diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 843a1276f..57fcb1d08 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -708,11 +708,11 @@ export class VectorQuery extends StandardQueryBase { * 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) { diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 095a7b5ff..481edc54b 100644 --- a/python/python/lancedb/query.py +++ b/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. diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 6840be052..89d146736 100644 --- a/python/python/tests/test_query.py +++ b/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") diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index b76865043..c64ec223c 100644 --- a/rust/lancedb/src/query.rs +++ b/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 { @@ -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::>() + .await + .unwrap(); + assert_eq!( + offset_results + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 2 + ); } #[tokio::test] diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 9feb9d5ab..5d8e9aea5 100644 --- a/rust/lancedb/src/table/query.rs +++ b/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> + 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::>(); 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::>(); - 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}" ); }