diff --git a/vortex-tensor/benches/cosine_similarity.rs b/vortex-tensor/benches/cosine_similarity.rs index 5b800ff0c92..df41b4ffc11 100644 --- a/vortex-tensor/benches/cosine_similarity.rs +++ b/vortex-tensor/benches/cosine_similarity.rs @@ -28,6 +28,8 @@ use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_tensor::encodings::normalized::Normalized; +use vortex_tensor::scalar_fns::NormMode; use vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity; use vortex_tensor::vector::Vector; @@ -80,12 +82,37 @@ fn constant_vector(width: usize) -> ArrayRef { ConstantArray::new(ext_scalar, ELEMENTS / width).into_array() } -fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { +fn normalized_vectors(width: usize, offset: usize) -> ArrayRef { + let row_count = ELEMENTS / width; + let elements: Buffer = (0..ELEMENTS) + .map(|i| { + if i % width == offset % width { + 1.0 + } else { + 0.0 + } + }) + .collect(); + let storage = FixedSizeListArray::new( + elements.into_array(), + u32::try_from(width).unwrap(), + Validity::NonNullable, + row_count, + ) + .into_array(); + let direction = Vector::try_new_vector_array(storage).unwrap(); + let norms = PrimitiveArray::from_iter((1..=row_count).map(|norm| norm as f64)).into_array(); + + // SAFETY: Every direction row is unit length, and the non-negative norms have matching length. + unsafe { Normalized::new_unchecked(direction, norms, Validity::NonNullable) }.into_array() +} + +fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, mode: NormMode) { let session = vortex_array::array_session(); bencher .with_inputs(|| { ( - CosineSimilarity::try_new(lhs.clone(), rhs.clone()) + CosineSimilarity::try_new(lhs.clone(), rhs.clone(), mode) .unwrap() .into_array(), session.create_execution_ctx(), @@ -97,13 +124,23 @@ fn bench_cosine(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) { /// The control: both operands vary by row, so every norm must be computed in the row loop. #[divan::bench(args = WIDTHS)] fn column_x_column(bencher: Bencher, width: usize) { - bench_cosine(bencher, vectors(width, 0), vectors(width, 31)); + bench_cosine( + bencher, + vectors(width, 0), + vectors(width, 31), + NormMode::Exact, + ); } /// The rhs is a broadcast query vector, whose norm is the same in every row. #[divan::bench(args = WIDTHS)] fn column_x_constant(bencher: Bencher, width: usize) { - bench_cosine(bencher, vectors(width, 0), constant_vector(width)); + bench_cosine( + bencher, + vectors(width, 0), + constant_vector(width), + NormMode::Exact, + ); } /// One query vector represented as an extension array over constant storage. @@ -124,5 +161,32 @@ fn extension_constant_vector(width: usize) -> ArrayRef { /// The rhs is the same broadcast query represented as extension-wrapped constant storage. #[divan::bench(args = WIDTHS)] fn column_x_extension_constant(bencher: Bencher, width: usize) { - bench_cosine(bencher, vectors(width, 0), extension_constant_vector(width)); + bench_cosine( + bencher, + vectors(width, 0), + extension_constant_vector(width), + NormMode::Exact, + ); +} + +/// Measures exact cosine similarity over two [`Normalized`] inputs. +#[divan::bench(args = WIDTHS)] +fn normalized_exact(bencher: Bencher, width: usize) { + bench_cosine( + bencher, + normalized_vectors(width, 0), + normalized_vectors(width, 1), + NormMode::Exact, + ); +} + +/// Measures cosine similarity while trusting both normalized-direction claims. +#[divan::bench(args = WIDTHS)] +fn normalized_assume(bencher: Bencher, width: usize) { + bench_cosine( + bencher, + normalized_vectors(width, 0), + normalized_vectors(width, 1), + NormMode::AssumeNormalized, + ); } diff --git a/vortex-tensor/benches/l2_norm.rs b/vortex-tensor/benches/l2_norm.rs index bb5dc0dd0ab..83a9e6de277 100644 --- a/vortex-tensor/benches/l2_norm.rs +++ b/vortex-tensor/benches/l2_norm.rs @@ -22,6 +22,8 @@ use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_tensor::encodings::normalized::Normalized; +use vortex_tensor::scalar_fns::NormMode; use vortex_tensor::scalar_fns::l2_norm::L2Norm; use vortex_tensor::vector::Vector; @@ -54,13 +56,32 @@ fn vectors(width: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } -fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { +fn normalized_vectors(width: usize) -> ArrayRef { + let row_count = ELEMENTS / width; + let elements: Buffer = (0..ELEMENTS) + .map(|i| if i % width == 0 { 1.0 } else { 0.0 }) + .collect(); + let storage = FixedSizeListArray::new( + elements.into_array(), + u32::try_from(width).unwrap(), + Validity::NonNullable, + row_count, + ) + .into_array(); + let direction = Vector::try_new_vector_array(storage).unwrap(); + let norms = PrimitiveArray::from_iter((1..=row_count).map(|norm| norm as f64)).into_array(); + + // SAFETY: Every direction row is unit length, and the non-negative norms have matching length. + unsafe { Normalized::new_unchecked(direction, norms, Validity::NonNullable) }.into_array() +} + +fn bench_l2_norm(bencher: Bencher, input: ArrayRef, mode: NormMode) { let session = vortex_array::array_session(); bencher .counter(ItemsCount::new(input.len())) .with_inputs(|| { ( - L2Norm::try_new(input.clone()).unwrap().into_array(), + L2Norm::try_new(input.clone(), mode).unwrap().into_array(), session.create_execution_ctx(), ) }) @@ -69,7 +90,7 @@ fn bench_l2_norm(bencher: Bencher, input: ArrayRef) { #[divan::bench(args = WIDTHS)] fn non_nullable(bencher: Bencher, width: usize) { - bench_l2_norm(bencher, vectors(width)); + bench_l2_norm(bencher, vectors(width), NormMode::Exact); } #[divan::bench(args = WIDTHS)] @@ -78,5 +99,21 @@ fn nullable(bencher: Bencher, width: usize) { let input = MaskedArray::try_new(vectors(width), validity) .unwrap() .into_array(); - bench_l2_norm(bencher, input); + bench_l2_norm(bencher, input, NormMode::Exact); +} + +/// Measures the physical norm of a [`Normalized`] input. +#[divan::bench(args = WIDTHS)] +fn normalized_exact(bencher: Bencher, width: usize) { + bench_l2_norm(bencher, normalized_vectors(width), NormMode::Exact); +} + +/// Reads the stored norm while trusting the normalized-direction claim. +#[divan::bench(args = WIDTHS)] +fn normalized_assume(bencher: Bencher, width: usize) { + bench_l2_norm( + bencher, + normalized_vectors(width), + NormMode::AssumeNormalized, + ); } diff --git a/vortex-tensor/src/encodings/normalized/array.rs b/vortex-tensor/src/encodings/normalized/array.rs index 35c8bdc4ef4..1e995ca9402 100644 --- a/vortex-tensor/src/encodings/normalized/array.rs +++ b/vortex-tensor/src/encodings/normalized/array.rs @@ -45,6 +45,10 @@ pub type NormalizedArray = Array; /// Row `i` decodes to `normalized[i] * norms[i]`. The encoding supports [`Vector`] and /// [`FixedShapeTensor`] columns with float elements. /// +/// `Normalized` is a physical representation of those logical tensor types. The direction is +/// evidence that norm-based scalar functions can use when their [`NormMode`] permits an +/// approximation; it is not a logical refinement of [`Vector`] or [`FixedShapeTensor`]. +/// /// # Invariants /// /// Every [`NormalizedArray`] has three slots. @@ -56,7 +60,7 @@ pub type NormalizedArray = Array; /// Both data children have the array's length and element ptype. A missing validity slot means /// either non-nullable or nullable-all-valid data, as determined by the parent dtype. /// -/// [`try_new`](Self::try_new) also enforces the invariants that make the split lossless: +/// [`try_new`](Self::try_new) also validates the direction-and-norm relationship: /// /// - Each normalized row has L2 norm `1.0` or `0.0`, within the tolerance for its precision and /// width. @@ -68,9 +72,10 @@ pub type NormalizedArray = Array; /// # Lossy normalized children /// /// [`new_unchecked`](Self::new_unchecked) permits an approximate normalized child, such as a -/// quantized direction. The stored norms remain authoritative. [`L2Norm`], [`InnerProduct`], and -/// [`CosineSimilarity`] therefore operate on the stored children and can differ slightly from -/// decoding and recomputing. +/// quantized direction. [`NormMode::Exact`](crate::scalar_fns::NormMode::Exact) measures that +/// physical direction. Only +/// [`NormMode::AssumeNormalized`](crate::scalar_fns::NormMode::AssumeNormalized) trusts its norm as +/// one, which can differ from decoding and recomputing and does not carry an error bound. /// /// [`Vector`]: crate::vector::Vector /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor @@ -78,6 +83,7 @@ pub type NormalizedArray = Array; /// [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm /// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct /// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity +/// [`NormMode`]: crate::scalar_fns::NormMode #[derive(Clone, Debug)] pub struct Normalized; diff --git a/vortex-tensor/src/encodings/normalized/compress.rs b/vortex-tensor/src/encodings/normalized/compress.rs index 914ee55b89c..a807ee4cf6f 100644 --- a/vortex-tensor/src/encodings/normalized/compress.rs +++ b/vortex-tensor/src/encodings/normalized/compress.rs @@ -41,6 +41,7 @@ use crate::encodings::normalized::NormalizedArraySlotsExt; use crate::encodings::normalized::NormalizedSlots; use crate::encodings::normalized::array::DATA_CHILDREN; use crate::matcher::AnyTensor; +use crate::scalar_fns::NormMode; use crate::scalar_fns::l2_norm::L2Norm; use crate::utils::extract_constant_flat_row; use crate::utils::extract_flat_elements; @@ -137,7 +138,9 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult VortexResult { - ScalarFnArray::try_new(CosineSimilarity.bind(EmptyOptions), vec![lhs, rhs]) + pub fn try_new(lhs: ArrayRef, rhs: ArrayRef, mode: NormMode) -> VortexResult { + ScalarFnArray::try_new(CosineSimilarity.bind(mode), vec![lhs, rhs]) } } impl ScalarFnVTable for CosineSimilarity { - type Options = EmptyOptions; + type Options = NormMode; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.cosine_similarity"); *ID } + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NormMode::deserialize(metadata) + } + fn arity(&self, _options: &Self::Options) -> Arity { Arity::Exact(2) } @@ -103,7 +119,7 @@ impl ScalarFnVTable for CosineSimilarity { fn execute( &self, - _options: &Self::Options, + options: &Self::Options, args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -122,13 +138,13 @@ impl ScalarFnVTable for CosineSimilarity { // 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); + return self.execute_both_normalized(lhs, rhs, *options, len, ctx); } NormalizedOrientation::One { normalized_array, plain, } => { - return self.execute_one_normalized(normalized_array, plain, len, ctx); + return self.execute_one_normalized(normalized_array, plain, *options, len, ctx); } NormalizedOrientation::Neither => {} } @@ -136,9 +152,9 @@ impl ScalarFnVTable for CosineSimilarity { // 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(lhs_ref.clone())?; - let norm_rhs_arr = L2Norm::try_new(rhs_ref.clone())?; + // Ordinary inputs carry no normalized claim, so both modes measure their physical norms. + let norm_lhs_arr = L2Norm::try_new(lhs_ref.clone(), NormMode::Exact)?; + let norm_rhs_arr = L2Norm::try_new(rhs_ref.clone(), NormMode::Exact)?; let dot_arr = InnerProduct::try_new(lhs_ref, rhs_ref)?; // Execute to get the inner product and norms of the arrays. We only fully decompress @@ -147,8 +163,8 @@ impl ScalarFnVTable for CosineSimilarity { 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. + // TODO(connor)[Tensor]: Replace this loop after binary numeric operations support + // zero-denominator division. The branch currently prevents a direct binary expression. match_each_float_ptype!(dot.ptype(), |T| { let dots = dot.as_slice::(); let norms_l = norm_l.as_slice::(); @@ -188,13 +204,40 @@ impl ScalarFnVTable for CosineSimilarity { } } +/// Metadata for a serialized [`CosineSimilarity`] array. +/// +/// `assume_normalized` is optional so metadata written before [`NormMode`] remains readable. +#[derive(Clone, prost::Message)] +struct CosineSimilarityMetadata { + /// The left input dtype needed to deserialize its child. + #[prost(message, optional, tag = "1")] + lhs_dtype: Option, + + /// The right input dtype needed to deserialize its child. + #[prost(message, optional, tag = "2")] + rhs_dtype: Option, + + /// Whether execution trusts normalized encoding evidence. + /// + /// `None` preserves the behavior of metadata written before [`NormMode`] was explicit. + #[prost(bool, optional, tag = "3")] + assume_normalized: Option, +} + impl ScalarFnArrayVTable for CosineSimilarity { fn serialize( &self, view: &ScalarFnArrayView, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some(BinaryTensorOpMetadata::encode_from_view(view)?)) + let array = view.as_::(); + let metadata = CosineSimilarityMetadata { + lhs_dtype: Some(array.child_at(0).dtype().try_into()?), + rhs_dtype: Some(array.child_at(1).dtype().try_into()?), + assume_normalized: Some(view.options.assumes_normalized()), + }; + + Ok(Some(metadata.encode_to_vec())) } fn deserialize( @@ -205,24 +248,39 @@ impl ScalarFnArrayVTable for CosineSimilarity { children: &dyn ArrayChildren, session: &VortexSession, ) -> VortexResult> { - let reconstructed = - BinaryTensorOpMetadata::decode_children(metadata, len, children, session)?; + let metadata = CosineSimilarityMetadata::decode(metadata) + .map_err(|error| vortex_err!("Failed to decode CosineSimilarity metadata: {error}"))?; + let lhs_dtype = metadata + .lhs_dtype + .as_ref() + .ok_or_else(|| vortex_err!("CosineSimilarity metadata missing lhs_dtype"))?; + let rhs_dtype = metadata + .rhs_dtype + .as_ref() + .ok_or_else(|| vortex_err!("CosineSimilarity metadata missing rhs_dtype"))?; + let lhs_dtype = DType::from_proto(lhs_dtype, session)?; + let rhs_dtype = DType::from_proto(rhs_dtype, session)?; + validate_binary_tensor_float_inputs(&lhs_dtype, &rhs_dtype)?; + + let lhs = children.get(0, &lhs_dtype, len)?; + let rhs = children.get(1, &rhs_dtype, len)?; + Ok(ScalarFnArrayParts { - options: EmptyOptions, - children: reconstructed, + options: NormMode::from_serialized(metadata.assume_normalized), + children: vec![lhs, rhs], }) } } impl CosineSimilarity { - /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so - /// `cosine_similarity = dot(n_l, n_r)`. + /// Computes cosine similarity from two [`Normalized`] direction-and-norm pairs. /// /// [`Normalized`]: crate::encodings::normalized::Normalized fn execute_both_normalized( &self, lhs_ref: &ArrayRef, rhs_ref: &ArrayRef, + mode: NormMode, len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -231,9 +289,17 @@ impl CosineSimilarity { 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 direction_norms = if mode.assumes_normalized() { + None + } else { + let lhs: PrimitiveArray = L2Norm::try_new(normalized_l.clone(), NormMode::Exact)? + .into_array() + .execute(ctx)?; + let rhs: PrimitiveArray = L2Norm::try_new(normalized_r.clone(), NormMode::Exact)? + .into_array() + .execute(ctx)?; + Some((lhs, rhs)) + }; let dot: PrimitiveArray = InnerProduct::try_new(normalized_l, normalized_r)? .into_array() .execute(ctx)?; @@ -244,12 +310,24 @@ impl CosineSimilarity { let dots = dot.as_slice::(); let norms_l = norms_l.as_slice::(); let norms_r = norms_r.as_slice::(); + let direction_norms = direction_norms + .as_ref() + .map(|(lhs, rhs)| (lhs.as_slice::(), rhs.as_slice::())); let buffer: Buffer = (0..len) .map(|i| { if norms_l[i] == T::zero() || norms_r[i] == T::zero() { + return T::zero(); + } + + let Some((direction_norms_l, direction_norms_r)) = direction_norms else { + return dots[i]; + }; + let denominator = direction_norms_l[i] * direction_norms_r[i]; + + if denominator == T::zero() { T::zero() } else { - dots[i] + dots[i] / denominator } }) .collect(); @@ -259,16 +337,17 @@ impl CosineSimilarity { }) } - /// One side is [`Normalized`]-encoded: treat the normalized child as authoritative, so - /// `cosine_similarity = dot(n, b) / ||b||`. + /// Computes cosine similarity when one side is [`Normalized`]-encoded. /// /// [`Normalized`]: crate::encodings::normalized::Normalized /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. + /// 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, + mode: NormMode, len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -276,26 +355,43 @@ impl CosineSimilarity { let (normalized, normalized_norms) = extract_normalized_children(normalized_ref); + let direction_norm = if mode.assumes_normalized() { + None + } else { + Some( + L2Norm::try_new(normalized.clone(), NormMode::Exact)? + .into_array() + .execute::(ctx)?, + ) + }; let dot_arr = InnerProduct::try_new(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(plain_ref.clone())?; + let norm_arr = L2Norm::try_new(plain_ref.clone(), NormMode::Exact)?; 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. + // TODO(connor)[Tensor]: Replace this loop after binary numeric operations support + // zero-denominator division. The branch currently prevents a direct binary expression. 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 direction_norms = direction_norm.as_ref().map(|norm| norm.as_slice::()); let buffer: Buffer = (0..len) .map(|i| { if normalized_norms[i] == T::zero() || plain_norms[i] == T::zero() { + return T::zero(); + } + + let denominator = + direction_norms.map_or(plain_norms[i], |norms| norms[i] * plain_norms[i]); + + if denominator == T::zero() { T::zero() } else { - dots[i] / plain_norms[i] + dots[i] / denominator } }) .collect(); @@ -309,6 +405,8 @@ impl CosineSimilarity { #[cfg(test)] mod tests { + use half::f16; + use prost::Message; use rstest::rstest; use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; @@ -316,11 +414,15 @@ mod tests { use vortex_array::VortexSessionExecute; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; + use vortex_array::matcher::Matcher; use vortex_array::validity::Validity; use vortex_error::VortexResult; + use super::CosineSimilarityMetadata; use crate::encodings::normalized::Normalized; + use crate::scalar_fns::NormMode; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::tests::SESSION; use crate::types::vector::Vector; @@ -332,7 +434,15 @@ mod tests { /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef) -> VortexResult> { - let result = CosineSimilarity::try_new(lhs, rhs)?; + eval_cosine_similarity_with_mode(lhs, rhs, NormMode::Exact) + } + + fn eval_cosine_similarity_with_mode( + lhs: ArrayRef, + rhs: ArrayRef, + mode: NormMode, + ) -> VortexResult> { + let result = CosineSimilarity::try_new(lhs, rhs, mode)?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; Ok(prim.as_slice::().to_vec()) @@ -496,7 +606,7 @@ mod tests { 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 result = CosineSimilarity::try_new(lhs, rhs)?; + let result = CosineSimilarity::try_new(lhs, rhs, NormMode::Exact)?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; @@ -544,6 +654,63 @@ mod tests { Ok(()) } + #[test] + fn mode_controls_near_unit_f16_shortcut() -> VortexResult<()> { + let element = f16::from_f32(1.009_765_6); + let direction = vector_array(2, &[element, f16::ZERO])?; + let norms = PrimitiveArray::from_iter([f16::ONE]).into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let vector = + Normalized::try_new(direction, norms, Validity::NonNullable, &mut ctx)?.into_array(); + + let exact: PrimitiveArray = + CosineSimilarity::try_new(vector.clone(), vector.clone(), NormMode::Exact)? + .into_array() + .execute(&mut ctx)?; + let assumed: PrimitiveArray = + CosineSimilarity::try_new(vector.clone(), vector, NormMode::AssumeNormalized)? + .into_array() + .execute(&mut ctx)?; + + assert_eq!(exact.as_slice::()[0], f16::ONE); + assert!(assumed.as_slice::()[0] > f16::ONE); + Ok(()) + } + + #[test] + fn exact_mode_measures_lossy_direction() -> VortexResult<()> { + let direction = tensor_array(&[2], &[0.8f64, 0.0])?; + let stored_norms = PrimitiveArray::from_iter([5.0f64]).into_array(); + // SAFETY: The children satisfy the structural requirements. This fixture intentionally + // supplies a lossy direction whose physical norm is 0.8. + let normalized = + unsafe { Normalized::new_unchecked(direction, stored_norms, Validity::NonNullable) } + .into_array(); + let plain = tensor_array(&[2], &[1.0f64, 0.0])?; + + assert_close( + &eval_cosine_similarity_with_mode(normalized.clone(), plain.clone(), NormMode::Exact)?, + &[1.0], + ); + assert_close( + &eval_cosine_similarity_with_mode(normalized, plain, NormMode::AssumeNormalized)?, + &[0.8], + ); + Ok(()) + } + + #[test] + fn mode_does_not_change_ordinary_inputs() -> VortexResult<()> { + let lhs = tensor_array(&[2], &[3.0f64, 4.0])?; + let rhs = tensor_array(&[2], &[4.0f64, 3.0])?; + + let exact = eval_cosine_similarity_with_mode(lhs.clone(), rhs.clone(), NormMode::Exact)?; + let assumed = eval_cosine_similarity_with_mode(lhs, rhs, NormMode::AssumeNormalized)?; + + assert_close(&exact, &assumed); + Ok(()) + } + #[test] fn one_side_normalized_lhs() -> VortexResult<()> { // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. @@ -579,7 +746,7 @@ mod tests { let validity = Validity::from_iter([true, false]); let rhs = Normalized::try_new(normalized_r, norms_r, validity, &mut ctx)?.into_array(); - let result = CosineSimilarity::try_new(lhs, rhs)?; + let result = CosineSimilarity::try_new(lhs, rhs, NormMode::Exact)?; let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; assert!(prim.is_valid(0, &mut ctx)?); @@ -752,26 +919,58 @@ mod tests { #[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(lhs.clone(), rhs.clone())?.into_array(); + let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); + for mode in [NormMode::Exact, NormMode::AssumeNormalized] { + let original = CosineSimilarity::try_new(lhs.clone(), rhs.clone(), mode)?.into_array(); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("CosineSimilarity serialize must produce metadata"); + let children = vec![lhs.clone(), rhs.clone()]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + let recovered_view = ExactScalarFn::::try_match(&recovered) + .expect("deserialized array must retain its scalar function"); + assert_eq!(*recovered_view.options, mode); + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + } + + Ok(()) + } + #[test] + fn legacy_metadata_uses_assume_normalized_mode() -> VortexResult<()> { + let lhs = cosine_vector_lhs(); + let rhs = cosine_vector_rhs(); + let original = + CosineSimilarity::try_new(lhs.clone(), rhs.clone(), NormMode::Exact)?.into_array(); let plugin = ScalarFnArrayPlugin::new(CosineSimilarity); let metadata = plugin .serialize(&original, &SESSION)? .expect("CosineSimilarity serialize must produce metadata"); + let mut metadata = CosineSimilarityMetadata::decode(metadata.as_slice())?; + metadata.assume_normalized = None; - let children = vec![lhs, rhs]; let recovered = plugin.deserialize( original.dtype(), original.len(), - &metadata, + &metadata.encode_to_vec(), &[], - &children, + &[lhs, rhs], &SESSION, )?; - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); + let recovered_view = ExactScalarFn::::try_match(&recovered) + .expect("deserialized array must retain its scalar function"); + assert_eq!(*recovered_view.options, NormMode::AssumeNormalized); Ok(()) } diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 8781cbc4abe..1e65226c049 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -29,7 +29,6 @@ use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; 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::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; @@ -45,6 +44,7 @@ use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; use crate::matcher::AnyTensor; +use crate::scalar_fns::NormMode; use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; use crate::utils::reattach_validity; @@ -57,10 +57,10 @@ use crate::utils::validate_tensor_float_input; /// The input must be a tensor-like extension array with a float element type. The output is a float /// column of the same float type. /// -/// When the input is [`Normalized`]-encoded, this operator treats the stored norms as -/// authoritative. For lossy normalized children, that means `L2Norm` intentionally reads the -/// stored norms instead of re-deriving them from fully decoded coordinates. That behavior is part -/// of the storage contract, not a separate lossy-compute mode. +/// [`NormMode::Exact`] measures the physical direction stored by a [`Normalized`] encoding and +/// multiplies that result by the stored norm. [`NormMode::AssumeNormalized`] instead trusts that +/// direction as unit length and returns the stored norm directly. The approximate mode does not +/// provide an error bound for unchecked or lossy encodings. /// /// [`Normalized`]: crate::encodings::normalized::Normalized #[derive(Clone)] @@ -71,21 +71,33 @@ impl L2Norm { /// /// # Errors /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches). - pub fn try_new(child: ArrayRef) -> VortexResult { - ScalarFnArray::try_new(L2Norm.bind(EmptyOptions), vec![child]) + /// Returns an error if the [`ScalarFnArray`] cannot be constructed because its input is not a + /// float tensor. + pub fn try_new(child: ArrayRef, mode: NormMode) -> VortexResult { + ScalarFnArray::try_new(L2Norm.bind(mode), vec![child]) } } impl ScalarFnVTable for L2Norm { - type Options = EmptyOptions; + type Options = NormMode; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("vortex.tensor.l2_norm"); *ID } + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + Ok(Some(options.serialize())) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + NormMode::deserialize(metadata) + } + fn arity(&self, _options: &Self::Options) -> Arity { Arity::Exact(1) } @@ -108,7 +120,7 @@ impl ScalarFnVTable for L2Norm { fn execute( &self, - _options: &Self::Options, + options: &Self::Options, args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -118,19 +130,39 @@ impl ScalarFnVTable for L2Norm { let ext = input_ref.dtype().as_extension(); let tensor_match = ext .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); + .vortex_expect("the input dtype was validated in `return_dtype`"); let tensor_flat_size = tensor_match.list_size() as usize; let element_ptype = tensor_match.element_ptype(); let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); - // Stored norms are authoritative. Reattach the parent validity because the child is - // non-nullable. if input_ref.is::() { - let (_, norms) = extract_normalized_children(&input_ref); - let norms = reattach_validity(norms, input_ref.validity()?)?; - vortex_ensure_eq!(norms.dtype(), &norm_dtype); - return Ok(norms); + let (direction, stored_norms) = extract_normalized_children(&input_ref); + if options.assumes_normalized() { + let norms = reattach_validity(stored_norms, input_ref.validity()?)?; + vortex_ensure_eq!(norms.dtype(), &norm_dtype); + + return Ok(norms); + } + + let direction_norms: PrimitiveArray = L2Norm::try_new(direction, NormMode::Exact)? + .into_array() + .execute(ctx)?; + let stored_norms: PrimitiveArray = stored_norms.execute(ctx)?; + let validity = input_ref.validity()?; + + return match_each_float_ptype!(element_ptype, |T| { + let direction_norms = direction_norms.as_slice::(); + let stored_norms = stored_norms.as_slice::(); + let norms: Buffer = direction_norms + .iter() + .zip(stored_norms) + .map(|(&direction_norm, &stored_norm)| direction_norm * stored_norm) + .collect(); + + // SAFETY: The buffer has one value per input row, matching `validity`. + Ok(unsafe { PrimitiveArray::new_unchecked(norms, validity) }.into_array()) + }); } // Optimize for the constant array case. @@ -144,10 +176,11 @@ impl ScalarFnVTable for L2Norm { let norm_scalar = match_each_float_ptype!(element_ptype, |T| { let values: Vec = elements .iter() - .map(|s| { - s.as_primitive() + .map(|scalar| { + scalar + .as_primitive() .as_::() - .vortex_expect("element was somehow not the correct float") + .vortex_expect("the input dtype established float elements") }) .collect(); let norm = l2_norm_row::(&values); @@ -194,13 +227,20 @@ impl ScalarFnVTable for L2Norm { } } -/// Metadata for a serialized [`L2Norm`] array: the single `input` child's [`DType`], which carries -/// the extension type (`FixedShapeTensor` vs `Vector`), dimension, and nullability that are not -/// recoverable from the parent's primitive-float output. +/// Metadata for a serialized [`L2Norm`] array. +/// +/// `assume_normalized` is optional so metadata written before [`NormMode`] remains readable. #[derive(Clone, prost::Message)] pub(super) struct L2NormMetadata { + /// The input dtype needed to deserialize the child. #[prost(message, optional, tag = "1")] input_dtype: Option, + + /// Whether execution trusts normalized encoding evidence. + /// + /// `None` preserves the behavior of metadata written before [`NormMode`] was explicit. + #[prost(bool, optional, tag = "2")] + assume_normalized: Option, } impl ScalarFnArrayVTable for L2Norm { @@ -211,7 +251,12 @@ impl ScalarFnArrayVTable for L2Norm { ) -> VortexResult>> { let scalar_fn_array = view.as_::(); let input_dtype = Some(scalar_fn_array.child_at(0).dtype().try_into()?); - Ok(Some(L2NormMetadata { input_dtype }.encode_to_vec())) + let metadata = L2NormMetadata { + input_dtype, + assume_normalized: Some(view.options.assumes_normalized()), + }; + + Ok(Some(metadata.encode_to_vec())) } fn deserialize( @@ -230,8 +275,9 @@ impl ScalarFnArrayVTable for L2Norm { .ok_or_else(|| vortex_err!("L2NormMetadata missing input_dtype"))?; let input_dtype = DType::from_proto(input_pb, session)?; let child = children.get(0, &input_dtype, len)?; + Ok(ScalarFnArrayParts { - options: EmptyOptions, + options: NormMode::from_serialized(metadata.assume_normalized), children: vec![child], }) } @@ -251,6 +297,7 @@ fn l2_norm_row(v: &[T]) -> T { #[cfg(test)] mod tests { + use prost::Message; use rstest::rstest; use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; @@ -261,16 +308,20 @@ mod tests { use vortex_array::arrays::ConstantArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDType; + use vortex_array::matcher::Matcher; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_error::VortexResult; + use super::L2NormMetadata; use crate::encodings::normalized::Normalized; + use crate::scalar_fns::NormMode; use crate::scalar_fns::l2_norm::L2Norm; use crate::tests::SESSION; use crate::types::vector::Vector; @@ -281,7 +332,11 @@ mod tests { /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. fn eval_l2_norm(input: ArrayRef) -> VortexResult> { - let result = L2Norm::try_new(input)?; + eval_l2_norm_with_mode(input, NormMode::Exact) + } + + fn eval_l2_norm_with_mode(input: ArrayRef, mode: NormMode) -> VortexResult> { + let result = L2Norm::try_new(input, mode)?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; Ok(prim.as_slice::().to_vec()) @@ -335,7 +390,7 @@ mod tests { let arr = tensor_array(&[2], &[3.0, 4.0, 0.0, 0.0])?; let arr = MaskedArray::try_new(arr, Validity::from_iter([true, false]))?.into_array(); - let result = L2Norm::try_new(arr)?; + let result = L2Norm::try_new(arr, NormMode::Exact)?; let mut ctx = SESSION.create_execution_ctx(); let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?; @@ -354,7 +409,7 @@ mod tests { fn constant_non_null_input_yields_constant_output() -> VortexResult<()> { let input = literal_vector_array(&[3.0f64, 4.0], 4); - let result = L2Norm::try_new(input)?.into_array(); + let result = L2Norm::try_new(input, NormMode::Exact)?.into_array(); let mut ctx = SESSION.create_execution_ctx(); let output = result.execute_until::(&mut ctx)?; @@ -384,7 +439,7 @@ mod tests { let null_scalar = Scalar::null(DType::Extension(ext_dtype)); let input = ConstantArray::new(null_scalar, 3).into_array(); - let result = L2Norm::try_new(input)?.into_array(); + let result = L2Norm::try_new(input, NormMode::Exact)?.into_array(); let mut ctx = SESSION.create_execution_ctx(); let output = result.execute_until::(&mut ctx)?; @@ -409,7 +464,7 @@ mod tests { let validity = Validity::from_iter([true, false]); let input = Normalized::try_new(normalized, norms, validity, &mut ctx)?.into_array(); - let result = L2Norm::try_new(input)?.into_array(); + let result = L2Norm::try_new(input, NormMode::Exact)?.into_array(); let prim: PrimitiveArray = result.execute(&mut ctx)?; assert_eq!( @@ -423,30 +478,81 @@ mod tests { Ok(()) } + #[test] + fn mode_controls_lossy_direction_norm() -> VortexResult<()> { + let direction = vector_array(2, &[0.8f64, 0.0])?; + let stored_norms = PrimitiveArray::from_iter([5.0f64]).into_array(); + // SAFETY: The children satisfy the structural requirements. This fixture intentionally + // supplies a lossy direction whose physical norm is 0.8. + let input = + unsafe { Normalized::new_unchecked(direction, stored_norms, Validity::NonNullable) } + .into_array(); + + assert_close( + &eval_l2_norm_with_mode(input.clone(), NormMode::Exact)?, + &[4.0], + ); + assert_close( + &eval_l2_norm_with_mode(input, NormMode::AssumeNormalized)?, + &[5.0], + ); + Ok(()) + } + #[rstest] #[case::fixed_shape_tensor(l2_norm_tensor_child())] #[case::vector(l2_norm_vector_child())] fn serde_round_trip(#[case] child: ArrayRef) -> VortexResult<()> { - let original = L2Norm::try_new(child.clone())?.into_array(); + let plugin = ScalarFnArrayPlugin::new(L2Norm); + for mode in [NormMode::Exact, NormMode::AssumeNormalized] { + let original = L2Norm::try_new(child.clone(), mode)?.into_array(); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("L2Norm serialize must produce metadata"); + let children = vec![child.clone()]; + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + let recovered_view = ExactScalarFn::::try_match(&recovered) + .expect("deserialized array must retain its scalar function"); + assert_eq!(*recovered_view.options, mode); + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + } + Ok(()) + } + + #[test] + fn legacy_metadata_uses_assume_normalized_mode() -> VortexResult<()> { + let child = l2_norm_vector_child(); + let original = L2Norm::try_new(child.clone(), NormMode::Exact)?.into_array(); let plugin = ScalarFnArrayPlugin::new(L2Norm); let metadata = plugin .serialize(&original, &SESSION)? .expect("L2Norm serialize must produce metadata"); + let mut metadata = L2NormMetadata::decode(metadata.as_slice())?; + metadata.assume_normalized = None; - let children = vec![child]; let recovered = plugin.deserialize( original.dtype(), original.len(), - &metadata, + &metadata.encode_to_vec(), &[], - &children, + &[child], &SESSION, )?; - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); + let recovered_view = ExactScalarFn::::try_match(&recovered) + .expect("deserialized array must retain its scalar function"); + assert_eq!(*recovered_view.options, NormMode::AssumeNormalized); Ok(()) } diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 68f10ca6b01..e87491eb3b7 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -2,7 +2,79 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Scalar function expressions defined on tensor and tensor-like extension types. +//! +//! Each child module owns one expression. [`NormMode`] defines whether norm-based expressions +//! measure physical coordinates or trust normalized encoding evidence. + +use std::fmt::Display; +use std::fmt::Formatter; + +use prost::Message; +use vortex_error::VortexResult; +use vortex_error::vortex_err; pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; + +/// Controls whether norm-based functions may trust [`Normalized`] encoding evidence. +/// +/// This policy belongs to the scalar function rather than the logical tensor dtype. It only +/// changes execution for [`Normalized`]-encoded inputs; ordinary tensors use their physical +/// coordinates in both modes. +/// +/// [`Normalized`]: crate::encodings::normalized::Normalized +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum NormMode { + /// Compute physical direction norms instead of assuming that they are exactly one. + Exact, + + /// Trust each `Normalized` direction as unit length and omit its norm computation. + /// + /// Checked arrays satisfy the encoding's documented tolerance. Unchecked or lossy arrays do + /// not carry an error bound, so this mode can produce values outside the mathematical range. + AssumeNormalized, +} + +impl NormMode { + pub(crate) fn assumes_normalized(self) -> bool { + matches!(self, Self::AssumeNormalized) + } + + pub(crate) fn from_serialized(assume_normalized: Option) -> Self { + match assume_normalized { + Some(false) => Self::Exact, + Some(true) | None => Self::AssumeNormalized, + } + } + + pub(crate) fn serialize(self) -> Vec { + NormModeMetadata { + assume_normalized: Some(self.assumes_normalized()), + } + .encode_to_vec() + } + + pub(crate) fn deserialize(metadata: &[u8]) -> VortexResult { + let metadata = NormModeMetadata::decode(metadata) + .map_err(|error| vortex_err!("Failed to decode NormMode metadata: {error}"))?; + + Ok(Self::from_serialized(metadata.assume_normalized)) + } +} + +impl Display for NormMode { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::Exact => f.write_str("exact"), + Self::AssumeNormalized => f.write_str("assume_normalized"), + } + } +} + +#[derive(Clone, prost::Message)] +struct NormModeMetadata { + /// Whether execution trusts normalized encoding evidence. + #[prost(bool, optional, tag = "1")] + assume_normalized: Option, +} diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 3e33fe20db9..553f3f98d94 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -65,10 +65,10 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { /// # Panics /// /// Panics if `array` is not [`Normalized`]-encoded. Callers reach this through -/// [`NormalizedOrientation::classify`], which has already matched on the encoding. +/// [`NormalizedOrientation`], which has already matched on the encoding. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -/// [`NormalizedOrientation::classify`]: crate::encodings::normalized::NormalizedOrientation::classify +/// [`NormalizedOrientation`]: crate::encodings::normalized::NormalizedOrientation pub fn extract_normalized_children(array: &ArrayRef) -> (ArrayRef, ArrayRef) { let normalized_array = array .as_opt::() @@ -234,12 +234,10 @@ pub fn extract_constant_flat_row( Ok(FlatRow { elems }) } -/// Metadata for a serialized binary tensor-op array (shared by [`InnerProduct`] and -/// [`CosineSimilarity`]). Both operands share the same extension dtype up to nullability -/// (enforced by their `return_dtype` checks), but their individual nullabilities are lost in the -/// parent's unioned output, so both are persisted. +/// Metadata for a serialized [`InnerProduct`] array. Both operands share the same extension dtype +/// up to nullability (enforced by its `return_dtype` checks), but their individual nullabilities +/// are lost in the parent's unioned output, so both are persisted. /// -/// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity /// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct #[derive(Clone, prost::Message)] pub(crate) struct BinaryTensorOpMetadata { diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index 1356a27168f..1d9248a02d8 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -18,11 +18,12 @@ //! use vortex_array::{ArrayRef, VortexSessionExecute}; //! use vortex_array::arrays::BoolArray; //! use vortex_session::VortexSession; +//! use vortex_tensor::scalar_fns::NormMode; //! use vortex_tensor::vector_search::build_similarity_search_tree; //! //! fn run(session: &VortexSession, data: ArrayRef, query: &[f32]) -> anyhow::Result<()> { //! let mut ctx = session.create_execution_ctx(); -//! let tree = build_similarity_search_tree(data, query, 0.8)?; +//! let tree = build_similarity_search_tree(data, query, 0.8, NormMode::Exact)?; //! let _matches: BoolArray = tree.execute(&mut ctx)?; //! Ok(()) //! } @@ -43,6 +44,7 @@ use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; +use crate::scalar_fns::NormMode; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::types::vector::Vector; @@ -63,7 +65,8 @@ use crate::types::vector::Vector; /// ``` /// /// The element type is inferred from `T` and must match the element type of `data`'s -/// [`Vector`] extension dtype. +/// [`Vector`] extension dtype. `mode` controls whether cosine similarity trusts normalized +/// encoding evidence. /// /// This function performs no execution; it is safe to call inside a benchmark setup closure. /// @@ -75,11 +78,12 @@ pub fn build_similarity_search_tree>( data: ArrayRef, query: &[T], threshold: T, + mode: NormMode, ) -> VortexResult { let num_rows = data.len(); let query_vec = Vector::constant_array(query, num_rows)?; - let cosine = CosineSimilarity::try_new(data, query_vec)?.into_array(); + let cosine = CosineSimilarity::try_new(data, query_vec, mode)?.into_array(); let threshold_scalar = Scalar::primitive(threshold, Nullability::NonNullable); let threshold_array = ConstantArray::new(threshold_scalar, num_rows).into_array(); @@ -95,6 +99,7 @@ mod tests { use vortex_error::VortexResult; use super::build_similarity_search_tree; + use crate::scalar_fns::NormMode; use crate::tests::SESSION; use crate::utils::test_helpers::vector_array; @@ -112,7 +117,7 @@ mod tests { )?; let query = [1.0f32, 0.0, 0.0]; - let tree = build_similarity_search_tree(data, &query, 0.5)?; + let tree = build_similarity_search_tree(data, &query, 0.5, NormMode::Exact)?; let mut ctx = SESSION.create_execution_ctx(); let result: BoolArray = tree.execute(&mut ctx)?;