fix: share scans across batched vector queries (#3805)

<!-- lance-gatekeeper-fix:v1 agent=d30696bc46eb32f04c9927b0792e35d3
generation=1 -->

## Summary

- use the Lance native batch KNN path so fixed-size batch vector
searches share one flat table scan
- validate consistent query-vector dimensions and retain the per-vector
plan when offsets require its existing semantics
- add Rust and Python regressions and update Rust, Python, and
TypeScript API documentation

## Root cause

LanceDB expanded every vector in a batch into a separate scan plan and
joined the plans with `UnionExec`. For unindexed tables on S3, a batch
of ten vectors therefore ran ten concurrent full scans, amplifying CPU
and retained data enough to produce the reported memory spike.

The native Lance batch KNN path performs bounded-memory selection for
all query vectors over one flat scan. LanceDB now supplies the vectors
as a batch and avoids applying a global scanner limit to the combined
per-query results. Batch queries with a nonzero offset keep the previous
plan because the native batch API does not support per-query offsets.

## Validation

- targeted Rust batch-query plan and execution tests
- `cargo check --quiet --features remote --tests --examples`
- `cargo clippy --quiet --features remote --tests --examples`
- `cargo fmt --all -- --check`
- targeted Python batch-vector regression after rebuilding the extension
- Ruff formatting/checks for the touched Python files
- Node.js build, lint, docs generation, and targeted batch-vector Jest
test
- `git diff --check`

Fixes #2468

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>

nodejs/lancedb/query.ts

@@ -727,11 +727,11 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
    * Add a query vector to the search
    *
    * This method can be called multiple times to add multiple query vectors
-   * to the search. If multiple query vectors are added, then they will be searched
-   * in parallel, and the results will be concatenated. A column called `query_index`
-   * will be added to indicate the index of the query vector that produced the result.
-   *
-   * Performance wise, this is equivalent to running multiple queries concurrently.
+   * to the search. A column called `query_index` will be added to indicate the index
+   * of the query vector that produced the result. Flat searches share one table scan
+   * across the query vectors, avoiding the scan and memory amplification of running
+   * multiple queries concurrently. Indexed searches may still perform per-vector
+   * index work.
    */
   addQueryVector(vector: IntoVector): VectorQuery {
     if (vector instanceof Promise) {

python/python/lancedb/query.py

@@ -3401,9 +3401,10 @@ class AsyncQuery(AsyncStandardQuery):
         pass in multiple vectors. When multiple vectors are passed in, if the vector
         column is with multivector type, then the vectors will be treated as a single
         query. Or the vectors will be treated as multiple queries, this can be useful
-        if you want to find the nearest vectors to multiple query vectors.
-        This is not expected to be faster than making multiple queries concurrently;
-        it is just a convenience method. If multiple vectors are passed in then
+        if you want to find the nearest vectors to multiple query vectors. Flat
+        searches share one table scan across the query vectors, avoiding the scan
+        and memory amplification of making multiple queries concurrently. If
+        multiple vectors are passed in then
         an additional column `query_index` will be added to the results. This column
         will contain the index of the query vector that the result is nearest to.
         """
@@ -3532,8 +3533,8 @@ class AsyncFTSQuery(AsyncStandardQuery):
 
         Typically, a single vector is passed in as the query. However, you can also
         pass in multiple vectors.  This can be useful if you want to find the nearest
-        vectors to multiple query vectors. This is not expected to be faster than
-        making multiple queries concurrently; it is just a convenience method.
+        vectors to multiple query vectors. Flat searches share one table scan across
+        the query vectors instead of issuing concurrent full scans.
         If multiple vectors are passed in then an additional column `query_index`
         will be added to the results.  This column will contain the index of the
         query vector that the result is nearest to.

python/python/tests/test_query.py

@@ -897,6 +897,23 @@ def test_query_builder_batches(table):
     assert rs_list["id"][1] == 2
 
 
+def test_batch_vector_query_shares_filtered_flat_scan(table):
+    query = (
+        table.search([[1.0, 2.0], [3.0, 4.0]])
+        .where("id > 0", prefilter=True)
+        .limit(1)
+        .select(["id"])
+    )
+
+    plan = query.explain_plan(verbose=True)
+    assert "KNNVectorDistance: queries=2" in plan
+    assert "UnionExec" not in plan
+
+    results = query.to_arrow()
+    assert len(results) == 2
+    assert results["query_index"].to_pylist() == [0, 1]
+
+
 def test_dynamic_projection(table):
     rs = (
         LanceVectorQueryBuilder(table, [0, 0], "vector")

rust/lancedb/src/query.rs

@@ -1174,12 +1174,12 @@ impl VectorQuery {
 
     /// Add another query vector to the search.
     ///
-    /// Multiple searches will be dispatched as part of the query.
-    /// This is a convenience method for adding multiple query vectors
-    /// to the search. It is not expected to be faster than issuing
-    /// multiple queries concurrently.
+    /// Multiple searches will be dispatched as a batch. Flat searches share
+    /// one table scan across the query vectors, avoiding the scan and memory
+    /// amplification of issuing the searches concurrently. Indexed searches
+    /// may still perform per-vector index work.
     ///
-    /// The output data will contain an additional columns `query_index` which
+    /// The output data will contain an additional column `query_index` which
     /// will contain the index of the query vector that was used to generate the
     /// result.
     pub fn add_query_vector(mut self, vector: impl IntoQueryVector) -> Result<Self> {
@@ -1646,7 +1646,11 @@ mod tests {
     use std::{collections::HashSet, sync::Arc};
 
     use super::*;
-    use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
+    use arrow::{
+        array::downcast_array,
+        compute::concat_batches,
+        datatypes::{Int32Type, UInt8Type},
+    };
     use arrow_array::{
         FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray,
         types::Float32Type,
@@ -2334,7 +2338,8 @@ mod tests {
             .limit(1);
 
         let plan = query.explain_plan(true).await.unwrap();
-        assert!(plan.contains("UnionExec"));
+        assert!(plan.contains("KNNVectorDistance: queries=2"));
+        assert!(!plan.contains("UnionExec"));
 
         let results = query
             .execute()
@@ -2349,6 +2354,100 @@ mod tests {
         // We don't guarantee order.
         assert!(query_index.values().contains(&0));
         assert!(query_index.values().contains(&1));
+
+        // Batch KNN does not support a per-query offset, so offset queries keep
+        // the legacy per-vector plan to preserve their result semantics.
+        let offset_query = table
+            .query()
+            .nearest_to(&[0.1, 0.2, 0.3, 0.4])
+            .unwrap()
+            .add_query_vector(&[0.5, 0.6, 0.7, 0.8])
+            .unwrap()
+            .limit(1)
+            .offset(1);
+        assert!(
+            offset_query
+                .explain_plan(true)
+                .await
+                .unwrap()
+                .contains("UnionExec")
+        );
+        let offset_results = offset_query
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        assert_eq!(
+            offset_results
+                .iter()
+                .map(RecordBatch::num_rows)
+                .sum::<usize>(),
+            2
+        );
+    }
+
+    #[tokio::test]
+    async fn test_multiple_binary_query_vectors() {
+        let vectors = FixedSizeListArray::from_iter_primitive::<UInt8Type, _, _>(
+            vec![
+                Some(vec![Some(0), Some(0)]),
+                Some(vec![Some(255), Some(255)]),
+            ],
+            2,
+        );
+        let schema = Arc::new(ArrowSchema::new(vec![
+            ArrowField::new("id", DataType::Int32, false),
+            ArrowField::new("vector", vectors.data_type().clone(), false),
+        ]));
+        let batch = RecordBatch::try_new(
+            schema,
+            vec![Arc::new(Int32Array::from(vec![0, 1])), Arc::new(vectors)],
+        )
+        .unwrap();
+
+        let conn = connect("memory://").execute().await.unwrap();
+        let table = conn
+            .create_table("binary_batch", batch)
+            .execute()
+            .await
+            .unwrap();
+        let query = table
+            .query()
+            .nearest_to(&[0.0, 0.0])
+            .unwrap()
+            .add_query_vector(&[255.0, 255.0])
+            .unwrap()
+            .distance_type(DistanceType::Hamming)
+            .limit(1);
+
+        // Binary queries retain the per-vector plan because Lance's binary
+        // nearest path requires primitive UInt8 query arrays.
+        assert!(
+            query
+                .explain_plan(true)
+                .await
+                .unwrap()
+                .contains("UnionExec")
+        );
+
+        let results = query
+            .execute()
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let results = concat_batches(&results[0].schema(), &results).unwrap();
+        assert_eq!(results.num_rows(), 2);
+
+        let ids = results["id"].as_primitive::<Int32Type>();
+        assert!(ids.values().contains(&0));
+        assert!(ids.values().contains(&1));
+        let query_index = results["query_index"].as_primitive::<Int32Type>();
+        assert!(query_index.values().contains(&0));
+        assert!(query_index.values().contains(&1));
     }
 
     #[tokio::test]

rust/lancedb/src/table/query.rs

@@ -21,7 +21,6 @@ use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::projection::ProjectionExec;
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::union::UnionExec;
-use futures::future::try_join_all;
 use lance::dataset::mem_wal::DatasetMemWalExt;
 use lance::dataset::scanner::DatasetRecordBatchStream;
 use lance::dataset::scanner::Scanner;
@@ -170,6 +169,7 @@ pub async fn create_plan(
     let mut column = query.column.clone();
 
     let mut query_vector = query.query_vector.first().cloned();
+    let mut is_batch_query = false;
     if query.query_vector.len() > 1 {
         if column.is_none() {
             // Infer a vector column with the same dimension of the query vector.
@@ -180,16 +180,37 @@ pub async fn create_plan(
             )?);
         }
         let vector_field = schema.field(column.as_ref().unwrap()).unwrap();
-        if let DataType::List(_) = vector_field.data_type() {
-            // Multivector handling: concatenate into FixedSizeList<FixedSizeList<_>>
+        let (_, element_type) =
+            lance::index::vector::utils::get_vector_type(schema, column.as_ref().unwrap())?;
+        let is_binary = matches!(element_type, DataType::UInt8);
+        if matches!(vector_field.data_type(), DataType::List(_))
+            || (query.base.offset.unwrap_or(0) == 0 && !is_binary)
+        {
+            // Lance distinguishes these cases from the vector column type: a
+            // list-like query against a List column is one multivector query,
+            // while the same query against a FixedSizeList column is a batch of
+            // independent queries. The batch path shares a single flat scan and
+            // bounds retained candidate data instead of running one scan per
+            // query vector.
             let vectors = query
                 .query_vector
                 .iter()
                 .map(|arr| arr.as_ref())
                 .collect::<Vec<_>>();
             let dim = vectors[0].len();
+            if let Some((query_index, actual_dim)) = vectors
+                .iter()
+                .enumerate()
+                .find_map(|(index, vector)| (vector.len() != dim).then_some((index, vector.len())))
+            {
+                return Err(Error::InvalidInput {
+                    message: format!(
+                        "query vector at index {query_index} has dimension {actual_dim}, expected {dim}"
+                    ),
+                });
+            }
             let mut fsl_builder = FixedSizeListBuilder::with_capacity(
-                Float32Builder::with_capacity(dim),
+                Float32Builder::with_capacity(dim * vectors.len()),
                 dim as i32,
                 vectors.len(),
             );
@@ -200,8 +221,12 @@ pub async fn create_plan(
                 fsl_builder.append(true);
             }
             query_vector = Some(Arc::new(fsl_builder.finish()));
+            is_batch_query = !matches!(vector_field.data_type(), DataType::List(_));
         } else {
-            // Multiple query vectors: create a plan for each and union them
+            // Lance's batch path has no per-query offset, and its binary path
+            // requires primitive UInt8 queries rather than a fixed-size list.
+            // Keep the prior plan shape for these cases so offsets are applied
+            // per query and binary query vectors retain their primitive shape.
             let query_vecs = query.query_vector.clone();
             let plan_futures = query_vecs
                 .into_iter()
@@ -214,7 +239,7 @@ pub async fn create_plan(
                     }
                 })
                 .collect::<Vec<_>>();
-            let plans = try_join_all(plan_futures).await?;
+            let plans = futures::future::try_join_all(plan_futures).await?;
             return create_multi_vector_plan(plans);
         }
     }
@@ -251,10 +276,14 @@ pub async fn create_plan(
         }
     }
 
-    scanner.limit(
-        query.base.limit.map(|limit| limit as i64),
-        query.base.offset.map(|offset| offset as i64),
-    )?;
+    // For a batch query, `nearest` already applies k to each query vector.
+    // Adding Scanner's global limit would truncate the combined result to k rows.
+    if !is_batch_query {
+        scanner.limit(
+            query.base.limit.map(|limit| limit as i64),
+            query.base.offset.map(|offset| offset as i64),
+        )?;
+    }
 
     if let Some(ef) = query.ef {
         scanner.ef(ef);
@@ -1088,7 +1117,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn test_create_plan_multivector_structure() {
+    async fn test_create_plan_batch_vector_uses_shared_scan() {
         use arrow_array::{Float32Array, RecordBatch};
         use arrow_schema::{DataType, Field, Schema};
         use datafusion_physical_plan::display::DisplayableExecutionPlan;
@@ -1115,11 +1144,18 @@ mod tests {
             .unwrap();
         let native_table = table.as_native().unwrap();
 
-        // This triggers the "create_multi_vector_plan" logic branch
+        // A batch of vectors against a fixed-size vector column should use
+        // Lance's native batch KNN path instead of independent scan plans.
         let q1 = Arc::new(Float32Array::from(vec![1.0, 2.0]));
         let q2 = Arc::new(Float32Array::from(vec![3.0, 4.0]));
 
         let req = VectorQueryRequest {
+            base: QueryRequest {
+                filter: Some(QueryFilter::Sql("id >= 0".to_string())),
+                limit: Some(1),
+                select: Select::Columns(vec!["id".to_string()]),
+                ..Default::default()
+            },
             column: Some("vector".to_string()),
             query_vector: vec![q1, q2],
             ..Default::default()
@@ -1136,19 +1172,17 @@ mod tests {
             .indent(true)
             .to_string();
 
-        // We expect a RepartitionExec wrapping a UnionExec
         assert!(
-            display.contains("RepartitionExec"),
-            "Plan should include Repartitioning"
+            display.contains("KNNVectorDistance: queries=2"),
+            "plan should use native batch KNN, got:\n{display}"
         );
         assert!(
-            display.contains("UnionExec"),
-            "Plan should include a Union of multiple searches"
+            !display.contains("UnionExec"),
+            "flat batch KNN should share one scan, got:\n{display}"
         );
-        // We expect the projection to add the 'query_index' column (logic inside multi_vector_plan)
         assert!(
             display.contains("query_index"),
-            "Plan should add query_index column"
+            "plan should add query_index column, got:\n{display}"
         );
     }