chore: make Ivf not be restricted by arrow float type (#2352)

lancedb · May 19, 2024 · 7d9dbda · 7d9dbda
1 parent 6845edf
commit 7d9dbda
Show file tree

Hide file tree

Showing 11 changed files with 332 additions and 534 deletions.
diff --git a/rust/lance-index/benches/find_partitions.rs b/rust/lance-index/benches/find_partitions.rs
@@ -4,14 +4,13 @@
 use arrow_array::Float32Array;
 use arrow_array::{types::Float32Type, FixedSizeListArray};
 use lance_arrow::FixedSizeListArrayExt;
-use std::sync::Arc;
 
 use criterion::{criterion_group, criterion_main, Criterion};
 #[cfg(target_os = "linux")]
 use pprof::criterion::{Output, PProfProfiler};
 
-use lance_index::vector::ivf::{Ivf, IvfImpl};
-use lance_linalg::{distance::MetricType, MatrixView};
+use lance_index::vector::ivf::Ivf;
+use lance_linalg::distance::DistanceType;
 use lance_testing::datagen::generate_random_array_with_seed;
 
 fn bench_partitions(c: &mut Criterion) {
@@ -21,21 +20,19 @@ fn bench_partitions(c: &mut Criterion) {
     let query: Float32Array = generate_random_array_with_seed::<Float32Type>(DIMENSION, SEED);
 
     for num_centroids in &[10240, 65536] {
-        let centroids = Arc::new(generate_random_array_with_seed::<Float32Type>(
-            num_centroids * DIMENSION,
-            SEED,
-        ));
-        let matrix = MatrixView::<Float32Type>::new(centroids.clone(), DIMENSION);
+        let centroids =
+            generate_random_array_with_seed::<Float32Type>(num_centroids * DIMENSION, SEED);
+        let fsl = FixedSizeListArray::try_new_from_values(centroids, DIMENSION as i32).unwrap();
 
         for k in &[1, 10, 50] {
-            let ivf = IvfImpl::new(matrix.clone(), MetricType::L2, "vector", vec![], None);
+            let ivf = Ivf::new(fsl.clone(), DistanceType::L2, vec![]);
             c.bench_function(format!("IVF{},k={},L2", num_centroids, k).as_str(), |b| {
                 b.iter(|| {
                     let _ = ivf.find_partitions(&query, *k);
                 })
             });
 
-            let ivf = IvfImpl::new(matrix.clone(), MetricType::Cosine, "vector", vec![], None);
+            let ivf = Ivf::new(fsl.clone(), DistanceType::Cosine, vec![]);
             c.bench_function(
                 format!("IVF{},k={},Cosine", num_centroids, k).as_str(),
                 |b| {
@@ -46,7 +43,7 @@ fn bench_partitions(c: &mut Criterion) {
             );
         }
 
-        let ivf = IvfImpl::new(matrix.clone(), MetricType::L2, "vector", vec![], None);
+        let ivf = Ivf::new(fsl.clone(), DistanceType::L2, vec![]);
         let batch = generate_random_array_with_seed::<Float32Type>(DIMENSION * 4096, SEED);
         let fsl = FixedSizeListArray::try_new_from_values(batch, DIMENSION as i32).unwrap();
         c.bench_function(