fix(index): accumulate k-means centroids in a wider float

[FANNG1]

Fixes #8607.

Building any IVF index over a float16 vector column hangs forever at 100% CPU once a cluster's values sum past the f16 range. The issue has the reproduction and the full chain; the short version is that the centroid update accumulated in the storage type, non-finite centroids made argmin assign nothing, and split_clusters then spun in an unbounded loop that could never satisfy its exit condition.

Changes

Three changes, each sufficient on its own to prevent the hang. They are kept together because the first is the root cause and the other two make the failure mode unreachable rather than merely unlikely.

to_kmeans accumulates in f64. Summing, dividing and the split perturbation all happen there, and the centroids are narrowed back to the storage type once, at the end. f64 holds every sum an f16, f32 or f64 cluster can produce, so this closes the f32 overflow path as well as the f16 one. Narrowing clamps to the storage type's maximum, because the split perturbation scales a centroid by 1 + 1/1024, which would otherwise round a top-of-range value to inf.

split_clusters takes donors largest-first from a BinaryHeap instead of by unbounded probabilistic search, breaking ties on the smaller cluster id, and returns once no cluster has more than one vector to donate. This replaces the direct port of faiss's Clustering.cpp donor selection, so donor choice is now deterministic and the routine terminates regardless of its input — it no longer relies on a precondition held by its only caller. The n parameter is gone with the probability denominator.

train_kmeans rejects training data where fewer than k vectors could be assigned to any centroid. That threshold is exact rather than conservative: with at least k vectors assigned, pigeonhole guarantees that whenever an empty cluster exists some cluster has two or more members and can donate, so split_clusters always fills every cluster; below k, the old loop was guaranteed to reach a state with an empty cluster and all counts at most one, which is precisely where it hung. Small datasets, heavy duplication and the hierarchical path are unaffected.

Non-finite rows cannot trip it from the index-building path, because every training entry point runs filter_finite_training_data first (ivf.rs, builder.rs), so k-means never sees a NaN or infinite vector and kmeans_random_init cannot seed a centroid from one. test_create_ivf_flat_with_nan_rows pins that interaction — a column that is 99% NaN still indexes off the remaining rows, repeatedly, since the seeding is random. A caller reaching KMeans::new_with_params directly with unfiltered data and drawing an all-non-finite seed set now gets this error where it previously hung.

PQ codebook training reaches the same to_kmeans, so IVF_PQ with the dot metric — which skips residuals and therefore quantizes the raw vectors — was hanging for the same reason and is fixed by the same change.

Compatibility

No public signature changes. The accumulator is a private helper, the KMeansAlgo for KMeansAlgoFloat<T> bound stays at T: ArrowNumericType, and the update uses only bounds that implementation already required, so downstream code generic over T keeps compiling untouched.

Two behavior changes worth calling out:

• Donor selection is no longer faiss-compatible. Reseeding an empty cluster now always splits the current largest cluster rather than sampling one with probability proportional to its size. This is the same heuristic train_hierarchical_kmeans already uses, and it is also asymptotically cheaper — the old search made roughly n / mean_cluster_size random draws per empty cluster.
• Some previously-hanging builds now return an error instead. That is the intended outcome, and the message names the likely causes (non-finite values, zero-length vectors under a normalizing metric, distances that overflow).

f32 centroids shift in the last few bits, since their sums are now exact rather than rounded at every step. Nothing depends on the previous values: kmeans_random_init seeds from SmallRng::from_os_rng(), so training is already non-deterministic between runs.

Tests

New in lance-index: split_clusters filling empty clusters deterministically while preserving the vector count, stopping when donors run out ([3,0,0,0,0] -> [1,1,1,0,0]) and when there is no donor at all, the narrowing clamp in both directions, a to_kmeans case with a cluster larger than 65504, f16 training at magnitude 1000, and an all-NaN float32 case that must error rather than hang — the last one pins the unbounded loop without involving f16 at all.

test_create_ivf_flat_f16 is now parameterized over vector magnitude and checks distance finiteness and recall; a new test_create_ivf_pq_f16_dot_large_values covers the PQ codebook path. Both hang on the unfixed build and pass in about a second with the fix, verified by running them against a build with only kmeans.rs reverted.

The stale comment on the k * 512 training cap is updated: it existed to keep f16 centroid updates from underflowing, which the accumulator now handles, and what is left is a crude cost bound that keeps a prefix of the data rather than sampling.

cargo clippy -p lance-index -p lance --tests --benches -- -D warnings is clean; cargo test -p lance-index --lib vector:: is 370 passed and cargo test -p lance --lib index::vector is 265 passed.

https://claude.ai/code/session_01HkgaggGvsdGuUGgnwEgKke