Skip to content

Add distributed vector search with global top-k merging - #87

Open
everySympathy wants to merge 3 commits into
daft-engine:mainfrom
everySympathy:codex/distributed-vector-search
Open

everySympathy wants to merge 3 commits into
daft-engine:mainfrom
everySympathy:codex/distributed-vector-search

Conversation

@everySympathy

Copy link
Copy Markdown
Collaborator

An indexed nearest-neighbor query through read_lance currently executes in one scan task. This adds an opt-in vector_search API that assigns index segment groups to separate scan tasks and merges their candidates into a global top-k in the Daft execution plan. With the Ray runner, a single query can execute across workers; a query matrix receives an independent top-k for each query.

Behavior

  • Pin every task to the dataset version opened during planning and copy query vectors before lazy execution.
  • Require full index coverage and select one vector index by column, with an optional explicit index name.
  • Balance segment groups by live fragment row counts when grouping is needed, and perform one local k-nearest search per task.
  • Apply the API filter before nearest-neighbor selection, while keeping downstream DataFrame operations after global top-k selection.
  • Return _distance, _rowid, and _query_id. Results sort by query ID, distance, and row ID; equal-distance membership at the k boundary is selected by Lance and can vary with segment grouping. Approximate candidate refinement also runs per task and can differ from a monolithic scanner.

Benchmark

Adds a reproducible benchmark comparing native pylance, the existing Daft scan path, segment parallelism, and independent-query parallelism on the same snapshot and index. It controls local CPU/thread budgets and records raw durations, recall, environment, source hashes, and local process resource observations. Runs that fail the recall threshold do not write performance records.

Included local measurements are historical and precede the final candidate and grouping fixes. They do not measure the final implementation or establish a multi-node advantage. The benchmark guide and report document these limits and distinguish local process measurements from cluster resource costs.

Validation

  • Native full suite: 548 passed, 10 skipped, 2 xfailed, 2 xpassed.
  • Ray query, vector-index build, and concurrency suites: 73 passed.
  • Ruff format/check, strict mypy on daft_lance, and git diff --check passed.
  • Regression coverage includes bounded candidates for tied distances, query-buffer mutation before execution, uneven segment grouping, quoted field names, independent batch top-k, filters, namespace access, pinned snapshots, and incomplete coverage rejection.
  • The uneven-grouping regression fails when the old round-robin planner is substituted; a quoted-column benchmark smoke run also passed.

@FANNG1

FANNG1 commented Oct 13, 2026

Copy link
Copy Markdown
Contributor

Thanks for the PR, I will take a review

@FANNG1

FANNG1 commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

Performance gain

Segment splitting is not the best option in any case the benchmark covers: 1M single query 92ms vs 35ms for native pylance, batch 437ms vs 267ms for one task per query. And it is the most complex part of the change (balanced grouping, global top-k merge, coverage check, tie semantics). Which workload is it aimed at, and is there data backing it?

Splitting out batch queries to make review easier

Could the query matrix (multiple query vectors at once) move to a follow-up PR and leave only the single-query case here? WDYT?

Unindexed fragments

I would like to understand the trade-off behind it. lance itself permits partial results for a segment-only search.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants