From 8025e79dddaaaae316397d1ed9e8e574dcf95c59 Mon Sep 17 00:00:00 2001 From: Junrui Lee Date: Wed, 16 Sep 2026 14:32:22 +0800 Subject: [PATCH 1/2] feat: add IVF-SQ L2 distance range search --- core/README.md | 32 +- core/src/collect.rs | 13 +- core/src/index.rs | 24 +- core/src/ivfsq_io.rs | 494 +++++++++++++++++++++++++++--- core/src/range.rs | 11 +- core/src/topk.rs | 17 ++ core/tests/range_search.rs | 597 ++++++++++++++++++++++++++++++++++++- docs/api.html | 2 +- docs/index.html | 2 +- docs/ivf-sq.html | 1 + docs/range-search.html | 19 +- 11 files changed, 1144 insertions(+), 68 deletions(-) diff --git a/core/README.md b/core/README.md index 0ab6ecd..3a22afb 100644 --- a/core/README.md +++ b/core/README.md @@ -22,14 +22,15 @@ `paimon-vindex-core` contains the Rust implementations and seek-based readers for IVF-FLAT, IVF-SQ, IVF-PQ, IVF-RQ, and DiskANN. -The Rust reader supports distance range search for IVF-FLAT and IVF-RQ with +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. Both families support single and batch queries, +`RangeSearchResult` buffers. All three families support single and batch queries, with or without a serialized Roaring allow-list, and a fixed positive `nprobe`. IVF-FLAT tests exact distances; IVF-RQ tests its one-bit or full multi-bit -estimated distances. Results are uncapped and unordered. Probing every list -removes the IVF coverage gap, but not IVF-RQ's quantization error. The range -path does not change top-K search or the v1 storage format. +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 +search or the v1 storage format. See the [range search guide](../docs/range-search.html) for membership, validation, filtering, and statistics. C/JNI range bindings are not included. @@ -41,6 +42,27 @@ 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 + +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` +CSR buffers without sorting, padding, or a top-K cap. + +IVF-SQ uses the same blocked SIMD estimator as top-K, including the stored +per-list residual bounds. Raising `nprobe` visits more lists but does not remove +quantization error: even at `nprobe == nlist`, membership can differ from the +original vectors' distances at either boundary. Prefer IVF-FLAT when exact +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. +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. + The crate ships its [normative v1 storage-format specification](STORAGE_FORMAT.md) and byte-exact fixtures. Project documentation, language bindings, and contribution guidance live in the diff --git a/core/src/collect.rs b/core/src/collect.rs index 5bd57ab..1146f12 100644 --- a/core/src/collect.rs +++ b/core/src/collect.rs @@ -26,8 +26,8 @@ //! `scan_codes_range` taking a `RangeQueryResult&`; one collector serves both //! here, so there is a single kernel rather than a pair to keep in step. //! -//! `ivfflat_io::ReaderTopKHeap` and [`RangeCollector`] are the two -//! implementations. +//! `ivfflat_io::ReaderTopKHeap`, `topk::TopKHeap`, and [`RangeCollector`] +//! implement the collection policies. use std::io; @@ -64,7 +64,8 @@ pub(crate) trait Collector { fn cutoff(&self) -> f32; /// Delivers one row, with the value the family's scan computed for it. For - /// IVF-Flat that value is an exact distance; for IVF-RQ it is an estimate. + /// IVF-Flat that value is an exact distance; for IVF-RQ and IVF-SQ it is an + /// estimate. /// /// Fallible because a collector may own a resource the scan cannot see: the /// oversized-list path streams chunks through a callback, and without a @@ -130,6 +131,12 @@ impl RangeCollector { pub(crate) fn into_rows(self) -> Vec<(i64, f32)> { self.rows } + + pub(crate) fn merge(&mut self, mut other: Self) { + self.scanned += other.scanned; + self.early_abandoned += other.early_abandoned; + self.rows.append(&mut other.rows); + } } impl Collector for RangeCollector { diff --git a/core/src/index.rs b/core/src/index.rs index 0f23c09..8baa879 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -1866,9 +1866,9 @@ impl VectorIndexReader { } /// Distance range search. For the contract see - /// [`IVFFlatIndexReader::range_search`] and - /// [`IVFRQIndexReader::range_search`]. IVF-RQ membership uses estimated - /// distances rather than distances to the original vectors. + /// [`IVFFlatIndexReader::range_search`] (exact distances), + /// [`IVFRQIndexReader::range_search`] (RQ estimates), and + /// [`IVFSQIndexReader::range_search`] (SQ estimates). Only L2 is supported. /// /// The empty-band short-circuit lives **inside each family's reader**, so a /// family that cannot do range search at all still fails loud for every @@ -1883,15 +1883,16 @@ impl VectorIndexReader { match self { Self::IvfFlat(reader) => reader.range_search(query, params), Self::IvfRq(reader) => reader.range_search(query, params), - Self::IvfSq(_) => Err(range_unsupported("ivf_sq")), + Self::IvfSq(reader) => reader.range_search(query, params), Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), } } /// Range search restricted to a serialized Roaring allow-list. For the - /// contract see [`IVFFlatIndexReader::range_search_with_roaring_filter`] - /// and [`IVFRQIndexReader::range_search`]. + /// contract see [`IVFFlatIndexReader::range_search_with_roaring_filter`], + /// [`IVFRQIndexReader::range_search`], and + /// [`IVFSQIndexReader::range_search_with_roaring_filter`]. pub fn range_search_with_roaring_filter( &mut self, query: &[f32], @@ -1907,15 +1908,14 @@ impl VectorIndexReader { match self { 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(_) => Err(range_unsupported("ivf_sq")), + Self::IvfSq(reader) => reader.range_search_with_filter(query, params, Some(&filter)), Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), } } /// Batched distance range search. For the contract see - /// [`IVFFlatIndexReader::range_search`] and - /// [`IVFRQIndexReader::range_search`]. + /// [`Self::range_search`]. pub fn range_search_batch( &mut self, queries: &[f32], @@ -1927,7 +1927,7 @@ impl VectorIndexReader { match self { Self::IvfFlat(reader) => reader.range_search_batch(queries, query_count, params), Self::IvfRq(reader) => reader.range_search_batch(queries, query_count, params), - Self::IvfSq(_) => Err(range_unsupported("ivf_sq")), + Self::IvfSq(reader) => reader.range_search_batch(queries, query_count, params), Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), } @@ -1951,7 +1951,9 @@ impl VectorIndexReader { Self::IvfRq(reader) => { reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter)) } - Self::IvfSq(_) => Err(range_unsupported("ivf_sq")), + Self::IvfSq(reader) => { + reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter)) + } Self::IvfPq(_) => Err(range_unsupported("ivf_pq")), Self::DiskAnn(_) => Err(range_unsupported("diskann")), } diff --git a/core/src/ivfsq_io.rs b/core/src/ivfsq_io.rs index 5daf168..9a2b7c8 100644 --- a/core/src/ivfsq_io.rs +++ b/core/src/ivfsq_io.rs @@ -17,7 +17,9 @@ //! Stable v1 storage and positional-I/O search for IVF-SQ8. +use crate::collect::{Collector, RangeCollector}; use crate::distance::{preprocess_vectors, MetricType}; +use crate::index::validate_queries; use crate::index_io_util::{ bounded_ivf_payload_batch_end, bounded_ivf_stream_chunk_rows, bytes_to_f32_vec, checked_list_bytes, checked_list_offset, checked_section_size, decode_delta_varint_ids, @@ -30,6 +32,7 @@ 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::read_options::VectorIndexReaderOptions; use crate::sq::ScalarQuantizer; use crate::topk::TopKHeap; @@ -37,7 +40,7 @@ use rayon::prelude::*; use std::collections::VecDeque; use std::io; use std::mem::size_of; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; pub const IVF_SQ_MAGIC: u32 = 0x49565351; // "IVSQ" pub const IVF_SQ_VERSION: u32 = 1; @@ -478,10 +481,18 @@ impl IVFSQIndexReader { } fn read_scan_lists(&mut self, list_ids: &[usize]) -> io::Result>> { + self.read_scan_lists_with_count(list_ids) + .map(|(lists, _)| lists) + } + + fn read_scan_lists_with_count( + &mut self, + list_ids: &[usize], + ) -> io::Result<(Vec>, usize)> { if self.list_cache.is_none() { - return self - .read_inverted_lists(list_ids) - .map(|lists| lists.into_iter().map(Arc::new).collect()); + let lists = self.read_inverted_lists(list_ids)?; + let reads = lists.iter().filter(|list| !list.ids.is_empty()).count(); + return Ok((lists.into_iter().map(Arc::new).collect(), reads)); } let mut results = vec![None; list_ids.len()]; let mut misses = Vec::new(); @@ -509,9 +520,11 @@ impl IVFSQIndexReader { } misses.push((position, list_id)); } + let mut reads = 0; if !misses.is_empty() { let missing_ids = misses.iter().map(|&(_, id)| id).collect::>(); let loaded = self.read_inverted_lists(&missing_ids)?; + reads = loaded.iter().filter(|list| !list.ids.is_empty()).count(); for ((position, list_id), list) in misses.into_iter().zip(loaded) { let list = Arc::new(list); self.list_cache.as_mut().unwrap().insert(CachedSqList { @@ -524,7 +537,7 @@ impl IVFSQIndexReader { results[position] = Some(list); } } - Ok(results.into_iter().map(Option::unwrap).collect()) + Ok((results.into_iter().map(Option::unwrap).collect(), reads)) } fn batch_read_end(&self, list_ids: &[usize]) -> io::Result { @@ -556,7 +569,7 @@ impl IVFSQIndexReader { fn for_each_streamed_list_chunk( &mut self, list_id: usize, - mut consume: impl FnMut(&[i64], &[u8]), + mut consume: impl FnMut(&[i64], &[u8]) -> io::Result<()>, ) -> io::Result<()> { self.ensure_loaded()?; let count = self.list_counts[list_id] as usize; @@ -600,7 +613,7 @@ impl IVFSQIndexReader { self.reader .pread(&mut [ReadRequest::new(chunk_offset, &mut codes)])?; let row_end = row_start + chunk_rows; - consume(&ids[row_start..row_end], &codes); + consume(&ids[row_start..row_end], &codes)?; row_start = row_end; } Ok(()) @@ -651,10 +664,10 @@ impl IVFSQIndexReader { ¢roid, &sq, metric, - filter, + SqRowSelection::Filter(filter), &mut scratch, &mut heap, - ); + ) })?; batch_start += 1; continue; @@ -677,7 +690,7 @@ impl IVFSQIndexReader { filter, &mut SqScanScratch::default(), &mut heap, - ); + )?; let cutoff = heap.distance_limit(); let per_list_results = lists[1..] .par_iter() @@ -693,10 +706,10 @@ impl IVFSQIndexReader { filter, scratch, &mut local_heap, - ); - local_heap.into_sorted() + )?; + Ok(local_heap.into_sorted()) }) - .collect::>(); + .collect::>>()?; for results in per_list_results { for (distance, row_id) in results { heap.push(distance, row_id); @@ -715,7 +728,7 @@ impl IVFSQIndexReader { filter, &mut scratch, &mut heap, - ); + )?; } } batch_start = batch_end; @@ -733,6 +746,228 @@ impl IVFSQIndexReader { let filter = decode_roaring_filter(roaring_filter_bytes)?; self.search_with_filter(query, k, nprobe, Some(&filter)) } + + /// Returns every probed row whose SQ-estimated squared L2 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. + 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) + } + + /// Restricts membership to the serialized Roaring allow-list. Malformed + /// filters are rejected even for an empty band. + 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 SQ-estimate range search; shared probed lists are read once. + 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) + } + + 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); + let band = params.band(); + if band.is_empty() { + return Ok(builder.build()); + } + self.ensure_loaded()?; + let dimension = self.d; + let (probe_lists, _) = kmeans::find_topk_batch( + queries, + query_count, + &self.quantizer_centroids, + self.nlist, + dimension, + nprobe, + ); + let mut list_to_queries = vec![Vec::new(); self.nlist]; + let mut unique_lists = Vec::new(); + for (query_index, lists) in probe_lists.iter().enumerate() { + builder.record_lists_probed(query_index, lists.len()); + for &list_id in lists { + if list_to_queries[list_id].is_empty() { + unique_lists.push(list_id); + } + list_to_queries[list_id].push(query_index); + } + } + let mut collectors = (0..query_count) + .map(|_| RangeCollector::new(band)) + .collect::>(); + let mut scratch = SqScanScratch::default(); + let mut batch_start = 0; + while batch_start < unique_lists.len() { + let first_list = unique_lists[batch_start]; + if ivf_payload_is_oversized(self.list_payload_len(first_list)?) { + let centroid = self.quantizer_centroids + [first_list * dimension..(first_list + 1) * dimension] + .to_vec(); + let sq = self.list_sqs.get(first_list).unwrap_or(&self.sq).clone(); + builder.record_list_read(); + self.for_each_streamed_list_chunk(first_list, |ids, codes| { + let masks = filter.map(|filter| sq_filter_masks(ids, filter)); + let selection = SqRowSelection::from_masks(masks.as_deref()); + for &query_index in &list_to_queries[first_list] { + scan_sq_rows( + &queries[query_index * dimension..(query_index + 1) * dimension], + ids, + codes, + ¢roid, + &sq, + MetricType::L2, + selection, + &mut scratch, + &mut collectors[query_index], + )?; + } + Ok(()) + })?; + batch_start += 1; + continue; + } + let count = self.batch_read_end(&unique_lists[batch_start..])?.max(1); + let batch_end = (batch_start + count).min(unique_lists.len()); + let (lists, reads) = + self.read_scan_lists_with_count(&unique_lists[batch_start..batch_end])?; + for _ in 0..reads { + builder.record_list_read(); + } + let masks = filter.map(|filter| { + lists + .iter() + .map(|list| sq_filter_masks(&list.ids, filter)) + .collect::>() + }); + let candidates = lists.iter().fold(0usize, |total, list| { + total.saturating_add( + list.ids + .len() + .saturating_mul(list_to_queries[list.list_id].len()), + ) + }); + let scan_one = |query_index: usize, + position: usize, + scratch: &mut SqScanScratch, + collector: &mut RangeCollector| { + let list = &lists[position]; + let list_id = list.list_id; + let selection = SqRowSelection::from_masks( + masks.as_ref().map(|masks| masks[position].as_slice()), + ); + scan_sq_rows( + &queries[query_index * dimension..(query_index + 1) * dimension], + &list.ids, + &list.codes, + &self.quantizer_centroids[list_id * dimension..(list_id + 1) * dimension], + self.list_sqs.get(list_id).unwrap_or(&self.sq), + MetricType::L2, + selection, + scratch, + collector, + ) + }; + if query_count == 1 && lists.len() > 1 && candidates >= PARALLEL_SQ_SCAN_MIN_CANDIDATES + { + let output = Mutex::new(&mut collectors[0]); + lists.par_iter().enumerate().try_for_each_init( + SqScanScratch::default, + |scratch, (position, _)| { + let mut collector = RangeCollector::new(band); + scan_one(0, position, scratch, &mut collector)?; + output.lock().expect("range output lock").merge(collector); + Ok::<(), io::Error>(()) + }, + )?; + } else { + let mut positions = vec![None; self.nlist]; + for (position, list) in lists.iter().enumerate() { + positions[list.list_id] = Some(position); + } + let scan_query = + |scratch: &mut SqScanScratch, + (query_index, collector): (usize, &mut RangeCollector)| { + for &list_id in &probe_lists[query_index] { + if let Some(position) = positions[list_id] { + scan_one(query_index, position, scratch, collector)?; + } + } + Ok::<(), io::Error>(()) + }; + if query_count > 1 && candidates >= PARALLEL_SQ_SCAN_MIN_CANDIDATES { + collectors + .par_iter_mut() + .enumerate() + .try_for_each_init(SqScanScratch::default, scan_query)?; + } else { + for query in collectors.iter_mut().enumerate() { + scan_query(&mut scratch, query)?; + } + } + } + batch_start = batch_end; + } + for (query_index, collector) in collectors.into_iter().enumerate() { + builder.record_scanned(query_index, collector.scanned()); + builder.record_early_abandoned(query_index, collector.early_abandoned()); + builder.take_rows(query_index, collector.into_rows()); + } + Ok(builder.build()) + } } pub fn search_batch_ivfsq_reader( @@ -833,11 +1068,12 @@ pub(crate) fn search_batch_ivfsq_reader_filter_range( ¢roid, &sq, metric, - filter, + SqRowSelection::Filter(filter), &mut stream_scratch, &mut heaps[query_index], - ); + )?; } + Ok(()) })?; batch_start += 1; continue; @@ -854,7 +1090,7 @@ pub(crate) fn search_batch_ivfsq_reader_filter_range( } // Keep a query's heap across partitions. Besides avoiding nprobe // allocations and merges, this carries the current cutoff into later scans. - heaps.par_iter_mut().enumerate().for_each_init( + heaps.par_iter_mut().enumerate().try_for_each_init( SqScanScratch::default, |scratch, (query_index, heap)| { let query = &processed[query_index * d..(query_index + 1) * d]; @@ -869,11 +1105,12 @@ pub(crate) fn search_batch_ivfsq_reader_filter_range( filter, scratch, heap, - ); + )?; } } + Ok::<(), io::Error>(()) }, - ); + )?; batch_start = batch_end; } @@ -1060,12 +1297,34 @@ struct SqScanScratch { distances: Vec, } +#[derive(Clone, Copy)] +enum SqRowSelection<'a> { + Filter(Option<&'a dyn RowIdFilter>), + BlockMasks(&'a [u32]), +} + +impl<'a> SqRowSelection<'a> { + fn from_masks(masks: Option<&'a [u32]>) -> Self { + masks.map(Self::BlockMasks).unwrap_or(Self::Filter(None)) + } +} + +fn sq_filter_masks(ids: &[i64], filter: &dyn RowIdFilter) -> Vec { + ids.chunks(IVF_SQ_SCAN_BLOCK_SIZE) + .map(|block| { + block.iter().enumerate().fold(0, |mask, (lane, &id)| { + mask | (u32::from(filter.contains(id)) << lane) + }) + }) + .collect() +} + // Below this point Rayon task setup and per-list heap merging dominate the // blocked SQ arithmetic. Production-sized lists usually cross the threshold; // small indexes stay on the lower-overhead sequential path. const PARALLEL_SQ_SCAN_MIN_CANDIDATES: usize = 8 * 1024; -fn scan_sq_list( +fn scan_sq_list( query: &[f32], list: &SqListData, centroid: &[f32], @@ -1073,8 +1332,8 @@ fn scan_sq_list( metric: MetricType, filter: Option<&dyn RowIdFilter>, scratch: &mut SqScanScratch, - heap: &mut TopKHeap, -) { + collector: &mut C, +) -> io::Result<()> { scan_sq_rows( query, &list.ids, @@ -1082,23 +1341,28 @@ fn scan_sq_list( centroid, sq, metric, - filter, + SqRowSelection::Filter(filter), scratch, - heap, - ); + collector, + ) } -fn scan_sq_rows( +fn scan_sq_rows( query: &[f32], ids: &[i64], codes: &[u8], centroid: &[f32], sq: &ScalarQuantizer, metric: MetricType, - filter: Option<&dyn RowIdFilter>, + selection: SqRowSelection<'_>, scratch: &mut SqScanScratch, - heap: &mut TopKHeap, -) { + collector: &mut C, +) -> io::Result<()> { + if matches!(selection, SqRowSelection::BlockMasks(masks) if masks.iter().all(|&mask| mask == 0)) + { + return Ok(()); + } + let cutoff = collector.cutoff(); sq.distances_to_blocked_codes_with_offset( query, codes, @@ -1106,18 +1370,43 @@ fn scan_sq_rows( centroid, metric, IVF_SQ_SCAN_BLOCK_SIZE, - heap.distance_limit(), + cutoff, &mut scratch.parameters, &mut scratch.distances, ); - for (&row_id, &distance) in ids.iter().zip(&scratch.distances) { - if filter.map(|f| !f.contains(row_id)).unwrap_or(false) { - continue; + let mut collect_row = |row_id, distance: f32| { + if distance.is_finite() && distance >= cutoff { + collector.note_abandoned(); + } else { + collector.push(row_id, distance)?; } - if heap.should_consider(distance) { - heap.push(distance, row_id); + Ok::<(), io::Error>(()) + }; + match selection { + SqRowSelection::Filter(filter) => { + for (&row_id, &distance) in ids.iter().zip(&scratch.distances) { + if filter + .map(|filter| !filter.contains(row_id)) + .unwrap_or(false) + { + continue; + } + collect_row(row_id, distance)?; + } + } + SqRowSelection::BlockMasks(masks) => { + for (block, &mask) in masks.iter().enumerate() { + let mut remaining = mask; + while remaining != 0 { + let position = + block * IVF_SQ_SCAN_BLOCK_SIZE + remaining.trailing_zeros() as usize; + collect_row(ids[position], scratch.distances[position])?; + remaining &= remaining - 1; + } + } } } + Ok(()) } fn padded_results(heap: TopKHeap, k: usize) -> (Vec, Vec) { @@ -1299,11 +1588,145 @@ mod tests { use super::*; use crate::io::PosWriter; use crate::io::ReadRequest; + use crate::range::{Bound, DistanceBand}; use roaring::RoaringTreemap; use std::io::Cursor; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; + #[test] + fn ivfsq_range_batch_evaluates_filter_once_per_list_row() { + struct CountingFilter(AtomicUsize); + + impl RowIdFilter for CountingFilter { + fn contains(&self, id: i64) -> bool { + self.0.fetch_add(1, Ordering::Relaxed); + id % 3 == 0 + } + } + + let (index, data, ids) = build_index(37, 4, 1_024); + let mut reader = IVFSQIndexReader::open(Cursor::new(serialized_index(&index))).unwrap(); + let filter = CountingFilter(AtomicUsize::new(0)); + let band = DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(); + let result = reader + .range_search_batch_with_filter( + &data[..37 * 3], + 3, + VectorRangeSearchParams::new(band, 4), + Some(&filter), + ) + .unwrap(); + assert_eq!(filter.0.load(Ordering::Relaxed), ids.len()); + let expected = ids.iter().filter(|&&id| id % 3 == 0).count(); + for query_index in 0..3 { + assert_eq!(result.query(query_index).labels.len(), expected); + } + } + + #[test] + fn ivfsq_range_cutoff_preserves_block_and_tail_membership() { + for dimension in [1, 31, 32, 33, 64, 65, 128] { + for count in [1, 31, 32, 33, 64, 67] { + let sq = ScalarQuantizer::with_bounds(dimension, -1.3, 2.7); + let query = vec![0.2; dimension]; + let centroid = vec![0.7; dimension]; + let ids = (0..count as i64).collect::>(); + let codes = (0..count) + .flat_map(|row| { + (0..dimension) + .map(move |component| ((row * 17 + component * 7) % 256) as u8) + }) + .collect::>(); + let blocked = block_sorted_sq_codes( + &codes, + &(0..count).collect::>(), + dimension, + IVF_SQ_SCAN_BLOCK_SIZE, + ); + let mut full = SqScanScratch::default(); + sq.distances_to_blocked_codes_with_offset( + &query, + &blocked, + count, + ¢roid, + MetricType::L2, + IVF_SQ_SCAN_BLOCK_SIZE, + f32::INFINITY, + &mut full.parameters, + &mut full.distances, + ); + let allowed: RoaringTreemap = (0..count as u64).filter(|id| id % 3 == 0).collect(); + for filter in [None, Some(&allowed as &dyn RowIdFilter)] { + let masks = filter.map(|filter| sq_filter_masks(&ids, filter)); + for selection in [ + SqRowSelection::Filter(filter), + SqRowSelection::from_masks(masks.as_deref()), + ] { + for upper in [ + Bound::Unbounded, + Bound::Finite(0.0), + Bound::Finite(full.distances[count / 2]), + ] { + let band = DistanceBand::new(Bound::Finite(0.0), upper, MetricType::L2) + .unwrap(); + let mut collector = RangeCollector::new(band); + scan_sq_rows( + &query, + &ids, + &blocked, + ¢roid, + &sq, + MetricType::L2, + selection, + &mut SqScanScratch::default(), + &mut collector, + ) + .unwrap(); + let expected = ids + .iter() + .copied() + .zip(full.distances.iter().copied()) + .filter(|(id, distance)| { + filter.map(|filter| filter.contains(*id)).unwrap_or(true) + && band.admit(*distance) + }) + .map(|(id, distance)| (id, distance.to_bits())) + .collect::>(); + let scanned = ids + .iter() + .filter(|&&id| { + filter.map(|filter| filter.contains(id)).unwrap_or(true) + }) + .count(); + assert_eq!(collector.scanned(), scanned); + assert_eq!(collector.early_abandoned(), scanned - expected.len()); + assert_eq!( + collector + .into_rows() + .into_iter() + .map(|(id, distance)| (id, distance.to_bits())) + .collect::>(), + expected, + "dimension={dimension}, count={count}" + ); + } + } + } + } + } + } + + #[test] + fn ivfsq_streamed_list_propagates_consumer_failure() { + let (index, _, _) = build_index(8, 1, 257); + let mut reader = IVFSQIndexReader::open(Cursor::new(serialized_index(&index))).unwrap(); + let error = reader + .for_each_streamed_list_chunk(0, |_, _| Err(io::Error::other("collector failed"))) + .unwrap_err(); + assert_eq!(error.to_string(), "collector failed"); + } + #[test] fn ivfsq_partition_cache_reuses_payloads_and_keeps_filters_query_local() { let (index, data, _) = build_index(37, 8, 4_097); @@ -1453,6 +1876,7 @@ mod tests { .for_each_streamed_list_chunk(0, |ids, codes| { actual_ids.extend_from_slice(ids); actual_codes.extend_from_slice(codes); + Ok(()) }) .unwrap(); assert_eq!(actual_ids, expected.0); diff --git a/core/src/range.rs b/core/src/range.rs index 7f7bb50..4ed9639 100644 --- a/core/src/range.rs +++ b/core/src/range.rs @@ -345,8 +345,9 @@ impl RangeSearchStats { pub fn lists_probed(&self) -> usize { self.lists_probed } - /// Rows read and at least partially evaluated, **including** rows abandoned - /// early. + /// Allow-listed rows read and at least partially evaluated, **including** + /// rows abandoned early. Blocked SQ arithmetic may also evaluate excluded + /// lanes; those do not enter this logical counter. pub fn rows_scanned(&self) -> usize { self.rows_scanned } @@ -355,7 +356,8 @@ impl RangeSearchStats { } /// Rows the scan rejected against the abandon cutoff rather than evaluating /// into the band test, which for IVF-Flat under L2 means their distance is - /// above the band's upper cut. + /// above the band's upper cut. IVF-SQ also counts estimates equal to that + /// exclusive cut, using its blocked quantized-distance kernel. /// /// 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 @@ -379,7 +381,8 @@ impl RangeSearchCallStats { /// logical measure): empty lists do not count, because the existing reader /// issues no payload I/O for them, and the several chunks of an oversized /// list count once. There is no re-reading, so this is also the actual - /// number of list reads. + /// number of list reads. IVF-SQ partition-cache hits do not count, since + /// they issue no payload I/O. pub fn list_reads(&self) -> usize { self.list_reads } diff --git a/core/src/topk.rs b/core/src/topk.rs index c600af1..499f210 100644 --- a/core/src/topk.rs +++ b/core/src/topk.rs @@ -17,6 +17,8 @@ use std::collections::HashMap; +use crate::collect::Collector; + pub(crate) struct TopKHeap { k: usize, max_distance: f32, @@ -133,6 +135,21 @@ impl TopKHeap { } } +impl Collector for TopKHeap { + #[inline] + fn cutoff(&self) -> f32 { + self.distance_limit() + } + + #[inline] + fn push(&mut self, id: i64, value: f32) -> std::io::Result<()> { + if self.should_consider(value) { + self.push(value, id); + } + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/core/tests/range_search.rs b/core/tests/range_search.rs index 0918373..809e4d8 100644 --- a/core/tests/range_search.rs +++ b/core/tests/range_search.rs @@ -27,15 +27,20 @@ use paimon_vindex_core::distance::{fvec_l2sqr, MetricType}; use paimon_vindex_core::index::{VectorIndexReader, VectorSearchParams}; -use paimon_vindex_core::io::PosWriter; +use paimon_vindex_core::io::{PosWriter, ReadRequest, SeekRead}; use paimon_vindex_core::ivfflat::IVFFlatIndex; use paimon_vindex_core::ivfflat_io::write_ivfflat_index; use paimon_vindex_core::ivfrq::IVFRQIndex; use paimon_vindex_core::ivfrq_io::write_ivfrq_index; +use paimon_vindex_core::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::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::sync::{Arc, Mutex}; use roaring::RoaringTreemap; @@ -947,6 +952,596 @@ fn serialize_roaring(allowed: &HashSet) -> Vec { bytes } +fn build_sq_index(dimension: usize, rows_per_list: usize, nlist: usize) -> IVFSQIndex { + let mut index = IVFSQIndex::new(dimension, nlist, MetricType::L2); + index.set_quantizer_centroids( + (0..nlist) + .flat_map(|list_id| { + (0..dimension).map(move |component| list_id as f32 + component as f32 * 0.01) + }) + .collect(), + ); + index.sq = ScalarQuantizer::with_bounds(dimension, -10.0, 10.0); + for list_id in 0..nlist { + index.list_sqs[list_id] = ScalarQuantizer::with_dimension_bounds( + dimension, + (0..dimension) + .map(|component| -0.7 - component as f32 * 0.001) + .collect(), + vec![1.3 + list_id as f32 * 0.1; dimension], + ); + for row in (0..rows_per_list).rev() { + index.ids[list_id].push((1i64 << 33) + (list_id * rows_per_list + row) as i64); + index.codes[list_id] + .extend((0..dimension).map(|component| ((row * 17 + component * 13) % 256) as u8)); + } + } + index +} + +fn serialize_sq(index: &IVFSQIndex) -> Vec { + let mut bytes = Vec::new(); + write_ivfsq_index(index, &mut PosWriter::new(&mut bytes)).unwrap(); + bytes +} + +#[test] +fn ivf_sq_range_matches_sq_estimates_for_all_entry_points() { + let dimension = 65; + let nlist = 4; + let rows_per_list = 67; + let index = build_sq_index(dimension, rows_per_list, nlist); + let queries: Vec = [0.0, 1.5, 3.0] + .into_iter() + .flat_map(|offset| (0..dimension).map(move |component| offset + component as f32 * 0.01)) + .collect(); + let mut reader = VectorIndexReader::open(Cursor::new(serialize_sq(&index))).unwrap(); + let allowed: HashSet = index + .ids + .iter() + .flatten() + .copied() + .filter(|id| id % 3 == 0) + .collect(); + let filter = serialize_roaring(&allowed); + for nprobe in [1, 2, nlist, nlist + 10] { + let band = l2(10.0, 180.0); + let params = VectorRangeSearchParams::new(band, nprobe); + let batch = reader.range_search_batch(&queries, 3, params).unwrap(); + let filtered_batch = reader + .range_search_batch_with_roaring_filter(&queries, 3, params, &filter) + .unwrap(); + for (query_index, query) in queries.chunks_exact(dimension).enumerate() { + let (ids, distances) = top_k(&mut reader, query, rows_per_list * nlist, nprobe); + let expected = bits_of( + ids.into_iter() + .zip(distances) + .filter(|(id, distance)| *id != -1 && band.admit(*distance)) + .collect(), + ); + assert!(!expected.is_empty()); + assert_eq!(pairs_of(batch.query(query_index)), expected); + assert_eq!( + pairs_of(reader.range_search(query, params).unwrap().query(0)), + expected + ); + let filtered: Vec<_> = expected + .into_iter() + .filter(|(id, _)| allowed.contains(id)) + .collect(); + assert!(!filtered.is_empty()); + assert_eq!(pairs_of(filtered_batch.query(query_index)), filtered); + assert_eq!( + pairs_of( + reader + .range_search_with_roaring_filter(query, params, &filter) + .unwrap() + .query(0) + ), + filtered, + ); + } + } +} + +#[test] +fn ivf_sq_range_uses_estimated_not_original_distance() { + let mut index = IVFSQIndex::new(1, 1, MetricType::L2); + index.set_quantizer_centroids(vec![0.0]); + index.sq = ScalarQuantizer::with_bounds(1, 0.0, 255.0); + index.list_sqs[0] = index.sq.clone(); + let vectors = [0.49, 0.51]; + let ids = [10, 20]; + index.add(&vectors, &ids, 2); + let mut reader = VectorIndexReader::open(Cursor::new(serialize_sq(&index))).unwrap(); + for band in [l2(0.0, 0.1), l2(0.2, 0.3)] { + let result = reader + .range_search(&[0.0], VectorRangeSearchParams::new(band, 1)) + .unwrap(); + let (labels, distances) = top_k(&mut reader, &[0.0], 2, 1); + let expected = bits_of( + labels + .into_iter() + .zip(distances) + .filter(|(_, distance)| band.admit(*distance)) + .collect(), + ); + assert_eq!(pairs_of(result.query(0)), expected); + assert_ne!( + pairs_of(result.query(0)), + bits_of(brute_force_band(&[0.0], &vectors, &ids, 1, band)) + ); + } +} + +#[test] +fn ivf_sq_range_preserves_boundaries_and_has_no_top_k_cap() { + let index = build_sq_index(65, 67, 3); + let mut reader = VectorIndexReader::open(Cursor::new(serialize_sq(&index))).unwrap(); + let query = vec![0.3; 65]; + let (ids, distances) = top_k(&mut reader, &query, 201, 3); + let lower = distances[20]; + let upper = distances[80]; + assert!(lower < upper); + for band in [ + l2(lower, upper), + l2(lower, lower), + l2(0.0, 0.0), + DistanceBand::new(Bound::Unbounded, Bound::Finite(upper), MetricType::L2).unwrap(), + DistanceBand::new(Bound::Finite(lower), Bound::Unbounded, MetricType::L2).unwrap(), + DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(), + ] { + let result = reader + .range_search(&query, VectorRangeSearchParams::new(band, 10)) + .unwrap(); + let expected = bits_of( + ids.iter() + .copied() + .zip(distances.iter().copied()) + .filter(|(_, distance)| band.admit(*distance)) + .collect(), + ); + assert_eq!(pairs_of(result.query(0)), expected); + assert_eq!(result.query(0).stats.rows_committed(), expected.len()); + if band.is_empty() { + 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); + } else { + assert_eq!(result.query(0).stats.lists_probed(), 3); + assert_eq!(result.query(0).stats.rows_scanned(), 201); + } + } +} + +#[test] +fn ivf_sq_range_parallel_scans_preserve_query_order_and_membership() { + let index = build_sq_index(65, 2_101, 4); + let bytes = serialize_sq(&index); + let queries: Vec = [0.0, 0.5, 2.0] + .into_iter() + .flat_map(|value| vec![value; 65]) + .collect(); + let mut reader = VectorIndexReader::open(Cursor::new(bytes)).unwrap(); + let whole = DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(); + let all = reader + .range_search_batch(&queries, 3, VectorRangeSearchParams::new(whole, 4)) + .unwrap(); + let band = l2(15.0, 180.0); + let params = VectorRangeSearchParams::new(band, 4); + let expected: Vec<_> = (0..3) + .map(|query_index| { + bits_of( + all.query(query_index) + .labels + .iter() + .copied() + .zip(all.query(query_index).distances.iter().copied()) + .filter(|(_, distance)| band.admit(*distance)) + .collect(), + ) + }) + .collect(); + let mut permuted = queries[130..].to_vec(); + permuted.extend_from_slice(&queries[..130]); + let allowed: HashSet<_> = index + .ids + .iter() + .flatten() + .copied() + .filter(|id| id % 5 == 0) + .collect(); + for threads in [1, 4] { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + .install(|| { + let result = reader.range_search_batch(&permuted, 3, params).unwrap(); + let filtered = reader + .range_search_batch_with_roaring_filter( + &queries, + 3, + params, + &serialize_roaring(&allowed), + ) + .unwrap(); + for (query_index, &original_index) in [2, 0, 1].iter().enumerate() { + assert!(!expected[original_index].is_empty()); + assert_eq!( + pairs_of(result.query(query_index)), + expected[original_index] + ); + let single = reader + .range_search( + &queries[original_index * 65..(original_index + 1) * 65], + params, + ) + .unwrap(); + assert_eq!(pairs_of(single.query(0)), expected[original_index]); + assert_eq!( + pairs_of(filtered.query(original_index)), + expected[original_index] + .iter() + .copied() + .filter(|(id, _)| allowed.contains(id)) + .collect::>() + ); + } + }); + } +} + +#[derive(Default)] +struct SqReadTrace { + calls: usize, + ranges: usize, + max_bytes: usize, +} + +struct SqRecordingReader { + inner: Cursor>, + trace: Arc>, +} + +impl SeekRead for SqRecordingReader { + fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> std::io::Result<()> { + let mut trace = self.trace.lock().unwrap(); + trace.calls += 1; + trace.ranges += ranges.len(); + for request in ranges.iter() { + trace.max_bytes = trace.max_bytes.max(request.buf.len()); + } + self.inner.pread(ranges) + } +} + +#[test] +fn ivf_sq_range_reuses_shared_lists_and_cache_with_query_local_filters() { + let mut index = build_sq_index(33, 67, 4); + index.ids[3].clear(); + index.codes[3].clear(); + let bytes = serialize_sq(&index); + let whole = DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(); + let params = VectorRangeSearchParams::new(whole, 4); + for budget in [0, 4 * 1024 * 1024] { + let trace = Arc::new(Mutex::new(SqReadTrace::default())); + let source = SqRecordingReader { + inner: Cursor::new(bytes.clone()), + trace: Arc::clone(&trace), + }; + let mut reader = + IVFSQIndexReader::open_with_options(source, VectorIndexReaderOptions::new(budget)) + .unwrap(); + *trace.lock().unwrap() = SqReadTrace::default(); + let queries = vec![0.0; 3 * 33]; + let batch = reader.range_search_batch(&queries, 3, params).unwrap(); + assert_eq!(batch.call_stats().list_reads(), 3); + assert_eq!(trace.lock().unwrap().calls, 1); + assert_eq!(trace.lock().unwrap().ranges, 3); + assert_eq!(batch.lims(), &[0, 201, 402, 603]); + let empty_filter = serialize_roaring(&HashSet::new()); + let filtered = reader + .range_search_batch_with_roaring_filter(&queries, 3, params, &empty_filter) + .unwrap(); + assert!(filtered.labels().is_empty()); + let repeated = reader.range_search(&queries[..33], params).unwrap(); + assert_eq!(pairs_of(repeated.query(0)), pairs_of(batch.query(0))); + if budget > 0 { + assert_eq!(trace.lock().unwrap().calls, 1); + assert_eq!(filtered.call_stats().list_reads(), 0); + assert_eq!(repeated.call_stats().list_reads(), 0); + } else { + assert_eq!(trace.lock().unwrap().calls, 3); + assert_eq!(repeated.call_stats().list_reads(), 3); + } + } +} + +#[test] +fn ivf_sq_range_respects_bounded_reads_and_deduplicates_partial_probes() { + struct SingleRangeReader(SqRecordingReader); + + impl SeekRead for SingleRangeReader { + fn pread(&mut self, ranges: &mut [ReadRequest<'_>]) -> std::io::Result<()> { + assert!(ranges.len() <= 1); + self.0.pread(ranges) + } + + fn read_capabilities(&self) -> paimon_vindex_core::io::SeekReadCapabilities { + paimon_vindex_core::io::SeekReadCapabilities { + max_ranges_per_pread: 1, + ..Default::default() + } + } + } + + let index = build_sq_index(33, 67, 4); + let mut queries = index.quantizer_centroids()[..66].to_vec(); + queries.extend_from_slice(&index.quantizer_centroids()[..33]); + let trace = Arc::new(Mutex::new(SqReadTrace::default())); + let source = SingleRangeReader(SqRecordingReader { + inner: Cursor::new(serialize_sq(&index)), + trace: Arc::clone(&trace), + }); + let mut reader = IVFSQIndexReader::open(source).unwrap(); + let band = DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(); + for (nprobe, reads, rows) in [(1, 2, 67), (4, 4, 268)] { + *trace.lock().unwrap() = SqReadTrace::default(); + let result = reader + .range_search_batch(&queries, 3, VectorRangeSearchParams::new(band, nprobe)) + .unwrap(); + assert_eq!(result.call_stats().list_reads(), reads); + assert_eq!(trace.lock().unwrap().calls, reads); + for query_index in 0..3 { + assert_eq!(result.query(query_index).stats.lists_probed(), nprobe); + assert_eq!(result.query(query_index).stats.rows_scanned(), rows); + assert_eq!(result.query(query_index).labels.len(), rows); + } + assert_eq!(pairs_of(result.query(0)), pairs_of(result.query(2))); + } +} + +#[test] +fn ivf_sq_range_streams_oversized_lists_for_single_batch_and_filter() { + let dimension = 512; + let count = 131_073; + let mut index = IVFSQIndex::new(dimension, 1, MetricType::L2); + index.set_quantizer_centroids(vec![0.0; dimension]); + index.sq = ScalarQuantizer::with_bounds(dimension, 0.0, 1.0); + index.list_sqs[0] = index.sq.clone(); + index.ids[0] = (0..count as i64).collect(); + index.codes[0] = vec![255; count * dimension]; + index.codes[0][..32 * dimension].fill(0); + index.codes[0][32 * dimension..64 * dimension].fill(64); + index.codes[0][(count - 1) * dimension..].fill(0); + let trace = Arc::new(Mutex::new(SqReadTrace::default())); + let source = SqRecordingReader { + inner: Cursor::new(serialize_sq(&index)), + trace: Arc::clone(&trace), + }; + drop(index); + let mut reader = IVFSQIndexReader::open(source).unwrap(); + let mut queries = vec![0.0; dimension]; + queries.extend(vec![0.25; dimension]); + let params = VectorRangeSearchParams::new(l2(0.0, 1.0), 1); + *trace.lock().unwrap() = SqReadTrace::default(); + let result = reader.range_search_batch(&queries, 2, params).unwrap(); + assert_eq!(result.call_stats().list_reads(), 1); + assert!(trace.lock().unwrap().calls > 2); + assert!(trace.lock().unwrap().max_bytes < count * dimension); + assert!(trace.lock().unwrap().max_bytes <= 64 * 1024 * 1024); + let mut first_ids: Vec<_> = (0..32).collect(); + first_ids.push((count - 1) as i64); + assert_eq!(result.query(0).labels, first_ids); + assert_eq!(result.query(1).labels, (32..64).collect::>()); + let allowed: HashSet = (0..count as i64).filter(|id| id % 2 == 0).collect(); + let filter = serialize_roaring(&allowed); + let filtered = reader + .range_search_batch_with_roaring_filter(&queries, 2, params, &filter) + .unwrap(); + for (query_index, query) in queries.chunks_exact(dimension).enumerate() { + assert_eq!(result.query(query_index).stats.rows_scanned(), count); + assert!(result.query(query_index).stats.early_abandoned() > count - 100); + let single = reader.range_search(query, params).unwrap(); + assert_eq!( + pairs_of(single.query(0)), + pairs_of(result.query(query_index)) + ); + let expected: Vec<_> = pairs_of(result.query(query_index)) + .into_iter() + .filter(|(id, _)| allowed.contains(id)) + .collect(); + assert_eq!(pairs_of(filtered.query(query_index)), expected); + assert_eq!( + pairs_of( + reader + .range_search_with_roaring_filter(query, params, &filter) + .unwrap() + .query(0) + ), + expected + ); + } +} + +#[test] +fn ivf_sq_range_validates_all_entry_points_before_empty_band_shortcuts() { + use std::io::ErrorKind::{InvalidInput, Unsupported}; + + for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { + let mut index = build_sq_index(33, 35, 2); + index.metric = metric; + let bytes = serialize_sq(&index); + let mut unified = VectorIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let mut direct = IVFSQIndexReader::open(Cursor::new(bytes)).unwrap(); + let empty = DistanceBand::new(Bound::Finite(1.0), Bound::Finite(1.0), metric).unwrap(); + let valid_filter = serialize_roaring(&HashSet::new()); + let mismatched_metric = if metric == MetricType::L2 { + MetricType::Cosine + } else { + MetricType::L2 + }; + let mismatched = + DistanceBand::new(Bound::Finite(1.0), Bound::Finite(1.0), mismatched_metric).unwrap(); + for (queries, query_count, band, nprobe, filter, error) in [ + ( + vec![0.0; 32], + 1, + empty, + 2, + valid_filter.clone(), + Some(InvalidInput), + ), + ( + vec![f32::NAN; 33], + 1, + empty, + 2, + valid_filter.clone(), + Some(InvalidInput), + ), + ( + vec![f32::INFINITY; 33], + 1, + empty, + 2, + valid_filter.clone(), + Some(InvalidInput), + ), + ( + vec![0.0; 33], + 1, + empty, + 0, + valid_filter.clone(), + Some(InvalidInput), + ), + ( + vec![0.0; 33], + 1, + mismatched, + 2, + valid_filter.clone(), + Some(InvalidInput), + ), + ( + vec![0.0; 33], + 1, + empty, + 2, + vec![0xde, 0xad], + Some(InvalidInput), + ), + ( + vec![], + 0, + empty, + 2, + valid_filter.clone(), + Some(InvalidInput), + ), + ( + vec![0.0; 33], + usize::MAX, + empty, + 2, + valid_filter.clone(), + Some(InvalidInput), + ), + ( + vec![0.0; 33], + 1, + empty, + 2, + valid_filter.clone(), + if metric == MetricType::L2 { + None + } else { + Some(Unsupported) + }, + ), + ] { + let params = VectorRangeSearchParams::new(band, nprobe); + for low_level in [false, true] { + for entry in 0..4 { + if query_count != 1 && entry < 2 || filter != valid_filter && entry % 2 == 0 { + continue; + } + let result = match (low_level, entry) { + (false, 0) => unified.range_search(&queries, params), + (false, 1) => { + unified.range_search_with_roaring_filter(&queries, params, &filter) + } + (false, 2) => unified.range_search_batch(&queries, query_count, params), + (false, _) => unified.range_search_batch_with_roaring_filter( + &queries, + query_count, + params, + &filter, + ), + (true, 0) => direct.range_search(&queries, params), + (true, 1) => { + direct.range_search_with_roaring_filter(&queries, params, &filter) + } + (true, 2) => direct.range_search_batch(&queries, query_count, params), + (true, _) => direct.range_search_batch_with_roaring_filter( + &queries, + query_count, + params, + &filter, + ), + }; + if let Some(kind) = error { + assert_eq!( + result.unwrap_err().kind(), + kind, + "metric={metric:?}, direct={low_level}, entry={entry}" + ); + } else { + assert!(result.unwrap().labels().is_empty()); + } + } + } + } + } +} + +#[test] +fn ivf_sq_range_propagates_nonfinite_estimates_and_payload_errors() { + let index = build_sq_index(33, 35, 2); + let bytes = serialize_sq(&index); + let whole = DistanceBand::new(Bound::Unbounded, Bound::Unbounded, MetricType::L2).unwrap(); + let mut reader = VectorIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + for band in [whole, l2(0.0, 1.0)] { + let params = VectorRangeSearchParams::new(band, 2); + assert_eq!( + reader + .range_search(&[f32::MAX; 33], params) + .unwrap_err() + .kind(), + std::io::ErrorKind::InvalidData + ); + assert_eq!( + reader + .range_search_batch(&[f32::MAX; 66], 2, params) + .unwrap_err() + .kind(), + std::io::ErrorKind::InvalidData + ); + } + let mut truncated = bytes; + truncated.truncate(truncated.len() - 10); + let mut reader = VectorIndexReader::open(Cursor::new(truncated)).unwrap(); + assert_eq!( + reader + .range_search(&[0.0; 33], VectorRangeSearchParams::new(whole, 2)) + .unwrap_err() + .kind(), + std::io::ErrorKind::UnexpectedEof + ); +} + // --- Task 8: the engine ---------------------------------------------------- #[test] diff --git a/docs/api.html b/docs/api.html index bc9b7c2..d1daf06 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 and IVF-RQ with l2: IVF-FLAT tests exact distances, while IVF-RQ tests estimated distances. The other families and metrics report Unsupported; callers requiring exact, complete membership need an exhaustive scan, not top-K followed by filtering. See Range search for the full contract. C/JNI range bindings are not included.

+

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.

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

Distance range search support

- + diff --git a/docs/ivf-sq.html b/docs/ivf-sq.html index bc29e0c..6e3f3ff 100644 --- a/docs/ivf-sq.html +++ b/docs/ivf-sq.html @@ -24,6 +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.

Configuration

diff --git a/docs/range-search.html b/docs/range-search.html index 8a8f107..440bd83 100644 --- a/docs/range-search.html +++ b/docs/range-search.html @@ -11,7 +11,7 @@
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 computes estimates. Range search answers "which rows are within this distance", where Top-K answers "which rows are closest".

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

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

Units and ordering

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

Entry pointSignature
Single queryrange_search(query, params)
Single query, filteredrange_search_with_roaring_filter(query, params, filter_bytes)
IVF-FLATSupported, l2 onlyExact distanceA probed row's membership is exact, so at nprobe = nlist the in-band set is complete — among allow-listed rows, when a filter is supplied
IVF-RQSupported, l2 onlyOne-bit or full multi-bit estimateSingle/batch and Roaring-filtered variants; even nprobe = nlist does not remove quantization-induced membership errors
IVF-SQPlannedEstimateAs above
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
DiskANNNot planned—Graph traversal is inherently k-oriented and has no natural radius termination criterion
MetricInternal valueNote
l2Squared Euclidean distanceA caller holding a Euclidean radius must square it, or use endpoint derivation below
cosine1 - cos, same direction as L2Not yet certified for range search
inner_product-inner_product, direction reversedNot yet certified for range search

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

-
A row within a few ULP of the upper cutIVF-FLAT can abandon a row when its partially accumulated squared distance passes the upper cut: it uses the same non-decreasing accumulation for pruning and the committed distance. IVF-RQ does not use this rule or top-K's coarse/FastScan bounds. It evaluates the complete estimate from F32 lookup sums before testing either cut, so every eligible row reaches the band test. Both families use the supplied cuts without a margin.
+
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.

"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 IVF-RQ's quantization error. An estimate can lie on the other side of a cut from the exact distance, producing missing or extra rows relative to an exact-distance predicate even at nprobe = nlist. A filter additionally excludes rows that were never eligible. Neither family truncates the rows admitted by its own distance calculation. Top-K plus post-filtering is not an equivalent fallback; exact, complete membership requires full-probe IVF-FLAT or an exhaustive raw-vector scan.
+
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.
@@ -79,20 +79,25 @@

Usage

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

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

-

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

+

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

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

Choosing an index type

-

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

+

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.

Choose according to the membership requirement.IVF-FLAT tests full-vector distances. IVF-RQ uses RaBitQ, with a one-bit estimate for one-bit files and the full multi-bit estimate otherwise, not Faiss's residual/additive quantizer. Its band predicate is precise relative to that estimate, not to the raw vector. Tests cover an independent estimated-distance oracle, single/batch equivalence, filters, statistics, parallel scans, and non-finite inputs/data; they do not establish exact-distance recall guarantees. See IVF-RQ range semantics.
+
IVF-SQ membership uses an estimate.IVF-FLAT computes exact distances from stored f32 vectors. IVF-SQ instead reuses top-K's blocked scalar-quantized estimator, reconstructing residuals with each list's stored bounds and centroid. The same estimated value determines band membership and is returned in distances; there is no original-vector reranking and no top-K fallback. Prefer IVF-FLAT if original-distance membership must be exact.
+

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

+

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.

+

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.

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-SQ, IVF-PQ or DiskANNUnsupported
An endpoint with no representable cut, meaning do not push the predicate downUnsupported
A non-finite IVF-RQ centroid, consumed factor, or computed estimateInvalid data
+
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

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.

@@ -106,6 +111,6 @@

Why not DiskANN

-
Apache Paimon Vector IndexRange search · IVF-FLAT and IVF-RQ
+
Apache Paimon Vector IndexRange search · IVF-FLAT, IVF-RQ and IVF-SQ
From eff3d50fbab61d07aa941a686b121600e6d01df1 Mon Sep 17 00:00:00 2001 From: Junrui Lee Date: Wed, 16 Sep 2026 16:41:17 +0800 Subject: [PATCH 2/2] fix: address IVF-SQ range search review feedback --- core/src/ivfsq_io.rs | 197 +++++++++++++++++++++++++++- core/tests/range_search.rs | 260 ++++++++++++++----------------------- 2 files changed, 296 insertions(+), 161 deletions(-) diff --git a/core/src/ivfsq_io.rs b/core/src/ivfsq_io.rs index 9a2b7c8..fc3c06e 100644 --- a/core/src/ivfsq_io.rs +++ b/core/src/ivfsq_io.rs @@ -862,7 +862,9 @@ impl IVFSQIndexReader { self.for_each_streamed_list_chunk(first_list, |ids, codes| { let masks = filter.map(|filter| sq_filter_masks(ids, filter)); let selection = SqRowSelection::from_masks(masks.as_deref()); - for &query_index in &list_to_queries[first_list] { + let query_indices = &list_to_queries[first_list]; + if query_indices.len() == 1 { + let query_index = query_indices[0]; scan_sq_rows( &queries[query_index * dimension..(query_index + 1) * dimension], ids, @@ -874,6 +876,24 @@ impl IVFSQIndexReader { &mut scratch, &mut collectors[query_index], )?; + } else { + let mut chunk_collectors = query_indices + .iter() + .map(|&query_index| (query_index, RangeCollector::new(band))) + .collect::>(); + scan_sq_range_chunk( + queries, + ids, + codes, + ¢roid, + &sq, + selection, + &mut scratch, + &mut chunk_collectors, + )?; + for (query_index, collector) in chunk_collectors { + collectors[query_index].merge(collector); + } } Ok(()) })?; @@ -1324,6 +1344,44 @@ fn sq_filter_masks(ids: &[i64], filter: &dyn RowIdFilter) -> Vec { // small indexes stay on the lower-overhead sequential path. const PARALLEL_SQ_SCAN_MIN_CANDIDATES: usize = 8 * 1024; +fn scan_sq_range_chunk( + queries: &[f32], + ids: &[i64], + codes: &[u8], + centroid: &[f32], + sq: &ScalarQuantizer, + selection: SqRowSelection<'_>, + scratch: &mut SqScanScratch, + collectors: &mut [(usize, C)], +) -> io::Result<()> { + let dimension = centroid.len(); + let scan_query = |scratch: &mut SqScanScratch, (query_index, collector): &mut (usize, C)| { + scan_sq_rows( + &queries[*query_index * dimension..(*query_index + 1) * dimension], + ids, + codes, + centroid, + sq, + MetricType::L2, + selection, + scratch, + collector, + ) + }; + if collectors.len() > 1 + && ids.len().saturating_mul(collectors.len()) >= PARALLEL_SQ_SCAN_MIN_CANDIDATES + { + collectors + .par_iter_mut() + .try_for_each_init(SqScanScratch::default, scan_query)?; + } else { + for collector in collectors { + scan_query(scratch, collector)?; + } + } + Ok(()) +} + fn scan_sq_list( query: &[f32], list: &SqListData, @@ -1594,6 +1652,143 @@ mod tests { use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; + #[test] + fn ivfsq_range_streamed_chunk_scans_queries_on_multiple_workers() { + struct TrackingCollector<'a> { + inner: RangeCollector, + workers: &'a AtomicU64, + } + + impl Collector for TrackingCollector<'_> { + fn cutoff(&self) -> f32 { + let worker = rayon::current_thread_index().unwrap(); + self.workers.fetch_or(1 << worker, Ordering::Relaxed); + self.inner.cutoff() + } + + fn push(&mut self, id: i64, value: f32) -> io::Result<()> { + self.inner.push(id, value) + } + + fn note_abandoned(&mut self) { + self.inner.note_abandoned(); + } + } + + let dimension = 65; + let count = 8_193; + let ids = (0..count as i64).collect::>(); + let codes = vec![0; count * dimension]; + let centroid = vec![0.0; dimension]; + let sq = ScalarQuantizer::with_bounds(dimension, 0.0, 1.0); + 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 [ + SqRowSelection::Filter(None), + SqRowSelection::BlockMasks(&masks), + ] { + let workers = AtomicU64::new(0); + let query_indices = [14, 2, 12, 4, 10, 6, 8, 0]; + let mut collectors = query_indices + .iter() + .map(|&query_index| { + ( + query_index, + TrackingCollector { + inner: RangeCollector::new(band), + workers: &workers, + }, + ) + }) + .collect::>(); + pool.install(|| { + scan_sq_range_chunk( + &queries, + &ids, + &codes, + ¢roid, + &sq, + selection, + &mut SqScanScratch::default(), + &mut collectors, + ) + .unwrap(); + }); + assert!( + workers.load(Ordering::Relaxed).count_ones() > 1, + "streamed chunks must scan active queries on multiple Rayon workers" + ); + for (query_index, collector) in collectors { + let mut expected = RangeCollector::new(band); + scan_sq_rows( + &queries[query_index * dimension..(query_index + 1) * dimension], + &ids, + &codes, + ¢roid, + &sq, + MetricType::L2, + selection, + &mut SqScanScratch::default(), + &mut expected, + ) + .unwrap(); + assert_eq!(collector.inner.scanned(), expected.scanned()); + assert_eq!( + collector.inner.early_abandoned(), + expected.early_abandoned() + ); + assert_eq!(collector.inner.into_rows(), expected.into_rows()); + } + } + } + + #[test] + fn ivfsq_range_streamed_chunk_propagates_parallel_collector_failure() { + struct FailingCollector; + + impl Collector for FailingCollector { + fn cutoff(&self) -> f32 { + f32::INFINITY + } + + fn push(&mut self, _id: i64, _value: f32) -> io::Result<()> { + Err(io::Error::other("parallel collector failed")) + } + } + + let ids = (0..4_097).collect::>(); + let codes = vec![0; ids.len()]; + let sq = ScalarQuantizer::with_bounds(1, 0.0, 1.0); + let error = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() + .install(|| { + scan_sq_range_chunk( + &[0.0, 1.0], + &ids, + &codes, + &[0.0], + &sq, + SqRowSelection::Filter(None), + &mut SqScanScratch::default(), + &mut [(0, FailingCollector), (1, FailingCollector)], + ) + .unwrap_err() + }); + assert_eq!(error.kind(), io::ErrorKind::Other); + assert_eq!(error.to_string(), "parallel collector failed"); + } + #[test] fn ivfsq_range_batch_evaluates_filter_once_per_list_row() { struct CountingFilter(AtomicUsize); diff --git a/core/tests/range_search.rs b/core/tests/range_search.rs index 809e4d8..3afcfa0 100644 --- a/core/tests/range_search.rs +++ b/core/tests/range_search.rs @@ -1325,58 +1325,64 @@ fn ivf_sq_range_streams_oversized_lists_for_single_batch_and_filter() { let mut queries = vec![0.0; dimension]; queries.extend(vec![0.25; dimension]); let params = VectorRangeSearchParams::new(l2(0.0, 1.0), 1); - *trace.lock().unwrap() = SqReadTrace::default(); - let result = reader.range_search_batch(&queries, 2, params).unwrap(); - assert_eq!(result.call_stats().list_reads(), 1); - assert!(trace.lock().unwrap().calls > 2); - assert!(trace.lock().unwrap().max_bytes < count * dimension); - assert!(trace.lock().unwrap().max_bytes <= 64 * 1024 * 1024); - let mut first_ids: Vec<_> = (0..32).collect(); - first_ids.push((count - 1) as i64); - assert_eq!(result.query(0).labels, first_ids); - assert_eq!(result.query(1).labels, (32..64).collect::>()); - let allowed: HashSet = (0..count as i64).filter(|id| id % 2 == 0).collect(); - let filter = serialize_roaring(&allowed); - let filtered = reader - .range_search_batch_with_roaring_filter(&queries, 2, params, &filter) - .unwrap(); - for (query_index, query) in queries.chunks_exact(dimension).enumerate() { - assert_eq!(result.query(query_index).stats.rows_scanned(), count); - assert!(result.query(query_index).stats.early_abandoned() > count - 100); - let single = reader.range_search(query, params).unwrap(); - assert_eq!( - pairs_of(single.query(0)), - pairs_of(result.query(query_index)) - ); - let expected: Vec<_> = pairs_of(result.query(query_index)) - .into_iter() - .filter(|(id, _)| allowed.contains(id)) - .collect(); - assert_eq!(pairs_of(filtered.query(query_index)), expected); - assert_eq!( - pairs_of( - reader - .range_search_with_roaring_filter(query, params, &filter) - .unwrap() - .query(0) - ), - expected - ); + let mut check_streamed_queries = || { + *trace.lock().unwrap() = SqReadTrace::default(); + let result = reader.range_search_batch(&queries, 2, params).unwrap(); + assert_eq!(result.call_stats().list_reads(), 1); + assert!(trace.lock().unwrap().calls > 2); + assert!(trace.lock().unwrap().max_bytes < count * dimension); + assert!(trace.lock().unwrap().max_bytes <= 64 * 1024 * 1024); + let mut first_ids: Vec<_> = (0..32).collect(); + first_ids.push((count - 1) as i64); + assert_eq!(result.query(0).labels, first_ids); + assert_eq!(result.query(1).labels, (32..64).collect::>()); + let allowed: HashSet = (0..count as i64).filter(|id| id % 2 == 0).collect(); + let filter = serialize_roaring(&allowed); + let filtered = reader + .range_search_batch_with_roaring_filter(&queries, 2, params, &filter) + .unwrap(); + for (query_index, query) in queries.chunks_exact(dimension).enumerate() { + assert_eq!(result.query(query_index).stats.rows_scanned(), count); + assert!(result.query(query_index).stats.early_abandoned() > count - 100); + let single = reader.range_search(query, params).unwrap(); + assert_eq!( + pairs_of(single.query(0)), + pairs_of(result.query(query_index)) + ); + let expected: Vec<_> = pairs_of(result.query(query_index)) + .into_iter() + .filter(|(id, _)| allowed.contains(id)) + .collect(); + assert_eq!(pairs_of(filtered.query(query_index)), expected); + assert_eq!( + pairs_of( + reader + .range_search_with_roaring_filter(query, params, &filter) + .unwrap() + .query(0) + ), + expected + ); + } + }; + for threads in [1, 4] { + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap() + .install(&mut check_streamed_queries); } } #[test] -fn ivf_sq_range_validates_all_entry_points_before_empty_band_shortcuts() { +fn ivf_sq_range_batch_validates_inputs_before_empty_band_shortcuts() { use std::io::ErrorKind::{InvalidInput, Unsupported}; for metric in [MetricType::L2, MetricType::Cosine, MetricType::InnerProduct] { let mut index = build_sq_index(33, 35, 2); index.metric = metric; - let bytes = serialize_sq(&index); - let mut unified = VectorIndexReader::open(Cursor::new(bytes.clone())).unwrap(); - let mut direct = IVFSQIndexReader::open(Cursor::new(bytes)).unwrap(); + let mut reader = IVFSQIndexReader::open(Cursor::new(serialize_sq(&index))).unwrap(); let empty = DistanceBand::new(Bound::Finite(1.0), Bound::Finite(1.0), metric).unwrap(); - let valid_filter = serialize_roaring(&HashSet::new()); let mismatched_metric = if metric == MetricType::L2 { MetricType::Cosine } else { @@ -1384,126 +1390,60 @@ fn ivf_sq_range_validates_all_entry_points_before_empty_band_shortcuts() { }; let mismatched = DistanceBand::new(Bound::Finite(1.0), Bound::Finite(1.0), mismatched_metric).unwrap(); - for (queries, query_count, band, nprobe, filter, error) in [ - ( - vec![0.0; 32], - 1, - empty, - 2, - valid_filter.clone(), - Some(InvalidInput), - ), - ( - vec![f32::NAN; 33], - 1, - empty, - 2, - valid_filter.clone(), - Some(InvalidInput), - ), - ( - vec![f32::INFINITY; 33], - 1, - empty, - 2, - valid_filter.clone(), - Some(InvalidInput), - ), - ( - vec![0.0; 33], - 1, - empty, - 0, - valid_filter.clone(), - Some(InvalidInput), - ), - ( - vec![0.0; 33], - 1, - mismatched, - 2, - valid_filter.clone(), - Some(InvalidInput), - ), - ( - vec![0.0; 33], - 1, - empty, - 2, - vec![0xde, 0xad], - Some(InvalidInput), - ), - ( - vec![], - 0, - empty, - 2, - valid_filter.clone(), - Some(InvalidInput), - ), - ( - vec![0.0; 33], - usize::MAX, - empty, - 2, - valid_filter.clone(), - Some(InvalidInput), - ), - ( - vec![0.0; 33], - 1, - empty, - 2, - valid_filter.clone(), - if metric == MetricType::L2 { - None - } else { - Some(Unsupported) - }, - ), + let query = [0.0; 33]; + for (case, queries, query_count, band, nprobe) in [ + ("dimension", &query[..32], 1, empty, 2), + ("NaN", &[f32::NAN; 33], 1, empty, 2), + ("infinity", &[f32::INFINITY; 33], 1, empty, 2), + ("nprobe", &query, 1, empty, 0), + ("metric", &query, 1, mismatched, 2), + ("zero queries", &[], 0, empty, 2), + ("query count overflow", &query, usize::MAX, empty, 2), ] { - let params = VectorRangeSearchParams::new(band, nprobe); - for low_level in [false, true] { - for entry in 0..4 { - if query_count != 1 && entry < 2 || filter != valid_filter && entry % 2 == 0 { - continue; - } - let result = match (low_level, entry) { - (false, 0) => unified.range_search(&queries, params), - (false, 1) => { - unified.range_search_with_roaring_filter(&queries, params, &filter) - } - (false, 2) => unified.range_search_batch(&queries, query_count, params), - (false, _) => unified.range_search_batch_with_roaring_filter( - &queries, - query_count, - params, - &filter, - ), - (true, 0) => direct.range_search(&queries, params), - (true, 1) => { - direct.range_search_with_roaring_filter(&queries, params, &filter) - } - (true, 2) => direct.range_search_batch(&queries, query_count, params), - (true, _) => direct.range_search_batch_with_roaring_filter( - &queries, - query_count, - params, - &filter, - ), - }; - if let Some(kind) = error { - assert_eq!( - result.unwrap_err().kind(), - kind, - "metric={metric:?}, direct={low_level}, entry={entry}" - ); - } else { - assert!(result.unwrap().labels().is_empty()); - } - } - } + let error = reader + .range_search_batch( + queries, + query_count, + VectorRangeSearchParams::new(band, nprobe), + ) + .unwrap_err(); + assert_eq!(error.kind(), InvalidInput, "{case}, metric={metric:?}"); } + let params = VectorRangeSearchParams::new(empty, 2); + assert_eq!( + reader + .range_search_batch_with_roaring_filter(&query, 1, params, &[0xde, 0xad]) + .unwrap_err() + .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); + } + } +} + +#[test] +fn ivf_sq_range_wrappers_delegate_input_validation() { + let bytes = serialize_sq(&build_sq_index(33, 35, 2)); + let mut unified = VectorIndexReader::open(Cursor::new(bytes.clone())).unwrap(); + let mut direct = IVFSQIndexReader::open(Cursor::new(bytes)).unwrap(); + let query = [0.0; 32]; + let filter = serialize_roaring(&HashSet::new()); + let params = VectorRangeSearchParams::new(l2(1.0, 1.0), 2); + for result in [ + unified.range_search(&query, params), + unified.range_search_with_roaring_filter(&query, params, &filter), + unified.range_search_batch(&query, 1, params), + unified.range_search_batch_with_roaring_filter(&query, 1, params, &filter), + direct.range_search(&query, params), + direct.range_search_with_roaring_filter(&query, params, &filter), + direct.range_search_batch_with_roaring_filter(&query, 1, params, &filter), + ] { + assert_eq!(result.unwrap_err().kind(), std::io::ErrorKind::InvalidInput); } }