diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 1baadff51..2a1283f22 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -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, @@ -2384,6 +2388,68 @@ mod tests { ); } + #[tokio::test] + async fn test_multiple_binary_query_vectors() { + let vectors = FixedSizeListArray::from_iter_primitive::( + 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::>() + .await + .unwrap(); + let results = concat_batches(&results[0].schema(), &results).unwrap(); + assert_eq!(results.num_rows(), 2); + + let ids = results["id"].as_primitive::(); + assert!(ids.values().contains(&0)); + assert!(ids.values().contains(&1)); + let query_index = results["query_index"].as_primitive::(); + assert!(query_index.values().contains(&0)); + assert!(query_index.values().contains(&1)); + } + #[tokio::test] async fn test_hybrid_search() { let tmp_dir = tempdir().unwrap(); diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index c746fc44a..2684ac5e2 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -180,8 +180,11 @@ pub async fn create_plan( )?); } let vector_field = schema.field(column.as_ref().unwrap()).unwrap(); + 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 + || (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, @@ -220,9 +223,10 @@ pub async fn create_plan( query_vector = Some(Arc::new(fsl_builder.finish())); is_batch_query = !matches!(vector_field.data_type(), DataType::List(_)); } else { - // 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. + // 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()