diff --git a/core/README.md b/core/README.md index 3a22afb..5f297de 100644 --- a/core/README.md +++ b/core/README.md @@ -22,16 +22,28 @@ `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, IVF-RQ, and IVF-SQ with +The Rust reader supports distance range search for IVF-FLAT, IVF-RQ, IVF-SQ, and IVF-PQ with squared L2, using `DistanceBand`, `VectorRangeSearchParams`, and CSR -`RangeSearchResult` buffers. All three families support single and batch queries, +`RangeSearchResult` buffers. All four 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; IVF-SQ tests scalar-quantized estimates. Results are uncapped -and unordered. Probing every list removes the IVF coverage gap, but not the -quantization error of IVF-RQ or IVF-SQ. The range path does not change top-K +estimated distances; IVF-SQ tests scalar-quantized estimates; IVF-PQ tests +floating-point ADC estimates for 4-bit and +8-bit codes, with optional residual encoding and OPQ. Results are uncapped and +unordered. Probing every list removes the IVF coverage gap, but not the +compressed families' quantization error. The range path does not change top-K search or the v1 storage format. +PQ range uses direct squared-L2 subvector lookup tables and sums their selected +entries in subquantizer order. It does not use top-K's u8 FastScan tables or +precomputed norm identities; membership is independent of list size, batch +size, and `optimize_for_search`. Finite estimates can therefore differ from +top-K's distances. Every filter-eligible row is fully evaluated, with no early +abandonment. Non-finite consumed estimates, rotated queries, or coarse distances +return `InvalidData`, including overflow and distances to unselected centroids. +Unique non-empty lists are read once per call, with oversized lists streamed +through the existing bounded reader. DiskANN range remains unsupported. + See the [range search guide](../docs/range-search.html) for membership, validation, filtering, and statistics. C/JNI range bindings are not included. @@ -61,7 +73,7 @@ Large lists stream in bounded chunks; scan scratch is reused, and a finite upper allows entire SQ blocks to stop after their partial distances reach that cut. Result memory still grows with the number of hits. Cache hits are excluded from `call_stats().list_reads()`. Range support does not extend to other metrics, -IVF-PQ, DiskANN, or language bindings in this change. +DiskANN, or language bindings. The crate ships its [normative v1 storage-format specification](STORAGE_FORMAT.md) and byte-exact fixtures. Project documentation, language bindings, and diff --git a/core/src/collect.rs b/core/src/collect.rs index 1146f12..6647aeb 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; for IVF-RQ and IVF-SQ it is an + /// IVF-Flat that value is an exact distance; for IVF-RQ, IVF-SQ and IVF-PQ it is an /// estimate. /// /// Fallible because a collector may own a resource the scan cannot see: the diff --git a/core/src/index.rs b/core/src/index.rs index 8baa879..9598b80 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -1867,8 +1867,9 @@ impl VectorIndexReader { /// Distance range search. For the contract see /// [`IVFFlatIndexReader::range_search`] (exact distances), - /// [`IVFRQIndexReader::range_search`] (RQ estimates), and - /// [`IVFSQIndexReader::range_search`] (SQ estimates). Only L2 is supported. + /// [`IVFRQIndexReader::range_search`] (RQ estimates), + /// [`IVFSQIndexReader::range_search`] (SQ estimates), and + /// [`IVFPQIndexReader::range_search`] (PQ estimates). Only L2 is supported. /// /// 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 @@ -1884,15 +1885,16 @@ impl VectorIndexReader { Self::IvfFlat(reader) => reader.range_search(query, params), Self::IvfRq(reader) => reader.range_search(query, params), Self::IvfSq(reader) => reader.range_search(query, params), - Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), + Self::IvfPq(reader) => reader.range_search(query, params), Self::DiskAnn(_) => Err(range_unsupported("diskann")), } } /// Range search restricted to a serialized Roaring allow-list. For the /// contract see [`IVFFlatIndexReader::range_search_with_roaring_filter`], - /// [`IVFRQIndexReader::range_search`], and - /// [`IVFSQIndexReader::range_search_with_roaring_filter`]. + /// [`IVFRQIndexReader::range_search`], + /// [`IVFSQIndexReader::range_search_with_roaring_filter`], and + /// [`IVFPQIndexReader::range_search_with_roaring_filter`]. pub fn range_search_with_roaring_filter( &mut self, query: &[f32], @@ -1909,7 +1911,7 @@ impl VectorIndexReader { Self::IvfFlat(reader) => reader.range_search_with_filter(query, params, Some(&filter)), Self::IvfRq(reader) => reader.range_search_with_filter(query, params, Some(&filter)), Self::IvfSq(reader) => reader.range_search_with_filter(query, params, Some(&filter)), - Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), + Self::IvfPq(reader) => reader.range_search_with_filter(query, params, Some(&filter)), Self::DiskAnn(_) => Err(range_unsupported("diskann")), } } @@ -1928,7 +1930,7 @@ impl VectorIndexReader { Self::IvfFlat(reader) => reader.range_search_batch(queries, query_count, params), Self::IvfRq(reader) => reader.range_search_batch(queries, query_count, params), Self::IvfSq(reader) => reader.range_search_batch(queries, query_count, params), - Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), + Self::IvfPq(reader) => reader.range_search_batch(queries, query_count, params), Self::DiskAnn(_) => Err(range_unsupported("diskann")), } } @@ -1954,7 +1956,9 @@ impl VectorIndexReader { Self::IvfSq(reader) => { reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter)) } - Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), + Self::IvfPq(reader) => { + reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter)) + } Self::DiskAnn(_) => Err(range_unsupported("diskann")), } } diff --git a/core/src/io.rs b/core/src/io.rs index 3067be1..95c8c76 100644 --- a/core/src/io.rs +++ b/core/src/io.rs @@ -934,6 +934,17 @@ impl IVFPQIndexReader { &mut self, list_id: usize, mut consume: impl FnMut(&ProductQuantizer, &[i64], &[u8]), + ) -> io::Result<()> { + self.try_for_each_streamed_list_chunk(list_id, |pq, ids, codes| { + consume(pq, ids, codes); + Ok(()) + }) + } + + pub(crate) fn try_for_each_streamed_list_chunk( + &mut self, + list_id: usize, + mut consume: impl FnMut(&ProductQuantizer, &[i64], &[u8]) -> io::Result<()>, ) -> io::Result<()> { self.ensure_loaded()?; let count = self.list_counts[list_id] as usize; @@ -1004,7 +1015,7 @@ impl IVFPQIndexReader { .pread(&mut [ReadRequest::new(chunk_offset, payload.codes_mut())])?; } let row_end = row_start + chunk_rows; - consume(&self.pq, &ids[row_start..row_end], payload.codes()); + consume(&self.pq, &ids[row_start..row_end], payload.codes())?; row_start = row_end; } Ok(()) diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 60644e8..ea1f692 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -16,16 +16,19 @@ // under the License. use crate::coarse::CoarseAssignment; +use crate::collect::{Collector, RangeCollector}; use crate::distance::{ - fvec_inner_product, fvec_madd, fvec_normalize, pq_distance_four_codes, pq_distance_from_table, - MetricType, + fvec_inner_product, fvec_l2sqr, fvec_madd, fvec_normalize, pq_distance_four_codes, + pq_distance_from_table, MetricType, }; +use crate::index::validate_queries; use crate::index_io_util::ivf_payload_is_oversized; use crate::io::{IVFPQIndexReader, InvertedListPayload, SeekRead}; use crate::kmeans::{self, KMeansConfig}; use crate::logging::{emit_log, LogLevel}; use crate::opq::OPQMatrix; use crate::pq::ProductQuantizer; +use crate::range::{RangeResultBuilder, RangeSearchResult, VectorRangeSearchParams}; use crate::sparse_table::SparseTable; use rayon::prelude::*; use roaring::RoaringTreemap; @@ -1665,6 +1668,423 @@ pub fn search_with_reader_roaring_filter( search_with_reader_filter(reader, query, k, nprobe, Some(&filter)) } +impl IVFPQIndexReader { + /// Returns every eligible probed row whose PQ estimate is in the L2 band. + /// + /// Both code widths use floating-point ADC: direct squared-L2 distances + /// from the query (after OPQ and optional coarse residual subtraction) to + /// each selected PQ centroid, summed in subquantizer order. Range search + /// does not use top-K's u8 FastScan tables or precomputed norm identities, + /// so membership is independent of list size and `optimize_for_search`. + /// This is an estimate of the original vector's distance, even at full + /// nprobe. Results are uncapped and unordered; top-K is unchanged. + /// + /// Non-finite transformed queries, coarse distances, or consumed PQ + /// distances return `InvalidData`. Filtered-out rows are not evaluated. + /// Every eligible row is fully evaluated, without early abandonment. + 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], + query_count: usize, + params: VectorRangeSearchParams, + ) -> io::Result { + self.range_search_batch_with_filter(queries, query_count, params, None) + } + + /// Batched range search restricted to a serialized Roaring allow-list. + pub fn range_search_batch_with_roaring_filter( + &mut self, + queries: &[f32], + query_count: usize, + params: VectorRangeSearchParams, + roaring_filter_bytes: &[u8], + ) -> io::Result { + let filter = decode_roaring_filter(roaring_filter_bytes)?; + self.range_search_batch_with_filter(queries, query_count, params, Some(&filter)) + } + + pub fn range_search_batch_with_filter( + &mut self, + queries: &[f32], + query_count: usize, + params: VectorRangeSearchParams, + filter: Option<&dyn RowIdFilter>, + ) -> io::Result { + validate_queries(queries, query_count, self.d)?; + if params.band().metric() != self.metric { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "band metric {:?} does not match index metric {:?}", + params.band().metric(), + self.metric + ), + )); + } + let nprobe = params.validate(self.nlist)?; + let mut builder = RangeResultBuilder::new(query_count); + if params.band().is_empty() { + return Ok(builder.build()); + } + self.ensure_loaded()?; + let prepare_query = |(query_index, query): (usize, &[f32])| { + let mut prepared = query.to_vec(); + if let Some(opq) = &self.opq { + opq.apply(query, &mut prepared); + } + if prepared.iter().any(|value| !value.is_finite()) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "non-finite IVF-PQ rotated query", + )); + } + let probes = kmeans::find_topk_checked( + &prepared, + &self.quantizer_centroids, + self.nlist, + self.d, + nprobe, + ) + .map_err(|list| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("non-finite IVF-PQ query-centroid distance for list {list}"), + ) + })?; + Ok(( + PqRangeQuery::new(query_index, prepared, RangeCollector::new(params.band())), + probes, + )) + }; + let coarse_work = query_count + .saturating_mul(self.nlist) + .saturating_mul(self.d); + let prepared_queries = + if query_count > 1 && coarse_work >= PARALLEL_PQ_RANGE_MIN_COARSE_COMPONENTS { + queries + .par_chunks_exact(self.d) + .enumerate() + .map(prepare_query) + .collect::>>()? + } else { + queries + .chunks_exact(self.d) + .enumerate() + .map(prepare_query) + .collect::>>()? + }; + let mut list_to_queries = vec![Vec::new(); self.nlist]; + let mut query_states = prepared_queries + .into_iter() + .map(|(query, probes)| { + builder.record_lists_probed(query.query_index, probes.len()); + for (_, list) in probes { + list_to_queries[list].push(query.query_index); + } + Some(query) + }) + .collect::>(); + let mut list_ids = (0..self.nlist) + .filter(|&list| !list_to_queries[list].is_empty() && self.list_counts[list] > 0) + .collect::>(); + list_ids.sort_unstable_by_key(|&list| self.list_offsets[list]); + let cached_queries = pq_range_cache_query_limit(&self.pq, PQ_RANGE_TABLE_CACHE_BYTES); + let worker_count = if list_ids.iter().any(|&list| { + let queries = list_to_queries[list].len(); + queries > 1 + && (self.list_counts[list] as usize).saturating_mul(queries) + >= PARALLEL_PQ_RANGE_MIN_CANDIDATES + }) { + query_count.min(rayon::current_num_threads()) + } else { + 1 + }; + let mut scratch = (0..worker_count) + .map(|_| PqRangeScratch::default()) + .collect::>(); + let mut active_queries = Vec::new(); + let mut batch_start = 0; + while batch_start < list_ids.len() { + let list_id = list_ids[batch_start]; + if ivf_payload_is_oversized(self.list_payload_len(list_id)?) { + let centroid = self.by_residual.then(|| { + self.quantizer_centroids[list_id * self.d..(list_id + 1) * self.d].to_vec() + }); + let transposed = self.transposed_codes; + active_queries.extend( + list_to_queries[list_id] + .iter() + .map(|&query_index| query_states[query_index].take().unwrap()), + ); + self.try_for_each_streamed_list_chunk(list_id, |pq, ids, codes| { + scan_pq_range_list( + pq, + list_id, + centroid.as_deref(), + ids, + codes, + transposed, + filter, + cached_queries, + &mut scratch, + &mut active_queries, + ) + })?; + for query in active_queries.drain(..) { + let query_index = query.query_index; + query_states[query_index] = Some(query); + } + builder.record_list_read(); + batch_start += 1; + } else { + let batch_end = batch_start + self.batch_read_end(&list_ids[batch_start..])?; + let lists = self.read_inverted_list_payloads(&list_ids[batch_start..batch_end])?; + for list in lists { + builder.record_list_read(); + let centroid = self.by_residual.then(|| { + &self.quantizer_centroids + [list.list_id * self.d..(list.list_id + 1) * self.d] + }); + active_queries.extend( + list_to_queries[list.list_id] + .iter() + .map(|&query_index| query_states[query_index].take().unwrap()), + ); + scan_pq_range_list( + &self.pq, + list.list_id, + centroid, + &list.ids, + list.codes(), + self.transposed_codes, + filter, + if self.by_residual { 0 } else { cached_queries }, + &mut scratch, + &mut active_queries, + )?; + for query in active_queries.drain(..) { + let query_index = query.query_index; + query_states[query_index] = Some(query); + } + } + batch_start = batch_end; + } + } + for query in query_states.into_iter().map(Option::unwrap) { + builder.record_scanned(query.query_index, query.collector.scanned()); + builder.take_rows(query.query_index, query.collector.into_rows()); + } + Ok(builder.build()) + } +} + +const PQ_RANGE_TABLE_CACHE_BYTES: usize = 8 * 1024 * 1024; +const PARALLEL_PQ_RANGE_MIN_CANDIDATES: usize = 8 * 1024; +const PARALLEL_PQ_RANGE_MIN_COARSE_COMPONENTS: usize = 128 * 1024; + +#[derive(Default)] +struct PqRangeScratch { + residual_query: Vec, + table: Vec, +} + +struct PqRangeQuery { + query_index: usize, + query: Vec, + collector: C, + table: Vec, + table_list: Option, +} + +impl PqRangeQuery { + fn new(query_index: usize, query: Vec, collector: C) -> Self { + Self { + query_index, + query, + collector, + table: Vec::new(), + table_list: None, + } + } +} + +fn pq_range_cache_query_limit(pq: &ProductQuantizer, budget: usize) -> usize { + pq.m() + .checked_mul(pq.ksub()) + .and_then(|elements| elements.checked_mul(std::mem::size_of::())) + .map_or(0, |bytes| budget / bytes) +} + +fn scan_pq_range_list( + pq: &ProductQuantizer, + list_id: usize, + centroid: Option<&[f32]>, + ids: &[i64], + codes: &[u8], + transposed: bool, + filter: Option<&dyn RowIdFilter>, + cached_queries: usize, + scratch: &mut [PqRangeScratch], + queries: &mut [PqRangeQuery], +) -> io::Result<()> { + let positions = matching_rows(ids, filter); + if ids.is_empty() || positions.as_ref().is_some_and(MatchingRows::is_empty) { + return Ok(()); + } + let scan_query = |scratch: &mut PqRangeScratch, query: &mut PqRangeQuery| { + let table = if query.query_index < cached_queries { + if query.table_list.is_none() + || (centroid.is_some() && query.table_list != Some(list_id)) + { + build_pq_range_table( + pq, + &query.query, + centroid, + &mut scratch.residual_query, + &mut query.table, + ); + query.table_list = Some(list_id); + } + &query.table + } else { + build_pq_range_table( + pq, + &query.query, + centroid, + &mut scratch.residual_query, + &mut scratch.table, + ); + &scratch.table + }; + scan_pq_range_codes( + pq, + table, + ids, + codes, + transposed, + positions.as_ref(), + &mut query.collector, + ) + }; + let matching_count = positions.as_ref().map_or(ids.len(), MatchingRows::len); + if scratch.len() > 1 + && queries.len() > 1 + && matching_count.saturating_mul(queries.len()) >= PARALLEL_PQ_RANGE_MIN_CANDIDATES + { + let queries_per_worker = queries.len().div_ceil(scratch.len()); + queries + .par_chunks_mut(queries_per_worker) + .zip(scratch.par_iter_mut()) + .try_for_each(|(queries, scratch)| { + for query in queries { + scan_query(scratch, query)?; + } + Ok(()) + }) + } else { + for query in queries { + scan_query(&mut scratch[0], query)?; + } + Ok(()) + } +} + +fn build_pq_range_table( + pq: &ProductQuantizer, + query: &[f32], + centroid: Option<&[f32]>, + residual_query: &mut Vec, + table: &mut Vec, +) { + let query = if let Some(centroid) = centroid { + residual_query.resize(query.len(), 0.0); + for ((residual, value), coarse) in residual_query.iter_mut().zip(query).zip(centroid) { + *residual = value - coarse; + } + residual_query.as_slice() + } else { + query + }; + table.resize(pq.m() * pq.ksub(), 0.0); + for sub in 0..pq.m() { + let query_chunk = &query[sub * pq.dsub()..(sub + 1) * pq.dsub()]; + for code in 0..pq.ksub() { + let offset = (sub * pq.ksub() + code) * pq.dsub(); + table[sub * pq.ksub() + code] = + fvec_l2sqr(query_chunk, &pq.centroids()[offset..offset + pq.dsub()]); + } + } +} + +fn scan_pq_range_codes( + pq: &ProductQuantizer, + table: &[f32], + ids: &[i64], + codes: &[u8], + transposed: bool, + positions: Option<&MatchingRows>, + collector: &mut C, +) -> io::Result<()> { + let mut scan_row = |row: usize| { + let mut distance = 0.0; + for sub in 0..pq.m() { + let column = if pq.nbits() == 4 { sub / 2 } else { sub }; + let offset = if transposed { + column * ids.len() + row + } else { + row * pq.code_size() + column + }; + let code = if pq.nbits() == 4 { + (codes[offset] >> (4 * (sub % 2))) & 15 + } else { + codes[offset] + } as usize; + distance += table[sub * pq.ksub() + code]; + } + collector.push(ids[row], distance) + }; + if let Some(positions) = positions { + for row in positions.positions() { + scan_row(row)?; + } + } else { + for row in 0..ids.len() { + scan_row(row)?; + } + } + Ok(()) +} + fn scan_reader_list( entry: &InvertedListPayload, dis0: f32, @@ -2964,6 +3384,260 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; + #[test] + fn pq_range_scans_active_queries_on_multiple_workers() { + struct TrackingCollector<'a> { + inner: RangeCollector, + workers: &'a AtomicUsize, + } + + impl Collector for TrackingCollector<'_> { + fn cutoff(&self) -> f32 { + self.inner.cutoff() + } + + fn push(&mut self, id: i64, distance: f32) -> io::Result<()> { + if self.inner.scanned() == 0 { + let worker = rayon::current_thread_index().unwrap(); + self.workers.fetch_or(1 << worker, Ordering::Relaxed); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + self.inner.push(id, distance) + } + } + + let mut pq = ProductQuantizer::with_nbits(128, 16, 8); + pq.set_centroids(vec![0.0; 128 * pq.ksub()]); + let ids = (0..8193).collect::>(); + let codes = vec![0; ids.len() * pq.code_size()]; + let queries = (0..16) + .map(|query_index| vec![query_index as f32 * 0.25; 128]) + .collect::>(); + let band = crate::range::DistanceBand::new( + crate::range::Bound::Unbounded, + crate::range::Bound::Unbounded, + MetricType::L2, + ) + .unwrap(); + let workers = AtomicUsize::new(0); + let query_indices = [14, 2, 12, 4, 10, 6, 8, 0]; + let mut collectors = query_indices + .iter() + .map(|&query_index| { + PqRangeQuery::new( + query_index, + queries[query_index].clone(), + TrackingCollector { + inner: RangeCollector::new(band), + workers: &workers, + }, + ) + }) + .collect::>(); + let mut scratch = (0..4) + .map(|_| PqRangeScratch::default()) + .collect::>(); + rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() + .install(|| { + scan_pq_range_list( + &pq, + 0, + None, + &ids, + &codes, + true, + None, + 0, + &mut scratch, + &mut collectors, + ) + .unwrap(); + }); + assert!( + workers.load(Ordering::Relaxed).count_ones() > 1, + "large PQ range batches must scan queries on multiple workers" + ); + for query in collectors { + assert_eq!(query.collector.inner.scanned(), ids.len()); + let distance = 8.0 * (query.query_index * query.query_index) as f32; + for (row, (id, value)) in query.collector.inner.into_rows().into_iter().enumerate() { + assert_eq!(value, distance); + assert_eq!(id, row as i64); + } + } + } + + fn pq_range_test_queries(count: usize) -> Vec> { + let band = crate::range::DistanceBand::new( + crate::range::Bound::Unbounded, + crate::range::Bound::Unbounded, + MetricType::L2, + ) + .unwrap(); + (0..count) + .rev() + .map(|query_index| { + PqRangeQuery::new( + query_index, + vec![query_index as f32 * 0.25; 32], + RangeCollector::new(band), + ) + }) + .collect() + } + + #[test] + fn pq_range_reuses_bounded_tables_without_changing_estimates() { + for bits in [4, 8] { + let mut pq = ProductQuantizer::with_nbits(32, 4, bits); + pq.set_centroids(vec![0.125; 32 * pq.ksub()]); + let table_bytes = pq.m() * pq.ksub() * std::mem::size_of::(); + let budget = 2 * table_bytes; + assert_eq!(pq_range_cache_query_limit(&pq, table_bytes - 1), 0); + assert_eq!(pq_range_cache_query_limit(&pq, budget), 2); + for residual in [false, true] { + let mut cached = pq_range_test_queries(5); + let mut uncached = pq_range_test_queries(5); + let mut scratch = [PqRangeScratch::default()]; + let mut reference_scratch = [PqRangeScratch::default()]; + let mut table_addresses = Vec::new(); + let mut scratch_address = None; + for (chunk, list_id) in [0, 0, 1].into_iter().enumerate() { + let ids = (chunk * 16..(chunk + 1) * 16) + .map(|row| row as i64) + .collect::>(); + let codes = vec![0; ids.len() * pq.code_size()]; + let centroid = vec![list_id as f32 * 0.5; 32]; + for (queries, scratch, cache_limit) in [ + (&mut cached, &mut scratch, 2), + (&mut uncached, &mut reference_scratch, 0), + ] { + scan_pq_range_list( + &pq, + list_id, + residual.then_some(centroid.as_slice()), + &ids, + &codes, + true, + None, + cache_limit, + scratch, + queries, + ) + .unwrap(); + } + let addresses = cached + .iter() + .filter(|query| query.query_index < 2) + .map(|query| { + assert_eq!(query.table_list, Some(if residual { list_id } else { 0 })); + query.table.as_ptr() + }) + .collect::>(); + if chunk == 0 { + table_addresses = addresses; + scratch_address = Some(scratch[0].table.as_ptr()); + } else { + assert_eq!(addresses, table_addresses); + assert_eq!(Some(scratch[0].table.as_ptr()), scratch_address); + } + assert!(cached + .iter() + .filter(|query| query.query_index >= 2) + .all(|query| query.table.is_empty())); + assert!( + cached + .iter() + .map(|query| query.table.capacity() * std::mem::size_of::()) + .sum::() + <= budget + ); + } + for (cached, uncached) in cached.into_iter().zip(uncached) { + assert_eq!(cached.collector.scanned(), uncached.collector.scanned()); + assert_eq!(cached.collector.into_rows(), uncached.collector.into_rows()); + } + } + } + } + + #[test] + fn pq_range_excluded_rows_do_not_allocate_tables() { + let pq = ProductQuantizer::with_nbits(32, 4, 8); + let ids = (0..8193).collect::>(); + let codes = vec![0; ids.len() * pq.code_size()]; + let mut queries = pq_range_test_queries(4); + let mut scratch = [PqRangeScratch::default()]; + scan_pq_range_list( + &pq, + 0, + None, + &ids, + &codes, + true, + Some(&RoaringTreemap::new()), + 4, + &mut scratch, + &mut queries, + ) + .unwrap(); + assert!(queries.iter().all(|query| query.table.is_empty() + && query.table_list.is_none() + && query.collector.scanned() == 0)); + assert!(scratch[0].table.is_empty()); + assert!(scratch[0].residual_query.is_empty()); + } + + #[test] + fn pq_range_parallel_collector_errors_propagate() { + struct FailingCollector; + + impl Collector for FailingCollector { + fn cutoff(&self) -> f32 { + f32::INFINITY + } + + fn push(&mut self, _id: i64, _distance: f32) -> io::Result<()> { + Err(io::Error::other("parallel PQ collector failed")) + } + } + + let mut pq = ProductQuantizer::with_nbits(32, 4, 8); + pq.set_centroids(vec![0.0; 32 * pq.ksub()]); + let ids = (0..8193).collect::>(); + let codes = vec![0; ids.len() * pq.code_size()]; + let mut queries = (0..4) + .map(|query_index| PqRangeQuery::new(query_index, vec![0.0; 32], FailingCollector)) + .collect::>(); + let mut scratch = (0..4) + .map(|_| PqRangeScratch::default()) + .collect::>(); + let error = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() + .install(|| { + scan_pq_range_list( + &pq, + 0, + None, + &ids, + &codes, + true, + None, + 4, + &mut scratch, + &mut queries, + ) + .unwrap_err() + }); + assert_eq!(error.kind(), io::ErrorKind::Other); + assert_eq!(error.to_string(), "parallel PQ collector failed"); + } + struct CountingFilter { contains_calls: AtomicUsize, } diff --git a/core/src/range.rs b/core/src/range.rs index 4ed9639..ef9796a 100644 --- a/core/src/range.rs +++ b/core/src/range.rs @@ -362,8 +362,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. IVF-RQ evaluates the complete estimate of every eligible row - /// without early abandonment, so its count is always zero. + /// rows too. IVF-RQ and IVF-PQ evaluate the complete estimate of every + /// eligible row without early abandonment, so their 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 3afcfa0..461e726 100644 --- a/core/tests/range_search.rs +++ b/core/tests/range_search.rs @@ -27,25 +27,1127 @@ use paimon_vindex_core::distance::{fvec_l2sqr, MetricType}; use paimon_vindex_core::index::{VectorIndexReader, VectorSearchParams}; -use paimon_vindex_core::io::{PosWriter, ReadRequest, SeekRead}; +use paimon_vindex_core::io::{ + write_index, IVFPQIndexReader, PosWriter, ReadRequest, SeekRead, SeekReadCapabilities, +}; use paimon_vindex_core::ivfflat::IVFFlatIndex; use paimon_vindex_core::ivfflat_io::write_ivfflat_index; +use paimon_vindex_core::ivfpq::IVFPQIndex; use paimon_vindex_core::ivfrq::IVFRQIndex; use paimon_vindex_core::ivfrq_io::write_ivfrq_index; use paimon_vindex_core::ivfsq::IVFSQIndex; use paimon_vindex_core::ivfsq_io::{write_ivfsq_index, IVFSQIndexReader}; -use paimon_vindex_core::range::{Bound, DistanceBand, QueryResult, VectorRangeSearchParams}; +use paimon_vindex_core::range::{ + Bound, DistanceBand, QueryResult, RangeSearchResult, VectorRangeSearchParams, +}; use paimon_vindex_core::read_options::VectorIndexReaderOptions; use paimon_vindex_core::rq::RQRotation; use paimon_vindex_core::sq::ScalarQuantizer; use std::collections::HashSet; -use std::io::Cursor; +use std::io::{self, Cursor}; use std::sync::{Arc, Mutex}; use roaring::RoaringTreemap; type Reader = VectorIndexReader>>; +fn pq_fixture(bits: usize, residual: bool, opq: bool) -> IVFPQIndex { + let mut index = IVFPQIndex::with_nbits(8, 4, 4, bits, MetricType::L2, opq); + index.by_residual = residual; + index.set_quantizer_centroids( + (0..index.nlist) + .flat_map(|list| (0..index.d).map(move |dim| (list * 4 + dim) as f32 / 4.0)) + .collect(), + ); + index.pq.set_centroids( + (0..index.pq.m()) + .flat_map(|sub| { + (0..index.pq.ksub()).flat_map(move |code| { + (0..2).map(move |dim| ((code * 7 + sub * 3 + dim) % 31) as f32 / 16.0) + }) + }) + .collect(), + ); + if let Some(rotation) = &mut index.opq { + rotation.rotation = vec![0.0; index.d * index.d]; + for dim in 0..index.d { + rotation.rotation[dim * index.d + (dim + 1) % index.d] = 1.0; + } + rotation.is_trained = true; + } + for list in 0..3 { + for row in 0..273 { + index.ids[list].push(1000 + (list * 273 + row) as i64); + let codes = (0..index.pq.m()) + .map(|sub| ((row * (sub + 1) + list + sub * 5) % index.pq.ksub()) as u8) + .collect::>(); + if bits == 4 { + index.codes[list].extend(codes.chunks_exact(2).map(|pair| pair[0] | pair[1] << 4)); + } else { + index.codes[list].extend(codes); + } + } + } + index +} + +fn pq_dense_opq_fixture(bits: usize, residual: bool) -> IVFPQIndex { + let dimension = 64; + let rows_per_list = 1027; + let mut index = IVFPQIndex::with_nbits(dimension, 3, 8, bits, MetricType::L2, true); + index.by_residual = residual; + index.set_quantizer_centroids( + (0..index.nlist) + .flat_map(|list| { + (0..dimension).map(move |coordinate| { + (list as f32 - 1.0) * 1.5 + ((coordinate * 7) % 17) as f32 / 29.0 + }) + }) + .collect(), + ); + index.pq.set_centroids( + (0..index.pq.m()) + .flat_map(|sub| { + (0..index.pq.ksub()).flat_map(move |code| { + (0..8).map(move |coordinate| { + ((code * 13 + sub * 7 + coordinate * 11) % 97) as f32 / 23.0 - 2.0 + }) + }) + }) + .collect(), + ); + let rotation = index.opq.as_mut().unwrap(); + rotation.rotation = (0..dimension) + .flat_map(|row| { + (0..dimension).map(move |column| { + if (row & column).count_ones() % 2 == 0 { + 0.125 + } else { + -0.125 + } + }) + }) + .collect(); + rotation.is_trained = true; + for left in rotation.rotation.chunks_exact(dimension) { + assert!(left.iter().all(|value| value.abs() == 0.125)); + } + for (row, left) in rotation.rotation.chunks_exact(dimension).enumerate() { + for (column, right) in rotation.rotation.chunks_exact(dimension).enumerate() { + let dot = left + .iter() + .zip(right) + .map(|(left, right)| left * right) + .sum::(); + assert_eq!(dot, if row == column { 1.0 } else { 0.0 }); + } + } + for list in 0..index.nlist { + for row in 0..rows_per_list { + index.ids[list].push(10_000 + (list * rows_per_list + row) as i64); + let codes = (0..index.pq.m()) + .map(|sub| { + ((row * (2 * sub + 1) + row / 16 + list * 17 + sub * 11) % index.pq.ksub()) + as u8 + }) + .collect::>(); + if bits == 4 { + index.codes[list].extend(codes.chunks_exact(2).map(|pair| pair[0] | pair[1] << 4)); + } else { + index.codes[list].extend(codes); + } + } + } + index +} + +fn pq_coarse_batch_fixture() -> IVFPQIndex { + let fixture = pq_dense_opq_fixture(8, true); + let mut index = IVFPQIndex::with_nbits(64, 128, 8, 8, MetricType::L2, true); + let mut centroids = fixture.quantizer_centroids().to_vec(); + centroids.resize(index.nlist * index.d, 1000.0); + index.set_quantizer_centroids(centroids); + index.pq = fixture.pq; + index.opq = fixture.opq; + for list in 0..fixture.nlist { + index.ids[list] = fixture.ids[list][..7].to_vec(); + index.codes[list] = fixture.codes[list][..7 * index.pq.code_size()].to_vec(); + } + index +} + +fn pq_bytes(index: &IVFPQIndex) -> Vec { + let mut bytes = Vec::new(); + write_index(index, &mut PosWriter::new(&mut bytes)).unwrap(); + bytes +} + +fn pq_reader(index: &IVFPQIndex) -> Reader { + VectorIndexReader::open(Cursor::new(pq_bytes(index))).unwrap() +} + +fn pq_oracle(index: &IVFPQIndex, query: &[f32], nprobe: usize) -> Vec<(i64, f32)> { + let mut rotated = query.to_vec(); + if let Some(rotation) = &index.opq { + for (dim, value) in rotated.iter_mut().enumerate() { + *value = query + .iter() + .enumerate() + .map(|(column, value)| value * rotation.rotation[dim * index.d + column]) + .sum(); + } + } + let mut lists = (0..index.nlist) + .map(|list| { + let distance = rotated + .iter() + .zip(&index.quantizer_centroids()[list * index.d..(list + 1) * index.d]) + .map(|(value, centroid)| (value - centroid).powi(2)) + .sum::(); + (distance, list) + }) + .collect::>(); + lists.sort_by(|left, right| left.partial_cmp(right).unwrap()); + let mut rows = Vec::new(); + for &(_, list) in lists.iter().take(nprobe) { + for (row, &id) in index.ids[list].iter().enumerate() { + let code = + &index.codes[list][row * index.pq.code_size()..(row + 1) * index.pq.code_size()]; + let mut distance = 0.0; + for sub in 0..index.pq.m() { + let label = if index.pq.nbits() == 4 { + (code[sub / 2] >> (4 * (sub % 2))) & 15 + } else { + code[sub] + } as usize; + let mut term = 0.0; + for dim in 0..index.pq.dsub() { + let coordinate = sub * index.pq.dsub() + dim; + let coarse = if index.by_residual { + index.quantizer_centroids()[list * index.d + coordinate] + } else { + 0.0 + }; + let centroid = index.pq.centroids() + [(sub * index.pq.ksub() + label) * index.pq.dsub() + dim]; + term += (rotated[coordinate] - coarse - centroid).powi(2); + } + distance += term; + } + rows.push((id, distance)); + } + } + rows.sort_by_key(|&(id, _)| id); + rows +} + +fn pq_rows(query: QueryResult<'_>) -> Vec<(i64, f32)> { + let mut rows = query + .labels + .iter() + .copied() + .zip(query.distances.iter().copied()) + .collect::>(); + rows.sort_by_key(|&(id, _)| id); + rows +} + +fn pq_all_params(nprobe: usize) -> VectorRangeSearchParams { + VectorRangeSearchParams::new( + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(), + nprobe, + ) +} + +fn pq_entry_points( + reader: &mut Reader, + query: &[f32], + params: VectorRangeSearchParams, + filter: &[u8], +) -> [io::Result; 4] { + [ + reader.range_search(query, params), + reader.range_search_with_roaring_filter(query, params, filter), + reader.range_search_batch(query, 1, params), + reader.range_search_batch_with_roaring_filter(query, 1, params, filter), + ] +} + +#[test] +fn pq_range_matches_decoded_oracle_for_both_code_widths() { + for bits in [4, 8] { + for residual in [false, true] { + for opq in [false, true] { + let index = pq_fixture(bits, residual, opq); + let query = [0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0]; + for nprobe in [1, 2, 4, 20] { + let expected = pq_oracle(&index, &query, nprobe); + let band = l2(0.5, 30.0); + let result = pq_reader(&index) + .range_search(&query, VectorRangeSearchParams::new(band, nprobe)) + .unwrap(); + assert_eq!( + pq_rows(result.query(0)), + expected + .into_iter() + .filter(|&(_, value)| band.admit(value)) + .collect::>() + ); + } + } + } + } +} + +#[test] +fn pq_range_dense_opq_parallel_batches_preserve_float_adc_results() { + let pools = [1, 4].map(|threads| { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + }); + let query_count = 16; + for bits in [4, 8] { + for residual in [false, true] { + let index = pq_dense_opq_fixture(bits, residual); + assert!(index.pq.dsub() >= 8); + let dimension = index.d; + let queries = (0..query_count) + .flat_map(|query| { + (0..dimension).map(move |coordinate| { + if coordinate == 0 { + (query % 3) as f32 * 12.0 - 12.0 + } else { + ((query * 19 + coordinate * 11 + query * coordinate) % 101) as f32 + / 31.0 + - 1.5 + } + }) + }) + .collect::>(); + let allowed = index + .ids + .iter() + .flatten() + .copied() + .filter(|id| id % 2 == 0) + .collect::>(); + assert!(index.ids.iter().all(|ids| { + ids.iter().filter(|id| allowed.contains(id)).count() * query_count > 8192 + })); + let filter = serialize_roaring(&allowed); + let params = pq_all_params(index.nlist); + let bytes = pq_bytes(&index); + let mut reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let reference = pools[0].install(|| { + queries + .chunks_exact(dimension) + .map(|query| pq_rows(reader.range_search(query, params).unwrap().query(0))) + .collect::>() + }); + for (query_index, query) in queries.chunks_exact(dimension).enumerate() { + let expected = pq_oracle(&index, query, index.nlist); + assert_eq!(reference[query_index].len(), expected.len()); + for (&(id, distance), &(expected_id, expected_distance)) in + reference[query_index].iter().zip(&expected) + { + assert_eq!(id, expected_id); + assert!( + (distance - expected_distance).abs() + <= 1e-5 * expected_distance.abs().max(1.0), + "PQ{bits}, residual={residual}, query={query_index}, id={id}: \ + SIMD distance {distance}, scalar oracle {expected_distance}" + ); + } + } + let mut distances = reference[0] + .iter() + .map(|&(_, value)| value) + .collect::>(); + distances.sort_by(f32::total_cmp); + let cut = distances[distances.len() / 2]; + assert!(distances[0] < cut && cut < *distances.last().unwrap()); + let bands = [ + params.band(), + DistanceBand::new(Bound::Unbounded, Bound::Finite(cut), MetricType::L2).unwrap(), + DistanceBand::new(Bound::Finite(cut), Bound::Unbounded, MetricType::L2).unwrap(), + l2(cut, cut.next_up()), + l2(cut.next_down(), cut), + ]; + for pool in &pools { + pool.install(|| { + let mut reader = IVFPQIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + for optimized in [false, true] { + if optimized { + reader.optimize_for_search().unwrap(); + } + for band in bands { + let band_params = VectorRangeSearchParams::new(band, index.nlist); + let batch = reader + .range_search_batch(&queries, query_count, band_params) + .unwrap(); + let filtered = reader + .range_search_batch_with_roaring_filter( + &queries, + query_count, + band_params, + &filter, + ) + .unwrap(); + assert_eq!(batch.query_count(), query_count); + assert_eq!(filtered.query_count(), query_count); + assert_eq!(batch.call_stats().list_reads(), index.nlist); + assert_eq!(filtered.call_stats().list_reads(), index.nlist); + for (query_index, query) in queries.chunks_exact(dimension).enumerate() + { + let expected = reference[query_index] + .iter() + .copied() + .filter(|&(_, value)| band.admit(value)) + .collect::>(); + let expected_filtered = expected + .iter() + .copied() + .filter(|(id, _)| allowed.contains(id)) + .collect::>(); + assert_eq!(pq_rows(batch.query(query_index)), expected); + assert_eq!(pq_rows(filtered.query(query_index)), expected_filtered); + for (result, scanned, committed) in [ + (&batch, reference[query_index].len(), expected.len()), + (&filtered, allowed.len(), expected_filtered.len()), + ] { + let stats = result.query(query_index).stats; + assert_eq!(stats.rows_scanned(), scanned); + assert_eq!(stats.rows_committed(), committed); + assert_eq!(stats.lists_probed(), index.nlist); + assert_eq!(stats.early_abandoned(), 0); + } + if band == params.band() || query_index == 0 { + let single = reader.range_search(query, band_params).unwrap(); + let single_filtered = reader + .range_search_with_filter( + query, + band_params, + Some(&allowed), + ) + .unwrap(); + assert_eq!(pq_rows(single.query(0)), expected); + assert_eq!( + pq_rows(single_filtered.query(0)), + expected_filtered + ); + } + } + } + } + }); + } + } + } +} + +#[test] +fn pq_range_coarse_parallel_partial_probes_preserve_query_order() { + let index = pq_coarse_batch_fixture(); + let query_count = 16; + let queries = (0..query_count) + .flat_map(|query_index| { + (0..index.d).map(move |coordinate| { + if coordinate == 0 { + (query_index % 3) as f32 * 12.0 - 12.0 + } else { + (query_index * 17 + coordinate) as f32 / 64.0 - 2.0 + } + }) + }) + .collect::>(); + let reversed = queries + .chunks_exact(index.d) + .rev() + .flatten() + .copied() + .collect::>(); + let params = pq_all_params(2); + let snapshots = [1, 4].map(|threads| { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + .install(|| { + let mut reader = pq_reader(&index); + let batch = reader + .range_search_batch(&queries, query_count, params) + .unwrap(); + let reordered = reader + .range_search_batch(&reversed, query_count, params) + .unwrap(); + assert_eq!(batch.query_count(), query_count); + assert_eq!(reordered.query_count(), query_count); + let rows = (0..query_count) + .map(|query_index| pq_rows(batch.query(query_index))) + .collect::>(); + for (query_index, query) in queries.chunks_exact(index.d).enumerate() { + let single = reader.range_search(query, params).unwrap(); + assert_eq!(rows[query_index], pq_rows(single.query(0))); + assert_eq!( + rows[query_index], + pq_rows(reordered.query(query_count - 1 - query_index)) + ); + assert_eq!(batch.query(query_index).stats.lists_probed(), 2); + let expected = pq_oracle(&index, query, 2); + assert_eq!(rows[query_index].len(), 14); + assert_eq!(rows[query_index].len(), expected.len()); + for (&(id, distance), &(expected_id, expected_distance)) in + rows[query_index].iter().zip(&expected) + { + assert_eq!(id, expected_id); + assert!( + (distance - expected_distance).abs() + <= 1e-5 * expected_distance.abs().max(1.0) + ); + } + } + assert!(rows.windows(2).all(|pair| pair[0] != pair[1])); + rows + }) + }); + assert_eq!(snapshots[0], snapshots[1]); +} + +#[test] +fn pq_range_coarse_parallel_overflow_precedes_payload_io() { + let index = pq_coarse_batch_fixture(); + let query_count = 16; + let params = pq_all_params(1); + let empty_filter = serialize_roaring(&HashSet::new()); + for threads in [1, 4] { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + .install(|| { + let valid_prefix = vec![0.0; (query_count - 1) * index.d]; + pq_reader(&index) + .range_search_batch(&valid_prefix, query_count - 1, params) + .unwrap(); + for (late_value, far_value, diagnostic) in [ + (1e20, 1000.0, "query-centroid distance"), + (0.0, 1e20, "query-centroid distance for list 127"), + (f32::MAX, 1000.0, "rotated query"), + ] { + let mut queries = vec![0.0; query_count * index.d]; + queries[(query_count - 1) * index.d..].fill(late_value); + let trace = Arc::new(Mutex::new(SqReadTrace::default())); + let source = SqRecordingReader { + inner: Cursor::new(pq_bytes(&index)), + trace: Arc::clone(&trace), + }; + let mut reader = IVFPQIndexReader::open(source).unwrap(); + reader.ensure_loaded().unwrap(); + *reader.quantizer_centroids.last_mut().unwrap() = far_value; + assert!(queries.iter().all(|value| value.is_finite())); + assert!(reader + .quantizer_centroids + .iter() + .all(|value| value.is_finite())); + *trace.lock().unwrap() = SqReadTrace::default(); + for result in [ + reader.range_search_batch(&queries, query_count, params), + reader.range_search_batch_with_roaring_filter( + &queries, + query_count, + params, + &empty_filter, + ), + ] { + let error = result.unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + assert!(error.to_string().contains(diagnostic), "{error}"); + } + assert_eq!(trace.lock().unwrap().calls, 0); + } + }); + } +} + +#[test] +fn pq_range_boundaries_partition_the_full_estimated_result() { + for bits in [4, 8] { + let index = pq_fixture(bits, true, false); + let query = [0.25; 8]; + let expected = pq_oracle(&index, &query, 4); + let cut = expected[200].1; + let mut reader = pq_reader(&index); + let bands = [ + pq_all_params(4).band(), + DistanceBand::new(Bound::Unbounded, Bound::Finite(cut), MetricType::L2).unwrap(), + DistanceBand::new(Bound::Finite(cut), Bound::Unbounded, MetricType::L2).unwrap(), + l2(cut, cut.next_up()), + l2(cut.next_down(), cut), + l2(cut, cut), + ]; + for band in bands { + let result = reader + .range_search(&query, VectorRangeSearchParams::new(band, 4)) + .unwrap(); + assert_eq!( + pq_rows(result.query(0)), + expected + .iter() + .copied() + .filter(|&(_, value)| band.admit(value)) + .collect::>() + ); + assert_eq!(result.query(0).stats.early_abandoned(), 0); + } + assert!(expected.iter().any(|&(_, value)| value == cut)); + assert!(expected.len() > 200); + } +} + +#[test] +fn pq_range_batch_filters_stats_and_permutations_match_single_queries() { + for bits in [4, 8] { + for opq in [false, true] { + let index = pq_fixture(bits, true, opq); + let queries = [vec![0.25; 8], vec![1.5; 8], vec![4.0; 8], vec![0.25; 8]].concat(); + let mut reader = pq_reader(&index); + for step in [1, 2, 127, 10000] { + let allowed = (1000..1819).step_by(step).collect::>(); + let filter = serialize_roaring(&allowed); + for nprobe in [1, 2, 4, 20] { + let params = VectorRangeSearchParams::new(l2(0.5, 20.0), nprobe); + let batch = reader.range_search_batch(&queries, 4, params).unwrap(); + let filtered = reader + .range_search_batch_with_roaring_filter(&queries, 4, params, &filter) + .unwrap(); + for (query_index, query) in queries.chunks_exact(8).enumerate() { + let single = reader.range_search(query, params).unwrap(); + assert_eq!(pq_rows(batch.query(query_index)), pq_rows(single.query(0))); + let single_filtered = reader + .range_search_with_roaring_filter(query, params, &filter) + .unwrap(); + let expected = pq_rows(batch.query(query_index)) + .into_iter() + .filter(|&(id, _)| allowed.contains(&id)) + .collect::>(); + assert_eq!(pq_rows(filtered.query(query_index)), expected); + assert_eq!(pq_rows(single_filtered.query(0)), expected); + let scanned = pq_oracle(&index, query, nprobe) + .iter() + .filter(|&&(id, _)| allowed.contains(&id)) + .count(); + let stats = filtered.query(query_index).stats; + assert_eq!(stats.rows_scanned(), scanned); + assert_eq!(stats.rows_committed(), expected.len()); + assert_eq!(stats.lists_probed(), nprobe.min(index.nlist)); + assert_eq!(stats.early_abandoned(), 0); + } + assert!(filtered.call_stats().list_reads() <= 3); + assert_eq!( + filtered.call_stats().list_reads(), + batch.call_stats().list_reads() + ); + if nprobe >= 4 { + assert_eq!(batch.call_stats().list_reads(), 3); + } + let order = [2, 0, 3, 1]; + let permuted_queries = order + .iter() + .flat_map(|&position| { + queries[position * 8..(position + 1) * 8].iter().copied() + }) + .collect::>(); + let permuted = reader + .range_search_batch(&permuted_queries, 4, params) + .unwrap(); + for (position, &original) in order.iter().enumerate() { + assert_eq!( + pq_rows(permuted.query(position)), + pq_rows(batch.query(original)) + ); + } + } + } + for allowed in [HashSet::new(), HashSet::from([999999])] { + let filter = serialize_roaring(&allowed); + for result in pq_entry_points(&mut reader, &[0.25; 8], pq_all_params(4), &filter) + .into_iter() + .skip(1) + .step_by(2) + { + let result = result.unwrap(); + assert!(result.labels().is_empty()); + assert_eq!(result.query(0).stats.rows_scanned(), 0); + } + } + } + } +} + +#[test] +fn pq_range_native_entry_points_and_optimization_leave_top_k_unchanged() { + for bits in [4, 8] { + for residual in [false, true] { + for opq in [false, true] { + let index = pq_fixture(bits, residual, opq); + let mut reader = IVFPQIndexReader::open(Cursor::new(pq_bytes(&index))).unwrap(); + let query = [0.25; 8]; + let params = pq_all_params(4); + let expected = pq_oracle(&index, &query, 4); + let allowed = index.ids.iter().flatten().copied().collect::>(); + let filter = serialize_roaring(&allowed); + for optimized in [false, true] { + if optimized { + reader.optimize_for_search().unwrap(); + } + let before = reader.search(&query, 25, 4).unwrap(); + let precomputed = reader.precomputed_table.clone(); + for result in [ + reader.range_search(&query, params), + reader.range_search_with_roaring_filter(&query, params, &filter), + reader.range_search_batch(&query, 1, params), + reader.range_search_batch_with_roaring_filter(&query, 1, params, &filter), + reader.range_search_with_filter(&query, params, Some(&allowed)), + reader.range_search_batch_with_filter(&query, 1, params, Some(&allowed)), + ] { + let result = result.unwrap(); + assert_eq!(pq_rows(result.query(0)), expected); + assert_eq!(result.query(0).stats.rows_scanned(), expected.len()); + assert_eq!(result.query(0).stats.lists_probed(), 4); + assert_eq!(result.call_stats().list_reads(), 3); + } + assert_eq!(reader.search(&query, 25, 4).unwrap(), before); + assert_eq!(reader.precomputed_table, precomputed); + } + } + } + } +} + +#[test] +fn pq_range_empty_bands_validate_before_skipping_metadata_io() { + let bytes = pq_bytes(&pq_fixture(4, true, false)); + let mut reader = VectorIndexReader::open(Cursor::new(bytes[..64].to_vec())).unwrap(); + let mut native = IVFPQIndexReader::open(Cursor::new(bytes[..64].to_vec())).unwrap(); + let empty = VectorRangeSearchParams::new(l2(1.0, 1.0), 4); + let filter = serialize_roaring(&HashSet::new()); + let query = [0.0; 8]; + let unified = pq_entry_points(&mut reader, &query, empty, &filter); + let direct = [ + native.range_search(&query, empty), + native.range_search_with_roaring_filter(&query, empty, &filter), + native.range_search_batch(&query, 1, empty), + native.range_search_batch_with_roaring_filter(&query, 1, empty, &filter), + ]; + for result in unified.into_iter().chain(direct) { + let result = result.unwrap(); + assert_eq!(result.lims(), &[0, 0]); + assert_eq!(result.query(0).stats.lists_probed(), 0); + assert_eq!(result.query(0).stats.rows_scanned(), 0); + assert_eq!(result.call_stats().list_reads(), 0); + } + for bad_query in [ + vec![0.0; 7], + vec![f32::NAN; 8], + vec![f32::INFINITY; 8], + vec![f32::NEG_INFINITY; 8], + ] { + for result in pq_entry_points(&mut reader, &bad_query, empty, &filter) { + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidInput); + } + assert_eq!( + native.range_search(&bad_query, empty).unwrap_err().kind(), + io::ErrorKind::InvalidInput + ); + } + for params in [empty, pq_all_params(4)] { + assert_eq!( + reader + .range_search_with_roaring_filter(&query, params, b"invalid") + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + native + .range_search_batch_with_roaring_filter(&query, 1, params, b"invalid") + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + } + for count in [0, 2, usize::MAX] { + assert_eq!( + reader + .range_search_batch(&query, count, empty) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + native + .range_search_batch(&query, count, empty) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + } + for result in pq_entry_points(&mut reader, &query, pq_all_params(0), &filter) { + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidInput); + } +} + +#[test] +fn pq_range_metric_mismatch_and_uncertified_metrics_fail_loud() { + let filter = serialize_roaring(&HashSet::new()); + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + let mut index = pq_fixture(8, false, false); + index.metric = metric; + let mut reader = pq_reader(&index); + let mut native = IVFPQIndexReader::open(Cursor::new(pq_bytes(&index))).unwrap(); + for band_metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + for nprobe in [0, 4] { + let band = + DistanceBand::new(Bound::Finite(1.0), Bound::Finite(1.0), band_metric).unwrap(); + let params = VectorRangeSearchParams::new(band, nprobe); + for result in pq_entry_points(&mut reader, &[0.0; 8], params, &filter) + .into_iter() + .chain([native.range_search(&[0.0; 8], params)]) + { + if nprobe == 0 || metric != band_metric { + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidInput); + } else if metric != MetricType::L2 { + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::Unsupported); + } else { + assert!(result.is_ok()); + } + } + } + } + } +} + +#[test] +fn pq_range_rejects_nonfinite_unselected_coarse_data_and_rotation() { + let filter = serialize_roaring(&HashSet::new()); + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY, f32::MAX] { + for rotation in [false, true] { + let index = pq_fixture(8, true, true); + let mut reader = pq_reader(&index); + let VectorIndexReader::IvfPq(native) = &mut reader else { + unreachable!() + }; + native.ensure_loaded().unwrap(); + if rotation { + native.opq.as_mut().unwrap().rotation[0] = value; + } else { + let last = native.quantizer_centroids.len() - 1; + native.quantizer_centroids[last] = value; + } + for result in pq_entry_points(&mut reader, &[2.0; 8], pq_all_params(1), &filter) { + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidData); + } + } + } +} + +#[test] +fn pq_range_nonfinite_estimates_are_checked_only_for_eligible_rows() { + for bits in [4, 8] { + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY, f32::MAX] { + let mut index = pq_fixture(bits, false, false); + for list in 0..index.nlist { + index.ids[list].clear(); + index.codes[list].clear(); + } + index.ids[0] = vec![1, 2]; + index.codes[0] = vec![0; index.pq.code_size() * 2]; + index.codes[0][index.pq.code_size()] = 1; + let mut centroids = index.pq.centroids().to_vec(); + centroids[index.pq.dsub()] = value; + index.pq.set_centroids(centroids); + let mut reader = pq_reader(&index); + for band in [pq_all_params(4).band(), l2(0.0, 0.001)] { + let params = VectorRangeSearchParams::new(band, 4); + assert_eq!( + reader.range_search(&[0.0; 8], params).unwrap_err().kind(), + io::ErrorKind::InvalidData + ); + for allowed in [HashSet::new(), HashSet::from([1])] { + let filter = serialize_roaring(&allowed); + let filtered = reader + .range_search_with_roaring_filter(&[0.0; 8], params, &filter) + .unwrap(); + assert_eq!(filtered.query(0).stats.rows_scanned(), allowed.len()); + assert_eq!(filtered.query(0).stats.early_abandoned(), 0); + } + let filter = serialize_roaring(&HashSet::from([2])); + assert_eq!( + reader + .range_search_batch_with_roaring_filter(&[0.0; 16], 2, params, &filter) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + } + } + } +} + +#[test] +fn pq_range_uses_estimated_not_original_distances() { + for bits in [4, 8] { + let mut index = IVFPQIndex::with_nbits(8, 1, 4, bits, MetricType::L2, false); + index.set_quantizer_centroids(vec![0.0; 8]); + index.pq.set_centroids(vec![0.0; index.d * index.pq.ksub()]); + let original = [0.25; 8]; + index.add(&original, &[42], 1); + let band = l2(0.0, 0.25); + assert!(!band.admit(fvec_l2sqr(&original, &[0.0; 8]))); + let result = pq_reader(&index) + .range_search(&[0.0; 8], VectorRangeSearchParams::new(band, 1)) + .unwrap(); + assert_eq!(pq_rows(result.query(0)), vec![(42, 0.0)]); + } +} + +#[test] +fn pq_range_direct_adc_avoids_norm_cancellation_and_rejects_sum_overflow() { + for bits in [4, 8] { + let mut index = IVFPQIndex::with_nbits(8, 1, 4, bits, MetricType::L2, false); + index.by_residual = false; + index.set_quantizer_centroids(vec![0.0; 8]); + index + .pq + .set_centroids(vec![100_000_008.0; index.d * index.pq.ksub()]); + index.ids[0] = vec![42]; + index.codes[0] = vec![0; index.pq.code_size()]; + let mut reader = pq_reader(&index); + let result = reader + .range_search(&[100_000_000.0; 8], pq_all_params(1)) + .unwrap(); + assert_eq!(pq_rows(result.query(0)), vec![(42, 512.0)]); + let value = (f32::MAX / 4.0).sqrt(); + assert!((2.0 * value * value).is_finite()); + index + .pq + .set_centroids(vec![value; index.d * index.pq.ksub()]); + assert_eq!( + pq_reader(&index) + .range_search(&[0.0; 8], pq_all_params(1)) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + } +} + +struct PqStreamingSource { + prefix: Vec, + rows: usize, + code_size: usize, + tail_code: u8, + reads: std::sync::Arc>>, +} + +impl SeekRead for PqStreamingSource { + fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()> { + assert!(ranges.len() <= 3); + assert!( + ranges + .iter() + .map(|request| request.buf.len()) + .sum::() + <= 64 * 1024 * 1024 + ); + for request in ranges { + let start = request.pos as usize; + let end = start + request.buf.len(); + if end > self.prefix.len() + self.rows * self.code_size { + return Err(io::ErrorKind::UnexpectedEof.into()); + } + self.reads.lock().unwrap().push((start, request.buf.len())); + request.buf.fill(0); + if start < self.prefix.len() { + let prefix_end = end.min(self.prefix.len()); + request.buf[..prefix_end - start].copy_from_slice(&self.prefix[start..prefix_end]); + } + for column in 0..self.code_size { + let tail_position = self.prefix.len() + (column + 1) * self.rows - 1; + if (start..end).contains(&tail_position) { + request.buf[tail_position - start] = self.tail_code; + } + } + } + Ok(()) + } + + fn read_capabilities(&self) -> SeekReadCapabilities { + SeekReadCapabilities { + max_ranges_per_pread: 3, + ..Default::default() + } + } +} + +fn pq_streaming_source(bits: usize, corrupt_first_code: bool) -> PqStreamingSource { + let dimension = 256; + let mut index = IVFPQIndex::with_nbits(dimension, 1, dimension, bits, MetricType::L2, false); + index.by_residual = false; + index.set_quantizer_centroids(vec![0.0; dimension]); + index.pq.set_centroids( + (0..dimension * index.pq.ksub()) + .map(|position| (position % index.pq.ksub()) as f32) + .collect(), + ); + index.ids[0] = vec![0]; + index.codes[0] = vec![0; index.pq.code_size()]; + let mut bytes = pq_bytes(&index); + let rows = 64 * 1024 * 1024 / index.pq.code_size() + 1; + let offset_table = 64 + dimension * 4 + dimension * index.pq.ksub() * 4; + let list_offset = offset_table + 16; + bytes[32..40].copy_from_slice(&(rows as i64).to_le_bytes()); + bytes[offset_table + 8..offset_table + 12].copy_from_slice(&(rows as i32).to_le_bytes()); + bytes[offset_table + 12..offset_table + 16].copy_from_slice(&(rows as i32).to_le_bytes()); + bytes.truncate(list_offset); + bytes.extend_from_slice(&0i64.to_le_bytes()); + bytes.extend_from_slice(&(rows as i32).to_le_bytes()); + bytes.push(0); + bytes.resize(bytes.len() + rows - 1, 1); + if corrupt_first_code { + let codebook = 64 + dimension * 4; + bytes[codebook..codebook + 4].copy_from_slice(&f32::NAN.to_le_bytes()); + } + PqStreamingSource { + prefix: bytes, + rows, + code_size: index.pq.code_size(), + tail_code: if bits == 4 { 0x11 } else { 1 }, + reads: Default::default(), + } +} + +#[test] +fn pq_range_streams_oversized_lists_once_and_propagates_collector_errors() { + let pools = [1, 4].map(|threads| { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + }); + let query_count = 33; + let selected_prefix_rows = 257; + assert!(query_count * selected_prefix_rows > 8192); + assert!(query_count * 256 * 256 * size_of::() > 8 * 1024 * 1024); + let queries = (0..query_count) + .flat_map(|query| vec![(query + 1) as f32 / 64.0; 256]) + .collect::>(); + for bits in [4, 8] { + for corrupt in [false, true] { + let source = pq_streaming_source(bits, corrupt); + let rows = source.rows; + let code_start = source.prefix.len(); + let code_bytes = rows * source.code_size; + let reads = source.reads.clone(); + let mut reader = VectorIndexReader::open(source).unwrap(); + let allowed = (0..selected_prefix_rows as i64) + .chain(std::iter::once(rows as i64 - 1)) + .collect::>(); + let filter = serialize_roaring(&allowed); + let mut reference_reads = None; + for pool in &pools { + pool.install(|| { + reads.lock().unwrap().clear(); + let outcome = reader.range_search_batch_with_roaring_filter( + &queries, + query_count, + pq_all_params(1), + &filter, + ); + let payload_reads = reads + .lock() + .unwrap() + .iter() + .copied() + .filter(|&(start, _)| start >= code_start) + .collect::>(); + let read_bytes = payload_reads + .iter() + .map(|&(_, length)| length) + .sum::(); + if let Some(reference) = &reference_reads { + assert_eq!( + &payload_reads, reference, + "worker count must not change I/O" + ); + } else { + reference_reads = Some(payload_reads); + } + if corrupt { + assert_eq!(outcome.unwrap_err().kind(), io::ErrorKind::InvalidData); + assert!( + read_bytes < code_bytes, + "the collector error must stop further chunks" + ); + } else { + let result = outcome.unwrap(); + assert_eq!( + read_bytes, code_bytes, + "shared list payload must be read exactly once" + ); + assert_eq!(result.call_stats().list_reads(), 1); + assert_eq!(result.query_count(), query_count); + for (query_index, query) in queries.chunks_exact(256).enumerate() { + let mut expected = (0..selected_prefix_rows as i64) + .map(|id| (id, 256.0 * query[0].powi(2))) + .collect::>(); + expected.push((rows as i64 - 1, 256.0 * (query[0] - 1.0).powi(2))); + assert_eq!(pq_rows(result.query(query_index)), expected); + let stats = result.query(query_index).stats; + assert_eq!(stats.rows_scanned(), allowed.len()); + assert_eq!(stats.rows_committed(), allowed.len()); + assert_eq!(stats.lists_probed(), 1); + assert_eq!(stats.early_abandoned(), 0); + } + } + }); + } + } + } +} + +#[test] +fn pq_range_empty_lists_and_shared_probe_reads_have_exact_counts() { + let index = pq_fixture(4, true, false); + let queries = index.quantizer_centroids().to_vec(); + let result = pq_reader(&index) + .range_search_batch(&queries, 4, pq_all_params(1)) + .unwrap(); + assert_eq!(result.call_stats().list_reads(), 3); + for query in 0..4 { + let expected = if query < 3 { 273 } else { 0 }; + assert_eq!(result.query(query).stats.rows_scanned(), expected); + assert_eq!(result.query(query).stats.lists_probed(), 1); + assert_eq!(result.query(query).labels.len(), expected); + } + let mut empty = pq_fixture(8, true, false); + for list in 0..empty.nlist { + empty.ids[list].clear(); + empty.codes[list].clear(); + } + let result = pq_reader(&empty) + .range_search_batch(&queries, 4, pq_all_params(4)) + .unwrap(); + assert_eq!(result.call_stats().list_reads(), 0); + assert_eq!(result.lims(), &[0, 0, 0, 0, 0]); + for query in 0..4 { + assert_eq!(result.query(query).stats.rows_scanned(), 0); + assert_eq!(result.query(query).stats.lists_probed(), 4); + } +} + 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( diff --git a/docs/api.html b/docs/api.html index d1daf06..915cb78 100644 --- a/docs/api.html +++ b/docs/api.html @@ -76,7 +76,7 @@

Shared search parameters

Range search parameters and results

-

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, IVF-RQ, and IVF-SQ with l2: IVF-FLAT tests exact distances, while IVF-RQ and IVF-SQ test quantized estimates. 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. Range bindings are not included.

+

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, IVF-RQ, IVF-SQ, and IVF-PQ with l2: IVF-FLAT tests exact distances, while IVF-RQ, IVF-SQ, and IVF-PQ test quantized estimates. 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. Range bindings are not included.

diff --git a/docs/index.html b/docs/index.html index 074193d..c98ab3f 100644 --- a/docs/index.html +++ b/docs/index.html @@ -126,7 +126,7 @@

Distance range search support

- +
Entry pointSignature
Single queryrange_search(query, params)
Single query, filteredrange_search_with_roaring_filter(query, params, filter_bytes)
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-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-SQSupported in Rust, l2 onlySQ estimateSingle/batch, with or without a Roaring filter; even full probing cannot guarantee membership under the original vectors' distances
IVF-PQPlannedEstimateAs above; would use the float-LUT path rather than fastscan
IVF-PQSupported, l2 onlyFloating-point ADC estimate4-bit/8-bit, single/batch and Roaring-filtered variants; no FastScan or top-K truncation, but quantization error remains
DiskANNNot planned—Graph traversal is inherently k-oriented and has no natural radius termination criterion
diff --git a/docs/ivf-pq.html b/docs/ivf-pq.html index 86deb93..7c60d9d 100644 --- a/docs/ivf-pq.html +++ b/docs/ivf-pq.html @@ -17,6 +17,7 @@

Positioning and trade-offs

+

Rust distance range search supports squared L2 for 4-bit and 8-bit PQ, with or without residual encoding and OPQ. All four single/batch and Roaring-filtered entry points use floating-point ADC estimates rather than FastScan or top-K truncation. Even probing every list does not make membership exact relative to the original vectors.

Vector payloadm bytes
Sub-codebook256 centroids per subspace
SearchADC distance lookup
Extra trainingPQ; optional OPQ

Good fit

  • Million-scale or larger collections where raw vectors are too large.
  • A mature and tunable accuracy–space trade-off.
  • Cache locality and scan throughput matter more than exact distances.
  • Long-lived Readers serve many repeated queries.

Poor fit

  • Quantization-induced rank changes are unacceptable.
  • The dimension has no useful divisor for m.
  • Training samples are sparse or the distribution changes frequently.
  • The simplest possible build pipeline is required.
Current public configurationThe unified VectorIndexConfig::IvfPq and options API always create 8-bit PQ. The core and v1 format contain a 4-bit path, but no public pq.nbits option currently exposes it.
diff --git a/docs/range-search.html b/docs/range-search.html index 440bd83..c56081c 100644 --- a/docs/range-search.html +++ b/docs/range-search.html @@ -11,13 +11,13 @@
VIPaimon Vector Index
-

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 and IVF-SQ compute 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, IVF-RQ and IVF-SQL2 only · Rust API
+

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, IVF-SQ and IVF-PQ compute 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, IVF-RQ, IVF-SQ and IVF-PQL2 only · Rust API

Semantic contract

-
IntervalHalf-open [lower, upper)
Result sizeUnbounded
Row orderUnspecified
MembershipExact (FLAT), estimated (RQ)
+
IntervalHalf-open [lower, upper)
Result sizeUnbounded
Row orderUnspecified
MembershipExact (FLAT), estimated (RQ/SQ/PQ)

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 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.
@@ -29,14 +29,14 @@

Units and ordering

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 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-SQ similarly uses its blocked estimated squared distances to prune at the exclusive upper cut. 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. All three families use the supplied cuts without a margin.
+
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-SQ similarly uses its blocked estimated squared distances to prune at the exclusive upper cut. IVF-RQ and IVF-PQ do not use this rule or top-K's coarse/FastScan bounds. They evaluate the complete estimate from F32 lookup sums before testing either cut, so every eligible row reaches the band test. All four 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.

-
Coverage and estimation are different gapsRaising nprobe improves list coverage; it does not remove the quantization error of IVF-RQ or IVF-SQ. 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. None of these families 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.
+
Coverage and estimation are different gapsRaising nprobe improves list coverage; it does not remove the quantization error of IVF-RQ, IVF-SQ or IVF-PQ. 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. None of these families 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.
@@ -79,13 +79,14 @@

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().

-

All three 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 at most once per call and shared across queries; IVF-SQ cache hits require no payload read.

-

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.

+

All four 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 at most once per call and shared across queries; IVF-SQ cache hits require no payload read.

+

For IVF-RQ and IVF-PQ, 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, IVF-RQ, and IVF-SQ. IVF-PQ and DiskANN still return Unsupported. Only L2, fixed probe widths, and Rust entry points are included.

+

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

+

IVF-PQ uses direct squared-L2 lookup tables for 4-bit and 8-bit codes, after optional OPQ rotation and coarse residual subtraction. The selected table entries are summed in subquantizer order without early abandonment. Neither u8 FastScan nor precomputed norm identities participate, so membership does not depend on list size, batch size, or optimize_for_search. These estimates may differ from top-K's distances. They are not distances to the original vectors; probing every list does not eliminate quantization error. Oversized lists use the existing bounded streaming reader.

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.
IVF-SQ membership uses an estimate.IVF-FLAT computes exact distances from stored f32 vectors. IVF-SQ instead reuses top-K's blocked scalar-quantized estimator, reconstructing residuals with each list's stored bounds and centroid. The same estimated value determines band membership and is returned in distances; there is no original-vector reranking and no top-K fallback. Prefer IVF-FLAT if original-distance membership must be exact.

A reproducible boundary example is in core/tests/range_search.rs: with a one-dimensional centroid of zero and SQ bounds [0, 255], inputs 0.49 and 0.51 quantize to 0 and 1. For query zero, band [0, 0.1) includes the first estimate despite its true squared distance being outside; band [0.2, 0.3) misses both although both true squared distances lie inside. This demonstrates the membership gap, not a general recall estimate.

@@ -97,8 +98,9 @@

Choosing an index type

Fail-loud combinations

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-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
+
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
DiskANNUnsupported
An endpoint with no representable cut, meaning do not push the predicate downUnsupported
A non-finite consumed IVF-RQ/PQ value 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.

+

IVF-PQ likewise checks the transformed query and every direct coarse distance, including unselected centroids. Every consumed PQ estimate must be finite, including for rows outside the requested band; a NaN, infinity, or overflowing accumulation returns InvalidData without a partial result. Filtered-out and unprobed rows are not evaluated, so unused non-finite codebook entries do not by themselves fail the scan.

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.

@@ -111,6 +113,6 @@

Why not DiskANN

-
Apache Paimon Vector IndexRange search · IVF-FLAT, IVF-RQ and IVF-SQ
+