diff --git a/rust/lance-linalg/src/distance.rs b/rust/lance-linalg/src/distance.rs index 81a80aaacd4..0852f450886 100644 --- a/rust/lance-linalg/src/distance.rs +++ b/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,24 @@ 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, so it is rejected rather than widened. +fn int8_query_to_f32(query: &PrimitiveArray) -> Result { + 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::>(), + )) +} + #[inline] fn assert_equal_lengths(left_len: usize, right_len: usize) { assert_eq!( @@ -501,6 +521,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 +601,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 = 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 = 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 = 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 = 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] diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index 1512571f6fd..692683f5893 100644 --- a/rust/lance-linalg/src/distance/cosine.rs +++ b/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::(from.as_primitive(), to), DataType::Float64 => do_cosine_distance_arrow_batch::(from.as_primitive(), to), DataType::Int8 => do_cosine_distance_arrow_batch::( - &from - .as_primitive::() - .into_iter() - .map(|x| x.unwrap() as f32) - .collect(), + &int8_query_to_f32(from.as_primitive::())?, &to.convert_to_floating_point()?, ), _ => Err(Error::InvalidArgumentError(format!( diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 22274a074d2..fba1aa72695 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/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,6 +804,12 @@ 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`. @@ -811,19 +817,12 @@ pub fn dot_distance_arrow_batch( from: &dyn Array, to: &FixedSizeListArray, ) -> Result> { - let dimension = to.value_length() as usize; - debug_assert_eq!(from.len(), dimension); - match *from.data_type() { DataType::Float16 => do_dot_distance_arrow_batch::(from.as_primitive(), to), DataType::Float32 => do_dot_distance_arrow_batch::(from.as_primitive(), to), DataType::Float64 => do_dot_distance_arrow_batch::(from.as_primitive(), to), DataType::Int8 => do_dot_distance_arrow_batch::( - &from - .as_primitive::() - .into_iter() - .map(|x| x.unwrap() as f32) - .collect(), + &int8_query_to_f32(from.as_primitive::())?, &to.convert_to_floating_point()?, ), _ => Err(Error::InvalidArgumentError(format!( diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 3af13d58840..6defb60a378 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/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,6 +969,12 @@ 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`. @@ -981,11 +987,7 @@ pub fn l2_distance_arrow_batch( DataType::Float32 => do_l2_distance_arrow_batch::(from.as_primitive(), to), DataType::Float64 => do_l2_distance_arrow_batch::(from.as_primitive(), to), DataType::Int8 => do_l2_distance_arrow_batch::( - &from - .as_primitive::() - .into_iter() - .map(|x| x.unwrap() as f32) - .collect(), + &int8_query_to_f32(from.as_primitive::())?, &to.convert_to_floating_point()?, ), _ => Err(Error::ComputeError(format!(