diff --git a/vortex-tensor/src/encodings/normalized/array.rs b/vortex-tensor/src/encodings/normalized/array.rs index c60bf556b02..64c63339464 100644 --- a/vortex-tensor/src/encodings/normalized/array.rs +++ b/vortex-tensor/src/encodings/normalized/array.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use prost::Message; use vortex_array::Array; use vortex_array::ArrayId; use vortex_array::ArrayParts; @@ -15,26 +14,30 @@ use vortex_array::array_slots; use vortex_array::arrays::ConstantArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; use vortex_array::serde::ArrayChildren; use vortex_array::validity::Validity; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityVTable; +use vortex_array::vtable::child_to_validity; +use vortex_array::vtable::validity_to_child; use vortex_array::vtable::with_empty_buffers; use vortex_error::VortexResult; -use vortex_error::vortex_err; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::execute::denormalize; use crate::encodings::normalized::rules::RULES; -use crate::encodings::normalized::validate::validate_l2_normalized_rows_against_norms; use crate::encodings::normalized::validate::validate_normalized_children; +use crate::encodings::normalized::validate::validate_normalized_rows; use crate::utils::validate_tensor_float_input; -/// An [`Normalized`]-encoded Vortex array. +/// A [`Normalized`]-encoded Vortex array. pub type NormalizedArray = Array; /// The norm-split encoding for tensor-like columns. @@ -47,18 +50,29 @@ pub type NormalizedArray = Array; /// /// Every [`NormalizedArray`] structurally guarantees, via [`VTable::validate`]: /// -/// - `normalized` is a tensor-like extension array with a float element type. -/// - `norms` is a primitive column whose ptype equals the tensor element ptype. +/// - `normalized` is a non-nullable tensor-like extension array with a float element type, whose +/// dtype is the array's own dtype with nullability stripped. +/// - `norms` is a non-nullable primitive column whose ptype equals the tensor element ptype. /// - both children have the array's length. -/// - the array dtype is `normalized.dtype().union_nullability(norms.nullability())`. +/// - the `validity` slot is present only when the array's dtype is nullable, in which case it is a +/// non-nullable boolean column of the array's length. +/// +/// Nulls therefore live on the array itself rather than in either child, which is what keeps the +/// two children free to be reshaped independently: neither the decode path nor the read-through +/// operators ever have to widen a child's dtype to match the parent's. /// /// On top of that, [`try_new`](Self::try_new) enforces the semantic invariants that make the split /// lossless: /// -/// - every valid row of `normalized` has L2 norm `1.0` or `0.0`, within the tolerance implied by -/// the element precision. +/// - every row of `normalized` has L2 norm `1.0` or `0.0`, within the tolerance implied by the +/// element precision. /// - every stored norm is non-negative. -/// - a stored norm of `0.0` is paired with an all-zero normalized row. +/// - a stored norm of `0.0` is paired with an all-zero normalized row, and an all-zero normalized +/// row is paired with a stored norm of `0.0`. +/// +/// Those checks run over every row, including rows the `validity` marks null. [`normalize`] zeroes +/// both children at null positions, which satisfies them, so callers building a nullable column +/// should go through `normalize` rather than pairing raw children with a mask. /// /// # Lossy normalized children /// @@ -71,26 +85,38 @@ pub type NormalizedArray = Array; /// storage contract, not a separate lossy-compute mode. /// /// [`AnyTensor`]: crate::matcher::AnyTensor +/// [`normalize`]: crate::encodings::normalized::normalize /// [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm /// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct /// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity #[derive(Clone, Debug)] pub struct Normalized; -/// The two child arrays of an [`NormalizedArray`]. +/// The slots of a [`NormalizedArray`]: its two children plus its validity. #[array_slots(Normalized)] pub struct NormalizedSlots { - /// The unit-norm (or zero) direction of each row, as a tensor-like extension array. + /// The unit-norm (or zero) direction of each row, as a non-nullable tensor-like extension + /// array. #[slot(0)] pub normalized: ArrayRef, - /// The authoritative L2 norm of each row, as a primitive float column. + /// The authoritative L2 norm of each row, as a non-nullable primitive float column. #[slot(1)] pub norms: ArrayRef, + + /// The validity / null map of the array. + /// + /// Both children are non-nullable, so this is the column's only record of which rows are null. + #[slot(2)] + pub validity: Option, } +/// The number of required slots: `normalized` and `norms`. The `validity` slot is optional, so a +/// serialized array has either this many children or one more. +pub(super) const DATA_CHILDREN: usize = NormalizedSlots::COUNT - 1; + impl Normalized { - /// Builds an [`NormalizedArray`], validating that `normalized` really is row-wise L2-normalized + /// Builds a [`NormalizedArray`], validating that `normalized` really is row-wise L2-normalized /// against `norms`. /// /// This is the constructor for exact norm splits. It scans both children, so it costs @@ -103,68 +129,72 @@ impl Normalized { pub fn try_new( normalized: ArrayRef, norms: ArrayRef, + validity: Validity, ctx: &mut ExecutionCtx, ) -> VortexResult { - let len = normalized.len(); - let dtype = normalized - .dtype() - .union_nullability(norms.dtype().nullability()); - let slots = NormalizedSlots { normalized, norms }.into_slots(); - // Structural validation has to come first: the row scan walks both children in lockstep // and assumes they are a matching-length tensor/float pair. - let normalized_array = Array::try_from_parts( - ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots), - )?; - validate_l2_normalized_rows_against_norms( - normalized_array.normalized(), - Some(normalized_array.norms()), - ctx, - )?; + let array = Array::try_from_parts(normalized_parts(normalized, norms, validity))?; + validate_normalized_rows(array.normalized(), Some(array.norms()), ctx)?; - Ok(normalized_array) + Ok(array) } - /// Builds an [`NormalizedArray`] without validation. + /// Builds a [`NormalizedArray`] without validation. /// /// # Safety /// /// The caller must uphold the structural invariants listed on [`Normalized`]. In particular, - /// both children must have the same length, `normalized` must be a float tensor, and `norms` - /// must be a primitive column with the same element ptype. + /// both children must be non-nullable and have the same length, `normalized` must be a float + /// tensor, and `norms` must be a primitive column with the same element ptype. /// /// This does not check the unit-norm relationship. Violating it can produce wrong answers but /// not memory unsafety. - pub unsafe fn new_unchecked(normalized: ArrayRef, norms: ArrayRef) -> NormalizedArray { - let len = normalized.len(); - let dtype = normalized - .dtype() - .union_nullability(norms.dtype().nullability()); - let slots = NormalizedSlots { normalized, norms }.into_slots(); - - unsafe { - Array::from_parts_unchecked( - ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots), - ) - } + pub unsafe fn new_unchecked( + normalized: ArrayRef, + norms: ArrayRef, + validity: Validity, + ) -> NormalizedArray { + unsafe { Array::from_parts_unchecked(normalized_parts(normalized, norms, validity)) } } } -/// Metadata for a serialized [`NormalizedArray`]: its children's nullabilities. +/// Assembles the [`ArrayParts`] shared by both constructors and by deserialization. /// -/// The parent dtype supplies the tensor shape and element ptype. Its nullability is the union of -/// the children, so it cannot identify which child is nullable. -#[derive(Clone, prost::Message)] -pub struct NormalizedMetadata { - /// Whether the `normalized` child is nullable. - #[prost(bool, tag = "1")] - pub normalized_is_nullable: bool, - - /// Whether the `norms` child is nullable. - #[prost(bool, tag = "2")] - pub norms_is_nullable: bool, +/// The array's dtype is the `normalized` child's dtype widened by whatever nullability `validity` +/// implies, which is the only place the parent's nullability comes from. +fn normalized_parts( + normalized: ArrayRef, + norms: ArrayRef, + validity: Validity, +) -> ArrayParts { + let len = normalized.len(); + let dtype = normalized.dtype().union_nullability(validity.nullability()); + let slots = NormalizedSlots { + normalized, + norms, + validity: validity_to_child(&validity, len), + } + .into_slots(); + + ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots) +} + +/// Accessors for [`NormalizedArray`] that are derived from its slots rather than being one. +pub trait NormalizedArrayExt: NormalizedArraySlotsExt { + /// The column's validity. + /// + /// Both children are non-nullable, so this is the array's complete null information. + fn normalized_validity(&self) -> Validity { + child_to_validity( + self.as_ref().slots()[NormalizedSlots::VALIDITY].as_ref(), + self.as_ref().dtype().nullability(), + ) + } } +impl NormalizedArrayExt for T {} + impl VTable for Normalized { type TypedArrayData = EmptyArrayData; @@ -185,7 +215,7 @@ impl VTable for Normalized { ) -> VortexResult<()> { let slots = NormalizedSlotsView::from_slots(slots); - validate_normalized_children(slots.normalized, slots.norms, dtype, len) + validate_normalized_children(slots.normalized, slots.norms, slots.validity, dtype, len) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -208,17 +238,13 @@ impl VTable for Normalized { with_empty_buffers(self, array, buffers) } + /// The array carries no metadata: the parent dtype supplies the tensor shape, element ptype, + /// and nullability, and both children's dtypes follow from it. fn serialize( - array: ArrayView<'_, Self>, + _array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - Ok(Some( - NormalizedMetadata { - normalized_is_nullable: array.normalized().dtype().is_nullable(), - norms_is_nullable: array.norms().dtype().is_nullable(), - } - .encode_to_vec(), - )) + Ok(Some(vec![])) } fn deserialize( @@ -230,18 +256,35 @@ impl VTable for Normalized { children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - let metadata = NormalizedMetadata::decode(metadata) - .map_err(|e| vortex_err!("Failed to decode NormalizedMetadata: {e}"))?; + vortex_ensure!( + metadata.is_empty(), + "NormalizedArray expects empty metadata, got {} bytes", + metadata.len(), + ); let element_ptype = validate_tensor_float_input(dtype)?.element_ptype(); - let normalized_dtype = dtype.with_nullability(metadata.normalized_is_nullable.into()); - let norms_dtype = DType::Primitive(element_ptype, metadata.norms_is_nullable.into()); + let normalized_dtype = dtype.as_nonnullable(); + let norms_dtype = DType::Primitive(element_ptype, Nullability::NonNullable); let normalized = children.get(0, &normalized_dtype, len)?; let norms = children.get(1, &norms_dtype, len)?; - let slots = NormalizedSlots { normalized, norms }.into_slots(); - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) + // An absent validity child means "no nulls"; the parent's nullability is what distinguishes + // `NonNullable` from `AllValid`. + let validity = if children.len() == NormalizedSlots::COUNT { + Validity::Array(children.get(NormalizedSlots::VALIDITY, &Validity::DTYPE, len)?) + } else { + vortex_ensure_eq!( + children.len(), + DATA_CHILDREN, + "NormalizedArray expects {DATA_CHILDREN} or {} children, got {}", + NormalizedSlots::COUNT, + children.len(), + ); + Validity::from(dtype.nullability()) + }; + + Ok(normalized_parts(normalized, norms, validity)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -250,10 +293,19 @@ impl VTable for Normalized { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let dtype = array.dtype().clone(); + let validity = array.normalized_validity(); + let row_count = array.len(); let slots = array.slots_view(); - denormalize(slots.normalized, slots.norms, array.len(), dtype, ctx) - .map(ExecutionResult::done) + denormalize( + slots.normalized, + slots.norms, + validity, + row_count, + dtype, + ctx, + ) + .map(ExecutionResult::done) } fn reduce_parent( @@ -267,10 +319,7 @@ impl VTable for Normalized { impl ValidityVTable for Normalized { fn validity(array: ArrayView<'_, Normalized>) -> VortexResult { - array - .normalized() - .validity()? - .and(array.norms().validity()?) + Ok(array.normalized_validity()) } } @@ -284,12 +333,18 @@ impl OperationsVTable for Normalized { // one-row constants, which also lets the constant-norms fast path do the multiply. let normalized = array.normalized().execute_scalar(index, ctx)?; let norms = array.norms().execute_scalar(index, ctx)?; + let dtype = array.dtype().clone(); + + // `Array::execute_scalar` resolves null rows before dispatching here, so this row is valid + // and only the parent's nullability has to be reproduced. + let validity = Validity::from(dtype.nullability()); let row = denormalize( &ConstantArray::new(normalized, 1).into_array(), &ConstantArray::new(norms, 1).into_array(), + validity, 1, - array.dtype().clone(), + dtype, ctx, )?; diff --git a/vortex-tensor/src/encodings/normalized/compress.rs b/vortex-tensor/src/encodings/normalized/compress.rs index 727483296e4..da696a0d8c3 100644 --- a/vortex-tensor/src/encodings/normalized/compress.rs +++ b/vortex-tensor/src/encodings/normalized/compress.rs @@ -38,8 +38,10 @@ use vortex_error::VortexResult; use crate::encodings::normalized::Normalized; use crate::encodings::normalized::NormalizedArray; +use crate::encodings::normalized::NormalizedArrayExt; use crate::encodings::normalized::NormalizedArraySlotsExt; use crate::encodings::normalized::NormalizedSlots; +use crate::encodings::normalized::array::DATA_CHILDREN; use crate::matcher::AnyTensor; use crate::scalar_fns::l2_norm::L2Norm; use crate::utils::extract_constant_flat_row; @@ -55,20 +57,28 @@ impl Scheme for NormalizedScheme { "vortex.tensor.normalized" } + /// Matching has to be as narrow as [`compress`](Self::compress) is: this scheme reports + /// [`EstimateVerdict::AlwaysUse`], so a canonical array it claims is never offered to another + /// scheme. Claiming an integer tensor here would abort the whole column's compression on the + /// float-only gate in `compress` rather than falling through. fn matches(&self, canonical: &Canonical) -> bool { - matches!( - canonical, - Canonical::Extension(ext) if ext.ext_dtype().is::() - ) + let Canonical::Extension(ext) = canonical else { + return false; + }; + + ext.ext_dtype() + .metadata_opt::() + .is_some_and(|tensor| tensor.element_ptype().is_float()) } fn produced_encodings(&self) -> Vec { vec![Normalized.id()] } - /// Children: normalized=0, norms=1. + /// Children: normalized=0, norms=1. The validity slot is passed through uncompressed, matching + /// how the compressor treats `FixedSizeListArray` and `StructArray` validity. fn num_children(&self) -> usize { - NormalizedSlots::COUNT + DATA_CHILDREN } fn expected_compression_ratio( @@ -107,37 +117,41 @@ impl Scheme for NormalizedScheme { )?; // SAFETY: Cascading preserves the split's child lengths and dtypes. - Ok(unsafe { Normalized::new_unchecked(normalized, norms) }.into_array()) + Ok(unsafe { + Normalized::new_unchecked(normalized, norms, normalized_array.normalized_validity()) + } + .into_array()) } } /// Splits a tensor-like column into its exact [`Normalized`] representation. /// -/// # Normalized child +/// # Children /// -/// The normalized child is always **non-nullable**. Every non-null row with a positive L2 norm is -/// divided by its norm to produce a unit-norm row. +/// Both children are **non-nullable**. Every non-null row with a positive L2 norm is divided by its +/// norm to produce a unit-norm row. /// -/// Rows that are null in the original input are **zeroed out** in the normalized output. Null rows -/// may carry undefined physical storage values, and we do not want that garbage propagating into -/// downstream lossy encodings of the normalized child. +/// Rows that are null in the original input are **zeroed out** in both children. Null rows may +/// carry undefined physical storage values, and we do not want that garbage propagating into +/// downstream lossy encodings of the normalized child — nor into the read-through operators, which +/// consume the norms buffer densely. /// /// # Nullability /// -/// Nullability is tracked entirely by the norms child, which inherits the input's nulls through -/// [`L2Norm`]'s validity propagation. The [`Normalized`] array's validity is the `and` of both -/// children, so an all-valid normalized child plus a nullable norms child reproduces the input's -/// validity exactly. +/// The input's nulls move onto the [`Normalized`] array's own validity, which it takes from +/// [`L2Norm`]'s validity propagation. Because the children carry no nulls of their own, that +/// validity is the reconstructed column's validity exactly. /// /// Because this computes exact norms first and then divides by them, the returned `normalized` -/// child satisfies the strict unit-norm invariant. +/// child satisfies the strict unit-norm invariant, and zeroing null rows satisfies both directions +/// of the zero-norm rule. pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let row_count = input.len(); let tensor_match = validate_tensor_float_input(input.dtype())?; let tensor_flat_size = tensor_match.list_size() as usize; // Constant fast path: if the input is a constant-backed extension, normalize the single stored - // row once and return an `Normalized` whose children are both `ConstantArray`s. + // row once and return a `Normalized` whose children are both `ConstantArray`s. if let Some(wrapped) = try_build_constant_normalized(&input, row_count, ctx)? { return Ok(wrapped); } @@ -145,8 +159,10 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult VortexResult(); let total_elements = row_count * tensor_flat_size; let mut elements = BufferMut::::with_capacity(total_elements); + let mut norms = BufferMut::::with_capacity(row_count); for i in 0..row_count { - let is_valid = norms_valid.value(i); - let norm = norm_values[i]; + // A null row's stored values are undefined, so its computed norm is meaningless. Zero + // both children there instead of storing whatever the garbage happened to produce. + let norm = if valid.value(i) { + norm_values[i] + } else { + T::zero() + }; + norms.push(norm); // SAFETY: We allocated `row_count * tensor_flat_size` capacity and push exactly // `tensor_flat_size` elements per row. - - // Null rows must be explicitly zeroed out. - if !is_valid || norm == T::zero() { + if norm == T::zero() { unsafe { elements.push_n_unchecked(T::zero(), tensor_flat_size) }; } else { for &x in flat.row::(i) { @@ -178,21 +199,24 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult() .vortex_expect("caller validated input has AnyTensor metadata"); let list_size = tensor_match.list_size() as usize; - let original_nullability = input.dtype().nullability(); - let ext_dtype = input.dtype().as_extension().clone(); - let storage_fsl_nullability = storage.dtype().nullability(); + + // The stored row is non-null, so every row is valid; the input's nullability only decides + // whether the column *can* hold nulls. Both children drop it: they are always non-nullable. + let validity = Validity::from(input.dtype().nullability()); + let normalized_ext_dtype = input.dtype().as_nonnullable().as_extension().clone(); // Materialize just the single stored row; this does not expand the constant to the full column // length. @@ -255,20 +281,20 @@ pub(crate) fn try_build_constant_normalized( .collect() }; - // The rebuilt FSL scalar preserves the original storage FSL's nullability so the resulting - // `ExtensionArray::new` call accepts the same extension dtype. - let fsl_scalar = Scalar::fixed_size_list(element_dtype, children, storage_fsl_nullability); - let norms_scalar = Scalar::primitive(norm_t, original_nullability); + // Both scalars are non-nullable, matching the non-nullable extension dtype the normalized + // child is rebuilt under. + let fsl_scalar = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable); + let norms_scalar = Scalar::primitive(norm_t, Nullability::NonNullable); (fsl_scalar, norms_scalar) }); let normalized_storage = ConstantArray::new(normalized_fsl_scalar, len).into_array(); - let normalized = ExtensionArray::new(ext_dtype, normalized_storage).into_array(); + let normalized = ExtensionArray::new(normalized_ext_dtype, normalized_storage).into_array(); let norms = ConstantArray::new(norms_scalar, len).into_array(); // SAFETY: The constant children have matching lengths and element ptypes. Ok(Some(unsafe { - Normalized::new_unchecked(normalized, norms) + Normalized::new_unchecked(normalized, norms, validity) })) } diff --git a/vortex-tensor/src/encodings/normalized/execute.rs b/vortex-tensor/src/encodings/normalized/execute.rs index 637c8c07117..33333111baf 100644 --- a/vortex-tensor/src/encodings/normalized/execute.rs +++ b/vortex-tensor/src/encodings/normalized/execute.rs @@ -8,6 +8,7 @@ use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; 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::FixedSizeListArrayExt; @@ -25,30 +26,26 @@ use vortex_error::VortexResult; use crate::matcher::AnyTensor; use crate::utils::extract_flat_elements; -use crate::utils::unit_norm_tolerance; /// Reconstructs the original tensor column by scaling each normalized row by its stored norm. /// -/// `dtype` is the parent [`NormalizedArray`]'s dtype, so the reconstructed column carries the -/// unioned nullability of both children. +/// `dtype` is the parent [`NormalizedArray`]'s dtype and `validity` its null map, so the +/// reconstructed column carries the parent's nullability rather than either child's — both children +/// are non-nullable. /// /// [`NormalizedArray`]: crate::encodings::normalized::NormalizedArray pub(super) fn denormalize( normalized: &ArrayRef, norms: &ArrayRef, + validity: Validity, row_count: usize, dtype: DType, ctx: &mut ExecutionCtx, ) -> VortexResult { - let validity = normalized.validity()?.and(norms.validity()?)?; - // Constant norms let us scale the whole backing buffer at once, or skip the multiply entirely - // when every norm is already 1. The nullability guard keeps us on the general path when the - // constant is a non-null value inside a nullable column, since the fast path cannot widen the - // normalized child's dtype to match the parent's. + // when every norm is exactly 1. if let Some(constant) = norms.as_opt::() && constant.scalar().value().is_some() - && normalized.dtype() == &dtype { return denormalize_constant_norms(normalized, constant.scalar(), dtype, validity, ctx); } @@ -76,9 +73,9 @@ pub(super) fn denormalize( /// Scales every row by the same stored norm. /// -/// Two things make this cheaper than the general path: a norm of `1.0` is the identity, so the -/// normalized child is already the answer; and otherwise the scale factor applies uniformly to the -/// flat backing buffer, so it becomes one lazy multiply over the elements array instead of a +/// Two things make this cheaper than the general path: a norm of exactly `1.0` is the identity, so +/// the normalized child is already the answer; and otherwise the scale factor applies uniformly to +/// the flat backing buffer, so it becomes one lazy multiply over the elements array instead of a /// per-row loop. fn denormalize_constant_norms( normalized: &ArrayRef, @@ -87,17 +84,19 @@ fn denormalize_constant_norms( validity: Validity, ctx: &mut ExecutionCtx, ) -> VortexResult { - let tensor_flat_size = tensor_flat_size(normalized.dtype()); - let error = norm + let norm_value = norm .value() .vortex_expect("the caller only takes this path for a non-null constant norm") .as_primitive() .as_f64() - .vortex_expect("norms are validated to be a float column, so the scalar fits in f64") - - 1.0f64; - - if error.abs() < unit_norm_tolerance(norm.dtype().as_ptype(), tensor_flat_size) { - return Ok(normalized.clone()); + .vortex_expect("norms are validated to be a float column, so the scalar fits in f64"); + + // Only an exact `1.0` is the identity. Skipping the multiply for a merely *near*-unit norm + // would leave this path disagreeing with the general one in the last bits, and `scalar_at` + // routes every row through here — so a per-row read would answer differently than a bulk decode + // of the same column. + if norm_value == 1.0 { + return apply_validity(normalized.clone(), validity); } let normalized: ExtensionArray = normalized.clone().execute(ctx)?; @@ -106,8 +105,8 @@ fn denormalize_constant_norms( let scale = ConstantArray::new(norm.clone(), storage.elements().len()).into_array(); let elements = storage.elements().clone().binary(scale, Operator::Mul)?; - // SAFETY: Only the element values changed; the list size, validity, and row count are carried - // over from the storage array we just executed. + // SAFETY: Only the element values changed; the list size and row count are carried over from + // the storage array we just executed, and the validity is the parent's. let storage = unsafe { FixedSizeListArray::new_unchecked(elements, storage.list_size(), validity, storage.len()) }; @@ -115,6 +114,17 @@ fn denormalize_constant_norms( Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) } +/// Attaches the parent's `validity` to an array that stands in for the decoded column unchanged. +/// +/// The `normalized` child is non-nullable, so a nullable parent needs its null map put back before +/// the child can be returned as-is. +fn apply_validity(array: ArrayRef, validity: Validity) -> VortexResult { + match validity { + Validity::NonNullable => Ok(array), + validity => Ok(MaskedArray::try_new(array, validity)?.into_array()), + } +} + /// Rebuilds a tensor-like extension array from flat primitive elements. fn build_tensor_array( dtype: DType, diff --git a/vortex-tensor/src/encodings/normalized/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs index 545236bba7d..d9ea43cd55f 100644 --- a/vortex-tensor/src/encodings/normalized/mod.rs +++ b/vortex-tensor/src/encodings/normalized/mod.rs @@ -3,10 +3,12 @@ //! The [`Normalized`] encoding: a norm-split physical layout for tensor-like columns. //! -//! An [`Normalized`] array stores a tensor or vector column as two children: +//! A [`Normalized`] array stores a tensor or vector column as two non-nullable children plus its +//! own validity: //! -//! - `normalized`, a tensor-like column whose valid rows are unit-norm (or zero), and -//! - `norms`, a primitive float column holding the authoritative L2 norm of each row. +//! - `normalized`, a tensor-like column whose rows are unit-norm (or zero), +//! - `norms`, a primitive float column holding the authoritative L2 norm of each row, and +//! - `validity`, the column's null map. //! //! The logical value of row `i` is `normalized[i] * norms[i]`, so canonicalizing the array //! reconstructs the original tensor column. Splitting magnitude away from direction is what makes @@ -14,6 +16,10 @@ //! value range, and quantizing it only perturbs direction while the exact magnitude survives in //! `norms`. //! +//! Keeping nulls on the array rather than in either child means neither the decode path nor the +//! read-through operators have to widen a child's dtype to reach the parent's, and it leaves both +//! children free to be reshaped independently. +//! //! Because the split is physical rather than logical, [`L2Norm`], [`InnerProduct`], and //! [`CosineSimilarity`] can read straight through it instead of decoding first. //! @@ -24,8 +30,8 @@ mod array; pub use array::Normalized; pub use array::NormalizedArray; +pub use array::NormalizedArrayExt; pub use array::NormalizedArraySlotsExt; -pub use array::NormalizedMetadata; pub use array::NormalizedSlots; mod compress; @@ -41,7 +47,7 @@ pub(crate) use orientation::NormalizedOrientation; mod rules; mod validate; -pub use validate::validate_l2_normalized_rows_against_norms; +pub use validate::validate_normalized_rows; #[cfg(test)] mod tests; diff --git a/vortex-tensor/src/encodings/normalized/rules.rs b/vortex-tensor/src/encodings/normalized/rules.rs index 7db946c25b0..718db9c37d2 100644 --- a/vortex-tensor/src/encodings/normalized/rules.rs +++ b/vortex-tensor/src/encodings/normalized/rules.rs @@ -11,6 +11,7 @@ use vortex_array::optimizer::rules::ParentRuleSet; use vortex_error::VortexResult; use crate::encodings::normalized::Normalized; +use crate::encodings::normalized::array::NormalizedArrayExt; use crate::encodings::normalized::array::NormalizedArraySlotsExt; pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ @@ -37,12 +38,13 @@ impl ArrayParentReduceRule for NormalizedSliceRule { ) -> VortexResult> { let range = parent.slice_range(); - // SAFETY: Slicing both children preserves their structure. + // SAFETY: Slicing both children and the validity preserves their structure. Ok(Some( unsafe { Normalized::new_unchecked( array.normalized().slice(range.clone())?, array.norms().slice(range.clone())?, + array.normalized_validity().slice(range.clone())?, ) } .into_array(), @@ -69,12 +71,14 @@ impl ArrayParentReduceRule for NormalizedFilterRule { ) -> VortexResult> { let mask = parent.filter_mask(); - // SAFETY: Filtering both children with the same mask preserves their structure. + // SAFETY: Filtering both children and the validity with the same mask preserves their + // structure. Ok(Some( unsafe { Normalized::new_unchecked( array.normalized().filter(mask.clone())?, array.norms().filter(mask.clone())?, + array.normalized_validity().filter(mask)?, ) } .into_array(), diff --git a/vortex-tensor/src/encodings/normalized/tests.rs b/vortex-tensor/src/encodings/normalized/tests.rs index a3dd603a7c3..536f359adb9 100644 --- a/vortex-tensor/src/encodings/normalized/tests.rs +++ b/vortex-tensor/src/encodings/normalized/tests.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use prost::Message; use rstest::rstest; use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; @@ -17,6 +16,7 @@ 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::FixedSizeListArrayExt; use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -32,11 +32,12 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use crate::encodings::normalized::Normalized; +use crate::encodings::normalized::NormalizedArrayExt; use crate::encodings::normalized::NormalizedArraySlotsExt; -use crate::encodings::normalized::NormalizedMetadata; use crate::encodings::normalized::NormalizedScheme; +use crate::encodings::normalized::NormalizedSlots; use crate::encodings::normalized::normalize; -use crate::encodings::normalized::validate_l2_normalized_rows_against_norms; +use crate::encodings::normalized::validate_normalized_rows; use crate::tests::SESSION; use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; @@ -46,16 +47,23 @@ use crate::utils::test_helpers::vector_array; /// Builds a [`Normalized`] array through the checked constructor and executes it, which is the /// end-to-end path every decode test cares about. -fn eval_normalized(normalized: ArrayRef, norms: ArrayRef) -> VortexResult { +fn eval_normalized( + normalized: ArrayRef, + norms: ArrayRef, + validity: Validity, +) -> VortexResult { let mut ctx = SESSION.create_execution_ctx(); - let normalized_array = Normalized::try_new(normalized, norms, &mut ctx)?; + let normalized_array = Normalized::try_new(normalized, norms, validity, &mut ctx)?; normalized_array.into_array().execute(&mut ctx) } -/// Snapshots a tensor-like array as `(dtype, per-row validity, flat elements)` so two columns can -/// be compared without depending on their physical encoding. -fn tensor_snapshot(array: ArrayRef) -> VortexResult<(DType, Vec, Vec)> { +/// Snapshots a tensor-like array as `(dtype, per-row validity, per-element values)` so two columns +/// can be compared without depending on their physical encoding. +/// +/// Elements belonging to null rows come back as `None`. A null row's physical storage values are +/// unspecified, so comparing them would pin an implementation detail rather than the column. +fn tensor_snapshot(array: ArrayRef) -> VortexResult<(DType, Vec, Vec>)> { let mut ctx = SESSION.create_execution_ctx(); let ext: ExtensionArray = array.execute(&mut ctx)?; let validity = (0..ext.len()) @@ -64,21 +72,33 @@ fn tensor_snapshot(array: ArrayRef) -> VortexResult<(DType, Vec, Vec) let storage: FixedSizeListArray = ext.storage_array().clone().execute(&mut ctx)?; let elements: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; - Ok(( - ext.dtype().clone(), - validity, - elements.as_slice::().to_vec(), - )) + let list_size = storage.list_size() as usize; + let values = elements + .as_slice::() + .iter() + .enumerate() + .map(|(i, &value)| validity[i / list_size].then_some(value)) + .collect(); + + Ok((ext.dtype().clone(), validity, values)) } #[track_caller] fn assert_tensor_arrays_eq(actual: ArrayRef, expected: ArrayRef) -> VortexResult<()> { - let (actual_dtype, actual_validity, actual_elements) = tensor_snapshot(actual)?; - let (expected_dtype, expected_validity, expected_elements) = tensor_snapshot(expected)?; + let (actual_dtype, actual_validity, actual_values) = tensor_snapshot(actual)?; + let (expected_dtype, expected_validity, expected_values) = tensor_snapshot(expected)?; assert_eq!(actual_dtype, expected_dtype); assert_eq!(actual_validity, expected_validity); - assert_close(&actual_elements, &expected_elements); + assert_eq!(actual_values.len(), expected_values.len()); + + for (i, (actual, expected)) in actual_values.iter().zip(&expected_values).enumerate() { + match (actual, expected) { + (None, None) => {} + (Some(actual), Some(expected)) => assert_close(&[*actual], &[*expected]), + _ => panic!("element {i}: got {actual:?}, expected {expected:?}"), + } + } Ok(()) } @@ -95,6 +115,12 @@ fn constant_f64_norms(value: f64, len: usize) -> ArrayRef { ConstantArray::new(Scalar::primitive(value, Nullability::NonNullable), len).into_array() } +fn nullable_vector_input() -> VortexResult { + let vectors = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0])?; + + Ok(MaskedArray::try_new(vectors, Validity::from_iter([true, false, true]))?.into_array()) +} + // ============================================================================= // Decoding // ============================================================================= @@ -104,7 +130,7 @@ fn decodes_vectors() -> VortexResult<()> { let normalized = vector_array(3, &[0.6, 0.8, 0.0, 0.0, 0.0, 0.0])?; let norms = PrimitiveArray::from_iter([5.0f64, 0.0]).into_array(); - let actual = eval_normalized(normalized, norms)?; + let actual = eval_normalized(normalized, norms, Validity::NonNullable)?; let expected = vector_array(3, &[3.0, 4.0, 0.0, 0.0, 0.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) @@ -115,43 +141,46 @@ fn decodes_fixed_shape_tensors() -> VortexResult<()> { let normalized = tensor_array(&[2, 2], &[0.5, 0.5, 0.5, 0.5, 1.0, 0.0, 0.0, 0.0])?; let norms = PrimitiveArray::from_iter([4.0f64, 2.0]).into_array(); - let actual = eval_normalized(normalized, norms)?; + let actual = eval_normalized(normalized, norms, Validity::NonNullable)?; let expected = tensor_array(&[2, 2], &[2.0, 2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) } #[test] -fn decodes_null_rows_from_either_child() -> VortexResult<()> { - let normalized = vector_array(2, &[0.6, 0.8, 1.0, 0.0, 0.0, 0.0])?; - let normalized = - MaskedArray::try_new(normalized, Validity::from_iter([true, false, true]))?.into_array(); - let norms = PrimitiveArray::from_option_iter([Some(5.0f64), Some(2.0), None]).into_array(); +fn decodes_null_rows_from_the_stored_validity() -> VortexResult<()> { + let normalized = vector_array(2, &[0.6, 0.8, 0.0, 0.0, 1.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 0.0, 2.0]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let actual: ExtensionArray = eval_normalized(normalized, norms)?.execute(&mut ctx)?; + let validity = Validity::from_iter([true, false, true]); + let actual: ExtensionArray = eval_normalized(normalized, norms, validity)?.execute(&mut ctx)?; let storage: FixedSizeListArray = actual.storage_array().clone().execute(&mut ctx)?; let elements: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; assert!(actual.is_valid(0, &mut ctx)?); assert!(!actual.is_valid(1, &mut ctx)?); - assert!(!actual.is_valid(2, &mut ctx)?); + assert!(actual.is_valid(2, &mut ctx)?); assert_close(&elements.as_slice::()[..2], &[3.0, 4.0]); + assert_close(&elements.as_slice::()[4..], &[2.0, 0.0]); Ok(()) } +/// Both children are non-nullable, so the stored validity is the column's only null record. #[test] -fn validity_is_the_intersection_of_both_children() -> VortexResult<()> { +fn validity_comes_from_the_stored_null_map() -> VortexResult<()> { let normalized = vector_array(2, &[1.0, 0.0, 1.0, 0.0, 1.0, 0.0])?; - let normalized = - MaskedArray::try_new(normalized, Validity::from_iter([true, false, true]))?.into_array(); - let norms = PrimitiveArray::from_option_iter([Some(1.0f64), Some(1.0), None]).into_array(); + let norms = PrimitiveArray::from_iter([1.0f64, 1.0, 1.0]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let normalized_array = Normalized::try_new(normalized, norms, &mut ctx)?; + let validity = Validity::from_iter([true, false, false]); + let normalized_array = Normalized::try_new(normalized, norms, validity, &mut ctx)?; assert!(normalized_array.dtype().is_nullable()); + assert!(!normalized_array.normalized().dtype().is_nullable()); + assert!(!normalized_array.norms().dtype().is_nullable()); + let mask = normalized_array .as_ref() .validity()? @@ -174,21 +203,29 @@ fn constant_unit_norms_decode_to_the_normalized_child() -> VortexResult<()> { let normalized = vector_array(3, &[1.0, 0.0, 0.0, 0.0, 1.0, 0.0])?; let norms = constant_f64_norms(1.0, 2); - let actual = eval_normalized(normalized.clone(), norms)?; + let actual = eval_normalized(normalized.clone(), norms, Validity::NonNullable)?; assert_tensor_arrays_eq(actual, normalized) } #[test] -fn constant_near_unit_norms_decode_to_the_normalized_child() -> VortexResult<()> { - // A norm that differs from 1.0 by less than the f64 unit-norm tolerance must still hit the - // identity fast path. +fn constant_near_unit_norms_are_still_multiplied() -> VortexResult<()> { + // Only an exact 1.0 is the identity. A norm that merely differs from 1.0 by less than the + // unit-norm tolerance must still be applied, so that a per-row `scalar_at` cannot answer + // differently than a bulk decode of the same column. + let near_unit = 1.0f64 + 2.0 * f64::EPSILON; let normalized = vector_array(3, &[1.0, 0.0, 0.0, 0.0, 1.0, 0.0])?; - let norms = constant_f64_norms(1.0 + 1e-12, 2); + let norms = constant_f64_norms(near_unit, 2); + + let mut ctx = SESSION.create_execution_ctx(); + let decoded = eval_normalized(normalized, norms, Validity::NonNullable)?; + let ext: ExtensionArray = decoded.execute(&mut ctx)?; + let storage: FixedSizeListArray = ext.storage_array().clone().execute(&mut ctx)?; + let elements: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; - let actual = eval_normalized(normalized.clone(), norms)?; + assert_eq!(elements.as_slice::()[0], near_unit); - assert_tensor_arrays_eq(actual, normalized) + Ok(()) } #[test] @@ -196,7 +233,7 @@ fn constant_nonunit_norms_scale_vectors() -> VortexResult<()> { let normalized = vector_array(3, &[0.6, 0.8, 0.0, 1.0, 0.0, 0.0])?; let norms = constant_f64_norms(5.0, 2); - let actual = eval_normalized(normalized, norms)?; + let actual = eval_normalized(normalized, norms, Validity::NonNullable)?; let expected = vector_array(3, &[3.0, 4.0, 0.0, 5.0, 0.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) @@ -209,27 +246,33 @@ fn constant_nonunit_norms_scale_fixed_shape_tensors() -> VortexResult<()> { let normalized = tensor_array(&[2, 2], &[0.5, 0.5, 0.5, 0.5, 1.0, 0.0, 0.0, 0.0])?; let norms = constant_f64_norms(4.0, 2); - let actual = eval_normalized(normalized, norms)?; + let actual = eval_normalized(normalized, norms, Validity::NonNullable)?; let expected = tensor_array(&[2, 2], &[2.0, 2.0, 2.0, 2.0, 4.0, 0.0, 0.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) } -#[test] -fn nullable_constant_norms_widen_the_decoded_dtype() -> VortexResult<()> { - // A non-null constant inside a *nullable* norms column cannot take the identity fast path: - // the parent dtype is nullable while the normalized child is not. - let normalized = vector_array(2, &[1.0, 0.0, 0.0, 1.0])?; - let norms = - ConstantArray::new(Scalar::primitive(1.0f64, Nullability::Nullable), 2).into_array(); +/// Regression: the constant-norms paths have to reach the array's nullability starting from a +/// non-nullable `normalized` child. The identity path (`norm == 1.0`) and the bulk-multiply path +/// get there by different routes, so both need covering — the multiply path used to widen the FSL +/// elements and panic on `ExtensionArray::new`. +#[rstest] +#[case::unit_norm(1.0)] +#[case::non_unit_norm(5.0)] +fn nullable_constant_norms_decode_to_the_nullable_dtype(#[case] norm: f64) -> VortexResult<()> { + let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; + let norms = constant_f64_norms(norm, 2); let mut ctx = SESSION.create_execution_ctx(); - let normalized_array = Normalized::try_new(normalized, norms, &mut ctx)?; + let validity = Validity::from_iter([true, false]); + let normalized_array = Normalized::try_new(normalized, norms, validity, &mut ctx)?; let dtype = normalized_array.dtype().clone(); let decoded: ArrayRef = normalized_array.into_array().execute(&mut ctx)?; assert!(dtype.is_nullable()); assert_eq!(decoded.dtype(), &dtype); + assert!(decoded.is_valid(0, &mut ctx)?); + assert!(!decoded.is_valid(1, &mut ctx)?); Ok(()) } @@ -263,13 +306,40 @@ fn nullable_constant_norms_widen_the_decoded_dtype() -> VortexResult<()> { vector_array(2, &[1.0f64, 0.0, 0.0, 1.0]).expect("valid vector array"), PrimitiveArray::from_iter([1.0f64]).into_array(), )] +#[case::nullable_normalized( + nullable_unit_vectors().expect("valid masked array"), + PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(), +)] +#[case::nullable_norms( + vector_array(2, &[1.0f64, 0.0, 0.0, 1.0]).expect("valid vector array"), + PrimitiveArray::from_option_iter([Some(1.0f64), None]).into_array(), +)] fn rejects_structurally_invalid_children( #[case] normalized: ArrayRef, #[case] norms: ArrayRef, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - assert!(Normalized::try_new(normalized, norms, &mut ctx).is_err()); + assert!(Normalized::try_new(normalized, norms, Validity::NonNullable, &mut ctx).is_err()); + + Ok(()) +} + +fn nullable_unit_vectors() -> VortexResult { + let vectors = vector_array(2, &[1.0f64, 0.0, 0.0, 1.0])?; + + Ok(MaskedArray::try_new(vectors, Validity::AllValid)?.into_array()) +} + +#[test] +fn rejects_a_validity_of_the_wrong_length() -> VortexResult<()> { + let normalized = vector_array(2, &[1.0f64, 0.0, 0.0, 1.0])?; + let norms = PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + let validity = Validity::from_iter([true, false, true]); + + assert!(Normalized::try_new(normalized, norms, validity, &mut ctx).is_err()); Ok(()) } @@ -287,13 +357,19 @@ fn rejects_structurally_invalid_children( vector_array(2, &[1.0f64, 0.0, 0.0, 0.0]).expect("valid vector array"), PrimitiveArray::from_iter([0.0f64, 0.0]).into_array(), )] +// The mirror image of the case above: it decodes to `[0.0, 0.0]` while `L2Norm` reads the stored +// `5.0` straight back, so the split is not lossless. +#[case::zero_row_with_nonzero_norm( + vector_array(2, &[0.0f64, 0.0]).expect("valid vector array"), + PrimitiveArray::from_iter([5.0f64]).into_array(), +)] fn checked_construction_rejects_semantic_violations( #[case] normalized: ArrayRef, #[case] norms: ArrayRef, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - assert!(Normalized::try_new(normalized, norms, &mut ctx).is_err()); + assert!(Normalized::try_new(normalized, norms, Validity::NonNullable, &mut ctx).is_err()); Ok(()) } @@ -303,7 +379,7 @@ fn accepts_zero_vectors_paired_with_zero_norms() -> VortexResult<()> { let normalized = vector_array(2, &[0.0, 0.0, 1.0, 0.0])?; let norms = PrimitiveArray::from_iter([0.0f64, 3.0]).into_array(); - let actual = eval_normalized(normalized, norms)?; + let actual = eval_normalized(normalized, norms, Validity::NonNullable)?; let expected = vector_array(2, &[0.0, 0.0, 3.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) @@ -315,11 +391,7 @@ fn validate_accepts_normalized_f16_rows() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let normalized_array = normalize(input, &mut ctx)?; - validate_l2_normalized_rows_against_norms( - &normalized_array.normalized().clone(), - None, - &mut ctx, - ) + validate_normalized_rows(&normalized_array.normalized().clone(), None, &mut ctx) } #[test] @@ -327,7 +399,7 @@ fn validate_rejects_unnormalized_rows() -> VortexResult<()> { let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0])?; let mut ctx = SESSION.create_execution_ctx(); - assert!(validate_l2_normalized_rows_against_norms(&input, None, &mut ctx).is_err()); + assert!(validate_normalized_rows(&input, None, &mut ctx).is_err()); Ok(()) } @@ -343,6 +415,7 @@ fn validate_rejects_unnormalized_rows() -> VortexResult<()> { )] #[case::constant_tensor(constant_tensor_array(&[2], &[3.0, 4.0], 3).expect("valid tensor array"))] #[case::constant_vector(Vector::constant_array(&[3.0, 4.0], 2).expect("valid vector array"))] +#[case::nullable_vector(nullable_vector_input().expect("valid vector array"))] fn normalize_round_trips(#[case] input: ArrayRef) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let normalized_array = normalize(input.clone(), &mut ctx)?; @@ -401,24 +474,36 @@ fn normalize_zeroes_rows_with_zero_norms() -> VortexResult<()> { } #[test] -fn normalize_preserves_nulls_through_the_norms_child() -> VortexResult<()> { +fn normalize_moves_input_nulls_onto_the_array() -> VortexResult<()> { + // Row 1 is masked out but physically holds the unit vector `[1.0, 0.0]`, so a norm of 1.0 would + // survive into the norms child if the null were not applied. let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 1.0])?; let input = MaskedArray::try_new(input, Validity::from_iter([true, false, true]))?.into_array(); let mut ctx = SESSION.create_execution_ctx(); let normalized_array = normalize(input, &mut ctx)?; + assert!(normalized_array.dtype().is_nullable()); assert!(!normalized_array.normalized().dtype().is_nullable()); - assert!(normalized_array.norms().dtype().is_nullable()); + assert!(!normalized_array.norms().dtype().is_nullable()); let mask = normalized_array - .as_ref() - .validity()? + .normalized_validity() .execute_mask(3, &mut ctx)?; assert!(mask.value(0)); assert!(!mask.value(1)); assert!(mask.value(2)); + // Both children are zeroed at the null row rather than carrying whatever the masked-out storage + // happened to hold, so no garbage reaches a downstream lossy encoding. + let norms: PrimitiveArray = normalized_array.norms().clone().execute(&mut ctx)?; + assert_close(&norms.as_slice::()[1..2], &[0.0]); + + let normalized: ExtensionArray = normalized_array.normalized().clone().execute(&mut ctx)?; + let storage: FixedSizeListArray = normalized.storage_array().clone().execute(&mut ctx)?; + let elements: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; + assert_close(&elements.as_slice::()[2..4], &[0.0, 0.0]); + Ok(()) } @@ -465,6 +550,34 @@ fn filter_stays_encoded_and_decodes_correctly() -> VortexResult<()> { assert_tensor_arrays_eq(filtered, expected) } +/// The push-down rules rebuild the array from sliced/filtered children, so they have to carry the +/// validity along with them. +#[test] +fn slice_and_filter_carry_the_validity() -> VortexResult<()> { + let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0, 5.0, 12.0])?; + let input = + MaskedArray::try_new(input, Validity::from_iter([true, false, true, false]))?.into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + let normalized_array = normalize(input, &mut ctx)?.into_array(); + + let sliced = normalized_array + .slice(1..3)? + .execute_until::(&mut ctx)?; + assert!(sliced.is::()); + assert!(!sliced.is_valid(0, &mut ctx)?); + assert!(sliced.is_valid(1, &mut ctx)?); + + let filtered = normalized_array + .filter(Mask::from_iter([true, true, false, false]))? + .execute_until::(&mut ctx)?; + assert!(filtered.is::()); + assert!(filtered.is_valid(0, &mut ctx)?); + assert!(!filtered.is_valid(1, &mut ctx)?); + + Ok(()) +} + #[test] fn take_decodes_correctly() -> VortexResult<()> { let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0, 5.0, 12.0])?; @@ -496,13 +609,53 @@ fn scalar_at_reads_a_single_denormalized_row() -> VortexResult<()> { Ok(()) } +/// `scalar_at` collapses the row's norm to a one-element constant, which routes it through the +/// constant-norms path. A norm within `unit_norm_tolerance` of 1.0 must not be treated as the +/// identity there, or a per-row read would disagree with a bulk decode in the last bits. +#[test] +fn scalar_at_matches_bulk_decode_for_near_unit_norms() -> VortexResult<()> { + let near_unit = 1.0f64 + 2.0 * f64::EPSILON; + let normalized = vector_array(2, &[1.0f64, 0.0])?; + let norms = PrimitiveArray::from_iter([near_unit]).into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + let normalized_array = + Normalized::try_new(normalized, norms, Validity::NonNullable, &mut ctx)?.into_array(); + let bulk: ArrayRef = normalized_array.clone().execute(&mut ctx)?; + + let row = normalized_array.execute_scalar(0, &mut ctx)?; + assert_eq!(row, bulk.execute_scalar(0, &mut ctx)?); + + // Both paths must have actually applied the norm, not just agreed on skipping it. + let ext: ExtensionArray = bulk.execute(&mut ctx)?; + let storage: FixedSizeListArray = ext.storage_array().clone().execute(&mut ctx)?; + let elements: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; + assert_eq!(elements.as_slice::()[0], near_unit); + + Ok(()) +} + +#[test] +fn scalar_at_reads_a_nullable_column() -> VortexResult<()> { + let input = nullable_vector_input()?; + let mut ctx = SESSION.create_execution_ctx(); + let normalized_array = normalize(input.clone(), &mut ctx)?.into_array(); + + for i in 0..input.len() { + assert_eq!( + normalized_array.execute_scalar(i, &mut ctx)?, + input.execute_scalar(i, &mut ctx)?, + ); + } + + Ok(()) +} + // ============================================================================= // Serialization // ============================================================================= /// Round-trips through the array plugin registry, which is the same path a Vortex file takes. -/// `normalize` leaves the normalized child non-nullable and the norms child nullable -/// whenever the input is, so this exercises two different per-child nullabilities. #[rstest] #[case::vector(vector_array(3, &[3.0, 4.0, 0.0, 0.0, 0.0, 0.0]).expect("valid vector array"))] #[case::fixed_shape_tensor( @@ -533,52 +686,36 @@ fn serde_round_trip(#[case] input: ArrayRef) -> VortexResult<()> { assert_tensor_arrays_eq(recovered, original) } -fn nullable_vector_input() -> VortexResult { - let vectors = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0])?; - - Ok(MaskedArray::try_new(vectors, Validity::from_iter([true, false, true]))?.into_array()) -} - -/// The parent dtype supplies the tensor shape and element ptype, while metadata records the two -/// independently nullable children. +/// The array carries no metadata: the parent dtype supplies the tensor shape, element ptype, and +/// nullability, both children's dtypes follow from it, and the validity child shows up in the child +/// count. #[test] -fn serialized_metadata_pins_child_nullabilities() -> VortexResult<()> { +fn serialization_carries_no_metadata() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let input = MaskedArray::try_new( - vector_array(2, &[3.0, 4.0, 1.0, 0.0])?, - Validity::from_iter([true, false]), - )? - .into_array(); - let normalized_array = normalize(input, &mut ctx)?; - - let bytes = SESSION - .array_serialize(&normalized_array.clone().into_array())? - .expect("Normalized must serialize"); - let metadata = NormalizedMetadata::decode(bytes.as_slice())?; + let nullable = normalize(nullable_vector_input()?, &mut ctx)?.into_array(); + let non_nullable = normalize(vector_array(2, &[3.0, 4.0, 1.0, 0.0])?, &mut ctx)?.into_array(); + + for array in [&nullable, &non_nullable] { + let bytes = SESSION + .array_serialize(array)? + .expect("Normalized must serialize"); + assert!(bytes.is_empty(), "Normalized must not serialize metadata"); + } - assert_eq!( - metadata.normalized_is_nullable, - normalized_array.normalized().dtype().is_nullable(), - ); - assert_eq!( - metadata.norms_is_nullable, - normalized_array.norms().dtype().is_nullable(), - ); + assert_eq!(nullable.nchildren(), NormalizedSlots::COUNT); + assert_eq!(non_nullable.nchildren(), NormalizedSlots::COUNT - 1); Ok(()) } #[test] -fn serde_round_trip_preserves_normalized_nullability() -> VortexResult<()> { - let normalized = MaskedArray::try_new( - vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?, - Validity::from_iter([true, false]), - )? - .into_array(); +fn serde_round_trip_preserves_the_stored_validity() -> VortexResult<()> { + let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let original = Normalized::try_new(normalized, norms, &mut ctx)?.into_array(); + let validity = Validity::from_iter([true, false]); + let original = Normalized::try_new(normalized, norms, validity, &mut ctx)?.into_array(); let children: Vec = original.children(); let metadata = SESSION .array_serialize(&original)? @@ -594,10 +731,16 @@ fn serde_round_trip_preserves_normalized_nullability() -> VortexResult<()> { &SESSION, )?; + assert_eq!(recovered.dtype(), original.dtype()); + let recovered = recovered.as_::(); - assert!(recovered.normalized().dtype().is_nullable()); + assert!(!recovered.normalized().dtype().is_nullable()); assert!(!recovered.norms().dtype().is_nullable()); + let mask = recovered.normalized_validity().execute_mask(2, &mut ctx)?; + assert!(mask.value(0)); + assert!(!mask.value(1)); + Ok(()) } @@ -645,9 +788,32 @@ fn scheme_matches_tensor_columns(#[case] input: ArrayRef) -> VortexResult<()> { Ok(()) } -#[test] -fn compressor_emits_the_dedicated_encoding() -> VortexResult<()> { - let input = collinear_vectors(1024)?; +/// The scheme reports `AlwaysUse`, so a canonical array it claims is never offered to another +/// scheme. Claiming a non-float tensor would abort the whole column's compression on the float-only +/// gate in `compress` rather than falling through. +#[rstest] +#[case::integer_tensor(tensor_array(&[2], &[1i32, 2, 3, 4]).expect("valid tensor array"))] +#[case::non_tensor_extension(non_tensor_extension_array().expect("valid date array"))] +fn scheme_does_not_match_non_float_tensors(#[case] input: ArrayRef) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let canonical: Canonical = input.clone().execute(&mut ctx)?; + + assert!(!NormalizedScheme.matches(&canonical)); + + let compressor = BtrBlocksCompressorBuilder::default() + .with_new_scheme(&NormalizedScheme) + .build(); + let compressed = compressor.compress(&input, &mut ctx)?; + + assert_ne!(compressed.encoding_id(), ArrayVTable::id(&Normalized)); + + Ok(()) +} + +#[rstest] +#[case::non_nullable(collinear_vectors(1024).expect("valid vector array"))] +#[case::nullable(nullable_collinear_vectors(1024).expect("valid vector array"))] +fn compressor_emits_the_dedicated_encoding(#[case] input: ArrayRef) -> VortexResult<()> { let compressor = BtrBlocksCompressorBuilder::default() .with_new_scheme(&NormalizedScheme) .build(); @@ -659,3 +825,10 @@ fn compressor_emits_the_dedicated_encoding() -> VortexResult<()> { assert!(compressed.nbytes() < input.nbytes()); assert_tensor_arrays_eq(compressed, input) } + +fn nullable_collinear_vectors(rows: usize) -> VortexResult { + let vectors = collinear_vectors(rows)?; + let validity = Validity::from_iter((0..rows).map(|i| i % 8 != 0)); + + Ok(MaskedArray::try_new(vectors, validity)?.into_array()) +} diff --git a/vortex-tensor/src/encodings/normalized/validate.rs b/vortex-tensor/src/encodings/normalized/validate.rs index 3c27231b639..76196011e73 100644 --- a/vortex-tensor/src/encodings/normalized/validate.rs +++ b/vortex-tensor/src/encodings/normalized/validate.rs @@ -8,7 +8,9 @@ use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; use vortex_array::match_each_float_ptype; +use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -18,7 +20,7 @@ use crate::utils::extract_flat_elements; use crate::utils::unit_norm_tolerance; use crate::utils::validate_tensor_float_input; -/// Validates the structural invariants of a [`Normalized`] array's children. +/// Validates the structural invariants of a [`Normalized`] array's slots. /// /// These are the cheap, dtype-and-length checks that every [`NormalizedArray`] upholds, whichever /// constructor built it. They run on construction and on deserialization. @@ -28,6 +30,7 @@ use crate::utils::validate_tensor_float_input; pub(super) fn validate_normalized_children( normalized: &ArrayRef, norms: &ArrayRef, + validity: Option<&ArrayRef>, dtype: &DType, len: usize, ) -> VortexResult<()> { @@ -47,44 +50,71 @@ pub(super) fn validate_normalized_children( let tensor_match = validate_tensor_float_input(normalized.dtype())?; let element_ptype = tensor_match.element_ptype(); - let DType::Primitive(norms_ptype, _) = norms.dtype() else { - vortex_bail!( - "Normalized norms must be a primitive float array, got {}", - norms.dtype(), - ); - }; + // Both children are non-nullable so that the array's validity is the column's only null + // record, which is what lets the decode and read-through paths skip dtype widening entirely. vortex_ensure_eq!( - *norms_ptype, - element_ptype, - "Normalized norms dtype must match the normalized element dtype ({element_ptype}), \ - got {norms_ptype}", + *normalized.dtype(), + dtype.as_nonnullable(), + "Normalized normalized child must be the non-nullable array dtype ({}), got {}", + dtype.as_nonnullable(), + normalized.dtype(), ); - let expected = normalized - .dtype() - .union_nullability(norms.dtype().nullability()); + let expected_norms_dtype = DType::Primitive(element_ptype, Nullability::NonNullable); vortex_ensure_eq!( - *dtype, - expected, - "Normalized dtype must be the union of its children's nullability ({expected}), got {dtype}", + *norms.dtype(), + expected_norms_dtype, + "Normalized norms must be a non-nullable {element_ptype} column ({expected_norms_dtype}), \ + got {}", + norms.dtype(), ); + if let Some(validity) = validity { + vortex_ensure!( + dtype.is_nullable(), + "Normalized must not carry a validity slot when its dtype is non-nullable ({dtype})", + ); + vortex_ensure_eq!( + *validity.dtype(), + Validity::DTYPE, + "Normalized validity must be a {} column, got {}", + Validity::DTYPE, + validity.dtype(), + ); + vortex_ensure_eq!( + validity.len(), + len, + "Normalized validity must have the array length ({len}), got {}", + validity.len(), + ); + } + Ok(()) } /// Validates that `normalized` and (when supplied) the matching `norms` jointly satisfy the /// semantic [`Normalized`] invariants: /// -/// - Every valid row of `normalized` has L2 norm `1.0` or `0.0`, within the tolerance implied by -/// the element precision. -/// - When `norms` is supplied, every stored norm is non-negative and any row whose stored norm is -/// `0.0` is exactly the zero vector in `normalized`. +/// - Every row of `normalized` has L2 norm `1.0` or `0.0`, within the tolerance implied by the +/// element precision. +/// - When `norms` is supplied, every stored norm is non-negative, and a row is the zero vector in +/// `normalized` exactly when its stored norm is `0.0`. /// -/// This costs `O(len * list_size)`, which is why it is a separate step rather than part of the -/// encoding's structural validation. +/// The second half is symmetric on purpose. Checking only one direction would accept +/// `normalized = [0.0, 0.0]` paired with `norms = [5.0]`, which decodes to `[0.0, 0.0]` while +/// [`L2Norm`] reads the stored `5.0` straight back — precisely the split that +/// [`Normalized::try_new`] promises is lossless. +/// +/// This scans every row, so it costs `O(len * list_size)`, which is why it is a separate step +/// rather than part of the encoding's structural validation. Rows a caller intends to be null are +/// scanned too; [`normalize`] zeroes both children at null positions, which satisfies both +/// directions of the zero-norm rule. /// /// [`Normalized`]: crate::encodings::normalized::Normalized -pub fn validate_l2_normalized_rows_against_norms( +/// [`Normalized::try_new`]: crate::encodings::normalized::Normalized::try_new +/// [`normalize`]: crate::encodings::normalized::normalize +/// [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm +pub fn validate_normalized_rows( normalized: &ArrayRef, norms: Option<&ArrayRef>, ctx: &mut ExecutionCtx, @@ -123,29 +153,15 @@ pub fn validate_l2_normalized_rows_against_norms( } let normalized: ExtensionArray = normalized.clone().execute(ctx)?; - let normalized_validity = normalized.as_ref().validity()?; - let flat = extract_flat_elements(normalized.storage_array(), tensor_flat_size, ctx)?; let norms = norms .map(|norms| norms.clone().execute::(ctx)) .transpose()?; - let combined_validity = match &norms { - Some(norms) => normalized_validity.and(norms.validity()?)?, - None => normalized_validity, - }; - - // Resolve validity to a mask once rather than probing it per row. - let combined_valid = combined_validity.execute_mask(row_count, ctx)?; - match_each_float_ptype!(element_ptype, |T| { let stored_norms = norms.as_ref().map(|norms| norms.as_slice::()); for i in 0..row_count { - if !combined_valid.value(i) { - continue; - } - let (row_norm_sq, is_zero_row) = flat.row::(i) .iter() @@ -168,12 +184,13 @@ pub fn validate_l2_normalized_rows_against_norms( "Normalized norms must be non-negative, but row {i} has {stored_norm_f64:.6}", ); - if stored_norm_f64 == 0.0 { - vortex_ensure!( - is_zero_row, - "Normalized normalized child must be all zeros when norms row {i} is 0.0", - ); - } + vortex_ensure!( + is_zero_row == (stored_norm_f64 == 0.0), + "Normalized normalized child must be all zeros exactly when its stored norm is \ + 0.0, but row {i} pairs a {} normalized row with a stored norm of \ + {stored_norm_f64:.6}", + if is_zero_row { "zero" } else { "nonzero" }, + ); } } }); diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index ca8fcf0efd4..34f4cee8ca3 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -117,7 +117,7 @@ impl ScalarFnVTable for CosineSimilarity { let len = args.row_count(); // If either side is a constant tensor-like extension array, eagerly normalize the single - // stored row and re-encode it as an `Normalized` whose children are both `ConstantArray`s. + // stored row and re-encode it as a `Normalized` whose children are both `ConstantArray`s. // The `Normalized` fast path below then picks it up. if let Some(normalized_array) = try_build_constant_normalized(&lhs_ref, len, ctx)? { lhs_ref = normalized_array.into_array(); @@ -580,14 +580,15 @@ mod tests { } #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on rhs). + fn both_normalized_null_rows() -> VortexResult<()> { + // Row 0: valid, row 1: null (via the stored validity on rhs). let mut ctx = SESSION.create_execution_ctx(); let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; let normalized_r = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_r = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); - let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); + let norms_r = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + let validity = Validity::from_iter([true, false]); + let rhs = Normalized::try_new(normalized_r, norms_r, validity, &mut ctx)?.into_array(); let scalar_fn = CosineSimilarity::new().erased(); let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; @@ -610,13 +611,17 @@ mod tests { // Intentionally violates the unit-norm invariant by pairing a nonzero normalized row // with a stored norm of `0.0`, mimicking lossy storage. // SAFETY: The children are structurally valid. - let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.into_array(); + let lhs = + unsafe { Normalized::new_unchecked(normalized_l, norms_l, Validity::NonNullable) } + .into_array(); let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); // Same as above for the rhs operand. // SAFETY: The children are structurally valid. - let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.into_array(); + let rhs = + unsafe { Normalized::new_unchecked(normalized_r, norms_r, Validity::NonNullable) } + .into_array(); // `dot(normalized_l, normalized_r) = 1.0`, but the authoritative stored norms are both // `0.0`, so cosine similarity must be `0.0`. @@ -635,7 +640,9 @@ mod tests { // Intentionally pairs a nonzero normalized row with a stored norm of `0.0`, mimicking // lossy storage where the stored norm is authoritative. // SAFETY: The children are structurally valid. - let normalized_array = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); + let normalized_array = + unsafe { Normalized::new_unchecked(normalized, norms, Validity::NonNullable) } + .into_array(); let plain = tensor_array(&[2], &[1.0, 0.0])?; diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 53ae82eb4a2..1dbaa11c88d 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -468,13 +468,14 @@ mod tests { } #[test] - fn both_normalized_null_norms() -> VortexResult<()> { - // Row 0: valid, row 1: null (via nullable norms on lhs). + fn both_normalized_null_rows() -> VortexResult<()> { + // Row 0: valid, row 1: null (via the stored validity on lhs). let normalized_l = tensor_array(&[2], &[0.6, 0.8, 1.0, 0.0])?; - let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); + let norms_l = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let lhs = Normalized::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); + let validity = Validity::from_iter([true, false]); + let lhs = Normalized::try_new(normalized_l, norms_l, validity, &mut ctx)?.into_array(); let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; let scalar_fn = InnerProduct::new().erased(); diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index b7e9060ed3f..80def79633c 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -11,6 +11,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; use vortex_array::arrays::ScalarFnArray; @@ -35,6 +36,7 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -44,6 +46,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::encodings::normalized::Normalized; +use crate::encodings::normalized::NormalizedArrayExt; use crate::matcher::AnyTensor; use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; @@ -131,8 +134,15 @@ impl ScalarFnVTable for L2Norm { // L2Norm over a `Normalized`-encoded column is defined to read back the authoritative stored // norms. Callers of lossy encodings opt into that storage semantics instead of forcing a // decode-and-recompute path here. - if input_ref.is::() { + // + // The stored norms are non-nullable — nulls live on the `Normalized` array itself — so a + // nullable input needs its null map reattached to reach `norm_dtype`. + if let Some(normalized_array) = input_ref.as_opt::() { let (_, norms) = extract_normalized_children(&input_ref); + let norms = match normalized_array.normalized_validity() { + Validity::NonNullable => norms, + validity => MaskedArray::try_new(norms, validity)?.into_array(), + }; vortex_ensure_eq!(norms.dtype(), &norm_dtype); return Ok(norms); } @@ -275,11 +285,13 @@ mod tests { use vortex_array::validity::Validity; use vortex_error::VortexResult; + use crate::encodings::normalized::Normalized; use crate::scalar_fns::l2_norm::L2Norm; use crate::tests::SESSION; use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::literal_vector_array; + use crate::utils::test_helpers::normalized_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; @@ -408,6 +420,44 @@ mod tests { Ok(()) } + /// The read-through returns the stored norms child, which is always non-nullable — nulls live + /// on the [`Normalized`] array itself. A nullable input therefore needs its null map reattached + /// to reach the declared return dtype, which used to be an assertion failure instead. + #[test] + fn reads_through_a_nullable_normalized_column() -> VortexResult<()> { + let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; + let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + let validity = Validity::from_iter([true, false]); + let input = Normalized::try_new(normalized, norms, validity, &mut ctx)?.into_array(); + + let result = ScalarFnArray::try_new(L2Norm::new().erased(), vec![input])?.into_array(); + let prim: PrimitiveArray = result.execute(&mut ctx)?; + + assert_eq!( + prim.dtype(), + &DType::Primitive(PType::F64, Nullability::Nullable) + ); + assert!(prim.is_valid(0, &mut ctx)?); + assert!(!prim.is_valid(1, &mut ctx)?); + assert_close(&[prim.as_slice::()[0]], &[5.0]); + + Ok(()) + } + + /// A non-nullable [`Normalized`] column reads straight back as the stored norms child, with no + /// masking wrapper in the way. + #[test] + fn reads_through_a_non_nullable_normalized_column() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 3.0], &mut ctx)?; + + assert_close(&eval_l2_norm(input)?, &[5.0, 3.0]); + + Ok(()) + } + #[rstest] #[case::fixed_shape_tensor(l2_norm_tensor_child())] #[case::vector(l2_norm_vector_child())] diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 488694bd47f..d1fe6a2382f 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -370,9 +370,9 @@ pub mod test_helpers { ConstantArray::new(ext_scalar, len).into_array() } - /// Creates a [`Normalized`] array from pre-normalized tensor elements and matching norms. The - /// caller must ensure every row of `normalized_elements` is unit-norm or zero, since this - /// goes through the checked constructor. + /// Creates a non-nullable [`Normalized`] array from pre-normalized tensor elements and matching + /// norms. The caller must ensure every row of `normalized_elements` is unit-norm or zero, since + /// this goes through the checked constructor. pub fn normalized_array( shape: &[usize], normalized_elements: &[T], @@ -382,7 +382,8 @@ pub mod test_helpers { let normalized = tensor_array(shape, normalized_elements)?; let norms = PrimitiveArray::new(Buffer::copy_from(norms), Validity::NonNullable).into_array(); - Ok(Normalized::try_new(normalized, norms, ctx)?.into_array()) + + Ok(Normalized::try_new(normalized, norms, Validity::NonNullable, ctx)?.into_array()) } /// Asserts that each element in `actual` is within `1e-10` of the corresponding `expected` diff --git a/vortex/src/editions/unstable/v2026_04.rs b/vortex/src/editions/unstable/v2026_04.rs index 90955f01d04..f6e302e83e9 100644 --- a/vortex/src/editions/unstable/v2026_04.rs +++ b/vortex/src/editions/unstable/v2026_04.rs @@ -21,7 +21,7 @@ pub static DECLARATION: EditionDeclaration = EditionDeclaration { &"vortex.patched", &"vortex.tensor.cosine_similarity", &"vortex.tensor.inner_product", - &"vortex.tensor.normalized", &"vortex.tensor.l2_norm", + &"vortex.tensor.normalized", ], };