diff --git a/vortex-tensor/src/encodings/normalized/compress.rs b/vortex-tensor/src/encodings/normalized/compress.rs index a807ee4cf6f..c342c9bf594 100644 --- a/vortex-tensor/src/encodings/normalized/compress.rs +++ b/vortex-tensor/src/encodings/normalized/compress.rs @@ -34,6 +34,7 @@ use vortex_compressor::scheme::SchemeExt; use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::encodings::normalized::Normalized; use crate::encodings::normalized::NormalizedArray; @@ -43,11 +44,12 @@ use crate::encodings::normalized::array::DATA_CHILDREN; use crate::matcher::AnyTensor; use crate::scalar_fns::NormMode; use crate::scalar_fns::l2_norm::L2Norm; +use crate::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. +/// The compression scheme that rewrites an ordinary tensor column into [`Normalized`] storage. #[derive(Debug)] pub struct NormalizedScheme; @@ -63,9 +65,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 { @@ -120,15 +124,29 @@ impl Scheme for NormalizedScheme { } } -/// Splits a tensor-like column into its exact [`Normalized`] representation. +/// Splits a [`Vector`] or [`FixedShapeTensor`] column into its exact [`Normalized`] representation. /// /// The children are non-nullable, and the input validity moves to the parent. Both children are /// zero at null rows so masked physical values cannot reach downstream encodings. /// /// # Errors /// -/// Returns an error if `input` is not a float tensor column or if execution fails. +/// Returns an error if `input` is a [`UnitVector`], is not a float tensor column, or cannot be +/// executed. +/// +/// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor +/// [`UnitVector`]: crate::unit_vector::UnitVector +/// [`Vector`]: crate::vector::Vector pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + vortex_ensure!( + !input + .dtype() + .as_extension_opt() + .is_some_and(|dtype| dtype.is::()), + InvalidArgument: "Normalized input must not already be a UnitVector, got {}", + input.dtype(), + ); + let row_count = input.len(); let tensor_match = validate_tensor_float_input(input.dtype())?; let tensor_flat_size = tensor_match.list_size() as usize; @@ -200,12 +218,23 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult VortexResult> { + if input + .dtype() + .as_extension_opt() + .is_some_and(|dtype| dtype.is::()) + { + return Ok(None); + } + let Some(ext) = input.as_opt::() else { return Ok(None); }; diff --git a/vortex-tensor/src/encodings/normalized/tests.rs b/vortex-tensor/src/encodings/normalized/tests.rs index 0687e1ef750..9ad9583ae66 100644 --- a/vortex-tensor/src/encodings/normalized/tests.rs +++ b/vortex-tensor/src/encodings/normalized/tests.rs @@ -43,6 +43,7 @@ use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::constant_tensor_array; use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::unit_vector_array; use crate::utils::test_helpers::vector_array; fn eval_normalized( @@ -806,6 +807,17 @@ fn scheme_does_not_match_non_float_tensors(#[case] input: ArrayRef) -> VortexRes Ok(()) } +#[test] +fn scheme_does_not_match_unit_vectors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = unit_vector_array(2, &[0.6f64, 0.8], &mut ctx)?; + let canonical: Canonical = input.clone().execute(&mut ctx)?; + + assert!(!NormalizedScheme.matches(&canonical)); + assert!(normalize(input, &mut ctx).is_err()); + Ok(()) +} + #[rstest] #[case::non_nullable(collinear_vectors(1024).expect("valid vector array"))] #[case::nullable(nullable_collinear_vectors(1024).expect("valid vector array"))] diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index fe56827c15b..9543df928d6 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -4,6 +4,14 @@ //! Types and functionality for working with tensors, vectors, and related mathematical constructs //! including unit vectors, spherical coordinates, and similarity measures such as cosine //! similarity. +//! +//! [`Vector`], [`UnitVector`], and [`FixedShapeTensor`] define logical tensor dtypes. The +//! [`scalar_fns`] module provides tensor operations, while [`encodings`] contains their physical +//! representations. +//! +//! [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor +//! [`UnitVector`]: crate::unit_vector::UnitVector +//! [`Vector`]: crate::vector::Vector #![cfg_attr( test, @@ -23,7 +31,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 +42,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 +53,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 +65,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,17 +81,19 @@ 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 // to persist these encodings do not pay the registry cost or widen their stable-encoding // surface unintentionally. - if std::env::var_os(SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV).is_some_and(|v| !v.is_empty()) { + if std::env::var_os(SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV).is_some_and(|value| !value.is_empty()) { let session_arrays = session.arrays(); 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..d23febeb95c 100644 --- a/vortex-tensor/src/matcher.rs +++ b/vortex-tensor/src/matcher.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Matcher for tensor-like extension types. +//! Matchers for tensor-like extension types. +//! +//! [`AnyTensor`] recognizes fixed-shape tensors and both vector refinements. [`TensorMatch`] +//! exposes their shared float element type and flattened row width without erasing the distinction +//! between fixed-shape tensors and vectors. use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDTypeRef; @@ -14,10 +18,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 +28,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 +58,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 4fe4892be40..294b43fc661 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -48,16 +48,18 @@ 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 of that element ptype. /// -/// [`NormMode::Exact`] measures the physical direction stored by a [`Normalized`] encoding before -/// applying the cosine formula. [`NormMode::AssumeNormalized`] instead trusts that direction as -/// unit length and omits its norm computation. The approximate mode does not clamp its output or -/// provide an error bound for unchecked or lossy encodings. +/// [`NormMode::Exact`] measures all physical coordinates before applying the cosine formula. +/// [`NormMode::AssumeNormalized`] omits norms claimed by a [`UnitVector`] dtype or [`Normalized`] +/// encoding. The approximate mode does not clamp its output or provide an error bound for unchecked +/// or lossy claims. /// /// [`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; @@ -152,9 +154,10 @@ impl ScalarFnVTable for CosineSimilarity { // Compute combined validity. let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - // Ordinary inputs carry no normalized claim, so both modes measure their physical norms. - let norm_lhs_arr = L2Norm::try_new(lhs_ref.clone(), NormMode::Exact)?; - let norm_rhs_arr = L2Norm::try_new(rhs_ref.clone(), NormMode::Exact)?; + // UnitVector inputs can carry a normalized claim. Ordinary tensors measure physical norms + // in both modes. + let norm_lhs_arr = L2Norm::try_new(lhs_ref.clone(), *options)?; + let norm_rhs_arr = L2Norm::try_new(rhs_ref.clone(), *options)?; let dot_arr = InnerProduct::try_new(lhs_ref, rhs_ref)?; // Execute to get the inner product and norms of the arrays. We only fully decompress @@ -369,7 +372,7 @@ impl CosineSimilarity { let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; - let norm_arr = L2Norm::try_new(plain_ref.clone(), NormMode::Exact)?; + let norm_arr = L2Norm::try_new(plain_ref.clone(), mode)?; let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; // TODO(connor)[Tensor]: Replace this loop after binary numeric operations support @@ -430,6 +433,7 @@ mod tests { use crate::utils::test_helpers::constant_tensor_array; use crate::utils::test_helpers::normalized_array; use crate::utils::test_helpers::tensor_array; + use crate::utils::test_helpers::unit_vector_array; use crate::utils::test_helpers::vector_array; /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. @@ -582,6 +586,35 @@ mod tests { Ok(()) } + #[test] + fn mode_controls_unit_vector_norms() -> VortexResult<()> { + let value = 1.0 + 8.0 * f64::EPSILON; + let mut ctx = SESSION.create_execution_ctx(); + let input = unit_vector_array(2, &[value, 0.0], &mut ctx)?; + + assert_eq!( + eval_cosine_similarity_with_mode(input.clone(), input.clone(), NormMode::Exact)?, + vec![1.0], + ); + let assumed = + eval_cosine_similarity_with_mode(input.clone(), input, NormMode::AssumeNormalized)?; + assert!(assumed[0] > 1.0); + Ok(()) + } + + #[test] + fn mixes_unit_and_ordinary_vectors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let unit = unit_vector_array(2, &[0.6f64, 0.8], &mut ctx)?; + let ordinary = vector_array(2, &[3.0f64, 4.0])?; + + assert_close( + &eval_cosine_similarity_with_mode(unit, ordinary, NormMode::AssumeNormalized)?, + &[1.0], + ); + Ok(()) + } + #[test] fn vector_constant_query() -> VortexResult<()> { let data = vector_array( diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 111b11dfcaa..cae80ea21df 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 of that element ptype. /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector +/// [`UnitVector`]: crate::unit_vector::UnitVector #[derive(Clone)] pub struct InnerProduct; @@ -92,7 +94,8 @@ impl ScalarFnVTable for InnerProduct { let lhs = &arg_dtypes[0]; let rhs = &arg_dtypes[1]; - // TODO(connor): relax the float-only gate once integer tensors are supported. + // TODO(connor)[Tensor]: Add integer products after their result dtype and overflow + // semantics are defined. The current kernels and return dtype support only floats. let tensor_match = validate_binary_tensor_float_inputs(lhs, rhs)?; let ptype = tensor_match.element_ptype(); let nullability = Nullability::from(lhs.is_nullable() || rhs.is_nullable()); @@ -238,7 +241,8 @@ impl InnerProduct { /// /// [`Normalized`]: crate::encodings::normalized::Normalized /// - /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. + /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as + /// `plain_ref`. fn execute_one_normalized( &self, normalized_ref: &ArrayRef, @@ -296,6 +300,7 @@ mod tests { use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::normalized_array; use crate::utils::test_helpers::tensor_array; + use crate::utils::test_helpers::unit_vector_array; use crate::utils::test_helpers::vector_array; /// Evaluates inner product between two tensor arrays and returns the result as `Vec`. @@ -370,6 +375,16 @@ mod tests { Ok(()) } + #[test] + fn mixes_unit_and_ordinary_vectors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let unit = unit_vector_array(2, &[0.6f64, 0.8], &mut ctx)?; + let ordinary = vector_array(2, &[3.0f64, 4.0])?; + + assert_close(&eval_inner_product(unit, ordinary)?, &[5.0]); + Ok(()) + } + #[test] fn null_input_row() -> VortexResult<()> { // 3 rows of dim-2 vectors. Row 1 of lhs is masked as null. diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 1e65226c049..57d3349e53c 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -4,6 +4,7 @@ //! L2 norm expression for tensor-like types. use num_traits::Float; +use num_traits::One; use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -45,6 +46,7 @@ use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; use crate::matcher::AnyTensor; use crate::scalar_fns::NormMode; +use crate::types::unit_vector::UnitVector; use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; use crate::utils::reattach_validity; @@ -57,12 +59,13 @@ use crate::utils::validate_tensor_float_input; /// The input must be a tensor-like extension array with a float element type. The output is a float /// column of the same float type. /// -/// [`NormMode::Exact`] measures the physical direction stored by a [`Normalized`] encoding and -/// multiplies that result by the stored norm. [`NormMode::AssumeNormalized`] instead trusts that -/// direction as unit length and returns the stored norm directly. The approximate mode does not -/// provide an error bound for unchecked or lossy encodings. +/// [`NormMode::Exact`] measures physical coordinates, including a [`UnitVector`] value or the +/// direction stored by a [`Normalized`] encoding. [`NormMode::AssumeNormalized`] returns one for a +/// [`UnitVector`] and returns a [`Normalized`] value's stored norm directly. The approximate mode +/// does not provide an error bound for unchecked or lossy claims. /// /// [`Normalized`]: crate::encodings::normalized::Normalized +/// [`UnitVector`]: crate::unit_vector::UnitVector #[derive(Clone)] pub struct L2Norm; @@ -136,6 +139,15 @@ impl ScalarFnVTable for L2Norm { let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); + if options.assumes_normalized() && ext.is::() { + let one = match_each_float_ptype!(element_ptype, |T| { + Scalar::primitive(T::one(), Nullability::NonNullable) + }); + let ones = ConstantArray::new(one, row_count).into_array(); + + return reattach_validity(ones, input_ref.validity()?); + } + if input_ref.is::() { let (direction, stored_norms) = extract_normalized_children(&input_ref); if options.assumes_normalized() { @@ -328,6 +340,7 @@ mod tests { use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::literal_vector_array; use crate::utils::test_helpers::tensor_array; + use crate::utils::test_helpers::unit_vector_array; use crate::utils::test_helpers::vector_array; /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. @@ -384,6 +397,23 @@ mod tests { Ok(()) } + #[test] + fn mode_controls_unit_vector_claim() -> VortexResult<()> { + let value = 1.0 + 8.0 * f64::EPSILON; + let mut ctx = SESSION.create_execution_ctx(); + let input = unit_vector_array(2, &[value, 0.0], &mut ctx)?; + + assert_eq!( + eval_l2_norm_with_mode(input.clone(), NormMode::Exact)?, + vec![value], + ); + assert_eq!( + eval_l2_norm_with_mode(input, NormMode::AssumeNormalized)?, + vec![1.0], + ); + Ok(()) + } + #[test] fn null_input_row() -> VortexResult<()> { // 2 rows of dim-2 vectors. Row 1 is masked as null. 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..9c9bea3d4cc --- /dev/null +++ b/vortex-tensor/src/scalar_fns/l2_normalize.rs @@ -0,0 +1,486 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! L2 normalization for vector columns. +//! +//! [`L2Normalize`] converts each finite, nonzero [`Vector`] row into a [`UnitVector`]. Exact-zero +//! rows become null because they have no direction, and input nulls remain null. An existing +//! [`UnitVector`] is returned unchanged. Execution returns an error if a valid input row contains +//! a non-finite coordinate. +//! +//! [`Vector`]: crate::vector::Vector +//! [`UnitVector`]: crate::unit_vector::UnitVector + +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::BoolArray; +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::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::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::proto::dtype as pb; +use vortex_array::expr::Expression; +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::ScalarFnVTableExt; +use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; +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::types::vector::Vector; +use crate::utils::extract_flat_elements; + +/// Normalizes vector rows into [`UnitVector`] values. +/// +/// Ordinary [`Vector`] inputs produce a nullable result because exact-zero rows have no direction. +/// Existing [`UnitVector`] inputs retain their dtype and nullability. +/// +/// [`Vector`]: crate::vector::Vector +#[derive(Clone)] +pub struct L2Normalize; + +impl L2Normalize { + /// Constructs a [`ScalarFnArray`] that lazily normalizes `child`. + /// + /// # Errors + /// + /// Returns an error if `child` is not a float [`Vector`] or [`UnitVector`], or if the scalar + /// function array cannot be constructed. + /// + /// [`Vector`]: crate::vector::Vector + /// [`UnitVector`]: crate::unit_vector::UnitVector + pub fn try_new(child: ArrayRef) -> VortexResult { + ScalarFnArray::try_new(L2Normalize.bind(EmptyOptions), 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 { + normalized_dtype(&arg_dtypes[0]) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + if input.dtype().as_extension().is::() { + return Ok(input); + } + + normalize_vector(input, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + _expression: &Expression, + ) -> VortexResult> { + // A valid zero row becomes null, so validity requires evaluating the values. + Ok(None) + } + + 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 L2Normalize metadata: {error}"))?; + let input_dtype = metadata + .input_dtype + .as_ref() + .ok_or_else(|| vortex_err!("L2Normalize metadata missing input_dtype"))?; + let input_dtype = DType::from_proto(input_dtype, session)?; + normalized_dtype(&input_dtype)?; + let child = children.get(0, &input_dtype, len)?; + + Ok(ScalarFnArrayParts { + options: EmptyOptions, + children: vec![child], + }) + } +} + +fn normalized_dtype(input_dtype: &DType) -> VortexResult { + let ext_dtype = input_dtype.as_extension_opt().ok_or_else(|| { + vortex_err!("L2Normalize input must be a Vector or UnitVector, got {input_dtype}") + })?; + let metadata = ext_dtype.metadata_opt::().ok_or_else(|| { + vortex_err!("L2Normalize input must be a Vector or UnitVector, got {input_dtype}") + })?; + + if ext_dtype.is::() { + return Ok(input_dtype.clone()); + } + + vortex_ensure!( + ext_dtype.is::(), + "L2Normalize input must be a Vector or UnitVector, got {input_dtype}", + ); + let storage_dtype = DType::FixedSizeList( + DType::Primitive(metadata.element_ptype(), Nullability::NonNullable).into(), + metadata.dimensions(), + Nullability::Nullable, + ); + let output_dtype = ExtDType::::try_new(EmptyMetadata, storage_dtype)?; + + Ok(DType::Extension(output_dtype.erased())) +} + +fn normalize_vector(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let row_count = input.len(); + let metadata = input.dtype().as_extension().metadata::(); + let dimensions = metadata.dimensions() as usize; + + let input: ExtensionArray = input.execute(ctx)?; + let input_validity = input.as_ref().validity()?; + let valid_rows = input_validity + .nullability() + .is_nullable() + .then(|| input_validity.execute_mask(row_count, ctx)) + .transpose()?; + let flat = extract_flat_elements(input.storage_array(), dimensions, ctx)?; + + match_each_float_ptype!(metadata.element_ptype(), |T| { + let mut elements = BufferMut::::with_capacity(row_count * dimensions); + let mut output_validity = Vec::with_capacity(row_count); + + for row_idx in 0..row_count { + if valid_rows + .as_ref() + .is_some_and(|valid_rows| !valid_rows.value(row_idx)) + { + elements.push_n(T::zero(), dimensions); + output_validity.push(false); + continue; + } + + let is_nonzero = normalize_row_into(flat.row::(row_idx), &mut elements, row_idx)?; + output_validity.push(is_nonzero); + } + + let validity = Validity::Array(BoolArray::from_iter(output_validity).into_array()); + // SAFETY: The loop writes exactly `row_count * dimensions` non-nullable elements. + let elements = + unsafe { PrimitiveArray::new_unchecked(elements.freeze(), Validity::NonNullable) }; + let storage = FixedSizeListArray::try_new( + elements.into_array(), + metadata.dimensions(), + validity, + row_count, + )?; + + // SAFETY: `normalize_row_into` emits a finite unit direction for every valid row. Zero and + // input-null rows are marked null, so their payloads do not participate in the invariant. + unsafe { UnitVector::new_unchecked(storage.into_array()) } + }) +} + +/// Writes one normalized row and returns whether its direction is defined. +fn normalize_row_into( + row: &[T], + output: &mut BufferMut, + row_idx: usize, +) -> VortexResult { + let mut scale = 0.0f64; + for value in row { + let value = ToPrimitive::to_f64(value) + .vortex_expect("float NativePType values must convert to f64"); + vortex_ensure!( + value.is_finite(), + InvalidArgument: "L2Normalize input row {row_idx} must be finite, got {value}", + ); + scale = scale.max(value.abs()); + } + + if scale == 0.0 { + output.push_n(T::zero(), row.len()); + return Ok(false); + } + + let scaled_norm = row + .iter() + .map(|value| { + ToPrimitive::to_f64(value).vortex_expect("float NativePType values must convert to f64") + / scale + }) + .map(|value| value * value) + .sum::() + .sqrt(); + + for value in row { + let value = ToPrimitive::to_f64(value) + .vortex_expect("float NativePType values must convert to f64"); + let normalized = T::from_f64((value / scale) / scaled_norm) + .vortex_expect("normalized float coordinates must fit their input ptype"); + + output.push(normalized); + } + + Ok(true) +} + +#[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::FixedSizeListArray; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::extension::ExtensionArrayExt; + use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; + use vortex_array::arrays::scalar_fn::ExactScalarFn; + use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; + use vortex_array::dtype::NativePType; + use vortex_array::matcher::Matcher; + use vortex_array::validity::Validity; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + + use crate::scalar_fns::l2_normalize::L2Normalize; + use crate::tests::SESSION; + use crate::types::unit_vector::UnitVector; + use crate::types::vector::Vector; + use crate::utils::test_helpers::assert_close; + use crate::utils::test_helpers::tensor_array; + use crate::utils::test_helpers::unit_vector_array; + use crate::utils::test_helpers::vector_array; + + fn evaluate(input: ArrayRef) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + L2Normalize::try_new(input)?.into_array().execute(&mut ctx) + } + + fn values(array: &ExtensionArray) -> VortexResult> { + let mut ctx = SESSION.create_execution_ctx(); + let storage: FixedSizeListArray = array.storage_array().clone().execute(&mut ctx)?; + let values: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; + + Ok(values.as_slice::().to_vec()) + } + + #[test] + fn normalizes_vector_rows() -> VortexResult<()> { + let output = evaluate(vector_array(2, &[3.0f64, 4.0, 1.0, 0.0])?)?; + + assert!(output.dtype().as_extension().is::()); + assert_close(&values::(&output)?, &[0.6, 0.8, 1.0, 0.0]); + Ok(()) + } + + #[test] + fn zero_rows_become_null() -> VortexResult<()> { + let output = evaluate(vector_array(2, &[3.0f64, 4.0, 0.0, 0.0])?)?; + let mut ctx = SESSION.create_execution_ctx(); + + assert!(output.is_valid(0, &mut ctx)?); + assert!(!output.is_valid(1, &mut ctx)?); + assert_close(&values::(&output)?, &[0.6, 0.8, 0.0, 0.0]); + Ok(()) + } + + #[test] + fn input_nulls_remain_null() -> VortexResult<()> { + let input = vector_array(2, &[3.0f64, 4.0, 1.0, 0.0])?; + let input = MaskedArray::try_new(input, Validity::from_iter([false, true]))?.into_array(); + let output = evaluate(input)?; + let mut ctx = SESSION.create_execution_ctx(); + + assert!(!output.is_valid(0, &mut ctx)?); + assert!(output.is_valid(1, &mut ctx)?); + Ok(()) + } + + #[test] + fn normalizes_constant_storage() -> VortexResult<()> { + let input = Vector::constant_array(&[3.0f64, 4.0], 3)?; + let output = evaluate(input)?; + + assert_close(&values::(&output)?, &[0.6, 0.8, 0.6, 0.8, 0.6, 0.8]); + Ok(()) + } + + #[test] + fn normalizes_empty_input() -> VortexResult<()> { + let output = evaluate(vector_array::(2, &[])?)?; + + assert_eq!(output.len(), 0); + assert!(output.dtype().as_extension().is::()); + Ok(()) + } + + #[test] + fn rejects_non_finite_rows() -> VortexResult<()> { + let input = vector_array(2, &[f64::INFINITY, 1.0])?; + let mut ctx = SESSION.create_execution_ctx(); + + assert!( + L2Normalize::try_new(input)? + .into_array() + .execute::(&mut ctx) + .is_err() + ); + Ok(()) + } + + #[test] + fn handles_extreme_f64_values() -> VortexResult<()> { + let tiny = f64::from_bits(1); + let output = evaluate(vector_array(2, &[f64::MAX, f64::MAX, tiny, 0.0])?)?; + let values = values::(&output)?; + + assert_close(&values[..2], &[2.0f64.sqrt().recip(); 2]); + assert_eq!(&values[2..], &[1.0, 0.0]); + Ok(()) + } + + #[test] + fn f16_output_satisfies_unit_vector_validation() -> VortexResult<()> { + let dimensions = 768; + let output = evaluate(vector_array( + dimensions, + &vec![f16::ONE; dimensions as usize], + )?)?; + let mut ctx = SESSION.create_execution_ctx(); + + UnitVector::try_new_unit_vector_array(output.storage_array().clone(), &mut ctx)?; + Ok(()) + } + + #[test] + fn unit_vector_input_is_identity() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = unit_vector_array(2, &[0.6f64, 0.8], &mut ctx)?; + let output: ExtensionArray = L2Normalize::try_new(input.clone())? + .into_array() + .execute(&mut ctx)?; + + vortex_array::assert_arrays_eq!(output.into_array(), input, &mut ctx); + Ok(()) + } + + #[test] + fn serde_round_trip() -> VortexResult<()> { + let child = vector_array(2, &[3.0f64, 4.0])?; + let original = L2Normalize::try_new(child.clone())?.into_array(); + let plugin = ScalarFnArrayPlugin::new(L2Normalize); + let metadata = plugin + .serialize(&original, &SESSION)? + .vortex_expect("L2Normalize serialization must produce metadata"); + let recovered = plugin.deserialize( + original.dtype(), + original.len(), + &metadata, + &[], + &[child], + &SESSION, + )?; + + assert!(ExactScalarFn::::try_match(&recovered).is_some()); + assert_eq!(recovered.dtype(), original.dtype()); + Ok(()) + } + + #[test] + fn rejects_fixed_shape_tensor() -> VortexResult<()> { + let input = tensor_array(&[2], &[1.0f64, 0.0])?; + + assert!(L2Normalize::try_new(input).is_err()); + Ok(()) + } + + #[test] + fn return_dtype_is_nullable_for_ordinary_vectors() -> VortexResult<()> { + let input = vector_array(2, &[1.0f64, 0.0])?; + let output = L2Normalize::try_new(input)?.into_array(); + + assert!(output.dtype().is_nullable()); + assert!(output.dtype().as_extension().is::()); + assert!(!output.dtype().as_extension().is::()); + Ok(()) + } +} diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index e87491eb3b7..e892766a4f4 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -4,7 +4,7 @@ //! Scalar function expressions defined on tensor and tensor-like extension types. //! //! Each child module owns one expression. [`NormMode`] defines whether norm-based expressions -//! measure physical coordinates or trust normalized encoding evidence. +//! measure physical coordinates or trust normalized dtype and encoding claims. use std::fmt::Display; use std::fmt::Formatter; @@ -16,20 +16,22 @@ use vortex_error::vortex_err; pub mod cosine_similarity; pub mod inner_product; pub mod l2_norm; +pub mod l2_normalize; -/// Controls whether norm-based functions may trust [`Normalized`] encoding evidence. +/// Controls whether norm-based functions may trust normalized-value claims. /// -/// This policy belongs to the scalar function rather than the logical tensor dtype. It only -/// changes execution for [`Normalized`]-encoded inputs; ordinary tensors use their physical -/// coordinates in both modes. +/// [`Normalized`] encodings claim that their direction child is normalized, while [`UnitVector`] +/// dtypes claim that each non-null value is unit length. Ordinary tensors carry neither claim and +/// use their physical coordinates in both modes. /// /// [`Normalized`]: crate::encodings::normalized::Normalized +/// [`UnitVector`]: crate::unit_vector::UnitVector #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum NormMode { /// Compute physical direction norms instead of assuming that they are exactly one. Exact, - /// Trust each `Normalized` direction as unit length and omit its norm computation. + /// Trust normalized encoding and dtype claims and omit their norm computation. /// /// Checked arrays satisfy the encoding's documented tolerance. Unchecked or lossy arrays do /// not carry an error bound, so this mode can produce values outside the mathematical range. diff --git a/vortex-tensor/src/types/mod.rs b/vortex-tensor/src/types/mod.rs index 47dcabdb36d..f29d7a57074 100644 --- a/vortex-tensor/src/types/mod.rs +++ b/vortex-tensor/src/types/mod.rs @@ -2,6 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Internal homes for tensor extension types. +//! +//! Each child module owns one extension dtype, its validation, and its interchange support. The +//! crate root re-exports these modules as [`fixed_shape_tensor`], [`unit_vector`], and [`vector`]. 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..19e312ccc0f --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/arrow.rs @@ -0,0 +1,241 @@ +// 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-vector 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::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(|dtype| dtype.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() || *list_size <= 0 { + 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())?; + // SAFETY: The Arrow extension marker is the producer's UnitVector claim. Exact operations + // still measure the imported coordinates; approximate operations explicitly trust it. + let imported = unsafe { UnitVector::new_unchecked(storage)? }; + + Ok(ArrowImport::Imported(imported)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::ArrayRef as ArrowArrayRef; + use arrow_array::FixedSizeListArray as ArrowFixedSizeListArray; + use arrow_array::Float32Array; + use arrow_schema::DataType; + use arrow_schema::Field; + use vortex_array::EmptyMetadata; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_arrow::ArrowSession; + use vortex_error::VortexResult; + + use super::ARROW_UNIT_VECTOR_EXTENSION_NAME; + use crate::types::unit_vector::UnitVector; + + 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..b7914db5091 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/mod.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Unit-vector extension type for fixed-length float vectors. +//! +//! A [`UnitVector`] has the same storage layout as [`Vector`], but each non-null row is finite, +//! nonzero, and has an L2 norm within [`unit_norm_tolerance`] of one. The refinement is approximate +//! because its coordinates use finite-precision floating-point values. +//! +//! Use [`UnitVector::try_new_unit_vector_array`] when the storage values are not already trusted. +//! [`UnitVector::new_unchecked`] preserves claims from trusted producers and interchange metadata. +//! Norm-based operations only replace the physical norm with one when the caller selects +//! [`NormMode::AssumeNormalized`]. +//! +//! [`NormMode::AssumeNormalized`]: crate::scalar_fns::NormMode::AssumeNormalized +//! [`Vector`]: crate::vector::Vector +//! [`unit_norm_tolerance`]: crate::unit_norm_tolerance + +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; + +mod arrow; +pub use arrow::ARROW_UNIT_VECTOR_EXTENSION_NAME; + +mod matcher; +pub use matcher::AnyUnitVector; + +mod validate; + +mod vtable; + +/// A fixed-length float vector whose non-null rows are approximately unit length. +#[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 a non-null row is zero, + /// non-finite, or outside [`unit_norm_tolerance`](crate::unit_norm_tolerance) of unit length. + pub fn try_new_unit_vector_array( + storage: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // SAFETY: `validate_unit_vector_rows` validates every non-null row before this array is + // returned. + let array = unsafe { Self::new_unchecked(storage)? }; + validate::validate_unit_vector_rows(&array, ctx)?; + + Ok(array) + } + + /// Constructs a [`UnitVector`] array without validating its row values. + /// + /// # Safety + /// + /// Every non-null row must be finite, nonzero, and have an L2 norm within + /// [`unit_norm_tolerance`](crate::unit_norm_tolerance) of one. Violating this contract can + /// produce incorrect results from operations using [`NormMode::AssumeNormalized`], but it + /// cannot cause memory unsafety. + /// + /// [`NormMode::AssumeNormalized`]: crate::scalar_fns::NormMode::AssumeNormalized + 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..f7ffde72a98 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/tests.rs @@ -0,0 +1,110 @@ +// 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::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::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 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_rows() -> VortexResult<()> { + let array = unit_vector(2, &[0.6, 0.8, 1.0, 0.0])?; + + assert!(array.dtype().as_extension().is::()); + Ok(()) +} + +#[test] +fn checked_constructor_rejects_zero_row() { + assert!(unit_vector(2, &[0.0, 0.0]).is_err()); +} + +#[test] +fn checked_constructor_rejects_non_unit_row() { + assert!(unit_vector(2, &[3.0, 4.0]).is_err()); +} + +#[test] +fn checked_constructor_rejects_non_finite_row() { + assert!(unit_vector(2, &[f64::NAN, 0.0]).is_err()); +} + +#[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 scalar_constructor_rejects_zero_value() -> VortexResult<()> { + let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + let storage = Scalar::fixed_size_list( + element_dtype, + vec![ + Scalar::primitive(0.0f64, Nullability::NonNullable), + Scalar::primitive(0.0f64, Nullability::NonNullable), + ], + Nullability::NonNullable, + ); + + assert!(Scalar::try_new(unit_dtype(PType::F64, 2)?, storage.into_value()).is_err()); + Ok(()) +} + +#[test] +fn dtype_rejects_zero_dimensions() { + assert!(ExtDType::::try_new(EmptyMetadata, storage_dtype(PType::F32, 0)).is_err()); +} + +#[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); +} diff --git a/vortex-tensor/src/types/unit_vector/validate.rs b/vortex-tensor/src/types/unit_vector/validate.rs new file mode 100644 index 00000000000..fd2a9a9d323 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/validate.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::ToPrimitive; +use num_traits::Zero; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::match_each_float_ptype; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::types::unit_vector::AnyUnitVector; +use crate::utils::extract_flat_elements; +use crate::utils::unit_norm_tolerance; + +pub(super) fn validate_unit_vector_rows( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let metadata = array.dtype().as_extension().metadata::(); + let row_count = array.len(); + if row_count == 0 { + return Ok(()); + } + + let array: ExtensionArray = array.clone().execute(ctx)?; + let validity = array.as_ref().validity()?; + let valid_rows = validity + .nullability() + .is_nullable() + .then(|| validity.execute_mask(row_count, ctx)) + .transpose()?; + let flat = extract_flat_elements(array.storage_array(), metadata.dimensions() as usize, ctx)?; + let tolerance = unit_norm_tolerance(metadata.element_ptype(), metadata.dimensions() as usize); + + match_each_float_ptype!(metadata.element_ptype(), |T| { + for row_idx in 0..row_count { + if valid_rows + .as_ref() + .is_some_and(|valid_rows| !valid_rows.value(row_idx)) + { + continue; + } + + let (sum_squares, is_zero) = flat.row::(row_idx).iter().fold( + (0.0f64, true), + |(sum_squares, is_zero), value| { + let value_f64 = ToPrimitive::to_f64(value) + .vortex_expect("UnitVector dtype validation established float elements"); + ( + sum_squares + value_f64 * value_f64, + is_zero && value.is_zero(), + ) + }, + ); + let norm = sum_squares.sqrt(); + + vortex_ensure!( + !is_zero && norm.is_finite() && (norm - 1.0).abs() <= tolerance, + "UnitVector row must be finite, nonzero, and have L2 norm within {tolerance:.6} \ + of 1.0, got row {row_idx} with norm {norm:.6}", + ); + } + }); + + 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..c9c801b9247 --- /dev/null +++ b/vortex-tensor/src/types/unit_vector/vtable.rs @@ -0,0 +1,90 @@ +// 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::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 validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + validate_vector_storage_dtype(ext_dtype.storage_dtype())?; + + let DType::FixedSizeList(_, dimensions, _) = ext_dtype.storage_dtype() else { + unreachable!("UnitVector storage validation established FixedSizeList storage") + }; + vortex_ensure!( + *dimensions > 0, + "UnitVector dimensions must be greater than zero, got {dimensions}", + ); + + Ok(()) + } + + fn unpack_native<'a>( + ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult> { + let elements = storage_value.as_list(); + let DType::FixedSizeList(element_dtype, dimensions, _) = ext_dtype.storage_dtype() else { + unreachable!("UnitVector dtype validation established FixedSizeList storage") + }; + let tolerance = unit_norm_tolerance(element_dtype.as_ptype(), *dimensions as usize); + + let (sum_squares, is_zero) = elements.iter().try_fold( + (0.0f64, true), + |(sum_squares, is_zero), element| -> VortexResult<_> { + let value = element + .as_ref() + .ok_or_else(|| { + vortex_err!("UnitVector scalar elements must be non-null, got 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_squares + value * value, is_zero && value == 0.0)) + }, + )?; + let norm = sum_squares.sqrt(); + + vortex_ensure!( + !is_zero && norm.is_finite() && (norm - 1.0).abs() <= tolerance, + "UnitVector scalar must be finite, nonzero, and have L2 norm within {tolerance:.6} \ + of 1.0, 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..f86e705d5dd 100644 --- a/vortex-tensor/src/types/vector/matcher.rs +++ b/vortex-tensor/src/types/vector/matcher.rs @@ -10,20 +10,13 @@ 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 [`Vector`] and [`UnitVector`] dtypes. 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 a vector dtype's 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. @@ -39,28 +32,21 @@ 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, dimensions, _) = ext_dtype.storage_dtype() else { + vortex_panic!("vector dtype must have FixedSizeList storage") + }; - Some(vector_metadata) - } + VectorMatcherMetadata::try_new(element_dtype.as_ptype(), *dimensions) + .vortex_expect("vector dtype validation established float elements") } impl VectorMatcherMetadata { @@ -70,7 +56,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 a float, got {element_ptype}", + ); Ok(Self { element_ptype, @@ -100,9 +89,11 @@ mod tests { use vortex_array::dtype::extension::ExtDType; use vortex_error::VortexResult; - use super::*; + use super::AnyVector; use crate::types::fixed_shape_tensor::FixedShapeTensor; use crate::types::fixed_shape_tensor::FixedShapeTensorMetadata; + use crate::types::unit_vector::UnitVector; + use crate::types::vector::Vector; fn vector_storage_dtype(element_ptype: PType, dimensions: u32) -> DType { DType::FixedSizeList( @@ -124,6 +115,18 @@ mod tests { Ok(()) } + #[test] + fn matches_unit_vector_dtype_metadata() -> VortexResult<()> { + let ext_dtype = + ExtDType::::try_new(EmptyMetadata, vector_storage_dtype(PType::F32, 256))? + .erased(); + + let metadata = ext_dtype.metadata::(); + assert_eq!(metadata.element_ptype(), PType::F32); + assert_eq!(metadata.dimensions(), 256); + Ok(()) + } + #[test] fn does_not_match_fixed_shape_tensor() -> VortexResult<()> { let ext_dtype = ExtDType::::try_new( diff --git a/vortex-tensor/src/types/vector/mod.rs b/vortex-tensor/src/types/vector/mod.rs index af424b9cc41..c7f62e5f754 100644 --- a/vortex-tensor/src/types/vector/mod.rs +++ b/vortex-tensor/src/types/vector/mod.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Vector extension type for fixed-length float vectors (e.g., embeddings). +//! Vector extension type for fixed-length float vectors, such as embeddings. +//! +//! [`Vector`] establishes the fixed-size, non-nullable-float element layout. [`AnyVector`] matches +//! both ordinary vectors and the [`UnitVector`](crate::unit_vector::UnitVector) refinement. use vortex_array::ArrayRef; use vortex_array::EmptyMetadata; @@ -17,6 +20,16 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +mod arrow; +pub use arrow::ARROW_VECTOR_EXTENSION_NAME; + +mod matcher; +pub use matcher::AnyVector; +pub use matcher::VectorMatcherMetadata; +pub(crate) use matcher::match_vector_storage; + +mod vtable; + /// Validates that `storage` is a valid storage dtype for a [`Vector`]. /// /// The storage must be a `FixedSizeList` with non-nullable float @@ -32,18 +45,18 @@ pub(crate) fn validate_vector_storage_dtype(storage: &DType) -> VortexResult<()> ); vortex_ensure!( !element_dtype.is_nullable(), - "Vector element dtype must be non-nullable" + "Vector element dtype must be non-nullable, got {element_dtype}", ); Ok(()) } -/// The Vector extension type. +/// A fixed-length vector with non-nullable floating-point elements. #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] pub struct Vector; impl Vector { - /// Helper function for creating a new [`Vector`] [`ExtensionArray`]. + /// Constructs a [`Vector`] [`ExtensionArray`] from its storage array. /// /// # Errors /// @@ -53,8 +66,8 @@ impl Vector { .map(|ext| ext.into_array()) } - /// Helper function to build a [`Vector`] [`ExtensionArray`] whose storage is a - /// [`ConstantArray`], broadcasting a single vector `elements` across `len` rows. + /// Constructs a [`Vector`] [`ExtensionArray`] whose [`ConstantArray`] storage broadcasts + /// `elements` across `len` rows. /// /// # Errors /// @@ -66,19 +79,10 @@ impl Vector { let element_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable); let children: Vec = elements .iter() - .map(|&v| Scalar::primitive(v, Nullability::NonNullable)) + .map(|&value| Scalar::primitive(value, Nullability::NonNullable)) .collect(); let storage_scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); Self::try_new_vector_array(ConstantArray::new(storage_scalar, len).into_array()) } } - -mod arrow; -mod matcher; - -pub use arrow::ARROW_VECTOR_EXTENSION_NAME; -pub use matcher::AnyVector; -pub use matcher::VectorMatcherMetadata; - -mod vtable; diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 553f3f98d94..1650e7dc407 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -35,14 +35,14 @@ 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 so the refinement remains useful for large vectors. /// /// Reference: Croci, Fasi, Higham, Mary, Mikaitis (2022). "Stochastic rounding: implementation, /// error analysis and applications." Royal Society Open Science, 9: 211631, §6.1 "Probabilistic @@ -57,7 +57,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. @@ -113,11 +118,21 @@ pub fn validate_binary_tensor_float_inputs<'a>( lhs: &'a DType, rhs: &DType, ) -> VortexResult> { + let lhs_match = validate_tensor_float_input(lhs)?; + if lhs.eq_ignore_nullability(rhs) { + return Ok(lhs_match); + } + + let rhs_match = validate_tensor_float_input(rhs)?; vortex_ensure!( - lhs.eq_ignore_nullability(rhs), - "binary tensor expression expects inputs to have the same dtype, got {lhs} and {rhs}" + matches!( + (lhs_match, rhs_match), + (TensorMatch::Vector(lhs), TensorMatch::Vector(rhs)) if lhs == rhs + ), + "binary tensor expression expects compatible inputs, 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. @@ -300,6 +315,7 @@ pub mod test_helpers { use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; @@ -313,6 +329,7 @@ pub mod test_helpers { use crate::encodings::normalized::Normalized; use crate::types::fixed_shape_tensor::FixedShapeTensor; use crate::types::fixed_shape_tensor::FixedShapeTensorMetadata; + use crate::types::unit_vector::UnitVector; use crate::types::vector::Vector; /// Builds a `FixedSizeList` storage array from flat `elements`. The row count is @@ -351,6 +368,18 @@ pub mod test_helpers { Vector::try_new_vector_array(flat_fsl(elements, dim)) } + /// Builds and validates a [`UnitVector`] extension array from flat `elements`. + pub fn unit_vector_array( + dim: u32, + elements: &[T], + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let vector = vector_array(dim, elements)?; + let vector: ExtensionArray = vector.execute(ctx)?; + + UnitVector::try_new_unit_vector_array(vector.storage_array().clone(), ctx) + } + /// Builds a [`FixedShapeTensor`] extension array whose storage is a [`ConstantArray`], /// representing a single query tensor broadcast to `len` rows. pub fn constant_tensor_array>( diff --git a/vortex/src/editions/preview/v2026_04.rs b/vortex/src/editions/preview/v2026_04.rs index 4d1a5ede4b8..2087637e9ec 100644 --- a/vortex/src/editions/preview/v2026_04.rs +++ b/vortex/src/editions/preview/v2026_04.rs @@ -23,7 +23,9 @@ 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"), EditionMember::dtype(&"vortex.tensor.fixed_shape_tensor"), + EditionMember::dtype(&"vortex.tensor.unit_vector"), EditionMember::dtype(&"vortex.tensor.vector"), ], }; diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 4ba621baf7a..053e06b1f69 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -6,16 +6,24 @@ use std::sync::Arc; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +#[cfg(feature = "unstable_encodings")] +use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; +#[cfg(feature = "unstable_encodings")] +use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +#[cfg(feature = "unstable_encodings")] +use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::field_path; use vortex_array::session::ArraySessionExt; use vortex_array::stream::ArrayStreamExt; +#[cfg(feature = "unstable_encodings")] +use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_buffer::ByteBufferMut; use vortex_edition::ComponentKind; @@ -41,6 +49,12 @@ use vortex_layout::session::LayoutSession; use vortex_sequence::Sequence; use vortex_session::VortexSession; use vortex_session::registry::Id; +#[cfg(feature = "unstable_encodings")] +use vortex_tensor::scalar_fns::l2_normalize::L2Normalize; +#[cfg(feature = "unstable_encodings")] +use vortex_tensor::unit_vector::UnitVector; +#[cfg(feature = "unstable_encodings")] +use vortex_tensor::vector::Vector; use vortex_utils::aliases::hash_set::HashSet; use super::CORE_2025_05_0; @@ -52,6 +66,8 @@ use super::DEFAULT_CORE_EDITION; use super::DEFAULT_PREVIEW_EDITION; use super::EDITION_DECLARATIONS; use super::PREVIEW_2026_06_0; +#[cfg(feature = "unstable_encodings")] +use crate::VortexSessionDefault; fn session() -> Result { let session = EditionSession::empty(); @@ -499,6 +515,34 @@ async fn write_with(session: &VortexSession, array: ArrayRef) -> VortexResult VortexResult<()> { + let session = VortexSession::default(); + session + .arrays() + .register(ScalarFnArrayPlugin::new(L2Normalize)); + + let elements = PrimitiveArray::from_iter([0.6f64, 0.8]).into_array(); + let storage = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 1)?.into_array(); + let mut ctx = session.create_execution_ctx(); + let unit = UnitVector::try_new_unit_vector_array(storage, &mut ctx)?; + write_with(&session, unit).await?; + + let elements = PrimitiveArray::from_iter([3.0f64, 4.0]).into_array(); + let storage = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 1)?.into_array(); + let vector = Vector::try_new_vector_array(storage)?; + let lazy_normalized = L2Normalize::try_new(vector)?.into_array(); + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut buffer, lazy_normalized.to_array_stream()) + .await?; + + Ok(()) +} + /// The layout encodings a written file actually contains, depth first. fn written_layout_ids(session: &VortexSession, buffer: ByteBufferMut) -> VortexResult> { let file = session.open_options().open_buffer(buffer)?;