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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
