diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 28f4d341470..ac461997454 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -7596,6 +7596,9 @@ mod test { use crate::dataset::scanner::test_dataset::TestVectorDataset; use crate::dataset::{NewColumnTransform, WriteParams}; use crate::index::vector::{StageParams, VectorIndexParams}; + // Imported through the public `io::exec` re-export rather than the crate-private + // `knn` module, so the tests below cover that surface too. + use crate::io::exec::{ANNIvfBatchExec, QUERY_INDEX_COL}; use crate::utils::test::{ DatagenExt, FragmentCount, FragmentRowCount, ThrottledStoreWrapper, assert_plan_node_equals, }; @@ -10478,6 +10481,92 @@ mod test { } } + /// Finds the batch vector-search node in a physical plan. + fn find_ann_ivf_batch_exec(plan: &dyn ExecutionPlan) -> Option<&ANNIvfBatchExec> { + if let Some(batch_exec) = plan.downcast_ref::() { + return Some(batch_exec); + } + plan.children() + .into_iter() + .find_map(|child| find_ann_ivf_batch_exec(child.as_ref())) + } + + /// A caller that matches the batch node in a plan reads the search back out + /// of it through the public `io::exec` surface, so the accessors must return + /// what the scanner fed the constructor. + #[rstest] + #[case::no_prefilter(None)] + #[case::prefilter(Some("i > 100"))] + #[tokio::test] + async fn test_batch_knn_indexed_exposes_plan_inputs(#[case] filter: Option<&str>) { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + test_ds.make_vector_index().await.unwrap(); + let dataset = &test_ds.dataset; + let (queries, query_values) = batch_knn_two_queries(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &queries, 2).unwrap(); + scan.nprobes(2); + if let Some(filter) = filter { + scan.filter(filter).unwrap(); + scan.prefilter(true); + } + scan.project(&["i"]).unwrap(); + + let plan = scan.create_plan().await.unwrap(); + let batch_exec = find_ann_ivf_batch_exec(plan.as_ref()) + .expect("indexed batch KNN should plan an ANNIvfBatchExec"); + + let query = batch_exec.query(); + assert_eq!(query.column, "vec"); + assert_eq!(query.k, 2); + assert_eq!(query.minimum_nprobes, 2); + assert_eq!(query.maximum_nprobes, Some(2)); + assert_eq!(query.metric_type, Some(DistanceType::L2)); + assert_eq!( + query.key.as_primitive::().values(), + query_values.as_slice(), + "query key must hold both query vectors concatenated" + ); + + assert_eq!(batch_exec.query_count(), 2); + assert_eq!( + query.key.len() / batch_exec.query_count(), + 32, + "query count must divide the key into the column's vectors" + ); + + assert_eq!(batch_exec.dataset().uri(), dataset.uri()); + assert_eq!( + batch_exec.dataset().version().version, + dataset.version().version + ); + + let expected_indices = dataset.load_indices_by_name("idx").await.unwrap(); + assert!(!expected_indices.is_empty()); + assert_eq!( + batch_exec + .indices() + .iter() + .map(|index| index.uuid) + .collect::>(), + expected_indices + .iter() + .map(|index| index.uuid) + .collect::>() + ); + + match (filter, batch_exec.prefilter_source()) { + (None, PreFilterSource::None) => {} + (Some(_), PreFilterSource::FilteredRowIds(_)) => {} + (_, source) => panic!("unexpected prefilter source {source:?} for filter {filter:?}"), + } + + assert_eq!(batch_exec.schema().field(0).name(), QUERY_INDEX_COL); + } + /// Batch indexed search must merge each query's top-k across multiple delta /// indices, not just within a single delta. #[tokio::test] diff --git a/rust/lance/src/io/exec.rs b/rust/lance/src/io/exec.rs index d37b58a238e..b14148b7952 100644 --- a/rust/lance/src/io/exec.rs +++ b/rust/lance/src/io/exec.rs @@ -30,7 +30,10 @@ pub mod testing; pub mod utils; pub use filter::LanceFilterExec; -pub use knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNVectorDistanceExec}; +pub use knn::{ + ANNIvfBatchExec, ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNVectorDistanceExec, + QUERY_INDEX_COL, +}; pub use lance_datafusion::planner::Planner; pub use lance_index::scalar::expression::FilterPlan; pub use optimizer::get_physical_optimizer; diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 76fcd62f5c9..3f10653cf32 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -2475,7 +2475,7 @@ pub fn new_knn_batch_exec( /// because the two-node pipeline streams one partition-list per delta through a /// per-query top-k, whereas the shared scan must invert queries onto partitions /// and keep one heap per query in a single pass. It still reuses the underlying -/// primitives (partition load, prefilter wiring via [`build_dataset_prefilter`], +/// primitives (partition load, prefilter wiring via `build_dataset_prefilter`, /// and the per-partition accumulate the index performs). /// /// Output schema: `{query_index: Int32, _distance: Float32, _rowid: UInt64}`, @@ -2531,6 +2531,32 @@ impl ANNIvfBatchExec { metrics: ExecutionPlanMetricsSet::new(), }) } + + /// Returns a reference to the vector query. Its `key` holds all + /// [`Self::query_count`] vectors concatenated. + pub fn query(&self) -> &Query { + &self.query + } + + /// Returns the number of query vectors packed into [`Self::query`]. + pub fn query_count(&self) -> usize { + self.query_count + } + + /// Returns a reference to the dataset. + pub fn dataset(&self) -> &Arc { + &self.dataset + } + + /// Returns a reference to the index metadata. + pub fn indices(&self) -> &[IndexMetadata] { + &self.indices + } + + /// Returns a reference to the prefilter source. + pub fn prefilter_source(&self) -> &PreFilterSource { + &self.prefilter_source + } } impl DisplayAs for ANNIvfBatchExec {