[LuciferYang]

Closes #9345

train_weighted_hierarchical_kmeans, the trainer behind the streaming coreset path for both f32 and f64 centroids, padded the cluster list with clones of the heaviest cluster whenever it formed fewer than target_k. A clone is bit-identical to the cluster it came from, and assignment resolves a tie to the incumbent, so a clone starts with no rows and gains some in the refinement passes that follow only if the original drifts away from it. The build succeeded with partitions that were empty or nearly so.

It now returns invalid_input naming the coreset row count and, when the coreset is not smaller than num_partitions, how many clusters were formed. The flat hierarchical trainer already rejects the same request (lance-index/src/vector/kmeans.rs:1561). The two are not the same error kind, though: the flat path raises ArrowError::InvalidArgumentError, which surfaces in Python as RuntimeError, while this one is Error::InvalidInput and surfaces as ValueError. Aligning them is a separate change. A coreset that has fewer rows than num_partitions gets its own message suggesting a larger sample_rate, since more diverse data would not help there.

A k-means split can also leave every row on one side when the rows differ. Under Dot this happens often enough that main's test_build_ivf_model_streaming_dot_coreset_training, 257 partitions over 1,024 distinct rows, came up 9 to 24 clusters short in 4 of 25 local runs with this rejection and no bisection; main padded those shortfalls with clones. Before rejecting, the trainer now bisects the remaining multi-row clusters by L2 around the row farthest from its centroid and the row farthest from that one, which always yields two non-empty halves when the rows differ. The trainer's output therefore falls short only when the coreset holds fewer distinct rows than num_partitions, and the bisection runs only where main would have padded.

That guarantee covers the trainer's output. The refinement passes after it run under the build's metric, and under Dot they can still pull distinct centroids for collinear positive data back onto the same point, as they could with the clones. That is left for a follow-up.

clusters.truncate(target_k) after the check would be a no-op once the shortfall is an error, because the loop stops at target_k and every step it takes is bounded by what is left. It is a debug_assert_eq! on that invariant instead.

How was this patch tested?

test_weighted_hierarchical_rejects_padding_with_duplicates is an rstest pair over the two shapes that reach the shortfall: fewer_rows_than_partitions, 6 rows against target_k 16, expects the coreset-size message, and enough_rows_all_identical, 32 identical rows against 16 partitions, expects the one about rows that cannot be split. Both assert the coreset row count appears in the message. Restoring the padding fails both cases.

test_weighted_hierarchical_bisects_clusters_kmeans_cannot_split trains 4 partitions from 8 distinct positive 1-D rows under Dot, where every k-means split sends all rows to the largest centroid, and asserts the trainer returns 4 distinct centroids, with one case each for f32 and f64. It fails without the bisection.

[LuciferYang]

Reworked the loop rather than the wording — you were pointing at the same hazard my own review round flagged and I had left as a note.

The loop abandoned itself when the cluster at the top of the heap was a singleton. Priority there is loss, not size (Ord only stratifies on finalized, then compares loss), so a heavy singleton could outrank clusters that still had rows to give: the build was rejected while target_k was reachable, and the Reduce num_partitions to <= N in the message was lower than what the trainer could actually have formed. Only the singleton case changed. A finalized cluster at the top genuinely does mean nothing splittable is left, because Ord sorts finalized last, so breaking there is still right; the singleton is now finalized, pushed back so it sinks, and the loop keeps splitting the rest.

One thing I could not do, so it is not claimed anywhere: there is no test pinning that guard. Heap priority comes from the initial clustering's cluster_losses rather than from the losses argument a test supplies, so I could not force a heavy singleton to the top deterministically. A fixture that looked convincing passed with the old break as well, so I deleted it instead of shipping a test with no discriminating power. The two shortfall cases remain covered and both fail if the padding loop comes back.

[LuciferYang]

Pushed the regression guard, and one honest correction about the case you named.

I added the 257-of-258 shape and then mutation-tested it: reverting the singleton branch from continue back to break leaves it green, so it does not distinguish the two heads on this branch. The reason is the heap ordering itself — for non-finalized clusters Ord compares loss, and a singleton's cluster loss is ~0, so singletons already sink below anything splittable without needing the finalized flag. A singleton only reaches the top when a per-row base loss spike inflates its cluster loss (cluster_losses[c] += base_losses[row] + weight * distance), and that value is produced by the initial clustering rather than by the losses argument a test passes in, so I could not force it deterministically. Four fixtures failed that way before I stopped.

The test is in anyway, because it does guard something real: that the shortfall check cannot over-reject an attainable count. Changing if clusters.len() < target_k to <= fails it immediately with "could only form 257 … Reduce num_partitions to <= 257". So the singleton repair itself rests on the ordering argument plus the flat trainer's precedent, not on a test, and I would rather say that than imply coverage I do not have. Making it testable means injecting heap priority, which is a separate change.

On CI: linux-arm is unrelated. It fails in query::primitives::test_query_string::case_2_large_utf8 with Resources exhausted: Additional allocation failed for ExternalSorterMerge[0] while building a scalar index on a string column — several 40 MB sorter reservations against the runner's pool. PR #9004, an unrelated residual-partition-id change, fails the same test at the same panic site (rust/lance/tests/utils/mod.rs:247) with the same allocation error, differing only in which index type it happened to be building when memory ran out. Nothing in this PR allocates; it touches one function in the IVF partition trainer.

[LuciferYang]

Added your fixture, and I was wrong about why my own attempts failed — worth correcting explicitly since I put it in the PR body.

I claimed the supplied losses cannot raise heap priority because priority comes from the initial clustering. That is wrong: assign_weighted_f32_points folds each base_losses[row] into initial.cluster_losses, which is exactly the ordering key. My fixtures failed for a duller reason — I used a 1e3 spike against a 0.1 baseline, nowhere near enough to dominate, and a 4- or 8-dimensional layout where the outlier did not end up alone. Your 1e20 against 0.0 at dimension = 1 makes it unconditional.

test_weighted_hierarchical_exhausts_splittable_clusters is in as written. Verified both ways: it passes on this head, and reverting the singleton continue to break fails it with could only form 16 non-empty clusters … Reduce num_partitions to <= 16 — the 257-of-258 case you described. The three weighted_hierarchical cases run in 0.03s.

I also dropped the earlier over-rejection guard I had added in its place; this fixture covers the repair directly, and that one only pinned that the shortfall check does not fire early.

[ddupg]

We've also encountered this issue with streaming IVF training, resulting in many duplicate centroids and empty partitions. Thanks for the fix—looking forward to seeing it merged!

[LuciferYang]

@BubbleCal could you take a look, since the streaming k-means trainer is from your #6913? When the weighted hierarchical trainer forms fewer clusters than num_partitions, it pads with clones of the heaviest one, and those partitions start empty (ddupg reported hitting this above). It now bisects the clusters k-means could not split and rejects only when the coreset has too few distinct rows. It is ported onto the generic f32/f64 trainer from #8662.

[lance-gatekeeper]

⚠️ Gate recommendation: approve with a non-blocking risk.

The exhausted-split expectation is corrected without changing production code. The native f32/f64 fallback still reaches attainable partition counts without duplicate padding and rejects unattainable counts.

One existing Dot limitation remains: positive collinear data can collapse back into duplicate centroids during subsequent refinement, leaving most partitions empty. This also occurs with the former padding path, so no additional change is requested here. For such workloads, verify final partition occupancy; bisection is an initialization fallback, not a guarantee of non-empty final partitions.

[LuciferYang]

Thank you @Xuanwo