This is an automated email from the ASF dual-hosted git repository.
jerry-024 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
The following commit(s) were added to refs/heads/main by this push:
new 643b4ed feat: add IVF-SQ L2 distance range search (#106)
643b4ed is described below
commit 643b4edb24a90ce3f2b60d0f600029b9f4a93479
Author: Junrui Lee <[email protected]>
AuthorDate: Thu Sep 17 15:24:11 2026 +0800
feat: add IVF-SQ L2 distance range search (#106)
---
core/README.md | 32 ++-
core/src/collect.rs | 13 +-
core/src/index.rs | 24 +-
core/src/ivfsq_io.rs | 689 ++++++++++++++++++++++++++++++++++++++++++---
core/src/range.rs | 11 +-
core/src/topk.rs | 17 ++
core/tests/range_search.rs | 537 ++++++++++++++++++++++++++++++++++-
docs/api.html | 2 +-
docs/index.html | 2 +-
docs/ivf-sq.html | 1 +
docs/range-search.html | 19 +-
11 files changed, 1279 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<R: SeekRead> VectorIndexReader<R> {
}
/// 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<R: SeekRead> VectorIndexReader<R> {
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<R: SeekRead> VectorIndexReader<R> {
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<R: SeekRead> VectorIndexReader<R> {
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<R: SeekRead> VectorIndexReader<R> {
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..fc3c06e 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<R: SeekRead> IVFSQIndexReader<R> {
}
fn read_scan_lists(&mut self, list_ids: &[usize]) ->
io::Result<Vec<Arc<SqListData>>> {
+ 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<Arc<SqListData>>, 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<R: SeekRead> IVFSQIndexReader<R> {
}
misses.push((position, list_id));
}
+ let mut reads = 0;
if !misses.is_empty() {
let missing_ids = misses.iter().map(|&(_, id)|
id).collect::<Vec<_>>();
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<R: SeekRead> IVFSQIndexReader<R> {
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<usize> {
@@ -556,7 +569,7 @@ impl<R: SeekRead> IVFSQIndexReader<R> {
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<R: SeekRead> IVFSQIndexReader<R> {
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<R: SeekRead> IVFSQIndexReader<R> {
¢roid,
&sq,
metric,
- filter,
+ SqRowSelection::Filter(filter),
&mut scratch,
&mut heap,
- );
+ )
})?;
batch_start += 1;
continue;
@@ -677,7 +690,7 @@ impl<R: SeekRead> IVFSQIndexReader<R> {
filter,
&mut SqScanScratch::default(),
&mut heap,
- );
+ )?;
let cutoff = heap.distance_limit();
let per_list_results = lists[1..]
.par_iter()
@@ -693,10 +706,10 @@ impl<R: SeekRead> IVFSQIndexReader<R> {
filter,
scratch,
&mut local_heap,
- );
- local_heap.into_sorted()
+ )?;
+ Ok(local_heap.into_sorted())
})
- .collect::<Vec<_>>();
+ .collect::<io::Result<Vec<_>>>()?;
for results in per_list_results {
for (distance, row_id) in results {
heap.push(distance, row_id);
@@ -715,7 +728,7 @@ impl<R: SeekRead> IVFSQIndexReader<R> {
filter,
&mut scratch,
&mut heap,
- );
+ )?;
}
}
batch_start = batch_end;
@@ -733,6 +746,248 @@ impl<R: SeekRead> IVFSQIndexReader<R> {
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<RangeSearchResult> {
+ 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<RangeSearchResult> {
+ 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<RangeSearchResult> {
+ 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<RangeSearchResult> {
+ 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<RangeSearchResult> {
+ 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<RangeSearchResult> {
+ 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::<Vec<_>>();
+ 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());
+ 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,
+ codes,
+ ¢roid,
+ &sq,
+ MetricType::L2,
+ selection,
+ &mut scratch,
+ &mut collectors[query_index],
+ )?;
+ } else {
+ let mut chunk_collectors = query_indices
+ .iter()
+ .map(|&query_index| (query_index,
RangeCollector::new(band)))
+ .collect::<Vec<_>>();
+ 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(())
+ })?;
+ 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::<Vec<_>>()
+ });
+ 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<R: SeekRead>(
@@ -833,11 +1088,12 @@ pub(crate) fn search_batch_ivfsq_reader_filter_range<R:
SeekRead>(
¢roid,
&sq,
metric,
- filter,
+ SqRowSelection::Filter(filter),
&mut stream_scratch,
&mut heaps[query_index],
- );
+ )?;
}
+ Ok(())
})?;
batch_start += 1;
continue;
@@ -854,7 +1110,7 @@ pub(crate) fn search_batch_ivfsq_reader_filter_range<R:
SeekRead>(
}
// 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 +1125,12 @@ pub(crate) fn search_batch_ivfsq_reader_filter_range<R:
SeekRead>(
filter,
scratch,
heap,
- );
+ )?;
}
}
+ Ok::<(), io::Error>(())
},
- );
+ )?;
batch_start = batch_end;
}
@@ -1060,12 +1317,72 @@ struct SqScanScratch {
distances: Vec<f32>,
}
+#[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<u32> {
+ 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_range_chunk<C: Collector + Send>(
+ 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<C: Collector>(
query: &[f32],
list: &SqListData,
centroid: &[f32],
@@ -1073,8 +1390,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 +1399,28 @@ fn scan_sq_list(
centroid,
sq,
metric,
- filter,
+ SqRowSelection::Filter(filter),
scratch,
- heap,
- );
+ collector,
+ )
}
-fn scan_sq_rows(
+fn scan_sq_rows<C: Collector>(
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 +1428,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<i64>, Vec<f32>) {
@@ -1299,11 +1646,282 @@ 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_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::<Vec<_>>();
+ 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::<Vec<_>>();
+ 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::<Vec<_>>();
+ 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::<Vec<_>>();
+ 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);
+
+ 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::<Vec<_>>();
+ let codes = (0..count)
+ .flat_map(|row| {
+ (0..dimension)
+ .map(move |component| ((row * 17 + component * 7)
% 256) as u8)
+ })
+ .collect::<Vec<_>>();
+ let blocked = block_sorted_sq_codes(
+ &codes,
+ &(0..count).collect::<Vec<_>>(),
+ 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::<Vec<_>>();
+ 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::<Vec<_>>(),
+ 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 +2071,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..3afcfa0 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,536 @@ fn serialize_roaring(allowed: &HashSet<i64>) -> Vec<u8> {
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<u8> {
+ 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<f32> = [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<i64> = 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<f32> = [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::<Vec<_>>()
+ );
+ }
+ });
+ }
+}
+
+#[derive(Default)]
+struct SqReadTrace {
+ calls: usize,
+ ranges: usize,
+ max_bytes: usize,
+}
+
+struct SqRecordingReader {
+ inner: Cursor<Vec<u8>>,
+ trace: Arc<Mutex<SqReadTrace>>,
+}
+
+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);
+ 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::<Vec<i64>>());
+ let allowed: HashSet<i64> = (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_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 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 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();
+ 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 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);
+ }
+}
+
+#[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 @@
<section class="article-section" id="range">
<h2>Range search parameters and results</h2>
- <p>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 <code>l2</code>: IVF-FLAT
tests exact distances, while IVF-RQ tests estimated distances. The other
families and metrics report <code>Unsupported</code>; callers requiring exact,
complete membership need an exhaustive scan, not top-K followed by filtering.
See <a href="range-search.html">Ra [...]
+ <p>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 <code>l2</code>:
IVF-FLAT tests exact distances, while IVF-RQ and IVF-SQ test quantized
estimates. The other families and metrics report <code>Unsupported</code>;
callers requiring exact, complete membership need an exhaustive scan, not top-K
followed by filtering. See <a href="ra [...]
<div class="table-wrap"><table><thead><tr><th>Entry
point</th><th>Signature</th></tr></thead><tbody>
<tr><td>Single query</td><td><code>range_search(query,
params)</code></td></tr>
<tr><td>Single query,
filtered</td><td><code>range_search_with_roaring_filter(query, params,
filter_bytes)</code></td></tr>
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 @@
<tbody>
<tr><td><a href="ivf-flat.html">IVF-FLAT</a></td><td>Supported,
<code>l2</code> only</td><td>Exact distance</td><td>A probed row's membership
is exact, so at <code>nprobe = nlist</code> the in-band set is complete — among
allow-listed rows, when a filter is supplied</td></tr>
<tr><td><a
href="ivf-rq.html#range">IVF-RQ</a></td><td>Supported, <code>l2</code>
only</td><td>One-bit or full multi-bit estimate</td><td>Single/batch and
Roaring-filtered variants; even <code>nprobe = nlist</code> does not remove
quantization-induced membership errors</td></tr>
- <tr><td><a
href="ivf-sq.html">IVF-SQ</a></td><td>Planned</td><td>Estimate</td><td>As
above</td></tr>
+ <tr><td><a href="ivf-sq.html">IVF-SQ</a></td><td>Supported in
Rust, <code>l2</code> only</td><td>SQ estimate</td><td>Single/batch, with or
without a Roaring filter; even full probing cannot guarantee membership under
the original vectors' distances</td></tr>
<tr><td><a
href="ivf-pq.html">IVF-PQ</a></td><td>Planned</td><td>Estimate</td><td>As
above; would use the float-LUT path rather than fastscan</td></tr>
<tr><td><a href="diskann.html">DiskANN</a></td><td>Not
planned</td><td>—</td><td>Graph traversal is inherently <em>k</em>-oriented and
has no natural radius termination criterion</td></tr>
</tbody>
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 @@
<h2>Build and search</h2>
<ol><li>Train the IVF coarse centroids and assign training vectors
to lists.</li><li>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.</li><li>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.</li><l [...]
<p>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.</p>
+ <p>The Rust reader also supports <a href="range-search.html">L2
distance range search</a>, 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 oversiz [...]
</section>
<section class="article-section" id="usage">
<h2>Configuration</h2>
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 @@
<a class="skip-link" href="#main">Skip to content</a>
<header class="site-header"><div class="header-inner"><a class="brand"
href="index.html" aria-label="Paimon Vector Index documentation home"><span
class="brand-mark">VI</span><span>Paimon Vector Index</span></a><nav
class="site-nav" data-site-nav aria-label="Documentation"><a
href="index.html">Overview</a><a href="api.html">API</a><a
href="development.html">Development</a><a href="range-search.html"
aria-current="page">Range search</a><a href="ivf-flat.html">IVF-FLAT</a><a
href="ivf-pq [...]
<main id="main"><div class="page-shell">
- <section class="hero detail-hero"><p class="eyebrow">Every row inside a
distance band</p><h1>Range search</h1><p class="hero-lead">Return every
eligible probed row whose family-specific distance falls inside a half-open
band <code>[lower, upper)</code>, 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".</p><div
class="badge-row"><span class="badge [...]
+ <section class="hero detail-hero"><p class="eyebrow">Every row inside a
distance band</p><h1>Range search</h1><p class="hero-lead">Return every
eligible probed row whose family-specific distance falls inside a half-open
band <code>[lower, upper)</code>, 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".</p><div class="badge-row"><span cl [...]
<div class="doc-layout">
<aside class="toc" aria-label="On this page"><strong>On this
page</strong><a href="#contract">Semantic contract</a><a href="#units">Units
and ordering</a><a href="#completeness">"No cap" is not "complete"</a><a
href="#endpoints">Endpoints and cuts</a><a href="#usage">Usage</a><a
href="#selection">Choosing an index type</a><a href="#errors">Fail-loud
combinations</a><a href="#diskann">Why not DiskANN</a></aside>
<article class="article">
@@ -29,14 +29,14 @@
<p>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.</p>
<div
class="table-wrap"><table><thead><tr><th>Metric</th><th>Internal
value</th><th>Note</th></tr></thead><tbody><tr><td><code>l2</code></td><td>Squared
Euclidean distance</td><td>A caller holding a Euclidean radius must square it,
or use endpoint derivation
below</td></tr><tr><td><code>cosine</code></td><td><code>1 - cos</code>, same
direction as L2</td><td>Not yet certified for range
search</td></tr><tr><td><code>inner_product</code></td><td><code>-inner_product</code>,
direc [...]
<p><strong>Row order is not part of the contract.</strong> 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
<code>ORDER BY</code>'s. Results are neither padded nor sorted, which is how
range search differs from Top-K.</p>
- <div class="callout"><strong>A row within a few ULP of the upper
cut</strong>IVF-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 [...]
+ <div class="callout"><strong>A row within a few ULP of the upper
cut</strong>IVF-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 bef [...]
</section>
<section class="article-section" id="completeness">
<h2>"No cap" is not "complete"</h2>
<p>Range search never truncates its result. That is a promise about
<em>not dropping rows it found</em>, and it is <strong>not</strong> a promise
that it found every in-band row in the file.</p>
<p>Only the <code>nprobe</code> nearest lists are probed, so a row
lying inside the band but in an unprobed list is not returned. A smaller
<code>nprobe</code> returns <strong>no more</strong> in-band rows, and
potentially fewer: if every matching row already lies in the lists it still
probes, the result is unchanged. At <code>nprobe == nlist</code> every list is
probed and, because IVF-FLAT computes exact distances, the result is then the
complete in-band set.</p>
- <div class="callout"><strong>Coverage and estimation are different
gaps</strong>Raising <code>nprobe</code> 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 <code>nprobe = nlist</code>. A filter
additionally excludes rows that were never eligible. Neither family truncates
the rows admitted by its own dista [...]
+ <div class="callout"><strong>Coverage and estimation are different
gaps</strong>Raising <code>nprobe</code> 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 <code>nprobe = nlist</code>. A
filter additionally excludes rows that were never eligible. None of these
families truncates the rows ad [...]
</section>
<section class="article-section" id="endpoints">
@@ -79,20 +79,25 @@ let band = DistanceBand::from_endpoints(
// 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)?;</code></pre></div>
<p>Results use a CSR layout, so a batch of queries shares three
contiguous buffers. <code>lims</code> holds <code>query_count + 1</code>
offsets, and query <code>i</code> owns <code>labels[lims[i]..lims[i+1]]</code>
together with the matching slice of <code>distances</code>. Per-query counters
are available through <code>query(i).stats</code>, and counters covering the
whole call through <code>call_stats()</code>.</p>
- <p>Both families expose <code>range_search</code>,
<code>range_search_batch</code>, and their <code>_with_roaring_filter</code>
variants in Rust. The filter is an allow-list and does not widen the fixed
<code>nprobe</code>. 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.</p>
+ <p>All three families expose <code>range_search</code>,
<code>range_search_batch</code>, and their <code>_with_roaring_filter</code>
variants in Rust. The filter is an allow-list and does not widen the fixed
<code>nprobe</code>. 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.</p>
<p>For IVF-RQ, <code>lists_probed</code> includes empty selected
lists; <code>rows_scanned</code> counts filter-eligible rows evaluated;
<code>rows_committed</code> counts returned rows; and
<code>early_abandoned</code> is zero. Call-level <code>list_reads</code> counts
unique non-empty lists, not query/list pairs or storage read rounds. These
result-owned counters leave the last top-K statistics unchanged.</p>
</section>
<section class="article-section" id="selection">
<h2>Choosing an index type</h2>
- <p>This version implements range search for <strong>IVF-FLAT and
IVF-RQ</strong>. IVF-SQ, IVF-PQ and DiskANN still return
<code>Unsupported</code>. Only L2, fixed probe widths, and Rust entry points
are included.</p>
+ <p>This version implements range search for <strong>IVF-FLAT,
IVF-RQ, and IVF-SQ</strong>. IVF-PQ and DiskANN still return
<code>Unsupported</code>. Only L2, fixed probe widths, and Rust entry points
are included.</p>
<div class="callout"><strong>Choose according to the membership
requirement.</strong>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 i [...]
+ <div class="callout"><strong>IVF-SQ membership uses an
estimate.</strong>IVF-FLAT computes exact distances from stored
<code>f32</code> 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 <code>distances</code>; there is no original-vector reranking and
no top-K fallback. Prefer IVF-FLAT if original-distance membersh [...]
+ <p>A reproducible boundary example is in
<code>core/tests/range_search.rs</code>: with a one-dimensional centroid of
zero and SQ bounds <code>[0, 255]</code>, inputs <code>0.49</code> and
<code>0.51</code> quantize to <code>0</code> and <code>1</code>. For query
zero, band <code>[0, 0.1)</code> includes the first estimate despite its true
squared distance being outside; band <code>[0.2, 0.3)</code> misses both
although both true squared distances lie inside. This demonstrates t [...]
+ <p><strong>Performance and memory:</strong> 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. [...]
+ <p>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.</p>
+ <p><code>call_stats().list_reads()</code> excludes empty lists and
IVF-SQ cache hits and counts a streamed list once. <code>rows_scanned()</code>
counts allow-listed rows reaching collection or cutoff rejection; blocked
arithmetic can also evaluate excluded lanes. For IVF-SQ,
<code>early_abandoned()</code> includes estimates equal to or above the upper
cut and is diagnostic, not a count of operations saved.</p>
</section>
<section class="article-section" id="errors">
<h2>Fail-loud combinations</h2>
<p><em>Invalid input</em> means the call itself is wrong.
<em>Unsupported</em> means the request cannot be served. <em>Invalid data</em>
covers corrupt consumed index data and non-finite computed distances; no
partial result is returned.</p>
- <div
class="table-wrap"><table><thead><tr><th>Situation</th><th>Class</th></tr></thead><tbody><tr><td>Inverted
band, where <code>lower > upper</code></td><td>Invalid
input</td></tr><tr><td>A non-finite cut</td><td>Invalid
input</td></tr><tr><td>A negative cut under squared L2</td><td>Invalid
input</td></tr><tr><td>An operator that does not match its side</td><td>Invalid
input</td></tr><tr><td>A band whose metric differs from the index's
metric</td><td>Invalid input</td></tr> [...]
+ <div
class="table-wrap"><table><thead><tr><th>Situation</th><th>Class</th></tr></thead><tbody><tr><td>Inverted
band, where <code>lower > upper</code></td><td>Invalid
input</td></tr><tr><td>A non-finite cut</td><td>Invalid
input</td></tr><tr><td>A negative cut under squared L2</td><td>Invalid
input</td></tr><tr><td>An operator that does not match its side</td><td>Invalid
input</td></tr><tr><td>A band whose metric differs from the index's
metric</td><td>Invalid input</td></tr> [...]
<p>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 <code>f_add</code> and
<code>f_rescale</code> for the estimate it consumes: coarse for one-bit codes,
full for multi-bit codes. Multi-bit coarse factors, including
<code>f_error</code>, are not used by range search and are not validated on
this path. Filtered-out and unprobed rows are not [...]
<p>An <strong>empty band is not an error</strong>: 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.</p>
</section>
@@ -106,6 +111,6 @@ let tail = DistanceBand::new(Bound::Finite(2.0),
Bound::Unbounded, MetricType::L
</article>
</div>
</div></main>
- <footer class="site-footer"><div class="footer-inner"><span>Apache Paimon
Vector Index</span><span>Range search · IVF-FLAT and
IVF-RQ</span></div></footer>
+ <footer class="site-footer"><div class="footer-inner"><span>Apache Paimon
Vector Index</span><span>Range search · IVF-FLAT, IVF-RQ and
IVF-SQ</span></div></footer>
</body>
</html>
