perf: do HNSW search with threads of CPU runtime (#2251)

Signed-off-by: BubbleCal <[email protected]>

rust/lance/src/index/vector/hnsw.rs

@@ -10,6 +10,7 @@ use std::{
 
 use arrow_array::{Float32Array, RecordBatch, UInt64Array};
 use async_trait::async_trait;
+use lance_core::utils::tokio::spawn_cpu;
 use lance_core::{datatypes::Schema, Error, Result};
 use lance_file::reader::FileReader;
 use lance_index::vector::quantizer::Quantizer;
@@ -46,7 +47,7 @@ pub(crate) struct HNSWIndexOptions {
 
 #[derive(Clone)]
 pub(crate) struct HNSWIndex<Q: Quantization> {
-    hnsw: HNSW,
+    hnsw: Arc<HNSW>,
 
     // TODO: move these into IVFIndex after the refactor is complete
     partition_storage: IvfQuantizationStorage<Q>,
@@ -80,7 +81,7 @@ impl<Q: Quantization> HNSWIndex<Q> {
 
         let ivf_store = IvfQuantizationStorage::open(aux_reader).await?;
         Ok(Self {
-            hnsw,
+            hnsw: Arc::new(hnsw),
             partition_storage: ivf_store,
             partition_metadata,
             options,
@@ -176,7 +177,9 @@ impl<Q: Quantization + Send + Sync + 'static> VectorIndex for HNSWIndex<Q> {
             });
         }
 
-        let results = self.hnsw.search(query.key.clone(), k, ef, bitmap)?;
+        let hnsw = self.hnsw.clone();
+        let key = query.key.clone();
+        let results = spawn_cpu(move || hnsw.search(key, k, ef, bitmap)).await?;
 
         let row_ids = UInt64Array::from_iter_values(results.iter().map(|x| row_ids[x.id as usize]));
         let distances = Arc::new(Float32Array::from_iter_values(
@@ -231,7 +234,7 @@ impl<Q: Quantization + Send + Sync + 'static> VectorIndex for HNSWIndex<Q> {
         .await?;
 
         Ok(Box::new(Self {
-            hnsw,
+            hnsw: Arc::new(hnsw),
             partition_storage: self.partition_storage.clone(),
             partition_metadata: self.partition_metadata.clone(),
             options: self.options.clone(),
@@ -258,7 +261,7 @@ impl<Q: Quantization + Send + Sync + 'static> VectorIndex for HNSWIndex<Q> {
         .await?;
 
         Ok(Box::new(Self {
-            hnsw,
+            hnsw: Arc::new(hnsw),
             partition_storage: self.partition_storage.clone(),
             partition_metadata: self.partition_metadata.clone(),
             options: self.options.clone(),

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

@@ -2750,25 +2750,37 @@ mod tests {
         assert_eq!(1, results.len());
         assert_eq!(k, results[0].num_rows());
 
-        let results = results[0]
+        let row_ids = results[0]
             .column_by_name(ROW_ID)
             .unwrap()
             .as_any()
             .downcast_ref::<UInt64Array>()
             .unwrap()
             .iter()
             .map(|v| v.unwrap() as u32)
-            .collect::<HashSet<_>>();
+            .collect::<Vec<_>>();
+        let dists = results[0]
+            .column_by_name("_distance")
+            .unwrap()
+            .as_any()
+            .downcast_ref::<Float32Array>()
+            .unwrap()
+            .values()
+            .to_vec();
 
-        let gt = ground_truth(&mat, query.values(), k, distance_type);
+        let results = dists.into_iter().zip(row_ids.into_iter()).collect_vec();
+        let gt = ground_truth(&mat, query.values(), k, DistanceType::L2);
+
+        let results_set = results.iter().map(|r| r.1).collect::<HashSet<_>>();
         let gt_set = gt.iter().map(|r| r.1).collect::<HashSet<_>>();
-        let recall = results.intersection(&gt_set).count() as f32 / k as f32;
+
+        let recall = results_set.intersection(&gt_set).count() as f32 / k as f32;
         assert!(
             recall >= 0.9,
             "recall: {}\n results: {:?}\n\ngt: {:?}",
             recall,
-            results.iter().sorted().collect_vec(),
-            gt_set.iter().sorted().collect_vec()
+            results,
+            gt,
         );
     }