JunRuiLee commented on code in PR #107:
URL:
https://github.com/apache/paimon-vector-index/pull/107#discussion_r4043694751
##########
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:
Fixed in 6e657d0. Larger batches now use parallel `find_topk_checked`,
preserving non-finite checks for unselected centroids. Added regression tests
for query ordering and overflow before payload I/O.
--
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]
