feat: add vector_norm, vector_dims and normalize for all vector types (#485)

fix: remove linking hacks in CI

Signed-off-by: usamoi <usamoi@outlook.com>

Author: usamoi

.github/workflows/psql.yml

@@ -95,7 +95,7 @@ jobs:
           echo PGRX_TARGET_INFO_PATH_PG$VERSION=$HOME/.pgrx_binding >> "$GITHUB_ENV"
       - name: Build Release
         run: |
-          cargo build --no-default-features --features "pg$VERSION" --release
+          cargo build --lib --no-default-features --features "pg$VERSION" --release
           ./tools/schema.sh --no-default-features --features "pg$VERSION" --release | expand -t 4 > ./target/vectors--$SEMVER.sql
       - name: Set up PostgreSQL
         run: |

.github/workflows/release.yml

@@ -92,7 +92,7 @@ jobs:
           echo PGRX_TARGET_INFO_PATH_PG$VERSION=$HOME/.pgrx_binding >> "$GITHUB_ENV"
       - name: Build
         run: |
-          cargo build --no-default-features --features pg$VERSION --release --target $ARCH-unknown-linux-gnu
+          cargo build --lib --no-default-features --features pg$VERSION --release --target $ARCH-unknown-linux-gnu
           ./tools/schema.sh --no-default-features --features pg$VERSION --release --target $ARCH-unknown-linux-gnu | expand -t 4 > ./target/vectors--$SEMVER.sql
       - name: Package
         run: |

.github/workflows/rust.yml

@@ -100,7 +100,7 @@ jobs:
       - name: Clippy
         run: cargo clippy --no-default-features --features "pg$VERSION" --target $ARCH-unknown-linux-gnu
       - name: Build
-        run: cargo build --no-default-features --features "pg$VERSION" --target $ARCH-unknown-linux-gnu
+        run: cargo build --lib --no-default-features --features "pg$VERSION" --target $ARCH-unknown-linux-gnu
       - name: Post Set up Cache
         uses: actions/cache/save@v4
         if: ${{ !steps.cache.outputs.cache-hit }}

crates/base/src/scalar/f32.rs

@@ -1,4 +1,5 @@
 use super::ScalarLike;
+use num_traits::Zero;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
 use std::fmt::{Debug, Display};
@@ -65,7 +66,7 @@ unsafe impl bytemuck::Zeroable for F32 {}
 
 unsafe impl bytemuck::Pod for F32 {}
 
-impl num_traits::Zero for F32 {
+impl Zero for F32 {
     fn zero() -> Self {
         Self(f32::zero())
     }
@@ -468,7 +469,7 @@ impl AddAssign<F32> for F32 {
 
 impl Sum for F32 {
     fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
-        iter.fold(F32(0.0), Add::add)
+        iter.fold(F32::zero(), Add::add)
     }
 }
 

crates/base/src/scalar/half_f16.rs

@@ -1,9 +1,11 @@
 use super::ScalarLike;
 use crate::scalar::F32;
 use half::f16;
+use num_traits::Zero;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
 use std::fmt::{Debug, Display};
+use std::iter::Sum;
 use std::num::ParseFloatError;
 use std::ops::*;
 use std::str::FromStr;
@@ -51,7 +53,7 @@ unsafe impl bytemuck::Zeroable for F16 {}
 
 unsafe impl bytemuck::Pod for F16 {}
 
-impl num_traits::Zero for F16 {
+impl Zero for F16 {
     fn zero() -> Self {
         Self(f16::zero())
     }
@@ -452,6 +454,12 @@ impl AddAssign<F16> for F16 {
     }
 }
 
+impl Sum for F16 {
+    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
+        iter.fold(F16::zero(), Add::add)
+    }
+}
+
 impl Sub<F16> for F16 {
     type Output = F16;
 

crates/base/src/vector/bvecf32.rs

@@ -17,6 +17,7 @@ impl BVecf32Owned {
     pub fn new(dims: u16, data: Vec<usize>) -> Self {
         Self::new_checked(dims, data).unwrap()
     }
+
     #[inline(always)]
     pub fn new_checked(dims: u16, data: Vec<usize>) -> Option<Self> {
         if dims == 0 {
@@ -31,6 +32,7 @@ impl BVecf32Owned {
         }
         unsafe { Some(Self::new_unchecked(dims, data)) }
     }
+
     /// # Safety
     ///
     /// * `dims` must be in `1..=65535`.
@@ -50,24 +52,6 @@ impl BVecf32Owned {
             data: vec![0; size],
         }
     }
-
-    #[inline(always)]
-    pub fn set(&mut self, index: usize, value: bool) {
-        assert!(index < self.dims as usize);
-        if value {
-            self.data[index / BVEC_WIDTH] |= 1 << (index % BVEC_WIDTH);
-        } else {
-            self.data[index / BVEC_WIDTH] &= !(1 << (index % BVEC_WIDTH));
-        }
-    }
-
-    /// # Safety
-    ///
-    /// The caller must ensure that it won't modify the padding bits
-    #[inline(always)]
-    pub unsafe fn data_mut(&mut self) -> &mut [usize] {
-        &mut self.data
-    }
 }
 
 impl VectorOwned for BVecf32Owned {
@@ -81,13 +65,15 @@ impl VectorOwned for BVecf32Owned {
         self.dims as u32
     }
 
+    #[inline(always)]
     fn for_borrow(&self) -> BVecf32Borrowed<'_> {
         BVecf32Borrowed {
             dims: self.dims,
             data: &self.data,
         }
     }
 
+    #[inline(always)]
     fn to_vec(&self) -> Vec<F32> {
         self.for_borrow().to_vec()
     }
@@ -104,6 +90,7 @@ impl<'a> BVecf32Borrowed<'a> {
     pub fn new(dims: u16, data: &'a [usize]) -> Self {
         Self::new_checked(dims, data).unwrap()
     }
+
     #[inline(always)]
     pub fn new_checked(dims: u16, data: &'a [usize]) -> Option<Self> {
         if dims == 0 {
@@ -118,6 +105,7 @@ impl<'a> BVecf32Borrowed<'a> {
         }
         unsafe { Some(Self::new_unchecked(dims, data)) }
     }
+
     /// # Safety
     ///
     /// * `dims` must be in `1..=65535`.
@@ -128,6 +116,18 @@ impl<'a> BVecf32Borrowed<'a> {
         Self { dims, data }
     }
 
+    #[inline(always)]
+    pub fn data(&self) -> &'a [usize] {
+        self.data
+    }
+
+    #[inline(always)]
+    pub fn get(&self, index: usize) -> bool {
+        assert!(index < self.dims as usize);
+        self.data[index / BVEC_WIDTH] & (1 << (index % BVEC_WIDTH)) != 0
+    }
+
+    #[inline(always)]
     pub fn iter(self) -> impl Iterator<Item = bool> + 'a {
         let mut index = 0;
         std::iter::from_fn(move || {
@@ -140,16 +140,6 @@ impl<'a> BVecf32Borrowed<'a> {
             }
         })
     }
-
-    pub fn get(&self, index: usize) -> bool {
-        assert!(index < self.dims as usize);
-        self.data[index / BVEC_WIDTH] & (1 << (index % BVEC_WIDTH)) != 0
-    }
-
-    #[inline(always)]
-    pub fn data(&self) -> &'a [usize] {
-        self.data
-    }
 }
 
 impl<'a> VectorBorrowed for BVecf32Borrowed<'a> {
@@ -168,9 +158,20 @@ impl<'a> VectorBorrowed for BVecf32Borrowed<'a> {
         }
     }
 
+    #[inline(always)]
     fn to_vec(&self) -> Vec<F32> {
         self.iter().map(|i| F32(i as u32 as f32)).collect()
     }
+
+    #[inline(always)]
+    fn length(&self) -> F32 {
+        length(*self)
+    }
+
+    #[inline(always)]
+    fn normalize(&self) -> BVecf32Owned {
+        unimplemented!()
+    }
 }
 
 impl<'a> Ord for BVecf32Borrowed<'a> {

crates/base/src/vector/mod.rs

@@ -11,6 +11,7 @@ pub use vecf32::{Vecf32Borrowed, Vecf32Owned};
 pub use veci8::{Veci8Borrowed, Veci8Owned};
 
 use crate::scalar::ScalarLike;
+use crate::scalar::F32;
 use serde::{Deserialize, Serialize};
 
 #[repr(u8)]
@@ -45,6 +46,10 @@ pub trait VectorBorrowed: Copy {
     fn dims(&self) -> u32;
 
     fn to_vec(&self) -> Vec<Self::Scalar>;
+
+    fn length(&self) -> F32;
+
+    fn normalize(&self) -> Self::Owned;
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

crates/base/src/vector/svecf32.rs

@@ -15,6 +15,7 @@ impl SVecf32Owned {
     pub fn new(dims: u32, indexes: Vec<u32>, values: Vec<F32>) -> Self {
         Self::new_checked(dims, indexes, values).unwrap()
     }
+
     #[inline(always)]
     pub fn new_checked(dims: u32, indexes: Vec<u32>, values: Vec<F32>) -> Option<Self> {
         if !(1..=1_048_575).contains(&dims) {
@@ -39,6 +40,7 @@ impl SVecf32Owned {
         }
         unsafe { Some(Self::new_unchecked(dims, indexes, values)) }
     }
+
     /// # Safety
     ///
     /// * `dims` must be in `1..=1_048_575`.
@@ -53,10 +55,12 @@ impl SVecf32Owned {
             values,
         }
     }
+
     #[inline(always)]
     pub fn indexes(&self) -> &[u32] {
         &self.indexes
     }
+
     #[inline(always)]
     pub fn values(&self) -> &[F32] {
         &self.values
@@ -74,6 +78,7 @@ impl VectorOwned for SVecf32Owned {
         self.dims
     }
 
+    #[inline(always)]
     fn for_borrow(&self) -> SVecf32Borrowed<'_> {
         SVecf32Borrowed {
             dims: self.dims,
@@ -82,6 +87,7 @@ impl VectorOwned for SVecf32Owned {
         }
     }
 
+    #[inline(always)]
     fn to_vec(&self) -> Vec<F32> {
         let mut dense = vec![F32::zero(); self.dims as usize];
         for (&index, &value) in self.indexes.iter().zip(self.values.iter()) {
@@ -103,6 +109,7 @@ impl<'a> SVecf32Borrowed<'a> {
     pub fn new(dims: u32, indexes: &'a [u32], values: &'a [F32]) -> Self {
         Self::new_checked(dims, indexes, values).unwrap()
     }
+
     #[inline(always)]
     pub fn new_checked(dims: u32, indexes: &'a [u32], values: &'a [F32]) -> Option<Self> {
         if !(1..=1_048_575).contains(&dims) {
@@ -127,6 +134,7 @@ impl<'a> SVecf32Borrowed<'a> {
         }
         unsafe { Some(Self::new_unchecked(dims, indexes, values)) }
     }
+
     /// # Safety
     ///
     /// * `dims` must be in `1..=1_048_575`.
@@ -141,14 +149,21 @@ impl<'a> SVecf32Borrowed<'a> {
             values,
         }
     }
+
     #[inline(always)]
     pub fn indexes(&self) -> &[u32] {
         self.indexes
     }
+
     #[inline(always)]
     pub fn values(&self) -> &[F32] {
         self.values
     }
+
+    #[inline(always)]
+    pub fn len(&self) -> u32 {
+        self.indexes.len().try_into().unwrap()
+    }
 }
 
 impl<'a> VectorBorrowed for SVecf32Borrowed<'a> {
@@ -160,6 +175,7 @@ impl<'a> VectorBorrowed for SVecf32Borrowed<'a> {
         self.dims
     }
 
+    #[inline(always)]
     fn for_own(&self) -> SVecf32Owned {
         SVecf32Owned {
             dims: self.dims,
@@ -168,19 +184,25 @@ impl<'a> VectorBorrowed for SVecf32Borrowed<'a> {
         }
     }
 
+    #[inline(always)]
     fn to_vec(&self) -> Vec<F32> {
         let mut dense = vec![F32::zero(); self.dims as usize];
         for (&index, &value) in self.indexes.iter().zip(self.values.iter()) {
             dense[index as usize] = value;
         }
         dense
     }
-}
 
-impl<'a> SVecf32Borrowed<'a> {
     #[inline(always)]
-    pub fn len(&self) -> u32 {
-        self.indexes.len().try_into().unwrap()
+    fn length(&self) -> F32 {
+        length(*self)
+    }
+
+    #[inline(always)]
+    fn normalize(&self) -> SVecf32Owned {
+        let mut own = self.for_own();
+        l2_normalize(&mut own);
+        own
     }
 }