Skip to content

feat(index): share IVF partition scans across batch vector queries#2

Open
sezruby wants to merge 1 commit into
mainfrom
knn-batch-6822
Open

feat(index): share IVF partition scans across batch vector queries#2
sezruby wants to merge 1 commit into
mainfrom
knn-batch-6822

Conversation

@sezruby

@sezruby sezruby commented Jun 15, 2026

Copy link
Copy Markdown
Owner

Implements #6822: extend batch vector queries to indexed/ANN search. Rebased on latest main.

Summary

Batch vector search (#6821, PR lance-format#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 each time) and unioning the results. This PR makes the indexed/ANN path share index-level state across the batch: it reads each probed IVF partition's storage once and scores every query that probes it, with the prefilter built once and shared.

Approach

  • VectorIndex trait (lance-index): defaulted supports_batch_partition_search() + search_partitions_batch(...) (default returns not_supported), so non-IVF indices are explicitly unsupported.
  • 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, reusing accumulate_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 via build_dataset_prefilter.
  • Each query vector is normalized independently for cosine (normalize_batch_query_for_index).

Design notes (pre-empting review questions)

  • Why a new exec node instead of extending KNNVectorDistanceExec / the two ANN nodes? The two-node single-query pipeline streams one partition-list per delta through a per-query top-k. Sharing the scan requires inverting queries onto partitions and keeping one heap per query in a single pass — a different dataflow. The new node still reuses the underlying primitives (partition load, build_dataset_prefilter, and the index's per-partition accumulate), and the single-query nodes are untouched. Happy to fold it in differently if you'd prefer.
  • Why gate on the index-type string, not 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_type reads metadata with no I/O. The opened index re-checks the trait as a defensive invariant.
  • nprobes gate (correctness). The shared path searches exactly minimum_nprobes partitions/query. The single-query path is adaptive (early_pruning floor + late-search expansion), so it only matches when nprobes is fixed. The fast path is therefore gated to minimum_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 for you: fixed-nprobes-first with batched early/late as a follow-up, or the full adaptive path in one PR?
  • Memory. Peak = the union of probed partitions held during scoring — the same buffering the existing single-query global-heap path uses (search_partitions), widened to the batch's partition union. Per-delta output is k-bounded, so cross-delta accumulation is O(deltas × k), not O(nprobes × rows).

Fallback matrix (no regression)

Case Behavior
IVF_FLAT/PQ/SQ/RQ, fixed nprobes, fully indexed shared-scan fast path
adaptive nprobes / refine_factor / IVF_HNSW_* / mixed indexed+unindexed per-query indexed loop (exact)

Test plan

  • cargo test -p lance --lib test_batch_knn15 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: "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 shared build_dataset_prefilter).
  • cargo fmt --all && cargo clippy -p lance -p lance-index --tests --benches -- -D warnings.
  • Python: pytest -k batch (L2 + cosine × three/single queries); ruff clean; pyright clean on changed lines.
  • Benchmark (benchmarks/test_search.py): batch vs repeated-single ANN; standalone timing (50k rows, dim 128, IVF_PQ 64 partitions, m=32, k=10, nprobes=10) → 2.48× speedup.

Closes lance-format#6822

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>
@sezruby sezruby closed this Jul 6, 2026
@sezruby sezruby reopened this Jul 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Extend batch vector queries to ANN and indexed search

1 participant