feat: vchordrqfscan (#105)

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

src/lib.rs

@@ -7,9 +7,10 @@
 mod datatype;
 mod postgres;
 mod projection;
-mod types;
 mod upgrade;
+mod utils;
 mod vchordrq;
+mod vchordrqfscan;
 
 pgrx::pg_module_magic!();
 pgrx::extension_sql_file!("./sql/bootstrap.sql", bootstrap);
@@ -23,6 +24,7 @@ unsafe extern "C" fn _PG_init() {
     detect::init();
     unsafe {
         vchordrq::init();
+        vchordrqfscan::init();
 
         #[cfg(any(feature = "pg13", feature = "pg14"))]
         pgrx::pg_sys::EmitWarningsOnPlaceholders(c"vchord".as_ptr());

src/sql/finalize.sql

@@ -39,16 +39,28 @@ IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_amhan
 CREATE FUNCTION vchordrq_prewarm(regclass, integer default 0) RETURNS TEXT
 STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrq_prewarm_wrapper';
 
+CREATE FUNCTION vchordrqfscan_amhandler(internal) RETURNS index_am_handler
+IMMUTABLE STRICT PARALLEL SAFE LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrqfscan_amhandler_wrapper';
+
+CREATE FUNCTION vchordrqfscan_prewarm(regclass, integer default 0) RETURNS TEXT
+STRICT LANGUAGE c AS 'MODULE_PATHNAME', '_vchordrqfscan_prewarm_wrapper';
+
 -- List of access methods
 
 CREATE ACCESS METHOD vchordrq TYPE INDEX HANDLER vchordrq_amhandler;
 
+CREATE ACCESS METHOD Vchordrqfscan TYPE INDEX HANDLER Vchordrqfscan_amhandler;
+
 -- List of operator families
 
 CREATE OPERATOR FAMILY vector_l2_ops USING vchordrq;
 CREATE OPERATOR FAMILY vector_ip_ops USING vchordrq;
 CREATE OPERATOR FAMILY vector_cosine_ops USING vchordrq;
 
+CREATE OPERATOR FAMILY vector_l2_ops USING Vchordrqfscan;
+CREATE OPERATOR FAMILY vector_ip_ops USING Vchordrqfscan;
+CREATE OPERATOR FAMILY vector_cosine_ops USING Vchordrqfscan;
+
 -- List of operator classes
 
 CREATE OPERATOR CLASS vector_l2_ops
@@ -68,3 +80,21 @@ CREATE OPERATOR CLASS vector_cosine_ops
     OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
     OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH,
     FUNCTION 1 _vchordrq_support_vector_cosine_ops();
+
+CREATE OPERATOR CLASS vector_l2_ops
+    FOR TYPE vector USING Vchordrqfscan FAMILY vector_l2_ops AS
+    OPERATOR 1 <-> (vector, vector) FOR ORDER BY float_ops,
+    OPERATOR 2 <<->> (vector, sphere_vector) FOR SEARCH,
+    FUNCTION 1 _Vchordrqfscan_support_vector_l2_ops();
+
+CREATE OPERATOR CLASS vector_ip_ops
+    FOR TYPE vector USING Vchordrqfscan FAMILY vector_ip_ops AS
+    OPERATOR 1 <#> (vector, vector) FOR ORDER BY float_ops,
+    OPERATOR 2 <<#>> (vector, sphere_vector) FOR SEARCH,
+    FUNCTION 1 _Vchordrqfscan_support_vector_ip_ops();
+
+CREATE OPERATOR CLASS vector_cosine_ops
+    FOR TYPE vector USING Vchordrqfscan FAMILY vector_cosine_ops AS
+    OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops,
+    OPERATOR 2 <<=>> (vector, sphere_vector) FOR SEARCH,
+    FUNCTION 1 _Vchordrqfscan_support_vector_cosine_ops();

src/vchordrq/algorithm/k_means.rs → src/utils/k_means.rs

@@ -1,4 +1,4 @@
-use crate::vchordrq::algorithm::parallelism::{ParallelIterator, Parallelism};
+use super::parallelism::{ParallelIterator, Parallelism};
 use base::scalar::*;
 use half::f16;
 use rand::rngs::StdRng;

src/utils/mod.rs

@@ -0,0 +1,2 @@
+pub mod k_means;
+pub mod parallelism;

src/vchordrq/algorithm/build.rs

@@ -1,14 +1,13 @@
 use crate::postgres::BufferWriteGuard;
 use crate::postgres::Relation;
-use crate::types::VchordrqBuildOptions;
-use crate::types::VchordrqExternalBuildOptions;
-use crate::types::VchordrqIndexingOptions;
-use crate::types::VchordrqInternalBuildOptions;
-use crate::vchordrq::algorithm::k_means;
 use crate::vchordrq::algorithm::rabitq;
 use crate::vchordrq::algorithm::tuples::*;
 use crate::vchordrq::algorithm::vectors;
 use crate::vchordrq::index::am_options::Opfamily;
+use crate::vchordrq::types::VchordrqBuildOptions;
+use crate::vchordrq::types::VchordrqExternalBuildOptions;
+use crate::vchordrq::types::VchordrqIndexingOptions;
+use crate::vchordrq::types::VchordrqInternalBuildOptions;
 use base::distance::DistanceKind;
 use base::index::VectorOptions;
 use base::scalar::ScalarLike;
@@ -158,13 +157,13 @@ impl Structure {
         }
         let mut result = Vec::<Self>::new();
         for w in internal_build.lists.iter().rev().copied().chain(once(1)) {
-            let means = crate::vchordrq::algorithm::parallelism::RayonParallelism::scoped(
+            let means = crate::utils::parallelism::RayonParallelism::scoped(
                 internal_build.build_threads as _,
                 Arc::new(|| {
                     pgrx::check_for_interrupts!();
                 }),
                 |parallelism| {
-                    k_means::k_means(
+                    crate::utils::k_means::k_means(
                         parallelism,
                         w as usize,
                         vector_options.dims as usize,
@@ -182,7 +181,8 @@ impl Structure {
             if let Some(structure) = result.last() {
                 let mut children = vec![Vec::new(); means.len()];
                 for i in 0..structure.len() as u32 {
-                    let target = k_means::k_means_lookup(&structure.means[i as usize], &means);
+                    let target =
+                        crate::utils::k_means::k_means_lookup(&structure.means[i as usize], &means);
                     children[target].push(i);
                 }
                 let (means, children) = std::iter::zip(means, children)

src/vchordrq/algorithm/mod.rs

@@ -1,7 +1,5 @@
 pub mod build;
 pub mod insert;
-pub mod k_means;
-pub mod parallelism;
 pub mod prewarm;
 pub mod rabitq;
 pub mod scan;

src/vchordrq/algorithm/rabitq.rs

@@ -141,3 +141,12 @@ fn asymmetric_binary_dot_product(x: &[u64], y: &(Vec<u64>, Vec<u64>, Vec<u64>, V
     }
     (t0 << 0) + (t1 << 1) + (t2 << 2) + (t3 << 3)
 }
+
+pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance {
+    match d {
+        DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)),
+        DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)),
+        DistanceKind::Hamming => unimplemented!(),
+        DistanceKind::Jaccard => unimplemented!(),
+    }
+}

src/vchordrq/algorithm/vectors.rs

@@ -1,6 +1,6 @@
 use crate::postgres::Relation;
+use crate::vchordrq::algorithm::rabitq::distance;
 use crate::vchordrq::algorithm::tuples::VectorTuple;
-use crate::vchordrq::index::utils::distance;
 use base::distance::Distance;
 use base::distance::DistanceKind;
 

src/vchordrq/index/am_options.rs

@@ -1,7 +1,7 @@
 use crate::datatype::memory_pgvector_vector::PgvectorVectorInput;
 use crate::datatype::memory_pgvector_vector::PgvectorVectorOutput;
 use crate::datatype::typmod::Typmod;
-use crate::types::VchordrqIndexingOptions;
+use crate::vchordrq::types::VchordrqIndexingOptions;
 use base::distance::*;
 use base::index::*;
 use base::vector::*;

src/vchordrq/index/utils.rs

@@ -1,5 +1,3 @@
-use base::distance::{Distance, DistanceKind};
-use base::scalar::ScalarLike;
 use base::search::*;
 
 pub fn pointer_to_ctid(pointer: Pointer) -> pgrx::pg_sys::ItemPointerData {
@@ -20,12 +18,3 @@ pub fn ctid_to_pointer(ctid: pgrx::pg_sys::ItemPointerData) -> Pointer {
     value |= ctid.ip_posid as u64;
     Pointer::new(value)
 }
-
-pub fn distance(d: DistanceKind, lhs: &[f32], rhs: &[f32]) -> Distance {
-    match d {
-        DistanceKind::L2 => Distance::from_f32(f32::reduce_sum_of_d2(lhs, rhs)),
-        DistanceKind::Dot => Distance::from_f32(-f32::reduce_sum_of_xy(lhs, rhs)),
-        DistanceKind::Hamming => unimplemented!(),
-        DistanceKind::Jaccard => unimplemented!(),
-    }
-}

src/vchordrq/mod.rs

@@ -1,6 +1,7 @@
 mod algorithm;
 mod gucs;
 mod index;
+mod types;
 
 pub unsafe fn init() {
     unsafe {