Add UnitVector tensor refinement

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::<AnyTensor>()
-            .is_some_and(|tensor| tensor.element_ptype().is_float())
+        !ext.ext_dtype().is::<UnitVector>()
+            && ext
+                .ext_dtype()
+                .metadata_opt::<AnyTensor>()
+                .is_some_and(|tensor| tensor.element_ptype().is_float())
     }
 
     fn produced_encodings(&self) -> Vec<ArrayId> {
@@ -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<NormalizedArray> {
+    vortex_ensure!(
+        !input
+            .dtype()
+            .as_extension_opt()
+            .is_some_and(|dtype| dtype.is::<UnitVector>()),
+        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<Normal
 /// 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.
+/// [`UnitVector`] input also returns `Ok(None)`. Any other matching input produces constant
+/// normalized and norms children.
+///
+/// [`UnitVector`]: crate::unit_vector::UnitVector
 pub(crate) fn try_build_constant_normalized(
     input: &ArrayRef,
     len: usize,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<Option<NormalizedArray>> {
+    if input
+        .dtype()
+        .as_extension_opt()
+        .is_some_and(|dtype| dtype.is::<UnitVector>())
+    {
+        return Ok(None);
+    }
+
     let Some(ext) = input.as_opt::<Extension>() else {
         return Ok(None);
     };

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"))]

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));
     }
 }
 

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::<AnyVector>() {
             return Some(TensorMatch::Vector(metadata));
         }

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<f64>`.
@@ -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(