fix: share scans across batched vector queries

nodejs/lancedb/query.ts

@@ -708,11 +708,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

@@ -3384,9 +3384,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.
         """
@@ -3515,8 +3516,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

@@ -888,6 +888,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> {
@@ -2355,7 +2355,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()
@@ -2370,6 +2371,38 @@ 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]

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;
@@ -167,6 +166,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.
@@ -177,16 +177,34 @@ 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<_>>
+        if matches!(vector_field.data_type(), DataType::List(_))
+            || query.base.offset.unwrap_or(0) == 0
+        {
+            // 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(),
             );
@@ -197,8 +215,11 @@ 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. Keep the prior plan
+            // shape for offset queries so the offset is applied to each query,
+            // rather than globally across the combined results.
             let query_vecs = query.query_vector.clone();
             let plan_futures = query_vecs
                 .into_iter()
@@ -211,7 +232,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);
         }
     }
@@ -248,10 +269,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);
@@ -1007,7 +1032,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;
@@ -1034,11 +1059,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()
@@ -1055,19 +1087,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}"
         );
     }