feat: support distance pushdown

Signed-off-by: cutecutecat <[email protected]>

crates/base/src/vector/bvecf32.rs

@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
 pub const BVEC_WIDTH: usize = usize::BITS as usize;
 
 // When using binary vector, please ensure that the padding bits are always zero.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct BVecf32Owned {
     dims: u16,
     data: Vec<usize>,

crates/base/src/vector/mod.rs

@@ -52,7 +52,7 @@ pub trait VectorBorrowed: Copy {
     fn normalize(&self) -> Self::Owned;
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub enum OwnedVector {
     Vecf32(Vecf32Owned),
     Vecf16(Vecf16Owned),

crates/base/src/vector/svecf32.rs

@@ -3,7 +3,7 @@ use crate::vector::{VectorBorrowed, VectorKind, VectorOwned};
 use num_traits::{Float, Zero};
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct SVecf32Owned {
     dims: u32,
     indexes: Vec<u32>,

crates/base/src/vector/vecf16.rs

@@ -3,7 +3,7 @@ use crate::scalar::{ScalarLike, F16, F32};
 use num_traits::{Float, Zero};
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 #[repr(transparent)]
 pub struct Vecf16Owned(Vec<F16>);
 

crates/base/src/vector/vecf32.rs

@@ -3,7 +3,7 @@ use crate::scalar::F32;
 use num_traits::{Float, Zero};
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 #[repr(transparent)]
 pub struct Vecf32Owned(Vec<F32>);
 

crates/base/src/vector/veci8.rs

@@ -3,7 +3,7 @@ use crate::scalar::{F32, I8};
 use num_traits::Float;
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct Veci8Owned {
     dims: u32,
     data: Vec<I8>,

crates/base/src/worker.rs

@@ -1,4 +1,5 @@
 use crate::index::*;
+use crate::scalar::F32;
 use crate::search::*;
 use crate::vector::*;
 
@@ -32,15 +33,15 @@ pub trait ViewBasicOperations {
         &'a self,
         vector: &'a OwnedVector,
         opts: &'a SearchOptions,
-    ) -> Result<Box<dyn Iterator<Item = Pointer> + 'a>, BasicError>;
+    ) -> Result<Box<dyn Iterator<Item = (F32, Pointer)> + 'a>, BasicError>;
 }
 
 pub trait ViewVbaseOperations {
     fn vbase<'a>(
         &'a self,
         vector: &'a OwnedVector,
         opts: &'a SearchOptions,
-    ) -> Result<Box<dyn Iterator<Item = Pointer> + 'a>, VbaseError>;
+    ) -> Result<Box<dyn Iterator<Item = (F32, Pointer)> + 'a>, VbaseError>;
 }
 
 pub trait ViewListOperations {

crates/hnsw/Cargo.toml

@@ -4,6 +4,7 @@ version.workspace = true
 edition.workspace = true
 
 [dependencies]
+log.workspace = true
 bytemuck.workspace = true
 parking_lot.workspace = true
 

crates/index/src/lib.rs

@@ -17,6 +17,7 @@ use arc_swap::ArcSwap;
 pub use base::distance::*;
 pub use base::index::*;
 use base::operator::*;
+use base::scalar::F32;
 pub use base::search::*;
 pub use base::vector::*;
 use common::clean::clean;
@@ -311,7 +312,7 @@ impl<O: Op> IndexView<O> {
         &'a self,
         vector: Borrowed<'_, O>,
         opts: &'a SearchOptions,
-    ) -> Result<impl Iterator<Item = Pointer> + 'a, BasicError> {
+    ) -> Result<impl Iterator<Item = (F32, Pointer)> + 'a, BasicError> {
         if self.options.vector.dims != vector.dims() {
             return Err(BasicError::InvalidVector);
         }
@@ -348,7 +349,7 @@ impl<O: Op> IndexView<O> {
         let loser = LoserTree::new(heaps);
         Ok(loser.filter_map(|x| {
             if self.delete.check(x.payload).is_some() {
-                Some(x.payload.pointer())
+                Some((x.distance, x.payload.pointer()))
             } else {
                 None
             }
@@ -358,7 +359,7 @@ impl<O: Op> IndexView<O> {
         &'a self,
         vector: Borrowed<'a, O>,
         opts: &'a SearchOptions,
-    ) -> Result<impl Iterator<Item = Pointer> + 'a, VbaseError> {
+    ) -> Result<impl Iterator<Item = (F32, Pointer)> + 'a, VbaseError> {
         if self.options.vector.dims != vector.dims() {
             return Err(VbaseError::InvalidVector);
         }
@@ -391,7 +392,7 @@ impl<O: Op> IndexView<O> {
         let loser = LoserTree::new(beta);
         Ok(loser.filter_map(|x| {
             if self.delete.check(x.payload).is_some() {
-                Some(x.payload.pointer())
+                Some((x.distance, x.payload.pointer()))
             } else {
                 None
             }

crates/service/src/instance.rs

@@ -1,6 +1,7 @@
 use base::distance::*;
 use base::index::*;
 use base::operator::*;
+use base::scalar::F32;
 use base::search::*;
 use base::vector::*;
 use base::worker::*;
@@ -297,11 +298,11 @@ impl ViewBasicOperations for InstanceView {
         &'a self,
         vector: &'a OwnedVector,
         opts: &'a SearchOptions,
-    ) -> Result<Box<dyn Iterator<Item = Pointer> + 'a>, BasicError> {
+    ) -> Result<Box<dyn Iterator<Item = (F32, Pointer)> + 'a>, BasicError> {
         match (self, vector) {
             (InstanceView::Vecf32Cos(x), OwnedVector::Vecf32(vector)) => {
                 Ok(Box::new(x.basic(vector.for_borrow(), opts)?)
-                    as Box<dyn Iterator<Item = Pointer>>)
+                    as Box<dyn Iterator<Item = (F32, Pointer)>>)
             }
             (InstanceView::Vecf32Dot(x), OwnedVector::Vecf32(vector)) => {
                 Ok(Box::new(x.basic(vector.for_borrow(), opts)?))
@@ -358,11 +359,11 @@ impl ViewVbaseOperations for InstanceView {
         &'a self,
         vector: &'a OwnedVector,
         opts: &'a SearchOptions,
-    ) -> Result<Box<dyn Iterator<Item = Pointer> + 'a>, VbaseError> {
+    ) -> Result<Box<dyn Iterator<Item = (F32, Pointer)> + 'a>, VbaseError> {
         match (self, vector) {
             (InstanceView::Vecf32Cos(x), OwnedVector::Vecf32(vector)) => {
                 Ok(Box::new(x.vbase(vector.for_borrow(), opts)?)
-                    as Box<dyn Iterator<Item = Pointer>>)
+                    as Box<dyn Iterator<Item = (F32, Pointer)>>)
             }
             (InstanceView::Vecf32Dot(x), OwnedVector::Vecf32(vector)) => {
                 Ok(Box::new(x.vbase(vector.for_borrow(), opts)?))

src/datatype/memory_bvecf32.rs

@@ -8,6 +8,7 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping;
 use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable;
 use pgrx::FromDatum;
 use pgrx::IntoDatum;
+use pgrx::UnboxDatum;
 use std::alloc::Layout;
 use std::ops::Deref;
 use std::ops::DerefMut;
@@ -137,6 +138,32 @@ impl<'a> FromDatum for BVecf32Input<'a> {
     }
 }
 
+impl<'a> IntoDatum for BVecf32Input<'a> {
+    // This is a fake IntoDatum, used for try_from_datum -> is_binary_coercible
+    // Don't call it directly!
+    fn into_datum(self) -> Option<Datum> {
+        match self {
+            BVecf32Input::Owned(o) => o.into_datum(),
+            BVecf32Input::Borrowed(_) => None,
+        }
+    }
+
+    fn type_oid() -> Oid {
+        BVecf32Output::type_oid()
+    }
+}
+
+unsafe impl<'a> UnboxDatum for BVecf32Input<'a> {
+    type As<'src> = BVecf32Input<'src>  where Self: 'src;
+
+    unsafe fn unbox<'src>(d: pgrx::Datum<'src>) -> Self::As<'src>
+    where
+        Self: 'src,
+    {
+        unsafe { BVecf32Input::from_datum(d.sans_lifetime(), false).unwrap() }
+    }
+}
+
 impl IntoDatum for BVecf32Output {
     fn into_datum(self) -> Option<Datum> {
         Some(Datum::from(self.into_raw() as *mut ()))

src/datatype/memory_svecf32.rs

@@ -9,6 +9,7 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping;
 use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable;
 use pgrx::FromDatum;
 use pgrx::IntoDatum;
+use pgrx::UnboxDatum;
 use std::alloc::Layout;
 use std::ops::Deref;
 use std::ptr::NonNull;
@@ -153,6 +154,32 @@ impl<'a> FromDatum for SVecf32Input<'a> {
     }
 }
 
+impl<'a> IntoDatum for SVecf32Input<'a> {
+    // This is a fake IntoDatum, used for try_from_datum -> is_binary_coercible
+    // Don't call it directly!
+    fn into_datum(self) -> Option<Datum> {
+        match self {
+            SVecf32Input::Owned(o) => o.into_datum(),
+            SVecf32Input::Borrowed(_) => None,
+        }
+    }
+
+    fn type_oid() -> Oid {
+        SVecf32Output::type_oid()
+    }
+}
+
+unsafe impl<'a> UnboxDatum for SVecf32Input<'a> {
+    type As<'src> = SVecf32Input<'src>  where Self: 'src;
+
+    unsafe fn unbox<'src>(d: pgrx::Datum<'src>) -> Self::As<'src>
+    where
+        Self: 'src,
+    {
+        unsafe { SVecf32Input::from_datum(d.sans_lifetime(), false).unwrap() }
+    }
+}
+
 impl IntoDatum for SVecf32Output {
     fn into_datum(self) -> Option<Datum> {
         Some(Datum::from(self.into_raw() as *mut ()))

src/datatype/memory_vecf16.rs

@@ -9,6 +9,7 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping;
 use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable;
 use pgrx::FromDatum;
 use pgrx::IntoDatum;
+use pgrx::UnboxDatum;
 use std::alloc::Layout;
 use std::ops::Deref;
 use std::ptr::NonNull;
@@ -130,6 +131,32 @@ impl<'a> FromDatum for Vecf16Input<'a> {
     }
 }
 
+impl<'a> IntoDatum for Vecf16Input<'a> {
+    // This is a fake IntoDatum, used for try_from_datum -> is_binary_coercible
+    // Don't call it directly!
+    fn into_datum(self) -> Option<Datum> {
+        match self {
+            Vecf16Input::Owned(o) => o.into_datum(),
+            Vecf16Input::Borrowed(_) => None,
+        }
+    }
+
+    fn type_oid() -> Oid {
+        Vecf16Output::type_oid()
+    }
+}
+
+unsafe impl<'a> UnboxDatum for Vecf16Input<'a> {
+    type As<'src> = Vecf16Input<'src>  where Self: 'src;
+
+    unsafe fn unbox<'src>(d: pgrx::Datum<'src>) -> Self::As<'src>
+    where
+        Self: 'src,
+    {
+        unsafe { Vecf16Input::from_datum(d.sans_lifetime(), false).unwrap() }
+    }
+}
+
 impl IntoDatum for Vecf16Output {
     fn into_datum(self) -> Option<Datum> {
         Some(Datum::from(self.into_raw() as *mut ()))

src/datatype/memory_vecf32.rs

@@ -9,6 +9,7 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping;
 use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable;
 use pgrx::FromDatum;
 use pgrx::IntoDatum;
+use pgrx::UnboxDatum;
 use std::alloc::Layout;
 use std::ops::Deref;
 use std::ptr::NonNull;
@@ -130,6 +131,32 @@ impl<'a> FromDatum for Vecf32Input<'a> {
     }
 }
 
+impl<'a> IntoDatum for Vecf32Input<'a> {
+    // This is a fake IntoDatum, used for try_from_datum -> is_binary_coercible
+    // Don't call it directly!
+    fn into_datum(self) -> Option<Datum> {
+        match self {
+            Vecf32Input::Owned(o) => o.into_datum(),
+            Vecf32Input::Borrowed(_) => None,
+        }
+    }
+
+    fn type_oid() -> Oid {
+        Vecf32Output::type_oid()
+    }
+}
+
+unsafe impl<'a> UnboxDatum for Vecf32Input<'a> {
+    type As<'src> = Vecf32Input<'src>  where Self: 'src;
+
+    unsafe fn unbox<'src>(d: pgrx::Datum<'src>) -> Self::As<'src>
+    where
+        Self: 'src,
+    {
+        unsafe { Vecf32Input::from_datum(d.sans_lifetime(), false).unwrap() }
+    }
+}
+
 impl IntoDatum for Vecf32Output {
     fn into_datum(self) -> Option<Datum> {
         Some(Datum::from(self.into_raw() as *mut ()))

src/datatype/memory_veci8.rs

@@ -9,6 +9,7 @@ use pgrx::pgrx_sql_entity_graph::metadata::SqlMapping;
 use pgrx::pgrx_sql_entity_graph::metadata::SqlTranslatable;
 use pgrx::FromDatum;
 use pgrx::IntoDatum;
+use pgrx::UnboxDatum;
 use std::alloc::Layout;
 use std::ops::Deref;
 use std::ops::DerefMut;
@@ -192,6 +193,32 @@ impl<'a> FromDatum for Veci8Input<'a> {
     }
 }
 
+impl<'a> IntoDatum for Veci8Input<'a> {
+    // This is a fake IntoDatum, used for try_from_datum -> is_binary_coercible
+    // Don't call it directly!
+    fn into_datum(self) -> Option<Datum> {
+        match self {
+            Veci8Input::Owned(o) => o.into_datum(),
+            Veci8Input::Borrowed(_) => None,
+        }
+    }
+
+    fn type_oid() -> Oid {
+        Veci8Output::type_oid()
+    }
+}
+
+unsafe impl<'a> UnboxDatum for Veci8Input<'a> {
+    type As<'src> = Veci8Input<'src>  where Self: 'src;
+
+    unsafe fn unbox<'src>(d: pgrx::Datum<'src>) -> Self::As<'src>
+    where
+        Self: 'src,
+    {
+        unsafe { Veci8Input::from_datum(d.sans_lifetime(), false).unwrap() }
+    }
+}
+
 impl IntoDatum for Veci8Output {
     fn into_datum(self) -> Option<Datum> {
         Some(Datum::from(self.into_raw() as *mut ()))