feat!: support hamming distance & binary vector

Cargo.toml

@@ -21,7 +21,7 @@ exclude = ["python"]
 resolver = "2"
 
 [workspace.package]
-version = "0.20.1"
+version = "0.21.0"
 edition = "2021"
 authors = ["Lance Devs <dev@lancedb.com>"]
 license = "Apache-2.0"
@@ -44,21 +44,21 @@ categories = [
 rust-version = "1.78"
 
 [workspace.dependencies]
-lance = { version = "=0.20.1", path = "./rust/lance" }
-lance-arrow = { version = "=0.20.1", path = "./rust/lance-arrow" }
-lance-core = { version = "=0.20.1", path = "./rust/lance-core" }
-lance-datafusion = { version = "=0.20.1", path = "./rust/lance-datafusion" }
-lance-datagen = { version = "=0.20.1", path = "./rust/lance-datagen" }
-lance-encoding = { version = "=0.20.1", path = "./rust/lance-encoding" }
-lance-encoding-datafusion = { version = "=0.20.1", path = "./rust/lance-encoding-datafusion" }
-lance-file = { version = "=0.20.1", path = "./rust/lance-file" }
-lance-index = { version = "=0.20.1", path = "./rust/lance-index" }
-lance-io = { version = "=0.20.1", path = "./rust/lance-io" }
-lance-jni = { version = "=0.20.1", path = "./java/core/lance-jni" }
-lance-linalg = { version = "=0.20.1", path = "./rust/lance-linalg" }
-lance-table = { version = "=0.20.1", path = "./rust/lance-table" }
-lance-test-macros = { version = "=0.20.1", path = "./rust/lance-test-macros" }
-lance-testing = { version = "=0.20.1", path = "./rust/lance-testing" }
+lance = { version = "=0.21.0", path = "./rust/lance" }
+lance-arrow = { version = "=0.21.0", path = "./rust/lance-arrow" }
+lance-core = { version = "=0.21.0", path = "./rust/lance-core" }
+lance-datafusion = { version = "=0.21.0", path = "./rust/lance-datafusion" }
+lance-datagen = { version = "=0.21.0", path = "./rust/lance-datagen" }
+lance-encoding = { version = "=0.21.0", path = "./rust/lance-encoding" }
+lance-encoding-datafusion = { version = "=0.21.0", path = "./rust/lance-encoding-datafusion" }
+lance-file = { version = "=0.21.0", path = "./rust/lance-file" }
+lance-index = { version = "=0.21.0", path = "./rust/lance-index" }
+lance-io = { version = "=0.21.0", path = "./rust/lance-io" }
+lance-jni = { version = "=0.21.0", path = "./java/core/lance-jni" }
+lance-linalg = { version = "=0.21.0", path = "./rust/lance-linalg" }
+lance-table = { version = "=0.21.0", path = "./rust/lance-table" }
+lance-test-macros = { version = "=0.21.0", path = "./rust/lance-test-macros" }
+lance-testing = { version = "=0.21.0", path = "./rust/lance-testing" }
 approx = "0.5.1"
 # Note that this one does not include pyarrow
 arrow = { version = "53.2", optional = false, features = ["prettyprint"] }
@@ -111,7 +111,7 @@ datafusion-physical-expr = { version = "42.0", features = [
 ] }
 deepsize = "0.2.0"
 either = "1.0"
-fsst = { version = "=0.20.1", path = "./rust/lance-encoding/src/compression_algo/fsst" }
+fsst = { version = "=0.21.0", path = "./rust/lance-encoding/src/compression_algo/fsst" }
 futures = "0.3"
 http = "1.1.0"
 hyperloglogplus = { version = "0.4.1", features = ["const-loop"] }

python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pylance"
-version = "0.20.1"
+version = "0.21.0"
 edition = "2021"
 authors = ["Lance Devs <dev@lancedb.com>"]
 rust-version = "1.65"

python/src/utils.rs

@@ -15,6 +15,7 @@
 use std::sync::Arc;
 
 use arrow::compute::concat;
+use arrow::datatypes::Float32Type;
 use arrow::pyarrow::{FromPyArrow, ToPyArrow};
 use arrow_array::{cast::AsArray, Array, FixedSizeListArray, Float32Array, UInt32Array};
 use arrow_data::ArrayData;
@@ -26,7 +27,7 @@ use lance_file::writer::FileWriter;
 use lance_index::scalar::IndexWriter;
 use lance_index::vector::hnsw::{builder::HnswBuildParams, HNSW};
 use lance_index::vector::v3::subindex::IvfSubIndex;
-use lance_linalg::kmeans::compute_partitions;
+use lance_linalg::kmeans::{compute_partitions, KMeansAlgoFloat};
 use lance_linalg::{
     distance::DistanceType,
     kmeans::{KMeans as LanceKMeans, KMeansParams},
@@ -132,14 +133,15 @@ impl KMeans {
         if !matches!(fixed_size_arr.value_type(), DataType::Float32) {
             return Err(PyValueError::new_err("Must be a FixedSizeList of Float32"));
         };
-        let values: Arc<Float32Array> = fixed_size_arr.values().as_primitive().clone().into();
-        let centroids: &Float32Array = kmeans.centroids.as_primitive();
-        let cluster_ids = UInt32Array::from(compute_partitions(
-            centroids.values(),
-            values.values(),
-            kmeans.dimension,
-            kmeans.distance_type,
-        ));
+        let values = fixed_size_arr.values().as_primitive();
+        let centroids = kmeans.centroids.as_primitive();
+        let cluster_ids =
+            UInt32Array::from(compute_partitions::<
+                Float32Type,
+                KMeansAlgoFloat<Float32Type>,
+            >(
+                centroids, values, kmeans.dimension, kmeans.distance_type
+            ));
         cluster_ids.into_data().to_pyarrow(py)
     }
 

rust/lance-index/benches/hnsw.rs

@@ -15,7 +15,7 @@ use lance_index::vector::v3::subindex::IvfSubIndex;
 use pprof::criterion::{Output, PProfProfiler};
 
 use lance_index::vector::{
-    flat::storage::FlatStorage,
+    flat::storage::FlatFloatStorage,
     hnsw::builder::{HnswBuildParams, HNSW},
 };
 use lance_linalg::distance::DistanceType;
@@ -31,7 +31,7 @@ fn bench_hnsw(c: &mut Criterion) {
 
     let data = generate_random_array_with_seed::<Float32Type>(TOTAL * DIMENSION, SEED);
     let fsl = FixedSizeListArray::try_new_from_values(data, DIMENSION as i32).unwrap();
-    let vectors = Arc::new(FlatStorage::new(fsl.clone(), DistanceType::L2));
+    let vectors = Arc::new(FlatFloatStorage::new(fsl.clone(), DistanceType::L2));
 
     let query = fsl.value(0);
     c.bench_function(

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

@@ -28,7 +28,7 @@ use crate::{
     },
 };
 
-use super::storage::{FlatStorage, FLAT_COLUMN};
+use super::storage::{FlatBinStorage, FlatFloatStorage, FLAT_COLUMN};
 
 /// A Flat index is any index that stores no metadata, and
 /// during query, it simply scans over the storage and returns the top k results
@@ -166,7 +166,7 @@ impl FlatQuantizer {
 impl Quantization for FlatQuantizer {
     type BuildParams = ();
     type Metadata = FlatMetadata;
-    type Storage = FlatStorage;
+    type Storage = FlatFloatStorage;
 
     fn build(data: &dyn Array, distance_type: DistanceType, _: &Self::BuildParams) -> Result<Self> {
         let dim = data.as_fixed_size_list().value_length();
@@ -228,3 +228,81 @@ impl TryFrom<Quantizer> for FlatQuantizer {
         }
     }
 }
+
+#[derive(Debug, Clone, DeepSizeOf)]
+pub struct FlatBinQuantizer {
+    dim: usize,
+    distance_type: DistanceType,
+}
+
+impl FlatBinQuantizer {
+    pub fn new(dim: usize, distance_type: DistanceType) -> Self {
+        Self { dim, distance_type }
+    }
+}
+
+impl Quantization for FlatBinQuantizer {
+    type BuildParams = ();
+    type Metadata = FlatMetadata;
+    type Storage = FlatBinStorage;
+
+    fn build(data: &dyn Array, distance_type: DistanceType, _: &Self::BuildParams) -> Result<Self> {
+        let dim = data.as_fixed_size_list().value_length();
+        Ok(Self::new(dim as usize, distance_type))
+    }
+
+    fn code_dim(&self) -> usize {
+        self.dim
+    }
+
+    fn column(&self) -> &'static str {
+        FLAT_COLUMN
+    }
+
+    fn from_metadata(metadata: &Self::Metadata, distance_type: DistanceType) -> Result<Quantizer> {
+        Ok(Quantizer::FlatBin(Self {
+            dim: metadata.dim,
+            distance_type,
+        }))
+    }
+
+    fn metadata(
+        &self,
+        _: Option<crate::vector::quantizer::QuantizationMetadata>,
+    ) -> Result<serde_json::Value> {
+        let metadata = FlatMetadata { dim: self.dim };
+        Ok(serde_json::to_value(metadata)?)
+    }
+
+    fn metadata_key() -> &'static str {
+        "flat"
+    }
+
+    fn quantization_type() -> QuantizationType {
+        QuantizationType::Flat
+    }
+
+    fn quantize(&self, vectors: &dyn Array) -> Result<ArrayRef> {
+        Ok(vectors.slice(0, vectors.len()))
+    }
+}
+
+impl From<FlatBinQuantizer> for Quantizer {
+    fn from(value: FlatBinQuantizer) -> Self {
+        Self::FlatBin(value)
+    }
+}
+
+impl TryFrom<Quantizer> for FlatBinQuantizer {
+    type Error = Error;
+
+    fn try_from(value: Quantizer) -> Result<Self> {
+        match value {
+            Quantizer::FlatBin(quantizer) => Ok(quantizer),
+            _ => Err(Error::invalid_input(
+                "quantizer is not FlatBinQuantizer",
+                location!(),
+            )),
+        }
+    }
+}