JunRuiLee commented on code in PR #104:
URL:
https://github.com/apache/paimon-vector-index/pull/104#discussion_r4013089462
##########
core/src/ivfrq_io.rs:
##########
@@ -609,6 +613,217 @@ impl<R: SeekRead> IVFRQIndexReader<R> {
let filter = decode_roaring_filter(roaring_filter_bytes)?;
self.search_with_filter(query, k, nprobe, Some(&filter))
}
+
+ /// Returns every eligible row in the probed lists whose IVF-RQ estimated
+ /// distance is in the requested band. Only squared L2 is supported.
+ ///
+ /// Membership uses the one-bit estimate or, for multi-bit codes, the full
+ /// estimate. It does not use top-K's coarse lower-bound or FastScan
pruning:
+ /// those bounds do not certify membership in a band of estimated
distances.
+ /// Estimates are not clamped to zero. Even probing all lists does not make
+ /// membership exact with respect to the original vectors.
+ ///
+ /// Results are uncapped and unordered. Non-finite centroids, consumed
+ /// factors or computed distances return `InvalidData`. Filtered-out rows
+ /// are not evaluated. Range statistics live in the result; the last top-K
+ /// statistics are left unchanged.
+ 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)
+ }
+
+ /// Range search restricted to a serialized Roaring allow-list.
+ pub fn range_search_with_roaring_filter(
+ &mut self,
+ query: &[f32],
+ params: VectorRangeSearchParams,
+ roaring_filter_bytes: &[u8],
+ ) -> io::Result<RangeSearchResult> {
+ let filter = decode_roaring_filter(roaring_filter_bytes)?;
+ self.range_search_with_filter(query, params, Some(&filter))
+ }
+
+ /// Batched range search with the same estimates and membership as single
+ /// queries. Each unique non-empty probed list is read once per call.
+ pub fn range_search_batch(
+ &mut self,
+ queries: &[f32],
+ nq: usize,
+ params: VectorRangeSearchParams,
+ ) -> io::Result<RangeSearchResult> {
+ self.range_search_batch_with_filter(queries, nq, params, None)
+ }
+
+ /// Batched range search restricted to a serialized Roaring allow-list.
+ pub fn range_search_batch_with_roaring_filter(
+ &mut self,
+ queries: &[f32],
+ nq: 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, nq, params, Some(&filter))
+ }
+
+ pub fn range_search_batch_with_filter(
+ &mut self,
+ queries: &[f32],
+ nq: usize,
+ params: VectorRangeSearchParams,
+ filter: Option<&dyn RowIdFilter>,
+ ) -> io::Result<RangeSearchResult> {
+ validate_queries(queries, nq, self.d)?;
+ if params.band().metric() != self.metric {
+ return Err(invalid_input(format!(
+ "band metric {:?} does not match index metric {:?}",
+ params.band().metric(),
+ self.metric
+ )));
+ }
+ let nprobe = params.validate(self.nlist)?;
+ let mut builder = RangeResultBuilder::new(nq);
+ if params.band().is_empty() {
+ return Ok(builder.build());
+ }
+ self.ensure_loaded()?;
+ if self
+ .quantizer_centroids
+ .iter()
+ .any(|value| !value.is_finite())
+ {
+ return Err(invalid_data("non-finite IVF-RQ centroid"));
+ }
+ let probe_lists = queries
+ .par_chunks_exact(self.d)
+ .map(|query| {
+ let mut distances = self
+ .quantizer_centroids
+ .chunks_exact(self.d)
+ .enumerate()
+ .map(|(list_id, centroid)| {
+ let distance = fvec_l2sqr(query, centroid);
+ if !distance.is_finite() {
+ return Err(invalid_data(format!(
+ "non-finite IVF-RQ query-centroid distance for
list {list_id}"
+ )));
+ }
+ Ok((list_id, distance))
+ })
+ .collect::<io::Result<Vec<_>>>()?;
+ let compare = |left: &(usize, f32), right: &(usize, f32)| {
+ left.1.total_cmp(&right.1).then(left.0.cmp(&right.0))
+ };
+ if nprobe < distances.len() {
+ distances.select_nth_unstable_by(nprobe - 1, compare);
+ distances.truncate(nprobe);
+ distances.shrink_to_fit();
+ }
+ distances.sort_unstable_by(compare);
+ Ok(distances)
+ })
+ .collect::<io::Result<Vec<_>>>()?;
Review Comment:
Thanks for calling this out. Addressed in
7a9547e18ccba82ee6b2ae5d9e4aeb01f75dbab5, now pushed to this PR.
The range path now calls `kmeans::find_topk_checked`, which keeps at most
`min(nlist, 2 * nprobe)` centroid candidates. It periodically partitions that
buffer back to `nprobe` using the existing distance-then-list-ID ordering,
rather than collecting every centroid distance. Final sorting is in-place
(`sort_unstable_by`), so wide probes do not allocate another sorting buffer. It
still evaluates and validates every direct L2 distance before discarding a
candidate, and passes the selected distances into `RQQueryTerms`. Query-level
Rayon parallelism and the SIMD distance kernel are preserved; existing top-K
behavior is unchanged. The buffer bound applies only to centroid selection, not
to the number of returned range rows.
A local allocation probe with 65,536 lists and `nprobe=16` measured the
selector's largest allocation dropping from **1,048,576 bytes to 512 bytes**
(not a claim about total reader memory). Peak-live-allocation checks also pass
for wide and full probes, including 1,024 lists with `nprobe=512` and
`nprobe=1024`.
I did not directly substitute `find_topk_batch_with_centroid_norms` for two
reasons:
1. **Error semantics:** that helper can discard a non-finite distance when
selecting top-K. For example, two queries `[0]`, centroids `[0]` and `[1e20]`,
and `nprobe=1` return the finite selected distance without an error. This PR
explicitly requires `InvalidData` for overflow even in an unselected centroid.
The checked selector preserves that contract without adding another full
distance pass.
2. **Performance is workload-dependent:** in a local ARM64 probe-only
benchmark (`nq=100`, `nlist=65536`, `nprobe=16`, 12 Rayon threads; nine-run
medians with the three implementations interleaved in one process), the old/new
paths took **15.7/14.3 ms at d=128** and **155.2/153.8 ms at d=768**. The
existing batch helper took about **44.5/174.6 ms**, respectively. It *was*
faster with one Rayon thread, so this is not a claim that direct scanning
always wins; the bounded change fixes candidate memory without unconditionally
switching the execution strategy.
Validation: the full core test suite passes, including **548 library tests
and 42 range-search tests**, with existing ignores unchanged. Added coverage
checks bounded candidate capacity, repeated compaction, ties, probe widths,
large offsets, and multi-list single/batch/filter equivalence; the
unselected-overflow regression now also covers 65 lists. Formatting and
workspace/all-target Clippy with `-D warnings` pass.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
