From b278f3a786e5c42e232e8241ac74d725105637c3 Mon Sep 17 00:00:00 2001 From: Sergey Troshkov Date: Tue, 25 Aug 2026 11:33:47 +0700 Subject: [PATCH 1/4] feat: add quantized vector results cache --- Cargo.lock | 1 + java/lance-jni/Cargo.lock | 1 + python/Cargo.lock | 1 + rust/lance-index/src/vector.rs | 51 + rust/lance-index/src/vector/flat/index.rs | 266 +++-- rust/lance-index/src/vector/v3/subindex.rs | 26 +- rust/lance/Cargo.toml | 1 + rust/lance/src/index/vector/ivf/v2.rs | 588 +++++++++- rust/lance/src/io/exec.rs | 1 + rust/lance/src/io/exec/knn.rs | 1154 +++++++++++++++++- rust/lance/src/io/exec/knn_results_cache.rs | 1162 +++++++++++++++++++ rust/lance/src/session/index_caches.rs | 861 +++++++++++++- 12 files changed, 4020 insertions(+), 93 deletions(-) create mode 100644 rust/lance/src/io/exec/knn_results_cache.rs diff --git a/Cargo.lock b/Cargo.lock index f420aebfb36..484cd942489 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4422,6 +4422,7 @@ dependencies = [ "aws-credential-types", "aws-sdk-dynamodb", "aws-sdk-s3", + "blake3", "byteorder", "bytes", "chrono", diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 96c25e25691..b2c93243673 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -3805,6 +3805,7 @@ dependencies = [ "async-recursion", "async-trait", "async_cell", + "blake3", "byteorder", "bytes", "chrono", diff --git a/python/Cargo.lock b/python/Cargo.lock index 0dfc7892cf9..cb333d4326a 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -3973,6 +3973,7 @@ dependencies = [ "async-trait", "async_cell", "aws-sdk-dynamodb", + "blake3", "byteorder", "bytes", "chrono", diff --git a/rust/lance-index/src/vector.rs b/rust/lance-index/src/vector.rs index 1192b2f5afb..35324413323 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,38 @@ 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. The initial results-cache + /// integration uses it only for flat IVF_RQ and IVF_SQ 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 cached 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/Cargo.toml b/rust/lance/Cargo.toml index 417d41a5fe1..d4b7fc4e342 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -45,6 +45,7 @@ arrow-schema = { workspace = true } arrow-select = { workspace = true } async-recursion.workspace = true async-trait.workspace = true +blake3.workspace = true byteorder.workspace = true bytes.workspace = true chrono.workspace = true diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index 730386bc756..3d3e9a6d8a6 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,186 @@ 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" + || !matches!( + Q::quantization_type(), + QuantizationType::Rabit | QuantizationType::Scalar + ) + { + 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" + || !matches!( + Q::quantization_type(), + QuantizationType::Rabit | QuantizationType::Scalar + ) + { + 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 +3157,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 +3175,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 +3188,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 +3211,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 +3576,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, @@ -6088,6 +6331,329 @@ mod tests { test_index_impl::(params, nlist, 0.75, -1.0..1.0, None).await; } + #[rstest] + #[case::rq4_l2(DistanceType::L2, Some(4), false, DIM)] + #[case::rq4_cosine(DistanceType::Cosine, Some(4), false, DIM)] + #[case::rq4_dot(DistanceType::Dot, Some(4), false, DIM)] + #[case::rq8_l2(DistanceType::L2, Some(8), false, DIM)] + #[case::rq8_cosine(DistanceType::Cosine, Some(8), false, DIM)] + #[case::rq8_dot(DistanceType::Dot, Some(8), false, DIM)] + #[case::sq8_l2(DistanceType::L2, None, true, DIM)] + #[case::sq8_cosine(DistanceType::Cosine, None, true, DIM)] + #[case::sq8_dot(DistanceType::Dot, None, true, DIM)] + #[case::rq4_dot_768(DistanceType::Dot, Some(4), false, 768)] + #[case::sq8_l2_3072(DistanceType::L2, None, true, 3072)] + #[case::rq8_cosine_4096(DistanceType::Cosine, Some(8), false, 4096)] + #[tokio::test] + async fn test_score_partition_candidates_matches_accurate_native_search( + #[case] distance_type: DistanceType, + #[case] rq_num_bits: Option, + #[case] has_negative_values: bool, + #[case] dimension: usize, + ) { + let test_dir = TempStrDir::default(); + let test_uri = test_dir.as_str(); + let range = if has_negative_values { + -1.0..1.0 + } else { + 0.0..1.0 + }; + 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, range).await; + let ivf_params = IvfBuildParams::new(4); + let params = if let Some(num_bits) = rq_num_bits { + VectorIndexParams::with_ivf_rq_params( + distance_type, + ivf_params, + RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), + ) + } else { + VectorIndexParams::with_ivf_sq_params( + distance_type, + ivf_params, + SQBuildParams::default(), + ) + }; + dataset + .create_index(&["vector"], IndexType::Vector, None, ¶ms, true) + .await + .unwrap(); + + let indices = dataset.load_indices().await.unwrap(); + if rq_num_bits.is_none() { + let obj_store = Arc::new(ObjectStore::local()); + let scheduler = ScanScheduler::new(obj_store, SchedulerConfig::default_for_testing()); + let metadata = get_sq_metadata(&dataset, scheduler, &indices[0].uuid.to_string()).await; + assert!( + metadata.bounds.start < 0.0 && metadata.bounds.end > 0.0, + "SQ bounds should be learned from both negative and positive values, got {:?}", + metadata.bounds + ); + } + 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::sq8(None)] + #[case::rq4(Some(4))] + #[case::rq8(Some(8))] + #[tokio::test] + async fn test_selected_scoring_excludes_null_and_non_finite_stored_vectors( + #[case] rq_num_bits: Option, + ) { + 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 = if let Some(num_bits) = rq_num_bits { + VectorIndexParams::with_ivf_rq_params( + DistanceType::L2, + ivf_params, + RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), + ) + } else { + VectorIndexParams::with_ivf_sq_params( + DistanceType::L2, + ivf_params, + SQBuildParams::default(), + ) + }; + 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.rs b/rust/lance/src/io/exec.rs index d37b58a238e..f610f1b7f55 100644 --- a/rust/lance/src/io/exec.rs +++ b/rust/lance/src/io/exec.rs @@ -15,6 +15,7 @@ pub mod filtered_read; pub mod filtered_read_proto; pub mod fts; pub(crate) mod knn; +mod knn_results_cache; mod optimizer; mod projection; mod pushdown_scan; diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 76fcd62f5c9..c38f1982ef7 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -69,6 +69,7 @@ use crate::index::vector::utils::{get_vector_type, validate_distance_type_for}; use crate::{Error, Result}; use lance_arrow::*; +use super::knn_results_cache; use super::row_addr_mask::MaskAndLoader; use super::utils::{ FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedRecordBatchStreamAdapter, PreFilterSource, @@ -108,14 +109,21 @@ impl AnnPartitionMetrics { pub struct AnnIndexMetrics { index_metrics: IndexMetrics, partitions_searched: Count, + results_cache_hits: Count, + results_cache_misses: Count, baseline_metrics: BaselineMetrics, } +const VECTOR_RESULTS_CACHE_HITS_METRIC: &str = "vector_results_cache_hits"; +const VECTOR_RESULTS_CACHE_MISSES_METRIC: &str = "vector_results_cache_misses"; + impl AnnIndexMetrics { pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { Self { index_metrics: IndexMetrics::new(metrics, partition), partitions_searched: metrics.new_count(PARTITIONS_SEARCHED_METRIC, partition), + results_cache_hits: metrics.new_count(VECTOR_RESULTS_CACHE_HITS_METRIC, partition), + results_cache_misses: metrics.new_count(VECTOR_RESULTS_CACHE_MISSES_METRIC, partition), baseline_metrics: BaselineMetrics::new(metrics, partition), } } @@ -1532,6 +1540,9 @@ pub struct ANNIvfSubIndexExec { properties: Arc, metrics: ExecutionPlanMetricsSet, + + #[cfg(test)] + results_cache_enabled_override: Option, } impl ANNIvfSubIndexExec { @@ -1564,6 +1575,8 @@ impl ANNIvfSubIndexExec { external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), + #[cfg(test)] + results_cache_enabled_override: None, }) } @@ -1600,6 +1613,29 @@ impl ANNIvfSubIndexExec { pub fn prefilter_source(&self) -> &PreFilterSource { &self.prefilter_source } + + fn is_results_cache_enabled(&self) -> bool { + #[cfg(test)] + if let Some(is_enabled) = self.results_cache_enabled_override { + return is_enabled; + } + knn_results_cache::is_enabled() + } + + #[cfg(test)] + fn copy_with_results_cache_enabled_for_testing(&self, is_enabled: bool) -> Result { + let mut copy = Self::try_new( + self.input.clone(), + self.dataset.clone(), + self.indices.clone(), + self.query.clone(), + self.prefilter_source.clone(), + )?; + copy.overlay_block = self.overlay_block.clone(); + copy.external_mask = self.external_mask.clone(); + copy.results_cache_enabled_override = Some(is_enabled); + Ok(copy) + } } impl DisplayAs for ANNIvfSubIndexExec { @@ -2214,6 +2250,8 @@ impl ExecutionPlan for ANNIvfSubIndexExec { external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), + #[cfg(test)] + results_cache_enabled_override: self.results_cache_enabled_override, } } else { return Err(DataFusionError::Internal( @@ -2250,6 +2288,19 @@ impl ExecutionPlan for ANNIvfSubIndexExec { HashMap::new() }); let prefilter_source = self.prefilter_source.clone(); + let has_prefilter = + !matches!(prefilter_source, PreFilterSource::None) || self.external_mask.is_some(); + let has_overlay = self.overlay_block.is_some(); + let results_cache_enabled = self.is_results_cache_enabled(); + let index_metadata_by_uuid = results_cache_enabled.then(|| { + Arc::new( + indices + .iter() + .cloned() + .map(|metadata| (metadata.uuid, metadata)) + .collect::>(), + ) + }); let metrics = Arc::new(AnnIndexMetrics::new(&self.metrics, partition)); let metrics_clone = metrics.clone(); let timer = Instant::now(); @@ -2324,6 +2375,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { let indices_by_uuid = indices_by_uuid.clone(); let state = state.clone(); let segment_bitmaps = segment_bitmaps.clone(); + let index_metadata_by_uuid = index_metadata_by_uuid.clone(); let mut query = query.clone(); async move { let index_metadata = indices_by_uuid.get(&index_uuid).ok_or_else(|| { @@ -2366,6 +2418,68 @@ impl ExecutionPlan for ANNIvfSubIndexExec { } None => None, }; + if results_cache_enabled + && let Some(index_metadata) = index_metadata_by_uuid + .as_ref() + .and_then(|metadata| metadata.get(&index_uuid)) + && let Some(identity) = knn_results_cache::identity_for_query( + ds.as_ref(), + index_metadata, + raw_index.as_ref(), + &query, + part_ids.values(), + has_prefilter, + has_overlay, + ) + .await + { + let parallelism = effective_query_parallelism( + &query, + raw_index.as_ref(), + target_partitions, + ); + let index_metrics: Arc = + Arc::new(metrics.index_metrics.clone()); + match knn_results_cache::search( + knn_results_cache::ResultsCacheSearchParams { + cache: &ds.index_cache.0, + identity, + index: raw_index.clone(), + query: &query, + partitions: part_ids.clone(), + centroid_distances: q_c_dists.clone(), + prefilter: pre_filter.clone(), + segment_mask: seg_mask.clone(), + metrics: index_metrics, + parallelism, + }, + ) + .await + { + Ok(cache_result) => { + if cache_result.was_hit { + metrics.results_cache_hits.add(1); + } else { + metrics.results_cache_misses.add(1); + metrics.partitions_searched.add(part_ids.len()); + } + metrics + .baseline_metrics + .record_output(cache_result.batch.num_rows()); + return DataFusionResult::Ok( + stream::once(async move { Ok(cache_result.batch) }).boxed(), + ); + } + Err(error) => { + // The feature is experimental and best-effort. Any cache-only + // failure falls through to the existing ANN path. + metrics.results_cache_misses.add(1); + log::debug!( + "falling back after vector results cache failure: {error}" + ); + } + } + } let early_search = Self::initial_search( raw_index.clone(), @@ -3005,19 +3119,27 @@ mod tests { }; use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; use async_trait::async_trait; + use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::error::Result as DataFusionResult; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::tempfile::TempStrDir; - use lance_datafusion::exec::{ExecutionStatsCallback, ExecutionSummaryCounts}; + use lance_datafusion::exec::{ExecutionStatsCallback, ExecutionSummaryCounts, execute_plan}; use lance_datafusion::utils::FIND_PARTITIONS_ELAPSED_METRIC; use lance_datagen::{BatchCount, RowCount, array}; + use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; + use lance_index::metrics::NoOpMetricsCollector; use lance_index::optimize::OptimizeOptions; + use lance_index::scalar::ScalarIndexParams; + use lance_index::vector::bq::{RQBuildParams, RQRotationType}; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; use lance_index::vector::quantizer::QuantizationType; + use lance_index::vector::sq::builder::SQBuildParams; use lance_index::vector::v3::subindex::SubIndexType; - use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, PreparedPartitionSearchHandle}; + use lance_index::vector::{ + ApproxMode, DEFAULT_QUERY_PARALLELISM, PreparedPartitionSearchHandle, + }; use lance_index::{Index, IndexType}; use lance_io::traits::Reader; use lance_linalg::distance::MetricType; @@ -3027,10 +3149,12 @@ mod tests { use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; + use crate::dataset::optimize::{CompactionOptions, compact_files}; use crate::dataset::{WriteMode, WriteParams}; use crate::index::vector::VectorIndexParams; use crate::index::vector::ivf::v2::STREAMING_SEARCH_BATCH_SIZE; use crate::io::exec::testing::TestingExec; + use crate::session::index_caches::{CachedVectorCandidate, VectorResultsCacheEntry}; fn base_query() -> Query { Query { @@ -4789,6 +4913,288 @@ mod tests { } } + struct ResultsCacheTestFixture { + dataset: Dataset, + centroids: ArrayRef, + query: ArrayRef, + metric_type: MetricType, + index_params: VectorIndexParams, + _tmp_dir: TempStrDir, + } + + impl ResultsCacheTestFixture { + const NUM_PARTITIONS: usize = 4; + const K: usize = 8; + const ROWS_PER_FRAGMENT: usize = 32; + const NUM_FRAGMENTS: usize = 4; + const INITIAL_ROWS: usize = Self::ROWS_PER_FRAGMENT * Self::NUM_FRAGMENTS; + const RANK_SCALE_STEP: f32 = 0.01; + const RANK_ORTHOGONAL_STEP: f32 = 0.005; + + fn data_batch(start_row: usize, num_rows: usize) -> RecordBatch { + let mut vectors = Vec::with_capacity(num_rows * 8); + for row in start_row..start_row + num_rows { + let centroid_id = row % Self::NUM_PARTITIONS; + let rank = row / Self::NUM_PARTITIONS; + // Both changes make larger ranks worse under L2, cosine, and dot, + // giving the quantized recall checks a reproducible top-k boundary. + let scale = 1.0 - rank as f32 * Self::RANK_SCALE_STEP; + let orthogonal_offset = rank as f32 * Self::RANK_ORTHOGONAL_STEP; + let mut vector = [0.0; 8]; + match centroid_id { + 0 => vector[0] = scale, + 1 => vector[1] = scale, + 2 => vector[0] = -scale, + 3 => vector[1] = -scale, + _ => unreachable!(), + } + vector[centroid_id + 2] = orthogonal_offset; + vectors.extend_from_slice(&vector); + } + let vectors = Arc::new( + FixedSizeListArray::try_new_from_values(Float32Array::from(vectors), 8).unwrap(), + ); + let row_ids = Arc::new(UInt64Array::from_iter_values( + start_row as u64..(start_row + num_rows) as u64, + )); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("vector", vectors.data_type().clone(), false), + ArrowField::new("row", DataType::UInt64, false), + ])); + RecordBatch::try_new(schema, vec![vectors, row_ids]).unwrap() + } + + async fn new(rq_num_bits: Option) -> Self { + Self::with_metric(rq_num_bits, MetricType::L2).await + } + + async fn with_metric(rq_num_bits: Option, metric_type: MetricType) -> Self { + let tmp_dir = TempStrDir::default(); + let centroids: ArrayRef = 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, // +x + 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // +y + -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // -x + 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // -y + ]), + 8, + ) + .unwrap(), + ); + let batches = (0..Self::NUM_FRAGMENTS) + .map(|fragment| { + Ok(Self::data_batch( + fragment * Self::ROWS_PER_FRAGMENT, + Self::ROWS_PER_FRAGMENT, + )) + }) + .collect::>(); + let schema = batches[0].as_ref().unwrap().schema(); + let reader = RecordBatchIterator::new(batches, schema); + let mut dataset = Dataset::write( + reader, + tmp_dir.as_str(), + Some(WriteParams { + max_rows_per_file: Self::ROWS_PER_FRAGMENT, + max_rows_per_group: Self::ROWS_PER_FRAGMENT, + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!(dataset.get_fragments().len(), Self::NUM_FRAGMENTS); + + let ivf_params = IvfBuildParams::try_with_centroids( + Self::NUM_PARTITIONS, + Arc::new(centroids.as_fixed_size_list().clone()), + ) + .unwrap(); + let index_params = if let Some(num_bits) = rq_num_bits { + VectorIndexParams::with_ivf_rq_params( + metric_type, + ivf_params, + RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), + ) + } else { + VectorIndexParams::with_ivf_sq_params( + metric_type, + ivf_params, + SQBuildParams::default(), + ) + }; + dataset + .create_index(&["vector"], IndexType::Vector, None, &index_params, false) + .await + .unwrap(); + + let dataset = Dataset::open(tmp_dir.as_str()).await.unwrap(); + let query = centroids.as_fixed_size_list().value(0); + Self { + dataset, + centroids, + query, + metric_type, + index_params, + _tmp_dir: tmp_dir, + } + } + + async fn vector_index_uuids(&self) -> Vec { + self.dataset + .load_indices() + .await + .unwrap() + .iter() + .filter(|index| index.name != FRAG_REUSE_INDEX_NAME) + .map(|index| index.uuid) + .collect() + } + + async fn fragment_reuse_uuid(&self) -> Option { + self.dataset + .load_indices() + .await + .unwrap() + .iter() + .find(|index| index.name == FRAG_REUSE_INDEX_NAME) + .map(|index| index.uuid) + } + + async fn append_index_segment(&mut self) -> Vec { + let batch = Self::data_batch(Self::INITIAL_ROWS, Self::ROWS_PER_FRAGMENT); + let schema = batch.schema(); + self.dataset + .append(RecordBatchIterator::new([Ok(batch)], schema), None) + .await + .unwrap(); + self.dataset + .optimize_indices(&OptimizeOptions::append()) + .await + .unwrap(); + self.vector_index_uuids().await + } + + async fn delete_query_centroid_rows(&mut self) { + self.dataset.delete("row % 4 = 0").await.unwrap(); + } + + async fn compact_with_fragment_reuse(&mut self) -> Uuid { + let metrics = compact_files( + &mut self.dataset, + CompactionOptions { + target_rows_per_fragment: Self::ROWS_PER_FRAGMENT * 2, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert!(metrics.fragments_removed > 0); + assert!(metrics.fragments_added > 0); + self.fragment_reuse_uuid() + .await + .expect("deferred compaction should create a fragment-reuse index") + } + + async fn replace_with_unsupported_flat_index(&mut self) { + let ivf_params = IvfBuildParams::try_with_centroids( + Self::NUM_PARTITIONS, + Arc::new(self.centroids.as_fixed_size_list().clone()), + ) + .unwrap(); + let index_params = VectorIndexParams::with_ivf_flat_params(MetricType::L2, ivf_params); + self.dataset + .create_index(&["vector"], IndexType::Vector, None, &index_params, true) + .await + .unwrap(); + } + + async fn create_row_scalar_index(&mut self) { + self.dataset + .create_index( + &["row"], + IndexType::Scalar, + None, + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + } + + async fn insert_out_of_range_results_cache_entry(&self) { + let indices = self.dataset.load_indices().await.unwrap(); + let index_metadata = indices + .iter() + .find(|index| index.name != FRAG_REUSE_INDEX_NAME) + .unwrap(); + let index = self + .dataset + .open_vector_index("vector", &index_metadata.uuid, &NoOpMetricsCollector) + .await + .unwrap(); + let query = Query { + column: "vector".to_string(), + key: self.query.clone(), + k: Self::K, + lower_bound: None, + upper_bound: None, + minimum_nprobes: Self::NUM_PARTITIONS, + maximum_nprobes: Some(Self::NUM_PARTITIONS), + ef: None, + refine_factor: None, + metric_type: None, + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: ApproxMode::Accurate, + }; + let (partitions, _) = index.find_partitions(&query).unwrap(); + let identity = knn_results_cache::identity_for_query( + &self.dataset, + index_metadata, + index.as_ref(), + &query, + partitions.values(), + false, + false, + ) + .await + .unwrap(); + let entry = VectorResultsCacheEntry::try_new( + identity.clone(), + vec![CachedVectorCandidate::new( + identity.partition_ids()[0], + u32::MAX, + )], + ) + .unwrap(); + self.dataset + .index_cache + .0 + .insert_with_key(&identity, Arc::new(entry)) + .await; + } + + async fn replace_index(&mut self) -> (Uuid, Uuid) { + let old_uuid = self.dataset.load_indices().await.unwrap()[0].uuid; + self.dataset + .create_index( + &["vector"], + IndexType::Vector, + None, + &self.index_params, + true, + ) + .await + .unwrap(); + let new_uuid = self.dataset.load_indices().await.unwrap()[0].uuid; + (old_uuid, new_uuid) + } + } + #[derive(Default)] struct StatsHolder { pub collected_stats: Arc>>, @@ -4818,6 +5224,750 @@ mod tests { ); } + #[derive(Clone, Copy, Debug)] + enum ResultsCacheQueryMode { + Supported, + Prefilter, + ScalarPrefilter, + DistanceRange, + AdaptiveNprobes, + Overlay, + } + + fn with_results_cache_test_options( + plan: Arc, + is_enabled: bool, + has_overlay: bool, + ) -> Arc { + plan.transform_down(|node| { + let Some(ann) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let mut replacement = ann + .copy_with_results_cache_enabled_for_testing(is_enabled) + .unwrap(); + if has_overlay { + replacement = replacement.with_overlay_block(RowAddrMask::all_rows()); + } + let replacement: Arc = Arc::new(replacement); + Ok(Transformed::yes(replacement)) + }) + .unwrap() + .data + } + + async fn run_results_cache_query( + fixture: &ResultsCacheTestFixture, + is_cache_enabled: bool, + refine_factor: Option, + ) -> (Vec, ExecutionSummaryCounts) { + run_results_cache_query_with_mode( + fixture, + is_cache_enabled, + refine_factor, + ResultsCacheQueryMode::Supported, + ) + .await + } + + async fn run_results_cache_query_with_mode( + fixture: &ResultsCacheTestFixture, + is_cache_enabled: bool, + refine_factor: Option, + query_mode: ResultsCacheQueryMode, + ) -> (Vec, ExecutionSummaryCounts) { + let (batch, stats) = run_results_cache_batch_with_mode( + fixture, + is_cache_enabled, + refine_factor, + query_mode, + ApproxMode::Accurate, + ) + .await; + ( + batch[ROW_ID].as_primitive::().values().to_vec(), + stats, + ) + } + + async fn run_results_cache_batch_with_mode( + fixture: &ResultsCacheTestFixture, + is_cache_enabled: bool, + refine_factor: Option, + query_mode: ResultsCacheQueryMode, + approx_mode: ApproxMode, + ) -> (RecordBatch, ExecutionSummaryCounts) { + let stats_holder = StatsHolder::default(); + let mut scanner = fixture.dataset.scan(); + scanner + .nearest("vector", fixture.query.as_ref(), ResultsCacheTestFixture::K) + .unwrap() + .distance_metric(fixture.metric_type) + .approx_mode(approx_mode); + match query_mode { + ResultsCacheQueryMode::Supported | ResultsCacheQueryMode::Overlay => { + scanner.nprobes(ResultsCacheTestFixture::NUM_PARTITIONS); + } + ResultsCacheQueryMode::Prefilter | ResultsCacheQueryMode::ScalarPrefilter => { + scanner + .nprobes(ResultsCacheTestFixture::NUM_PARTITIONS) + .filter("row >= 0") + .unwrap() + .prefilter(true); + } + ResultsCacheQueryMode::DistanceRange => { + scanner + .nprobes(ResultsCacheTestFixture::NUM_PARTITIONS) + .distance_range(None, Some(f32::MAX)); + } + ResultsCacheQueryMode::AdaptiveNprobes => { + scanner + .minimum_nprobes(1) + .maximum_nprobes(ResultsCacheTestFixture::NUM_PARTITIONS); + } + } + scanner + .project(&Vec::::new()) + .unwrap() + .with_row_id() + .scan_stats_callback(stats_holder.get_setter()); + if let Some(refine_factor) = refine_factor { + scanner.refine(refine_factor); + } + + let plan = scanner.create_plan().await.unwrap(); + if matches!(query_mode, ResultsCacheQueryMode::ScalarPrefilter) { + let rendered = format!( + "{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ); + assert!( + rendered.contains("ScalarIndexQuery"), + "expected a scalar-index prefilter plan, got:\n{rendered}" + ); + } + let plan = with_results_cache_test_options( + plan, + is_cache_enabled, + matches!(query_mode, ResultsCacheQueryMode::Overlay), + ); + let batches = execute_plan(plan, scanner.execution_options()) + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); + (batch, stats_holder.consume()) + } + + async fn exact_result_batch(fixture: &ResultsCacheTestFixture) -> RecordBatch { + fixture + .dataset + .scan() + .nearest("vector", fixture.query.as_ref(), ResultsCacheTestFixture::K) + .unwrap() + .distance_metric(fixture.metric_type) + .use_index(false) + .project(&Vec::::new()) + .unwrap() + .with_row_id() + .try_into_batch() + .await + .unwrap() + } + + async fn exact_results(fixture: &ResultsCacheTestFixture) -> Vec { + let batch = exact_result_batch(fixture).await; + batch[ROW_ID].as_primitive::().values().to_vec() + } + + fn count_metric(stats: &ExecutionSummaryCounts, name: &str) -> usize { + stats.all_counts.get(name).copied().unwrap_or_default() + } + + fn recall(actual: &[u64], expected: &[u64]) -> f32 { + actual + .iter() + .filter(|row_id| expected.contains(row_id)) + .count() as f32 + / expected.len() as f32 + } + + fn assert_valid_result_set(row_ids: &[u64]) { + assert_eq!(row_ids.len(), ResultsCacheTestFixture::K); + for (position, row_id) in row_ids.iter().enumerate() { + assert!( + !row_ids[..position].contains(row_id), + "result row id {row_id} occurs more than once in {row_ids:?}" + ); + } + } + + fn assert_result_batches_match(actual: &RecordBatch, expected: &RecordBatch) { + assert_eq!(actual.num_rows(), expected.num_rows()); + let actual_row_ids = actual[ROW_ID].as_primitive::(); + let expected_row_ids = expected[ROW_ID].as_primitive::(); + let actual_distances = actual[DIST_COL].as_primitive::(); + let expected_distances = expected[DIST_COL].as_primitive::(); + for position in 0..actual.num_rows() { + assert_eq!( + actual_row_ids.value(position), + expected_row_ids.value(position) + ); + let actual_distance = actual_distances.value(position); + let expected_distance = expected_distances.value(position); + let tolerance = 1.0e-5 * (1.0 + expected_distance.abs()); + assert!( + actual_distance.is_finite() + && (actual_distance - expected_distance).abs() <= tolerance, + "distance mismatch at position {position}: actual={actual_distance}, expected={expected_distance}" + ); + } + } + + fn assert_quantized_result_batches_agree(actual: &RecordBatch, expected: &RecordBatch) { + assert_eq!(actual.num_rows(), expected.num_rows()); + let actual_row_ids = actual[ROW_ID] + .as_primitive::() + .values() + .to_vec(); + let expected_row_ids = expected[ROW_ID] + .as_primitive::() + .values() + .to_vec(); + assert_valid_result_set(&actual_row_ids); + assert_valid_result_set(&expected_row_ids); + assert!(recall(&actual_row_ids, &expected_row_ids) >= 0.5); + + let actual_distances = actual[DIST_COL].as_primitive::(); + let expected_distances = expected[DIST_COL].as_primitive::(); + for position in 0..actual.num_rows() { + let actual_distance = actual_distances.value(position); + let expected_distance = expected_distances.value(position); + let tolerance = 1.0e-4 * (1.0 + expected_distance.abs()); + assert!( + actual_distance.is_finite() + && (actual_distance - expected_distance).abs() <= tolerance, + "quantized distance mismatch at position {position}: actual={actual_distance}, expected={expected_distance}" + ); + } + } + + fn complete_refine_factor(row_count: usize) -> u32 { + u32::try_from(row_count.div_ceil(ResultsCacheTestFixture::K)).unwrap() + } + + #[rstest] + #[case::sq8(None)] + #[case::rq4(Some(4))] + #[case::rq8(Some(8))] + #[tokio::test] + async fn test_results_cache_end_to_end_and_index_replacement_invalidation( + #[case] rq_num_bits: Option, + ) { + let mut fixture = ResultsCacheTestFixture::new(rq_num_bits).await; + let expected = exact_results(&fixture).await; + + let (cold, cold_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&cold_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert!(cold_stats.bytes_read > 0); + assert_valid_result_set(&cold); + + let (warm, warm_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + assert_eq!( + count_metric(&warm_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + assert_eq!(warm_stats.bytes_read, 0); + assert_eq!(warm_stats.parts_loaded, 0); + assert_valid_result_set(&warm); + + let refine_factor = Some( + (ResultsCacheTestFixture::ROWS_PER_FRAGMENT * ResultsCacheTestFixture::NUM_FRAGMENTS + / ResultsCacheTestFixture::K) as u32, + ); + let (exact_cold, exact_cold_stats) = + run_results_cache_query(&fixture, true, refine_factor).await; + assert_eq!( + count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert_eq!(recall(&exact_cold, &expected), 1.0); + + let (exact_warm, exact_warm_stats) = + run_results_cache_query(&fixture, true, refine_factor).await; + assert_eq!( + count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + assert_eq!(recall(&exact_warm, &expected), 1.0); + + let (disabled, disabled_stats) = + run_results_cache_query(&fixture, false, refine_factor).await; + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + assert_eq!(recall(&disabled, &expected), 1.0); + + let (old_uuid, new_uuid) = fixture.replace_index().await; + assert_ne!(old_uuid, new_uuid); + let (after_replacement, replacement_stats) = + run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&replacement_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&replacement_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert_valid_result_set(&after_replacement); + } + + #[rstest] + #[case::sq8_l2(None, MetricType::L2)] + #[case::sq8_cosine(None, MetricType::Cosine)] + #[case::sq8_dot(None, MetricType::Dot)] + #[case::rq4_l2(Some(4), MetricType::L2)] + #[case::rq4_cosine(Some(4), MetricType::Cosine)] + #[case::rq4_dot(Some(4), MetricType::Dot)] + #[case::rq8_l2(Some(8), MetricType::L2)] + #[case::rq8_cosine(Some(8), MetricType::Cosine)] + #[case::rq8_dot(Some(8), MetricType::Dot)] + #[tokio::test] + async fn test_results_cache_scoring_matrix_matches_direct_native_and_exact_refinement( + #[case] rq_num_bits: Option, + #[case] metric_type: MetricType, + ) { + let fixture = ResultsCacheTestFixture::with_metric(rq_num_bits, metric_type).await; + let exact = exact_result_batch(&fixture).await; + let exact_row_ids = exact[ROW_ID].as_primitive::().values().to_vec(); + + let (_, cold_stats) = run_results_cache_batch_with_mode( + &fixture, + true, + None, + ResultsCacheQueryMode::Supported, + ApproxMode::Accurate, + ) + .await; + assert_eq!( + count_metric(&cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + let (warm, warm_stats) = run_results_cache_batch_with_mode( + &fixture, + true, + None, + ResultsCacheQueryMode::Supported, + ApproxMode::Accurate, + ) + .await; + assert_eq!( + count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + let direct = run_results_cache_batch_with_mode( + &fixture, + false, + None, + ResultsCacheQueryMode::Supported, + ApproxMode::Accurate, + ) + .await + .0; + assert_quantized_result_batches_agree(&warm, &direct); + let warm_row_ids = warm[ROW_ID].as_primitive::().values().to_vec(); + assert_valid_result_set(&warm_row_ids); + assert!( + recall(&warm_row_ids, &exact_row_ids) >= 0.5, + "{metric_type} warm-cache recall fell below 0.5: warm={warm_row_ids:?}, exact={exact_row_ids:?}" + ); + + let refine_factor = Some(complete_refine_factor( + ResultsCacheTestFixture::INITIAL_ROWS, + )); + let (exact_cold, exact_cold_stats) = run_results_cache_batch_with_mode( + &fixture, + true, + refine_factor, + ResultsCacheQueryMode::Supported, + ApproxMode::Accurate, + ) + .await; + assert_eq!( + count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert_result_batches_match(&exact_cold, &exact); + let (exact_warm, exact_warm_stats) = run_results_cache_batch_with_mode( + &fixture, + true, + refine_factor, + ResultsCacheQueryMode::Supported, + ApproxMode::Accurate, + ) + .await; + assert_eq!( + count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + assert_result_batches_match(&exact_warm, &exact); + } + + #[rstest] + #[case::sq8(None)] + #[case::rq4(Some(4))] + #[case::rq8(Some(8))] + #[tokio::test] + async fn test_results_cache_dataset_append_and_multi_segment_invalidation( + #[case] rq_num_bits: Option, + ) { + let mut fixture = ResultsCacheTestFixture::new(rq_num_bits).await; + let original_version = fixture.dataset.version_id(); + let original_index_uuids = fixture.vector_index_uuids().await; + assert_eq!(original_index_uuids.len(), 1); + + let (_, cold_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + let (_, warm_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + + let appended_index_uuids = fixture.append_index_segment().await; + assert!(fixture.dataset.version_id() > original_version); + assert_eq!(appended_index_uuids.len(), 2); + assert!(appended_index_uuids.contains(&original_index_uuids[0])); + + let (after_append, after_append_stats) = + run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&after_append_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + appended_index_uuids.len() + ); + assert_eq!( + count_metric(&after_append_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_valid_result_set(&after_append); + + let (after_append_warm, after_append_warm_stats) = + run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&after_append_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + appended_index_uuids.len() + ); + assert_eq!( + count_metric(&after_append_warm_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + assert_valid_result_set(&after_append_warm); + + let expected = exact_results(&fixture).await; + let row_count = fixture.dataset.count_rows(None).await.unwrap(); + let refine_factor = Some(complete_refine_factor(row_count)); + let (exact_cold, exact_cold_stats) = + run_results_cache_query(&fixture, true, refine_factor).await; + assert_eq!( + count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + appended_index_uuids.len() + ); + assert_eq!(recall(&exact_cold, &expected), 1.0); + + let (exact_warm, exact_warm_stats) = + run_results_cache_query(&fixture, true, refine_factor).await; + assert_eq!( + count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + appended_index_uuids.len() + ); + assert_eq!(recall(&exact_warm, &expected), 1.0); + } + + #[rstest] + #[case::sq8(None)] + #[case::rq4(Some(4))] + #[case::rq8(Some(8))] + #[tokio::test] + async fn test_results_cache_deletion_and_fragment_reuse_invalidation( + #[case] rq_num_bits: Option, + ) { + let mut fixture = ResultsCacheTestFixture::new(rq_num_bits).await; + let deleted_row_ids = exact_results(&fixture).await; + + let (_, initial_cold_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&initial_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + let (_, initial_warm_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&initial_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + + let version_before_delete = fixture.dataset.version_id(); + fixture.delete_query_centroid_rows().await; + assert!(fixture.dataset.version_id() > version_before_delete); + + let (after_delete, after_delete_stats) = + run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&after_delete_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert_eq!( + count_metric(&after_delete_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_valid_result_set(&after_delete); + assert!( + after_delete + .iter() + .all(|row_id| !deleted_row_ids.contains(row_id)) + ); + + let (_, after_delete_warm_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&after_delete_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + + let vector_index_uuid = fixture.vector_index_uuids().await[0]; + assert_eq!(fixture.fragment_reuse_uuid().await, None); + let version_before_compaction = fixture.dataset.version_id(); + let fragment_reuse_uuid = fixture.compact_with_fragment_reuse().await; + assert!(fixture.dataset.version_id() > version_before_compaction); + assert_eq!(fixture.vector_index_uuids().await, vec![vector_index_uuid]); + assert_eq!( + fixture.fragment_reuse_uuid().await, + Some(fragment_reuse_uuid) + ); + + let (after_compaction, after_compaction_stats) = + run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&after_compaction_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert_eq!( + count_metric(&after_compaction_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_valid_result_set(&after_compaction); + + let (_, after_compaction_warm_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric( + &after_compaction_warm_stats, + VECTOR_RESULTS_CACHE_HITS_METRIC + ), + 1 + ); + + let expected = exact_results(&fixture).await; + let row_count = fixture.dataset.count_rows(None).await.unwrap(); + let refine_factor = Some(complete_refine_factor(row_count)); + let (exact_cold, exact_cold_stats) = + run_results_cache_query(&fixture, true, refine_factor).await; + assert_eq!( + count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert_eq!(recall(&exact_cold, &expected), 1.0); + + let (exact_warm, exact_warm_stats) = + run_results_cache_query(&fixture, true, refine_factor).await; + assert_eq!( + count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + assert_eq!(recall(&exact_warm, &expected), 1.0); + } + + #[rstest] + #[case::prefilter(ResultsCacheQueryMode::Prefilter)] + #[case::scalar_prefilter(ResultsCacheQueryMode::ScalarPrefilter)] + #[case::distance_range(ResultsCacheQueryMode::DistanceRange)] + #[case::adaptive_nprobes(ResultsCacheQueryMode::AdaptiveNprobes)] + #[case::overlay(ResultsCacheQueryMode::Overlay)] + #[tokio::test] + async fn test_results_cache_bypasses_unsupported_query_shapes( + #[case] query_mode: ResultsCacheQueryMode, + ) { + let mut fixture = ResultsCacheTestFixture::new(None).await; + if matches!(query_mode, ResultsCacheQueryMode::ScalarPrefilter) { + fixture.create_row_scalar_index().await; + } + let row_count = fixture.dataset.count_rows(None).await.unwrap(); + let refine_factor = Some(complete_refine_factor(row_count)); + let (expected, disabled_stats) = + run_results_cache_query_with_mode(&fixture, false, refine_factor, query_mode).await; + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + assert_valid_result_set(&expected); + let exact = exact_results(&fixture).await; + assert!(recall(&expected, &exact) >= 0.5); + for _ in 0..2 { + let (row_ids, stats) = + run_results_cache_query_with_mode(&fixture, true, refine_factor, query_mode).await; + assert_eq!( + count_metric(&stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0, + "unsupported query mode {query_mode:?} recorded a cache hit" + ); + assert_eq!( + count_metric(&stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0, + "unsupported query mode {query_mode:?} recorded a cache miss" + ); + assert_eq!(row_ids, expected); + } + } + + #[rstest] + #[case::rq4_fast(4, ApproxMode::Fast)] + #[case::rq4_normal(4, ApproxMode::Normal)] + #[case::rq8_fast(8, ApproxMode::Fast)] + #[case::rq8_normal(8, ApproxMode::Normal)] + #[tokio::test] + async fn test_results_cache_bypasses_non_accurate_rq( + #[case] rq_num_bits: u8, + #[case] approx_mode: ApproxMode, + ) { + let fixture = ResultsCacheTestFixture::new(Some(rq_num_bits)).await; + let (expected, disabled_stats) = run_results_cache_batch_with_mode( + &fixture, + false, + None, + ResultsCacheQueryMode::Supported, + approx_mode, + ) + .await; + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + + for _ in 0..2 { + let (actual, stats) = run_results_cache_batch_with_mode( + &fixture, + true, + None, + ResultsCacheQueryMode::Supported, + approx_mode, + ) + .await; + assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_HITS_METRIC), 0); + assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), 0); + assert_result_batches_match(&actual, &expected); + } + } + + #[tokio::test] + async fn test_results_cache_bypasses_unsupported_index_type() { + let mut fixture = ResultsCacheTestFixture::new(None).await; + fixture.replace_with_unsupported_flat_index().await; + + let row_count = fixture.dataset.count_rows(None).await.unwrap(); + let refine_factor = Some(complete_refine_factor(row_count)); + let (expected, disabled_stats) = + run_results_cache_query(&fixture, false, refine_factor).await; + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + assert_valid_result_set(&expected); + let exact = exact_results(&fixture).await; + assert!(recall(&expected, &exact) >= 0.5); + for _ in 0..2 { + let (row_ids, stats) = run_results_cache_query(&fixture, true, refine_factor).await; + assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_HITS_METRIC), 0); + assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), 0); + assert_eq!(row_ids, expected); + } + } + + #[rstest] + #[case::sq8(None)] + #[case::rq4(Some(4))] + #[case::rq8(Some(8))] + #[tokio::test] + async fn test_results_cache_unusable_entry_falls_back_and_is_replaced( + #[case] rq_num_bits: Option, + ) { + let fixture = ResultsCacheTestFixture::new(rq_num_bits).await; + let (ordinary, disabled_stats) = run_results_cache_query(&fixture, false, None).await; + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + assert_valid_result_set(&ordinary); + fixture.insert_out_of_range_results_cache_entry().await; + + let (fallback, fallback_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&fallback_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 0 + ); + assert_eq!( + count_metric(&fallback_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 1 + ); + assert_valid_result_set(&fallback); + assert!(recall(&fallback, &ordinary) >= 0.5); + + let (warm, warm_stats) = run_results_cache_query(&fixture, true, None).await; + assert_eq!( + count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), + 1 + ); + assert_eq!( + count_metric(&warm_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), + 0 + ); + assert_valid_result_set(&warm); + assert!(recall(&warm, &ordinary) >= 0.5); + } + #[rstest] #[tokio::test] async fn test_no_max_nprobes(#[values(1, 20)] num_deltas: usize) { diff --git a/rust/lance/src/io/exec/knn_results_cache.rs b/rust/lance/src/io/exec/knn_results_cache.rs new file mode 100644 index 00000000000..285faec9ac0 --- /dev/null +++ b/rust/lance/src/io/exec/knn_results_cache.rs @@ -0,0 +1,1162 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Experimental cache execution for reusable IVF candidate pools. + +use std::collections::{BinaryHeap, HashMap}; +use std::sync::Arc; + +use arrow::datatypes::{Float32Type, UInt64Type}; +use arrow_array::cast::AsArray; +use arrow_array::{Float32Array, RecordBatch, UInt32Array, UInt64Array}; +use futures::{StreamExt, TryStreamExt}; +use lance_core::cache::LanceCache; +use lance_core::{Error, ROW_ID, Result}; +use lance_index::IndexType; +use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; +use lance_index::metrics::MetricsCollector; +use lance_index::prefilter::PreFilter; +use lance_index::vector::graph::{OrderedFloat, OrderedNode}; +use lance_index::vector::quantizer::QuantizationType; +use lance_index::vector::v3::subindex::SubIndexType; +use lance_index::vector::{DIST_COL, PartitionSearchResult, Query, VectorIndex}; +use lance_select::RowAddrMask; +use lance_table::format::IndexMetadata; + +use crate::dataset::Dataset; +use crate::index::DatasetIndexExt; +use crate::session::index_caches::{ + CachedVectorCandidate, VectorResultsCacheEntry, VectorResultsCacheIdentity, +}; + +use super::knn::KNN_INDEX_SCHEMA; + +/// Maximum reusable candidate pool retained for one exact query/search shape. +/// +/// This is deliberately fixed while the feature is experimental. It is part of +/// the cache identity, so changing it cannot reuse entries created under another +/// limit. +pub(super) const RESULTS_CACHE_CANDIDATE_LIMIT: usize = 1000; + +const VECTOR_RESULTS_CACHE_ENV: &str = "LANCE_EXPERIMENTAL_VECTOR_RESULTS_CACHE"; + +pub(super) struct ResultsCacheSearch { + pub batch: RecordBatch, + /// True for both a warm entry and a concurrent load coalesced behind its owner. + pub was_hit: bool, +} + +pub(super) struct ResultsCacheSearchParams<'a> { + pub cache: &'a LanceCache, + pub identity: VectorResultsCacheIdentity, + pub index: Arc, + pub query: &'a Query, + pub partitions: Arc, + pub centroid_distances: Arc, + pub prefilter: Arc, + pub segment_mask: Option>, + pub metrics: Arc, + pub parallelism: usize, +} + +pub(super) fn is_enabled() -> bool { + std::env::var(VECTOR_RESULTS_CACHE_ENV) + .ok() + .is_some_and(|value| value == "1") +} + +pub(super) async fn identity_for_query( + dataset: &Dataset, + index_metadata: &IndexMetadata, + index: &dyn VectorIndex, + query: &Query, + partition_ids: &[u32], + has_prefilter: bool, + has_overlay: bool, +) -> Option { + let (sub_index_type, quantization_type) = match index.index_type() { + IndexType::IvfRq => (SubIndexType::Flat, QuantizationType::Rabit), + IndexType::IvfSq => (SubIndexType::Flat, QuantizationType::Scalar), + _ => return None, + }; + let index_store = dataset.object_store_for_index(index_metadata).await.ok()?; + let fragment_reuse_uuid = dataset + .load_indices() + .await + .ok()? + .iter() + .find(|metadata| metadata.name == FRAG_REUSE_INDEX_NAME) + .map(|metadata| metadata.uuid); + let manifest_location = dataset.manifest_location(); + let manifest_path = manifest_location.path.as_ref(); + let manifest_etag = manifest_location.e_tag.as_deref().unwrap_or_default(); + let manifest_size = manifest_location + .size + .map_or_else(|| "none".to_owned(), |size| format!("some:{size}")); + let dataset_read_identity = format!( + "{}:{}/{}:{}/{}:{}", + manifest_path.len(), + manifest_path, + manifest_etag.len(), + manifest_etag, + manifest_size.len(), + manifest_size + ); + + VectorResultsCacheIdentity::try_new( + &index_store.store_prefix, + &dataset_read_identity, + dataset.version_id(), + index_metadata, + fragment_reuse_uuid, + query, + index.metric_type(), + sub_index_type, + quantization_type, + partition_ids, + RESULTS_CACHE_CANDIDATE_LIMIT, + has_prefilter, + has_overlay, + ) + .ok() +} + +pub(super) async fn search(params: ResultsCacheSearchParams<'_>) -> Result { + let ResultsCacheSearchParams { + cache, + identity, + index, + query, + partitions, + centroid_distances, + prefilter, + segment_mask, + metrics, + parallelism, + } = params; + let mut populated_batch = None; + let populated_batch_slot = &mut populated_batch; + let loader_identity = identity.clone(); + let loader_index = index.clone(); + let loader_partitions = partitions.clone(); + let loader_centroid_distances = centroid_distances.clone(); + let loader_prefilter = prefilter.clone(); + let loader_segment_mask = segment_mask.clone(); + let loader_metrics = metrics.clone(); + let cache_lookup = cache + .get_or_insert_with_key_hit(identity.clone(), move || async move { + let (candidates, batch) = populate_candidates( + loader_index, + query, + loader_partitions, + loader_centroid_distances, + loader_prefilter, + loader_segment_mask, + loader_metrics, + parallelism, + loader_identity.result_limit(), + ) + .await?; + *populated_batch_slot = Some(batch); + VectorResultsCacheEntry::try_new(loader_identity, candidates) + }) + .await; + + let (entry, was_cached) = match cache_lookup { + Ok(result) => result, + Err(error) => { + if let Some(batch) = populated_batch { + // The candidate search succeeded, so an invalid cache entry must not + // discard its query result or cause the ordinary path to repeat work. + log::debug!("not storing invalid vector results cache entry: {error}"); + return Ok(ResultsCacheSearch { + batch, + was_hit: false, + }); + } + return Err(error); + } + }; + + if !was_cached { + let batch = populated_batch.ok_or_else(|| { + Error::internal( + "vector results cache loader completed without its candidate-search batch", + ) + })?; + return Ok(ResultsCacheSearch { + batch, + was_hit: false, + }); + } + + if entry.is_compatible_with(&identity) { + match replay_candidates( + index.clone(), + query, + entry.candidates(), + identity.result_limit(), + metrics.clone(), + parallelism, + ) + .await + { + Ok(batch) => { + return Ok(ResultsCacheSearch { + batch, + was_hit: true, + }); + } + Err(error) => { + // A stale or malformed in-memory entry must never fail the query. + // Re-run the candidate-producing path and replace it. + log::debug!("ignoring unusable vector results cache entry: {error}"); + } + } + } + + // An unusable existing entry cannot enter get-or-insert's loader, so replace it + // directly. Normal cold misses are single-flighted by the lookup above. + let (candidates, batch) = populate_candidates( + index, + query, + partitions, + centroid_distances, + prefilter, + segment_mask, + metrics, + parallelism, + identity.result_limit(), + ) + .await?; + match VectorResultsCacheEntry::try_new(identity.clone(), candidates) { + Ok(entry) => cache.insert_with_key(&identity, Arc::new(entry)).await, + Err(error) => { + // Cache population is best-effort and cannot change query success. + log::debug!("not storing invalid vector results cache entry: {error}"); + } + } + Ok(ResultsCacheSearch { + batch, + was_hit: false, + }) +} + +#[allow(clippy::too_many_arguments)] +async fn populate_candidates( + index: Arc, + query: &Query, + partitions: Arc, + centroid_distances: Arc, + prefilter: Arc, + segment_mask: Option>, + metrics: Arc, + parallelism: usize, + result_limit: usize, +) -> Result<(Vec, RecordBatch)> { + if partitions.len() != centroid_distances.len() { + return Err(Error::invalid_input(format!( + "partition count {} does not match centroid distance count {} for vector results cache", + partitions.len(), + centroid_distances.len() + ))); + } + let accumulated = futures::stream::iter(0..partitions.len()) + .map(|partition_index| { + let index = index.clone(); + let prefilter = prefilter.clone(); + let metrics = metrics.clone(); + let mut partition_query = query.clone(); + let partition_id = partitions.value(partition_index); + partition_query.dist_q_c = centroid_distances.value(partition_index); + async move { + index + .search_in_partition_with_candidates( + partition_id as usize, + &partition_query, + prefilter, + metrics.as_ref(), + RESULTS_CACHE_CANDIDATE_LIMIT, + ) + .await + } + }) + .buffered(parallelism.max(1)) + .try_fold( + SearchAccumulator::new(RESULTS_CACHE_CANDIDATE_LIMIT, result_limit), + move |mut accumulated, partition_result| { + let segment_mask = segment_mask.clone(); + async move { + accumulated.add(partition_result, segment_mask.as_deref())?; + Ok(accumulated) + } + }, + ) + .await?; + accumulated.finish() +} + +struct SearchAccumulator { + candidate_limit: usize, + result_limit: usize, + candidates: BinaryHeap>, + results: BinaryHeap>, +} + +impl SearchAccumulator { + fn new(candidate_limit: usize, result_limit: usize) -> Self { + Self { + candidate_limit, + result_limit, + candidates: BinaryHeap::with_capacity(candidate_limit), + results: BinaryHeap::with_capacity(result_limit), + } + } + + fn add( + &mut self, + search_result: PartitionSearchResult, + segment_mask: Option<&RowAddrMask>, + ) -> Result<()> { + if search_result.candidates.len() != search_result.batch.num_rows() { + return Err(Error::internal(format!( + "partition search returned {} candidate identities for {} rows", + search_result.candidates.len(), + search_result.batch.num_rows() + ))); + } + let row_ids = search_result + .batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("partition search result has no row-id column"))? + .as_primitive::(); + let distances = search_result + .batch + .column_by_name(DIST_COL) + .ok_or_else(|| Error::internal("partition search result has no distance column"))? + .as_primitive::(); + if row_ids.len() != distances.len() { + return Err(Error::internal(format!( + "partition search returned {} row ids and {} distances", + row_ids.len(), + distances.len() + ))); + } + for ((candidate, &row_id), &distance) in search_result + .candidates + .iter() + .zip(row_ids.values()) + .zip(distances.values()) + { + if distance.is_nan() { + continue; + } + if segment_mask.is_some_and(|mask| !mask.selected(row_id)) { + continue; + } + let candidate = + CachedVectorCandidate::new(candidate.partition_id, candidate.offset_in_partition); + push_top_candidate( + &mut self.candidates, + self.candidate_limit, + candidate, + distance, + ); + push_top_candidate(&mut self.results, self.result_limit, row_id, distance); + } + Ok(()) + } + + fn finish(self) -> Result<(Vec, RecordBatch)> { + let mut ordered_candidates = self.candidates.into_vec(); + ordered_candidates.sort_by(|left, right| { + left.dist + .cmp(&right.dist) + .then_with(|| left.id.partition_id().cmp(&right.id.partition_id())) + .then_with(|| { + left.id + .offset_in_partition() + .cmp(&right.id.offset_in_partition()) + }) + }); + let candidates = ordered_candidates.into_iter().map(|node| node.id).collect(); + Ok((candidates, batch_from_heap(self.results)?)) + } +} + +fn push_top_candidate( + heap: &mut BinaryHeap>, + limit: usize, + id: T, + distance: f32, +) { + if limit == 0 { + return; + } + let node = OrderedNode::new(id, OrderedFloat(distance)); + if heap.len() < limit { + heap.push(node); + } else if heap + .peek() + .is_some_and(|farthest| farthest.dist > node.dist) + { + heap.pop(); + heap.push(node); + } +} + +async fn replay_candidates( + index: Arc, + query: &Query, + candidates: &[CachedVectorCandidate], + result_limit: usize, + metrics: Arc, + parallelism: usize, +) -> Result { + let mut offsets_by_partition = HashMap::>::new(); + for candidate in candidates { + offsets_by_partition + .entry(candidate.partition_id()) + .or_default() + .push(candidate.offset_in_partition()); + } + + let scored_partitions = futures::stream::iter(offsets_by_partition) + .map(|(partition_id, offsets)| { + let index = index.clone(); + let query = query.clone(); + let metrics = metrics.clone(); + async move { + let batch = index + .score_partition_candidates( + partition_id as usize, + &query, + &offsets, + metrics.as_ref(), + ) + .await?; + Result::Ok((partition_id, offsets.len(), batch)) + } + }) + .buffered(parallelism.max(1)) + .try_collect::>() + .await?; + + let mut top_results = BinaryHeap::with_capacity(result_limit); + for (partition_id, expected_rows, batch) in scored_partitions { + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("candidate scoring result has no row-id column"))? + .as_primitive::(); + let distances = batch + .column_by_name(DIST_COL) + .ok_or_else(|| Error::internal("candidate scoring result has no distance column"))? + .as_primitive::(); + if row_ids.len() != expected_rows || distances.len() != expected_rows { + return Err(Error::internal(format!( + "candidate scoring for partition {partition_id} returned {} row ids and {} distances for {} offsets", + row_ids.len(), + distances.len(), + expected_rows + ))); + } + for (&row_id, &distance) in row_ids.values().iter().zip(distances.values()) { + if distance.is_nan() { + continue; + } + push_top_candidate(&mut top_results, result_limit, row_id, distance); + } + } + batch_from_heap(top_results) +} + +fn batch_from_heap(heap: BinaryHeap>) -> Result { + let mut ordered = heap.into_vec(); + ordered.sort_by(|left, right| { + left.dist + .cmp(&right.dist) + .then_with(|| left.id.cmp(&right.id)) + }); + let mut row_ids = Vec::with_capacity(ordered.len()); + let mut distances = Vec::with_capacity(ordered.len()); + for result in ordered { + row_ids.push(result.id); + distances.push(result.dist.0); + } + Ok(RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(distances)), + Arc::new(UInt64Array::from(row_ids)), + ], + )?) +} + +#[cfg(test)] +mod tests { + use std::any::Any; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use async_trait::async_trait; + use lance_core::cache::QuickCacheBackend; + use lance_core::deepsize::DeepSizeOf; + use lance_core::utils::row_addr_remap::RowAddrRemap; + use lance_index::Index; + use lance_index::metrics::NoOpMetricsCollector; + use lance_index::prefilter::NoFilter; + use lance_index::vector::ivf::storage::IvfModel; + use lance_index::vector::quantizer::Quantizer; + use lance_index::vector::v3::subindex::SubIndexType; + use lance_index::vector::{ApproxMode, DEFAULT_QUERY_PARALLELISM, PartitionSearchCandidate}; + use lance_io::traits::Reader; + use roaring::RoaringBitmap; + use uuid::Uuid; + + use super::*; + + #[derive(Debug, DeepSizeOf)] + struct TestCandidateIndex { + miss_partition_searches: AtomicUsize, + miss_search_yields: AtomicUsize, + block_miss_searches: AtomicUsize, + hit_partition_scores: AtomicUsize, + active_hit_partition_scores: AtomicUsize, + max_active_hit_partition_scores: AtomicUsize, + return_malformed_hit_shape: AtomicUsize, + row_ids: Vec, + } + + impl TestCandidateIndex { + fn new() -> Self { + Self { + miss_partition_searches: AtomicUsize::new(0), + miss_search_yields: AtomicUsize::new(0), + block_miss_searches: AtomicUsize::new(0), + hit_partition_scores: AtomicUsize::new(0), + active_hit_partition_scores: AtomicUsize::new(0), + max_active_hit_partition_scores: AtomicUsize::new(0), + return_malformed_hit_shape: AtomicUsize::new(0), + row_ids: Vec::new(), + } + } + } + + #[async_trait] + impl Index for TestCandidateIndex { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_index(self: Arc) -> Arc { + self + } + + fn statistics(&self) -> Result { + Ok(serde_json::json!({})) + } + + async fn prewarm(&self) -> Result<()> { + Ok(()) + } + + fn index_type(&self) -> IndexType { + IndexType::IvfRq + } + + async fn calculate_included_frags(&self) -> Result { + Ok(RoaringBitmap::new()) + } + } + + #[async_trait] + impl VectorIndex for TestCandidateIndex { + async fn search( + &self, + _query: &Query, + _pre_filter: Arc, + _metrics: &dyn MetricsCollector, + ) -> Result { + Err(Error::not_supported("test whole-index search")) + } + + fn find_partitions(&self, _query: &Query) -> Result<(UInt32Array, Float32Array)> { + Ok(( + UInt32Array::from(vec![0, 1]), + Float32Array::from(vec![0.0, 0.0]), + )) + } + + fn total_partitions(&self) -> usize { + 2 + } + + async fn search_in_partition( + &self, + _partition_id: usize, + _query: &Query, + _pre_filter: Arc, + _metrics: &dyn MetricsCollector, + ) -> Result { + Err(Error::not_supported("test partition search")) + } + + async fn search_in_partition_with_candidates( + &self, + partition_id: usize, + _query: &Query, + _pre_filter: Arc, + _metrics: &dyn MetricsCollector, + _candidate_limit: usize, + ) -> Result { + self.miss_partition_searches.fetch_add(1, Ordering::Relaxed); + for _ in 0..self.miss_search_yields.load(Ordering::Relaxed) { + tokio::task::yield_now().await; + } + while self.block_miss_searches.load(Ordering::Relaxed) != 0 { + tokio::task::yield_now().await; + } + let (row_ids, distances) = match partition_id { + 0 => (vec![0, 1], vec![0.4, 0.2]), + 1 => (vec![10, 11], vec![0.3, 0.1]), + _ => { + return Err(Error::invalid_input(format!( + "unexpected test partition {partition_id}" + ))); + } + }; + let candidates = (0..row_ids.len()) + .map(|offset| PartitionSearchCandidate { + partition_id: partition_id as u32, + offset_in_partition: offset as u32, + }) + .collect(); + let batch = RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(distances)), + Arc::new(UInt64Array::from(row_ids)), + ], + )?; + Ok(PartitionSearchResult { batch, candidates }) + } + + async fn score_partition_candidates( + &self, + partition_id: usize, + _query: &Query, + offsets_in_partition: &[u32], + _metrics: &dyn MetricsCollector, + ) -> Result { + self.hit_partition_scores.fetch_add(1, Ordering::Relaxed); + let active = self + .active_hit_partition_scores + .fetch_add(1, Ordering::Relaxed) + + 1; + self.max_active_hit_partition_scores + .fetch_max(active, Ordering::Relaxed); + tokio::task::yield_now().await; + let mut scored = offsets_in_partition + .iter() + .map(|offset| match (partition_id, *offset) { + (0, 0) => Ok((0, 0.05)), + (0, 1) => Ok((1, 0.8)), + (1, 0) => Ok((10, 0.6)), + (1, 1) => Ok((11, 0.1)), + _ => Err(Error::invalid_input(format!( + "unexpected test candidate ({partition_id}, {offset})" + ))), + }) + .collect::>>()?; + if self.return_malformed_hit_shape.load(Ordering::Relaxed) != 0 { + scored.pop(); + } + self.active_hit_partition_scores + .fetch_sub(1, Ordering::Relaxed); + let (row_ids, distances): (Vec<_>, Vec<_>) = scored.into_iter().unzip(); + Ok(RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(distances)), + Arc::new(UInt64Array::from(row_ids)), + ], + )?) + } + + fn is_loadable(&self) -> bool { + false + } + + fn use_residual(&self) -> bool { + false + } + + async fn load( + &self, + _reader: Arc, + _offset: usize, + _length: usize, + ) -> Result> { + Err(Error::not_supported("test index load")) + } + + fn num_rows(&self) -> u64 { + 0 + } + + fn row_ids(&self) -> Box + '_> { + Box::new(self.row_ids.iter()) + } + + async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { + Ok(()) + } + + async fn to_batch_stream( + &self, + _with_vector: bool, + ) -> Result { + Err(Error::not_supported("test batch stream")) + } + + fn ivf_model(&self) -> &IvfModel { + unreachable!("test cache path does not inspect the IVF model") + } + + fn quantizer(&self) -> Quantizer { + unreachable!("test cache path does not inspect the quantizer") + } + + fn partition_size(&self, _part_id: usize) -> usize { + 2 + } + + fn sub_index_type(&self) -> (SubIndexType, QuantizationType) { + (SubIndexType::Flat, QuantizationType::Rabit) + } + + fn metric_type(&self) -> lance_linalg::distance::DistanceType { + lance_linalg::distance::DistanceType::L2 + } + } + + fn test_query() -> Query { + Query { + column: "vector".to_string(), + key: Arc::new(Float32Array::from(vec![0.0, 1.0])), + k: 2, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 2, + maximum_nprobes: Some(2), + ef: None, + refine_factor: None, + metric_type: Some(lance_linalg::distance::DistanceType::L2), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + approx_mode: ApproxMode::Accurate, + } + } + + fn test_identity(query: &Query) -> VectorResultsCacheIdentity { + let metadata = IndexMetadata { + uuid: Uuid::from_u128(1), + fields: vec![0], + covering_fields: vec![], + name: "vector_idx".to_string(), + dataset_version: 1, + fragment_bitmap: Some([0_u32].into_iter().collect()), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.table.VectorIndexDetails".to_string(), + value: vec![1], + })), + index_version: 1, + created_at: None, + base_id: None, + files: None, + }; + VectorResultsCacheIdentity::try_new( + "memory", + "manifest-1", + 1, + &metadata, + None, + query, + lance_linalg::distance::DistanceType::L2, + SubIndexType::Flat, + QuantizationType::Rabit, + &[0, 1], + RESULTS_CACHE_CANDIDATE_LIMIT, + false, + false, + ) + .unwrap() + } + + fn row_ids(batch: &RecordBatch) -> Vec { + batch[ROW_ID].as_primitive::().values().to_vec() + } + + #[test] + fn accumulator_filters_nan_and_orders_ties_with_aligned_rows() { + let mut accumulator = SearchAccumulator::new(4, 4); + accumulator + .add( + PartitionSearchResult { + batch: RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(vec![0.2, f32::NAN, 0.2])), + Arc::new(UInt64Array::from(vec![30, 20, 10])), + ], + ) + .unwrap(), + candidates: vec![ + PartitionSearchCandidate { + partition_id: 0, + offset_in_partition: 2, + }, + PartitionSearchCandidate { + partition_id: 0, + offset_in_partition: 1, + }, + PartitionSearchCandidate { + partition_id: 0, + offset_in_partition: 0, + }, + ], + }, + None, + ) + .unwrap(); + + let (candidates, batch) = accumulator.finish().unwrap(); + assert_eq!( + candidates, + vec![ + CachedVectorCandidate::new(0, 0), + CachedVectorCandidate::new(0, 2), + ] + ); + assert_eq!(row_ids(&batch), vec![10, 30]); + assert_eq!( + batch[DIST_COL].as_primitive::().values(), + &[0.2, 0.2] + ); + } + + #[test] + fn accumulator_filters_candidates_outside_segment() { + let mut accumulator = SearchAccumulator::new(4, 4); + let segment_mask = + RowAddrMask::from_allowed(lance_select::RowAddrTreeMap::from_iter([10_u64, 30])); + accumulator + .add( + PartitionSearchResult { + batch: RecordBatch::try_new( + KNN_INDEX_SCHEMA.clone(), + vec![ + Arc::new(Float32Array::from(vec![0.3, 0.2, 0.1])), + Arc::new(UInt64Array::from(vec![30, 20, 10])), + ], + ) + .unwrap(), + candidates: vec![ + PartitionSearchCandidate { + partition_id: 0, + offset_in_partition: 2, + }, + PartitionSearchCandidate { + partition_id: 0, + offset_in_partition: 1, + }, + PartitionSearchCandidate { + partition_id: 0, + offset_in_partition: 0, + }, + ], + }, + Some(&segment_mask), + ) + .unwrap(); + + let (candidates, batch) = accumulator.finish().unwrap(); + assert_eq!( + candidates, + vec![ + CachedVectorCandidate::new(0, 0), + CachedVectorCandidate::new(0, 2), + ] + ); + assert_eq!(row_ids(&batch), vec![10, 30]); + } + + #[tokio::test] + async fn miss_populates_candidates_and_hit_rescores_them() { + let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(64 * 1024))); + let index = Arc::new(TestCandidateIndex::new()); + let query = test_query(); + let identity = test_identity(&query); + let partitions = Arc::new(UInt32Array::from(vec![0, 1])); + let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); + + let miss = search(ResultsCacheSearchParams { + cache: &cache, + identity: identity.clone(), + index: index.clone(), + query: &query, + partitions: partitions.clone(), + centroid_distances: centroid_distances.clone(), + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + .await + .unwrap(); + assert!(!miss.was_hit); + assert_eq!(row_ids(&miss.batch), vec![11, 1]); + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); + assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 0); + + let hit = search(ResultsCacheSearchParams { + cache: &cache, + identity, + index: index.clone(), + query: &query, + partitions, + centroid_distances, + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + .await + .unwrap(); + assert!(hit.was_hit); + assert_eq!(row_ids(&hit.batch), vec![0, 11]); + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); + assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 2); + assert_eq!( + index + .max_active_hit_partition_scores + .load(Ordering::Relaxed), + 2 + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn concurrent_cold_searches_populate_candidates_once() { + const SEARCHES: usize = 4; + + let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(64 * 1024))); + let index = Arc::new(TestCandidateIndex::new()); + index.miss_search_yields.store(16, Ordering::Relaxed); + let query = test_query(); + let identity = test_identity(&query); + let partitions = Arc::new(UInt32Array::from(vec![0, 1])); + let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); + + let results = futures::future::join_all((0..SEARCHES).map(|_| { + search(ResultsCacheSearchParams { + cache: &cache, + identity: identity.clone(), + index: index.clone(), + query: &query, + partitions: partitions.clone(), + centroid_distances: centroid_distances.clone(), + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + })) + .await + .into_iter() + .collect::>>() + .unwrap(); + + assert_eq!(results.iter().filter(|result| !result.was_hit).count(), 1); + assert_eq!(results.iter().filter(|result| result.was_hit).count(), 3); + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); + assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 6); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn concurrent_cold_search_retries_after_population_is_cancelled() { + let cache = Arc::new(LanceCache::with_backend(Arc::new( + QuickCacheBackend::with_capacity(64 * 1024), + ))); + let index = Arc::new(TestCandidateIndex::new()); + index.block_miss_searches.store(1, Ordering::Relaxed); + + let owner = { + let cache = cache.clone(); + let index = index.clone(); + tokio::spawn(async move { + let query = test_query(); + let identity = test_identity(&query); + search(ResultsCacheSearchParams { + cache: cache.as_ref(), + identity, + index, + query: &query, + partitions: Arc::new(UInt32Array::from(vec![0, 1])), + centroid_distances: Arc::new(Float32Array::from(vec![0.0, 0.0])), + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + .await + }) + }; + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while index.miss_partition_searches.load(Ordering::Relaxed) < 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("population owner did not start both partition searches"); + + let contender = { + let cache = cache.clone(); + let index = index.clone(); + tokio::spawn(async move { + let query = test_query(); + let identity = test_identity(&query); + search(ResultsCacheSearchParams { + cache: cache.as_ref(), + identity, + index, + query: &query, + partitions: Arc::new(UInt32Array::from(vec![0, 1])), + centroid_distances: Arc::new(Float32Array::from(vec![0.0, 0.0])), + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + .await + }) + }; + for _ in 0..32 { + tokio::task::yield_now().await; + } + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); + assert!(!contender.is_finished()); + + owner.abort(); + assert!(matches!(owner.await, Err(error) if error.is_cancelled())); + index.block_miss_searches.store(0, Ordering::Relaxed); + let result = tokio::time::timeout(std::time::Duration::from_secs(5), contender) + .await + .expect("contender remained parked after population cancellation") + .unwrap() + .unwrap(); + assert!(!result.was_hit); + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); + } + + #[tokio::test] + async fn different_queries_with_same_partitions_do_not_share_candidates() { + let cache = LanceCache::with_capacity(64 * 1024); + let index = Arc::new(TestCandidateIndex::new()); + let first_query = test_query(); + let first_identity = test_identity(&first_query); + let partitions = Arc::new(UInt32Array::from(vec![0, 1])); + let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); + + let first = search(ResultsCacheSearchParams { + cache: &cache, + identity: first_identity, + index: index.clone(), + query: &first_query, + partitions: partitions.clone(), + centroid_distances: centroid_distances.clone(), + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + .await + .unwrap(); + assert!(!first.was_hit); + + let mut second_query = first_query.clone(); + second_query.key = Arc::new(Float32Array::from(vec![1.0, 0.0])); + let second_identity = test_identity(&second_query); + let second = search(ResultsCacheSearchParams { + cache: &cache, + identity: second_identity.clone(), + index: index.clone(), + query: &second_query, + partitions: partitions.clone(), + centroid_distances: centroid_distances.clone(), + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + .await + .unwrap(); + assert!(!second.was_hit); + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); + assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 0); + + let repeated_second = search(ResultsCacheSearchParams { + cache: &cache, + identity: second_identity, + index: index.clone(), + query: &second_query, + partitions, + centroid_distances, + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }) + .await + .unwrap(); + assert!(repeated_second.was_hit); + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); + assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 2); + } + + #[tokio::test] + async fn results_cache_malformed_hit_shape_falls_back_and_replaces_entry() { + let cache = LanceCache::with_capacity(64 * 1024); + let index = Arc::new(TestCandidateIndex::new()); + let query = test_query(); + let identity = test_identity(&query); + let partitions = Arc::new(UInt32Array::from(vec![0, 1])); + let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); + let params = || ResultsCacheSearchParams { + cache: &cache, + identity: identity.clone(), + index: index.clone(), + query: &query, + partitions: partitions.clone(), + centroid_distances: centroid_distances.clone(), + prefilter: Arc::new(NoFilter), + segment_mask: None, + metrics: Arc::new(NoOpMetricsCollector), + parallelism: 2, + }; + + let cold = search(params()).await.unwrap(); + assert!(!cold.was_hit); + assert_eq!(row_ids(&cold.batch), vec![11, 1]); + + index.return_malformed_hit_shape.store(1, Ordering::Relaxed); + let fallback = search(params()).await.unwrap(); + assert!(!fallback.was_hit); + assert_eq!(row_ids(&fallback.batch), vec![11, 1]); + assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); + + index.return_malformed_hit_shape.store(0, Ordering::Relaxed); + let replaced = search(params()).await.unwrap(); + assert!(replaced.was_hit); + assert_eq!(row_ids(&replaced.batch), vec![0, 11]); + } +} diff --git a/rust/lance/src/session/index_caches.rs b/rust/lance/src/session/index_caches.rs index 23922c7a4b0..86c59cede6a 100644 --- a/rust/lance/src/session/index_caches.rs +++ b/rust/lance/src/session/index_caches.rs @@ -10,14 +10,25 @@ //! │ │ //! └────┴──► Index-specific cache (prefixed by index UUID and FRI UUID) -use std::{borrow::Cow, ops::Deref, sync::Arc}; +use std::{borrow::Cow, collections::HashSet, ops::Deref, sync::Arc}; +use arrow_array::{ + cast::AsArray, + types::{Float16Type, Float32Type, Float64Type}, +}; +use arrow_schema::DataType; use lance_core::cache::{CacheKey, CacheKeySchema, KeyBuilder, LanceCache}; use lance_core::deepsize::{Context, DeepSizeOf}; use lance_index::frag_reuse::CompactFragReuseIndex; +use lance_index::vector::quantizer::QuantizationType; +use lance_index::vector::v3::subindex::SubIndexType; +use lance_index::vector::{ApproxMode, Query}; +use lance_linalg::distance::DistanceType; use lance_table::format::IndexMetadata; use uuid::Uuid; +use crate::{Error, Result}; + /// A type-safe wrapper around a LanceCache that enforces namespaces for index data. pub struct GlobalIndexCache(pub(super) LanceCache); @@ -89,6 +100,491 @@ pub(crate) fn write_index_identity(builder: &mut KeyBuilder, uuid: &Uuid, fri_uu } } +/// One candidate stored by the experimental vector-results cache. +/// +/// The offset is meaningful only inside `partition_id` of the exact index +/// segment identified by [`VectorResultsCacheIdentity`]. It must still be +/// bounds-checked against the loaded partition before use. +#[derive(Clone, Copy, Debug, DeepSizeOf, Eq, Hash, PartialEq)] +pub struct CachedVectorCandidate { + partition_id: u32, + offset_in_partition: u32, +} + +impl CachedVectorCandidate { + /// Create a partition-local candidate identity. + pub fn new(partition_id: u32, offset_in_partition: u32) -> Self { + Self { + partition_id, + offset_in_partition, + } + } + + /// Return the IVF partition containing this candidate. + pub fn partition_id(&self) -> u32 { + self.partition_id + } + + /// Return the candidate's local offset inside its IVF partition. + pub fn offset_in_partition(&self) -> u32 { + self.offset_in_partition + } +} + +/// Domain-separates exact query fingerprints from other BLAKE3 uses. +const VECTOR_RESULTS_QUERY_FINGERPRINT_CONTEXT: &str = + "lance.vector-results-cache-query-fingerprint.v1"; + +/// Complete compatibility identity for one reusable vector candidate pool. +/// +/// The exact query values are represented by a cryptographic fingerprint. A +/// candidate pool can therefore be reused only by the same query bit pattern; +/// the ordered IVF partition list remains in the identity to bind the pool to +/// its search shape. Unsupported query forms are rejected by [`Self::try_new`] +/// instead of being represented by a partial key. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct VectorResultsCacheIdentity { + store_identity: String, + dataset_read_identity: String, + dataset_version: u64, + index_uuid: Uuid, + frag_reuse_uuid: Option, + index_dataset_version: u64, + index_version: i32, + index_base_id: Option, + index_fields: Vec, + column: String, + metric_variant: u32, + sub_index_variant: u32, + quantization_variant: u32, + approx_mode_variant: u32, + vector_type_variant: u32, + dimension: u32, + query_fingerprint: [u8; 32], + nprobes: u32, + partition_ids: Vec, + result_limit: u32, + candidate_limit: u32, +} + +impl DeepSizeOf for VectorResultsCacheIdentity { + fn deep_size_of_children(&self, context: &mut Context) -> usize { + self.store_identity.deep_size_of_children(context) + + self.dataset_read_identity.deep_size_of_children(context) + + self.index_fields.deep_size_of_children(context) + + self.column.deep_size_of_children(context) + + self.partition_ids.deep_size_of_children(context) + } +} + +impl VectorResultsCacheIdentity { + /// Build a fail-closed cache identity for the initial supported query scope. + /// + /// `dataset_read_identity` must identify the exact manifest being read, not + /// merely the dataset URI. `has_prefilter` represents a user/scalar-index + /// prefilter; ordinary deletion filtering is safe because the exact dataset + /// read identity and version are included in the key. + #[allow(clippy::too_many_arguments)] + pub fn try_new( + store_identity: &str, + dataset_read_identity: &str, + dataset_version: u64, + index: &IndexMetadata, + frag_reuse_uuid: Option, + query: &Query, + index_metric: DistanceType, + sub_index_type: SubIndexType, + quantization_type: QuantizationType, + partition_ids: &[u32], + candidate_limit: usize, + has_prefilter: bool, + has_overlay: bool, + ) -> Result { + if store_identity.is_empty() { + return Err(Error::invalid_input( + "vector results cache requires a non-empty object-store identity", + )); + } + if dataset_read_identity.is_empty() { + return Err(Error::invalid_input( + "vector results cache requires an exact dataset read identity", + )); + } + if index.fragment_bitmap.is_none() { + return Err(Error::not_supported(format!( + "vector results cache requires known fragment coverage for index {}", + index.uuid + ))); + } + if index.index_details.is_none() { + return Err(Error::not_supported(format!( + "vector results cache requires current index details for index {}", + index.uuid + ))); + } + if index.fields.is_empty() { + return Err(Error::invalid_input(format!( + "vector results cache cannot identify an index with no fields: {}", + index.uuid + ))); + } + if has_prefilter { + return Err(Error::not_supported( + "vector results cache does not support prefiltered queries", + )); + } + if has_overlay { + return Err(Error::not_supported( + "vector results cache does not support data overlays", + )); + } + if query.lower_bound.is_some() || query.upper_bound.is_some() { + return Err(Error::not_supported( + "vector results cache does not support distance bounds", + )); + } + if query.ef.is_some() { + return Err(Error::not_supported( + "vector results cache does not support HNSW query parameters", + )); + } + if !query.use_index { + return Err(Error::invalid_input( + "vector results cache requires index search to be enabled", + )); + } + if query.column.is_empty() { + return Err(Error::invalid_input( + "vector results cache requires a non-empty vector column", + )); + } + if query.key.is_empty() || query.key.null_count() != 0 { + return Err(Error::invalid_input(format!( + "vector results cache requires a non-empty, non-null query vector, got length {} with {} nulls", + query.key.len(), + query.key.null_count() + ))); + } + + let mut query_hasher = + blake3::Hasher::new_derive_key(VECTOR_RESULTS_QUERY_FINGERPRINT_CONTEXT); + let (vector_type_variant, is_finite) = match query.key.data_type() { + DataType::Float16 => { + let values = query.key.as_primitive::().values(); + query_hasher.update(values.inner().as_slice()); + (0, values.iter().all(|value| value.is_finite())) + } + DataType::Float32 => { + let values = query.key.as_primitive::().values(); + query_hasher.update(values.inner().as_slice()); + (1, values.iter().all(|value| value.is_finite())) + } + DataType::Float64 => { + let values = query.key.as_primitive::().values(); + query_hasher.update(values.inner().as_slice()); + (2, values.iter().all(|value| value.is_finite())) + } + data_type => { + return Err(Error::not_supported(format!( + "vector results cache supports one float vector, got query type {data_type}" + ))); + } + }; + if !is_finite { + return Err(Error::invalid_input( + "vector results cache requires a finite query vector", + )); + } + let dimension = u32::try_from(query.key.len()).map_err(|_| { + Error::invalid_input(format!( + "query vector dimension {} exceeds the vector results cache limit", + query.key.len() + )) + })?; + let query_fingerprint = *query_hasher.finalize().as_bytes(); + + let metric_type = query.metric_type.unwrap_or(index_metric); + if metric_type != index_metric { + return Err(Error::invalid_input(format!( + "query metric {metric_type} does not match index metric {index_metric} for vector results cache" + ))); + } + let metric_variant = match metric_type { + DistanceType::L2 => 0, + DistanceType::Cosine => 1, + DistanceType::Dot => 2, + DistanceType::Hamming => { + return Err(Error::not_supported( + "vector results cache does not support Hamming distance", + )); + } + }; + let sub_index_variant = match sub_index_type { + SubIndexType::Flat => 0, + SubIndexType::Hnsw => { + return Err(Error::not_supported( + "vector results cache initially supports only flat IVF sub-indices", + )); + } + }; + let quantization_variant = match quantization_type { + QuantizationType::Scalar => 0, + QuantizationType::Rabit => { + if query.approx_mode != ApproxMode::Accurate { + return Err(Error::not_supported(format!( + "vector results cache supports RQ only with ApproxMode::Accurate; got {:?}", + query.approx_mode + ))); + } + 1 + } + other => { + return Err(Error::not_supported(format!( + "vector results cache does not support {other} quantization" + ))); + } + }; + let approx_mode_variant = match query.approx_mode { + ApproxMode::Fast => 0, + ApproxMode::Normal => 1, + ApproxMode::Accurate => 2, + }; + + let Some(maximum_nprobes) = query.maximum_nprobes else { + return Err(Error::not_supported( + "vector results cache requires a fixed maximum_nprobes", + )); + }; + if query.minimum_nprobes != maximum_nprobes { + return Err(Error::not_supported(format!( + "vector results cache requires fixed nprobes, got minimum_nprobes={} and maximum_nprobes={maximum_nprobes}", + query.minimum_nprobes + ))); + } + if maximum_nprobes == 0 || maximum_nprobes != partition_ids.len() { + return Err(Error::invalid_input(format!( + "vector results cache nprobes {maximum_nprobes} does not match {} searched partitions", + partition_ids.len() + ))); + } + if partition_ids.iter().copied().collect::>().len() != partition_ids.len() { + return Err(Error::invalid_input(format!( + "vector results cache requires unique partition ids, got {partition_ids:?}" + ))); + } + let nprobes = u32::try_from(maximum_nprobes).map_err(|_| { + Error::invalid_input(format!( + "nprobes {maximum_nprobes} exceeds the vector results cache limit" + )) + })?; + + let refine_factor = query.refine_factor.unwrap_or(1) as usize; + let result_limit = query.k.checked_mul(refine_factor).ok_or_else(|| { + Error::invalid_input(format!( + "vector results cache result limit overflows: k={} refine_factor={refine_factor}", + query.k + )) + })?; + if result_limit == 0 || candidate_limit < result_limit { + return Err(Error::invalid_input(format!( + "vector results cache candidate limit {candidate_limit} must be at least the requested result limit {result_limit}" + ))); + } + let result_limit = u32::try_from(result_limit).map_err(|_| { + Error::invalid_input(format!( + "result limit {result_limit} exceeds the vector results cache limit" + )) + })?; + let candidate_limit = u32::try_from(candidate_limit).map_err(|_| { + Error::invalid_input(format!( + "candidate limit {candidate_limit} exceeds the vector results cache limit" + )) + })?; + + Ok(Self { + store_identity: store_identity.to_owned(), + dataset_read_identity: dataset_read_identity.to_owned(), + dataset_version, + index_uuid: index.uuid, + frag_reuse_uuid, + index_dataset_version: index.dataset_version, + index_version: index.index_version, + index_base_id: index.base_id, + index_fields: index.fields.clone(), + column: query.column.clone(), + metric_variant, + sub_index_variant, + quantization_variant, + approx_mode_variant, + vector_type_variant, + dimension, + query_fingerprint, + nprobes, + partition_ids: partition_ids.to_vec(), + result_limit, + candidate_limit, + }) + } + + /// Return the exact ordered partition list represented by this cache key. + pub fn partition_ids(&self) -> &[u32] { + &self.partition_ids + } + + /// Return the maximum number of candidates stored in a compatible entry. + pub fn candidate_limit(&self) -> usize { + self.candidate_limit as usize + } + + /// Return the number of scored candidates emitted for this query shape. + pub fn result_limit(&self) -> usize { + self.result_limit as usize + } +} + +/// In-memory candidate pool stored under a [`VectorResultsCacheIdentity`]. +/// +/// No codec is registered yet, so this does not create a persistent cache or a +/// file-format compatibility surface. +#[derive(Clone, Debug, DeepSizeOf)] +pub struct VectorResultsCacheEntry { + identity: VectorResultsCacheIdentity, + candidates: Vec, +} + +impl VectorResultsCacheEntry { + /// Create an entry after validating its candidate count and partitions. + pub fn try_new( + identity: VectorResultsCacheIdentity, + candidates: Vec, + ) -> Result { + let entry = Self { + identity, + candidates, + }; + if !entry.has_valid_shape() { + return Err(Error::invalid_input( + "vector results cache entry contains incompatible or duplicate candidates", + )); + } + Ok(entry) + } + + /// Return true only if the entry exactly matches the requested identity and + /// all candidates satisfy the identity's structural constraints. + pub fn is_compatible_with(&self, identity: &VectorResultsCacheIdentity) -> bool { + self.identity == *identity && self.has_valid_shape() + } + + /// Return the validated partition-local candidates. + pub fn candidates(&self) -> &[CachedVectorCandidate] { + &self.candidates + } + + fn has_valid_shape(&self) -> bool { + if self.candidates.len() > self.identity.candidate_limit() { + return false; + } + let valid_partitions = self + .identity + .partition_ids() + .iter() + .copied() + .collect::>(); + let unique_candidates = self.candidates.iter().copied().collect::>(); + unique_candidates.len() == self.candidates.len() + && self + .candidates + .iter() + .all(|candidate| valid_partitions.contains(&candidate.partition_id())) + } +} + +impl CacheKey for VectorResultsCacheIdentity { + type ValueType = VectorResultsCacheEntry; + + fn key(&self) -> Cow<'_, str> { + let partition_ids = self + .partition_ids + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let base_id = self.index_base_id.map_or(-1, i64::from); + Cow::Owned(format!( + "{store_len}:{store}/{read_len}:{read}/{dataset_version}/{index_uuid}/{frag_reuse_uuid:?}/{index_dataset_version}/{index_version}/{base_id}/{index_fields:?}/{column_len}:{column}/{metric_variant}/{sub_index_variant}/{quantization_variant}/{approx_mode_variant}/{vector_type_variant}/{dimension}/{query_fingerprint:?}/{nprobes}/{partition_ids}/{result_limit}/{candidate_limit}", + store_len = self.store_identity.len(), + store = self.store_identity, + read_len = self.dataset_read_identity.len(), + read = self.dataset_read_identity, + dataset_version = self.dataset_version, + index_uuid = self.index_uuid, + frag_reuse_uuid = self.frag_reuse_uuid, + index_dataset_version = self.index_dataset_version, + index_version = self.index_version, + index_fields = self.index_fields, + column_len = self.column.len(), + column = self.column, + metric_variant = self.metric_variant, + sub_index_variant = self.sub_index_variant, + quantization_variant = self.quantization_variant, + approx_mode_variant = self.approx_mode_variant, + vector_type_variant = self.vector_type_variant, + dimension = self.dimension, + query_fingerprint = self.query_fingerprint, + nprobes = self.nprobes, + result_limit = self.result_limit, + candidate_limit = self.candidate_limit, + )) + } + + fn type_name() -> &'static str { + "VectorResultsCacheEntry" + } + + fn stable_type_id() -> &'static str { + "lance.VectorResultsCacheEntry" + } + + fn schema() -> CacheKeySchema { + CacheKeySchema::new("lance.vector-results-cache-key", 2) + } + + fn write_key(&self, builder: &mut KeyBuilder) { + builder.write_str(&self.store_identity); + builder.write_str(&self.dataset_read_identity); + builder.write_u64(self.dataset_version); + write_index_identity(builder, &self.index_uuid, self.frag_reuse_uuid.as_ref()); + builder.write_u64(self.index_dataset_version); + builder.write_i32(self.index_version); + if let Some(base_id) = self.index_base_id { + builder.write_some(); + builder.write_u32(base_id); + } else { + builder.write_none(); + } + builder.write_sequence_len(self.index_fields.len() as u64); + for field_id in &self.index_fields { + builder.write_i32(*field_id); + } + builder.write_str(&self.column); + builder.write_variant(self.metric_variant); + builder.write_variant(self.sub_index_variant); + builder.write_variant(self.quantization_variant); + builder.write_variant(self.approx_mode_variant); + builder.write_variant(self.vector_type_variant); + builder.write_u32(self.dimension); + builder.write_fixed_bytes(&self.query_fingerprint); + builder.write_u32(self.nprobes); + builder.write_sequence_len(self.partition_ids.len() as u64); + for partition_id in &self.partition_ids { + builder.write_u32(*partition_id); + } + builder.write_u32(self.result_limit); + builder.write_u32(self.candidate_limit); + } +} + // Cache key types for type-safe cache access #[derive(Debug)] @@ -207,6 +703,89 @@ impl CacheKey for ScalarIndexDetailsKey<'_> { #[cfg(test)] mod tests { use super::*; + use arrow_array::{ArrayRef, Float32Array}; + use lance_core::cache::{CacheNamespace, InternalCacheKey}; + use lance_index::vector::DEFAULT_QUERY_PARALLELISM; + + fn vector_index_metadata() -> IndexMetadata { + IndexMetadata { + uuid: Uuid::from_u128(0x11111111_2222_3333_4444_555555555555), + fields: vec![7], + covering_fields: vec![], + name: "vector_idx".to_string(), + dataset_version: 5, + fragment_bitmap: Some([1_u32, 3, 8].into_iter().collect()), + index_details: Some(Arc::new(prost_types::Any { + type_url: "type.googleapis.com/lance.table.VectorIndexDetails".to_string(), + value: vec![1], + })), + index_version: 3, + created_at: None, + base_id: Some(2), + files: None, + } + } + + fn vector_query() -> Query { + Query { + column: "vector".to_string(), + key: Arc::new(Float32Array::from(vec![0.25, 0.5, 0.75])) as ArrayRef, + k: 10, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 3, + maximum_nprobes: Some(3), + ef: None, + refine_factor: Some(2), + metric_type: Some(DistanceType::Cosine), + use_index: true, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: f32::NAN, + approx_mode: ApproxMode::Accurate, + } + } + + fn vector_results_identity_for(query: &Query) -> VectorResultsCacheIdentity { + VectorResultsCacheIdentity::try_new( + "s3$account-a", + "_versions/17.manifest", + 17, + &vector_index_metadata(), + Some(Uuid::from_u128(0xaaaaaaaa_bbbb_cccc_dddd_eeeeeeeeeeee)), + query, + DistanceType::Cosine, + SubIndexType::Flat, + QuantizationType::Rabit, + &[2, 5, 9], + 100, + false, + false, + ) + .unwrap() + } + + fn vector_results_identity() -> VectorResultsCacheIdentity { + vector_results_identity_for(&vector_query()) + } + + fn physical_key(identity: &VectorResultsCacheIdentity) -> InternalCacheKey { + let mut builder = KeyBuilder::new( + CacheNamespace::root(), + VectorResultsCacheIdentity::stable_type_id(), + VectorResultsCacheIdentity::schema(), + ); + identity.write_key(&mut builder); + builder.finish() + } + + fn assert_identity_field_isolated( + base: &VectorResultsCacheIdentity, + change: impl FnOnce(&mut VectorResultsCacheIdentity), + ) { + let mut changed = base.clone(); + change(&mut changed); + assert_ne!(physical_key(base), physical_key(&changed)); + } #[test] fn index_metadata_key_isolates_object_store_identity() { @@ -239,4 +818,284 @@ mod tests { assert_ne!(first.key(), second.key()); } + + #[test] + fn vector_results_cache_key_isolates_every_identity_axis() { + let base = vector_results_identity(); + + assert_identity_field_isolated(&base, |identity| { + identity.store_identity.push_str("-rotated") + }); + assert_identity_field_isolated(&base, |identity| { + identity.dataset_read_identity.push_str("-detached") + }); + assert_identity_field_isolated(&base, |identity| identity.dataset_version += 1); + assert_identity_field_isolated(&base, |identity| identity.index_uuid = Uuid::new_v4()); + assert_identity_field_isolated(&base, |identity| identity.frag_reuse_uuid = None); + assert_identity_field_isolated(&base, |identity| identity.index_dataset_version += 1); + assert_identity_field_isolated(&base, |identity| identity.index_version += 1); + assert_identity_field_isolated(&base, |identity| identity.index_base_id = None); + assert_identity_field_isolated(&base, |identity| identity.index_fields.push(8)); + assert_identity_field_isolated(&base, |identity| identity.column.push_str("_new")); + assert_identity_field_isolated(&base, |identity| identity.metric_variant += 1); + assert_identity_field_isolated(&base, |identity| identity.sub_index_variant += 1); + assert_identity_field_isolated(&base, |identity| identity.quantization_variant += 1); + assert_identity_field_isolated(&base, |identity| identity.approx_mode_variant += 1); + assert_identity_field_isolated(&base, |identity| identity.vector_type_variant += 1); + assert_identity_field_isolated(&base, |identity| identity.dimension += 1); + assert_identity_field_isolated(&base, |identity| identity.query_fingerprint[0] ^= 1); + assert_identity_field_isolated(&base, |identity| identity.nprobes += 1); + assert_identity_field_isolated(&base, |identity| identity.partition_ids.swap(0, 1)); + assert_identity_field_isolated(&base, |identity| identity.result_limit += 1); + assert_identity_field_isolated(&base, |identity| identity.candidate_limit += 1); + } + + #[test] + fn vector_results_cache_key_isolates_exact_query_values() { + let first_query = vector_query(); + let mut second_query = first_query.clone(); + second_query.key = Arc::new(Float32Array::from(vec![0.25, 0.5, 0.76])); + + let first = vector_results_identity_for(&first_query); + let second = vector_results_identity_for(&second_query); + + assert_eq!(first.partition_ids(), second.partition_ids()); + assert_ne!(first.query_fingerprint, second.query_fingerprint); + assert_ne!(physical_key(&first), physical_key(&second)); + } + + #[test] + fn vector_results_cache_identity_rejects_unsafe_query_forms() { + let index = vector_index_metadata(); + let make_identity = |query: &Query, + index: &IndexMetadata, + sub_index_type, + quantization_type, + partition_ids: &[u32], + has_prefilter, + has_overlay| { + VectorResultsCacheIdentity::try_new( + "s3$account-a", + "_versions/17.manifest", + 17, + index, + None, + query, + DistanceType::Cosine, + sub_index_type, + quantization_type, + partition_ids, + 100, + has_prefilter, + has_overlay, + ) + }; + + let query = vector_query(); + assert!( + make_identity( + &query, + &index, + SubIndexType::Flat, + QuantizationType::Rabit, + &[2, 5, 9], + true, + false, + ) + .is_err() + ); + assert!( + make_identity( + &query, + &index, + SubIndexType::Flat, + QuantizationType::Rabit, + &[2, 5, 9], + false, + true, + ) + .is_err() + ); + + let mut bounded = query.clone(); + bounded.lower_bound = Some(0.1); + let mut adaptive = query.clone(); + adaptive.maximum_nprobes = None; + let mut mismatched_metric = query.clone(); + mismatched_metric.metric_type = Some(DistanceType::Dot); + let mut hnsw_query = query.clone(); + hnsw_query.ef = Some(64); + let mut empty_query = query.clone(); + empty_query.key = Arc::new(Float32Array::from(Vec::::new())); + let mut null_query = query.clone(); + null_query.key = Arc::new(Float32Array::from(vec![Some(0.25), None, Some(0.75)])); + let mut nan_query = query.clone(); + nan_query.key = Arc::new(Float32Array::from(vec![0.25, f32::NAN, 0.75])); + let mut infinite_query = query.clone(); + infinite_query.key = Arc::new(Float32Array::from(vec![0.25, f32::INFINITY, 0.75])); + for unsupported_query in [ + bounded, + adaptive, + mismatched_metric, + hnsw_query, + empty_query, + null_query, + nan_query, + infinite_query, + ] { + assert!( + make_identity( + &unsupported_query, + &index, + SubIndexType::Flat, + QuantizationType::Rabit, + &[2, 5, 9], + false, + false, + ) + .is_err() + ); + } + + for approx_mode in [ApproxMode::Fast, ApproxMode::Normal] { + let mut rq_query = query.clone(); + rq_query.approx_mode = approx_mode; + let error = make_identity( + &rq_query, + &index, + SubIndexType::Flat, + QuantizationType::Rabit, + &[2, 5, 9], + false, + false, + ) + .unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. })); + assert!( + error + .to_string() + .contains("supports RQ only with ApproxMode::Accurate") + ); + assert!( + make_identity( + &rq_query, + &index, + SubIndexType::Flat, + QuantizationType::Scalar, + &[2, 5, 9], + false, + false, + ) + .is_ok(), + "SQ ignores the RQ-specific approximation policy" + ); + } + + assert!( + make_identity( + &query, + &index, + SubIndexType::Hnsw, + QuantizationType::Scalar, + &[2, 5, 9], + false, + false, + ) + .is_err() + ); + assert!( + make_identity( + &query, + &index, + SubIndexType::Flat, + QuantizationType::Product, + &[2, 5, 9], + false, + false, + ) + .is_err() + ); + assert!( + make_identity( + &query, + &index, + SubIndexType::Flat, + QuantizationType::Rabit, + &[2, 2, 9], + false, + false, + ) + .is_err() + ); + + let mut unknown_coverage = index; + unknown_coverage.fragment_bitmap = None; + assert!( + make_identity( + &query, + &unknown_coverage, + SubIndexType::Flat, + QuantizationType::Rabit, + &[2, 5, 9], + false, + false, + ) + .is_err() + ); + } + + #[test] + fn vector_results_cache_entry_validates_candidate_shape() { + let identity = vector_results_identity(); + let candidates = vec![ + CachedVectorCandidate::new(2, 10), + CachedVectorCandidate::new(5, 20), + CachedVectorCandidate::new(9, 30), + ]; + let entry = VectorResultsCacheEntry::try_new(identity.clone(), candidates).unwrap(); + assert!(entry.is_compatible_with(&identity)); + assert_eq!(entry.candidates()[1].partition_id(), 5); + assert_eq!(entry.candidates()[1].offset_in_partition(), 20); + + let mut different_identity = identity.clone(); + different_identity.dataset_version += 1; + assert!(!entry.is_compatible_with(&different_identity)); + assert!( + VectorResultsCacheEntry::try_new( + identity.clone(), + vec![CachedVectorCandidate::new(7, 10)], + ) + .is_err() + ); + assert!( + VectorResultsCacheEntry::try_new( + identity, + vec![ + CachedVectorCandidate::new(2, 10), + CachedVectorCandidate::new(2, 10), + ], + ) + .is_err() + ); + } + + #[tokio::test] + async fn vector_results_cache_identity_produces_typed_cache_misses() { + let cache = LanceCache::with_capacity(4096); + let identity = vector_results_identity(); + let entry = Arc::new( + VectorResultsCacheEntry::try_new( + identity.clone(), + vec![CachedVectorCandidate::new(2, 10)], + ) + .unwrap(), + ); + cache.insert_with_key(&identity, entry.clone()).await; + let cached = cache.get_with_key(&identity).await.unwrap(); + assert!(cached.is_compatible_with(&identity)); + + let mut other_version = identity; + other_version.dataset_version += 1; + assert!(cache.get_with_key(&other_version).await.is_none()); + } } From c29a19478ef59dc466bc82f4e8507660c9865639 Mon Sep 17 00:00:00 2001 From: Sergey Troshkov Date: Tue, 25 Aug 2026 16:43:46 +0700 Subject: [PATCH 2/4] feat: add opt-in quantized vector refinement --- Cargo.lock | 1 - java/lance-jni/Cargo.lock | 1 - java/lance-jni/src/blocking_scanner.rs | 6 + java/lance-jni/src/utils.rs | 4 + java/src/main/java/org/lance/ipc/Query.java | 25 + java/src/test/java/org/lance/JNITest.java | 37 +- python/Cargo.lock | 1 - python/python/lance/dataset.py | 18 + python/python/tests/test_vector_index.py | 23 + python/src/dataset.rs | 30 + rust/lance-index/src/vector.rs | 7 +- rust/lance/Cargo.toml | 1 - rust/lance/src/dataset/scanner.rs | 233 +++- rust/lance/src/index/vector/ivf/v2.rs | 97 +- rust/lance/src/io/exec.rs | 1 - rust/lance/src/io/exec/knn.rs | 1366 +++---------------- rust/lance/src/io/exec/knn_results_cache.rs | 1162 ---------------- rust/lance/src/session/index_caches.rs | 857 +----------- 18 files changed, 589 insertions(+), 3281 deletions(-) delete mode 100644 rust/lance/src/io/exec/knn_results_cache.rs diff --git a/Cargo.lock b/Cargo.lock index 484cd942489..f420aebfb36 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4422,7 +4422,6 @@ dependencies = [ "aws-credential-types", "aws-sdk-dynamodb", "aws-sdk-s3", - "blake3", "byteorder", "bytes", "chrono", diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index b2c93243673..96c25e25691 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -3805,7 +3805,6 @@ dependencies = [ "async-recursion", "async-trait", "async_cell", - "blake3", "byteorder", "bytes", "chrono", 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/Cargo.lock b/python/Cargo.lock index cb333d4326a..0dfc7892cf9 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -3973,7 +3973,6 @@ dependencies = [ "async-trait", "async_cell", "aws-sdk-dynamodb", - "blake3", "byteorder", "bytes", "chrono", diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index c70f845bb38..b4383a5218d 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -7088,6 +7088,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, @@ -7122,6 +7123,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, @@ -7133,6 +7140,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, @@ -8341,6 +8349,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, @@ -8376,6 +8385,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 @@ -8456,6 +8469,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 @@ -8486,6 +8503,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 ab5f26379b6..49e2d9976a6 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") @@ -2626,6 +2637,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 35324413323..41115cebb06 100644 --- a/rust/lance-index/src/vector.rs +++ b/rust/lance-index/src/vector.rs @@ -269,8 +269,8 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index { /// Search one partition and preserve the local offsets of emitted rows. /// - /// Implementations may reject this operation. The initial results-cache - /// integration uses it only for flat IVF_RQ and IVF_SQ sub-indices. + /// 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, @@ -287,8 +287,7 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug + Index { /// /// The returned rows preserve the order of `offsets_in_partition`. This /// operation does not apply a pre-filter or distance bounds; callers must - /// validate cached candidate identity and apply query result shaping around - /// it. + /// validate candidate identity and apply query result shaping around it. async fn score_partition_candidates( &self, _partition_id: usize, diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index d4b7fc4e342..417d41a5fe1 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -45,7 +45,6 @@ arrow-schema = { workspace = true } arrow-select = { workspace = true } async-recursion.workspace = true async-trait.workspace = true -blake3.workspace = true byteorder.workspace = true bytes.workspace = true chrono.workspace = true diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 28f4d341470..468c13ab6a1 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, @@ -5806,9 +5872,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 +5956,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 +6904,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 +6967,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())?, @@ -7266,6 +7361,7 @@ pub mod test_dataset { IndexType, scalar::{ScalarIndexParams, inverted::tokenizer::InvertedIndexParams}, vector::{ + bq::{RQBuildParams, RQRotationType}, hnsw::builder::HnswBuildParams, ivf::IvfBuildParams, kmeans::{KMeansParams, train_kmeans}, @@ -7407,6 +7503,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 @@ -16283,6 +16397,113 @@ 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_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 3d3e9a6d8a6..0218d2bf56b 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -2270,12 +2270,7 @@ impl VectorIndex for IVFInd metrics: &dyn MetricsCollector, candidate_limit: usize, ) -> Result { - if S::name() != "FLAT" - || !matches!( - Q::quantization_type(), - QuantizationType::Rabit | QuantizationType::Scalar - ) - { + 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(), @@ -2350,12 +2345,7 @@ impl VectorIndex for IVFInd offsets_in_partition: &[u32], metrics: &dyn MetricsCollector, ) -> Result { - if S::name() != "FLAT" - || !matches!( - Q::quantization_type(), - QuantizationType::Rabit | QuantizationType::Scalar - ) - { + 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(), @@ -6332,65 +6322,37 @@ mod tests { } #[rstest] - #[case::rq4_l2(DistanceType::L2, Some(4), false, DIM)] - #[case::rq4_cosine(DistanceType::Cosine, Some(4), false, DIM)] - #[case::rq4_dot(DistanceType::Dot, Some(4), false, DIM)] - #[case::rq8_l2(DistanceType::L2, Some(8), false, DIM)] - #[case::rq8_cosine(DistanceType::Cosine, Some(8), false, DIM)] - #[case::rq8_dot(DistanceType::Dot, Some(8), false, DIM)] - #[case::sq8_l2(DistanceType::L2, None, true, DIM)] - #[case::sq8_cosine(DistanceType::Cosine, None, true, DIM)] - #[case::sq8_dot(DistanceType::Dot, None, true, DIM)] - #[case::rq4_dot_768(DistanceType::Dot, Some(4), false, 768)] - #[case::sq8_l2_3072(DistanceType::L2, None, true, 3072)] - #[case::rq8_cosine_4096(DistanceType::Cosine, Some(8), false, 4096)] + #[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] rq_num_bits: Option, - #[case] has_negative_values: bool, + #[case] num_bits: u8, #[case] dimension: usize, ) { let test_dir = TempStrDir::default(); let test_uri = test_dir.as_str(); - let range = if has_negative_values { - -1.0..1.0 - } else { - 0.0..1.0 - }; 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, range).await; + generate_f32_test_dataset_with_shape(test_uri, num_rows, dimension, 0.0..1.0).await; let ivf_params = IvfBuildParams::new(4); - let params = if let Some(num_bits) = rq_num_bits { - VectorIndexParams::with_ivf_rq_params( - distance_type, - ivf_params, - RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), - ) - } else { - VectorIndexParams::with_ivf_sq_params( - distance_type, - ivf_params, - SQBuildParams::default(), - ) - }; + 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(); - if rq_num_bits.is_none() { - let obj_store = Arc::new(ObjectStore::local()); - let scheduler = ScanScheduler::new(obj_store, SchedulerConfig::default_for_testing()); - let metadata = get_sq_metadata(&dataset, scheduler, &indices[0].uuid.to_string()).await; - assert!( - metadata.bounds.start < 0.0 && metadata.bounds.end > 0.0, - "SQ bounds should be learned from both negative and positive values, got {:?}", - metadata.bounds - ); - } let index = dataset .open_vector_index("vector", &indices[0].uuid, &NoOpMetricsCollector) .await @@ -6532,12 +6494,11 @@ mod tests { } #[rstest] - #[case::sq8(None)] - #[case::rq4(Some(4))] - #[case::rq8(Some(8))] + #[case::rq4(4)] + #[case::rq8(8)] #[tokio::test] async fn test_selected_scoring_excludes_null_and_non_finite_stored_vectors( - #[case] rq_num_bits: Option, + #[case] num_bits: u8, ) { const EDGE_DIM: usize = 8; let test_dir = TempStrDir::default(); @@ -6588,19 +6549,11 @@ mod tests { .unwrap(), ); let ivf_params = IvfBuildParams::try_with_centroids(1, centroids).unwrap(); - let params = if let Some(num_bits) = rq_num_bits { - VectorIndexParams::with_ivf_rq_params( - DistanceType::L2, - ivf_params, - RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), - ) - } else { - VectorIndexParams::with_ivf_sq_params( - DistanceType::L2, - ivf_params, - SQBuildParams::default(), - ) - }; + 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 diff --git a/rust/lance/src/io/exec.rs b/rust/lance/src/io/exec.rs index f610f1b7f55..d37b58a238e 100644 --- a/rust/lance/src/io/exec.rs +++ b/rust/lance/src/io/exec.rs @@ -15,7 +15,6 @@ pub mod filtered_read; pub mod filtered_read_proto; pub mod fts; pub(crate) mod knn; -mod knn_results_cache; mod optimizer; mod projection; mod pushdown_scan; diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index c38f1982ef7..2e8ac22e482 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; @@ -69,7 +71,6 @@ use crate::index::vector::utils::{get_vector_type, validate_distance_type_for}; use crate::{Error, Result}; use lance_arrow::*; -use super::knn_results_cache; use super::row_addr_mask::MaskAndLoader; use super::utils::{ FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedRecordBatchStreamAdapter, PreFilterSource, @@ -109,21 +110,14 @@ impl AnnPartitionMetrics { pub struct AnnIndexMetrics { index_metrics: IndexMetrics, partitions_searched: Count, - results_cache_hits: Count, - results_cache_misses: Count, baseline_metrics: BaselineMetrics, } -const VECTOR_RESULTS_CACHE_HITS_METRIC: &str = "vector_results_cache_hits"; -const VECTOR_RESULTS_CACHE_MISSES_METRIC: &str = "vector_results_cache_misses"; - impl AnnIndexMetrics { pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { Self { index_metrics: IndexMetrics::new(metrics, partition), partitions_searched: metrics.new_count(PARTITIONS_SEARCHED_METRIC, partition), - results_cache_hits: metrics.new_count(VECTOR_RESULTS_CACHE_HITS_METRIC, partition), - results_cache_misses: metrics.new_count(VECTOR_RESULTS_CACHE_MISSES_METRIC, partition), baseline_metrics: BaselineMetrics::new(metrics, partition), } } @@ -1249,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(), @@ -1269,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)) } @@ -1541,8 +1539,8 @@ pub struct ANNIvfSubIndexExec { metrics: ExecutionPlanMetricsSet, - #[cfg(test)] - results_cache_enabled_override: Option, + /// Coarse-candidate overfetch factor for RQ Fast-to-Accurate refinement. + quantized_refine_factor: Option, } impl ANNIvfSubIndexExec { @@ -1575,8 +1573,7 @@ impl ANNIvfSubIndexExec { external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), - #[cfg(test)] - results_cache_enabled_override: None, + quantized_refine_factor: None, }) } @@ -1594,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 @@ -1613,29 +1621,6 @@ impl ANNIvfSubIndexExec { pub fn prefilter_source(&self) -> &PreFilterSource { &self.prefilter_source } - - fn is_results_cache_enabled(&self) -> bool { - #[cfg(test)] - if let Some(is_enabled) = self.results_cache_enabled_override { - return is_enabled; - } - knn_results_cache::is_enabled() - } - - #[cfg(test)] - fn copy_with_results_cache_enabled_for_testing(&self, is_enabled: bool) -> Result { - let mut copy = Self::try_new( - self.input.clone(), - self.dataset.clone(), - self.indices.clone(), - self.query.clone(), - self.prefilter_source.clone(), - )?; - copy.overlay_block = self.overlay_block.clone(); - copy.external_mask = self.external_mask.clone(); - copy.results_cache_enabled_override = Some(is_enabled); - Ok(copy) - } } impl DisplayAs for ANNIvfSubIndexExec { @@ -1653,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!( @@ -1663,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(()) } } } @@ -1864,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, @@ -1901,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) } @@ -1946,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 @@ -2039,7 +2172,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 = @@ -2088,6 +2221,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( @@ -2097,6 +2231,7 @@ impl ANNIvfSubIndexExec { pre_filter, metrics, seg_mask, + quantized_refine_factor, ) .await?; state.record_late_batch(batch.num_rows()); @@ -2124,12 +2259,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; @@ -2176,6 +2312,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, @@ -2184,6 +2321,7 @@ impl ANNIvfSubIndexExec { pre_filter, metrics, seg_mask, + quantized_refine_factor, ) .await?; state.record_batch(&batch); @@ -2250,8 +2388,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), - #[cfg(test)] - results_cache_enabled_override: self.results_cache_enabled_override, + quantized_refine_factor: self.quantized_refine_factor, } } else { return Err(DataFusionError::Internal( @@ -2270,6 +2407,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(); @@ -2288,19 +2426,6 @@ impl ExecutionPlan for ANNIvfSubIndexExec { HashMap::new() }); let prefilter_source = self.prefilter_source.clone(); - let has_prefilter = - !matches!(prefilter_source, PreFilterSource::None) || self.external_mask.is_some(); - let has_overlay = self.overlay_block.is_some(); - let results_cache_enabled = self.is_results_cache_enabled(); - let index_metadata_by_uuid = results_cache_enabled.then(|| { - Arc::new( - indices - .iter() - .cloned() - .map(|metadata| (metadata.uuid, metadata)) - .collect::>(), - ) - }); let metrics = Arc::new(AnnIndexMetrics::new(&self.metrics, partition)); let metrics_clone = metrics.clone(); let timer = Instant::now(); @@ -2375,7 +2500,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { let indices_by_uuid = indices_by_uuid.clone(); let state = state.clone(); let segment_bitmaps = segment_bitmaps.clone(); - let index_metadata_by_uuid = index_metadata_by_uuid.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(|| { @@ -2418,68 +2543,6 @@ impl ExecutionPlan for ANNIvfSubIndexExec { } None => None, }; - if results_cache_enabled - && let Some(index_metadata) = index_metadata_by_uuid - .as_ref() - .and_then(|metadata| metadata.get(&index_uuid)) - && let Some(identity) = knn_results_cache::identity_for_query( - ds.as_ref(), - index_metadata, - raw_index.as_ref(), - &query, - part_ids.values(), - has_prefilter, - has_overlay, - ) - .await - { - let parallelism = effective_query_parallelism( - &query, - raw_index.as_ref(), - target_partitions, - ); - let index_metrics: Arc = - Arc::new(metrics.index_metrics.clone()); - match knn_results_cache::search( - knn_results_cache::ResultsCacheSearchParams { - cache: &ds.index_cache.0, - identity, - index: raw_index.clone(), - query: &query, - partitions: part_ids.clone(), - centroid_distances: q_c_dists.clone(), - prefilter: pre_filter.clone(), - segment_mask: seg_mask.clone(), - metrics: index_metrics, - parallelism, - }, - ) - .await - { - Ok(cache_result) => { - if cache_result.was_hit { - metrics.results_cache_hits.add(1); - } else { - metrics.results_cache_misses.add(1); - metrics.partitions_searched.add(part_ids.len()); - } - metrics - .baseline_metrics - .record_output(cache_result.batch.num_rows()); - return DataFusionResult::Ok( - stream::once(async move { Ok(cache_result.batch) }).boxed(), - ); - } - Err(error) => { - // The feature is experimental and best-effort. Any cache-only - // failure falls through to the existing ANN path. - metrics.results_cache_misses.add(1); - log::debug!( - "falling back after vector results cache failure: {error}" - ); - } - } - } let early_search = Self::initial_search( raw_index.clone(), @@ -2491,6 +2554,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { state.clone(), target_partitions, seg_mask.clone(), + quantized_refine_factor, ); let late_search = Self::late_search( raw_index.clone(), @@ -2503,6 +2567,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { state, target_partitions, seg_mask, + quantized_refine_factor, ); DataFusionResult::Ok(early_search.chain(late_search).boxed()) } @@ -3119,27 +3184,19 @@ mod tests { }; use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; use async_trait::async_trait; - use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::error::Result as DataFusionResult; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use lance_core::deepsize::DeepSizeOf; use lance_core::utils::tempfile::TempStrDir; - use lance_datafusion::exec::{ExecutionStatsCallback, ExecutionSummaryCounts, execute_plan}; + use lance_datafusion::exec::{ExecutionStatsCallback, ExecutionSummaryCounts}; use lance_datafusion::utils::FIND_PARTITIONS_ELAPSED_METRIC; use lance_datagen::{BatchCount, RowCount, array}; - use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; - use lance_index::metrics::NoOpMetricsCollector; use lance_index::optimize::OptimizeOptions; - use lance_index::scalar::ScalarIndexParams; - use lance_index::vector::bq::{RQBuildParams, RQRotationType}; use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; use lance_index::vector::quantizer::QuantizationType; - use lance_index::vector::sq::builder::SQBuildParams; use lance_index::vector::v3::subindex::SubIndexType; - use lance_index::vector::{ - ApproxMode, DEFAULT_QUERY_PARALLELISM, PreparedPartitionSearchHandle, - }; + use lance_index::vector::{DEFAULT_QUERY_PARALLELISM, PreparedPartitionSearchHandle}; use lance_index::{Index, IndexType}; use lance_io::traits::Reader; use lance_linalg::distance::MetricType; @@ -3149,12 +3206,10 @@ mod tests { use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; - use crate::dataset::optimize::{CompactionOptions, compact_files}; use crate::dataset::{WriteMode, WriteParams}; use crate::index::vector::VectorIndexParams; use crate::index::vector::ivf::v2::STREAMING_SEARCH_BATCH_SIZE; use crate::io::exec::testing::TestingExec; - use crate::session::index_caches::{CachedVectorCandidate, VectorResultsCacheEntry}; fn base_query() -> Query { Query { @@ -3981,6 +4036,7 @@ mod tests { state, usize::MAX, None, + None, ) .try_collect::>() .await @@ -4032,6 +4088,7 @@ mod tests { state, usize::MAX, None, + None, ) .try_collect::>() .await @@ -4096,6 +4153,7 @@ mod tests { state.clone(), usize::MAX, None, + None, ) .try_collect::>() .await @@ -4185,6 +4243,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask.clone()), + None, ) .try_collect::>() .await @@ -4218,6 +4277,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + None, ) .try_collect::>() .await @@ -4280,6 +4340,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask.clone()), + None, ) .try_collect::>() .await @@ -4297,6 +4358,7 @@ mod tests { state.clone(), usize::MAX, Some(seg_mask), + None, ) .try_collect::>() .await @@ -4330,6 +4392,7 @@ mod tests { state.clone(), usize::MAX, None, + None, ) .try_collect::>(); @@ -4349,6 +4412,7 @@ mod tests { state, usize::MAX, None, + None, ) .try_collect::>(); @@ -4913,288 +4977,6 @@ mod tests { } } - struct ResultsCacheTestFixture { - dataset: Dataset, - centroids: ArrayRef, - query: ArrayRef, - metric_type: MetricType, - index_params: VectorIndexParams, - _tmp_dir: TempStrDir, - } - - impl ResultsCacheTestFixture { - const NUM_PARTITIONS: usize = 4; - const K: usize = 8; - const ROWS_PER_FRAGMENT: usize = 32; - const NUM_FRAGMENTS: usize = 4; - const INITIAL_ROWS: usize = Self::ROWS_PER_FRAGMENT * Self::NUM_FRAGMENTS; - const RANK_SCALE_STEP: f32 = 0.01; - const RANK_ORTHOGONAL_STEP: f32 = 0.005; - - fn data_batch(start_row: usize, num_rows: usize) -> RecordBatch { - let mut vectors = Vec::with_capacity(num_rows * 8); - for row in start_row..start_row + num_rows { - let centroid_id = row % Self::NUM_PARTITIONS; - let rank = row / Self::NUM_PARTITIONS; - // Both changes make larger ranks worse under L2, cosine, and dot, - // giving the quantized recall checks a reproducible top-k boundary. - let scale = 1.0 - rank as f32 * Self::RANK_SCALE_STEP; - let orthogonal_offset = rank as f32 * Self::RANK_ORTHOGONAL_STEP; - let mut vector = [0.0; 8]; - match centroid_id { - 0 => vector[0] = scale, - 1 => vector[1] = scale, - 2 => vector[0] = -scale, - 3 => vector[1] = -scale, - _ => unreachable!(), - } - vector[centroid_id + 2] = orthogonal_offset; - vectors.extend_from_slice(&vector); - } - let vectors = Arc::new( - FixedSizeListArray::try_new_from_values(Float32Array::from(vectors), 8).unwrap(), - ); - let row_ids = Arc::new(UInt64Array::from_iter_values( - start_row as u64..(start_row + num_rows) as u64, - )); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("vector", vectors.data_type().clone(), false), - ArrowField::new("row", DataType::UInt64, false), - ])); - RecordBatch::try_new(schema, vec![vectors, row_ids]).unwrap() - } - - async fn new(rq_num_bits: Option) -> Self { - Self::with_metric(rq_num_bits, MetricType::L2).await - } - - async fn with_metric(rq_num_bits: Option, metric_type: MetricType) -> Self { - let tmp_dir = TempStrDir::default(); - let centroids: ArrayRef = 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, // +x - 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // +y - -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // -x - 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // -y - ]), - 8, - ) - .unwrap(), - ); - let batches = (0..Self::NUM_FRAGMENTS) - .map(|fragment| { - Ok(Self::data_batch( - fragment * Self::ROWS_PER_FRAGMENT, - Self::ROWS_PER_FRAGMENT, - )) - }) - .collect::>(); - let schema = batches[0].as_ref().unwrap().schema(); - let reader = RecordBatchIterator::new(batches, schema); - let mut dataset = Dataset::write( - reader, - tmp_dir.as_str(), - Some(WriteParams { - max_rows_per_file: Self::ROWS_PER_FRAGMENT, - max_rows_per_group: Self::ROWS_PER_FRAGMENT, - ..Default::default() - }), - ) - .await - .unwrap(); - assert_eq!(dataset.get_fragments().len(), Self::NUM_FRAGMENTS); - - let ivf_params = IvfBuildParams::try_with_centroids( - Self::NUM_PARTITIONS, - Arc::new(centroids.as_fixed_size_list().clone()), - ) - .unwrap(); - let index_params = if let Some(num_bits) = rq_num_bits { - VectorIndexParams::with_ivf_rq_params( - metric_type, - ivf_params, - RQBuildParams::with_rotation_type(num_bits, RQRotationType::Fast), - ) - } else { - VectorIndexParams::with_ivf_sq_params( - metric_type, - ivf_params, - SQBuildParams::default(), - ) - }; - dataset - .create_index(&["vector"], IndexType::Vector, None, &index_params, false) - .await - .unwrap(); - - let dataset = Dataset::open(tmp_dir.as_str()).await.unwrap(); - let query = centroids.as_fixed_size_list().value(0); - Self { - dataset, - centroids, - query, - metric_type, - index_params, - _tmp_dir: tmp_dir, - } - } - - async fn vector_index_uuids(&self) -> Vec { - self.dataset - .load_indices() - .await - .unwrap() - .iter() - .filter(|index| index.name != FRAG_REUSE_INDEX_NAME) - .map(|index| index.uuid) - .collect() - } - - async fn fragment_reuse_uuid(&self) -> Option { - self.dataset - .load_indices() - .await - .unwrap() - .iter() - .find(|index| index.name == FRAG_REUSE_INDEX_NAME) - .map(|index| index.uuid) - } - - async fn append_index_segment(&mut self) -> Vec { - let batch = Self::data_batch(Self::INITIAL_ROWS, Self::ROWS_PER_FRAGMENT); - let schema = batch.schema(); - self.dataset - .append(RecordBatchIterator::new([Ok(batch)], schema), None) - .await - .unwrap(); - self.dataset - .optimize_indices(&OptimizeOptions::append()) - .await - .unwrap(); - self.vector_index_uuids().await - } - - async fn delete_query_centroid_rows(&mut self) { - self.dataset.delete("row % 4 = 0").await.unwrap(); - } - - async fn compact_with_fragment_reuse(&mut self) -> Uuid { - let metrics = compact_files( - &mut self.dataset, - CompactionOptions { - target_rows_per_fragment: Self::ROWS_PER_FRAGMENT * 2, - defer_index_remap: true, - ..Default::default() - }, - None, - ) - .await - .unwrap(); - assert!(metrics.fragments_removed > 0); - assert!(metrics.fragments_added > 0); - self.fragment_reuse_uuid() - .await - .expect("deferred compaction should create a fragment-reuse index") - } - - async fn replace_with_unsupported_flat_index(&mut self) { - let ivf_params = IvfBuildParams::try_with_centroids( - Self::NUM_PARTITIONS, - Arc::new(self.centroids.as_fixed_size_list().clone()), - ) - .unwrap(); - let index_params = VectorIndexParams::with_ivf_flat_params(MetricType::L2, ivf_params); - self.dataset - .create_index(&["vector"], IndexType::Vector, None, &index_params, true) - .await - .unwrap(); - } - - async fn create_row_scalar_index(&mut self) { - self.dataset - .create_index( - &["row"], - IndexType::Scalar, - None, - &ScalarIndexParams::default(), - false, - ) - .await - .unwrap(); - } - - async fn insert_out_of_range_results_cache_entry(&self) { - let indices = self.dataset.load_indices().await.unwrap(); - let index_metadata = indices - .iter() - .find(|index| index.name != FRAG_REUSE_INDEX_NAME) - .unwrap(); - let index = self - .dataset - .open_vector_index("vector", &index_metadata.uuid, &NoOpMetricsCollector) - .await - .unwrap(); - let query = Query { - column: "vector".to_string(), - key: self.query.clone(), - k: Self::K, - lower_bound: None, - upper_bound: None, - minimum_nprobes: Self::NUM_PARTITIONS, - maximum_nprobes: Some(Self::NUM_PARTITIONS), - ef: None, - refine_factor: None, - metric_type: None, - use_index: true, - query_parallelism: DEFAULT_QUERY_PARALLELISM, - dist_q_c: 0.0, - approx_mode: ApproxMode::Accurate, - }; - let (partitions, _) = index.find_partitions(&query).unwrap(); - let identity = knn_results_cache::identity_for_query( - &self.dataset, - index_metadata, - index.as_ref(), - &query, - partitions.values(), - false, - false, - ) - .await - .unwrap(); - let entry = VectorResultsCacheEntry::try_new( - identity.clone(), - vec![CachedVectorCandidate::new( - identity.partition_ids()[0], - u32::MAX, - )], - ) - .unwrap(); - self.dataset - .index_cache - .0 - .insert_with_key(&identity, Arc::new(entry)) - .await; - } - - async fn replace_index(&mut self) -> (Uuid, Uuid) { - let old_uuid = self.dataset.load_indices().await.unwrap()[0].uuid; - self.dataset - .create_index( - &["vector"], - IndexType::Vector, - None, - &self.index_params, - true, - ) - .await - .unwrap(); - let new_uuid = self.dataset.load_indices().await.unwrap()[0].uuid; - (old_uuid, new_uuid) - } - } - #[derive(Default)] struct StatsHolder { pub collected_stats: Arc>>, @@ -5224,750 +5006,6 @@ mod tests { ); } - #[derive(Clone, Copy, Debug)] - enum ResultsCacheQueryMode { - Supported, - Prefilter, - ScalarPrefilter, - DistanceRange, - AdaptiveNprobes, - Overlay, - } - - fn with_results_cache_test_options( - plan: Arc, - is_enabled: bool, - has_overlay: bool, - ) -> Arc { - plan.transform_down(|node| { - let Some(ann) = node.downcast_ref::() else { - return Ok(Transformed::no(node)); - }; - let mut replacement = ann - .copy_with_results_cache_enabled_for_testing(is_enabled) - .unwrap(); - if has_overlay { - replacement = replacement.with_overlay_block(RowAddrMask::all_rows()); - } - let replacement: Arc = Arc::new(replacement); - Ok(Transformed::yes(replacement)) - }) - .unwrap() - .data - } - - async fn run_results_cache_query( - fixture: &ResultsCacheTestFixture, - is_cache_enabled: bool, - refine_factor: Option, - ) -> (Vec, ExecutionSummaryCounts) { - run_results_cache_query_with_mode( - fixture, - is_cache_enabled, - refine_factor, - ResultsCacheQueryMode::Supported, - ) - .await - } - - async fn run_results_cache_query_with_mode( - fixture: &ResultsCacheTestFixture, - is_cache_enabled: bool, - refine_factor: Option, - query_mode: ResultsCacheQueryMode, - ) -> (Vec, ExecutionSummaryCounts) { - let (batch, stats) = run_results_cache_batch_with_mode( - fixture, - is_cache_enabled, - refine_factor, - query_mode, - ApproxMode::Accurate, - ) - .await; - ( - batch[ROW_ID].as_primitive::().values().to_vec(), - stats, - ) - } - - async fn run_results_cache_batch_with_mode( - fixture: &ResultsCacheTestFixture, - is_cache_enabled: bool, - refine_factor: Option, - query_mode: ResultsCacheQueryMode, - approx_mode: ApproxMode, - ) -> (RecordBatch, ExecutionSummaryCounts) { - let stats_holder = StatsHolder::default(); - let mut scanner = fixture.dataset.scan(); - scanner - .nearest("vector", fixture.query.as_ref(), ResultsCacheTestFixture::K) - .unwrap() - .distance_metric(fixture.metric_type) - .approx_mode(approx_mode); - match query_mode { - ResultsCacheQueryMode::Supported | ResultsCacheQueryMode::Overlay => { - scanner.nprobes(ResultsCacheTestFixture::NUM_PARTITIONS); - } - ResultsCacheQueryMode::Prefilter | ResultsCacheQueryMode::ScalarPrefilter => { - scanner - .nprobes(ResultsCacheTestFixture::NUM_PARTITIONS) - .filter("row >= 0") - .unwrap() - .prefilter(true); - } - ResultsCacheQueryMode::DistanceRange => { - scanner - .nprobes(ResultsCacheTestFixture::NUM_PARTITIONS) - .distance_range(None, Some(f32::MAX)); - } - ResultsCacheQueryMode::AdaptiveNprobes => { - scanner - .minimum_nprobes(1) - .maximum_nprobes(ResultsCacheTestFixture::NUM_PARTITIONS); - } - } - scanner - .project(&Vec::::new()) - .unwrap() - .with_row_id() - .scan_stats_callback(stats_holder.get_setter()); - if let Some(refine_factor) = refine_factor { - scanner.refine(refine_factor); - } - - let plan = scanner.create_plan().await.unwrap(); - if matches!(query_mode, ResultsCacheQueryMode::ScalarPrefilter) { - let rendered = format!( - "{}", - datafusion::physical_plan::displayable(plan.as_ref()).indent(true) - ); - assert!( - rendered.contains("ScalarIndexQuery"), - "expected a scalar-index prefilter plan, got:\n{rendered}" - ); - } - let plan = with_results_cache_test_options( - plan, - is_cache_enabled, - matches!(query_mode, ResultsCacheQueryMode::Overlay), - ); - let batches = execute_plan(plan, scanner.execution_options()) - .unwrap() - .try_collect::>() - .await - .unwrap(); - let batch = concat_batches(&batches[0].schema(), &batches).unwrap(); - (batch, stats_holder.consume()) - } - - async fn exact_result_batch(fixture: &ResultsCacheTestFixture) -> RecordBatch { - fixture - .dataset - .scan() - .nearest("vector", fixture.query.as_ref(), ResultsCacheTestFixture::K) - .unwrap() - .distance_metric(fixture.metric_type) - .use_index(false) - .project(&Vec::::new()) - .unwrap() - .with_row_id() - .try_into_batch() - .await - .unwrap() - } - - async fn exact_results(fixture: &ResultsCacheTestFixture) -> Vec { - let batch = exact_result_batch(fixture).await; - batch[ROW_ID].as_primitive::().values().to_vec() - } - - fn count_metric(stats: &ExecutionSummaryCounts, name: &str) -> usize { - stats.all_counts.get(name).copied().unwrap_or_default() - } - - fn recall(actual: &[u64], expected: &[u64]) -> f32 { - actual - .iter() - .filter(|row_id| expected.contains(row_id)) - .count() as f32 - / expected.len() as f32 - } - - fn assert_valid_result_set(row_ids: &[u64]) { - assert_eq!(row_ids.len(), ResultsCacheTestFixture::K); - for (position, row_id) in row_ids.iter().enumerate() { - assert!( - !row_ids[..position].contains(row_id), - "result row id {row_id} occurs more than once in {row_ids:?}" - ); - } - } - - fn assert_result_batches_match(actual: &RecordBatch, expected: &RecordBatch) { - assert_eq!(actual.num_rows(), expected.num_rows()); - let actual_row_ids = actual[ROW_ID].as_primitive::(); - let expected_row_ids = expected[ROW_ID].as_primitive::(); - let actual_distances = actual[DIST_COL].as_primitive::(); - let expected_distances = expected[DIST_COL].as_primitive::(); - for position in 0..actual.num_rows() { - assert_eq!( - actual_row_ids.value(position), - expected_row_ids.value(position) - ); - let actual_distance = actual_distances.value(position); - let expected_distance = expected_distances.value(position); - let tolerance = 1.0e-5 * (1.0 + expected_distance.abs()); - assert!( - actual_distance.is_finite() - && (actual_distance - expected_distance).abs() <= tolerance, - "distance mismatch at position {position}: actual={actual_distance}, expected={expected_distance}" - ); - } - } - - fn assert_quantized_result_batches_agree(actual: &RecordBatch, expected: &RecordBatch) { - assert_eq!(actual.num_rows(), expected.num_rows()); - let actual_row_ids = actual[ROW_ID] - .as_primitive::() - .values() - .to_vec(); - let expected_row_ids = expected[ROW_ID] - .as_primitive::() - .values() - .to_vec(); - assert_valid_result_set(&actual_row_ids); - assert_valid_result_set(&expected_row_ids); - assert!(recall(&actual_row_ids, &expected_row_ids) >= 0.5); - - let actual_distances = actual[DIST_COL].as_primitive::(); - let expected_distances = expected[DIST_COL].as_primitive::(); - for position in 0..actual.num_rows() { - let actual_distance = actual_distances.value(position); - let expected_distance = expected_distances.value(position); - let tolerance = 1.0e-4 * (1.0 + expected_distance.abs()); - assert!( - actual_distance.is_finite() - && (actual_distance - expected_distance).abs() <= tolerance, - "quantized distance mismatch at position {position}: actual={actual_distance}, expected={expected_distance}" - ); - } - } - - fn complete_refine_factor(row_count: usize) -> u32 { - u32::try_from(row_count.div_ceil(ResultsCacheTestFixture::K)).unwrap() - } - - #[rstest] - #[case::sq8(None)] - #[case::rq4(Some(4))] - #[case::rq8(Some(8))] - #[tokio::test] - async fn test_results_cache_end_to_end_and_index_replacement_invalidation( - #[case] rq_num_bits: Option, - ) { - let mut fixture = ResultsCacheTestFixture::new(rq_num_bits).await; - let expected = exact_results(&fixture).await; - - let (cold, cold_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&cold_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert!(cold_stats.bytes_read > 0); - assert_valid_result_set(&cold); - - let (warm, warm_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - assert_eq!( - count_metric(&warm_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - assert_eq!(warm_stats.bytes_read, 0); - assert_eq!(warm_stats.parts_loaded, 0); - assert_valid_result_set(&warm); - - let refine_factor = Some( - (ResultsCacheTestFixture::ROWS_PER_FRAGMENT * ResultsCacheTestFixture::NUM_FRAGMENTS - / ResultsCacheTestFixture::K) as u32, - ); - let (exact_cold, exact_cold_stats) = - run_results_cache_query(&fixture, true, refine_factor).await; - assert_eq!( - count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert_eq!(recall(&exact_cold, &expected), 1.0); - - let (exact_warm, exact_warm_stats) = - run_results_cache_query(&fixture, true, refine_factor).await; - assert_eq!( - count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - assert_eq!(recall(&exact_warm, &expected), 1.0); - - let (disabled, disabled_stats) = - run_results_cache_query(&fixture, false, refine_factor).await; - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - assert_eq!(recall(&disabled, &expected), 1.0); - - let (old_uuid, new_uuid) = fixture.replace_index().await; - assert_ne!(old_uuid, new_uuid); - let (after_replacement, replacement_stats) = - run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&replacement_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&replacement_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert_valid_result_set(&after_replacement); - } - - #[rstest] - #[case::sq8_l2(None, MetricType::L2)] - #[case::sq8_cosine(None, MetricType::Cosine)] - #[case::sq8_dot(None, MetricType::Dot)] - #[case::rq4_l2(Some(4), MetricType::L2)] - #[case::rq4_cosine(Some(4), MetricType::Cosine)] - #[case::rq4_dot(Some(4), MetricType::Dot)] - #[case::rq8_l2(Some(8), MetricType::L2)] - #[case::rq8_cosine(Some(8), MetricType::Cosine)] - #[case::rq8_dot(Some(8), MetricType::Dot)] - #[tokio::test] - async fn test_results_cache_scoring_matrix_matches_direct_native_and_exact_refinement( - #[case] rq_num_bits: Option, - #[case] metric_type: MetricType, - ) { - let fixture = ResultsCacheTestFixture::with_metric(rq_num_bits, metric_type).await; - let exact = exact_result_batch(&fixture).await; - let exact_row_ids = exact[ROW_ID].as_primitive::().values().to_vec(); - - let (_, cold_stats) = run_results_cache_batch_with_mode( - &fixture, - true, - None, - ResultsCacheQueryMode::Supported, - ApproxMode::Accurate, - ) - .await; - assert_eq!( - count_metric(&cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - let (warm, warm_stats) = run_results_cache_batch_with_mode( - &fixture, - true, - None, - ResultsCacheQueryMode::Supported, - ApproxMode::Accurate, - ) - .await; - assert_eq!( - count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - let direct = run_results_cache_batch_with_mode( - &fixture, - false, - None, - ResultsCacheQueryMode::Supported, - ApproxMode::Accurate, - ) - .await - .0; - assert_quantized_result_batches_agree(&warm, &direct); - let warm_row_ids = warm[ROW_ID].as_primitive::().values().to_vec(); - assert_valid_result_set(&warm_row_ids); - assert!( - recall(&warm_row_ids, &exact_row_ids) >= 0.5, - "{metric_type} warm-cache recall fell below 0.5: warm={warm_row_ids:?}, exact={exact_row_ids:?}" - ); - - let refine_factor = Some(complete_refine_factor( - ResultsCacheTestFixture::INITIAL_ROWS, - )); - let (exact_cold, exact_cold_stats) = run_results_cache_batch_with_mode( - &fixture, - true, - refine_factor, - ResultsCacheQueryMode::Supported, - ApproxMode::Accurate, - ) - .await; - assert_eq!( - count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert_result_batches_match(&exact_cold, &exact); - let (exact_warm, exact_warm_stats) = run_results_cache_batch_with_mode( - &fixture, - true, - refine_factor, - ResultsCacheQueryMode::Supported, - ApproxMode::Accurate, - ) - .await; - assert_eq!( - count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - assert_result_batches_match(&exact_warm, &exact); - } - - #[rstest] - #[case::sq8(None)] - #[case::rq4(Some(4))] - #[case::rq8(Some(8))] - #[tokio::test] - async fn test_results_cache_dataset_append_and_multi_segment_invalidation( - #[case] rq_num_bits: Option, - ) { - let mut fixture = ResultsCacheTestFixture::new(rq_num_bits).await; - let original_version = fixture.dataset.version_id(); - let original_index_uuids = fixture.vector_index_uuids().await; - assert_eq!(original_index_uuids.len(), 1); - - let (_, cold_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - let (_, warm_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - - let appended_index_uuids = fixture.append_index_segment().await; - assert!(fixture.dataset.version_id() > original_version); - assert_eq!(appended_index_uuids.len(), 2); - assert!(appended_index_uuids.contains(&original_index_uuids[0])); - - let (after_append, after_append_stats) = - run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&after_append_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - appended_index_uuids.len() - ); - assert_eq!( - count_metric(&after_append_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_valid_result_set(&after_append); - - let (after_append_warm, after_append_warm_stats) = - run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&after_append_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - appended_index_uuids.len() - ); - assert_eq!( - count_metric(&after_append_warm_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - assert_valid_result_set(&after_append_warm); - - let expected = exact_results(&fixture).await; - let row_count = fixture.dataset.count_rows(None).await.unwrap(); - let refine_factor = Some(complete_refine_factor(row_count)); - let (exact_cold, exact_cold_stats) = - run_results_cache_query(&fixture, true, refine_factor).await; - assert_eq!( - count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - appended_index_uuids.len() - ); - assert_eq!(recall(&exact_cold, &expected), 1.0); - - let (exact_warm, exact_warm_stats) = - run_results_cache_query(&fixture, true, refine_factor).await; - assert_eq!( - count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - appended_index_uuids.len() - ); - assert_eq!(recall(&exact_warm, &expected), 1.0); - } - - #[rstest] - #[case::sq8(None)] - #[case::rq4(Some(4))] - #[case::rq8(Some(8))] - #[tokio::test] - async fn test_results_cache_deletion_and_fragment_reuse_invalidation( - #[case] rq_num_bits: Option, - ) { - let mut fixture = ResultsCacheTestFixture::new(rq_num_bits).await; - let deleted_row_ids = exact_results(&fixture).await; - - let (_, initial_cold_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&initial_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - let (_, initial_warm_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&initial_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - - let version_before_delete = fixture.dataset.version_id(); - fixture.delete_query_centroid_rows().await; - assert!(fixture.dataset.version_id() > version_before_delete); - - let (after_delete, after_delete_stats) = - run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&after_delete_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert_eq!( - count_metric(&after_delete_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_valid_result_set(&after_delete); - assert!( - after_delete - .iter() - .all(|row_id| !deleted_row_ids.contains(row_id)) - ); - - let (_, after_delete_warm_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&after_delete_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - - let vector_index_uuid = fixture.vector_index_uuids().await[0]; - assert_eq!(fixture.fragment_reuse_uuid().await, None); - let version_before_compaction = fixture.dataset.version_id(); - let fragment_reuse_uuid = fixture.compact_with_fragment_reuse().await; - assert!(fixture.dataset.version_id() > version_before_compaction); - assert_eq!(fixture.vector_index_uuids().await, vec![vector_index_uuid]); - assert_eq!( - fixture.fragment_reuse_uuid().await, - Some(fragment_reuse_uuid) - ); - - let (after_compaction, after_compaction_stats) = - run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&after_compaction_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert_eq!( - count_metric(&after_compaction_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_valid_result_set(&after_compaction); - - let (_, after_compaction_warm_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric( - &after_compaction_warm_stats, - VECTOR_RESULTS_CACHE_HITS_METRIC - ), - 1 - ); - - let expected = exact_results(&fixture).await; - let row_count = fixture.dataset.count_rows(None).await.unwrap(); - let refine_factor = Some(complete_refine_factor(row_count)); - let (exact_cold, exact_cold_stats) = - run_results_cache_query(&fixture, true, refine_factor).await; - assert_eq!( - count_metric(&exact_cold_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert_eq!(recall(&exact_cold, &expected), 1.0); - - let (exact_warm, exact_warm_stats) = - run_results_cache_query(&fixture, true, refine_factor).await; - assert_eq!( - count_metric(&exact_warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - assert_eq!(recall(&exact_warm, &expected), 1.0); - } - - #[rstest] - #[case::prefilter(ResultsCacheQueryMode::Prefilter)] - #[case::scalar_prefilter(ResultsCacheQueryMode::ScalarPrefilter)] - #[case::distance_range(ResultsCacheQueryMode::DistanceRange)] - #[case::adaptive_nprobes(ResultsCacheQueryMode::AdaptiveNprobes)] - #[case::overlay(ResultsCacheQueryMode::Overlay)] - #[tokio::test] - async fn test_results_cache_bypasses_unsupported_query_shapes( - #[case] query_mode: ResultsCacheQueryMode, - ) { - let mut fixture = ResultsCacheTestFixture::new(None).await; - if matches!(query_mode, ResultsCacheQueryMode::ScalarPrefilter) { - fixture.create_row_scalar_index().await; - } - let row_count = fixture.dataset.count_rows(None).await.unwrap(); - let refine_factor = Some(complete_refine_factor(row_count)); - let (expected, disabled_stats) = - run_results_cache_query_with_mode(&fixture, false, refine_factor, query_mode).await; - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - assert_valid_result_set(&expected); - let exact = exact_results(&fixture).await; - assert!(recall(&expected, &exact) >= 0.5); - for _ in 0..2 { - let (row_ids, stats) = - run_results_cache_query_with_mode(&fixture, true, refine_factor, query_mode).await; - assert_eq!( - count_metric(&stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0, - "unsupported query mode {query_mode:?} recorded a cache hit" - ); - assert_eq!( - count_metric(&stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0, - "unsupported query mode {query_mode:?} recorded a cache miss" - ); - assert_eq!(row_ids, expected); - } - } - - #[rstest] - #[case::rq4_fast(4, ApproxMode::Fast)] - #[case::rq4_normal(4, ApproxMode::Normal)] - #[case::rq8_fast(8, ApproxMode::Fast)] - #[case::rq8_normal(8, ApproxMode::Normal)] - #[tokio::test] - async fn test_results_cache_bypasses_non_accurate_rq( - #[case] rq_num_bits: u8, - #[case] approx_mode: ApproxMode, - ) { - let fixture = ResultsCacheTestFixture::new(Some(rq_num_bits)).await; - let (expected, disabled_stats) = run_results_cache_batch_with_mode( - &fixture, - false, - None, - ResultsCacheQueryMode::Supported, - approx_mode, - ) - .await; - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - - for _ in 0..2 { - let (actual, stats) = run_results_cache_batch_with_mode( - &fixture, - true, - None, - ResultsCacheQueryMode::Supported, - approx_mode, - ) - .await; - assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_HITS_METRIC), 0); - assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), 0); - assert_result_batches_match(&actual, &expected); - } - } - - #[tokio::test] - async fn test_results_cache_bypasses_unsupported_index_type() { - let mut fixture = ResultsCacheTestFixture::new(None).await; - fixture.replace_with_unsupported_flat_index().await; - - let row_count = fixture.dataset.count_rows(None).await.unwrap(); - let refine_factor = Some(complete_refine_factor(row_count)); - let (expected, disabled_stats) = - run_results_cache_query(&fixture, false, refine_factor).await; - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - assert_valid_result_set(&expected); - let exact = exact_results(&fixture).await; - assert!(recall(&expected, &exact) >= 0.5); - for _ in 0..2 { - let (row_ids, stats) = run_results_cache_query(&fixture, true, refine_factor).await; - assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_HITS_METRIC), 0); - assert_eq!(count_metric(&stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), 0); - assert_eq!(row_ids, expected); - } - } - - #[rstest] - #[case::sq8(None)] - #[case::rq4(Some(4))] - #[case::rq8(Some(8))] - #[tokio::test] - async fn test_results_cache_unusable_entry_falls_back_and_is_replaced( - #[case] rq_num_bits: Option, - ) { - let fixture = ResultsCacheTestFixture::new(rq_num_bits).await; - let (ordinary, disabled_stats) = run_results_cache_query(&fixture, false, None).await; - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&disabled_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - assert_valid_result_set(&ordinary); - fixture.insert_out_of_range_results_cache_entry().await; - - let (fallback, fallback_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&fallback_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 0 - ); - assert_eq!( - count_metric(&fallback_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 1 - ); - assert_valid_result_set(&fallback); - assert!(recall(&fallback, &ordinary) >= 0.5); - - let (warm, warm_stats) = run_results_cache_query(&fixture, true, None).await; - assert_eq!( - count_metric(&warm_stats, VECTOR_RESULTS_CACHE_HITS_METRIC), - 1 - ); - assert_eq!( - count_metric(&warm_stats, VECTOR_RESULTS_CACHE_MISSES_METRIC), - 0 - ); - assert_valid_result_set(&warm); - assert!(recall(&warm, &ordinary) >= 0.5); - } - #[rstest] #[tokio::test] async fn test_no_max_nprobes(#[values(1, 20)] num_deltas: usize) { diff --git a/rust/lance/src/io/exec/knn_results_cache.rs b/rust/lance/src/io/exec/knn_results_cache.rs deleted file mode 100644 index 285faec9ac0..00000000000 --- a/rust/lance/src/io/exec/knn_results_cache.rs +++ /dev/null @@ -1,1162 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The Lance Authors - -//! Experimental cache execution for reusable IVF candidate pools. - -use std::collections::{BinaryHeap, HashMap}; -use std::sync::Arc; - -use arrow::datatypes::{Float32Type, UInt64Type}; -use arrow_array::cast::AsArray; -use arrow_array::{Float32Array, RecordBatch, UInt32Array, UInt64Array}; -use futures::{StreamExt, TryStreamExt}; -use lance_core::cache::LanceCache; -use lance_core::{Error, ROW_ID, Result}; -use lance_index::IndexType; -use lance_index::frag_reuse::FRAG_REUSE_INDEX_NAME; -use lance_index::metrics::MetricsCollector; -use lance_index::prefilter::PreFilter; -use lance_index::vector::graph::{OrderedFloat, OrderedNode}; -use lance_index::vector::quantizer::QuantizationType; -use lance_index::vector::v3::subindex::SubIndexType; -use lance_index::vector::{DIST_COL, PartitionSearchResult, Query, VectorIndex}; -use lance_select::RowAddrMask; -use lance_table::format::IndexMetadata; - -use crate::dataset::Dataset; -use crate::index::DatasetIndexExt; -use crate::session::index_caches::{ - CachedVectorCandidate, VectorResultsCacheEntry, VectorResultsCacheIdentity, -}; - -use super::knn::KNN_INDEX_SCHEMA; - -/// Maximum reusable candidate pool retained for one exact query/search shape. -/// -/// This is deliberately fixed while the feature is experimental. It is part of -/// the cache identity, so changing it cannot reuse entries created under another -/// limit. -pub(super) const RESULTS_CACHE_CANDIDATE_LIMIT: usize = 1000; - -const VECTOR_RESULTS_CACHE_ENV: &str = "LANCE_EXPERIMENTAL_VECTOR_RESULTS_CACHE"; - -pub(super) struct ResultsCacheSearch { - pub batch: RecordBatch, - /// True for both a warm entry and a concurrent load coalesced behind its owner. - pub was_hit: bool, -} - -pub(super) struct ResultsCacheSearchParams<'a> { - pub cache: &'a LanceCache, - pub identity: VectorResultsCacheIdentity, - pub index: Arc, - pub query: &'a Query, - pub partitions: Arc, - pub centroid_distances: Arc, - pub prefilter: Arc, - pub segment_mask: Option>, - pub metrics: Arc, - pub parallelism: usize, -} - -pub(super) fn is_enabled() -> bool { - std::env::var(VECTOR_RESULTS_CACHE_ENV) - .ok() - .is_some_and(|value| value == "1") -} - -pub(super) async fn identity_for_query( - dataset: &Dataset, - index_metadata: &IndexMetadata, - index: &dyn VectorIndex, - query: &Query, - partition_ids: &[u32], - has_prefilter: bool, - has_overlay: bool, -) -> Option { - let (sub_index_type, quantization_type) = match index.index_type() { - IndexType::IvfRq => (SubIndexType::Flat, QuantizationType::Rabit), - IndexType::IvfSq => (SubIndexType::Flat, QuantizationType::Scalar), - _ => return None, - }; - let index_store = dataset.object_store_for_index(index_metadata).await.ok()?; - let fragment_reuse_uuid = dataset - .load_indices() - .await - .ok()? - .iter() - .find(|metadata| metadata.name == FRAG_REUSE_INDEX_NAME) - .map(|metadata| metadata.uuid); - let manifest_location = dataset.manifest_location(); - let manifest_path = manifest_location.path.as_ref(); - let manifest_etag = manifest_location.e_tag.as_deref().unwrap_or_default(); - let manifest_size = manifest_location - .size - .map_or_else(|| "none".to_owned(), |size| format!("some:{size}")); - let dataset_read_identity = format!( - "{}:{}/{}:{}/{}:{}", - manifest_path.len(), - manifest_path, - manifest_etag.len(), - manifest_etag, - manifest_size.len(), - manifest_size - ); - - VectorResultsCacheIdentity::try_new( - &index_store.store_prefix, - &dataset_read_identity, - dataset.version_id(), - index_metadata, - fragment_reuse_uuid, - query, - index.metric_type(), - sub_index_type, - quantization_type, - partition_ids, - RESULTS_CACHE_CANDIDATE_LIMIT, - has_prefilter, - has_overlay, - ) - .ok() -} - -pub(super) async fn search(params: ResultsCacheSearchParams<'_>) -> Result { - let ResultsCacheSearchParams { - cache, - identity, - index, - query, - partitions, - centroid_distances, - prefilter, - segment_mask, - metrics, - parallelism, - } = params; - let mut populated_batch = None; - let populated_batch_slot = &mut populated_batch; - let loader_identity = identity.clone(); - let loader_index = index.clone(); - let loader_partitions = partitions.clone(); - let loader_centroid_distances = centroid_distances.clone(); - let loader_prefilter = prefilter.clone(); - let loader_segment_mask = segment_mask.clone(); - let loader_metrics = metrics.clone(); - let cache_lookup = cache - .get_or_insert_with_key_hit(identity.clone(), move || async move { - let (candidates, batch) = populate_candidates( - loader_index, - query, - loader_partitions, - loader_centroid_distances, - loader_prefilter, - loader_segment_mask, - loader_metrics, - parallelism, - loader_identity.result_limit(), - ) - .await?; - *populated_batch_slot = Some(batch); - VectorResultsCacheEntry::try_new(loader_identity, candidates) - }) - .await; - - let (entry, was_cached) = match cache_lookup { - Ok(result) => result, - Err(error) => { - if let Some(batch) = populated_batch { - // The candidate search succeeded, so an invalid cache entry must not - // discard its query result or cause the ordinary path to repeat work. - log::debug!("not storing invalid vector results cache entry: {error}"); - return Ok(ResultsCacheSearch { - batch, - was_hit: false, - }); - } - return Err(error); - } - }; - - if !was_cached { - let batch = populated_batch.ok_or_else(|| { - Error::internal( - "vector results cache loader completed without its candidate-search batch", - ) - })?; - return Ok(ResultsCacheSearch { - batch, - was_hit: false, - }); - } - - if entry.is_compatible_with(&identity) { - match replay_candidates( - index.clone(), - query, - entry.candidates(), - identity.result_limit(), - metrics.clone(), - parallelism, - ) - .await - { - Ok(batch) => { - return Ok(ResultsCacheSearch { - batch, - was_hit: true, - }); - } - Err(error) => { - // A stale or malformed in-memory entry must never fail the query. - // Re-run the candidate-producing path and replace it. - log::debug!("ignoring unusable vector results cache entry: {error}"); - } - } - } - - // An unusable existing entry cannot enter get-or-insert's loader, so replace it - // directly. Normal cold misses are single-flighted by the lookup above. - let (candidates, batch) = populate_candidates( - index, - query, - partitions, - centroid_distances, - prefilter, - segment_mask, - metrics, - parallelism, - identity.result_limit(), - ) - .await?; - match VectorResultsCacheEntry::try_new(identity.clone(), candidates) { - Ok(entry) => cache.insert_with_key(&identity, Arc::new(entry)).await, - Err(error) => { - // Cache population is best-effort and cannot change query success. - log::debug!("not storing invalid vector results cache entry: {error}"); - } - } - Ok(ResultsCacheSearch { - batch, - was_hit: false, - }) -} - -#[allow(clippy::too_many_arguments)] -async fn populate_candidates( - index: Arc, - query: &Query, - partitions: Arc, - centroid_distances: Arc, - prefilter: Arc, - segment_mask: Option>, - metrics: Arc, - parallelism: usize, - result_limit: usize, -) -> Result<(Vec, RecordBatch)> { - if partitions.len() != centroid_distances.len() { - return Err(Error::invalid_input(format!( - "partition count {} does not match centroid distance count {} for vector results cache", - partitions.len(), - centroid_distances.len() - ))); - } - let accumulated = futures::stream::iter(0..partitions.len()) - .map(|partition_index| { - let index = index.clone(); - let prefilter = prefilter.clone(); - let metrics = metrics.clone(); - let mut partition_query = query.clone(); - let partition_id = partitions.value(partition_index); - partition_query.dist_q_c = centroid_distances.value(partition_index); - async move { - index - .search_in_partition_with_candidates( - partition_id as usize, - &partition_query, - prefilter, - metrics.as_ref(), - RESULTS_CACHE_CANDIDATE_LIMIT, - ) - .await - } - }) - .buffered(parallelism.max(1)) - .try_fold( - SearchAccumulator::new(RESULTS_CACHE_CANDIDATE_LIMIT, result_limit), - move |mut accumulated, partition_result| { - let segment_mask = segment_mask.clone(); - async move { - accumulated.add(partition_result, segment_mask.as_deref())?; - Ok(accumulated) - } - }, - ) - .await?; - accumulated.finish() -} - -struct SearchAccumulator { - candidate_limit: usize, - result_limit: usize, - candidates: BinaryHeap>, - results: BinaryHeap>, -} - -impl SearchAccumulator { - fn new(candidate_limit: usize, result_limit: usize) -> Self { - Self { - candidate_limit, - result_limit, - candidates: BinaryHeap::with_capacity(candidate_limit), - results: BinaryHeap::with_capacity(result_limit), - } - } - - fn add( - &mut self, - search_result: PartitionSearchResult, - segment_mask: Option<&RowAddrMask>, - ) -> Result<()> { - if search_result.candidates.len() != search_result.batch.num_rows() { - return Err(Error::internal(format!( - "partition search returned {} candidate identities for {} rows", - search_result.candidates.len(), - search_result.batch.num_rows() - ))); - } - let row_ids = search_result - .batch - .column_by_name(ROW_ID) - .ok_or_else(|| Error::internal("partition search result has no row-id column"))? - .as_primitive::(); - let distances = search_result - .batch - .column_by_name(DIST_COL) - .ok_or_else(|| Error::internal("partition search result has no distance column"))? - .as_primitive::(); - if row_ids.len() != distances.len() { - return Err(Error::internal(format!( - "partition search returned {} row ids and {} distances", - row_ids.len(), - distances.len() - ))); - } - for ((candidate, &row_id), &distance) in search_result - .candidates - .iter() - .zip(row_ids.values()) - .zip(distances.values()) - { - if distance.is_nan() { - continue; - } - if segment_mask.is_some_and(|mask| !mask.selected(row_id)) { - continue; - } - let candidate = - CachedVectorCandidate::new(candidate.partition_id, candidate.offset_in_partition); - push_top_candidate( - &mut self.candidates, - self.candidate_limit, - candidate, - distance, - ); - push_top_candidate(&mut self.results, self.result_limit, row_id, distance); - } - Ok(()) - } - - fn finish(self) -> Result<(Vec, RecordBatch)> { - let mut ordered_candidates = self.candidates.into_vec(); - ordered_candidates.sort_by(|left, right| { - left.dist - .cmp(&right.dist) - .then_with(|| left.id.partition_id().cmp(&right.id.partition_id())) - .then_with(|| { - left.id - .offset_in_partition() - .cmp(&right.id.offset_in_partition()) - }) - }); - let candidates = ordered_candidates.into_iter().map(|node| node.id).collect(); - Ok((candidates, batch_from_heap(self.results)?)) - } -} - -fn push_top_candidate( - heap: &mut BinaryHeap>, - limit: usize, - id: T, - distance: f32, -) { - if limit == 0 { - return; - } - let node = OrderedNode::new(id, OrderedFloat(distance)); - if heap.len() < limit { - heap.push(node); - } else if heap - .peek() - .is_some_and(|farthest| farthest.dist > node.dist) - { - heap.pop(); - heap.push(node); - } -} - -async fn replay_candidates( - index: Arc, - query: &Query, - candidates: &[CachedVectorCandidate], - result_limit: usize, - metrics: Arc, - parallelism: usize, -) -> Result { - let mut offsets_by_partition = HashMap::>::new(); - for candidate in candidates { - offsets_by_partition - .entry(candidate.partition_id()) - .or_default() - .push(candidate.offset_in_partition()); - } - - let scored_partitions = futures::stream::iter(offsets_by_partition) - .map(|(partition_id, offsets)| { - let index = index.clone(); - let query = query.clone(); - let metrics = metrics.clone(); - async move { - let batch = index - .score_partition_candidates( - partition_id as usize, - &query, - &offsets, - metrics.as_ref(), - ) - .await?; - Result::Ok((partition_id, offsets.len(), batch)) - } - }) - .buffered(parallelism.max(1)) - .try_collect::>() - .await?; - - let mut top_results = BinaryHeap::with_capacity(result_limit); - for (partition_id, expected_rows, batch) in scored_partitions { - let row_ids = batch - .column_by_name(ROW_ID) - .ok_or_else(|| Error::internal("candidate scoring result has no row-id column"))? - .as_primitive::(); - let distances = batch - .column_by_name(DIST_COL) - .ok_or_else(|| Error::internal("candidate scoring result has no distance column"))? - .as_primitive::(); - if row_ids.len() != expected_rows || distances.len() != expected_rows { - return Err(Error::internal(format!( - "candidate scoring for partition {partition_id} returned {} row ids and {} distances for {} offsets", - row_ids.len(), - distances.len(), - expected_rows - ))); - } - for (&row_id, &distance) in row_ids.values().iter().zip(distances.values()) { - if distance.is_nan() { - continue; - } - push_top_candidate(&mut top_results, result_limit, row_id, distance); - } - } - batch_from_heap(top_results) -} - -fn batch_from_heap(heap: BinaryHeap>) -> Result { - let mut ordered = heap.into_vec(); - ordered.sort_by(|left, right| { - left.dist - .cmp(&right.dist) - .then_with(|| left.id.cmp(&right.id)) - }); - let mut row_ids = Vec::with_capacity(ordered.len()); - let mut distances = Vec::with_capacity(ordered.len()); - for result in ordered { - row_ids.push(result.id); - distances.push(result.dist.0); - } - Ok(RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(distances)), - Arc::new(UInt64Array::from(row_ids)), - ], - )?) -} - -#[cfg(test)] -mod tests { - use std::any::Any; - use std::sync::atomic::{AtomicUsize, Ordering}; - - use async_trait::async_trait; - use lance_core::cache::QuickCacheBackend; - use lance_core::deepsize::DeepSizeOf; - use lance_core::utils::row_addr_remap::RowAddrRemap; - use lance_index::Index; - use lance_index::metrics::NoOpMetricsCollector; - use lance_index::prefilter::NoFilter; - use lance_index::vector::ivf::storage::IvfModel; - use lance_index::vector::quantizer::Quantizer; - use lance_index::vector::v3::subindex::SubIndexType; - use lance_index::vector::{ApproxMode, DEFAULT_QUERY_PARALLELISM, PartitionSearchCandidate}; - use lance_io::traits::Reader; - use roaring::RoaringBitmap; - use uuid::Uuid; - - use super::*; - - #[derive(Debug, DeepSizeOf)] - struct TestCandidateIndex { - miss_partition_searches: AtomicUsize, - miss_search_yields: AtomicUsize, - block_miss_searches: AtomicUsize, - hit_partition_scores: AtomicUsize, - active_hit_partition_scores: AtomicUsize, - max_active_hit_partition_scores: AtomicUsize, - return_malformed_hit_shape: AtomicUsize, - row_ids: Vec, - } - - impl TestCandidateIndex { - fn new() -> Self { - Self { - miss_partition_searches: AtomicUsize::new(0), - miss_search_yields: AtomicUsize::new(0), - block_miss_searches: AtomicUsize::new(0), - hit_partition_scores: AtomicUsize::new(0), - active_hit_partition_scores: AtomicUsize::new(0), - max_active_hit_partition_scores: AtomicUsize::new(0), - return_malformed_hit_shape: AtomicUsize::new(0), - row_ids: Vec::new(), - } - } - } - - #[async_trait] - impl Index for TestCandidateIndex { - fn as_any(&self) -> &dyn Any { - self - } - - fn as_index(self: Arc) -> Arc { - self - } - - fn statistics(&self) -> Result { - Ok(serde_json::json!({})) - } - - async fn prewarm(&self) -> Result<()> { - Ok(()) - } - - fn index_type(&self) -> IndexType { - IndexType::IvfRq - } - - async fn calculate_included_frags(&self) -> Result { - Ok(RoaringBitmap::new()) - } - } - - #[async_trait] - impl VectorIndex for TestCandidateIndex { - async fn search( - &self, - _query: &Query, - _pre_filter: Arc, - _metrics: &dyn MetricsCollector, - ) -> Result { - Err(Error::not_supported("test whole-index search")) - } - - fn find_partitions(&self, _query: &Query) -> Result<(UInt32Array, Float32Array)> { - Ok(( - UInt32Array::from(vec![0, 1]), - Float32Array::from(vec![0.0, 0.0]), - )) - } - - fn total_partitions(&self) -> usize { - 2 - } - - async fn search_in_partition( - &self, - _partition_id: usize, - _query: &Query, - _pre_filter: Arc, - _metrics: &dyn MetricsCollector, - ) -> Result { - Err(Error::not_supported("test partition search")) - } - - async fn search_in_partition_with_candidates( - &self, - partition_id: usize, - _query: &Query, - _pre_filter: Arc, - _metrics: &dyn MetricsCollector, - _candidate_limit: usize, - ) -> Result { - self.miss_partition_searches.fetch_add(1, Ordering::Relaxed); - for _ in 0..self.miss_search_yields.load(Ordering::Relaxed) { - tokio::task::yield_now().await; - } - while self.block_miss_searches.load(Ordering::Relaxed) != 0 { - tokio::task::yield_now().await; - } - let (row_ids, distances) = match partition_id { - 0 => (vec![0, 1], vec![0.4, 0.2]), - 1 => (vec![10, 11], vec![0.3, 0.1]), - _ => { - return Err(Error::invalid_input(format!( - "unexpected test partition {partition_id}" - ))); - } - }; - let candidates = (0..row_ids.len()) - .map(|offset| PartitionSearchCandidate { - partition_id: partition_id as u32, - offset_in_partition: offset as u32, - }) - .collect(); - let batch = RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(distances)), - Arc::new(UInt64Array::from(row_ids)), - ], - )?; - Ok(PartitionSearchResult { batch, candidates }) - } - - async fn score_partition_candidates( - &self, - partition_id: usize, - _query: &Query, - offsets_in_partition: &[u32], - _metrics: &dyn MetricsCollector, - ) -> Result { - self.hit_partition_scores.fetch_add(1, Ordering::Relaxed); - let active = self - .active_hit_partition_scores - .fetch_add(1, Ordering::Relaxed) - + 1; - self.max_active_hit_partition_scores - .fetch_max(active, Ordering::Relaxed); - tokio::task::yield_now().await; - let mut scored = offsets_in_partition - .iter() - .map(|offset| match (partition_id, *offset) { - (0, 0) => Ok((0, 0.05)), - (0, 1) => Ok((1, 0.8)), - (1, 0) => Ok((10, 0.6)), - (1, 1) => Ok((11, 0.1)), - _ => Err(Error::invalid_input(format!( - "unexpected test candidate ({partition_id}, {offset})" - ))), - }) - .collect::>>()?; - if self.return_malformed_hit_shape.load(Ordering::Relaxed) != 0 { - scored.pop(); - } - self.active_hit_partition_scores - .fetch_sub(1, Ordering::Relaxed); - let (row_ids, distances): (Vec<_>, Vec<_>) = scored.into_iter().unzip(); - Ok(RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(distances)), - Arc::new(UInt64Array::from(row_ids)), - ], - )?) - } - - fn is_loadable(&self) -> bool { - false - } - - fn use_residual(&self) -> bool { - false - } - - async fn load( - &self, - _reader: Arc, - _offset: usize, - _length: usize, - ) -> Result> { - Err(Error::not_supported("test index load")) - } - - fn num_rows(&self) -> u64 { - 0 - } - - fn row_ids(&self) -> Box + '_> { - Box::new(self.row_ids.iter()) - } - - async fn remap(&mut self, _mapping: &RowAddrRemap) -> Result<()> { - Ok(()) - } - - async fn to_batch_stream( - &self, - _with_vector: bool, - ) -> Result { - Err(Error::not_supported("test batch stream")) - } - - fn ivf_model(&self) -> &IvfModel { - unreachable!("test cache path does not inspect the IVF model") - } - - fn quantizer(&self) -> Quantizer { - unreachable!("test cache path does not inspect the quantizer") - } - - fn partition_size(&self, _part_id: usize) -> usize { - 2 - } - - fn sub_index_type(&self) -> (SubIndexType, QuantizationType) { - (SubIndexType::Flat, QuantizationType::Rabit) - } - - fn metric_type(&self) -> lance_linalg::distance::DistanceType { - lance_linalg::distance::DistanceType::L2 - } - } - - fn test_query() -> Query { - Query { - column: "vector".to_string(), - key: Arc::new(Float32Array::from(vec![0.0, 1.0])), - k: 2, - lower_bound: None, - upper_bound: None, - minimum_nprobes: 2, - maximum_nprobes: Some(2), - ef: None, - refine_factor: None, - metric_type: Some(lance_linalg::distance::DistanceType::L2), - use_index: true, - query_parallelism: DEFAULT_QUERY_PARALLELISM, - dist_q_c: 0.0, - approx_mode: ApproxMode::Accurate, - } - } - - fn test_identity(query: &Query) -> VectorResultsCacheIdentity { - let metadata = IndexMetadata { - uuid: Uuid::from_u128(1), - fields: vec![0], - covering_fields: vec![], - name: "vector_idx".to_string(), - dataset_version: 1, - fragment_bitmap: Some([0_u32].into_iter().collect()), - index_details: Some(Arc::new(prost_types::Any { - type_url: "type.googleapis.com/lance.table.VectorIndexDetails".to_string(), - value: vec![1], - })), - index_version: 1, - created_at: None, - base_id: None, - files: None, - }; - VectorResultsCacheIdentity::try_new( - "memory", - "manifest-1", - 1, - &metadata, - None, - query, - lance_linalg::distance::DistanceType::L2, - SubIndexType::Flat, - QuantizationType::Rabit, - &[0, 1], - RESULTS_CACHE_CANDIDATE_LIMIT, - false, - false, - ) - .unwrap() - } - - fn row_ids(batch: &RecordBatch) -> Vec { - batch[ROW_ID].as_primitive::().values().to_vec() - } - - #[test] - fn accumulator_filters_nan_and_orders_ties_with_aligned_rows() { - let mut accumulator = SearchAccumulator::new(4, 4); - accumulator - .add( - PartitionSearchResult { - batch: RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(vec![0.2, f32::NAN, 0.2])), - Arc::new(UInt64Array::from(vec![30, 20, 10])), - ], - ) - .unwrap(), - candidates: vec![ - PartitionSearchCandidate { - partition_id: 0, - offset_in_partition: 2, - }, - PartitionSearchCandidate { - partition_id: 0, - offset_in_partition: 1, - }, - PartitionSearchCandidate { - partition_id: 0, - offset_in_partition: 0, - }, - ], - }, - None, - ) - .unwrap(); - - let (candidates, batch) = accumulator.finish().unwrap(); - assert_eq!( - candidates, - vec![ - CachedVectorCandidate::new(0, 0), - CachedVectorCandidate::new(0, 2), - ] - ); - assert_eq!(row_ids(&batch), vec![10, 30]); - assert_eq!( - batch[DIST_COL].as_primitive::().values(), - &[0.2, 0.2] - ); - } - - #[test] - fn accumulator_filters_candidates_outside_segment() { - let mut accumulator = SearchAccumulator::new(4, 4); - let segment_mask = - RowAddrMask::from_allowed(lance_select::RowAddrTreeMap::from_iter([10_u64, 30])); - accumulator - .add( - PartitionSearchResult { - batch: RecordBatch::try_new( - KNN_INDEX_SCHEMA.clone(), - vec![ - Arc::new(Float32Array::from(vec![0.3, 0.2, 0.1])), - Arc::new(UInt64Array::from(vec![30, 20, 10])), - ], - ) - .unwrap(), - candidates: vec![ - PartitionSearchCandidate { - partition_id: 0, - offset_in_partition: 2, - }, - PartitionSearchCandidate { - partition_id: 0, - offset_in_partition: 1, - }, - PartitionSearchCandidate { - partition_id: 0, - offset_in_partition: 0, - }, - ], - }, - Some(&segment_mask), - ) - .unwrap(); - - let (candidates, batch) = accumulator.finish().unwrap(); - assert_eq!( - candidates, - vec![ - CachedVectorCandidate::new(0, 0), - CachedVectorCandidate::new(0, 2), - ] - ); - assert_eq!(row_ids(&batch), vec![10, 30]); - } - - #[tokio::test] - async fn miss_populates_candidates_and_hit_rescores_them() { - let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(64 * 1024))); - let index = Arc::new(TestCandidateIndex::new()); - let query = test_query(); - let identity = test_identity(&query); - let partitions = Arc::new(UInt32Array::from(vec![0, 1])); - let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); - - let miss = search(ResultsCacheSearchParams { - cache: &cache, - identity: identity.clone(), - index: index.clone(), - query: &query, - partitions: partitions.clone(), - centroid_distances: centroid_distances.clone(), - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - .await - .unwrap(); - assert!(!miss.was_hit); - assert_eq!(row_ids(&miss.batch), vec![11, 1]); - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); - assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 0); - - let hit = search(ResultsCacheSearchParams { - cache: &cache, - identity, - index: index.clone(), - query: &query, - partitions, - centroid_distances, - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - .await - .unwrap(); - assert!(hit.was_hit); - assert_eq!(row_ids(&hit.batch), vec![0, 11]); - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); - assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 2); - assert_eq!( - index - .max_active_hit_partition_scores - .load(Ordering::Relaxed), - 2 - ); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn concurrent_cold_searches_populate_candidates_once() { - const SEARCHES: usize = 4; - - let cache = LanceCache::with_backend(Arc::new(QuickCacheBackend::with_capacity(64 * 1024))); - let index = Arc::new(TestCandidateIndex::new()); - index.miss_search_yields.store(16, Ordering::Relaxed); - let query = test_query(); - let identity = test_identity(&query); - let partitions = Arc::new(UInt32Array::from(vec![0, 1])); - let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); - - let results = futures::future::join_all((0..SEARCHES).map(|_| { - search(ResultsCacheSearchParams { - cache: &cache, - identity: identity.clone(), - index: index.clone(), - query: &query, - partitions: partitions.clone(), - centroid_distances: centroid_distances.clone(), - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - })) - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert_eq!(results.iter().filter(|result| !result.was_hit).count(), 1); - assert_eq!(results.iter().filter(|result| result.was_hit).count(), 3); - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); - assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 6); - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn concurrent_cold_search_retries_after_population_is_cancelled() { - let cache = Arc::new(LanceCache::with_backend(Arc::new( - QuickCacheBackend::with_capacity(64 * 1024), - ))); - let index = Arc::new(TestCandidateIndex::new()); - index.block_miss_searches.store(1, Ordering::Relaxed); - - let owner = { - let cache = cache.clone(); - let index = index.clone(); - tokio::spawn(async move { - let query = test_query(); - let identity = test_identity(&query); - search(ResultsCacheSearchParams { - cache: cache.as_ref(), - identity, - index, - query: &query, - partitions: Arc::new(UInt32Array::from(vec![0, 1])), - centroid_distances: Arc::new(Float32Array::from(vec![0.0, 0.0])), - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - .await - }) - }; - tokio::time::timeout(std::time::Duration::from_secs(5), async { - while index.miss_partition_searches.load(Ordering::Relaxed) < 2 { - tokio::task::yield_now().await; - } - }) - .await - .expect("population owner did not start both partition searches"); - - let contender = { - let cache = cache.clone(); - let index = index.clone(); - tokio::spawn(async move { - let query = test_query(); - let identity = test_identity(&query); - search(ResultsCacheSearchParams { - cache: cache.as_ref(), - identity, - index, - query: &query, - partitions: Arc::new(UInt32Array::from(vec![0, 1])), - centroid_distances: Arc::new(Float32Array::from(vec![0.0, 0.0])), - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - .await - }) - }; - for _ in 0..32 { - tokio::task::yield_now().await; - } - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 2); - assert!(!contender.is_finished()); - - owner.abort(); - assert!(matches!(owner.await, Err(error) if error.is_cancelled())); - index.block_miss_searches.store(0, Ordering::Relaxed); - let result = tokio::time::timeout(std::time::Duration::from_secs(5), contender) - .await - .expect("contender remained parked after population cancellation") - .unwrap() - .unwrap(); - assert!(!result.was_hit); - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); - } - - #[tokio::test] - async fn different_queries_with_same_partitions_do_not_share_candidates() { - let cache = LanceCache::with_capacity(64 * 1024); - let index = Arc::new(TestCandidateIndex::new()); - let first_query = test_query(); - let first_identity = test_identity(&first_query); - let partitions = Arc::new(UInt32Array::from(vec![0, 1])); - let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); - - let first = search(ResultsCacheSearchParams { - cache: &cache, - identity: first_identity, - index: index.clone(), - query: &first_query, - partitions: partitions.clone(), - centroid_distances: centroid_distances.clone(), - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - .await - .unwrap(); - assert!(!first.was_hit); - - let mut second_query = first_query.clone(); - second_query.key = Arc::new(Float32Array::from(vec![1.0, 0.0])); - let second_identity = test_identity(&second_query); - let second = search(ResultsCacheSearchParams { - cache: &cache, - identity: second_identity.clone(), - index: index.clone(), - query: &second_query, - partitions: partitions.clone(), - centroid_distances: centroid_distances.clone(), - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - .await - .unwrap(); - assert!(!second.was_hit); - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); - assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 0); - - let repeated_second = search(ResultsCacheSearchParams { - cache: &cache, - identity: second_identity, - index: index.clone(), - query: &second_query, - partitions, - centroid_distances, - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }) - .await - .unwrap(); - assert!(repeated_second.was_hit); - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); - assert_eq!(index.hit_partition_scores.load(Ordering::Relaxed), 2); - } - - #[tokio::test] - async fn results_cache_malformed_hit_shape_falls_back_and_replaces_entry() { - let cache = LanceCache::with_capacity(64 * 1024); - let index = Arc::new(TestCandidateIndex::new()); - let query = test_query(); - let identity = test_identity(&query); - let partitions = Arc::new(UInt32Array::from(vec![0, 1])); - let centroid_distances = Arc::new(Float32Array::from(vec![0.0, 0.0])); - let params = || ResultsCacheSearchParams { - cache: &cache, - identity: identity.clone(), - index: index.clone(), - query: &query, - partitions: partitions.clone(), - centroid_distances: centroid_distances.clone(), - prefilter: Arc::new(NoFilter), - segment_mask: None, - metrics: Arc::new(NoOpMetricsCollector), - parallelism: 2, - }; - - let cold = search(params()).await.unwrap(); - assert!(!cold.was_hit); - assert_eq!(row_ids(&cold.batch), vec![11, 1]); - - index.return_malformed_hit_shape.store(1, Ordering::Relaxed); - let fallback = search(params()).await.unwrap(); - assert!(!fallback.was_hit); - assert_eq!(row_ids(&fallback.batch), vec![11, 1]); - assert_eq!(index.miss_partition_searches.load(Ordering::Relaxed), 4); - - index.return_malformed_hit_shape.store(0, Ordering::Relaxed); - let replaced = search(params()).await.unwrap(); - assert!(replaced.was_hit); - assert_eq!(row_ids(&replaced.batch), vec![0, 11]); - } -} diff --git a/rust/lance/src/session/index_caches.rs b/rust/lance/src/session/index_caches.rs index 86c59cede6a..2f17f35d165 100644 --- a/rust/lance/src/session/index_caches.rs +++ b/rust/lance/src/session/index_caches.rs @@ -10,13 +10,8 @@ //! │ │ //! └────┴──► Index-specific cache (prefixed by index UUID and FRI UUID) -use std::{borrow::Cow, collections::HashSet, ops::Deref, sync::Arc}; +use std::{borrow::Cow, ops::Deref, sync::Arc}; -use arrow_array::{ - cast::AsArray, - types::{Float16Type, Float32Type, Float64Type}, -}; -use arrow_schema::DataType; use lance_core::cache::{CacheKey, CacheKeySchema, KeyBuilder, LanceCache}; use lance_core::deepsize::{Context, DeepSizeOf}; use lance_index::frag_reuse::CompactFragReuseIndex; @@ -27,8 +22,6 @@ use lance_linalg::distance::DistanceType; use lance_table::format::IndexMetadata; use uuid::Uuid; -use crate::{Error, Result}; - /// A type-safe wrapper around a LanceCache that enforces namespaces for index data. pub struct GlobalIndexCache(pub(super) LanceCache); @@ -100,491 +93,6 @@ pub(crate) fn write_index_identity(builder: &mut KeyBuilder, uuid: &Uuid, fri_uu } } -/// One candidate stored by the experimental vector-results cache. -/// -/// The offset is meaningful only inside `partition_id` of the exact index -/// segment identified by [`VectorResultsCacheIdentity`]. It must still be -/// bounds-checked against the loaded partition before use. -#[derive(Clone, Copy, Debug, DeepSizeOf, Eq, Hash, PartialEq)] -pub struct CachedVectorCandidate { - partition_id: u32, - offset_in_partition: u32, -} - -impl CachedVectorCandidate { - /// Create a partition-local candidate identity. - pub fn new(partition_id: u32, offset_in_partition: u32) -> Self { - Self { - partition_id, - offset_in_partition, - } - } - - /// Return the IVF partition containing this candidate. - pub fn partition_id(&self) -> u32 { - self.partition_id - } - - /// Return the candidate's local offset inside its IVF partition. - pub fn offset_in_partition(&self) -> u32 { - self.offset_in_partition - } -} - -/// Domain-separates exact query fingerprints from other BLAKE3 uses. -const VECTOR_RESULTS_QUERY_FINGERPRINT_CONTEXT: &str = - "lance.vector-results-cache-query-fingerprint.v1"; - -/// Complete compatibility identity for one reusable vector candidate pool. -/// -/// The exact query values are represented by a cryptographic fingerprint. A -/// candidate pool can therefore be reused only by the same query bit pattern; -/// the ordered IVF partition list remains in the identity to bind the pool to -/// its search shape. Unsupported query forms are rejected by [`Self::try_new`] -/// instead of being represented by a partial key. -#[derive(Clone, Debug, Eq, PartialEq)] -pub struct VectorResultsCacheIdentity { - store_identity: String, - dataset_read_identity: String, - dataset_version: u64, - index_uuid: Uuid, - frag_reuse_uuid: Option, - index_dataset_version: u64, - index_version: i32, - index_base_id: Option, - index_fields: Vec, - column: String, - metric_variant: u32, - sub_index_variant: u32, - quantization_variant: u32, - approx_mode_variant: u32, - vector_type_variant: u32, - dimension: u32, - query_fingerprint: [u8; 32], - nprobes: u32, - partition_ids: Vec, - result_limit: u32, - candidate_limit: u32, -} - -impl DeepSizeOf for VectorResultsCacheIdentity { - fn deep_size_of_children(&self, context: &mut Context) -> usize { - self.store_identity.deep_size_of_children(context) - + self.dataset_read_identity.deep_size_of_children(context) - + self.index_fields.deep_size_of_children(context) - + self.column.deep_size_of_children(context) - + self.partition_ids.deep_size_of_children(context) - } -} - -impl VectorResultsCacheIdentity { - /// Build a fail-closed cache identity for the initial supported query scope. - /// - /// `dataset_read_identity` must identify the exact manifest being read, not - /// merely the dataset URI. `has_prefilter` represents a user/scalar-index - /// prefilter; ordinary deletion filtering is safe because the exact dataset - /// read identity and version are included in the key. - #[allow(clippy::too_many_arguments)] - pub fn try_new( - store_identity: &str, - dataset_read_identity: &str, - dataset_version: u64, - index: &IndexMetadata, - frag_reuse_uuid: Option, - query: &Query, - index_metric: DistanceType, - sub_index_type: SubIndexType, - quantization_type: QuantizationType, - partition_ids: &[u32], - candidate_limit: usize, - has_prefilter: bool, - has_overlay: bool, - ) -> Result { - if store_identity.is_empty() { - return Err(Error::invalid_input( - "vector results cache requires a non-empty object-store identity", - )); - } - if dataset_read_identity.is_empty() { - return Err(Error::invalid_input( - "vector results cache requires an exact dataset read identity", - )); - } - if index.fragment_bitmap.is_none() { - return Err(Error::not_supported(format!( - "vector results cache requires known fragment coverage for index {}", - index.uuid - ))); - } - if index.index_details.is_none() { - return Err(Error::not_supported(format!( - "vector results cache requires current index details for index {}", - index.uuid - ))); - } - if index.fields.is_empty() { - return Err(Error::invalid_input(format!( - "vector results cache cannot identify an index with no fields: {}", - index.uuid - ))); - } - if has_prefilter { - return Err(Error::not_supported( - "vector results cache does not support prefiltered queries", - )); - } - if has_overlay { - return Err(Error::not_supported( - "vector results cache does not support data overlays", - )); - } - if query.lower_bound.is_some() || query.upper_bound.is_some() { - return Err(Error::not_supported( - "vector results cache does not support distance bounds", - )); - } - if query.ef.is_some() { - return Err(Error::not_supported( - "vector results cache does not support HNSW query parameters", - )); - } - if !query.use_index { - return Err(Error::invalid_input( - "vector results cache requires index search to be enabled", - )); - } - if query.column.is_empty() { - return Err(Error::invalid_input( - "vector results cache requires a non-empty vector column", - )); - } - if query.key.is_empty() || query.key.null_count() != 0 { - return Err(Error::invalid_input(format!( - "vector results cache requires a non-empty, non-null query vector, got length {} with {} nulls", - query.key.len(), - query.key.null_count() - ))); - } - - let mut query_hasher = - blake3::Hasher::new_derive_key(VECTOR_RESULTS_QUERY_FINGERPRINT_CONTEXT); - let (vector_type_variant, is_finite) = match query.key.data_type() { - DataType::Float16 => { - let values = query.key.as_primitive::().values(); - query_hasher.update(values.inner().as_slice()); - (0, values.iter().all(|value| value.is_finite())) - } - DataType::Float32 => { - let values = query.key.as_primitive::().values(); - query_hasher.update(values.inner().as_slice()); - (1, values.iter().all(|value| value.is_finite())) - } - DataType::Float64 => { - let values = query.key.as_primitive::().values(); - query_hasher.update(values.inner().as_slice()); - (2, values.iter().all(|value| value.is_finite())) - } - data_type => { - return Err(Error::not_supported(format!( - "vector results cache supports one float vector, got query type {data_type}" - ))); - } - }; - if !is_finite { - return Err(Error::invalid_input( - "vector results cache requires a finite query vector", - )); - } - let dimension = u32::try_from(query.key.len()).map_err(|_| { - Error::invalid_input(format!( - "query vector dimension {} exceeds the vector results cache limit", - query.key.len() - )) - })?; - let query_fingerprint = *query_hasher.finalize().as_bytes(); - - let metric_type = query.metric_type.unwrap_or(index_metric); - if metric_type != index_metric { - return Err(Error::invalid_input(format!( - "query metric {metric_type} does not match index metric {index_metric} for vector results cache" - ))); - } - let metric_variant = match metric_type { - DistanceType::L2 => 0, - DistanceType::Cosine => 1, - DistanceType::Dot => 2, - DistanceType::Hamming => { - return Err(Error::not_supported( - "vector results cache does not support Hamming distance", - )); - } - }; - let sub_index_variant = match sub_index_type { - SubIndexType::Flat => 0, - SubIndexType::Hnsw => { - return Err(Error::not_supported( - "vector results cache initially supports only flat IVF sub-indices", - )); - } - }; - let quantization_variant = match quantization_type { - QuantizationType::Scalar => 0, - QuantizationType::Rabit => { - if query.approx_mode != ApproxMode::Accurate { - return Err(Error::not_supported(format!( - "vector results cache supports RQ only with ApproxMode::Accurate; got {:?}", - query.approx_mode - ))); - } - 1 - } - other => { - return Err(Error::not_supported(format!( - "vector results cache does not support {other} quantization" - ))); - } - }; - let approx_mode_variant = match query.approx_mode { - ApproxMode::Fast => 0, - ApproxMode::Normal => 1, - ApproxMode::Accurate => 2, - }; - - let Some(maximum_nprobes) = query.maximum_nprobes else { - return Err(Error::not_supported( - "vector results cache requires a fixed maximum_nprobes", - )); - }; - if query.minimum_nprobes != maximum_nprobes { - return Err(Error::not_supported(format!( - "vector results cache requires fixed nprobes, got minimum_nprobes={} and maximum_nprobes={maximum_nprobes}", - query.minimum_nprobes - ))); - } - if maximum_nprobes == 0 || maximum_nprobes != partition_ids.len() { - return Err(Error::invalid_input(format!( - "vector results cache nprobes {maximum_nprobes} does not match {} searched partitions", - partition_ids.len() - ))); - } - if partition_ids.iter().copied().collect::>().len() != partition_ids.len() { - return Err(Error::invalid_input(format!( - "vector results cache requires unique partition ids, got {partition_ids:?}" - ))); - } - let nprobes = u32::try_from(maximum_nprobes).map_err(|_| { - Error::invalid_input(format!( - "nprobes {maximum_nprobes} exceeds the vector results cache limit" - )) - })?; - - let refine_factor = query.refine_factor.unwrap_or(1) as usize; - let result_limit = query.k.checked_mul(refine_factor).ok_or_else(|| { - Error::invalid_input(format!( - "vector results cache result limit overflows: k={} refine_factor={refine_factor}", - query.k - )) - })?; - if result_limit == 0 || candidate_limit < result_limit { - return Err(Error::invalid_input(format!( - "vector results cache candidate limit {candidate_limit} must be at least the requested result limit {result_limit}" - ))); - } - let result_limit = u32::try_from(result_limit).map_err(|_| { - Error::invalid_input(format!( - "result limit {result_limit} exceeds the vector results cache limit" - )) - })?; - let candidate_limit = u32::try_from(candidate_limit).map_err(|_| { - Error::invalid_input(format!( - "candidate limit {candidate_limit} exceeds the vector results cache limit" - )) - })?; - - Ok(Self { - store_identity: store_identity.to_owned(), - dataset_read_identity: dataset_read_identity.to_owned(), - dataset_version, - index_uuid: index.uuid, - frag_reuse_uuid, - index_dataset_version: index.dataset_version, - index_version: index.index_version, - index_base_id: index.base_id, - index_fields: index.fields.clone(), - column: query.column.clone(), - metric_variant, - sub_index_variant, - quantization_variant, - approx_mode_variant, - vector_type_variant, - dimension, - query_fingerprint, - nprobes, - partition_ids: partition_ids.to_vec(), - result_limit, - candidate_limit, - }) - } - - /// Return the exact ordered partition list represented by this cache key. - pub fn partition_ids(&self) -> &[u32] { - &self.partition_ids - } - - /// Return the maximum number of candidates stored in a compatible entry. - pub fn candidate_limit(&self) -> usize { - self.candidate_limit as usize - } - - /// Return the number of scored candidates emitted for this query shape. - pub fn result_limit(&self) -> usize { - self.result_limit as usize - } -} - -/// In-memory candidate pool stored under a [`VectorResultsCacheIdentity`]. -/// -/// No codec is registered yet, so this does not create a persistent cache or a -/// file-format compatibility surface. -#[derive(Clone, Debug, DeepSizeOf)] -pub struct VectorResultsCacheEntry { - identity: VectorResultsCacheIdentity, - candidates: Vec, -} - -impl VectorResultsCacheEntry { - /// Create an entry after validating its candidate count and partitions. - pub fn try_new( - identity: VectorResultsCacheIdentity, - candidates: Vec, - ) -> Result { - let entry = Self { - identity, - candidates, - }; - if !entry.has_valid_shape() { - return Err(Error::invalid_input( - "vector results cache entry contains incompatible or duplicate candidates", - )); - } - Ok(entry) - } - - /// Return true only if the entry exactly matches the requested identity and - /// all candidates satisfy the identity's structural constraints. - pub fn is_compatible_with(&self, identity: &VectorResultsCacheIdentity) -> bool { - self.identity == *identity && self.has_valid_shape() - } - - /// Return the validated partition-local candidates. - pub fn candidates(&self) -> &[CachedVectorCandidate] { - &self.candidates - } - - fn has_valid_shape(&self) -> bool { - if self.candidates.len() > self.identity.candidate_limit() { - return false; - } - let valid_partitions = self - .identity - .partition_ids() - .iter() - .copied() - .collect::>(); - let unique_candidates = self.candidates.iter().copied().collect::>(); - unique_candidates.len() == self.candidates.len() - && self - .candidates - .iter() - .all(|candidate| valid_partitions.contains(&candidate.partition_id())) - } -} - -impl CacheKey for VectorResultsCacheIdentity { - type ValueType = VectorResultsCacheEntry; - - fn key(&self) -> Cow<'_, str> { - let partition_ids = self - .partition_ids - .iter() - .map(u32::to_string) - .collect::>() - .join(","); - let base_id = self.index_base_id.map_or(-1, i64::from); - Cow::Owned(format!( - "{store_len}:{store}/{read_len}:{read}/{dataset_version}/{index_uuid}/{frag_reuse_uuid:?}/{index_dataset_version}/{index_version}/{base_id}/{index_fields:?}/{column_len}:{column}/{metric_variant}/{sub_index_variant}/{quantization_variant}/{approx_mode_variant}/{vector_type_variant}/{dimension}/{query_fingerprint:?}/{nprobes}/{partition_ids}/{result_limit}/{candidate_limit}", - store_len = self.store_identity.len(), - store = self.store_identity, - read_len = self.dataset_read_identity.len(), - read = self.dataset_read_identity, - dataset_version = self.dataset_version, - index_uuid = self.index_uuid, - frag_reuse_uuid = self.frag_reuse_uuid, - index_dataset_version = self.index_dataset_version, - index_version = self.index_version, - index_fields = self.index_fields, - column_len = self.column.len(), - column = self.column, - metric_variant = self.metric_variant, - sub_index_variant = self.sub_index_variant, - quantization_variant = self.quantization_variant, - approx_mode_variant = self.approx_mode_variant, - vector_type_variant = self.vector_type_variant, - dimension = self.dimension, - query_fingerprint = self.query_fingerprint, - nprobes = self.nprobes, - result_limit = self.result_limit, - candidate_limit = self.candidate_limit, - )) - } - - fn type_name() -> &'static str { - "VectorResultsCacheEntry" - } - - fn stable_type_id() -> &'static str { - "lance.VectorResultsCacheEntry" - } - - fn schema() -> CacheKeySchema { - CacheKeySchema::new("lance.vector-results-cache-key", 2) - } - - fn write_key(&self, builder: &mut KeyBuilder) { - builder.write_str(&self.store_identity); - builder.write_str(&self.dataset_read_identity); - builder.write_u64(self.dataset_version); - write_index_identity(builder, &self.index_uuid, self.frag_reuse_uuid.as_ref()); - builder.write_u64(self.index_dataset_version); - builder.write_i32(self.index_version); - if let Some(base_id) = self.index_base_id { - builder.write_some(); - builder.write_u32(base_id); - } else { - builder.write_none(); - } - builder.write_sequence_len(self.index_fields.len() as u64); - for field_id in &self.index_fields { - builder.write_i32(*field_id); - } - builder.write_str(&self.column); - builder.write_variant(self.metric_variant); - builder.write_variant(self.sub_index_variant); - builder.write_variant(self.quantization_variant); - builder.write_variant(self.approx_mode_variant); - builder.write_variant(self.vector_type_variant); - builder.write_u32(self.dimension); - builder.write_fixed_bytes(&self.query_fingerprint); - builder.write_u32(self.nprobes); - builder.write_sequence_len(self.partition_ids.len() as u64); - for partition_id in &self.partition_ids { - builder.write_u32(*partition_id); - } - builder.write_u32(self.result_limit); - builder.write_u32(self.candidate_limit); - } -} - // Cache key types for type-safe cache access #[derive(Debug)] @@ -703,89 +211,6 @@ impl CacheKey for ScalarIndexDetailsKey<'_> { #[cfg(test)] mod tests { use super::*; - use arrow_array::{ArrayRef, Float32Array}; - use lance_core::cache::{CacheNamespace, InternalCacheKey}; - use lance_index::vector::DEFAULT_QUERY_PARALLELISM; - - fn vector_index_metadata() -> IndexMetadata { - IndexMetadata { - uuid: Uuid::from_u128(0x11111111_2222_3333_4444_555555555555), - fields: vec![7], - covering_fields: vec![], - name: "vector_idx".to_string(), - dataset_version: 5, - fragment_bitmap: Some([1_u32, 3, 8].into_iter().collect()), - index_details: Some(Arc::new(prost_types::Any { - type_url: "type.googleapis.com/lance.table.VectorIndexDetails".to_string(), - value: vec![1], - })), - index_version: 3, - created_at: None, - base_id: Some(2), - files: None, - } - } - - fn vector_query() -> Query { - Query { - column: "vector".to_string(), - key: Arc::new(Float32Array::from(vec![0.25, 0.5, 0.75])) as ArrayRef, - k: 10, - lower_bound: None, - upper_bound: None, - minimum_nprobes: 3, - maximum_nprobes: Some(3), - ef: None, - refine_factor: Some(2), - metric_type: Some(DistanceType::Cosine), - use_index: true, - query_parallelism: DEFAULT_QUERY_PARALLELISM, - dist_q_c: f32::NAN, - approx_mode: ApproxMode::Accurate, - } - } - - fn vector_results_identity_for(query: &Query) -> VectorResultsCacheIdentity { - VectorResultsCacheIdentity::try_new( - "s3$account-a", - "_versions/17.manifest", - 17, - &vector_index_metadata(), - Some(Uuid::from_u128(0xaaaaaaaa_bbbb_cccc_dddd_eeeeeeeeeeee)), - query, - DistanceType::Cosine, - SubIndexType::Flat, - QuantizationType::Rabit, - &[2, 5, 9], - 100, - false, - false, - ) - .unwrap() - } - - fn vector_results_identity() -> VectorResultsCacheIdentity { - vector_results_identity_for(&vector_query()) - } - - fn physical_key(identity: &VectorResultsCacheIdentity) -> InternalCacheKey { - let mut builder = KeyBuilder::new( - CacheNamespace::root(), - VectorResultsCacheIdentity::stable_type_id(), - VectorResultsCacheIdentity::schema(), - ); - identity.write_key(&mut builder); - builder.finish() - } - - fn assert_identity_field_isolated( - base: &VectorResultsCacheIdentity, - change: impl FnOnce(&mut VectorResultsCacheIdentity), - ) { - let mut changed = base.clone(); - change(&mut changed); - assert_ne!(physical_key(base), physical_key(&changed)); - } #[test] fn index_metadata_key_isolates_object_store_identity() { @@ -818,284 +243,4 @@ mod tests { assert_ne!(first.key(), second.key()); } - - #[test] - fn vector_results_cache_key_isolates_every_identity_axis() { - let base = vector_results_identity(); - - assert_identity_field_isolated(&base, |identity| { - identity.store_identity.push_str("-rotated") - }); - assert_identity_field_isolated(&base, |identity| { - identity.dataset_read_identity.push_str("-detached") - }); - assert_identity_field_isolated(&base, |identity| identity.dataset_version += 1); - assert_identity_field_isolated(&base, |identity| identity.index_uuid = Uuid::new_v4()); - assert_identity_field_isolated(&base, |identity| identity.frag_reuse_uuid = None); - assert_identity_field_isolated(&base, |identity| identity.index_dataset_version += 1); - assert_identity_field_isolated(&base, |identity| identity.index_version += 1); - assert_identity_field_isolated(&base, |identity| identity.index_base_id = None); - assert_identity_field_isolated(&base, |identity| identity.index_fields.push(8)); - assert_identity_field_isolated(&base, |identity| identity.column.push_str("_new")); - assert_identity_field_isolated(&base, |identity| identity.metric_variant += 1); - assert_identity_field_isolated(&base, |identity| identity.sub_index_variant += 1); - assert_identity_field_isolated(&base, |identity| identity.quantization_variant += 1); - assert_identity_field_isolated(&base, |identity| identity.approx_mode_variant += 1); - assert_identity_field_isolated(&base, |identity| identity.vector_type_variant += 1); - assert_identity_field_isolated(&base, |identity| identity.dimension += 1); - assert_identity_field_isolated(&base, |identity| identity.query_fingerprint[0] ^= 1); - assert_identity_field_isolated(&base, |identity| identity.nprobes += 1); - assert_identity_field_isolated(&base, |identity| identity.partition_ids.swap(0, 1)); - assert_identity_field_isolated(&base, |identity| identity.result_limit += 1); - assert_identity_field_isolated(&base, |identity| identity.candidate_limit += 1); - } - - #[test] - fn vector_results_cache_key_isolates_exact_query_values() { - let first_query = vector_query(); - let mut second_query = first_query.clone(); - second_query.key = Arc::new(Float32Array::from(vec![0.25, 0.5, 0.76])); - - let first = vector_results_identity_for(&first_query); - let second = vector_results_identity_for(&second_query); - - assert_eq!(first.partition_ids(), second.partition_ids()); - assert_ne!(first.query_fingerprint, second.query_fingerprint); - assert_ne!(physical_key(&first), physical_key(&second)); - } - - #[test] - fn vector_results_cache_identity_rejects_unsafe_query_forms() { - let index = vector_index_metadata(); - let make_identity = |query: &Query, - index: &IndexMetadata, - sub_index_type, - quantization_type, - partition_ids: &[u32], - has_prefilter, - has_overlay| { - VectorResultsCacheIdentity::try_new( - "s3$account-a", - "_versions/17.manifest", - 17, - index, - None, - query, - DistanceType::Cosine, - sub_index_type, - quantization_type, - partition_ids, - 100, - has_prefilter, - has_overlay, - ) - }; - - let query = vector_query(); - assert!( - make_identity( - &query, - &index, - SubIndexType::Flat, - QuantizationType::Rabit, - &[2, 5, 9], - true, - false, - ) - .is_err() - ); - assert!( - make_identity( - &query, - &index, - SubIndexType::Flat, - QuantizationType::Rabit, - &[2, 5, 9], - false, - true, - ) - .is_err() - ); - - let mut bounded = query.clone(); - bounded.lower_bound = Some(0.1); - let mut adaptive = query.clone(); - adaptive.maximum_nprobes = None; - let mut mismatched_metric = query.clone(); - mismatched_metric.metric_type = Some(DistanceType::Dot); - let mut hnsw_query = query.clone(); - hnsw_query.ef = Some(64); - let mut empty_query = query.clone(); - empty_query.key = Arc::new(Float32Array::from(Vec::::new())); - let mut null_query = query.clone(); - null_query.key = Arc::new(Float32Array::from(vec![Some(0.25), None, Some(0.75)])); - let mut nan_query = query.clone(); - nan_query.key = Arc::new(Float32Array::from(vec![0.25, f32::NAN, 0.75])); - let mut infinite_query = query.clone(); - infinite_query.key = Arc::new(Float32Array::from(vec![0.25, f32::INFINITY, 0.75])); - for unsupported_query in [ - bounded, - adaptive, - mismatched_metric, - hnsw_query, - empty_query, - null_query, - nan_query, - infinite_query, - ] { - assert!( - make_identity( - &unsupported_query, - &index, - SubIndexType::Flat, - QuantizationType::Rabit, - &[2, 5, 9], - false, - false, - ) - .is_err() - ); - } - - for approx_mode in [ApproxMode::Fast, ApproxMode::Normal] { - let mut rq_query = query.clone(); - rq_query.approx_mode = approx_mode; - let error = make_identity( - &rq_query, - &index, - SubIndexType::Flat, - QuantizationType::Rabit, - &[2, 5, 9], - false, - false, - ) - .unwrap_err(); - assert!(matches!(error, Error::NotSupported { .. })); - assert!( - error - .to_string() - .contains("supports RQ only with ApproxMode::Accurate") - ); - assert!( - make_identity( - &rq_query, - &index, - SubIndexType::Flat, - QuantizationType::Scalar, - &[2, 5, 9], - false, - false, - ) - .is_ok(), - "SQ ignores the RQ-specific approximation policy" - ); - } - - assert!( - make_identity( - &query, - &index, - SubIndexType::Hnsw, - QuantizationType::Scalar, - &[2, 5, 9], - false, - false, - ) - .is_err() - ); - assert!( - make_identity( - &query, - &index, - SubIndexType::Flat, - QuantizationType::Product, - &[2, 5, 9], - false, - false, - ) - .is_err() - ); - assert!( - make_identity( - &query, - &index, - SubIndexType::Flat, - QuantizationType::Rabit, - &[2, 2, 9], - false, - false, - ) - .is_err() - ); - - let mut unknown_coverage = index; - unknown_coverage.fragment_bitmap = None; - assert!( - make_identity( - &query, - &unknown_coverage, - SubIndexType::Flat, - QuantizationType::Rabit, - &[2, 5, 9], - false, - false, - ) - .is_err() - ); - } - - #[test] - fn vector_results_cache_entry_validates_candidate_shape() { - let identity = vector_results_identity(); - let candidates = vec![ - CachedVectorCandidate::new(2, 10), - CachedVectorCandidate::new(5, 20), - CachedVectorCandidate::new(9, 30), - ]; - let entry = VectorResultsCacheEntry::try_new(identity.clone(), candidates).unwrap(); - assert!(entry.is_compatible_with(&identity)); - assert_eq!(entry.candidates()[1].partition_id(), 5); - assert_eq!(entry.candidates()[1].offset_in_partition(), 20); - - let mut different_identity = identity.clone(); - different_identity.dataset_version += 1; - assert!(!entry.is_compatible_with(&different_identity)); - assert!( - VectorResultsCacheEntry::try_new( - identity.clone(), - vec![CachedVectorCandidate::new(7, 10)], - ) - .is_err() - ); - assert!( - VectorResultsCacheEntry::try_new( - identity, - vec![ - CachedVectorCandidate::new(2, 10), - CachedVectorCandidate::new(2, 10), - ], - ) - .is_err() - ); - } - - #[tokio::test] - async fn vector_results_cache_identity_produces_typed_cache_misses() { - let cache = LanceCache::with_capacity(4096); - let identity = vector_results_identity(); - let entry = Arc::new( - VectorResultsCacheEntry::try_new( - identity.clone(), - vec![CachedVectorCandidate::new(2, 10)], - ) - .unwrap(), - ); - cache.insert_with_key(&identity, entry.clone()).await; - let cached = cache.get_with_key(&identity).await.unwrap(); - assert!(cached.is_compatible_with(&identity)); - - let mut other_version = identity; - other_version.dataset_version += 1; - assert!(cache.get_with_key(&other_version).await.is_none()); - } } From 5a8b979114a90941fc3e77d9f3a8c0a2b70226f5 Mon Sep 17 00:00:00 2001 From: Sergey Troshkov Date: Tue, 25 Aug 2026 17:28:13 +0700 Subject: [PATCH 3/4] fix(index): score selective rows during quantized refinement --- rust/lance/src/dataset/scanner.rs | 31 ++++++++++++++++++++++++++ rust/lance/src/io/exec/knn.rs | 3 +++ rust/lance/src/session/index_caches.rs | 4 ---- 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 468c13ab6a1..6a0223b9007 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -16490,6 +16490,37 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") 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) diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 2e8ac22e482..6436708c9fc 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -2119,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 diff --git a/rust/lance/src/session/index_caches.rs b/rust/lance/src/session/index_caches.rs index 2f17f35d165..23922c7a4b0 100644 --- a/rust/lance/src/session/index_caches.rs +++ b/rust/lance/src/session/index_caches.rs @@ -15,10 +15,6 @@ use std::{borrow::Cow, ops::Deref, sync::Arc}; use lance_core::cache::{CacheKey, CacheKeySchema, KeyBuilder, LanceCache}; use lance_core::deepsize::{Context, DeepSizeOf}; use lance_index::frag_reuse::CompactFragReuseIndex; -use lance_index::vector::quantizer::QuantizationType; -use lance_index::vector::v3::subindex::SubIndexType; -use lance_index::vector::{ApproxMode, Query}; -use lance_linalg::distance::DistanceType; use lance_table::format::IndexMetadata; use uuid::Uuid; From 1d567fadcef8537dbbd0e0b916606bb8377ce0d1 Mon Sep 17 00:00:00 2001 From: Sergey Troshkov Date: Mon, 14 Sep 2026 13:03:55 +0700 Subject: [PATCH 4/4] fix(index): rerank batched quantized queries --- rust/lance/src/dataset/scanner.rs | 36 +++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 6a0223b9007..44cfca90de6 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -5609,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 @@ -5634,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 @@ -10131,6 +10132,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