feat: support all float types in MatrixView

rust/lance-arrow/Cargo.toml

@@ -18,3 +18,4 @@ arrow-schema = { workspace = true }
 serde_json = { workspace = true }
 serde = { workspace = true }
 half = { workspace = true }
+num-traits = { workspace = true }

rust/lance-arrow/src/bfloat16.rs

@@ -25,6 +25,8 @@ use arrow_data::ArrayData;
 use arrow_schema::{ArrowError, DataType};
 use half::bf16;
 
+pub struct BFloat16Type {}
+
 pub struct BFloat16Array {
     inner: FixedSizeBinaryArray,
 }

rust/lance-arrow/src/floats.rs

@@ -0,0 +1,62 @@
+// Copyright 2023 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Floats Array
+
+use arrow_array::{Array, Float16Array, Float32Array, Float64Array};
+use half::{bf16, f16};
+use num_traits::{Float, FromPrimitive};
+
+use super::bfloat16::BFloat16Array;
+
+/// [FloatArray] is a trait that is implemented by all float type arrays.
+pub trait FloatArray: Array + From<Vec<Self::Native>> {
+    type Native: Float + FromPrimitive;
+
+    /// Returns a reference to the underlying data as a slice.
+    fn as_slice(&self) -> &[Self::Native];
+}
+
+impl FloatArray for BFloat16Array {
+    type Native = bf16;
+
+    fn as_slice(&self) -> &[Self::Native] {
+        // TODO: apache/arrow-rs#4820
+        todo!()
+    }
+}
+
+impl FloatArray for Float16Array {
+    type Native = f16;
+
+    fn as_slice(&self) -> &[Self::Native] {
+        self.values()
+    }
+}
+
+impl FloatArray for Float32Array {
+    type Native = f32;
+
+    fn as_slice(&self) -> &[Self::Native] {
+        self.values()
+    }
+}
+
+impl FloatArray for Float64Array {
+    type Native = f64;
+
+    fn as_slice(&self) -> &[Self::Native] {
+        self.values()
+    }
+}

rust/lance-arrow/src/lib.rs

@@ -28,6 +28,8 @@ use arrow_schema::{ArrowError, DataType, Field, FieldRef, Fields, Schema};
 pub mod schema;
 pub use schema::*;
 pub mod bfloat16;
+pub mod floats;
+pub use floats::*;
 
 type Result<T> = std::result::Result<T, ArrowError>;
 

rust/lance-linalg/src/distance.rs

@@ -94,10 +94,9 @@ impl TryFrom<&str> for DistanceType {
             "l2" | "euclidean" => Ok(Self::L2),
             "cosine" => Ok(Self::Cosine),
             "dot" => Ok(Self::Dot),
-            _ => Err(ArrowError::InvalidArgumentError(
-                format!("Metric type '{s}' is not supported,")
-                    + "only 'l2'/'euclidean', 'cosine', and 'dot' are supported.",
-            )),
+            _ => Err(ArrowError::InvalidArgumentError(format!(
+                "Metric type '{s}' is not supported"
+            ))),
         }
     }
 }

rust/lance-linalg/src/kmeans.rs

@@ -16,10 +16,10 @@ use std::cmp::min;
 use std::collections::HashSet;
 use std::sync::Arc;
 
-use arrow_array::cast::AsArray;
-use arrow_array::FixedSizeListArray;
+use arrow_array::types::Float32Type;
 use arrow_array::{
-    builder::Float32Builder, cast::as_primitive_array, new_empty_array, Array, Float32Array,
+    cast::{as_primitive_array, AsArray},
+    new_empty_array, Array, FixedSizeListArray, Float32Array,
 };
 use arrow_schema::{ArrowError, DataType};
 use arrow_select::concat::concat;
@@ -154,12 +154,12 @@ async fn kmeans_random_init(
     let chosen = (0..data.len() / dimension)
         .choose_multiple(&mut rng, k)
         .to_vec();
-    let mut builder = Float32Builder::with_capacity(k * dimension);
+    let mut builder: Vec<f32> = Vec::with_capacity(k * dimension);
     for i in chosen {
-        builder.append_slice(&data.values()[i * dimension..(i + 1) * dimension]);
+        builder.extend(data.values()[i * dimension..(i + 1) * dimension].iter());
     }
     let mut kmeans = KMeans::empty(k, dimension, metric_type);
-    kmeans.centroids = Arc::new(builder.finish());
+    kmeans.centroids = Arc::new(builder.into());
     Ok(kmeans)
 }
 
@@ -246,7 +246,7 @@ impl KMeanMembership {
         }
         let centroids = concat(&mean_refs).unwrap();
         Ok(KMeans {
-            centroids: Arc::new(as_primitive_array(centroids.as_ref()).clone()),
+            centroids: Arc::new(centroids.as_ref().as_primitive::<Float32Type>().clone()),
             dimension,
             k: self.k,
             metric_type: self.metric_type,
@@ -319,7 +319,7 @@ impl KMeans {
     /// - *metric_type*: the metric type to calculate distance.
     /// - *rng*: random generator.
     pub async fn init_random(
-        data: &MatrixView,
+        data: &MatrixView<Float32Array>,
         k: usize,
         metric_type: MetricType,
         rng: impl Rng,
@@ -431,7 +431,7 @@ impl KMeans {
     ///   let kmeans = membership.to_kmeans();
     /// }
     /// ```
-    pub async fn train_once(&self, data: &MatrixView) -> KMeanMembership {
+    pub async fn train_once(&self, data: &MatrixView<Float32Array>) -> KMeanMembership {
         self.compute_membership(data.data().clone()).await
     }