feat: record loss for IVF and KMeans

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

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

@@ -34,6 +34,9 @@ pub struct IvfModel {
 
     /// Number of vectors in each partition.
     pub lengths: Vec<u32>,
+
+    /// Kmeans loss
+    pub loss: Option<f64>,
 }
 
 impl DeepSizeOf for IvfModel {
@@ -53,6 +56,7 @@ impl IvfModel {
             centroids: None,
             offsets: vec![],
             lengths: vec![],
+            loss: None,
         }
     }
 
@@ -61,6 +65,7 @@ impl IvfModel {
             centroids: Some(centroids),
             offsets: vec![],
             lengths: vec![],
+            loss: None,
         }
     }
 
@@ -88,6 +93,14 @@ impl IvfModel {
         self.lengths[part] as usize
     }
 
+    pub fn num_rows(&self) -> u64 {
+        self.lengths.iter().map(|x| *x as u64).sum()
+    }
+
+    pub fn avg_loss(&self) -> Option<f64> {
+        self.loss.map(|loss| loss / self.num_rows() as f64)
+    }
+
     /// Use the query vector to find `nprobes` closest partitions.
     pub fn find_partitions(
         &self,
@@ -215,6 +228,7 @@ impl TryFrom<PbIvf> for IvfModel {
             centroids,
             offsets,
             lengths: proto.lengths,
+            loss: None,
         })
     }
 }

rust/lance-linalg/src/kmeans.rs

@@ -103,6 +103,9 @@ pub struct KMeans {
 
     /// How to calculate distance between two vectors.
     pub distance_type: DistanceType,
+
+    /// The loss of the last training.
+    pub loss: f64,
 }
 
 /// Randomly initialize kmeans centroids.
@@ -127,6 +130,7 @@ fn kmeans_random_init<T: ArrowPrimitiveType>(
         centroids: Arc::new(centroids),
         dimension,
         distance_type,
+        loss: f64::MAX,
     }
 }
 
@@ -191,6 +195,7 @@ pub trait KMeansAlgo<T: Num> {
         k: usize,
         membership: &[Option<u32>],
         distance_type: DistanceType,
+        loss: f64,
     ) -> KMeans;
 }
 
@@ -245,6 +250,7 @@ where
         k: usize,
         membership: &[Option<u32>],
         distance_type: DistanceType,
+        loss: f64,
     ) -> KMeans {
         let mut cluster_cnts = vec![0_u64; k];
         let mut new_centroids = vec![T::Native::zero(); k * dimension];
@@ -293,6 +299,7 @@ where
             centroids: Arc::new(PrimitiveArray::<T>::from_iter_values(new_centroids)),
             dimension,
             distance_type,
+            loss,
         }
     }
 }
@@ -337,6 +344,7 @@ impl KMeansAlgo<u8> for KModeAlgo {
         k: usize,
         membership: &[Option<u32>],
         distance_type: DistanceType,
+        loss: f64,
     ) -> KMeans {
         assert_eq!(distance_type, DistanceType::Hamming);
 
@@ -379,6 +387,7 @@ impl KMeansAlgo<u8> for KModeAlgo {
             centroids: Arc::new(UInt8Array::from(centroids)),
             dimension,
             distance_type,
+            loss,
         }
     }
 }
@@ -389,6 +398,7 @@ impl KMeans {
             centroids: arrow_array::array::new_empty_array(&DataType::Float32),
             dimension,
             distance_type,
+            loss: f64::MAX,
         }
     }
 
@@ -398,6 +408,7 @@ impl KMeans {
         centroids: ArrayRef,
         dimension: usize,
         distance_type: DistanceType,
+        loss: f64,
     ) -> Self {
         assert!(matches!(
             centroids.data_type(),
@@ -407,6 +418,7 @@ impl KMeans {
             centroids,
             dimension,
             distance_type,
+            loss,
         }
     }
 
@@ -496,6 +508,7 @@ impl KMeans {
                     k,
                     &membership,
                     params.distance_type,
+                    last_loss,
                 );
                 last_membership = Some(membership);
                 if (loss - last_loss).abs() / last_loss < params.tolerance {

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

@@ -781,6 +781,7 @@ mod tests {
         let gt_set = gt.iter().map(|r| r.1).collect::<HashSet<_>>();
 
         let recall = row_ids.intersection(&gt_set).count() as f32 / k as f32;
+        println!("recall: {}", recall);
         assert!(
             recall >= recall_requirement,
             "recall: {}\n results: {:?}\n\ngt: {:?}",