perf(index): parallelize kmeans centroid recomputation

[isaac-dasari]

Summary

• recompute KMeans centroids by partitioning input vectors across workers so each input vector is scanned once
• merge thread-local centroid accumulators in parallel and cap their aggregate allocation at 64 MiB
• add correctness coverage for f16, f32, f64, ignored memberships, empty-cluster splitting, and the memory cap
• add a focused Criterion benchmark and include it in the existing Rust benchmark workflow

The previous implementation divided centroids among workers, but every worker scanned the complete input. That made centroid recomputation perform O(num_vectors * workers) input reads. This changes the input scan to O(num_vectors) while avoiding shared-write contention.

Closes #6369.

Validation

Passed locally:

• cargo fmt --all -- --check
• cargo check -p lance-index --tests --benches --profile bench -j 1
• cargo clippy -p lance-index --tests --benches --profile bench -j 1 -- -D warnings
• git diff --check

The optimized test/benchmark linker exceeded the available resources on the local 7.6 GiB development machine, so executable test and benchmark runs are intentionally delegated to GitHub Actions. The PR is draft pending those results.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The revision fixes the prior large-grid regression, but the memory-only strategy selector still chooses dense input partitioning for supported low-sample or high-dimensional workloads where zeroing and merging the private grids costs much more than the useful centroid work.

Use a workload-aware fallback—or a stable membership-to-row index followed by centroid-owner summation—and add representative base/head benchmarks for these shapes.

[lance-gatekeeper]

This memory-only cutoff still selects InputPartitioned when zeroing and merging one dense k × dimension grid per worker costs much more than the useful vector work. Optimized direct base/head harnesses on 62 threads measured:

• N=512, k=256, dimension=1024 (the supported sample_rate=2 shape): head 3.9–4.6 ms versus base 0.45–0.69 ms, roughly 9× slower.
• N=16,384, k=64, dimension=1024 (the default 256 samples per centroid): head 1.30–1.39 ms versus base 0.385–0.583 ms, roughly 3× slower.

Make the selector account for workload relative to private-grid initialization and reduction, or build a stable membership-to-row index and keep centroid-owner summation. Add representative base/head benchmarks so these supported shapes cannot regress.

Reproducer

For each shape, identical optimized harnesses at the base and head called KMeansAlgoFloat::<Float32Type>::to_kmeans with data = vec![1.0_f32; N * dimension], membership[row] = Some((row % k) as u32), and balanced cluster_sizes. Ten iterations verified every centroid value and reported elapsed time:

LANCE_CPU_THREADS=62 RAYON_NUM_THREADS=62 cargo run --release --quiet

Expected: this performance-only change should not materially regress supported centroid-recomputation workloads. Observed: the ranges above after warm-up.