fix(datafusion): return vector_search rows in best-first relevance order#613
Open
XiaoHongbo-Hope wants to merge 6 commits into
Open
fix(datafusion): return vector_search rows in best-first relevance order#613XiaoHongbo-Hope wants to merge 6 commits into
XiaoHongbo-Hope wants to merge 6 commits into
Conversation
The `vector_search` table-valued function materialized its results with a
plain row-range scan, so `SELECT * FROM vector_search(...)` returned the
matching rows in physical file order instead of by relevance. This
contradicts the documented contract ("returns the top-k rows ordered by
relevance score") and means the first returned row is often not the most
similar one, especially at larger `limit`s.
Materialize the scored row-ids (best-first) via `execute_scored`, read the
matching rows (with projection pushdown, plus the internal `_ROW_ID`
column), then realign them to the index's relevance rank via a gather and
drop `_ROW_ID`. A scored row-id that is not materialized fails loud rather
than silently shrinking the top-k. The output schema is unchanged (no score
column exposed). `execute_read` is intentionally not reused: it is
statically `!Send` due to the primary-key-vector path, which the async
`TableProvider::scan` cannot hold.
Add a regression test that builds an ivf-flat vindex table and asserts the
returned order is `[2, 5, 1]` (relevance order), which differs from the
ascending-id file order `[1, 2, 5]` the old code produced, covering both the
projected (`SELECT id`) and full (`SELECT *`) paths.
XiaoHongbo-Hope
force-pushed
the
fix-vector-search-tvf-ordering
branch
from
July 26, 2026 13:00
11afe92 to
78913fd
Compare
The gather assembled the reordered result with the provider's Arrow schema (where string columns are DataFusion `Utf8View`) but kept the Paimon read's column types (`Utf8`), so a `vector_search` query returning a string column could fail on the type mismatch. Cast each column to the output field type when it differs, matching the normal scan path's `to_datafusion_batch`. Add a regression test with a `VARCHAR` column asserting the query executes and stays relevance-ordered.
This was referenced Jul 26, 2026
…dent test_vector_search_casts_string_column asserted a specific relevance order for a 3-row fixture, which is not stable across environments for such a tiny index (the exact-match hit can rank either first or second). Its purpose is the string-column schema cast, not ordering (covered by test_vector_search_orders_by_relevance), so assert the hit set order-independently.
The scan previously ignored DataFusion's pushed-down limit, so an outer `LIMIT n` over a large top-k (e.g. `vector_search(..., 1_000_000) LIMIT 1`) still searched, read, and materialized every row before truncating — a coordinator OOM risk. Bound the search by `min(top_k, outer_limit)` so the read and in-memory materialization only cover the rows the query can return. Add a test asserting an outer LIMIT returns the correctly-ranked top-n.
…correctly Two fixes to the vector_search scan: - The previous commit applied the outer LIMIT by shrinking the search's top-k (`with_limit(min(k, outer))`). That changes the ANN search width (list size, refine candidates), so `vector_search(..., 100) LIMIT 1` is not the first row of the top-100 search. Instead keep the full top-k for the search and truncate the already-ranked result before reading rows — the recall is unchanged and a large top-k with a small outer LIMIT only reads/materializes what it returns. - Move the search + read + rank-order gather out of `scan()` (planning time) into a custom single-partition `VectorSearchExec` whose stream does the work at execution time. Planning and `EXPLAIN` no longer read or materialize anything, and the work runs under DataFusion's `TaskContext`.
Reject non-empty children in with_new_children (the exec is a leaf) and reject a non-zero partition in execute (it has a single partition), instead of silently ignoring them.
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.
Makes the DataFusion
vector_searchtable function return rows in the index'sbest-first relevance order (previously it re-read rows in physical file order, so the
first row was often not the most similar).
What it does
execute_scored, reads the matchingrows (projection pushdown + internal
_ROW_ID), then gathers them back into relevancerank and drops
_ROW_ID. Output schema is unchanged.Utf8→ DataFusionUtf8View) so stringcolumns round-trip.
VectorSearchExec(not atplanning time), so
EXPLAINstays cheap and the work runs under theTaskContext.LIMITby truncating thealready-ranked result before reading, so recall is unchanged and a large top-k with a
small outer
LIMITonly reads what it returns.Dependency / known limitation — needs #614
This PR makes the scan honor the order of
SearchResult.row_ids, but that order is onlyfully correct once #614 is merged. On current
main,SearchResult::top_kdoes notsort when
candidate_count <= k(i.e.limit >= number of matched rows), so for thatcase the index returns rows in scored-map order and this scan faithfully reproduces it.
With #614, the ranking is correct for all limits.
Because of that, the string-column test asserts its hit set order-independently (the cast
is what it verifies); deterministic relevance ordering is covered by
test_vector_search_orders_by_relevance, which uses a top-k smaller than the candidatecount (the sorted/truncation path). Merge order: #614 first, then rebase this PR for
the full-ordering guarantee across all limits.