add normalized vector type

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-array/src/arrays/extension/array.rs

@@ -6,6 +6,7 @@ use std::fmt::Formatter;
 
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 
 use crate::ArrayRef;
 use crate::array::Array;
@@ -90,10 +91,18 @@ impl ExtensionData {
     pub fn try_new(ext_dtype: ExtDTypeRef, storage_dtype: &DType) -> VortexResult<Self> {
         // TODO(connor): Replace these statements once we add `validate_storage_array`.
         // ext_dtype.validate_storage_array(&storage_array)?;
-        assert_eq!(
+        //
+        // The storage array's outer nullability is allowed to differ from the extension's declared
+        // storage outer nullability. Nested storage nullability must still match exactly.
+        vortex_ensure!(
+            storage_dtypes_match_ignoring_outer_nullability(
+                ext_dtype.storage_dtype(),
+                storage_dtype
+            ),
+            "ExtensionArray: storage_dtype must match storage array DType (ignoring outer \
+             nullability only), got extension storage {} and array storage {}",
             ext_dtype.storage_dtype(),
             storage_dtype,
-            "ExtensionArray: storage_dtype must match storage array DType",
         );
 
         // SAFETY: we validate that the inputs are valid above.
@@ -113,10 +122,18 @@ impl ExtensionData {
         // ext_dtype
         //     .validate_storage_array(&storage_array)
         //     .vortex_expect("[Debug Assertion]: Invalid storage array for `ExtensionArray`");
-        debug_assert_eq!(
+        //
+        // Match the contract of [`Self::try_new`]: the storage dtype must match the extension's
+        // declared storage dtype ignoring only outer nullability.
+        debug_assert!(
+            storage_dtypes_match_ignoring_outer_nullability(
+                ext_dtype.storage_dtype(),
+                storage_dtype
+            ),
+            "ExtensionArray: storage_dtype must match storage array DType (ignoring outer \
+             nullability only), got extension storage {} and array storage {}",
             ext_dtype.storage_dtype(),
             storage_dtype,
-            "ExtensionArray: storage_dtype must match storage array DType",
         );
 
         Self { ext_dtype }
@@ -128,6 +145,13 @@ impl ExtensionData {
     }
 }
 
+fn storage_dtypes_match_ignoring_outer_nullability(
+    ext_storage_dtype: &DType,
+    storage_dtype: &DType,
+) -> bool {
+    ext_storage_dtype.with_nullability(storage_dtype.nullability()) == *storage_dtype
+}
+
 pub trait ExtensionArrayExt: TypedArrayRef<Extension> {
     fn storage_array(&self) -> &ArrayRef {
         self.as_ref().slots()[STORAGE_SLOT]
@@ -179,3 +203,93 @@ impl Array<Extension> {
         Self::try_new(ext_dtype, storage_array)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use vortex_buffer::Buffer;
+
+    use super::*;
+    use crate::IntoArray;
+    use crate::arrays::ExtensionArray;
+    use crate::arrays::FixedSizeListArray;
+    use crate::arrays::PrimitiveArray;
+    use crate::dtype::Nullability;
+    use crate::dtype::PType;
+    use crate::dtype::extension::ExtId;
+    use crate::extension::EmptyMetadata;
+    use crate::scalar::ScalarValue;
+    use crate::validity::Validity;
+
+    #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
+    struct TestExt;
+
+    impl ExtVTable for TestExt {
+        type Metadata = EmptyMetadata;
+        type NativeValue<'a> = &'a ScalarValue;
+
+        fn id(&self) -> ExtId {
+            ExtId::new("vortex.test.extension")
+        }
+
+        fn serialize_metadata(&self, _metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
+            Ok(Vec::new())
+        }
+
+        fn deserialize_metadata(&self, _metadata: &[u8]) -> VortexResult<Self::Metadata> {
+            Ok(EmptyMetadata)
+        }
+
+        fn validate_dtype(_ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
+            Ok(())
+        }
+
+        fn unpack_native<'a>(
+            _ext_dtype: &'a ExtDType<Self>,
+            storage_value: &'a ScalarValue,
+        ) -> VortexResult<Self::NativeValue<'a>> {
+            Ok(storage_value)
+        }
+    }
+
+    fn fsl_dtype(element_nullability: Nullability, list_nullability: Nullability) -> DType {
+        DType::FixedSizeList(
+            Arc::new(DType::Primitive(PType::F32, element_nullability)),
+            2,
+            list_nullability,
+        )
+    }
+
+    #[test]
+    fn extension_storage_allows_outer_nullability_mismatch() -> VortexResult<()> {
+        let ext_dtype = ExtDType::<TestExt>::try_new(
+            EmptyMetadata,
+            fsl_dtype(Nullability::NonNullable, Nullability::NonNullable),
+        )?
+        .erased();
+
+        let elements = PrimitiveArray::from_iter([1.0f32, 0.0]).into_array();
+        let storage = FixedSizeListArray::try_new(elements, 2, Validity::AllValid, 1)?.into_array();
+
+        ExtensionArray::try_new(ext_dtype, storage)?;
+        Ok(())
+    }
+
+    #[test]
+    fn extension_storage_rejects_nested_nullability_mismatch() -> VortexResult<()> {
+        let ext_dtype = ExtDType::<TestExt>::try_new(
+            EmptyMetadata,
+            fsl_dtype(Nullability::NonNullable, Nullability::NonNullable),
+        )?
+        .erased();
+
+        let elements =
+            PrimitiveArray::new(Buffer::copy_from([1.0f32, 0.0]), Validity::AllValid).into_array();
+        let storage =
+            FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 1)?.into_array();
+
+        assert!(ExtensionArray::try_new(ext_dtype, storage).is_err());
+        Ok(())
+    }
+}

vortex-tensor/public-api.lock

@@ -250,6 +250,64 @@ pub type vortex_tensor::matcher::AnyTensor::Match<'a> = vortex_tensor::matcher::
 
 pub fn vortex_tensor::matcher::AnyTensor::try_match<'a>(ext_dtype: &'a vortex_array::dtype::extension::erased::ExtDTypeRef) -> core::option::Option<Self::Match>
 
+pub mod vortex_tensor::normalized_vector
+
+pub struct vortex_tensor::normalized_vector::AnyNormalizedVector
+
+impl vortex_array::dtype::extension::matcher::Matcher for vortex_tensor::normalized_vector::AnyNormalizedVector
+
+pub type vortex_tensor::normalized_vector::AnyNormalizedVector::Match<'a> = vortex_tensor::vector::VectorMatcherMetadata
+
+pub fn vortex_tensor::normalized_vector::AnyNormalizedVector::try_match<'a>(ext_dtype: &'a vortex_array::dtype::extension::erased::ExtDTypeRef) -> core::option::Option<Self::Match>
+
+pub struct vortex_tensor::normalized_vector::NormalizedVector
+
+impl vortex_tensor::normalized_vector::NormalizedVector
+
+pub unsafe fn vortex_tensor::normalized_vector::NormalizedVector::new_unchecked(storage: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::try_new(storage: vortex_array::array::erased::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
+impl core::clone::Clone for vortex_tensor::normalized_vector::NormalizedVector
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::clone(&self) -> vortex_tensor::normalized_vector::NormalizedVector
+
+impl core::cmp::Eq for vortex_tensor::normalized_vector::NormalizedVector
+
+impl core::cmp::PartialEq for vortex_tensor::normalized_vector::NormalizedVector
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::eq(&self, other: &vortex_tensor::normalized_vector::NormalizedVector) -> bool
+
+impl core::default::Default for vortex_tensor::normalized_vector::NormalizedVector
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::default() -> vortex_tensor::normalized_vector::NormalizedVector
+
+impl core::fmt::Debug for vortex_tensor::normalized_vector::NormalizedVector
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::hash::Hash for vortex_tensor::normalized_vector::NormalizedVector
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::hash<__H: core::hash::Hasher>(&self, state: &mut __H)
+
+impl core::marker::StructuralPartialEq for vortex_tensor::normalized_vector::NormalizedVector
+
+impl vortex_array::dtype::extension::vtable::ExtVTable for vortex_tensor::normalized_vector::NormalizedVector
+
+pub type vortex_tensor::normalized_vector::NormalizedVector::Metadata = vortex_array::extension::EmptyMetadata
+
+pub type vortex_tensor::normalized_vector::NormalizedVector::NativeValue<'a> = &'a vortex_array::scalar::scalar_value::ScalarValue
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::deserialize_metadata(&self, _metadata: &[u8]) -> vortex_error::VortexResult<Self::Metadata>
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::id(&self) -> vortex_array::dtype::extension::ExtId
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::serialize_metadata(&self, _metadata: &Self::Metadata) -> vortex_error::VortexResult<alloc::vec::Vec<u8>>
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::unpack_native<'a>(_ext_dtype: &'a vortex_array::dtype::extension::typed::ExtDType<Self>, storage_value: &'a vortex_array::scalar::scalar_value::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
+
+pub fn vortex_tensor::normalized_vector::NormalizedVector::validate_dtype(ext_dtype: &vortex_array::dtype::extension::typed::ExtDType<Self>) -> vortex_error::VortexResult<()>
+
 pub mod vortex_tensor::scalar_fns
 
 pub mod vortex_tensor::scalar_fns::cosine_similarity
@@ -382,8 +440,6 @@ pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::validity(&self, _options:
 
 pub fn vortex_tensor::scalar_fns::l2_denorm::normalize_as_l2_denorm(input: vortex_array::array::erased::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::vtable::ScalarFnArray>
 
-pub fn vortex_tensor::scalar_fns::l2_denorm::validate_l2_normalized_rows_against_norms(normalized: &vortex_array::array::erased::ArrayRef, norms: core::option::Option<&vortex_array::array::erased::ArrayRef>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<()>
-
 pub mod vortex_tensor::scalar_fns::l2_norm
 
 pub struct vortex_tensor::scalar_fns::l2_norm::L2Norm
@@ -574,7 +630,9 @@ pub fn vortex_tensor::vector::VectorMatcherMetadata::dimensions(&self) -> u32
 
 pub fn vortex_tensor::vector::VectorMatcherMetadata::element_ptype(&self) -> vortex_array::dtype::ptype::PType
 
-pub fn vortex_tensor::vector::VectorMatcherMetadata::try_new(element_ptype: vortex_array::dtype::ptype::PType, dimensions: u32) -> vortex_error::VortexResult<Self>
+pub fn vortex_tensor::vector::VectorMatcherMetadata::is_normalized(&self) -> bool
+
+pub fn vortex_tensor::vector::VectorMatcherMetadata::try_new(element_ptype: vortex_array::dtype::ptype::PType, dimensions: u32, is_normalized: bool) -> vortex_error::VortexResult<Self>
 
 impl core::clone::Clone for vortex_tensor::vector::VectorMatcherMetadata
 

vortex-tensor/src/encodings/l2_denorm.rs

@@ -13,8 +13,9 @@ use vortex_compressor::scheme::Scheme;
 use vortex_compressor::stats::ArrayAndStats;
 use vortex_error::VortexResult;
 
-use crate::matcher::AnyTensor;
+use crate::normalized_vector::AnyNormalizedVector;
 use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm;
+use crate::types::vector::AnyVector;
 
 #[derive(Debug)]
 pub struct L2DenormScheme;
@@ -25,10 +26,11 @@ impl Scheme for L2DenormScheme {
     }
 
     fn matches(&self, canonical: &Canonical) -> bool {
-        matches!(
-            canonical,
-            Canonical::Extension(ext) if ext.ext_dtype().is::<AnyTensor>()
-        )
+        let Canonical::Extension(ext) = canonical else {
+            return false;
+        };
+
+        ext.ext_dtype().is::<AnyVector>() && !ext.ext_dtype().is::<AnyNormalizedVector>()
     }
 
     fn expected_compression_ratio(
@@ -37,6 +39,7 @@ impl Scheme for L2DenormScheme {
         _compress_ctx: CompressorContext,
         _exec_ctx: &mut ExecutionCtx,
     ) -> CompressionEstimate {
+        // We almost always want to pre-normalize our data if the vector is not already normalized.
         CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse)
     }
 
@@ -51,3 +54,62 @@ impl Scheme for L2DenormScheme {
         Ok(l2_denorm.into_array())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use vortex_array::Canonical;
+    use vortex_array::IntoArray;
+    use vortex_array::arrays::ExtensionArray;
+    use vortex_array::arrays::FixedSizeListArray;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+    use vortex_array::dtype::extension::ExtDType;
+    use vortex_array::extension::EmptyMetadata;
+    use vortex_array::validity::Validity;
+    use vortex_compressor::scheme::Scheme;
+    use vortex_error::VortexResult;
+
+    use super::L2DenormScheme;
+    use crate::types::fixed_shape::FixedShapeTensor;
+    use crate::types::fixed_shape::FixedShapeTensorMetadata;
+    use crate::types::vector::Vector;
+
+    fn fsl_storage(elements: &[f32], list_size: u32) -> VortexResult<FixedSizeListArray> {
+        let len = elements.len() / list_size as usize;
+        let elements = PrimitiveArray::from_iter(elements.iter().copied()).into_array();
+        FixedSizeListArray::try_new(elements, list_size, Validity::NonNullable, len)
+    }
+
+    #[test]
+    fn matches_vector() -> VortexResult<()> {
+        let fsl = fsl_storage(&[1.0, 0.0], 2)?;
+        let ext_dtype = ExtDType::<Vector>::try_new(EmptyMetadata, fsl.dtype().clone())?.erased();
+        let canonical = Canonical::Extension(ExtensionArray::new(ext_dtype, fsl.into_array()));
+
+        assert!(L2DenormScheme.matches(&canonical));
+        Ok(())
+    }
+
+    #[test]
+    fn rejects_fixed_shape_tensor() -> VortexResult<()> {
+        let fsl = fsl_storage(&[1.0, 0.0, 0.0, 1.0], 4)?;
+        let storage_dtype = DType::FixedSizeList(
+            Arc::new(DType::Primitive(PType::F32, Nullability::NonNullable)),
+            4,
+            Nullability::NonNullable,
+        );
+        let ext_dtype = ExtDType::<FixedShapeTensor>::try_new(
+            FixedShapeTensorMetadata::new(vec![2, 2]),
+            storage_dtype,
+        )?
+        .erased();
+        let canonical = Canonical::Extension(ExtensionArray::new(ext_dtype, fsl.into_array()));
+
+        assert!(!L2DenormScheme.matches(&canonical));
+        Ok(())
+    }
+}

vortex-tensor/src/encodings/turboquant/compress.rs

@@ -39,8 +39,8 @@ use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm;
 use crate::scalar_fns::sorf_transform::SorfMatrix;
 use crate::scalar_fns::sorf_transform::SorfOptions;
 use crate::scalar_fns::sorf_transform::SorfTransform;
+use crate::types::normalized_vector::NormalizedVector;
 use crate::types::vector::AnyVector;
-use crate::types::vector::Vector;
 use crate::utils::cast_to_f32;
 
 /// Configuration for TurboQuant encoding.
@@ -80,8 +80,8 @@ impl Default for TurboQuantConfig {
 ///
 /// # Errors
 ///
-/// Returns an error if `input` is not a tensor-like extension array, if normalization fails, or
-/// if [`turboquant_encode_unchecked`] rejects the input shape.
+/// Returns an error if `input` is not a tensor-like extension array, if normalization fails, or if
+/// [`turboquant_encode_unchecked`] rejects the input shape.
 pub fn turboquant_encode(
     input: ArrayRef,
     config: &TurboQuantConfig,
@@ -101,7 +101,7 @@ pub fn turboquant_encode(
     let tq = unsafe { turboquant_encode_unchecked(normalized_ext, config, ctx) }?;
 
     // SAFETY: TurboQuant is a lossy approximation of the normalized child, so we intentionally
-    // bypass the strict normalized-row validation when reattaching the stored norms.
+    // bypass the strict normalized-row and zero-row validation when reattaching the stored norms.
     Ok(unsafe { L2Denorm::new_array_unchecked(tq, norms, num_rows) }?.into_array())
 }
 
@@ -156,7 +156,10 @@ pub unsafe fn turboquant_encode_unchecked(
             Validity::NonNullable,
             0,
         )?;
-        let empty_padded_vector = Vector::try_new_vector_array(empty_fsl.into_array())?;
+        // SAFETY: An empty FSL contains no rows, so the unit-norm-or-zero invariant holds
+        // vacuously.
+        let empty_padded_vector =
+            unsafe { NormalizedVector::new_unchecked(empty_fsl.into_array()) }?;
 
         let sorf_options = SorfOptions {
             seed,
@@ -172,7 +175,11 @@ pub unsafe fn turboquant_encode_unchecked(
     let core = turboquant_quantize_core(&fsl, seed, config.bit_width, config.num_rounds, ctx)?;
     let quantized_fsl =
         build_quantized_fsl(num_rows, core.all_indices, core.centroids, core.padded_dim)?;
-    let padded_vector = Vector::try_new_vector_array(quantized_fsl)?;
+    // SAFETY: TurboQuant is a lossy approximation of the already-unit-norm input. The
+    // quantized rows are approximately unit-norm by construction; downstream callers
+    // (notably the enclosing `L2Denorm` wrapper) treat the stored-norm + NormalizedVector
+    // claim as authoritative rather than decode-verified.
+    let padded_vector = unsafe { NormalizedVector::new_unchecked(quantized_fsl) }?;
 
     let sorf_options = SorfOptions {
         seed,