You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@paimon.apache.org by "JunRuiLee (via GitHub)" <gi...@apache.org> on 2026/09/16 06:32:54 UTC
[PR] feat: add IVF-SQ L2 distance range search [paimon-vector-index]
JunRuiLee opened a new pull request, #106:
URL: https://github.com/apache/paimon-vector-index/pull/106
## Summary
Part of #97. Adds IVF-SQ **L2** distance range search on top of current `main`, including the already-merged IVF-RQ support from #104.
- Support single and batch queries, each with or without a serialized Roaring allow-list, through both the unified and direct SQ readers.
- Reuse `DistanceBand`, `VectorRangeSearchParams`, `RangeSearchResult`, and the generic `Collector` scan seam. Membership is decided by the same blocked SQ-estimated squared-L2 values used by the SQ scanner, not original-vector distances or a capped top-K fallback.
- Preserve the existing top-K behavior, other metrics, bindings, and on-disk format. Existing RQ range support remains enabled.
- Document quantization-induced membership errors, half-open bounds, statistics, and the supported-family matrix.
## Performance
- Retain blocked SIMD evaluation and finite-upper-cut early abandonment.
- Read each unique probed list once per batch, reuse the existing partition cache, and stream oversized lists in bounded chunks.
- Reuse scan scratch and query-owned collectors; large single queries merge once per list rather than once per row.
- Build compact, query-local filter masks once per list/chunk and share them across queries. Masks are not retained in the partition cache; fully excluded lists/chunks skip distance evaluation.
A local release-mode synthetic check (macOS arm64, Rust 1.95.0, four Rayon threads; 32,768 rows × 128 dimensions; `nlist=16`, `nprobe=8`, batch size 16, warm 64 MiB cache, 10% Roaring filter) compared the initial range path with repeated filter lookups against shared masks:
| Band | Repeated filter lookups | Shared masks | Reduction |
| --- | ---: | ---: | ---: |
| Narrow `[0, 150)` | 1.594 ms/batch | 0.980 ms/batch | 38.5% |
| Dense `[0, 10000)` | 1.775 ms/batch | 1.176 ms/batch | 33.7% |
| Empty-hit `[0, 1)` | 1.432 ms/batch | 0.755 ms/batch | 47.3% |
These are local microbenchmarks, not production or cross-platform guarantees. Result checksums match before/after mask reuse, and the existing top-K comparison showed identical checksums without a material local slowdown.
## Coverage
- SQ-estimate oracle across all four entry points and multiple probe widths.
- Exact boundary inclusion/exclusion, unbounded and empty bands, no result cap, and a concrete distinction from original-vector membership.
- SIMD blocks/tails, query permutation and parallel scans, cache reuse, query-local filtering, partial-probe deduplication, and bounded multi-range reads.
- Oversized-list single/batch/filter paths, invalid queries and metrics, non-finite estimates, payload errors, and fallible streaming.
- A lookup-count regression test proves that three queries over 1,024 rows perform 1,024 filter lookups instead of 3,072.
## Validation
Executed locally after rebasing onto `c2d9436`:
- `cargo fmt --all -- --check`
- `cargo clippy --all-targets --workspace -- -D warnings`
- `cargo test -p paimon-vindex-core --lib`: 551 passed, 1 existing ignored test.
- `cargo test -p paimon-vindex-core --test range_search`: 51 passed, including the existing Flat and RQ cases.
- `cargo test -p paimon-vindex-core --release --test range_search`: 51 passed.
- `cargo test -p paimon-vindex-core --test storage_format_fixtures`: 2 passed, 1 existing ignored fixture printer.
The exact CI Clippy invocation also supplies `-A clippy::chunks-exact-to-as-chunks`, which local Rust 1.95.0 does not recognize. The workspace Clippy command above passes without that lint allowance.
Linux/x86_64 and full cross-language CI remain for the repository workflows. There is no storage migration or binding change; reverting this additive PR does not require rewriting index files.
## Known review follow-up
Independent review found no blocking issues. One non-blocking P2 hardening follow-up remains: SQ range coarse-probe selection uses the existing unchecked centroid ranking, so a non-finite centroid or an overflowed coarse squared distance in an unselected list can escape validation. This is a source-traced corruption/extreme-value case, not a newly reproduced ordinary-input failure. A follow-up should add SQ-range-only checked probe validation and tests for unselected NaN/overflow cases; existing top-K and RQ behavior remain unchanged here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add IVF-SQ L2 distance range search [paimon-vector-index]
Posted by "jerry-024 (via GitHub)" <gi...@apache.org>.
jerry-024 commented on code in PR #106:
URL: https://github.com/apache/paimon-vector-index/pull/106#discussion_r4023815981
##########
core/src/ivfsq_io.rs:
##########
@@ -733,6 +746,228 @@ 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());
+ 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],
+ )?;
+ }
Review Comment:
<!-- dlf-review -->
**[Major]** The oversized-list batch path serializes the per-query SQ scans.
For every streamed chunk, this loop walks all entries in `list_to_queries[first_list]` sequentially with a single `SqScanScratch`, while the normal-list path parallelizes the same per-query work. Because this branch only handles list payloads above 64 MiB, multiple queries probing the same list add substantial CPU time serially and leave Rayon workers idle. Please keep the shared chunk read, but parallelize the per-query scans with worker-local scratch and query-local collectors, then merge the chunk results.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add IVF-SQ L2 distance range search [paimon-vector-index]
Posted by "jerry-024 (via GitHub)" <gi...@apache.org>.
jerry-024 commented on code in PR #106:
URL: https://github.com/apache/paimon-vector-index/pull/106#discussion_r4023897141
##########
core/tests/range_search.rs:
##########
@@ -947,6 +952,596 @@ 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);
+ *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
+ );
+ }
+}
+
+#[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] {
Review Comment:
<!-- dlf-review -->
**[Minor]** This test repeats the same validation rules across every wrapper.
The unified and direct entry points ultimately exercise the same batch range-search validation, so the metric × reader × entry-point matrix rechecks the dimension, non-finite, `nprobe`, and metric errors many times. Could we keep one invalid-input matrix against the direct batch path, cover malformed filters through the filtered path, and retain one delegation smoke case for the remaining wrappers? That would preserve the boundary coverage while removing most of this ~140-line test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add IVF-SQ L2 distance range search [paimon-vector-index]
Posted by "JunRuiLee (via GitHub)" <gi...@apache.org>.
JunRuiLee commented on code in PR #106:
URL: https://github.com/apache/paimon-vector-index/pull/106#discussion_r4024117294
##########
core/src/ivfsq_io.rs:
##########
@@ -733,6 +746,228 @@ 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());
+ 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],
+ )?;
+ }
Review Comment:
Addressed in `eff3d50`.
Kept the shared chunk read and filter mask, and parallelized active queries with worker-local `SqScanScratch` and query-local collectors. Chunk results and counters are merged after the parallel scan completes; the single-query path and small-tail fallback remain sequential.
Added a regression that checks actual multi-worker execution, query mapping, rows and statistics, plus parallel collector-error propagation. The oversized-list integration now runs with both one and four threads. Formatting, workspace Clippy, 553 core unit tests, and all 52 range integration tests in debug and release pass locally.
##########
core/tests/range_search.rs:
##########
@@ -947,6 +952,596 @@ 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);
+ *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
+ );
+ }
+}
+
+#[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] {
Review Comment:
Addressed in `eff3d50`.
Consolidated the invalid-input matrix into the direct batch path, checked malformed filters through the filtered batch path, and kept one invalid-dimension delegation smoke case per remaining wrapper. This removes the reader/entry-point cross product while retaining empty-band validation ordering and unsupported-metric coverage.
All 52 range integration tests pass in both debug and release.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: add IVF-SQ L2 distance range search [paimon-vector-index]
Posted by "jerry-024 (via GitHub)" <gi...@apache.org>.
jerry-024 merged PR #106:
URL: https://github.com/apache/paimon-vector-index/pull/106
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org