fix: cosine kmeans (#1723)

Use L2 over normalized vector to run cosine for kmeans.

Author: eddyxu

python/src/utils.rs

@@ -101,7 +101,7 @@ impl KMeans {
             return Err(PyValueError::new_err("Must be a FixedSizeList of Float32"));
         };
         let values: Arc<Float32Array> = fixed_size_arr.values().as_primitive().clone().into();
-        let membership = RT.block_on(Some(py), kmeans.compute_membership(values, None))?;
+        let membership = RT.block_on(Some(py), kmeans.compute_membership(values))?;
         let cluster_ids: UInt32Array = membership
             .cluster_id_and_distances
             .iter()

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

@@ -22,20 +22,16 @@ use arrow_array::builder::UInt32Builder;
 use arrow_array::types::{Float16Type, Float32Type, Float64Type};
 use arrow_array::UInt64Array;
 use arrow_array::{
-    cast::AsArray, types::UInt32Type, Array, FixedSizeListArray, Float32Array, RecordBatch,
-    UInt32Array,
+    cast::AsArray, types::UInt32Type, Array, FixedSizeListArray, RecordBatch, UInt32Array,
 };
-use arrow_ord::sort::sort_to_indices;
 use arrow_schema::{DataType, Field};
 use arrow_select::take::take;
 use async_trait::async_trait;
-use futures::{stream, StreamExt, TryStreamExt};
+use futures::{stream, StreamExt};
 use lance_arrow::*;
 use lance_core::{Error, Result, ROW_ID};
 use lance_linalg::{
-    distance::{
-        cosine_distance_batch, dot_distance_batch, l2_distance_batch, Cosine, Dot, MetricType, L2,
-    },
+    distance::{Cosine, Dot, MetricType, L2},
     MatrixView,
 };
 use log::{debug, info};
@@ -51,6 +47,7 @@ use crate::vector::{
     transform::Transformer,
 };
 pub use builder::IvfBuildParams;
+use lance_linalg::kmeans::KMeans;
 
 fn new_ivf_impl<T: ArrowFloatType + Dot + Cosine + L2 + 'static>(
     centroids: &T::ArrayType,
@@ -349,23 +346,22 @@ impl<T: ArrowFloatType + Dot + L2 + Cosine + 'static> IvfImpl<T> {
             // wouldn't cover all rows but 13-element chunks (13 * 32 = 416) would
             // have one empty chunk at the end. This filter removes those empty chunks.
             .filter(|range| futures::future::ready(range.start < range.end))
-            .map(|range| {
+            .map(|range| async {
                 let centroids = centroids.clone();
                 let data = data.clone();
-                tokio::task::spawn_blocking(move || {
-                    compute_partitions::<T>(
-                        centroids.as_slice(),
-                        &data.as_slice()[range],
-                        dimension,
-                        metric_type,
-                    )
-                })
+
+                compute_partitions::<T>(
+                    centroids.as_slice(),
+                    &data.as_slice()[range],
+                    dimension,
+                    metric_type,
+                )
                 .in_current_span()
+                .await
             })
             .buffered(chunks)
-            .try_collect()
-            .await
-            .expect("compute_partitions: schedule CPU task");
+            .collect::<Vec<_>>()
+            .await;
 
         UInt32Array::from_iter(result.iter().flatten().copied())
     }
@@ -428,16 +424,6 @@ impl<T: ArrowFloatType + Dot + L2 + Cosine + 'static> Ivf for IvfImpl<T> {
     }
 
     fn find_partitions(&self, query: &dyn Array, nprobes: usize) -> Result<UInt32Array> {
-        if query.len() != self.dimension() {
-            return Err(Error::IO {
-                message: format!(
-                    "Ivf::find_partition: dimension mismatch: {} != {}",
-                    query.len(),
-                    self.dimension()
-                ),
-                location: location!(),
-            });
-        }
         let query = query
             .as_any()
             .downcast_ref::<T::ArrayType>()
@@ -449,23 +435,13 @@ impl<T: ArrowFloatType + Dot + L2 + Cosine + 'static> Ivf for IvfImpl<T> {
                 ),
                 location: Default::default(),
             })?;
-        let centroid_values = self.centroids.data();
-        let centroids = centroid_values.as_slice();
-        let dim = query.len();
-        let distances = Float32Array::from_iter_values(match self.metric_type {
-            lance_linalg::distance::DistanceType::L2 => {
-                l2_distance_batch(query.as_slice(), centroids, dim)
-            }
-            lance_linalg::distance::DistanceType::Cosine => {
-                cosine_distance_batch(query.as_slice(), centroids, dim)
-            }
-            lance_linalg::distance::DistanceType::Dot => {
-                dot_distance_batch(query.as_slice(), centroids, dim)
-            }
-        });
-
-        let top_k_partitions = sort_to_indices(&distances, None, Some(nprobes))?;
-        Ok(top_k_partitions)
+        // TODO: hold kmeans in this struct.
+        let kmeans = KMeans::<T>::with_centroids(
+            self.centroids.data().clone(),
+            self.dimension(),
+            self.metric_type,
+        );
+        Ok(kmeans.find_partitions(query.as_slice(), nprobes)?)
     }
 
     async fn partition_transform(&self, batch: &RecordBatch, column: &str) -> Result<RecordBatch> {

rust/lance-linalg/Cargo.toml

@@ -11,6 +11,7 @@ categories = { workspace = true }
 
 [dependencies]
 arrow-array = { workspace = true }
+arrow-ord = { workspace = true }
 arrow-schema = { workspace = true }
 arrow-select = { workspace = true }
 futures = { workspace = true }