diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index c41ba501..b1c7ce02 100644 --- a/crates/vchordrq/src/types.rs +++ b/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 { + 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>), diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 8c0e9fd3..d0b2da9e 100644 --- a/src/index/gucs.rs +++ b/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 = GucSetting::::new(false); static VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS: GucSetting = GucSetting::::new(0); @@ -75,6 +83,13 @@ static mut VCHORDRQ_MAXSIM_REFINE_CONFIG: *mut pgrx::pg_sys::config_generic = co static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); +static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting = GucSetting::::new(-1); + +const VCHORDRQ_MAXSIM_CANDIDATE_LIMIT_MAX: i32 = 65_536; + +static VCHORDRQ_MAXSIM_BACKEND: GucSetting = + GucSetting::::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 { + 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() } diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 5040b587..ff0a5601 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/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(), diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 1c51c97f..40699b68 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/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>; }; + 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, _, _)>>, )| (rough, payload); - let iter: Box> = match opfamily.vector_kind() { + let coarse = match opfamily.vector_kind() { VectorKind::Vecf32 => { + let fetcher = &mut fetcher; type Op = vchordrq::operator::Op, Dot>; let unprojected = vectors .into_iter() @@ -141,7 +154,7 @@ impl SearchBuilder for MaxsimBuilder { .iter() .map(|vector| RandomProject::project(vector.as_borrowed())) .collect::>(); - 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, Dot>; let unprojected = vectors .into_iter() @@ -285,7 +300,7 @@ impl SearchBuilder for MaxsimBuilder { .iter() .map(|vector| RandomProject::project(vector.as_borrowed())) .collect::>(); - 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; let unprojected = vectors .into_iter() @@ -425,7 +442,7 @@ impl SearchBuilder for MaxsimBuilder { } }) .collect::>(); - 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; let unprojected = vectors .into_iter() @@ -565,7 +584,7 @@ impl SearchBuilder for MaxsimBuilder { } }) .collect::>(); - 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> = 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::>() - .into_iter_sorted_polyfill() - .map(|(Reverse(distance), AlwaysEqual(key))| { - let distance = distance.to_f32(); - let recheck = false; - (distance, key, recheck) - }); - let iter: Box> = 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)>, + Vec<(Distance, NonZero)>, + Distance, +); + +fn aggregate_token_searches( + iter: impl Iterator, + query_count: usize, +) -> BinaryHeap<(Reverse, 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. diff --git a/src/index/vchordrq/scanners/maxsim/rerank.rs b/src/index/vchordrq/scanners/maxsim/rerank.rs new file mode 100644 index 00000000..3f7aa26e --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/rerank.rs @@ -0,0 +1,215 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use crate::index::vchordrq::opclass::Opfamily; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use vchordrq::types::OwnedVector; + +pub(super) type HeapKey = [u16; 3]; + +#[derive(Clone, Copy, Debug)] +pub(super) struct Candidate { + pub distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) struct CandidateTensor { + pub candidate: Candidate, + pub vectors: Vec, +} + +/// Supplies the full tensor for a candidate selected by the index. +/// +/// Keeping this boundary separate from scoring lets future backends consume a +/// different source without changing candidate generation. The reference +/// source below reads the indexed column from the PostgreSQL heap. +pub(super) trait CandidateTensorSource { + fn fetch(&mut self, candidate: Candidate) -> Result, RerankError>; +} + +pub(super) struct HeapTensorSource<'a, F> { + fetcher: &'a mut F, + opfamily: Opfamily, +} + +impl<'a, F> HeapTensorSource<'a, F> { + pub fn new(fetcher: &'a mut F, opfamily: Opfamily) -> Self { + Self { fetcher, opfamily } + } +} + +impl CandidateTensorSource for HeapTensorSource<'_, F> { + fn fetch(&mut self, candidate: Candidate) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + if !tuple.filter() { + return Ok(None); + } + let (values, is_nulls) = tuple.build(); + if is_nulls[0] { + return Err(RerankError::TensorMismatch); + } + let vectors = + unsafe { self.opfamily.input_vectors(values[0]) }.ok_or(RerankError::TensorMismatch)?; + Ok(Some(CandidateTensor { candidate, vectors })) + } +} + +/// Exact MaxSim scorer independent of candidate generation and tensor storage. +pub(super) trait ExactMaxsimBackend { + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result, RerankError>; +} + +#[derive(Debug)] +pub(super) enum RerankError { + TensorMismatch, +} + +impl std::fmt::Display for RerankError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::TensorMismatch => write!(f, "MaxSim tensor kind or dimension is not matched"), + } + } +} + +#[derive(Default)] +pub(super) struct CpuExactMaxsimBackend; + +impl ExactMaxsimBackend for CpuExactMaxsimBackend { + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result, RerankError> { + let mut results = BinaryHeap::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let distance = + exact_maxsim_distance(query, &tensor.vectors).ok_or(RerankError::TensorMismatch)?; + results.push((Reverse(distance), AlwaysEqual(tensor.candidate.heap_key))); + } + Ok(results + .into_iter_sorted_polyfill() + .map(|(Reverse(distance), AlwaysEqual(heap_key))| Candidate { distance, heap_key }) + .collect()) + } +} + +fn exact_maxsim_distance(query: &[OwnedVector], document: &[OwnedVector]) -> Option { + if query.is_empty() || document.is_empty() { + return None; + } + let mut maxsim = 0.0f32; + for query_vector in query { + let mut best = Distance::INFINITY; + for document_vector in document { + best = std::cmp::min(best, document_vector.operator_dot(query_vector)?); + } + maxsim += best.to_f32(); + } + Some(Distance::from_f32(maxsim)) +} + +// Emulate unstable library feature `binary_heap_into_iter_sorted`. +trait IntoIterSortedPolyfill { + fn into_iter_sorted_polyfill(self) -> IntoIterSorted; +} + +impl IntoIterSortedPolyfill for BinaryHeap { + fn into_iter_sorted_polyfill(self) -> IntoIterSorted { + IntoIterSorted(self) + } +} + +struct IntoIterSorted(BinaryHeap); + +impl Iterator for IntoIterSorted { + type Item = T; + + fn next(&mut self) -> Option { + self.0.pop() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockSource(BTreeMap>); + + impl CandidateTensorSource for MockSource { + fn fetch(&mut self, candidate: Candidate) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + #[test] + fn cpu_backend_orders_candidates_by_exact_maxsim() { + let first = [0, 0, 1]; + let second = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]; + let mut candidates = vec![ + Candidate { + distance: Distance::from_f32(-2.0), + heap_key: second, + }, + Candidate { + distance: Distance::from_f32(-1.0), + heap_key: first, + }, + ] + .into_iter(); + let mut source = MockSource(BTreeMap::from([ + (first, vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]), + (second, vec![vector(&[0.5, 0.5])]), + ])); + + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap(); + + assert_eq!( + results.iter().map(|x| x.heap_key).collect::>(), + vec![first, second] + ); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].distance.to_f32(), -1.0); + } +} diff --git a/src/index/vchordrq/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs index b345da37..822fb23a 100644 --- a/src/index/vchordrq/scanners/mod.rs +++ b/src/index/vchordrq/scanners/mod.rs @@ -15,6 +15,7 @@ mod default; mod maxsim; +use crate::index::gucs::PostgresMaxsimBackend; use crate::index::scanners::Io; pub use default::DefaultBuilder; @@ -27,6 +28,8 @@ pub struct SearchOptions { pub max_scan_tuples: Option, pub maxsim_refine: u32, pub maxsim_threshold: u32, + pub maxsim_candidate_limit: Option, + pub maxsim_backend: PostgresMaxsimBackend, pub io_search: Io, pub io_rerank: Io, pub prefilter: bool, diff --git a/tests/vchordrq/maxsim_cpu_exact.slt b/tests/vchordrq/maxsim_cpu_exact.slt new file mode 100644 index 00000000..6851162b --- /dev/null +++ b/tests/vchordrq/maxsim_cpu_exact.slt @@ -0,0 +1,48 @@ +statement ok +CREATE TABLE t (id integer, val vector(2)[]); + +statement ok +INSERT INTO t VALUES + (1, ARRAY['[1,0]'::vector, '[0,1]'::vector]), + (2, ARRAY['[0.5,0.5]'::vector]), + (3, ARRAY['[-1,0]'::vector, '[0,-1]'::vector]); + +statement ok +CREATE INDEX t_val_idx ON t USING vchordrq (val vector_maxsim_ops) +WITH (options = $$ +build.internal.lists = [] +$$); + +statement ok +SET enable_seqscan = off; + +statement ok +SET vchordrq.probes = ''; + +statement ok +SET vchordrq.maxsim_refine = 16; + +statement ok +SET vchordrq.maxsim_candidate_limit = 3; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +query I +SELECT id +FROM t +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 3; +---- +1 +2 +3 + +statement ok +RESET vchordrq.maxsim_backend; + +statement ok +RESET vchordrq.maxsim_candidate_limit; + +statement ok +DROP TABLE t;