feat(index): split IVF partitions to target size and join undersized ones in one optimize pass - #9051
Conversation
…nt in one optimize pass Incremental IVF maintenance (the split / join / reassign pass of `optimize_indices`) now converges in one pass and reads far fewer raw vectors: - An oversized partition is split `ceil(rows / target)` ways by one balanced k-means on `256 x ways` sampled rows instead of being halved once per call, so a partition many times the target no longer needs one optimize per halving. Splits whose sampled rows are too alike to separate are skipped instead of producing empty partitions. - Of the 64 neighbor partitions of a split, only those where a sampled row is now closer to a new centroid than to its own (SPFresh's necessary condition for a row to move, with a 5% margin) are re-read and reassigned; the others keep their rows and codes. - Every undersized partition is joined away in one pass, each vector going to the nearest remaining partition among the 64 nearest neighbors of its old centroid. In steady state, joins are decided on partition sizes summed over all segments, so a delta segment's naturally small partitions are no longer join candidates on their own. - Split and join thresholds use the `target_partition_size` the index was created with, read from the persisted index details, instead of the index type's default. - Raw vectors are fetched with `io_parallelism` concurrent `take_rows` chunks, split samples are uniform and seeded per partition, and an optimize whose split candidates turned out empty no longer forces a merge. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
47a0702 to
a888540
Compare
…vy hubs safe Address the review of the one-pass split/join optimize: - Joins keep enough partitions for the rows to average the target size, so joining every undersized partition into the largest one cannot rebuild a partition above the split threshold (20 x 24 rows at target 100 made 480). - Joined partitions are loaded one at a time and quantized before the next is read, and the quantized rows of one pass are capped by JOIN_BYTES_BUDGET; memory no longer scales with the total of every removed partition. - A split of more than 256 ways retries with flat k-means when hierarchical training refuses a duplicate-heavy sample, instead of failing the optimize. - Steady-state joins need every segment merged into one, so they are only offered when the segments share an IVF model. The template append helper in the v2 tests now drifts rows like the other fixtures: a ceil(rows / target)-way split of identical rows left pieces empty at random. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The minimum-kept rule only bounded the global average: joined rows go to the nearest surviving centroid, so many small partitions around one survivor could still rebuild a partition above the split threshold (thirty coinciding 24-row partitions made 744 rows at threshold 400). The planner now projects each candidate onto its nearest survivor and skips it when that destination would exceed the threshold. A destination that is still undersized may be joined away in a later round with its received rows re-projected exactly, so a handful of tiny partitions still converges to one in a single optimize. A split's new centroid is now kept only if a sampled row is nearer to it than to its siblings and the nearby old centroids, with the transformer's own assignment routine; a centroid that only won rows a neighbor keeps produced an empty partition. The optimize tests build their index on fixed centroids so their layout no longer depends on k-means seeding. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…by row The planner reserved a joined partition's rows against the survivor nearest its centroid, but rows are assigned one by one and can land on a different survivor, which could still end up above four times the target. The join now gives every row to its nearest neighbor that has room below the split threshold, from the partition sizes the pass was planned on; only when all neighbors are full does a row go to the nearest one, with a warning. The planner's projection is an estimate on top of that. Each partition's nearest surviving neighbors are cached and refetched only once all of them are gone, so a chain of joins (512 empty partitions) plans in milliseconds instead of recomputing distances to every centroid at every step. The rows projected onto a partition move along when it is joined away, replacing the exact re-projection of each contributor. The old per-row reassign helper and its enum had no callers left. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…full A joined partition's rows only looked at the 64 nearest survivors; an earlier join could fill that whole window, and the fallback then knowingly placed rows in a full partition (a 410-row survivor at threshold 400). Rows now fall back to the nearest surviving partition anywhere with room, and the planner keeps at least ceil(rows / threshold) survivors so that room always exists; only an index already at the threshold everywhere lands a row in a full partition, which the next optimize splits. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
A multivector row can have vectors in several partitions. Joining one of them loaded the whole row from the dataset and reassigned every vector, so the vectors that surviving partitions keep were added a second time, and the room and byte accounting counted fewer entries than the join emitted. The loader now reports how many entries of each row the partition held, and for a row loaded with more vectors than that only the ones this centroid wins against its 64 nearest neighbors by the largest margin are reassigned. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ined bytes Inferring which vectors of a multivector row a partition held from distances to a bounded centroid window could pick the wrong one, dropping a removed vector and duplicating a surviving one. A join now reindexes every logical row with an entry in a joined partition exactly once: all of its vectors are placed anew and its old entries leave every partition through a shared set carried by the join adjustment, so identity and count are both exact for any quantizer. The planner's per-row byte estimate summed only the code column, which for RaBitQ omits the factor and extra-code columns (20 of 65 bytes). It now sums the quantizer's output fields with saturating arithmetic, and the join also measures the bytes it actually retains and stops joining further partitions once the budget is reached; the remaining planned joins wait for the next optimize. Rows are quantized under the old centroid numbering so this can happen one partition at a time. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…them Whole-row reindexing was inconsistent with the state it ran against: room ignored the entries the reindexed rows free in survivors, a byte-budget stop after routing left candidate exclusion and deletion disagreeing with the kept set, one partition's expansion to every vector of its rows was unbounded, and RaBitQ's partition transformer reassigned rows instead of honoring the chosen destination. The join now gathers the rows to reindex from every joined partition first, counts survivor room after those rows' entries are gone (reading survivors only on multivector indexes, where a joined row has entries elsewhere), and routes and quantizes chunk by chunk against that one state, streaming the batches through a shuffler to temp files like the split path; the byte budget and its estimate are gone. The partition transformer keeps a supplied partition id and measures the distance to it rather than reassign, so RaBitQ's residual and factors match the partition the row is stored in. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…wn window A multivector row reindexed under one joined partition had all of its vectors routed through that partition's 64-neighbor window, so a vector from a distant region landed in a wrong partition whenever the window had room, and a query probing its exact centroid could miss it. Each vector of a multivector row now gets the window of its own nearest surviving centroids, the same cost as indexing it in the first place; single-vector rows keep their partition's window. Multivector rows are also fetched 64 at a time with one chunk in flight, and the flattened vectors are routed and quantized 1,024 at a time, so the join's working set no longer grows with how many vectors the fetched rows expand to; a single oversized row is the one irreducible case. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
A join fetched 64 multivector rows at a time, so several rows with many vectors each could exhaust memory before the 1,024-vector routing cap or the shuffler's backpressure took effect. Rows are now regrouped by their decoded size: a fixed-size vector column is chunked to a 32 MiB budget from its known row size, a multivector column is fetched one row at a time with four fetches in flight and regrouped by measured bytes, and only a single row larger than the budget ever exceeds it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…othing An ablation of every optimization in the one-pass split/join (each switched off on its own, two hub scenarios, three seeds of the split's k-means) left only the k-way split and the batched join moving the outcome: - Sampled neighbor pruning skipped 9-16% of partition re-reads but the 512-row samples cost more time than they saved (hub 15.5 s vs 14.5 s without), with identical recall. Every neighbor of a split is reassigned again, as before. - The planner's projection of join destinations (nearest surviving neighbor, cached, multi-round) produced the same joins as the plain planner in every run; the row-level room check in the join already carries the bound. The planner now joins the smallest undersized partitions first, as many as leave ceil(rows / threshold) survivors. - Seeded uniform sampling of the split's training rows and the balance factor on its k-means were within noise or slightly worse (fewer undersized pieces without the balance factor). The split's k-means is seeded per partition instead, so a rerun trains the same split and the results above are reproducible. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
Ablation update (be8df61): I switched off each optimization in this PR on its own (env-gated build, not part of the PR) and measured two hub scenarios on the 6M-row DEEP base with three seeds of the split's k-means each. Only the k-way split and the batched join move the outcome; four parts were removed:
The split's k-means is now seeded per partition, so the branch is deterministic across runs. The Results section of the description has the new main-vs-PR numbers on this harness (hub: not converged after 6 calls / 60.7 s on main vs 2 calls / 14.2 s here, recall unchanged) and the full ablation table. PR title adjusted since reassignment is no longer pruned. |
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The ablation update shows that only k-way splitting and batched joining move the measured outcome, and this revision removes four non-contributing mechanisms. The remaining implementation still enforces join capacity during row routing, routes multivectors independently, bounds decoded fetches before flattening, and seeds k-means per partition, preserving the accepted convergence, recall, and working-set contract with less planning complexity.
What changed?
Incremental IVF maintenance (
optimize_indiceson an IVF_FLAT/SQ/PQ/RQ index, the SPFresh-style split / join / reassign pass inrust/lance/src/index/vector/builder.rs) now converges in one pass and reads far fewer raw vectors:4 x targetis splitceil(rows / target)ways (capped at 1024) by one k-means on256 x wayssampled rows (seeded per partition, so a rerun trains the same split), instead of one 2-way split per optimize call. A 1.5M-row partition with a 4096-row target needed ~7 optimize passes (each re-reading raw vectors); it now needs one. Above 256 ways the trainer is hierarchical, which refuses a sample with fewer distinct rows than pieces; the split then retries with flat k-means and keeps only the centroids that win sampled rows against their siblings and the 64 nearest old centroids (the same assignment the IVF transformer makes, ties to the old centroid), so a duplicate-heavy hub still splits as far as its distinct rows allow (or is left alone when fewer than two centroids win rows) and no new partition comes out empty.join_partitions,PartitionAdjustment::Join { kept_partitions }), smallest first, each vector going to the nearest of the 64 nearest remaining partitions, instead of one join per optimize call. The join is executed against one finalized state. The rows to reindex are gathered first from every joined partition, their old entries are dropped from every partition through a shared set carried by the join adjustment (reindexed_row_ids), and survivor room below the split threshold is counted after those entries are gone (survivors are read for that only on multivector indexes, the only ones where a joined row has entries elsewhere). Each vector then goes to its nearest neighbor with room (choose_join_destination) among the 64 nearest survivors of the joined partition's centroid for a single-vector row, or of the vector itself for a multivector row (whose vectors may lie far from the partition the row was found under), or to the nearest surviving partition anywhere with room when all 64 are full; the planner keeps at leastceil(rows / threshold)survivors so that room always exists, and only an index already at the threshold everywhere lands a row in a full partition (logged, and split by the next optimize). Rows are fetched in batches admitted by decoded size (regroup_by_bytes, 32 MiB): a fixed-size vector column is chunked to that budget from its known row size, a multivector column is fetched one row at a time with four fetches in flight and regrouped by measured bytes, and only a single row larger than the budget ever exceeds it; each batch is then routed and quantized 1,024 vectors at a time with the chosen partition forced (the partition transformer now keeps a supplied partition id and measures the distance to it instead of reassigning, which RaBitQ's residual needs) and streamed through a shuffler to temp files like the split path, so memory is bounded by one chunk regardless of how many partitions are joined or how many vectors a reindexed row has elsewhere; there is no byte budget. The planner joins the smallest undersized partitions first, as many as leaveceil(rows / threshold)survivors; where each row lands is decided when the join runs.target_partition_sizedrives the thresholds. Split and join thresholds used the index type's default (8192 / 4096 rows) even when the index was created with another target; the persisted value is now read from the index details (target_partition_size_from_details) both in the builder and in the steady-state segment selector. In steady state, joins are decided on sizes summed over all segments and only when the segments share one IVF model; segments with different centroids are left alone.io_parallelismconcurrenttake_rowschunks instead of one at a time; an optimize whose split candidates turned out empty no longer forces a full merge.Why?
Partition skew is the direct cause of IVF query tail latency, and repairing it was expensive: every optimize halved one oversized partition, re-read the raw vectors of it and its 64 neighbors, and re-quantized all of them, so a hub partition took many passes and each pass cost O(65 partitions) of dataset reads. Joins were one per pass as well.
Results
Measured with the ablation harness (
ablate.py: every optimize and every recall evaluation runs in its own process). Dataset: the first 6M rows of DEEP-10M (96-d, L2) with an IVF_PQ index of 1024 partitions and 32 sub-vectors (target partition size 8192, so the split threshold is 32,768 rows and the join threshold 2,048), built once and copied for every run. Two scenarios append near-duplicate hubs on top of it and then calloptimize_indicesuntil it is a no-op (at most 6 calls):hub= 8 hubs x 150k rows (18x the target),smallhub= 8 hubs x 30k rows. Machine: Mac mini M4 (10 cores, 26 GB), local NVMe,maturin develop --releasebuilds ofmain(4157766) and this PR's head; every number is the mean over three runs (the branch is deterministic across them, std 0). Optimize time, partitions re-read and largest partition: lower is better; recall@10 (1,000 DEEP queries, nprobes 16, refine 10, against exact ground truth over the live rows): higher is better.On the near-duplicate hubs
main's 2-way splits peel a few hundred rows off each hub per call, so the hubs stay at ~150k rows while undersized pieces pile up (47 after six calls); the k-way split takes each hub to the target size in one call and the batched join removes the pieces in the next.Ablation
Each optimization in the PR was switched off on its own (an env-gated build that is not part of the PR) and measured on both scenarios, three seeds of the split's k-means each. Only the two core changes move the outcome; the rest were removed from the PR.
take_rowschunksTests
builder.rs:optimize_splits_oversized_partition_to_target_in_one_pass,optimize_joins_all_undersized_partitions_in_one_pass,optimize_split_threshold_honors_persisted_target_partition_size,optimize_join_leaves_no_partition_above_the_split_threshold(20 x 24-row partitions at target 100; the optimize tests build their index on fixed centroids so the layout does not depend on k-means seeding),optimize_splits_duplicate_heavy_partition_as_far_as_its_distinct_rows_allow(a 303-way split of five duplicate groups, built on fixed centroids),optimize_join_keeps_every_destination_below_the_split_threshold(24 rows nearest to a 390-row partition at threshold 400: ten fit, fourteen spill to the next neighbor),join_destination_spills_to_the_next_neighbor,join_destination_searches_beyond_a_full_window(a source whose 64-neighbor window is full after an earlier join reaches a partition outside it),plan_partition_adjustment_keeps_room_for_every_joined_row,plan_partition_adjustment_joins_all_undersized_partitions_but_one,optimize_join_counts_room_after_the_reindexed_rows_leave(multivector rows whose 360 entries leave a survivor before their 380 vectors are placed),join_window_follows_each_vector_of_a_multivector_row(a vector at 100 found under a partition near 0 reaches its own surviving centroid),regroup_by_bytes_bounds_each_batch_to_the_budget;apply_centroid_splits_correct_count_and_orderingcovers a 3-way split;select_reassign_candidates_skips_deleted_partitioncovers the exclusion set.test_join_partition_on_delete_multivecnow retains a row with one vector in each partition and checks the joined index holds exactly three entries per row.test_keeps_chosen_partitions_and_measures_their_distanceinlance-indexcovers the forced-partition path of the partition transformer.ivf/v2.rsandindex.rsupdated to the new granularity (ceil(rows / target)pieces, batched joins); two fixtures that appended identical rows now add a tiny per-row drift, since identical rows cannot be split by clustering (the split is skipped instead of producing empty partitions).optimize_steady_state_keeps_small_delta_partitionsguards the steady-state rule that joins are decided on sizes summed over all segments;test_vector_append_is_segment_set_native_with_distinct_modelsnow uses two partitions per segment so a steady-state optimize of segments with different models is exercised as a no-op.cargo test -p lance --lib -- index::: 870 passed, 1 failed once (dataset::index::frag_reuse::tests::test_cleanup_frag_reuse_index, a scalar-index compaction cleanup test); it passed on three reruns (twice on this branch, once on a tree without these changes) and does not touch the vector rebalance path, so it is classified as flaky under the parallel run.cargo clippy -p lance --tests --benches -- -D warningsandcargo fmt --allare clean.Limitations / follow-ups
🤖 Generated with Claude Code