diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 1f5be08b38b..30e96820e18 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -72,6 +72,10 @@ fn dot_scalar< } /// Dot product. +/// +/// # Panics +/// +/// Panics under the conditions [`Dot::dot`] documents. #[inline] pub fn dot(from: &[T], to: &[T]) -> f32 { T::dot(from, to) @@ -80,12 +84,20 @@ pub fn dot(from: &[T], to: &[T]) -> f32 { /// Dot product between two f32 slices, dispatched to the widest SIMD backend /// available at runtime. See [`crate::distance::l2::l2_f32`] for why this is /// needed on top of the generic [`dot`]. +/// +/// # Panics +/// +/// Panics if `x` and `y` have different lengths. #[inline] pub fn dot_f32(x: &[f32], y: &[f32]) -> f32 { f32::dot(x, y) } /// Negative [Dot] distance. +/// +/// # Panics +/// +/// Panics under the conditions [`Dot::dot`] documents. #[inline] pub fn dot_distance(from: &[T], to: &[T]) -> f32 { 1.0 - T::dot(from, to) @@ -94,6 +106,12 @@ pub fn dot_distance(from: &[T], to: &[T]) -> f32 { /// Dot product pub trait Dot: Num { /// Dot product. + /// + /// # Panics + /// + /// `x` and `y` must have the same length. An implementation is required to + /// reject a mismatch rather than read past the shorter slice; the five in + /// this crate do it by panicking. fn dot(x: &[Self], y: &[Self]) -> f32; /// Dot product of `x` against each `dimension`-sized vector in `batch`. @@ -106,6 +124,11 @@ pub trait Dot: Num { /// Returns `impl Iterator` rather than a trait object: hot consumers drive /// this one element at a time, so a `Box` would cost a /// virtual call per element and an allocation per batch. + /// + /// # Panics + /// + /// Panics unless `dimension` is non-zero, `x.len()` equals `dimension`, and + /// `batch.len()` is a multiple of `dimension`. fn dot_batch<'a>( x: &'a [Self], batch: &'a [Self], @@ -756,6 +779,10 @@ impl Dot for u8 { } /// Negative dot product, to present the relative order of dot distance. +/// +/// # Panics +/// +/// Panics under the conditions [`Dot::dot_batch`] documents. pub fn dot_distance_batch<'a, T: Dot>( from: &'a [T], to: &'a [T], diff --git a/rust/lance-linalg/src/distance/dot_u8.rs b/rust/lance-linalg/src/distance/dot_u8.rs index 32033f336e4..3cd5cdbb117 100644 --- a/rust/lance-linalg/src/distance/dot_u8.rs +++ b/rust/lance-linalg/src/distance/dot_u8.rs @@ -36,6 +36,10 @@ use super::{U8_U32_ACCUMULATOR_MAX_LEN, assert_equal_lengths}; /// /// The result is the low 32 bits of the exact dot product. Use /// [`dot_u8_u64`] when the full result is required. +/// +/// # Panics +/// +/// Panics if `a` and `b` have different lengths. #[inline] pub fn dot_u8_scalar(a: &[u8], b: &[u8]) -> u32 { assert_equal_lengths(a.len(), b.len()); @@ -153,6 +157,10 @@ fn select_backend() -> DotU8Fn { /// /// The result is the low 32 bits of the exact dot product. Use /// [`dot_u8_u64`] when the full result is required. +/// +/// # Panics +/// +/// Panics if `a` and `b` have different lengths. #[inline] pub fn dot_u8(a: &[u8], b: &[u8]) -> u32 { assert_equal_lengths(a.len(), b.len()); diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 16f4ef194c8..3fb21af012b 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -47,6 +47,12 @@ use crate::simd::x86::hsum256_ps; /// pub trait L2: Num { /// Calculate the L2 distance between two vectors. + /// + /// # Panics + /// + /// `x` and `y` must have the same length. An implementation is required to + /// reject a mismatch rather than read past the shorter slice; the five in + /// this crate do it by panicking. fn l2(x: &[Self], y: &[Self]) -> f32; /// L2 distance from `x` to each `dimension`-sized vector in `y`. @@ -60,6 +66,11 @@ pub trait L2: Num { /// assignment loop drives this one element at a time, so a /// `Box` would cost a virtual call per element and an /// allocation per batch. + /// + /// # Panics + /// + /// Panics unless `dimension` is non-zero, `x.len()` equals `dimension`, and + /// `y.len()` is a multiple of `dimension`. fn l2_batch<'a>( x: &'a [Self], y: &'a [Self], @@ -70,6 +81,11 @@ pub trait L2: Num { } } +/// Squared L2 distance between two vectors of any [`L2`] element type. +/// +/// # Panics +/// +/// Panics under the conditions [`L2::l2`] documents. #[inline] pub fn l2(from: &[T], to: &[T]) -> f32 { T::l2(from, to) @@ -82,12 +98,20 @@ pub fn l2(from: &[T], to: &[T]) -> f32 { /// to [`l2`], whose x86_64 implementation selects the best runtime-supported /// kernel. This entry point gives hot-path callers such as the in-memory HNSW /// index an explicit f32 API. +/// +/// # Panics +/// +/// Panics if `x` and `y` have different lengths. #[inline] pub fn l2_f32(x: &[f32], y: &[f32]) -> f32 { f32::l2(x, y) } /// Calculate L2 distance between two uint8 slices. +/// +/// # Panics +/// +/// Panics if `key` and `target` have different lengths. #[inline] pub fn l2_distance_uint_scalar(key: &[u8], target: &[u8]) -> f32 { assert_equal_lengths(key.len(), target.len()); @@ -102,6 +126,11 @@ pub fn l2_distance_uint_scalar(key: &[u8], target: &[u8]) -> f32 { /// It relies on LLVM for auto-vectorization and unrolling. /// /// This is pub for test/benchmark only. use [l2] instead. +/// +/// # Panics +/// +/// Panics if `from` and `to` have different lengths, and separately if `LANES` +/// is zero, which `chunks_exact` rejects. #[inline] pub fn l2_scalar< T: AsPrimitive, @@ -839,6 +868,14 @@ pub struct L2Prepared { impl L2Prepared { /// Transpose `targets` from AoS `[num_targets][dimension]` to SoA layout. + /// + /// `targets.len()` must be a multiple of `dimension`. + /// + /// # Panics + /// + /// Panics if `dimension` is zero. With debug assertions on, also panics if + /// `targets.len()` is not a multiple of `dimension`; without them the + /// trailing partial vector is dropped. pub fn new(targets: &[f32], dimension: usize) -> Self { let num_targets = targets.len() / dimension; debug_assert_eq!(targets.len(), num_targets * dimension); @@ -859,7 +896,13 @@ impl L2Prepared { /// Compute L2 distances from `query` to every target, writing into `out`. /// - /// `out` must have length `num_targets`. It will be zeroed before accumulation. + /// `query` must have length `dimension` and `out` must have length + /// `num_targets`. `out` will be zeroed before accumulation. + /// + /// # Panics + /// + /// With debug assertions on, panics unless both lengths match. A `query` + /// longer than `dimension` can also panic on a slice range without them. pub fn distances_into(&self, query: &[f32], out: &mut [f32]) { debug_assert_eq!(query.len(), self.dimension); debug_assert_eq!(out.len(), self.num_targets); @@ -908,6 +951,10 @@ impl L2Prepared { } /// Compute L2 distance between two vectors. +/// +/// # Panics +/// +/// Panics if `from` and `to` have different lengths. #[inline] pub fn l2_distance(from: &[f32], to: &[f32]) -> f32 { l2(from, to) @@ -924,6 +971,10 @@ pub fn l2_distance(from: &[f32], to: &[f32]) -> f32 { /// Returns /// /// An iterator of pair-wise distance between `from` vector to each vector in the batch. +/// +/// # Panics +/// +/// Panics under the conditions [`L2::l2_batch`] documents. pub fn l2_distance_batch<'a, T: L2>( from: &'a [T], to: &'a [T], diff --git a/rust/lance-linalg/src/distance/l2_u8.rs b/rust/lance-linalg/src/distance/l2_u8.rs index efcff21dcb4..2242b219a51 100644 --- a/rust/lance-linalg/src/distance/l2_u8.rs +++ b/rust/lance-linalg/src/distance/l2_u8.rs @@ -29,6 +29,10 @@ use super::{U8_U32_ACCUMULATOR_MAX_LEN, assert_equal_lengths}; /// /// The result is the low 32 bits of the exact squared distance. Use /// [`l2_u8_u64`] when the full result is required. +/// +/// # Panics +/// +/// Panics if `a` and `b` have different lengths. #[inline] pub fn l2_u8_scalar(a: &[u8], b: &[u8]) -> u32 { assert_equal_lengths(a.len(), b.len()); @@ -161,6 +165,10 @@ fn select_backend() -> L2U8Fn { /// /// The result is the low 32 bits of the exact squared distance. Use /// [`l2_u8_u64`] when the full result is required. +/// +/// # Panics +/// +/// Panics if `a` and `b` have different lengths. #[inline] pub fn l2_u8(a: &[u8], b: &[u8]) -> u32 { assert_equal_lengths(a.len(), b.len());