From e149ff344c9e07b2baaae3ffed049460d16d5714 Mon Sep 17 00:00:00 2001 From: my-vegetable-has-exploded Date: Fri, 17 May 2024 12:11:40 +0000 Subject: [PATCH 1/4] feat: support element-wise multiplication. Signed-off-by: my-vegetable-has-exploded --- src/datatype/operators_svecf32.rs | 41 +++++++++++++++++++++++++++++++ src/datatype/operators_vecf32.rs | 11 +++++++++ src/sql/finalize.sql | 14 +++++++++++ tests/sqllogictest/sparse.slt | 5 ++++ tests/sqllogictest/vector.slt | 5 ++++ 5 files changed, 76 insertions(+) diff --git a/src/datatype/operators_svecf32.rs b/src/datatype/operators_svecf32.rs index 413cd71dd..ed8f1b9f7 100644 --- a/src/datatype/operators_svecf32.rs +++ b/src/datatype/operators_svecf32.rs @@ -88,6 +88,47 @@ fn _vectors_svecf32_operator_minus(lhs: SVecf32Input<'_>, rhs: SVecf32Input<'_>) SVecf32Output::new(SVecf32Borrowed::new(lhs.dims(), &indexes, &values)) } +/// Calculate the element-wise multiplication of two sparse vectors. +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vectors_svecf32_operator_mul(lhs: SVecf32Input<'_>, rhs: SVecf32Input<'_>) -> SVecf32Output { + check_matched_dims(lhs.dims() as _, rhs.dims() as _); + + let size1 = lhs.len(); + let size2 = rhs.len(); + let mut pos1 = 0; + let mut pos2 = 0; + let mut pos = 0; + let mut indexes = vec![0; std::cmp::min(size1, size2)]; + let mut values = vec![F32::zero(); std::cmp::min(size1, size2)]; + let lhs = lhs.for_borrow(); + let rhs = rhs.for_borrow(); + while pos1 < size1 && pos2 < size2 { + let lhs_index = lhs.indexes()[pos1]; + let rhs_index = rhs.indexes()[pos2]; + match lhs_index.cmp(&rhs_index) { + std::cmp::Ordering::Less => { + pos1 += 1; + } + std::cmp::Ordering::Equal => { + let lhs_value = lhs.values()[pos1]; + let rhs_value = rhs.values()[pos2]; + indexes[pos] = lhs_index; + values[pos] = lhs_value * rhs_value; + pos1 += 1; + pos2 += 1; + pos += 1; + } + std::cmp::Ordering::Greater => { + pos2 += 1; + } + } + } + indexes.truncate(pos); + values.truncate(pos); + + SVecf32Output::new(SVecf32Borrowed::new(lhs.dims(), &indexes, &values)) +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vectors_svecf32_operator_lt(lhs: SVecf32Input<'_>, rhs: SVecf32Input<'_>) -> bool { check_matched_dims(lhs.dims() as _, rhs.dims() as _); diff --git a/src/datatype/operators_vecf32.rs b/src/datatype/operators_vecf32.rs index 136616366..935fabdab 100644 --- a/src/datatype/operators_vecf32.rs +++ b/src/datatype/operators_vecf32.rs @@ -26,6 +26,17 @@ fn _vectors_vecf32_operator_minus(lhs: Vecf32Input<'_>, rhs: Vecf32Input<'_>) -> Vecf32Output::new(Vecf32Borrowed::new(&v)) } +/// Calculate the element-wise multiplication of two vectors. +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vectors_vecf32_operator_mul(lhs: Vecf32Input<'_>, rhs: Vecf32Input<'_>) -> Vecf32Output { + let n = check_matched_dims(lhs.dims(), rhs.dims()); + let mut v = vec![F32::zero(); n]; + for i in 0..n { + v[i] = lhs[i] * rhs[i]; + } + Vecf32Output::new(Vecf32Borrowed::new(&v)) +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vectors_vecf32_operator_lt(lhs: Vecf32Input<'_>, rhs: Vecf32Input<'_>) -> bool { check_matched_dims(lhs.dims(), rhs.dims()); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index b5ce4a22b..70382dacc 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -140,6 +140,20 @@ CREATE OPERATOR - ( RIGHTARG = veci8 ); +CREATE OPERATOR * ( + PROCEDURE = _vectors_vecf32_operator_mul, + LEFTARG = vector, + RIGHTARG = vector, + COMMUTATOR = * +); + +CREATE OPERATOR * ( + PROCEDURE = _vectors_svecf32_operator_mul, + LEFTARG = svector, + RIGHTARG = svector, + COMMUTATOR = * +); + CREATE OPERATOR & ( PROCEDURE = _vectors_bvecf32_operator_and, LEFTARG = bvector, diff --git a/tests/sqllogictest/sparse.slt b/tests/sqllogictest/sparse.slt index 018b37423..07513d465 100644 --- a/tests/sqllogictest/sparse.slt +++ b/tests/sqllogictest/sparse.slt @@ -42,6 +42,11 @@ SELECT to_svector(5, '{1,2}', '{1,2}'); ---- [0, 1, 2, 0, 0] +query I +SELECT to_svector(5, '{1,2}', '{1,1}') * to_svector(5, '{1,3}', '{2,2}'); +---- +[0, 2, 0, 0, 0] + statement error Lengths of index and value are not matched. SELECT to_svector(5, '{1,2,3}', '{1,2}'); diff --git a/tests/sqllogictest/vector.slt b/tests/sqllogictest/vector.slt index 1b5d77f8d..f89d64c4b 100644 --- a/tests/sqllogictest/vector.slt +++ b/tests/sqllogictest/vector.slt @@ -41,6 +41,11 @@ SELECT vector_dims(v) FROM unnest(ARRAY['[1,2]'::vector, '[3]']) v; 2 1 +query I +SELECT '[1,2,3]'::vector * '[4,5,6]'::vector; +---- +[4, 10, 18] + query ? SELECT avg(v) FROM unnest(ARRAY['[1,2,3]'::vector, '[3,5,7]']) v; ---- From 29448caa7f494349c7c2a5bbb955a882d1084e43 Mon Sep 17 00:00:00 2001 From: my-vegetable-has-exploded Date: Sat, 18 May 2024 06:01:42 +0000 Subject: [PATCH 2/4] fix tap. Signed-off-by: my-vegetable-has-exploded --- src/sql/finalize.sql | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 70382dacc..3f411414a 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -141,17 +141,17 @@ CREATE OPERATOR - ( ); CREATE OPERATOR * ( - PROCEDURE = _vectors_vecf32_operator_mul, - LEFTARG = vector, - RIGHTARG = vector, - COMMUTATOR = * + PROCEDURE = _vectors_vecf32_operator_mul, + LEFTARG = vector, + RIGHTARG = vector, + COMMUTATOR = * ); CREATE OPERATOR * ( - PROCEDURE = _vectors_svecf32_operator_mul, - LEFTARG = svector, - RIGHTARG = svector, - COMMUTATOR = * + PROCEDURE = _vectors_svecf32_operator_mul, + LEFTARG = svector, + RIGHTARG = svector, + COMMUTATOR = * ); CREATE OPERATOR & ( From 09699cb48aceaa03325adca52d8e94d727fb437e Mon Sep 17 00:00:00 2001 From: my-vegetable-has-exploded Date: Sun, 19 May 2024 11:28:23 +0000 Subject: [PATCH 3/4] Implement mul for vecf16 & veci8. Signed-off-by: my-vegetable-has-exploded --- src/datatype/operators_vecf16.rs | 11 +++++++++++ src/datatype/operators_veci8.rs | 14 ++++++++++++++ src/sql/finalize.sql | 14 ++++++++++++++ tests/sqllogictest/fp16.slt | 5 +++++ tests/sqllogictest/int8.slt | 5 +++++ 5 files changed, 49 insertions(+) diff --git a/src/datatype/operators_vecf16.rs b/src/datatype/operators_vecf16.rs index d2ea2d8c2..c3b7c8808 100644 --- a/src/datatype/operators_vecf16.rs +++ b/src/datatype/operators_vecf16.rs @@ -26,6 +26,17 @@ fn _vectors_vecf16_operator_minus(lhs: Vecf16Input<'_>, rhs: Vecf16Input<'_>) -> Vecf16Output::new(Vecf16Borrowed::new(&v)) } +/// Calculate the element-wise multiplication of two f16 vectors. +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vectors_vecf16_operator_mul(lhs: Vecf16Input<'_>, rhs: Vecf16Input<'_>) -> Vecf16Output { + let n = check_matched_dims(lhs.dims(), rhs.dims()); + let mut v = vec![F16::zero(); n]; + for i in 0..n { + v[i] = lhs[i] * rhs[i]; + } + Vecf16Output::new(Vecf16Borrowed::new(&v)) +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vectors_vecf16_operator_lt(lhs: Vecf16Input<'_>, rhs: Vecf16Input<'_>) -> bool { check_matched_dims(lhs.dims(), rhs.dims()); diff --git a/src/datatype/operators_veci8.rs b/src/datatype/operators_veci8.rs index 46c894649..74b3e1bde 100644 --- a/src/datatype/operators_veci8.rs +++ b/src/datatype/operators_veci8.rs @@ -31,6 +31,20 @@ fn _vectors_veci8_operator_minus(lhs: Veci8Input<'_>, rhs: Veci8Input<'_>) -> Ve ) } +/// Calculate the element-wise multiplication of two i8 vectors. +#[pgrx::pg_extern(immutable, strict, parallel_safe)] +fn _vectors_veci8_operator_mul(lhs: Veci8Input<'_>, rhs: Veci8Input<'_>) -> Veci8Output { + check_matched_dims(lhs.len(), rhs.len()); + let data = (0..lhs.len()) + .map(|i| lhs.index(i) * rhs.index(i)) + .collect::>(); + let (vector, alpha, offset) = veci8::i8_quantization(&data); + let (sum, l2_norm) = veci8::i8_precompute(&vector, alpha, offset); + Veci8Output::new( + Veci8Borrowed::new_checked(lhs.len() as u32, &vector, alpha, offset, sum, l2_norm).unwrap(), + ) +} + #[pgrx::pg_extern(immutable, strict, parallel_safe)] fn _vectors_veci8_operator_lt(lhs: Veci8Input<'_>, rhs: Veci8Input<'_>) -> bool { check_matched_dims(lhs.len(), rhs.len()); diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index 3f411414a..c3c54fae3 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -147,6 +147,13 @@ CREATE OPERATOR * ( COMMUTATOR = * ); +CREATE OPERATOR * ( + PROCEDURE = _vectors_vecf16_operator_mul, + LEFTARG = vecf16, + RIGHTARG = vecf16, + COMMUTATOR = * +); + CREATE OPERATOR * ( PROCEDURE = _vectors_svecf32_operator_mul, LEFTARG = svector, @@ -154,6 +161,13 @@ CREATE OPERATOR * ( COMMUTATOR = * ); +CREATE OPERATOR * ( + PROCEDURE = _vectors_veci8_operator_mul, + LEFTARG = veci8, + RIGHTARG = veci8, + COMMUTATOR = * +); + CREATE OPERATOR & ( PROCEDURE = _vectors_bvecf32_operator_and, LEFTARG = bvector, diff --git a/tests/sqllogictest/fp16.slt b/tests/sqllogictest/fp16.slt index e7cdf6046..8289c1782 100644 --- a/tests/sqllogictest/fp16.slt +++ b/tests/sqllogictest/fp16.slt @@ -35,5 +35,10 @@ SELECT COUNT(1) FROM (SELECT 1 FROM t ORDER BY val <=> '[0.5,0.5,0.5]'::vecf16 l ---- 10 +query I +SELECT '[1,2,3]'::vecf16 * '[4,5,6]'::vecf16; +---- +[4, 10, 18] + statement ok DROP TABLE t; \ No newline at end of file diff --git a/tests/sqllogictest/int8.slt b/tests/sqllogictest/int8.slt index caac864d2..2aa25a5bd 100644 --- a/tests/sqllogictest/int8.slt +++ b/tests/sqllogictest/int8.slt @@ -43,6 +43,11 @@ SELECT to_veci8(5, 1, 0, '{0,1,2,0,0}'); ---- [0, 1, 2, 0, 0] +query I +SELECT '[2,2,2]'::veci8 * '[2,2,2]'::veci8; +---- +[4, 4, 4] + statement error Lengths of values and len are not matched. SELECT to_veci8(5, 1, 0, '{0,1,2,0,0,0}'); From e62f6a6cce4953a013d57d35a63dd00282c16595 Mon Sep 17 00:00:00 2001 From: my-vegetable-has-exploded Date: Mon, 20 May 2024 08:49:30 +0000 Subject: [PATCH 4/4] fix zero. Signed-off-by: my-vegetable-has-exploded --- src/datatype/operators_svecf32.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/datatype/operators_svecf32.rs b/src/datatype/operators_svecf32.rs index ed8f1b9f7..be4e39272 100644 --- a/src/datatype/operators_svecf32.rs +++ b/src/datatype/operators_svecf32.rs @@ -110,13 +110,15 @@ fn _vectors_svecf32_operator_mul(lhs: SVecf32Input<'_>, rhs: SVecf32Input<'_>) - pos1 += 1; } std::cmp::Ordering::Equal => { + // only both indexes are not zero, values are multiplied let lhs_value = lhs.values()[pos1]; let rhs_value = rhs.values()[pos2]; indexes[pos] = lhs_index; values[pos] = lhs_value * rhs_value; pos1 += 1; pos2 += 1; - pos += 1; + // only increment pos if the value is not zero + pos += (!values[pos].is_zero()) as usize; } std::cmp::Ordering::Greater => { pos2 += 1;