fix(fts): split tail-partition merge by the worker memory budget - #7601
Conversation
merge_all_tail_partitions folded every leftover worker builder into a single partition unconditionally; on a build whose workers never hit the flush threshold (large LANCE_FTS_PARTITION_SIZE), the entire index collapsed into one partition and queries lost all intra-partition parallelism. Fold with the same memory/doc-count checks as the partition-merge path instead, so the final partition count converges to total builder memory / budget.
Tail partitions hold most of the index when workers rarely hit the flush threshold; writing them sequentially serialized the posting-list compression of nearly the whole index behind one producer thread at the end of the build. Write them through buffer_unordered instead.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Xuanwo
left a comment
There was a problem hiding this comment.
Before merging, please make the partition-shape change safe for fuzzy queries. This PR increases the number of partitions in tail-heavy builds, but fuzzy expansion currently applies max_expansions inside each partition and then unions the expanded terms across partitions. That means the same query can expand to more terms, and return different results, solely because the tail was split.
We should either enforce the cap globally after combining partition expansions or explicitly change the contract, and add a regression test that builds multiple tail partitions and verifies fuzzy query behavior.
Review follow-up: build the same corpus as one partition and as a tail-split multi-partition layout through the real worker/budget path, and assert fuzzy OR/AND queries return identical (row_id, score) sets while max_expansions is not the binding constraint.
|
Thanks for the catch — the mechanism is exactly as you described: Handled in two parts:
|
…ns (#7634) Follow-up to the review discussion on #7601: fuzzy expansion previously ran inside every partition with a per-partition `max_expansions` cap, and the per-partition picks were unioned afterwards. The effective cap was `num_partitions × max_expansions`, so the same query could match more terms — and return different results — purely because of how the corpus happened to be partitioned (e.g. after a tail-heavy build splits its leftovers). ## What - **Expansion now runs once per query**, in `InvertedIndex::bm25_search`, under the query-wide `max_expansions` budget. For each query token the per-partition FST candidates merge into one lexicographically ordered set and the remaining budget takes a prefix of it. Each partition is only asked for its `remaining` lexicographically smallest candidates, which is lossless for that selection: any term among the merged lex-smallest `remaining` is also among its own partition's smallest `remaining`. - **Partitions receive the final token list** and no longer expand (`load_posting_lists` drops its `expand_fuzzy` call); `params.fuzziness` still drives the grouped fuzzy dedup/scoring semantics. - **The BM25 scorer reuses the same expansion** (`bm25_scorer_for_final_tokens`), so scorer terms and searched terms stay in lockstep and a query pays for one expansion instead of one for the scorer plus one per partition. - **Fuzzy AND/phrase keep their contract** — every original token position must retain at least one expansion — via an index-level check (previously implicit in each partition's own expansion). - `InvertedPartition::expand_fuzzy` stays for API compatibility, reimplemented on the shared per-token candidate collector; single-partition behavior is unchanged (same FST lexicographic order, same budget fill). ## Semantics Single-partition indexes behave exactly as before. Multi-partition indexes previously over-expanded in proportion to their partition count; they now honor the documented cap, and the expansion (hence the result set) is a pure function of the segment's vocabulary, independent of partition shape. ## Tests - `test_fuzzy_expansion_cap_is_global_across_partitions`: two partitions with disjoint variants, binding cap → exactly the 3 lexicographically smallest terms across both (fails on main, which returns 4). - `test_fuzzy_results_independent_of_partition_shape`: the same four docs built as one partition and as two return identical `(row_id, score)` sets under a binding cap (fails on main). - Existing fuzzy suite (grouped AND, position grouping, whole-query cap) unchanged and passing. The non-binding-cap regression test over the real tail-split builder path lives in #7601. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Yang Cen <yang@lancedb.com>
Two fixes for FTS tail partitions:
merge_all_tail_partitionsfolded every leftover worker builder into asingle partition unconditionally. On builds whose workers never hit the
flush threshold (large
LANCE_FTS_PARTITION_SIZE), the entire indexcollapsed into one partition and queries lost all intra-query
parallelism. Fold with the same memory/doc-count budget as the
partition-merge path, so the final partition count converges to
total builder memory / budget regardless of worker layout.
writing them sequentially serialized posting-list compression behind one
producer at the end of the build. Write them through buffer_unordered.
Standalone verification vs its merge-base
5M-doc controlled build,
LANCE_FTS_PARTITION_SIZE=128, 96 workers,per-branch-tip wheels:
tail partition (25 partitions total including flushed ones); with this fix
the same tails split into 23 budget-sized partitions (47 total).
Posting bytes essentially unchanged (3.59 vs 3.60 GB).
at this small scale; on the 200M-doc build, where tails hold most of the
data, the build went 2.5h → 2.0h.
Independent of the block-size/impact PR stack.
🤖 Generated with Claude Code