perf(index): parallelize kmeans centroid recomputation

.github/workflows/rust-benchmark.yml

@@ -50,7 +50,7 @@ jobs:
         working-directory: ./rust/lance-index
         run: |
           # TODO: a few benchmarks are failing. Re-enable everything once they are fixed.
-          cargo bench --bench sq --bench hnsw --bench inverted --bench pq_dist_table --bench pq_assignment -- --output-format bencher | tee -a ../../output.txt
+          cargo bench --bench sq --bench hnsw --bench inverted --bench pq_dist_table --bench pq_assignment --bench kmeans_recompute -- --output-format bencher | tee -a ../../output.txt
       - name: Store benchmark result
         if: github.event_name != 'pull_request'
         uses: benchmark-action/github-action-benchmark@a7bc2366eda11037936ea57d811a43b3418d3073 # v1.21.0

rust/lance-index/Cargo.toml

@@ -143,6 +143,10 @@ harness = false
 name = "kmeans"
 harness = false
 
+[[bench]]
+name = "kmeans_recompute"
+harness = false
+
 [[bench]]
 name = "compute_partition"
 harness = false

rust/lance-index/benches/kmeans_recompute.rs

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use std::hint::black_box;
+
+use arrow_array::types::Float32Type;
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use lance_index::vector::kmeans::{KMeansAlgo, KMeansAlgoFloat};
+use lance_linalg::distance::DistanceType;
+
+fn bench_recompute_centroids(c: &mut Criterion) {
+    let mut group = c.benchmark_group("kmeans_recompute_centroids");
+
+    let cases = [
+        ("input_partitioned", 128 * 1024, 128, 256),
+        ("large_centroid_grid", 65_536, 1024, 4096),
+    ];
+
+    for (name, num_vectors, dimension, k) in cases {
+        let data = vec![1.0_f32; num_vectors * dimension];
+        let membership = (0..num_vectors)
+            .map(|row| Some((row % k) as u32))
+            .collect::<Vec<_>>();
+        let cluster_sizes = vec![num_vectors / k; k];
+
+        group.bench_with_input(
+            BenchmarkId::new(name, format!("{num_vectors}x{dimension}d_{k}k")),
+            &num_vectors,
+            |b, _| {
+                b.iter(|| {
+                    let mut cluster_sizes = cluster_sizes.clone();
+                    black_box(KMeansAlgoFloat::<Float32Type>::to_kmeans(
+                        black_box(&data),
+                        dimension,
+                        k,
+                        black_box(&membership),
+                        &mut cluster_sizes,
+                        DistanceType::L2,
+                        0.0,
+                    ))
+                });
+            },
+        );
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_recompute_centroids);
+criterion_main!(benches);