From 9b87776d619a42c3355aa59d9d0b9ea8a7cf33ed Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 18 May 2026 16:37:43 +0800 Subject: [PATCH 01/32] feat: support batch flat vector queries Add a flat KNN batch query path so callers can submit multiple query vectors and share scan work while preserving per-query top-k results. Co-authored-by: Cursor --- python/python/lance/dataset.py | 17 +- python/python/tests/test_vector_index.py | 35 +++ python/src/dataset.rs | 17 +- rust/lance/benches/vector_index.rs | 187 +++++++++++- rust/lance/src/dataset/scanner.rs | 226 +++++++++++++- rust/lance/src/io/exec.rs | 4 +- rust/lance/src/io/exec/knn.rs | 358 ++++++++++++++++++++++- 7 files changed, 826 insertions(+), 18 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 5737ec013b5..936f5812113 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -92,7 +92,7 @@ pa.Array, pa.Scalar, np.ndarray, - Iterable[float], + Iterable[Union[float, Iterable[float]]], ] LANCE_COMMIT_MESSAGE_KEY = "__lance_commit_message" _BLOB_PANDAS_MODE_LAZY = "lazy" @@ -1099,6 +1099,12 @@ def scanner( "distance_range": (0.0, 1.0), } + ``q`` may also be a 2-D array-like value for fixed-size vector columns. + In that case Lance runs a flat batch KNN query, returns up to ``k`` rows + for each query vector, and adds ``_query_index`` to identify the source + query for each result row. Indexed/ANN batch search is not used in this + first implementation. + batch_size: int, default None The maximum number of rows per batch. In some cases batches can be smaller than this size. Note: this can be overridden by @@ -5989,6 +5995,10 @@ def nearest( Parameters ---------- + q: QueryVectorLike + A single query vector or, for fixed-size vector columns, a 2-D array-like + batch of query vectors. Batch queries return up to ``k`` rows per query + and include ``_query_index`` in the output. query_parallelism: int, optional Maximum partition-search concurrency for a single vector query. The default is 0. Value 0 uses the automatic policy, which @@ -7137,7 +7147,10 @@ def _build_vector_search_query( column: str The name of the vector column to search. q: QueryVectorLike - The query vector. + The query vector. For fixed-size vector columns, this may be a 2-D + array-like batch of query vectors. Batch queries run flat KNN, apply + ``k`` per query vector, and add ``_query_index`` to the result so + callers can split rows by input query. k: int, optional The number of nearest neighbors to return. metric: str, optional diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 356f72a5e66..bc448cf8e29 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -180,6 +180,41 @@ def test_flat(dataset): run(dataset) +def test_batch_flat_query_matches_repeated_single_queries(dataset): + queries = np.random.randn(2, 128).astype(np.float32) + k = 5 + + batch = dataset.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": queries, + "k": k, + }, + ) + + assert batch.num_rows == queries.shape[0] * k + assert batch.column_names == ["id", "_distance", "_query_index"] + assert batch["_query_index"].to_pylist() == [0] * k + [1] * k + + for query_index, query in enumerate(queries): + single = dataset.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": query, + "k": k, + "use_index": False, + }, + ) + batch_slice = batch.filter(pc.field("_query_index") == query_index) + assert batch_slice["id"].to_pylist() == single["id"].to_pylist() + np.testing.assert_allclose( + batch_slice["_distance"].to_numpy(), + single["_distance"].to_numpy(), + ) + + def test_ann(indexed_dataset): run(indexed_dataset) diff --git a/python/src/dataset.rs b/python/src/dataset.rs index c868504e87c..1a59348fdc6 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1368,17 +1368,26 @@ impl Dataset { query_parallelism, ) = vector_query_params_from_dict(nearest, default_k)?; - let (_, element_type) = get_vector_type(self_.ds.schema(), &column) + let (vector_type, element_type) = get_vector_type(self_.ds.schema(), &column) .map_err(|e| PyValueError::new_err(e.to_string()))?; - let scanner = match element_type { - DataType::UInt8 => { + let is_batch_query = + matches!(q.data_type(), DataType::List(_) | DataType::FixedSizeList(_, _)) + && matches!(vector_type, DataType::FixedSizeList(_, _)); + let scanner = match (is_batch_query, element_type) { + (true, DataType::UInt8) => { + return Err(PyValueError::new_err( + "Batch nearest is not supported for binary vector columns", + )); + } + (false, DataType::UInt8) => { let q = arrow::compute::cast(&q, &DataType::UInt8).map_err(|e| { PyValueError::new_err(format!("Failed to cast q to binary vector: {}", e)) })?; let q = q.as_primitive::(); scanner.nearest(&column, q, k) } - _ => scanner.nearest(&column, &q, k), + (true, _) => scanner.nearest_batch(&column, &q, k), + (false, _) => scanner.nearest(&column, &q, k), }; let distance_range: Option<(Option, Option)> = if let Some(dr) = nearest.get_item("distance_range")? { diff --git a/rust/lance/benches/vector_index.rs b/rust/lance/benches/vector_index.rs index 21c9aa4e4aa..8004b08f9f5 100644 --- a/rust/lance/benches/vector_index.rs +++ b/rust/lance/benches/vector_index.rs @@ -3,18 +3,19 @@ #![allow(clippy::print_stdout)] use std::sync::Arc; +use std::time::Duration; use arrow_array::{ FixedSizeListArray, Float32Array, RecordBatch, RecordBatchIterator, cast::as_primitive_array, }; use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema}; -use criterion::{Criterion, criterion_group, criterion_main}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use futures::TryStreamExt; #[cfg(target_os = "linux")] use pprof::criterion::{Output, PProfProfiler}; use rand::Rng; -use lance::dataset::{Dataset, WriteMode, WriteParams, builder::DatasetBuilder}; +use lance::dataset::{Dataset, ReadParams, WriteMode, WriteParams, builder::DatasetBuilder}; use lance::index::DatasetIndexExt; use lance::index::vector::VectorIndexParams; use lance_arrow::{FixedSizeListArrayExt, as_fixed_size_list_array}; @@ -22,7 +23,23 @@ use lance_index::{ IndexType, vector::{ivf::IvfBuildParams, pq::PQBuildParams}, }; +use lance_io::object_store::{ObjectStoreParams, WrappingObjectStore}; use lance_linalg::distance::MetricType; +use object_store::{ + ObjectStore, + throttle::{ThrottleConfig, ThrottledStore}, +}; + +#[derive(Debug)] +struct ThrottledStoreWrapper { + config: ThrottleConfig, +} + +impl WrappingObjectStore for ThrottledStoreWrapper { + fn wrap(&self, _prefix: &str, original: Arc) -> Arc { + Arc::new(ThrottledStore::new(original, self.config.clone())) + } +} fn bench_ivf_pq_index(c: &mut Criterion) { // default tokio runtime @@ -124,6 +141,94 @@ fn bench_ivf_pq_index(c: &mut Criterion) { ); } +fn bench_batch_flat_knn(c: &mut Criterion) { + const DIM: i32 = 4; + const K: usize = 10; + const QUERY_COUNT: usize = 8; + + let rt = tokio::runtime::Runtime::new().unwrap(); + let uri = format!("memory://batch_flat_vec_data_{}", rand::random::()); + let dataset = rt.block_on(async { + create_flat_file( + &uri, + WriteMode::Create, + 50_000, + 5_000, + DIM, + Duration::from_millis(5), + ) + .await + }); + let first_batch = rt.block_on(async { + dataset + .scan() + .try_into_stream() + .await + .unwrap() + .try_next() + .await + .unwrap() + .unwrap() + }); + let vector_column = first_batch.column_by_name("vector").unwrap(); + let vectors = as_fixed_size_list_array(vector_column); + let query_values = (0..QUERY_COUNT) + .flat_map(|query_index| { + let values = vectors.value(query_index); + as_primitive_array::(&values) + .values() + .to_vec() + }) + .collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), DIM) + .unwrap(); + + let mut group = c.benchmark_group("batch_flat_knn"); + group.bench_function(BenchmarkId::new("separate_queries", QUERY_COUNT), |b| { + b.to_async(&rt).iter(|| async { + for query_index in 0..QUERY_COUNT { + let query = Float32Array::from( + query_values[query_index * DIM as usize..(query_index + 1) * DIM as usize] + .to_vec(), + ); + let results = dataset + .scan() + .nearest("vector", &query, K) + .unwrap() + .use_index(false) + .project::<&str>(&[]) + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(!results.is_empty()); + } + }) + }); + group.bench_function(BenchmarkId::new("batch_query", QUERY_COUNT), |b| { + b.to_async(&rt).iter(|| async { + let results = dataset + .scan() + .nearest_batch("vector", &queries, K) + .unwrap() + .project::<&str>(&[]) + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(!results.is_empty()); + }) + }); + group.finish(); +} + async fn create_file(path: &std::path::Path, mode: WriteMode) { let schema = Arc::new(ArrowSchema::new(vec![Field::new( "vector", @@ -187,6 +292,80 @@ async fn create_file(path: &std::path::Path, mode: WriteMode) { .unwrap(); } +async fn create_flat_file( + uri: &str, + mode: WriteMode, + num_rows: i32, + batch_size: i32, + dim: i32, + wait_get_per_call: Duration, +) -> Dataset { + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "vector", + DataType::FixedSizeList( + FieldRef::new(Field::new("item", DataType::Float32, true)), + dim, + ), + false, + )])); + + let batches: Vec = (0..(num_rows / batch_size)) + .map(|_| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new( + FixedSizeListArray::try_new_from_values( + create_float32_array(batch_size * dim), + dim, + ) + .unwrap(), + )], + ) + .unwrap() + }) + .collect(); + + if !uri.starts_with("memory://") { + std::fs::remove_dir_all(uri).map_or_else(|_| println!("{} not exists", uri), |_| {}); + } + let store_params = if wait_get_per_call.is_zero() { + None + } else { + Some(ObjectStoreParams { + object_store_wrapper: Some(Arc::new(ThrottledStoreWrapper { + config: ThrottleConfig { + wait_get_per_call, + ..Default::default() + }, + })), + ..Default::default() + }) + }; + let write_params = WriteParams { + max_rows_per_file: num_rows as usize, + max_rows_per_group: batch_size as usize, + store_params: store_params.clone(), + mode, + ..Default::default() + }; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); + let dataset = Dataset::write(reader, uri, Some(write_params)) + .await + .unwrap(); + if uri.starts_with("memory://") { + dataset + } else { + DatasetBuilder::from_uri(uri) + .with_read_params(ReadParams { + store_options: store_params, + ..Default::default() + }) + .load() + .await + .unwrap() + } +} + fn create_float32_array(num_elements: i32) -> Float32Array { // generate an Arrow Float32Array with 10000*128 elements randomly let mut rng = rand::rng(); @@ -202,12 +381,12 @@ criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10) .with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); - targets = bench_ivf_pq_index); + targets = bench_ivf_pq_index, bench_batch_flat_knn); // Non-linux version does not support pprof. #[cfg(not(target_os = "linux"))] criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10); - targets = bench_ivf_pq_index); + targets = bench_ivf_pq_index, bench_batch_flat_knn); criterion_main!(benches); diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 547707affcb..8cce12f2c45 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -22,6 +22,7 @@ use datafusion::logical_expr::{Expr, ScalarUDF, col, lit}; use datafusion::physical_expr::PhysicalSortExpr; #[allow(deprecated)] use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::expressions; use datafusion::physical_plan::projection::ProjectionExec as DFProjectionExec; use datafusion::physical_plan::sorts::sort::SortExec; @@ -98,9 +99,10 @@ use crate::io::exec::fts::{ use crate::io::exec::knn::MultivectorScoringExec; use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; use crate::io::exec::{ - AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNVectorDistanceExec, LancePushdownScanExec, - LanceScanExec, Planner, PreFilterSource, ScanConfig, TakeExec, - knn::{KNN_INDEX_SCHEMA, new_knn_exec}, + AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNBatchVectorDistanceExec, + KNNVectorDistanceExec, LancePushdownScanExec, LanceScanExec, Planner, PreFilterSource, + ScanConfig, TakeExec, + knn::{KNN_INDEX_SCHEMA, QUERY_INDEX_COL, new_knn_exec}, project, }; use crate::io::exec::{AddRowOffsetExec, LanceFilterExec, LanceScanConfig, get_physical_optimizer}; @@ -768,6 +770,7 @@ pub struct Scanner { ordering: Option>, nearest: Option, + nearest_query_count: usize, /// If false, do not use any scalar indices for the scan /// @@ -1023,6 +1026,7 @@ impl Scanner { offset: None, ordering: None, nearest: None, + nearest_query_count: 1, use_stats: true, ordered: true, fragments: None, @@ -1528,6 +1532,120 @@ impl Scanner { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, }); + self.nearest_query_count = 1; + Ok(self) + } + + /// Find the k-nearest neighbors for each query vector in a batch. + /// + /// This first implementation uses the flat scan path. It reads candidate + /// vectors once and computes distances for all query vectors, returning up + /// to `k` rows per query with an additional `_query_index` column. + pub fn nearest_batch(&mut self, column: &str, q: &dyn Array, k: usize) -> Result<&mut Self> { + if !self.prefilter { + // We can allow fragment scan if the input to nearest is a prefilter. + // The fragment scan will be performed by the prefilter. + self.ensure_not_fragment_scan()?; + } + + if k == 0 { + return Err(Error::invalid_input("k must be positive".to_string())); + } + if q.is_empty() { + return Err(Error::invalid_input( + "Query vector batch must have at least one query".to_string(), + )); + } + + let (vector_type, element_type) = get_vector_type(self.dataset.schema(), column)?; + if matches!(vector_type, DataType::List(_)) { + return Err(Error::not_supported( + "Batch nearest is not supported for multivector columns".to_string(), + )); + } + let dim = get_vector_dim(self.dataset.schema(), column)?; + + let (q, query_count) = match q.data_type() { + DataType::FixedSizeList(_, _) => { + let fsl = q.as_fixed_size_list(); + if fsl.value_length() as usize != dim { + return Err(Error::invalid_input(format!( + "query dim({}) doesn't match the column {} vector dim({})", + fsl.value_length(), + column, + dim, + ))); + } + (fsl.values().clone(), fsl.len()) + } + DataType::List(_) => { + let list_array = q.as_list::(); + for i in 0..list_array.len() { + let vec = list_array.value(i); + if vec.len() != dim { + return Err(Error::invalid_input(format!( + "query dim({}) doesn't match the column {} vector dim({})", + vec.len(), + column, + dim, + ))); + } + } + (list_array.values().clone(), list_array.len()) + } + _ => { + if q.len() % dim != 0 { + return Err(Error::invalid_input(format!( + "query batch len({}) must be a multiple of column {} vector dim({})", + q.len(), + column, + dim, + ))); + } + (q.slice(0, q.len()), q.len() / dim) + } + }; + + if query_count == 0 { + return Err(Error::invalid_input( + "Query vector batch must have at least one query".to_string(), + )); + } + + let key = match &element_type { + dt if dt == q.data_type() => q, + dt if dt.is_floating() => coerce_float_vector( + q.as_any().downcast_ref::().unwrap(), + FloatType::try_from(dt)?, + )?, + _ => { + return Err(Error::invalid_input(format!( + "Column {} has element type {} and the query vector batch is {}", + column, + element_type, + q.data_type(), + ))); + } + }; + + self.nearest = Some(Query { + column: column.to_string(), + key, + k, + lower_bound: None, + upper_bound: None, + minimum_nprobes: 1, + maximum_nprobes: None, + ef: None, + refine_factor: None, + metric_type: None, + // Batch KNN is flat-only for now. ANN batching needs per-query + // partition/result grouping instead of the single-query top-k plan. + use_index: false, + query_parallelism: DEFAULT_QUERY_PARALLELISM, + dist_q_c: 0.0, + }); + self.nearest_query_count = query_count; Ok(self) } @@ -1628,7 +1746,12 @@ impl Scanner { /// This is essentially a weak consistency search, only on the indexed data. pub fn fast_search(&mut self) -> &mut Self { if let Some(q) = self.nearest.as_mut() { - q.use_index = true; + if self.nearest_query_count > 1 { + log::warn!("fast_search is ignored for batch nearest queries"); + return self; + } else { + q.use_index = true; + } } self.fast_search = true; self.projection_plan.include_row_id(); // fast search requires _rowid @@ -1688,7 +1811,12 @@ impl Scanner { /// Set whether to use the index if available pub fn use_index(&mut self, use_index: bool) -> &mut Self { if let Some(q) = self.nearest.as_mut() { - q.use_index = use_index + if self.nearest_query_count > 1 && use_index { + log::warn!("use_index(true) is ignored for batch nearest queries"); + q.use_index = false; + } else { + q.use_index = use_index; + } } self } @@ -1858,6 +1986,9 @@ impl Scanner { if self.nearest.as_ref().is_some() { extra_columns.push(ArrowField::new(DIST_COL, DataType::Float32, true)); + if self.nearest_query_count > 1 { + extra_columns.push(ArrowField::new(QUERY_INDEX_COL, DataType::UInt32, true)); + } }; if self.full_text_query.is_some() { @@ -1906,6 +2037,12 @@ impl Scanner { let vector_expr = expressions::col(DIST_COL, current_schema)?; output_expr.push((vector_expr, DIST_COL.to_string())); } + if self.nearest_query_count > 1 + && output_expr.iter().all(|(_, name)| name != QUERY_INDEX_COL) + { + let query_index_expr = expressions::col(QUERY_INDEX_COL, current_schema)?; + output_expr.push((query_index_expr, QUERY_INDEX_COL.to_string())); + } if self.full_text_query.is_some() && output_expr.iter().all(|(_, name)| name != SCORE_COL) { @@ -4267,6 +4404,18 @@ impl Scanner { default_distance_type_for(&element_type) } }; + if self.nearest_query_count > 1 { + let input = Arc::new(CoalescePartitionsExec::new(input)); + return KNNBatchVectorDistanceExec::try_new( + input, + &q.column, + q.key.clone(), + self.nearest_query_count, + q.k, + metric_type, + ) + .map(|exec| Arc::new(exec) as Arc); + } let flat_dist = Arc::new(KNNVectorDistanceExec::try_new( input, &q.column, @@ -4991,7 +5140,7 @@ mod test { use arrow::array::as_primitive_array; use arrow::datatypes::{Float64Type, Int32Type, Int64Type}; use arrow_array::cast::AsArray; - use arrow_array::types::{Float32Type, UInt64Type}; + use arrow_array::types::{Float32Type, UInt32Type, UInt64Type}; use arrow_array::{ ArrayRef, FixedSizeListArray, Float16Array, Int32Array, LargeStringArray, PrimitiveArray, RecordBatchIterator, StringArray, StructArray, UInt8Array, @@ -5637,6 +5786,71 @@ mod test { assert_eq!(expected_i, actual_i); } + #[tokio::test] + async fn test_batch_knn_flat_results_include_query_index() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + test_ds.make_vector_index().await.unwrap(); + let dataset = &test_ds.dataset; + + let query_values = (32..96).map(|v| v as f32).collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) + .unwrap(); + + let mut scan = dataset.scan(); + scan.nearest_batch("vec", &queries, 2).unwrap(); + scan.project(&["i"]).unwrap(); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("KNNBatchVectorDistance"), + "expected flat batch KNN plan, got:\n{}", + plan + ); + assert!( + !plan.contains("ANNSubIndex"), + "batch KNN should not use ANN index yet, got:\n{}", + plan + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 4); + + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + assert_eq!(query_indices.values(), &[0, 0, 1, 1]); + + let batch_ids = batch["i"].as_primitive::(); + let batch_distances = batch[DIST_COL].as_primitive::(); + + for query_index in 0..2 { + let query = + Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); + let single = dataset + .scan() + .nearest("vec", &query, 2) + .unwrap() + .use_index(false) + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let single_ids = single["i"].as_primitive::(); + let single_distances = single[DIST_COL].as_primitive::(); + + for result_index in 0..2 { + let batch_index = query_index * 2 + result_index; + assert_eq!(batch_ids.value(batch_index), single_ids.value(result_index)); + assert_eq!( + batch_distances.value(batch_index), + single_distances.value(result_index) + ); + } + } + } + #[rstest] #[tokio::test] async fn test_can_project_distance() { diff --git a/rust/lance/src/io/exec.rs b/rust/lance/src/io/exec.rs index 0b93e3c2834..4889b8fcd65 100644 --- a/rust/lance/src/io/exec.rs +++ b/rust/lance/src/io/exec.rs @@ -27,7 +27,9 @@ pub mod testing; pub mod utils; pub use filter::LanceFilterExec; -pub use knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNVectorDistanceExec}; +pub use knn::{ + ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNBatchVectorDistanceExec, KNNVectorDistanceExec, +}; 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 49ee7be86bc..2643e2de42b 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use std::any::Any; -use std::collections::{HashMap, HashSet}; +use std::cmp::Ordering as CmpOrdering; +use std::collections::{BinaryHeap, HashMap, HashSet}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, Mutex}; use std::time::Instant; @@ -16,6 +17,7 @@ use arrow_array::{ cast::AsArray, }; use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use arrow_select::concat::concat_batches; use datafusion::physical_plan::PlanProperties; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ @@ -69,6 +71,8 @@ use super::utils::{ SelectionVectorToPrefilter, }; +pub const QUERY_INDEX_COL: &str = "_query_index"; + pub struct AnnPartitionMetrics { index_metrics: IndexMetrics, partitions_ranked: Count, @@ -348,6 +352,358 @@ impl ExecutionPlan for KNNVectorDistanceExec { } } +#[derive(Clone)] +struct BatchKnnCandidate { + query_index: u32, + distance: f32, + row_id: u64, + batch: Arc, + row_index: u32, +} + +impl PartialEq for BatchKnnCandidate { + fn eq(&self, other: &Self) -> bool { + self.query_index == other.query_index + && self.distance == other.distance + && self.row_id == other.row_id + && self.row_index == other.row_index + } +} + +impl Eq for BatchKnnCandidate {} + +impl PartialOrd for BatchKnnCandidate { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for BatchKnnCandidate { + fn cmp(&self, other: &Self) -> CmpOrdering { + self.distance + .total_cmp(&other.distance) + .then_with(|| self.row_id.cmp(&other.row_id)) + .then_with(|| self.query_index.cmp(&other.query_index)) + .then_with(|| self.row_index.cmp(&other.row_index)) + } +} + +/// [ExecutionPlan] that computes flat KNN for a batch of query vectors. +/// +/// This node consumes the input once, computes distances from each input vector +/// to each query vector, and keeps the top-k rows independently for every query. +#[derive(Debug)] +pub struct KNNBatchVectorDistanceExec { + pub input: Arc, + pub query: ArrayRef, + pub query_count: usize, + pub k: usize, + pub column: String, + pub distance_type: DistanceType, + + input_schema: SchemaRef, + output_schema: SchemaRef, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl DisplayAs for KNNBatchVectorDistanceExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "KNNBatchVectorDistance: queries={}, k={}, metric={}", + self.query_count, self.k, self.distance_type + ), + DisplayFormatType::TreeRender => write!( + f, + "KNNBatchVectorDistance\nqueries={}\nk={}\nmetric={}", + self.query_count, self.k, self.distance_type + ), + } + } +} + +impl KNNBatchVectorDistanceExec { + pub fn try_new( + input: Arc, + column: &str, + query: ArrayRef, + query_count: usize, + k: usize, + distance_type: DistanceType, + ) -> Result { + if query_count == 0 { + return Err(Error::invalid_input( + "query_count must be positive for batch KNN".to_string(), + )); + } + if k == 0 { + return Err(Error::invalid_input( + "k must be positive for batch KNN".to_string(), + )); + } + if query.len() % query_count != 0 { + return Err(Error::invalid_input(format!( + "query length ({}) must be divisible by query_count ({})", + query.len(), + query_count + ))); + } + + let mut input_schema = input.schema().as_ref().clone(); + let (_, element_type) = get_vector_type(&(&input_schema).try_into()?, column)?; + validate_distance_type_for(distance_type, &element_type)?; + + if input_schema.column_with_name(DIST_COL).is_some() { + input_schema = input_schema.without_column(DIST_COL); + } + if input_schema.column_with_name(QUERY_INDEX_COL).is_some() { + input_schema = input_schema.without_column(QUERY_INDEX_COL); + } + let input_schema = Arc::new(input_schema); + let output_schema = Arc::new( + input_schema + .as_ref() + .clone() + .try_with_column(Field::new(QUERY_INDEX_COL, DataType::UInt32, true))? + .try_with_column(Field::new(DIST_COL, DataType::Float32, true))?, + ); + + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + + Ok(Self { + input, + query, + query_count, + k, + column: column.to_string(), + distance_type, + input_schema, + output_schema, + properties, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + async fn execute_batch( + input: SendableRecordBatchStream, + input_schema: SchemaRef, + output_schema: SchemaRef, + column: String, + query: ArrayRef, + query_count: usize, + k: usize, + distance_type: DistanceType, + ) -> DataFusionResult { + let query_dim = query.len() / query_count; + let mut heaps = (0..query_count) + .map(|_| BinaryHeap::::with_capacity(k + 1)) + .collect::>(); + let mut input = input; + let mut fallback_row_id = 0_u64; + + while let Some(batch) = input.next().await { + let mut batch = batch?; + if batch.column_by_name(DIST_COL).is_some() { + batch = batch.drop_column(DIST_COL).map_err(|e| { + DataFusionError::ArrowError(Box::new(e), Some("drop _distance".to_string())) + })?; + } + if batch.column_by_name(QUERY_INDEX_COL).is_some() { + batch = batch.drop_column(QUERY_INDEX_COL).map_err(|e| { + DataFusionError::ArrowError(Box::new(e), Some("drop _query_index".to_string())) + })?; + } + if batch.num_rows() == 0 { + continue; + } + + let row_ids = batch + .column_by_name(ROW_ID) + .map(|row_ids| row_ids.as_primitive::().clone()); + let batch = Arc::new(batch); + + for query_index in 0..query_count { + let key = query.slice(query_index * query_dim, query_dim); + let with_distances = + compute_distance(key, distance_type, &column, batch.as_ref().clone()) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let distances = with_distances[DIST_COL].as_primitive::(); + let heap = &mut heaps[query_index]; + for row_index in 0..batch.num_rows() { + if !distances.is_valid(row_index) { + continue; + } + let distance = distances.value(row_index); + if !distance.is_finite() { + continue; + }; + let row_id = row_ids + .as_ref() + .map(|row_ids| row_ids.value(row_index)) + .unwrap_or(fallback_row_id + row_index as u64); + let candidate = BatchKnnCandidate { + query_index: query_index as u32, + distance, + row_id, + batch: batch.clone(), + row_index: row_index as u32, + }; + if heap.len() < k { + heap.push(candidate); + } else if heap + .peek() + .is_some_and(|worst| candidate.cmp(worst).is_lt()) + { + heap.pop(); + heap.push(candidate); + } + } + } + + fallback_row_id += batch.num_rows() as u64; + } + + let mut results = heaps + .into_iter() + .flat_map(BinaryHeap::into_vec) + .collect::>(); + results.sort_by(|left, right| { + left.query_index + .cmp(&right.query_index) + .then_with(|| left.distance.total_cmp(&right.distance)) + .then_with(|| left.row_id.cmp(&right.row_id)) + }); + + if results.is_empty() { + return Ok(RecordBatch::new_empty(output_schema)); + } + + let mut query_indices = UInt32Builder::with_capacity(results.len()); + let mut distances = Float32Builder::with_capacity(results.len()); + let mut row_batches = Vec::with_capacity(results.len()); + for result in results { + query_indices.append_value(result.query_index); + distances.append_value(result.distance); + let indices = UInt32Array::from(vec![result.row_index]); + row_batches.push( + arrow_select::take::take_record_batch(result.batch.as_ref(), &indices).map_err( + |e| { + DataFusionError::ArrowError(Box::new(e), Some("take top-k row".to_string())) + }, + )?, + ); + } + + let output = concat_batches(&input_schema, &row_batches) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + output + .try_with_column( + Field::new(QUERY_INDEX_COL, DataType::UInt32, true), + Arc::new(query_indices.finish()), + ) + .and_then(|batch| { + batch.try_with_column( + Field::new(DIST_COL, DataType::Float32, true), + Arc::new(distances.finish()), + ) + }) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) + } +} + +impl ExecutionPlan for KNNBatchVectorDistanceExec { + fn name(&self) -> &str { + "KNNBatchVectorDistanceExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.output_schema.clone() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::SinglePartition] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return Err(DataFusionError::Internal( + "KNNBatchVectorDistanceExec node must have exactly one child".to_string(), + )); + } + + Ok(Arc::new(Self::try_new( + children.pop().expect("length checked"), + &self.column, + self.query.clone(), + self.query_count, + self.k, + self.distance_type, + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let input = self.input.execute(partition, context)?; + let stream = stream::once(Self::execute_batch( + input, + self.input_schema.clone(), + self.output_schema.clone(), + self.column.clone(), + self.query.clone(), + self.query_count, + self.k, + self.distance_type, + )); + Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new( + self.output_schema.clone(), + stream.boxed(), + partition, + &self.metrics, + ))) + } + + fn partition_statistics(&self, _partition: Option) -> DataFusionResult { + Ok(Statistics { + num_rows: Precision::Inexact(self.query_count * self.k), + ..Statistics::new_unknown(self.schema().as_ref()) + }) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn supports_limit_pushdown(&self) -> bool { + false + } +} + pub static KNN_INDEX_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![ Field::new(DIST_COL, DataType::Float32, true), From 264e0b23a0ddd78d844512825ab1d201bd1ba691 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 18 May 2026 18:00:54 +0800 Subject: [PATCH 02/32] refactor: align batch vector query with nearest API Fold batch flat KNN into the existing nearest and KNN execution paths so the public API and plan nodes stay consistent with reviewer feedback. Co-authored-by: Cursor --- python/src/dataset.rs | 3 +- rust/lance/benches/vector_index.rs | 77 +---- rust/lance/src/dataset/scanner.rs | 200 +++---------- rust/lance/src/io/exec.rs | 4 +- rust/lance/src/io/exec/knn.rs | 460 ++++++++++++----------------- 5 files changed, 246 insertions(+), 498 deletions(-) diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 1a59348fdc6..c242545d636 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1386,8 +1386,7 @@ impl Dataset { let q = q.as_primitive::(); scanner.nearest(&column, q, k) } - (true, _) => scanner.nearest_batch(&column, &q, k), - (false, _) => scanner.nearest(&column, &q, k), + (_, _) => scanner.nearest(&column, &q, k), }; let distance_range: Option<(Option, Option)> = if let Some(dr) = nearest.get_item("distance_range")? { diff --git a/rust/lance/benches/vector_index.rs b/rust/lance/benches/vector_index.rs index 8004b08f9f5..b6379e4bdb4 100644 --- a/rust/lance/benches/vector_index.rs +++ b/rust/lance/benches/vector_index.rs @@ -3,7 +3,6 @@ #![allow(clippy::print_stdout)] use std::sync::Arc; -use std::time::Duration; use arrow_array::{ FixedSizeListArray, Float32Array, RecordBatch, RecordBatchIterator, cast::as_primitive_array, @@ -15,7 +14,7 @@ use futures::TryStreamExt; use pprof::criterion::{Output, PProfProfiler}; use rand::Rng; -use lance::dataset::{Dataset, ReadParams, WriteMode, WriteParams, builder::DatasetBuilder}; +use lance::dataset::{Dataset, WriteMode, WriteParams, builder::DatasetBuilder}; use lance::index::DatasetIndexExt; use lance::index::vector::VectorIndexParams; use lance_arrow::{FixedSizeListArrayExt, as_fixed_size_list_array}; @@ -23,23 +22,7 @@ use lance_index::{ IndexType, vector::{ivf::IvfBuildParams, pq::PQBuildParams}, }; -use lance_io::object_store::{ObjectStoreParams, WrappingObjectStore}; use lance_linalg::distance::MetricType; -use object_store::{ - ObjectStore, - throttle::{ThrottleConfig, ThrottledStore}, -}; - -#[derive(Debug)] -struct ThrottledStoreWrapper { - config: ThrottleConfig, -} - -impl WrappingObjectStore for ThrottledStoreWrapper { - fn wrap(&self, _prefix: &str, original: Arc) -> Arc { - Arc::new(ThrottledStore::new(original, self.config.clone())) - } -} fn bench_ivf_pq_index(c: &mut Criterion) { // default tokio runtime @@ -147,18 +130,15 @@ fn bench_batch_flat_knn(c: &mut Criterion) { const QUERY_COUNT: usize = 8; let rt = tokio::runtime::Runtime::new().unwrap(); - let uri = format!("memory://batch_flat_vec_data_{}", rand::random::()); - let dataset = rt.block_on(async { - create_flat_file( - &uri, - WriteMode::Create, - 50_000, - 5_000, - DIM, - Duration::from_millis(5), - ) - .await - }); + let uri = std::env::temp_dir() + .join(format!( + "batch_flat_vec_data_{}.lance", + rand::random::() + )) + .to_string_lossy() + .to_string(); + let dataset = + rt.block_on(async { create_flat_file(&uri, WriteMode::Create, 50_000, 5_000, DIM).await }); let first_batch = rt.block_on(async { dataset .scan() @@ -213,7 +193,7 @@ fn bench_batch_flat_knn(c: &mut Criterion) { b.to_async(&rt).iter(|| async { let results = dataset .scan() - .nearest_batch("vector", &queries, K) + .nearest("vector", &queries, K) .unwrap() .project::<&str>(&[]) .unwrap() @@ -298,7 +278,6 @@ async fn create_flat_file( num_rows: i32, batch_size: i32, dim: i32, - wait_get_per_call: Duration, ) -> Dataset { let schema = Arc::new(ArrowSchema::new(vec![Field::new( "vector", @@ -325,45 +304,17 @@ async fn create_flat_file( }) .collect(); - if !uri.starts_with("memory://") { - std::fs::remove_dir_all(uri).map_or_else(|_| println!("{} not exists", uri), |_| {}); - } - let store_params = if wait_get_per_call.is_zero() { - None - } else { - Some(ObjectStoreParams { - object_store_wrapper: Some(Arc::new(ThrottledStoreWrapper { - config: ThrottleConfig { - wait_get_per_call, - ..Default::default() - }, - })), - ..Default::default() - }) - }; + std::fs::remove_dir_all(uri).map_or_else(|_| println!("{} not exists", uri), |_| {}); let write_params = WriteParams { max_rows_per_file: num_rows as usize, max_rows_per_group: batch_size as usize, - store_params: store_params.clone(), mode, ..Default::default() }; let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); - let dataset = Dataset::write(reader, uri, Some(write_params)) + Dataset::write(reader, uri, Some(write_params)) .await - .unwrap(); - if uri.starts_with("memory://") { - dataset - } else { - DatasetBuilder::from_uri(uri) - .with_read_params(ReadParams { - store_options: store_params, - ..Default::default() - }) - .load() - .await - .unwrap() - } + .unwrap() } fn create_float32_array(num_elements: i32) -> Float32Array { diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 8cce12f2c45..10d1bffaf2a 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -99,9 +99,8 @@ use crate::io::exec::fts::{ use crate::io::exec::knn::MultivectorScoringExec; use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; use crate::io::exec::{ - AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNBatchVectorDistanceExec, - KNNVectorDistanceExec, LancePushdownScanExec, LanceScanExec, Planner, PreFilterSource, - ScanConfig, TakeExec, + AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNVectorDistanceExec, LancePushdownScanExec, + LanceScanExec, Planner, PreFilterSource, ScanConfig, TakeExec, knn::{KNN_INDEX_SCHEMA, QUERY_INDEX_COL, new_knn_exec}, project, }; @@ -1453,15 +1452,8 @@ impl Scanner { let (vector_type, element_type) = get_vector_type(self.dataset.schema(), column)?; let dim = get_vector_dim(self.dataset.schema(), column)?; - let q = match q.data_type() { + let (q, query_count) = match q.data_type() { DataType::List(_) | DataType::FixedSizeList(_, _) => { - if !matches!(vector_type, DataType::List(_)) { - return Err(Error::invalid_input(format!( - "Query is multivector but column {}({})is not multivector", - column, vector_type, - ))); - } - if let Some(list_array) = q.as_list_opt::() { for i in 0..list_array.len() { let vec = list_array.value(i); @@ -1474,7 +1466,12 @@ impl Scanner { ))); } } - list_array.values().clone() + let query_count = if matches!(vector_type, DataType::List(_)) { + 1 + } else { + list_array.len() + }; + (list_array.values().clone(), query_count) } else { let fsl = q.as_fixed_size_list(); if fsl.value_length() as usize != dim { @@ -1485,130 +1482,34 @@ impl Scanner { dim, ))); } - fsl.values().clone() - } - } - _ => { - if q.len() != dim { - return Err(Error::invalid_input(format!( - "query dim({}) doesn't match the column {} vector dim({})", - q.len(), - column, - dim, - ))); + let query_count = if matches!(vector_type, DataType::List(_)) { + 1 + } else { + fsl.len() + }; + (fsl.values().clone(), query_count) } - q.slice(0, q.len()) } - }; - - let key = match &element_type { - dt if dt == q.data_type() => q, - dt if dt.is_floating() => coerce_float_vector( - q.as_any().downcast_ref::().unwrap(), - FloatType::try_from(dt)?, - )?, _ => { - return Err(Error::invalid_input(format!( - "Column {} has element type {} and the query vector is {}", - column, - element_type, - q.data_type(), - ))); - } - }; - - self.nearest = Some(Query { - column: column.to_string(), - key, - k, - lower_bound: None, - upper_bound: None, - minimum_nprobes: 1, - maximum_nprobes: None, - ef: None, - refine_factor: None, - metric_type: None, - use_index: true, - query_parallelism: DEFAULT_QUERY_PARALLELISM, - dist_q_c: 0.0, - }); - self.nearest_query_count = 1; - Ok(self) - } - - /// Find the k-nearest neighbors for each query vector in a batch. - /// - /// This first implementation uses the flat scan path. It reads candidate - /// vectors once and computes distances for all query vectors, returning up - /// to `k` rows per query with an additional `_query_index` column. - pub fn nearest_batch(&mut self, column: &str, q: &dyn Array, k: usize) -> Result<&mut Self> { - if !self.prefilter { - // We can allow fragment scan if the input to nearest is a prefilter. - // The fragment scan will be performed by the prefilter. - self.ensure_not_fragment_scan()?; - } - - if k == 0 { - return Err(Error::invalid_input("k must be positive".to_string())); - } - if q.is_empty() { - return Err(Error::invalid_input( - "Query vector batch must have at least one query".to_string(), - )); - } - - let (vector_type, element_type) = get_vector_type(self.dataset.schema(), column)?; - if matches!(vector_type, DataType::List(_)) { - return Err(Error::not_supported( - "Batch nearest is not supported for multivector columns".to_string(), - )); - } - let dim = get_vector_dim(self.dataset.schema(), column)?; - - let (q, query_count) = match q.data_type() { - DataType::FixedSizeList(_, _) => { - let fsl = q.as_fixed_size_list(); - if fsl.value_length() as usize != dim { + if q.len() != dim + && (!matches!(vector_type, DataType::FixedSizeList(_, _)) + || !q.len().is_multiple_of(dim)) + { return Err(Error::invalid_input(format!( "query dim({}) doesn't match the column {} vector dim({})", - fsl.value_length(), - column, - dim, - ))); - } - (fsl.values().clone(), fsl.len()) - } - DataType::List(_) => { - let list_array = q.as_list::(); - for i in 0..list_array.len() { - let vec = list_array.value(i); - if vec.len() != dim { - return Err(Error::invalid_input(format!( - "query dim({}) doesn't match the column {} vector dim({})", - vec.len(), - column, - dim, - ))); - } - } - (list_array.values().clone(), list_array.len()) - } - _ => { - if q.len() % dim != 0 { - return Err(Error::invalid_input(format!( - "query batch len({}) must be a multiple of column {} vector dim({})", q.len(), column, dim, ))); } - (q.slice(0, q.len()), q.len() / dim) + let query_count = if q.len() == dim { 1 } else { q.len() / dim }; + (q.slice(0, q.len()), query_count) } }; if query_count == 0 { return Err(Error::invalid_input( - "Query vector batch must have at least one query".to_string(), + "Query vector must have non-zero length".to_string(), )); } @@ -1620,7 +1521,7 @@ impl Scanner { )?, _ => { return Err(Error::invalid_input(format!( - "Column {} has element type {} and the query vector batch is {}", + "Column {} has element type {} and the query vector is {}", column, element_type, q.data_type(), @@ -1639,9 +1540,7 @@ impl Scanner { ef: None, refine_factor: None, metric_type: None, - // Batch KNN is flat-only for now. ANN batching needs per-query - // partition/result grouping instead of the single-query top-k plan. - use_index: false, + use_index: true, query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, }); @@ -1746,12 +1645,7 @@ impl Scanner { /// This is essentially a weak consistency search, only on the indexed data. pub fn fast_search(&mut self) -> &mut Self { if let Some(q) = self.nearest.as_mut() { - if self.nearest_query_count > 1 { - log::warn!("fast_search is ignored for batch nearest queries"); - return self; - } else { - q.use_index = true; - } + q.use_index = true; } self.fast_search = true; self.projection_plan.include_row_id(); // fast search requires _rowid @@ -1811,12 +1705,7 @@ impl Scanner { /// Set whether to use the index if available pub fn use_index(&mut self, use_index: bool) -> &mut Self { if let Some(q) = self.nearest.as_mut() { - if self.nearest_query_count > 1 && use_index { - log::warn!("use_index(true) is ignored for batch nearest queries"); - q.use_index = false; - } else { - q.use_index = use_index; - } + q.use_index = use_index } self } @@ -3647,8 +3536,14 @@ impl Scanner { // Sanity check let (vector_type, element_type) = get_vector_type(self.dataset.schema(), &q.column)?; + if self.nearest_query_count > 1 && self.fast_search { + return Err(Error::not_supported( + "fast_search is not supported for batch nearest queries".to_string(), + )); + } + let column_id = self.dataset.schema().field_id(q.column.as_str())?; - let use_index = q.use_index; + let use_index = q.use_index && self.nearest_query_count == 1; let indices = if use_index { self.dataset.load_indices().await? } else { @@ -4404,22 +4299,17 @@ impl Scanner { default_distance_type_for(&element_type) } }; - if self.nearest_query_count > 1 { - let input = Arc::new(CoalescePartitionsExec::new(input)); - return KNNBatchVectorDistanceExec::try_new( - input, - &q.column, - q.key.clone(), - self.nearest_query_count, - q.k, - metric_type, - ) - .map(|exec| Arc::new(exec) as Arc); - } - let flat_dist = Arc::new(KNNVectorDistanceExec::try_new( + let input = if self.nearest_query_count > 1 { + Arc::new(CoalescePartitionsExec::new(input)) as Arc + } else { + input + }; + let flat_dist = Arc::new(KNNVectorDistanceExec::try_new_batch( input, &q.column, q.key.clone(), + self.nearest_query_count, + q.k, metric_type, )?); @@ -4464,6 +4354,10 @@ impl Scanner { flat_dist }; + if self.nearest_query_count > 1 { + return Ok(knn_plan); + } + // Use DataFusion's [SortExec] for Top-K search let sort = SortExec::new( [ @@ -5800,12 +5694,12 @@ mod test { .unwrap(); let mut scan = dataset.scan(); - scan.nearest_batch("vec", &queries, 2).unwrap(); + scan.nearest("vec", &queries, 2).unwrap(); scan.project(&["i"]).unwrap(); let plan = scan.explain_plan(false).await.unwrap(); assert!( - plan.contains("KNNBatchVectorDistance"), + plan.contains("KNNVectorDistance: queries=2"), "expected flat batch KNN plan, got:\n{}", plan ); diff --git a/rust/lance/src/io/exec.rs b/rust/lance/src/io/exec.rs index 4889b8fcd65..0b93e3c2834 100644 --- a/rust/lance/src/io/exec.rs +++ b/rust/lance/src/io/exec.rs @@ -27,9 +27,7 @@ pub mod testing; pub mod utils; pub use filter::LanceFilterExec; -pub use knn::{ - ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNBatchVectorDistanceExec, KNNVectorDistanceExec, -}; +pub use knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNVectorDistanceExec}; 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 2643e2de42b..f8581e60c82 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -145,23 +145,52 @@ pub struct KNNVectorDistanceExec { /// The vector query to execute. pub query: ArrayRef, + pub query_count: usize, + pub k: usize, pub column: String, pub distance_type: DistanceType, + input_schema: SchemaRef, output_schema: SchemaRef, properties: Arc, metrics: ExecutionPlanMetricsSet, } +struct BatchKnnConfig { + input_schema: SchemaRef, + output_schema: SchemaRef, + column: String, + query: ArrayRef, + query_count: usize, + k: usize, + distance_type: DistanceType, +} + impl DisplayAs for KNNVectorDistanceExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "KNNVectorDistance: metric={}", self.distance_type,) + if self.query_count > 1 { + write!( + f, + "KNNVectorDistance: queries={}, k={}, metric={}", + self.query_count, self.k, self.distance_type, + ) + } else { + write!(f, "KNNVectorDistance: metric={}", self.distance_type,) + } } DisplayFormatType::TreeRender => { - write!(f, "KNNVectorDistance\nmetric={}", self.distance_type,) + if self.query_count > 1 { + write!( + f, + "KNNVectorDistance\nqueries={}\nk={}\nmetric={}", + self.query_count, self.k, self.distance_type, + ) + } else { + write!(f, "KNNVectorDistance\nmetric={}", self.distance_type,) + } } } } @@ -177,181 +206,9 @@ impl KNNVectorDistanceExec { query: ArrayRef, distance_type: DistanceType, ) -> Result { - let mut output_schema = input.schema().as_ref().clone(); - let (_, element_type) = get_vector_type(&(&output_schema).try_into()?, column)?; - validate_distance_type_for(distance_type, &element_type)?; - - // FlatExec appends a distance column to the input schema. The input - // may already have a distance column (possibly in the wrong position), so - // we need to remove it before adding a new one. - if output_schema.column_with_name(DIST_COL).is_some() { - output_schema = output_schema.without_column(DIST_COL); - } - let output_schema = Arc::new(output_schema.try_with_column(Field::new( - DIST_COL, - DataType::Float32, - true, - ))?); - - // This node has the same partitioning & boundedness as the input node - // but it destroys any ordering. - let properties = Arc::new( - input - .properties() - .as_ref() - .clone() - .with_eq_properties(EquivalenceProperties::new(output_schema.clone())), - ); - - Ok(Self { - input, - query, - column: column.to_string(), - distance_type, - output_schema, - properties, - metrics: ExecutionPlanMetricsSet::new(), - }) - } -} - -impl ExecutionPlan for KNNVectorDistanceExec { - fn name(&self) -> &str { - "KNNVectorDistanceExec" - } - - fn as_any(&self) -> &dyn Any { - self - } - - /// Flat KNN inherits the schema from input node, and add one distance column. - fn schema(&self) -> arrow_schema::SchemaRef { - self.output_schema.clone() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn with_new_children( - self: Arc, - mut children: Vec>, - ) -> DataFusionResult> { - if children.len() != 1 { - return Err(DataFusionError::Internal( - "KNNVectorDistanceExec node must have exactly one child".to_string(), - )); - } - - Ok(Arc::new(Self::try_new( - children.pop().expect("length checked"), - &self.column, - self.query.clone(), - self.distance_type, - )?)) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> DataFusionResult { - let input_stream = self.input.execute(partition, context)?; - let input_schema = input_stream.schema(); - let key = self.query.clone(); - let column = self.column.clone(); - let dt = self.distance_type; - let schema = self.schema(); - - // Empty batches don't have a vector column to score; filter them out - // before reaching the helper so the transform always sees real work. - let filtered_input = Box::pin(RecordBatchStreamAdapter::new( - input_schema, - input_stream.try_filter(|batch| future::ready(batch.num_rows() > 0)), - )) as SendableRecordBatchStream; - - // Mirror of the helper's elapsed_compute counter; used to attribute - // wall-clock from the spawn_blocking distance kernel back onto the - // node's `elapsed_compute` metric. - let elapsed_compute = BaselineMetrics::new(&self.metrics, partition) - .elapsed_compute() - .clone(); - - let stream = InstrumentedChildInputStream::new( - filtered_input, - schema, - move |batch| { - let key = key.clone(); - let column = column.clone(); - let elapsed_compute = elapsed_compute.clone(); - async move { - // Time around the .await to capture the spawn_blocking - // distance work, which otherwise runs while this future is - // Pending and is missed by the helper's own poll timer. - let start = std::time::Instant::now(); - let batch = compute_distance(key, dt, &column, batch) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; - elapsed_compute.add_duration(start.elapsed()); - - let distances = batch[DIST_COL].as_primitive::(); - let mask = BooleanArray::from_iter( - distances - .iter() - .map(|v| Some(v.map(|v| !v.is_nan()).unwrap_or(false))), - ); - arrow::compute::filter_record_batch(&batch, &mask) - .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) - } - }, - get_num_compute_intensive_cpus(), - partition, - &self.metrics, - ); - Ok(Box::pin(stream) as SendableRecordBatchStream) - } - - fn partition_statistics(&self, partition: Option) -> DataFusionResult { - let inner_stats = self.input.partition_statistics(partition)?; - let schema = self.input.schema(); - let dist_stats = inner_stats - .column_statistics - .iter() - .zip(schema.fields()) - .find(|(_, field)| field.name() == &self.column) - .map(|(stats, _)| ColumnStatistics { - null_count: stats.null_count, - ..Default::default() - }) - .unwrap_or_default(); - let column_statistics = inner_stats - .column_statistics - .into_iter() - .zip(schema.fields()) - .filter(|(_, field)| field.name() != DIST_COL) - .map(|(stats, _)| stats) - .chain(std::iter::once(dist_stats)) - .collect::>(); - Ok(Statistics { - num_rows: inner_stats.num_rows, - column_statistics, - ..Statistics::new_unknown(self.schema().as_ref()) - }) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn properties(&self) -> &Arc { - &self.properties + Self::try_new_batch(input, column, query, 1, 0, distance_type) } - fn supports_limit_pushdown(&self) -> bool { - false - } -} - #[derive(Clone)] struct BatchKnnCandidate { query_index: u32, @@ -388,44 +245,7 @@ impl Ord for BatchKnnCandidate { } } -/// [ExecutionPlan] that computes flat KNN for a batch of query vectors. -/// -/// This node consumes the input once, computes distances from each input vector -/// to each query vector, and keeps the top-k rows independently for every query. -#[derive(Debug)] -pub struct KNNBatchVectorDistanceExec { - pub input: Arc, - pub query: ArrayRef, - pub query_count: usize, - pub k: usize, - pub column: String, - pub distance_type: DistanceType, - - input_schema: SchemaRef, - output_schema: SchemaRef, - properties: Arc, - metrics: ExecutionPlanMetricsSet, -} - -impl DisplayAs for KNNBatchVectorDistanceExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => write!( - f, - "KNNBatchVectorDistance: queries={}, k={}, metric={}", - self.query_count, self.k, self.distance_type - ), - DisplayFormatType::TreeRender => write!( - f, - "KNNBatchVectorDistance\nqueries={}\nk={}\nmetric={}", - self.query_count, self.k, self.distance_type - ), - } - } -} - -impl KNNBatchVectorDistanceExec { - pub fn try_new( + pub fn try_new_batch( input: Arc, column: &str, query: ArrayRef, @@ -435,26 +255,29 @@ impl KNNBatchVectorDistanceExec { ) -> Result { if query_count == 0 { return Err(Error::invalid_input( - "query_count must be positive for batch KNN".to_string(), + "query_count must be positive for KNN".to_string(), )); } - if k == 0 { - return Err(Error::invalid_input( - "k must be positive for batch KNN".to_string(), - )); - } - if query.len() % query_count != 0 { + if !query.len().is_multiple_of(query_count) { return Err(Error::invalid_input(format!( "query length ({}) must be divisible by query_count ({})", query.len(), query_count ))); } + if query_count > 1 && k == 0 { + return Err(Error::invalid_input( + "k must be positive for batch KNN".to_string(), + )); + } let mut input_schema = input.schema().as_ref().clone(); let (_, element_type) = get_vector_type(&(&input_schema).try_into()?, column)?; validate_distance_type_for(distance_type, &element_type)?; + // FlatExec appends a distance column to the input schema. The input + // may already have a distance column (possibly in the wrong position), so + // we need to remove it before adding a new one. if input_schema.column_with_name(DIST_COL).is_some() { input_schema = input_schema.without_column(DIST_COL); } @@ -462,20 +285,38 @@ impl KNNBatchVectorDistanceExec { input_schema = input_schema.without_column(QUERY_INDEX_COL); } let input_schema = Arc::new(input_schema); - let output_schema = Arc::new( - input_schema - .as_ref() - .clone() - .try_with_column(Field::new(QUERY_INDEX_COL, DataType::UInt32, true))? - .try_with_column(Field::new(DIST_COL, DataType::Float32, true))?, - ); + let mut output_schema = input_schema.as_ref().clone(); + if query_count > 1 { + output_schema = output_schema.try_with_column(Field::new( + QUERY_INDEX_COL, + DataType::UInt32, + true, + ))?; + } + let output_schema = Arc::new(output_schema.try_with_column(Field::new( + DIST_COL, + DataType::Float32, + true, + ))?); - let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(output_schema.clone()), - Partitioning::UnknownPartitioning(1), - EmissionType::Final, - Boundedness::Bounded, - )); + // This node has the same partitioning & boundedness as the input node + // but it destroys any ordering. + let properties = if query_count > 1 { + Arc::new(PlanProperties::new( + EquivalenceProperties::new(output_schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )) + } else { + Arc::new( + input + .properties() + .as_ref() + .clone() + .with_eq_properties(EquivalenceProperties::new(output_schema.clone())), + ) + }; Ok(Self { input, @@ -493,14 +334,17 @@ impl KNNBatchVectorDistanceExec { async fn execute_batch( input: SendableRecordBatchStream, - input_schema: SchemaRef, - output_schema: SchemaRef, - column: String, - query: ArrayRef, - query_count: usize, - k: usize, - distance_type: DistanceType, + config: BatchKnnConfig, ) -> DataFusionResult { + let BatchKnnConfig { + input_schema, + output_schema, + column, + query, + query_count, + k, + distance_type, + } = config; let query_dim = query.len() / query_count; let mut heaps = (0..query_count) .map(|_| BinaryHeap::::with_capacity(k + 1)) @@ -509,17 +353,7 @@ impl KNNBatchVectorDistanceExec { let mut fallback_row_id = 0_u64; while let Some(batch) = input.next().await { - let mut batch = batch?; - if batch.column_by_name(DIST_COL).is_some() { - batch = batch.drop_column(DIST_COL).map_err(|e| { - DataFusionError::ArrowError(Box::new(e), Some("drop _distance".to_string())) - })?; - } - if batch.column_by_name(QUERY_INDEX_COL).is_some() { - batch = batch.drop_column(QUERY_INDEX_COL).map_err(|e| { - DataFusionError::ArrowError(Box::new(e), Some("drop _query_index".to_string())) - })?; - } + let batch = batch?; if batch.num_rows() == 0 { continue; } @@ -529,14 +363,13 @@ impl KNNBatchVectorDistanceExec { .map(|row_ids| row_ids.as_primitive::().clone()); let batch = Arc::new(batch); - for query_index in 0..query_count { + for (query_index, heap) in heaps.iter_mut().enumerate().take(query_count) { let key = query.slice(query_index * query_dim, query_dim); let with_distances = compute_distance(key, distance_type, &column, batch.as_ref().clone()) .await .map_err(|e| DataFusionError::External(Box::new(e)))?; let distances = with_distances[DIST_COL].as_primitive::(); - let heap = &mut heaps[query_index]; for row_index in 0..batch.num_rows() { if !distances.is_valid(row_index) { continue; @@ -619,16 +452,17 @@ impl KNNBatchVectorDistanceExec { } } -impl ExecutionPlan for KNNBatchVectorDistanceExec { +impl ExecutionPlan for KNNVectorDistanceExec { fn name(&self) -> &str { - "KNNBatchVectorDistanceExec" + "KNNVectorDistanceExec" } fn as_any(&self) -> &dyn Any { self } - fn schema(&self) -> SchemaRef { + /// Flat KNN inherits the schema from input node, and add one distance column. + fn schema(&self) -> arrow_schema::SchemaRef { self.output_schema.clone() } @@ -636,21 +470,17 @@ impl ExecutionPlan for KNNBatchVectorDistanceExec { vec![&self.input] } - fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] - } - fn with_new_children( self: Arc, mut children: Vec>, ) -> DataFusionResult> { if children.len() != 1 { return Err(DataFusionError::Internal( - "KNNBatchVectorDistanceExec node must have exactly one child".to_string(), + "KNNVectorDistanceExec node must have exactly one child".to_string(), )); } - Ok(Arc::new(Self::try_new( + Ok(Arc::new(Self::try_new_batch( children.pop().expect("length checked"), &self.column, self.query.clone(), @@ -665,28 +495,96 @@ impl ExecutionPlan for KNNBatchVectorDistanceExec { partition: usize, context: Arc, ) -> DataFusionResult { - let input = self.input.execute(partition, context)?; - let stream = stream::once(Self::execute_batch( - input, - self.input_schema.clone(), - self.output_schema.clone(), - self.column.clone(), - self.query.clone(), - self.query_count, - self.k, - self.distance_type, - )); + let input_stream = self.input.execute(partition, context)?; + if self.query_count > 1 { + let stream = stream::once(Self::execute_batch( + input_stream, + BatchKnnConfig { + input_schema: self.input_schema.clone(), + output_schema: self.output_schema.clone(), + column: self.column.clone(), + query: self.query.clone(), + query_count: self.query_count, + k: self.k, + distance_type: self.distance_type, + }, + )); + let schema = self.schema(); + return Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new( + schema, + stream.boxed(), + partition, + &self.metrics, + )) as SendableRecordBatchStream); + } + let key = self.query.clone(); + let column = self.column.clone(); + let dt = self.distance_type; + let stream = input_stream + .try_filter(|batch| future::ready(batch.num_rows() > 0)) + .map(move |batch| { + let key = key.clone(); + let column = column.clone(); + async move { + let batch = compute_distance(key, dt, &column, batch?) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + let distances = batch[DIST_COL].as_primitive::(); + let mask = BooleanArray::from_iter( + distances + .iter() + .map(|v| Some(v.map(|v| !v.is_nan()).unwrap_or(false))), + ); + arrow::compute::filter_record_batch(&batch, &mask) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) + } + }) + .buffer_unordered(get_num_compute_intensive_cpus()); + let schema = self.schema(); Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new( - self.output_schema.clone(), + schema, stream.boxed(), partition, &self.metrics, - ))) + )) as SendableRecordBatchStream) } - fn partition_statistics(&self, _partition: Option) -> DataFusionResult { + fn partition_statistics(&self, partition: Option) -> DataFusionResult { + let inner_stats = self.input.partition_statistics(partition)?; + let schema = self.input.schema(); + let dist_stats = inner_stats + .column_statistics + .iter() + .zip(schema.fields()) + .find(|(_, field)| field.name() == &self.column) + .map(|(stats, _)| ColumnStatistics { + null_count: stats.null_count, + ..Default::default() + }) + .unwrap_or_default(); + let column_statistics = inner_stats + .column_statistics + .into_iter() + .zip(schema.fields()) + .filter(|(_, field)| field.name() != DIST_COL) + .map(|(stats, _)| stats) + .collect::>(); + let column_statistics = if self.query_count > 1 { + column_statistics + .into_iter() + .chain(std::iter::once(ColumnStatistics::default())) + .chain(std::iter::once(dist_stats)) + .collect::>() + } else { + column_statistics + .into_iter() + .chain(std::iter::once(dist_stats)) + .collect::>() + }; Ok(Statistics { - num_rows: Precision::Inexact(self.query_count * self.k), + num_rows: inner_stats.num_rows, + column_statistics, ..Statistics::new_unknown(self.schema().as_ref()) }) } @@ -702,6 +600,14 @@ impl ExecutionPlan for KNNBatchVectorDistanceExec { fn supports_limit_pushdown(&self) -> bool { false } + + fn required_input_distribution(&self) -> Vec { + if self.query_count > 1 { + vec![Distribution::SinglePartition] + } else { + vec![Distribution::UnspecifiedDistribution] + } + } } pub static KNN_INDEX_SCHEMA: LazyLock = LazyLock::new(|| { From 10a56874dcc7acd4a4d9ba4d5b8b8146e898d871 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 18 May 2026 18:16:19 +0800 Subject: [PATCH 03/32] bench: scale batch flat vector query benchmark Use a larger local-disk dataset and stream benchmark data generation so batch query gains are measured under a more realistic scan workload. Co-authored-by: Cursor --- rust/lance/benches/vector_index.rs | 80 ++++++++++++++++++++---------- 1 file changed, 53 insertions(+), 27 deletions(-) diff --git a/rust/lance/benches/vector_index.rs b/rust/lance/benches/vector_index.rs index b6379e4bdb4..09fafa776d9 100644 --- a/rust/lance/benches/vector_index.rs +++ b/rust/lance/benches/vector_index.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use arrow_array::{ FixedSizeListArray, Float32Array, RecordBatch, RecordBatchIterator, cast::as_primitive_array, }; -use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema}; +use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema as ArrowSchema}; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use futures::TryStreamExt; #[cfg(target_os = "linux")] @@ -125,9 +125,11 @@ fn bench_ivf_pq_index(c: &mut Criterion) { } fn bench_batch_flat_knn(c: &mut Criterion) { - const DIM: i32 = 4; + const DIM: i32 = 512; const K: usize = 10; - const QUERY_COUNT: usize = 8; + const NUM_ROWS: usize = 1_000_000; + const BATCH_SIZE: usize = 10_000; + const QUERY_COUNT: usize = 10; let rt = tokio::runtime::Runtime::new().unwrap(); let uri = std::env::temp_dir() @@ -137,8 +139,9 @@ fn bench_batch_flat_knn(c: &mut Criterion) { )) .to_string_lossy() .to_string(); - let dataset = - rt.block_on(async { create_flat_file(&uri, WriteMode::Create, 50_000, 5_000, DIM).await }); + let dataset = rt.block_on(async { + create_flat_file(&uri, WriteMode::Create, NUM_ROWS, BATCH_SIZE, DIM).await + }); let first_batch = rt.block_on(async { dataset .scan() @@ -275,8 +278,8 @@ async fn create_file(path: &std::path::Path, mode: WriteMode) { async fn create_flat_file( uri: &str, mode: WriteMode, - num_rows: i32, - batch_size: i32, + num_rows: usize, + batch_size: usize, dim: i32, ) -> Dataset { let schema = Arc::new(ArrowSchema::new(vec![Field::new( @@ -288,39 +291,62 @@ async fn create_flat_file( false, )])); - let batches: Vec = (0..(num_rows / batch_size)) - .map(|_| { - RecordBatch::try_new( - schema.clone(), - vec![Arc::new( - FixedSizeListArray::try_new_from_values( - create_float32_array(batch_size * dim), - dim, - ) - .unwrap(), - )], + struct FlatVectorBatchIter { + schema: Arc, + remaining_rows: usize, + batch_size: usize, + dim: i32, + } + + impl Iterator for FlatVectorBatchIter { + type Item = Result; + + fn next(&mut self) -> Option { + if self.remaining_rows == 0 { + return None; + } + let rows = self.remaining_rows.min(self.batch_size); + self.remaining_rows -= rows; + + let values = create_float32_array(rows * self.dim as usize); + Some( + RecordBatch::try_new( + self.schema.clone(), + vec![Arc::new( + FixedSizeListArray::try_new_from_values(values, self.dim).unwrap(), + )], + ) + .map_err(Into::into), ) - .unwrap() - }) - .collect(); + } + } std::fs::remove_dir_all(uri).map_or_else(|_| println!("{} not exists", uri), |_| {}); let write_params = WriteParams { - max_rows_per_file: num_rows as usize, - max_rows_per_group: batch_size as usize, + max_rows_per_file: num_rows, + max_rows_per_group: batch_size, mode, ..Default::default() }; - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); + let reader = RecordBatchIterator::new( + FlatVectorBatchIter { + schema: schema.clone(), + remaining_rows: num_rows, + batch_size, + dim, + }, + schema, + ); Dataset::write(reader, uri, Some(write_params)) .await .unwrap() } -fn create_float32_array(num_elements: i32) -> Float32Array { - // generate an Arrow Float32Array with 10000*128 elements randomly +fn create_float32_array(num_elements: usize) -> Float32Array { + // Generate random values on demand so large benchmark datasets do not need + // to be fully materialized in memory before writing. let mut rng = rand::rng(); - let mut values = Vec::with_capacity(num_elements as usize); + let mut values = Vec::with_capacity(num_elements); for _ in 0..num_elements { values.push(rng.random_range(0.0..1.0)); } From beddd3d2c30983a6cb254e96cca5c578da7efb02 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 18 May 2026 18:31:20 +0800 Subject: [PATCH 04/32] bench: parameterize batch vector query benchmark Allow the local-disk batch KNN benchmark to vary row count, dimensionality, and query count so PR results can show scaling trends. Co-authored-by: Cursor --- rust/lance/benches/vector_index.rs | 120 ++++++++++++++++++----------- 1 file changed, 73 insertions(+), 47 deletions(-) diff --git a/rust/lance/benches/vector_index.rs b/rust/lance/benches/vector_index.rs index 09fafa776d9..52767b0ca9b 100644 --- a/rust/lance/benches/vector_index.rs +++ b/rust/lance/benches/vector_index.rs @@ -125,11 +125,16 @@ fn bench_ivf_pq_index(c: &mut Criterion) { } fn bench_batch_flat_knn(c: &mut Criterion) { - const DIM: i32 = 512; + const DEFAULT_DIM: i32 = 512; const K: usize = 10; - const NUM_ROWS: usize = 1_000_000; - const BATCH_SIZE: usize = 10_000; - const QUERY_COUNT: usize = 10; + const DEFAULT_NUM_ROWS: usize = 1_000_000; + const DEFAULT_BATCH_SIZE: usize = 10_000; + const DEFAULT_QUERY_COUNT: usize = 10; + + let dim = bench_env("LANCE_BATCH_KNN_DIM", DEFAULT_DIM); + let num_rows = bench_env("LANCE_BATCH_KNN_ROWS", DEFAULT_NUM_ROWS); + let batch_size = bench_env("LANCE_BATCH_KNN_BATCH_SIZE", DEFAULT_BATCH_SIZE); + let query_count = bench_env("LANCE_BATCH_KNN_QUERY_COUNT", DEFAULT_QUERY_COUNT); let rt = tokio::runtime::Runtime::new().unwrap(); let uri = std::env::temp_dir() @@ -140,7 +145,7 @@ fn bench_batch_flat_knn(c: &mut Criterion) { .to_string_lossy() .to_string(); let dataset = rt.block_on(async { - create_flat_file(&uri, WriteMode::Create, NUM_ROWS, BATCH_SIZE, DIM).await + create_flat_file(&uri, WriteMode::Create, num_rows, batch_size, dim).await }); let first_batch = rt.block_on(async { dataset @@ -155,7 +160,7 @@ fn bench_batch_flat_knn(c: &mut Criterion) { }); let vector_column = first_batch.column_by_name("vector").unwrap(); let vectors = as_fixed_size_list_array(vector_column); - let query_values = (0..QUERY_COUNT) + let query_values = (0..query_count) .flat_map(|query_index| { let values = vectors.value(query_index); as_primitive_array::(&values) @@ -164,22 +169,51 @@ fn bench_batch_flat_knn(c: &mut Criterion) { }) .collect::>(); let queries = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), DIM) + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), dim) .unwrap(); let mut group = c.benchmark_group("batch_flat_knn"); - group.bench_function(BenchmarkId::new("separate_queries", QUERY_COUNT), |b| { - b.to_async(&rt).iter(|| async { - for query_index in 0..QUERY_COUNT { - let query = Float32Array::from( - query_values[query_index * DIM as usize..(query_index + 1) * DIM as usize] - .to_vec(), - ); + group.bench_function( + BenchmarkId::new( + "separate_queries", + format!("rows={num_rows},dim={dim},queries={query_count}"), + ), + |b| { + b.to_async(&rt).iter(|| async { + for query_index in 0..query_count { + let query = Float32Array::from( + query_values[query_index * dim as usize..(query_index + 1) * dim as usize] + .to_vec(), + ); + let results = dataset + .scan() + .nearest("vector", &query, K) + .unwrap() + .use_index(false) + .project::<&str>(&[]) + .unwrap() + .try_into_stream() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(!results.is_empty()); + } + }) + }, + ); + group.bench_function( + BenchmarkId::new( + "batch_query", + format!("rows={num_rows},dim={dim},queries={query_count}"), + ), + |b| { + b.to_async(&rt).iter(|| async { let results = dataset .scan() - .nearest("vector", &query, K) + .nearest("vector", &queries, K) .unwrap() - .use_index(false) .project::<&str>(&[]) .unwrap() .try_into_stream() @@ -189,27 +223,22 @@ fn bench_batch_flat_knn(c: &mut Criterion) { .await .unwrap(); assert!(!results.is_empty()); - } - }) - }); - group.bench_function(BenchmarkId::new("batch_query", QUERY_COUNT), |b| { - b.to_async(&rt).iter(|| async { - let results = dataset - .scan() - .nearest("vector", &queries, K) - .unwrap() - .project::<&str>(&[]) - .unwrap() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert!(!results.is_empty()); - }) - }); + }) + }, + ); group.finish(); + drop(dataset); + let _ = std::fs::remove_dir_all(uri); +} + +fn bench_env(key: &str, default: T) -> T +where + T: std::str::FromStr, +{ + std::env::var(key) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(default) } async fn create_file(path: &std::path::Path, mode: WriteMode) { @@ -243,8 +272,8 @@ async fn create_file(path: &std::path::Path, mode: WriteMode) { let test_uri = path.to_str().unwrap(); std::fs::remove_dir_all(test_uri).map_or_else(|_| println!("{} not exists", test_uri), |_| {}); let write_params = WriteParams { - max_rows_per_file: num_rows as usize, - max_rows_per_group: batch_size as usize, + max_rows_per_file: num_rows, + max_rows_per_group: batch_size, mode, ..Default::default() }; @@ -309,15 +338,12 @@ async fn create_flat_file( self.remaining_rows -= rows; let values = create_float32_array(rows * self.dim as usize); - Some( - RecordBatch::try_new( - self.schema.clone(), - vec![Arc::new( - FixedSizeListArray::try_new_from_values(values, self.dim).unwrap(), - )], - ) - .map_err(Into::into), - ) + Some(RecordBatch::try_new( + self.schema.clone(), + vec![Arc::new( + FixedSizeListArray::try_new_from_values(values, self.dim).unwrap(), + )], + )) } } From e1a66c3ed87866a545283182f6e2aeea7b5b80d8 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 18 May 2026 19:08:48 +0800 Subject: [PATCH 05/32] fix: align batch query review feedback Use the LanceDB-compatible query_index result column and move the batch flat KNN benchmark to Python so benchmark scaling can be reproduced from the binding API. Co-authored-by: Cursor --- python/python/benchmarks/test_search.py | 100 ++++++++++++ python/python/lance/dataset.py | 6 +- python/python/tests/test_vector_index.py | 6 +- rust/lance/benches/vector_index.rs | 200 +---------------------- rust/lance/src/dataset/scanner.rs | 6 + rust/lance/src/io/exec/knn.rs | 21 ++- 6 files changed, 131 insertions(+), 208 deletions(-) diff --git a/python/python/benchmarks/test_search.py b/python/python/benchmarks/test_search.py index 61076e61687..783319160e1 100644 --- a/python/python/benchmarks/test_search.py +++ b/python/python/benchmarks/test_search.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors +import os import shutil from pathlib import Path from typing import NamedTuple, Union @@ -13,6 +14,11 @@ N_DIMS = 768 NUM_ROWS = 100_000 NEW_ROWS = 10_000 +BATCH_FLAT_KNN_DIM = 512 +BATCH_FLAT_KNN_K = 10 +BATCH_FLAT_KNN_BATCH_SIZE = 10_000 +BATCH_FLAT_KNN_QUERY_COUNT = 10 +BATCH_FLAT_KNN_ROWS = 1_000_000 def find_or_clean(dataset_path: Path) -> Union[lance.LanceDataset, None]: @@ -64,6 +70,49 @@ def create_table(num_rows, offset) -> pa.Table: ) +def get_benchmark_env(key: str, default: int) -> int: + value = os.environ.get(key) + if value is None: + return default + return int(value) + + +def create_flat_vector_table(num_rows: int, offset: int, dim: int) -> pa.Table: + rng = np.random.default_rng(seed=offset) + values = rng.random((num_rows, dim), dtype=np.float32) + vectors = pa.FixedSizeListArray.from_arrays(pa.array(values.ravel()), dim) + ids = pa.array(range(offset, offset + num_rows)) + return pa.table({"vector": vectors, "id": ids}) + + +def create_batch_flat_knn_dataset( + data_dir: Path, num_rows: int, batch_size: int, dim: int +) -> lance.LanceDataset: + tmp_path = data_dir / f"batch_flat_knn_{num_rows}_{batch_size}_{dim}" + dataset = find_or_clean(tmp_path) + if dataset: + return dataset + + rows_remaining = num_rows + offset = 0 + dataset = None + while rows_remaining > 0: + next_batch_length = min(rows_remaining, batch_size) + rows_remaining -= next_batch_length + table = create_flat_vector_table(next_batch_length, offset, dim) + if offset == 0: + dataset = lance.write_dataset( + table, tmp_path, data_storage_version="stable" + ) + else: + dataset = lance.write_dataset( + table, tmp_path, mode="append", data_storage_version="stable" + ) + offset += next_batch_length + + return dataset + + def create_base_dataset(data_dir: Path) -> lance.LanceDataset: tmp_path = data_dir / "search_dataset" dataset = find_or_clean(tmp_path) @@ -173,6 +222,57 @@ def test_knn_search(test_dataset, benchmark): assert result.num_rows > 0 +@pytest.mark.benchmark(group="batch_flat_knn") +@pytest.mark.parametrize("mode", ["separate", "batch"]) +def test_batch_flat_knn(data_dir: Path, benchmark, mode: str): + dim = get_benchmark_env("LANCE_BATCH_KNN_DIM", BATCH_FLAT_KNN_DIM) + num_rows = get_benchmark_env("LANCE_BATCH_KNN_ROWS", BATCH_FLAT_KNN_ROWS) + batch_size = get_benchmark_env( + "LANCE_BATCH_KNN_BATCH_SIZE", BATCH_FLAT_KNN_BATCH_SIZE + ) + query_count = get_benchmark_env( + "LANCE_BATCH_KNN_QUERY_COUNT", BATCH_FLAT_KNN_QUERY_COUNT + ) + rounds = get_benchmark_env("LANCE_BATCH_KNN_ROUNDS", 10) + dataset = create_batch_flat_knn_dataset(data_dir, num_rows, batch_size, dim) + query_table = dataset.to_table(columns=["vector"], limit=query_count) + query_values = np.asarray( + query_table["vector"].combine_chunks().values, dtype=np.float32 + ).reshape(query_count, dim) + + def separate_queries(): + total_rows = 0 + for query in query_values: + total_rows += dataset.to_table( + columns=[], + nearest={ + "column": "vector", + "q": query, + "k": BATCH_FLAT_KNN_K, + "use_index": False, + }, + ).num_rows + return total_rows + + def batch_query(): + return dataset.to_table( + columns=[], + nearest={ + "column": "vector", + "q": query_values, + "k": BATCH_FLAT_KNN_K, + "use_index": False, + }, + ).num_rows + + if mode == "separate": + result = benchmark.pedantic(separate_queries, rounds=rounds, iterations=1) + else: + result = benchmark.pedantic(batch_query, rounds=rounds, iterations=1) + + assert result == query_count * BATCH_FLAT_KNN_K + + @pytest.mark.benchmark(group="query_ann") def test_ann_no_refine(test_dataset, benchmark): q = pc.random(N_DIMS).cast(pa.float32()) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 936f5812113..788a58ae24b 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -1101,7 +1101,7 @@ def scanner( ``q`` may also be a 2-D array-like value for fixed-size vector columns. In that case Lance runs a flat batch KNN query, returns up to ``k`` rows - for each query vector, and adds ``_query_index`` to identify the source + for each query vector, and adds ``query_index`` to identify the source query for each result row. Indexed/ANN batch search is not used in this first implementation. @@ -5998,7 +5998,7 @@ def nearest( q: QueryVectorLike A single query vector or, for fixed-size vector columns, a 2-D array-like batch of query vectors. Batch queries return up to ``k`` rows per query - and include ``_query_index`` in the output. + and include ``query_index`` in the output. query_parallelism: int, optional Maximum partition-search concurrency for a single vector query. The default is 0. Value 0 uses the automatic policy, which @@ -7149,7 +7149,7 @@ def _build_vector_search_query( q: QueryVectorLike The query vector. For fixed-size vector columns, this may be a 2-D array-like batch of query vectors. Batch queries run flat KNN, apply - ``k`` per query vector, and add ``_query_index`` to the result so + ``k`` per query vector, and add ``query_index`` to the result so callers can split rows by input query. k: int, optional The number of nearest neighbors to return. diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index bc448cf8e29..61848c9667b 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -194,8 +194,8 @@ def test_batch_flat_query_matches_repeated_single_queries(dataset): ) assert batch.num_rows == queries.shape[0] * k - assert batch.column_names == ["id", "_distance", "_query_index"] - assert batch["_query_index"].to_pylist() == [0] * k + [1] * k + assert batch.column_names == ["id", "_distance", "query_index"] + assert batch["query_index"].to_pylist() == [0] * k + [1] * k for query_index, query in enumerate(queries): single = dataset.to_table( @@ -207,7 +207,7 @@ def test_batch_flat_query_matches_repeated_single_queries(dataset): "use_index": False, }, ) - batch_slice = batch.filter(pc.field("_query_index") == query_index) + batch_slice = batch.filter(pc.field("query_index") == query_index) assert batch_slice["id"].to_pylist() == single["id"].to_pylist() np.testing.assert_allclose( batch_slice["_distance"].to_numpy(), diff --git a/rust/lance/benches/vector_index.rs b/rust/lance/benches/vector_index.rs index 52767b0ca9b..21c9aa4e4aa 100644 --- a/rust/lance/benches/vector_index.rs +++ b/rust/lance/benches/vector_index.rs @@ -7,8 +7,8 @@ use std::sync::Arc; use arrow_array::{ FixedSizeListArray, Float32Array, RecordBatch, RecordBatchIterator, cast::as_primitive_array, }; -use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema as ArrowSchema}; -use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use arrow_schema::{DataType, Field, FieldRef, Schema as ArrowSchema}; +use criterion::{Criterion, criterion_group, criterion_main}; use futures::TryStreamExt; #[cfg(target_os = "linux")] use pprof::criterion::{Output, PProfProfiler}; @@ -124,123 +124,6 @@ fn bench_ivf_pq_index(c: &mut Criterion) { ); } -fn bench_batch_flat_knn(c: &mut Criterion) { - const DEFAULT_DIM: i32 = 512; - const K: usize = 10; - const DEFAULT_NUM_ROWS: usize = 1_000_000; - const DEFAULT_BATCH_SIZE: usize = 10_000; - const DEFAULT_QUERY_COUNT: usize = 10; - - let dim = bench_env("LANCE_BATCH_KNN_DIM", DEFAULT_DIM); - let num_rows = bench_env("LANCE_BATCH_KNN_ROWS", DEFAULT_NUM_ROWS); - let batch_size = bench_env("LANCE_BATCH_KNN_BATCH_SIZE", DEFAULT_BATCH_SIZE); - let query_count = bench_env("LANCE_BATCH_KNN_QUERY_COUNT", DEFAULT_QUERY_COUNT); - - let rt = tokio::runtime::Runtime::new().unwrap(); - let uri = std::env::temp_dir() - .join(format!( - "batch_flat_vec_data_{}.lance", - rand::random::() - )) - .to_string_lossy() - .to_string(); - let dataset = rt.block_on(async { - create_flat_file(&uri, WriteMode::Create, num_rows, batch_size, dim).await - }); - let first_batch = rt.block_on(async { - dataset - .scan() - .try_into_stream() - .await - .unwrap() - .try_next() - .await - .unwrap() - .unwrap() - }); - let vector_column = first_batch.column_by_name("vector").unwrap(); - let vectors = as_fixed_size_list_array(vector_column); - let query_values = (0..query_count) - .flat_map(|query_index| { - let values = vectors.value(query_index); - as_primitive_array::(&values) - .values() - .to_vec() - }) - .collect::>(); - let queries = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), dim) - .unwrap(); - - let mut group = c.benchmark_group("batch_flat_knn"); - group.bench_function( - BenchmarkId::new( - "separate_queries", - format!("rows={num_rows},dim={dim},queries={query_count}"), - ), - |b| { - b.to_async(&rt).iter(|| async { - for query_index in 0..query_count { - let query = Float32Array::from( - query_values[query_index * dim as usize..(query_index + 1) * dim as usize] - .to_vec(), - ); - let results = dataset - .scan() - .nearest("vector", &query, K) - .unwrap() - .use_index(false) - .project::<&str>(&[]) - .unwrap() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert!(!results.is_empty()); - } - }) - }, - ); - group.bench_function( - BenchmarkId::new( - "batch_query", - format!("rows={num_rows},dim={dim},queries={query_count}"), - ), - |b| { - b.to_async(&rt).iter(|| async { - let results = dataset - .scan() - .nearest("vector", &queries, K) - .unwrap() - .project::<&str>(&[]) - .unwrap() - .try_into_stream() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert!(!results.is_empty()); - }) - }, - ); - group.finish(); - drop(dataset); - let _ = std::fs::remove_dir_all(uri); -} - -fn bench_env(key: &str, default: T) -> T -where - T: std::str::FromStr, -{ - std::env::var(key) - .ok() - .and_then(|value| value.parse().ok()) - .unwrap_or(default) -} - async fn create_file(path: &std::path::Path, mode: WriteMode) { let schema = Arc::new(ArrowSchema::new(vec![Field::new( "vector", @@ -272,8 +155,8 @@ async fn create_file(path: &std::path::Path, mode: WriteMode) { let test_uri = path.to_str().unwrap(); std::fs::remove_dir_all(test_uri).map_or_else(|_| println!("{} not exists", test_uri), |_| {}); let write_params = WriteParams { - max_rows_per_file: num_rows, - max_rows_per_group: batch_size, + max_rows_per_file: num_rows as usize, + max_rows_per_group: batch_size as usize, mode, ..Default::default() }; @@ -304,75 +187,10 @@ async fn create_file(path: &std::path::Path, mode: WriteMode) { .unwrap(); } -async fn create_flat_file( - uri: &str, - mode: WriteMode, - num_rows: usize, - batch_size: usize, - dim: i32, -) -> Dataset { - let schema = Arc::new(ArrowSchema::new(vec![Field::new( - "vector", - DataType::FixedSizeList( - FieldRef::new(Field::new("item", DataType::Float32, true)), - dim, - ), - false, - )])); - - struct FlatVectorBatchIter { - schema: Arc, - remaining_rows: usize, - batch_size: usize, - dim: i32, - } - - impl Iterator for FlatVectorBatchIter { - type Item = Result; - - fn next(&mut self) -> Option { - if self.remaining_rows == 0 { - return None; - } - let rows = self.remaining_rows.min(self.batch_size); - self.remaining_rows -= rows; - - let values = create_float32_array(rows * self.dim as usize); - Some(RecordBatch::try_new( - self.schema.clone(), - vec![Arc::new( - FixedSizeListArray::try_new_from_values(values, self.dim).unwrap(), - )], - )) - } - } - - std::fs::remove_dir_all(uri).map_or_else(|_| println!("{} not exists", uri), |_| {}); - let write_params = WriteParams { - max_rows_per_file: num_rows, - max_rows_per_group: batch_size, - mode, - ..Default::default() - }; - let reader = RecordBatchIterator::new( - FlatVectorBatchIter { - schema: schema.clone(), - remaining_rows: num_rows, - batch_size, - dim, - }, - schema, - ); - Dataset::write(reader, uri, Some(write_params)) - .await - .unwrap() -} - -fn create_float32_array(num_elements: usize) -> Float32Array { - // Generate random values on demand so large benchmark datasets do not need - // to be fully materialized in memory before writing. +fn create_float32_array(num_elements: i32) -> Float32Array { + // generate an Arrow Float32Array with 10000*128 elements randomly let mut rng = rand::rng(); - let mut values = Vec::with_capacity(num_elements); + let mut values = Vec::with_capacity(num_elements as usize); for _ in 0..num_elements { values.push(rng.random_range(0.0..1.0)); } @@ -384,12 +202,12 @@ criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10) .with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); - targets = bench_ivf_pq_index, bench_batch_flat_knn); + targets = bench_ivf_pq_index); // Non-linux version does not support pprof. #[cfg(not(target_os = "linux"))] criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10); - targets = bench_ivf_pq_index, bench_batch_flat_knn); + targets = bench_ivf_pq_index); criterion_main!(benches); diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 10d1bffaf2a..876cb155ad1 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1466,6 +1466,9 @@ impl Scanner { ))); } } + // A list-like query against a multivector column is one multivector query. + // The same list-like query against a fixed-size vector column is a batch + // of single-vector queries. let query_count = if matches!(vector_type, DataType::List(_)) { 1 } else { @@ -1482,6 +1485,9 @@ impl Scanner { dim, ))); } + // A list-like query against a multivector column is one multivector query. + // The same list-like query against a fixed-size vector column is a batch + // of single-vector queries. let query_count = if matches!(vector_type, DataType::List(_)) { 1 } else { diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index f8581e60c82..c1d0b82373b 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -71,7 +71,7 @@ use super::utils::{ SelectionVectorToPrefilter, }; -pub const QUERY_INDEX_COL: &str = "_query_index"; +pub const QUERY_INDEX_COL: &str = "query_index"; pub struct AnnPartitionMetrics { index_metrics: IndexMetrics, @@ -347,10 +347,9 @@ impl Ord for BatchKnnCandidate { } = config; let query_dim = query.len() / query_count; let mut heaps = (0..query_count) - .map(|_| BinaryHeap::::with_capacity(k + 1)) + .map(|_| BinaryHeap::::with_capacity(k)) .collect::>(); let mut input = input; - let mut fallback_row_id = 0_u64; while let Some(batch) = input.next().await { let batch = batch?; @@ -360,7 +359,13 @@ impl Ord for BatchKnnCandidate { let row_ids = batch .column_by_name(ROW_ID) - .map(|row_ids| row_ids.as_primitive::().clone()); + .ok_or_else(|| { + DataFusionError::Internal( + "KNNVectorDistanceExec batch mode requires _rowid in input".to_string(), + ) + })? + .as_primitive::() + .clone(); let batch = Arc::new(batch); for (query_index, heap) in heaps.iter_mut().enumerate().take(query_count) { @@ -378,14 +383,10 @@ impl Ord for BatchKnnCandidate { if !distance.is_finite() { continue; }; - let row_id = row_ids - .as_ref() - .map(|row_ids| row_ids.value(row_index)) - .unwrap_or(fallback_row_id + row_index as u64); let candidate = BatchKnnCandidate { query_index: query_index as u32, distance, - row_id, + row_id: row_ids.value(row_index), batch: batch.clone(), row_index: row_index as u32, }; @@ -400,8 +401,6 @@ impl Ord for BatchKnnCandidate { } } } - - fallback_row_id += batch.num_rows() as u64; } let mut results = heaps From 42fe7b90603a00a5c6db2ec9be7046399a84c42a Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 18 May 2026 19:51:53 +0800 Subject: [PATCH 06/32] fix: format python dataset binding Apply rustfmt output expected by CI for the batch query binding change. Co-authored-by: Cursor --- python/src/dataset.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/src/dataset.rs b/python/src/dataset.rs index c242545d636..82947e11902 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1370,9 +1370,10 @@ impl Dataset { let (vector_type, element_type) = get_vector_type(self_.ds.schema(), &column) .map_err(|e| PyValueError::new_err(e.to_string()))?; - let is_batch_query = - matches!(q.data_type(), DataType::List(_) | DataType::FixedSizeList(_, _)) - && matches!(vector_type, DataType::FixedSizeList(_, _)); + let is_batch_query = matches!( + q.data_type(), + DataType::List(_) | DataType::FixedSizeList(_, _) + ) && matches!(vector_type, DataType::FixedSizeList(_, _)); let scanner = match (is_batch_query, element_type) { (true, DataType::UInt8) => { return Err(PyValueError::new_err( From 4217c77221c5aa673fbfeafe833c9d3bd7a9865e Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 18 May 2026 22:13:29 +0800 Subject: [PATCH 07/32] bench: use pytest params for batch knn benchmark Move batch flat KNN benchmark configuration into pytest parameters so review and reproduction do not rely on environment variables. Co-authored-by: Cursor --- python/python/benchmarks/test_search.py | 41 ++++++++++++++----------- 1 file changed, 23 insertions(+), 18 deletions(-) diff --git a/python/python/benchmarks/test_search.py b/python/python/benchmarks/test_search.py index 783319160e1..3750bdf4396 100644 --- a/python/python/benchmarks/test_search.py +++ b/python/python/benchmarks/test_search.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors -import os import shutil from pathlib import Path from typing import NamedTuple, Union @@ -70,13 +69,6 @@ def create_table(num_rows, offset) -> pa.Table: ) -def get_benchmark_env(key: str, default: int) -> int: - value = os.environ.get(key) - if value is None: - return default - return int(value) - - def create_flat_vector_table(num_rows: int, offset: int, dim: int) -> pa.Table: rng = np.random.default_rng(seed=offset) values = rng.random((num_rows, dim), dtype=np.float32) @@ -224,16 +216,29 @@ def test_knn_search(test_dataset, benchmark): @pytest.mark.benchmark(group="batch_flat_knn") @pytest.mark.parametrize("mode", ["separate", "batch"]) -def test_batch_flat_knn(data_dir: Path, benchmark, mode: str): - dim = get_benchmark_env("LANCE_BATCH_KNN_DIM", BATCH_FLAT_KNN_DIM) - num_rows = get_benchmark_env("LANCE_BATCH_KNN_ROWS", BATCH_FLAT_KNN_ROWS) - batch_size = get_benchmark_env( - "LANCE_BATCH_KNN_BATCH_SIZE", BATCH_FLAT_KNN_BATCH_SIZE - ) - query_count = get_benchmark_env( - "LANCE_BATCH_KNN_QUERY_COUNT", BATCH_FLAT_KNN_QUERY_COUNT - ) - rounds = get_benchmark_env("LANCE_BATCH_KNN_ROUNDS", 10) +@pytest.mark.parametrize( + ("dim", "num_rows", "batch_size", "query_count", "rounds"), + [ + ( + BATCH_FLAT_KNN_DIM, + BATCH_FLAT_KNN_ROWS, + BATCH_FLAT_KNN_BATCH_SIZE, + BATCH_FLAT_KNN_QUERY_COUNT, + 10, + ) + ], + ids=["1m_rows_512d_m10"], +) +def test_batch_flat_knn( + data_dir: Path, + benchmark, + mode: str, + dim: int, + num_rows: int, + batch_size: int, + query_count: int, + rounds: int, +): dataset = create_batch_flat_knn_dataset(data_dir, num_rows, batch_size, dim) query_table = dataset.to_table(columns=["vector"], limit=query_count) query_values = np.asarray( From f4b3f2589fc14744fc29a1430f97b4c4d360ee9c Mon Sep 17 00:00:00 2001 From: zoey Date: Tue, 19 May 2026 00:37:38 +0800 Subject: [PATCH 08/32] fix: respect batch vector query parameters Route batched queries through vector indices when available and apply distance range bounds before per-query top-k selection on the flat path. Co-authored-by: Cursor --- rust/lance/src/dataset/scanner.rs | 210 ++++++++++++++++++++++++++++-- rust/lance/src/io/exec/knn.rs | 131 ++++++++++++------- 2 files changed, 283 insertions(+), 58 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 876cb155ad1..6a491957687 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -38,7 +38,7 @@ use datafusion::scalar::ScalarValue; use datafusion_expr::ExprSchemable; use datafusion_expr::execution_props::ExecutionProps; use datafusion_functions::core::getfield::GetFieldFunc; -use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::expressions::{Column, Literal}; use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalExpr, create_physical_expr}; use datafusion_physical_plan::joins::PartitionMode; use datafusion_physical_plan::projection::ProjectionExec; @@ -101,7 +101,7 @@ use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; use crate::io::exec::{ AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNVectorDistanceExec, LancePushdownScanExec, LanceScanExec, Planner, PreFilterSource, ScanConfig, TakeExec, - knn::{KNN_INDEX_SCHEMA, QUERY_INDEX_COL, new_knn_exec}, + knn::{KNN_INDEX_SCHEMA, KnnBatchParams, QUERY_INDEX_COL, new_knn_exec}, project, }; use crate::io::exec::{AddRowOffsetExec, LanceFilterExec, LanceScanConfig, get_physical_optimizer}; @@ -3532,6 +3532,7 @@ impl Scanner { } // ANN/KNN search execution node with optional prefilter + #[async_recursion] async fn vector_search( &self, filter_plan: &ExprFilterPlan, @@ -3542,14 +3543,8 @@ impl Scanner { // Sanity check let (vector_type, element_type) = get_vector_type(self.dataset.schema(), &q.column)?; - if self.nearest_query_count > 1 && self.fast_search { - return Err(Error::not_supported( - "fast_search is not supported for batch nearest queries".to_string(), - )); - } - let column_id = self.dataset.schema().field_id(q.column.as_str())?; - let use_index = q.use_index && self.nearest_query_count == 1; + let use_index = q.use_index; let indices = if use_index { self.dataset.load_indices().await? } else { @@ -3689,6 +3684,10 @@ impl Scanner { }; if let Some((index_name, index_segments, index_metric)) = index_and_segments { + if self.nearest_query_count > 1 { + return self.batch_indexed_vector_search(filter_plan, &q).await; + } + log::trace!("index found for vector search"); // Use the index's metric type q.metric_type = Some(index_metric); @@ -3766,6 +3765,85 @@ impl Scanner { } } + async fn batch_indexed_vector_search( + &self, + filter_plan: &ExprFilterPlan, + q: &Query, + ) -> Result> { + let query_dim = q.key.len() / self.nearest_query_count; + let mut query_plans = Vec::with_capacity(self.nearest_query_count); + + for query_index in 0..self.nearest_query_count { + let mut single_query = q.clone(); + single_query.key = q.key.slice(query_index * query_dim, query_dim); + + let mut single_scanner = self.clone(); + single_scanner.nearest_query_count = 1; + single_scanner.nearest = Some(single_query.clone()); + + let single_plan = single_scanner + .vector_search(filter_plan, &single_query) + .await?; + query_plans.push(Self::add_query_index_column( + single_plan, + query_index as u32, + )?); + } + + let unioned = UnionExec::try_new(query_plans)?; + let unioned = Arc::new(RepartitionExec::try_new( + unioned, + Partitioning::RoundRobinBatch(1), + )?) as Arc; + + let query_index_sort = PhysicalSortExpr { + expr: expressions::col(QUERY_INDEX_COL, unioned.schema().as_ref())?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }; + let distance_sort = PhysicalSortExpr { + expr: expressions::col(DIST_COL, unioned.schema().as_ref())?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }; + let row_id_sort = PhysicalSortExpr { + expr: expressions::col(ROW_ID, unioned.schema().as_ref())?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }; + + Ok(Arc::new(SortExec::new( + [query_index_sort, distance_sort, row_id_sort].into(), + unioned, + ))) + } + + fn add_query_index_column( + plan: Arc, + query_index: u32, + ) -> Result> { + let schema = plan.schema(); + let mut projection_exprs = Vec::with_capacity(schema.fields().len() + 1); + for field in schema.fields() { + projection_exprs.push(( + Arc::new(Column::new_with_schema(field.name(), schema.as_ref())?) + as Arc, + field.name().clone(), + )); + } + projection_exprs.push(( + Arc::new(Literal::new(ScalarValue::UInt32(Some(query_index)))) as Arc, + QUERY_INDEX_COL.to_string(), + )); + Ok(Arc::new(ProjectionExec::try_new(projection_exprs, plan)?)) + } + /// Combine ANN results with KNN results for data appended after index creation async fn knn_combined( &self, @@ -4314,9 +4392,13 @@ impl Scanner { input, &q.column, q.key.clone(), - self.nearest_query_count, - q.k, - metric_type, + KnnBatchParams { + query_count: self.nearest_query_count, + k: q.k, + lower_bound: q.lower_bound, + upper_bound: q.upper_bound, + distance_type: metric_type, + }, )?); let lower: Option<(Expr, Arc)> = q @@ -5042,8 +5124,8 @@ mod test { use arrow_array::cast::AsArray; use arrow_array::types::{Float32Type, UInt32Type, UInt64Type}; use arrow_array::{ - ArrayRef, FixedSizeListArray, Float16Array, Int32Array, LargeStringArray, PrimitiveArray, - RecordBatchIterator, StringArray, StructArray, UInt8Array, + ArrayRef, BooleanArray, FixedSizeListArray, Float16Array, Int32Array, LargeStringArray, + PrimitiveArray, RecordBatchIterator, StringArray, StructArray, UInt8Array, }; use arrow_ord::sort::sort_to_indices; @@ -5701,6 +5783,7 @@ mod test { let mut scan = dataset.scan(); scan.nearest("vec", &queries, 2).unwrap(); + scan.use_index(false); scan.project(&["i"]).unwrap(); let plan = scan.explain_plan(false).await.unwrap(); @@ -5751,6 +5834,105 @@ mod test { } } + #[tokio::test] + async fn test_batch_knn_flat_respects_distance_range() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let dataset = &test_ds.dataset; + + let query_values = (32..96).map(|v| v as f32).collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) + .unwrap(); + + let batch = dataset + .scan() + .nearest("vec", &queries, 2) + .unwrap() + .use_index(false) + .distance_range(Some(1.0), None) + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + + assert_eq!(batch.num_rows(), 4); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0, 1, 1] + ); + + for query_index in 0..2 { + let query = + Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); + let single = dataset + .scan() + .nearest("vec", &query, 2) + .unwrap() + .use_index(false) + .distance_range(Some(1.0), None) + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + let mask = BooleanArray::from_iter( + query_indices + .iter() + .map(|value| value.map(|value| value == query_index as u32)), + ); + let batch_slice = arrow::compute::filter_record_batch(&batch, &mask).unwrap(); + assert_eq!( + batch_slice["i"].as_primitive::().values(), + single["i"].as_primitive::().values() + ); + assert_eq!( + batch_slice[DIST_COL].as_primitive::().values(), + single[DIST_COL].as_primitive::().values() + ); + } + } + + #[tokio::test] + async fn test_batch_knn_uses_index_when_available() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + test_ds.make_vector_index().await.unwrap(); + let dataset = &test_ds.dataset; + + let query_values = (32..96).map(|v| v as f32).collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) + .unwrap(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &queries, 2).unwrap(); + scan.project(&["i"]).unwrap(); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "batch KNN should use the vector index when available, got:\n{}", + plan + ); + assert!( + !plan.contains("KNNVectorDistance: queries=2"), + "indexed batch KNN should not force the flat batch path, got:\n{}", + plan + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 4); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0, 1, 1] + ); + } + #[rstest] #[tokio::test] async fn test_can_project_distance() { diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index c1d0b82373b..2adc196c7cb 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -147,6 +147,8 @@ pub struct KNNVectorDistanceExec { pub query: ArrayRef, pub query_count: usize, pub k: usize, + pub lower_bound: Option, + pub upper_bound: Option, pub column: String, pub distance_type: DistanceType, @@ -157,6 +159,14 @@ pub struct KNNVectorDistanceExec { metrics: ExecutionPlanMetricsSet, } +pub struct KnnBatchParams { + pub query_count: usize, + pub k: usize, + pub lower_bound: Option, + pub upper_bound: Option, + pub distance_type: DistanceType, +} + struct BatchKnnConfig { input_schema: SchemaRef, output_schema: SchemaRef, @@ -164,6 +174,8 @@ struct BatchKnnConfig { query: ArrayRef, query_count: usize, k: usize, + lower_bound: Option, + upper_bound: Option, distance_type: DistanceType, } @@ -206,53 +218,33 @@ impl KNNVectorDistanceExec { query: ArrayRef, distance_type: DistanceType, ) -> Result { - Self::try_new_batch(input, column, query, 1, 0, distance_type) - } - -#[derive(Clone)] -struct BatchKnnCandidate { - query_index: u32, - distance: f32, - row_id: u64, - batch: Arc, - row_index: u32, -} - -impl PartialEq for BatchKnnCandidate { - fn eq(&self, other: &Self) -> bool { - self.query_index == other.query_index - && self.distance == other.distance - && self.row_id == other.row_id - && self.row_index == other.row_index - } -} - -impl Eq for BatchKnnCandidate {} - -impl PartialOrd for BatchKnnCandidate { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for BatchKnnCandidate { - fn cmp(&self, other: &Self) -> CmpOrdering { - self.distance - .total_cmp(&other.distance) - .then_with(|| self.row_id.cmp(&other.row_id)) - .then_with(|| self.query_index.cmp(&other.query_index)) - .then_with(|| self.row_index.cmp(&other.row_index)) + Self::try_new_batch( + input, + column, + query, + KnnBatchParams { + query_count: 1, + k: 0, + lower_bound: None, + upper_bound: None, + distance_type, + }, + ) } -} - pub fn try_new_batch( + pub(crate) fn try_new_batch( input: Arc, column: &str, query: ArrayRef, - query_count: usize, - k: usize, - distance_type: DistanceType, + params: KnnBatchParams, ) -> Result { + let KnnBatchParams { + query_count, + k, + lower_bound, + upper_bound, + distance_type, + } = params; if query_count == 0 { return Err(Error::invalid_input( "query_count must be positive for KNN".to_string(), @@ -323,6 +315,8 @@ impl Ord for BatchKnnCandidate { query, query_count, k, + lower_bound, + upper_bound, column: column.to_string(), distance_type, input_schema, @@ -343,6 +337,8 @@ impl Ord for BatchKnnCandidate { query, query_count, k, + lower_bound, + upper_bound, distance_type, } = config; let query_dim = query.len() / query_count; @@ -383,6 +379,11 @@ impl Ord for BatchKnnCandidate { if !distance.is_finite() { continue; }; + if lower_bound.is_some_and(|lower_bound| distance < lower_bound) + || upper_bound.is_some_and(|upper_bound| distance >= upper_bound) + { + continue; + } let candidate = BatchKnnCandidate { query_index: query_index as u32, distance, @@ -483,9 +484,13 @@ impl ExecutionPlan for KNNVectorDistanceExec { children.pop().expect("length checked"), &self.column, self.query.clone(), - self.query_count, - self.k, - self.distance_type, + KnnBatchParams { + query_count: self.query_count, + k: self.k, + lower_bound: self.lower_bound, + upper_bound: self.upper_bound, + distance_type: self.distance_type, + }, )?)) } @@ -505,6 +510,8 @@ impl ExecutionPlan for KNNVectorDistanceExec { query: self.query.clone(), query_count: self.query_count, k: self.k, + lower_bound: self.lower_bound, + upper_bound: self.upper_bound, distance_type: self.distance_type, }, )); @@ -609,6 +616,42 @@ impl ExecutionPlan for KNNVectorDistanceExec { } } +#[derive(Clone)] +struct BatchKnnCandidate { + query_index: u32, + distance: f32, + row_id: u64, + batch: Arc, + row_index: u32, +} + +impl PartialEq for BatchKnnCandidate { + fn eq(&self, other: &Self) -> bool { + self.query_index == other.query_index + && self.distance == other.distance + && self.row_id == other.row_id + && self.row_index == other.row_index + } +} + +impl Eq for BatchKnnCandidate {} + +impl PartialOrd for BatchKnnCandidate { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for BatchKnnCandidate { + fn cmp(&self, other: &Self) -> CmpOrdering { + self.distance + .total_cmp(&other.distance) + .then_with(|| self.row_id.cmp(&other.row_id)) + .then_with(|| self.query_index.cmp(&other.query_index)) + .then_with(|| self.row_index.cmp(&other.row_index)) + } +} + pub static KNN_INDEX_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![ Field::new(DIST_COL, DataType::Float32, true), From 82ab93762ce3df1123ed3b912105274deba4bad4 Mon Sep 17 00:00:00 2001 From: zoey Date: Tue, 19 May 2026 01:02:24 +0800 Subject: [PATCH 09/32] test: assert indexed batch KNN matches single-query distance_range Co-authored-by: Cursor --- rust/lance/src/dataset/scanner.rs | 63 +++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 6a491957687..2228c24f039 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -5933,6 +5933,69 @@ mod test { ); } + #[tokio::test] + async fn test_batch_knn_indexed_respects_distance_range() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + test_ds.make_vector_index().await.unwrap(); + let dataset = &test_ds.dataset; + + let query_values = (32..96).map(|v| v as f32).collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) + .unwrap(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &queries, 2).unwrap(); + scan.distance_range(Some(1.0), None); + scan.project(&["i"]).unwrap(); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("ANNSubIndex"), + "indexed batch KNN should use the vector index, got:\n{}", + plan + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 4); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0, 1, 1] + ); + + for query_index in 0..2 { + let query = + Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); + let single = dataset + .scan() + .nearest("vec", &query, 2) + .unwrap() + .distance_range(Some(1.0), None) + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + let mask = BooleanArray::from_iter( + query_indices + .iter() + .map(|value| value.map(|value| value == query_index as u32)), + ); + let batch_slice = arrow::compute::filter_record_batch(&batch, &mask).unwrap(); + assert_eq!( + batch_slice["i"].as_primitive::().values(), + single["i"].as_primitive::().values() + ); + assert_eq!( + batch_slice[DIST_COL].as_primitive::().values(), + single[DIST_COL].as_primitive::().values() + ); + } + } + #[rstest] #[tokio::test] async fn test_can_project_distance() { From dfd4532fcd15b1a012c763746a4f9e9fffb5eb27 Mon Sep 17 00:00:00 2001 From: zoey Date: Thu, 21 May 2026 07:15:58 +0800 Subject: [PATCH 10/32] docs: align batch vector query Python docs with implementation Update nearest/search docstrings to describe indexed batch queries and add Python tests that batch distance_range matches per-query searches. Co-authored-by: Cursor --- python/python/lance/dataset.py | 20 +++++---- python/python/tests/test_vector_index.py | 53 ++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 8 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 788a58ae24b..a4a028ed7f6 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -1100,10 +1100,11 @@ def scanner( } ``q`` may also be a 2-D array-like value for fixed-size vector columns. - In that case Lance runs a flat batch KNN query, returns up to ``k`` rows - for each query vector, and adds ``query_index`` to identify the source - query for each result row. Indexed/ANN batch search is not used in this - first implementation. + In that case Lance runs a batch nearest-neighbor query, returns up to + ``k`` rows for each query vector, and adds ``query_index`` to identify the + source query for each result row. When ``use_index`` is true and a vector + index is available, each query vector is searched through the index + path; otherwise the flat batch path is used. batch_size: int, default None The maximum number of rows per batch. In some cases batches can be @@ -5998,7 +5999,9 @@ def nearest( q: QueryVectorLike A single query vector or, for fixed-size vector columns, a 2-D array-like batch of query vectors. Batch queries return up to ``k`` rows per query - and include ``query_index`` in the output. + and include ``query_index`` in the output. When ``use_index`` is true and + a vector index is available, each query vector is searched through the + index path; otherwise the flat batch path is used. query_parallelism: int, optional Maximum partition-search concurrency for a single vector query. The default is 0. Value 0 uses the automatic policy, which @@ -7148,9 +7151,10 @@ def _build_vector_search_query( The name of the vector column to search. q: QueryVectorLike The query vector. For fixed-size vector columns, this may be a 2-D - array-like batch of query vectors. Batch queries run flat KNN, apply - ``k`` per query vector, and add ``query_index`` to the result so - callers can split rows by input query. + array-like batch of query vectors. Batch queries return up to ``k`` rows per + query vector and include ``query_index`` in the output. When ``use_index`` + is true and a vector index is available, each query vector is searched + through the index path; otherwise the flat batch path is used. k: int, optional The number of nearest neighbors to return. metric: str, optional diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 61848c9667b..a02bef3c8a3 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -215,6 +215,59 @@ def test_batch_flat_query_matches_repeated_single_queries(dataset): ) +def _assert_batch_matches_single_queries(ds, queries, k, nearest_kwargs): + batch = ds.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": queries, + "k": k, + **nearest_kwargs, + }, + ) + assert all( + nearest_kwargs["distance_range"][0] <= d < nearest_kwargs["distance_range"][1] + for d in batch["_distance"].to_pylist() + ) + + for query_index, query in enumerate(queries): + single = ds.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": query, + "k": k, + **nearest_kwargs, + }, + ) + batch_slice = batch.filter(pc.field("query_index") == query_index) + assert batch_slice["id"].to_pylist() == single["id"].to_pylist() + np.testing.assert_allclose( + batch_slice["_distance"].to_numpy(), + single["_distance"].to_numpy(), + ) + + +def test_batch_flat_respects_distance_range(dataset): + queries = np.random.randn(2, 128).astype(np.float32) + _assert_batch_matches_single_queries( + dataset, + queries, + k=5, + nearest_kwargs={"use_index": False, "distance_range": (0.0, 50.0)}, + ) + + +def test_batch_indexed_respects_distance_range(indexed_dataset): + queries = np.random.randn(2, 128).astype(np.float32) + _assert_batch_matches_single_queries( + indexed_dataset, + queries, + k=5, + nearest_kwargs={"distance_range": (0.0, 50.0)}, + ) + + def test_ann(indexed_dataset): run(indexed_dataset) From 69013a007fa5223692cf18872ccec77cce19b5b7 Mon Sep 17 00:00:00 2001 From: zoey Date: Thu, 21 May 2026 07:36:44 +0800 Subject: [PATCH 11/32] fix: include query_index in empty batch fast_search results When fast_search is used with a batch nearest query and no vector index, return an empty result whose schema still contains query_index. Co-authored-by: Cursor --- python/python/tests/test_vector_index.py | 15 ++++++++++++++ rust/lance/src/dataset/scanner.rs | 26 ++++++++++++++++++++++-- rust/lance/src/io/exec/knn.rs | 15 ++++++++++---- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index a02bef3c8a3..4fc2c034635 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -268,6 +268,21 @@ def test_batch_indexed_respects_distance_range(indexed_dataset): ) +def test_batch_fast_search_without_index_returns_empty_with_query_index(dataset): + queries = np.random.randn(2, 128).astype(np.float32) + batch = dataset.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": queries, + "k": 5, + }, + fast_search=True, + ) + assert batch.num_rows == 0 + assert "query_index" in batch.column_names + + def test_ann(indexed_dataset): run(indexed_dataset) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 2228c24f039..cfd731749be 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -101,7 +101,7 @@ use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; use crate::io::exec::{ AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNVectorDistanceExec, LancePushdownScanExec, LanceScanExec, Planner, PreFilterSource, ScanConfig, TakeExec, - knn::{KNN_INDEX_SCHEMA, KnnBatchParams, QUERY_INDEX_COL, new_knn_exec}, + knn::{KnnBatchParams, QUERY_INDEX_COL, knn_empty_result_schema, new_knn_exec}, project, }; use crate::io::exec::{AddRowOffsetExec, LanceFilterExec, LanceScanConfig, get_physical_optimizer}; @@ -3725,7 +3725,9 @@ impl Scanner { Ok(knn_node) } else { if self.fast_search { - return Ok(Arc::new(EmptyExec::new(KNN_INDEX_SCHEMA.clone()))); + return Ok(Arc::new(EmptyExec::new(knn_empty_result_schema( + self.nearest_query_count > 1, + )))); } // Resolve metric type for flat search (use default if not specified) let metric = q @@ -10105,6 +10107,26 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_eq!(fast_rows, 0); } + #[tokio::test] + async fn test_batch_fast_search_without_index_returns_empty_with_query_index() { + let dataset = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let query_values = (32..96).map(|v| v as f32).collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values), 32).unwrap(); + + let mut scanner = dataset.dataset.scan(); + scanner.nearest("vec", &queries, 2).unwrap().fast_search(); + let batch = scanner.try_into_batch().await.unwrap(); + + assert_eq!(batch.num_rows(), 0); + assert!( + batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), + "batch fast_search without index should still expose query_index in schema" + ); + } + #[rstest] #[tokio::test] pub async fn test_scan_planning_io( diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 2adc196c7cb..b799efeddc6 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -652,12 +652,19 @@ impl Ord for BatchKnnCandidate { } } -pub static KNN_INDEX_SCHEMA: LazyLock = LazyLock::new(|| { - Arc::new(Schema::new(vec![ +pub static KNN_INDEX_SCHEMA: LazyLock = LazyLock::new(|| knn_empty_result_schema(false)); + +/// Schema for empty vector-search results (e.g. `fast_search` with no index). +pub fn knn_empty_result_schema(include_query_index: bool) -> SchemaRef { + let mut fields = vec![ Field::new(DIST_COL, DataType::Float32, true), ROW_ID_FIELD.clone(), - ])) -}); + ]; + if include_query_index { + fields.push(Field::new(QUERY_INDEX_COL, DataType::UInt32, true)); + } + Arc::new(Schema::new(fields)) +} pub static KNN_PARTITION_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![ From f64e3d6e8f5e9716e1751cfd5150b25f29ded7e5 Mon Sep 17 00:00:00 2001 From: zoey Date: Thu, 21 May 2026 17:01:14 +0800 Subject: [PATCH 12/32] fix: treat batch nearest by query shape and skip SortExec top-k Use is_batch_nearest based on list-like queries on fixed-size vector columns instead of query_count > 1, so single-vector batch queries still get query_index and avoid SortExec TopK(fetch=k) truncating m*k results to k rows. Co-authored-by: Cursor --- python/python/tests/test_vector_index.py | 19 ++++ rust/lance/src/dataset/scanner.rs | 112 +++++++++++++++++++++-- rust/lance/src/io/exec/knn.rs | 22 +++-- 3 files changed, 137 insertions(+), 16 deletions(-) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 4fc2c034635..d12adeb4467 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -268,6 +268,25 @@ def test_batch_indexed_respects_distance_range(indexed_dataset): ) +def test_batch_single_vector_list_query_includes_query_index(dataset): + query = np.random.randn(1, 128).astype(np.float32) + k = 5 + + batch = dataset.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": query, + "k": k, + "use_index": False, + }, + ) + + assert batch.num_rows == k + assert "query_index" in batch.column_names + assert batch["query_index"].to_pylist() == [0] * k + + def test_batch_fast_search_without_index_returns_empty_with_query_index(dataset): queries = np.random.randn(2, 128).astype(np.float32) batch = dataset.to_table( diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index cfd731749be..d7815bffafe 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -770,6 +770,9 @@ pub struct Scanner { nearest: 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). + is_batch_nearest: bool, /// If false, do not use any scalar indices for the scan /// @@ -1026,6 +1029,7 @@ impl Scanner { ordering: None, nearest: None, nearest_query_count: 1, + is_batch_nearest: false, use_stats: true, ordered: true, fragments: None, @@ -1430,6 +1434,23 @@ impl Scanner { Ok(self) } + fn is_batch_nearest_query( + vector_type: &DataType, + query_type: &DataType, + query_len: usize, + dim: usize, + ) -> bool { + match vector_type { + DataType::FixedSizeList(_, _) => { + matches!( + query_type, + DataType::List(_) | DataType::FixedSizeList(_, _) + ) || (query_len > dim && query_len.is_multiple_of(dim)) + } + _ => false, + } + } + /// Find k-nearest neighbor within the vector column. /// the query can be a Float16Array, Float32Array, Float64Array, UInt8Array, /// or a ListArray/FixedSizeListArray of the above types. @@ -1451,8 +1472,10 @@ impl Scanner { // make sure the field exists let (vector_type, element_type) = get_vector_type(self.dataset.schema(), column)?; let dim = get_vector_dim(self.dataset.schema(), column)?; + let query_type = q.data_type().clone(); + let query_len = q.len(); - let (q, query_count) = match q.data_type() { + let (q, query_count) = match &query_type { DataType::List(_) | DataType::FixedSizeList(_, _) => { if let Some(list_array) = q.as_list_opt::() { for i in 0..list_array.len() { @@ -1551,6 +1574,8 @@ impl Scanner { dist_q_c: 0.0, }); self.nearest_query_count = query_count; + self.is_batch_nearest = + Self::is_batch_nearest_query(&vector_type, &query_type, query_len, dim); Ok(self) } @@ -1881,7 +1906,7 @@ impl Scanner { if self.nearest.as_ref().is_some() { extra_columns.push(ArrowField::new(DIST_COL, DataType::Float32, true)); - if self.nearest_query_count > 1 { + if self.is_batch_nearest { extra_columns.push(ArrowField::new(QUERY_INDEX_COL, DataType::UInt32, true)); } }; @@ -1932,8 +1957,7 @@ impl Scanner { let vector_expr = expressions::col(DIST_COL, current_schema)?; output_expr.push((vector_expr, DIST_COL.to_string())); } - if self.nearest_query_count > 1 - && output_expr.iter().all(|(_, name)| name != QUERY_INDEX_COL) + if self.is_batch_nearest && output_expr.iter().all(|(_, name)| name != QUERY_INDEX_COL) { let query_index_expr = expressions::col(QUERY_INDEX_COL, current_schema)?; output_expr.push((query_index_expr, QUERY_INDEX_COL.to_string())); @@ -3684,7 +3708,7 @@ impl Scanner { }; if let Some((index_name, index_segments, index_metric)) = index_and_segments { - if self.nearest_query_count > 1 { + if self.is_batch_nearest { return self.batch_indexed_vector_search(filter_plan, &q).await; } @@ -3726,7 +3750,7 @@ impl Scanner { } else { if self.fast_search { return Ok(Arc::new(EmptyExec::new(knn_empty_result_schema( - self.nearest_query_count > 1, + self.is_batch_nearest, )))); } // Resolve metric type for flat search (use default if not specified) @@ -3781,6 +3805,7 @@ impl Scanner { let mut single_scanner = self.clone(); single_scanner.nearest_query_count = 1; + single_scanner.is_batch_nearest = false; single_scanner.nearest = Some(single_query.clone()); let single_plan = single_scanner @@ -4385,7 +4410,7 @@ impl Scanner { default_distance_type_for(&element_type) } }; - let input = if self.nearest_query_count > 1 { + let input = if self.is_batch_nearest { Arc::new(CoalescePartitionsExec::new(input)) as Arc } else { input @@ -4395,6 +4420,7 @@ impl Scanner { &q.column, q.key.clone(), KnnBatchParams { + is_batch: self.is_batch_nearest, query_count: self.nearest_query_count, k: q.k, lower_bound: q.lower_bound, @@ -4444,7 +4470,7 @@ impl Scanner { flat_dist }; - if self.nearest_query_count > 1 { + if self.is_batch_nearest { return Ok(knn_plan); } @@ -10107,6 +10133,76 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_eq!(fast_rows, 0); } + #[tokio::test] + async fn test_batch_single_vector_list_query_includes_query_index() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let dataset = &test_ds.dataset; + + let query_values = (32..64).map(|v| v as f32).collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values), 32).unwrap(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &queries, 2).unwrap().use_index(false); + scan.project(&["i"]).unwrap(); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("KNNVectorDistance: queries=1"), + "single-vector batch query should use batch KNN path, got:\n{}", + plan + ); + assert!( + !plan.contains("SortExec: TopK(fetch="), + "batch KNN must not apply per-query SortExec top-k, got:\n{}", + plan + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 2); + assert!( + batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), + "batch-shaped query with one vector should still return query_index" + ); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0] + ); + } + + #[tokio::test] + async fn test_batch_flat_plan_returns_m_times_k_without_sort_topk() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let dataset = &test_ds.dataset; + + let query_values = (32..96).map(|v| v as f32).collect::>(); + let queries = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values), 32).unwrap(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &queries, 2).unwrap().use_index(false); + scan.project(&["i"]).unwrap(); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("KNNVectorDistance: queries=2"), + "expected batch flat KNN plan, got:\n{}", + plan + ); + assert!( + !plan.contains("SortExec: TopK(fetch="), + "batch flat KNN must not truncate to k rows globally, got:\n{}", + plan + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 4); + } + #[tokio::test] async fn test_batch_fast_search_without_index_returns_empty_with_query_index() { let dataset = TestVectorDataset::new(LanceFileVersion::Stable, true) diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index b799efeddc6..83d12f0240f 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -145,6 +145,7 @@ pub struct KNNVectorDistanceExec { /// The vector query to execute. pub query: ArrayRef, + pub is_batch: bool, pub query_count: usize, pub k: usize, pub lower_bound: Option, @@ -160,6 +161,7 @@ pub struct KNNVectorDistanceExec { } pub struct KnnBatchParams { + pub is_batch: bool, pub query_count: usize, pub k: usize, pub lower_bound: Option, @@ -183,7 +185,7 @@ impl DisplayAs for KNNVectorDistanceExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - if self.query_count > 1 { + if self.is_batch { write!( f, "KNNVectorDistance: queries={}, k={}, metric={}", @@ -194,7 +196,7 @@ impl DisplayAs for KNNVectorDistanceExec { } } DisplayFormatType::TreeRender => { - if self.query_count > 1 { + if self.is_batch { write!( f, "KNNVectorDistance\nqueries={}\nk={}\nmetric={}", @@ -223,6 +225,7 @@ impl KNNVectorDistanceExec { column, query, KnnBatchParams { + is_batch: false, query_count: 1, k: 0, lower_bound: None, @@ -239,6 +242,7 @@ impl KNNVectorDistanceExec { params: KnnBatchParams, ) -> Result { let KnnBatchParams { + is_batch, query_count, k, lower_bound, @@ -257,7 +261,7 @@ impl KNNVectorDistanceExec { query_count ))); } - if query_count > 1 && k == 0 { + if is_batch && k == 0 { return Err(Error::invalid_input( "k must be positive for batch KNN".to_string(), )); @@ -278,7 +282,7 @@ impl KNNVectorDistanceExec { } let input_schema = Arc::new(input_schema); let mut output_schema = input_schema.as_ref().clone(); - if query_count > 1 { + if is_batch { output_schema = output_schema.try_with_column(Field::new( QUERY_INDEX_COL, DataType::UInt32, @@ -293,7 +297,7 @@ impl KNNVectorDistanceExec { // This node has the same partitioning & boundedness as the input node // but it destroys any ordering. - let properties = if query_count > 1 { + let properties = if is_batch { Arc::new(PlanProperties::new( EquivalenceProperties::new(output_schema.clone()), Partitioning::UnknownPartitioning(1), @@ -313,6 +317,7 @@ impl KNNVectorDistanceExec { Ok(Self { input, query, + is_batch, query_count, k, lower_bound, @@ -485,6 +490,7 @@ impl ExecutionPlan for KNNVectorDistanceExec { &self.column, self.query.clone(), KnnBatchParams { + is_batch: self.is_batch, query_count: self.query_count, k: self.k, lower_bound: self.lower_bound, @@ -500,7 +506,7 @@ impl ExecutionPlan for KNNVectorDistanceExec { context: Arc, ) -> DataFusionResult { let input_stream = self.input.execute(partition, context)?; - if self.query_count > 1 { + if self.is_batch { let stream = stream::once(Self::execute_batch( input_stream, BatchKnnConfig { @@ -576,7 +582,7 @@ impl ExecutionPlan for KNNVectorDistanceExec { .filter(|(_, field)| field.name() != DIST_COL) .map(|(stats, _)| stats) .collect::>(); - let column_statistics = if self.query_count > 1 { + let column_statistics = if self.is_batch { column_statistics .into_iter() .chain(std::iter::once(ColumnStatistics::default())) @@ -608,7 +614,7 @@ impl ExecutionPlan for KNNVectorDistanceExec { } fn required_input_distribution(&self) -> Vec { - if self.query_count > 1 { + if self.is_batch { vec![Distribution::SinglePartition] } else { vec![Distribution::UnspecifiedDistribution] From 584c07be308a5557b516ca76b95c33da18058c9c Mon Sep 17 00:00:00 2001 From: zoey Date: Fri, 22 May 2026 07:36:55 +0800 Subject: [PATCH 13/32] test: deduplicate batch KNN tests Consolidate overlapping Rust/Python batch nearest tests via shared helpers. No production changes; merge with main deferred. Co-authored-by: Cursor --- python/python/tests/test_vector_index.py | 66 ++--- rust/lance/src/dataset/scanner.rs | 339 ++++++++--------------- 2 files changed, 145 insertions(+), 260 deletions(-) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index d12adeb4467..195cfa35de2 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -180,9 +180,17 @@ def test_flat(dataset): run(dataset) -def test_batch_flat_query_matches_repeated_single_queries(dataset): - queries = np.random.randn(2, 128).astype(np.float32) +@pytest.mark.parametrize( + "queries", + [ + np.random.randn(2, 128).astype(np.float32), + np.random.randn(1, 128).astype(np.float32), + ], + ids=["two_queries", "single_query"], +) +def test_batch_flat_query_matches_repeated_single_queries(dataset, queries): k = 5 + query_count = queries.shape[0] batch = dataset.to_table( columns=["id"], @@ -190,29 +198,21 @@ def test_batch_flat_query_matches_repeated_single_queries(dataset): "column": "vector", "q": queries, "k": k, + "use_index": False, }, ) - assert batch.num_rows == queries.shape[0] * k + assert batch.num_rows == query_count * k assert batch.column_names == ["id", "_distance", "query_index"] - assert batch["query_index"].to_pylist() == [0] * k + [1] * k + expected_query_index = sum([[i] * k for i in range(query_count)], []) + assert batch["query_index"].to_pylist() == expected_query_index - for query_index, query in enumerate(queries): - single = dataset.to_table( - columns=["id"], - nearest={ - "column": "vector", - "q": query, - "k": k, - "use_index": False, - }, - ) - batch_slice = batch.filter(pc.field("query_index") == query_index) - assert batch_slice["id"].to_pylist() == single["id"].to_pylist() - np.testing.assert_allclose( - batch_slice["_distance"].to_numpy(), - single["_distance"].to_numpy(), - ) + _assert_batch_matches_single_queries( + dataset, + queries, + k=k, + nearest_kwargs={"use_index": False}, + ) def _assert_batch_matches_single_queries(ds, queries, k, nearest_kwargs): @@ -225,10 +225,9 @@ def _assert_batch_matches_single_queries(ds, queries, k, nearest_kwargs): **nearest_kwargs, }, ) - assert all( - nearest_kwargs["distance_range"][0] <= d < nearest_kwargs["distance_range"][1] - for d in batch["_distance"].to_pylist() - ) + if "distance_range" in nearest_kwargs: + lo, hi = nearest_kwargs["distance_range"] + assert all(lo <= d < hi for d in batch["_distance"].to_pylist()) for query_index, query in enumerate(queries): single = ds.to_table( @@ -268,25 +267,6 @@ def test_batch_indexed_respects_distance_range(indexed_dataset): ) -def test_batch_single_vector_list_query_includes_query_index(dataset): - query = np.random.randn(1, 128).astype(np.float32) - k = 5 - - batch = dataset.to_table( - columns=["id"], - nearest={ - "column": "vector", - "q": query, - "k": k, - "use_index": False, - }, - ) - - assert batch.num_rows == k - assert "query_index" in batch.column_names - assert batch["query_index"].to_pylist() == [0] * k - - def test_batch_fast_search_without_index_returns_empty_with_query_index(dataset): queries = np.random.randn(2, 128).astype(np.float32) batch = dataset.to_table( diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index d7815bffafe..b1692dbf5b2 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -5796,21 +5796,66 @@ mod test { assert_eq!(expected_i, actual_i); } - #[tokio::test] - async fn test_batch_knn_flat_results_include_query_index() { - let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) - .await - .unwrap(); - test_ds.make_vector_index().await.unwrap(); - let dataset = &test_ds.dataset; - + fn batch_knn_two_queries() -> (FixedSizeListArray, Vec) { let query_values = (32..96).map(|v| v as f32).collect::>(); let queries = FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) .unwrap(); + (queries, query_values) + } + + async fn assert_batch_matches_single_queries( + dataset: &Dataset, + batch: &RecordBatch, + query_values: &[f32], + k: usize, + use_index: bool, + distance_range: Option<(Option, Option)>, + ) { + let query_count = query_values.len() / 32; + assert_eq!(batch.num_rows(), query_count * k); + + for query_index in 0..query_count { + let query = + Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); + let mut scan = dataset.scan(); + scan.nearest("vec", &query, k).unwrap(); + scan.use_index(use_index); + if let Some((lower, upper)) = distance_range { + scan.distance_range(lower, upper); + } + scan.project(&["i"]).unwrap(); + let single = scan.try_into_batch().await.unwrap(); + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + let mask = BooleanArray::from_iter( + query_indices + .iter() + .map(|value| value.map(|value| value == query_index as u32)), + ); + let batch_slice = arrow::compute::filter_record_batch(batch, &mask).unwrap(); + assert_eq!( + batch_slice["i"].as_primitive::().values(), + single["i"].as_primitive::().values() + ); + assert_eq!( + batch_slice[DIST_COL].as_primitive::().values(), + single[DIST_COL].as_primitive::().values() + ); + } + } + + #[tokio::test] + async fn test_batch_knn_flat() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let dataset = &test_ds.dataset; + let k = 2; + + let (queries, query_values) = batch_knn_two_queries(); let mut scan = dataset.scan(); - scan.nearest("vec", &queries, 2).unwrap(); + scan.nearest("vec", &queries, k).unwrap(); scan.use_index(false); scan.project(&["i"]).unwrap(); @@ -5822,44 +5867,52 @@ mod test { ); assert!( !plan.contains("ANNSubIndex"), - "batch KNN should not use ANN index yet, got:\n{}", + "flat batch KNN should not use ANN index, got:\n{}", + plan + ); + assert!( + !plan.contains("SortExec: TopK(fetch="), + "batch flat KNN must not truncate to k rows globally, got:\n{}", plan ); let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 4); - - let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); - assert_eq!(query_indices.values(), &[0, 0, 1, 1]); - - let batch_ids = batch["i"].as_primitive::(); - let batch_distances = batch[DIST_COL].as_primitive::(); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0, 1, 1] + ); + assert_batch_matches_single_queries(dataset, &batch, &query_values, k, false, None).await; - for query_index in 0..2 { - let query = - Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); - let single = dataset - .scan() - .nearest("vec", &query, 2) - .unwrap() - .use_index(false) - .project(&["i"]) - .unwrap() - .try_into_batch() - .await + let query_values_one = (32..64).map(|v| v as f32).collect::>(); + let queries_one = + FixedSizeListArray::try_new_from_values(Float32Array::from(query_values_one.clone()), 32) .unwrap(); - let single_ids = single["i"].as_primitive::(); - let single_distances = single[DIST_COL].as_primitive::(); + let mut scan = dataset.scan(); + scan.nearest("vec", &queries_one, k).unwrap(); + scan.use_index(false); + scan.project(&["i"]).unwrap(); - for result_index in 0..2 { - let batch_index = query_index * 2 + result_index; - assert_eq!(batch_ids.value(batch_index), single_ids.value(result_index)); - assert_eq!( - batch_distances.value(batch_index), - single_distances.value(result_index) - ); - } - } + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("KNNVectorDistance: queries=1"), + "single-vector batch query should use batch KNN path, got:\n{}", + plan + ); + assert!( + !plan.contains("SortExec: TopK(fetch="), + "batch KNN must not apply per-query SortExec top-k, got:\n{}", + plan + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert!( + batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), + "batch-shaped query with one vector should still return query_index" + ); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0] + ); } #[tokio::test] @@ -5868,11 +5921,7 @@ mod test { .await .unwrap(); let dataset = &test_ds.dataset; - - let query_values = (32..96).map(|v| v as f32).collect::>(); - let queries = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) - .unwrap(); + let (queries, query_values) = batch_knn_two_queries(); let batch = dataset .scan() @@ -5886,56 +5935,29 @@ mod test { .await .unwrap(); - assert_eq!(batch.num_rows(), 4); assert_eq!( batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0, 1, 1] ); - - for query_index in 0..2 { - let query = - Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); - let single = dataset - .scan() - .nearest("vec", &query, 2) - .unwrap() - .use_index(false) - .distance_range(Some(1.0), None) - .project(&["i"]) - .unwrap() - .try_into_batch() - .await - .unwrap(); - let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); - let mask = BooleanArray::from_iter( - query_indices - .iter() - .map(|value| value.map(|value| value == query_index as u32)), - ); - let batch_slice = arrow::compute::filter_record_batch(&batch, &mask).unwrap(); - assert_eq!( - batch_slice["i"].as_primitive::().values(), - single["i"].as_primitive::().values() - ); - assert_eq!( - batch_slice[DIST_COL].as_primitive::().values(), - single[DIST_COL].as_primitive::().values() - ); - } + assert_batch_matches_single_queries( + dataset, + &batch, + &query_values, + 2, + false, + Some((Some(1.0), None)), + ) + .await; } #[tokio::test] - async fn test_batch_knn_uses_index_when_available() { + async fn test_batch_knn_indexed() { let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) .await .unwrap(); test_ds.make_vector_index().await.unwrap(); let dataset = &test_ds.dataset; - - let query_values = (32..96).map(|v| v as f32).collect::>(); - let queries = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) - .unwrap(); + let (queries, query_values) = batch_knn_two_queries(); let mut scan = dataset.scan(); scan.nearest("vec", &queries, 2).unwrap(); @@ -5954,78 +5976,31 @@ mod test { ); let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 4); assert_eq!( batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0, 1, 1] ); - } - #[tokio::test] - async fn test_batch_knn_indexed_respects_distance_range() { - let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + let batch = dataset + .scan() + .nearest("vec", &queries, 2) + .unwrap() + .distance_range(Some(1.0), None) + .project(&["i"]) + .unwrap() + .try_into_batch() .await .unwrap(); - test_ds.make_vector_index().await.unwrap(); - let dataset = &test_ds.dataset; - - let query_values = (32..96).map(|v| v as f32).collect::>(); - let queries = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values.clone()), 32) - .unwrap(); - - let mut scan = dataset.scan(); - scan.nearest("vec", &queries, 2).unwrap(); - scan.distance_range(Some(1.0), None); - scan.project(&["i"]).unwrap(); - - let plan = scan.explain_plan(false).await.unwrap(); - assert!( - plan.contains("ANNSubIndex"), - "indexed batch KNN should use the vector index, got:\n{}", - plan - ); - - let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 4); - assert_eq!( - batch[QUERY_INDEX_COL].as_primitive::().values(), - &[0, 0, 1, 1] - ); - - for query_index in 0..2 { - let query = - Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); - let single = dataset - .scan() - .nearest("vec", &query, 2) - .unwrap() - .distance_range(Some(1.0), None) - .project(&["i"]) - .unwrap() - .try_into_batch() - .await - .unwrap(); - let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); - let mask = BooleanArray::from_iter( - query_indices - .iter() - .map(|value| value.map(|value| value == query_index as u32)), - ); - let batch_slice = arrow::compute::filter_record_batch(&batch, &mask).unwrap(); - assert_eq!( - batch_slice["i"].as_primitive::().values(), - single["i"].as_primitive::().values() - ); - assert_eq!( - batch_slice[DIST_COL].as_primitive::().values(), - single[DIST_COL].as_primitive::().values() - ); - } + assert_batch_matches_single_queries( + dataset, + &batch, + &query_values, + 2, + true, + Some((Some(1.0), None)), + ) + .await; } - - #[rstest] - #[tokio::test] async fn test_can_project_distance() { let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) .await @@ -10133,76 +10108,6 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") assert_eq!(fast_rows, 0); } - #[tokio::test] - async fn test_batch_single_vector_list_query_includes_query_index() { - let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) - .await - .unwrap(); - let dataset = &test_ds.dataset; - - let query_values = (32..64).map(|v| v as f32).collect::>(); - let queries = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values), 32).unwrap(); - - let mut scan = dataset.scan(); - scan.nearest("vec", &queries, 2).unwrap().use_index(false); - scan.project(&["i"]).unwrap(); - - let plan = scan.explain_plan(false).await.unwrap(); - assert!( - plan.contains("KNNVectorDistance: queries=1"), - "single-vector batch query should use batch KNN path, got:\n{}", - plan - ); - assert!( - !plan.contains("SortExec: TopK(fetch="), - "batch KNN must not apply per-query SortExec top-k, got:\n{}", - plan - ); - - let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 2); - assert!( - batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), - "batch-shaped query with one vector should still return query_index" - ); - assert_eq!( - batch[QUERY_INDEX_COL].as_primitive::().values(), - &[0, 0] - ); - } - - #[tokio::test] - async fn test_batch_flat_plan_returns_m_times_k_without_sort_topk() { - let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) - .await - .unwrap(); - let dataset = &test_ds.dataset; - - let query_values = (32..96).map(|v| v as f32).collect::>(); - let queries = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values), 32).unwrap(); - - let mut scan = dataset.scan(); - scan.nearest("vec", &queries, 2).unwrap().use_index(false); - scan.project(&["i"]).unwrap(); - - let plan = scan.explain_plan(false).await.unwrap(); - assert!( - plan.contains("KNNVectorDistance: queries=2"), - "expected batch flat KNN plan, got:\n{}", - plan - ); - assert!( - !plan.contains("SortExec: TopK(fetch="), - "batch flat KNN must not truncate to k rows globally, got:\n{}", - plan - ); - - let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.num_rows(), 4); - } - #[tokio::test] async fn test_batch_fast_search_without_index_returns_empty_with_query_index() { let dataset = TestVectorDataset::new(LanceFileVersion::Stable, true) From fc0e7f0ce7f4b5938ac62505cf1858eafe10fae3 Mon Sep 17 00:00:00 2001 From: zoey Date: Fri, 22 May 2026 15:48:37 +0800 Subject: [PATCH 14/32] fix: align batch query branch with main Keep the main-based batch vector query branch compiling cleanly after conflict resolution. Co-authored-by: Cursor --- rust/lance/src/dataset/scanner.rs | 10 +++++++--- rust/lance/src/io/exec/knn.rs | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index b1692dbf5b2..04ddc986c38 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -5884,9 +5884,11 @@ mod test { assert_batch_matches_single_queries(dataset, &batch, &query_values, k, false, None).await; let query_values_one = (32..64).map(|v| v as f32).collect::>(); - let queries_one = - FixedSizeListArray::try_new_from_values(Float32Array::from(query_values_one.clone()), 32) - .unwrap(); + let queries_one = FixedSizeListArray::try_new_from_values( + Float32Array::from(query_values_one.clone()), + 32, + ) + .unwrap(); let mut scan = dataset.scan(); scan.nearest("vec", &queries_one, k).unwrap(); scan.use_index(false); @@ -6001,6 +6003,8 @@ mod test { ) .await; } + + #[tokio::test] async fn test_can_project_distance() { let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) .await diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 83d12f0240f..167ff0a4298 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -67,7 +67,7 @@ use crate::{Error, Result}; use lance_arrow::*; use super::utils::{ - FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedChildInputStream, PreFilterSource, + FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedRecordBatchStreamAdapter, PreFilterSource, SelectionVectorToPrefilter, }; From e14e04a001bb2a999405055b6c2e19a32641ea08 Mon Sep 17 00:00:00 2001 From: zoey Date: Fri, 22 May 2026 19:32:41 +0800 Subject: [PATCH 15/32] fix: address latest batch KNN review feedback Add a regression test that batch flat KNN returns k rows per query instead of being truncated by SortExec, and keep query_index autoprojection separate from scoring-column autoprojection. Co-authored-by: Cursor --- rust/lance/src/dataset/scanner.rs | 53 ++++++++++++++++++++++++++++--- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 04ddc986c38..57e8bf58102 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1957,11 +1957,6 @@ impl Scanner { let vector_expr = expressions::col(DIST_COL, current_schema)?; output_expr.push((vector_expr, DIST_COL.to_string())); } - if self.is_batch_nearest && output_expr.iter().all(|(_, name)| name != QUERY_INDEX_COL) - { - let query_index_expr = expressions::col(QUERY_INDEX_COL, current_schema)?; - output_expr.push((query_index_expr, QUERY_INDEX_COL.to_string())); - } if self.full_text_query.is_some() && output_expr.iter().all(|(_, name)| name != SCORE_COL) { @@ -1975,6 +1970,13 @@ impl Scanner { } } + // Batch nearest queries always expose query_index when the caller requested an + // explicit projection but omitted this discriminator column. + if self.is_batch_nearest && output_expr.iter().all(|(_, name)| name != QUERY_INDEX_COL) { + let query_index_expr = expressions::col(QUERY_INDEX_COL, current_schema)?; + output_expr.push((query_index_expr, QUERY_INDEX_COL.to_string())); + } + if self.legacy_with_row_id { let row_id_pos = output_expr .iter() @@ -5917,6 +5919,47 @@ mod test { ); } + #[tokio::test] + async fn test_batch_knn_flat_returns_top_k_per_query() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let dataset = &test_ds.dataset; + let k = 10; + let (queries, _) = batch_knn_two_queries(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &queries, k).unwrap(); + scan.use_index(false); + scan.project(&["i"]).unwrap(); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + !plan.contains("SortExec: TopK(fetch="), + "batch flat KNN must not apply global SortExec top-k truncation, got:\n{}", + plan + ); + + let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + 2 * k, + "batch flat KNN must return k rows per query vector" + ); + + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + for query_index in 0..2 { + let rows_for_query = query_indices + .iter() + .filter(|value| *value == Some(query_index as u32)) + .count(); + assert_eq!( + rows_for_query, k, + "query_index {query_index} should have exactly {k} rows" + ); + } + } + #[tokio::test] async fn test_batch_knn_flat_respects_distance_range() { let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) From 05a08242fa5aacc6dceeaecaf8acf882c77af90f Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 18:41:27 +0800 Subject: [PATCH 16/32] fix: address batch KNN review threads and harden tests Stop cloning full input batches in flat batch KNN heaps, restrict batch detection to list-shaped queries, reject reserved query_index conflicts, and add regression tests for BubbleCal's projection and Codex feedback. Co-authored-by: Cursor --- java/src/main/java/org/lance/ipc/Query.java | 5 + python/python/lance/dataset.py | 35 ++-- python/python/tests/test_vector_index.py | 36 +++++ rust/lance/src/dataset/scanner.rs | 167 ++++++++++++++++---- rust/lance/src/io/exec/knn.rs | 58 +++++-- 5 files changed, 246 insertions(+), 55 deletions(-) diff --git a/java/src/main/java/org/lance/ipc/Query.java b/java/src/main/java/org/lance/ipc/Query.java index 3ad1301ee59..7e682c7fd1a 100644 --- a/java/src/main/java/org/lance/ipc/Query.java +++ b/java/src/main/java/org/lance/ipc/Query.java @@ -140,6 +140,11 @@ public Builder setColumn(String column) { /** * Sets the vector to be searched. * + *

This API accepts a single query vector. The array length must match the target + * vector column dimension. Batch nearest-neighbor search with multiple query vectors + * requires a list-shaped query input and is not available through this {@code float[]} + * entry point. + * * @param key The search vector. * @return The Builder instance for method chaining. */ diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index a4a028ed7f6..8ed60319bed 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -1099,11 +1099,14 @@ def scanner( "distance_range": (0.0, 1.0), } - ``q`` may also be a 2-D array-like value for fixed-size vector columns. - In that case Lance runs a batch nearest-neighbor query, returns up to - ``k`` rows for each query vector, and adds ``query_index`` to identify the - source query for each result row. When ``use_index`` is true and a vector - index is available, each query vector is searched through the index + ``q`` may also be a 2-D array-like value, or a list of vectors, for + fixed-size vector columns. In that case Lance runs a batch nearest-neighbor + query, returns up to ``k`` rows for each query vector, and adds + ``query_index`` to identify the source query for each result row. + Flattened 1-D arrays whose length is a multiple of the vector dimension are + rejected. Datasets that already contain a ``query_index`` column cannot be + used for batch nearest-neighbor search. When ``use_index`` is true and a + vector index is available, each query vector is searched through the index path; otherwise the flat batch path is used. batch_size: int, default None @@ -5998,10 +6001,13 @@ def nearest( ---------- q: QueryVectorLike A single query vector or, for fixed-size vector columns, a 2-D array-like - batch of query vectors. Batch queries return up to ``k`` rows per query - and include ``query_index`` in the output. When ``use_index`` is true and - a vector index is available, each query vector is searched through the - index path; otherwise the flat batch path is used. + or list-shaped batch of query vectors. Batch queries return up to ``k`` rows + per query and include ``query_index`` in the output. Flattened 1-D inputs + whose length is a multiple of the vector dimension are rejected. Datasets + with an existing ``query_index`` column cannot be used for batch search. + When ``use_index`` is true and a vector index is available, each query + vector is searched through the index path; otherwise the flat batch path + is used. query_parallelism: int, optional Maximum partition-search concurrency for a single vector query. The default is 0. Value 0 uses the automatic policy, which @@ -7151,10 +7157,13 @@ def _build_vector_search_query( The name of the vector column to search. q: QueryVectorLike The query vector. For fixed-size vector columns, this may be a 2-D - array-like batch of query vectors. Batch queries return up to ``k`` rows per - query vector and include ``query_index`` in the output. When ``use_index`` - is true and a vector index is available, each query vector is searched - through the index path; otherwise the flat batch path is used. + array-like or list-shaped batch of query vectors. Batch queries return up to + ``k`` rows per query vector and include ``query_index`` in the output. + Flattened 1-D inputs whose length is a multiple of the vector dimension are + rejected. Datasets with an existing ``query_index`` column cannot be used for + batch search. When ``use_index`` is true and a vector index is available, + each query vector is searched through the index path; otherwise the flat batch + path is used. k: int, optional The number of nearest neighbors to return. metric: str, optional diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 195cfa35de2..2a5b218f293 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -247,6 +247,42 @@ def _assert_batch_matches_single_queries(ds, queries, k, nearest_kwargs): ) +def test_batch_vector_search_rejects_dataset_query_index_column(tmp_path): + dim = 128 + table = create_table(nvec=80, ndim=dim) + table = table.append_column( + "query_index", + pa.array(range(80), type=pa.uint32()), + ) + ds = lance.write_dataset(table, tmp_path / "with_query_index") + + queries = np.random.randn(2, dim).astype(np.float32) + with pytest.raises(Exception, match="query_index"): + ds.to_table( + columns=["id", "query_index"], + nearest={ + "column": "vector", + "q": queries, + "k": 5, + "use_index": False, + }, + ) + + +def test_flat_1d_query_length_multiple_of_dim_is_rejected(dataset): + q = np.random.randn(256).astype(np.float32) + with pytest.raises(ValueError, match=r"256.*128"): + dataset.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": q, + "k": 5, + "use_index": False, + }, + ) + + def test_batch_flat_respects_distance_range(dataset): queries = np.random.randn(2, 128).astype(np.float32) _assert_batch_matches_single_queries( diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 57e8bf58102..6b3c4069a34 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1434,21 +1434,17 @@ impl Scanner { Ok(self) } - fn is_batch_nearest_query( - vector_type: &DataType, - query_type: &DataType, - query_len: usize, - dim: usize, - ) -> bool { - match vector_type { - DataType::FixedSizeList(_, _) => { - matches!( - query_type, - DataType::List(_) | DataType::FixedSizeList(_, _) - ) || (query_len > dim && query_len.is_multiple_of(dim)) - } - _ => false, - } + /// Returns true when `q` is a batch of single-vector queries. + /// + /// List-like queries against a [`DataType::List`] vector column are treated as one + /// multivector query. The same list-like query against a fixed-size vector column is + /// treated as a batch of single-vector queries. + fn is_batch_nearest_query(vector_type: &DataType, query_type: &DataType) -> bool { + matches!(vector_type, DataType::FixedSizeList(_, _)) + && matches!( + query_type, + DataType::List(_) | DataType::FixedSizeList(_, _) + ) } /// Find k-nearest neighbor within the vector column. @@ -1473,7 +1469,6 @@ impl Scanner { let (vector_type, element_type) = get_vector_type(self.dataset.schema(), column)?; let dim = get_vector_dim(self.dataset.schema(), column)?; let query_type = q.data_type().clone(); - let query_len = q.len(); let (q, query_count) = match &query_type { DataType::List(_) | DataType::FixedSizeList(_, _) => { @@ -1520,10 +1515,7 @@ impl Scanner { } } _ => { - if q.len() != dim - && (!matches!(vector_type, DataType::FixedSizeList(_, _)) - || !q.len().is_multiple_of(dim)) - { + if q.len() != dim { return Err(Error::invalid_input(format!( "query dim({}) doesn't match the column {} vector dim({})", q.len(), @@ -1531,8 +1523,7 @@ impl Scanner { dim, ))); } - let query_count = if q.len() == dim { 1 } else { q.len() / dim }; - (q.slice(0, q.len()), query_count) + (q.slice(0, q.len()), 1) } }; @@ -1542,6 +1533,14 @@ impl Scanner { )); } + if Self::is_batch_nearest_query(&vector_type, &query_type) + && self.dataset.schema().field(QUERY_INDEX_COL).is_some() + { + return Err(Error::invalid_input(format!( + "batch nearest neighbor search cannot be used on datasets with column '{QUERY_INDEX_COL}'" + ))); + } + let key = match &element_type { dt if dt == q.data_type() => q, dt if dt.is_floating() => coerce_float_vector( @@ -1574,8 +1573,7 @@ impl Scanner { dist_q_c: 0.0, }); self.nearest_query_count = query_count; - self.is_batch_nearest = - Self::is_batch_nearest_query(&vector_type, &query_type, query_len, dim); + self.is_batch_nearest = Self::is_batch_nearest_query(&vector_type, &query_type); Ok(self) } @@ -1970,8 +1968,8 @@ impl Scanner { } } - // Batch nearest queries always expose query_index when the caller requested an - // explicit projection but omitted this discriminator column. + // Batch nearest queries expose the synthetic `query_index` discriminator separately + // from scoring-column autoprojection (`_distance` / `_score` above). if self.is_batch_nearest && output_expr.iter().all(|(_, name)| name != QUERY_INDEX_COL) { let query_index_expr = expressions::col(QUERY_INDEX_COL, current_schema)?; output_expr.push((query_index_expr, QUERY_INDEX_COL.to_string())); @@ -5155,7 +5153,7 @@ mod test { use arrow_array::types::{Float32Type, UInt32Type, UInt64Type}; use arrow_array::{ ArrayRef, BooleanArray, FixedSizeListArray, Float16Array, Int32Array, LargeStringArray, - PrimitiveArray, RecordBatchIterator, StringArray, StructArray, UInt8Array, + PrimitiveArray, RecordBatchIterator, StringArray, StructArray, UInt8Array, UInt32Array, }; use arrow_ord::sort::sort_to_indices; @@ -5919,6 +5917,121 @@ mod test { ); } + #[tokio::test] + async fn test_primitive_query_length_multiple_of_dim_is_rejected() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let dataset = &test_ds.dataset; + let q: Float32Array = (32..96).map(|v| v as f32).collect(); + + let err = match dataset.scan().nearest("vec", &q, 2) { + Err(err) => err.to_string(), + Ok(_) => panic!("expected primitive query length mismatch error"), + }; + assert!( + err.contains("query dim(64) doesn't match the column vec vector dim(32)"), + "unexpected error: {err}" + ); + } + + async fn dataset_with_query_index_column() -> (TempStrDir, Dataset) { + let path = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i", DataType::Int32, true), + ArrowField::new( + "vec", + DataType::FixedSizeList( + Arc::new(ArrowField::new("item", DataType::Float32, true)), + 32, + ), + true, + ), + ArrowField::new(QUERY_INDEX_COL, DataType::UInt32, true), + ])); + let vector_values: Float32Array = (0..32 * 80).map(|v| v as f32).collect(); + let vectors = FixedSizeListArray::try_new_from_values(vector_values, 32).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..80)), + Arc::new(vectors), + Arc::new(UInt32Array::from_iter((0..80).map(|v| v as u32))), + ], + ) + .unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(std::iter::once(Ok(batch)), schema.clone()), + &path, + None, + ) + .await + .unwrap(); + (path, dataset) + } + + #[tokio::test] + async fn test_batch_knn_rejects_dataset_query_index_column() { + let (_tmp, dataset) = dataset_with_query_index_column().await; + let (queries, _) = batch_knn_two_queries(); + let err = match dataset.scan().nearest("vec", &queries, 2) { + Err(err) => err.to_string(), + Ok(_) => panic!("expected reserved query_index column error"), + }; + assert!(err.contains(QUERY_INDEX_COL), "unexpected error: {err}"); + } + + #[tokio::test] + async fn test_batch_knn_rejects_selecting_dataset_query_index_column() { + let (_tmp, dataset) = dataset_with_query_index_column().await; + let (queries, _) = batch_knn_two_queries(); + let err = match dataset.scan().nearest("vec", &queries, 2) { + Err(err) => err.to_string(), + Ok(scan) => match scan.project(&[QUERY_INDEX_COL]) { + Err(err) => err.to_string(), + Ok(scan) => match scan.try_into_batch().await { + Err(err) => err.to_string(), + Ok(_) => { + panic!("expected error when batch nearest selects query_index column") + } + }, + }, + }; + assert!(err.contains(QUERY_INDEX_COL), "unexpected error: {err}"); + } + + #[tokio::test] + async fn test_single_knn_projects_dataset_query_index_column() { + let (_tmp, dataset) = dataset_with_query_index_column().await; + let q: Float32Array = (32..64).map(|v| v as f32).collect(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &q, 2).unwrap(); + scan.use_index(false); + scan.project(&["i"]).unwrap(); + let without_query_index = scan.try_into_batch().await.unwrap(); + + let mut scan = dataset.scan(); + scan.nearest("vec", &q, 2).unwrap(); + scan.use_index(false); + scan.project(&["i", QUERY_INDEX_COL]).unwrap(); + let with_query_index = scan.try_into_batch().await.unwrap(); + + assert_eq!(without_query_index.num_rows(), 2); + assert_eq!( + without_query_index["i"] + .as_primitive::() + .values(), + with_query_index["i"].as_primitive::().values() + ); + assert_eq!( + with_query_index[QUERY_INDEX_COL] + .as_primitive::() + .null_count(), + 0 + ); + } + #[tokio::test] async fn test_batch_knn_flat_returns_top_k_per_query() { let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 167ff0a4298..4aae0c61ee8 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -3,7 +3,7 @@ use std::any::Any; use std::cmp::Ordering as CmpOrdering; -use std::collections::{BinaryHeap, HashMap, HashSet}; +use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, Mutex}; use std::time::Instant; @@ -277,8 +277,10 @@ impl KNNVectorDistanceExec { if input_schema.column_with_name(DIST_COL).is_some() { input_schema = input_schema.without_column(DIST_COL); } - if input_schema.column_with_name(QUERY_INDEX_COL).is_some() { - input_schema = input_schema.without_column(QUERY_INDEX_COL); + if is_batch && input_schema.column_with_name(QUERY_INDEX_COL).is_some() { + return Err(Error::invalid_input(format!( + "batch KNN cannot run when the input already contains reserved column '{QUERY_INDEX_COL}'" + ))); } let input_schema = Arc::new(input_schema); let mut output_schema = input_schema.as_ref().clone(); @@ -351,6 +353,7 @@ impl KNNVectorDistanceExec { .map(|_| BinaryHeap::::with_capacity(k)) .collect::>(); let mut input = input; + let mut source_batches = Vec::>::new(); while let Some(batch) = input.next().await { let batch = batch?; @@ -367,7 +370,9 @@ impl KNNVectorDistanceExec { })? .as_primitive::() .clone(); + let batch_index = source_batches.len() as u32; let batch = Arc::new(batch); + source_batches.push(batch.clone()); for (query_index, heap) in heaps.iter_mut().enumerate().take(query_count) { let key = query.slice(query_index * query_dim, query_dim); @@ -393,7 +398,7 @@ impl KNNVectorDistanceExec { query_index: query_index as u32, distance, row_id: row_ids.value(row_index), - batch: batch.clone(), + batch_index, row_index: row_index as u32, }; if heap.len() < k { @@ -426,19 +431,42 @@ impl KNNVectorDistanceExec { let mut query_indices = UInt32Builder::with_capacity(results.len()); let mut distances = Float32Builder::with_capacity(results.len()); - let mut row_batches = Vec::with_capacity(results.len()); - for result in results { + let mut row_batches = vec![None; results.len()]; + let mut groups: BTreeMap> = BTreeMap::new(); + for (out_idx, result) in results.iter().enumerate() { + groups + .entry(result.batch_index) + .or_default() + .push((out_idx, result.row_index)); + } + for (batch_index, entries) in groups { + let row_indices = + UInt32Array::from_iter(entries.iter().map(|(_, row_index)| *row_index)); + let taken = arrow_select::take::take_record_batch( + source_batches[batch_index as usize].as_ref(), + &row_indices, + ) + .map_err(|e| { + DataFusionError::ArrowError(Box::new(e), Some("take top-k rows".to_string())) + })?; + for (slice_idx, (out_idx, _)) in entries.iter().enumerate() { + row_batches[*out_idx] = Some(taken.slice(slice_idx, 1)); + } + } + for result in &results { query_indices.append_value(result.query_index); distances.append_value(result.distance); - let indices = UInt32Array::from(vec![result.row_index]); - row_batches.push( - arrow_select::take::take_record_batch(result.batch.as_ref(), &indices).map_err( - |e| { - DataFusionError::ArrowError(Box::new(e), Some("take top-k row".to_string())) - }, - )?, - ); } + let row_batches = row_batches + .into_iter() + .map(|batch| { + batch.ok_or_else(|| { + DataFusionError::Internal( + "missing materialized row for batch KNN result".to_string(), + ) + }) + }) + .collect::>>()?; let output = concat_batches(&input_schema, &row_batches) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; @@ -627,7 +655,7 @@ struct BatchKnnCandidate { query_index: u32, distance: f32, row_id: u64, - batch: Arc, + batch_index: u32, row_index: u32, } From 5184b3b21504350b5462359c3defbef75cd07a34 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 18:47:42 +0800 Subject: [PATCH 17/32] style(java): fix Query.java spotless javadoc wrapping Co-authored-by: Cursor --- java/src/main/java/org/lance/ipc/Query.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/java/src/main/java/org/lance/ipc/Query.java b/java/src/main/java/org/lance/ipc/Query.java index 7e682c7fd1a..48013b375ee 100644 --- a/java/src/main/java/org/lance/ipc/Query.java +++ b/java/src/main/java/org/lance/ipc/Query.java @@ -140,10 +140,9 @@ public Builder setColumn(String column) { /** * Sets the vector to be searched. * - *

This API accepts a single query vector. The array length must match the target - * vector column dimension. Batch nearest-neighbor search with multiple query vectors - * requires a list-shaped query input and is not available through this {@code float[]} - * entry point. + *

This API accepts a single query vector. The array length must match the target vector + * column dimension. Batch nearest-neighbor search with multiple query vectors requires a + * list-shaped query input and is not available through this {@code float[]} entry point. * * @param key The search vector. * @return The Builder instance for method chaining. From 20f70a15b65682b73dd4651e04923dff1013f770 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 19:44:57 +0800 Subject: [PATCH 18/32] fix: retain only heap-referenced batches in batch flat KNN Cache input RecordBatches only when a row enters a per-query top-k heap and prune unreferenced batches after each scan batch so memory stays bounded by O(query_count * k) batches instead of total scanned data. Co-authored-by: Cursor --- rust/lance/src/io/exec/knn.rs | 42 ++++++++++++++++++++++++++--------- 1 file changed, 31 insertions(+), 11 deletions(-) diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 4aae0c61ee8..083204ffbf6 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -353,7 +353,10 @@ impl KNNVectorDistanceExec { .map(|_| BinaryHeap::::with_capacity(k)) .collect::>(); let mut input = input; - let mut source_batches = Vec::>::new(); + // Retain only input batches that still have a row in a query heap (at most + // `query_count * k` batches), instead of every scanned batch. + let mut cached_batches: HashMap> = HashMap::new(); + let mut batch_counter: u32 = 0; while let Some(batch) = input.next().await { let batch = batch?; @@ -370,9 +373,9 @@ impl KNNVectorDistanceExec { })? .as_primitive::() .clone(); - let batch_index = source_batches.len() as u32; + let batch_index = batch_counter; + batch_counter += 1; let batch = Arc::new(batch); - source_batches.push(batch.clone()); for (query_index, heap) in heaps.iter_mut().enumerate().take(query_count) { let key = query.slice(query_index * query_dim, query_dim); @@ -401,17 +404,32 @@ impl KNNVectorDistanceExec { batch_index, row_index: row_index as u32, }; - if heap.len() < k { + let entered_heap = if heap.len() < k { heap.push(candidate); + true } else if heap .peek() .is_some_and(|worst| candidate.cmp(worst).is_lt()) { heap.pop(); heap.push(candidate); + true + } else { + false + }; + if entered_heap { + cached_batches + .entry(batch_index) + .or_insert_with(|| batch.clone()); } } } + + let referenced_batch_indices: HashSet = heaps + .iter() + .flat_map(|heap| heap.iter().map(|candidate| candidate.batch_index)) + .collect(); + cached_batches.retain(|batch_index, _| referenced_batch_indices.contains(batch_index)); } let mut results = heaps @@ -440,15 +458,17 @@ impl KNNVectorDistanceExec { .push((out_idx, result.row_index)); } for (batch_index, entries) in groups { + let source_batch = cached_batches.get(&batch_index).ok_or_else(|| { + DataFusionError::Internal(format!( + "batch KNN missing cached input batch for index {batch_index}" + )) + })?; let row_indices = UInt32Array::from_iter(entries.iter().map(|(_, row_index)| *row_index)); - let taken = arrow_select::take::take_record_batch( - source_batches[batch_index as usize].as_ref(), - &row_indices, - ) - .map_err(|e| { - DataFusionError::ArrowError(Box::new(e), Some("take top-k rows".to_string())) - })?; + let taken = arrow_select::take::take_record_batch(source_batch.as_ref(), &row_indices) + .map_err(|e| { + DataFusionError::ArrowError(Box::new(e), Some("take top-k rows".to_string())) + })?; for (slice_idx, (out_idx, _)) in entries.iter().enumerate() { row_batches[*out_idx] = Some(taken.slice(slice_idx, 1)); } From 2e8049271c56186341a1b73d3f3e9604f15b3433 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 19:55:19 +0800 Subject: [PATCH 19/32] ci: install cargo-deny binary directly in rust workflow Avoid flaky Docker builds for cargo-deny-action when GitHub release downloads fail under curl --silent. Co-authored-by: Cursor --- .github/workflows/rust.yml | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index e3c17671ce3..2063378f79d 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -78,10 +78,19 @@ jobs: runs-on: ubuntu-24.04 steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 - - uses: EmbarkStudios/cargo-deny-action@3fd3802e88374d3fe9159b834c7714ec57d6c979 # v2 - with: - log-level: warn - command: check + - name: Install cargo-deny + run: | + version=0.19.0 + arch="$(uname -m)" + case "${arch}" in + x86_64) target=x86_64-unknown-linux-musl ;; + aarch64) target=aarch64-unknown-linux-musl ;; + *) echo "unsupported arch: ${arch}" >&2; exit 1 ;; + esac + url="https://github.com/EmbarkStudios/cargo-deny/releases/download/${version}/cargo-deny-${version}-${target}.tar.gz" + curl -fsSL "${url}" | tar -xzf - -C /usr/local/bin --strip-components=1 + - name: Run cargo-deny + run: cargo deny check linux-build: runs-on: "ubuntu-24.04-8x" From 8645a8610c30b987dfa311344152815bae7d1b64 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 20:19:58 +0800 Subject: [PATCH 20/32] fix: align batch and single flat KNN distance filtering Share flat KNN distance validity rules (drop null/NaN, keep infinity), apply the same _distance not-null filter on batch output, and add a regression test for infinite distance handling. Co-authored-by: Cursor --- rust/lance/src/dataset/scanner.rs | 15 +++++++--- rust/lance/src/io/exec/knn.rs | 50 ++++++++++++++++++++++++------- 2 files changed, 51 insertions(+), 14 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 6b3c4069a34..66bc92ad1f0 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4471,7 +4471,7 @@ impl Scanner { }; if self.is_batch_nearest { - return Ok(knn_plan); + return Self::flat_knn_not_null_filter(knn_plan); } // Use DataFusion's [SortExec] for Top-K search @@ -4497,10 +4497,17 @@ impl Scanner { ) .with_fetch(Some(q.k)); - let logical_not_null = col(DIST_COL).is_not_null(); - let not_nulls = Arc::new(LanceFilterExec::try_new(logical_not_null, Arc::new(sort))?); + Self::flat_knn_not_null_filter(Arc::new(sort)) + } - Ok(not_nulls) + fn flat_knn_not_null_filter( + knn_plan: Arc, + ) -> Result> { + let logical_not_null = col(DIST_COL).is_not_null(); + Ok(Arc::new(LanceFilterExec::try_new( + logical_not_null, + knn_plan, + )?)) } fn get_fragments_as_bitmap(&self) -> RoaringBitmap { diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 083204ffbf6..527e1dff9c3 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -73,6 +73,24 @@ use super::utils::{ pub const QUERY_INDEX_COL: &str = "query_index"; +/// Whether a computed flat KNN distance should participate in top-k selection. +/// +/// Matches the single-query flat path: keep infinite distances and drop only +/// null/NaN values. +fn flat_knn_distance_is_candidate( + distances: &arrow::array::PrimitiveArray, + row_index: usize, +) -> Option { + if !distances.is_valid(row_index) { + return None; + } + let distance = distances.value(row_index); + if distance.is_nan() { + return None; + } + Some(distance) +} + pub struct AnnPartitionMetrics { index_metrics: IndexMetrics, partitions_ranked: Count, @@ -385,11 +403,8 @@ impl KNNVectorDistanceExec { .map_err(|e| DataFusionError::External(Box::new(e)))?; let distances = with_distances[DIST_COL].as_primitive::(); for row_index in 0..batch.num_rows() { - if !distances.is_valid(row_index) { - continue; - } - let distance = distances.value(row_index); - if !distance.is_finite() { + let Some(distance) = flat_knn_distance_is_candidate(distances, row_index) + else { continue; }; if lower_bound.is_some_and(|lower_bound| distance < lower_bound) @@ -591,11 +606,9 @@ impl ExecutionPlan for KNNVectorDistanceExec { .map_err(|e| DataFusionError::External(Box::new(e)))?; let distances = batch[DIST_COL].as_primitive::(); - let mask = BooleanArray::from_iter( - distances - .iter() - .map(|v| Some(v.map(|v| !v.is_nan()).unwrap_or(false))), - ); + let mask = BooleanArray::from_iter((0..distances.len()).map(|row_index| { + Some(flat_knn_distance_is_candidate(distances, row_index).is_some()) + })); arrow::compute::filter_record_batch(&batch, &mask) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) } @@ -2610,6 +2623,23 @@ mod tests { assert_eq!(expected, results[0]); } + #[test] + fn flat_knn_distance_keeps_infinity() { + let distances = Float32Array::from(vec![ + Some(f32::NAN), + Some(f32::INFINITY), + None, + Some(1.0), + ]); + assert!(flat_knn_distance_is_candidate(&distances, 0).is_none()); + assert_eq!( + flat_knn_distance_is_candidate(&distances, 1), + Some(f32::INFINITY) + ); + assert!(flat_knn_distance_is_candidate(&distances, 2).is_none()); + assert_eq!(flat_knn_distance_is_candidate(&distances, 3), Some(1.0)); + } + #[test] fn test_create_knn_flat() { let dim: usize = 128; From d28627e5359e99bbd4be3eab7529a811105fa298 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 20:27:33 +0800 Subject: [PATCH 21/32] style: apply rustfmt to flat_knn_distance_keeps_infinity test Co-authored-by: Cursor --- rust/lance/src/io/exec/knn.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 527e1dff9c3..682cc6a3970 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -2625,12 +2625,8 @@ mod tests { #[test] fn flat_knn_distance_keeps_infinity() { - let distances = Float32Array::from(vec![ - Some(f32::NAN), - Some(f32::INFINITY), - None, - Some(1.0), - ]); + let distances = + Float32Array::from(vec![Some(f32::NAN), Some(f32::INFINITY), None, Some(1.0)]); assert!(flat_knn_distance_is_candidate(&distances, 0).is_none()); assert_eq!( flat_knn_distance_is_candidate(&distances, 1), From 63ab0aca2b1803668c9c7cb453b505fee4730e00 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 21:36:58 +0800 Subject: [PATCH 22/32] ci: install cargo-deny into user-writable bin directory Co-authored-by: Cursor --- .github/workflows/rust.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 2063378f79d..3103af2e4a1 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -88,7 +88,10 @@ jobs: *) echo "unsupported arch: ${arch}" >&2; exit 1 ;; esac url="https://github.com/EmbarkStudios/cargo-deny/releases/download/${version}/cargo-deny-${version}-${target}.tar.gz" - curl -fsSL "${url}" | tar -xzf - -C /usr/local/bin --strip-components=1 + install_dir="${HOME}/.local/bin" + mkdir -p "${install_dir}" + curl -fsSL "${url}" | tar -xzf - -C "${install_dir}" --strip-components=1 + echo "${install_dir}" >> "${GITHUB_PATH}" - name: Run cargo-deny run: cargo deny check From f3ffa4aeb1df870c2fbdbf062f12c126833061de Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 21:39:24 +0800 Subject: [PATCH 23/32] revert: restore main cargo-deny workflow in rust.yml Keep this PR focused on batch KNN changes only; the prior direct-install workflow edits are unrelated to the feature. Co-authored-by: Cursor --- .github/workflows/rust.yml | 20 ++++---------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 3103af2e4a1..e3c17671ce3 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -78,22 +78,10 @@ jobs: runs-on: ubuntu-24.04 steps: - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 - - name: Install cargo-deny - run: | - version=0.19.0 - arch="$(uname -m)" - case "${arch}" in - x86_64) target=x86_64-unknown-linux-musl ;; - aarch64) target=aarch64-unknown-linux-musl ;; - *) echo "unsupported arch: ${arch}" >&2; exit 1 ;; - esac - url="https://github.com/EmbarkStudios/cargo-deny/releases/download/${version}/cargo-deny-${version}-${target}.tar.gz" - install_dir="${HOME}/.local/bin" - mkdir -p "${install_dir}" - curl -fsSL "${url}" | tar -xzf - -C "${install_dir}" --strip-components=1 - echo "${install_dir}" >> "${GITHUB_PATH}" - - name: Run cargo-deny - run: cargo deny check + - uses: EmbarkStudios/cargo-deny-action@3fd3802e88374d3fe9159b834c7714ec57d6c979 # v2 + with: + log-level: warn + command: check linux-build: runs-on: "ubuntu-24.04-8x" From 196db0f04bbfcba68bb389f4c6b8bbc70afe04c9 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 22:17:01 +0800 Subject: [PATCH 24/32] test: stabilize delta-index HNSW recall in append test Use refine(4) so ANN over-fetches candidates before flat re-ranking, avoiding flaky mac-build failures when HNSW returns approximate neighbors. Co-authored-by: Cursor --- rust/lance/src/index/append.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 16b17752ef4..ed5fc4bd509 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -979,7 +979,8 @@ mod tests { .unwrap() .nearest("vector", array.value(0).as_primitive::(), 2) .unwrap() - .refine(1) + // Over-fetch ANN candidates so delta-index HNSW recall is stable in CI. + .refine(4) .try_into_batch() .await .unwrap(); From 10378135253657df61c21eb14b780dec33f23bd7 Mon Sep 17 00:00:00 2001 From: zoey Date: Mon, 25 May 2026 22:24:30 +0800 Subject: [PATCH 25/32] revert: drop out-of-scope append.rs CI test tweak The mac-build failure was a flaky single-query HNSW test unrelated to batch KNN; keep append.rs unchanged from the batch feature commits. Co-authored-by: Cursor --- rust/lance/src/index/append.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index ed5fc4bd509..16b17752ef4 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -979,8 +979,7 @@ mod tests { .unwrap() .nearest("vector", array.value(0).as_primitive::(), 2) .unwrap() - // Over-fetch ANN candidates so delta-index HNSW recall is stable in CI. - .refine(4) + .refine(1) .try_into_batch() .await .unwrap(); From 079ac74b366817787223487fab4de28c08db3a46 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Mon, 25 May 2026 23:15:24 +0800 Subject: [PATCH 26/32] fix(vector): simplify batch knn handling --- java/lance-jni/src/blocking_scanner.rs | 4 +- .../test/java/org/lance/VectorSearchTest.java | 60 +++---- rust/lance/src/dataset/scanner.rs | 27 +-- rust/lance/src/io/exec/knn.rs | 162 ++++++++---------- 4 files changed, 111 insertions(+), 142 deletions(-) diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index b77c170d497..1a8249fdbcf 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -317,7 +317,9 @@ pub(crate) fn build_scanner_with_options<'a>( let key_array = env.get_vec_f32_from_method(&java_obj, "getKey")?; let key = Float32Array::from(key_array); let k = env.get_int_as_usize_from_method(&java_obj, "getK")?; - let _ = scanner.nearest(&column, &key, k); + scanner + .nearest(&column, &key, k) + .map_err(|err| Error::input_error(err.to_string()))?; let minimum_nprobes = env.get_int_as_usize_from_method(&java_obj, "getMinimumNprobes")?; scanner.minimum_nprobes(minimum_nprobes); diff --git a/java/src/test/java/org/lance/VectorSearchTest.java b/java/src/test/java/org/lance/VectorSearchTest.java index 0a34640da7e..8a82ecb2849 100644 --- a/java/src/test/java/org/lance/VectorSearchTest.java +++ b/java/src/test/java/org/lance/VectorSearchTest.java @@ -63,35 +63,37 @@ void test_create_index() throws Exception { } } - // rust/lance-linalg/src/distance/l2.rs:256:5: - // 5assertion `left == right` failed - // Directly panic instead of throwing an exception - // @Test - // void search_invalid_vector() throws Exception { - // try (TestVectorDataset testVectorDataset = new - // TestVectorDataset(tempDir.resolve("test_create_index"))) { - // try (Dataset dataset = testVectorDataset.create()) { - // float[] key = new float[30]; - // for (int i = 0; i < 30; i++) { - // key[i] = (float) (i + 30); - // } - // ScanOptions options = new ScanOptions.Builder() - // .nearest(new Query.Builder() - // .setColumn(TestVectorDataset.vectorColumnName) - // .setKey(key) - // .setK(5) - // .setUseIndex(false) - // .build()) - // .build(); - // assertThrows(IllegalArgumentException.class, () -> { - // try (Scanner scanner = dataset.newScan(options)) { - // try (ArrowReader reader = scanner.scanBatches()) { - // } - // } - // }); - // } - // } - // } + @Test + void search_invalid_vector() throws Exception { + try (TestVectorDataset testVectorDataset = + new TestVectorDataset(tempDir.resolve("search_invalid_vector"))) { + try (Dataset dataset = testVectorDataset.create()) { + float[] key = new float[30]; + for (int i = 0; i < 30; i++) { + key[i] = (float) (i + 30); + } + ScanOptions options = + new ScanOptions.Builder() + .nearest( + new Query.Builder() + .setColumn(TestVectorDataset.vectorColumnName) + .setKey(key) + .setK(5) + .setUseIndex(false) + .build()) + .build(); + assertThrows( + IllegalArgumentException.class, + () -> { + try (Scanner scanner = dataset.newScan(options)) { + try (ArrowReader reader = scanner.scanBatches()) { + reader.loadNextBatch(); + } + } + }); + } + } + } @ParameterizedTest @ValueSource(booleans = {false, true}) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 66bc92ad1f0..edabfa18402 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4429,6 +4429,10 @@ impl Scanner { }, )?); + if self.is_batch_nearest { + return Ok(flat_dist); + } + let lower: Option<(Expr, Arc)> = q .lower_bound .map(|v| -> Result<(Expr, Arc)> { @@ -4470,10 +4474,6 @@ impl Scanner { flat_dist }; - if self.is_batch_nearest { - return Self::flat_knn_not_null_filter(knn_plan); - } - // Use DataFusion's [SortExec] for Top-K search let sort = SortExec::new( [ @@ -5988,25 +5988,6 @@ mod test { assert!(err.contains(QUERY_INDEX_COL), "unexpected error: {err}"); } - #[tokio::test] - async fn test_batch_knn_rejects_selecting_dataset_query_index_column() { - let (_tmp, dataset) = dataset_with_query_index_column().await; - let (queries, _) = batch_knn_two_queries(); - let err = match dataset.scan().nearest("vec", &queries, 2) { - Err(err) => err.to_string(), - Ok(scan) => match scan.project(&[QUERY_INDEX_COL]) { - Err(err) => err.to_string(), - Ok(scan) => match scan.try_into_batch().await { - Err(err) => err.to_string(), - Ok(_) => { - panic!("expected error when batch nearest selects query_index column") - } - }, - }, - }; - assert!(err.contains(QUERY_INDEX_COL), "unexpected error: {err}"); - } - #[tokio::test] async fn test_single_knn_projects_dataset_query_index_column() { let (_tmp, dataset) = dataset_with_query_index_column().await; diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 682cc6a3970..049db093cac 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -3,7 +3,7 @@ use std::any::Any; use std::cmp::Ordering as CmpOrdering; -use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet}; +use std::collections::{BinaryHeap, HashMap, HashSet}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, Mutex}; use std::time::Instant; @@ -67,8 +67,8 @@ use crate::{Error, Result}; use lance_arrow::*; use super::utils::{ - FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedRecordBatchStreamAdapter, PreFilterSource, - SelectionVectorToPrefilter, + FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedChildInputStream, + InstrumentedRecordBatchStreamAdapter, PreFilterSource, SelectionVectorToPrefilter, }; pub const QUERY_INDEX_COL: &str = "query_index"; @@ -371,10 +371,6 @@ impl KNNVectorDistanceExec { .map(|_| BinaryHeap::::with_capacity(k)) .collect::>(); let mut input = input; - // Retain only input batches that still have a row in a query heap (at most - // `query_count * k` batches), instead of every scanned batch. - let mut cached_batches: HashMap> = HashMap::new(); - let mut batch_counter: u32 = 0; while let Some(batch) = input.next().await { let batch = batch?; @@ -391,16 +387,12 @@ impl KNNVectorDistanceExec { })? .as_primitive::() .clone(); - let batch_index = batch_counter; - batch_counter += 1; - let batch = Arc::new(batch); for (query_index, heap) in heaps.iter_mut().enumerate().take(query_count) { let key = query.slice(query_index * query_dim, query_dim); - let with_distances = - compute_distance(key, distance_type, &column, batch.as_ref().clone()) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; + let with_distances = compute_distance(key, distance_type, &column, batch.clone()) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; let distances = with_distances[DIST_COL].as_primitive::(); for row_index in 0..batch.num_rows() { let Some(distance) = flat_knn_distance_is_candidate(distances, row_index) @@ -412,39 +404,37 @@ impl KNNVectorDistanceExec { { continue; } - let candidate = BatchKnnCandidate { - query_index: query_index as u32, - distance, - row_id: row_ids.value(row_index), - batch_index, - row_index: row_index as u32, + let query_index = query_index as u32; + let row_id = row_ids.value(row_index); + let row_index = row_index as u32; + let candidate_is_better = |worst: &BatchKnnCandidate| { + distance + .total_cmp(&worst.distance) + .then_with(|| row_id.cmp(&worst.row_id)) + .then_with(|| query_index.cmp(&worst.query_index)) + .then_with(|| row_index.cmp(&worst.row_index)) + .is_lt() }; - let entered_heap = if heap.len() < k { - heap.push(candidate); - true - } else if heap - .peek() - .is_some_and(|worst| candidate.cmp(worst).is_lt()) - { + if heap.len() < k { + heap.push(BatchKnnCandidate { + query_index, + distance, + row_id, + batch: batch.clone(), + row_index, + }); + } else if heap.peek().is_some_and(candidate_is_better) { heap.pop(); - heap.push(candidate); - true - } else { - false - }; - if entered_heap { - cached_batches - .entry(batch_index) - .or_insert_with(|| batch.clone()); + heap.push(BatchKnnCandidate { + query_index, + distance, + row_id, + batch: batch.clone(), + row_index, + }); } } } - - let referenced_batch_indices: HashSet = heaps - .iter() - .flat_map(|heap| heap.iter().map(|candidate| candidate.batch_index)) - .collect(); - cached_batches.retain(|batch_index, _| referenced_batch_indices.contains(batch_index)); } let mut results = heaps @@ -464,44 +454,17 @@ impl KNNVectorDistanceExec { let mut query_indices = UInt32Builder::with_capacity(results.len()); let mut distances = Float32Builder::with_capacity(results.len()); - let mut row_batches = vec![None; results.len()]; - let mut groups: BTreeMap> = BTreeMap::new(); - for (out_idx, result) in results.iter().enumerate() { - groups - .entry(result.batch_index) - .or_default() - .push((out_idx, result.row_index)); - } - for (batch_index, entries) in groups { - let source_batch = cached_batches.get(&batch_index).ok_or_else(|| { - DataFusionError::Internal(format!( - "batch KNN missing cached input batch for index {batch_index}" - )) - })?; - let row_indices = - UInt32Array::from_iter(entries.iter().map(|(_, row_index)| *row_index)); - let taken = arrow_select::take::take_record_batch(source_batch.as_ref(), &row_indices) - .map_err(|e| { - DataFusionError::ArrowError(Box::new(e), Some("take top-k rows".to_string())) - })?; - for (slice_idx, (out_idx, _)) in entries.iter().enumerate() { - row_batches[*out_idx] = Some(taken.slice(slice_idx, 1)); - } - } - for result in &results { + let mut row_batches = Vec::with_capacity(results.len()); + for result in results { query_indices.append_value(result.query_index); distances.append_value(result.distance); + let indices = UInt32Array::from(vec![result.row_index]); + row_batches.push( + arrow_select::take::take_record_batch(&result.batch, &indices).map_err(|e| { + DataFusionError::ArrowError(Box::new(e), Some("take top-k row".to_string())) + })?, + ); } - let row_batches = row_batches - .into_iter() - .map(|batch| { - batch.ok_or_else(|| { - DataFusionError::Internal( - "missing materialized row for batch KNN result".to_string(), - ) - }) - }) - .collect::>>()?; let output = concat_batches(&input_schema, &row_batches) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; @@ -592,18 +555,42 @@ impl ExecutionPlan for KNNVectorDistanceExec { &self.metrics, )) as SendableRecordBatchStream); } + let input_schema = self.input.schema(); let key = self.query.clone(); let column = self.column.clone(); let dt = self.distance_type; - let stream = input_stream - .try_filter(|batch| future::ready(batch.num_rows() > 0)) - .map(move |batch| { + let schema = self.schema(); + + // Empty batches don't have a vector column to score; filter them out + // before reaching the helper so the transform always sees real work. + let filtered_input = Box::pin(RecordBatchStreamAdapter::new( + input_schema, + input_stream.try_filter(|batch| future::ready(batch.num_rows() > 0)), + )) as SendableRecordBatchStream; + + // Mirror of the helper's elapsed_compute counter; used to attribute + // wall-clock from the spawn_blocking distance kernel back onto the + // node's `elapsed_compute` metric. + let elapsed_compute = BaselineMetrics::new(&self.metrics, partition) + .elapsed_compute() + .clone(); + + let stream = InstrumentedChildInputStream::new( + filtered_input, + schema, + move |batch| { let key = key.clone(); let column = column.clone(); + let elapsed_compute = elapsed_compute.clone(); async move { - let batch = compute_distance(key, dt, &column, batch?) + // Time around the .await to capture the spawn_blocking + // distance work, which otherwise runs while this future is + // Pending and is missed by the helper's own poll timer. + let start = Instant::now(); + let batch = compute_distance(key, dt, &column, batch) .await .map_err(|e| DataFusionError::External(Box::new(e)))?; + elapsed_compute.add_duration(start.elapsed()); let distances = batch[DIST_COL].as_primitive::(); let mask = BooleanArray::from_iter((0..distances.len()).map(|row_index| { @@ -612,15 +599,12 @@ impl ExecutionPlan for KNNVectorDistanceExec { arrow::compute::filter_record_batch(&batch, &mask) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) } - }) - .buffer_unordered(get_num_compute_intensive_cpus()); - let schema = self.schema(); - Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new( - schema, - stream.boxed(), + }, + get_num_compute_intensive_cpus(), partition, &self.metrics, - )) as SendableRecordBatchStream) + ); + Ok(Box::pin(stream) as SendableRecordBatchStream) } fn partition_statistics(&self, partition: Option) -> DataFusionResult { @@ -688,7 +672,7 @@ struct BatchKnnCandidate { query_index: u32, distance: f32, row_id: u64, - batch_index: u32, + batch: RecordBatch, row_index: u32, } From dd36022b3992eeacedd403af3fd4e1e311034c4c Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Mon, 25 May 2026 23:32:29 +0800 Subject: [PATCH 27/32] refactor(vector): trim batch knn checks and tests --- python/python/tests/test_vector_index.py | 20 -------- python/src/dataset.rs | 22 ++++----- rust/lance/src/dataset/scanner.rs | 63 ++++++------------------ 3 files changed, 25 insertions(+), 80 deletions(-) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 2a5b218f293..d4c606cbe44 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -283,26 +283,6 @@ def test_flat_1d_query_length_multiple_of_dim_is_rejected(dataset): ) -def test_batch_flat_respects_distance_range(dataset): - queries = np.random.randn(2, 128).astype(np.float32) - _assert_batch_matches_single_queries( - dataset, - queries, - k=5, - nearest_kwargs={"use_index": False, "distance_range": (0.0, 50.0)}, - ) - - -def test_batch_indexed_respects_distance_range(indexed_dataset): - queries = np.random.randn(2, 128).astype(np.float32) - _assert_batch_matches_single_queries( - indexed_dataset, - queries, - k=5, - nearest_kwargs={"distance_range": (0.0, 50.0)}, - ) - - def test_batch_fast_search_without_index_returns_empty_with_query_index(dataset): queries = np.random.randn(2, 128).astype(np.float32) batch = dataset.to_table( diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 82947e11902..62ef1670518 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1368,26 +1368,22 @@ impl Dataset { query_parallelism, ) = vector_query_params_from_dict(nearest, default_k)?; - let (vector_type, element_type) = get_vector_type(self_.ds.schema(), &column) + let (_, element_type) = get_vector_type(self_.ds.schema(), &column) .map_err(|e| PyValueError::new_err(e.to_string()))?; - let is_batch_query = matches!( - q.data_type(), - DataType::List(_) | DataType::FixedSizeList(_, _) - ) && matches!(vector_type, DataType::FixedSizeList(_, _)); - let scanner = match (is_batch_query, element_type) { - (true, DataType::UInt8) => { - return Err(PyValueError::new_err( - "Batch nearest is not supported for binary vector columns", - )); - } - (false, DataType::UInt8) => { + let scanner = match element_type { + DataType::UInt8 + if !matches!( + q.data_type(), + DataType::List(_) | DataType::FixedSizeList(_, _) + ) => + { let q = arrow::compute::cast(&q, &DataType::UInt8).map_err(|e| { PyValueError::new_err(format!("Failed to cast q to binary vector: {}", e)) })?; let q = q.as_primitive::(); scanner.nearest(&column, q, k) } - (_, _) => scanner.nearest(&column, &q, k), + _ => scanner.nearest(&column, &q, k), }; let distance_range: Option<(Option, Option)> = if let Some(dr) = nearest.get_item("distance_range")? { diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index edabfa18402..cb61070de14 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1527,12 +1527,6 @@ impl Scanner { } }; - if query_count == 0 { - return Err(Error::invalid_input( - "Query vector must have non-zero length".to_string(), - )); - } - if Self::is_batch_nearest_query(&vector_type, &query_type) && self.dataset.schema().field(QUERY_INDEX_COL).is_some() { @@ -5884,10 +5878,26 @@ mod test { ); let batch = scan.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + 2 * k, + "batch flat KNN must return k rows per query vector" + ); assert_eq!( batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0, 1, 1] ); + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + for query_index in 0..2 { + let rows_for_query = query_indices + .iter() + .filter(|value| *value == Some(query_index as u32)) + .count(); + assert_eq!( + rows_for_query, k, + "query_index {query_index} should have exactly {k} rows" + ); + } assert_batch_matches_single_queries(dataset, &batch, &query_values, k, false, None).await; let query_values_one = (32..64).map(|v| v as f32).collect::>(); @@ -6020,47 +6030,6 @@ mod test { ); } - #[tokio::test] - async fn test_batch_knn_flat_returns_top_k_per_query() { - let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) - .await - .unwrap(); - let dataset = &test_ds.dataset; - let k = 10; - let (queries, _) = batch_knn_two_queries(); - - let mut scan = dataset.scan(); - scan.nearest("vec", &queries, k).unwrap(); - scan.use_index(false); - scan.project(&["i"]).unwrap(); - - let plan = scan.explain_plan(false).await.unwrap(); - assert!( - !plan.contains("SortExec: TopK(fetch="), - "batch flat KNN must not apply global SortExec top-k truncation, got:\n{}", - plan - ); - - let batch = scan.try_into_batch().await.unwrap(); - assert_eq!( - batch.num_rows(), - 2 * k, - "batch flat KNN must return k rows per query vector" - ); - - let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); - for query_index in 0..2 { - let rows_for_query = query_indices - .iter() - .filter(|value| *value == Some(query_index as u32)) - .count(); - assert_eq!( - rows_for_query, k, - "query_index {query_index} should have exactly {k} rows" - ); - } - } - #[tokio::test] async fn test_batch_knn_flat_respects_distance_range() { let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) From 3b98b5380db0a95b8e4180dfc8e088f3be654ac5 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 26 May 2026 00:48:14 +0800 Subject: [PATCH 28/32] chore(vector): remove batch knn benchmark --- python/python/benchmarks/test_search.py | 105 ------------------------ 1 file changed, 105 deletions(-) diff --git a/python/python/benchmarks/test_search.py b/python/python/benchmarks/test_search.py index 3750bdf4396..61076e61687 100644 --- a/python/python/benchmarks/test_search.py +++ b/python/python/benchmarks/test_search.py @@ -13,11 +13,6 @@ N_DIMS = 768 NUM_ROWS = 100_000 NEW_ROWS = 10_000 -BATCH_FLAT_KNN_DIM = 512 -BATCH_FLAT_KNN_K = 10 -BATCH_FLAT_KNN_BATCH_SIZE = 10_000 -BATCH_FLAT_KNN_QUERY_COUNT = 10 -BATCH_FLAT_KNN_ROWS = 1_000_000 def find_or_clean(dataset_path: Path) -> Union[lance.LanceDataset, None]: @@ -69,42 +64,6 @@ def create_table(num_rows, offset) -> pa.Table: ) -def create_flat_vector_table(num_rows: int, offset: int, dim: int) -> pa.Table: - rng = np.random.default_rng(seed=offset) - values = rng.random((num_rows, dim), dtype=np.float32) - vectors = pa.FixedSizeListArray.from_arrays(pa.array(values.ravel()), dim) - ids = pa.array(range(offset, offset + num_rows)) - return pa.table({"vector": vectors, "id": ids}) - - -def create_batch_flat_knn_dataset( - data_dir: Path, num_rows: int, batch_size: int, dim: int -) -> lance.LanceDataset: - tmp_path = data_dir / f"batch_flat_knn_{num_rows}_{batch_size}_{dim}" - dataset = find_or_clean(tmp_path) - if dataset: - return dataset - - rows_remaining = num_rows - offset = 0 - dataset = None - while rows_remaining > 0: - next_batch_length = min(rows_remaining, batch_size) - rows_remaining -= next_batch_length - table = create_flat_vector_table(next_batch_length, offset, dim) - if offset == 0: - dataset = lance.write_dataset( - table, tmp_path, data_storage_version="stable" - ) - else: - dataset = lance.write_dataset( - table, tmp_path, mode="append", data_storage_version="stable" - ) - offset += next_batch_length - - return dataset - - def create_base_dataset(data_dir: Path) -> lance.LanceDataset: tmp_path = data_dir / "search_dataset" dataset = find_or_clean(tmp_path) @@ -214,70 +173,6 @@ def test_knn_search(test_dataset, benchmark): assert result.num_rows > 0 -@pytest.mark.benchmark(group="batch_flat_knn") -@pytest.mark.parametrize("mode", ["separate", "batch"]) -@pytest.mark.parametrize( - ("dim", "num_rows", "batch_size", "query_count", "rounds"), - [ - ( - BATCH_FLAT_KNN_DIM, - BATCH_FLAT_KNN_ROWS, - BATCH_FLAT_KNN_BATCH_SIZE, - BATCH_FLAT_KNN_QUERY_COUNT, - 10, - ) - ], - ids=["1m_rows_512d_m10"], -) -def test_batch_flat_knn( - data_dir: Path, - benchmark, - mode: str, - dim: int, - num_rows: int, - batch_size: int, - query_count: int, - rounds: int, -): - dataset = create_batch_flat_knn_dataset(data_dir, num_rows, batch_size, dim) - query_table = dataset.to_table(columns=["vector"], limit=query_count) - query_values = np.asarray( - query_table["vector"].combine_chunks().values, dtype=np.float32 - ).reshape(query_count, dim) - - def separate_queries(): - total_rows = 0 - for query in query_values: - total_rows += dataset.to_table( - columns=[], - nearest={ - "column": "vector", - "q": query, - "k": BATCH_FLAT_KNN_K, - "use_index": False, - }, - ).num_rows - return total_rows - - def batch_query(): - return dataset.to_table( - columns=[], - nearest={ - "column": "vector", - "q": query_values, - "k": BATCH_FLAT_KNN_K, - "use_index": False, - }, - ).num_rows - - if mode == "separate": - result = benchmark.pedantic(separate_queries, rounds=rounds, iterations=1) - else: - result = benchmark.pedantic(batch_query, rounds=rounds, iterations=1) - - assert result == query_count * BATCH_FLAT_KNN_K - - @pytest.mark.benchmark(group="query_ann") def test_ann_no_refine(test_dataset, benchmark): q = pc.random(N_DIMS).cast(pa.float32()) From b6ae4afebfec6fec7997421544c7363a3f43a532 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 26 May 2026 01:06:16 +0800 Subject: [PATCH 29/32] fix(vector): align batch query index output --- python/python/lance/dataset.py | 13 +++-- python/python/tests/test_vector_index.py | 5 +- rust/lance/src/dataset/scanner.rs | 66 ++++++++++++++++-------- rust/lance/src/io/exec/knn.rs | 39 +++++++------- 4 files changed, 76 insertions(+), 47 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 18ad79b21a1..eeb2dacff6a 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -1102,7 +1102,8 @@ def scanner( ``q`` may also be a 2-D array-like value, or a list of vectors, for fixed-size vector columns. In that case Lance runs a batch nearest-neighbor query, returns up to ``k`` rows for each query vector, and adds - ``query_index`` to identify the source query for each result row. + an Int32 non-null ``query_index`` as the first output column to identify + the source query for each result row. Flattened 1-D arrays whose length is a multiple of the vector dimension are rejected. Datasets that already contain a ``query_index`` column cannot be used for batch nearest-neighbor search. When ``use_index`` is true and a @@ -6005,9 +6006,10 @@ def nearest( q: QueryVectorLike A single query vector or, for fixed-size vector columns, a 2-D array-like or list-shaped batch of query vectors. Batch queries return up to ``k`` rows - per query and include ``query_index`` in the output. Flattened 1-D inputs - whose length is a multiple of the vector dimension are rejected. Datasets - with an existing ``query_index`` column cannot be used for batch search. + per query and include Int32 non-null ``query_index`` as the first output + column. Flattened 1-D inputs whose length is a multiple of the vector + dimension are rejected. Datasets with an existing ``query_index`` column + cannot be used for batch search. When ``use_index`` is true and a vector index is available, each query vector is searched through the index path; otherwise the flat batch path is used. @@ -7161,7 +7163,8 @@ def _build_vector_search_query( q: QueryVectorLike The query vector. For fixed-size vector columns, this may be a 2-D array-like or list-shaped batch of query vectors. Batch queries return up to - ``k`` rows per query vector and include ``query_index`` in the output. + ``k`` rows per query vector and include Int32 non-null ``query_index`` as + the first output column. Flattened 1-D inputs whose length is a multiple of the vector dimension are rejected. Datasets with an existing ``query_index`` column cannot be used for batch search. When ``use_index`` is true and a vector index is available, diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index d4c606cbe44..6537dcb49b7 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -203,7 +203,10 @@ def test_batch_flat_query_matches_repeated_single_queries(dataset, queries): ) assert batch.num_rows == query_count * k - assert batch.column_names == ["id", "_distance", "query_index"] + assert batch.column_names == ["query_index", "id", "_distance"] + query_index_field = batch.schema.field("query_index") + assert query_index_field.type == pa.int32() + assert not query_index_field.nullable expected_query_index = sum([[i] * k for i in range(query_count)], []) assert batch["query_index"].to_pylist() == expected_query_index diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 0c59991aaed..dd03d83f4f8 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -101,7 +101,9 @@ use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; use crate::io::exec::{ AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNVectorDistanceExec, LancePushdownScanExec, LanceScanExec, Planner, PreFilterSource, ScanConfig, TakeExec, - knn::{KnnBatchParams, QUERY_INDEX_COL, knn_empty_result_schema, new_knn_exec}, + knn::{ + KnnBatchParams, QUERY_INDEX_COL, knn_empty_result_schema, new_knn_exec, query_index_field, + }, project, }; use crate::io::exec::{AddRowOffsetExec, LanceFilterExec, LanceScanConfig, get_physical_optimizer}; @@ -1899,7 +1901,7 @@ impl Scanner { if self.nearest.as_ref().is_some() { extra_columns.push(ArrowField::new(DIST_COL, DataType::Float32, true)); if self.is_batch_nearest { - extra_columns.push(ArrowField::new(QUERY_INDEX_COL, DataType::UInt32, true)); + extra_columns.push(query_index_field()); } }; @@ -1962,11 +1964,21 @@ impl Scanner { } } - // Batch nearest queries expose the synthetic `query_index` discriminator separately - // from scoring-column autoprojection (`_distance` / `_score` above). - if self.is_batch_nearest && output_expr.iter().all(|(_, name)| name != QUERY_INDEX_COL) { - let query_index_expr = expressions::col(QUERY_INDEX_COL, current_schema)?; - output_expr.push((query_index_expr, QUERY_INDEX_COL.to_string())); + // Batch nearest queries expose the synthetic `query_index` discriminator as + // the first output column for compatibility with LanceDB batch vector search. + if self.is_batch_nearest { + let query_index_expr = if let Some(pos) = output_expr + .iter() + .position(|(_, name)| name == QUERY_INDEX_COL) + { + output_expr.remove(pos) + } else { + ( + expressions::col(QUERY_INDEX_COL, current_schema)?, + QUERY_INDEX_COL.to_string(), + ) + }; + output_expr.insert(0, query_index_expr); } if self.legacy_with_row_id { @@ -3811,7 +3823,7 @@ impl Scanner { .await?; query_plans.push(Self::add_query_index_column( single_plan, - query_index as u32, + query_index as i32, )?); } @@ -3851,10 +3863,14 @@ impl Scanner { fn add_query_index_column( plan: Arc, - query_index: u32, + query_index: i32, ) -> Result> { let schema = plan.schema(); let mut projection_exprs = Vec::with_capacity(schema.fields().len() + 1); + projection_exprs.push(( + Arc::new(Literal::new(ScalarValue::Int32(Some(query_index)))) as Arc, + QUERY_INDEX_COL.to_string(), + )); for field in schema.fields() { projection_exprs.push(( Arc::new(Column::new_with_schema(field.name(), schema.as_ref())?) @@ -3862,10 +3878,6 @@ impl Scanner { field.name().clone(), )); } - projection_exprs.push(( - Arc::new(Literal::new(ScalarValue::UInt32(Some(query_index)))) as Arc, - QUERY_INDEX_COL.to_string(), - )); Ok(Arc::new(ProjectionExec::try_new(projection_exprs, plan)?)) } @@ -5833,11 +5845,11 @@ mod test { scan.project(&["i"]).unwrap(); let single = scan.try_into_batch().await.unwrap(); - let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); let mask = BooleanArray::from_iter( query_indices .iter() - .map(|value| value.map(|value| value == query_index as u32)), + .map(|value| value.map(|value| value == query_index as i32)), ); let batch_slice = arrow::compute::filter_record_batch(batch, &mask).unwrap(); assert_eq!( @@ -5883,20 +5895,23 @@ mod test { ); let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); + assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); + assert!(!batch.schema().field(0).is_nullable()); assert_eq!( batch.num_rows(), 2 * k, "batch flat KNN must return k rows per query vector" ); assert_eq!( - batch[QUERY_INDEX_COL].as_primitive::().values(), + batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0, 1, 1] ); - let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); + let query_indices = batch[QUERY_INDEX_COL].as_primitive::(); for query_index in 0..2 { let rows_for_query = query_indices .iter() - .filter(|value| *value == Some(query_index as u32)) + .filter(|value| *value == Some(query_index)) .count(); assert_eq!( rows_for_query, k, @@ -5933,8 +5948,11 @@ mod test { batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), "batch-shaped query with one vector should still return query_index" ); + assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); + assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); + assert!(!batch.schema().field(0).is_nullable()); assert_eq!( - batch[QUERY_INDEX_COL].as_primitive::().values(), + batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0] ); } @@ -6056,7 +6074,7 @@ mod test { .unwrap(); assert_eq!( - batch[QUERY_INDEX_COL].as_primitive::().values(), + batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0, 1, 1] ); assert_batch_matches_single_queries( @@ -6096,8 +6114,11 @@ mod test { ); let batch = scan.try_into_batch().await.unwrap(); + assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); + assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); + assert!(!batch.schema().field(0).is_nullable()); assert_eq!( - batch[QUERY_INDEX_COL].as_primitive::().values(), + batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0, 1, 1] ); @@ -10244,6 +10265,9 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") let batch = scanner.try_into_batch().await.unwrap(); assert_eq!(batch.num_rows(), 0); + assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); + assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); + assert!(!batch.schema().field(0).is_nullable()); assert!( batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), "batch fast_search without index should still expose query_index in schema" diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 049db093cac..27e1f8b7602 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -8,7 +8,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, LazyLock, Mutex}; use std::time::Instant; -use arrow::array::Float32Builder; +use arrow::array::{Float32Builder, Int32Builder}; use arrow::datatypes::{Float32Type, UInt32Type, UInt64Type}; use arrow_array::{Array, Float32Array, UInt32Array, UInt64Array}; use arrow_array::{ @@ -73,6 +73,10 @@ use super::utils::{ pub const QUERY_INDEX_COL: &str = "query_index"; +pub(crate) fn query_index_field() -> Field { + Field::new(QUERY_INDEX_COL, DataType::Int32, false) +} + /// Whether a computed flat KNN distance should participate in top-k selection. /// /// Matches the single-query flat path: keep infinite distances and drop only @@ -301,14 +305,13 @@ impl KNNVectorDistanceExec { ))); } let input_schema = Arc::new(input_schema); - let mut output_schema = input_schema.as_ref().clone(); - if is_batch { - output_schema = output_schema.try_with_column(Field::new( - QUERY_INDEX_COL, - DataType::UInt32, - true, - ))?; - } + let output_schema = if is_batch { + input_schema + .as_ref() + .try_with_column_at(0, query_index_field())? + } else { + input_schema.as_ref().clone() + }; let output_schema = Arc::new(output_schema.try_with_column(Field::new( DIST_COL, DataType::Float32, @@ -404,7 +407,7 @@ impl KNNVectorDistanceExec { { continue; } - let query_index = query_index as u32; + let query_index = query_index as i32; let row_id = row_ids.value(row_index); let row_index = row_index as u32; let candidate_is_better = |worst: &BatchKnnCandidate| { @@ -452,7 +455,7 @@ impl KNNVectorDistanceExec { return Ok(RecordBatch::new_empty(output_schema)); } - let mut query_indices = UInt32Builder::with_capacity(results.len()); + let mut query_indices = Int32Builder::with_capacity(results.len()); let mut distances = Float32Builder::with_capacity(results.len()); let mut row_batches = Vec::with_capacity(results.len()); for result in results { @@ -469,10 +472,7 @@ impl KNNVectorDistanceExec { let output = concat_batches(&input_schema, &row_batches) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; output - .try_with_column( - Field::new(QUERY_INDEX_COL, DataType::UInt32, true), - Arc::new(query_indices.finish()), - ) + .try_with_column_at(0, query_index_field(), Arc::new(query_indices.finish())) .and_then(|batch| { batch.try_with_column( Field::new(DIST_COL, DataType::Float32, true), @@ -628,9 +628,8 @@ impl ExecutionPlan for KNNVectorDistanceExec { .map(|(stats, _)| stats) .collect::>(); let column_statistics = if self.is_batch { - column_statistics - .into_iter() - .chain(std::iter::once(ColumnStatistics::default())) + std::iter::once(ColumnStatistics::default()) + .chain(column_statistics) .chain(std::iter::once(dist_stats)) .collect::>() } else { @@ -669,7 +668,7 @@ impl ExecutionPlan for KNNVectorDistanceExec { #[derive(Clone)] struct BatchKnnCandidate { - query_index: u32, + query_index: i32, distance: f32, row_id: u64, batch: RecordBatch, @@ -712,7 +711,7 @@ pub fn knn_empty_result_schema(include_query_index: bool) -> SchemaRef { ROW_ID_FIELD.clone(), ]; if include_query_index { - fields.push(Field::new(QUERY_INDEX_COL, DataType::UInt32, true)); + fields.insert(0, query_index_field()); } Arc::new(Schema::new(fields)) } From 15e8801a9b786f761e16d3436ae9999f06e1690c Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 26 May 2026 01:11:42 +0800 Subject: [PATCH 30/32] refactor(vector): simplify batch knn candidate filtering --- rust/lance/src/io/exec/knn.rs | 82 ++++++++++------------------------- 1 file changed, 24 insertions(+), 58 deletions(-) diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 27e1f8b7602..7068e1e3b27 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -77,24 +77,6 @@ pub(crate) fn query_index_field() -> Field { Field::new(QUERY_INDEX_COL, DataType::Int32, false) } -/// Whether a computed flat KNN distance should participate in top-k selection. -/// -/// Matches the single-query flat path: keep infinite distances and drop only -/// null/NaN values. -fn flat_knn_distance_is_candidate( - distances: &arrow::array::PrimitiveArray, - row_index: usize, -) -> Option { - if !distances.is_valid(row_index) { - return None; - } - let distance = distances.value(row_index); - if distance.is_nan() { - return None; - } - Some(distance) -} - pub struct AnnPartitionMetrics { index_metrics: IndexMetrics, partitions_ranked: Count, @@ -397,11 +379,17 @@ impl KNNVectorDistanceExec { .await .map_err(|e| DataFusionError::External(Box::new(e)))?; let distances = with_distances[DIST_COL].as_primitive::(); - for row_index in 0..batch.num_rows() { - let Some(distance) = flat_knn_distance_is_candidate(distances, row_index) - else { + let distance_values = distances.values(); + for row_index in 0..distances.len() { + if !distances.is_valid(row_index) { continue; - }; + } + let distance = distance_values[row_index]; + if distance.is_nan() { + continue; + } + // Single-query flat KNN applies distance_range as a plan filter. + // Batch mode filters before insertion so top-k stays per query. if lower_bound.is_some_and(|lower_bound| distance < lower_bound) || upper_bound.is_some_and(|upper_bound| distance >= upper_bound) { @@ -410,31 +398,21 @@ impl KNNVectorDistanceExec { let query_index = query_index as i32; let row_id = row_ids.value(row_index); let row_index = row_index as u32; - let candidate_is_better = |worst: &BatchKnnCandidate| { - distance - .total_cmp(&worst.distance) - .then_with(|| row_id.cmp(&worst.row_id)) - .then_with(|| query_index.cmp(&worst.query_index)) - .then_with(|| row_index.cmp(&worst.row_index)) - .is_lt() + let candidate = BatchKnnCandidate { + query_index, + distance, + row_id, + batch: batch.clone(), + row_index, }; if heap.len() < k { - heap.push(BatchKnnCandidate { - query_index, - distance, - row_id, - batch: batch.clone(), - row_index, - }); - } else if heap.peek().is_some_and(candidate_is_better) { + heap.push(candidate); + } else if heap + .peek() + .is_some_and(|worst| candidate.cmp(worst).is_lt()) + { heap.pop(); - heap.push(BatchKnnCandidate { - query_index, - distance, - row_id, - batch: batch.clone(), - row_index, - }); + heap.push(candidate); } } } @@ -593,8 +571,9 @@ impl ExecutionPlan for KNNVectorDistanceExec { elapsed_compute.add_duration(start.elapsed()); let distances = batch[DIST_COL].as_primitive::(); + let distance_values = distances.values(); let mask = BooleanArray::from_iter((0..distances.len()).map(|row_index| { - Some(flat_knn_distance_is_candidate(distances, row_index).is_some()) + Some(distances.is_valid(row_index) && !distance_values[row_index].is_nan()) })); arrow::compute::filter_record_batch(&batch, &mask) .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) @@ -2606,19 +2585,6 @@ mod tests { assert_eq!(expected, results[0]); } - #[test] - fn flat_knn_distance_keeps_infinity() { - let distances = - Float32Array::from(vec![Some(f32::NAN), Some(f32::INFINITY), None, Some(1.0)]); - assert!(flat_knn_distance_is_candidate(&distances, 0).is_none()); - assert_eq!( - flat_knn_distance_is_candidate(&distances, 1), - Some(f32::INFINITY) - ); - assert!(flat_knn_distance_is_candidate(&distances, 2).is_none()); - assert_eq!(flat_knn_distance_is_candidate(&distances, 3), Some(1.0)); - } - #[test] fn test_create_knn_flat() { let dim: usize = 128; From 8ede30156fadf45059cd11f523842bc18c9ee728 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 26 May 2026 01:19:35 +0800 Subject: [PATCH 31/32] refactor(vector): simplify batch knn scanner checks --- rust/lance/src/dataset/scanner.rs | 39 ++++++++++++------------------- 1 file changed, 15 insertions(+), 24 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index dd03d83f4f8..a2f6a6115dc 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1529,9 +1529,8 @@ impl Scanner { } }; - if Self::is_batch_nearest_query(&vector_type, &query_type) - && self.dataset.schema().field(QUERY_INDEX_COL).is_some() - { + let is_batch_nearest = Self::is_batch_nearest_query(&vector_type, &query_type); + if is_batch_nearest && self.dataset.schema().field(QUERY_INDEX_COL).is_some() { return Err(Error::invalid_input(format!( "batch nearest neighbor search cannot be used on datasets with column '{QUERY_INDEX_COL}'" ))); @@ -1569,7 +1568,7 @@ impl Scanner { dist_q_c: 0.0, }); self.nearest_query_count = query_count; - self.is_batch_nearest = Self::is_batch_nearest_query(&vector_type, &query_type); + self.is_batch_nearest = is_batch_nearest; Ok(self) } @@ -5822,6 +5821,14 @@ mod test { (queries, query_values) } + fn assert_query_index_field(batch: &RecordBatch) { + let schema = batch.schema(); + let field = schema.field(0); + assert_eq!(field.name(), QUERY_INDEX_COL); + assert_eq!(field.data_type(), &DataType::Int32); + assert!(!field.is_nullable()); + } + async fn assert_batch_matches_single_queries( dataset: &Dataset, batch: &RecordBatch, @@ -5895,9 +5902,7 @@ mod test { ); let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); - assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); - assert!(!batch.schema().field(0).is_nullable()); + assert_query_index_field(&batch); assert_eq!( batch.num_rows(), 2 * k, @@ -5944,13 +5949,7 @@ mod test { ); let batch = scan.try_into_batch().await.unwrap(); - assert!( - batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), - "batch-shaped query with one vector should still return query_index" - ); - assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); - assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); - assert!(!batch.schema().field(0).is_nullable()); + assert_query_index_field(&batch); assert_eq!( batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0] @@ -6114,9 +6113,7 @@ mod test { ); let batch = scan.try_into_batch().await.unwrap(); - assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); - assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); - assert!(!batch.schema().field(0).is_nullable()); + assert_query_index_field(&batch); assert_eq!( batch[QUERY_INDEX_COL].as_primitive::().values(), &[0, 0, 1, 1] @@ -10265,13 +10262,7 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") let batch = scanner.try_into_batch().await.unwrap(); assert_eq!(batch.num_rows(), 0); - assert_eq!(batch.schema().field(0).name(), QUERY_INDEX_COL); - assert_eq!(batch.schema().field(0).data_type(), &DataType::Int32); - assert!(!batch.schema().field(0).is_nullable()); - assert!( - batch.schema().column_with_name(QUERY_INDEX_COL).is_some(), - "batch fast_search without index should still expose query_index in schema" - ); + assert_query_index_field(&batch); } #[rstest] From 7c21a0a9ab62a943a60c2dea7ff6373714426443 Mon Sep 17 00:00:00 2001 From: BubbleCal Date: Tue, 26 May 2026 01:52:58 +0800 Subject: [PATCH 32/32] fix(vector): satisfy batch knn clippy lint --- rust/lance/src/io/exec/knn.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 7068e1e3b27..71239b4e34b 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -73,7 +73,7 @@ use super::utils::{ pub const QUERY_INDEX_COL: &str = "query_index"; -pub(crate) fn query_index_field() -> Field { +pub fn query_index_field() -> Field { Field::new(QUERY_INDEX_COL, DataType::Int32, false) }