diff --git a/core/README.md b/core/README.md index 3a22afb..00958a2 100644 --- a/core/README.md +++ b/core/README.md @@ -22,14 +22,17 @@ `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 -squared L2, using `DistanceBand`, `VectorRangeSearchParams`, and CSR -`RangeSearchResult` buffers. All three families support single and batch queries, +The Rust reader supports distance range search for IVF-FLAT, IVF-SQ, IVF-PQ and +IVF-RQ with L2, cosine and inner product, using `DistanceBand`, +`VectorRangeSearchParams`, and CSR `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. See the [range search guide](../docs/range-search.html) for membership, @@ -42,11 +45,36 @@ abstractions. It does not incorporate source code from Microsoft's The implementation supports L2, inner-product, and cosine search with the same lower-is-better distance semantics as the IVF indexes. -## IVF-SQ range search +## Distance semantics + +`DistanceBand::new` takes internal, lower-is-better f32 scores: squared L2, +cosine distance (`1 - cos`, not clamped), or negative inner product. +`MetricType::public_distance` converts a score to its public f64 predicate value: +L2 takes the f32 square root before widening, cosine widens unchanged, and inner +product negates. Returned `distances` always remain in internal units. + +Use `DistanceBand::from_endpoints` for public f64 predicates. `Ge`/`Gt` belong on +the lower side and `Le`/`Lt` on the upper side. L2 resolves square-root rounding; +cosine searches all finite f32 values (including negative roundoff and signed +zero); inner product reverses both the side and comparison when negating. +Endpoints are never rounded to f32 first. Missing sides are structurally +unbounded, so an inclusive endpoint at `f32::MAX` does not lose that value. +Out-of-domain cosine/IP predicates become empty or unbounded bands; unrepresentable +L2 cuts retain the existing `Unsupported` response. + +`IndexType::supports_range_search(metric)` and `reader.supports_range_search()` +report capability without reading list payloads. DiskANN and C/JNI range APIs +remain unsupported. Bad queries, mismatched metrics, non-finite endpoints, +malformed filters and zero `nprobe` are errors even for an empty band. +Non-finite consumed distances or cosine norms return `InvalidData`, not partial +results. Cosine queries use the existing normalization, including leaving zero +vectors at zero; normalization overflow is an error on this path. + +## IVF-SQ and IVF-PQ range search Use `DistanceBand` and `VectorRangeSearchParams` with `range_search`, `range_search_batch`, or their `*_with_roaring_filter` variants. Bands are -half-open `[lower, upper)` in squared-L2 units; results use `RangeSearchResult` +half-open `[lower, upper)` in internal metric units; results use `RangeSearchResult` CSR buffers without sorting, padding, or a top-K cap. IVF-SQ uses the same blocked SIMD estimator as top-K, including the stored @@ -57,11 +85,32 @@ membership is required. There is no original-vector reranking or top-K fallback. Batch queries share list reads, reuse the existing partition cache, and evaluate the Roaring allow-list once per list row using query-local one-bit-per-row masks. -Large lists stream in bounded chunks; scan scratch is reused, and a finite upper cut -allows entire SQ blocks to stop after their partial distances reach that cut. +Large lists stream in bounded chunks; scan scratch is reused, and under L2 a +finite upper cut 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. +`call_stats().list_reads()`. Cosine and IP never prune partial sums and report zero +`early_abandoned`; L2 pruning is unchanged. Filters exclude negative row IDs in +the Roaring API, while direct `RowIdFilter` implementations may admit them. + +IVF-PQ sums full f32 subvector distances in subquantizer order rather than using +top-K's quantized FastScan tables. +L2 uses ADC squared distances; cosine uses half the ADC squared distance after +query normalization; IP uses negative estimated dot product. The cosine score +is a unit-vector surrogate, not a re-normalized exact distance to decoded codes. +Like IVF-RQ's cosine estimator, it can differ from exact cosine even for zero +queries. IVF-FLAT and IVF-SQ define cosine distance involving a zero vector as 1. +PQ and RQ always evaluate complete estimates, without early abandonment. +Shared lists are read once per call, oversized PQ lists stream in bounded chunks, +and the PQ batch filter is evaluated once per list row, not once per query. +Query lookup tables are allocated lazily within an 8 MiB cache cap, with reusable +worker scratch beyond that budget. Residual tables are reused across chunks of +the same list and rebuilt for a different list. Large batches scan queries in +parallel; membership does not depend on worker count, list size, batch size or +`optimize_for_search`. Non-finite transformed queries, coarse distances (even to +unselected lists), or consumed estimates return `InvalidData` without partial results. +PQ range scores need not be bit-identical to top-K scores: the range path does +not use top-K table quantization, L2 table expansion, or its cosine score scale. 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..bd3c829 100644 --- a/core/src/collect.rs +++ b/core/src/collect.rs @@ -36,6 +36,8 @@ use crate::range::{Bound, DistanceBand}; /// The interface a scan kernel uses to hand candidate rows to a collector. pub(crate) trait Collector { + const VALIDATE_COSINE_INPUTS: bool = false; + /// Called when the scan kernel rejects a row against [`cutoff`] instead of /// delivering it. /// @@ -140,6 +142,8 @@ impl RangeCollector { } impl Collector for RangeCollector { + const VALIDATE_COSINE_INPUTS: bool = true; + #[inline] fn note_abandoned(&mut self) { self.scanned += 1; @@ -222,7 +226,7 @@ mod tests { } #[test] - fn an_uncertified_metric_never_prunes_on_a_partial_sum() { + fn a_non_l2_metric_never_prunes_on_a_partial_sum() { // A partial cosine or inner-product accumulation does not bound the full // value, so the cutoff must stay infinite no matter what the band says. for metric in [MetricType::Cosine, MetricType::InnerProduct] { diff --git a/core/src/distance.rs b/core/src/distance.rs index 7dadf93..b0179a3 100644 --- a/core/src/distance.rs +++ b/core/src/distance.rs @@ -27,6 +27,17 @@ pub enum MetricType { } impl MetricType { + /// Converts a finite internal score to the public predicate value. + /// L2 takes an f32 square root before widening; cosine is `1 - cos` + /// without clamping; inner product reverses the internal score's sign. + pub fn public_distance(self, distance: f32) -> f64 { + f64::from(match self { + Self::L2 => distance.sqrt(), + Self::Cosine => distance, + Self::InnerProduct => -distance, + }) + } + pub fn from_code(code: u32) -> Option { match code { 0 => Some(MetricType::L2), diff --git a/core/src/index.rs b/core/src/index.rs index 8baa879..7640e3c 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -121,6 +121,18 @@ pub enum IndexType { } impl IndexType { + /// Range membership is exact for IVF-FLAT and estimated for SQ, PQ and RQ. + /// DiskANN does not support range search. This does not imply full IVF coverage. + pub fn supports_range_search(self, metric: MetricType) -> bool { + matches!( + self, + Self::IvfFlat | Self::IvfSq | Self::IvfPq | Self::IvfRq + ) && matches!( + metric, + MetricType::L2 | MetricType::Cosine | MetricType::InnerProduct + ) + } + pub fn from_code(code: u32) -> Option { match code { 0 => Some(Self::IvfFlat), @@ -1567,6 +1579,12 @@ pub enum VectorIndexReader { } impl VectorIndexReader { + /// Reports family/metric capability without loading list payloads. + pub fn supports_range_search(&self) -> bool { + let metadata = self.metadata(); + metadata.index_type.supports_range_search(metadata.metric) + } + pub fn open(reader: R) -> io::Result { Self::open_with_options(reader, VectorIndexReaderOptions::default()) } @@ -1868,7 +1886,8 @@ 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. + /// [`IVFSQIndexReader::range_search`] (SQ estimates), and + /// [`IVFPQIndexReader::range_search`] (PQ estimates). All three metrics are 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,7 +1903,7 @@ 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")), } } @@ -1909,7 +1928,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 +1947,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 +1973,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")), } } @@ -2677,7 +2698,7 @@ fn validate_query(query: &[f32], dimension: usize) -> io::Result<()> { validate_finite_values(query, dimension, "query") } -/// IVF-Flat and IVF-RQ implement range search. The other families return +/// The IVF families implement range search. DiskANN returns /// `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/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/ivfflat_io.rs b/core/src/ivfflat_io.rs index 335d2e2..cadf1d8 100644 --- a/core/src/ivfflat_io.rs +++ b/core/src/ivfflat_io.rs @@ -26,7 +26,10 @@ use crate::io::{ReadRequest, SeekRead, SeekWrite}; use crate::ivfflat::IVFFlatIndex; use crate::ivfpq::RowIdFilter; use crate::kmeans; -use crate::range::{RangeResultBuilder, RangeSearchResult, VectorRangeSearchParams}; +use crate::range::{ + checked_cosine_norm, prepare_range_queries, range_probe_lists, RangeResultBuilder, + RangeSearchResult, VectorRangeSearchParams, +}; use rayon::prelude::*; use roaring::RoaringTreemap; use std::io; @@ -882,12 +885,8 @@ impl IVFFlatIndexReader { // centroid array. self.ensure_loaded()?; - // Note for metric certification: unlike the top-K path this does not - // apply `fvec_normalize` for cosine. That is currently unreachable, - // because `params.validate` rejects every non-L2 metric above, but - // whoever certifies cosine must add the normalization here as well as - // relaxing `ensure_certified_metric`. - // + let processed = prepare_range_queries(queries, self.d, self.metric)?; + let queries = processed.as_ref(); // Ranked probe selection, one group per query. // // This relies on `find_topk_batch` selecting the same centroids for a @@ -899,14 +898,14 @@ impl IVFFlatIndexReader { // its error bound leaves the ranking ambiguous, which is what makes the // two agree. `kmeans` pins that property with a test; if it is ever // relaxed, this caller needs an exact helper of its own again. - let (probe_lists, _coarse_distances) = kmeans::find_topk_batch( + let probe_lists = range_probe_lists( queries, - nq, &self.quantizer_centroids, - self.nlist, self.d, + self.nlist, nprobe, - ); + self.metric, + )?; for (qi, lists) in probe_lists.iter().enumerate() { builder.record_lists_probed(qi, lists.len()); } @@ -1344,6 +1343,11 @@ fn scan_flat_rows( continue; } let vector = &vectors[local_idx * d..(local_idx + 1) * d]; + let vector_norm = if C::VALIDATE_COSINE_INPUTS && metric == MetricType::Cosine { + Some(checked_cosine_norm(vector)?) + } else { + None + }; let distance = if metric == MetricType::L2 { // One traversal, whatever the outcome: the kernel abandons against // the collector's cutoff and otherwise hands back the value @@ -1359,7 +1363,7 @@ fn scan_flat_rows( } } } else { - distance_context.distance_to(vector, None) + distance_context.distance_to(vector, vector_norm) }; collector.push(id, distance)?; } diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index 60644e8..2394219 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -16,16 +16,21 @@ // 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::{ + prepare_range_queries, RangeResultBuilder, RangeSearchResult, VectorRangeSearchParams, +}; use crate::sparse_table::SparseTable; use rayon::prelude::*; use roaring::RoaringTreemap; @@ -1665,6 +1670,430 @@ 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 band. + /// + /// Both code widths use floating-point ADC, summed in subquantizer order. + /// L2 uses direct squared distances from the query (after OPQ and optional + /// coarse residual subtraction) to each selected PQ centroid. Range search + /// does not use top-K's u8 FastScan tables or precomputed norm identities, + /// so membership is independent of list size, batch 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. + /// Cosine normalizes the query before OPQ and returns half the squared-L2 + /// ADC estimate. Inner product uses negative dot-product tables, including + /// the coarse-centroid contribution for residual codes. Neither estimator + /// establishes exact membership over the original vectors. + /// + /// 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 mut query_states = Vec::with_capacity(query_count); + let mut list_to_queries = vec![Vec::new(); self.nlist]; + for (query_index, query) in queries.chunks_exact(self.d).enumerate() { + let normalized = prepare_range_queries(query, self.d, self.metric)?; + let query = normalized.as_ref(); + 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}"), + ) + })?; + builder.record_lists_probed(query_index, probes.len()); + for (_, list) in probes { + list_to_queries[list].push(query_index); + } + query_states.push(Some(PqRangeQuery::new( + query_index, + prepared, + self.metric, + RangeCollector::new(params.band()), + ))); + } + 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; + +#[derive(Default)] +struct PqRangeScratch { + residual_query: Vec, + table: Vec, +} + +struct PqRangeQuery { + query_index: usize, + query: Vec, + metric: MetricType, + collector: C, + table: Vec, + table_list: Option, +} + +impl PqRangeQuery { + fn new(query_index: usize, query: Vec, metric: MetricType, collector: C) -> Self { + Self { + query_index, + query, + metric, + 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, + query.metric, + centroid, + &mut scratch.residual_query, + &mut query.table, + ); + query.table_list = Some(list_id); + } + &query.table + } else { + build_pq_range_table( + pq, + &query.query, + query.metric, + centroid, + &mut scratch.residual_query, + &mut scratch.table, + ); + &scratch.table + }; + scan_pq_range_codes( + pq, + table, + query.metric, + 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], + metric: MetricType, + centroid: Option<&[f32]>, + residual_query: &mut Vec, + table: &mut Vec, +) { + let residual_centroid = centroid.filter(|_| metric != MetricType::InnerProduct); + let query = if let Some(centroid) = residual_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(); + let codeword = &pq.centroids()[offset..offset + pq.dsub()]; + table[sub * pq.ksub() + code] = match metric { + MetricType::L2 | MetricType::Cosine => fvec_l2sqr(query_chunk, codeword), + MetricType::InnerProduct => -fvec_inner_product(query_chunk, codeword), + }; + } + } + if metric == MetricType::InnerProduct { + if let Some(centroid) = centroid { + let offset = -fvec_inner_product(query, centroid); + for entry in &mut table[..pq.ksub()] { + *entry += offset; + } + } + } +} + +fn scan_pq_range_codes( + pq: &ProductQuantizer, + table: &[f32], + metric: MetricType, + 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]; + } + if metric == MetricType::Cosine { + distance *= 0.5; + } + 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 +3393,296 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; + #[test] + fn pq_range_scans_active_queries_on_multiple_workers() { + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + check_pq_range_parallel_workers(metric); + } + } + + fn check_pq_range_parallel_workers(metric: MetricType) { + 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, + metric, + ) + .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(), + metric, + 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 squared = 8.0 * (query.query_index * query.query_index) as f32; + let distance = match metric { + MetricType::L2 => squared, + MetricType::Cosine => squared * 0.5, + MetricType::InnerProduct => 0.0, + }; + 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, + metric: MetricType, + ) -> Vec> { + let band = crate::range::DistanceBand::new( + crate::range::Bound::Unbounded, + crate::range::Bound::Unbounded, + metric, + ) + .unwrap(); + (0..count) + .rev() + .map(|query_index| { + PqRangeQuery::new( + query_index, + vec![query_index as f32 * 0.25; 32], + metric, + RangeCollector::new(band), + ) + }) + .collect() + } + + #[test] + fn pq_range_reuses_bounded_tables_without_changing_estimates() { + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + check_pq_range_cached_estimates(metric); + } + } + + fn check_pq_range_cached_estimates(metric: MetricType) { + 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, metric); + let mut uncached = pq_range_test_queries(5, metric); + 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() { + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + check_pq_range_excluded_rows(metric); + } + } + + fn check_pq_range_excluded_rows(metric: MetricType) { + 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, metric); + 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() { + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + check_pq_range_collector_errors(metric); + } + } + + fn check_pq_range_collector_errors(metric: MetricType) { + 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], metric, 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/ivfrq_io.rs b/core/src/ivfrq_io.rs index fc9b866..3bde51f 100644 --- a/core/src/ivfrq_io.rs +++ b/core/src/ivfrq_io.rs @@ -26,7 +26,9 @@ 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::range::{ + prepare_range_queries, 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, @@ -615,7 +617,7 @@ impl IVFRQIndexReader { } /// Returns every eligible row in the probed lists whose IVF-RQ estimated - /// distance is in the requested band. Only squared L2 is supported. + /// distance is in the requested band, for L2, cosine or inner product. /// /// 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: @@ -706,6 +708,12 @@ impl IVFRQIndexReader { { return Err(invalid_data("non-finite IVF-RQ centroid")); } + let processed = prepare_range_queries(queries, self.d, self.metric)?; + let queries = processed.as_ref(); + let query_norms = queries + .chunks_exact(self.d) + .map(fvec_norm_l2sqr) + .collect::>(); let probe_lists = queries .par_chunks_exact(self.d) .map(|query| { @@ -764,10 +772,22 @@ impl IVFRQIndexReader { } let scan_one = |list: &RQReadList, query_index: usize, distance: f32| -> io::Result<()> { - let terms = RQQueryTerms { - g_add: distance, - g_error: distance.sqrt(), + let terms = if self.metric == MetricType::L2 { + RQQueryTerms { + g_add: distance, + g_error: distance.sqrt(), + } + } else { + self.quantizer.query_terms_from_coarse_distance( + distance, + query_norms[query_index], + self.quantizer_centroid_norms[list.list_id], + self.metric, + ) }; + if !terms.g_add.is_finite() || !terms.g_error.is_finite() { + return Err(invalid_data("non-finite IVF-RQ query terms")); + } let mut collector = RangeCollector::new(params.band()); scan_range_blocked_list( list, diff --git a/core/src/ivfsq_io.rs b/core/src/ivfsq_io.rs index fc3c06e..4bd0883 100644 --- a/core/src/ivfsq_io.rs +++ b/core/src/ivfsq_io.rs @@ -32,7 +32,10 @@ use crate::io::{ReadRequest, SeekRead, SeekWrite}; use crate::ivfpq::RowIdFilter; use crate::ivfsq::IVFSQIndex; use crate::kmeans; -use crate::range::{RangeResultBuilder, RangeSearchResult, VectorRangeSearchParams}; +use crate::range::{ + prepare_range_queries, range_probe_lists, RangeResultBuilder, RangeSearchResult, + VectorRangeSearchParams, +}; use crate::read_options::VectorIndexReaderOptions; use crate::sq::ScalarQuantizer; use crate::topk::TopKHeap; @@ -747,7 +750,7 @@ impl IVFSQIndexReader { self.search_with_filter(query, k, nprobe, Some(&filter)) } - /// Returns every probed row whose SQ-estimated squared L2 distance is in + /// Returns every probed row whose SQ-estimated internal distance is in /// the half-open band. Results are unsorted, unpadded, and never truncated. /// Even probing every list does not guarantee membership under the original /// vectors' distances: scalar quantization can move a row across either cut. @@ -827,14 +830,17 @@ impl IVFSQIndexReader { } self.ensure_loaded()?; let dimension = self.d; - let (probe_lists, _) = kmeans::find_topk_batch( + let processed = prepare_range_queries(queries, dimension, self.metric)?; + let queries = processed.as_ref(); + let metric = self.metric; + let probe_lists = range_probe_lists( queries, - query_count, &self.quantizer_centroids, - self.nlist, dimension, + self.nlist, nprobe, - ); + metric, + )?; 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() { @@ -871,7 +877,7 @@ impl IVFSQIndexReader { codes, ¢roid, &sq, - MetricType::L2, + metric, selection, &mut scratch, &mut collectors[query_index], @@ -887,6 +893,7 @@ impl IVFSQIndexReader { codes, ¢roid, &sq, + metric, selection, &mut scratch, &mut chunk_collectors, @@ -935,7 +942,7 @@ impl IVFSQIndexReader { &list.codes, &self.quantizer_centroids[list_id * dimension..(list_id + 1) * dimension], self.list_sqs.get(list_id).unwrap_or(&self.sq), - MetricType::L2, + metric, selection, scratch, collector, @@ -1350,6 +1357,7 @@ fn scan_sq_range_chunk( codes: &[u8], centroid: &[f32], sq: &ScalarQuantizer, + metric: MetricType, selection: SqRowSelection<'_>, scratch: &mut SqScanScratch, collectors: &mut [(usize, C)], @@ -1362,7 +1370,7 @@ fn scan_sq_range_chunk( codes, centroid, sq, - MetricType::L2, + metric, selection, scratch, collector, @@ -1421,17 +1429,32 @@ fn scan_sq_rows( return Ok(()); } let cutoff = collector.cutoff(); - sq.distances_to_blocked_codes_with_offset( - query, - codes, - ids.len(), - centroid, - metric, - IVF_SQ_SCAN_BLOCK_SIZE, - cutoff, - &mut scratch.parameters, - &mut scratch.distances, - ); + if C::VALIDATE_COSINE_INPUTS { + sq.distances_to_blocked_codes_with_offset_checked( + query, + codes, + ids.len(), + centroid, + metric, + IVF_SQ_SCAN_BLOCK_SIZE, + cutoff, + true, + &mut scratch.parameters, + &mut scratch.distances, + ); + } else { + sq.distances_to_blocked_codes_with_offset( + query, + codes, + ids.len(), + centroid, + metric, + IVF_SQ_SCAN_BLOCK_SIZE, + cutoff, + &mut scratch.parameters, + &mut scratch.distances, + ); + } let mut collect_row = |row_id, distance: f32| { if distance.is_finite() && distance >= cutoff { collector.note_abandoned(); @@ -1660,9 +1683,18 @@ mod tests { } impl Collector for TrackingCollector<'_> { + const VALIDATE_COSINE_INPUTS: bool = true; + fn cutoff(&self) -> f32 { let worker = rayon::current_thread_index().unwrap(); - self.workers.fetch_or(1 << worker, Ordering::Relaxed); + if self.workers.fetch_or(1 << worker, Ordering::Relaxed) == 0 { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(1); + while self.workers.load(Ordering::Relaxed).count_ones() < 2 + && std::time::Instant::now() < deadline + { + std::thread::sleep(std::time::Duration::from_millis(1)); + } + } self.inner.cutoff() } @@ -1684,18 +1716,22 @@ mod tests { let queries = (0..16) .flat_map(|query_index| vec![query_index as f32 * 0.25; dimension]) .collect::>(); - let band = - DistanceBand::new(Bound::Finite(1.0), Bound::Finite(200.0), MetricType::L2).unwrap(); let pool = rayon::ThreadPoolBuilder::new() .num_threads(4) .build() .unwrap(); let allowed: RoaringTreemap = (0..count as u64).filter(|id| id % 3 == 0).collect(); let masks = sq_filter_masks(&ids, &allowed); - for selection in [ + for (selection, metric) in [ SqRowSelection::Filter(None), SqRowSelection::BlockMasks(&masks), - ] { + ] + .into_iter() + .flat_map(|selection| { + [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] + .map(|metric| (selection, metric)) + }) { + let band = DistanceBand::new(Bound::Finite(1.0), Bound::Finite(200.0), metric).unwrap(); let workers = AtomicU64::new(0); let query_indices = [14, 2, 12, 4, 10, 6, 8, 0]; let mut collectors = query_indices @@ -1717,6 +1753,7 @@ mod tests { &codes, ¢roid, &sq, + metric, selection, &mut SqScanScratch::default(), &mut collectors, @@ -1735,7 +1772,7 @@ mod tests { &codes, ¢roid, &sq, - MetricType::L2, + metric, selection, &mut SqScanScratch::default(), &mut expected, @@ -1779,6 +1816,7 @@ mod tests { &codes, &[0.0], &sq, + MetricType::L2, SqRowSelection::Filter(None), &mut SqScanScratch::default(), &mut [(0, FailingCollector), (1, FailingCollector)], diff --git a/core/src/range.rs b/core/src/range.rs index 4ed9639..4aa3d90 100644 --- a/core/src/range.rs +++ b/core/src/range.rs @@ -32,9 +32,10 @@ //! to be safe because squared distances are non-negative, but that is a //! coincidence of one metric and should not become the representation for three. -use std::io; +use std::{borrow::Cow, io}; -use crate::distance::MetricType; +use crate::distance::{fvec_norm_l2sqr, fvec_normalize, MetricType}; +use crate::kmeans; /// One side's bound on an interval. #[derive(Debug, Clone, Copy, PartialEq)] @@ -135,17 +136,76 @@ impl DistanceBand { } } -/// Only L2 is certified so far. Cosine and inner-product bands return -/// `Unsupported`, which means "we cannot serve this request, please fall back", -/// not "the call has a bug". -pub(crate) fn ensure_certified_metric(metric: MetricType) -> io::Result<()> { +pub(crate) fn prepare_range_queries( + queries: &[f32], + dimension: usize, + metric: MetricType, +) -> io::Result> { + if metric != MetricType::Cosine { + return Ok(Cow::Borrowed(queries)); + } + let mut normalized = queries.to_vec(); + for query in normalized.chunks_exact_mut(dimension) { + let norm = fvec_normalize(query); + if !norm.is_finite() || query.iter().any(|value| !value.is_finite()) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "non-finite normalized query", + )); + } + } + Ok(Cow::Owned(normalized)) +} + +pub(crate) fn range_probe_lists( + queries: &[f32], + centroids: &[f32], + dimension: usize, + nlist: usize, + nprobe: usize, + metric: MetricType, +) -> io::Result>> { if metric == MetricType::L2 { - return Ok(()); + return Ok(kmeans::find_topk_batch( + queries, + queries.len() / dimension, + centroids, + nlist, + dimension, + nprobe, + ) + .0); + } + if centroids.iter().any(|value| !value.is_finite()) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "non-finite IVF centroid", + )); + } + queries + .chunks_exact(dimension) + .map(|query| { + kmeans::find_topk_checked(query, centroids, nlist, dimension, nprobe) + .map(|lists| lists.into_iter().map(|(_, list)| list).collect()) + .map_err(|list| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("non-finite query-centroid distance for list {list}"), + ) + }) + }) + .collect() +} + +pub(crate) fn checked_cosine_norm(vector: &[f32]) -> io::Result { + let norm_squared = fvec_norm_l2sqr(vector); + if vector.iter().any(|value| !value.is_finite()) || !norm_squared.is_finite() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "non-finite cosine vector or norm", + )); } - Err(io::Error::new( - io::ErrorKind::Unsupported, - format!("range search is not certified for metric {metric:?} yet"), - )) + Ok(norm_squared.sqrt()) } /// The comparison operator for one endpoint of a distance predicate. @@ -173,12 +233,41 @@ pub struct DistanceEndpoint { /// What a stored squared-L2 f32 distance displays as in SQL: `sqrtf`, then /// widened to double. -/// -/// **L2 only.** The conversions for the other two metrics are not implemented -/// here; see [`DistanceBand::from_endpoints`]. #[inline] fn l2_public_value(stored_bits: u32) -> f64 { - f64::from(f32::from_bits(stored_bits).sqrt()) + MetricType::L2.public_distance(f32::from_bits(stored_bits)) +} + +fn ordered_value(key: u32) -> f32 { + f32::from_bits(if key & 0x8000_0000 != 0 { + key ^ 0x8000_0000 + } else { + !key + }) +} + +fn first_linear_cut(endpoint: DistanceEndpoint) -> Option { + let admits = |key| { + let value = f64::from(ordered_value(key)); + match endpoint.op { + CutOperator::Ge | CutOperator::Lt => value >= endpoint.value, + CutOperator::Gt | CutOperator::Le => value > endpoint.value, + } + }; + let mut low = 0x0080_0000; + let mut high = 0xff7f_ffff; + if !admits(high) { + return None; + } + while low < high { + let middle = low + (high - low) / 2; + if admits(middle) { + high = middle; + } else { + low = middle + 1; + } + } + Some(ordered_value(low)) } /// The first bit pattern satisfying `l2_public_value(bits) >= endpoint`. @@ -238,36 +327,17 @@ impl DistanceBand { /// "the lower bound is `<`" is meaningless, and guessing the intent would /// only mask a dispatch bug in the caller. /// - /// # L2 only - /// - /// Cosine and inner product return `Unsupported`. Do **not** wave them - /// through the L2 path; both conversions differ: - /// - /// * **cosine**: the internal value `1 - cos` can be slightly below zero - /// because `distance.rs` does not clamp it, and `first_ge`/`first_gt` - /// search only the non-negative f32 bit range, where a negative value - /// simply cannot be found. An ordered key spanning the **whole** finite - /// f32 axis, negatives included, is required to binary search it. - /// * **inner product**: the internal distance is `-inner_product`, so the - /// public value **decreases** as the internal one increases. The endpoint - /// must first be negated, and **the side and its open/closed sense - /// flipped together** (`>= e` becomes `<= -e` on the internal value), or - /// the two bounds end up completely reversed. - /// - /// A metric-certification change must implement both conversions and add the - /// "three metrics × four operators × boundary ULP" tests, rather than merely - /// relaxing [`ensure_certified_metric`]. + /// Endpoints use [`MetricType::public_distance`]. Cosine searches the whole + /// finite f32 axis, including negative roundoff. Inner product negates the + /// endpoint and reverses both its side and comparison. Neither path rounds + /// the f64 endpoint to f32 before deciding membership. Out-of-domain linear + /// cuts become empty or structurally unbounded bands. L2 retains its + /// `Unsupported` result when no representable square-root cut exists. pub fn from_endpoints( lower: Option, upper: Option, metric: MetricType, ) -> io::Result { - // Order matters, and the rule is that **caller bugs outrank capability - // gaps**. A non-finite endpoint and an operator on the wrong side are - // both caller bugs (`InvalidInput`); an uncertified metric is a - // capability gap (`Unsupported`, so the caller should fall back). - // Reporting either bug as `Unsupported` would let an FFI caller - // silently fall back and never see it. for ep in [lower, upper].into_iter().flatten() { if !ep.value.is_finite() { return Err(invalid(format!("endpoint is not finite: {}", ep.value))); @@ -311,10 +381,36 @@ impl DistanceBand { )), }; if metric != MetricType::L2 { - return Err(io::Error::new( - io::ErrorKind::Unsupported, - format!("endpoint derivation is not implemented for metric {metric:?}"), - )); + if matches!((lower, upper), (Some(low), Some(high)) if low.value > high.value) { + return Err(invalid("inverted public band")); + } + let reverse = |endpoint: DistanceEndpoint| DistanceEndpoint { + value: -endpoint.value, + op: match endpoint.op { + CutOperator::Ge => CutOperator::Le, + CutOperator::Gt => CutOperator::Lt, + CutOperator::Le => CutOperator::Ge, + CutOperator::Lt => CutOperator::Gt, + }, + }; + let (lower, upper) = if metric == MetricType::InnerProduct { + (upper.map(reverse), lower.map(reverse)) + } else { + (lower, upper) + }; + let lower_cut = lower.map(first_linear_cut); + let upper_cut = upper.and_then(first_linear_cut); + if lower_cut == Some(None) + || upper_cut == Some(-f32::MAX) + || matches!((lower_cut, upper_cut), (Some(Some(low)), Some(high)) if low > high) + { + return Self::new(Bound::Finite(0.0), Bound::Finite(0.0), metric); + } + return Self::new( + lower_cut.flatten().map_or(Bound::Unbounded, Bound::Finite), + upper_cut.map_or(Bound::Unbounded, Bound::Finite), + metric, + ); } let resolve = |side: Option<(DistanceEndpoint, CutFinder)>| -> io::Result { match side { @@ -362,8 +458,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. Cosine and inner product never abandon partial sums. IVF-PQ + /// and IVF-RQ always evaluate complete estimates, so their counts are zero. pub fn early_abandoned(&self) -> usize { self.early_abandoned } @@ -590,12 +686,7 @@ impl VectorRangeSearchParams { /// equivalent resolution private for the same reason. Callers get /// validation implicitly, by calling a search method. pub(crate) fn validate(&self, nlist: usize) -> io::Result { - // A caller bug outranks a capability gap, so the shape check runs before - // the metric check. Reversing the two would report a zero nprobe on an - // uncertified metric as `Unsupported`, and an FFI caller would silently - // fall back to a scan instead of surfacing the bug. self.validate_shape()?; - ensure_certified_metric(self.band.metric())?; Ok(self.nprobe.min(nlist)) } } @@ -695,11 +786,13 @@ mod tests { } #[test] - fn only_l2_is_certified_in_this_pr() { - assert!(ensure_certified_metric(MetricType::L2).is_ok()); - for metric in [MetricType::Cosine, MetricType::InnerProduct] { - let err = ensure_certified_metric(metric).unwrap_err(); - assert_eq!(err.kind(), std::io::ErrorKind::Unsupported); + fn all_metrics_accept_positive_probe_widths() { + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + let band = DistanceBand::new(Bound::Unbounded, Bound::Unbounded, metric).unwrap(); + assert_eq!( + VectorRangeSearchParams::new(band, 3).validate(2).unwrap(), + 2 + ); } } @@ -721,11 +814,7 @@ mod tests { } #[test] - fn a_wrong_side_operator_outranks_the_metric_check() { - // A caller bug must not be reported as a capability gap. An operator on - // the wrong side is a dispatch bug worth fixing; if it were reported as - // Unsupported on an uncertified metric, an FFI caller would silently - // fall back to a scan and never see it. + fn a_wrong_side_operator_is_invalid_for_every_metric() { for metric in [MetricType::Cosine, MetricType::InnerProduct] { let ep = DistanceEndpoint { value: 0.5, @@ -766,21 +855,19 @@ mod tests { } #[test] - fn endpoint_derivation_is_unsupported_for_uncertified_metrics() { - // ensure_certified_metric alone cannot be the gate: from_endpoints is a - // public function, and silently deriving a band by the L2 rule for - // cosine or inner product would return a wrong answer. + fn linear_endpoint_derivation_preserves_public_comparisons() { for metric in [MetricType::Cosine, MetricType::InnerProduct] { let ep = DistanceEndpoint { value: 0.5, op: CutOperator::Ge, }; - let err = DistanceBand::from_endpoints(Some(ep), None, metric).unwrap_err(); - assert_eq!( - err.kind(), - std::io::ErrorKind::Unsupported, - "metric {metric:?}" - ); + let band = DistanceBand::from_endpoints(Some(ep), None, metric).unwrap(); + for distance in [-1.0, -0.5, 0.0, 0.5, 1.0] { + assert_eq!( + band.admit(distance), + metric.public_distance(distance) >= 0.5 + ); + } } } @@ -1027,9 +1114,7 @@ mod tests { } #[test] - fn a_zero_nprobe_outranks_the_metric_check() { - // Same rule as above, on the params path: nprobe == 0 is a caller bug - // and must stay InvalidInput even when the metric is not certified. + fn a_zero_nprobe_is_invalid_for_every_metric() { for metric in [MetricType::Cosine, MetricType::InnerProduct] { let band = DistanceBand::new(Bound::Finite(0.0), Bound::Finite(1.0), metric).unwrap(); let err = VectorRangeSearchParams::new(band, 0) diff --git a/core/src/sq.rs b/core/src/sq.rs index ffe58be..8baa222 100644 --- a/core/src/sq.rs +++ b/core/src/sq.rs @@ -260,6 +260,24 @@ impl ScalarQuantizer { cutoff: f32, parameters: &mut Vec, distances: &mut Vec, + ) { + self.distances_to_blocked_codes_with_offset_checked( + query, codes, count, offset, metric, block_size, cutoff, false, parameters, distances, + ); + } + + pub(crate) fn distances_to_blocked_codes_with_offset_checked( + &self, + query: &[f32], + codes: &[u8], + count: usize, + offset: &[f32], + metric: MetricType, + block_size: usize, + cutoff: f32, + check_finite: bool, + parameters: &mut Vec, + distances: &mut Vec, ) { debug_assert!(query.len() >= self.d); debug_assert!(offset.len() >= self.d); @@ -318,7 +336,11 @@ impl ScalarQuantizer { MetricType::Cosine => { for (distance, norm) in block_distances.iter_mut().zip(norms.unwrap()) { let denominator = query_norm * norm.sqrt(); - *distance = if denominator > 0.0 { + *distance = if check_finite + && (!distance.is_finite() || !denominator.is_finite()) + { + f32::NAN + } else if denominator > 0.0 { 1.0 - *distance / denominator } else { 1.0 diff --git a/core/tests/range_metrics.rs b/core/tests/range_metrics.rs new file mode 100644 index 0000000..b5abb84 --- /dev/null +++ b/core/tests/range_metrics.rs @@ -0,0 +1,1069 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use paimon_vindex_core::distance::{ + fvec_distance, fvec_inner_product, fvec_l2sqr, fvec_norm_l2sqr, fvec_normalize, MetricType, +}; +use paimon_vindex_core::index::{IndexType, VectorIndexReader, VectorSearchParams}; +use paimon_vindex_core::io::{write_index, IVFPQIndexReader, PosWriter}; +use paimon_vindex_core::ivfflat::IVFFlatIndex; +use paimon_vindex_core::ivfflat_io::{write_ivfflat_index, IVFFlatIndexReader}; +use paimon_vindex_core::ivfpq::IVFPQIndex; +use paimon_vindex_core::ivfrq::IVFRQIndex; +use paimon_vindex_core::ivfrq_io::{write_ivfrq_index, IVFRQIndexReader}; +use paimon_vindex_core::ivfsq::IVFSQIndex; +use paimon_vindex_core::ivfsq_io::{write_ivfsq_index, IVFSQIndexReader}; +use paimon_vindex_core::range::{ + Bound, CutOperator, DistanceBand, DistanceEndpoint, RangeSearchResult, VectorRangeSearchParams, +}; +use paimon_vindex_core::rq::RQRotation; +use paimon_vindex_core::sq::ScalarQuantizer; +use roaring::RoaringTreemap; +use std::collections::BTreeSet; +use std::io::{self, Cursor}; + +const DIMENSION: usize = 16; +const LISTS: usize = 3; +const ROWS: usize = 37; +const METRICS: [MetricType; 3] = [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct]; + +#[derive(Clone, Copy, Debug)] +enum Family { + Flat, + Sq, + Pq(usize, bool, bool), + Rq(usize), +} + +const FAMILIES: [Family; 13] = [ + Family::Flat, + Family::Sq, + Family::Pq(4, false, false), + Family::Pq(4, false, true), + Family::Pq(4, true, false), + Family::Pq(4, true, true), + Family::Pq(8, false, false), + Family::Pq(8, false, true), + Family::Pq(8, true, false), + Family::Pq(8, true, true), + Family::Rq(1), + Family::Rq(4), + Family::Rq(8), +]; + +enum Source { + Flat(IVFFlatIndex), + Sq(IVFSQIndex), + Pq(IVFPQIndex), + Rq(IVFRQIndex), +} + +fn normalized(mut vector: Vec, metric: MetricType) -> Vec { + if metric == MetricType::Cosine { + fvec_normalize(&mut vector); + } + vector +} + +fn fixture(family: Family, metric: MetricType) -> Source { + let centroids = (0..LISTS) + .flat_map(|list| { + normalized( + (0..DIMENSION) + .map(|dimension| if dimension == list { 4.0 } else { 0.0 }) + .collect(), + metric, + ) + }) + .collect::>(); + let vectors = (0..LISTS) + .flat_map(|list| { + (0..ROWS).flat_map(move |row| { + normalized( + (0..DIMENSION) + .map(|dimension| { + if dimension == list { + 4.0 + } else { + ((row * 13 + dimension * 7) % 31) as f32 / 31.0 - 0.5 + } + }) + .collect(), + metric, + ) + }) + }) + .collect::>(); + let ids = (0..LISTS * ROWS) + .map(|row| { + if row % 5 == 0 { + -1000 - row as i64 + } else { + 1000 + row as i64 + } + }) + .collect::>(); + match family { + Family::Flat => { + let mut index = IVFFlatIndex::new(DIMENSION, LISTS, metric); + index.set_quantizer_centroids(centroids); + index.add(&vectors, &ids, ids.len()); + Source::Flat(index) + } + Family::Sq => { + let mut index = IVFSQIndex::new(DIMENSION, LISTS, metric); + index.set_quantizer_centroids(centroids); + index.sq = ScalarQuantizer::with_bounds(DIMENSION, -0.6, 0.6); + index.list_sqs = (0..LISTS).map(|_| index.sq.clone()).collect(); + index.add(&vectors, &ids, ids.len()); + Source::Sq(index) + } + Family::Pq(bits, opq, residual) => { + let mut index = IVFPQIndex::with_nbits(DIMENSION, LISTS, 4, bits, metric, opq); + index.by_residual = residual; + index.set_quantizer_centroids(centroids); + let codebook = (0..4 * index.pq.ksub() * 4) + .map(|position| ((position * 17 + position / 7) % 101) as f32 / 100.0 - 0.5) + .collect(); + index.pq.set_centroids(codebook); + if let Some(rotation) = &mut index.opq { + rotation.rotation = vec![0.0; DIMENSION * DIMENSION]; + for dimension in 0..DIMENSION { + rotation.rotation[dimension * DIMENSION + (dimension + 4) % DIMENSION] = + if dimension % 2 == 0 { 1.0 } else { -1.0 }; + } + rotation.is_trained = true; + } + for list in 0..LISTS { + index.ids[list] = ids[list * ROWS..(list + 1) * ROWS].to_vec(); + index.codes[list] = (0..ROWS * index.pq.code_size()) + .map(|position| (position * 31 + list * 19) as u8) + .collect(); + } + Source::Pq(index) + } + Family::Rq(bits) => { + let mut index = IVFRQIndex::with_bits(DIMENSION, LISTS, bits, metric); + index.set_quantizer_centroids(centroids); + index.add(&vectors, &ids, ids.len()); + Source::Rq(index) + } + } +} + +impl Source { + fn bytes(&self) -> Vec { + let mut bytes = Vec::new(); + let mut writer = PosWriter::new(&mut bytes); + match self { + Self::Flat(index) => write_ivfflat_index(index, &mut writer), + Self::Sq(index) => write_ivfsq_index(index, &mut writer), + Self::Pq(index) => write_index(index, &mut writer), + Self::Rq(index) => write_ivfrq_index(index, &mut writer), + } + .unwrap(); + bytes + } + + fn centroids(&self) -> &[f32] { + match self { + Self::Flat(index) => index.quantizer_centroids(), + Self::Sq(index) => index.quantizer_centroids(), + Self::Pq(index) => index.quantizer_centroids(), + Self::Rq(index) => index.quantizer_centroids(), + } + } + + fn query(&self, query: &[f32], metric: MetricType) -> Vec { + let mut query = normalized(query.to_vec(), metric); + if let Self::Pq(index) = self { + if let Some(opq) = &index.opq { + let mut rotated = vec![0.0; DIMENSION]; + opq.apply(&query, &mut rotated); + query = rotated; + } + } + query + } + + fn lists(&self, query: &[f32], nprobe: usize) -> Vec { + let mut lists = (0..LISTS).collect::>(); + lists.sort_by(|&left, &right| { + fvec_l2sqr( + query, + &self.centroids()[left * DIMENSION..(left + 1) * DIMENSION], + ) + .total_cmp(&fvec_l2sqr( + query, + &self.centroids()[right * DIMENSION..(right + 1) * DIMENSION], + )) + .then(left.cmp(&right)) + }); + lists.truncate(nprobe.min(LISTS)); + lists + } + + fn oracle(&self, raw_query: &[f32], metric: MetricType, nprobe: usize) -> Vec<(i64, f32)> { + let query = self.query(raw_query, metric); + let mut rows = Vec::new(); + for list in self.lists(&query, nprobe) { + let centroid = &self.centroids()[list * DIMENSION..(list + 1) * DIMENSION]; + match self { + Self::Flat(index) => { + for (&id, vector) in index.ids[list] + .iter() + .zip(index.vectors[list].chunks_exact(DIMENSION)) + { + rows.push((id, fvec_distance(&query, vector, metric))); + } + } + Self::Sq(index) => { + let sq = &index.list_sqs[list]; + for (&id, code) in index.ids[list] + .iter() + .zip(index.codes[list].chunks_exact(DIMENSION)) + { + let mut dot = 0.0; + let mut norm: f32 = 0.0; + let mut squared = 0.0; + for dimension in 0..DIMENSION { + let decoded = (centroid[dimension] + sq.mins[dimension]) + + code[dimension] as f32 + * ((sq.maxs[dimension] - sq.mins[dimension]) * (1.0 / 255.0)); + dot += query[dimension] * decoded; + norm += decoded * decoded; + squared += (query[dimension] - decoded).powi(2); + } + let denominator = fvec_norm_l2sqr(&query).sqrt() * norm.sqrt(); + rows.push(( + id, + match metric { + MetricType::L2 => squared, + MetricType::InnerProduct => -dot, + MetricType::Cosine => { + if denominator > 0.0 { + 1.0 - dot / denominator + } else { + 1.0 + } + } + }, + )); + } + } + Self::Pq(index) => { + for (&id, code) in index.ids[list] + .iter() + .zip(index.codes[list].chunks_exact(index.pq.code_size())) + { + let mut decoded = vec![0.0; DIMENSION]; + index.pq.decode(code, &mut decoded); + let mut distance = 0.0; + for sub in 0..index.pq.m() { + let range = index.pq.chunk_range(sub); + let sub_query = query[range.clone()] + .iter() + .zip(¢roid[range.clone()]) + .map(|(&value, ¢er)| { + if index.by_residual && metric != MetricType::InnerProduct { + value - center + } else { + value + } + }) + .collect::>(); + let mut sub_distance = if metric == MetricType::InnerProduct { + -fvec_inner_product(&sub_query, &decoded[range]) + } else { + fvec_l2sqr(&sub_query, &decoded[range]) + }; + if sub == 0 && index.by_residual && metric == MetricType::InnerProduct { + sub_distance -= fvec_inner_product(&query, centroid); + } + distance += sub_distance; + } + if metric == MetricType::Cosine { + distance *= 0.5; + } + rows.push((id, distance)); + } + } + Self::Rq(index) => { + let rotation = + RQRotation::new(DIMENSION, 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 sum = rotated.iter().sum::(); + let coarse = fvec_l2sqr(&query, centroid); + let add = match metric { + MetricType::L2 => coarse, + MetricType::Cosine => 0.5 * coarse, + MetricType::InnerProduct => { + -0.5 * (fvec_norm_l2sqr(&query) + fvec_norm_l2sqr(centroid) - coarse) + } + }; + for (position, &id) in index.ids[list].iter().enumerate() { + let code = &index.codes[list] + [position * index.code_size()..(position + 1) * index.code_size()]; + let subset = |plane: usize, byte: usize| { + (0..8) + .rev() + .filter(|bit| { + code[plane * index.plane_size() + byte] & (1 << bit) != 0 + }) + .fold(0.0, |total, bit| total + rotated[byte * 8 + bit]) + }; + let mut unsigned = (0..index.plane_size()) + .map(|byte| subset(0, byte)) + .sum::() + * (1usize << (index.bits - 1)) as f32; + for plane in 1..index.bits { + for byte in 0..index.plane_size() { + unsigned += (1usize << (index.bits - 1 - plane)) as f32 + * subset(plane, byte); + } + } + let factors = index.factors[list][position].full; + let distance = factors.f_add + + add + + factors.f_rescale + * (unsigned - ((1usize << index.bits) - 1) as f32 * 0.5 * sum); + rows.push((id, distance)); + } + } + } + } + rows + } +} + +fn direct( + bytes: &[u8], + family: Family, + queries: &[f32], + batch: bool, + params: VectorRangeSearchParams, + filter: Option<&[u8]>, +) -> io::Result { + macro_rules! run { + ($reader:ty) => {{ + let mut reader = <$reader>::open(Cursor::new(bytes.to_vec()))?; + match (batch, filter) { + (false, None) => reader.range_search(queries, params), + (false, Some(filter)) => { + reader.range_search_with_roaring_filter(queries, params, filter) + } + (true, None) => { + reader.range_search_batch(queries, queries.len() / DIMENSION, params) + } + (true, Some(filter)) => reader.range_search_batch_with_roaring_filter( + queries, + queries.len() / DIMENSION, + params, + filter, + ), + } + }}; + } + match family { + Family::Flat => run!(IVFFlatIndexReader>>), + Family::Sq => run!(IVFSQIndexReader>>), + Family::Pq(..) => run!(IVFPQIndexReader>>), + Family::Rq(..) => run!(IVFRQIndexReader>>), + } +} + +fn pairs(rows: impl IntoIterator) -> Vec<(i64, u32)> { + let mut rows = rows + .into_iter() + .map(|(id, distance)| (id, distance.to_bits())) + .collect::>(); + rows.sort_unstable(); + rows +} + +fn result_pairs(result: &RangeSearchResult, query: usize) -> Vec<(i64, u32)> { + pairs( + result + .query(query) + .labels + .iter() + .copied() + .zip(result.query(query).distances.iter().copied()), + ) +} + +fn band(metric: MetricType) -> DistanceBand { + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, metric).unwrap() +} + +fn filter_bytes(ids: impl IntoIterator) -> Vec { + let filter = ids + .into_iter() + .map(|id| id as u64) + .collect::(); + let mut bytes = Vec::new(); + filter.serialize_into(&mut bytes).unwrap(); + bytes +} + +#[test] +fn all_families_match_independent_metric_oracles_and_four_entry_points() { + for metric in METRICS { + for family in FAMILIES { + if metric == MetricType::L2 && matches!(family, Family::Sq) { + continue; + } + let source = fixture(family, metric); + let bytes = source.bytes(); + let queries = (0..3 * DIMENSION) + .map(|position| ((position * 13) % 37) as f32 / 17.0 - 0.9) + .collect::>(); + for nprobe in [1, LISTS] { + let oracles = queries + .chunks_exact(DIMENSION) + .map(|query| source.oracle(query, metric, nprobe)) + .collect::>(); + let mut distances = oracles + .iter() + .flatten() + .map(|row| row.1) + .collect::>(); + distances.sort_by(f32::total_cmp); + let lower = distances[distances.len() / 4]; + let upper = distances[distances.len() * 3 / 4]; + let cuts = + DistanceBand::new(Bound::Finite(lower), Bound::Finite(upper), metric).unwrap(); + for band in [band(metric), cuts] { + let params = VectorRangeSearchParams::new(band, nprobe); + let allowed = oracles + .iter() + .flatten() + .filter(|row| row.0 >= 0 && row.0 % 3 != 0) + .map(|row| row.0) + .collect::>(); + let filter = filter_bytes(allowed.iter().copied()); + for selected in [None, Some(filter.as_slice())] { + let mut reader = + VectorIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + assert!(reader.supports_range_search()); + let batch = if let Some(filter) = selected { + reader + .range_search_batch_with_roaring_filter(&queries, 3, params, filter) + } else { + reader.range_search_batch(&queries, 3, params) + } + .unwrap(); + let typed = + direct(&bytes, family, &queries, true, params, selected).unwrap(); + let unique = queries + .chunks_exact(DIMENSION) + .flat_map(|query| source.lists(&source.query(query, metric), nprobe)) + .collect::>(); + assert_eq!(batch.call_stats().list_reads(), unique.len()); + for (query_index, query) in queries.chunks_exact(DIMENSION).enumerate() { + let expected = pairs(oracles[query_index].iter().copied().filter( + |(id, distance)| { + band.admit(*distance) + && (selected.is_none() || allowed.contains(id)) + }, + )); + assert_eq!( + result_pairs(&batch, query_index), + expected, + "{family:?} {metric:?} query={query_index}" + ); + assert_eq!(result_pairs(&typed, query_index), expected); + let single = if let Some(filter) = selected { + reader.range_search_with_roaring_filter(query, params, filter) + } else { + reader.range_search(query, params) + } + .unwrap(); + assert_eq!(result_pairs(&single, 0), expected); + assert_eq!( + result_pairs( + &direct(&bytes, family, query, false, params, selected) + .unwrap(), + 0 + ), + expected + ); + let stats = batch.query(query_index).stats; + assert_eq!(stats.rows_committed(), expected.len()); + assert_eq!(stats.lists_probed(), nprobe); + assert_eq!( + stats.rows_scanned(), + oracles[query_index] + .iter() + .filter(|row| selected.is_none() || allowed.contains(&row.0)) + .count() + ); + if metric != MetricType::L2 + || matches!(family, Family::Rq(_) | Family::Pq(..)) + { + assert_eq!(stats.early_abandoned(), 0); + } + } + } + } + } + } + } +} + +#[test] +fn endpoints_match_public_predicates_at_ulps_zeros_and_extremes() { + let mut values = vec![ + -f32::MAX, + -f32::MIN_POSITIVE, + -f32::from_bits(1), + -0.0, + 0.0, + f32::from_bits(1), + f32::MIN_POSITIVE, + f32::MAX, + ]; + for center in [-2.0f32, -0.5, 0.5, 1.0, 2.0] { + values.extend([ + f32::from_bits(center.to_bits() - 1), + center, + f32::from_bits(center.to_bits() + 1), + ]); + } + for metric in METRICS { + for &stored in &values { + if metric == MetricType::L2 && stored < 0.0 { + continue; + } + let public = match metric { + MetricType::L2 => f64::from(stored.sqrt()), + MetricType::Cosine => f64::from(stored), + MetricType::InnerProduct => -f64::from(stored), + }; + assert_eq!(metric.public_distance(stored), public); + let epsilon = public.abs().max(f64::MIN_POSITIVE) * f64::EPSILON; + for value in [public - epsilon, public, public + epsilon] { + for op in [ + CutOperator::Ge, + CutOperator::Gt, + CutOperator::Le, + CutOperator::Lt, + ] { + let endpoint = Some(DistanceEndpoint { value, op }); + let (lower, upper) = if matches!(op, CutOperator::Ge | CutOperator::Gt) { + (endpoint, None) + } else { + (None, endpoint) + }; + let derived = DistanceBand::from_endpoints(lower, upper, metric); + let band = match derived { + Ok(band) => band, + Err(error) => { + assert_eq!(metric, MetricType::L2); + assert_eq!(error.kind(), io::ErrorKind::Unsupported); + continue; + } + }; + for &candidate in &values { + if metric == MetricType::L2 && candidate < 0.0 { + continue; + } + let displayed = metric.public_distance(candidate); + let expected = match op { + CutOperator::Ge => displayed >= value, + CutOperator::Gt => displayed > value, + CutOperator::Le => displayed <= value, + CutOperator::Lt => displayed < value, + }; + assert_eq!( + band.admit(candidate), + expected, + "{metric:?} {op:?} {value} candidate={candidate}" + ); + } + } + } + } + } + for metric in [MetricType::Cosine, MetricType::InnerProduct] { + for &stored in &values { + let value = metric.public_distance(stored); + let singleton = DistanceBand::from_endpoints( + Some(DistanceEndpoint { + value, + op: CutOperator::Ge, + }), + Some(DistanceEndpoint { + value, + op: CutOperator::Le, + }), + metric, + ) + .unwrap(); + for &candidate in &values { + assert_eq!( + singleton.admit(candidate), + metric.public_distance(candidate) == value, + "{metric:?} singleton={value} candidate={candidate}" + ); + } + } + for (value, op) in [(-f64::MAX, CutOperator::Lt), (f64::MAX, CutOperator::Gt)] { + let endpoint = Some(DistanceEndpoint { value, op }); + let (lower, upper) = if op == CutOperator::Gt { + (endpoint, None) + } else { + (None, endpoint) + }; + assert!(DistanceBand::from_endpoints(lower, upper, metric) + .unwrap() + .is_empty()); + } + for value in [-f64::MAX, 0.0, f64::MAX] { + let empty = DistanceBand::from_endpoints( + Some(DistanceEndpoint { + value, + op: CutOperator::Gt, + }), + Some(DistanceEndpoint { + value, + op: CutOperator::Lt, + }), + metric, + ) + .unwrap(); + assert!(values.iter().all(|&value| !empty.admit(value))); + } + } +} + +#[test] +fn capability_validation_empty_bands_and_topk_are_preserved() { + for metric in METRICS { + assert!(!IndexType::DiskAnn.supports_range_search(metric)); + for family in FAMILIES { + let bytes = fixture(family, metric).bytes(); + let query = vec![0.2; DIMENSION]; + let mut reader = VectorIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let params = VectorRangeSearchParams::new(band(metric), LISTS); + let topk = reader + .search(&query, VectorSearchParams::new(7, LISTS)) + .unwrap(); + reader.range_search(&query, params).unwrap(); + assert_eq!( + reader + .search(&query, VectorSearchParams::new(7, LISTS)) + .unwrap(), + topk + ); + let empty = VectorRangeSearchParams::new( + DistanceBand::new(Bound::Finite(0.5), Bound::Finite(0.5), metric).unwrap(), + LISTS, + ); + assert_eq!( + direct(&bytes, family, &query[..DIMENSION - 1], false, empty, None) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + let wrong_metric = if metric == MetricType::L2 { + MetricType::Cosine + } else { + MetricType::L2 + }; + assert_eq!( + direct( + &bytes, + family, + &query, + false, + VectorRangeSearchParams::new(band(wrong_metric), 1), + None + ) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + for batch in [false, true] { + let result = direct(&bytes, family, &query, batch, empty, None).unwrap(); + assert!(result.labels().is_empty()); + assert_eq!(result.call_stats().list_reads(), 0); + assert_eq!(result.query(0).stats.lists_probed(), 0); + for invalid in [ + vec![f32::NAN; DIMENSION], + vec![f32::INFINITY; DIMENSION], + vec![f32::NEG_INFINITY; DIMENSION], + ] { + assert_eq!( + direct(&bytes, family, &invalid, batch, empty, None) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + } + assert_eq!( + direct(&bytes, family, &query, batch, empty, Some(&[255])) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + assert_eq!( + direct( + &bytes, + family, + &query, + batch, + VectorRangeSearchParams::new(band(metric), 0), + None + ) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidInput + ); + } + } + } +} + +#[test] +fn nonfinite_consumed_data_is_rejected_without_poisoning_filtered_rows() { + for metric in [MetricType::Cosine, MetricType::InnerProduct] { + for family in FAMILIES { + for invalid in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + let mut source = fixture(family, metric); + match &mut source { + Source::Flat(index) => index.vectors[0][0] = invalid, + Source::Sq(index) => { + index.list_sqs[0] = ScalarQuantizer::with_bounds(DIMENSION, 1.0e38, 1.0e38); + } + Source::Pq(index) => { + let mut centroids = index.pq.centroids().to_vec(); + let code = index.codes[0][0] as usize; + centroids[code * index.pq.dsub()] = invalid; + index.pq.set_centroids(centroids); + } + Source::Rq(index) => { + index.factors[0][0].coarse.f_add = invalid; + index.factors[0][0].full.f_add = invalid; + } + } + let query = vec![4.0; DIMENSION]; + let bytes = source.bytes(); + let params = VectorRangeSearchParams::new(band(metric), LISTS); + let empty_filter = filter_bytes([]); + for batch in [false, true] { + let queries = if batch { + query.repeat(2) + } else { + query.clone() + }; + assert_eq!( + direct(&bytes, family, &queries, batch, params, None) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData, + "{family:?} {metric:?}" + ); + let empty = + direct(&bytes, family, &queries, batch, params, Some(&empty_filter)) + .unwrap(); + assert!(empty.labels().is_empty()); + assert_eq!(empty.query(0).stats.rows_scanned(), 0); + } + let good = source + .oracle(&query, metric, LISTS) + .into_iter() + .find(|&(id, distance)| id >= 0 && distance.is_finite()) + .unwrap(); + let only_good = filter_bytes([good.0]); + let result = + direct(&bytes, family, &query, false, params, Some(&only_good)).unwrap(); + assert_eq!(result.labels(), &[good.0]); + } + } + } +} + +#[test] +fn cosine_zero_queries_still_validate_consumed_row_norms() { + let metric = MetricType::Cosine; + let mut flat = IVFFlatIndex::new(DIMENSION, 1, metric); + flat.set_quantizer_centroids(vec![0.0; DIMENSION]); + flat.ids[0] = vec![1, 2]; + flat.vectors[0] = vec![0.0; 2 * DIMENSION]; + flat.vectors[0][DIMENSION..].fill(1.0e20); + + let mut sq = IVFSQIndex::new(DIMENSION, 1, metric); + sq.set_quantizer_centroids(vec![0.0; DIMENSION]); + sq.sq = ScalarQuantizer::with_bounds(DIMENSION, 0.0, 1.0e20); + sq.list_sqs = vec![sq.sq.clone()]; + sq.ids[0] = vec![1, 2]; + sq.codes[0] = vec![0; 2 * DIMENSION]; + sq.codes[0][DIMENSION..].fill(255); + + for (family, source) in [ + (Family::Flat, Source::Flat(flat)), + (Family::Sq, Source::Sq(sq)), + ] { + let bytes = source.bytes(); + let params = VectorRangeSearchParams::new(band(metric), 1); + let bad = filter_bytes([2]); + let good = filter_bytes([1]); + for count in [1, 2] { + let queries = vec![0.0; count * DIMENSION]; + for filter in [None, Some(bad.as_slice())] { + assert_eq!( + direct(&bytes, family, &queries, count > 1, params, filter) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + } + let result = direct(&bytes, family, &queries, count > 1, params, Some(&good)).unwrap(); + for query in 0..count { + assert_eq!(result.query(query).labels, &[1]); + assert_eq!(result.query(query).distances, &[1.0]); + assert_eq!(result.query(query).stats.rows_scanned(), 1); + } + } + } +} + +#[test] +fn nonfinite_coarse_data_and_finite_query_overflow_fail_loud() { + for metric in [MetricType::Cosine, MetricType::InnerProduct] { + for family in FAMILIES { + for invalid in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY, 1.0e20] { + let source = fixture(family, metric); + let mut reader = VectorIndexReader::open(Cursor::new(source.bytes())).unwrap(); + macro_rules! corrupt { + ($reader:expr) => {{ + $reader.ensure_loaded().unwrap(); + $reader.quantizer_centroids[DIMENSION * (LISTS - 1)] = invalid; + }}; + } + match &mut reader { + VectorIndexReader::IvfFlat(reader) => corrupt!(reader), + VectorIndexReader::IvfSq(reader) => corrupt!(reader), + VectorIndexReader::IvfPq(reader) => corrupt!(reader), + VectorIndexReader::IvfRq(reader) => corrupt!(reader), + VectorIndexReader::DiskAnn(_) => unreachable!(), + } + let params = VectorRangeSearchParams::new(band(metric), 1); + assert_eq!( + reader + .range_search(&[0.0; DIMENSION], params) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + } + let source = fixture(family, metric); + assert_eq!( + direct( + &source.bytes(), + family, + &[f32::MAX; DIMENSION], + false, + VectorRangeSearchParams::new(band(metric), 1), + None + ) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + } + } +} + +#[test] +fn flat_zero_vectors_and_inner_product_extrema_use_public_semantics() { + for metric in [MetricType::Cosine, MetricType::InnerProduct] { + let mut index = IVFFlatIndex::new(DIMENSION, 1, metric); + index.set_quantizer_centroids(vec![0.0; DIMENSION]); + let values = if metric == MetricType::Cosine { + vec![0.0, -1.0, 1.0] + } else { + vec![-f32::MAX, -1.0, 0.0, 1.0, f32::MAX] + }; + for (position, &value) in values.iter().enumerate() { + index.ids[0].push(position as i64); + let mut vector = vec![0.0; DIMENSION]; + vector[0] = value; + index.vectors[0].extend(vector); + } + let bytes = Source::Flat(index).bytes(); + for first in [0.0, 1.0] { + let mut query = vec![0.0; DIMENSION]; + query[0] = first; + let params = VectorRangeSearchParams::new(band(metric), 1); + let result = direct(&bytes, Family::Flat, &query, false, params, None).unwrap(); + assert_eq!(result.labels().len(), values.len()); + for (&id, &distance) in result.labels().iter().zip(result.distances()) { + let value = values[id as usize]; + let expected = if metric == MetricType::Cosine { + if value == 0.0 || first == 0.0 { + 1.0 + } else { + 1.0 - value + } + } else { + -(first * value) + }; + assert_eq!(distance, expected); + } + for public in [-f64::from(f32::MAX), 0.0, 1.0, f64::from(f32::MAX)] { + let selected = DistanceBand::from_endpoints( + Some(DistanceEndpoint { + value: public, + op: CutOperator::Ge, + }), + None, + metric, + ) + .unwrap(); + let expected = pairs( + result + .labels() + .iter() + .copied() + .zip(result.distances().iter().copied()) + .filter(|&(_, distance)| metric.public_distance(distance) >= public), + ); + let found = direct( + &bytes, + Family::Flat, + &query, + false, + VectorRangeSearchParams::new(selected, 1), + None, + ) + .unwrap(); + assert_eq!(result_pairs(&found, 0), expected); + } + } + } +} + +#[test] +fn quantized_zero_queries_keep_estimator_membership() { + for family in FAMILIES { + if matches!(family, Family::Flat) { + continue; + } + let source = fixture(family, MetricType::Cosine); + let query = [0.0; DIMENSION]; + let expected = source.oracle(&query, MetricType::Cosine, LISTS); + let result = direct( + &source.bytes(), + family, + &query, + false, + VectorRangeSearchParams::new(band(MetricType::Cosine), LISTS), + None, + ) + .unwrap(); + assert_eq!(result_pairs(&result, 0), pairs(expected)); + } +} + +#[test] +fn pq_streaming_shares_reads_filters_and_propagates_errors() { + use paimon_vindex_core::io::{ReadRequest, SeekRead}; + use paimon_vindex_core::ivfpq::RowIdFilter; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + + struct LimitedReader { + data: Cursor>, + calls: Arc, + } + impl SeekRead for LimitedReader { + fn pread(&mut self, requests: &mut [ReadRequest<'_>]) -> io::Result<()> { + assert!( + requests + .iter() + .map(|request| request.buf.len()) + .sum::() + <= 64 * 1024 * 1024 + ); + self.calls.fetch_add(1, Ordering::Relaxed); + self.data.pread(requests) + } + } + struct CountingFilter { + calls: AtomicUsize, + last: i64, + } + impl RowIdFilter for CountingFilter { + fn contains(&self, id: i64) -> bool { + self.calls.fetch_add(1, Ordering::Relaxed); + id == 0 || id == self.last + } + } + for bits in [4, 8] { + let dimension = 256; + let code_size = dimension * bits / 8; + let count = 64 * 1024 * 1024 / code_size + 1; + let mut index = + IVFPQIndex::with_nbits(dimension, 1, dimension, bits, MetricType::Cosine, false); + index.set_quantizer_centroids(vec![0.0; dimension]); + index + .pq + .set_centroids(vec![0.25; dimension * index.pq.ksub()]); + index.ids[0] = (0..count as i64).collect(); + index.codes[0] = vec![0; count * code_size]; + let bytes = Source::Pq(index).bytes(); + let calls = Arc::new(AtomicUsize::new(0)); + let mut reader = IVFPQIndexReader::open(LimitedReader { + data: Cursor::new(bytes), + calls: calls.clone(), + }) + .unwrap(); + let query = vec![1.0; dimension]; + let filter = CountingFilter { + calls: AtomicUsize::new(0), + last: count as i64 - 1, + }; + let params = VectorRangeSearchParams::new(band(MetricType::Cosine), 1); + let result = reader + .range_search_batch_with_filter(&query.repeat(2), 2, params, Some(&filter)) + .unwrap(); + assert_eq!(result.call_stats().list_reads(), 1); + assert_eq!(filter.calls.load(Ordering::Relaxed), count); + assert!(calls.load(Ordering::Relaxed) > 3); + for query_index in 0..2 { + assert_eq!(result.query(query_index).labels, &[0, count as i64 - 1]); + assert_eq!(result.query(query_index).stats.rows_scanned(), 2); + assert_eq!(result.query(query_index).stats.early_abandoned(), 0); + assert_eq!(result.query(query_index).distances, &[4.5, 4.5]); + } + let mut centroids = reader.pq.centroids().to_vec(); + centroids[0] = f32::NAN; + reader.pq.set_centroids(centroids); + assert_eq!( + reader + .range_search_with_filter(&query, params, Some(&filter)) + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + } +} diff --git a/core/tests/range_search.rs b/core/tests/range_search.rs index 3afcfa0..9936796 100644 --- a/core/tests/range_search.rs +++ b/core/tests/range_search.rs @@ -27,25 +27,1061 @@ 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, metric: MetricType) -> IVFPQIndex { + let dimension = 64; + let rows_per_list = 1027; + let mut index = IVFPQIndex::with_nbits(dimension, 3, 8, bits, metric, 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_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)> { + pq_scalar_oracle(index, query, nprobe) + .into_iter() + .map(|(id, distance, _)| (id, distance)) + .collect() +} + +fn pq_scalar_oracle(index: &IVFPQIndex, query: &[f32], nprobe: usize) -> Vec<(i64, f32, f32)> { + let mut prepared = query.to_vec(); + if index.metric == MetricType::Cosine { + let norm = prepared + .iter() + .map(|value| value * value) + .sum::() + .sqrt(); + if norm > 0.0 { + for value in &mut prepared { + *value /= norm; + } + } + } + let query = prepared.as_slice(); + 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; + let mut absolute_terms = 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]; + let contribution = match index.metric { + MetricType::L2 | MetricType::Cosine => { + (rotated[coordinate] - coarse - centroid).powi(2) + } + MetricType::InnerProduct => -rotated[coordinate] * centroid, + }; + term += contribution; + absolute_terms += contribution.abs(); + } + if sub == 0 && index.by_residual && index.metric == MetricType::InnerProduct { + term -= rotated + .iter() + .zip(&index.quantizer_centroids()[list * index.d..(list + 1) * index.d]) + .map(|(value, coarse)| { + let contribution = value * coarse; + absolute_terms += contribution.abs(); + contribution + }) + .sum::(); + } + distance += term; + } + if index.metric == MetricType::Cosine { + distance *= 0.5; + absolute_terms *= 0.5; + } + rows.push((id, distance, absolute_terms)); + } + } + 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() { + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + check_pq_range_dense_opq(metric); + } +} + +fn check_pq_range_dense_opq(metric: MetricType) { + 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, metric); + 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 = VectorRangeSearchParams::new( + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, metric).unwrap(), + 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_scalar_oracle(&index, query, index.nlist); + assert_eq!(reference[query_index].len(), expected.len()); + for (&(id, distance), &(expected_id, expected_distance, absolute_terms)) in + reference[query_index].iter().zip(&expected) + { + assert_eq!(id, expected_id); + assert!( + (distance - expected_distance).abs() <= 1e-5 * absolute_terms.max(1.0), + "PQ{bits}, {metric:?}, 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), metric).unwrap(), + DistanceBand::new(Bound::Finite(cut), Bound::Unbounded, metric).unwrap(), + DistanceBand::new(Bound::Finite(cut), Bound::Finite(cut.next_up()), metric) + .unwrap(), + DistanceBand::new(Bound::Finite(cut.next_down()), Bound::Finite(cut), metric) + .unwrap(), + ]; + 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_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_validates_metric_and_nprobe_even_for_empty_bands() { + 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 { + 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, + metric: MetricType, +) -> PqStreamingSource { + let dimension = 256; + let mut index = IVFPQIndex::with_nbits(dimension, 1, dimension, bits, metric, 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() { + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + check_pq_range_streaming(metric); + } +} + +fn check_pq_range_streaming(metric: MetricType) { + 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); + let queries = (0..query_count) + .flat_map(|query| vec![(query + 1) as f32 / 64.0; 256]) + .collect::>(); + for bits in [4, 8] { + assert_eq!( + query_count * 256 * (1 << bits) * size_of::() > 8 * 1024 * 1024, + bits == 8 + ); + for corrupt in [false, true] { + let source = pq_streaming_source(bits, corrupt, metric); + 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, + VectorRangeSearchParams::new( + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, metric).unwrap(), + 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 estimate = |codeword: f32| match metric { + MetricType::L2 => 256.0 * (query[0] - codeword).powi(2), + MetricType::Cosine => 128.0 * (0.0625 - codeword).powi(2), + MetricType::InnerProduct => -256.0 * query[0] * codeword, + }; + let mut expected = (0..selected_prefix_rows as i64) + .map(|id| (id, estimate(0.0))) + .collect::>(); + expected.push((rows as i64 - 1, estimate(1.0))); + 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( @@ -629,19 +1665,17 @@ fn rq_range_unified_metric_capability_and_validation_precedence() { 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(), + for result in [ + reader.range_search(&[1.0; 13], params).unwrap(), + reader.range_search_batch(&[1.0; 26], 2, params).unwrap(), reader .range_search_with_roaring_filter(&[1.0; 13], params, &filter) - .unwrap_err(), + .unwrap(), reader .range_search_batch_with_roaring_filter(&[1.0; 26], 2, params, &filter) - .unwrap_err(), + .unwrap(), ] { - assert_eq!(error.kind(), ErrorKind::Unsupported); + assert_eq!(result.query(0).labels.len(), usize::from(!band.is_empty())); } assert_eq!( reader @@ -1376,7 +2410,7 @@ fn ivf_sq_range_streams_oversized_lists_for_single_batch_and_filter() { #[test] fn ivf_sq_range_batch_validates_inputs_before_empty_band_shortcuts() { - use std::io::ErrorKind::{InvalidInput, Unsupported}; + use std::io::ErrorKind::InvalidInput; for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { let mut index = build_sq_index(33, 35, 2); @@ -1417,12 +2451,8 @@ fn ivf_sq_range_batch_validates_inputs_before_empty_band_shortcuts() { .kind(), InvalidInput ); - let result = reader.range_search_batch(&query, 1, params); - if metric == MetricType::L2 { - assert!(result.unwrap().labels().is_empty()); - } else { - assert_eq!(result.unwrap_err().kind(), Unsupported); - } + let result = reader.range_search_batch(&query, 1, params).unwrap(); + assert!(result.labels().is_empty()); } } @@ -1534,15 +2564,8 @@ fn a_band_whose_metric_disagrees_with_the_index_is_rejected() { std::io::ErrorKind::InvalidInput, "index {index_metric:?} / band {band_metric:?} must report InvalidInput" ); - } else if band_metric != MetricType::L2 { - // Matching but not yet certified. - assert_eq!( - outcome.unwrap_err().kind(), - std::io::ErrorKind::Unsupported, - "index {index_metric:?} / band {band_metric:?}" - ); } else { - assert!(outcome.is_ok(), "L2/L2 must succeed"); + assert!(outcome.is_ok(), "matching metrics must succeed"); } } } @@ -2444,11 +3467,10 @@ fn a_zero_nprobe_outranks_the_metric_capability_gap() { "nprobe == 0 must be reported before the uncertified-metric gap" ); - // With a valid nprobe the metric gap is what is left to report. - let err = cosine_reader + let result = cosine_reader .range_search(&[0.0; 8], VectorRangeSearchParams::new(cosine_band, 4)) - .unwrap_err(); - assert_eq!(err.kind(), std::io::ErrorKind::Unsupported); + .unwrap(); + assert!(result.labels().is_empty()); } /// Builds a single-list index whose payload exceeds the 64 MiB threshold that diff --git a/docs/api.html b/docs/api.html index d1daf06..56bd9f9 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-SQ, IVF-PQ and IVF-RQ with L2, cosine and inner product. FLAT tests exact distances; SQ, PQ and RQ test quantized estimates. Query capability with IndexType::supports_range_search(metric) or reader.supports_range_search(). DiskANN reports Unsupported; exact, complete membership requires full-probe FLAT or an exhaustive scan, not top-K followed by filtering. See Range search for public endpoint conversion and the full contract. Range bindings are not included.

diff --git a/docs/index.html b/docs/index.html index 074193d..860e274 100644 --- a/docs/index.html +++ b/docs/index.html @@ -123,10 +123,10 @@

Distance range search support

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-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-FLATSupported, l2 / cosine / ipExact 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 / cosine / ipOne-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 / cosine / ipSQ estimateSingle/batch, with or without a Roaring filter; even full probing cannot guarantee membership under the original vectors' distances
IVF-PQSupported, l2 / cosine / ipFloating-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..971948e 100644 --- a/docs/ivf-pq.html +++ b/docs/ivf-pq.html @@ -110,6 +110,8 @@

Tuning order

Implementation boundaries

+

Rust readers support range search for L2, cosine and inner product, with single/batch and Roaring-filter variants. Both 4-bit and 8-bit files, residual encoding and OPQ use complete f32 ADC estimates: squared distance for L2, half squared distance after query normalization for cosine, and negative dot product for IP. The cosine estimate is a unit-vector surrogate, not exact cosine of a re-normalized reconstruction. Range search does not use top-K's quantized lookup tables or result cap; its estimated membership is not exact membership over original vectors, even at full probe. Existing top-K scores and behavior are unchanged.

+

Range lookup tables are allocated lazily within an 8 MiB cache cap; larger batches reuse per-worker scratch. Shared lists are read once, oversized lists stream in bounded chunks, and residual tables are reused within each list. Parallel query scans preserve membership across list sizes, batch sizes, worker counts and optimize_for_search.

  • The public options API has no pq.nbits; unified construction creates 8-bit PQ.
  • L2 uses residual PQ. Inner product and cosine currently use by_residual=false.
  • The OPQ rotation, IVF centroids, and PQ codebooks are all stored in the file; Readers need no external model.
  • Training samples should represent production data. Distribution drift can degrade both IVF assignment and PQ codebooks.
diff --git a/docs/ivf-rq.html b/docs/ivf-rq.html index f1e233c..a4a3fd3 100644 --- a/docs/ivf-rq.html +++ b/docs/ivf-rq.html @@ -63,7 +63,7 @@

Usage

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.

+

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. L2, cosine and inner product are supported with fixed positive nprobe. Cosine normalizes queries before probe selection and rotation; each metric uses its own query terms and stored factors. Returned scores remain internal distances, not public inner-product similarities. 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.

diff --git a/docs/ivf-sq.html b/docs/ivf-sq.html index 6e3f3ff..f2366f1 100644 --- a/docs/ivf-sq.html +++ b/docs/ivf-sq.html @@ -24,7 +24,7 @@

Position

Build and search

  1. Train the IVF coarse centroids and assign training vectors to lists.
  2. Compute per-dimension residual extrema using partition-local reductions, then pool them across the training sample. This avoids clipping unseen vectors to the narrow or constant bounds of sparsely sampled partitions.
  3. Encode every residual coordinate to an unsigned byte. New indexes use pooled residual bounds for every list; existing v1 files retain their recorded per-list bounds.
  4. At query time, select nprobe lists, reuse cached partitions, and load missing sorted row IDs and codes in bounded multi-range batches. Scan the codes with the metric-specific kernel and retain the top K. Blocked L2 scans use SIMD and a conservative partial-distance cutoff.

Cosine input is normalized through the shared metric preprocessing path. Filters are checked while scanning, so excluded rows do not enter the top-K heap.

-

The Rust reader also supports L2 distance range search, for single and batch queries with or without a Roaring filter. It reuses this blocked SQ scanner and decides membership on its estimated squared distance, not on the original vector. Results have no top-K cap; quantization can move rows across either band boundary even when every list is probed. Range search retains shared batch reads, the partition cache, and bounded streaming of oversized lists.

+

The Rust reader also supports L2, cosine and inner-product range search, for single and batch queries with or without a Roaring filter. It reuses this blocked SQ scanner and decides membership on its estimated internal distance, not on the original vector. Only L2 uses partial-distance pruning; cosine/IP evaluate complete estimates and report zero early abandonment. Results have no top-K cap; quantization can move rows across either band boundary even when every list is probed. Range search retains shared batch reads, the partition cache, and bounded streaming of oversized lists.

Configuration

diff --git a/docs/range-search.html b/docs/range-search.html index 440bd83..9adb219 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-SQ, IVF-PQ and IVF-RQ 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 / SQ / PQ / RQL2 / cosine / inner product · 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 (SQ / PQ / 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 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.
@@ -26,17 +26,18 @@

Semantic contract

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
+

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

+
MetricInternal valuePublic predicate value
l2Squared Euclidean distance or estimatef32 square root, then widened to f64
cosine1 - cos or family estimate; no clampingInternal value widened to f64
inner_product-inner_product or family estimateNegated internal value widened to f64
+

MetricType::public_distance defines this conversion. Returned distances and DistanceBand::new always use internal units; DistanceBand::from_endpoints takes public f64 endpoints. Cosine queries are normalized before both probe selection and scanning; zero queries remain zero. FLAT and SQ return cosine distance 1 when either vector has zero norm. PQ and RQ retain their unit-vector estimators even for zero queries, rather than claiming exact cosine membership.

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 cutUnder L2, IVF-FLAT can abandon a row when its partially accumulated squared distance passes the upper cut, using the same accumulation for pruning and collection. IVF-SQ similarly prunes blocked squared estimates at the exclusive upper cut. Cosine and inner product never use partial-sum pruning. PQ and RQ always evaluate complete f32 estimates, without top-K's FastScan or coarse-bound pruning. 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-SQ, IVF-PQ or IVF-RQ. 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.
@@ -44,8 +45,8 @@

Endpoints and cuts

A predicate such as distance(v, query) < 0.5 compares against the value the engine displays, which for an L2 index is sqrt of the stored squared f32 distance, widened to double. Deriving a cut from that endpoint therefore belongs in this library rather than in the caller.

Why not just square the endpointendpoint * endpoint can land one or two ULP away from the correct cut, because it does not absorb the rounding that the displayed sqrt introduced. Some endpoints square exactly, so the discrepancy is data-dependent rather than universal, which is what makes it easy to miss. One ULP is enough to move a row sitting exactly on a bucket boundary into the neighbouring bucket. Cut derivation instead binary searches the f32 bit patterns for the first value that the predicate admits, which absorbs that rounding exactly.

Pass the already-folded literal from the right-hand side of the predicate as-is. No squaring, no square root, and no binary search on the caller's part. Each endpoint carries its comparison operator, and the operator must match the side it is on:

-
SideAccepted operatorsMeaning of the derived cut
lowerGe, GtThe first value that is admitted, since the lower end is closed
upperLe, LtThe first value that is excluded, since the upper end is open
-

A mismatch, such as a Lt operator on the lower side, is rejected rather than reinterpreted: "the lower bound is <" has no meaning, and guessing the intent would hide a dispatch bug in the caller. An endpoint too large to be represented as a cut reports that the predicate cannot be expressed as a band, which means it should not be pushed down.

+
Public predicate sideAccepted operatorsPredicate
lowerGe, GtPublic value is at least, or strictly greater than, the endpoint
upperLe, LtPublic value is at most, or strictly less than, the endpoint
+

A mismatch, such as a Lt operator on the lower side, is rejected rather than reinterpreted. Cosine searches the entire finite f32 axis, including negative values and signed zeros. Inner product negates endpoints and reverses their sides and comparisons: public ip >= e becomes internal distance <= -e. Endpoints are not rounded to f32 first. Out-of-domain cosine/IP cuts become empty or structurally unbounded bands, preserving inclusive membership at f32::MAX. L2 retains Unsupported when no representable square-root cut exists.

Displayed values have plateausSeveral adjacent squared f32 values round to the same displayed value. So > and >= can derive cuts several bit patterns apart when the endpoint sits exactly on such a plateau, and can derive the same cut when the endpoint falls between two displayed values. What holds in every case is that each cut is the smallest one satisfying its own predicate.
@@ -78,26 +79,40 @@

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)?;
+
Rust · cosine distance and inner-product similarity
let near_cosine = DistanceBand::from_endpoints(
+    None,
+    Some(DistanceEndpoint { value: 0.2, op: CutOperator::Le }),
+    MetricType::Cosine,
+)?;
+let high_similarity = DistanceBand::from_endpoints(
+    Some(DistanceEndpoint { value: 0.8, op: CutOperator::Ge }),
+    None,
+    MetricType::InnerProduct,
+)?;
+

The inner-product predicate is a lower bound on public similarity, not on the returned negative-dot score. Endpoint conversion reverses the internal cut automatically.

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.

+

All four families expose range_search, range_search_batch, and their _with_roaring_filter variants through both typed readers and VectorIndexReader. The filter is an allow-list and does not widen the fixed nprobe. Roaring filters admit only non-negative row IDs; a direct RowIdFilter may admit signed IDs. 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.

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.

+

Range search supports IVF-FLAT, IVF-SQ, IVF-PQ and IVF-RQ under L2, cosine and inner product, with fixed positive probe widths and Rust entry points. Query capability through IndexType::supports_range_search(metric) or reader.supports_range_search(). DiskANN remains unsupported; no storage-format, C/JNI binding or top-K behavior changes 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.
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.
+
IVF-PQ uses complete floating-point ADC estimates.Both packed 4-bit and 8-bit codes, residual encoding and OPQ are supported. L2 sums direct squared subvector distances. Cosine uses half the ADC squared distance after query normalization, a unit-vector surrogate rather than exact cosine of a re-normalized reconstruction. Inner product uses negative estimated dot product. The range path does not reuse top-K's quantized FastScan tables, expanded L2 tables or cosine score scale, so scores need not be bit-identical to top-K. It never truncates or reranks raw vectors. Shared lists are read once, oversized lists stream in bounded chunks, and the allow-list is evaluated once per list row across the batch.

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.

-

Performance and memory: IVF-SQ keeps the existing SIMD block layout and uses a finite upper cut to abandon a block once all partial squared distances reach that exclusive cut. A lower cut alone cannot prune a partial sum. Batch queries read each unique list once, reuse cached partitions, and keep query-owned collectors instead of materializing a list-by-query result matrix. Large single queries can scan lists in parallel and merge once per list, not per row. Oversized lists stream in bounded chunks, with reusable scan scratch. Output memory still grows with all admitted rows; there is no result cap.

+

Performance and memory: Under L2, IVF-SQ keeps the existing SIMD block layout and uses a finite upper cut to abandon a block once all partial squared distances reach that exclusive cut. A lower cut alone cannot prune a partial sum; cosine and IP do not prune partial sums. Batch queries read each unique list once, reuse cached partitions, and keep query-owned collectors instead of materializing a list-by-query result matrix. Large single queries can scan lists in parallel and merge once per list, not per row. Oversized lists stream in bounded chunks, with reusable scan scratch. Output memory still grows with all admitted rows; there is no result cap.

+

IVF-PQ lazily caches query lookup tables within an 8 MiB cap, falling back to reusable worker scratch beyond that budget. Residual tables are reused across chunks of one list and refreshed when the list changes. Large batches scan queries in parallel. Membership is independent of list size, batch size, worker count and optimize_for_search.

Filtered SQ batches evaluate the allow-list once per list row and share compact, query-local block masks (one bit per row) across queries, including streamed chunks. These masks never enter the partition cache. An entirely excluded list or chunk skips distance evaluation; partially selected blocks retain the same SIMD arithmetic as unfiltered search.

-

call_stats().list_reads() excludes empty lists and IVF-SQ cache hits and counts a streamed list once. rows_scanned() counts allow-listed rows reaching collection or cutoff rejection; blocked arithmetic can also evaluate excluded lanes. For IVF-SQ, early_abandoned() includes estimates equal to or above the upper cut and is diagnostic, not a count of operations saved.

+

call_stats().list_reads() excludes empty lists and IVF-SQ cache hits and counts a streamed list once. rows_scanned() counts allow-listed rows reaching collection or cutoff rejection; blocked arithmetic can also evaluate excluded lanes. Under L2, SQ's early_abandoned() includes estimates equal to or above the upper cut. It is zero for cosine/IP and for every PQ/RQ range search.

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, non-finite cut, or negative squared-L2 cutInvalid input
Operator on the wrong side or mismatched index/band metricsInvalid input
Zero nprobe, wrong query dimensions, non-finite query values, or malformed Roaring filterInvalid input
DiskANNUnsupported
L2 endpoint with no representable square-root cutUnsupported
Non-finite consumed distance, factor, vector or cosine norm; non-finite normalization or rotationInvalid data
+

For cosine/IP every family validates all centroids and direct query-centroid distances before choosing lists, including unselected lists. Non-finite cosine vectors and norms cannot silently become distance 1 via zero-norm handling. Consumed PQ entries and RQ factors must produce finite estimates; unused codebook entries and filtered-out row estimates do not poison valid results. SQ bounds retain their existing metadata validation. No error returns a partial result.

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.

@@ -111,6 +126,6 @@

Why not DiskANN

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