perf: use faster kmean find partition routing for pq assignment

rust/lance-index/Cargo.toml

@@ -76,6 +76,10 @@ harness = false
 name = "pq_dist_table"
 harness = false
 
+[[bench]]
+name = "pq_assignment"
+harness = false
+
 [[bench]]
 name = "hnsw"
 harness = false

rust/lance-index/benches/pq_assignment.rs

@@ -0,0 +1,44 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+//! Benchmark of Building PQ code from Dense Vectors.
+
+use std::sync::Arc;
+
+use arrow_array::{types::Float32Type, FixedSizeListArray};
+use criterion::{criterion_group, criterion_main, Criterion};
+use lance_arrow::FixedSizeListArrayExt;
+use lance_index::vector::pq::{ProductQuantizer, ProductQuantizerImpl};
+use lance_linalg::distance::DistanceType;
+use lance_testing::datagen::generate_random_array_with_seed;
+
+const PQ: usize = 96;
+const DIM: usize = 1536;
+const TOTAL: usize = 32 * 1024;
+
+fn pq_transform(c: &mut Criterion) {
+    let codebook = Arc::new(generate_random_array_with_seed::<Float32Type>(
+        256 * DIM,
+        [88; 32],
+    ));
+
+    let vectors = generate_random_array_with_seed::<Float32Type>(DIM * TOTAL, [3; 32]);
+    let fsl = FixedSizeListArray::try_new_from_values(vectors, DIM as i32).unwrap();
+
+    for dt in [DistanceType::L2, DistanceType::Dot].iter() {
+        let pq = ProductQuantizerImpl::<Float32Type>::new(PQ, 8, DIM, codebook.clone(), *dt);
+
+        c.bench_function(format!("{},{}", dt, TOTAL).as_str(), |b| {
+            b.iter(|| {
+                let _ = pq.transform(&fsl).unwrap();
+            })
+        });
+    }
+}
+
+criterion_group!(
+    name=benches;
+    config = Criterion::default().significance_level(0.1).sample_size(10);
+    targets = pq_transform);
+
+criterion_main!(benches);

rust/lance-index/src/vector/pq.rs

@@ -14,7 +14,7 @@ use lance_arrow::*;
 use lance_core::{Error, Result};
 use lance_linalg::distance::{dot_distance_batch, l2_distance_batch, DistanceType, Dot, L2};
 use lance_linalg::kernels::argmin_value_float;
-use lance_linalg::kmeans::kmeans_find_partitions;
+use lance_linalg::kmeans::compute_partition;
 use lance_linalg::{distance::MetricType, MatrixView};
 use rayon::prelude::*;
 use snafu::{location, Location};
@@ -360,24 +360,23 @@ where
                     location: location!(),
                 })?;
 
+        let sub_dim = dim / num_sub_vectors;
         let values = flatten_data
             .as_slice()
             .par_chunks(dim)
             .map(|vector| {
                 vector
-                    .chunks_exact(dim / num_sub_vectors)
+                    .chunks_exact(sub_dim)
                     .enumerate()
-                    .flat_map(|(sub_idx, sub_vector)| {
+                    .map(|(sub_idx, sub_vector)| {
                         let centroids = get_sub_vector_centroids(
                             codebook.as_slice(),
                             dim,
                             num_bits,
                             num_sub_vectors,
                             sub_idx,
                         );
-                        let parts = kmeans_find_partitions(centroids, sub_vector, 1, distance_type)
-                            .expect("kmeans_find_partitions failed");
-                        parts.values().iter().map(|v| *v as u8).collect::<Vec<_>>()
+                        compute_partition(centroids, sub_vector, distance_type).map(|v| v as u8)
                     })
                     .collect::<Vec<_>>()
             })

rust/lance-index/src/vector/pq/utils.rs

@@ -54,6 +54,7 @@ pub fn num_centroids(num_bits: impl Into<u32>) -> usize {
     2_usize.pow(num_bits.into())
 }
 
+#[inline]
 pub fn get_sub_vector_centroids<T>(
     codebook: &[T],
     dimension: usize,

rust/lance-linalg/src/kmeans.rs

@@ -685,22 +685,32 @@ pub fn compute_partitions<T: Float + L2 + Dot + Sync>(
     let dimension = dimension.as_();
     vectors
         .par_chunks(dimension)
-        .map(|vec| {
-            argmin_value(match distance_type {
-                DistanceType::L2 => l2_distance_batch(vec, centroids, dimension),
-                DistanceType::Dot => dot_distance_batch(vec, centroids, dimension),
-                _ => {
-                    panic!(
-                        "KMeans::find_partitions: {} is not supported",
-                        distance_type
-                    );
-                }
-            })
-            .map(|(idx, _)| idx)
-        })
+        .map(|vec| compute_partition(centroids, vec, distance_type))
         .collect::<Vec<_>>()
 }
 
+#[inline]
+pub fn compute_partition<T: Float + L2 + Dot>(
+    centroids: &[T],
+    vector: &[T],
+    distance_type: DistanceType,
+) -> Option<u32> {
+    match distance_type {
+        DistanceType::L2 => {
+            argmin_value_float(l2_distance_batch(vector, centroids, vector.len())).map(|c| c.0)
+        }
+        DistanceType::Dot => {
+            argmin_value_float(dot_distance_batch(vector, centroids, vector.len())).map(|c| c.0)
+        }
+        _ => {
+            panic!(
+                "KMeans::compute_partition: distance type {} is not supported",
+                distance_type
+            );
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::iter::repeat;