perf(index): parallelize kmeans centroid recomputation

[isaac-dasari]

Summary

• replace repeated KMeans centroid membership rescans with a workload-aware hybrid recomputation path
• keep the existing centroid-owner rescan for workloads where indexing does not amortize
• use a parallel stable owner index for workloads where removing repeated membership scans is measurably beneficial
• preserve disjoint centroid-owner writes and ascending input-row accumulation order
• add representative Criterion coverage and include the KMeans recomputation benchmark in the Rust benchmark workflow

Closes #6369.

Motivation

The existing parallel recomputation partitions centroids across workers. Each centroid owner then scans the full membership array and only accumulates vectors belonging to its centroid range. This preserves locality and avoids shared writes, but repeats membership scanning across owners.

A naive input-partitioned reduction removes those rescans but requires private k × dimension accumulators and can reduce effective parallelism or spend more time initializing and merging dense grids than doing useful centroid work. A serial membership index also regresses wall-clock time because its setup is not parallelized.

This revision keeps the good properties of the existing owner-based algorithm and only removes repeated membership scans when the complete workload makes that worthwhile.

Implementation

Workload-aware strategy selection

The recomputation path considers the actual workload rather than using a row-count-only threshold. The selector accounts for worker count, centroid count, vector dimension/value size, owner count, index traffic, and the amount of work available to amortize index construction.

High-dimensional, low-sample, and low-parallelism workloads remain on the original centroid-owner rescan path. Low-dimensional workloads with enough parallel owner work use the indexed path.

Parallel stable owner index

For the indexed path, membership rows are grouped by centroid owner using contiguous input partitions processed in parallel. Each partition preserves source row order; owner slices are combined in partition order, so every owner sees globally ascending input rows.

Centroid owners then process only their indexed rows while retaining the same disjoint centroid-chunk ownership as the original implementation. This avoids shared centroid writes and preserves accumulation order and numerical results.

The index is proportional to the number of participating rows plus owner metadata; it does not allocate a dense k × dimension accumulator per worker.

Correctness coverage

Focused tests cover:

• stable owner indexing and globally ascending row order
• indexed recomputation matching the serial/reference accumulation exactly
• owner-rescan fallback matching the reference path
• workload selection across the representative high-dimensional, low-dimensional, and low-parallelism cases
• existing empty-cluster and centroid recomputation behavior

Benchmark coverage

The committed Criterion benchmark covers representative KMeans recomputation shapes:

• low_sample_high_dim: N=512, dimension=1024, k=256
• default_ivf_high_dim: N=16,384, dimension=1024, k=64
• large_incremental_ivf: N=65,536, dimension=1024, k=4096
• default_pq_subvector: N=65,536, dimension=64, k=256
• max_sample_low_dim: N=131,072, dimension=128, k=256

The latest Gatekeeper current-base/current-head verification is approved (K-approved). It reports exact centroid bit patterns, cluster sizes, and row order preserved across the representative envelope. The indexed low-dimensional cases remain about 1.8–1.9× faster, while workloads that select the owner-rescan fallback show no material regression.

Validation

Passed locally on the current base:

• 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 PR is ready for maintainer review.

[lance-gatekeeper]

✅ Gate recommendation: approve.

The workload-aware hybrid remains the right trade-off on the updated base: the KMeans patch, its numerical and concurrency dependencies, and direct training callers are unchanged. Fresh current-base/current-head checks preserve every centroid bit pattern, cluster size, and row order. The indexed PQ and low-dimensional max-sample cases remain materially faster, while repeated alternating measurements show no stable regression on fallback or two-thread paths.

[isaac-dasari]

@Xuanwo kindly review the changes and merge the PR