add normalized vector type

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

Author: connortsui20

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

@@ -78,3 +78,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

@@ -14,8 +14,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;
@@ -26,10 +27,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(
@@ -38,6 +40,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)
     }
 
@@ -52,3 +55,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/centroids.rs

@@ -36,7 +36,7 @@ static CENTROID_CACHE: LazyLock<DashMap<(u32, u8), Buffer<f32>>> = LazyLock::new
 /// Returns `2^bit_width` centroids sorted in ascending order, representing optimal scalar
 /// quantization levels for the coordinate distribution after random rotation in
 /// `dimension`-dimensional space.
-pub fn get_centroids(dimension: u32, bit_width: u8) -> VortexResult<Buffer<f32>> {
+pub fn compute_or_get_centroids(dimension: u32, bit_width: u8) -> VortexResult<Buffer<f32>> {
     vortex_ensure!(
         (1..=MAX_BIT_WIDTH).contains(&bit_width),
         "TurboQuant bit_width must be 1-{}, got {bit_width}",
@@ -239,7 +239,7 @@ mod tests {
         #[case] bits: u8,
         #[case] expected: usize,
     ) -> VortexResult<()> {
-        let centroids = get_centroids(dim, bits)?;
+        let centroids = compute_or_get_centroids(dim, bits)?;
         assert_eq!(centroids.len(), expected);
         Ok(())
     }
@@ -251,7 +251,7 @@ mod tests {
     #[case(128, 4)]
     #[case(768, 2)]
     fn centroids_are_sorted(#[case] dim: u32, #[case] bits: u8) -> VortexResult<()> {
-        let centroids = get_centroids(dim, bits)?;
+        let centroids = compute_or_get_centroids(dim, bits)?;
         for window in centroids.windows(2) {
             assert!(
                 window[0] < window[1],
@@ -268,7 +268,7 @@ mod tests {
     #[case(256, 2)]
     #[case(768, 2)]
     fn centroids_are_symmetric(#[case] dim: u32, #[case] bits: u8) -> VortexResult<()> {
-        let centroids = get_centroids(dim, bits)?;
+        let centroids = compute_or_get_centroids(dim, bits)?;
         let count = centroids.len();
         for idx in 0..count / 2 {
             let diff = (centroids[idx] + centroids[count - 1 - idx]).abs();
@@ -287,7 +287,7 @@ mod tests {
     #[case(128, 1)]
     #[case(128, 4)]
     fn centroids_within_bounds(#[case] dim: u32, #[case] bits: u8) -> VortexResult<()> {
-        let centroids = get_centroids(dim, bits)?;
+        let centroids = compute_or_get_centroids(dim, bits)?;
         for &val in centroids.iter() {
             assert!(
                 (-1.0..=1.0).contains(&val),
@@ -299,15 +299,15 @@ mod tests {
 
     #[test]
     fn centroids_cached() -> VortexResult<()> {
-        let c1 = get_centroids(128, 2)?;
-        let c2 = get_centroids(128, 2)?;
+        let c1 = compute_or_get_centroids(128, 2)?;
+        let c2 = compute_or_get_centroids(128, 2)?;
         assert_eq!(c1, c2);
         Ok(())
     }
 
     #[test]
     fn find_nearest_basic() -> VortexResult<()> {
-        let centroids = get_centroids(128, 2)?;
+        let centroids = compute_or_get_centroids(128, 2)?;
         let boundaries = compute_centroid_boundaries(&centroids);
         assert_eq!(find_nearest_centroid(-1.0, &boundaries), 0);
 
@@ -324,9 +324,9 @@ mod tests {
 
     #[test]
     fn rejects_invalid_params() {
-        assert!(get_centroids(128, 0).is_err());
-        assert!(get_centroids(128, 9).is_err());
-        assert!(get_centroids(1, 2).is_err());
-        assert!(get_centroids(127, 2).is_err());
+        assert!(compute_or_get_centroids(128, 0).is_err());
+        assert!(compute_or_get_centroids(128, 9).is_err());
+        assert!(compute_or_get_centroids(1, 2).is_err());
+        assert!(compute_or_get_centroids(127, 2).is_err());
     }
 }