From 9959f3ad64dd98f0c35f48bdaa211fe804c3cca5 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Thu, 17 Sep 2026 22:29:38 +0800 Subject: [PATCH] perf(knn): share candidate reads for batched refinement --- docs/src/quickstart/vector-search.md | 22 + python/python/benchmarks/batch_refinement.py | 237 ++++++++ python/python/tests/test_vector_index.py | 5 +- rust/lance/src/dataset/scanner.rs | 226 ++++++-- rust/lance/src/io/exec/knn.rs | 3 + rust/lance/src/io/exec/knn/refine.rs | 571 +++++++++++++++++++ 6 files changed, 1001 insertions(+), 63 deletions(-) create mode 100644 python/python/benchmarks/batch_refinement.py create mode 100644 rust/lance/src/io/exec/knn/refine.rs diff --git a/docs/src/quickstart/vector-search.md b/docs/src/quickstart/vector-search.md index ac77e12b752..cf45c864c0a 100644 --- a/docs/src/quickstart/vector-search.md +++ b/docs/src/quickstart/vector-search.md @@ -227,6 +227,28 @@ sift1m.to_table( - `nprobes` => how many partitions (in the coarse quantizer) to probe - `refine_factor` => controls "re-ranking". If k=10 and refine_factor=5 then retrieve 50 nearest neighbors by ANN and re-sort using actual distances then return top 10. This improves recall without sacrificing performance too much +Pass a two-dimensional query array to refine several queries together: + +```python +sift1m.to_table( + nearest={ + "column": "vector", + "q": samples[:8], + "k": 10, + "nprobes": 10, + "refine_factor": 5, + } +) +``` + +Results include a `query_index` column identifying each query. With fixed, +positive `nprobes` and fully indexed IVF_FLAT/PQ/SQ/RQ data, Lance shares index +partition scans and reads duplicate refinement candidates once per input batch. +Each query still ranks only its own candidates. Vector reads are split into +chunks targeting 32 MiB of decoded vector data; candidate metadata and per-query +top-k results consume additional memory. Adaptive probes, HNSW, incomplete index +coverage, stale vector overlays, and external row masks use per-query execution. + !!! note "Memory Usage" The latencies above include file I/O as Lance currently doesn't hold anything in memory. Along with index building speed, creating a purely in-memory version of the dataset would make the biggest impact on performance. diff --git a/python/python/benchmarks/batch_refinement.py b/python/python/benchmarks/batch_refinement.py new file mode 100644 index 00000000000..c402056882a --- /dev/null +++ b/python/python/benchmarks/batch_refinement.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Compare native batch refinement with serial and concurrent single queries. + +Build with release-with-debug, then run from python/: + uv run --no-sync python python/benchmarks/batch_refinement.py \ + --directory /tmp/batch-refinement --output /tmp/refinement.json + +Reuse the directory across revisions to preserve the dataset, index and queries. +The index is warmed; the OS page cache is not flushed. This is not a cold-storage +benchmark. Dataset creation and validation are excluded from measured latency. +""" + +import argparse +import json +import platform +import random +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import lance +import numpy as np +import pyarrow as pa + + +def prepare(args): + args.directory.mkdir(parents=True, exist_ok=True) + dataset_path = args.directory / "vectors.lance" + query_path = args.directory / "queries.npz" + config_path = args.directory / "config.json" + config = { + "rows": args.rows, + "dimension": args.dimension, + "partitions": args.partitions, + "seed": args.seed, + } + if dataset_path.exists(): + if json.loads(config_path.read_text()) != config: + raise ValueError("Existing dataset configuration differs from arguments") + else: + rng = np.random.default_rng(args.seed) + scale = np.exp(-np.arange(args.dimension) / 128).astype(np.float32) + vectors = rng.standard_normal((args.rows, args.dimension), dtype=np.float32) + vectors *= scale + vectors /= np.linalg.norm(vectors, axis=1, keepdims=True) + table = pa.table( + { + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(vectors.reshape(-1)), args.dimension + ), + "id": pa.array(np.arange(args.rows)), + } + ) + ds = lance.write_dataset(table, dataset_path, max_rows_per_file=4096) + ds.create_index( + "vector", + index_type="IVF_SQ", + metric="l2", + num_partitions=args.partitions, + ) + independent = rng.standard_normal((64, args.dimension), dtype=np.float32) + independent *= scale + independent /= np.linalg.norm(independent, axis=1, keepdims=True) + anchor = rng.standard_normal(args.dimension, dtype=np.float32) * scale + anchor /= np.linalg.norm(anchor) + nearby = ( + anchor + + 0.005 + * rng.standard_normal((64, args.dimension), dtype=np.float32) + * scale + ) + nearby /= np.linalg.norm(nearby, axis=1, keepdims=True) + np.savez(query_path, independent=independent, nearby=nearby) + config_path.write_text(json.dumps(config)) + with np.load(query_path) as data: + queries = {name: data[name] for name in data.files} + return dataset_path, queries + + +def unpack(table, count): + ids = table["id"].to_numpy() + distances = table["_distance"].to_numpy() + if "query_index" not in table.column_names: + return [(ids, distances)] + indices = table["query_index"].to_numpy() + return [(ids[indices == i], distances[indices == i]) for i in range(count)] + + +def validate(actual, expected): + assert len(actual) == len(expected) + for (ids, distances), (expected_ids, expected_distances) in zip(actual, expected): + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_allclose(distances, expected_distances, atol=2e-6, rtol=2e-5) + + +def run(args): + dataset_path, query_sets = prepare(args) + ds = lance.dataset(dataset_path, index_cache_size_bytes=512 * 1024**2) + parameters = { + "column": "vector", + "k": args.k, + "nprobes": args.nprobes, + "refine_factor": args.refine_factor, + "query_parallelism": 1, + } + + def single(query): + return unpack( + ds.to_table( + columns=["id", "_distance"], nearest={**parameters, "q": query} + ), + 1, + )[0] + + records = [] + plans = {} + checksums = {} + recall = {} + rng = random.Random(args.seed) + with ThreadPoolExecutor(max_workers=args.workers) as executor: + for distribution, all_queries in query_sets.items(): + hits = 0 + for query in all_queries[:8]: + actual_ids, _ = single(query) + exact = ds.to_table( + columns=["id", "_distance"], + nearest={ + "column": "vector", + "q": query, + "k": args.k, + "use_index": False, + }, + )["id"].to_numpy() + hits += len(np.intersect1d(actual_ids, exact)) + recall[distribution] = hits / (8 * args.k) + for width in args.batch_sizes: + queries = all_queries[:width] + expected = [single(query) for query in queries] + checksums[f"{distribution}/{width}"] = [ + ids.tolist() for ids, _ in expected + ] + nearest = {**parameters, "q": queries} + scanner = ds.scanner(columns=["id", "_distance"], nearest=nearest) + plans[f"{distribution}/{width}"] = scanner.analyze_plan() + + def serial(): + return [single(query) for query in queries] + + def parallel(): + return list(executor.map(single, queries)) + + def batch(): + return unpack( + ds.to_table(columns=["id", "_distance"], nearest=nearest), width + ) + + methods = {"serial": serial, "parallel": parallel, "batch": batch} + for method in methods.values(): + validate(method(), expected) + for repeat in range(args.repeats): + order = list(methods) + rng.shuffle(order) + for name in order: + start = time.perf_counter() + result = methods[name]() + elapsed_ms = (time.perf_counter() - start) * 1000 + validate(result, expected) + records.append( + { + "distribution": distribution, + "batch_size": width, + "method": name, + "repeat": repeat, + "elapsed_ms": elapsed_ms, + } + ) + medians = [ + { + "distribution": distribution, + "batch_size": width, + **{ + name: float( + np.median( + [ + record["elapsed_ms"] + for record in records + if record["distribution"] == distribution + and record["batch_size"] == width + and record["method"] == name + ] + ) + ) + for name in ["serial", "parallel", "batch"] + }, + } + for distribution in query_sets + for width in args.batch_sizes + ] + output = { + "label": args.label, + "platform": platform.platform(), + "lance_version": lance.__version__, + "parameters": { + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + }, + "dataset_version": ds.version, + "recall_at_k_over_eight_queries": recall, + "plans": plans, + "result_ids": checksums, + "medians_ms": medians, + "samples": records, + } + args.output.write_text(json.dumps(output, indent=2)) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--directory", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--label", default="release-with-debug") + parser.add_argument("--rows", type=int, default=32768) + parser.add_argument("--dimension", type=int, default=1024) + parser.add_argument("--partitions", type=int, default=64) + parser.add_argument("--nprobes", type=int, default=16) + parser.add_argument("--k", type=int, default=10) + parser.add_argument("--refine-factor", type=int, default=10) + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 8, 32, 64]) + parser.add_argument("--repeats", type=int, default=11) + parser.add_argument("--workers", type=int, default=8) + parser.add_argument("--seed", type=int, default=42) + arguments = parser.parse_args() + if any(width < 1 or width > 64 for width in arguments.batch_sizes): + parser.error("batch sizes must be between 1 and 64") + run(arguments) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index eb2c036f76b..a7d0ff5fc4b 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -243,8 +243,9 @@ def test_batch_flat_query_matches_repeated_single_queries(dataset, queries): @pytest.mark.parametrize("metric", ["l2", "cosine"]) @pytest.mark.parametrize("query_count", [3, 1], ids=["three_queries", "single_query"]) +@pytest.mark.parametrize("refine_factor", [None, 1, 4]) def test_batch_indexed_query_matches_repeated_single_queries( - dataset, metric, query_count + dataset, metric, query_count, refine_factor ): indexed = dataset.create_index( "vector", @@ -263,6 +264,8 @@ def test_batch_indexed_query_matches_repeated_single_queries( # nprobes covers every partition so the shared-scan batch path and the # repeated single-query path search the same partitions deterministically. nearest_kwargs = {"use_index": True, "nprobes": 4} + if refine_factor is not None: + nearest_kwargs["refine_factor"] = refine_factor batch = indexed.to_table( columns=["id"], nearest={"column": "vector", "q": queries, "k": k, **nearest_kwargs}, diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 8733b167263..094114fcab1 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -123,8 +123,8 @@ use crate::io::exec::{ AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNVectorDistanceExec, LancePushdownScanExec, LanceScanExec, Planner, PreFilterSource, RowAddrMaskFilterExec, ScanConfig, TakeExec, knn::{ - KnnBatchParams, QUERY_INDEX_COL, knn_empty_result_schema, new_knn_batch_exec, new_knn_exec, - query_index_field, + BatchRefineExec, KnnBatchParams, QUERY_INDEX_COL, knn_empty_result_schema, + new_knn_batch_exec, new_knn_exec, query_index_field, }, project, }; @@ -5543,7 +5543,6 @@ impl Scanner { /// search per query vector. /// /// Requires all of: - /// - no refine step (the batch path does not yet rerank); /// - fixed nprobes (`minimum_nprobes == maximum_nprobes`) — see below; /// - every segment an IVF index with a flat-style sub-index (i.e. not HNSW); /// - every target fragment covered by the *selected* `index_segments` (or @@ -5568,14 +5567,6 @@ impl Scanner { index_segments: &[IndexMetadata], q: &Query, ) -> Result { - // Any refine factor sends the query onto a reranking path that the - // shared batch scan does not implement: the single-query path reranks - // with the original vectors even when the factor is 1, and rejects a - // factor of 0 outright (`Refine factor cannot be zero`). The batch path - // does neither, so fall back to the per-query loop for every `Some(_)`. - if q.refine_factor.is_some() { - return Ok(false); - } // Only fixed nprobes is provably equivalent to single-query search; see // the method docs. Adaptive nprobes falls back to the per-query loop. if q.maximum_nprobes != Some(q.minimum_nprobes) { @@ -5688,24 +5679,48 @@ impl Scanner { // Fast path: when every index segment is an IVF index with a flat-style // sub-index (IVF_FLAT/PQ/SQ/RQ), search all query vectors in a single // pass that reads each partition's storage once and shares the prefilter - // across the batch. HNSW, refine, and mixed indexed/unindexed scans fall + // across the batch. HNSW and mixed indexed/unindexed scans fall // back to the per-query loop below, which never regresses behavior. + if q.refine_factor == Some(0) { + return Err(Error::invalid_input("Refine factor cannot be zero")); + } if self .batch_index_search_supported(index_name, index_segments, q) .await? { let mut batch_query = q.clone(); batch_query.metric_type = Some(index_metric); + batch_query.k = + q.k.checked_mul(q.refine_factor.unwrap_or(1) as usize) + .ok_or_else(|| { + Error::invalid_input(format!( + "Candidate count overflows usize: k={}, refine_factor={:?}", + q.k, q.refine_factor + )) + })?; + batch_query.refine_factor = None; let prefilter_source = self .prefilter_source(filter_plan, self.get_indexed_frags(index_segments)) .await?; - return new_knn_batch_exec( + let input = new_knn_batch_exec( self.dataset.clone(), index_segments, &batch_query, self.nearest_query_count, prefilter_source, - ); + )?; + return if q.refine_factor.is_some() { + let mut refine_query = q.clone(); + refine_query.metric_type = Some(index_metric); + Ok(Arc::new(BatchRefineExec::try_new( + self.dataset.clone(), + input, + refine_query, + self.nearest_query_count, + )?)) + } else { + Ok(input) + }; } let query_dim = q.key.len() / self.nearest_query_count; @@ -10124,62 +10139,139 @@ mod test { ); } - /// Any `refine_factor` sends the query onto a reranking path the shared-scan - /// batch node does not implement, so the scanner must fall back to the - /// per-query indexed loop and still produce correctly grouped results. - /// - /// All of `refine(0)`, `refine(1)`, and `refine(2)` must fall back: - /// `refine(1)` still reranks on the single-query path (a factor of 1 is not - /// a no-op), and `refine(0)` is rejected there with `Refine factor cannot be - /// zero` — the batch path would instead return empty results. Covering the - /// boundary factors guards the `refine_factor.is_some()` gate against - /// regressing back to a `> 1` check. + #[rstest] + #[case::factor_one(1, true, MetricType::L2, "none")] + #[case::shared(2, true, MetricType::L2, "none")] + #[case::cosine(2, true, MetricType::Cosine, "none")] + #[case::dot(2, true, MetricType::Dot, "none")] + #[case::adaptive(2, false, MetricType::L2, "none")] + #[case::prefilter(2, true, MetricType::L2, "prefilter")] + #[case::postfilter(2, true, MetricType::L2, "postfilter")] + #[case::empty(2, true, MetricType::L2, "empty")] + #[case::deleted(2, true, MetricType::L2, "deleted")] + #[case::bounded(2, true, MetricType::L2, "bounded")] #[tokio::test] - async fn test_batch_knn_indexed_refine_falls_back() { + async fn test_batch_knn_indexed_refine( + #[case] factor: u32, + #[case] fixed_probes: bool, + #[case] metric: MetricType, + #[case] scenario: &str, + ) { let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) .await .unwrap(); - test_ds.make_vector_index().await.unwrap(); + test_ds.make_vector_index_with_metric(metric).await.unwrap(); + if scenario == "deleted" { + test_ds.dataset.delete("i % 2 = 0").await.unwrap(); + } let dataset = &test_ds.dataset; - let (queries, _query_values) = batch_knn_two_queries(); - - for refine_factor in [1u32, 2] { - let mut scan = dataset.scan(); - scan.nearest("vec", &queries, 2).unwrap(); - scan.refine(refine_factor); - scan.project(&["i"]).unwrap(); - - let plan = scan.explain_plan(false).await.unwrap(); - assert!( - !plan.contains("ANNIvfBatch"), - "refine({refine_factor}) must not use the shared-scan batch node, got:\n{plan}" + let (queries, query_values) = batch_knn_two_queries(); + let mut scan = dataset.scan(); + scan.nearest("vec", &queries, 10).unwrap(); + scan.refine(factor).project(&["i", "vec"]).unwrap(); + if fixed_probes { + scan.nprobes(2); + } + match scenario { + "prefilter" | "postfilter" => { + scan.filter("i > 100") + .unwrap() + .prefilter(scenario == "prefilter"); + } + "empty" => { + scan.filter("i < 0").unwrap().prefilter(true); + } + "bounded" => { + scan.distance_range(Some(1.0), Some(1_000_000.0)); + } + _ => {} + } + let plan = scan.explain_plan(false).await.unwrap(); + assert_eq!(plan.contains("BatchRefine"), fixed_probes, "{plan}"); + assert_eq!(plan.contains("ANNIvfBatch"), fixed_probes, "{plan}"); + let batch = scan.try_into_batch().await.unwrap(); + assert_query_index_field(&batch); + for query_index in 0..2 { + let key = + Float32Array::from(query_values[query_index * 32..(query_index + 1) * 32].to_vec()); + let mut single = dataset.scan(); + single.nearest("vec", &key, 10).unwrap(); + single.refine(factor).project(&["i", "vec"]).unwrap(); + if fixed_probes { + single.nprobes(2); + } + match scenario { + "prefilter" | "postfilter" => { + single + .filter("i > 100") + .unwrap() + .prefilter(scenario == "prefilter"); + } + "empty" => { + single.filter("i < 0").unwrap().prefilter(true); + } + "bounded" => { + single.distance_range(Some(1.0), Some(1_000_000.0)); + } + _ => {} + } + let single = single.try_into_batch().await.unwrap(); + let mask = BooleanArray::from_iter( + batch[QUERY_INDEX_COL] + .as_primitive::() + .iter() + .map(|value| value.map(|value| value == query_index as i32)), ); + let actual = arrow::compute::filter_record_batch(&batch, &mask).unwrap(); + assert_eq!(actual["i"].as_ref(), single["i"].as_ref()); + assert_eq!(actual["vec"].as_ref(), single["vec"].as_ref()); + assert_eq!(actual[DIST_COL].as_ref(), single[DIST_COL].as_ref()); + if scenario != "none" { + continue; + } + let exact = dataset + .scan() + .nearest("vec", &key, 10) + .unwrap() + .use_index(false) + .distance_metric(metric) + .project(&["i"]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let expected_ids = exact["i"].as_primitive::(); + let hits = actual["i"] + .as_primitive::() + .values() + .iter() + .filter(|id| expected_ids.values().contains(id)) + .count(); assert!( - plan.contains("ANNSubIndex"), - "refine({refine_factor}) batch search should fall back to the per-query \ - indexed loop, 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], - "refine({refine_factor}) should still group results per query" + hits as f32 / 10.0 >= 0.5, + "recall={hits}/10, metric={metric:?}" ); } + } - // refine(0) is rejected on the fallback (per-query) path; the batch path - // must not silently accept it and return empty results instead. - let mut scan = dataset.scan(); - scan.nearest("vec", &queries, 2).unwrap(); - scan.refine(0); - scan.project(&["i"]).unwrap(); - let result = scan.try_into_batch().await; - assert!( - result.is_err(), - "refine(0) must error rather than fall through to an empty batch result" - ); + #[rstest] + #[case::fixed(true)] + #[case::adaptive(false)] + #[tokio::test] + async fn test_batch_knn_indexed_refine_zero(#[case] fixed_probes: bool) { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) + .await + .unwrap(); + test_ds.make_vector_index().await.unwrap(); + let (queries, _) = batch_knn_two_queries(); + let mut scan = test_ds.dataset.scan(); + scan.nearest("vec", &queries, 2).unwrap().refine(0); + if fixed_probes { + scan.nprobes(2); + } + let error = scan.try_into_batch().await.unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!(error.to_string().contains("Refine factor cannot be zero")); } /// Without pinned nprobes the shared-scan fast path is not equivalent to @@ -10290,8 +10382,9 @@ mod test { /// query must therefore fall back to the per-query loop when a mask is present, /// and every returned row must honor the mask — otherwise the batch path would /// silently return masked-out rows. + #[rstest] #[tokio::test] - async fn test_batch_knn_indexed_external_mask_falls_back() { + async fn test_batch_knn_indexed_external_mask_falls_back(#[values(false, true)] refine: bool) { let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, true) .await .unwrap(); @@ -10323,6 +10416,9 @@ mod test { let mut scan = dataset.scan(); scan.nearest("vec", &queries, k).unwrap(); scan.nprobes(2); + if refine { + scan.refine(2); + } scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( allow.iter().copied(), ))); @@ -10362,8 +10458,11 @@ mod test { /// selected segments' coverage (as `knn_combined` does), not the whole /// logical index, so the scanner falls back to the per-query loop, which /// re-scores the uncovered fragment on the flat path and returns every row. + #[rstest] #[tokio::test] - async fn test_batch_knn_indexed_partial_segment_selection_falls_back() { + async fn test_batch_knn_indexed_partial_segment_selection_falls_back( + #[values(false, true)] refine: bool, + ) { let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) .await .unwrap(); @@ -10382,6 +10481,9 @@ mod test { let mut scan = dataset.scan(); scan.nearest("vec", &queries, k).unwrap(); scan.nprobes(2); + if refine { + scan.refine(2); + } // Request both indexed fragments but select only the segment covering // fragment 0; fragment 1 is covered only by the unselected segment. scan.with_fragments(vec![fragments[0].clone(), fragments[1].clone()]); diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 3f10653cf32..3f9f0585b29 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -76,6 +76,9 @@ use super::utils::{ }; mod adaptive_probe; +mod refine; + +pub use refine::BatchRefineExec; use adaptive_probe::AutoProbePolicy; diff --git a/rust/lance/src/io/exec/knn/refine.rs b/rust/lance/src/io/exec/knn/refine.rs new file mode 100644 index 00000000000..191913fd8df --- /dev/null +++ b/rust/lance/src/io/exec/knn/refine.rs @@ -0,0 +1,571 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Read shared candidates once, then rank only each query's own candidates. + +use std::collections::{BinaryHeap, HashMap}; +use std::sync::Arc; + +use arrow_array::cast::AsArray; +use arrow_array::types::{Int32Type, UInt64Type}; +use arrow_array::{Array, Float32Array, Int32Array, RecordBatch, UInt32Array, UInt64Array}; +use arrow_schema::{DataType, SchemaRef}; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{Count, ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + SendableRecordBatchStream, +}; +use datafusion_physical_expr::{Distribution, EquivalenceProperties}; +use futures::{StreamExt, stream}; +use lance_core::ROW_ID; +use lance_core::utils::tokio::spawn_cpu; +use lance_datafusion::utils::ExecutionPlanMetricsSetExt; +use lance_index::vector::Query; + +use crate::dataset::{Dataset, ProjectionRequest}; +use crate::index::vector::utils::get_vector_type; +use crate::io::exec::utils::InstrumentedRecordBatchStreamAdapter; +use crate::{Error, Result}; + +use super::{ + BatchKnnCandidate, BatchKnnExtra, KNNVectorDistanceExec, QUERY_INDEX_COL, + knn_empty_result_schema, would_enter_heap, +}; + +/// Limit materialized vectors independently of batch width and refine factor. +/// Scoring gathers at most another chunk's worth of vectors for one query at a +/// time. A single vector wider than the budget is still read on its own. +const VECTOR_BATCH_BYTES: usize = 32 * 1024 * 1024; + +#[derive(Debug)] +pub struct BatchRefineExec { + dataset: Arc, + input: Arc, + query: Query, + query_count: usize, + rows_per_take: usize, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl BatchRefineExec { + pub(crate) fn try_new( + dataset: Arc, + input: Arc, + query: Query, + query_count: usize, + ) -> Result { + if query_count == 0 || query.k == 0 || !query.key.len().is_multiple_of(query_count) { + return Err(Error::invalid_input(format!( + "Batch refinement requires positive k and a key divisible by query_count: k={}, key_length={}, query_count={query_count}", + query.k, + query.key.len() + ))); + } + let (vector_type, element_type) = get_vector_type(dataset.schema(), &query.column)?; + let DataType::FixedSizeList(_, dim) = vector_type else { + return Err(Error::invalid_input(format!( + "Batch refinement requires fixed-size vectors, got {vector_type} for '{}'", + query.column + ))); + }; + let width = element_type.primitive_width().ok_or_else(|| { + Error::invalid_input(format!( + "Batch refinement requires primitive vector elements, got {element_type}" + )) + })?; + let row_bytes = (dim as usize) + .checked_mul(width) + .and_then(|bytes| bytes.checked_add(1)) + .ok_or_else(|| { + Error::invalid_input(format!( + "Vector byte size overflows: dimension={dim}, element_width={width}" + )) + })?; + if dim <= 0 || query.key.len() / query_count != dim as usize || query.metric_type.is_none() + { + return Err(Error::invalid_input(format!( + "Invalid batch refinement query: dimension={dim}, key_length={}, query_count={query_count}, metric={:?}", + query.key.len(), + query.metric_type + ))); + } + for (name, expected) in [ + (ROW_ID, DataType::UInt64), + (QUERY_INDEX_COL, DataType::Int32), + ] { + let schema = input.schema(); + let field = schema.field_with_name(name)?; + if field.data_type() != &expected { + return Err(Error::invalid_input(format!( + "Batch refinement requires {name}: {expected}, got {}", + field.data_type() + ))); + } + } + Ok(Self { + dataset, + input, + query, + query_count, + rows_per_take: (VECTOR_BATCH_BYTES / row_bytes).max(1), + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(knn_empty_result_schema(true)), + Partitioning::RoundRobinBatch(1), + EmissionType::Final, + Boundedness::Bounded, + )), + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + async fn refine( + self: Arc, + mut input: SendableRecordBatchStream, + vectors_read: Count, + read_batches: Count, + peak_vector_bytes: Count, + ) -> Result { + let mut heaps = (0..self.query_count) + .map(|_| BinaryHeap::::new()) + .collect::>(); + let projection = ProjectionRequest::from_columns( + [self.query.column.as_str(), ROW_ID], + self.dataset.schema(), + ); + let mut peak_bytes = 0; + while let Some(batch) = input.next().await { + let batch = batch?; + let query_count = self.query_count; + let candidates = spawn_cpu(move || -> Result> { + let ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("Batch refinement input missing _rowid"))? + .as_primitive::(); + let queries = batch + .column_by_name(QUERY_INDEX_COL) + .ok_or_else(|| Error::internal("Batch refinement input missing query_index"))? + .as_primitive::(); + let mut candidates = Vec::with_capacity(batch.num_rows()); + for (id, query) in ids.iter().zip(queries.iter()) { + let (Some(id), Some(query)) = (id, query) else { + return Err(Error::internal( + "Batch refinement input has NULL candidate identity", + )); + }; + if query < 0 || query as usize >= query_count { + return Err(Error::internal(format!( + "Candidate query_index={query} is outside query_count={query_count}" + ))); + } + candidates.push((id, query as usize)); + } + candidates.sort_unstable(); + Ok(candidates) + }) + .await?; + let mut offset = 0; + while offset < candidates.len() { + let start = offset; + let mut row_ids = + Vec::with_capacity(self.rows_per_take.min(candidates.len() - offset)); + while offset < candidates.len() { + let row_id = candidates[offset].0; + if row_ids.last() != Some(&row_id) { + if row_ids.len() == self.rows_per_take { + break; + } + row_ids.push(row_id); + } + offset += 1; + } + let batch = self.dataset.take_rows(&row_ids, projection.clone()).await?; + vectors_read.add(batch.num_rows()); + read_batches.add(1); + let vectors = + KNNVectorDistanceExec::resolve_vector_column(&batch, &self.query.column)?; + let bytes = vectors.get_array_memory_size(); + if bytes > peak_bytes { + peak_vector_bytes.add(bytes - peak_bytes); + peak_bytes = bytes; + } + // The ANN prefilter excludes deleted rows. Stable-ID lookup + // can also omit missing IDs: join by returned identity so an + // omitted row never shifts a candidate onto another vector. + let returned_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| Error::internal("Batch refinement take missing _rowid"))? + .as_primitive::(); + let positions: HashMap = returned_ids + .values() + .iter() + .enumerate() + .map(|(index, id)| (*id, index as u32)) + .collect(); + let mut groups = vec![Vec::new(); self.query_count]; + for &(id, query_index) in &candidates[start..offset] { + if let Some(&position) = positions.get(&id) { + groups[query_index].push(position); + } + } + let query = self.query.clone(); + let query_count = self.query_count; + let returned_ids = returned_ids.clone(); + let rows_per_take = self.rows_per_take; + heaps = spawn_cpu(move || -> Result<_> { + let metric = query + .metric_type + .ok_or_else(|| Error::internal("Batch refinement metric is missing"))?; + let dim = query.key.len() / query_count; + for (query_index, indices) in groups.into_iter().enumerate() { + if indices.is_empty() { + continue; + } + // Do not score the union against every query: only this + // query's original ANN candidates are eligible for top-k. + // Repeated candidates within one query must not expand + // the scoring buffer beyond the vector read budget. + for indices in indices.chunks(rows_per_take) { + let indices = UInt32Array::from(indices.to_vec()); + let selected = + arrow_select::take::take(vectors.as_ref(), &indices, None)?; + let key = query.key.slice(query_index * dim, dim); + let distances = metric.arrow_batch_func()( + key.as_ref(), + selected.as_fixed_size_list(), + )?; + let heap = &mut heaps[query_index]; + for (position, distance) in distances.iter().enumerate() { + let Some(distance) = distance else { continue }; + if distance.is_nan() + || query.lower_bound.is_some_and(|bound| distance < bound) + || query.upper_bound.is_some_and(|bound| distance >= bound) + { + continue; + } + let row_id = returned_ids.value(indices.value(position) as usize); + if would_enter_heap( + heap, + query.k, + distance, + row_id, + query_index as i32, + ) { + if heap.len() == query.k { + heap.pop(); + } + heap.push(BatchKnnCandidate { + query_index: query_index as i32, + distance, + row_id, + extra: BatchKnnExtra::RowIdOnly, + }); + } + } + } + } + Ok(heaps) + }) + .await?; + } + } + let mut results = heaps + .into_iter() + .flat_map(BinaryHeap::into_vec) + .collect::>(); + results.sort_unstable_by(|left, right| { + left.query_index + .cmp(&right.query_index) + .then_with(|| left.cmp(right)) + }); + Ok(RecordBatch::try_new( + self.schema(), + vec![ + Arc::new(Int32Array::from( + results + .iter() + .map(|candidate| candidate.query_index) + .collect::>(), + )), + Arc::new(Float32Array::from( + results + .iter() + .map(|candidate| candidate.distance) + .collect::>(), + )), + Arc::new(UInt64Array::from( + results + .iter() + .map(|candidate| candidate.row_id) + .collect::>(), + )), + ], + )?) + } +} + +impl DisplayAs for BatchRefineExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "BatchRefine: query_count={}, k={}, rows_per_take={}", + self.query_count, self.query.k, self.rows_per_take + ) + } +} + +impl ExecutionPlan for BatchRefineExec { + fn name(&self) -> &str { + "BatchRefineExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn schema(&self) -> SchemaRef { + knn_empty_result_schema(true) + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn required_input_distribution(&self) -> Vec { + vec![Distribution::SinglePartition] + } + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return Err(DataFusionError::Internal(format!( + "BatchRefineExec requires one child, got {}", + children.len() + ))); + } + let input = children.remove(0); + let mut node = Self::try_new( + self.dataset.clone(), + input, + self.query.clone(), + self.query_count, + )?; + node.rows_per_take = self.rows_per_take; + Ok(Arc::new(node)) + } + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let input = self.input.execute(partition, context)?; + let node = Arc::new(Self { + dataset: self.dataset.clone(), + input: self.input.clone(), + query: self.query.clone(), + query_count: self.query_count, + rows_per_take: self.rows_per_take, + properties: self.properties.clone(), + metrics: self.metrics.clone(), + }); + let vectors_read = self.metrics.new_count("refine_vectors_read", partition); + let read_batches = self.metrics.new_count("refine_read_batches", partition); + let peak_vector_bytes = self + .metrics + .new_count("refine_peak_vector_bytes", partition); + let result = async move { + node.refine(input, vectors_read, read_batches, peak_vector_bytes) + .await + .map_err(DataFusionError::from) + }; + Ok(Box::pin(InstrumentedRecordBatchStreamAdapter::new( + self.schema(), + stream::once(result).boxed(), + partition, + &self.metrics, + ))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::record_batch; + use arrow_array::types::Float32Type; + use lance_datagen::{ArrayGeneratorExt, BatchCount, Dimension, RowCount, array, gen_batch}; + use lance_linalg::distance::DistanceType; + use rstest::rstest; + + use crate::dataset::WriteParams; + use crate::io::exec::testing::TestingExec; + + #[rstest] + #[tokio::test] + async fn test_batch_refine_shared_candidates( + #[values(1, 100)] rows_per_take: usize, + #[values(false, true)] stable_row_ids: bool, + #[values("normal", "deleted", "bounded", "all_null", "empty")] scenario: &str, + ) { + let nulls = if scenario == "all_null" { + vec![true; 8] + } else { + vec![false, false, false, false, false, false, false, true] + }; + let data = gen_batch() + .col( + "vec", + array::cycle_vec( + array::cycle::(vec![ + 0.0, + 2.0, + 3.0, + 10.0, + 15.0, + f32::NAN, + 30.0, + 0.0, + ]), + Dimension::from(1), + ) + .with_nulls(&nulls), + ) + .col("id", array::step::()) + .into_reader_rows(RowCount::from(8), BatchCount::from(1)); + let mut dataset = Dataset::write( + data, + "memory://", + Some(WriteParams { + max_rows_per_file: 4, + enable_stable_row_ids: stable_row_ids, + ..Default::default() + }), + ) + .await + .unwrap(); + let ids = dataset + .scan() + .with_row_id() + .project::<&str>(&[]) + .unwrap() + .try_into_batch() + .await + .unwrap(); + let ids = ids[ROW_ID].as_primitive::().values(); + // Both queries share invalid rows. Each also has a better neighbor in + // the other query's candidates, which must never become eligible. + let mut candidates = record_batch!( + (QUERY_INDEX_COL, Int32, [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1]), + ( + ROW_ID, + UInt64, + [ + ids[0], ids[0], ids[2], ids[4], ids[5], ids[7], ids[1], ids[3], ids[6], ids[5], + ids[7] + ] + ) + ) + .unwrap(); + if scenario == "deleted" { + dataset.delete("id = 2").await.unwrap(); + // Physical-ID takes require live rows, as guaranteed by the ANN + // prefilter. Stable-ID takes can additionally omit missing IDs. + if !stable_row_ids { + let mask = arrow_array::BooleanArray::from_iter( + candidates[ROW_ID] + .as_primitive::() + .iter() + .map(|id| id.map(|id| id != ids[2])), + ); + candidates = arrow_select::filter::filter_record_batch(&candidates, &mask).unwrap(); + } + } + let query = Query { + column: "vec".to_string(), + key: Arc::new(Float32Array::from(vec![9.0, 0.0])), + k: 1, + lower_bound: (scenario == "bounded").then_some(5.0), + upper_bound: (scenario == "bounded").then_some(101.0), + minimum_nprobes: 1, + maximum_nprobes: Some(1), + ef: None, + refine_factor: Some(2), + metric_type: Some(DistanceType::L2), + use_index: true, + query_parallelism: 1, + dist_q_c: 0.0, + approx_mode: Default::default(), + }; + let input = if scenario == "empty" { + candidates.slice(0, 0) + } else { + candidates + }; + let input = Arc::new(TestingExec::new(vec![input])); + let mut node = BatchRefineExec::try_new(Arc::new(dataset), input, query, 2).unwrap(); + node.rows_per_take = rows_per_take; + let node = Arc::new(node); + let output = node + .execute(0, Arc::new(TaskContext::default())) + .unwrap() + .next() + .await + .unwrap() + .unwrap(); + if matches!(scenario, "empty" | "all_null") { + assert_eq!(output.num_rows(), 0); + } else { + assert_eq!( + output[QUERY_INDEX_COL].as_primitive::().values(), + &[0, 1] + ); + let expected = [ + ids[if scenario == "deleted" { 4 } else { 2 }], + ids[if scenario == "bounded" { 3 } else { 1 }], + ]; + assert_eq!( + output[ROW_ID].as_primitive::().values(), + &expected + ); + let distances = [36.0, if scenario == "bounded" { 100.0 } else { 4.0 }]; + assert_eq!( + output["_distance"].as_primitive::().values(), + &distances + ); + } + let metrics = node.metrics().unwrap(); + let read = metrics + .sum_by_name("refine_vectors_read") + .unwrap() + .as_usize(); + assert_eq!( + read, + if scenario == "empty" { + 0 + } else if scenario == "deleted" { + 7 + } else { + 8 + } + ); + let batches = metrics + .sum_by_name("refine_read_batches") + .unwrap() + .as_usize(); + assert_eq!( + batches, + if scenario == "empty" { + 0 + } else { + let candidates = if scenario == "deleted" && !stable_row_ids { + 7_usize + } else { + 8 + }; + candidates.div_ceil(rows_per_take) + } + ); + } +}