diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 9d65220aa53..efaefa21c4e 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -403,6 +403,12 @@ pub(crate) fn build_scanner_with_options<'a>( scanner.refine(refine_factor); } + if let Some(factor) = + env.get_optional_u32_from_method(&java_obj, "getQuantizedRefineFactor")? + { + scanner.quantized_refine(factor); + } + if let Some(distance_type_str) = env.get_optional_string_from_method(&java_obj, "getDistanceTypeString")? { diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index 5a3f795880c..40efd15b9a5 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -300,6 +300,10 @@ pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> let ef = env.get_optional_usize_from_method(&java_obj, "getEf")?; let refine_factor = env.get_optional_u32_from_method(&java_obj, "getRefineFactor")?; + // Query filters cannot represent scanner execution stages, but read the + // property here so JNI contract tests cover the Java method signature. + let _quantized_refine_factor = + env.get_optional_u32_from_method(&java_obj, "getQuantizedRefineFactor")?; let distance_type = if let Some(distance_type_str) = env.get_optional_string_from_method(&java_obj, "getDistanceTypeString")? diff --git a/java/src/main/java/org/lance/ipc/Query.java b/java/src/main/java/org/lance/ipc/Query.java index 215865310df..d508a21520d 100644 --- a/java/src/main/java/org/lance/ipc/Query.java +++ b/java/src/main/java/org/lance/ipc/Query.java @@ -29,6 +29,7 @@ public class Query { private final Optional maximumNprobes; private final Optional ef; private final Optional refineFactor; + private final Optional quantizedRefineFactor; private final Optional distanceType; private final boolean useIndex; private final int queryParallelism; @@ -50,6 +51,10 @@ private Query(Builder builder) { this.maximumNprobes = builder.maximumNprobes; this.ef = builder.ef; this.refineFactor = builder.refineFactor; + Preconditions.checkArgument( + !builder.quantizedRefineFactor.isPresent() || builder.quantizedRefineFactor.get() > 0, + "Quantized refine factor must be greater than 0"); + this.quantizedRefineFactor = builder.quantizedRefineFactor; this.distanceType = builder.distanceType; this.useIndex = builder.useIndex; this.queryParallelism = builder.queryParallelism; @@ -84,6 +89,10 @@ public Optional getRefineFactor() { return refineFactor; } + public Optional getQuantizedRefineFactor() { + return quantizedRefineFactor; + } + public Optional getDistanceType() { return distanceType; } @@ -118,6 +127,7 @@ public String toString() { .add("maximumNprobes", maximumNprobes.orElse(null)) .add("ef", ef.orElse(null)) .add("refineFactor", refineFactor.orElse(null)) + .add("quantizedRefineFactor", quantizedRefineFactor.orElse(null)) .add("distanceType", distanceType.orElse(null)) .add("useIndex", useIndex) .add("queryParallelism", queryParallelism) @@ -133,6 +143,7 @@ public static class Builder { private Optional maximumNprobes = Optional.empty(); private Optional ef = Optional.empty(); private Optional refineFactor = Optional.empty(); + private Optional quantizedRefineFactor = Optional.empty(); private Optional distanceType = Optional.empty(); private boolean useIndex = true; private int queryParallelism = 0; @@ -244,6 +255,20 @@ public Builder setRefineFactor(int refineFactor) { return this; } + /** + * Reranks overfetched IVF_RQ candidates with the index's most accurate stored RQ data. + * + *

This stage does not read original vectors. It can be combined with {@link + * #setRefineFactor(int)}, in which case exact refinement remains the final stage. + * + * @param quantizedRefineFactor One-bit candidate overfetch factor; must be greater than zero. + * @return The Builder instance for method chaining. + */ + public Builder setQuantizedRefineFactor(int quantizedRefineFactor) { + this.quantizedRefineFactor = Optional.of(quantizedRefineFactor); + return this; + } + /** * Sets the distance metric type. * diff --git a/java/src/test/java/org/lance/JNITest.java b/java/src/test/java/org/lance/JNITest.java index 94db13d6dea..02396f176f9 100644 --- a/java/src/test/java/org/lance/JNITest.java +++ b/java/src/test/java/org/lance/JNITest.java @@ -31,6 +31,7 @@ import java.util.Optional; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; public class JNITest { @@ -54,21 +55,33 @@ public void testQuery() { Query defaultQuery = new Query.Builder().setColumn("column").setKey(new float[] {1.0f, 2.0f, 3.0f}).build(); assertEquals(ApproxMode.NORMAL, defaultQuery.getApproxMode()); + assertFalse(defaultQuery.getQuantizedRefineFactor().isPresent()); + + Query configuredQuery = + new Query.Builder() + .setColumn("column") + .setKey(new float[] {1.0f, 2.0f, 3.0f}) + .setK(10) + .setNprobes(20) + .setEf(30) + .setRefineFactor(40) + .setQuantizedRefineFactor(4) + .setDistanceType(DistanceType.L2) + .setUseIndex(true) + .setQueryParallelism(-1) + .setApproxMode(ApproxMode.ACCURATE) + .build(); + assertEquals(Optional.of(4), configuredQuery.getQuantizedRefineFactor()); + JniTestHelper.parseQuery(Optional.of(configuredQuery)); - JniTestHelper.parseQuery( - Optional.of( + assertThrows( + IllegalArgumentException.class, + () -> new Query.Builder() .setColumn("column") - .setKey(new float[] {1.0f, 2.0f, 3.0f}) - .setK(10) - .setNprobes(20) - .setEf(30) - .setRefineFactor(40) - .setDistanceType(DistanceType.L2) - .setUseIndex(true) - .setQueryParallelism(-1) - .setApproxMode(ApproxMode.ACCURATE) - .build())); + .setKey(new float[] {1.0f}) + .setQuantizedRefineFactor(0) + .build()); } @Test diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 1221d9d1930..296dab1c59c 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -7111,6 +7111,7 @@ def nearest( minimum_nprobes: Optional[int] = None, maximum_nprobes: Optional[int] = None, refine_factor: Optional[int] = None, + quantized_refine_factor: Optional[int] = None, use_index: bool = True, ef: Optional[int] = None, query_parallelism: Optional[int] = None, @@ -7145,6 +7146,12 @@ def nearest( setting. ``fast`` favors lower latency and may reduce recall, ``normal`` uses the default balance, and ``accurate`` favors higher recall and may increase latency. + quantized_refine_factor: int, optional + For an IVF_RQ flat index, overfetch this many one-bit candidates per + requested result and rerank them with the index's most accurate stored + RQ representation. Unlike ``refine_factor``, this does not read the + original vectors. The two stages can be combined; exact refinement + remains the final stage. """ self._nearest = _build_vector_search_query( column, @@ -7156,6 +7163,7 @@ def nearest( minimum_nprobes=minimum_nprobes, maximum_nprobes=maximum_nprobes, refine_factor=refine_factor, + quantized_refine_factor=quantized_refine_factor, use_index=use_index, ef=ef, query_parallelism=query_parallelism, @@ -8364,6 +8372,7 @@ def _build_vector_search_query( minimum_nprobes: Optional[int] = None, maximum_nprobes: Optional[int] = None, refine_factor: Optional[int] = None, + quantized_refine_factor: Optional[int] = None, use_index: bool = True, ef: Optional[int] = None, query_parallelism: Optional[int] = None, @@ -8399,6 +8408,10 @@ def _build_vector_search_query( The maximum number of partitions to search. refine_factor: int, optional The refine factor for the search. + quantized_refine_factor: int, optional + For an IVF_RQ flat index, overfetch this many one-bit candidates per + requested result and rerank them using the most accurate stored RQ data. + This does not read original vectors. use_index: bool, default True Whether to use the index for the search. ef: int, optional @@ -8479,6 +8492,10 @@ def _build_vector_search_query( raise ValueError("minimum_nprobes must be <= maximum_nprobes") if refine_factor is not None and int(refine_factor) < 1: raise ValueError(f"Refine factor must be 1 or more got {refine_factor}") + if quantized_refine_factor is not None and int(quantized_refine_factor) < 1: + raise ValueError( + f"Quantized refine factor must be 1 or more got {quantized_refine_factor}" + ) if ef is not None and int(ef) <= 0: # `ef` should be >= `k`, but `k` could be None so we can't check it here # the rust code will check it @@ -8509,6 +8526,7 @@ def _build_vector_search_query( "minimum_nprobes": minimum_nprobes, "maximum_nprobes": maximum_nprobes, "refine_factor": refine_factor, + "quantized_refine_factor": quantized_refine_factor, "use_index": use_index, "ef": ef, "query_parallelism": query_parallelism, diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index eb2c036f76b..b9ea10ecc09 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -1341,6 +1341,17 @@ def test_create_ivf_rq_multi_bit_searches_l2_and_cosine(): ) assert result.num_rows == 10 + quantized_refined = ds.to_table( + nearest={ + "column": "vector", + "q": mat[0], + "k": 10, + "quantized_refine_factor": 4, + }, + columns=["id"], + ) + assert quantized_refined.num_rows == 10 + cosine_ds = lance.write_dataset(tbl, "memory://") cosine_ds = _assert_recall_at_least(cosine_ds, mat[1], metric="cosine") cosine_stats = cosine_ds.stats.index_stats("vector_idx") @@ -2669,6 +2680,18 @@ def test_vector_index_invalid_approx_mode(indexed_dataset): ) +def test_vector_index_invalid_quantized_refine_factor(indexed_dataset): + with pytest.raises(ValueError, match="Quantized refine factor"): + indexed_dataset.scanner( + nearest={ + "column": "vector", + "q": np.random.randn(128), + "k": 10, + "quantized_refine_factor": 0, + } + ) + + def test_knn_deleted_rows(tmp_path): data = create_table() ds = lance.write_dataset(data, tmp_path) diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 6b442fb3027..e634ef1c4cc 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1534,6 +1534,7 @@ impl Dataset { maximum_nprobes, metric_type, refine_factor, + quantized_refine_factor, use_index, ef, query_parallelism, @@ -1597,6 +1598,9 @@ impl Dataset { if let Some(factor) = refine_factor { s = s.refine(factor); } + if let Some(factor) = quantized_refine_factor { + s = s.quantized_refine(factor); + } if let Some(m) = metric_type { s = s.distance_metric(m); } @@ -5512,6 +5516,7 @@ type VectorQueryParams = ( Option, Option, Option, + Option, bool, Option, i32, @@ -5643,6 +5648,23 @@ fn vector_query_params_from_dict( None }; + let quantized_refine_factor: Option = + if let Some(factor) = dict.get_item("quantized_refine_factor")? { + if factor.is_none() { + None + } else { + let factor: u32 = factor.extract()?; + if factor == 0 { + return Err(PyValueError::new_err( + "quantized_refine_factor must be greater than zero", + )); + } + Some(factor) + } + } else { + None + }; + let use_index: bool = if let Some(idx) = dict.get_item("use_index")? { idx.extract()? } else { @@ -5670,6 +5692,7 @@ fn vector_query_params_from_dict( maximum_nprobes, metric_type, refine_factor, + quantized_refine_factor, use_index, ef, query_parallelism, @@ -5707,12 +5730,19 @@ impl PySearchFilter { maximum_nprobes, metric_type_opt, refine_factor, + quantized_refine_factor, use_index, ef, query_parallelism, approx_mode, ) = vector_query_params_from_dict(query, default_k)?; + if quantized_refine_factor.is_some() { + return Err(PyValueError::new_err( + "quantized_refine_factor is only supported on dataset scanners", + )); + } + let metric_type = Some(metric_type_opt.unwrap_or(MetricType::L2)); let vector_query = VectorQuery { diff --git a/rust/lance-index/src/vector.rs b/rust/lance-index/src/vector.rs index 1192b2f5afb..41115cebb06 100644 --- a/rust/lance-index/src/vector.rs +++ b/rust/lance-index/src/vector.rs @@ -168,6 +168,25 @@ pub struct Query { pub approx_mode: ApproxMode, } +/// Stable partition-local identity for one candidate emitted by an IVF search. +#[derive(Debug, Clone, Copy, PartialEq, Eq, DeepSizeOf)] +pub struct PartitionSearchCandidate { + /// IVF partition containing the candidate. + pub partition_id: u32, + /// Zero-based offset inside the partition's vector storage. + pub offset_in_partition: u32, +} + +/// Result of a partition search that also exposes identities suitable for +/// selected-candidate re-scoring. +#[derive(Debug, Clone)] +pub struct PartitionSearchResult { + /// Ordinary `(distance, row id)` result batch. + pub batch: RecordBatch, + /// Candidate identities in the same order as rows in `batch`. + pub candidates: Vec, +} + impl From for DistanceType { fn from(proto: pb::VectorMetricType) -> Self { match proto { @@ -248,6 +267,37 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index { metrics: &dyn MetricsCollector, ) -> Result; + /// Search one partition and preserve the local offsets of emitted rows. + /// + /// Implementations may reject this operation. Quantized refinement uses it + /// only for flat IVF_RQ sub-indices. + async fn search_in_partition_with_candidates( + &self, + _partition_id: usize, + _query: &Query, + _pre_filter: Arc, + _metrics: &dyn MetricsCollector, + _candidate_limit: usize, + ) -> Result { + Err(Error::not_supported("search_in_partition_with_candidates")) + } + + /// Score partition-local candidate offsets with the index's most accurate + /// native storage representation. + /// + /// The returned rows preserve the order of `offsets_in_partition`. This + /// operation does not apply a pre-filter or distance bounds; callers must + /// validate candidate identity and apply query result shaping around it. + async fn score_partition_candidates( + &self, + _partition_id: usize, + _query: &Query, + _offsets_in_partition: &[u32], + _metrics: &dyn MetricsCollector, + ) -> Result { + Err(Error::not_supported("score_partition_candidates")) + } + /// Asynchronously prepare a single-partition search so the CPU-heavy portion /// can be executed separately. async fn prepare_partition_search( diff --git a/rust/lance-index/src/vector/flat/index.rs b/rust/lance-index/src/vector/flat/index.rs index 647d375d7ae..d4428eaa949 100644 --- a/rust/lance-index/src/vector/flat/index.rs +++ b/rust/lance-index/src/vector/flat/index.rs @@ -21,7 +21,7 @@ use crate::{ metrics::MetricsCollector, prefilter::PreFilter, vector::{ - ApproxMode, DIST_COL, Query, + ApproxMode, DIST_COL, PartitionSearchCandidate, PartitionSearchResult, Query, graph::{OrderedFloat, OrderedNode}, quantizer::{Quantization, QuantizationType, Quantizer, QuantizerMetadata}, storage::{ @@ -34,20 +34,20 @@ use crate::{ use super::storage::{FLAT_COLUMN, FlatBinStorage, FlatFloatStorage}; #[inline(always)] -fn push_candidate_local( - res: &mut BinaryHeap>, +fn push_candidate_local( + res: &mut BinaryHeap>, k: usize, - row_id: u64, + candidate: T, dist: OrderedFloat, ) { if k == 0 { return; } if res.len() < k { - res.push(OrderedNode::new(row_id, dist)); + res.push(OrderedNode::new(candidate, dist)); } else if res.peek().is_some_and(|node| node.dist > dist) { res.pop(); - res.push(OrderedNode::new(row_id, dist)); + res.push(OrderedNode::new(candidate, dist)); } } @@ -90,6 +90,110 @@ impl From<&Query> for FlatQueryParams { } } +#[allow(clippy::too_many_arguments)] +fn search_candidates_with_scratch( + query: ArrayRef, + candidate_limit: usize, + params: &FlatQueryParams, + storage: &impl VectorStore, + prefilter: Arc, + metrics: &dyn MetricsCollector, + residual: Option>, + scratch: &mut QueryScratch, +) -> Result>> { + if storage.len() > u32::MAX as usize { + return Err(Error::index(format!( + "flat partition contains {} rows, exceeding the u32 local-offset limit", + storage.len() + ))); + } + + let is_range_query = params.lower_bound.is_some() || params.upper_bound.is_some(); + let row_ids = storage.row_ids(); + let dist_calc = storage.dist_calculator_with_scratch( + query, + params.dist_q_c, + residual, + &mut scratch.query_f32, + DistanceCalculatorOptions { + approx_mode: params.approx_mode, + }, + ); + let mut candidates = BinaryHeap::with_capacity(candidate_limit); + metrics.record_comparisons(storage.len()); + + if prefilter.is_empty() { + dist_calc.distance_all_with_scratch( + candidate_limit, + &mut scratch.distances, + &mut scratch.u16, + &mut scratch.u8, + &mut scratch.u32, + ); + let distances = scratch.distances.iter().copied(); + if is_range_query { + let lower_bound = params.lower_bound.unwrap_or(f32::MIN).into(); + let upper_bound = params.upper_bound.unwrap_or(f32::MAX).into(); + for (offset, (&row_id, distance)) in row_ids.zip(distances).enumerate() { + let distance = distance.into(); + if distance < lower_bound || distance >= upper_bound { + continue; + } + push_candidate_local( + &mut candidates, + candidate_limit, + (offset as u32, row_id), + distance, + ); + } + } else { + for (offset, (&row_id, distance)) in row_ids.zip(distances).enumerate() { + push_candidate_local( + &mut candidates, + candidate_limit, + (offset as u32, row_id), + distance.into(), + ); + } + } + } else { + let row_addr_mask = prefilter.mask(); + if is_range_query { + let lower_bound = params.lower_bound.unwrap_or(f32::MIN).into(); + let upper_bound = params.upper_bound.unwrap_or(f32::MAX).into(); + for (offset, &row_id) in row_ids.enumerate() { + if !row_addr_mask.selected(row_id) { + continue; + } + let distance = dist_calc.distance(offset as u32).into(); + if distance < lower_bound || distance >= upper_bound { + continue; + } + push_candidate_local( + &mut candidates, + candidate_limit, + (offset as u32, row_id), + distance, + ); + } + } else { + for (offset, &row_id) in row_ids.enumerate() { + if !row_addr_mask.selected(row_id) { + continue; + } + push_candidate_local( + &mut candidates, + candidate_limit, + (offset as u32, row_id), + dist_calc.distance(offset as u32).into(), + ); + } + } + } + + Ok(candidates) +} + impl IvfSubIndex for FlatIndex { type QueryParams = FlatQueryParams; type BuildParams = (); @@ -139,81 +243,16 @@ impl IvfSubIndex for FlatIndex { residual: Option>, scratch: &mut QueryScratch, ) -> Result { - let is_range_query = params.lower_bound.is_some() || params.upper_bound.is_some(); - let row_ids = storage.row_ids(); - let dist_calc = storage.dist_calculator_with_scratch( - query, - params.dist_q_c, - residual, - &mut scratch.query_f32, - DistanceCalculatorOptions { - approx_mode: params.approx_mode, - }, - ); - let mut res = BinaryHeap::with_capacity(k); - metrics.record_comparisons(storage.len()); - - match prefilter.is_empty() { - true => { - dist_calc.distance_all_with_scratch( - k, - &mut scratch.distances, - &mut scratch.u16, - &mut scratch.u8, - &mut scratch.u32, - ); - let dists = scratch.distances.iter().copied(); - - if is_range_query { - let lower_bound = params.lower_bound.unwrap_or(f32::MIN).into(); - let upper_bound = params.upper_bound.unwrap_or(f32::MAX).into(); - - for (&row_id, dist) in row_ids.zip(dists) { - let dist = dist.into(); - if dist < lower_bound || dist >= upper_bound { - continue; - } - push_candidate_local(&mut res, k, row_id, dist); - } - } else { - for (&row_id, dist) in row_ids.zip(dists) { - let dist = dist.into(); - push_candidate_local(&mut res, k, row_id, dist); - } - } - } - false => { - let row_addr_mask = prefilter.mask(); - if is_range_query { - let lower_bound = params.lower_bound.unwrap_or(f32::MIN).into(); - let upper_bound = params.upper_bound.unwrap_or(f32::MAX).into(); - for (id, &row_addr) in row_ids.enumerate() { - if !row_addr_mask.selected(row_addr) { - continue; - } - let dist = dist_calc.distance(id as u32).into(); - if dist < lower_bound || dist >= upper_bound { - continue; - } - - push_candidate_local(&mut res, k, row_addr, dist); - } - } else { - for (id, &row_addr) in row_ids.enumerate() { - if !row_addr_mask.selected(row_addr) { - continue; - } - - let dist = dist_calc.distance(id as u32).into(); - push_candidate_local(&mut res, k, row_addr, dist); - } - } - } - }; + let candidates = search_candidates_with_scratch( + query, k, ¶ms, storage, prefilter, metrics, residual, scratch, + )?; // we don't need to sort the results by distances here // because there's a SortExec node in the query plan which sorts the results from all partitions - let (row_ids, dists): (Vec<_>, Vec<_>) = res.into_iter().map(|r| (r.id, r.dist.0)).unzip(); + let (row_ids, dists): (Vec<_>, Vec<_>) = candidates + .into_iter() + .map(|candidate| (candidate.id.1, candidate.dist.0)) + .unzip(); let (row_ids, dists) = (UInt64Array::from(row_ids), Float32Array::from(dists)); Ok(RecordBatch::try_new( @@ -222,6 +261,49 @@ impl IvfSubIndex for FlatIndex { )?) } + fn search_with_candidates_with_scratch( + &self, + query: ArrayRef, + candidate_limit: usize, + params: Self::QueryParams, + storage: &impl VectorStore, + prefilter: Arc, + metrics: &dyn MetricsCollector, + partition_id: u32, + residual: Option>, + scratch: &mut QueryScratch, + ) -> Result { + let results = search_candidates_with_scratch( + query, + candidate_limit, + ¶ms, + storage, + prefilter, + metrics, + residual, + scratch, + )?; + let mut candidates = Vec::with_capacity(results.len()); + let mut row_ids = Vec::with_capacity(results.len()); + let mut distances = Vec::with_capacity(results.len()); + for result in results { + candidates.push(PartitionSearchCandidate { + partition_id, + offset_in_partition: result.id.0, + }); + row_ids.push(result.id.1); + distances.push(result.dist.0); + } + let batch = RecordBatch::try_new( + ANN_SEARCH_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(distances)), + Arc::new(UInt64Array::from(row_ids)), + ], + )?; + Ok(PartitionSearchResult { batch, candidates }) + } + fn supports_global_topk_heap() -> bool { true } @@ -673,4 +755,32 @@ mod tests { vec![0, 3] ); } + + #[test] + fn test_flat_candidate_search_preserves_partition_offsets() { + let index = FlatIndex::default(); + let storage = test_storage(); + let mut scratch = QueryScratch::new(); + let result = index + .search_with_candidates_with_scratch( + query(), + 3, + FlatQueryParams::default(), + &storage, + Arc::new(NoFilter), + &NoOpMetricsCollector, + 7, + None, + &mut scratch, + ) + .unwrap(); + + assert_eq!(result.batch.num_rows(), result.candidates.len()); + let row_ids = + result.batch[lance_core::ROW_ID].as_primitive::(); + for (&row_id, candidate) in row_ids.values().iter().zip(&result.candidates) { + assert_eq!(candidate.partition_id, 7); + assert_eq!(row_id, candidate.offset_in_partition as u64); + } + } } diff --git a/rust/lance-index/src/vector/v3/subindex.rs b/rust/lance-index/src/vector/v3/subindex.rs index cf712c1b246..a0d19996fc3 100644 --- a/rust/lance-index/src/vector/v3/subindex.rs +++ b/rust/lance-index/src/vector/v3/subindex.rs @@ -14,7 +14,10 @@ use crate::metrics::MetricsCollector; use crate::vector::graph::OrderedNode; use crate::vector::storage::{QueryResidual, QueryScratch, VectorStore}; use crate::vector::{flat, hnsw}; -use crate::{prefilter::PreFilter, vector::Query}; +use crate::{ + prefilter::PreFilter, + vector::{PartitionSearchResult, Query}, +}; /// A sub index for IVF index pub trait IvfSubIndex: Send + Sync + Debug + DeepSizeOf { type QueryParams: Send + Sync + for<'a> From<&'a Query>; @@ -73,6 +76,27 @@ pub trait IvfSubIndex: Send + Sync + Debug + DeepSizeOf { self.search(query, k, params, storage, prefilter, metrics) } + /// Search the sub-index while preserving each result's partition-local + /// storage offset. Implementations that cannot expose stable offsets return + /// [`Error::NotSupported`]. + #[allow(clippy::too_many_arguments)] + fn search_with_candidates_with_scratch( + &self, + _query: ArrayRef, + _candidate_limit: usize, + _params: Self::QueryParams, + _storage: &impl VectorStore, + _prefilter: Arc, + _metrics: &dyn MetricsCollector, + _partition_id: u32, + _residual: Option>, + _scratch: &mut QueryScratch, + ) -> Result { + Err(Error::not_supported( + "partition-local candidate identities are not supported for this sub-index", + )) + } + /// Return true if this sub-index can accumulate candidates into a caller-owned heap. fn supports_global_topk_heap() -> bool { false diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 8733b167263..c63d89cb937 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1135,6 +1135,8 @@ pub struct Scanner { ordering: Option>, nearest: Option, + /// Candidate overfetch factor for RQ Fast-to-Accurate refinement. + quantized_refine_factor: Option, nearest_query_count: usize, /// True when the query shape represents a batch of single-vector queries /// (list-like query on a fixed-size vector column, or multiple concatenated vectors). @@ -1409,6 +1411,7 @@ impl Scanner { offset: None, ordering: None, nearest: None, + quantized_refine_factor: None, nearest_query_count: 1, is_batch_nearest: false, use_stats: true, @@ -2024,6 +2027,7 @@ impl Scanner { dist_q_c: 0.0, approx_mode: Default::default(), }); + self.quantized_refine_factor = None; self.nearest_query_count = query_count; self.is_batch_nearest = is_batch_nearest; Ok(self) @@ -2149,6 +2153,32 @@ impl Scanner { self } + /// Re-rank an over-fetched IVF_RQ candidate set with its multi-bit index data. + /// + /// Unlike [`Self::refine`], this does not read original vectors from the dataset. + /// It first discovers `factor * k` candidates with one-bit RQ scores and then + /// re-ranks them with the index's most accurate stored RQ representation. If + /// exact refinement is also configured, the quantized stage emits enough + /// candidates for that exact stage. + /// + /// ``` + /// # use arrow_array::Float32Array; + /// # use lance::{Dataset, Result}; + /// # async fn search(dataset: &Dataset) -> Result<()> { + /// let query = Float32Array::from(vec![0.0_f32; 128]); + /// let mut scanner = dataset.scan(); + /// scanner.nearest("vector", &query, 10)?; + /// scanner.quantized_refine(4); + /// let results = scanner.try_into_batch().await?; + /// # let _ = results; + /// # Ok(()) + /// # } + /// ``` + pub fn quantized_refine(&mut self, factor: u32) -> &mut Self { + self.quantized_refine_factor = Some(factor); + self + } + /// Change the distance [MetricType], i.e, L2 or Cosine distance. pub fn distance_metric(&mut self, metric_type: MetricType) -> &mut Self { if let Some(q) = self.nearest.as_mut() { @@ -5269,6 +5299,12 @@ impl Scanner { ) -> Result> { let mut q = q.clone(); + if matches!(self.quantized_refine_factor, Some(0)) { + return Err(Error::invalid_input( + "Quantized refine factor cannot be zero".to_string(), + )); + } + // Sanity check let (vector_type, element_type) = get_vector_type(self.dataset.schema(), &q.column)?; @@ -5450,12 +5486,37 @@ impl Scanner { // Build a prefilter block mask for stale rows (empty = no-op fast path). let overlay_block = self.stale_rows_block_mask(&stale_rows).await?; + let ann_query = if self.quantized_refine_factor.is_some() { + let exact_factor = q.refine_factor.unwrap_or(1) as usize; + let mut ann_query = q.clone(); + ann_query.k = q.k.checked_mul(exact_factor).ok_or_else(|| { + Error::invalid_input(format!( + "quantized refinement result count overflows: k={} exact_refine_factor={exact_factor}", + q.k + )) + })?; + ann_query.refine_factor = None; + ann_query + } else { + q.clone() + }; + let ann_node = match vector_type { DataType::FixedSizeList(_, _) => { - self.ann(&q, &index_segments, filter_plan, overlay_block.clone()) - .await? + self.ann( + &ann_query, + &index_segments, + filter_plan, + overlay_block.clone(), + ) + .await? } DataType::List(_) => { + if self.quantized_refine_factor.is_some() { + return Err(Error::not_supported( + "quantized refinement is not supported for multivector columns", + )); + } self.multivec_ann(&q, &index_segments, filter_plan, overlay_block.clone()) .await? } @@ -5489,6 +5550,11 @@ impl Scanner { Ok(knn_node) } else { + if self.quantized_refine_factor.is_some() { + return Err(Error::not_supported( + "quantized refinement requires an IVF_RQ vector index", + )); + } if self.fast_search { return Ok(Arc::new(EmptyExec::new(knn_empty_result_schema( self.is_batch_nearest, @@ -5543,7 +5609,7 @@ impl Scanner { /// search per query vector. /// /// Requires all of: - /// - no refine step (the batch path does not yet rerank); + /// - no exact or quantized refine step (the batch path does not yet rerank); /// - fixed nprobes (`minimum_nprobes == maximum_nprobes`) — see below; /// - every segment an IVF index with a flat-style sub-index (i.e. not HNSW); /// - every target fragment covered by the *selected* `index_segments` (or @@ -5568,12 +5634,13 @@ impl Scanner { index_segments: &[IndexMetadata], q: &Query, ) -> Result { - // Any refine factor sends the query onto a reranking path that the - // shared batch scan does not implement: the single-query path reranks - // with the original vectors even when the factor is 1, and rejects a - // factor of 0 outright (`Refine factor cannot be zero`). The batch path - // does neither, so fall back to the per-query loop for every `Some(_)`. - if q.refine_factor.is_some() { + // Any exact or quantized refine factor sends the query onto a reranking + // path that the shared batch scan does not implement. Exact refinement + // reranks with the original vectors even when the factor is 1, while + // quantized refinement overfetches and reranks with Accurate RQ scores. + // The batch path does neither, so fall back to the per-query loop for + // every configured refinement. + if q.refine_factor.is_some() || self.quantized_refine_factor.is_some() { return Ok(false); } // Only fixed nprobes is provably equivalent to single-query search; see @@ -5806,9 +5873,11 @@ impl Scanner { let q = q.clone(); debug_assert!(q.metric_type.is_some()); + let keep_quantized_distances = + self.quantized_refine_factor.is_some() && q.refine_factor.is_none(); // Ensure the vector column is present for distance computation. - if knn_node.schema().column_with_name(&q.column).is_none() { + if !keep_quantized_distances && knn_node.schema().column_with_name(&q.column).is_none() { let vector_projection = self .dataset .empty_projection() @@ -5888,11 +5957,36 @@ impl Scanner { // Union: flat paths first (matching original order), then ANN results. flat_inputs.push(knn_node); let unioned = UnionExec::try_new(flat_inputs)?; - let unioned = RepartitionExec::try_new( + let unioned: Arc = Arc::new(RepartitionExec::try_new( unioned, datafusion::physical_plan::Partitioning::RoundRobinBatch(1), - )?; - self.flat_knn(Arc::new(unioned), &q) + )?); + if keep_quantized_distances { + let sort = SortExec::new( + [ + PhysicalSortExpr { + expr: expressions::col(DIST_COL, unioned.schema().as_ref())?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }, + PhysicalSortExpr { + expr: expressions::col(ROW_ID, unioned.schema().as_ref())?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }, + ] + .into(), + unioned, + ) + .with_fetch(Some(q.k)); + Self::flat_knn_not_null_filter(Arc::new(sort)) + } else { + self.flat_knn(unioned, &q) + } } #[async_recursion] @@ -6811,6 +6905,7 @@ impl Scanner { prefilter_source, overlay_block, self.external_row_mask.clone(), + self.quantized_refine_factor, )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, inner_fanout_search.schema().as_ref())?, @@ -6873,6 +6968,7 @@ impl Scanner { prefilter_source.clone(), overlay_block.clone(), self.external_row_mask.clone(), + None, )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, ann_node.schema().as_ref())?, @@ -7270,6 +7366,7 @@ pub mod test_dataset { IndexType, scalar::{ScalarIndexParams, inverted::tokenizer::InvertedIndexParams}, vector::{ + bq::{RQBuildParams, RQRotationType}, hnsw::builder::HnswBuildParams, ivf::IvfBuildParams, kmeans::{KMeansParams, train_kmeans}, @@ -7411,6 +7508,24 @@ pub mod test_dataset { Ok(()) } + pub async fn make_rq_vector_index(&mut self, num_bits: u8) -> Result<()> { + let params = VectorIndexParams::with_ivf_rq_params( + MetricType::L2, + IvfBuildParams::new(2), + RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), + ); + self.dataset + .create_index( + &["vec"], + IndexType::Vector, + Some("idx".to_string()), + ¶ms, + true, + ) + .await?; + Ok(()) + } + pub async fn make_segmented_vector_index(&mut self) -> Result> { let batch = self .dataset @@ -10182,6 +10297,27 @@ mod test { ); } + #[tokio::test] + async fn test_batch_knn_quantized_refinement_uses_rerank_path() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + test_ds.make_rq_vector_index(4).await.unwrap(); + let (queries, _) = batch_knn_two_queries(); + + let mut scan = test_ds.dataset.scan(); + scan.nearest("vec", &queries, 10).unwrap(); + scan.nprobes(2); + scan.quantized_refine(40); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + !plan.contains("ANNIvfBatch"), + "quantized refinement must not use the shared-scan batch node, got:\n{plan}" + ); + assert!(plan.contains("quantized_refine_factor=40"), "{plan}"); + } + /// Without pinned nprobes the shared-scan fast path is not equivalent to /// single-query search (the single-query path applies an adaptive /// `early_pruning` floor and late-search expansion that the batch path does @@ -16534,6 +16670,144 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") scanner.nearest_mut().unwrap().approx_mode, ApproxMode::Accurate ); + + assert_eq!(scanner.quantized_refine_factor, None); + scanner.quantized_refine(4); + assert_eq!(scanner.quantized_refine_factor, Some(4)); + + scanner.nearest("vec", &query_vector, 5).unwrap(); + assert_eq!(scanner.quantized_refine_factor, None); + + scanner.quantized_refine(0); + let error = scanner.explain_plan(false).await.unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert!(error.to_string().contains("cannot be zero")); + } + + #[rstest] + #[case::rq4(4)] + #[case::rq8(8)] + #[tokio::test] + async fn test_ivf_rq_quantized_refinement(#[case] num_bits: u8) { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + test_ds.make_rq_vector_index(num_bits).await.unwrap(); + let query_vector = Float32Array::from_iter_values((0..32).map(|value| value as f32)); + + let mut accurate = test_ds.dataset.scan(); + accurate + .nearest("vec", &query_vector, 10) + .unwrap() + .nprobes(2) + .approx_mode(ApproxMode::Accurate) + .with_row_id() + .project(&["i"]) + .unwrap(); + let accurate_results = accurate.try_into_batch().await.unwrap(); + + // The fixture has 400 rows across multiple fragments. This factor makes + // the Fast candidate pass exhaustive, so selected-candidate scoring + // must produce the same top-k set as native Accurate RQ search. Ties in + // approximate scores may appear in a different order. + let mut quantized = test_ds.dataset.scan(); + quantized + .nearest("vec", &query_vector, 10) + .unwrap() + .nprobes(2) + .quantized_refine(40) + .with_row_id() + .project(&["i"]) + .unwrap(); + let plan = quantized.explain_plan(false).await.unwrap(); + assert!(plan.contains("quantized_refine_factor=40"), "{plan}"); + assert!(!plan.contains("KNNVectorDistance"), "{plan}"); + let quantized_results = quantized.try_into_batch().await.unwrap(); + let quantized_ids = BTreeSet::from_iter(batch_row_ids(&quantized_results)); + let accurate_ids = BTreeSet::from_iter(batch_row_ids(&accurate_results)); + assert_eq!(quantized_ids, accurate_ids); + assert!( + quantized_results[DIST_COL] + .as_primitive::() + .values() + .iter() + .all(|distance| distance.is_finite()) + ); + + // Exact refinement remains a separate final stage and still reads the + // original vectors only when the caller explicitly asks for it. + let mut exact = test_ds.dataset.scan(); + exact + .nearest("vec", &query_vector, 10) + .unwrap() + .use_index(false) + .with_row_id() + .project(&["i"]) + .unwrap(); + let exact_results = exact.try_into_batch().await.unwrap(); + + let mut quantized_then_exact = test_ds.dataset.scan(); + quantized_then_exact + .nearest("vec", &query_vector, 10) + .unwrap() + .nprobes(2) + .quantized_refine(40) + .refine(2) + .with_row_id() + .project(&["i"]) + .unwrap(); + let plan = quantized_then_exact.explain_plan(false).await.unwrap(); + assert!(plan.contains("quantized_refine_factor=40"), "{plan}"); + assert!(plan.contains("KNNVectorDistance"), "{plan}"); + let quantized_then_exact_results = quantized_then_exact.try_into_batch().await.unwrap(); + assert_eq!(quantized_then_exact_results, exact_results); + } + + #[tokio::test] + async fn test_quantized_refinement_scores_selective_prefilter_shortcut_rows() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + test_ds.make_rq_vector_index(4).await.unwrap(); + let query_vector = Float32Array::from_iter_values((0..32).map(|value| value as f32)); + + let mut scanner = test_ds.dataset.scan(); + scanner + .nearest("vec", &query_vector, 10) + .unwrap() + .minimum_nprobes(1) + .maximum_nprobes(2) + .prefilter(true) + .filter("i = 0 OR i = 399") + .unwrap() + .quantized_refine(2) + .distance_range(None, Some(f32::MAX)); + let results = scanner.try_into_batch().await.unwrap(); + + assert_eq!(results.num_rows(), 2); + assert!( + results[DIST_COL] + .as_primitive::() + .values() + .iter() + .all(|distance| distance.is_finite()) + ); + } + + #[tokio::test] + async fn test_quantized_refinement_rejects_non_rq_index() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + test_ds.make_vector_index().await.unwrap(); + let query_vector = Float32Array::from(vec![0.0; 32]); + let mut scanner = test_ds.dataset.scan(); + scanner + .nearest("vec", &query_vector, 5) + .unwrap() + .quantized_refine(2); + let error = scanner.try_into_batch().await.unwrap_err(); + assert!(error.to_string().contains("requires an IVF_RQ flat index")); } #[tokio::test] diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index d4f3134c2d4..08d2b8b971e 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -20,7 +20,9 @@ use crate::index::vector::{IndexFileVersion, builder::index_type_string}; use crate::index::{PreFilter, vector::VectorIndex}; use arrow::compute::concat_batches; use arrow_arith::numeric::sub; -use arrow_array::{ArrayRef, Float32Array, RecordBatch, UInt32Array, UInt64Array}; +use arrow_array::{ + ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, UInt32Array, UInt64Array, +}; use arrow_schema::DataType; use async_trait::async_trait; use datafusion::error::{DataFusionError, Result as DataFusionResult}; @@ -30,7 +32,7 @@ use futures::future::BoxFuture; use futures::prelude::stream::{self, TryStreamExt}; use futures::stream::FuturesUnordered; use futures::{Stream, StreamExt}; -use lance_arrow::RecordBatchExt; +use lance_arrow::{FixedSizeListArrayExt, RecordBatchExt}; use lance_core::cache::{ CacheCodec, CacheCodecImpl, CacheEntryReader, CacheEntryWriter, CacheKey, CacheKeySchema, KeyBuilder, LanceCache, WeakLanceCache, @@ -62,16 +64,16 @@ use lance_index::vector::quantizer::{ }; use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::storage::{ - QueryResidual, QueryScratch, QueryScratchCapacity, QueryScratchPool, RabitRawQueryContext, - VectorStore, + DistCalculator, DistanceCalculatorOptions, QueryResidual, QueryScratch, QueryScratchCapacity, + QueryScratchPool, RabitRawQueryContext, VectorStore, }; use lance_index::vector::v3::subindex::SubIndexType; use lance_index::{ INDEX_AUXILIARY_FILE_NAME, INDEX_FILE_NAME, Index, IndexType, pb, vector::{ - DISTANCE_TYPE_KEY, PartitionSearchControl, PreparedPartitionSearchHandle, Query, - VECTOR_RESULT_SCHEMA, ivf::storage::IVF_METADATA_KEY, quantizer::Quantization, - storage::IvfQuantizationStorage, v3::subindex::IvfSubIndex, + ApproxMode, DISTANCE_TYPE_KEY, PartitionSearchControl, PartitionSearchResult, + PreparedPartitionSearchHandle, Query, VECTOR_RESULT_SCHEMA, ivf::storage::IVF_METADATA_KEY, + quantizer::Quantization, storage::IvfQuantizationStorage, v3::subindex::IvfSubIndex, }, }; use lance_index::{INDEX_METADATA_SCHEMA_KEY, IndexMetadata}; @@ -1474,6 +1476,34 @@ impl IVFIndex { Ok(query) } + fn query_centroid_distance(&self, partition_id: usize, query: &Query) -> Result { + if partition_id >= self.ivf.num_partitions() { + return Err(Error::index(format!( + "partition {partition_id} is out of range for an IVF index with {} partitions", + self.ivf.num_partitions() + ))); + } + let centroid = self.ivf.centroid(partition_id).ok_or_else(|| { + Error::index(format!("partition centroid {partition_id} does not exist")) + })?; + let centroid = FixedSizeListArray::try_new_from_values( + centroid.clone(), + i32::try_from(centroid.len()).map_err(|_| { + Error::index(format!( + "partition centroid {partition_id} has unsupported dimension {}", + centroid.len() + )) + })?, + )?; + let distance_type = if self.distance_type == DistanceType::Cosine { + DistanceType::L2 + } else { + self.distance_type + }; + let distances = distance_type.arrow_batch_func()(query.key.as_ref(), ¢roid)?; + Ok(distances.value(0)) + } + fn query_scratch_capacity( ivf: &IvfModel, storage: &IvfQuantizationStorage, @@ -2232,6 +2262,176 @@ impl VectorIndex for IVFInd Ok(batch) } + async fn search_in_partition_with_candidates( + &self, + partition_id: usize, + query: &Query, + pre_filter: Arc, + metrics: &dyn MetricsCollector, + candidate_limit: usize, + ) -> Result { + if S::name() != "FLAT" || Q::quantization_type() != QuantizationType::Rabit { + return Err(Error::not_supported(format!( + "candidate-producing search is not supported for IVF_{}_{}", + S::name(), + Q::quantization_type() + ))); + } + let partition_id_u32 = u32::try_from(partition_id).map_err(|_| { + Error::invalid_input(format!( + "partition id {partition_id} exceeds the candidate identity limit" + )) + })?; + let part_entry = self.load_partition(partition_id, true, metrics).await?; + pre_filter.wait_for_ready().await?; + + let partition_centroid = self.ivf.centroid(partition_id); + let rq_search_cache = self.rq_search_cache.clone(); + let raw_query_context = self.prepare_rq_raw_query_context(&query.key)?; + let query = Self::preprocess_partition_query( + self.use_query_residual, + self.use_residual_scratch, + partition_id, + partition_centroid.as_ref(), + query, + )?; + let scratch_pool = self.scratch_pool.clone(); + let use_query_residual = self.use_query_residual; + let use_residual_scratch = self.use_residual_scratch; + let (result, local_metrics) = spawn_cpu(move || { + let param = (&query).into(); + let local_metrics = LocalMetricsCollector::default(); + let part = part_entry + .as_any() + .downcast_ref::>() + .ok_or(Error::internal( + "failed to downcast partition entry".to_string(), + ))?; + let rotated_partition_centroid = + rotated_partition_centroid_slice(rq_search_cache.as_deref(), partition_id); + let residual = Self::query_context_for_scratch( + use_query_residual, + use_residual_scratch, + partition_id, + partition_centroid.as_ref(), + rotated_partition_centroid, + raw_query_context.as_deref(), + )?; + let result = scratch_pool.with_scratch(|scratch| { + part.index.search_with_candidates_with_scratch( + query.key, + candidate_limit, + param, + &part.storage, + pre_filter, + &local_metrics, + partition_id_u32, + residual, + scratch, + ) + })?; + Result::Ok((result, local_metrics)) + }) + .await?; + + local_metrics.dump_into(metrics); + Ok(result) + } + + async fn score_partition_candidates( + &self, + partition_id: usize, + query: &Query, + offsets_in_partition: &[u32], + metrics: &dyn MetricsCollector, + ) -> Result { + if S::name() != "FLAT" || Q::quantization_type() != QuantizationType::Rabit { + return Err(Error::not_supported(format!( + "selected-candidate scoring is not supported for IVF_{}_{}", + S::name(), + Q::quantization_type() + ))); + } + + let mut query = query.clone(); + if Q::quantization_type() == QuantizationType::Rabit { + query.dist_q_c = self.query_centroid_distance(partition_id, &query)?; + } + let part_entry = self.load_partition(partition_id, true, metrics).await?; + let partition_centroid = self.ivf.centroid(partition_id); + let rq_search_cache = self.rq_search_cache.clone(); + let raw_query_context = self.prepare_rq_raw_query_context(&query.key)?; + let query = Self::preprocess_partition_query_owned( + self.use_query_residual, + self.use_residual_scratch, + partition_id, + partition_centroid.as_ref(), + query, + )?; + let offsets_in_partition = offsets_in_partition.to_vec(); + let scratch_pool = self.scratch_pool.clone(); + let use_query_residual = self.use_query_residual; + let use_residual_scratch = self.use_residual_scratch; + let (batch, local_metrics) = spawn_cpu(move || { + let local_metrics = LocalMetricsCollector::default(); + let part = part_entry + .as_any() + .downcast_ref::>() + .ok_or(Error::internal( + "failed to downcast partition entry".to_string(), + ))?; + let partition_len = part.storage.len(); + for &offset in &offsets_in_partition { + if offset as usize >= partition_len { + return Err(Error::index(format!( + "partition offset {offset} is out of range for partition {partition_id} with {partition_len} rows" + ))); + } + } + + let rotated_partition_centroid = + rotated_partition_centroid_slice(rq_search_cache.as_deref(), partition_id); + let residual = Self::query_context_for_scratch( + use_query_residual, + use_residual_scratch, + partition_id, + partition_centroid.as_ref(), + rotated_partition_centroid, + raw_query_context.as_deref(), + )?; + let batch = scratch_pool.with_scratch(|scratch| { + let dist_calc = part.storage.dist_calculator_with_scratch( + query.key, + query.dist_q_c, + residual, + &mut scratch.query_f32, + DistanceCalculatorOptions { + approx_mode: ApproxMode::Accurate, + }, + ); + let mut row_ids = Vec::with_capacity(offsets_in_partition.len()); + let mut distances = Vec::with_capacity(offsets_in_partition.len()); + for &offset in &offsets_in_partition { + row_ids.push(part.storage.row_id(offset)); + distances.push(dist_calc.distance(offset)); + } + Result::Ok(RecordBatch::try_new( + VECTOR_RESULT_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(distances)), + Arc::new(UInt64Array::from(row_ids)), + ], + )?) + })?; + local_metrics.record_comparisons(offsets_in_partition.len()); + Result::Ok((batch, local_metrics)) + }) + .await?; + + local_metrics.dump_into(metrics); + Ok(batch) + } + async fn prepare_partition_search( &self, partition_id: usize, @@ -2947,8 +3147,9 @@ mod tests { Array, ArrayRef, ArrowPrimitiveType, FixedSizeListArray, Float32Array, Int64Array, ListArray, PrimitiveArray, RecordBatch, RecordBatchIterator, UInt64Array, }; - use arrow_buffer::OffsetBuffer; + use arrow_buffer::{NullBuffer, OffsetBuffer}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use futures::TryStreamExt; use itertools::Itertools; use lance_arrow::FixedSizeListArrayExt; use lance_index::vector::bq::{ @@ -2964,6 +3165,7 @@ mod tests { use crate::dataset::{InsertBuilder, UpdateBuilder, WriteMode, WriteParams}; use crate::index::DatasetIndexExt; use crate::index::DatasetIndexInternalExt; + use crate::index::prefilter::DatasetPreFilter; use crate::index::vector::ivf::v2::{ IVFPartitionKey, IvfFlatIndex, IvfHnswSqIndex, IvfPq, IvfStateEntryBox, PartitionEntry, }; @@ -2976,7 +3178,6 @@ mod tests { dataset::optimize::{CompactionOptions, compact_files}, index::vector::IndexFileVersion, }; - use futures::TryStreamExt; use lance_core::cache::{CacheBackend, CacheCodecImpl, LanceCache, WeakLanceCache}; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::tempfile::TempStrDir; @@ -3000,13 +3201,16 @@ mod tests { use lance_index::vector::quantizer::QuantizerMetadata; use lance_index::vector::sq::ScalarQuantizer; use lance_index::vector::sq::builder::SQBuildParams; - use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, Query}; use lance_index::vector::{ + ApproxMode, DEFAULT_QUERY_PARALLELISM, Query, pq::storage::ProductQuantizationMetadata, sq::storage::{SQ_METADATA_KEY, ScalarQuantizationMetadata}, storage::STORAGE_METADATA_KEY, }; - use lance_index::{INDEX_AUXILIARY_FILE_NAME, metrics::NoOpMetricsCollector}; + use lance_index::{ + INDEX_AUXILIARY_FILE_NAME, + metrics::{LocalMetricsCollector, NoOpMetricsCollector}, + }; use lance_io::{ object_store::{ObjectStore, ObjectStoreParams, StorageOptionsAccessor}, scheduler::{ScanScheduler, SchedulerConfig}, @@ -3362,6 +3566,35 @@ mod tests { (dataset, Arc::new(vectors.as_fixed_size_list().clone())) } + async fn generate_f32_test_dataset_with_shape( + test_uri: &str, + num_rows: usize, + dimension: usize, + range: Range, + ) -> (Dataset, Arc) { + let ids = Arc::new(UInt64Array::from_iter_values(0..num_rows as u64)); + let values = generate_random_array_with_range::(num_rows * dimension, range); + let vectors = + Arc::new(FixedSizeListArray::try_new_from_values(values, dimension as i32).unwrap()); + let vectors = Arc::new(normalize_fsl(vectors.as_ref()).unwrap()); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("vector", vectors.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new(schema.clone(), vec![ids, vectors.clone()]).unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + test_uri, + Some(WriteParams { + mode: crate::dataset::WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); + (dataset, vectors) + } + async fn generate_multivec_test_dataset( test_uri: &str, range: Range, @@ -6118,6 +6351,292 @@ mod tests { test_index_impl::(params, nlist, 0.75, -1.0..1.0, None).await; } + #[rstest] + #[case::rq4_l2(DistanceType::L2, 4, DIM)] + #[case::rq4_cosine(DistanceType::Cosine, 4, DIM)] + #[case::rq4_dot(DistanceType::Dot, 4, DIM)] + #[case::rq8_l2(DistanceType::L2, 8, DIM)] + #[case::rq8_cosine(DistanceType::Cosine, 8, DIM)] + #[case::rq8_dot(DistanceType::Dot, 8, DIM)] + #[case::rq4_dot_768(DistanceType::Dot, 4, 768)] + #[case::rq8_cosine_4096(DistanceType::Cosine, 8, 4096)] + #[tokio::test] + async fn test_score_partition_candidates_matches_accurate_native_search( + #[case] distance_type: DistanceType, + #[case] num_bits: u8, + #[case] dimension: usize, + ) { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let num_rows = if dimension == DIM { NUM_ROWS } else { 128 }; + let (mut dataset, vectors) = + generate_f32_test_dataset_with_shape(test_uri, num_rows, dimension, 0.0..1.0).await; + let ivf_params = IvfBuildParams::new(4); + let params = VectorIndexParams::with_ivf_rq_params( + distance_type, + ivf_params, + RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + let index = dataset + .open_vector_index("vector", &indices[0].uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let mut query = Query { + column: "vector".to_string(), + key: vectors.value(0), + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 1, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: Some(distance_type), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: f32::NAN, + approx_mode: ApproxMode::Fast, + }; + let (partition_ids, centroid_distances) = index.find_partitions(&query).unwrap(); + let partition_position = (0..partition_ids.len()) + .find(|&position| index.partition_size(partition_ids.value(position) as usize) > 0) + .expect("the test index should contain a non-empty partition"); + let partition_id = partition_ids.value(partition_position) as usize; + let partition_len = index.partition_size(partition_id); + assert!(partition_len >= 3); + + let partition_batches = index + .partition_reader(partition_id, false, &NoOpMetricsCollector) + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let partition_row_ids = partition_batches + .iter() + .flat_map(|batch| { + batch[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(partition_row_ids.len(), partition_len); + + query.dist_q_c = centroid_distances.value(partition_position); + let prefilter = Arc::new(DatasetPreFilter::new( + Arc::new(dataset.clone()), + &indices, + None, + )); + let candidate_limit = partition_len.min(5); + let partition_search = index + .search_in_partition_with_candidates( + partition_id, + &query, + prefilter.clone(), + &NoOpMetricsCollector, + candidate_limit, + ) + .await + .unwrap(); + assert_eq!(partition_search.batch.num_rows(), candidate_limit); + assert_eq!(partition_search.candidates.len(), candidate_limit); + let candidate_row_ids = partition_search.batch[ROW_ID].as_primitive::(); + for (position, candidate) in partition_search.candidates.iter().enumerate() { + assert_eq!(candidate.partition_id as usize, partition_id); + assert_eq!( + candidate_row_ids.value(position), + partition_row_ids[candidate.offset_in_partition as usize] + ); + } + + let offsets = vec![0, (partition_len / 2) as u32, (partition_len - 1) as u32]; + let metrics = LocalMetricsCollector::default(); + let scored = index + .score_partition_candidates(partition_id, &query, &offsets, &metrics) + .await + .unwrap(); + assert_eq!(scored.num_rows(), offsets.len()); + assert_eq!(metrics.comparisons.load(Ordering::Relaxed), offsets.len()); + let empty = index + .score_partition_candidates(partition_id, &query, &[], &NoOpMetricsCollector) + .await + .unwrap(); + assert_eq!(empty.num_rows(), 0); + + query.k = partition_len; + query.approx_mode = ApproxMode::Accurate; + let searched = index + .search_in_partition(partition_id, &query, prefilter, &NoOpMetricsCollector) + .await + .unwrap(); + let expected_distances = searched[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .zip( + searched[DIST_COL] + .as_primitive::() + .values() + .iter() + .copied(), + ) + .collect::>(); + let scored_row_ids = scored[ROW_ID].as_primitive::(); + let scored_distances = scored[DIST_COL].as_primitive::(); + for (position, &offset) in offsets.iter().enumerate() { + let row_id = partition_row_ids[offset as usize]; + assert_eq!(scored_row_ids.value(position), row_id); + let expected_distance = expected_distances[&row_id]; + let actual_distance = scored_distances.value(position); + assert!( + actual_distance.is_finite(), + "selected score for row {row_id} should be finite" + ); + let tolerance = 1.0e-4 * (1.0 + expected_distance.abs()); + assert!( + (actual_distance - expected_distance).abs() <= tolerance, + "selected score {actual_distance} differs from accurate native score {expected_distance} for row {row_id}" + ); + } + + let error = index + .score_partition_candidates( + partition_id, + &query, + &[partition_len as u32], + &NoOpMetricsCollector, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains(&format!( + "partition offset {partition_len} is out of range for partition {partition_id}" + ))); + } + + #[rstest] + #[case::rq4(4)] + #[case::rq8(8)] + #[tokio::test] + async fn test_selected_scoring_excludes_null_and_non_finite_stored_vectors( + #[case] num_bits: u8, + ) { + const EDGE_DIM: usize = 8; + let test_dir = TempStrDir::default(); + let values = [ + [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0; EDGE_DIM], + [f32::NAN, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ] + .into_iter() + .flatten() + .collect::>(); + let vectors = Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::Float32, true)), + EDGE_DIM as i32, + Arc::new(Float32Array::from(values)), + Some(NullBuffer::from(vec![true, false, true, true, true, true])), + ) + .unwrap(), + ); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt64, false), + Field::new("vector", vectors.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(UInt64Array::from_iter_values(0..vectors.len() as u64)), + vectors, + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + test_dir.as_str(), + None, + ) + .await + .unwrap(); + let centroids = Arc::new( + FixedSizeListArray::try_new_from_values( + Float32Array::from(vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + EDGE_DIM as i32, + ) + .unwrap(), + ); + let ivf_params = IvfBuildParams::try_with_centroids(1, centroids).unwrap(); + let params = VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + ivf_params, + RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), + ); + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + let index = dataset + .open_vector_index("vector", &indices[0].uuid, &NoOpMetricsCollector) + .await + .unwrap(); + assert_eq!(index.partition_size(0), 4); + let query = Query { + column: "vector".to_string(), + key: Arc::new(Float32Array::from(vec![ + 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + ])), + k: 4, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 1, + maximum_nprobes: Some(1), + ef: None, + refine_factor: None, + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: ApproxMode::Accurate, + }; + let scored = index + .score_partition_candidates(0, &query, &[0, 1, 2, 3], &NoOpMetricsCollector) + .await + .unwrap(); + assert_eq!(scored.num_rows(), 4); + assert!( + scored[DIST_COL] + .as_primitive::() + .values() + .iter() + .all(|distance| distance.is_finite()) + ); + assert_eq!( + scored[ROW_ID] + .as_primitive::() + .values() + .iter() + .copied() + .collect::>() + .len(), + 4 + ); + } + // These queries probe every partition, so recall here measures RaBitQ quantization // error alone. At 1 bit per dimension it averages ~0.67 on this uniformly random, // L2-normalized data, and each build draws a fresh random rotation, so no bar worth diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 3f10653cf32..0e9f94f7c8e 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -47,12 +47,14 @@ use lance_datafusion::utils::{ DELTAS_SEARCHED_METRIC, ExecutionPlanMetricsSetExt, FIND_PARTITIONS_ELAPSED_METRIC, PARTITIONS_RANKED_METRIC, PARTITIONS_SEARCHED_METRIC, }; +use lance_index::IndexType; use lance_index::metrics::MetricsCollector; use lance_index::prefilter::PreFilter; use lance_index::vector::DIST_Q_C_COLUMN; +use lance_index::vector::graph::OrderedFloat; use lance_index::vector::{ - DIST_COL, INDEX_UUID_COLUMN, PART_ID_COLUMN, PartitionSearchControl, Query, VectorIndex, - flat::compute_distance, + ApproxMode, DIST_COL, INDEX_UUID_COLUMN, PART_ID_COLUMN, PartitionSearchControl, Query, + VectorIndex, flat::compute_distance, }; use lance_linalg::distance::DistanceType; use lance_linalg::kernels::normalize_arrow; @@ -1241,6 +1243,7 @@ pub fn new_knn_exec( prefilter_source: PreFilterSource, overlay_block: Option, external_mask: Option>, + quantized_refine_factor: Option, ) -> Result> { let ivf_node = ANNIvfPartitionExec::try_new( dataset.clone(), @@ -1261,6 +1264,9 @@ pub fn new_knn_exec( if external_mask.is_some() { sub_index = sub_index.with_external_mask(external_mask); } + if let Some(factor) = quantized_refine_factor { + sub_index = sub_index.with_quantized_refine_factor(factor)?; + } Ok(Arc::new(sub_index)) } @@ -1532,6 +1538,9 @@ pub struct ANNIvfSubIndexExec { properties: Arc, metrics: ExecutionPlanMetricsSet, + + /// Coarse-candidate overfetch factor for RQ Fast-to-Accurate refinement. + quantized_refine_factor: Option, } impl ANNIvfSubIndexExec { @@ -1564,6 +1573,7 @@ impl ANNIvfSubIndexExec { external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), + quantized_refine_factor: None, }) } @@ -1581,6 +1591,17 @@ impl ANNIvfSubIndexExec { self } + /// Refine RQ candidates from one-bit Fast scores to multi-bit Accurate scores. + pub fn with_quantized_refine_factor(mut self, factor: u32) -> Result { + if factor == 0 { + return Err(Error::invalid_input( + "quantized refine factor must be greater than zero", + )); + } + self.quantized_refine_factor = Some(factor); + Ok(self) + } + /// Returns a reference to the vector query. pub fn query(&self) -> &Query { &self.query @@ -1617,8 +1638,12 @@ impl DisplayAs for ANNIvfSubIndexExec { self.indices[0].name, self.query.k * self.query.refine_factor.unwrap_or(1) as usize, self.indices.len(), - metric_str - ) + metric_str, + )?; + if let Some(factor) = self.quantized_refine_factor { + write!(f, ", quantized_refine_factor={factor}")?; + } + Ok(()) } DisplayFormatType::TreeRender => { write!( @@ -1627,8 +1652,12 @@ impl DisplayAs for ANNIvfSubIndexExec { self.indices[0].name, self.query.k * self.query.refine_factor.unwrap_or(1) as usize, self.indices.len(), - metric_str - ) + metric_str, + )?; + if let Some(factor) = self.quantized_refine_factor { + write!(f, "\nquantized_refine_factor={factor}")?; + } + Ok(()) } } } @@ -1828,6 +1857,84 @@ fn restrict_to_segment( .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) } +struct RestrictedPreFilter { + inner: Arc, + restriction: Arc, +} + +#[async_trait::async_trait] +impl PreFilter for RestrictedPreFilter { + async fn wait_for_ready(&self) -> Result<()> { + self.inner.wait_for_ready().await + } + + fn is_empty(&self) -> bool { + false + } + + fn mask(&self) -> Arc { + Arc::new(self.inner.mask().as_ref().clone() & self.restriction.as_ref().clone()) + } + + fn filter_row_ids<'a>(&self, row_ids: Box + 'a>) -> Vec { + self.mask().selected_indices(row_ids) + } +} + +fn quantized_top_k(batch: &RecordBatch, query: &Query) -> Result { + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("quantized refinement result has no row-id column"))? + .as_primitive::(); + let distances = batch + .column_by_name(DIST_COL) + .ok_or_else(|| Error::internal("quantized refinement result has no distance column"))? + .as_primitive::(); + if row_ids.len() != distances.len() { + return Err(Error::internal(format!( + "quantized refinement returned {} row ids and {} distances", + row_ids.len(), + distances.len() + ))); + } + + let lower_bound = query.lower_bound.unwrap_or(f32::MIN); + let upper_bound = query.upper_bound.unwrap_or(f32::MAX); + let mut top_results = BinaryHeap::with_capacity(query.k); + for row_idx in 0..row_ids.len() { + if row_ids.is_null(row_idx) || distances.is_null(row_idx) { + continue; + } + let distance = distances.value(row_idx); + if distance.is_nan() || distance < lower_bound || distance >= upper_bound { + continue; + } + let node = (OrderedFloat(distance), row_ids.value(row_idx)); + if top_results.len() < query.k { + top_results.push(node); + } else if top_results.peek().is_some_and(|farthest| farthest > &node) { + top_results.pop(); + top_results.push(node); + } + } + + let mut ordered = top_results.into_vec(); + ordered.sort_unstable(); + let mut row_ids = Vec::with_capacity(ordered.len()); + let mut distances = Vec::with_capacity(ordered.len()); + for (distance, row_id) in ordered { + row_ids.push(row_id); + distances.push(distance.0); + } + Ok(RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(distances)), + Arc::new(UInt64Array::from(row_ids)), + ], + )?) +} + fn effective_query_parallelism( query: &Query, index: &dyn VectorIndex, @@ -1865,12 +1972,73 @@ impl ANNIvfSubIndexExec { pre_filter: Arc, metrics: Arc, seg_mask: Option>, + quantized_refine_factor: Option, ) -> DataFusionResult { - let batch = index - .search_in_partition(part_id, &query, pre_filter, &metrics.index_metrics) - .map_err(|e| DataFusionError::Execution(format!("Failed to calculate KNN: {}", e))) - .await?; - let batch = restrict_to_segment(batch, seg_mask.as_deref())?; + let batch = if let Some(factor) = quantized_refine_factor { + if index.index_type() != IndexType::IvfRq { + return Err(DataFusionError::NotImplemented(format!( + "quantized refinement requires an IVF_RQ flat index, got {}", + index.index_type() + ))); + } + let candidate_limit = query.k.checked_mul(factor as usize).ok_or_else(|| { + DataFusionError::Execution(format!( + "quantized refinement candidate count overflows: k={} factor={factor}", + query.k + )) + })?; + let mut coarse_query = query.clone(); + coarse_query.approx_mode = ApproxMode::Fast; + coarse_query.refine_factor = None; + coarse_query.lower_bound = None; + coarse_query.upper_bound = None; + let candidate_prefilter: Arc = match &seg_mask { + Some(restriction) => Arc::new(RestrictedPreFilter { + inner: pre_filter, + restriction: restriction.clone(), + }), + None => pre_filter, + }; + let candidates = index + .search_in_partition_with_candidates( + part_id, + &coarse_query, + candidate_prefilter, + &metrics.index_metrics, + candidate_limit, + ) + .await + .map_err(|e| { + DataFusionError::Execution(format!( + "failed to discover quantized refinement candidates: {e}" + )) + })?; + let offsets = candidates + .candidates + .iter() + .map(|candidate| candidate.offset_in_partition) + .collect::>(); + let scored = index + .score_partition_candidates(part_id, &query, &offsets, &metrics.index_metrics) + .await + .map_err(|e| { + DataFusionError::Execution(format!( + "failed to score quantized refinement candidates: {e}" + )) + })?; + let scored = restrict_to_segment(scored, seg_mask.as_deref())?; + quantized_top_k(&scored, &query).map_err(|e| { + DataFusionError::Execution(format!( + "failed to select quantized refinement results: {e}" + )) + })? + } else { + let batch = index + .search_in_partition(part_id, &query, pre_filter, &metrics.index_metrics) + .map_err(|e| DataFusionError::Execution(format!("Failed to calculate KNN: {e}"))) + .await?; + restrict_to_segment(batch, seg_mask.as_deref())? + }; metrics.baseline_metrics.record_output(batch.num_rows()); Ok(batch) } @@ -1910,6 +2078,7 @@ impl ANNIvfSubIndexExec { state: Arc, target_partitions: usize, seg_mask: Option>, + quantized_refine_factor: Option, ) -> impl Stream> { let stream = futures::stream::once(async move { let max_nprobes = query @@ -1950,6 +2119,9 @@ impl ANNIvfSubIndexExec { if let Some(max_results) = max_results && found_so_far < max_results && max_results <= query.k + // Quantized refinement must score every returned row and apply its + // distance bounds, so it cannot emit unscored rows from this shortcut. + && quantized_refine_factor.is_none() { // In this case there are fewer than k results matching the prefilter so // just return the prefilter ids and don't bother searching any further @@ -2003,7 +2175,7 @@ impl ANNIvfSubIndexExec { let query_parallelism = effective_query_parallelism(&query, index.as_ref(), target_partitions); - if query_parallelism <= 1 { + if query_parallelism <= 1 && quantized_refine_factor.is_none() { return stream::once(async move { let prefilter: Arc = prefilter; let index_metrics: Arc = @@ -2052,6 +2224,7 @@ impl ANNIvfSubIndexExec { let state = state.clone(); let index = index.clone(); let seg_mask = seg_mask.clone(); + let quantized_refine_factor = quantized_refine_factor; async move { metrics.partitions_searched.add(1); let batch = Self::search_partition( @@ -2061,6 +2234,7 @@ impl ANNIvfSubIndexExec { pre_filter, metrics, seg_mask, + quantized_refine_factor, ) .await?; state.record_late_batch(batch.num_rows()); @@ -2088,12 +2262,13 @@ impl ANNIvfSubIndexExec { state: Arc, target_partitions: usize, seg_mask: Option>, + quantized_refine_factor: Option, ) -> impl Stream> { let minimum_nprobes = query.minimum_nprobes.min(partitions.len()); let query_parallelism = effective_query_parallelism(&query, index.as_ref(), target_partitions); - if query_parallelism <= 1 { + if query_parallelism <= 1 && quantized_refine_factor.is_none() { metrics.partitions_searched.add(minimum_nprobes); return stream::once(async move { let prefilter: Arc = prefilter; @@ -2140,6 +2315,7 @@ impl ANNIvfSubIndexExec { let pre_filter = prefilter.clone(); let state = state.clone(); let seg_mask = seg_mask.clone(); + let quantized_refine_factor = quantized_refine_factor; async move { let batch = Self::search_partition( index, @@ -2148,6 +2324,7 @@ impl ANNIvfSubIndexExec { pre_filter, metrics, seg_mask, + quantized_refine_factor, ) .await?; state.record_batch(&batch); @@ -2214,6 +2391,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), + quantized_refine_factor: self.quantized_refine_factor, } } else { return Err(DataFusionError::Internal( @@ -2232,6 +2410,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { let schema = self.schema(); let target_partitions = context.session_config().target_partitions(); let query = self.query.clone(); + let quantized_refine_factor = self.quantized_refine_factor; let ds = self.dataset.clone(); let column = self.query.column.clone(); let indices = self.indices.clone(); @@ -2324,6 +2503,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { let indices_by_uuid = indices_by_uuid.clone(); let state = state.clone(); let segment_bitmaps = segment_bitmaps.clone(); + let quantized_refine_factor = quantized_refine_factor; let mut query = query.clone(); async move { let index_metadata = indices_by_uuid.get(&index_uuid).ok_or_else(|| { @@ -2377,6 +2557,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { state.clone(), target_partitions, seg_mask.clone(), + quantized_refine_factor, ); let late_search = Self::late_search( raw_index.clone(), @@ -2389,6 +2570,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { state, target_partitions, seg_mask, + quantized_refine_factor, ); DataFusionResult::Ok(early_search.chain(late_search).boxed()) } @@ -3883,6 +4065,7 @@ mod tests { state, usize::MAX, None, + None, ) .try_collect::>() .await @@ -3934,6 +4117,7 @@ mod tests { state, usize::MAX, None, + None, ) .try_collect::>() .await @@ -3998,6 +4182,7 @@ mod tests { state.clone(), usize::MAX, None, + None, ) .try_collect::>() .await @@ -4087,6 +4272,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask.clone()), + None, ) .try_collect::>() .await @@ -4120,6 +4306,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + None, ) .try_collect::>() .await @@ -4182,6 +4369,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask.clone()), + None, ) .try_collect::>() .await @@ -4199,6 +4387,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + None, ) .try_collect::>() .await @@ -4232,6 +4421,7 @@ mod tests { state.clone(), usize::MAX, None, + None, ) .try_collect::>(); @@ -4251,6 +4441,7 @@ mod tests { state, usize::MAX, None, + None, ) .try_collect::>();