feat(index): share IVF partition scans across batch vector queries#7640
Open
sezruby wants to merge 1 commit into
Open
feat(index): share IVF partition scans across batch vector queries#7640sezruby wants to merge 1 commit into
sezruby wants to merge 1 commit into
Conversation
github-actions
Bot
added
A-python
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
wjones127
self-requested a review
July 23, 2026 21:46
Extend batch vector search (lance-format#6821) to the indexed/ANN path so a single multi-query request reads each IVF partition's storage once and scores every query that probes it, instead of re-running a full single-query plan per vector and unioning the results (which re-opens the index and rebuilds the prefilter for each query). - Add `VectorIndex::search_partitions_batch` + `supports_batch_partition_search` (defaulted so non-IVF indices stay explicitly unsupported). - Implement them for `IVFIndex` with a flat-style sub-index (IVF_FLAT/PQ/SQ/RQ): load each distinct partition once and accumulate one top-k heap per query, sharing the prefilter across the whole batch. - Add `ANNIvfBatchExec`, which ranks every query against the centroids, runs the shared-scan batch search, merges per-query top-k across deltas, and emits `query_index`-tagged results; route to it from `Scanner::batch_indexed_vector_search` when the gate below holds. - Normalize each query vector independently for cosine (`normalize_batch_query_for_index`): normalizing the concatenated batch key with one global norm would scale each vector by a batch-composition-dependent factor and break equivalence with single-query search. The shared-scan fast path is gated to cases that are provably equivalent to repeated single-query search: fixed nprobes (`minimum_nprobes == maximum_nprobes`), no refine step, an IVF flat-style index, and fully-indexed fragments. With adaptive nprobes the single-query path applies an `early_pruning` floor and late-search expansion that the batch path does not, so those queries fall back to the per-query loop, which stays exact. HNSW, refine, and mixed indexed/unindexed scans also fall back. Tests: plan shape; exact batch-vs-repeated-single equivalence (nprobes pinned); cosine regression; shared prefilter; multi-delta cross-delta merge; and fallbacks for refine and adaptive nprobes. Python parametrized over L2 + cosine; a batch-vs-repeated-single ANN benchmark. Closes lance-format#6822 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
📝 WalkthroughWalkthroughAdds indexed ANN batch vector search with shared IVF partition reads, per-query result grouping, cosine-safe normalization, correctness-gated scanner planning, fallback coverage, equivalence tests, and batch-versus-repeated-search benchmarks. ChangesIndexed ANN batch search
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Scanner
participant ANNIvfBatchExec
participant IVFIndex
Client->>Scanner: batch nearest query
Scanner->>ANNIvfBatchExec: plan indexed batch search
ANNIvfBatchExec->>IVFIndex: search multiple queries across partitions
IVFIndex-->>ANNIvfBatchExec: per-query top-k results
ANNIvfBatchExec-->>Scanner: grouped query_index results
Scanner-->>Client: batch result table
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
python/python/tests/test_vector_index.py (1)
240-241: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSeed the RNG so equivalence failures are reproducible.
This test asserts exact
idequality between the shared-scan batch path and repeated single-query search over an approximate IVF_PQ index. With unseedednp.random.randn, any tie-break or ordering discrepancy (see theHashMapscoring-order note inrust/lance/src/index/vector/ivf/v2.rs) surfaces as an intermittent failure that cannot be replayed.♻️ Suggested change
+ rng = np.random.default_rng(42) scales = np.linspace(0.1, 10.0, query_count).reshape(-1, 1) - queries = (np.random.randn(query_count, 128) * scales).astype(np.float32) + queries = (rng.standard_normal((query_count, 128)) * scales).astype(np.float32)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/python/tests/test_vector_index.py` around lines 240 - 241, Seed NumPy’s random number generator before the randomized query construction in the test containing the shared-scan and repeated single-query equivalence assertions. Use a fixed, explicit seed so failures involving exact id ordering are deterministic and reproducible, while leaving the existing query generation and assertions unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/lance/src/dataset/scanner.rs`:
- Around line 4089-4096: Update the refine-factor gate in the batch
indexed-search path to reject any Some value, not only factors greater than one,
so refine(1) and refine(0) fall back to the per-query loop for reranking and
validation. Extend test_batch_knn_indexed_refine_falls_back to cover both refine
factors while preserving the existing refine(2) coverage.
- Around line 4137-4158: Update the batch fast-path gate in the surrounding
scanner method to fall back whenever
self.overlay_stale_vector_rows(index_segments)? is non-empty, before
constructing the batch executor. Preserve the existing batch behavior when no
stale overlay rows exist, and add coverage using an overlaid fragment that
verifies execution takes the per-query fallback path.
In `@rust/lance/src/index/vector/ivf/v2.rs`:
- Around line 1920-1938: The batch search flow around load_parallelism and the
try_collect::<Vec<_>>() currently retains every probed partition before scoring.
Replace this all-at-once collection with bounded-window processing, using
STREAMING_SEARCH_BATCH_SIZE or the existing bounded-channel pattern: load one
window, dispatch its scoring through spawn_cpu, release partition entries, then
continue until all assignments are processed while preserving result ordering
and error propagation.
- Around line 1902-1916: Update OrderedNode::cmp to use row_id as the final
comparison key after distance, ensuring equal-distance candidates are
deterministically selected and ordered. Preserve the existing heap behavior and
make the resulting global_heap_to_batch output match repeated single-query
ordering.
In `@rust/lance/src/io/exec/knn.rs`:
- Around line 2401-2411: Update the per-query batch handling around the loop
over per_query: first validate that the returned batch count matches the query
count, returning a diagnosable Internal error on mismatch before indexing
candidates. Then replace positional batch.column access with column_by_name
using the schema’s distance and row-id field names, preserving the existing
candidate extension behavior.
---
Nitpick comments:
In `@python/python/tests/test_vector_index.py`:
- Around line 240-241: Seed NumPy’s random number generator before the
randomized query construction in the test containing the shared-scan and
repeated single-query equivalence assertions. Use a fixed, explicit seed so
failures involving exact id ordering are deterministic and reproducible, while
leaving the existing query generation and assertions unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: QUIET
Plan: Pro Plus
Run ID: c6efabac-70c1-4356-8665-8f6423a9f771
📒 Files selected for processing (6)
python/python/benchmarks/test_search.pypython/python/tests/test_vector_index.pyrust/lance-index/src/vector.rsrust/lance/src/dataset/scanner.rsrust/lance/src/index/vector/ivf/v2.rsrust/lance/src/io/exec/knn.rs
Comment on lines
+4089
to
+4096
| if matches!(q.refine_factor, Some(rf) if rf > 1) { | ||
| 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) { | ||
| return Ok(false); | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
The refine gate lets refine(1) and refine(0) onto the non-reranking batch path.
The doc above states the requirement as "no refine step (the batch path does not yet rerank)", but the predicate only rejects rf > 1:
refine(1): the single-query path still re-ranks with the original vectors (if q.refine_factor.is_some()at Line 3989, andScanner::refinedocuments "even if the factor is 1, the results will still be re-ranked"). The batch path performs no rerank, so distances/order diverge from repeated single-query search on quantized indexes (IVF_PQ/SQ/RQ) — exactly what this gate exists to prevent.refine(0): the single-query path rejects it withRefine factor cannot be zero(Line 3963). The batch path skips that validation and reaches the index withheap_capacity = query.k * 0, silently returning empty results instead of an error.
Gating on is_some() fixes both: zero/one refine factors fall back to the per-query loop, which reranks and validates.
test_batch_knn_indexed_refine_falls_back only exercises refine(2), so neither case is covered — worth extending it with refine(1) and refine(0).
🐛 Proposed fix
- if matches!(q.refine_factor, Some(rf) if rf > 1) {
+ // Any refine factor (including 1) means the single-query path reranks
+ // with the original vectors, and 0 is rejected outright there; the
+ // shared-scan path does neither, so fall back.
+ if q.refine_factor.is_some() {
return Ok(false);
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if matches!(q.refine_factor, Some(rf) if rf > 1) { | |
| 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) { | |
| return Ok(false); | |
| } | |
| // Any refine factor (including 1) means the single-query path reranks | |
| // with the original vectors, and 0 is rejected outright there; the | |
| // shared-scan path does neither, so fall back. | |
| 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) { | |
| return Ok(false); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/scanner.rs` around lines 4089 - 4096, Update the
refine-factor gate in the batch indexed-search path to reject any Some value,
not only factors greater than one, so refine(1) and refine(0) fall back to the
per-query loop for reranking and validation. Extend
test_batch_knn_indexed_refine_falls_back to cover both refine factors while
preserving the existing refine(2) coverage.
Comment on lines
+4137
to
+4158
| // 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 | ||
| // back to the per-query loop below, which never regresses behavior. | ||
| 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); | ||
| let prefilter_source = self | ||
| .prefilter_source(filter_plan, self.get_indexed_frags(index_segments)) | ||
| .await?; | ||
| return new_knn_batch_exec( | ||
| self.dataset.clone(), | ||
| index_segments, | ||
| &batch_query, | ||
| self.nearest_query_count, | ||
| prefilter_source, | ||
| ); | ||
| } |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Batch path skips data-overlay stale-row reconciliation.
For the per-query path, vector_search computes overlay_stale_vector_rows / stale_rows_block_mask and passes an overlay_block into the ANN node, then re-scores those rows via knn_combined (Lines 3968-4012). The batch fast path builds its prefilter without an overlay block, and ANNIvfBatchExec::execute passes None explicitly (rust/lance/src/io/exec/knn.rs Lines 2334-2336, "The batch node has no data overlay to reconcile against"). Neither batch_index_search_supported nor this method checks for overlays, so a dataset whose fragments carry overlays committed after the vector index will get stale index entries from the batch path while the single-query path masks and re-scores them.
Simplest correctness-preserving fix: make the gate fall back when self.overlay_stale_vector_rows(index_segments)? is non-empty (it short-circuits cheaply when no fragment has overlays), and add a test with an overlaid fragment asserting the fallback.
🛡️ Suggested gate addition
if self.fast_search {
return Ok(true);
}
+ // Overlay-stale rows must be blocked from the index and re-scored on a
+ // flat path; the batch node does neither, so fall back when any exist.
+ if !self.overlay_stale_vector_rows(index_segments)?.is_empty() {
+ return Ok(false);
+ }
// The batch node only searches indexed partitions, so any unindexed🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/dataset/scanner.rs` around lines 4137 - 4158, Update the batch
fast-path gate in the surrounding scanner method to fall back whenever
self.overlay_stale_vector_rows(index_segments)? is non-empty, before
constructing the batch executor. Preserve the existing batch behavior when no
stale overlay rows exist, and add coverage using an overlaid fragment that
verifies execution takes the per-query fallback path.
Comment on lines
+1902
to
+1916
| // Invert the per-query partition lists so each distinct partition is | ||
| // loaded once and scored against every query that probes it. | ||
| let mut assignments: HashMap<u32, Vec<(usize, f32)>> = HashMap::new(); | ||
| for (query_index, (parts, dists)) in partitions_per_query | ||
| .iter() | ||
| .zip(q_c_dists_per_query.iter()) | ||
| .enumerate() | ||
| { | ||
| for (part_id, dist_q_c) in parts.values().iter().zip(dists.values().iter()) { | ||
| assignments | ||
| .entry(*part_id) | ||
| .or_default() | ||
| .push((query_index, *dist_q_c)); | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Does the global top-k heap impose a total order that includes the row id?
rg -nP -C6 'struct OrderedNode|impl .*Ord.* for OrderedNode' --type=rust
ast-grep run --pattern 'fn global_heap_to_batch($$$) { $$$ }' --lang rust rust/lance/src/index/vector/ivf/v2.rsRepository: lance-format/lance
Length of output: 156
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files =="
git ls-files 'rust/lance/src/index/vector/ivf/v2.rs' 'rust/lance/src/dataset/scanner.rs' | sed -n '1,20p'
echo "== outline around top of v2.rs =="
wc -l rust/lance/src/index/vector/ivf/v2.rs rust/lance/src/dataset/scanner.rs
ast-grep outline rust/lance/src/index/vector/ivf/v2.rs --length 200 | sed -n '1,200p'
echo "== search relevant terms =="
rg -n "assignments|partitions_per_query|global_heap_to_batch|TestVectorDataset|assert_batch_matches_single_queries|OrderedNode|OrderedCandidate|global_top|topk" rust/lance/src -S
echo "== lines 1860-1945 in ivf/v2.rs =="
sed -n '1860,1945p' rust/lance/src/index/vector/ivf/v2.rs
echo "== lines around scanner tests =="
rg -n "assert_batch_matches_single_queries|TestVectorDataset" rust/lance/src/dataset/scanner.rs -A 80 -B 20Repository: lance-format/lance
Length of output: 548
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== search relevant terms =="
rg -n "assignments|partitions_per_query|global_heap_to_batch|TestVectorDataset|assert_batch_matches_single_queries|OrderedNode|OrderedCandidate|global_top|topk" rust/lance/src -S || true
echo "== lines 1860-1945 in ivf/v2.rs =="
sed -n '1860,1945p' rust/lance/src/index/vector/ivf/v2.rs
echo "== lines around scanner tests =="
rg -n "assert_batch_matches_single_queries|TestVectorDataset" rust/lance/src/dataset/scanner.rs -A 100 -B 20 || true
echo "== import headers for v2.rs =="
sed -n '1,120p' rust/lance/src/index/vector/ivf/v2.rsRepository: lance-format/lance
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== ivf/v2.rs relevant scoring/heap/export lines =="
sed -n '800,935p' rust/lance/src/index/vector/ivf/v2.rs
sed -n '1945,2025p' rust/lance/src/index/vector/ivf/v2.rs
echo "== OrderedNode definition and Ord impl =="
rg -n -C 12 'struct OrderedNode|BinaryHeap<OrderedNode|Ord for Ordered|impl PartialOrd|impl PartialEq' rust/lance/src -S | sed -n '1,240p'
echo "== relevant tests around batch indexed equivalence =="
sed -n '6950,7825p' rust/lance/src/dataset/scanner.rsRepository: lance-format/lance
Length of output: 50374
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== exact OrderedNode references in graph submodule =="
rg -n -C 20 'OrderedNode' rust/lance/src/index/vector/graph.rs rust/lance/src/index/vector -S | sed -n '1,260p'
echo "== ivf.rs OrderedNode/topk area =="
sed -n '3980,4095p' rust/lance/src/index/vector/ivf.rs
echo "== Behavioral probe: Rust BinaryHeap Ord ties (if cargo availability) =="
if cargo --version >/dev/null 2>&1 && rustc --version >/dev/null 2>&1; then
tmp="$(mktemp -d)"
cat > "$tmp/heap_order_probe.rs" <<'RS'
fn main() {
let mut heap = std::collections::BinaryHeap::new();
for row_id in 0u64..10 {
heap.push((0.0f32, row_id));
}
let out: Vec<_> = heap.into_iter().map(|(d, id)| format!("({}, {})", d, id)).collect();
for (i, item) in out.iter().enumerate() {
println!("heap_order_probe.{}. {}", i, item);
}
}
RS
cargo run --manifest-path "$tmp/HeapOrderProbeCargo.toml" 2>/dev/null || (
cat > "$tmp/Cargo.toml" <<'TOML'
[package]
name = "heap_order_probe"
version = "0.0.0"
edition = "2021"
TOML
mkdir -p "$tmp/src"
mv "$tmp/heap_order_probe.rs" "$tmp/src/main.rs"
cargo run --manifest-path "$tmp/Cargo.toml"
)
else
echo "cargo/rustc unavailable"
fiRepository: lance-format/lance
Length of output: 14303
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate graph.rs =="
git ls-files | rg '(^|/)graph\.rs$|lance_index/src|index/vector' | sed -n '1,120p'
echo "== graph.rs OrderedNode definitions/usages =="
while IFS= read -r f; do
rg -n -C 8 'struct OrderedNode|impl .*OrderedNode|derive\(.*Ord.*\)|impl PartialOrd|impl Ord|BinaryHeap<OrderedNode' "$f"
done < <(git ls-files | rg '(^|/)graph\.rs$')
echo "== ivf storage topk calls into graph =="
rg -n -C 12 'accumulate_topk_with_scratch|topk|OrderedNode' rust/lance/src/index/vector -S | sed -n '1,260p'Repository: lance-format/lance
Length of output: 1335
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== rust/lance-index/src/vector/graph.rs OrderedNode area =="
sed -n '1,180p' rust/lance-index/src/vector/graph.rs
echo "== binary heap tie ordering probe =="
python3 - <<'PY'
# Rust std::collections::BinaryHeap iterates from the maximum heap element;
# equal-order elements can emerge in different relative orders.
import random
heap = []
for row_id in range(9):
random.shuffle(heap)
heap.append((0.0, row_id))
heap.sort(reverse=True)
seen = set()
for _ in range(200):
q = heap.copy()
out = []
while q:
_, row_id = q.pop()
out.append(row_id)
seen.add(tuple(out))
print("seen orderings:", len(seen))
PYRepository: lance-format/lance
Length of output: 4451
Use a deterministic tie-break for batch top-k results.
OrderedNode only orders by distance, so equal-distance candidates are not deterministically selected or ordered before global_heap_to_batch emits them. With ties present, the shared batch path can produce the same top-k set in a different order than repeated single-query search, causing the indexed equivalence assertions to become flaky. Add the row id as the final comparison key in OrderedNode::cmp, or sort the heap output by (distance, row_id) before returning the batch.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/index/vector/ivf/v2.rs` around lines 1902 - 1916, Update
OrderedNode::cmp to use row_id as the final comparison key after distance,
ensuring equal-distance candidates are deterministically selected and ordered.
Preserve the existing heap behavior and make the resulting global_heap_to_batch
output match repeated single-query ordering.
Comment on lines
+1920
to
+1938
| // Load each distinct partition's storage exactly once. This shared I/O | ||
| // is the whole point of batch search versus repeated single queries. | ||
| let load_parallelism = get_num_compute_intensive_cpus().max(1); | ||
| let load_index = self.clone(); | ||
| let load_metrics = metrics.clone(); | ||
| let loaded = stream::iter(assignments) | ||
| .map(move |(part_id, probing_queries)| { | ||
| let index = load_index.clone(); | ||
| let metrics = load_metrics.clone(); | ||
| async move { | ||
| let part_entry = index | ||
| .load_partition(part_id as usize, true, metrics.as_ref()) | ||
| .await?; | ||
| Result::Ok((part_id as usize, part_entry, probing_queries)) | ||
| } | ||
| }) | ||
| .buffered(load_parallelism) | ||
| .try_collect::<Vec<_>>() | ||
| .await?; |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
All probed partitions are materialized before any scoring — unbounded memory for wide batches.
try_collect::<Vec<_>>() holds every distinct probed partition's storage alive at once. For a batch of q queries with nprobes probes each, that is up to min(q * nprobes, num_partitions) partitions resident, and the Arc clones also pin those entries in the index cache. The single-query path deliberately avoids this by streaming prepared partitions through a bounded channel (STREAMING_SEARCH_BATCH_SIZE, Line 1663) after #7642.
Consider processing the load stream in chunks (load a bounded window, dispatch scoring for it via spawn_cpu, drop the entries, repeat) so peak memory is bounded by the window rather than by the probe fan-out.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/index/vector/ivf/v2.rs` around lines 1920 - 1938, The batch
search flow around load_parallelism and the try_collect::<Vec<_>>() currently
retains every probed partition before scoring. Replace this all-at-once
collection with bounded-window processing, using STREAMING_SEARCH_BATCH_SIZE or
the existing bounded-channel pattern: load one window, dispatch its scoring
through spawn_cpu, release partition entries, then continue until all
assignments are processed while preserving result ordering and error
propagation.
Comment on lines
+2401
to
+2411
| for (query_index, batch) in per_query.into_iter().enumerate() { | ||
| let dists = batch.column(0).as_primitive::<Float32Type>(); | ||
| let row_ids = batch.column(1).as_primitive::<UInt64Type>(); | ||
| candidates[query_index].extend( | ||
| dists | ||
| .values() | ||
| .iter() | ||
| .copied() | ||
| .zip(row_ids.values().iter().copied()), | ||
| ); | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Use column_by_name and validate the returned batch count.
Two issues in this loop:
batch.column(0)/batch.column(1)bind result columns by position. As per coding guidelines,Use column_by_name() for RecordBatch column access in production code; positional access silently mis-reads distances as row ids ifVECTOR_RESULT_SCHEMAfield order ever changes.per_querycomes from a trait method whose contract is "one batch per query". If an implementation returns more batches thanquery_count,candidates[query_index]panics inside the executor. A length check with anInternalerror keeps the failure diagnosable.
♻️ Proposed fix
+ if per_query.len() != query_count {
+ return Err(DataFusionError::Internal(format!(
+ "ANNIvfBatchExec: index {} returned {} result batches for {query_count} queries",
+ index_meta.uuid,
+ per_query.len()
+ )));
+ }
for (query_index, batch) in per_query.into_iter().enumerate() {
- let dists = batch.column(0).as_primitive::<Float32Type>();
- let row_ids = batch.column(1).as_primitive::<UInt64Type>();
+ let dists = batch
+ .column_by_name(DIST_COL)
+ .ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "ANNIvfBatchExec: batch partition search result missing {DIST_COL}"
+ ))
+ })?
+ .as_primitive::<Float32Type>();
+ let row_ids = batch
+ .column_by_name(ROW_ID)
+ .ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "ANNIvfBatchExec: batch partition search result missing {ROW_ID}"
+ ))
+ })?
+ .as_primitive::<UInt64Type>();
candidates[query_index].extend(📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| for (query_index, batch) in per_query.into_iter().enumerate() { | |
| let dists = batch.column(0).as_primitive::<Float32Type>(); | |
| let row_ids = batch.column(1).as_primitive::<UInt64Type>(); | |
| candidates[query_index].extend( | |
| dists | |
| .values() | |
| .iter() | |
| .copied() | |
| .zip(row_ids.values().iter().copied()), | |
| ); | |
| } | |
| if per_query.len() != query_count { | |
| return Err(DataFusionError::Internal(format!( | |
| "ANNIvfBatchExec: index {} returned {} result batches for {query_count} queries", | |
| index_meta.uuid, | |
| per_query.len() | |
| ))); | |
| } | |
| for (query_index, batch) in per_query.into_iter().enumerate() { | |
| let dists = batch | |
| .column_by_name(DIST_COL) | |
| .ok_or_else(|| { | |
| DataFusionError::Internal(format!( | |
| "ANNIvfBatchExec: batch partition search result missing {DIST_COL}" | |
| )) | |
| })? | |
| .as_primitive::<Float32Type>(); | |
| let row_ids = batch | |
| .column_by_name(ROW_ID) | |
| .ok_or_else(|| { | |
| DataFusionError::Internal(format!( | |
| "ANNIvfBatchExec: batch partition search result missing {ROW_ID}" | |
| )) | |
| })? | |
| .as_primitive::<UInt64Type>(); | |
| candidates[query_index].extend( | |
| dists | |
| .values() | |
| .iter() | |
| .copied() | |
| .zip(row_ids.values().iter().copied()), | |
| ); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/lance/src/io/exec/knn.rs` around lines 2401 - 2411, Update the per-query
batch handling around the loop over per_query: first validate that the returned
batch count matches the query count, returning a diagnosable Internal error on
mismatch before indexing candidates. Then replace positional batch.column access
with column_by_name using the schema’s distance and row-id field names,
preserving the existing candidate extension behavior.
Source: Coding guidelines
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Closes #6822.
Rationale for this change
Batch vector search (#6821, PR #6828) made indexed multi-query search work by looping the full single-query plan once per query vector — re-opening the index and rebuilding the prefilter for every query — and unioning the results. That throws away the main opportunity of a batch: query vectors mostly probe overlapping IVF partitions, so the same partition storage is read and decoded many times. This PR shares that index-level state across the batch.
What changes are included in this PR?
The indexed/ANN path now reads each probed IVF partition's storage once and scores every query that probes it, with the prefilter built once and shared.
VectorIndextrait (lance-index): defaultedsupports_batch_partition_search()+search_partitions_batch(...)(default returnsnot_supported). Additive — no breaking change.IVFIndex(ivf/v2.rs): batch search for flat-style sub-indices (IVF_FLAT/PQ/SQ/RQ). Invert per-query partition lists, load each distinct partition once, accumulate one top-k heap per query, reusingaccumulate_prepared_partition_search/global_heap_to_batch.ANNIvfBatchExec(io/exec/knn.rs): ranks each query against the centroids, runs the shared-scan batch search per delta, merges per-query top-k across deltas, emits{query_index, _distance, _rowid}. Prefilter wiring shared with the single-query node viabuild_dataset_prefilter.normalize_batch_query_for_index).Fallback matrix (no regression):
refine_factor/ IVF_HNSW_* / mixed indexed+unindexedDesign notes (anticipating review):
build_dataset_prefilter, the index's per-partition accumulate) and leaves the single-query nodes untouched.supports_batch_partition_search()? The gate is a planning-time decision and the single-query path likewise doesn't open the index there;derive_vector_index_typereads metadata with no I/O. The opened index re-checks the trait as a defensive invariant.minimum_nprobespartitions/query, but the single-query path is adaptive (early_pruningfloor + late-search expansion), so the two only match when nprobes is fixed. The fast path is gated tominimum_nprobes == maximum_nprobes; adaptive nprobes falls back to the per-query loop (verified: an unpinned batch diverged on every query before the gate, 0 divergence after). Open question: fixed-nprobes-first with batched early/late as a follow-up, or the full adaptive path in one PR?search_partitions) uses, widened to the batch's partition union; per-delta output is k-bounded, so cross-delta accumulation is O(deltas × k).Are these changes tested?
Yes.
cargo test -p lance --lib test_batch_knn— 15 tests: plan shape, exact batch-vs-repeated-single equivalence (nprobes pinned), cosine regression, shared prefilter, multi-delta cross-delta merge, and explicit fallbacks for refine, adaptive nprobes, and IVF_HNSW (acceptance criterion: "unsupported index types have explicit behavior and tests").cargo test -p lance --lib dataset::scanner::test::test_knn(29) — no single-query regression (exercises the sharedbuild_dataset_prefilter).cargo fmt --all&&cargo clippy -p lance -p lance-index --tests --benches -- -D warnings.pytest -k batch(L2 + cosine × three/single queries); index-type sweep (IVF_FLAT/PQ/SQ/HNSW_PQ/HNSW_SQ) matches repeated single-query;ruffclean;pyrightclean on changed lines.benchmarks/test_search.py): batch vs repeated-single ANN; local run (50k rows, dim 128, IVF_PQ 64 partitions, m=32, k=10, nprobes=10) → ~2.5× faster (ratio; absolute ms machine-dependent).Are there any user-facing changes?
No API changes. The batch-query surface (2-D / list query with a
query_indexoutput column) already shipped in #6828; this PR only changes how indexed batch queries execute, so Python and Java bindings are unaffected. It is a performance improvement for fixed-nprobes indexed batch search; results are identical to issuing the queries one at a time. No format or compatibility change.