From 4ddec3c9f7ffc6b57c6c37d555fc1bff647ad0b9 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 10 Aug 2026 21:00:46 -0400 Subject: [PATCH 1/4] Execute tensor product functions with RowFn Signed-off-by: Connor Tsui --- vortex-tensor/benches/cosine_similarity.rs | 35 +- vortex-tensor/benches/inner_product.rs | 43 +- vortex-tensor/src/encodings/normalized/mod.rs | 1 - .../src/scalar_fns/cosine_similarity.rs | 921 +++++------------- vortex-tensor/src/scalar_fns/inner_product.rs | 497 ++-------- vortex-tensor/src/scalar_fns/row.rs | 17 + .../src/scalar_fns/tests/cosine_similarity.rs | 607 ++++++++++++ .../src/scalar_fns/tests/inner_product.rs | 282 ++++++ vortex-tensor/src/scalar_fns/tests/mod.rs | 2 + vortex-tensor/src/utils.rs | 18 +- vortex-tensor/src/vector_search.rs | 4 +- 11 files changed, 1297 insertions(+), 1130 deletions(-) create mode 100644 vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs create mode 100644 vortex-tensor/src/scalar_fns/tests/inner_product.rs diff --git a/vortex-tensor/benches/cosine_similarity.rs b/vortex-tensor/benches/cosine_similarity.rs index 6cc5eb867ef..49551fbf701 100644 --- a/vortex-tensor/benches/cosine_similarity.rs +++ b/vortex-tensor/benches/cosine_similarity.rs @@ -21,11 +21,14 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity; @@ -41,15 +44,11 @@ fn main() { } /// Total `f64` elements per operand, held constant across widths: the row count is -/// `ELEMENTS / width`. This budget is a quarter of the one the other tensor benches use, because -/// the constant arms recompute the broadcast vector's norm per row and cost roughly ten times the -/// column arms per element. It is what keeps every arm inside the 1 ms per-iteration limit from -/// `docs/developer-guide/benchmarking.md`, measured against CodSpeed's CPU simulation. +/// `ELEMENTS / width`. The smaller budget keeps the wider cosine kernels inside the 1 ms +/// per-iteration limit from `docs/developer-guide/benchmarking.md` under CodSpeed simulation. const ELEMENTS: usize = 2_048; -/// Widths chosen to separate the two costs, as in `l2_norm.rs`: the redundant norm pass is -/// `O(rows * width)`, one third of the closure's arithmetic, so wide tensors show the hoist -/// while a narrow one is dominated by per-row framework costs. +/// Widths that expose both fixed row-framework costs and the `O(width)` kernel work. const WIDTHS: &[usize] = &[2, 32, 256]; /// `ELEMENTS / width` vectors of `width` `f64` elements, non-nullable. `seed` offsets the values so @@ -85,9 +84,9 @@ fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { bencher .with_inputs(|| { ( - CosineSimilarity::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + CosineSimilarity + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -106,6 +105,22 @@ fn column_x_constant(bencher: Bencher, width: usize) { bench_cosine(bencher, vectors(width, 0), constant_vector(width)); } +/// The lhs is a broadcast query vector, whose norm is the same in every row. +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_cosine(bencher, constant_vector(width), vectors(width, 31)); +} + +/// A nullable broadcast rhs exercises constant preparation and output validity together. +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_cosine(bencher, vectors(width, 0), rhs); +} + /// One query vector represented as an extension array over constant storage. fn extension_constant_vector(width: usize) -> ArrayRef { let ext_dtype = vectors(width, 0).dtype().as_extension().clone(); diff --git a/vortex-tensor/benches/inner_product.rs b/vortex-tensor/benches/inner_product.rs index 796e9b648d6..5c4adf1c7ec 100644 --- a/vortex-tensor/benches/inner_product.rs +++ b/vortex-tensor/benches/inner_product.rs @@ -15,11 +15,19 @@ use divan::Bencher; use divan::counter::ItemsCount; use mimalloc::MiMalloc; use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::scalar_fns::inner_product::InnerProduct; @@ -56,15 +64,25 @@ fn vectors(width: usize, seed: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +fn constant_vector(width: usize) -> ArrayRef { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let children = (0..width) + .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable)) + .collect(); + let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let vector = Scalar::extension::(EmptyMetadata, storage); + ConstantArray::new(vector, ELEMENTS / width).into_array() +} + fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(lhs.len())) .with_inputs(|| { ( - InnerProduct::try_new_array(lhs.clone(), rhs.clone()) - .unwrap() - .into_array(), + InnerProduct + .try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()]) + .unwrap(), session.create_execution_ctx(), ) }) @@ -84,3 +102,22 @@ fn nullable(bencher: Bencher, width: usize) { .into_array(); bench_inner_product(bencher, lhs, vectors(width, 31)); } + +#[divan::bench(args = WIDTHS)] +fn column_x_constant(bencher: Bencher, width: usize) { + bench_inner_product(bencher, vectors(width, 0), constant_vector(width)); +} + +#[divan::bench(args = WIDTHS)] +fn constant_x_column(bencher: Bencher, width: usize) { + bench_inner_product(bencher, constant_vector(width), vectors(width, 31)); +} + +#[divan::bench(args = WIDTHS)] +fn column_x_nullable_constant(bencher: Bencher, width: usize) { + let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); + let rhs = MaskedArray::try_new(constant_vector(width), validity) + .unwrap() + .into_array(); + bench_inner_product(bencher, vectors(width, 0), rhs); +} diff --git a/vortex-tensor/src/encodings/normalized/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs index 545236bba7d..2c9a72d2988 100644 --- a/vortex-tensor/src/encodings/normalized/mod.rs +++ b/vortex-tensor/src/encodings/normalized/mod.rs @@ -31,7 +31,6 @@ pub use array::NormalizedSlots; mod compress; pub use compress::NormalizedScheme; pub use compress::normalize; -pub(crate) use compress::try_build_constant_normalized; mod execute; diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 76d266a2721..6de8240d41d 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Cosine similarity expression for tensor-like types. +//! Cosine similarity between two tensor columns. +use num_traits::Float; use num_traits::Zero; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -14,36 +15,39 @@ use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; +use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::encodings::normalized::try_build_constant_normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; +use crate::scalar_fns::row::TensorRow; +#[cfg(test)] +use crate::scalar_fns::row::probe; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; +use crate::utils::l2_norm_row; /// Cosine similarity between two columns. /// /// Computes `dot(a, b) / (||a|| * ||b||)` over the flat backing buffer of each tensor or vector. /// The shape and permutation do not affect the result because cosine similarity only depends on the -/// element values, not their logical arrangement. +/// element values, not their logical arrangement. A zero norm on either side yields `0.0`. /// /// Both inputs must be tensor-like extension arrays ([`FixedShapeTensor`] or [`Vector`]) with the /// same dtype and a float element type. The output is a float column of the same float type. @@ -56,13 +60,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector /// [`Normalized`]: crate::encodings::normalized::Normalized -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct CosineSimilarity; impl CosineSimilarity { /// Creates a new [`TypedScalarFnInstance`] wrapping the cosine similarity operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(CosineSimilarity, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the cosine similarity between `lhs` and @@ -70,132 +74,94 @@ impl CosineSimilarity { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(CosineSimilarity::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for CosineSimilarity { +impl RowFn for CosineSimilarity { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("CosineSimilarity must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_prepared_into::<(TensorRow, TensorRow), UninitElementSink, _, _>( + |(lhs, rhs)| { + #[cfg(test)] + probe::record(lhs.is_some(), rhs.is_some()); + ConstantNorms { + lhs: lhs.map(l2_norm_row), + rhs: rhs.map(l2_norm_row), + } + }, + |norms, (lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { + InitializedElement::write( + output, + cosine_similarity_row_prepared(norms, lhs, rhs), + ) + } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands make the _stored_ norms and normalized children + /// authoritative: `cos(D(x, s), D(y, t)) = dot(x, y)` and `cos(D(x, s), y) = dot(x, y) / + /// ||y||`, in both cases forced to `0.0` on rows where any authoritative norm is `0.0` (even + /// for lossy children whose decoded coordinates are nonzero). + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, + args: &[ArrayRef], ctx: &mut ExecutionCtx, - ) -> VortexResult { - let mut lhs_ref = args.get(0)?; - let mut rhs_ref = args.get(1)?; - let len = args.row_count(); - - // If either side is a constant tensor-like extension array, eagerly normalize the single - // stored row and re-encode it as an `Normalized` whose children are both `ConstantArray`s. - // The `Normalized` fast path below then picks it up. - if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, len, ctx)? { - lhs_ref = normalized_array.into_array(); - } - if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, len, ctx)? { - rhs_ref = normalized_array.into_array(); - } + ) -> VortexResult> { + let lhs = args[0].clone(); + let rhs = args[1].clone(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { - NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); - } + match NormalizedOrientation::classify(&lhs, &rhs) { + NormalizedOrientation::Both { lhs, rhs } => cosine_both_normalized(lhs, rhs, ctx) + .map(|output| Some(RowExecution::Output(output))), NormalizedOrientation::One { normalized_array, plain, - } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); - } - NormalizedOrientation::Neither => {} + } => cosine_one_normalized(normalized_array, plain, ctx) + .map(|output| Some(RowExecution::Output(output))), + NormalizedOrientation::Neither => Ok(None), } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Compute inner product and norms as columnar operations, and propagate the options. - let norm_lhs_arr = L2Norm.try_new_array(len, EmptyOptions, [lhs_ref.clone()])?; - let norm_rhs_arr = L2Norm.try_new_array(len, EmptyOptions, [rhs_ref.clone()])?; - let dot_arr = InnerProduct::try_new_array(lhs_ref, rhs_ref)?; - - // Execute to get the inner product and norms of the arrays. We only fully decompress - // because we need to perform special logic (guard against 0) during division. - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - let norm_l: PrimitiveArray = norm_lhs_arr.into_array().execute(ctx)?; - let norm_r: PrimitiveArray = norm_rhs_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norm_l.as_slice::(); - let norms_r = norm_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - let denom = norms_l[i] * norms_r[i]; - - if denom == T::zero() { - T::zero() - } else { - dots[i] / denom - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false } } +vortex_array::impl_row_fn_vtable!(CosineSimilarity); + impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( &self, @@ -222,578 +188,177 @@ impl ScalarFnArrayVTable for CosineSimilarity { } } -impl CosineSimilarity { - /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so - /// `cosine_similarity = dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - // `Normalized` makes the normalized children authoritative, so their dot product is the - // cosine similarity even for lossy storage wrappers, except that a zero stored norm still - // represents a zero vector. - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let norms_l = norms_l.as_slice::(); - let norms_r = norms_r.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if norms_l[i] == T::zero() || norms_r[i] == T::zero() { - T::zero() - } else { - dots[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: treat the normalized child as authoritative, so - /// `cosine_similarity = dot(n, b) / ||b||`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, normalized_norms) = extract_normalized_children(normalized_ref); - - let dot_arr = InnerProduct::try_new_array(normalized, plain_ref.clone())?; - let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - - let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - - let norm_arr = L2Norm.try_new_array(len, EmptyOptions, [plain_ref.clone()])?; - let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; - - // TODO(connor): Ideally we would have a `SafeDiv` binary numeric operation. - // TODO(connor): This can be written in a more SIMD-friendly manner. - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let normalized_norms = normalized_norms.as_slice::(); - let plain_norms = plain_norm.as_slice::(); - let buffer: Buffer = (0..len) - .map(|i| { - if normalized_norms[i] == T::zero() || plain_norms[i] == T::zero() { - T::zero() - } else { - dots[i] / plain_norms[i] - } - }) - .collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } +/// Per-batch state for the cosine row kernel: the L2 norm of each operand that is constant for +/// the batch. +/// +/// A broadcast query vector holds the same elements in every row, so its norm is the same in +/// every row too. Computing it in the prepare step hoists an `O(width)` pass and a `sqrt` per row +/// out of the row loop. `None` marks an operand that varies by row, whose norm the row closure +/// computes exactly as it did before the hoist. +struct ConstantNorms { + /// The norm of the lhs when it is batch-constant. + lhs: Option, + + /// The norm of the rhs when it is batch-constant. + rhs: Option, } -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::cosine_similarity::CosineSimilarity; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::constant_tensor_array; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. - fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - #[test] - fn unit_vectors_1d() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 0.0, 1.0, 0.0, // Tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // Tensor 1 - 1.0, 0.0, 0.0, // Tensor 2 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - /// Single-row cosine similarity for various vector pairs. - #[rstest] - // Antiparallel -> -1.0. - #[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] - // dot=24, both magnitudes=5 -> 24/25 = 0.96. - #[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] - // Zero vector -> guarded to 0.0. - #[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); - Ok(()) - } - - /// Self-similarity across various tensor shapes should always produce 1.0. - #[rstest] - // 2x3 matrix, flattened to 6 elements. - #[case::matrix_2d( - &[2, 3], - &[ - 1.0, 0.0, 0.0, // row 0 - 0.0, 0.0, 0.0, // row 1 - ], - )] - // 2x2x2 tensor, 8 elements. - #[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] - fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { - let lhs = tensor_array(shape, elements)?; - let rhs = tensor_array(shape, elements)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn scalar_0d() -> VortexResult<()> { - // 0-dimensional tensor: each "tensor" is a single scalar value. - let lhs = tensor_array(&[], &[5.0, 3.0])?; - let rhs = tensor_array(&[], &[5.0, -3.0])?; - - // Same sign -> 1.0, opposite sign -> -1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); - Ok(()) - } - - #[test] - fn many_rows() -> VortexResult<()> { - // 5 tensors of shape [4] compared against themselves -> all 1.0. - let lhs = tensor_array( - &[4], - &[ - 1.0, 2.0, 3.0, 4.0, // tensor 0 - 0.0, 1.0, 0.0, 0.0, // tensor 1 - 5.0, 0.0, 5.0, 0.0, // tensor 2 - 1.0, 1.0, 1.0, 1.0, // tensor 3 - 0.0, 0.0, 0.0, 7.0, // tensor 4 - ], - )?; - let rhs = lhs.clone(); - - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0, 1.0, 1.0, 1.0, 1.0], - ); - Ok(()) - } - - #[test] - fn constant_query_tensor() -> VortexResult<()> { - // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. - let data = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 0.0, 1.0, 0.0, // tensor 1 - 0.0, 0.0, 1.0, // tensor 2 - 1.0, 0.0, 0.0, // tensor 3 - ], - )?; - let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn vector_unit_vectors() -> VortexResult<()> { - let lhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 1.0, 0.0, 0.0, // vector 1 - ], - )?; - - // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn vector_constant_query() -> VortexResult<()> { - let data = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // vector 0 - 0.0, 1.0, 0.0, // vector 1 - 0.0, 0.0, 1.0, // vector 2 - 1.0, 0.0, 0.0, // vector 3 - ], - )?; - let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; - - assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. - let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; - let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; - let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: self-similarity = 1.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_self_similarity() -> VortexResult<()> { - // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. - // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Self-similarity should always be 1.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); - Ok(()) - } - - #[test] - fn both_normalized_orthogonal() -> VortexResult<()> { - // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. - // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn both_normalized_zero_norm() -> VortexResult<()> { - // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS is plain [3.0, 4.0]. - // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[3.0, 4.0])?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 0.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); - Ok(()) - } - - #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on rhs). - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); - let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); - - let scalar_fn = CosineSimilarity::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[1.0]); - Ok(()) - } - - #[test] - fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine - // similarity for that row must be `0.0` even though the dot product of the normalized - // children is nonzero. - let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; - let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally violates the unit-norm invariant by pairing a nonzero normalized row - // with a stored norm of `0.0`, mimicking lossy storage. - // SAFETY: The children are structurally valid. - let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.into_array(); - - let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; - let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Same as above for the rhs operand. - // SAFETY: The children are structurally valid. - let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.into_array(); - - // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both - // `0.0`, so cosine similarity must be `0.0`. - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { - // Mimics a lossy encoding where the stored norm is authoritative but - // the decoded normalized child is physically nonzero. The plain side is a normal nonzero - // tensor with positive norm. cosine similarity must still be `0.0` because the - // authoritative stored norm on the normalized_array side is `0.0`. - let normalized = tensor_array(&[2], &[0.6, 0.8])?; - let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); - // Intentionally pairs a nonzero normalized row with a stored norm of `0.0`, mimicking - // lossy storage where the stored norm is authoritative. - // SAFETY: The children are structurally valid. - let normalized_array = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); - - let plain = tensor_array(&[2], &[1.0, 0.0])?; - - // Normalized encoding on the lhs: `One { normalized_array: lhs, plain: rhs }`. - assert_close( - &eval_cosine_similarity(normalized_array.clone(), plain.clone())?, - &[0.0], - ); - - // Normalized encoding on the rhs: `One { normalized_array: rhs, plain: lhs }`. The same - // zero-norm guard must fire regardless of operand order. - assert_close(&eval_cosine_similarity(plain, normalized_array)?, &[0.0]); - Ok(()) - } - - #[test] - fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { - // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. - // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. - let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 - 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 - 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 - 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { - // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[test] - fn both_constant_tensors() -> VortexResult<()> { - // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). - let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; - let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; - let expected = 1.0 / 2.0_f64.sqrt(); - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[expected, expected, expected], - ); - Ok(()) - } - - #[test] - fn constant_zero_norm_query() -> VortexResult<()> { - // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in - // `execute_one_normalized` and `execute_both_normalized`. - let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; - let rhs = tensor_array( - &[3], - &[ - 1.0, 2.0, 3.0, // - 4.0, 5.0, 6.0, // - 7.0, 8.0, 9.0, // - ], - )?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); - Ok(()) - } - - #[test] - fn constant_self_similarity_nonunit() -> VortexResult<()> { - // A non-unit constant query compared to itself must produce `1.0`. This exercises the - // helper's division: after normalization, both sides must be exactly unit so the - // Normalized fast path's inner product yields 1. - let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; - assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); - Ok(()) - } - - #[test] - fn vector_constant_matches_plain() -> VortexResult<()> { - // Exercise the `Vector` extension variant through the new pre-pass. - let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; - let rhs = vector_array( - 3, - &[ - 1.0, 0.0, 0.0, // - 1.0, 2.0, 2.0, // - 0.0, 0.0, 1.0, // - 2.0, 1.0, 2.0, // - ], - )?; - assert_close( - &eval_cosine_similarity(lhs, rhs)?, - &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], - ); - Ok(()) - } - - #[rstest] - #[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] - #[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = CosineSimilarity::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("CosineSimilarity serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) +/// Computes the cosine similarity of one row, taking any hoisted norm from `norms` and computing +/// the rest exactly as [`cosine_similarity_row`] does. +/// +/// Each arm accumulates the same values in the same order as [`cosine_similarity_row`], and the +/// denominator keeps its lhs-times-rhs order, so the result is bit-identical whether a norm was +/// hoisted or not. The match costs one predictable branch per row: the arm is the same for the +/// whole batch. +fn cosine_similarity_row_prepared( + norms: &ConstantNorms, + lhs: &[T], + rhs: &[T], +) -> T { + match (norms.lhs, norms.rhs) { + (None, None) => cosine_similarity_row(lhs, rhs), + (Some(lhs_norm), None) => { + let mut dot = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm_squared.sqrt()) + } + (None, Some(rhs_norm)) => { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + } + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm) + } + (Some(lhs_norm), Some(rhs_norm)) => { + let mut dot = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + } + cosine_from_parts(dot, lhs_norm * rhs_norm) + } } +} - fn cosine_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") - } +/// Computes the cosine similarity of two equal-length float slices. +/// +/// Returns `dot(a, b) / (||a|| * ||b||)`, or `0.0` when either norm is zero. +fn cosine_similarity_row(lhs: &[T], rhs: &[T]) -> T { + let mut dot = T::zero(); + let mut lhs_norm_squared = T::zero(); + let mut rhs_norm_squared = T::zero(); + for (&lhs_element, &rhs_element) in lhs.iter().zip(rhs) { + dot = dot + lhs_element * rhs_element; + lhs_norm_squared = lhs_norm_squared + lhs_element * lhs_element; + rhs_norm_squared = rhs_norm_squared + rhs_element * rhs_element; + } + + cosine_from_parts(dot, lhs_norm_squared.sqrt() * rhs_norm_squared.sqrt()) +} - fn cosine_vector_rhs() -> ArrayRef { - vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +/// The shared tail of every cosine arm: `dot / denominator`, guarded to `0.0` when it is +/// zero. +fn cosine_from_parts(dot: T, denominator: T) -> T { + if denominator == T::zero() { + T::zero() + } else { + dot / denominator } +} - fn cosine_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") - } +/// Both sides are [`Normalized`]-encoded: the normalized children are authoritative, so their dot +/// product is the cosine similarity, except that a row with a zero _stored_ norm is a zero vector. +/// +/// Unlike [`InnerProduct::reduce_encoded`], which composes lazy `Mul` arrays over the norm columns, +/// this executes and materializes. The zero-norm guard is a conditional per row rather than an +/// arithmetic factor, so there is no lazy array that expresses it; the norm columns are one value +/// per row rather than one per coordinate, so materializing them is cheap next to the decode this +/// avoids. +/// +/// [`InnerProduct::reduce_encoded`]: InnerProduct::reduce_encoded +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_both_normalized( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized_l, normalized_r])? + .execute(ctx)?; + let norms_l: PrimitiveArray = norms_l.execute(ctx)?; + let norms_r: PrimitiveArray = norms_r.execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let norms_l = norms_l.as_slice::(); + let norms_r = norms_r.as_slice::(); + // Zipped rather than indexed by `0..len`: one bounds check per iterator instead of three + // per row. A length disagreement between the children shortens the result, which the + // lifting reports against the batch row count rather than panicking mid-loop. + let buffer: Buffer = dots + .iter() + .zip(norms_l) + .zip(norms_r) + .map(|((&dot, &norm_l), &norm_r)| { + if norm_l.is_zero() || norm_r.is_zero() { + T::zero() + } else { + dot + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) +} - fn cosine_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") - } +/// One side is [`Normalized`]-encoded: `cos = dot(normalized, plain) / ||plain||`, forced to `0.0` +/// on rows where the stored norm or the plain norm is `0.0`. +/// +/// [`Normalized`]: crate::encodings::normalized::Normalized +fn cosine_one_normalized( + normalized_array: &ArrayRef, + plain: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = normalized_array.len(); + let (normalized, normalized_norms) = extract_normalized_children(normalized_array); + + let dot: PrimitiveArray = InnerProduct + .try_new_array(len, EmptyOptions, [normalized, plain.clone()])? + .execute(ctx)?; + let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; + let plain_norm: PrimitiveArray = L2Norm + .try_new_array(len, EmptyOptions, [plain.clone()])? + .execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let normalized_norms = normalized_norms.as_slice::(); + let plain_norms = plain_norm.as_slice::(); + // Zipped for the same reason as [`cosine_both_normalized`]. + let buffer: Buffer = dots + .iter() + .zip(normalized_norms) + .zip(plain_norms) + .map(|((&dot, &stored_norm), &plain_norm)| { + if stored_norm.is_zero() || plain_norm.is_zero() { + T::zero() + } else { + dot / plain_norm + } + }) + .collect(); + + Ok(PrimitiveArray::new(buffer, Validity::NonNullable).into_array()) + }) } diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 53ae82eb4a2..a25e36f9aab 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -6,40 +6,34 @@ use num_traits::Float; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; -use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::InitializedElement; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::serde::ArrayChildren; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::NormalizedOrientation; -use crate::matcher::AnyTensor; +use crate::scalar_fns::row::TensorRow; +use crate::scalar_fns::row::tensor_element_ptype; use crate::utils::BinaryTensorOpMetadata; -use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; -use crate::utils::validate_binary_tensor_float_inputs; /// Inner product (dot product) between two columns. /// @@ -52,13 +46,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector -#[derive(Clone)] +#[derive(Clone, Debug, Default)] pub struct InnerProduct; impl InnerProduct { /// Creates a new [`TypedScalarFnInstance`] wrapping the inner product operation. - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(InnerProduct, EmptyOptions) + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(Self, EmptyOptions) } /// Constructs a [`ScalarFnArray`] that lazily computes the inner product between `lhs` and @@ -66,119 +60,93 @@ impl InnerProduct { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). + /// Returns an error if the array cannot be constructed, such as when the input dtypes are + /// unsupported. pub fn try_new_array(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(InnerProduct::new().erased(), vec![lhs, rhs]) + ScalarFnArray::try_new(Self::new().erased(), vec![lhs, rhs]) } } -impl ScalarFnVTable for InnerProduct { +impl RowFn for InnerProduct { type Options = EmptyOptions; + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.inner_product"); *ID } - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) + fn serialize(&self, _options: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) } - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("lhs"), - 1 => ChildName::from("rhs"), - _ => unreachable!("InnerProduct must have exactly two children"), - } + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + Ok(EmptyOptions) } - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let lhs = &arg_dtypes[0]; - let rhs = &arg_dtypes[1]; - - // TODO(connor): relax the float-only gate once integer tensors are supported. - let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; - let ptype = tensor_match.element_ptype(); - let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); - Ok(DType::Primitive(ptype, nullability)) + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + match_each_float_ptype!(tensor_element_ptype(args)?, |T| { + visitor.visit_into::<(TensorRow, TensorRow), UninitElementSink, _>( + |(lhs, rhs), output| { + // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + unsafe { InitializedElement::write(output, inner_product_row(lhs, rhs)) } + }, + ) + }) } - fn execute( + /// [`Normalized`]-encoded operands factor through their stored norms: with `D(x, s)` denoting + /// `x * s` rowwise, `dot(D(x, s), D(y, t)) = s * t * dot(x, y)` and + /// `dot(D(x, s), y) = s * dot(x, y)`. The rewrite is expressed with lazy [`Operator::Mul`] + /// arrays over the (much smaller) norm columns, so no denormalized coordinates are decoded. + /// + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn reduce_encoded( &self, _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let lhs_ref = args.get(0)?; - let rhs_ref = args.get(1)?; - let len = args.row_count(); + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let len = args[0].len(); - // Take any Normalized read-through fast path that applies. - match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + Ok(match NormalizedOrientation::classify(&args[0], &args[1]) { NormalizedOrientation::Both { lhs, rhs } => { - return self.execute_both_normalized(lhs, rhs, len, ctx); + let (normalized_l, norms_l) = extract_normalized_children(lhs); + let (normalized_r, norms_r) = extract_normalized_children(rhs); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized_l, normalized_r])?; + Some( + dot.binary(norms_l, Operator::Mul)? + .binary(norms_r, Operator::Mul)?, + ) } NormalizedOrientation::One { normalized_array, plain, } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); + let (normalized, norms) = extract_normalized_children(normalized_array); + let dot = + InnerProduct.try_new_array(len, EmptyOptions, [normalized, plain.clone()])?; + Some(dot.binary(norms, Operator::Mul)?) } - NormalizedOrientation::Neither => {} + NormalizedOrientation::Neither => None, } - - // Compute combined validity. - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - // Canonicalize so we can perform the math directly. - let lhs: ExtensionArray = lhs_ref.execute(ctx)?; - let rhs: ExtensionArray = rhs_ref.execute(ctx)?; - - // We validated that both inputs have the same type. - let ext = lhs.dtype().as_extension(); - let tensor_match = ext - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let dimensions = tensor_match.list_size() as usize; - - // Extract the storage array from each extension input. We pass the storage (FSL) rather - // than the extension array to avoid canonicalizing the extension wrapper. - let lhs_storage = lhs.storage_array(); - let rhs_storage = rhs.storage_array(); - - let lhs_flat = extract_flat_elements(lhs_storage, dimensions, ctx)?; - let rhs_flat = extract_flat_elements(rhs_storage, dimensions, ctx)?; - - match_each_float_ptype!(lhs_flat.ptype(), |T| { - let buffer: Buffer = (0..len) - .map(|i| inner_product_row(lhs_flat.row::(i), rhs_flat.row::(i))) - .collect(); - - // SAFETY: The buffer length equals `row_count`, which matches the source validity - // length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - // The result is null if either input tensor is null. - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false + .map(RowExecution::Output)) } } +vortex_array::impl_row_fn_vtable!(InnerProduct); + impl ScalarFnArrayVTable for InnerProduct { fn serialize( &self, @@ -205,329 +173,12 @@ impl ScalarFnArrayVTable for InnerProduct { } } -impl InnerProduct { - /// Both sides are [`Normalized`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - fn execute_both_normalized( - &self, - lhs_ref: &ArrayRef, - rhs_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - - let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); - let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - - let norms_l: PrimitiveArray = norms_l.execute(ctx)?; - let norms_r: PrimitiveArray = norms_r.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let nl = norms_l.as_slice::(); - let nr = norms_r.as_slice::(); - let buffer: Buffer = (0..len).map(|i| nl[i] * nr[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } - - /// One side is [`Normalized`]-encoded: `inner_product = s * dot(n, other)`. - /// - /// [`Normalized`]: crate::encodings::normalized::Normalized - /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. - fn execute_one_normalized( - &self, - normalized_ref: &ArrayRef, - plain_ref: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - - let (normalized, norms) = extract_normalized_children(normalized_ref); - let normalized_norms: PrimitiveArray = norms.execute(ctx)?; - - let dot: PrimitiveArray = InnerProduct::try_new_array(normalized, plain_ref.clone())? - .into_array() - .execute(ctx)?; - - match_each_float_ptype!(dot.ptype(), |T| { - let dots = dot.as_slice::(); - let ns = normalized_norms.as_slice::(); - let buffer: Buffer = (0..len).map(|i| ns[i] * dots[i]).collect(); - - // SAFETY: The buffer length equals `len`, which matches the source validity length. - Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) - }) - } -} - /// Computes the inner product (dot product) of two equal-length float slices. /// /// Returns `sum(a_i * b_i)`. -fn inner_product_row(a: &[T], b: &[T]) -> T { - a.iter() - .zip(b.iter()) - .map(|(&x, &y)| x * y) - .fold(T::zero(), |acc, v| acc + v) -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::MaskedArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::ScalarFnArray; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::encodings::normalized::Normalized; - use crate::scalar_fns::inner_product::InnerProduct; - use crate::tests::SESSION; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::normalized_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. - fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - Ok(prim.as_slice::().to_vec()) - } - - /// Single-row inner product for various vector pairs. - #[rstest] - // Orthogonal: [1, 0] . [0, 1] = 0. - #[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] - // Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. - #[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] - // Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. - #[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] - // Scaled: [2, 0] . [3, 0] = 6. - #[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] - fn single_row( - #[case] shape: &[usize], - #[case] lhs_elems: &[f64], - #[case] rhs_elems: &[f64], - #[case] expected: &[f64], - ) -> VortexResult<()> { - let lhs = tensor_array(shape, lhs_elems)?; - let rhs = tensor_array(shape, rhs_elems)?; - assert_close(&eval_inner_product(lhs, rhs)?, expected); - Ok(()) - } - - #[test] - fn multiple_rows() -> VortexResult<()> { - let lhs = tensor_array( - &[3], - &[ - 1.0, 0.0, 0.0, // tensor 0 - 3.0, 4.0, 0.0, // tensor 1 - 1.0, 1.0, 1.0, // tensor 2 - ], - )?; - let rhs = tensor_array( - &[3], - &[ - 0.0, 1.0, 0.0, // tensor 0: dot = 0 - 3.0, 4.0, 0.0, // tensor 1: dot = 25 - 2.0, 2.0, 2.0, // tensor 2: dot = 6 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); - Ok(()) - } - - #[test] - fn vector_inner_product() -> VortexResult<()> { - let lhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0 - 1.0, 0.0, // vector 1 - ], - )?; - let rhs = vector_array( - 2, - &[ - 3.0, 4.0, // vector 0: dot = 25 - 0.0, 1.0, // vector 1: dot = 0 - ], - )?; - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn null_input_row() -> VortexResult<()> { - // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. - let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; - let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; - let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let mut ctx = SESSION.create_execution_ctx(); - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert!(prim.is_valid(2, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[23.0]); - assert_close(&[prim.as_slice::()[2]], &[127.0]); - Ok(()) - } - - #[test] - fn rejects_non_extension_dtype() { - let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); - let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - } - - #[test] - fn rejects_mismatched_dtypes() -> VortexResult<()> { - let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; - let rhs = vector_array(2, &[3.0_f64, 4.0])?; - let result = InnerProduct::try_new_array(lhs, rhs); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn both_normalized() -> VortexResult<()> { - // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). - // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). - // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; - - // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. - assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); - Ok(()) - } - - #[test] - fn both_normalized_multiple_rows() -> VortexResult<()> { - // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. - // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_lhs() -> VortexResult<()> { - // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // RHS: plain [1.0, 2.0]. - // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = tensor_array(&[2], &[1.0, 2.0])?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn one_side_normalized_rhs() -> VortexResult<()> { - // LHS: plain [1.0, 2.0]. - // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. - // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. - let mut ctx = SESSION.create_execution_ctx(); - let lhs = tensor_array(&[2], &[1.0, 2.0])?; - let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - - assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); - Ok(()) - } - - #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on lhs). - let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let lhs = Normalized::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); - let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - - let scalar_fn = InnerProduct::new().erased(); - let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; - let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; - - // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. - assert!(prim.is_valid(0, &mut ctx)?); - assert!(!prim.is_valid(1, &mut ctx)?); - assert_close(&[prim.as_slice::()[0]], &[25.0]); - Ok(()) - } - - #[rstest] - #[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] - #[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] - fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { - let original = InnerProduct::try_new_array(lhs.clone(), rhs.clone())?.into_array(); - - let plugin = ScalarFnArrayPlugin::new(InnerProduct); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("InnerProduct serialize must produce metadata"); - - let children = vec![lhs, rhs]; - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn inner_product_vector_lhs() -> ArrayRef { - vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") - } - - fn inner_product_vector_rhs() -> ArrayRef { - vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") - } - - fn inner_product_tensor_lhs() -> ArrayRef { - tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") - } - - fn inner_product_tensor_rhs() -> ArrayRef { - tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") - } +fn inner_product_row(lhs: &[T], rhs: &[T]) -> T { + lhs.iter() + .zip(rhs) + .map(|(&lhs_element, &rhs_element)| lhs_element * rhs_element) + .fold(T::zero(), |sum, product| sum + product) } diff --git a/vortex-tensor/src/scalar_fns/row.rs b/vortex-tensor/src/scalar_fns/row.rs index a99ad76ea6e..0b3cf3634ee 100644 --- a/vortex-tensor/src/scalar_fns/row.rs +++ b/vortex-tensor/src/scalar_fns/row.rs @@ -164,3 +164,20 @@ unsafe impl InputElement for TensorRow { } } } + +/// Records which operands a test's prepare step received as batch constants. +#[cfg(test)] +pub(crate) mod probe { + use std::cell::Cell; + + thread_local! { + /// Bit 0 records the lhs and bit 1 records the rhs. Thread-local storage prevents + /// concurrent tests from racing; row execution remains on the calling thread. + pub(crate) static SEEN_CONSTANTS: Cell = const { Cell::new(u8::MAX) }; + } + + /// Records whether each operand was constant for the current batch. + pub(crate) fn record(lhs_constant: bool, rhs_constant: bool) { + SEEN_CONSTANTS.set(u8::from(lhs_constant) | (u8::from(rhs_constant) << 1)); + } +} diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs new file mode 100644 index 00000000000..0e2687690a9 --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -0,0 +1,607 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::assert_arrays_eq; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::cosine_similarity::CosineSimilarity; +use crate::scalar_fns::row::probe; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::constant_tensor_array; +use crate::utils::test_helpers::literal_vector_array; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. +fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = CosineSimilarity::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = CosineSimilarity::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_cosine_similarity(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Like [`eval_cosine_similarity`], but returns the executed array for exact array comparisons. +fn eval_cosine_similarity_array( + lhs: ArrayRef, + rhs: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + Ok(result + .into_array() + .execute::(ctx)? + .into_array()) +} + +#[test] +fn unit_vectors_1d() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 0.0, 1.0, 0.0, // Tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // Tensor 1 + 1.0, 0.0, 0.0, // Tensor 2 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +/// Single-row cosine similarity for various vector pairs. +#[rstest] +// Antiparallel -> -1.0. +#[case::opposite(&[3], &[1.0, 0.0, 0.0], &[-1.0, 0.0, 0.0], &[-1.0])] +// dot=24, both magnitudes=5 -> 24/25 = 0.96. +#[case::non_unit(&[2], &[3.0, 4.0], &[4.0, 3.0], &[0.96])] +// Zero vector -> guarded to 0.0. +#[case::zero_norm(&[2], &[0.0, 0.0], &[1.0, 0.0], &[0.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, expected); + Ok(()) +} + +/// Self-similarity across various tensor shapes should always produce 1.0. +#[rstest] +// 2x3 matrix, flattened to 6 elements. +#[case::matrix_2d( + &[2, 3], + &[ + 1.0, 0.0, 0.0, // row 0 + 0.0, 0.0, 0.0, // row 1 + ], +)] +// 2x2x2 tensor, 8 elements. +#[case::tensor_3d(&[2, 2, 2], &[1.0; 8])] +fn self_similarity(#[case] shape: &[usize], #[case] elements: &[f64]) -> VortexResult<()> { + let lhs = tensor_array(shape, elements)?; + let rhs = tensor_array(shape, elements)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn scalar_0d() -> VortexResult<()> { + // 0-dimensional tensor: each "tensor" is a single scalar value. + let lhs = tensor_array(&[], &[5.0, 3.0])?; + let rhs = tensor_array(&[], &[5.0, -3.0])?; + + // Same sign -> 1.0, opposite sign -> -1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, -1.0]); + Ok(()) +} + +#[test] +fn many_rows() -> VortexResult<()> { + // 5 tensors of shape [4] compared against themselves -> all 1.0. + let lhs = tensor_array( + &[4], + &[ + 1.0, 2.0, 3.0, 4.0, // tensor 0 + 0.0, 1.0, 0.0, 0.0, // tensor 1 + 5.0, 0.0, 5.0, 0.0, // tensor 2 + 1.0, 1.0, 1.0, 1.0, // tensor 3 + 0.0, 0.0, 0.0, 7.0, // tensor 4 + ], + )?; + let rhs = lhs.clone(); + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0, 1.0, 1.0, 1.0, 1.0], + ); + Ok(()) +} + +#[test] +fn constant_query_tensor() -> VortexResult<()> { + // Compare 4 tensors of shape [3] against a single constant query tensor [1,0,0]. + let data = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 0.0, 1.0, 0.0, // tensor 1 + 0.0, 0.0, 1.0, // tensor 2 + 1.0, 0.0, 0.0, // tensor 3 + ], + )?; + let query = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn vector_unit_vectors() -> VortexResult<()> { + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 1.0, 0.0, 0.0, // vector 1 + ], + )?; + + // Row 0: identical -> 1.0, row 1: orthogonal -> 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn vector_constant_query() -> VortexResult<()> { + let data = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // vector 0 + 0.0, 1.0, 0.0, // vector 1 + 0.0, 0.0, 1.0, // vector 2 + 1.0, 0.0, 0.0, // vector 3 + ], + )?; + let query = Vector::constant_array(&[1.0, 0.0, 0.0], 4)?; + + assert_close(&eval_cosine_similarity(data, query)?, &[1.0, 0.0, 0.0, 1.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 2 rows of dim-2 vectors. Row 1 of rhs is masked as null. + let lhs = tensor_array(&[2], &[3.0, 4.0, 1.0, 0.0])?; + let rhs = tensor_array(&[2], &[3.0, 4.0, 0.0, 1.0])?; + let rhs = MaskedArray::try_new(rhs, Validity::from_iter([true, false]))?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: self-similarity = 1.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_self_similarity() -> VortexResult<()> { + // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. + // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Self-similarity should always be 1.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); + Ok(()) +} + +#[test] +fn both_normalized_orthogonal() -> VortexResult<()> { + // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. + // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn both_normalized_zero_norm() -> VortexResult<()> { + // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS is plain [3.0, 4.0]. + // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[3.0, 4.0])?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 0.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); + Ok(()) +} + +#[test] +fn both_normalized_null_norms() -> VortexResult<()> { + // Row 0: valid, row 1: null (via nullable norms on rhs). + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); + + let scalar_fn = CosineSimilarity.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[1.0]); + Ok(()) +} + +#[test] +fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine + // similarity for that row must be `0.0` even though the dot product of the normalized + // children is nonzero. + let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; + let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally violates the unit-norm invariant by + // pairing a nonzero normalized row with a stored norm of `0.0`, mimicking lossy storage. + let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.into_array(); + + let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; + let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: Same as above for the rhs operand. + let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.into_array(); + + // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both + // `0.0`, so cosine similarity must be `0.0`. + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + // Mimics a lossy encoding where the stored norm is authoritative but + // the decoded normalized child is physically nonzero. The plain side is a normal nonzero + // tensor with positive norm. cosine similarity must still be `0.0` because the + // authoritative stored norm on the denorm side is `0.0`. + let normalized = tensor_array(&[2], &[0.6, 0.8])?; + let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); + // SAFETY: This is a focused test that intentionally pairs a nonzero normalized row with a + // stored norm of `0.0`, mimicking lossy storage where the stored norm is authoritative. + let denorm = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + + let plain = tensor_array(&[2], &[1.0, 0.0])?; + + // Denorm on the lhs: `One { denorm: lhs, plain: rhs }`. + assert_close( + &eval_cosine_similarity(denorm.clone(), plain.clone())?, + &[0.0], + ); + + // Denorm on the rhs: `One { denorm: rhs, plain: lhs }`. The same zero-norm guard must + // fire regardless of operand order. + assert_close(&eval_cosine_similarity(plain, denorm)?, &[0.0]); + Ok(()) +} + +#[test] +fn constant_lhs_matches_plain_tensor() -> VortexResult<()> { + // The constant query `[1, 2, 2]` has norm 3, so its normalized form is `[1/3, 2/3, 2/3]`. + // Expected cosine similarity against each row is `dot([1, 2, 2], row) / (3 * ||row||)`. + let lhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // dot=1, ||rhs||=1, expected=1/3 + 1.0, 2.0, 2.0, // dot=9, ||rhs||=3, expected=1 + 0.0, 0.0, 1.0, // dot=2, ||rhs||=1, expected=2/3 + 2.0, 1.0, 2.0, // dot=8, ||rhs||=3, expected=8/9 + ], + )?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn constant_rhs_matches_plain_tensor() -> VortexResult<()> { + // Mirror of `constant_lhs_matches_plain_tensor` with the constant on the right. + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let rhs = constant_tensor_array(&[3], &[1.0, 2.0, 2.0], 4)?; + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + Ok(()) +} + +#[test] +fn both_constant_tensors() -> VortexResult<()> { + // `[1, 0, 0]` vs `[1, 1, 0]`. dot=1, ||lhs||=1, ||rhs||=sqrt(2), expected=1/sqrt(2). + let lhs = constant_tensor_array(&[3], &[1.0, 0.0, 0.0], 3)?; + let rhs = constant_tensor_array(&[3], &[1.0, 1.0, 0.0], 3)?; + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[expected, expected, expected], + ); + Ok(()) +} + +#[test] +fn constant_zero_norm_query() -> VortexResult<()> { + // A zero-norm constant query must produce `0.0` through the prepared row kernel's + // zero-denominator guard. + let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; + let rhs = tensor_array( + &[3], + &[ + 1.0, 2.0, 3.0, // + 4.0, 5.0, 6.0, // + 7.0, 8.0, 9.0, // + ], + )?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +#[test] +fn constant_self_similarity_nonunit() -> VortexResult<()> { + // The prepared path hoists both norms and computes the same dot product for every row. + let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); + Ok(()) +} + +/// An extension array over constant storage (what [`Vector::constant_array`] builds) is a batch +/// constant like any other. The row layer sees through the wrapper, so `prepare` hoists its norm. +#[test] +fn vector_constant_matches_plain() -> VortexResult<()> { + let lhs = Vector::constant_array(&[1.0, 2.0, 2.0], 4)?; + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + + assert_close( + &eval_cosine_similarity(lhs, rhs)?, + &[1.0 / 3.0, 1.0, 2.0 / 3.0, 8.0 / 9.0], + ); + assert_eq!( + probe::SEEN_CONSTANTS.get(), + 0b01, + "the extension-over-constant lhs must reach prepare as a batch constant", + ); + Ok(()) +} + +/// Both literal and extension-wrapped constant storage reach the prepared row path. The probe +/// ensures that the literal query remains a batch constant instead of becoming a varying column. +/// +/// [`ConstantArray`]: vortex_array::arrays::ConstantArray +#[test] +fn literal_constant_rhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(lhs.clone(), literal_vector_array(&query, 4), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(lhs, vector_array(3, &query.repeat(4))?, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// The mirror of [`literal_constant_rhs_matches_expanded_column`], exercising the hoisted-lhs arm +/// of the prepared kernel. +#[test] +fn literal_constant_lhs_matches_expanded_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let rhs = vector_array( + 3, + &[ + 1.0, 0.0, 0.0, // + 1.0, 2.0, 2.0, // + 0.0, 0.0, 1.0, // + 2.0, 1.0, 2.0, // + ], + )?; + let query = [1.0, 2.0, 2.0]; + + let from_constant = + eval_cosine_similarity_array(literal_vector_array(&query, 4), rhs.clone(), &mut ctx)?; + let from_expanded = + eval_cosine_similarity_array(vector_array(3, &query.repeat(4))?, rhs, &mut ctx)?; + + assert_arrays_eq!(from_constant, from_expanded, &mut ctx); + Ok(()) +} + +/// A zero-norm literal constant query must be guarded to `0.0` on every row by the prepared row +/// kernel, exactly as the unprepared kernel guards it. +#[test] +fn literal_constant_zero_norm_query_yields_zero() -> VortexResult<()> { + let lhs = vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = literal_vector_array(&[0.0f64, 0.0, 0.0], 2); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0, 0.0]); + Ok(()) +} + +/// Two literal constants are folded to a single-row execution by the row lifting, and that row +/// still runs the prepared kernel with both norms hoisted. +#[test] +fn both_literal_constants() -> VortexResult<()> { + let lhs = literal_vector_array(&[1.0f64, 0.0, 0.0], 3); + let rhs = literal_vector_array(&[1.0f64, 1.0, 0.0], 3); + let expected = 1.0 / 2.0_f64.sqrt(); + assert_close(&eval_cosine_similarity(lhs, rhs)?, &[expected; 3]); + Ok(()) +} + +#[rstest] +#[case::vector(cosine_vector_lhs(), cosine_vector_rhs())] +#[case::fixed_shape_tensor(cosine_tensor_lhs(), cosine_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + CosineSimilarity.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("CosineSimilarity serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn cosine_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_vector_rhs() -> ArrayRef { + vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).expect("valid vector array") +} + +fn cosine_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn cosine_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/inner_product.rs b/vortex-tensor/src/scalar_fns/tests/inner_product.rs new file mode 100644 index 00000000000..9a5a60f5dda --- /dev/null +++ b/vortex-tensor/src/scalar_fns/tests/inner_product.rs @@ -0,0 +1,282 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::MaskedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::encodings::normalized::Normalized; +use crate::scalar_fns::inner_product::InnerProduct; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::normalized_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; +use crate::utils::test_helpers::zero_width_vector_array; + +/// Evaluates inner product between two tensor arrays and returns the result as `Vec`. +fn eval_inner_product(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + Ok(prim.as_slice::().to_vec()) +} + +#[test] +fn inherent_constructors_remain_available() -> VortexResult<()> { + let _scalar_fn = InnerProduct::new(); + let lhs = tensor_array(&[1], &[2.0])?; + let rhs = tensor_array(&[1], &[3.0])?; + let array = InnerProduct::try_new_array(lhs, rhs)?; + + assert_eq!(array.len(), 1); + Ok(()) +} + +#[test] +fn zero_width_and_empty_inputs() -> VortexResult<()> { + let lhs = zero_width_vector_array::(3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + + let lhs = vector_array(2, &[] as &[f64])?; + let rhs = vector_array(2, &[] as &[f64])?; + assert!(eval_inner_product(lhs, rhs)?.is_empty()); + + let lhs = Vector::constant_array::(&[], 3)?; + let rhs = zero_width_vector_array::(3)?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 0.0, 0.0]); + Ok(()) +} + +/// Single-row inner product for various vector pairs. +#[rstest] +// Orthogonal: [1, 0] . [0, 1] = 0. +#[case::orthogonal(&[2], &[1.0, 0.0], &[0.0, 1.0], &[0.0])] +// Parallel: [3, 4] . [3, 4] = 9 + 16 = 25. +#[case::parallel(&[2], &[3.0, 4.0], &[3.0, 4.0], &[25.0])] +// Antiparallel: [1, 2] . [-1, -2] = -1 + -4 = -5. +#[case::antiparallel(&[2], &[1.0, 2.0], &[-1.0, -2.0], &[-5.0])] +// Scaled: [2, 0] . [3, 0] = 6. +#[case::scaled(&[2], &[2.0, 0.0], &[3.0, 0.0], &[6.0])] +fn single_row( + #[case] shape: &[usize], + #[case] lhs_elems: &[f64], + #[case] rhs_elems: &[f64], + #[case] expected: &[f64], +) -> VortexResult<()> { + let lhs = tensor_array(shape, lhs_elems)?; + let rhs = tensor_array(shape, rhs_elems)?; + assert_close(&eval_inner_product(lhs, rhs)?, expected); + Ok(()) +} + +#[test] +fn multiple_rows() -> VortexResult<()> { + let lhs = tensor_array( + &[3], + &[ + 1.0, 0.0, 0.0, // tensor 0 + 3.0, 4.0, 0.0, // tensor 1 + 1.0, 1.0, 1.0, // tensor 2 + ], + )?; + let rhs = tensor_array( + &[3], + &[ + 0.0, 1.0, 0.0, // tensor 0: dot = 0 + 3.0, 4.0, 0.0, // tensor 1: dot = 25 + 2.0, 2.0, 2.0, // tensor 2: dot = 6 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[0.0, 25.0, 6.0]); + Ok(()) +} + +#[test] +fn vector_inner_product() -> VortexResult<()> { + let lhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0 + 1.0, 0.0, // vector 1 + ], + )?; + let rhs = vector_array( + 2, + &[ + 3.0, 4.0, // vector 0: dot = 25 + 0.0, 1.0, // vector 1: dot = 0 + ], + )?; + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn null_input_row() -> VortexResult<()> { + // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. + let lhs = tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0])?; + let rhs = tensor_array(&[2], &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0])?; + let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let mut ctx = SESSION.create_execution_ctx(); + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 1*7 + 2*8 = 23, row 1: null, row 2: 5*11 + 6*12 = 127. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert!(prim.is_valid(2, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[23.0]); + assert_close(&[prim.as_slice::()[2]], &[127.0]); + Ok(()) +} + +#[test] +fn rejects_non_extension_dtype() { + let lhs = PrimitiveArray::from_iter([1.0_f64, 2.0]).into_array(); + let rhs = PrimitiveArray::from_iter([3.0_f64, 4.0]).into_array(); + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); +} + +#[test] +fn rejects_mismatched_dtypes() -> VortexResult<()> { + let lhs = tensor_array(&[2], &[1.0_f64, 2.0])?; + let rhs = vector_array(2, &[3.0_f64, 4.0])?; + let result = InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs, rhs]); + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn both_normalized() -> VortexResult<()> { + // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). + // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). + // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; + + // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. + assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); + Ok(()) +} + +#[test] +fn both_normalized_multiple_rows() -> VortexResult<()> { + // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. + // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS: plain [1.0, 2.0]. + // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = tensor_array(&[2], &[1.0, 2.0])?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS: plain [1.0, 2.0]. + // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. + let mut ctx = SESSION.create_execution_ctx(); + let lhs = tensor_array(&[2], &[1.0, 2.0])?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + + assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); + Ok(()) +} + +#[test] +fn both_normalized_null_norms() -> VortexResult<()> { + // Row 0: valid, row 1: null (via nullable norms on lhs). + let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; + let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + let lhs = Normalized::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + + let scalar_fn = InnerProduct.bind(EmptyOptions); + let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; + let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; + + // Row 0: 5.0 * 5.0 * dot([0.6, 0.8], [0.6, 0.8]) = 25.0, row 1: null. + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[25.0]); + Ok(()) +} + +#[rstest] +#[case::vector(inner_product_vector_lhs(), inner_product_vector_rhs())] +#[case::fixed_shape_tensor(inner_product_tensor_lhs(), inner_product_tensor_rhs())] +fn serde_round_trip(#[case] lhs: ArrayRef, #[case] rhs: ArrayRef) -> VortexResult<()> { + let original = + InnerProduct.try_new_array(lhs.len(), EmptyOptions, [lhs.clone(), rhs.clone()])?; + + let plugin = ScalarFnArrayPlugin::new(InnerProduct); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("InnerProduct serialize must produce metadata"); + + let children = vec![lhs, rhs]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) +} + +fn inner_product_vector_lhs() -> ArrayRef { + vector_array(3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).expect("valid vector array") +} + +fn inner_product_vector_rhs() -> ArrayRef { + vector_array(3, &[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).expect("valid vector array") +} + +fn inner_product_tensor_lhs() -> ArrayRef { + tensor_array(&[2], &[1.0, 2.0, 3.0, 4.0]).expect("valid tensor array") +} + +fn inner_product_tensor_rhs() -> ArrayRef { + tensor_array(&[2], &[5.0, 6.0, 7.0, 8.0]).expect("valid tensor array") +} diff --git a/vortex-tensor/src/scalar_fns/tests/mod.rs b/vortex-tensor/src/scalar_fns/tests/mod.rs index 5447772adda..bb3726e9329 100644 --- a/vortex-tensor/src/scalar_fns/tests/mod.rs +++ b/vortex-tensor/src/scalar_fns/tests/mod.rs @@ -3,5 +3,7 @@ //! Tests for the tensor scalar functions. +mod cosine_similarity; +mod inner_product; mod l2_norm; mod row; diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index df40bdea44a..62f0713319b 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -63,6 +63,9 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { } /// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. +/// +/// L2 norm and cosine similarity share this implementation so prepared constant norms use the +/// same accumulation order as varying rows. pub(crate) fn l2_norm_row(row: &[T]) -> T { let mut sum_squared = T::zero(); for &element in row { @@ -111,19 +114,6 @@ pub fn validate_tensor_float_input(input_dtype: &DType) -> VortexResult( - lhs: &'a DType, - rhs: &DType, -) -> VortexResult> { - vortex_ensure!( - lhs.eq_ignore_nullability(rhs), - "binary tensor expression expects inputs to have the same dtype, got {lhs} and {rhs}" - ); - validate_tensor_float_input(lhs) -} - /// Validates that every argument has the same float tensor dtype, ignoring nullability. pub fn validate_tensor_float_inputs(args: &[DType]) -> VortexResult> { let (first, rest) = args @@ -322,7 +312,7 @@ impl BinaryTensorOpMetadata { let lhs_dtype = DType::from_proto(lhs_pb, session)?; let rhs_dtype = DType::from_proto(rhs_pb, session)?; - validate_binary_tensor_float_inputs(&lhs_dtype, &rhs_dtype)?; + validate_tensor_float_inputs(&[lhs_dtype.clone(), rhs_dtype.clone()])?; let lhs = children.get(0, &lhs_dtype, len)?; let rhs = children.get(1, &rhs_dtype, len)?; diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index ad3b96d1bff..492bc837b89 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -35,11 +35,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::scalar::PValue; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; @@ -79,7 +81,7 @@ pub fn build_similarity_search_tree>( let num_rows = data.len(); let query_vec = Vector::constant_array(query, num_rows)?; - let cosine = CosineSimilarity::try_new_array(data, query_vec)?.into_array(); + let cosine = CosineSimilarity.try_new_array(num_rows, EmptyOptions, [data, query_vec])?; let threshold_scalar = Scalar::primitive(threshold, Nullability::NonNullable); let threshold_array = ConstantArray::new(threshold_scalar, num_rows).into_array(); From 87bd382973efe4264d749e718cdbbda139f603a0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 11 Aug 2026 13:09:00 -0400 Subject: [PATCH 2/4] Update tensor products for the RowFn API Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 2 -- vortex-tensor/src/scalar_fns/inner_product.rs | 2 -- vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs | 2 +- vortex-tensor/src/utils.rs | 2 +- 4 files changed, 2 insertions(+), 6 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 6de8240d41d..a2788cdaae1 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -160,8 +160,6 @@ impl RowFn for CosineSimilarity { } } -vortex_array::impl_row_fn_vtable!(CosineSimilarity); - impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index a25e36f9aab..976fbe3cb18 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -145,8 +145,6 @@ impl RowFn for InnerProduct { } } -vortex_array::impl_row_fn_vtable!(InnerProduct); - impl ScalarFnArrayVTable for InnerProduct { fn serialize( &self, diff --git a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs index 0e2687690a9..ecdbe23f835 100644 --- a/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/tests/cosine_similarity.rs @@ -490,7 +490,7 @@ fn vector_constant_matches_plain() -> VortexResult<()> { } /// Both literal and extension-wrapped constant storage reach the prepared row path. The probe -/// ensures that the literal query remains a batch constant instead of becoming a varying column. +/// ensures that the literal query remains a batch constant instead of becoming a per-row column. /// /// [`ConstantArray`]: vortex_array::arrays::ConstantArray #[test] diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 62f0713319b..c05dcf5c1b6 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -65,7 +65,7 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { /// Computes `sqrt(sum(v_i^2))` for one row. An empty or all-zero row produces `0.0`. /// /// L2 norm and cosine similarity share this implementation so prepared constant norms use the -/// same accumulation order as varying rows. +/// same accumulation order as rows from per-row inputs. pub(crate) fn l2_norm_row(row: &[T]) -> T { let mut sum_squared = T::zero(); for &element in row { From 659aa973a47f2f501a464c59063da76b34f14772 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 13:38:43 -0400 Subject: [PATCH 3/4] Opt tensor products into unstable RowFn Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 10 +++++----- vortex-tensor/src/scalar_fns/inner_product.rs | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index a2788cdaae1..e3defc3432f 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -18,13 +18,13 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowExecution; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_buffer::Buffer; diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 976fbe3cb18..d6a8be4d333 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -16,14 +16,14 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::InitializedElement; -use vortex_array::scalar_fn::RowExecution; -use vortex_array::scalar_fn::RowFn; -use vortex_array::scalar_fn::RowVisitor; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::UninitElementSink; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::unstable::row::InitializedElement; +use vortex_array::scalar_fn::unstable::row::RowExecution; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::UninitElementSink; use vortex_array::serde::ArrayChildren; use vortex_error::VortexResult; use vortex_session::VortexSession; From c3ed052b5045ee1afb8d4e6679d05fe8cb978238 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:22:42 -0400 Subject: [PATCH 4/4] Declare tensor product fallibility Signed-off-by: Connor Tsui --- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 1 + vortex-tensor/src/scalar_fns/inner_product.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index e3defc3432f..1582ef9c7d0 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -85,6 +85,7 @@ impl RowFn for CosineSimilarity { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index d6a8be4d333..a31cce5ddda 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -71,6 +71,7 @@ impl RowFn for InnerProduct { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.inner_product");