jerry-024 commented on code in PR #104:
URL:
https://github.com/apache/paimon-vector-index/pull/104#discussion_r4012666418
##########
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:
<!-- dlf-review -->
**[MAJOR] Avoid materializing every centroid distance per query**
This batch path builds an nlist-sized (usize, f32) vector for every query
before truncating it to nprobe, so it retains O(nq * nlist) tuples and performs
scalar O(nq * nlist * d) work before any inverted-list I/O. At 100 queries and
65,536 lists, the tuples alone occupy about 100 MiB on a 64-bit target. Please
reuse kmeans::find_topk_batch_with_centroid_norms, as the existing top-K path
does, and pass its coarse distances into RQQueryTerms.
--
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]
