Make kmeans test less flaky (#170)

rs/index/src/hnsw/reader.rs

@@ -1,9 +1,9 @@
 use std::fs::File;
 
+use anyhow::Result;
 use byteorder::{ByteOrder, LittleEndian};
 use memmap2::Mmap;
 use quantization::quantization::Quantizer;
-use anyhow::Result;
 
 use crate::hnsw::index::Hnsw;
 use crate::hnsw::writer::{Header, Version};

rs/index/src/lib.rs

@@ -2,7 +2,7 @@ pub mod hnsw;
 pub mod index;
 pub mod ivf;
 pub mod posting_list;
+pub mod spann;
 pub mod traverse_state;
 pub mod utils;
 pub mod vector;
-pub mod spann;

rs/index/src/spann/mod.rs

@@ -1,2 +1,2 @@
-pub mod reader;
 pub mod index;
+pub mod reader;

rs/index_writer/src/index_writer.rs

@@ -465,4 +465,82 @@ mod tests {
         assert!(ivf_vector_storage.exists());
         assert!(ivf_index.exists());
     }
+
+    #[test]
+    fn test_index_writer_process_ivf_hnsw() {
+        // Setup test data
+        let mut rng = rand::thread_rng();
+        let dimension = 10;
+        let num_rows = 100;
+        let data: Vec<Vec<f32>> = (0..num_rows)
+            .map(|_| (0..dimension).map(|_| rng.gen::<f32>()).collect())
+            .collect();
+
+        let mut mock_input = MockInput::new(data);
+
+        // Create a temporary directory for output
+        let temp_dir = TempDir::new("test_index_writer_process_ivf_hnsw")
+            .expect("Failed to create temporary directory");
+        let base_directory = temp_dir
+            .path()
+            .to_str()
+            .expect("Failed to convert temporary directory path to string")
+            .to_string();
+
+        // Configure IndexWriter
+        let base_config = BaseConfig {
+            output_path: base_directory.clone(),
+            dimension,
+            max_memory_size: 1024 * 1024 * 1024, // 1 GB
+            file_size: 1024 * 1024 * 1024,       // 1 GB
+        };
+        let hnsw_config = HnswConfig {
+            num_layers: 2,
+            max_num_neighbors: 10,
+            ef_construction: 100,
+            reindex: false,
+            quantizer_type: QuantizerType::ProductQuantizer,
+            subvector_dimension: 2,
+            num_bits: 2,
+            num_training_rows: 50,
+
+            max_iteration: 10,
+            batch_size: 10,
+        };
+        let ivf_config = IvfConfig {
+            num_clusters: 2,
+            num_data_points: 100,
+            max_clusters_per_vector: 1,
+
+            max_iteration: 10,
+            batch_size: 10,
+            tolerance: 0.0,
+            max_posting_list_size: usize::MAX,
+        };
+        let config = IndexWriterConfig::HnswIvf(HnswIvfConfig {
+            base_config,
+            hnsw_config,
+            ivf_config,
+        });
+
+        let mut index_writer = IndexWriter::new(config);
+
+        // Process the input
+        assert!(index_writer.process(&mut mock_input).is_ok());
+
+        // Check if output directories and files exist
+        let quantizer_directory_path = format!("{}/centroid_quantizer", base_directory);
+        let pq_directory = Path::new(&quantizer_directory_path);
+        let centroids_directory_path = format!("{}/centroids", base_directory);
+        let centroids_directory = Path::new(&centroids_directory_path);
+        let hnsw_vector_storage_path =
+            format!("{}/vector_storage", centroids_directory.to_str().unwrap());
+        let hnsw_vector_storage = Path::new(&hnsw_vector_storage_path);
+        let hnsw_index_path = format!("{}/index", centroids_directory.to_str().unwrap());
+        let hnsw_index = Path::new(&hnsw_index_path);
+        assert!(pq_directory.exists());
+        assert!(centroids_directory.exists());
+        assert!(hnsw_vector_storage.exists());
+        assert!(hnsw_index.exists());
+    }
 }

rs/utils/src/kmeans_builder/kmeans_builder.rs

@@ -24,6 +24,8 @@ pub struct KMeansBuilder {
 
     // Variant for this algorithm. Currently only Lloyd is supported.
     pub variant: KMeansVariant,
+
+    pub cluster_init_values: Option<Vec<usize>>,
 }
 
 pub struct KMeansResult {
@@ -49,6 +51,25 @@ impl KMeansBuilder {
             tolerance,
             dimension,
             variant,
+            cluster_init_values: None,
+        }
+    }
+
+    pub fn new_with_cluster_init_values(
+        num_cluters: usize,
+        max_iter: usize,
+        tolerance: f32,
+        dimension: usize,
+        variant: KMeansVariant,
+        cluster_init_values: Vec<usize>,
+    ) -> Self {
+        Self {
+            num_cluters,
+            max_iter,
+            tolerance,
+            dimension,
+            variant,
+            cluster_init_values: Some(cluster_init_values),
         }
     }
 
@@ -120,8 +141,17 @@ impl KMeansBuilder {
         let mut cluster_labels = vec![0; num_data_points];
 
         // Random initialization of cluster labels
-        for i in 0..num_data_points {
-            cluster_labels[i] = rand::random::<usize>() % self.num_cluters;
+        match &self.cluster_init_values {
+            Some(values) => {
+                for i in 0..num_data_points {
+                    cluster_labels[i] = values[i % values.len()];
+                }
+            }
+            None => {
+                for i in 0..num_data_points {
+                    cluster_labels[i] = rand::random::<usize>() % self.num_cluters;
+                }
+            }
         }
 
         let mut final_centroids = vec![0.0; self.num_cluters * self.dimension];
@@ -174,8 +204,7 @@ impl KMeansBuilder {
                         let centroid = centroids
                             [centroid_id * self.dimension..(centroid_id + 1) * self.dimension]
                             .as_ref();
-                        let distance =
-                            T::calculate_squared(dp, centroid) + penalties[centroid_id];
+                        let distance = T::calculate_squared(dp, centroid) + penalties[centroid_id];
 
                         if distance < min_cost {
                             min_cost = distance;
@@ -224,7 +253,14 @@ mod tests {
             .cloned()
             .collect();
 
-        let kmeans = KMeansBuilder::new(3, 100, 1e-4, 2, KMeansVariant::Lloyd);
+        let kmeans = KMeansBuilder::new_with_cluster_init_values(
+            3,
+            100,
+            1e-4,
+            2,
+            KMeansVariant::Lloyd,
+            vec![0, 0, 0, 1, 1, 1, 2, 2, 2],
+        );
         let result = kmeans
             .fit(flattened_data)
             .expect("KMeans run should succeed");
@@ -267,7 +303,14 @@ mod tests {
             .flatten()
             .cloned()
             .collect();
-        let kmeans = KMeansBuilder::new(3, 100, 10000.0, 2, KMeansVariant::Lloyd);
+        let kmeans = KMeansBuilder::new_with_cluster_init_values(
+            3,
+            100,
+            10000.0,
+            2,
+            KMeansVariant::Lloyd,
+            vec![0, 0, 0, 1, 1, 1, 2, 2, 2],
+        );
         let result = kmeans
             .fit(flattened_data)
             .expect("KMeans run should succeed");
@@ -301,7 +344,15 @@ mod tests {
             .flatten()
             .cloned()
             .collect();
-        let kmeans = KMeansBuilder::new(3, 100, 0.0, 2, KMeansVariant::Lloyd);
+        let kmeans = KMeansBuilder::new_with_cluster_init_values(
+            3,
+            100,
+            0.0,
+            2,
+            KMeansVariant::Lloyd,
+            vec![0, 0, 0, 1, 1, 1, 2, 2, 2],
+        );
+
         let result = kmeans
             .fit(flattened_data)
             .expect("KMeans run should succeed");