feat(vchordrq): add pluggable exact MaxSim reranking

crates/vchordrq/src/types.rs

@@ -18,6 +18,7 @@ use validator::{Validate, ValidationError};
 use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned};
 use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned};
 use vector::vect::{VectBorrowed, VectOwned};
+use vector::{VectorBorrowed, VectorOwned};
 
 #[derive(Debug, Clone, Serialize, Deserialize, Validate)]
 #[serde(deny_unknown_fields)]
@@ -61,6 +62,26 @@ pub enum OwnedVector {
     Rabitq4(Rabitq4Owned),
 }
 
+impl OwnedVector {
+    pub fn operator_dot(&self, rhs: &Self) -> Option<distance::Distance> {
+        match (self, rhs) {
+            (Self::Vecf32(lhs), Self::Vecf32(rhs)) => {
+                Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed()))
+            }
+            (Self::Vecf16(lhs), Self::Vecf16(rhs)) => {
+                Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed()))
+            }
+            (Self::Rabitq8(lhs), Self::Rabitq8(rhs)) => {
+                Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed()))
+            }
+            (Self::Rabitq4(lhs), Self::Rabitq4(rhs)) => {
+                Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed()))
+            }
+            _ => None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy)]
 pub enum BorrowedVector<'a> {
     Vecf32(VectBorrowed<'a, f32>),

src/index/gucs.rs

@@ -27,6 +27,14 @@ pub enum PostgresIo {
     ReadStream,
 }
 
+#[derive(Debug, Clone, Copy, PostgresGucEnum)]
+pub enum PostgresMaxsimBackend {
+    #[name = c"coarse_only"]
+    CoarseOnly,
+    #[name = c"cpu_exact"]
+    CpuExact,
+}
+
 static VCHORDRQ_QUERY_SAMPLING_ENABLE: GucSetting<bool> = GucSetting::<bool>::new(false);
 
 static VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS: GucSetting<i32> = GucSetting::<i32>::new(0);
@@ -75,6 +83,13 @@ static mut VCHORDRQ_MAXSIM_REFINE_CONFIG: *mut pgrx::pg_sys::config_generic = co
 
 static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting<i32> = GucSetting::<i32>::new(0);
 
+static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting<i32> = GucSetting::<i32>::new(-1);
+
+const VCHORDRQ_MAXSIM_CANDIDATE_LIMIT_MAX: i32 = 65_536;
+
+static VCHORDRQ_MAXSIM_BACKEND: GucSetting<PostgresMaxsimBackend> =
+    GucSetting::<PostgresMaxsimBackend>::new(PostgresMaxsimBackend::CoarseOnly);
+
 static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic =
     core::ptr::null_mut();
 
@@ -151,6 +166,24 @@ pub fn init() {
         GucContext::Userset,
         GucFlags::default(),
     );
+    GucRegistry::define_int_guc(
+        c"vchordrq.maxsim_candidate_limit",
+        c"Maximum number of index candidates passed to exact MaxSim reranking.",
+        c"A positive value is required when maxsim_backend is cpu_exact.",
+        &VCHORDRQ_MAXSIM_CANDIDATE_LIMIT,
+        -1,
+        VCHORDRQ_MAXSIM_CANDIDATE_LIMIT_MAX,
+        GucContext::Userset,
+        GucFlags::default(),
+    );
+    GucRegistry::define_enum_guc(
+        c"vchordrq.maxsim_backend",
+        c"Backend used after MaxSim candidate generation.",
+        c"coarse_only preserves existing behavior; cpu_exact reads full tensors from the heap.",
+        &VCHORDRQ_MAXSIM_BACKEND,
+        GucContext::Userset,
+        GucFlags::default(),
+    );
     GucRegistry::define_bool_guc(
         c"vchordrq.prefilter",
         c"`prefilter` argument of vchordrq.",
@@ -472,6 +505,15 @@ pub fn vchordrq_maxsim_threshold(index: pgrx::pg_sys::Relation) -> u32 {
     }
 }
 
+pub fn vchordrq_maxsim_candidate_limit() -> Option<u32> {
+    let value = VCHORDRQ_MAXSIM_CANDIDATE_LIMIT.get();
+    if value < 0 { None } else { Some(value as u32) }
+}
+
+pub fn vchordrq_maxsim_backend() -> PostgresMaxsimBackend {
+    VCHORDRQ_MAXSIM_BACKEND.get()
+}
+
 pub fn vchordrq_prefilter() -> bool {
     VCHORDRQ_PREFILTER.get()
 }

src/index/vchordrq/am/mod.rs

@@ -514,6 +514,8 @@ pub unsafe extern "C-unwind" fn amrescan(
             max_scan_tuples: gucs::vchordrq_max_scan_tuples(),
             maxsim_refine: gucs::vchordrq_maxsim_refine((*scan).indexRelation),
             maxsim_threshold: gucs::vchordrq_maxsim_threshold((*scan).indexRelation),
+            maxsim_candidate_limit: gucs::vchordrq_maxsim_candidate_limit(),
+            maxsim_backend: gucs::vchordrq_maxsim_backend(),
             io_search: gucs::vchordrq_io_search(),
             io_rerank: gucs::vchordrq_io_rerank(),
             prefilter: gucs::vchordrq_prefilter(),

src/index/vchordrq/scanners/maxsim.rs

@@ -12,7 +12,11 @@
 //
 // Copyright (c) 2025-2026 TensorChord Inc.
 
+mod rerank;
+
+use self::rerank::{Candidate, CpuExactMaxsimBackend, ExactMaxsimBackend, HeapTensorSource};
 use crate::index::fetcher::*;
+use crate::index::gucs::PostgresMaxsimBackend;
 use crate::index::scanners::{Io, SearchBuilder};
 use crate::index::vchordrq::dispatch::*;
 use crate::index::vchordrq::filter::filter;
@@ -99,10 +103,18 @@ impl SearchBuilder for MaxsimBuilder {
         }
         let maxsim_refine = options.maxsim_refine;
         let maxsim_threshold = options.maxsim_threshold;
+        let maxsim_backend = options.maxsim_backend;
+        let maxsim_candidate_limit = options.maxsim_candidate_limit;
+        if matches!(maxsim_backend, PostgresMaxsimBackend::CpuExact)
+            && !matches!(maxsim_candidate_limit, Some(1..))
+        {
+            pgrx::error!("cpu_exact MaxSim requires a positive vchordrq.maxsim_candidate_limit");
+        }
         let opfamily = self.opfamily;
         let Some(vectors) = vectors else {
             return Box::new(std::iter::empty()) as Box<dyn Iterator<Item = (f32, [u16; 3], bool)>>;
         };
+        let exact_query = vectors.clone();
         let method = how(index);
         if !matches!(method, RerankMethod::Index) {
             pgrx::error!("maxsim search with rerank_in_table is not supported");
@@ -124,8 +136,9 @@ impl SearchBuilder for MaxsimBuilder {
                 _,
                 AlwaysEqual<PackedRefMut8<(NonZero<u64>, _, _)>>,
             )| (rough, payload);
-        let iter: Box<dyn Iterator<Item = _>> = match opfamily.vector_kind() {
+        let coarse = match opfamily.vector_kind() {
             VectorKind::Vecf32 => {
+                let fetcher = &mut fetcher;
                 type Op = vchordrq::operator::Op<VectOwned<f32>, Dot>;
                 let unprojected = vectors
                     .into_iter()
@@ -141,7 +154,7 @@ impl SearchBuilder for MaxsimBuilder {
                     .iter()
                     .map(|vector| RandomProject::project(vector.as_borrowed()))
                     .collect::<Vec<_>>();
-                Box::new((0..n).map(move |i| {
+                let token_searches = (0..n).map(move |i| {
                     let (results, estimation_by_threshold) = match options.io_search {
                         Io::Plain => maxsim_search::<_, Op>(
                             index,
@@ -267,9 +280,11 @@ impl SearchBuilder for MaxsimBuilder {
                         rough_set.extend(rough_iter.map(rough_map));
                     }
                     (accu_set, rough_set, estimation_by_threshold)
-                }))
+                });
+                aggregate_token_searches(token_searches, n)
             }
             VectorKind::Vecf16 => {
+                let fetcher = &mut fetcher;
                 type Op = vchordrq::operator::Op<VectOwned<f16>, Dot>;
                 let unprojected = vectors
                     .into_iter()
@@ -285,7 +300,7 @@ impl SearchBuilder for MaxsimBuilder {
                     .iter()
                     .map(|vector| RandomProject::project(vector.as_borrowed()))
                     .collect::<Vec<_>>();
-                Box::new((0..n).map(move |i| {
+                let token_searches = (0..n).map(move |i| {
                     let (results, estimation_by_threshold) = match options.io_search {
                         Io::Plain => maxsim_search::<_, Op>(
                             index,
@@ -411,9 +426,11 @@ impl SearchBuilder for MaxsimBuilder {
                         rough_set.extend(rough_iter.map(rough_map));
                     }
                     (accu_set, rough_set, estimation_by_threshold)
-                }))
+                });
+                aggregate_token_searches(token_searches, n)
             }
             VectorKind::Rabitq8 => {
+                let fetcher = &mut fetcher;
                 type Op = vchordrq::operator::Op<Rabitq8Owned, Dot>;
                 let unprojected = vectors
                     .into_iter()
@@ -425,7 +442,7 @@ impl SearchBuilder for MaxsimBuilder {
                         }
                     })
                     .collect::<Vec<_>>();
-                Box::new((0..n).map(move |i| {
+                let token_searches = (0..n).map(move |i| {
                     let (results, estimation_by_threshold) = match options.io_search {
                         Io::Plain => maxsim_search::<_, Op>(
                             index,
@@ -551,9 +568,11 @@ impl SearchBuilder for MaxsimBuilder {
                         rough_set.extend(rough_iter.map(rough_map));
                     }
                     (accu_set, rough_set, estimation_by_threshold)
-                }))
+                });
+                aggregate_token_searches(token_searches, n)
             }
             VectorKind::Rabitq4 => {
+                let fetcher = &mut fetcher;
                 type Op = vchordrq::operator::Op<Rabitq4Owned, Dot>;
                 let unprojected = vectors
                     .into_iter()
@@ -565,7 +584,7 @@ impl SearchBuilder for MaxsimBuilder {
                         }
                     })
                     .collect::<Vec<_>>();
-                Box::new((0..n).map(move |i| {
+                let token_searches = (0..n).map(move |i| {
                     let (results, estimation_by_threshold) = match options.io_search {
                         Io::Plain => maxsim_search::<_, Op>(
                             index,
@@ -691,55 +710,35 @@ impl SearchBuilder for MaxsimBuilder {
                         rough_set.extend(rough_iter.map(rough_map));
                     }
                     (accu_set, rough_set, estimation_by_threshold)
-                }))
+                });
+                aggregate_token_searches(token_searches, n)
             }
         };
-        let mut updates = Vec::new();
-        let mut estimations = Vec::new();
-        for (query_id, (accu_set, rough_set, estimation_by_threshold)) in iter.enumerate() {
-            updates.reserve(accu_set.len() + rough_set.len());
-            let is_empty = accu_set.is_empty() && rough_set.is_empty();
-            let mut estimation_by_scope = Distance::NEG_INFINITY;
-            for (distance, payload) in accu_set {
-                estimation_by_scope = std::cmp::max(estimation_by_scope, distance);
-                let (key, _) = pointer_to_kv(payload);
-                updates.push((key, query_id, distance));
+        let iter: Box<dyn Iterator<Item = _>> = match maxsim_backend {
+            PostgresMaxsimBackend::CoarseOnly => Box::new(
+                coarse
+                    .into_iter_sorted_polyfill()
+                    .map(|(Reverse(distance), AlwaysEqual(key))| (distance.to_f32(), key, false)),
+            ),
+            PostgresMaxsimBackend::CpuExact => {
+                let mut candidates = coarse
+                    .into_iter_sorted_polyfill()
+                    .take(maxsim_candidate_limit.unwrap() as usize)
+                    .map(|(Reverse(distance), AlwaysEqual(heap_key))| Candidate {
+                        distance,
+                        heap_key,
+                    });
+                let mut source = HeapTensorSource::new(&mut fetcher, opfamily);
+                let results = CpuExactMaxsimBackend
+                    .rerank(&exact_query, &mut candidates, &mut source)
+                    .unwrap_or_else(|error| pgrx::error!("{error}"));
+                Box::new(
+                    results
+                        .into_iter()
+                        .map(|candidate| (candidate.distance.to_f32(), candidate.heap_key, false)),
+                )
             }
-            for (distance, payload) in rough_set {
-                let (key, _) = pointer_to_kv(payload);
-                updates.push((key, query_id, distance));
-            }
-            estimations.push(if !is_empty {
-                std::cmp::max(estimation_by_scope, estimation_by_threshold)
-            } else {
-                Distance::ZERO
-            });
-        }
-        updates.sort_unstable_by_key(|&(key, ..)| key);
-        let iter = updates
-            .chunk_by(|(kl, ..), (kr, ..)| kl == kr)
-            .map(|chunk| {
-                let key = chunk[0].0;
-                let mut value = vec![None; n];
-                for &(_, query_id, distance) in chunk {
-                    let this = value[query_id].get_or_insert(Distance::INFINITY);
-                    *this = std::cmp::min(*this, distance);
-                }
-                let mut maxsim = 0.0f32;
-                for (query_id, distance) in value.into_iter().enumerate() {
-                    let d = distance.unwrap_or(estimations[query_id]);
-                    maxsim += Distance::to_f32(d);
-                }
-                (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key))
-            })
-            .collect::<BinaryHeap<_>>()
-            .into_iter_sorted_polyfill()
-            .map(|(Reverse(distance), AlwaysEqual(key))| {
-                let distance = distance.to_f32();
-                let recheck = false;
-                (distance, key, recheck)
-            });
-        let iter: Box<dyn Iterator<Item = _>> = Box::new(iter);
+        };
         let iter = if let Some(max_scan_tuples) = options.max_scan_tuples {
             Box::new(iter.take(max_scan_tuples as _))
         } else {
@@ -750,6 +749,57 @@ impl SearchBuilder for MaxsimBuilder {
     }
 }
 
+type TokenSearchResult = (
+    Vec<(Distance, NonZero<u64>)>,
+    Vec<(Distance, NonZero<u64>)>,
+    Distance,
+);
+
+fn aggregate_token_searches(
+    iter: impl Iterator<Item = TokenSearchResult>,
+    query_count: usize,
+) -> BinaryHeap<(Reverse<Distance>, AlwaysEqual<[u16; 3]>)> {
+    let mut updates = Vec::new();
+    let mut estimations = Vec::new();
+    for (query_id, (accu_set, rough_set, estimation_by_threshold)) in iter.enumerate() {
+        updates.reserve(accu_set.len() + rough_set.len());
+        let is_empty = accu_set.is_empty() && rough_set.is_empty();
+        let mut estimation_by_scope = Distance::NEG_INFINITY;
+        for (distance, payload) in accu_set {
+            estimation_by_scope = std::cmp::max(estimation_by_scope, distance);
+            let (key, _) = pointer_to_kv(payload);
+            updates.push((key, query_id, distance));
+        }
+        for (distance, payload) in rough_set {
+            let (key, _) = pointer_to_kv(payload);
+            updates.push((key, query_id, distance));
+        }
+        estimations.push(if !is_empty {
+            std::cmp::max(estimation_by_scope, estimation_by_threshold)
+        } else {
+            Distance::ZERO
+        });
+    }
+    updates.sort_unstable_by_key(|&(key, ..)| key);
+    updates
+        .chunk_by(|(left, ..), (right, ..)| left == right)
+        .map(|chunk| {
+            let key = chunk[0].0;
+            let mut value = vec![None; query_count];
+            for &(_, query_id, distance) in chunk {
+                let this = value[query_id].get_or_insert(Distance::INFINITY);
+                *this = std::cmp::min(*this, distance);
+            }
+            let maxsim = value
+                .into_iter()
+                .enumerate()
+                .map(|(query_id, distance)| distance.unwrap_or(estimations[query_id]).to_f32())
+                .sum();
+            (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key))
+        })
+        .collect()
+}
+
 // Emulate unstable library feature `binary_heap_into_iter_sorted`.
 // See https://github.com/rust-lang/rust/issues/59278.