From 973fdfac58e5f56217e1b9f64ca25a3af3acb316 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Tue, 15 Sep 2026 21:42:46 -0500 Subject: [PATCH 1/3] feat(knn): expose ANNIvfBatchExec's query, width and inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shared-scan batch node kept every field private, so a caller that matched it in a plan could not read the query, the number of vectors in it, or the dataset and index metadata behind it — the accessors ANNIvfSubIndexExec has had all along. Sophon's WAL union needs them to build a fresh-tier arm for a batch vector search and to rebuild the node with an inflated candidate k. Also re-export the node and QUERY_INDEX_COL from `lance::io::exec`, alongside the single-query nodes. Co-Authored-By: Claude Opus 5 (1M context) --- rust/lance/src/io/exec.rs | 5 ++++- rust/lance/src/io/exec/knn.rs | 26 ++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) 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..862f53fd2d2 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -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 { From eaacd1789bf484cd747b7d34e2e84e21af3e5728 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Wed, 16 Sep 2026 10:22:11 -0500 Subject: [PATCH 2/3] test(knn): cover the batch plan-inspection surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The accessors and the `io::exec` re-export had no repository coverage, so a regression in either would go unnoticed here. Plan a batch indexed search, downcast the node through the public path, and assert the five accessors hand back the scanner's inputs — with and without a prefilter, so `prefilter_source()` is exercised in both shapes. Co-Authored-By: Claude Opus 5 (1M context) --- rust/lance/src/dataset/scanner.rs | 89 +++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 28f4d341470..09e5349302c 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.as_any().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] From b0533e1a82a9192dce1fcaa1b86eb8f9ba79d35f Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Wed, 16 Sep 2026 10:46:33 -0500 Subject: [PATCH 3/3] fix(knn): repair the batch accessor build and docs Two breaks from this branch: - The plan walk used `as_any().downcast_ref()`, but DataFusion 54 puts `downcast_ref` on `dyn ExecutionPlan` itself, so the test target did not compile. Match `find_filtered_read`, which already walks plans this way. - Re-exporting `ANNIvfBatchExec` brought its doc comment into the public docs, where its intra-doc link to the private `build_dataset_prefilter` is a rustdoc error. Keep the reference, drop the link. Co-Authored-By: Claude Opus 5 (1M context) --- rust/lance/src/dataset/scanner.rs | 2 +- rust/lance/src/io/exec/knn.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 09e5349302c..ac461997454 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -10483,7 +10483,7 @@ 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.as_any().downcast_ref::() { + if let Some(batch_exec) = plan.downcast_ref::() { return Some(batch_exec); } plan.children() diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 862f53fd2d2..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}`,