diff --git a/vortex-tensor/src/encodings/normalized/array.rs b/vortex-tensor/src/encodings/normalized/array.rs index 35c8bdc4ef4..c0956143c32 100644 --- a/vortex-tensor/src/encodings/normalized/array.rs +++ b/vortex-tensor/src/encodings/normalized/array.rs @@ -1,20 +1,26 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use prost::Message; use vortex_array::Array; use vortex_array::ArrayId; use vortex_array::ArrayParts; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::EmptyArrayData; +use vortex_array::EmptyMetadata; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::array_slots; use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; @@ -24,9 +30,11 @@ use vortex_array::vtable::ValidityVTable; use vortex_array::vtable::child_to_validity; use vortex_array::vtable::validity_to_child; use vortex_array::vtable::with_empty_buffers; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -35,6 +43,10 @@ use crate::encodings::normalized::execute::denormalize; use crate::encodings::normalized::rules::RULES; use crate::encodings::normalized::validate::validate_normalized_children; use crate::encodings::normalized::validate::validate_normalized_rows; +use crate::types::unit_vector::AnyUnitVector; +use crate::types::unit_vector::UnitVector; +use crate::types::vector::AnyVector; +use crate::types::vector::Vector; use crate::utils::validate_tensor_float_input; /// A [`Normalized`]-encoded Vortex array. @@ -49,7 +61,10 @@ pub type NormalizedArray = Array; /// /// Every [`NormalizedArray`] has three slots. /// -/// - `normalized` is a non-nullable float tensor with the parent dtype's shape. +/// - For a [`Vector`] parent, `normalized` is a non-nullable +/// [`UnitVector`](crate::unit_vector::UnitVector). A documented lossy transform may instead +/// erase the refinement and store an ordinary non-nullable Vector. +/// - For a [`FixedShapeTensor`] parent, `normalized` is the corresponding non-nullable tensor. /// - `norms` is a non-nullable primitive column with the tensor element ptype. /// - `validity` optionally contains the parent validity as a non-nullable boolean column. /// @@ -67,10 +82,11 @@ 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. +/// Unchecked construction permits an approximate normalized child, such as a quantized direction. +/// If it can no longer prove the UnitVector tolerance, a Vector direction **must** erase that +/// refinement to ordinary Vector. The stored norms remain authoritative. [`L2Norm`], +/// [`InnerProduct`], and [`CosineSimilarity`] therefore operate on the stored children and can +/// differ slightly from decoding and recomputing. /// /// [`Vector`]: crate::vector::Vector /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor @@ -110,12 +126,38 @@ impl Normalized { /// Returns an error if the children are structurally incompatible, or if they violate any of /// the semantic invariants listed on [`Normalized`]. pub fn try_new( - normalized: ArrayRef, + mut normalized: ArrayRef, norms: ArrayRef, validity: Validity, ctx: &mut ExecutionCtx, ) -> VortexResult { - let array = Array::try_from_parts(normalized_parts(normalized, norms, validity))?; + let dtype = decoded_dtype_for_direction(normalized.dtype(), validity.nullability())?; + if dtype + .as_extension_opt() + .is_some_and(|dtype| dtype.is::()) + && !normalized + .dtype() + .as_extension_opt() + .is_some_and(|dtype| dtype.is::()) + { + let extension: ExtensionArray = normalized.execute(ctx)?; + // SAFETY: The semantic scan below validates the direction before it is returned. + normalized = unsafe { UnitVector::new_unchecked(extension.storage_array().clone())? }; + } + if dtype + .as_extension_opt() + .is_some_and(|dtype| dtype.is::()) + { + vortex_ensure!( + normalized + .dtype() + .as_extension_opt() + .is_some_and(|dtype| dtype.is::()), + "exact Normalized vector direction must be a UnitVector, got {}", + normalized.dtype(), + ); + } + let array = Array::try_from_parts(normalized_parts(dtype, normalized, norms, validity))?; validate_normalized_rows(array.normalized(), Some(array.norms()), ctx)?; Ok(array) @@ -138,17 +180,56 @@ impl Normalized { norms: ArrayRef, validity: Validity, ) -> NormalizedArray { - unsafe { Array::from_parts_unchecked(normalized_parts(normalized, norms, validity)) } + let dtype = decoded_dtype_for_direction(normalized.dtype(), validity.nullability()) + .vortex_expect("new_unchecked requires a valid tensor direction dtype"); + unsafe { Self::new_unchecked_with_dtype(dtype, normalized, norms, validity) } + } + + /// Builds a [`NormalizedArray`] with an explicit decoded dtype and without validation. + /// + /// # Safety + /// + /// The caller must uphold the invariants of [`new_unchecked`](Self::new_unchecked), and the + /// normalized child must be compatible with `dtype`. A vector parent accepts a UnitVector + /// child, or an ordinary Vector child when a documented lossy transform erased the + /// refinement. Violating these requirements can produce incorrect results but not memory + /// unsafety. + pub(crate) unsafe fn new_unchecked_with_dtype( + dtype: DType, + normalized: ArrayRef, + norms: ArrayRef, + validity: Validity, + ) -> NormalizedArray { + unsafe { Array::from_parts_unchecked(normalized_parts(dtype, normalized, norms, validity)) } } } +fn decoded_dtype_for_direction( + direction_dtype: &DType, + nullability: Nullability, +) -> VortexResult { + if direction_dtype + .as_extension_opt() + .is_some_and(|dtype| dtype.is::()) + { + let storage_dtype = direction_dtype + .as_extension() + .storage_dtype() + .with_nullability(nullability); + let dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?; + return Ok(DType::Extension(dtype.erased())); + } + + Ok(direction_dtype.with_nullability(nullability)) +} + fn normalized_parts( + dtype: DType, normalized: ArrayRef, norms: ArrayRef, validity: Validity, ) -> ArrayParts { let len = normalized.len(); - let dtype = normalized.dtype().with_nullability(validity.nullability()); let slots = NormalizedSlots { normalized, norms, @@ -159,6 +240,13 @@ fn normalized_parts( ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots) } +#[derive(Clone, prost::Message)] +struct NormalizedMetadata { + /// The direction dtype, absent from legacy metadata where it matched the parent dtype. + #[prost(message, optional, tag = "1")] + direction_dtype: Option, +} + impl VTable for Normalized { type TypedArrayData = EmptyArrayData; @@ -203,11 +291,11 @@ impl VTable for Normalized { } fn serialize( - _array: ArrayView<'_, Self>, + array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - // The parent dtype determines both child dtypes and the array nullability. - Ok(Some(vec![])) + let direction_dtype = Some(array.slots_view().normalized.dtype().try_into()?); + Ok(Some(NormalizedMetadata { direction_dtype }.encode_to_vec())) } fn deserialize( @@ -217,16 +305,17 @@ impl VTable for Normalized { metadata: &[u8], _buffers: &[BufferHandle], children: &dyn ArrayChildren, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { - vortex_ensure!( - metadata.is_empty(), - "NormalizedArray expects empty metadata, got {} bytes", - metadata.len(), - ); - let element_ptype = validate_tensor_float_input(dtype)?.element_ptype(); - let normalized_dtype = dtype.as_nonnullable(); + let metadata = NormalizedMetadata::decode(metadata) + .map_err(|error| vortex_err!("Failed to decode NormalizedMetadata: {error}"))?; + let normalized_dtype = metadata + .direction_dtype + .as_ref() + .map(|dtype| DType::from_proto(dtype, session)) + .transpose()? + .unwrap_or_else(|| dtype.as_nonnullable()); let norms_dtype = DType::Primitive(element_ptype, Nullability::NonNullable); let normalized = children.get(0, &normalized_dtype, len)?; @@ -249,7 +338,7 @@ impl VTable for Normalized { ), }; - Ok(normalized_parts(normalized, norms, validity)) + Ok(normalized_parts(dtype.clone(), normalized, norms, validity)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { diff --git a/vortex-tensor/src/encodings/normalized/compress.rs b/vortex-tensor/src/encodings/normalized/compress.rs index 643c6bd93bb..d1dfce61fbc 100644 --- a/vortex-tensor/src/encodings/normalized/compress.rs +++ b/vortex-tensor/src/encodings/normalized/compress.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use num_traits::Float; -use num_traits::Zero; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayVTable; @@ -13,19 +11,13 @@ use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Extension; use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::FixedSizeListArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::EmptyOptions; use vortex_array::validity::Validity; -use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; @@ -34,7 +26,7 @@ use vortex_compressor::scheme::EstimateVerdict; use vortex_compressor::scheme::Scheme; use vortex_compressor::scheme::SchemeExt; use vortex_compressor::stats::ArrayAndStats; -use vortex_error::VortexExpect; +use vortex_error::VortexError; use vortex_error::VortexResult; use crate::encodings::normalized::Normalized; @@ -43,10 +35,11 @@ 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::l2_norm::L2Norm; +use crate::scalar_fns::l2_normalize::normalize_children; +use crate::scalar_fns::l2_normalize::normalize_row_into; +use crate::scalar_fns::l2_normalize::normalized_output_dtype; +use crate::types::unit_vector::UnitVector; use crate::utils::extract_constant_flat_row; -use crate::utils::extract_flat_elements; -use crate::utils::validate_tensor_float_input; /// The compression scheme that rewrites a tensor-like column into the [`Normalized`] encoding. #[derive(Debug)] @@ -64,9 +57,11 @@ impl Scheme for NormalizedScheme { // `AlwaysUse` prevents later schemes from seeing a claimed array, so match only the float // tensor dtypes accepted by `compress`. - ext.ext_dtype() - .metadata_opt::() - .is_some_and(|tensor| tensor.element_ptype().is_float()) + !ext.ext_dtype().is::() + && ext + .ext_dtype() + .metadata_opt::() + .is_some_and(|tensor| tensor.element_ptype().is_float()) } fn produced_encodings(&self) -> Vec { @@ -94,6 +89,7 @@ impl Scheme for NormalizedScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { + let dtype = data.array().dtype().clone(); let normalized_array = normalize(data.array().clone(), exec_ctx)?; // Splitting magnitude out is only worth anything if the children then compress: the @@ -117,7 +113,10 @@ impl Scheme for NormalizedScheme { // SAFETY: Cascading preserves the split's child lengths and dtypes, and the validity is // carried over from the split unchanged. - Ok(unsafe { Normalized::new_unchecked(normalized, norms, validity) }.into_array()) + Ok( + unsafe { Normalized::new_unchecked_with_dtype(dtype, normalized, norms, validity) } + .into_array(), + ) } } @@ -130,81 +129,23 @@ impl Scheme for NormalizedScheme { /// /// Returns an error if `input` is not a float tensor column or if execution fails. pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let row_count = input.len(); - let tensor_match = validate_tensor_float_input(input.dtype())?; - let tensor_flat_size = tensor_match.list_size() as usize; - - // Normalize a constant's stored row without expanding it to the column length. - if let Some(wrapped) = try_build_constant_normalized(&input, row_count, ctx)? { - return Ok(wrapped); + if let Some(normalized) = try_build_constant_normalized(&input, ctx)? { + return Ok(normalized); } - let norms_array: ArrayRef = L2Norm - .try_new_array(row_count, EmptyOptions, [input.clone()])? - .execute(ctx)?; - - // Execute before reading validity. Reading validity from the lazy array would run `L2Norm` - // again when the values are requested. - let primitive_norms: PrimitiveArray = norms_array.execute(ctx)?; - - let validity = primitive_norms.validity()?; - - // Skip `fill_null` for non-nullable input because it still executes a cast. - let norms: PrimitiveArray = if validity.nullability().is_nullable() { - let element_dtype = - DType::Primitive(tensor_match.element_ptype(), Nullability::NonNullable); - - primitive_norms - .into_array() - .fill_null(Scalar::zero_value(&element_dtype))? - .execute(ctx)? - } else { - primitive_norms - }; + let dtype = input.dtype().clone(); + let (normalized, norms, validity) = normalize_children(input, ctx)?; - let input: ExtensionArray = input.execute(ctx)?; - let normalized_dtype = input.dtype().as_nonnullable(); - let flat = extract_flat_elements(input.storage_array(), tensor_flat_size, ctx)?; - - let normalized = match_each_float_ptype!(flat.ptype(), |T| { - let norm_values = norms.as_slice::(); - - let total_elements = row_count * tensor_flat_size; - let mut elements = BufferMut::::with_capacity(total_elements); - for i in 0..row_count { - let norm = norm_values[i]; - - // SAFETY: We allocated `row_count * tensor_flat_size` capacity and push exactly - // `tensor_flat_size` elements per row. - if norm.is_zero() { - unsafe { elements.push_n_unchecked(T::zero(), tensor_flat_size) }; - } else { - for &x in flat.row::(i) { - unsafe { elements.push_unchecked(x / norm) }; - } - } - } - - build_normalized( - normalized_dtype, - tensor_flat_size, - row_count, - elements.freeze(), - ) - })?; - - // SAFETY: This split creates non-nullable children with matching lengths and element ptypes. - // The captured validity describes the same input rows. - Ok(unsafe { Normalized::new_unchecked(normalized, norms.into_array(), validity) }) + // SAFETY: `normalize_children` constructs compatible non-nullable children and carries the + // input validity. + Ok(unsafe { Normalized::new_unchecked_with_dtype(dtype, normalized, norms, validity) }) } /// Normalizes a single constant row without expanding it to the column length. /// -/// Returns `Ok(None)` unless `input` has a non-null constant fixed-size-list storage scalar. A -/// matching input produces constant normalized and norms children. +/// Returns `Ok(None)` unless `input` has a non-null constant tensor row. pub(crate) fn try_build_constant_normalized( input: &ArrayRef, - len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { let Some(ext) = input.as_opt::() else { @@ -218,78 +159,59 @@ pub(crate) fn try_build_constant_normalized( return Ok(None); } - let tensor_match = input + if input .dtype() .as_extension() .metadata_opt::() - .vortex_expect("caller validated input has AnyTensor metadata"); - let list_size = tensor_match.list_size() as usize; + .is_none() + { + return Ok(None); + } // A non-null constant with a nullable dtype is all-valid. let validity = Validity::from(input.dtype().nullability()); - let normalized_ext_dtype = input.dtype().as_nonnullable().as_extension().clone(); + let normalized_ext_dtype = normalized_output_dtype(input.dtype())? + .as_extension() + .clone(); - // Materialize just the single stored row; this does not expand the constant to the full column - // length. let flat = extract_constant_flat_row(storage, ctx)?; - let (normalized_fsl_scalar, norms_scalar) = match_each_float_ptype!(flat.ptype(), |T| { + let scalars = match_each_float_ptype!(flat.ptype(), |T| { let row = flat.as_slice::(); - - let mut sum_sq = T::zero(); - for &x in row { - sum_sq += x * x; - } - let norm_t: T = sum_sq.sqrt(); - - // Zero-norm rows must be stored as all-zeros so the unit-norm-or-zero invariant holds. - // This mirrors the per-row logic in `normalize`. + let mut normalized = BufferMut::::with_capacity(row.len()); + // SAFETY: `normalized` reserves space for the entire row. + let norm = unsafe { normalize_row_into(row, &mut normalized)? }; let element_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable); - let children: Vec = if norm_t.is_zero() { - (0..list_size) - .map(|_| Scalar::zero_value(&element_dtype)) - .collect() - } else { - row.iter() - .map(|&v| Scalar::primitive(v / norm_t, Nullability::NonNullable)) - .collect() - }; + let children = normalized + .freeze() + .iter() + .copied() + .map(|value| Scalar::primitive(value, Nullability::NonNullable)) + .collect(); let fsl_scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); - let norms_scalar = Scalar::primitive(norm_t, Nullability::NonNullable); - (fsl_scalar, norms_scalar) + let norms_scalar = Scalar::primitive(norm, Nullability::NonNullable); + Ok::<_, VortexError>((fsl_scalar, norms_scalar)) }); + // This is also an optimization for infallible scalar functions. Value-dependent failure must + // fall back to their ordinary execution path. + let Ok((normalized_fsl_scalar, norms_scalar)) = scalars else { + return Ok(None); + }; + let len = input.len(); let normalized_storage = ConstantArray::new(normalized_fsl_scalar, len).into_array(); - let normalized = ExtensionArray::new(normalized_ext_dtype, normalized_storage).into_array(); + let normalized = if normalized_ext_dtype.is::() { + // SAFETY: The stored row was produced by `normalize_row_into` and the constant repeats it. + unsafe { UnitVector::new_unchecked(normalized_storage)? } + } else { + ExtensionArray::new(normalized_ext_dtype, normalized_storage).into_array() + }; let norms = ConstantArray::new(norms_scalar, len).into_array(); // SAFETY: Both constants use `len`, are non-nullable, and have the input element ptype. The // validity comes from the same input column. Ok(Some(unsafe { - Normalized::new_unchecked(normalized, norms, validity) + Normalized::new_unchecked_with_dtype(input.dtype().clone(), normalized, norms, validity) })) } - -/// Builds the non-nullable tensor-like extension array that becomes the `normalized` child. -fn build_normalized( - dtype: DType, - tensor_flat_size: usize, - row_count: usize, - elements: Buffer, -) -> VortexResult { - let list_size = - u32::try_from(tensor_flat_size).vortex_expect("tensor flat size must fit into `u32`"); - - // SAFETY: Tensor elements are always non-nullable, so the validity carries no length. - let elements = unsafe { PrimitiveArray::new_unchecked(elements, Validity::NonNullable) }; - - let storage = FixedSizeListArray::try_new( - elements.into_array(), - list_size, - Validity::NonNullable, - row_count, - )?; - - Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) -} diff --git a/vortex-tensor/src/encodings/normalized/execute.rs b/vortex-tensor/src/encodings/normalized/execute.rs index b960e28f1c4..0096516b670 100644 --- a/vortex-tensor/src/encodings/normalized/execute.rs +++ b/vortex-tensor/src/encodings/normalized/execute.rs @@ -83,7 +83,13 @@ fn denormalize_constant_norms( // A near-unit norm must still be multiplied, or `scalar_at` can disagree with bulk decoding. if norm_value == 1.0 { - return reattach_validity(normalized.clone(), validity); + let normalized: ExtensionArray = normalized.clone().execute(ctx)?; + let decoded = ExtensionArray::new( + dtype.as_nonnullable().as_extension().clone(), + normalized.storage_array().clone(), + ) + .into_array(); + return reattach_validity(decoded, validity); } let normalized: ExtensionArray = normalized.clone().execute(ctx)?; diff --git a/vortex-tensor/src/encodings/normalized/rules.rs b/vortex-tensor/src/encodings/normalized/rules.rs index 9a45963b37b..56942db5350 100644 --- a/vortex-tensor/src/encodings/normalized/rules.rs +++ b/vortex-tensor/src/encodings/normalized/rules.rs @@ -35,7 +35,8 @@ impl ArrayParentReduceRule for NormalizedSliceRule { // SAFETY: Slicing every slot with the same range preserves their dtypes and lengths. Ok(Some( unsafe { - Normalized::new_unchecked( + Normalized::new_unchecked_with_dtype( + array.dtype().clone(), array.normalized().slice(range.clone())?, array.norms().slice(range.clone())?, array.validity()?.slice(range.clone())?, @@ -63,7 +64,8 @@ impl ArrayParentReduceRule for NormalizedFilterRule { // SAFETY: Filtering every slot with the same mask preserves their dtypes and lengths. Ok(Some( unsafe { - Normalized::new_unchecked( + Normalized::new_unchecked_with_dtype( + array.dtype().clone(), array.normalized().filter(mask.clone())?, array.norms().filter(mask.clone())?, array.validity()?.filter(mask)?, diff --git a/vortex-tensor/src/encodings/normalized/tests.rs b/vortex-tensor/src/encodings/normalized/tests.rs index 0687e1ef750..05e719dc7cd 100644 --- a/vortex-tensor/src/encodings/normalized/tests.rs +++ b/vortex-tensor/src/encodings/normalized/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use half::f16; use rstest::rstest; use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; @@ -40,6 +41,7 @@ use crate::encodings::normalized::normalize; use crate::encodings::normalized::validate_normalized_rows; use crate::tests::SESSION; use crate::types::vector::Vector; +use crate::unit_vector::AnyUnitVector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::constant_tensor_array; use crate::utils::test_helpers::tensor_array; @@ -382,7 +384,7 @@ fn accepts_zero_vectors_paired_with_zero_norms() -> VortexResult<()> { #[test] fn validate_accepts_normalized_f16_rows() -> VortexResult<()> { - let input = vector_array(2, &[3.0f32, 4.0, 0.0, 0.0].map(half::f16::from_f32))?; + let input = vector_array(2, &[3.0f32, 4.0, 0.0, 0.0].map(f16::from_f32))?; let mut ctx = SESSION.create_execution_ctx(); let normalized_array = normalize(input, &mut ctx)?; @@ -393,9 +395,9 @@ fn validate_accepts_normalized_f16_rows() -> VortexResult<()> { fn checked_construction_accepts_dense_normalized_f16_row() -> VortexResult<()> { // Every coordinate is smaller than the unit-norm tolerance. Exact zero detection must not // misclassify this row as the zero vector. - let element = half::f16::from_f32(1.0 / 128.0_f32.sqrt()); + let element = f16::from_f32(1.0 / 128.0_f32.sqrt()); let normalized = vector_array(128, &[element; 128])?; - let norms = PrimitiveArray::from_iter([half::f16::from_f32(1.0)]).into_array(); + let norms = PrimitiveArray::from_iter([f16::from_f32(1.0)]).into_array(); let mut ctx = SESSION.create_execution_ctx(); Normalized::try_new(normalized, norms, Validity::NonNullable, &mut ctx)?; @@ -465,6 +467,42 @@ fn normalize_keeps_constant_input_children_constant() -> VortexResult<()> { Ok(()) } +#[test] +fn normalize_constant_f16_uses_wide_accumulation() -> VortexResult<()> { + let values = vec![f16::ONE; 4096]; + let input = Vector::constant_array(&values, 8)?; + let mut ctx = SESSION.create_execution_ctx(); + let normalized = normalize(input, &mut ctx)?; + + validate_normalized_rows(normalized.normalized(), None, &mut ctx) +} + +#[test] +fn normalize_vector_uses_a_unit_vector_direction() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let normalized = normalize(vector_array(2, &[3.0f64, 4.0])?, &mut ctx)?; + + assert!( + normalized + .normalized() + .dtype() + .as_extension() + .is::() + ); + assert!(normalized.dtype().as_extension().is::()); + Ok(()) +} + +#[test] +fn scheme_skips_unit_vectors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let normalized = normalize(vector_array(2, &[3.0f64, 4.0])?, &mut ctx)?; + let unit: Canonical = normalized.normalized().clone().execute(&mut ctx)?; + + assert!(!NormalizedScheme.matches(&unit)); + Ok(()) +} + #[test] fn normalize_zeroes_rows_with_zero_norms() -> VortexResult<()> { let input = vector_array(2, &[0.0, 0.0, 3.0, 4.0])?; @@ -672,7 +710,7 @@ fn serde_round_trip(#[case] input: ArrayRef) -> VortexResult<()> { } #[test] -fn serialization_carries_no_metadata() -> VortexResult<()> { +fn serialization_carries_the_direction_dtype() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let nullable = normalize(nullable_vector_input()?, &mut ctx)?.into_array(); let non_nullable = normalize(vector_array(2, &[3.0, 4.0, 1.0, 0.0])?, &mut ctx)?.into_array(); @@ -681,7 +719,10 @@ fn serialization_carries_no_metadata() -> VortexResult<()> { let bytes = SESSION .array_serialize(array)? .expect("Normalized must serialize"); - assert!(bytes.is_empty(), "Normalized must not serialize metadata"); + assert!( + !bytes.is_empty(), + "Normalized must serialize its direction dtype" + ); } assert_eq!(nullable.nchildren(), NormalizedSlots::COUNT); @@ -690,6 +731,21 @@ fn serialization_carries_no_metadata() -> VortexResult<()> { Ok(()) } +#[test] +fn legacy_empty_metadata_uses_the_parent_dtype_for_the_direction() -> VortexResult<()> { + let direction = vector_array(2, &[0.6f64, 0.8])?; + let dtype = direction.dtype().clone(); + let norms = PrimitiveArray::from_iter([5.0f64]).into_array(); + let children = vec![direction, norms]; + + let recovered = + ArrayPlugin::deserialize(&Normalized, &dtype, 1, &[], &[], &children, &SESSION)?; + let recovered = recovered.as_::(); + + assert!(recovered.normalized().dtype().as_extension().is::()); + Ok(()) +} + #[test] fn serde_round_trip_of_a_nullable_column_with_no_null_rows() -> VortexResult<()> { // AllValid omits the validity child, so deserialization must recover nullability from the @@ -724,18 +780,58 @@ fn serde_round_trip_of_a_nullable_column_with_no_null_rows() -> VortexResult<()> Ok(()) } +#[test] +fn lossy_vector_direction_dtype_round_trips() -> VortexResult<()> { + let direction = vector_array(2, &[0.61f64, 0.79])?; + let dtype = direction.dtype().clone(); + let norms = PrimitiveArray::from_iter([5.0f64]).into_array(); + // SAFETY: A plain Vector direction is the documented escape hatch for a lossy transform. The + // child dtype, length, ptype, and parent dtype are compatible. + let original = unsafe { + Normalized::new_unchecked_with_dtype(dtype, direction, norms, Validity::NonNullable) + } + .into_array(); + let metadata = SESSION + .array_serialize(&original)? + .expect("Normalized must serialize"); + let recovered = ArrayPlugin::deserialize( + &Normalized, + original.dtype(), + original.len(), + &metadata, + &[], + &original.children(), + &SESSION, + )?; + let recovered = recovered.as_::(); + + assert!(recovered.normalized().dtype().as_extension().is::()); + assert_eq!(recovered.dtype(), original.dtype()); + Ok(()) +} + #[test] fn deserialize_rejects_validity_child_for_non_nullable_dtype() -> VortexResult<()> { let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; let dtype = normalized.dtype().clone(); + let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + // SAFETY: The children are compatible and satisfy the exact normalized invariants. + let original = unsafe { + Normalized::new_unchecked(normalized.clone(), norms.clone(), Validity::NonNullable) + } + .into_array(); + let metadata = SESSION + .array_serialize(&original)? + .expect("Normalized must serialize"); let children = vec![ normalized, - PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(), + norms, BoolArray::from_iter([true, false]).into_array(), ]; - let error = ArrayPlugin::deserialize(&Normalized, &dtype, 2, &[], &[], &children, &SESSION) - .unwrap_err(); + let error = + ArrayPlugin::deserialize(&Normalized, &dtype, 2, &metadata, &[], &children, &SESSION) + .unwrap_err(); assert!( error diff --git a/vortex-tensor/src/encodings/normalized/validate.rs b/vortex-tensor/src/encodings/normalized/validate.rs index d82722dac67..bac73fb26a1 100644 --- a/vortex-tensor/src/encodings/normalized/validate.rs +++ b/vortex-tensor/src/encodings/normalized/validate.rs @@ -17,6 +17,9 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; +use crate::types::vector::AnyVector; +use crate::types::vector::Vector; +use crate::unit_vector::AnyUnitVector; use crate::utils::extract_flat_elements; use crate::utils::unit_norm_tolerance; use crate::utils::validate_tensor_float_input; @@ -34,7 +37,7 @@ pub(super) fn validate_normalized_children( vortex_ensure_eq!( normalized.len(), len, - "Normalized normalized child must have the array length ({len}), got {}", + "Normalized direction must have the array length ({len}), got {}", normalized.len(), ); vortex_ensure_eq!( @@ -44,13 +47,27 @@ pub(super) fn validate_normalized_children( norms.len(), ); - let tensor_match = validate_tensor_float_input(normalized.dtype())?; - let element_ptype = tensor_match.element_ptype(); + let parent_match = validate_tensor_float_input(dtype)?; + let child_match = validate_tensor_float_input(normalized.dtype())?; + let element_ptype = parent_match.element_ptype(); - vortex_ensure_eq!( - *normalized.dtype(), - dtype.as_nonnullable(), - "Normalized normalized child must be the non-nullable array dtype ({}), got {}", + vortex_ensure!( + !dtype.as_extension().is::(), + "Normalized parent dtype must be Vector or FixedShapeTensor, got {dtype}", + ); + let parent_is_vector = dtype.as_extension().is::(); + let child_is_vector = normalized.dtype().as_extension().is::(); + let compatible_child = if parent_is_vector { + child_is_vector + && parent_match.element_ptype() == child_match.element_ptype() + && parent_match.list_size() == child_match.list_size() + && !normalized.dtype().is_nullable() + } else { + *normalized.dtype() == dtype.as_nonnullable() + }; + vortex_ensure!( + compatible_child, + "Normalized direction must be compatible with the non-nullable parent dtype {}, got {}", dtype.as_nonnullable(), normalized.dtype(), ); @@ -89,17 +106,14 @@ pub(super) fn validate_normalized_children( /// Validates the semantic invariants documented by [`Normalized`]. /// -/// The zero relationship is checked in both directions. Otherwise, a zero row with a nonzero -/// stored norm would decode differently from [`L2Norm`]. This `O(len * list_size)` scan includes -/// rows that the parent might mark null. +/// When `norms` is present, the zero relationship is checked in both directions for every row, +/// including parent-null rows. Without norms, null direction rows are skipped. /// /// # Errors /// -/// Returns an error if either child has an incompatible dtype or length, or if a row violates the -/// semantic invariants. +/// Returns an error if a child is incompatible or a row violates the semantic invariants. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -/// [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm pub fn validate_normalized_rows( normalized: &ArrayRef, norms: Option<&ArrayRef>, @@ -139,6 +153,16 @@ pub fn validate_normalized_rows( } let normalized: ExtensionArray = normalized.clone().execute(ctx)?; + let valid_rows = if norms.is_none() { + Some( + normalized + .as_ref() + .validity()? + .execute_mask(row_count, ctx)?, + ) + } else { + None + }; let flat = extract_flat_elements(normalized.storage_array(), tensor_flat_size, ctx)?; let norms = norms .map(|norms| norms.clone().execute::(ctx)) @@ -148,6 +172,13 @@ pub fn validate_normalized_rows( let stored_norms = norms.as_ref().map(|norms| norms.as_slice::()); for i in 0..row_count { + if valid_rows + .as_ref() + .is_some_and(|validity| !validity.value(i)) + { + continue; + } + let (row_norm_sq, is_zero_row) = flat.row::(i) .iter() @@ -159,9 +190,9 @@ pub fn validate_normalized_rows( let row_norm = row_norm_sq.sqrt(); vortex_ensure!( - row_norm.is_zero() || (row_norm - 1.0).abs() <= tolerance, - "Normalized normalized child must have L2 norm 1.0 or 0.0, but row {i} has \ - {row_norm:.6}", + is_zero_row || (row_norm - 1.0).abs() <= tolerance, + "Normalized direction must have L2 norm 1.0 or be exactly zero, but row {i} has \ + norm {row_norm:.6}", ); if let Some(stored_norms) = stored_norms { @@ -173,9 +204,8 @@ pub fn validate_normalized_rows( vortex_ensure!( is_zero_row == stored_norm_f64.is_zero(), - "Normalized normalized child must be all zeros exactly when its stored norm is \ - 0.0, but row {i} pairs a {} normalized row with a stored norm of \ - {stored_norm_f64:.6}", + "Normalized direction must be exactly zero if and only if its stored norm is \ + 0.0, but row {i} pairs a {} direction with norm {stored_norm_f64:.6}", if is_zero_row { "zero" } else { "nonzero" }, ); } diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index fe56827c15b..ca40ac4cb24 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -23,7 +23,9 @@ use crate::encodings::normalized::Normalized; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; +use crate::scalar_fns::l2_normalize::L2Normalize; use crate::types::fixed_shape_tensor::FixedShapeTensor; +use crate::types::unit_vector::UnitVector; use crate::types::vector::Vector; pub mod matcher; @@ -32,7 +34,9 @@ pub mod scalar_fns; mod types; pub use types::fixed_shape_tensor; +pub use types::unit_vector; pub use types::vector; +pub use utils::unit_norm_tolerance; pub mod encodings; @@ -41,10 +45,10 @@ pub mod vector_search; mod utils; /// Environment variable that gates registration of the tensor scalar-fn array plugins (the array -/// encodings that let [`CosineSimilarity`], [`InnerProduct`], and [`L2Norm`] persist in a Vortex -/// file). When unset, only the scalar functions themselves are registered; readers of files -/// containing serialized tensor scalar-fn arrays will fail to deserialize. Opt-in by setting the -/// variable to any non-empty value. +/// encodings that let [`CosineSimilarity`], [`InnerProduct`], [`L2Norm`], and [`L2Normalize`] +/// persist in a Vortex file). When unset, only the scalar functions themselves are registered; +/// readers of files containing serialized tensor scalar-fn arrays will fail to deserialize. +/// Opt-in by setting the variable to any non-empty value. /// /// This does **not** gate [`Normalized`]. That is a real array encoding rather than a persisted /// scalar function, and the compressor can emit it, so it always registers. @@ -53,11 +57,14 @@ pub const SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str = "VX_SCALAR_FN_ARRAY_TENSOR_P /// Initialize the Vortex tensor library with a Vortex session. pub fn initialize(session: &VortexSession) { session.dtypes().register(Vector); + session.dtypes().register(UnitVector); session.dtypes().register(FixedShapeTensor); let arrow_session = session.arrow(); arrow_session.register_exporter(Arc::new(Vector)); arrow_session.register_importer(Arc::new(Vector)); + arrow_session.register_exporter(Arc::new(UnitVector)); + arrow_session.register_importer(Arc::new(UnitVector)); session.arrays().register(Normalized); @@ -66,6 +73,7 @@ pub fn initialize(session: &VortexSession) { session_fns.register(CosineSimilarity); session_fns.register(InnerProduct); session_fns.register(L2Norm); + session_fns.register(L2Normalize); // Registering the scalar-fn array plugins lets the tensor scalar fns be serialized as array // encodings inside Vortex files. Gate this on an env var so applications that do not intend @@ -77,6 +85,7 @@ pub fn initialize(session: &VortexSession) { session_arrays.register(ScalarFnArrayPlugin::new(CosineSimilarity)); session_arrays.register(ScalarFnArrayPlugin::new(InnerProduct)); session_arrays.register(ScalarFnArrayPlugin::new(L2Norm)); + session_arrays.register(ScalarFnArrayPlugin::new(L2Normalize)); } } diff --git a/vortex-tensor/src/matcher.rs b/vortex-tensor/src/matcher.rs index 10aa2581d03..13e02336c1f 100644 --- a/vortex-tensor/src/matcher.rs +++ b/vortex-tensor/src/matcher.rs @@ -14,10 +14,8 @@ use crate::types::vector::VectorMatcherMetadata; /// Matcher for any tensor-like extension type. /// -/// Currently the different kinds of tensors that are available are: -/// -/// - `FixedShapeTensor` -/// - `Vector` +/// Matches [`FixedShapeTensor`](crate::fixed_shape_tensor::FixedShapeTensor), +/// [`Vector`](crate::vector::Vector), and [`UnitVector`](crate::unit_vector::UnitVector). pub struct AnyTensor; /// The matched variant of a tensor-like extension type. @@ -26,9 +24,7 @@ pub enum TensorMatch<'a> { /// A [`FixedShapeTensor`](crate::fixed_shape_tensor::FixedShapeTensor) extension type. FixedShapeTensor(FixedShapeTensorMatcherMetadata<'a>), - /// A [`Vector`](crate::vector::Vector) extension type. - /// - /// Note that we store an owned type here wrapping (copyable) data from the dtype. + /// A [`Vector`](crate::vector::Vector) or [`UnitVector`](crate::unit_vector::UnitVector). Vector(VectorMatcherMetadata), } @@ -58,7 +54,6 @@ impl Matcher for AnyTensor { return Some(TensorMatch::FixedShapeTensor(metadata)); } - // Special logic for vectors to get convenience metadata (instead of `EmptyMetadata`). if let Some(metadata) = ext_dtype.metadata_opt::() { return Some(TensorMatch::Vector(metadata)); } diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index ef6eed69e94..0bb7958be87 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -34,6 +34,7 @@ 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::unit_vector::AnyUnitVector; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_normalized_children; use crate::utils::validate_binary_tensor_float_inputs; @@ -44,16 +45,22 @@ use crate::utils::validate_binary_tensor_float_inputs; /// The shape and permutation do not affect the result because cosine similarity only depends on the /// element values, not their logical arrangement. /// -/// 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. +/// Fixed-shape tensor inputs must have the same dtype, ignoring top-level nullability. Vector +/// inputs may mix [`Vector`] and [`UnitVector`] when their element ptype and dimensions match. The +/// output is a float column with that element ptype. /// /// When either input is [`Normalized`]-encoded, this operator treats the stored norms and /// normalized children as authoritative. For lossy normalized children, that means the optimized /// read-through path may intentionally differ slightly from decoding both sides to dense /// coordinates and recomputing cosine from scratch. /// +/// A [`UnitVector`] norm is treated as one, while [`L2Norm`] still measures its physical +/// coordinates. With tolerance `t`, omitting one norm adds at most approximately `t` absolute +/// error; omitting both adds at most `2t + t²`. +/// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector +/// [`UnitVector`]: crate::unit_vector::UnitVector /// [`Normalized`]: crate::encodings::normalized::Normalized #[derive(Clone)] pub struct CosineSimilarity; @@ -117,10 +124,10 @@ impl ScalarFnVTable for CosineSimilarity { let len = args.row_count(); // Normalize extension-level constants so the encoded fast path can use them. - if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, len, ctx)? { + if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, ctx)? { lhs_ref = normalized_array.into_array(); } - if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, len, ctx)? { + if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, ctx)? { rhs_ref = normalized_array.into_array(); } @@ -138,6 +145,23 @@ impl ScalarFnVTable for CosineSimilarity { NormalizedOrientation::Neither => {} } + let lhs_is_unit = lhs_ref.dtype().as_extension().is::(); + let rhs_is_unit = rhs_ref.dtype().as_extension().is::(); + match (lhs_is_unit, rhs_is_unit) { + (true, true) => { + return InnerProduct::try_new_array(lhs_ref, rhs_ref)? + .into_array() + .execute(ctx); + } + (true, false) => { + return self.execute_one_unit(lhs_ref, rhs_ref, len, ctx); + } + (false, true) => { + return self.execute_one_unit(rhs_ref, lhs_ref, len, ctx); + } + (false, false) => {} + } + // Compute combined validity. let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; @@ -220,6 +244,37 @@ impl ScalarFnArrayVTable for CosineSimilarity { } impl CosineSimilarity { + fn execute_one_unit( + &self, + unit: ArrayRef, + plain: ArrayRef, + len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let validity = unit.validity()?.and(plain.validity()?)?; + let dot: PrimitiveArray = InnerProduct::try_new_array(unit, plain.clone())? + .into_array() + .execute(ctx)?; + let plain_norm: PrimitiveArray = L2Norm::try_new_array(plain)?.into_array().execute(ctx)?; + + match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let norms = plain_norm.as_slice::(); + let buffer: Buffer = (0..len) + .map(|i| { + if norms[i] == T::zero() { + T::zero() + } else { + dots[i] / norms[i] + } + }) + .collect(); + + // SAFETY: The buffer length equals `len`, which matches the source validity length. + Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array()) + }) + } + /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so /// `cosine_similarity = dot(n_l, n_r)`. /// @@ -286,6 +341,26 @@ impl CosineSimilarity { let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; + if plain_ref.dtype().as_extension().is::() { + return match_each_float_ptype!(dot.ptype(), |T| { + let dots = dot.as_slice::(); + let normalized_norms = normalized_norms.as_slice::(); + let buffer: Buffer = (0..len) + .map(|i| { + if normalized_norms[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()) + }); + } + let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; @@ -319,9 +394,11 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; + use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::validity::Validity; use vortex_error::VortexResult; @@ -330,6 +407,7 @@ mod tests { use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::tests::SESSION; use crate::types::vector::Vector; + use crate::unit_vector::UnitVector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::constant_tensor_array; use crate::utils::test_helpers::normalized_array; @@ -345,6 +423,36 @@ mod tests { Ok(prim.as_slice::().to_vec()) } + fn checked_unit_vector(values: &[f32]) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + let dimensions = u32::try_from(values.len())?; + let vector: ExtensionArray = vector_array(dimensions, values)?.execute(&mut ctx)?; + UnitVector::try_new_unit_vector_array(vector.storage_array().clone(), &mut ctx) + } + + #[test] + fn both_unit_vectors_omit_norms() -> VortexResult<()> { + let lhs = checked_unit_vector(&[0.6000005, 0.8])?; + let rhs = checked_unit_vector(&[1.0, 0.0])?; + let result = CosineSimilarity::try_new_array(lhs, rhs)? + .into_array() + .execute::(&mut SESSION.create_execution_ctx())?; + + assert_eq!(result.as_slice::()[0], 0.6000005); + Ok(()) + } + + #[test] + fn one_unit_vector_omits_its_norm() -> VortexResult<()> { + let unit = checked_unit_vector(&[0.6000005, 0.8])?; + let plain = vector_array(2, &[3.0f32, 4.0])?; + let result = CosineSimilarity::try_new_array(unit, plain)? + .into_array() + .execute::(&mut SESSION.create_execution_ctx())?; + assert!((result.as_slice::()[0] - 1.0000004).abs() < 1e-6); + Ok(()) + } + #[test] fn unit_vectors_1d() -> VortexResult<()> { let lhs = tensor_array( diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 74eb184045f..bbc0e828ddb 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -47,11 +47,13 @@ use crate::utils::validate_binary_tensor_float_inputs; /// this is the standard dot product; for higher-rank ([`FixedShapeTensor`]) arrays this is the /// Frobenius inner product. /// -/// 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. +/// Fixed-shape tensor inputs must have the same dtype, ignoring top-level nullability. Vector +/// inputs may mix [`Vector`] and [`UnitVector`] when their element ptype and dimensions match. The +/// output is a float column with that element ptype. /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector +/// [`UnitVector`]: crate::unit_vector::UnitVector #[derive(Clone)] pub struct InnerProduct; diff --git a/vortex-tensor/src/scalar_fns/l2_normalize.rs b/vortex-tensor/src/scalar_fns/l2_normalize.rs new file mode 100644 index 00000000000..5cb942f66ad --- /dev/null +++ b/vortex-tensor/src/scalar_fns/l2_normalize.rs @@ -0,0 +1,579 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splits float tensor rows into L2-normalized directions and norms. +//! +//! [`L2Normalize`] returns a struct with non-nullable `normalized` and `norms` fields. The struct +//! carries the input nullability, and null rows use zero-valued child payloads. An exact-zero row +//! produces a zero direction and zero norm. +//! +//! A [`Vector`](crate::vector::Vector) direction is refined to +//! [`UnitVector`]. A +//! [`FixedShapeTensor`](crate::fixed_shape_tensor::FixedShapeTensor) direction keeps its ordinary +//! tensor dtype because Vortex does not define a unit-tensor refinement. Normalization accumulates +//! squares in f64 and uses scaled accumulation when the ordinary sum overflows or underflows. + +use num_traits::ToPrimitive; +use num_traits::Zero; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; +use vortex_array::arrays::scalar_fn::ScalarFnArrayView; +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::FieldNames; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::proto::dtype as pb; +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::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::types::unit_vector::UnitVector; +use crate::types::vector::AnyVector; +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; + +/// Splits each tensor-like row into its L2-normalized direction and physical norm. +/// +/// Vector inputs produce a [`UnitVector`] direction. Fixed-shape tensor inputs retain their +/// ordinary tensor dtype because Vortex does not define a unit-tensor refinement. The two fields +/// are non-nullable, and input nullability is carried by the returned struct. +#[derive(Clone)] +pub struct L2Normalize; + +impl L2Normalize { + /// Creates an [`L2Normalize`] scalar function instance. + pub fn new() -> TypedScalarFnInstance { + TypedScalarFnInstance::new(L2Normalize, EmptyOptions) + } + + /// Constructs a lazy [`ScalarFnArray`] that normalizes `child`. + /// + /// # Errors + /// + /// Returns an error if `child` is not a float tensor-like array or the scalar-function array + /// cannot be constructed. + pub fn try_new_array(child: ArrayRef) -> VortexResult { + ScalarFnArray::try_new(L2Normalize::new().erased(), vec![child]) + } +} + +impl ScalarFnVTable for L2Normalize { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.tensor.l2_normalize"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("input"), + _ => unreachable!("L2Normalize must have exactly one child"), + } + } + + fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { + let input_dtype = &arg_dtypes[0]; + let tensor_match = validate_tensor_float_input(input_dtype)?; + let normalized_dtype = normalized_output_dtype(input_dtype)?; + let norms_dtype = DType::Primitive(tensor_match.element_ptype(), Nullability::NonNullable); + + Ok(DType::Struct( + StructFields::new( + FieldNames::from(["normalized", "norms"]), + vec![normalized_dtype, norms_dtype], + ), + input_dtype.nullability(), + )) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + normalize_array(input, ctx).map(|array| array.into_array()) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + true + } +} + +#[derive(Clone, prost::Message)] +struct L2NormalizeMetadata { + /// The child dtype required before deserializing the child array. + #[prost(message, optional, tag = "1")] + input_dtype: Option, +} + +impl ScalarFnArrayVTable for L2Normalize { + fn serialize( + &self, + view: &ScalarFnArrayView, + _session: &VortexSession, + ) -> VortexResult>> { + let array = view.as_::(); + let input_dtype = Some(array.child_at(0).dtype().try_into()?); + + Ok(Some(L2NormalizeMetadata { input_dtype }.encode_to_vec())) + } + + fn deserialize( + &self, + _dtype: &DType, + len: usize, + metadata: &[u8], + children: &dyn ArrayChildren, + session: &VortexSession, + ) -> VortexResult> { + let metadata = L2NormalizeMetadata::decode(metadata) + .map_err(|error| vortex_err!("Failed to decode L2NormalizeMetadata: {error}"))?; + let input_dtype = metadata + .input_dtype + .as_ref() + .ok_or_else(|| vortex_err!("L2NormalizeMetadata missing input_dtype"))?; + let input_dtype = DType::from_proto(input_dtype, session)?; + let child = children.get(0, &input_dtype, len)?; + + Ok(ScalarFnArrayParts { + options: EmptyOptions, + children: vec![child], + }) + } +} + +pub(crate) fn normalize_array( + input: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let row_count = input.len(); + let (normalized, norms, validity) = normalize_children(input, ctx)?; + + StructArray::try_new( + FieldNames::from(["normalized", "norms"]), + vec![normalized, norms], + row_count, + validity, + ) +} + +pub(crate) fn normalize_children( + input: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ArrayRef, ArrayRef, Validity)> { + let row_count = input.len(); + let tensor_match = validate_tensor_float_input(input.dtype())?; + let tensor_flat_size = tensor_match.list_size() as usize; + let output_dtype = normalized_output_dtype(input.dtype())?; + + let input: ExtensionArray = input.execute(ctx)?; + let validity = input.as_ref().validity()?; + let valid_rows = validity + .nullability() + .is_nullable() + .then(|| validity.execute_mask(row_count, ctx)) + .transpose()?; + let flat = extract_flat_elements(input.storage_array(), tensor_flat_size, ctx)?; + let (normalized, norms) = match_each_float_ptype!(flat.ptype(), |T| { + let mut elements = BufferMut::::with_capacity(row_count * tensor_flat_size); + let mut norms = BufferMut::::with_capacity(row_count); + + if let Some(valid_rows) = &valid_rows { + for row_idx in 0..row_count { + if !valid_rows.value(row_idx) { + // SAFETY: The buffers reserve one direction row and one norm per input row. + unsafe { + elements.push_n_unchecked(T::zero(), tensor_flat_size); + norms.push_unchecked(T::zero()); + } + } else { + // SAFETY: `elements` reserves `tensor_flat_size` values per input row. + let norm = + unsafe { normalize_row_into(flat.row::(row_idx), &mut elements)? }; + // SAFETY: `norms` reserves one value per input row. + unsafe { norms.push_unchecked(norm) }; + } + } + } else { + for row_idx in 0..row_count { + // SAFETY: `elements` reserves `tensor_flat_size` values for every input row. + let norm = unsafe { normalize_row_into(flat.row::(row_idx), &mut elements)? }; + // SAFETY: `norms` reserves one value for every input row. + unsafe { norms.push_unchecked(norm) }; + } + } + + let normalized = + build_normalized_array(output_dtype, tensor_flat_size, row_count, elements.freeze())?; + // SAFETY: The loop writes exactly one norm for each input row. + let norms = unsafe { PrimitiveArray::new_unchecked(norms.freeze(), Validity::NonNullable) }; + + Ok::<_, vortex_error::VortexError>((normalized, norms)) + })?; + + Ok((normalized, norms.into_array(), validity)) +} + +/// Writes the normalized row to `output` and returns its physical L2 norm. +/// +/// # Safety +/// +/// `output` must have spare capacity for every value in `row`. +// This runs once per row; inlining avoids returning a large `VortexResult` from the hot loop. +#[inline(always)] +pub(crate) unsafe fn normalize_row_into( + row: &[T], + output: &mut BufferMut, +) -> VortexResult { + let sum_squares = row.iter().fold(0.0f64, |sum, value| { + let value = + ToPrimitive::to_f64(value).vortex_expect("float NativePType values convert to f64"); + sum + value * value + }); + + let (norm_f64, scaled_divisor) = if sum_squares.is_finite() && sum_squares != 0.0 { + (sum_squares.sqrt(), None) + } else if sum_squares == 0.0 && row.iter().all(Zero::is_zero) { + (0.0, None) + } else { + let (scale, scaled_norm) = scaled_l2_norm(row); + (scale * scaled_norm, Some((scale, scaled_norm))) + }; + + vortex_ensure!( + norm_f64.is_finite(), + "L2 norm must be finite, got {norm_f64}" + ); + let norm = T::from_f64(norm_f64).ok_or_else(|| { + vortex_err!( + "L2 norm must be representable as {}, got {norm_f64}", + T::PTYPE, + ) + })?; + vortex_ensure!( + ToPrimitive::to_f64(&norm).is_some_and(f64::is_finite), + "L2 norm must be representable as {}, got {norm_f64}", + T::PTYPE, + ); + if norm_f64 == 0.0 { + // SAFETY: The caller reserves space for the entire row. + unsafe { output.push_n_unchecked(T::zero(), row.len()) }; + return Ok(norm); + } + + if let Some((scale, scaled_norm)) = scaled_divisor { + for value in row { + let value = + ToPrimitive::to_f64(value).vortex_expect("float NativePType values convert to f64"); + let normalized = T::from_f64((value / scale) / scaled_norm) + .vortex_expect("float NativePType values can represent an f64 direction"); + // SAFETY: The caller reserves space for the entire row. + unsafe { output.push_unchecked(normalized) }; + } + } else { + for value in row { + let value = + ToPrimitive::to_f64(value).vortex_expect("float NativePType values convert to f64"); + let normalized = T::from_f64(value / norm_f64) + .vortex_expect("float NativePType values can represent an f64 direction"); + // SAFETY: The caller reserves space for the entire row. + unsafe { output.push_unchecked(normalized) }; + } + } + + Ok(norm) +} + +fn scaled_l2_norm(row: &[T]) -> (f64, f64) { + let mut scale = 0.0f64; + let mut sum_squares = 1.0f64; + + for value in row { + let absolute = ToPrimitive::to_f64(value) + .vortex_expect("float NativePType values convert to f64") + .abs(); + if absolute.is_nan() { + scale = f64::NAN; + break; + } + if absolute.is_infinite() { + scale = f64::INFINITY; + break; + } + if absolute == 0.0 { + continue; + } + + if scale < absolute { + let ratio = scale / absolute; + sum_squares = 1.0 + sum_squares * ratio * ratio; + scale = absolute; + } else { + let ratio = absolute / scale; + sum_squares += ratio * ratio; + } + } + + (scale, sum_squares.sqrt()) +} + +pub(crate) fn normalized_output_dtype(input_dtype: &DType) -> VortexResult { + let ext_dtype = input_dtype.as_extension(); + if ext_dtype.is::() { + let unit_dtype = ExtDType::::try_new( + EmptyMetadata, + ext_dtype.storage_dtype().as_nonnullable(), + )?; + return Ok(DType::Extension(unit_dtype.erased())); + } + + Ok(input_dtype.as_nonnullable()) +} + +fn build_normalized_array( + dtype: DType, + tensor_flat_size: usize, + row_count: usize, + elements: Buffer, +) -> VortexResult { + let list_size = + u32::try_from(tensor_flat_size).vortex_expect("tensor flat size must fit into `u32`"); + // SAFETY: Tensor elements are always non-nullable, so the validity carries no length. + let elements = unsafe { PrimitiveArray::new_unchecked(elements, Validity::NonNullable) }; + let storage = FixedSizeListArray::try_new( + elements.into_array(), + list_size, + Validity::NonNullable, + row_count, + )?; + + if dtype.as_extension().is::() { + // SAFETY: `normalize_row_into` produced every valid row, and null rows contain zeros. + return unsafe { UnitVector::new_unchecked(storage.into_array()) }; + } + + Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) +} + +#[cfg(test)] +mod tests { + use half::f16; + use vortex_array::ArrayPlugin; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::extension::ExtensionArrayExt; + use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; + use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_error::VortexResult; + + use crate::encodings::normalized::validate_normalized_rows; + use crate::scalar_fns::l2_normalize::L2Normalize; + use crate::tests::SESSION; + use crate::unit_vector::AnyUnitVector; + use crate::unit_vector::UnitVector; + use crate::utils::test_helpers::tensor_array; + use crate::utils::test_helpers::vector_array; + + fn evaluate(input: ArrayRef) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + L2Normalize::try_new_array(input)? + .into_array() + .execute(&mut ctx) + } + + #[test] + fn vector_returns_unit_vector_and_norms() -> VortexResult<()> { + let result = evaluate(vector_array(2, &[3.0f64, 4.0, 0.0, 0.0])?)?; + let normalized = result.unmasked_field_by_name("normalized")?; + let norms: PrimitiveArray = result + .unmasked_field_by_name("norms")? + .clone() + .execute(&mut SESSION.create_execution_ctx())?; + + assert!(normalized.dtype().as_extension().is::()); + assert_eq!(norms.as_slice::(), &[5.0, 0.0]); + Ok(()) + } + + #[test] + fn fixed_shape_tensor_keeps_its_dtype() -> VortexResult<()> { + let input = tensor_array(&[2], &[3.0f64, 4.0])?; + let expected = input.dtype().as_nonnullable(); + let result = evaluate(input)?; + + assert_eq!( + result.unmasked_field_by_name("normalized")?.dtype(), + &expected, + ); + Ok(()) + } + + #[test] + fn unit_vector_input_reports_physical_norm() -> VortexResult<()> { + let vector = vector_array(2, &[0.6000005f32, 0.8])?; + let mut ctx = SESSION.create_execution_ctx(); + let vector: ExtensionArray = vector.execute(&mut ctx)?; + let unit = UnitVector::try_new_unit_vector_array(vector.storage_array().clone(), &mut ctx)?; + let result = evaluate(unit)?; + let norms: PrimitiveArray = result + .unmasked_field_by_name("norms")? + .clone() + .execute(&mut ctx)?; + + assert_ne!(norms.as_slice::()[0], 1.0); + assert!( + result + .unmasked_field_by_name("normalized")? + .dtype() + .as_extension() + .is::() + ); + Ok(()) + } + + #[test] + fn serde_round_trip() -> VortexResult<()> { + let child = vector_array(2, &[3.0f64, 4.0])?; + let original = L2Normalize::try_new_array(child.clone())?.into_array(); + let plugin = ScalarFnArrayPlugin::new(L2Normalize); + let metadata = plugin + .serialize(&original, &SESSION)? + .expect("L2Normalize must serialize metadata"); + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &[child], + &SESSION, + )?; + + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.encoding_id(), original.encoding_id()); + Ok(()) + } + + #[test] + fn reports_value_dependent_errors() { + assert!(L2Normalize.is_fallible(&EmptyOptions)); + } + + #[test] + fn f16_output_satisfies_the_capped_unit_tolerance() -> VortexResult<()> { + for dimensions in [128_u32, 768, 4096] { + let length = usize::try_from(dimensions)?; + let values = vec![f16::ONE; length]; + let result = evaluate(vector_array(dimensions, &values)?)?; + let normalized = result.unmasked_field_by_name("normalized")?; + + assert!(normalized.dtype().as_extension().is::()); + validate_normalized_rows(normalized, None, &mut SESSION.create_execution_ctx())?; + } + Ok(()) + } + + #[test] + fn scaled_fallback_handles_overflowing_sum_of_squares() -> VortexResult<()> { + let value = f64::MAX / 2.0; + let result = evaluate(vector_array(2, &[value, value])?)?; + let normalized = result.unmasked_field_by_name("normalized")?; + let norms: PrimitiveArray = result + .unmasked_field_by_name("norms")? + .clone() + .execute(&mut SESSION.create_execution_ctx())?; + + validate_normalized_rows(normalized, None, &mut SESSION.create_execution_ctx())?; + let norm = norms.as_slice::()[0]; + assert!(norm.is_finite()); + assert!(norm > value); + Ok(()) + } + + #[test] + fn scaled_fallback_handles_underflowing_sum_of_squares() -> VortexResult<()> { + let value = f64::MIN_POSITIVE; + let result = evaluate(vector_array(2, &[value, value])?)?; + let normalized = result.unmasked_field_by_name("normalized")?; + let norms: PrimitiveArray = result + .unmasked_field_by_name("norms")? + .clone() + .execute(&mut SESSION.create_execution_ctx())?; + + validate_normalized_rows(normalized, None, &mut SESSION.create_execution_ctx())?; + assert!(norms.as_slice::()[0] > 0.0); + Ok(()) + } + + #[test] + fn rejects_non_finite_norms() -> VortexResult<()> { + let input = vector_array(2, &[f64::NAN, 0.0])?; + + assert!(evaluate(input).is_err()); + Ok(()) + } + + #[test] + fn rejects_norms_that_do_not_fit_the_input_ptype() -> VortexResult<()> { + let input = vector_array(2, &[f16::MAX, f16::MAX])?; + + assert!(evaluate(input).is_err()); + Ok(()) + } +} diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 68f10ca6b01..cd562f2d0e0 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -6,3 +6,4 @@ pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; +pub mod l2_normalize; diff --git a/vortex-tensor/src/types/mod.rs b/vortex-tensor/src/types/mod.rs index 47dcabdb36d..cbb2be3fa60 100644 --- a/vortex-tensor/src/types/mod.rs +++ b/vortex-tensor/src/types/mod.rs @@ -4,4 +4,5 @@ //! Internal homes for tensor extension types. pub mod fixed_shape_tensor; +pub mod unit_vector; pub mod vector; diff --git a/vortex-tensor/src/types/unit_vector/arrow.rs b/vortex-tensor/src/types/unit_vector/arrow.rs new file mode 100644 index 00000000000..5a7935978d8 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/arrow.rs @@ -0,0 +1,229 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow conversion for [`UnitVector`]. +//! +//! Arrow data carrying the `vortex.tensor.unit_vector` extension name is trusted to satisfy the +//! unit-norm refinement. This keeps import structural and zero-copy; callers handling untrusted +//! values must use [`UnitVector::try_new_unit_vector_array`] instead. + +use std::sync::Arc; + +use arrow_array::Array; +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use crate::types::unit_vector::UnitVector; + +/// Arrow extension name used to identify [`UnitVector`] fields on the wire. +pub const ARROW_UNIT_VECTOR_EXTENSION_NAME: &str = "vortex.tensor.unit_vector"; + +static ARROW_UNIT_VECTOR: CachedId = CachedId::new(ARROW_UNIT_VECTOR_EXTENSION_NAME); + +#[expect( + clippy::disallowed_types, + reason = "Arrow's Field::set_metadata requires std::collections::HashMap" +)] +fn unit_vector_extension_metadata() -> std::collections::HashMap { + [( + EXTENSION_TYPE_NAME_KEY.to_string(), + ARROW_UNIT_VECTOR_EXTENSION_NAME.to_string(), + )] + .into() +} + +fn is_supported_float(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Float16 | DataType::Float32 | DataType::Float64 + ) +} + +impl ArrowExportVTable for UnitVector { + fn arrow_ext_id(&self) -> Id { + *ARROW_UNIT_VECTOR + } + + fn vortex_id(&self) -> Id { + UnitVector.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let DType::Extension(dtype) = dtype else { + return Ok(None); + }; + if !dtype.is::() { + return Ok(None); + } + + let mut field = session.to_arrow_field(name, dtype.storage_dtype())?; + field.set_metadata(unit_vector_extension_metadata()); + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if !array + .dtype() + .as_extension_opt() + .is_some_and(|ext| ext.is::()) + { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + let session = ctx.session().clone(); + let arrow_storage = session.arrow().execute_arrow(storage, Some(target), ctx)?; + + Ok(ArrowExport::Exported(arrow_storage)) + } +} + +impl ArrowImportVTable for UnitVector { + fn arrow_ext_id(&self) -> Id { + *ARROW_UNIT_VECTOR + } + + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult> { + if field.extension_type_name() != Some(ARROW_UNIT_VECTOR_EXTENSION_NAME) { + return Ok(None); + } + let DataType::FixedSizeList(element, list_size) = field.data_type() else { + return Ok(None); + }; + if !is_supported_float(element.data_type()) || element.is_nullable() { + return Ok(None); + } + + let storage_dtype = DType::FixedSizeList( + Arc::new(session.from_arrow_field(element.as_ref())?), + *list_size as u32, + field.is_nullable().into(), + ); + let dtype = ExtDType::try_with_vtable(UnitVector, EmptyMetadata, storage_dtype)?; + + Ok(Some(DType::Extension(dtype.erased()))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + _field: &Field, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult { + let DType::Extension(dtype) = dtype else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !dtype.is::() { + return Ok(ArrowImport::Unsupported(array)); + } + let DataType::FixedSizeList(element, _) = array.data_type() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !is_supported_float(element.data_type()) { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = session.from_arrow_array(array, dtype.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::FixedSizeListArray as ArrowFixedSizeListArray; + use arrow_array::Float32Array; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + + use super::*; + + const DIMENSIONS: u32 = 2; + + fn unit_vector_dtype() -> VortexResult { + let storage = DType::FixedSizeList( + Arc::new(DType::Primitive(PType::F32, Nullability::NonNullable)), + DIMENSIONS, + Nullability::NonNullable, + ); + let dtype = ExtDType::::try_new(EmptyMetadata, storage)?; + + Ok(DType::Extension(dtype.erased())) + } + + fn session_with_unit_vector() -> ArrowSession { + let session = ArrowSession::default(); + session.register_exporter(Arc::new(UnitVector)); + session.register_importer(Arc::new(UnitVector)); + session + } + + #[test] + fn field_round_trip_preserves_unit_vector() -> VortexResult<()> { + let session = session_with_unit_vector(); + let dtype = unit_vector_dtype()?; + let field = session.to_arrow_field("embedding", &dtype)?; + + assert_eq!( + field.extension_type_name(), + Some(ARROW_UNIT_VECTOR_EXTENSION_NAME), + ); + assert_eq!(session.from_arrow_field(&field)?, dtype); + Ok(()) + } + + #[test] + fn tagged_import_trusts_the_refinement() -> VortexResult<()> { + let session = session_with_unit_vector(); + let field = session.to_arrow_field("embedding", &unit_vector_dtype()?)?; + let values = Arc::new(Float32Array::from(vec![3.0, 4.0])); + let element = Arc::new(Field::new("item", DataType::Float32, false)); + let arrow: ArrowArrayRef = Arc::new(ArrowFixedSizeListArray::new( + element, + DIMENSIONS as i32, + values, + None, + )); + + let imported = session.from_arrow_array(arrow, &field)?; + assert!(imported.dtype().as_extension().is::()); + Ok(()) + } +} diff --git a/vortex-tensor/src/types/unit_vector/matcher.rs b/vortex-tensor/src/types/unit_vector/matcher.rs new file mode 100644 index 00000000000..b5e56ffef77 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/matcher.rs @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::dtype::extension::ExtDTypeRef; +use vortex_array::dtype::extension::Matcher; + +use crate::types::unit_vector::UnitVector; +use crate::types::vector::VectorMatcherMetadata; +use crate::types::vector::match_vector_storage; + +/// Matches exactly the [`UnitVector`] extension type. +pub struct AnyUnitVector; + +impl Matcher for AnyUnitVector { + type Match<'a> = VectorMatcherMetadata; + + fn try_match<'a>(ext_dtype: &'a ExtDTypeRef) -> Option> { + ext_dtype + .is::() + .then(|| match_vector_storage(ext_dtype)) + } +} diff --git a/vortex-tensor/src/types/unit_vector/mod.rs b/vortex-tensor/src/types/unit_vector/mod.rs new file mode 100644 index 00000000000..bf8d54efb13 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/mod.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Unit-vector extension type for fixed-length float vectors. +//! +//! A [`UnitVector`] uses the same fixed-size-list storage as +//! [`Vector`](crate::vector::Vector), but every non-null row is either exactly zero or has an L2 +//! norm within [`unit_norm_tolerance`](crate::unit_norm_tolerance) of one. Use +//! [`try_new_unit_vector_array`](UnitVector::try_new_unit_vector_array) at untrusted construction +//! boundaries. + +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_error::VortexResult; + +use crate::encodings::normalized::validate_normalized_rows; + +mod arrow; +pub use arrow::ARROW_UNIT_VECTOR_EXTENSION_NAME; + +mod matcher; +pub use matcher::AnyUnitVector; + +mod vtable; + +/// A fixed-length float vector that is unit norm within the configured tolerance, or exactly zero. +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] +pub struct UnitVector; + +impl UnitVector { + /// Constructs a [`UnitVector`] array after validating every non-null row. + /// + /// # Errors + /// + /// Returns an error if the storage dtype is incompatible or any non-null row is neither + /// exactly zero nor unit norm within the configured tolerance. + pub fn try_new_unit_vector_array( + storage: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // SAFETY: The array is validated immediately below before it is returned to the caller. + let array = unsafe { Self::new_unchecked(storage)? }; + validate_normalized_rows(&array, None, ctx)?; + + Ok(array) + } + + /// Constructs a [`UnitVector`] array without validating its row values. + /// + /// # Safety + /// + /// Every non-null row must be exactly zero or have an L2 norm within + /// [`unit_norm_tolerance`](crate::unit_norm_tolerance) of one. Violating this contract can + /// produce incorrect results in operations that use the refinement for approximate compute + /// shortcuts; it does not cause memory unsafety. + pub unsafe fn new_unchecked(storage: ArrayRef) -> VortexResult { + ExtensionArray::try_new_from_vtable(UnitVector, EmptyMetadata, storage) + .map(|array| array.into_array()) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-tensor/src/types/unit_vector/tests.rs b/vortex-tensor/src/types/unit_vector/tests.rs new file mode 100644 index 00000000000..c7f5eae4ff9 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/tests.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use crate::tests::SESSION; +use crate::types::unit_vector::AnyUnitVector; +use crate::types::unit_vector::UnitVector; +use crate::types::vector::AnyVector; +use crate::types::vector::Vector; +use crate::utils::test_helpers::vector_array; +use crate::utils::unit_norm_tolerance; + +fn unit_vector(dimensions: u32, values: &[f64]) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + let vector = vector_array(dimensions, values)?; + let vector: ExtensionArray = vector.execute(&mut ctx)?; + + UnitVector::try_new_unit_vector_array(vector.storage_array().clone(), &mut ctx) +} + +fn storage_dtype(ptype: PType, dimensions: u32) -> DType { + DType::FixedSizeList( + Arc::new(DType::Primitive(ptype, Nullability::NonNullable)), + dimensions, + Nullability::NonNullable, + ) +} + +fn vector_dtype(ptype: PType, dimensions: u32) -> VortexResult { + let dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype(ptype, dimensions))?; + + Ok(DType::Extension(dtype.erased())) +} + +fn unit_dtype(ptype: PType, dimensions: u32) -> VortexResult { + let dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype(ptype, dimensions))?; + + Ok(DType::Extension(dtype.erased())) +} + +#[test] +fn checked_constructor_accepts_unit_and_zero_rows() -> VortexResult<()> { + let array = unit_vector(2, &[0.6, 0.8, 0.0, 0.0])?; + + assert!(array.dtype().as_extension().is::()); + assert!(array.dtype().as_extension().is::()); + Ok(()) +} + +#[test] +fn checked_constructor_rejects_non_unit_row() { + assert!(unit_vector(2, &[3.0, 4.0]).is_err()); +} + +#[test] +fn checked_constructor_rejects_nonzero_row_with_underflowed_norm() { + assert!(unit_vector(2, &[f64::from_bits(1), 0.0]).is_err()); +} + +#[test] +fn scalar_constructor_rejects_non_unit_value() -> VortexResult<()> { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let storage = Scalar::fixed_size_list( + element_dtype, + vec![ + Scalar::primitive(3.0f64, Nullability::NonNullable), + Scalar::primitive(4.0f64, Nullability::NonNullable), + ], + Nullability::NonNullable, + ); + + assert!(Scalar::try_new(unit_dtype(PType::F64, 2)?, storage.into_value()).is_err()); + Ok(()) +} + +#[test] +fn checked_constructor_ignores_null_row_payloads() -> VortexResult<()> { + let elements = buffer![3.0f64, 4.0, 0.6, 0.8].into_array(); + let validity = Validity::Array(BoolArray::from_iter([false, true]).into_array()); + let storage = FixedSizeListArray::try_new(elements, 2, validity, 2)?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + + UnitVector::try_new_unit_vector_array(storage, &mut ctx)?; + Ok(()) +} + +#[test] +fn f16_tolerance_is_capped() { + assert_eq!(unit_norm_tolerance(PType::F16, 768), 1e-2); + assert!(unit_norm_tolerance(PType::F32, 768) < 1e-2); +} + +#[test] +fn unit_vector_coerces_and_casts_to_vector() -> VortexResult<()> { + let array = unit_vector(2, &[0.6, 0.8])?; + let target = vector_dtype(PType::F64, 2)?; + + assert!(target.can_coerce_from(array.dtype())); + assert!(!array.dtype().can_coerce_from(&target)); + + let cast = array.cast(target.clone())?; + assert_eq!(cast.dtype(), &target); + Ok(()) +} + +#[test] +fn mixed_least_supertype_is_vector_and_symmetric() -> VortexResult<()> { + let unit = unit_dtype(PType::F32, 4)?; + let vector = vector_dtype(PType::F64, 4)?; + let expected = vector.clone(); + + assert_eq!(unit.least_supertype(&vector), Some(expected.clone())); + assert_eq!(vector.least_supertype(&unit), Some(expected)); + Ok(()) +} + +#[test] +fn precision_widening_erases_unit_refinement() -> VortexResult<()> { + let f32_unit = unit_dtype(PType::F32, 4)?; + let f64_unit = unit_dtype(PType::F64, 4)?; + let expected = vector_dtype(PType::F64, 4)?; + + assert_eq!(f32_unit.least_supertype(&f64_unit), Some(expected.clone())); + assert_eq!(f64_unit.least_supertype(&f32_unit), Some(expected)); + Ok(()) +} diff --git a/vortex-tensor/src/types/unit_vector/vtable.rs b/vortex-tensor/src/types/unit_vector/vtable.rs new file mode 100644 index 00000000000..2ff8658689e --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/vtable.rs @@ -0,0 +1,118 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::EmptyMetadata; +use vortex_array::dtype::DType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::PValue; +use vortex_array::scalar::ScalarValue; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use crate::types::unit_vector::UnitVector; +use crate::types::vector::Vector; +use crate::types::vector::validate_vector_storage_dtype; +use crate::utils::unit_norm_tolerance; + +impl ExtVTable for UnitVector { + type Metadata = EmptyMetadata; + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.tensor.unit_vector"); + *ID + } + + fn serialize_metadata(&self, _metadata: &Self::Metadata) -> VortexResult> { + Ok(Vec::new()) + } + + fn deserialize_metadata(&self, _metadata: &[u8]) -> VortexResult { + Ok(EmptyMetadata) + } + + fn can_coerce_to(source: &ExtDType, target: &DType) -> bool { + let DType::Extension(target) = target else { + return false; + }; + if !target.is::() { + return false; + } + + target + .storage_dtype() + .can_coerce_from(source.storage_dtype()) + } + + fn least_supertype(source: &ExtDType, other: &DType) -> Option { + let DType::Extension(other) = other else { + return None; + }; + let other_is_unit = other.is::(); + if !other_is_unit && !other.is::() { + return None; + } + + let storage = source + .storage_dtype() + .least_supertype(other.storage_dtype())?; + + if other_is_unit + && source + .storage_dtype() + .eq_ignore_nullability(other.storage_dtype()) + { + let unit = ExtDType::::try_new(EmptyMetadata, storage).ok()?; + return Some(DType::Extension(unit.erased())); + } + + let vector = ExtDType::::try_new(EmptyMetadata, storage).ok()?; + Some(DType::Extension(vector.erased())) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + validate_vector_storage_dtype(ext_dtype.storage_dtype()) + } + + fn unpack_native<'a>( + ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult> { + let elements = storage_value.as_list(); + let DType::FixedSizeList(element_dtype, ..) = ext_dtype.storage_dtype() else { + unreachable!("UnitVector dtype validation established fixed-size-list storage") + }; + let element_ptype = element_dtype.as_ptype(); + let tolerance = unit_norm_tolerance(element_ptype, elements.len()); + + let (norm_squared, is_zero) = elements.iter().try_fold( + (0.0, true), + |(sum_squared, is_zero), element| -> VortexResult<_> { + let value = element + .as_ref() + .ok_or_else(|| vortex_err!("UnitVector scalar elements must be non-null"))? + .as_primitive(); + let value = match value { + PValue::F16(value) => value.to_f64(), + PValue::F32(value) => *value as f64, + PValue::F64(value) => *value, + _ => unreachable!("UnitVector dtype validation established float elements"), + }; + + Ok((sum_squared + value * value, is_zero && value == 0.0)) + }, + )?; + let norm = norm_squared.sqrt(); + + vortex_ensure!( + is_zero || (norm - 1.0).abs() <= tolerance, + "UnitVector scalar must have L2 norm 1.0 or be exactly zero, got {norm:.6}", + ); + + Ok(storage_value) + } +} diff --git a/vortex-tensor/src/types/vector/matcher.rs b/vortex-tensor/src/types/vector/matcher.rs index 9f5b0037029..157d85bba6f 100644 --- a/vortex-tensor/src/types/vector/matcher.rs +++ b/vortex-tensor/src/types/vector/matcher.rs @@ -5,33 +5,23 @@ use vortex_array::dtype::DType; use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDTypeRef; use vortex_array::dtype::extension::Matcher; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; +use crate::types::unit_vector::UnitVector; use crate::types::vector::Vector; +/// Matches both [`Vector`] and [`UnitVector`]. pub struct AnyVector; -/// Convenience metadata for vectors. -/// -/// Unlike `FixedShapeTensor`, the [`Vector`] type has `EmptyMetadata` as its metadata because all -/// of the important information is already stored in the dtype. -/// -/// However, it is quite inconvenient to repeatedly unwrap the dtype to get the element type of the -/// vector and the number of dimensions. -/// -/// Thus, we allow the matcher to return this metadata so that we can access this information more -/// easily. +/// Shape metadata derived from [`Vector`] or [`UnitVector`] fixed-size-list storage. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct VectorMatcherMetadata { - /// The element type of the vectors. Note that vector elements are _always_ non-nullable. - /// - /// This MUST be a floating point type (f16, f32, f64). + /// The floating-point element ptype. element_ptype: PType, - /// The number of dimensions of the vector. This is always fixed. + /// The number of elements in each vector row. dimensions: u32, } @@ -39,27 +29,22 @@ impl Matcher for AnyVector { type Match<'a> = VectorMatcherMetadata; fn try_match<'a>(ext_dtype: &'a ExtDTypeRef) -> Option> { - if !ext_dtype.is::() { + if !ext_dtype.is::() && !ext_dtype.is::() { return None; } - let DType::FixedSizeList(element_dtype, list_size, _) = ext_dtype.storage_dtype() else { - vortex_panic!("`Vector` type somehow did not have a `FixedSizeList` storage type") - }; - - let dimensions = *list_size; - - assert!(element_dtype.is_float(), "element dtype must be float"); - assert!( - !element_dtype.is_nullable(), - "element dtype must be non-nullable" - ); - let element_ptype = element_dtype.as_ptype(); + Some(match_vector_storage(ext_dtype)) + } +} - let vector_metadata = VectorMatcherMetadata::try_new(element_ptype, dimensions) - .vortex_expect("`Vector` type somehow did not have float elements"); +pub(crate) fn match_vector_storage(ext_dtype: &ExtDTypeRef) -> VectorMatcherMetadata { + let DType::FixedSizeList(element_dtype, list_size, _) = ext_dtype.storage_dtype() else { + vortex_panic!("vector dtype somehow did not have a `FixedSizeList` storage type") + }; - Some(vector_metadata) + VectorMatcherMetadata { + element_ptype: element_dtype.as_ptype(), + dimensions: *list_size, } } @@ -70,7 +55,10 @@ impl VectorMatcherMetadata { /// /// Returns an error if the element type is not a float. pub fn try_new(element_ptype: PType, dimensions: u32) -> VortexResult { - vortex_ensure!(element_ptype.is_float()); + vortex_ensure!( + element_ptype.is_float(), + "vector element ptype must be floating point, got {element_ptype}", + ); Ok(Self { element_ptype, diff --git a/vortex-tensor/src/types/vector/mod.rs b/vortex-tensor/src/types/vector/mod.rs index af424b9cc41..b34b4856daa 100644 --- a/vortex-tensor/src/types/vector/mod.rs +++ b/vortex-tensor/src/types/vector/mod.rs @@ -80,5 +80,6 @@ mod matcher; pub use arrow::ARROW_VECTOR_EXTENSION_NAME; pub use matcher::AnyVector; pub use matcher::VectorMatcherMetadata; +pub(crate) use matcher::match_vector_storage; mod vtable; diff --git a/vortex-tensor/src/types/vector/vtable.rs b/vortex-tensor/src/types/vector/vtable.rs index e526bd9982b..793efa0d2da 100644 --- a/vortex-tensor/src/types/vector/vtable.rs +++ b/vortex-tensor/src/types/vector/vtable.rs @@ -10,6 +10,7 @@ use vortex_array::scalar::ScalarValue; use vortex_error::VortexResult; use vortex_session::registry::CachedId; +use crate::types::unit_vector::UnitVector; use crate::types::vector::Vector; use crate::types::vector::validate_vector_storage_dtype; @@ -36,7 +37,7 @@ impl ExtVTable for Vector { let DType::Extension(other_ext) = other else { return None; }; - if !other_ext.is::() { + if !other_ext.is::() && !other_ext.is::() { return None; } let widened = ext_dtype diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 3e33fe20db9..3692445f39b 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -33,20 +33,18 @@ use crate::encodings::normalized::NormalizedArraySlotsExt; use crate::matcher::AnyTensor; use crate::matcher::TensorMatch; -/// Safety factor for unit-norm tolerance. Applied as a constant multiplier on the probabilistic -/// `√d · ε` bound so that legitimate round-off noise clears the check with headroom. -pub(crate) const SAFETY_FACTOR: usize = 10; +const UNIT_NORM_SAFETY_FACTOR: f64 = 10.0; + +const F16_MAX_UNIT_NORM_DRIFT: f64 = 1e-2; /// Returns the acceptable unit-norm drift for the given element precision and dimension count. /// -/// Uses the `c · √d · ε` bound where ε is machine epsilon and d is the vector dimension. Under -/// IEEE 754 round-to-nearest the probabilistic (RMS-case) forward error for computing ‖x‖₂ grows -/// as `O(√d · ε)` rather than the worst-case `O(d · ε)` from the classical Wilkinson bound, -/// assuming near-independent rounding errors across the d-term summation. +/// Uses `10 · √d · ε`, where `ε` is machine epsilon and `d` is the dimension count. The f16 +/// tolerance is capped at one percent. +/// +/// # Panics /// -/// Reference: Croci, Fasi, Higham, Mary, Mikaitis (2022). "Stochastic rounding: implementation, -/// error analysis and applications." Royal Society Open Science, 9: 211631, §6.1 "Probabilistic -/// error analysis." +/// Panics if `element_ptype` is not floating point. pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { let machine_epsilon: f64 = match element_ptype { PType::F64 => f64::EPSILON, @@ -57,7 +55,12 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { let dimensions_root = (dimensions as f64).sqrt(); - SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root + let tolerance = UNIT_NORM_SAFETY_FACTOR * machine_epsilon * dimensions_root; + if element_ptype == PType::F16 { + tolerance.min(F16_MAX_UNIT_NORM_DRIFT) + } else { + tolerance + } } /// Extracts the `(normalized, norms)` children of a [`Normalized`]-encoded array. @@ -107,17 +110,33 @@ pub fn validate_tensor_float_input(input_dtype: &DType) -> VortexResult( lhs: &'a DType, rhs: &DType, ) -> VortexResult> { + let lhs_match = validate_tensor_float_input(lhs)?; + let rhs_match = validate_tensor_float_input(rhs)?; + let both_vectors = matches!( + (lhs_match, rhs_match), + (TensorMatch::Vector(_), TensorMatch::Vector(_)) + ); + let compatible = lhs.eq_ignore_nullability(rhs) + || (both_vectors + && lhs_match.element_ptype() == rhs_match.element_ptype() + && lhs_match.list_size() == rhs_match.list_size()); + vortex_ensure!( - lhs.eq_ignore_nullability(rhs), - "binary tensor expression expects inputs to have the same dtype, got {lhs} and {rhs}" + compatible, + "binary tensor expression expects compatible input dtypes, got {lhs} and {rhs}" ); - validate_tensor_float_input(lhs) + + Ok(lhs_match) } /// The flat primitive elements of a tensor storage array, with typed row access. diff --git a/vortex/src/editions/unstable/v2026_04.rs b/vortex/src/editions/unstable/v2026_04.rs index 992ddd49853..1e1173b89c2 100644 --- a/vortex/src/editions/unstable/v2026_04.rs +++ b/vortex/src/editions/unstable/v2026_04.rs @@ -24,5 +24,6 @@ pub static DECLARATION: EditionDeclaration = EditionDeclaration { EditionMember::array(&"vortex.tensor.inner_product"), EditionMember::array(&"vortex.tensor.normalized"), EditionMember::array(&"vortex.tensor.l2_norm"), + EditionMember::array(&"vortex.tensor.l2_normalize"), ], };