fix(linalg): reject a null Int8 query instead of panicking

rust/lance-linalg/src/distance.rs

@@ -12,8 +12,10 @@
 use std::sync::Arc;
 
 use arrow_array::cast::AsArray;
-use arrow_array::types::{Float16Type, Float32Type, Float64Type, UInt8Type};
-use arrow_array::{Array, ArrowPrimitiveType, FixedSizeListArray, Float32Array, ListArray};
+use arrow_array::types::{Float16Type, Float32Type, Float64Type, Int8Type, UInt8Type};
+use arrow_array::{
+    Array, ArrowPrimitiveType, FixedSizeListArray, Float32Array, ListArray, PrimitiveArray,
+};
 use arrow_schema::{ArrowError, DataType};
 
 pub mod cosine;
@@ -26,6 +28,25 @@ pub mod l2;
 pub mod l2_u8;
 pub mod norm_l2;
 
+/// Widens an `Int8` query vector to `f32`, rejecting nulls.
+///
+/// The three `_arrow_batch` entry points that accept `Int8` take the query as a
+/// `&dyn Array`, so it has to be widened before it reaches a kernel. A null element has no
+/// distance to compute, and the widening this replaced resolved it with an
+/// `unwrap`.
+fn int8_query_to_f32(query: &PrimitiveArray<Int8Type>) -> Result<Float32Array> {
+    if query.null_count() > 0 {
+        return Err(ArrowError::InvalidArgumentError(format!(
+            "Int8 query vector `from` must not contain nulls, found {} in {} values",
+            query.null_count(),
+            query.len()
+        )));
+    }
+    Ok(Float32Array::from(
+        query.values().iter().map(|&v| v as f32).collect::<Vec<_>>(),
+    ))
+}
+
 #[inline]
 fn assert_equal_lengths(left_len: usize, right_len: usize) {
     assert_eq!(
@@ -501,6 +522,7 @@ mod tests {
     use arrow_buffer::{OffsetBuffer, ScalarBuffer};
     use arrow_schema::Field;
     use half::f16;
+    use lance_arrow::FixedSizeListArrayExt;
 
     #[cfg(target_arch = "x86_64")]
     #[test]
@@ -580,6 +602,83 @@ mod tests {
         );
     }
 
+    /// The `_arrow_batch` entry points that accept `Int8` widen the query element
+    /// by element.
+    /// A null there used to reach an `unwrap`, so a query column with a null in
+    /// its values panicked instead of returning an error, on the metrics that
+    /// accept `Int8`.
+    #[test]
+    fn test_arrow_batch_rejects_null_int8_query() {
+        let targets =
+            FixedSizeListArray::try_new_from_values(Int8Array::from(vec![1_i8, 2, 3, 4]), 2)
+                .unwrap();
+        let query: Arc<dyn Array> = Arc::new(Int8Array::from(vec![Some(1_i8), None]));
+
+        for dt in [DistanceType::L2, DistanceType::Cosine, DistanceType::Dot] {
+            let err = dt.arrow_batch_func()(query.as_ref(), &targets).unwrap_err();
+            assert!(
+                matches!(&err, ArrowError::InvalidArgumentError(m)
+                    if m.contains("Int8 query vector `from`") && m.contains("found 1 in 2 values")),
+                "{dt} accepted a null Int8 query element, got: {err}"
+            );
+        }
+
+        // The same query without nulls goes through, so the guard is not
+        // rejecting every `Int8` query.
+        let query: Arc<dyn Array> = Arc::new(Int8Array::from(vec![1_i8, 2]));
+        for dt in [DistanceType::L2, DistanceType::Cosine, DistanceType::Dot] {
+            assert_eq!(
+                dt.arrow_batch_func()(query.as_ref(), &targets)
+                    .unwrap()
+                    .len(),
+                2,
+                "{dt} rejected a well-formed Int8 query"
+            );
+        }
+
+        // A sliced query reads through `values()`, which has to follow the slice:
+        // the window here holds no nulls while the full buffer does. The L2
+        // distances are asserted literally rather than against a second call,
+        // since computing the expected values by the same route would hide a bug
+        // that transformed both alike. Query [3, 4] against [[1, 2], [3, 4]]
+        // gives (3-1)^2 + (4-2)^2 = 8 and 0.
+        let sliced = Int8Array::from(vec![None, Some(3_i8), Some(4), None]).slice(1, 2);
+        let query: Arc<dyn Array> = Arc::new(sliced);
+        let got = DistanceType::L2.arrow_batch_func()(query.as_ref(), &targets).unwrap();
+        assert_eq!(
+            got.values(),
+            &[8.0_f32, 0.0],
+            "L2 did not follow the query slice"
+        );
+        for dt in [DistanceType::Cosine, DistanceType::Dot] {
+            let got = dt.arrow_batch_func()(query.as_ref(), &targets).unwrap();
+            let want =
+                dt.arrow_batch_func()(Arc::new(Int8Array::from(vec![3_i8, 4])).as_ref(), &targets)
+                    .unwrap();
+            assert_eq!(got, want, "{dt} did not follow the query slice");
+        }
+    }
+
+    /// An input that is both a length mismatch and a null query must reach the
+    /// null error on all three metrics. `dot` used to carry a second
+    /// `debug_assert_eq!` on the dimension in its public entry point, ahead of
+    /// the `Int8` arm's null guard.
+    #[test]
+    fn test_arrow_batch_null_and_length_mismatch_agree() {
+        let targets =
+            FixedSizeListArray::try_new_from_values(Int8Array::from(vec![1_i8, 2, 3, 4]), 2)
+                .unwrap();
+        let query: Arc<dyn Array> = Arc::new(Int8Array::from(vec![Some(1_i8), None, Some(2)]));
+
+        for dt in [DistanceType::L2, DistanceType::Cosine, DistanceType::Dot] {
+            let err = dt.arrow_batch_func()(query.as_ref(), &targets).unwrap_err();
+            assert!(
+                matches!(&err, ArrowError::InvalidArgumentError(m) if m.contains("must not contain nulls")),
+                "{dt} did not report the null query, got: {err}"
+            );
+        }
+    }
+
     /// `Int8` is a valid vector element type elsewhere in the crate but has no
     /// multivector kernel, so it must be rejected for the type, not the metric.
     #[test]

rust/lance-linalg/src/distance/cosine.rs

@@ -24,7 +24,7 @@ use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray};
 #[allow(unused_imports)]
 use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport};
 
-use super::{Dot, norm_l2::norm_l2};
+use super::{Dot, int8_query_to_f32, norm_l2::norm_l2};
 use super::{Normalize, dot::dot};
 #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))]
 use crate::distance::BatchKind;
@@ -1362,9 +1362,18 @@ where
 /// - `from`: the vector to compute distance from.
 /// - `to`: a list of vectors to compute distance to.
 ///
+/// # Errors
+///
+/// Returns an error if `from` is an `Int8` array containing nulls, since a null
+/// query element has no distance to compute. The unsupported-type and downcast
+/// paths return errors of their own; this list is not exhaustive.
+///
 /// # Panics
 ///
-/// Panics if the length of `from` is not equal to the dimension (value length) of `to`.
+/// With debug assertions on, panics if the length of `from` is not equal to the
+/// dimension (value length) of `to`. Without them the mismatch is not reliably
+/// caught, since cosine has no always-on layout assert of its own, unlike the l2
+/// and dot equivalents.
 pub fn cosine_distance_arrow_batch(
     from: &dyn Array,
     to: &FixedSizeListArray,
@@ -1374,11 +1383,7 @@ pub fn cosine_distance_arrow_batch(
         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(),
+            &int8_query_to_f32(from.as_primitive::<Int8Type>())?,
             &to.convert_to_floating_point()?,
         ),
         _ => Err(Error::InvalidArgumentError(format!(

rust/lance-linalg/src/distance/dot.rs

@@ -28,7 +28,7 @@ use crate::Result;
     not(all(target_feature = "avx2", target_feature = "fma"))
 ))]
 use crate::distance::{BatchIter, BatchKernel, BatchKind, BatchOperation};
-use crate::distance::{assert_batch_layout, assert_equal_lengths};
+use crate::distance::{assert_batch_layout, assert_equal_lengths, int8_query_to_f32};
 #[cfg(all(
     target_arch = "x86_64",
     not(all(target_feature = "avx2", target_feature = "fma"))
@@ -804,26 +804,29 @@ where
 /// - `from`: the vector to compute distance from.
 /// - `to`: a list of vectors to compute distance to.
 ///
+/// # Errors
+///
+/// Returns an error if `from` is an `Int8` array containing nulls, since a null
+/// query element has no distance to compute. The unsupported-type and downcast
+/// paths return errors of their own; this list is not exhaustive.
+///
 /// # Panics
 ///
-/// Panics if the length of `from` is not equal to the dimension (value length) of `to`.
+/// Panics if the length of `from` is not equal to the dimension (value length)
+/// of `to`.
 pub fn dot_distance_arrow_batch(
     from: &dyn Array,
     to: &FixedSizeListArray,
 ) -> Result<Arc<Float32Array>> {
-    let dimension = to.value_length() as usize;
-    debug_assert_eq!(from.len(), dimension);
-
+    // `do_dot_distance_arrow_batch` already carries this `debug_assert_eq!`, and
+    // the copy that was here sat ahead of the `Int8` arm's null guard, where l2
+    // and cosine have theirs after it.
     match *from.data_type() {
         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(),
+            &int8_query_to_f32(from.as_primitive::<Int8Type>())?,
             &to.convert_to_floating_point()?,
         ),
         _ => Err(Error::InvalidArgumentError(format!(

rust/lance-linalg/src/distance/l2.rs

@@ -30,7 +30,7 @@ use lance_core::utils::cpu::SIMD_SUPPORT;
 use lance_core::utils::cpu::SimdSupport;
 use num_traits::{AsPrimitive, Num};
 
-use crate::distance::{assert_batch_layout, assert_equal_lengths};
+use crate::distance::{assert_batch_layout, assert_equal_lengths, int8_query_to_f32};
 
 #[cfg(all(
     target_arch = "x86_64",
@@ -969,9 +969,16 @@ where
 /// - `from`: the vector to compute distance from.
 /// - `to`: a list of vectors to compute distance to.
 ///
+/// # Errors
+///
+/// Returns an error if `from` is an `Int8` array containing nulls, since a null
+/// query element has no distance to compute. The unsupported-type and downcast
+/// paths return errors of their own; this list is not exhaustive.
+///
 /// # Panics
 ///
-/// Panics if the length of `from` is not equal to the dimension (value length) of `to`.
+/// Panics if the length of `from` is not equal to the dimension (value length)
+/// of `to`.
 pub fn l2_distance_arrow_batch(
     from: &dyn Array,
     to: &FixedSizeListArray,
@@ -981,11 +988,7 @@ pub fn l2_distance_arrow_batch(
         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(),
+            &int8_query_to_f32(from.as_primitive::<Int8Type>())?,
             &to.convert_to_floating_point()?,
         ),
         _ => Err(Error::ComputeError(format!(