Pull out L2Denorm from TurboQuant

vortex-tensor/public-api.lock

@@ -16,7 +16,9 @@ pub const vortex_tensor::encodings::turboquant::TurboQuant::MAX_CENTROIDS: usize
 
 pub const vortex_tensor::encodings::turboquant::TurboQuant::MIN_DIMENSION: u32
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::try_new_array(dtype: vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_tensor::encodings::turboquant::TurboQuantArray>
+pub unsafe fn vortex_tensor::encodings::turboquant::TurboQuant::new_array_unchecked(dtype: vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_tensor::encodings::turboquant::TurboQuantArray
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::try_new_array(dtype: vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_tensor::encodings::turboquant::TurboQuantArray>
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::validate_dtype(dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<vortex_tensor::vector::VectorMatcherMetadata>
 
@@ -34,7 +36,7 @@ pub type vortex_tensor::encodings::turboquant::TurboQuant::ArrayData = vortex_te
 
 pub type vortex_tensor::encodings::turboquant::TurboQuant::OperationsVTable = vortex_tensor::encodings::turboquant::TurboQuant
 
-pub type vortex_tensor::encodings::turboquant::TurboQuant::ValidityVTable = vortex_array::array::vtable::validity::ValidityVTableFromChild
+pub type vortex_tensor::encodings::turboquant::TurboQuant::ValidityVTable = vortex_tensor::encodings::turboquant::TurboQuant
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::buffer(_array: vortex_array::array::view::ArrayView<'_, Self>, idx: usize) -> vortex_array::buffer::BufferHandle
 
@@ -62,9 +64,9 @@ impl vortex_array::array::vtable::operations::OperationsVTable<vortex_tensor::en
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuant::scalar_at(array: vortex_array::array::view::ArrayView<'_, vortex_tensor::encodings::turboquant::TurboQuant>, index: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
 
-impl vortex_array::array::vtable::validity::ValidityChild<vortex_tensor::encodings::turboquant::TurboQuant> for vortex_tensor::encodings::turboquant::TurboQuant
+impl vortex_array::array::vtable::validity::ValidityVTable<vortex_tensor::encodings::turboquant::TurboQuant> for vortex_tensor::encodings::turboquant::TurboQuant
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuant::validity_child(array: vortex_array::array::view::ArrayView<'_, vortex_tensor::encodings::turboquant::TurboQuant>) -> vortex_array::array::erased::ArrayRef
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::validity(_array: vortex_array::array::view::ArrayView<'_, vortex_tensor::encodings::turboquant::TurboQuant>) -> vortex_error::VortexResult<vortex_array::validity::Validity>
 
 impl vortex_array::arrays::dict::take::TakeExecute for vortex_tensor::encodings::turboquant::TurboQuant
 
@@ -110,7 +112,7 @@ pub fn vortex_tensor::encodings::turboquant::TurboQuantData::padded_dim(&self) -
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantData::try_new(dimension: u32, bit_width: u8, num_rounds: u8) -> vortex_error::VortexResult<Self>
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuantData::validate(dtype: &vortex_array::dtype::DType, codes: &vortex_array::array::erased::ArrayRef, norms: &vortex_array::array::erased::ArrayRef, centroids: &vortex_array::array::erased::ArrayRef, rotation_signs: &vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<()>
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::validate(dtype: &vortex_array::dtype::DType, codes: &vortex_array::array::erased::ArrayRef, centroids: &vortex_array::array::erased::ArrayRef, rotation_signs: &vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<()>
 
 impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuantData
 
@@ -164,8 +166,6 @@ pub fn vortex_tensor::encodings::turboquant::TurboQuantArrayExt::centroids(&self
 
 pub fn vortex_tensor::encodings::turboquant::TurboQuantArrayExt::codes(&self) -> &vortex_array::array::erased::ArrayRef
 
-pub fn vortex_tensor::encodings::turboquant::TurboQuantArrayExt::norms(&self) -> &vortex_array::array::erased::ArrayRef
-
 pub fn vortex_tensor::encodings::turboquant::TurboQuantArrayExt::rotation_signs(&self) -> &vortex_array::array::erased::ArrayRef
 
 impl<T: vortex_array::array::typed::TypedArrayRef<vortex_tensor::encodings::turboquant::TurboQuant>> vortex_tensor::encodings::turboquant::TurboQuantArrayExt for T
@@ -174,12 +174,12 @@ pub fn T::centroids(&self) -> &vortex_array::array::erased::ArrayRef
 
 pub fn T::codes(&self) -> &vortex_array::array::erased::ArrayRef
 
-pub fn T::norms(&self) -> &vortex_array::array::erased::ArrayRef
-
 pub fn T::rotation_signs(&self) -> &vortex_array::array::erased::ArrayRef
 
 pub fn vortex_tensor::encodings::turboquant::turboquant_encode(ext: vortex_array::array::view::ArrayView<'_, vortex_array::arrays::extension::vtable::Extension>, config: &vortex_tensor::encodings::turboquant::TurboQuantConfig, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
 
+pub unsafe fn vortex_tensor::encodings::turboquant::turboquant_encode_unchecked(ext: vortex_array::array::view::ArrayView<'_, vortex_array::arrays::extension::vtable::Extension>, config: &vortex_tensor::encodings::turboquant::TurboQuantConfig, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
 pub type vortex_tensor::encodings::turboquant::TurboQuantArray = vortex_array::array::typed::Array<vortex_tensor::encodings::turboquant::TurboQuant>
 
 pub mod vortex_tensor::fixed_shape
@@ -430,7 +430,9 @@ impl vortex_tensor::scalar_fns::l2_denorm::L2Denorm
 
 pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::new(options: &vortex_tensor::scalar_fns::ApproxOptions) -> vortex_array::scalar_fn::typed::ScalarFn<vortex_tensor::scalar_fns::l2_denorm::L2Denorm>
 
-pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::try_new_array(options: &vortex_tensor::scalar_fns::ApproxOptions, normalized: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, len: usize) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::vtable::ScalarFnArray>
+pub unsafe fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::new_array_unchecked(options: &vortex_tensor::scalar_fns::ApproxOptions, normalized: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, len: usize) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::vtable::ScalarFnArray>
+
+pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::try_new_array(options: &vortex_tensor::scalar_fns::ApproxOptions, normalized: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, len: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::vtable::ScalarFnArray>
 
 impl core::clone::Clone for vortex_tensor::scalar_fns::l2_denorm::L2Denorm
 
@@ -458,6 +460,10 @@ pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::return_dtype(&self, _opti
 
 pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::validity(&self, _options: &Self::Options, expression: &vortex_array::expr::expression::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::expression::Expression>>
 
+pub fn vortex_tensor::scalar_fns::l2_denorm::normalize_as_l2_denorm(options: &vortex_tensor::scalar_fns::ApproxOptions, 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(input: 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

vortex-tensor/src/encodings/turboquant/array/data.rs

@@ -19,12 +19,16 @@ use crate::encodings::turboquant::vtable::TurboQuant;
 /// TurboQuant array data.
 ///
 /// TurboQuant is a lossy vector quantization encoding for [`Vector`](crate::vector::Vector)
-/// extension arrays. It stores quantized coordinate codes and per-vector norms, along with shared
+/// extension arrays. It stores quantized coordinate codes for unit-norm vectors, along with shared
 /// codebook centroids and the parameters of the current structured rotation.
 ///
+/// Norms should be stored externally in the [`L2Denorm`](crate::scalar_fns::l2_denorm::L2Denorm)
+/// `ScalarFnArray` wrapper.
+///
 /// See the [module docs](crate::encodings::turboquant) for algorithmic details.
 ///
-/// A degenerate TurboQuant array has zero rows and `bit_width == 0`, with all slots empty.
+/// Note that degenerate TurboQuant arrays have zero rows and `bit_width == 0`, with all slots
+/// empty.
 #[derive(Clone, Debug)]
 pub struct TurboQuantData {
     /// The vector dimension `d`, cached from the `FixedSizeList` storage dtype's list size.
@@ -95,16 +99,21 @@ impl TurboQuantData {
     pub fn validate(
         dtype: &DType,
         codes: &ArrayRef,
-        norms: &ArrayRef,
         centroids: &ArrayRef,
         rotation_signs: &ArrayRef,
     ) -> VortexResult<()> {
         let vector_metadata = TurboQuant::validate_dtype(dtype)?;
         let dimension = vector_metadata.dimensions();
         let padded_dim = dimension.next_power_of_two();
 
+        // TurboQuant arrays are always non-nullable. Nullability should be handled by the external
+        // L2Denorm ScalarFnArray wrapper.
+        vortex_ensure!(
+            !dtype.is_nullable(),
+            "TurboQuant dtype must be non-nullable, got {dtype}",
+        );
+
         // Codes must be a non-nullable FixedSizeList<u8> with list_size == padded_dim.
-        // Null vectors are represented by all-zero codes since validity lives in the norms array.
         let expected_codes_dtype = DType::FixedSizeList(
             Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)),
             padded_dim,
@@ -116,23 +125,6 @@ impl TurboQuantData {
             "codes dtype does not match expected {expected_codes_dtype}",
         );
 
-        let num_rows = codes.len();
-        vortex_ensure_eq!(
-            norms.len(),
-            num_rows,
-            "norms length must match codes length",
-        );
-
-        // Norms dtype must match the element ptype of the Vector, with the parent's nullability.
-        // Norms carry the validity of the entire TurboQuant array.
-        let element_ptype = vector_metadata.element_ptype();
-        let expected_norms_dtype = DType::Primitive(element_ptype, dtype.nullability());
-        vortex_ensure_eq!(
-            *norms.dtype(),
-            expected_norms_dtype,
-            "norms dtype does not match expected {expected_norms_dtype}",
-        );
-
         // Centroids are always f32 regardless of element type.
         let centroids_dtype = DType::Primitive(PType::F32, Nullability::NonNullable);
         vortex_ensure_eq!(
@@ -154,6 +146,7 @@ impl TurboQuantData {
             "rotation_signs dtype does not match expected {expected_signs_dtype}",
         );
         // Degenerate (empty) case: all children must be empty, and bit_width is 0.
+        let num_rows = codes.len();
         if num_rows == 0 {
             vortex_ensure!(
                 centroids.is_empty(),
@@ -198,13 +191,11 @@ impl TurboQuantData {
 
     pub(crate) fn make_slots(
         codes: ArrayRef,
-        norms: ArrayRef,
         centroids: ArrayRef,
         rotation_signs: ArrayRef,
     ) -> Vec<Option<ArrayRef>> {
         let mut slots = vec![None; Slot::COUNT];
         slots[Slot::Codes as usize] = Some(codes);
-        slots[Slot::Norms as usize] = Some(norms);
         slots[Slot::Centroids as usize] = Some(centroids);
         slots[Slot::RotationSigns as usize] = Some(rotation_signs);
         slots
@@ -242,12 +233,6 @@ pub trait TurboQuantArrayExt: TypedArrayRef<TurboQuant> {
             .vortex_expect("TurboQuantArray codes slot")
     }
 
-    fn norms(&self) -> &ArrayRef {
-        self.as_ref().slots()[Slot::Norms as usize]
-            .as_ref()
-            .vortex_expect("TurboQuantArray norms slot")
-    }
-
     fn centroids(&self) -> &ArrayRef {
         self.as_ref().slots()[Slot::Centroids as usize]
             .as_ref()

vortex-tensor/src/encodings/turboquant/array/slots.rs

@@ -2,22 +2,26 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 /// Slot positions for TurboQuantArray children.
+///
+/// Norms are not stored in the TurboQuantArray. They live in the external [`L2Denorm`]
+/// ScalarFnArray wrapper returned by [`turboquant_encode`].
+///
+/// [`L2Denorm`]: crate::scalar_fns::l2_denorm::L2Denorm
+/// [`turboquant_encode`]: crate::encodings::turboquant::turboquant_encode
 #[repr(usize)]
 #[derive(Clone, Copy, Debug)]
 pub(crate) enum Slot {
     Codes = 0,
-    Norms = 1,
-    Centroids = 2,
-    RotationSigns = 3,
+    Centroids = 1,
+    RotationSigns = 2,
 }
 
 impl Slot {
-    pub(crate) const COUNT: usize = 4;
+    pub(crate) const COUNT: usize = 3;
 
     pub(crate) fn name(self) -> &'static str {
         match self {
             Self::Codes => "codes",
-            Self::Norms => "norms",
             Self::Centroids => "centroids",
             Self::RotationSigns => "rotation_signs",
         }
@@ -26,9 +30,8 @@ impl Slot {
     pub(crate) fn from_index(idx: usize) -> Self {
         match idx {
             0 => Self::Codes,
-            1 => Self::Norms,
-            2 => Self::Centroids,
-            3 => Self::RotationSigns,
+            1 => Self::Centroids,
+            2 => Self::RotationSigns,
             _ => vortex_error::vortex_panic!("invalid slot index {idx}"),
         }
     }

vortex-tensor/src/encodings/turboquant/compute/slice.rs

@@ -17,14 +17,12 @@ impl SliceReduce for TurboQuant {
         array: ArrayView<'_, TurboQuant>,
         range: Range<usize>,
     ) -> VortexResult<Option<ArrayRef>> {
-        let sliced_codes = array.codes().slice(range.clone())?;
-        let sliced_norms = array.norms().slice(range)?;
+        let sliced_codes = array.codes().slice(range)?;
 
         Ok(Some(
             TurboQuant::try_new_array(
                 array.dtype().clone(),
                 sliced_codes,
-                sliced_norms,
                 array.centroids().clone(),
                 array.rotation_signs().clone(),
             )?

vortex-tensor/src/encodings/turboquant/compute/take.rs

@@ -19,13 +19,11 @@ impl TakeExecute for TurboQuant {
     ) -> VortexResult<Option<ArrayRef>> {
         // FSL children handle per-row take natively.
         let taken_codes = array.codes().take(indices.clone())?;
-        let taken_norms = array.norms().take(indices.clone())?;
 
         Ok(Some(
             TurboQuant::try_new_array(
                 array.dtype().clone(),
                 taken_codes,
-                taken_norms,
                 array.centroids().clone(),
                 array.rotation_signs().clone(),
             )?

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

@@ -17,6 +17,20 @@
 //! TurboQuant minimizes mean-squared reconstruction error (1-8 bits per coordinate)
 //! using MSE-optimal scalar quantization on coordinates of a rotated unit vector.
 //!
+//! The `TurboQuantArray` stores only the quantized unit-norm vector data (codes, centroids,
+//! rotation signs). Per-vector L2 norms are stored separately in an [`L2Denorm`] ScalarFnArray
+//! wrapper. The [`turboquant_encode`] function returns this wrapper:
+//!
+//! ```text
+//! ScalarFnArray(L2Denorm, [TurboQuantArray, norms])
+//! ```
+//!
+//! When executed, the TQ array decompresses to unit-norm vectors, and the [`L2Denorm`] function
+//! lazily re-applies the stored norms to reconstruct the original magnitudes.
+//!
+//! [`L2Denorm`]: crate::scalar_fns::l2_denorm::L2Denorm
+//! [`turboquant_encode`]: crate::encodings::turboquant::turboquant_encode
+//!
 //! The TurboQuant paper analyzes a full random orthogonal rotation. The current Vortex
 //! implementation instead uses a fixed 3-round Walsh-Hadamard-based structured transform with
 //! random sign diagonals. This is a practical approximation chosen for encode/decode efficiency,
@@ -48,9 +62,10 @@
 //! # Compression ratios
 //!
 //! Each vector is stored as `padded_dim * bit_width / 8` bytes of quantized codes plus one stored
-//! norm. In the current implementation, that norm uses the vector's element float type, not a
-//! separate fixed storage precision. Non-power-of-2 dimensions are padded to the next power of 2
-//! for the structured rotation, which reduces the effective ratio for those dimensions.
+//! norm (in the [`L2Denorm`] wrapper). In the current implementation, that norm uses the vector's
+//! element float type, not a separate fixed storage precision. Non-power-of-2 dimensions are
+//! padded to the next power of 2 for the structured rotation, which reduces the effective ratio
+//! for those dimensions.
 //!
 //! The table below assumes f32 input, so the stored norm is 4 bytes.
 //!
@@ -71,13 +86,17 @@
 //! use vortex_array::arrays::ExtensionArray;
 //! use vortex_array::arrays::FixedSizeListArray;
 //! use vortex_array::arrays::PrimitiveArray;
+//! use vortex_array::arrays::Extension;
+//! use vortex_array::arrays::scalar_fn::ScalarFnArrayExt;
 //! use vortex_array::dtype::extension::ExtDType;
 //! use vortex_array::extension::EmptyMetadata;
 //! use vortex_array::validity::Validity;
 //! use vortex_buffer::BufferMut;
 //! use vortex_array::session::ArraySession;
 //! use vortex_session::VortexSession;
-//! use vortex_tensor::encodings::turboquant::{TurboQuantConfig, turboquant_encode};
+//! use vortex_tensor::encodings::turboquant::{TurboQuantConfig, turboquant_encode_unchecked};
+//! use vortex_tensor::scalar_fns::ApproxOptions;
+//! use vortex_tensor::scalar_fns::l2_denorm::normalize_as_l2_denorm;
 //! use vortex_tensor::vector::Vector;
 //!
 //! // Create a Vector extension array of 100 random 128-d vectors.
@@ -95,14 +114,23 @@
 //!     .unwrap().erased();
 //! let ext = ExtensionArray::new(ext_dtype, fsl.into_array());
 //!
-//! // Quantize at 2 bits per coordinate.
-//! let config = TurboQuantConfig { bit_width: 2, seed: Some(42), num_rounds: 3 };
+//! // Normalize, then quantize the normalized child at 2 bits per coordinate.
 //! let session = VortexSession::empty().with::<ArraySession>();
 //! let mut ctx = session.create_execution_ctx();
-//! let encoded = turboquant_encode(ext.as_view(), &config, &mut ctx).unwrap();
+//! let l2_denorm = normalize_as_l2_denorm(
+//!     &ApproxOptions::Exact, ext.into_array(), &mut ctx,
+//! ).unwrap();
+//! let normalized = l2_denorm.child_at(0).clone();
+//!
+//! let normalized_ext = normalized.as_opt::<Extension>().unwrap();
+//! let config = TurboQuantConfig { bit_width: 2, seed: Some(42), num_rounds: 3 };
+//! // SAFETY: We just normalized the input.
+//! let tq = unsafe {
+//!     turboquant_encode_unchecked(normalized_ext, &config, &mut ctx).unwrap()
+//! };
 //!
 //! // Verify compression: 100 vectors x 128 dims x 4 bytes = 51200 bytes input.
-//! assert!(encoded.nbytes() < 51200);
+//! assert!(tq.nbytes() < 51200);
 //! ```
 
 mod array;
@@ -122,6 +150,7 @@ mod scheme;
 pub use scheme::TurboQuantScheme;
 pub use scheme::compress::TurboQuantConfig;
 pub use scheme::compress::turboquant_encode;
+pub use scheme::compress::turboquant_encode_unchecked;
 
 #[cfg(test)]
 mod tests;