diff --git a/core/README.md b/core/README.md index 940d417..0ab6ecd 100644 --- a/core/README.md +++ b/core/README.md @@ -22,6 +22,18 @@ `paimon-vindex-core` contains the Rust implementations and seek-based readers for IVF-FLAT, IVF-SQ, IVF-PQ, IVF-RQ, and DiskANN. +The Rust reader supports distance range search for IVF-FLAT and IVF-RQ with +squared L2, using `DistanceBand`, `VectorRangeSearchParams`, and CSR +`RangeSearchResult` buffers. Both families support single and batch queries, +with or without a serialized Roaring allow-list, and a fixed positive `nprobe`. +IVF-FLAT tests exact distances; IVF-RQ tests its one-bit or full multi-bit +estimated distances. Results are uncapped and unordered. Probing every list +removes the IVF coverage gap, but not IVF-RQ's quantization error. The range +path does not change top-K search or the v1 storage format. + +See the [range search guide](../docs/range-search.html) for membership, +validation, filtering, and statistics. C/JNI range bindings are not included. + The DiskANN and Vamana code is an independent Apache-licensed implementation based on the published algorithms and this project's existing storage abstractions. It does not incorporate source code from Microsoft's diff --git a/core/src/collect.rs b/core/src/collect.rs index 6a354c8..5bd57ab 100644 --- a/core/src/collect.rs +++ b/core/src/collect.rs @@ -64,7 +64,7 @@ pub(crate) trait Collector { fn cutoff(&self) -> f32; /// Delivers one row, with the value the family's scan computed for it. For - /// IVF-Flat that value is an exact distance. + /// IVF-Flat that value is an exact distance; for IVF-RQ it is an estimate. /// /// Fallible because a collector may own a resource the scan cannot see: the /// oversized-list path streams chunks through a callback, and without a diff --git a/core/src/index.rs b/core/src/index.rs index 96a0a35..0f23c09 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -1846,7 +1846,7 @@ impl VectorIndexReader { /// Without this a family that cannot serve range search would report /// `Unsupported` for an invalid width or a mismatched metric, and a caller /// treating `Unsupported` as "fall back to a scan" would silently paper over - /// its own bug. The IVF-Flat reader repeats the metric check because it is a + /// its own bug. Each supporting reader repeats the metric check because it is a /// public entry point in its own right; the comparison is two enum reads, so /// the duplication costs nothing measurable. fn validate_range_request(&self, params: &VectorRangeSearchParams) -> io::Result<()> { @@ -1866,7 +1866,9 @@ impl VectorIndexReader { } /// Distance range search. For the contract see - /// [`IVFFlatIndexReader::range_search`]. + /// [`IVFFlatIndexReader::range_search`] and + /// [`IVFRQIndexReader::range_search`]. IVF-RQ membership uses estimated + /// distances rather than distances to the original vectors. /// /// The empty-band short-circuit lives **inside each family's reader**, so a /// family that cannot do range search at all still fails loud for every @@ -1880,7 +1882,7 @@ impl VectorIndexReader { self.validate_range_request(¶ms)?; match self { Self::IvfFlat(reader) => reader.range_search(query, params), - Self::IvfRq(_) => Err(range_unsupported("ivf_rq")), + Self::IvfRq(reader) => reader.range_search(query, params), Self::IvfSq(_) => Err(range_unsupported("ivf_sq")), Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), @@ -1888,7 +1890,8 @@ impl VectorIndexReader { } /// Range search restricted to a serialized Roaring allow-list. For the - /// contract see [`IVFFlatIndexReader::range_search_with_roaring_filter`]. + /// contract see [`IVFFlatIndexReader::range_search_with_roaring_filter`] + /// and [`IVFRQIndexReader::range_search`]. pub fn range_search_with_roaring_filter( &mut self, query: &[f32], @@ -1903,7 +1906,7 @@ impl VectorIndexReader { let filter = decode_roaring_filter(roaring_filter_bytes)?; match self { Self::IvfFlat(reader) => reader.range_search_with_filter(query, params, Some(&filter)), - Self::IvfRq(_) => Err(range_unsupported("ivf_rq")), + Self::IvfRq(reader) => reader.range_search_with_filter(query, params, Some(&filter)), Self::IvfSq(_) => Err(range_unsupported("ivf_sq")), Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), @@ -1911,7 +1914,8 @@ impl VectorIndexReader { } /// Batched distance range search. For the contract see - /// [`IVFFlatIndexReader::range_search`]. + /// [`IVFFlatIndexReader::range_search`] and + /// [`IVFRQIndexReader::range_search`]. pub fn range_search_batch( &mut self, queries: &[f32], @@ -1922,7 +1926,7 @@ impl VectorIndexReader { self.validate_range_request(¶ms)?; match self { Self::IvfFlat(reader) => reader.range_search_batch(queries, query_count, params), - Self::IvfRq(_) => Err(range_unsupported("ivf_rq")), + Self::IvfRq(reader) => reader.range_search_batch(queries, query_count, params), Self::IvfSq(_) => Err(range_unsupported("ivf_sq")), Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), @@ -1944,7 +1948,9 @@ impl VectorIndexReader { Self::IvfFlat(reader) => { reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter)) } - Self::IvfRq(_) => Err(range_unsupported("ivf_rq")), + Self::IvfRq(reader) => { + reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter)) + } Self::IvfSq(_) => Err(range_unsupported("ivf_sq")), Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), @@ -2669,7 +2675,7 @@ fn validate_query(query: &[f32], dimension: usize) -> io::Result<()> { validate_finite_values(query, dimension, "query") } -/// Only IVF-Flat implements range search so far. The other families return +/// IVF-Flat and IVF-RQ implement range search. The other families return /// `Unsupported`, meaning "we cannot serve this request, please fall back", /// rather than "the call has a bug". For DiskANN the reason is a lasting one: /// graph traversal is inherently k-oriented and has no natural radius diff --git a/core/src/ivfrq_io.rs b/core/src/ivfrq_io.rs index 1aee33a..fc9b866 100644 --- a/core/src/ivfrq_io.rs +++ b/core/src/ivfrq_io.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. +use crate::collect::{Collector, RangeCollector}; use crate::distance::{fvec_norm_l2sqr, preprocess_vectors, MetricType}; +use crate::index::validate_queries; use crate::index_io_util::{ decode_delta_varint_ids, encode_delta_varint_ids, pread_batched_payloads, validate_search_inputs, @@ -24,6 +26,7 @@ use crate::io::{PreadCursor, ReadRequest, SeekRead, SeekWrite}; use crate::ivfpq::RowIdFilter; use crate::ivfrq::{build_timing_enabled, log_build_elapsed, log_build_timing, IVFRQIndex}; use crate::kmeans; +use crate::range::{RangeResultBuilder, RangeSearchResult, VectorRangeSearchParams}; use crate::rq::{ is_supported_rq_bits, padded_dimension, RQCodeFactors, RQQueryContext, RQQueryTerms, RQRotation, RQVectorFactors, RaBitQuantizer, DEFAULT_RQ_ROTATION_ROUNDS, RQ_SCAN_BLOCK_SIZE, @@ -32,6 +35,7 @@ use crate::topk::TopKHeap; use rayon::prelude::*; use roaring::RoaringTreemap; use std::io; +use std::sync::Mutex; pub const IVF_RQ_MAGIC: u32 = 0x49565251; // "IVRQ" pub const IVF_RQ_VERSION: u32 = 1; @@ -609,6 +613,205 @@ impl IVFRQIndexReader { let filter = decode_roaring_filter(roaring_filter_bytes)?; self.search_with_filter(query, k, nprobe, Some(&filter)) } + + /// Returns every eligible row in the probed lists whose IVF-RQ estimated + /// distance is in the requested band. Only squared L2 is supported. + /// + /// Membership uses the one-bit estimate or, for multi-bit codes, the full + /// estimate. It does not use top-K's coarse lower-bound or FastScan pruning: + /// those bounds do not certify membership in a band of estimated distances. + /// Estimates are not clamped to zero. Even probing all lists does not make + /// membership exact with respect to the original vectors. + /// + /// Results are uncapped and unordered. Non-finite centroids, consumed + /// factors or computed distances return `InvalidData`. Filtered-out rows + /// are not evaluated. Range statistics live in the result; the last top-K + /// statistics are left unchanged. + pub fn range_search( + &mut self, + query: &[f32], + params: VectorRangeSearchParams, + ) -> io::Result { + self.range_search_with_filter(query, params, None) + } + + pub fn range_search_with_filter( + &mut self, + query: &[f32], + params: VectorRangeSearchParams, + filter: Option<&dyn RowIdFilter>, + ) -> io::Result { + self.range_search_batch_with_filter(query, 1, params, filter) + } + + /// Range search restricted to a serialized Roaring allow-list. + pub fn range_search_with_roaring_filter( + &mut self, + query: &[f32], + params: VectorRangeSearchParams, + roaring_filter_bytes: &[u8], + ) -> io::Result { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + self.range_search_with_filter(query, params, Some(&filter)) + } + + /// Batched range search with the same estimates and membership as single + /// queries. Each unique non-empty probed list is read once per call. + pub fn range_search_batch( + &mut self, + queries: &[f32], + nq: usize, + params: VectorRangeSearchParams, + ) -> io::Result { + self.range_search_batch_with_filter(queries, nq, params, None) + } + + /// Batched range search restricted to a serialized Roaring allow-list. + pub fn range_search_batch_with_roaring_filter( + &mut self, + queries: &[f32], + nq: usize, + params: VectorRangeSearchParams, + roaring_filter_bytes: &[u8], + ) -> io::Result { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + self.range_search_batch_with_filter(queries, nq, params, Some(&filter)) + } + + pub fn range_search_batch_with_filter( + &mut self, + queries: &[f32], + nq: usize, + params: VectorRangeSearchParams, + filter: Option<&dyn RowIdFilter>, + ) -> io::Result { + validate_queries(queries, nq, self.d)?; + if params.band().metric() != self.metric { + return Err(invalid_input(format!( + "band metric {:?} does not match index metric {:?}", + params.band().metric(), + self.metric + ))); + } + let nprobe = params.validate(self.nlist)?; + let mut builder = RangeResultBuilder::new(nq); + if params.band().is_empty() { + return Ok(builder.build()); + } + self.ensure_loaded()?; + if self + .quantizer_centroids + .iter() + .any(|value| !value.is_finite()) + { + return Err(invalid_data("non-finite IVF-RQ centroid")); + } + let probe_lists = queries + .par_chunks_exact(self.d) + .map(|query| { + kmeans::find_topk_checked( + query, + &self.quantizer_centroids, + self.nlist, + self.d, + nprobe, + ) + .map_err(|list_id| { + invalid_data(format!( + "non-finite IVF-RQ query-centroid distance for list {list_id}" + )) + }) + }) + .collect::>>()?; + let mut query_contexts = Vec::with_capacity(nq); + let mut scratch = vec![0.0; self.padded_d]; + for query in queries.chunks_exact(self.d) { + let mut rotated = vec![0.0; self.padded_d]; + self.rotation.rotate(query, &mut rotated, &mut scratch); + if rotated.iter().any(|value| !value.is_finite()) { + return Err(invalid_data("non-finite IVF-RQ rotated query")); + } + query_contexts.push(self.quantizer.prepare_query(rotated)); + } + let mut list_to_queries = vec![Vec::new(); self.nlist]; + let mut unique_lists = Vec::new(); + for (query_index, lists) in probe_lists.iter().enumerate() { + builder.record_lists_probed(query_index, lists.len()); + for &(distance, list_id) in lists { + if list_to_queries[list_id].is_empty() { + unique_lists.push(list_id); + } + list_to_queries[list_id].push((query_index, distance)); + } + } + + #[derive(Default)] + struct QueryMerge { + rows: Vec<(i64, f32)>, + scanned: usize, + } + let query_merges = (0..nq) + .map(|_| Mutex::new(QueryMerge::default())) + .collect::>(); + let mut batch_start = 0; + while batch_start < unique_lists.len() { + let batch_end = batch_start + self.batch_read_end(&unique_lists[batch_start..])?; + let loaded_lists = self.read_inverted_lists(&unique_lists[batch_start..batch_end])?; + for list in &loaded_lists { + if !list.ids.is_empty() { + builder.record_list_read(); + } + } + let scan_one = + |list: &RQReadList, query_index: usize, distance: f32| -> io::Result<()> { + let terms = RQQueryTerms { + g_add: distance, + g_error: distance.sqrt(), + }; + let mut collector = RangeCollector::new(params.band()); + scan_range_blocked_list( + list, + &self.quantizer, + &query_contexts[query_index], + terms, + filter, + &mut collector, + )?; + let mut output = query_merges[query_index].lock().expect("output lock"); + output.scanned += collector.scanned(); + output.rows.extend(collector.into_rows()); + Ok(()) + }; + let candidate_count = loaded_lists + .iter() + .map(|list| { + list.ids + .len() + .saturating_mul(list_to_queries[list.list_id].len()) + }) + .fold(0usize, usize::saturating_add); + if candidate_count >= PARALLEL_RQ_SCAN_MIN_CANDIDATES { + loaded_lists.par_iter().try_for_each(|list| { + list_to_queries[list.list_id].par_iter().try_for_each( + |&(query_index, distance)| scan_one(list, query_index, distance), + ) + })?; + } else { + for list in &loaded_lists { + for &(query_index, distance) in &list_to_queries[list.list_id] { + scan_one(list, query_index, distance)?; + } + } + } + batch_start = batch_end; + } + for (query_index, output) in query_merges.into_iter().enumerate() { + let output = output.into_inner().expect("output lock"); + builder.record_scanned(query_index, output.scanned); + builder.take_rows(query_index, output.rows); + } + Ok(builder.build()) + } } pub struct RQReadList { @@ -1136,6 +1339,80 @@ fn scan_blocked_list( } } +fn scan_range_blocked_list( + list: &RQReadList, + quantizer: &RaBitQuantizer, + query: &RQQueryContext, + query_terms: RQQueryTerms, + filter: Option<&dyn RowIdFilter>, + collector: &mut C, +) -> io::Result<()> { + let bits = quantizer.bits(); + let plane_size = quantizer.plane_size(); + let center = ((1usize << bits) - 1) as f32 * 0.5; + let query_sum = quantizer.query_sum(query); + let codes = list.blocked_codes(); + for block_start in (0..list.ids.len()).step_by(RQ_SCAN_BLOCK_SIZE) { + let lanes = (list.ids.len() - block_start).min(RQ_SCAN_BLOCK_SIZE); + let code_start = block_start * quantizer.code_size(); + let factor_start = block_start * quantizer.factor_fields(); + let mut allowed = [true; RQ_SCAN_BLOCK_SIZE]; + if let Some(filter) = filter { + for (lane, allowed) in allowed[..lanes].iter_mut().enumerate() { + *allowed = filter.contains(list.ids[block_start + lane]); + } + } + let mut unsigned = [0.0; RQ_SCAN_BLOCK_SIZE]; + for byte in 0..plane_size { + for lane in 0..lanes { + if allowed[lane] { + unsigned[lane] += quantizer.byte_subset_sum( + query, + byte, + codes[code_start + byte * lanes + lane], + ); + } + } + } + for value in &mut unsigned[..lanes] { + *value *= (1usize << (bits - 1)) as f32; + } + for plane in 1..bits { + let weight = (1usize << (bits - 1 - plane)) as f32; + let plane_start = code_start + plane * plane_size * lanes; + for byte in 0..plane_size { + for lane in 0..lanes { + if allowed[lane] { + unsigned[lane] += weight + * quantizer.byte_subset_sum( + query, + byte, + codes[plane_start + byte * lanes + lane], + ); + } + } + } + } + for lane in 0..lanes { + if !allowed[lane] { + continue; + } + let id = list.ids[block_start + lane]; + let field = if bits == 1 { 0 } else { 3 }; + let factors = read_block_factor(list, factor_start, lanes, field, lane, false); + if !factors.f_add.is_finite() || !factors.f_rescale.is_finite() { + return Err(invalid_data(format!( + "non-finite IVF-RQ factors for row {id}" + ))); + } + let estimate = + quantizer.estimate(unsigned[lane] - center * query_sum, factors, query_terms); + collector.push(id, estimate)?; + } + } + Ok(()) +} + fn read_block_factor( list: &RQReadList, block_start: usize, @@ -1386,6 +1663,307 @@ mod tests { } } + fn range_fixture(nlist: usize, per_list: usize) -> (IVFRQIndex, Vec) { + let dimension = 16; + let mut index = IVFRQIndex::with_bits(dimension, nlist, 4, MetricType::L2); + index.set_quantizer_centroids( + (0..nlist) + .flat_map(|list| vec![list as f32 * 16.0; dimension]) + .collect(), + ); + let vectors = (0..nlist * per_list) + .flat_map(|row| { + (0..dimension).map(move |dim| { + (row / per_list) as f32 * 16.0 + ((row * 13 + dim * 7) % 31) as f32 * 0.01 + }) + }) + .collect::>(); + let ids = (0..nlist * per_list) + .map(|row| row as i64) + .collect::>(); + index.add(&vectors, &ids, ids.len()); + assert!(index.ids.iter().all(|list| list.len() == per_list)); + let mut bytes = Vec::new(); + write_ivfrq_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + (index, bytes) + } + + fn all_range_params(nprobe: usize) -> VectorRangeSearchParams { + use crate::range::{Bound, DistanceBand}; + VectorRangeSearchParams::new( + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(), + nprobe, + ) + } + + fn range_pairs(result: &RangeSearchResult, query: usize) -> Vec<(i64, u32)> { + let query = result.query(query); + let mut pairs = query + .labels + .iter() + .zip(query.distances) + .map(|(&id, &distance)| (id, distance.to_bits())) + .collect::>(); + pairs.sort_unstable(); + pairs + } + + #[test] + fn ivfrq_range_direct_entry_points_and_empty_band_validation() { + use crate::range::{Bound, DistanceBand}; + let (index, bytes) = range_fixture(3, 37); + let stats = Arc::new(Mutex::new(ReaderStats::default())); + let mut reader = IVFRQIndexReader::open(CountingReader { + inner: Cursor::new(bytes), + stats: Arc::clone(&stats), + max_ranges_per_pread: 0, + }) + .unwrap(); + let query = vec![0.2; index.d]; + let queries = query.repeat(2); + let mut filter = RoaringTreemap::new(); + filter.insert(1); + filter.insert(50); + let mut filter_bytes = Vec::new(); + filter.serialize_into(&mut filter_bytes).unwrap(); + let empty = VectorRangeSearchParams::new( + DistanceBand::new(Bound::Finite(1.0), Bound::Finite(1.0), MetricType::L2).unwrap(), + 3, + ); + let calls = stats.lock().unwrap().calls; + for result in [ + reader.range_search(&query, empty).unwrap(), + reader + .range_search_with_roaring_filter(&query, empty, &filter_bytes) + .unwrap(), + reader.range_search_batch(&queries, 2, empty).unwrap(), + reader + .range_search_batch_with_roaring_filter(&queries, 2, empty, &filter_bytes) + .unwrap(), + ] { + assert!(result.labels().is_empty()); + assert_eq!(result.lims().len(), result.query_count() + 1); + assert_eq!(result.call_stats().list_reads(), 0); + assert_eq!(result.query(0).stats.lists_probed(), 0); + } + for params in [empty, all_range_params(3)] { + for query in [ + vec![0.0; index.d - 1], + vec![f32::NAN; index.d], + vec![f32::INFINITY; index.d], + vec![f32::NEG_INFINITY; index.d], + ] { + assert_eq!( + reader.range_search(&query, params).unwrap_err().kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + reader + .range_search_batch(&query, 1, params) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + } + assert_eq!( + reader + .range_search_with_roaring_filter(&query, params, &[255]) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + reader + .range_search_batch_with_roaring_filter(&queries, 2, params, &[255]) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + reader + .range_search_batch(&[], 0, params) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + reader + .range_search_batch(&[], usize::MAX, params) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + reader + .range_search_batch(&query, 2, params) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + } + assert_eq!( + reader + .range_search(&query, all_range_params(0)) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + let wrong_metric = VectorRangeSearchParams::new( + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::Cosine).unwrap(), + 3, + ); + assert_eq!( + reader + .range_search(&query, wrong_metric) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!(stats.lock().unwrap().calls, calls); + let params = all_range_params(99); + let single = reader.range_search(&query, params).unwrap(); + let batch = reader.range_search_batch(&queries, 2, params).unwrap(); + let filtered = reader + .range_search_with_roaring_filter(&query, params, &filter_bytes) + .unwrap(); + let filtered_batch = reader + .range_search_batch_with_roaring_filter(&queries, 2, params, &filter_bytes) + .unwrap(); + assert_eq!(single.labels().len(), 111); + assert_eq!(filtered.labels().len(), 2); + for query_index in 0..2 { + assert_eq!(range_pairs(&single, 0), range_pairs(&batch, query_index)); + assert_eq!( + range_pairs(&filtered, 0), + range_pairs(&filtered_batch, query_index) + ); + } + } + + #[test] + fn ivfrq_range_counts_empty_lists_and_honors_read_capabilities() { + let (mut index, _) = range_fixture(5, 37); + index.ids[2].clear(); + index.codes[2].clear(); + index.factors[2].clear(); + let mut bytes = Vec::new(); + write_ivfrq_index(&index, &mut PosWriter::new(&mut bytes)).unwrap(); + for (max_ranges, expected_calls) in [(0, 1), (1, 4), (2, 2)] { + let stats = Arc::new(Mutex::new(ReaderStats::default())); + let mut reader = IVFRQIndexReader::open(CountingReader { + inner: Cursor::new(bytes.clone()), + stats: Arc::clone(&stats), + max_ranges_per_pread: max_ranges, + }) + .unwrap(); + reader.ensure_loaded().unwrap(); + *stats.lock().unwrap() = ReaderStats::default(); + let result = reader + .range_search_batch(&vec![0.2; index.d * 3], 3, all_range_params(5)) + .unwrap(); + assert_eq!(stats.lock().unwrap().calls, expected_calls); + assert_eq!( + stats.lock().unwrap().max_ranges_per_batch, + if max_ranges == 0 { 4 } else { max_ranges } + ); + assert_eq!(result.call_stats().list_reads(), 4); + for query in 0..3 { + assert_eq!(result.query(query).stats.lists_probed(), 5); + assert_eq!(result.query(query).stats.rows_scanned(), 4 * 37); + assert_eq!(result.query(query).stats.rows_committed(), 4 * 37); + } + } + } + + #[test] + fn ivfrq_range_parallel_single_and_batch_match_one_worker() { + struct RendezvousFilter { + workers: AtomicU64, + parallel: bool, + gate: Mutex<()>, + ready: std::sync::Condvar, + } + + impl RowIdFilter for RendezvousFilter { + fn contains(&self, _id: i64) -> bool { + let worker = rayon::current_thread_index().expect("scan runs inside the pool"); + let mask = 1u64 << worker; + let previous = self.workers.fetch_or(mask, Ordering::Relaxed); + if self.parallel && previous & mask == 0 { + let guard = self.gate.lock().unwrap(); + self.ready.notify_all(); + let (_guard, timeout) = self + .ready + .wait_timeout_while(guard, std::time::Duration::from_secs(10), |_| { + self.workers.load(Ordering::Relaxed).count_ones() < 2 + }) + .unwrap(); + assert!( + !timeout.timed_out(), + "a second scan worker must reach the filter" + ); + } + true + } + } + + for (nlist, nq) in [(8, 1), (8, 3), (1, 8)] { + let (index, bytes) = range_fixture(nlist, 1057); + let queries = (0..nq) + .flat_map(|query| vec![query as f32 * 0.1; index.d]) + .collect::>(); + let mut expected = None; + for workers in [1, 4] { + let stats = Arc::new(Mutex::new(ReaderStats::default())); + let mut reader = IVFRQIndexReader::open(CountingReader { + inner: Cursor::new(bytes.clone()), + stats: Arc::clone(&stats), + max_ranges_per_pread: 0, + }) + .unwrap(); + reader.ensure_loaded().unwrap(); + *stats.lock().unwrap() = ReaderStats::default(); + let filter = RendezvousFilter { + workers: AtomicU64::new(0), + parallel: workers > 1, + gate: Mutex::new(()), + ready: std::sync::Condvar::new(), + }; + let result = rayon::ThreadPoolBuilder::new() + .num_threads(workers) + .build() + .unwrap() + .install(|| { + reader + .range_search_batch_with_filter( + &queries, + nq, + all_range_params(nlist), + Some(&filter), + ) + .unwrap() + }); + let pairs = (0..nq) + .map(|query| range_pairs(&result, query)) + .collect::>(); + if let Some(expected) = &expected { + assert_eq!(&pairs, expected); + assert!(filter.workers.load(Ordering::Relaxed).count_ones() > 1); + } else { + expected = Some(pairs); + } + assert_eq!(stats.lock().unwrap().calls, 1); + assert_eq!(result.call_stats().list_reads(), nlist); + for query in 0..nq { + assert_eq!(result.query(query).stats.rows_scanned(), nlist * 1057); + assert_eq!(result.query(query).stats.rows_committed(), nlist * 1057); + assert_eq!(result.query(query).stats.early_abandoned(), 0); + } + assert_eq!(reader.last_search_stats(), IVFRQSearchStats::default()); + } + } + } + #[test] fn ivfrq_four_bit_roundtrip_uses_blocked_layout() { let d = 13; diff --git a/core/src/kmeans.rs b/core/src/kmeans.rs index 6ec4926..a9352aa 100644 --- a/core/src/kmeans.rs +++ b/core/src/kmeans.rs @@ -770,6 +770,47 @@ pub fn find_topk( (indices, distances) } +pub(crate) fn find_topk_checked( + point: &[f32], + centroids: &[f32], + centroid_count: usize, + dimension: usize, + nprobe: usize, +) -> Result, usize> { + let nprobe = nprobe.min(centroid_count); + if nprobe == 0 { + return Ok(Vec::new()); + } + let capacity = nprobe.saturating_mul(2).min(centroid_count); + let mut candidates = Vec::with_capacity(capacity); + let mut cutoff = (f32::INFINITY, usize::MAX); + for centroid in 0..centroid_count { + let distance = fvec_l2sqr( + point, + ¢roids[centroid * dimension..(centroid + 1) * dimension], + ); + if !distance.is_finite() { + return Err(centroid); + } + let candidate = (distance, centroid); + if compare_distance_then_index(&candidate, &cutoff).is_ge() { + continue; + } + candidates.push(candidate); + if candidates.len() == capacity && nprobe < capacity { + candidates.select_nth_unstable_by(nprobe - 1, compare_distance_then_index); + cutoff = candidates[nprobe - 1]; + candidates.truncate(nprobe); + } + } + if nprobe < candidates.len() { + candidates.select_nth_unstable_by(nprobe - 1, compare_distance_then_index); + candidates.truncate(nprobe); + } + candidates.sort_unstable_by(compare_distance_then_index); + Ok(candidates) +} + /// Batch find top-nprobe nearest centroids using SGEMM, with direct L2 fallback when /// cancellation error can dominate the nearest computed distance. /// Returns (all_indices, all_distances) each of length nq * nprobe. @@ -1449,6 +1490,83 @@ mod tests { assert_eq!(indices[0], 0); } + #[test] + fn test_find_topk_checked_matches_full_sort_with_bounded_candidates() { + let mut random = StdRng::seed_from_u64(104); + let centroid_count = 257; + for dimension in [1, 13, 32, 129] { + let centroids = (0..centroid_count * dimension) + .map(|_| random.gen_range(-8.0f32..8.0)) + .collect::>(); + for _ in 0..4 { + let query = (0..dimension) + .map(|_| random.gen_range(-8.0f32..8.0)) + .collect::>(); + let mut expected = centroids + .chunks_exact(dimension) + .enumerate() + .map(|(centroid, values)| (fvec_l2sqr(&query, values), centroid)) + .collect::>(); + expected.sort_by(compare_distance_then_index); + for width in [0, 1, 2, 7, 16, 128, 129, 256, 257, usize::MAX] { + let nprobe = width.min(centroid_count); + let actual = + find_topk_checked(&query, ¢roids, centroid_count, dimension, width) + .unwrap(); + assert_eq!(actual, expected[..nprobe]); + assert!(actual.capacity() <= nprobe * 2); + } + } + } + } + + #[test] + fn test_find_topk_checked_preserves_ties_across_candidate_compactions() { + let centroids = (0..128) + .map(|centroid| ((127 - centroid) / 4) as f32) + .collect::>(); + let actual = find_topk_checked(&[0.0], ¢roids, centroids.len(), 1, 7).unwrap(); + assert_eq!( + actual, + vec![ + (0.0, 124), + (0.0, 125), + (0.0, 126), + (0.0, 127), + (1.0, 120), + (1.0, 121), + (1.0, 122), + ] + ); + let tied = find_topk_checked(&[0.0], &[1.0; 128], 128, 1, 3).unwrap(); + assert_eq!(tied, vec![(1.0, 0), (1.0, 1), (1.0, 2)]); + } + + #[test] + fn test_find_topk_checked_rejects_unselected_nonfinite_distances() { + for invalid in [1e20, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + for position in [0, 31, 63] { + let mut centroids = vec![1.0; 64]; + centroids[position] = invalid; + assert_eq!( + find_topk_checked(&[0.0], ¢roids, 64, 1, 1), + Err(position) + ); + } + } + assert_eq!(find_topk_checked(&[f32::MAX], &[0.0, 1.0], 2, 1, 1), Err(0)); + } + + #[test] + fn test_find_topk_checked_keeps_direct_distances_for_large_offsets() { + let query = [1e20f32]; + let centroids = [query[0], f32::from_bits(query[0].to_bits() + 1)]; + let actual = find_topk_checked(&query, ¢roids, 2, 1, 2).unwrap(); + assert_eq!(actual[0], (0.0, 0)); + assert_eq!(actual[1], (fvec_l2sqr(&query, ¢roids[1..]), 1)); + assert!(actual[1].0.is_finite()); + } + #[test] fn test_kmeans_rejects_invalid_data_shapes_before_early_return() { let config = KMeansConfig::default(); diff --git a/core/src/range.rs b/core/src/range.rs index c08c209..7f7bb50 100644 --- a/core/src/range.rs +++ b/core/src/range.rs @@ -360,7 +360,8 @@ impl RangeSearchStats { /// A diagnostic, not a work measure: a row is counted whether the kernel /// stopped at its first term or at its last, so this is not the number of /// rows whose evaluation was short-circuited. `rows_scanned` counts these - /// rows too. + /// rows too. IVF-RQ evaluates the complete estimate of every eligible row + /// without early abandonment, so its count is always zero. pub fn early_abandoned(&self) -> usize { self.early_abandoned } diff --git a/core/tests/range_search.rs b/core/tests/range_search.rs index 7600f83..0918373 100644 --- a/core/tests/range_search.rs +++ b/core/tests/range_search.rs @@ -30,7 +30,10 @@ use paimon_vindex_core::index::{VectorIndexReader, VectorSearchParams}; use paimon_vindex_core::io::PosWriter; use paimon_vindex_core::ivfflat::IVFFlatIndex; use paimon_vindex_core::ivfflat_io::write_ivfflat_index; +use paimon_vindex_core::ivfrq::IVFRQIndex; +use paimon_vindex_core::ivfrq_io::write_ivfrq_index; use paimon_vindex_core::range::{Bound, DistanceBand, QueryResult, VectorRangeSearchParams}; +use paimon_vindex_core::rq::RQRotation; use std::collections::HashSet; use std::io::Cursor; @@ -38,6 +41,628 @@ use roaring::RoaringTreemap; type Reader = VectorIndexReader>>; +fn rq_fixture(dimension: usize, bits: usize, nlist: usize, per_list: usize) -> IVFRQIndex { + let mut index = IVFRQIndex::with_bits(dimension, nlist, bits, MetricType::L2); + index.set_quantizer_centroids( + (0..nlist) + .flat_map(|list| (0..dimension).map(move |dim| list as f32 * 4.0 + dim as f32 * 0.01)) + .collect(), + ); + let mut state = 7181u64; + let mut vectors = Vec::new(); + for list in 0..nlist { + for _ in 0..per_list { + for dim in 0..dimension { + state = state.wrapping_mul(6364136223846793005).wrapping_add(1); + let noise = ((state >> 40) as f32 / (1u32 << 24) as f32 - 0.5) * 2.0; + vectors.push(index.quantizer_centroids()[list * dimension + dim] + noise); + } + } + } + let ids = (0..nlist * per_list) + .map(|row| 1000 + (nlist * per_list - row) as i64) + .collect::>(); + index.add(&vectors, &ids, ids.len()); + assert!(index.ids.iter().all(|ids| ids.len() == per_list)); + index +} + +fn rq_bytes(index: &IVFRQIndex) -> Vec { + let mut bytes = Vec::new(); + write_ivfrq_index(index, &mut PosWriter::new(&mut bytes)).unwrap(); + bytes +} + +fn rq_reader(index: &IVFRQIndex) -> Reader { + VectorIndexReader::open(Cursor::new(rq_bytes(index))).unwrap() +} + +fn rq_estimated_oracle(index: &IVFRQIndex, query: &[f32], nprobe: usize) -> Vec<(i64, f32)> { + let rotation = RQRotation::new(index.d, index.rotation_seed, index.rotation_rounds); + let mut rotated = vec![0.0; index.padded_d]; + rotation.rotate(query, &mut rotated, &mut vec![0.0; index.padded_d]); + let query_sum = rotated.iter().sum::(); + let mut lists = (0..index.nlist) + .map(|list| { + let distance = fvec_l2sqr( + query, + &index.quantizer_centroids()[list * index.d..(list + 1) * index.d], + ); + (list, distance) + }) + .collect::>(); + lists.sort_by(|left, right| left.1.total_cmp(&right.1).then(left.0.cmp(&right.0))); + let mut rows = Vec::new(); + for (list, coarse_distance) in lists.into_iter().take(nprobe) { + for (position, &id) in index.ids[list].iter().enumerate() { + let code = &index.codes[list] + [position * index.code_size()..(position + 1) * index.code_size()]; + let byte_sum = |plane: usize, byte: usize| { + let pattern = code[plane * index.plane_size() + byte]; + (0..8) + .rev() + .filter(|bit| pattern & (1 << bit) != 0) + .fold(0.0, |sum, bit| sum + rotated[byte * 8 + bit]) + }; + let mut unsigned = (0..index.plane_size()) + .map(|byte| byte_sum(0, byte)) + .sum::() + * (1usize << (index.bits - 1)) as f32; + for plane in 1..index.bits { + let weight = (1usize << (index.bits - 1 - plane)) as f32; + for byte in 0..index.plane_size() { + unsigned += weight * byte_sum(plane, byte); + } + } + let center = ((1usize << index.bits) - 1) as f32 * 0.5; + let factors = index.factors[list][position].full; + let estimate = factors.f_add + + coarse_distance + + factors.f_rescale * (unsigned - center * query_sum); + rows.push((id, estimate)); + } + } + rows +} + +#[test] +fn rq_range_matches_estimated_distance_oracle() { + for bits in 1..=8 { + for dimension in [13, 256] { + let index = rq_fixture(dimension, bits, 4, 37); + let query = (0..dimension) + .map(|dim| dim as f32 * 0.01 + 0.2) + .collect::>(); + for nprobe in [1, 3, 4] { + let oracle = rq_estimated_oracle(&index, &query, nprobe); + let mut values = oracle.iter().map(|row| row.1).collect::>(); + values.sort_by(f32::total_cmp); + let lower = values[values.len() / 4].max(0.0); + let upper = values[values.len() * 3 / 4].max(lower); + for band in [ + l2(lower, upper), + DistanceBand::new(Bound::Unbounded, Bound::Finite(upper), MetricType::L2) + .unwrap(), + DistanceBand::new(Bound::Finite(lower), Bound::Unbounded, MetricType::L2) + .unwrap(), + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(), + ] { + let result = rq_reader(&index) + .range_search(&query, VectorRangeSearchParams::new(band, nprobe)) + .unwrap(); + let expected = oracle + .iter() + .copied() + .filter(|row| band.admit(row.1)) + .collect(); + assert_eq!( + pairs_of(result.query(0)), + bits_of(expected), + "bits={bits}, dimension={dimension}" + ); + assert_eq!(result.query(0).stats.rows_scanned(), nprobe * 37); + assert_eq!(result.query(0).stats.early_abandoned(), 0); + assert_eq!( + result.query(0).stats.rows_committed(), + result.labels().len() + ); + assert_eq!(result.call_stats().list_reads(), nprobe); + } + } + } + } +} + +fn rq_all_distances() -> DistanceBand { + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap() +} + +#[test] +fn rq_range_bounded_probes_match_oracle_across_many_lists() { + let index = rq_fixture(13, 4, 65, 3); + let queries = [0.2, 126.0, 256.2] + .into_iter() + .flat_map(|base| (0..index.d).map(move |dimension| base + dimension as f32 * 0.01)) + .collect::>(); + let allowed = index + .ids + .iter() + .flatten() + .copied() + .filter(|row| row % 2 == 0) + .collect::>(); + let filter = serialize_roaring(&allowed); + for width in [1, 3, 16, 33, 65, usize::MAX] { + let nprobe = width.min(index.nlist); + for band in [rq_all_distances(), l2(0.0, 256.0)] { + let params = VectorRangeSearchParams::new(band, width); + let mut reader = rq_reader(&index); + let batch = reader.range_search_batch(&queries, 3, params).unwrap(); + let filtered = reader + .range_search_batch_with_roaring_filter(&queries, 3, params, &filter) + .unwrap(); + for (query_index, query) in queries.chunks_exact(index.d).enumerate() { + let expected = rq_estimated_oracle(&index, query, nprobe) + .into_iter() + .filter(|row| band.admit(row.1)) + .collect::>(); + assert_eq!( + pairs_of(batch.query(query_index)), + bits_of(expected.clone()) + ); + let single = reader.range_search(query, params).unwrap(); + assert_eq!(pairs_of(single.query(0)), bits_of(expected.clone())); + assert_eq!( + pairs_of(filtered.query(query_index)), + bits_of( + expected + .into_iter() + .filter(|row| allowed.contains(&row.0)) + .collect() + ) + ); + assert_eq!(batch.query(query_index).stats.rows_scanned(), nprobe * 3); + } + } + } +} + +#[test] +fn rq_range_batch_single_filters_and_statistics_agree() { + for bits in [1, 4, 8] { + let index = rq_fixture(256, bits, 4, 37); + let queries = [0.2, 4.2, 2.0, 8.7, 0.2] + .into_iter() + .flat_map(|offset| (0..index.d).map(move |dim| dim as f32 * 0.01 + offset)) + .collect::>(); + let all_ids = index.ids.iter().flatten().copied().collect::>(); + let sparse_ids = all_ids + .iter() + .copied() + .filter(|id| id % 3 == 0) + .collect::>(); + for nprobe in [1, 2, 9] { + let params = VectorRangeSearchParams::new(l2(10.0, 2000.0), nprobe); + let mut reader = rq_reader(&index); + let unfiltered = reader.range_search_batch(&queries, 5, params).unwrap(); + for allowed in [&all_ids, &sparse_ids, &HashSet::new()] { + let filter = serialize_roaring(allowed); + let batch = reader + .range_search_batch_with_roaring_filter(&queries, 5, params, &filter) + .unwrap(); + let mut probed_lists = HashSet::new(); + for (query_index, query) in queries.chunks_exact(index.d).enumerate() { + let oracle = rq_estimated_oracle(&index, query, nprobe); + for (list, ids) in index.ids.iter().enumerate() { + if oracle.iter().any(|row| ids.contains(&row.0)) { + probed_lists.insert(list); + } + } + let single = reader + .range_search_with_roaring_filter(query, params, &filter) + .unwrap(); + let unfiltered_single = reader.range_search(query, params).unwrap(); + assert_eq!( + pairs_of(single.query(0)), + pairs_of(batch.query(query_index)) + ); + assert_eq!( + pairs_of(unfiltered_single.query(0)), + pairs_of(unfiltered.query(query_index)) + ); + let expected = oracle + .iter() + .copied() + .filter(|row| allowed.contains(&row.0) && params.band().admit(row.1)) + .collect(); + assert_eq!(pairs_of(batch.query(query_index)), bits_of(expected)); + let stats = batch.query(query_index).stats; + assert_eq!(stats.lists_probed(), nprobe.min(index.nlist)); + assert_eq!( + stats.rows_scanned(), + oracle.iter().filter(|row| allowed.contains(&row.0)).count() + ); + assert_eq!( + stats.rows_committed(), + batch.query(query_index).labels.len() + ); + assert_eq!(stats.early_abandoned(), 0); + assert_eq!( + batch.lims()[query_index + 1] - batch.lims()[query_index], + stats.rows_committed() + ); + if allowed == &all_ids { + assert_eq!( + pairs_of(batch.query(query_index)), + pairs_of(unfiltered.query(query_index)) + ); + } + } + assert_eq!(batch.call_stats().list_reads(), probed_lists.len()); + } + } + } +} + +#[test] +fn rq_range_membership_is_estimated_not_exact() { + let dimension = 13; + let mut index = IVFRQIndex::with_bits(dimension, 1, 1, MetricType::L2); + index.set_quantizer_centroids(vec![0.0; dimension]); + let vectors = (0..96 * dimension) + .map(|value| ((value * 37 % 113) as f32 - 56.0) / 23.0) + .collect::>(); + let ids = (0..96).collect::>(); + index.add(&vectors, &ids, ids.len()); + let query = vec![0.3; dimension]; + let oracle = rq_estimated_oracle(&index, &query, 1); + let (witness, estimated, exact) = oracle + .iter() + .find_map(|&(id, estimate)| { + let exact = fvec_l2sqr( + &query, + &vectors[id as usize * dimension..(id as usize + 1) * dimension], + ); + (estimate > 0.0 && (estimate - exact).abs() > 0.01).then_some((id, estimate, exact)) + }) + .expect("fixture must distinguish estimated and exact distances"); + let band = l2(0.0, (estimated + exact) * 0.5); + let result = rq_reader(&index) + .range_search(&query, VectorRangeSearchParams::new(band, 1)) + .unwrap(); + assert_eq!(result.labels().contains(&witness), band.admit(estimated)); + assert_ne!(result.labels().contains(&witness), band.admit(exact)); + assert_eq!( + pairs_of(result.query(0)), + bits_of(oracle.into_iter().filter(|row| band.admit(row.1)).collect()) + ); +} + +#[test] +fn rq_range_does_not_prune_on_coarse_bounds_or_clamp_estimates() { + for bits in [1, 4, 8] { + let mut index = rq_fixture(13, bits, 1, 37); + let negative_id = index.ids[0][0]; + for (row, factors) in index.factors[0].iter_mut().enumerate() { + factors.coarse.f_add = 1000.0; + factors.coarse.f_rescale = 0.0; + factors.coarse.f_error = 0.0; + factors.full.f_add = if row == 0 { -1.0 } else { 0.5 }; + factors.full.f_rescale = 0.0; + if bits == 1 { + factors.coarse = factors.full; + } + } + let query = index.quantizer_centroids().to_vec(); + let mut reader = rq_reader(&index); + let all = reader + .range_search(&query, VectorRangeSearchParams::new(rq_all_distances(), 1)) + .unwrap(); + assert_eq!(all.labels().len(), 37); + assert!(pairs_of(all.query(0)).contains(&(negative_id, (-1.0f32).to_bits()))); + let finite = reader + .range_search(&query, VectorRangeSearchParams::new(l2(0.0, 1.0), 1)) + .unwrap(); + assert_eq!(finite.labels().len(), 36); + assert!(!finite.labels().contains(&negative_id)); + assert_eq!(finite.query(0).stats.early_abandoned(), 0); + } +} + +#[test] +fn rq_range_preserves_topk_results_and_last_search_stats() { + for bits in [1, 4] { + let index = rq_fixture(256, bits, 4, 65); + let query = vec![0.3; index.d]; + let filter = serialize_roaring( + &index + .ids + .iter() + .flatten() + .copied() + .filter(|id| id % 2 == 0) + .collect(), + ); + let mut reader = rq_reader(&index); + for filtered in [false, true] { + let params = VectorSearchParams::new(7, 3); + let before = if filtered { + reader + .search_with_roaring_filter(&query, params, &filter) + .unwrap() + } else { + reader.search(&query, params).unwrap() + }; + let stats = reader.ivfrq_search_stats(); + let range = reader + .range_search(&query, VectorRangeSearchParams::new(rq_all_distances(), 3)) + .unwrap(); + assert!(range.labels().len() > 7); + assert_eq!(reader.ivfrq_search_stats(), stats); + let after = if filtered { + reader + .search_with_roaring_filter(&query, params, &filter) + .unwrap() + } else { + reader.search(&query, params).unwrap() + }; + assert_eq!(before, after); + assert_eq!(reader.ivfrq_search_stats(), stats); + } + } +} + +#[test] +fn rq_range_nonfinite_factors_fail_only_when_consumed() { + use std::io::ErrorKind; + for bits in [1, 4] { + for nonfinite in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + for field in [0, 1] { + let mut index = rq_fixture(13, bits, 2, 37); + let bad_id = index.ids[0][36]; + let factors = if bits == 1 { + &mut index.factors[0][36].coarse + } else { + &mut index.factors[0][36].full + }; + if field == 0 { + factors.f_add = nonfinite; + } else { + factors.f_rescale = nonfinite; + } + let query = vec![0.2; index.d]; + let params = VectorRangeSearchParams::new(l2(0.0, 0.01), 2); + let mut reader = rq_reader(&index); + assert_eq!( + reader.range_search(&query, params).unwrap_err().kind(), + ErrorKind::InvalidData + ); + assert_eq!( + reader + .range_search_batch(&query.repeat(2), 2, params) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); + let only_bad = serialize_roaring(&HashSet::from([bad_id])); + assert_eq!( + reader + .range_search_with_roaring_filter(&query, params, &only_bad) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); + assert_eq!( + reader + .range_search_batch_with_roaring_filter( + &query.repeat(2), + 2, + params, + &only_bad + ) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); + let only_good = serialize_roaring( + &index + .ids + .iter() + .flatten() + .copied() + .filter(|&id| id != bad_id) + .collect(), + ); + let result = reader + .range_search_with_roaring_filter(&query, params, &only_good) + .unwrap(); + assert_eq!(result.query(0).stats.rows_scanned(), 73); + } + } + } + let mut index = rq_fixture(13, 4, 1, 37); + for factors in &mut index.factors[0] { + factors.coarse.f_add = f32::NAN; + factors.coarse.f_rescale = f32::NEG_INFINITY; + factors.coarse.f_error = f32::INFINITY; + } + let query = vec![0.2; index.d]; + let result = rq_reader(&index) + .range_search(&query, VectorRangeSearchParams::new(rq_all_distances(), 1)) + .unwrap(); + assert_eq!( + pairs_of(result.query(0)), + bits_of(rq_estimated_oracle(&index, &query, 1)) + ); +} + +#[test] +fn rq_range_rejects_nonfinite_centroids_and_arithmetic_overflow() { + use paimon_vindex_core::ivfrq_io::IVF_RQ_HEADER_SIZE; + use std::io::ErrorKind; + let index = rq_fixture(13, 4, 2, 37); + let params = VectorRangeSearchParams::new(rq_all_distances(), 1); + for nonfinite in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + let mut bytes = rq_bytes(&index); + let offset = IVF_RQ_HEADER_SIZE + index.d * 4; + bytes[offset..offset + 4].copy_from_slice(&nonfinite.to_le_bytes()); + let mut reader = VectorIndexReader::open(Cursor::new(bytes)).unwrap(); + assert_eq!( + reader + .range_search(&vec![0.0; index.d], params) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); + } + assert_eq!( + rq_reader(&index) + .range_search(&vec![f32::MAX; index.d], params) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); + assert_eq!( + rq_reader(&index) + .range_search(&vec![1e20; index.d], params) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); + let mut index = rq_fixture(13, 4, 1, 37); + for factors in &mut index.factors[0] { + factors.full.f_add = f32::MAX; + factors.full.f_rescale = f32::MAX; + } + let mut reader = rq_reader(&index); + assert_eq!( + reader + .range_search(&vec![1.0; index.d], params) + .unwrap_err() + .kind(), + ErrorKind::InvalidData + ); +} + +#[test] +fn rq_range_rejects_overflow_in_an_unselected_centroid() { + use paimon_vindex_core::ivfrq_io::IVF_RQ_HEADER_SIZE; + use std::io::ErrorKind; + + for nlist in [2, 65] { + let index = rq_fixture(13, 4, nlist, 37); + let mut bytes = rq_bytes(&index); + let offset = IVF_RQ_HEADER_SIZE + (nlist - 1) * index.d * 4; + bytes[offset..offset + 4].copy_from_slice(&1e20f32.to_le_bytes()); + let mut reader = VectorIndexReader::open(Cursor::new(bytes)).unwrap(); + let params = VectorRangeSearchParams::new(rq_all_distances(), 1); + let query = vec![0.0; index.d]; + let queries = query.repeat(2); + let filter = serialize_roaring(&index.ids[0].iter().copied().collect()); + for error in [ + reader.range_search(&query, params).unwrap_err(), + reader.range_search_batch(&queries, 2, params).unwrap_err(), + reader + .range_search_with_roaring_filter(&query, params, &filter) + .unwrap_err(), + reader + .range_search_batch_with_roaring_filter(&queries, 2, params, &filter) + .unwrap_err(), + ] { + assert_eq!(error.kind(), ErrorKind::InvalidData); + } + } +} + +#[test] +fn rq_range_filters_preserve_signed_ids_and_fixed_probe_width() { + let mut index = rq_fixture(13, 4, 2, 37); + index.ids[0][0] = i64::MIN; + index.ids[0][1] = -1; + index.ids[0][2] = 1i64 << 33; + index.ids[0][3] = i64::MAX; + let mut filter = RoaringTreemap::new(); + for id in [i64::MIN as u64, u64::MAX, 1u64 << 33, i64::MAX as u64] { + filter.insert(id); + } + let mut bytes = Vec::new(); + filter.serialize_into(&mut bytes).unwrap(); + let query = vec![0.2; index.d]; + let params = VectorRangeSearchParams::new(rq_all_distances(), 1); + let mut reader = rq_reader(&index); + let all = reader.range_search(&query, params).unwrap(); + assert!(all.labels().contains(&i64::MIN)); + assert!(all.labels().contains(&-1)); + let filtered = reader + .range_search_with_roaring_filter(&query, params, &bytes) + .unwrap(); + assert_eq!( + filtered.labels().iter().copied().collect::>(), + HashSet::from([1i64 << 33, i64::MAX]) + ); + let far_list = serialize_roaring(&index.ids[1].iter().copied().collect()); + let far = reader + .range_search_with_roaring_filter(&query, params, &far_list) + .unwrap(); + assert!(far.labels().is_empty()); + assert_eq!(far.query(0).stats.lists_probed(), 1); + assert_eq!(far.query(0).stats.rows_scanned(), 0); + assert_eq!(far.call_stats().list_reads(), 1); +} + +#[test] +fn rq_range_unified_metric_capability_and_validation_precedence() { + use std::io::ErrorKind; + for metric in [MetricType::Cosine, MetricType::InnerProduct] { + let mut index = IVFRQIndex::with_bits(13, 1, 4, metric); + index.set_quantizer_centroids(vec![0.0; 13]); + index.add(&[1.0; 13], &[1], 1); + let mut reader = rq_reader(&index); + let filter = serialize_roaring(&HashSet::from([1])); + for band in [ + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, metric).unwrap(), + DistanceBand::new(Bound::Finite(1.0), Bound::Finite(1.0), metric).unwrap(), + ] { + let params = VectorRangeSearchParams::new(band, 1); + for error in [ + reader.range_search(&[1.0; 13], params).unwrap_err(), + reader + .range_search_batch(&[1.0; 26], 2, params) + .unwrap_err(), + reader + .range_search_with_roaring_filter(&[1.0; 13], params, &filter) + .unwrap_err(), + reader + .range_search_batch_with_roaring_filter(&[1.0; 26], 2, params, &filter) + .unwrap_err(), + ] { + assert_eq!(error.kind(), ErrorKind::Unsupported); + } + assert_eq!( + reader + .range_search(&[f32::NAN; 13], params) + .unwrap_err() + .kind(), + ErrorKind::InvalidInput + ); + assert_eq!( + reader + .range_search_with_roaring_filter(&[1.0; 13], params, &[255]) + .unwrap_err() + .kind(), + ErrorKind::InvalidInput + ); + assert_eq!( + reader + .range_search(&[1.0; 13], VectorRangeSearchParams::new(band, 0)) + .unwrap_err() + .kind(), + ErrorKind::InvalidInput + ); + } + } +} + // --- fixtures -------------------------------------------------------------- // // Indexes are built by assigning rows to lists **explicitly** rather than by diff --git a/docs/api.html b/docs/api.html index 800df0a..bc9b7c2 100644 --- a/docs/api.html +++ b/docs/api.html @@ -76,7 +76,7 @@

Shared search parameters

Range search parameters and results

-

Range search returns every probed row inside a half-open distance band instead of a fixed number of nearest neighbours. It is currently available on IVF-FLAT with the l2 metric; the other families and metrics report the request as unsupported so a caller can fall back. See Range search for the full semantic contract.

+

Rust range search returns every eligible probed row inside a half-open distance band instead of a fixed number of nearest neighbours. All four entry points support IVF-FLAT and IVF-RQ with l2: IVF-FLAT tests exact distances, while IVF-RQ tests estimated distances. The other families and metrics report Unsupported; callers requiring exact, complete membership need an exhaustive scan, not top-K followed by filtering. See Range search for the full contract. C/JNI range bindings are not included.

diff --git a/docs/index.html b/docs/index.html index a458ad1..b82b871 100644 --- a/docs/index.html +++ b/docs/index.html @@ -124,7 +124,7 @@

Distance range search support

- + diff --git a/docs/ivf-rq.html b/docs/ivf-rq.html index 49cd529..f1e233c 100644 --- a/docs/ivf-rq.html +++ b/docs/ivf-rq.html @@ -13,7 +13,7 @@

Multi-bit rotated residual quantization

IVF-RQ

Partition vectors with IVF, spread each residual through a deterministic orthogonal transform, and store centered scalar levels as bit planes. A cheap sign-plane estimate rejects weak candidates before the remaining planes are evaluated.

4 bits / dimension by default1–8 build-time bitstwo-stage scanMagic: IVRQ
- +

Positioning and trade-offs

@@ -61,6 +61,15 @@

Usage

let params = VectorSearchParams::new(10, 64);
+
+

Estimated-distance range search

+

The Rust VectorIndexReader and IVFRQIndexReader support single/batch range search, with or without a serialized Roaring allow-list, using the shared DistanceBand, VectorRangeSearchParams, and RangeSearchResult. Only squared L2 and fixed positive nprobe are supported. See the shared usage examples.

+

Each eligible row is evaluated using its one-bit estimate or full multi-bit estimate from all stored bit planes. Membership is the half-open test lower <= estimate < upper, not a comparison against the original vector's exact distance. Negative estimates remain negative. Quantization can cause both missing rows and extra rows relative to an exact-distance band, even with every list probed.

+

Unlike top-K, range search does not use the coarse error-bound or quantized FastScan pruning stages: those do not certify membership in an estimate band. It uses F32 lookup sums for every allowed row, then hands the estimate to the shared collector. Probe selection computes direct L2 distances to every centroid, rejects non-finite values, and reuses selected distances in the estimate. A bounded candidate buffer uses O(nprobe) distance-selection storage per query, without skipping distance validation for unselected centroids. This range-only path uses the same calculations for single and batch queries, rather than top-K's SGEMM probe optimization. Lists are read once per call and shared across queries; sufficiently large scans run in parallel. Results have no fixed-K truncation, padding, or ordering guarantee.

+

Range statistics are returned with the result. rows_scanned counts filter-eligible rows, early_abandoned is zero, and list_reads counts unique non-empty lists. The last top-K statistics and all existing top-K behavior are unchanged. No storage-format change or C/JNI range binding is introduced.

+
Correctness is relative to the estimateTests compare range output against an independently accumulated estimated-distance oracle; they do not establish exact-distance recall guarantees. The Top-K recall and throughput measurements on this page are not range-search measurements. Use IVF-FLAT with full probing or an exhaustive raw-vector scan when exact, complete membership is required.
+
+

Parameters

Entry pointSignature
Single queryrange_search(query, params)
Single query, filteredrange_search_with_roaring_filter(query, params, filter_bytes)
IndexRange searchMembership decided onNote
IVF-FLATSupported, l2 onlyExact distanceA probed row's membership is exact, so at nprobe = nlist the in-band set is complete — among allow-listed rows, when a filter is supplied
IVF-RQPlannedEstimateAn estimate on the wrong side of a cut removes the row entirely, so prefer IVF-FLAT when completeness matters; accuracy will be measured in the contribution that enables it
IVF-RQSupported, l2 onlyOne-bit or full multi-bit estimateSingle/batch and Roaring-filtered variants; even nprobe = nlist does not remove quantization-induced membership errors
IVF-SQPlannedEstimateAs above
IVF-PQPlannedEstimateAs above; would use the float-LUT path rather than fastscan
DiskANNNot planned—Graph traversal is inherently k-oriented and has no natural radius termination criterion
ParameterRequirement / defaultPurposeGuidance
dimensionInferred by Java/Python one-shot training; otherwise > 0Logical vector dimensionStorage pads internally to a multiple of 64.
nlistAuto from expected-vector-count, or explicit > 0IVF partition countCompare the resolved value with the same IVF-FLAT baseline.
rq.bits1–8; auto from max-bytes-per-vector, otherwise 4Persisted residual level widthHigher values increase recall, file bytes, I/O, and scan work linearly.
metricRequiredL2 / inner product / cosineSemantic, not inferred; fixed in the file.
ivf.train.max-points-per-centroidPositive integer; default 256Maximum training vectors per coarse centroidCaps coarse K-means input at nlist × value vectors, including hierarchical clustering stages.
ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it.
nprobeAutomatic by default; explicit expert overrideLists probedAuto accounts for K, average list size, and filter selectivity.
diff --git a/docs/range-search.html b/docs/range-search.html index 8488229..8a8f107 100644 --- a/docs/range-search.html +++ b/docs/range-search.html @@ -11,31 +11,32 @@
VIPaimon Vector Index
-

Every row inside a distance band

Range search

Return every probed row whose distance to the query falls inside a half-open band [lower, upper), with no result limit. Range search answers "which rows are within this distance", where Top-K answers "which rows are closest".

No limit, no capHalf-open intervalIVF-FLAT onlyL2 only
+

Every row inside a distance band

Range search

Return every eligible probed row whose family-specific distance falls inside a half-open band [lower, upper), with no result limit. IVF-FLAT computes exact distances; IVF-RQ computes estimates. Range search answers "which rows are within this distance", where Top-K answers "which rows are closest".

No limit, no capHalf-open intervalIVF-FLAT and IVF-RQL2 only

Semantic contract

-
IntervalHalf-open [lower, upper)
Result sizeUnbounded
Row orderUnspecified
MembershipExact distance (IVF-FLAT)
+
IntervalHalf-open [lower, upper)
Result sizeUnbounded
Row orderUnspecified
MembershipExact (FLAT), estimated (RQ)

A band is left-closed and right-open. A row at exactly lower is returned; a row at exactly upper is not. This is what makes adjacent bands tile a range without overlapping or leaving gaps, so [a,b) and [b,c) together return exactly what [a,c) returns.

Either side may be unbounded, and unboundedness is a distinct state rather than a large or small number. A band with both sides unbounded is the whole space and is legal. An empty band where lower == upper is also legal and returns zero rows.

-
Why unboundedness is not a sentinel valueUsing a finite value to mean "unbounded" loses rows on two of the three metrics. Cosine distance is not clamped, so two identical normalized vectors can score slightly below zero, and a lower of 0.0 would exclude the most similar rows. An inner-product distance can be exactly f32::MAX, which a right-open interval ending at that value excludes precisely. Squared L2 happens to be safe with 0.0, but that is a coincidence of one metric.
+
Why unboundedness is not a sentinel valueA finite cut is not a substitute for an unbounded side. IVF-RQ squared-L2 estimates can be negative and are not clamped: a lower cut of 0.0 excludes them, while Bound::Unbounded includes them. Likewise, a finite upper cut excludes a value exactly at that cut. Use structural unboundedness to request every finite estimate.

Units and ordering

Cuts are expressed in the index's own distance space, not in whatever unit the caller happens to think in.

+

For IVF-RQ, this is the raw estimate in squared-L2 units. Negative estimates have no real-valued Euclidean radius. The endpoint examples describe non-negative squared distances; endpoint conversion does not turn estimated membership into an exact predicate over the original vectors.

MetricInternal valueNote
l2Squared Euclidean distanceA caller holding a Euclidean radius must square it, or use endpoint derivation below
cosine1 - cos, same direction as L2Not yet certified for range search
inner_product-inner_product, direction reversedNot yet certified for range search

Row order is not part of the contract. Within one list rows come back in physical order, but no order across lists is specified or promised, and two runs of the same query may differ. Do not depend on any observed order: sorting is the caller's job, and in SQL it is ORDER BY's. Results are neither padded nor sorted, which is how range search differs from Top-K.

-
A row within a few ULP of the upper cutThe scan abandons a row as soon as its partially accumulated distance passes the upper cut, and the cut is used exactly as given, with no margin. That is sound because the value compared while pruning is a prefix of the value that will be committed: one kernel accumulates the distance and reads its own running sum, rather than a second kernel accumulating in a different order. Adding a non-negative float is non-decreasing under IEEE rounding, so a partial sum above the cut proves the completed distance is above it. An earlier design did use two kernels whose sums disagreed by a few ULP once the dimension reached 128, and had to widen the cut to cover the divergence; sharing one accumulation removes the divergence instead of budgeting for it, and removes the second pass over an admitted row along with it.
+
A row within a few ULP of the upper cutIVF-FLAT can abandon a row when its partially accumulated squared distance passes the upper cut: it uses the same non-decreasing accumulation for pruning and the committed distance. IVF-RQ does not use this rule or top-K's coarse/FastScan bounds. It evaluates the complete estimate from F32 lookup sums before testing either cut, so every eligible row reaches the band test. Both families use the supplied cuts without a margin.

"No cap" is not "complete"

Range search never truncates its result. That is a promise about not dropping rows it found, and it is not a promise that it found every in-band row in the file.

Only the nprobe nearest lists are probed, so a row lying inside the band but in an unprobed list is not returned. A smaller nprobe returns no more in-band rows, and potentially fewer: if every matching row already lies in the lists it still probes, the result is unchanged. At nprobe == nlist every list is probed and, because IVF-FLAT computes exact distances, the result is then the complete in-band set.

-
The two are different knobsProbe coverage is a recall concern and is raised with nprobe. It is not the same as a result-count limit, which is a result-set concern. This version has no result-count limit at all, so for an unfiltered search probe coverage is the only reason a row can be missing. A filtered search has one more: a row the allow-list excludes was never eligible in the first place.
+
Coverage and estimation are different gapsRaising nprobe improves list coverage; it does not remove IVF-RQ's quantization error. An estimate can lie on the other side of a cut from the exact distance, producing missing or extra rows relative to an exact-distance predicate even at nprobe = nlist. A filter additionally excludes rows that were never eligible. Neither family truncates the rows admitted by its own distance calculation. Top-K plus post-filtering is not an equivalent fallback; exact, complete membership requires full-probe IVF-FLAT or an exhaustive raw-vector scan.
@@ -78,18 +79,21 @@

Usage

// A one-sided band: everything at or beyond 2.0, with no upper end. let tail = DistanceBand::new(Bound::Finite(2.0), Bound::Unbounded, MetricType::L2)?;

Results use a CSR layout, so a batch of queries shares three contiguous buffers. lims holds query_count + 1 offsets, and query i owns labels[lims[i]..lims[i+1]] together with the matching slice of distances. Per-query counters are available through query(i).stats, and counters covering the whole call through call_stats().

+

Both families expose range_search, range_search_batch, and their _with_roaring_filter variants in Rust. The filter is an allow-list and does not widen the fixed nprobe. A query has the same label/distance multiset alone or in a batch; order remains unspecified. Unique non-empty lists are read once per call and shared across queries.

+

For IVF-RQ, lists_probed includes empty selected lists; rows_scanned counts filter-eligible rows evaluated; rows_committed counts returned rows; and early_abandoned is zero. Call-level list_reads counts unique non-empty lists, not query/list pairs or storage read rounds. These result-owned counters leave the last top-K statistics unchanged.

Choosing an index type

-

This version implements range search for IVF-FLAT only. The other IVF families report that the request is unsupported so a caller can fall back, and they gain support in later work.

-
Compressed IVF families are not supported yet.IVF-FLAT stores full f32 vectors and computes the distance used for band membership directly, so a probed row's membership is exact. IVF-RQ, IVF-SQ and IVF-PQ store only codes, so their membership would be decided on an estimate, and an estimate on the wrong side of a cut removes the row from the result entirely rather than merely degrading an ordering as it would in Top-K. All three return Unsupported in this release. Their membership semantics and reproducible accuracy measurements will be documented in the contributions that enable them; this release makes no recall claim for unsupported families.
+

This version implements range search for IVF-FLAT and IVF-RQ. IVF-SQ, IVF-PQ and DiskANN still return Unsupported. Only L2, fixed probe widths, and Rust entry points are included.

+
Choose according to the membership requirement.IVF-FLAT tests full-vector distances. IVF-RQ uses RaBitQ, with a one-bit estimate for one-bit files and the full multi-bit estimate otherwise, not Faiss's residual/additive quantizer. Its band predicate is precise relative to that estimate, not to the raw vector. Tests cover an independent estimated-distance oracle, single/batch equivalence, filters, statistics, parallel scans, and non-finite inputs/data; they do not establish exact-distance recall guarantees. See IVF-RQ range semantics.

Fail-loud combinations

-

Two error classes are distinguished, and the difference is actionable. Invalid input means the call itself is wrong and is worth fixing as a bug. Unsupported means this request cannot be served and the caller should fall back to a scan.

-
SituationClass
Inverted band, where lower > upperInvalid input
A non-finite cutInvalid input
A negative cut under squared L2Invalid input
An operator that does not match its sideInvalid input
A band whose metric differs from the index's metricInvalid input
nprobe of 0Invalid input
A query whose dimension differs from the index's, or holds a non-finite valueInvalid input
A malformed Roaring filterInvalid input
cosine or inner_product, not yet certifiedUnsupported
IVF-RQ, IVF-SQ, IVF-PQ or DiskANNUnsupported
Automatic probe width, not yet implementedUnsupported
An endpoint with no representable cut, meaning do not push the predicate downUnsupported
+

Invalid input means the call itself is wrong. Unsupported means the request cannot be served. Invalid data covers corrupt consumed index data and non-finite computed distances; no partial result is returned.

+
SituationClass
Inverted band, where lower > upperInvalid input
A non-finite cutInvalid input
A negative cut under squared L2Invalid input
An operator that does not match its sideInvalid input
A band whose metric differs from the index's metricInvalid input
nprobe of 0Invalid input
A query whose dimension differs from the index's, or holds a non-finite valueInvalid input
A malformed Roaring filterInvalid input
cosine or inner_product, not yet certifiedUnsupported
IVF-SQ, IVF-PQ or DiskANNUnsupported
An endpoint with no representable cut, meaning do not push the predicate downUnsupported
A non-finite IVF-RQ centroid, consumed factor, or computed estimateInvalid data
+

IVF-RQ validates centroids and every direct query-centroid distance before selecting lists for a non-empty band, including distances to lists that would not be selected. It requires finite f_add and f_rescale for the estimate it consumes: coarse for one-bit codes, full for multi-bit codes. Multi-bit coarse factors, including f_error, are not used by range search and are not validated on this path. Filtered-out and unprobed rows are not evaluated. Finite inputs can still overflow during rotation, query-centroid distance calculation, or estimation, which also returns InvalidData.

An empty band is not an error: it returns zero rows. It also does not mask a bad call. The dimension, metric and width are all validated before the empty band takes its shortcut, and a family that cannot do range search at all rejects every band, the empty one included. A malformed Roaring filter is likewise rejected before that shortcut.

@@ -102,6 +106,6 @@

Why not DiskANN

-
Apache Paimon Vector IndexRange search · IVF-FLAT
+