diff --git a/java/lance-jni/src/blocking_scanner.rs b/java/lance-jni/src/blocking_scanner.rs index 128bb724a41..f18b0d92a27 100644 --- a/java/lance-jni/src/blocking_scanner.rs +++ b/java/lance-jni/src/blocking_scanner.rs @@ -318,7 +318,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/main/java/org/lance/ipc/Query.java b/java/src/main/java/org/lance/ipc/Query.java index 3ad1301ee59..48013b375ee 100644 --- a/java/src/main/java/org/lance/ipc/Query.java +++ b/java/src/main/java/org/lance/ipc/Query.java @@ -140,6 +140,10 @@ 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/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/python/python/lance/dataset.py b/python/python/lance/dataset.py index 3990509ee97..eeb2dacff6a 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,17 @@ def scanner( "distance_range": (0.0, 1.0), } + ``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 + 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 + 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 smaller than this size. Note: this can be overridden by @@ -5992,6 +6003,16 @@ def nearest( Parameters ---------- + 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 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. query_parallelism: int, optional Maximum partition-search concurrency for a single vector query. The default is 0. Value 0 uses the automatic policy, which @@ -7140,7 +7161,15 @@ 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 or list-shaped batch of query vectors. Batch queries return up to + ``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, + 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 356f72a5e66..6537dcb49b7 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -180,6 +180,127 @@ def test_flat(dataset): run(dataset) +@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"], + nearest={ + "column": "vector", + "q": queries, + "k": k, + "use_index": False, + }, + ) + + assert batch.num_rows == query_count * k + 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 + + _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): + batch = ds.to_table( + columns=["id"], + nearest={ + "column": "vector", + "q": queries, + "k": k, + **nearest_kwargs, + }, + ) + 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( + 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_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_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/python/src/dataset.rs b/python/src/dataset.rs index c868504e87c..62ef1670518 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -1371,7 +1371,12 @@ impl Dataset { let (_, 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 => { + 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)) })?; diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 21cfbd629c4..a2f6a6115dc 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; @@ -37,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; @@ -100,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::{KNN_INDEX_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}; @@ -768,6 +771,10 @@ pub struct Scanner { ordering: Option>, 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 /// @@ -1023,6 +1030,8 @@ impl Scanner { offset: None, ordering: None, nearest: None, + nearest_query_count: 1, + is_batch_nearest: false, use_stats: true, ordered: true, fragments: None, @@ -1427,6 +1436,19 @@ impl Scanner { Ok(self) } + /// 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. /// the query can be a Float16Array, Float32Array, Float64Array, UInt8Array, /// or a ListArray/FixedSizeListArray of the above types. @@ -1448,16 +1470,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 q = match q.data_type() { + let (q, query_count) = match &query_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); @@ -1470,7 +1486,15 @@ impl Scanner { ))); } } - list_array.values().clone() + // 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 { + list_array.len() + }; + (list_array.values().clone(), query_count) } else { let fsl = q.as_fixed_size_list(); if fsl.value_length() as usize != dim { @@ -1481,7 +1505,15 @@ impl Scanner { dim, ))); } - fsl.values().clone() + // 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 { + fsl.len() + }; + (fsl.values().clone(), query_count) } } _ => { @@ -1493,10 +1525,17 @@ impl Scanner { dim, ))); } - q.slice(0, q.len()) + (q.slice(0, q.len()), 1) } }; + 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}'" + ))); + } + let key = match &element_type { dt if dt == q.data_type() => q, dt if dt.is_floating() => coerce_float_vector( @@ -1528,6 +1567,8 @@ impl Scanner { query_parallelism: DEFAULT_QUERY_PARALLELISM, dist_q_c: 0.0, }); + self.nearest_query_count = query_count; + self.is_batch_nearest = is_batch_nearest; Ok(self) } @@ -1858,6 +1899,9 @@ 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(query_index_field()); + } }; if self.full_text_query.is_some() { @@ -1919,6 +1963,23 @@ impl Scanner { } } + // 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 { let row_id_pos = output_expr .iter() @@ -3504,6 +3565,7 @@ impl Scanner { } // ANN/KNN search execution node with optional prefilter + #[async_recursion] async fn vector_search( &self, filter_plan: &ExprFilterPlan, @@ -3655,6 +3717,10 @@ impl Scanner { }; if let Some((index_name, index_segments, index_metric)) = index_and_segments { + if self.is_batch_nearest { + 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); @@ -3692,7 +3758,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.is_batch_nearest, + )))); } // Resolve metric type for flat search (use default if not specified) let metric = q @@ -3732,6 +3800,86 @@ 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.is_batch_nearest = false; + 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 i32, + )?); + } + + 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: 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())?) + as Arc, + field.name().clone(), + )); + } + 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, @@ -4271,13 +4419,29 @@ impl Scanner { default_distance_type_for(&element_type) } }; - let flat_dist = Arc::new(KNNVectorDistanceExec::try_new( + let input = if self.is_batch_nearest { + Arc::new(CoalescePartitionsExec::new(input)) as Arc + } else { + input + }; + let flat_dist = Arc::new(KNNVectorDistanceExec::try_new_batch( input, &q.column, q.key.clone(), - metric_type, + KnnBatchParams { + is_batch: self.is_batch_nearest, + query_count: self.nearest_query_count, + k: q.k, + lower_bound: q.lower_bound, + upper_bound: q.upper_bound, + distance_type: metric_type, + }, )?); + if self.is_batch_nearest { + return Ok(flat_dist); + } + let lower: Option<(Expr, Arc)> = q .lower_bound .map(|v| -> Result<(Expr, Arc)> { @@ -4342,10 +4506,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 { @@ -4996,10 +5167,10 @@ 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, + ArrayRef, BooleanArray, FixedSizeListArray, Float16Array, Int32Array, LargeStringArray, + PrimitiveArray, RecordBatchIterator, StringArray, StructArray, UInt8Array, UInt32Array, }; use arrow_ord::sort::sort_to_indices; @@ -5642,7 +5813,333 @@ mod test { assert_eq!(expected_i, actual_i); } - #[rstest] + 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) + } + + 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, + 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 i32)), + ); + 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, k).unwrap(); + scan.use_index(false); + scan.project(&["i"]).unwrap(); + + let plan = scan.explain_plan(false).await.unwrap(); + assert!( + plan.contains("KNNVectorDistance: queries=2"), + "expected flat batch KNN plan, got:\n{}", + plan + ); + assert!( + !plan.contains("ANNSubIndex"), + "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_query_index_field(&batch); + 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)) + .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::>(); + 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); + 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_query_index_field(&batch); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0] + ); + } + + #[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_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_respects_distance_range() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + let dataset = &test_ds.dataset; + let (queries, query_values) = batch_knn_two_queries(); + + 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[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0, 1, 1] + ); + assert_batch_matches_single_queries( + dataset, + &batch, + &query_values, + 2, + false, + Some((Some(1.0), None)), + ) + .await; + } + + #[tokio::test] + 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 (queries, query_values) = batch_knn_two_queries(); + + 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_query_index_field(&batch); + assert_eq!( + batch[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 0, 1, 1] + ); + + let batch = dataset + .scan() + .nearest("vec", &queries, 2) + .unwrap() + .distance_range(Some(1.0), None) + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_batch_matches_single_queries( + dataset, + &batch, + &query_values, + 2, + true, + Some((Some(1.0), None)), + ) + .await; + } + #[tokio::test] async fn test_can_project_distance() { let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) @@ -9751,6 +10248,23 @@ 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_query_index_field(&batch); + } + #[rstest] #[tokio::test] async fn test_fast_search_scalar_index_skips_unindexed_fragments( diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 49ee7be86bc..71239b4e34b 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -2,12 +2,13 @@ // 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; -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::{ @@ -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::{ @@ -65,10 +67,16 @@ use crate::{Error, Result}; use lance_arrow::*; use super::utils::{ - FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedChildInputStream, PreFilterSource, - SelectionVectorToPrefilter, + FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedChildInputStream, + InstrumentedRecordBatchStreamAdapter, PreFilterSource, SelectionVectorToPrefilter, }; +pub const QUERY_INDEX_COL: &str = "query_index"; + +pub fn query_index_field() -> Field { + Field::new(QUERY_INDEX_COL, DataType::Int32, false) +} + pub struct AnnPartitionMetrics { index_metrics: IndexMetrics, partitions_ranked: Count, @@ -141,23 +149,66 @@ 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, + pub upper_bound: Option, pub column: String, pub distance_type: DistanceType, + input_schema: SchemaRef, output_schema: SchemaRef, properties: Arc, metrics: ExecutionPlanMetricsSet, } +pub struct KnnBatchParams { + pub is_batch: bool, + 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, + column: String, + query: ArrayRef, + query_count: usize, + k: usize, + lower_bound: Option, + upper_bound: Option, + 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.is_batch { + 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.is_batch { + write!( + f, + "KNNVectorDistance\nqueries={}\nk={}\nmetric={}", + self.query_count, self.k, self.distance_type, + ) + } else { + write!(f, "KNNVectorDistance\nmetric={}", self.distance_type,) + } } } } @@ -173,16 +224,76 @@ 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)?; + Self::try_new_batch( + input, + column, + query, + KnnBatchParams { + is_batch: false, + query_count: 1, + k: 0, + lower_bound: None, + upper_bound: None, + distance_type, + }, + ) + } + + pub(crate) fn try_new_batch( + input: Arc, + column: &str, + query: ArrayRef, + params: KnnBatchParams, + ) -> Result { + let KnnBatchParams { + is_batch, + 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(), + )); + } + 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 is_batch && 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 output_schema.column_with_name(DIST_COL).is_some() { - output_schema = output_schema.without_column(DIST_COL); + if input_schema.column_with_name(DIST_COL).is_some() { + input_schema = input_schema.without_column(DIST_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 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, @@ -191,24 +302,163 @@ impl KNNVectorDistanceExec { // 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())), - ); + let properties = if is_batch { + 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, query, + is_batch, + query_count, + k, + lower_bound, + upper_bound, column: column.to_string(), distance_type, + input_schema, output_schema, properties, metrics: ExecutionPlanMetricsSet::new(), }) } + + async fn execute_batch( + input: SendableRecordBatchStream, + config: BatchKnnConfig, + ) -> DataFusionResult { + let BatchKnnConfig { + input_schema, + output_schema, + column, + query, + query_count, + k, + lower_bound, + upper_bound, + distance_type, + } = config; + let query_dim = query.len() / query_count; + let mut heaps = (0..query_count) + .map(|_| BinaryHeap::::with_capacity(k)) + .collect::>(); + let mut input = input; + + while let Some(batch) = input.next().await { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| { + DataFusionError::Internal( + "KNNVectorDistanceExec batch mode requires _rowid in input".to_string(), + ) + })? + .as_primitive::() + .clone(); + + 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.clone()) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let distances = with_distances[DIST_COL].as_primitive::(); + 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) + { + continue; + } + let query_index = query_index as i32; + let row_id = row_ids.value(row_index); + let row_index = row_index as u32; + let candidate = BatchKnnCandidate { + query_index, + distance, + row_id, + batch: batch.clone(), + row_index, + }; + 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); + } + } + } + } + + 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 = 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 { + 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 output = concat_batches(&input_schema, &row_batches) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?; + output + .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), + Arc::new(distances.finish()), + ) + }) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) + } } impl ExecutionPlan for KNNVectorDistanceExec { @@ -239,11 +489,18 @@ impl ExecutionPlan for KNNVectorDistanceExec { )); } - Ok(Arc::new(Self::try_new( + Ok(Arc::new(Self::try_new_batch( children.pop().expect("length checked"), &self.column, self.query.clone(), - self.distance_type, + KnnBatchParams { + is_batch: self.is_batch, + query_count: self.query_count, + k: self.k, + lower_bound: self.lower_bound, + upper_bound: self.upper_bound, + distance_type: self.distance_type, + }, )?)) } @@ -253,7 +510,30 @@ impl ExecutionPlan for KNNVectorDistanceExec { context: Arc, ) -> DataFusionResult { let input_stream = self.input.execute(partition, context)?; - let input_schema = input_stream.schema(); + if self.is_batch { + 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, + lower_bound: self.lower_bound, + upper_bound: self.upper_bound, + distance_type: self.distance_type, + }, + )); + let schema = self.schema(); + return Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new( + schema, + stream.boxed(), + partition, + &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; @@ -284,18 +564,17 @@ impl ExecutionPlan for KNNVectorDistanceExec { // 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 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( - distances - .iter() - .map(|v| Some(v.map(|v| !v.is_nan()).unwrap_or(false))), - ); + let distance_values = distances.values(); + let mask = BooleanArray::from_iter((0..distances.len()).map(|row_index| { + 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)) } @@ -326,8 +605,18 @@ impl ExecutionPlan for KNNVectorDistanceExec { .zip(schema.fields()) .filter(|(_, field)| field.name() != DIST_COL) .map(|(stats, _)| stats) - .chain(std::iter::once(dist_stats)) .collect::>(); + let column_statistics = if self.is_batch { + std::iter::once(ColumnStatistics::default()) + .chain(column_statistics) + .chain(std::iter::once(dist_stats)) + .collect::>() + } else { + column_statistics + .into_iter() + .chain(std::iter::once(dist_stats)) + .collect::>() + }; Ok(Statistics { num_rows: inner_stats.num_rows, column_statistics, @@ -346,14 +635,65 @@ impl ExecutionPlan for KNNVectorDistanceExec { fn supports_limit_pushdown(&self) -> bool { false } + + fn required_input_distribution(&self) -> Vec { + if self.is_batch { + vec![Distribution::SinglePartition] + } else { + vec![Distribution::UnspecifiedDistribution] + } + } } -pub static KNN_INDEX_SCHEMA: LazyLock = LazyLock::new(|| { - Arc::new(Schema::new(vec![ +#[derive(Clone)] +struct BatchKnnCandidate { + query_index: i32, + distance: f32, + row_id: u64, + batch: RecordBatch, + 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(|| 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.insert(0, query_index_field()); + } + Arc::new(Schema::new(fields)) +} pub static KNN_PARTITION_SCHEMA: LazyLock = LazyLock::new(|| { Arc::new(Schema::new(vec![