jerry-024 commented on code in PR #107:
URL:
https://github.com/apache/paimon-vector-index/pull/107#discussion_r4043430732
##########
core/src/ivfpq.rs:
##########
@@ -1665,6 +1668,401 @@ pub fn search_with_reader_roaring_filter<R: SeekRead>(
search_with_reader_filter(reader, query, k, nprobe, Some(&filter))
}
+impl<R: SeekRead> IVFPQIndexReader<R> {
+ /// Returns every eligible probed row whose PQ estimate is in the L2 band.
+ ///
+ /// Both code widths use floating-point ADC: direct squared-L2 distances
+ /// from the query (after OPQ and optional coarse residual subtraction) to
+ /// each selected PQ centroid, summed in subquantizer order. Range search
+ /// does not use top-K's u8 FastScan tables or precomputed norm identities,
+ /// so membership is independent of list size and `optimize_for_search`.
+ /// This is an estimate of the original vector's distance, even at full
+ /// nprobe. Results are uncapped and unordered; top-K is unchanged.
+ ///
+ /// Non-finite transformed queries, coarse distances, or consumed PQ
+ /// distances return `InvalidData`. Filtered-out rows are not evaluated.
+ /// Every eligible row is fully evaluated, without early abandonment.
+ 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],
+ query_count: usize,
+ params: VectorRangeSearchParams,
+ ) -> io::Result<RangeSearchResult> {
+ self.range_search_batch_with_filter(queries, query_count, params, None)
+ }
+
+ /// Batched range search restricted to a serialized Roaring allow-list.
+ 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);
+ if params.band().is_empty() {
+ return Ok(builder.build());
+ }
+ self.ensure_loaded()?;
+ let mut query_states = Vec::with_capacity(query_count);
+ let mut list_to_queries = vec![Vec::new(); self.nlist];
+ for (query_index, query) in queries.chunks_exact(self.d).enumerate() {
+ let mut prepared = query.to_vec();
+ if let Some(opq) = &self.opq {
+ opq.apply(query, &mut prepared);
+ }
+ if prepared.iter().any(|value| !value.is_finite()) {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ "non-finite IVF-PQ rotated query",
+ ));
+ }
+ let probes = kmeans::find_topk_checked(
+ &prepared,
+ &self.quantizer_centroids,
+ self.nlist,
+ self.d,
+ nprobe,
+ )
+ .map_err(|list| {
+ io::Error::new(
+ io::ErrorKind::InvalidData,
+ format!("non-finite IVF-PQ query-centroid distance for
list {list}"),
+ )
+ })?;
Review Comment:
<!-- dlf-review -->
**[MAJOR] Batch coarse assignment instead of serializing it per query**
`range_search_batch_with_filter` calls `find_topk_checked` once per query,
and that helper scans all `nlist` centroids with individual distance calls. The
public batch path therefore performs `O(query_count * nlist * dimension)`
coarse-assignment work serially on the caller thread before any list I/O,
bypassing the existing batched SGEMM/parallel path used by sibling readers.
Please add or reuse a checked batched assignment helper (or at minimum the
parallel checked pattern used by IVF-RQ), while preserving the current contract
that rejects non-finite distances even for unselected centroids, then build
`list_to_queries` from the batched result.
--
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]
