feat: int8 support for distance functions (#3605)

triandco · LuQQiu · web-flow · commit 5338238195a2 · 2025-04-03T16:58:52.000+08:00
Using [AllMiniLM12V2](https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2) to generate embeddings, I found that with a tiny bit of transformation, I can store the embedding as int8 to massively reduce its storage requirement and turn it into f32 during distance calculation (using cosine and dot) with marginal change to the outcome. I would like to add support for int8 by converting FixedSizeList type of Int8Array to Float32Array during calculation. The conversion is done with a ``convert_to_floating_point`` on the ``FixedSizeListArrayExt`` trait. In theory, this can also work with other types like uint8, int32, int64. This would make a big difference in storage requirement on the compatible models. --------- Co-authored-by: LuQQiu <luqiujob@gmail.com>
diff --git a/rust/lance-arrow/src/lib.rs b/rust/lance-arrow/src/lib.rs
@@ -13,6 +13,7 @@ use arrow_array::{
     GenericListArray, OffsetSizeTrait, PrimitiveArray, RecordBatch, StructArray, UInt32Array,
     UInt8Array,
 };
+use arrow_array::{Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array};
 use arrow_buffer::MutableBuffer;
 use arrow_data::ArrayDataBuilder;
 use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields, IntervalUnit, Schema};
@@ -235,6 +236,10 @@ pub trait FixedSizeListArrayExt {
     /// assert_eq!(sampled.values().len(), 160);
     /// ```
     fn sample(&self, n: usize) -> Result<FixedSizeListArray>;
+
+    /// Ensure the [FixedSizeListArray] of Float16, Float32, Float64,
+    /// Int8, Int16, Int32, Int64, UInt8, UInt32 type to its closest floating point type.
+    fn convert_to_floating_point(&self) -> Result<FixedSizeListArray>;
 }
 
 impl FixedSizeListArrayExt for FixedSizeListArray {
@@ -253,6 +258,136 @@ impl FixedSizeListArrayExt for FixedSizeListArray {
         let chosen = (0..self.len() as u32).choose_multiple(&mut rng, n);
         take(self, &UInt32Array::from(chosen), None).map(|arr| arr.as_fixed_size_list().clone())
     }
+
+    fn convert_to_floating_point(&self) -> Result<FixedSizeListArray> {
+        match self.data_type() {
+            DataType::FixedSizeList(field, size) => match field.data_type() {
+                DataType::Float16 | DataType::Float32 | DataType::Float64 => Ok(self.clone()),
+                DataType::Int8 => Ok(Self::new(
+                    Arc::new(arrow_schema::Field::new(
+                        field.name(),
+                        DataType::Float32,
+                        field.is_nullable(),
+                    )),
+                    *size,
+                    Arc::new(Float32Array::from_iter_values(
+                        self.values()
+                            .as_any()
+                            .downcast_ref::<Int8Array>()
+                            .ok_or(ArrowError::ParseError(
+                                "Fail to cast primitive array to Int8Type".to_string(),
+                            ))?
+                            .into_iter()
+                            .filter_map(|x| x.map(|y| y as f32)),
+                    )),
+                    self.nulls().cloned(),
+                )),
+                DataType::Int16 => Ok(Self::new(
+                    Arc::new(arrow_schema::Field::new(
+                        field.name(),
+                        DataType::Float32,
+                        field.is_nullable(),
+                    )),
+                    *size,
+                    Arc::new(Float32Array::from_iter_values(
+                        self.values()
+                            .as_any()
+                            .downcast_ref::<Int16Array>()
+                            .ok_or(ArrowError::ParseError(
+                                "Fail to cast primitive array to Int8Type".to_string(),
+                            ))?
+                            .into_iter()
+                            .filter_map(|x| x.map(|y| y as f32)),
+                    )),
+                    self.nulls().cloned(),
+                )),
+                DataType::Int32 => Ok(Self::new(
+                    Arc::new(arrow_schema::Field::new(
+                        field.name(),
+                        DataType::Float32,
+                        field.is_nullable(),
+                    )),
+                    *size,
+                    Arc::new(Float32Array::from_iter_values(
+                        self.values()
+                            .as_any()
+                            .downcast_ref::<Int32Array>()
+                            .ok_or(ArrowError::ParseError(
+                                "Fail to cast primitive array to Int8Type".to_string(),
+                            ))?
+                            .into_iter()
+                            .filter_map(|x| x.map(|y| y as f32)),
+                    )),
+                    self.nulls().cloned(),
+                )),
+                DataType::Int64 => Ok(Self::new(
+                    Arc::new(arrow_schema::Field::new(
+                        field.name(),
+                        DataType::Float64,
+                        field.is_nullable(),
+                    )),
+                    *size,
+                    Arc::new(Float64Array::from_iter_values(
+                        self.values()
+                            .as_any()
+                            .downcast_ref::<Int64Array>()
+                            .ok_or(ArrowError::ParseError(
+                                "Fail to cast primitive array to Int8Type".to_string(),
+                            ))?
+                            .into_iter()
+                            .filter_map(|x| x.map(|y| y as f64)),
+                    )),
+                    self.nulls().cloned(),
+                )),
+                DataType::UInt8 => Ok(Self::new(
+                    Arc::new(arrow_schema::Field::new(
+                        field.name(),
+                        DataType::Float64,
+                        field.is_nullable(),
+                    )),
+                    *size,
+                    Arc::new(Float64Array::from_iter_values(
+                        self.values()
+                            .as_any()
+                            .downcast_ref::<UInt8Array>()
+                            .ok_or(ArrowError::ParseError(
+                                "Fail to cast primitive array to Int8Type".to_string(),
+                            ))?
+                            .into_iter()
+                            .filter_map(|x| x.map(|y| y as f64)),
+                    )),
+                    self.nulls().cloned(),
+                )),
+                DataType::UInt32 => Ok(Self::new(
+                    Arc::new(arrow_schema::Field::new(
+                        field.name(),
+                        DataType::Float64,
+                        field.is_nullable(),
+                    )),
+                    *size,
+                    Arc::new(Float64Array::from_iter_values(
+                        self.values()
+                            .as_any()
+                            .downcast_ref::<UInt32Array>()
+                            .ok_or(ArrowError::ParseError(
+                                "Fail to cast primitive array to Int8Type".to_string(),
+                            ))?
+                            .into_iter()
+                            .filter_map(|x| x.map(|y| y as f64)),
+                    )),
+                    self.nulls().cloned(),
+                )),
+                data_type => Err(ArrowError::ParseError(format!(
+                    "Expect either floating type or integer got {:?}",
+                    data_type
+                ))),
+            },
+            data_type => Err(ArrowError::ParseError(format!(
+                "Expect either FixedSizeList got {:?}",
+                data_type
+            ))),
+        }
+    }
 }
 
 /// Force downcast of an [`Array`], such as an [`ArrayRef`], to
@@ -412,6 +547,14 @@ pub trait RecordBatchExt {
     /// Replace a column (specified by name) and return the new [`RecordBatch`].
     fn replace_column_by_name(&self, name: &str, column: Arc<dyn Array>) -> Result<RecordBatch>;
 
+    /// Replace a column schema (specified by name) and return the new [`RecordBatch`].
+    fn replace_column_schema_by_name(
+        &self,
+        name: &str,
+        new_data_type: DataType,
+        column: Arc<dyn Array>,
+    ) -> Result<RecordBatch>;
+
     /// Get (potentially nested) column by qualified name.
     fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef>;
 
@@ -519,6 +662,37 @@ impl RecordBatchExt for RecordBatch {
         Self::try_new(self.schema(), columns)
     }
 
+    fn replace_column_schema_by_name(
+        &self,
+        name: &str,
+        new_data_type: DataType,
+        column: Arc<dyn Array>,
+    ) -> Result<RecordBatch> {
+        let fields = self
+            .schema()
+            .fields()
+            .iter()
+            .map(|x| {
+                if x.name() != name {
+                    x.clone()
+                } else {
+                    let new_field = Field::new(name, new_data_type.clone(), x.is_nullable());
+                    Arc::new(new_field)
+                }
+            })
+            .collect::<Vec<_>>();
+        let schema = Schema::new_with_metadata(fields, self.schema().metadata.clone());
+        let mut columns = self.columns().to_vec();
+        let field_i = self
+            .schema()
+            .fields()
+            .iter()
+            .position(|f| f.name() == name)
+            .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?;
+        columns[field_i] = column;
+        Self::try_new(Arc::new(schema), columns)
+    }
+
     fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef> {
         let split = name.split('.').collect::<Vec<_>>();
         if split.is_empty() {
diff --git a/rust/lance-index/src/vector/residual.rs b/rust/lance-index/src/vector/residual.rs
@@ -136,6 +136,13 @@ pub(crate) fn compute_residual(
         (DataType::Float64, DataType::Float64) => {
             do_compute_residual::<Float64Type>(centroids, vectors, distance_type, partitions)
         }
+        (DataType::Float32, DataType::Int8) => {
+            do_compute_residual::<Float32Type>(
+                centroids,
+                &vectors.convert_to_floating_point()?,
+                distance_type,
+                partitions)
+        }
         _ => Err(Error::Index {
             message: format!(
                 "Compute residual vector: centroids and vector type mismatch: centroid: {}, vector: {}",
@@ -181,7 +188,16 @@ impl Transformer for ResidualTransform {
             compute_residual(&self.centroids, original_vectors, None, Some(part_ids_ref))?;
 
         // Replace original column with residual column.
-        let batch = batch.replace_column_by_name(&self.vec_col, Arc::new(residual_arr))?;
+        let batch = if residual_arr.data_type() != original.data_type() {
+            batch.replace_column_schema_by_name(
+                &self.vec_col,
+                residual_arr.data_type().clone(),
+                Arc::new(residual_arr),
+            )?
+        } else {
+            batch.replace_column_by_name(&self.vec_col, Arc::new(residual_arr))?
+        };
+
         Ok(batch)
     }
 }
diff --git a/rust/lance-index/src/vector/transform.rs b/rust/lance-index/src/vector/transform.rs
@@ -142,6 +142,7 @@ impl Transformer for KeepFiniteVectors {
                     DataType::Float32 => is_all_finite::<Float32Type>(&data),
                     DataType::Float64 => is_all_finite::<Float64Type>(&data),
                     DataType::UInt8 => data.null_count() == 0,
+                    DataType::Int8 => data.null_count() == 0,
                     _ => false,
                 };
                 if is_valid {
diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs
@@ -11,12 +11,12 @@ use std::sync::Arc;
 
 use arrow_array::{
     cast::AsArray,
-    types::{Float16Type, Float32Type, Float64Type},
+    types::{Float16Type, Float32Type, Float64Type, Int8Type},
     Array, FixedSizeListArray, Float32Array,
 };
 use arrow_schema::DataType;
 use half::{bf16, f16};
-use lance_arrow::{ArrowFloatType, FloatArray};
+use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray};
 #[cfg(feature = "fp16kernels")]
 use lance_core::utils::cpu::SimdSupport;
 use lance_core::utils::cpu::FP16_SIMD_SUPPORT;
@@ -320,6 +320,14 @@ pub fn cosine_distance_arrow_batch(
         DataType::Float16 => do_cosine_distance_arrow_batch::<Float16Type>(from.as_primitive(), to),
         DataType::Float32 => do_cosine_distance_arrow_batch::<Float32Type>(from.as_primitive(), to),
         DataType::Float64 => do_cosine_distance_arrow_batch::<Float64Type>(from.as_primitive(), to),
+        DataType::Int8 => do_cosine_distance_arrow_batch::<Float32Type>(
+            &from
+                .as_primitive::<Int8Type>()
+                .into_iter()
+                .map(|x| x.unwrap() as f32)
+                .collect(),
+            &to.convert_to_floating_point()?,
+        ),
         _ => Err(Error::InvalidArgumentError(format!(
             "Unsupported data type {:?}",
             from.data_type()
diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs
@@ -8,11 +8,11 @@ use std::ops::AddAssign;
 use std::sync::Arc;
 
 use crate::Error;
-use arrow_array::types::{Float16Type, Float64Type};
+use arrow_array::types::{Float16Type, Float64Type, Int8Type};
 use arrow_array::{cast::AsArray, types::Float32Type, Array, FixedSizeListArray, Float32Array};
 use arrow_schema::DataType;
 use half::{bf16, f16};
-use lance_arrow::{ArrowFloatType, FloatArray};
+use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray};
 #[cfg(feature = "fp16kernels")]
 use lance_core::utils::cpu::SimdSupport;
 use lance_core::utils::cpu::FP16_SIMD_SUPPORT;
@@ -278,6 +278,14 @@ pub fn dot_distance_arrow_batch(
         DataType::Float16 => do_dot_distance_arrow_batch::<Float16Type>(from.as_primitive(), to),
         DataType::Float32 => do_dot_distance_arrow_batch::<Float32Type>(from.as_primitive(), to),
         DataType::Float64 => do_dot_distance_arrow_batch::<Float64Type>(from.as_primitive(), to),
+        DataType::Int8 => do_dot_distance_arrow_batch::<Float32Type>(
+            &from
+                .as_primitive::<Int8Type>()
+                .into_iter()
+                .map(|x| x.unwrap() as f32)
+                .collect(),
+            &to.convert_to_floating_point()?,
+        ),
         _ => Err(Error::InvalidArgumentError(format!(
             "Unsupported data type: {:?}",
             from.data_type()
diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs
@@ -10,12 +10,12 @@ use std::sync::Arc;
 
 use arrow_array::{
     cast::AsArray,
-    types::{Float16Type, Float32Type, Float64Type},
+    types::{Float16Type, Float32Type, Float64Type, Int8Type},
     Array, FixedSizeListArray, Float32Array,
 };
 use arrow_schema::DataType;
 use half::{bf16, f16};
-use lance_arrow::{ArrowFloatType, FloatArray};
+use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray};
 #[cfg(feature = "fp16kernels")]
 use lance_core::utils::cpu::SimdSupport;
 use lance_core::utils::cpu::FP16_SIMD_SUPPORT;
@@ -293,6 +293,14 @@ pub fn l2_distance_arrow_batch(
         DataType::Float16 => do_l2_distance_arrow_batch::<Float16Type>(from.as_primitive(), to),
         DataType::Float32 => do_l2_distance_arrow_batch::<Float32Type>(from.as_primitive(), to),
         DataType::Float64 => do_l2_distance_arrow_batch::<Float64Type>(from.as_primitive(), to),
+        DataType::Int8 => do_l2_distance_arrow_batch::<Float32Type>(
+            &from
+                .as_primitive::<Int8Type>()
+                .into_iter()
+                .map(|x| x.unwrap() as f32)
+                .collect(),
+            &to.convert_to_floating_point()?,
+        ),
         _ => Err(Error::ComputeError(format!(
             "Unsupported data type: {}",
             from.data_type()
diff --git a/rust/lance-linalg/src/kmeans.rs b/rust/lance-linalg/src/kmeans.rs
@@ -23,6 +23,7 @@ use arrow_array::{ArrowNumericType, UInt8Array};
 use arrow_ord::sort::sort_to_indices;
 use arrow_schema::{ArrowError, DataType};
 use bitvec::prelude::*;
+use lance_arrow::FixedSizeListArrayExt;
 use log::{info, warn};
 use num_traits::{AsPrimitive, Float, FromPrimitive, Num, Zero};
 use rand::prelude::*;
@@ -720,6 +721,15 @@ pub fn compute_partitions_arrow_array(
             centroids.value_length(),
             distance_type,
         )),
+        (DataType::Float32, DataType::Int8) => Ok(compute_partitions::<
+            Float32Type,
+            KMeansAlgoFloat<Float32Type>,
+        >(
+            centroids.values().as_primitive(),
+            vectors.convert_to_floating_point()?.values().as_primitive(),
+            centroids.value_length(),
+            distance_type,
+        )),
         (DataType::Float64, DataType::Float64) => Ok(compute_partitions::<
             Float64Type,
             KMeansAlgoFloat<Float64Type>,
@@ -736,7 +746,7 @@ pub fn compute_partitions_arrow_array(
             distance_type,
         )),
         _ => Err(ArrowError::InvalidArgumentError(
-            "Centroids and vectors have different types".to_string(),
+            "Centroids and vectors have incompatible types".to_string(),
         )),
     }
 }
diff --git a/rust/lance-testing/src/datagen.rs b/rust/lance-testing/src/datagen.rs
@@ -9,7 +9,8 @@ use std::{iter::repeat_with, ops::Range};
 
 use arrow_array::types::ArrowPrimitiveType;
 use arrow_array::{
-    Float32Array, Int32Array, PrimitiveArray, RecordBatch, RecordBatchIterator, RecordBatchReader,
+    Float32Array, Int32Array, Int8Array, PrimitiveArray, RecordBatch, RecordBatchIterator,
+    RecordBatchReader,
 };
 use arrow_schema::{DataType, Field, Schema as ArrowSchema};
 use lance_arrow::{fixed_size_list_type, ArrowFloatType, FixedSizeListArrayExt};
@@ -222,6 +223,13 @@ pub fn generate_random_array(n: usize) -> Float32Array {
     Float32Array::from_iter_values(repeat_with(|| rng.gen::<f32>()).take(n))
 }
 
+/// Create a random float32 array where each element is uniformly
+/// distributed between [0..1]
+pub fn generate_random_int8_array(n: usize) -> Int8Array {
+    let mut rng = rand::thread_rng();
+    Int8Array::from_iter_values(repeat_with(|| rng.gen::<i8>()).take(n))
+}
+
 /// Create a random primitive array where each element is uniformly distributed a
 /// given range.
 pub fn generate_random_array_with_range<T: ArrowPrimitiveType>(
diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs
diff --git a/rust/lance/src/index/vector/ivf.rs b/rust/lance/src/index/vector/ivf.rs
diff --git a/rust/lance/src/index/vector/utils.rs b/rust/lance/src/index/vector/utils.rs
