feat: support element-wise multiplication. (#480)

* feat: support element-wise multiplication.

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>

* fix tap.

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>

* Implement mul for vecf16 & veci8.

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>

* fix zero.

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>

---------

Signed-off-by: my-vegetable-has-exploded <wy1109468038@gmail.com>

Author: my-vegetable-has-exploded

src/datatype/operators_svecf32.rs

@@ -88,6 +88,49 @@ 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 => {
+                // 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;
+                // only increment pos if the value is not zero
+                pos += (!values[pos].is_zero()) as usize;
+            }
+            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 _);

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());

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());

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::<Vec<_>>();
+    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());

src/sql/finalize.sql

@@ -140,6 +140,34 @@ CREATE OPERATOR - (
     RIGHTARG = veci8
 );
 
+CREATE OPERATOR * (
+    PROCEDURE = _vectors_vecf32_operator_mul,
+    LEFTARG = vector,
+    RIGHTARG = vector,
+    COMMUTATOR = *
+);
+
+CREATE OPERATOR * (
+    PROCEDURE = _vectors_vecf16_operator_mul,
+    LEFTARG = vecf16,
+    RIGHTARG = vecf16,
+    COMMUTATOR = *
+);
+
+CREATE OPERATOR * (
+    PROCEDURE = _vectors_svecf32_operator_mul,
+    LEFTARG = svector,
+    RIGHTARG = svector,
+    COMMUTATOR = *
+);
+
+CREATE OPERATOR * (
+    PROCEDURE = _vectors_veci8_operator_mul,
+    LEFTARG = veci8,
+    RIGHTARG = veci8,
+    COMMUTATOR = *
+);
+
 CREATE OPERATOR & (
     PROCEDURE = _vectors_bvecf32_operator_and,
     LEFTARG = bvector,

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;

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}');
 

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}');
 

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;
 ----