diff --git a/vortex-tensor/benches/normalized.rs b/vortex-tensor/benches/normalized.rs index 39ef3396866..65aaef37128 100644 --- a/vortex-tensor/benches/normalized.rs +++ b/vortex-tensor/benches/normalized.rs @@ -1,13 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Baseline throughput for decoding the `Normalized` encoding over tensor columns. +//! Baseline throughput for the `Normalized` encoding over tensor columns, in both directions: +//! `non_nullable` and `nullable` decode, `encode_non_nullable` and `encode_nullable` split a plain +//! tensor column into the two children. //! //! The arms vary vector width and input nullability. Their names are intended to remain stable //! across implementation changes so CodSpeed can compare them against `develop`. //! //! Rows are derived from a fixed element budget rather than fixed per arm, so widening a vector //! trades rows for elements instead of multiplying the work. See [`ELEMENTS`]. +//! +//! Width is the axis that separates the two costs the encode path pays. A narrow vector leaves the +//! per-row work visible, while a wide one buries it under the per-element division. #![expect(clippy::unwrap_used)] @@ -24,10 +29,11 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_tensor::encodings::normalized::Normalized; +use vortex_tensor::encodings::normalized::normalize; use vortex_tensor::vector::Vector; -// Decoding allocates the output inside the timed region, so use the vendored allocator instead -// of measuring glibc differences between CodSpeed runner images. +// Both directions allocate their output inside the timed region, so use the vendored allocator +// instead of measuring glibc differences between CodSpeed runner images. #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; @@ -56,12 +62,34 @@ fn normalized_vectors(width: usize) -> ArrayRef { Vector::try_new_vector_array(storage).unwrap() } +/// Vectors of varying magnitude, so no row is a zero vector and every row takes the dividing +/// branch in [`normalize`]. +fn plain_vectors(width: usize) -> ArrayRef { + let rows = ELEMENTS / width; + let elements: Buffer = (0..rows) + .flat_map(|row| (0..width).map(move |i| (row + 1) as f64 * (i + 1) as f64)) + .collect(); + let storage = FixedSizeListArray::new( + elements.into_array(), + u32::try_from(width).unwrap(), + Validity::NonNullable, + rows, + ) + .into_array(); + Vector::try_new_vector_array(storage).unwrap() +} + +/// Masks every eighth row. The row count is `ELEMENTS / width`, matching both vector builders. +fn sparse_nulls(width: usize) -> Validity { + Validity::from_iter((0..ELEMENTS / width).map(|i| !i.is_multiple_of(8))) +} + fn norms(rows: usize) -> ArrayRef { let values: Buffer = (0..rows).map(|i| 1.0 + ((i % 13) as f64) / 13.0).collect(); PrimitiveArray::new(values, Validity::NonNullable).into_array() } -fn bench_normalized(bencher: Bencher, normalized: ArrayRef) { +fn bench_decode(bencher: Bencher, normalized: ArrayRef, validity: Validity) { let session = vortex_array::array_session(); let rows = normalized.len(); let norms = norms(rows); @@ -69,7 +97,13 @@ fn bench_normalized(bencher: Bencher, normalized: ArrayRef) { .counter(ItemsCount::new(rows)) .with_inputs(|| { let mut ctx = session.create_execution_ctx(); - let array = Normalized::try_new(normalized.clone(), norms.clone(), &mut ctx).unwrap(); + let array = Normalized::try_new( + normalized.clone(), + norms.clone(), + validity.clone(), + &mut ctx, + ) + .unwrap(); (array, ctx) }) .bench_values(|(array, mut ctx)| { @@ -80,16 +114,35 @@ fn bench_normalized(bencher: Bencher, normalized: ArrayRef) { }); } +fn bench_encode(bencher: Bencher, input: ArrayRef) { + let session = vortex_array::array_session(); + let rows = input.len(); + bencher + .counter(ItemsCount::new(rows)) + .with_inputs(|| (input.clone(), session.create_execution_ctx())) + .bench_values(|(input, mut ctx)| normalize(input, &mut ctx).unwrap()); +} + #[divan::bench(args = WIDTHS)] fn non_nullable(bencher: Bencher, width: usize) { - bench_normalized(bencher, normalized_vectors(width)); + bench_decode(bencher, normalized_vectors(width), Validity::NonNullable); } #[divan::bench(args = WIDTHS)] fn nullable(bencher: Bencher, width: usize) { - let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0)); - let normalized = MaskedArray::try_new(normalized_vectors(width), validity) + bench_decode(bencher, normalized_vectors(width), sparse_nulls(width)); +} + +#[divan::bench(args = WIDTHS)] +fn encode_non_nullable(bencher: Bencher, width: usize) { + bench_encode(bencher, plain_vectors(width)); +} + +#[divan::bench(args = WIDTHS)] +fn encode_nullable(bencher: Bencher, width: usize) { + let input = MaskedArray::try_new(plain_vectors(width), sparse_nulls(width)) .unwrap() .into_array(); - bench_normalized(bencher, normalized); + + bench_encode(bencher, input); } diff --git a/vortex-tensor/src/encodings/normalized/array.rs b/vortex-tensor/src/encodings/normalized/array.rs index c60bf556b02..fa357c2e2a9 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_bail; +use vortex_error::vortex_ensure; 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,23 +50,34 @@ 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 /// /// [`new_unchecked`](Self::new_unchecked) deliberately skips the semantic scan so that -/// `normalized` may be an *approximation* of the unit-norm direction, such as a quantized child. +/// `normalized` may be an _approximation_ of the unit-norm direction, such as a quantized child. /// The stored norms stay authoritative in that case, and the read-through rules in /// [`L2Norm`], [`InnerProduct`], and [`CosineSimilarity`] are defined against the stored children /// rather than against decoded coordinates. Those operators may therefore return slightly @@ -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,66 +129,56 @@ 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 carrying whatever nullability `validity` +/// implies. `validity` is the sole source of the array's nullability, as it is for every canonical +/// container, so a nullable child cannot quietly widen the parent. +fn normalized_parts( + normalized: ArrayRef, + norms: ArrayRef, + validity: Validity, +) -> ArrayParts { + let len = normalized.len(); + let dtype = normalized.dtype().with_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) } impl VTable for Normalized { @@ -185,7 +201,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 +224,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 +242,37 @@ 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 = match children.len() { + DATA_CHILDREN => Validity::from(dtype.nullability()), + NormalizedSlots::COUNT => { + vortex_ensure!( + dtype.is_nullable(), + "Normalized validity child requires a nullable dtype, got {dtype}", + ); + Validity::Array(children.get(NormalizedSlots::VALIDITY, &Validity::DTYPE, len)?) + } + other => vortex_bail!( + "Normalized expects {DATA_CHILDREN} or {} children, got {other}", + NormalizedSlots::COUNT, + ), + }; + + Ok(normalized_parts(normalized, norms, validity)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -250,10 +281,10 @@ impl VTable for Normalized { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { let dtype = array.dtype().clone(); + let validity = array.validity()?; let slots = array.slots_view(); - denormalize(slots.normalized, slots.norms, array.len(), dtype, ctx) - .map(ExecutionResult::done) + denormalize(slots.normalized, slots.norms, validity, dtype, ctx).map(ExecutionResult::done) } fn reduce_parent( @@ -267,10 +298,11 @@ impl VTable for Normalized { impl ValidityVTable for Normalized { fn validity(array: ArrayView<'_, Normalized>) -> VortexResult { - array - .normalized() - .validity()? - .and(array.norms().validity()?) + // Both children are non-nullable, so the slot is the column's complete null information. + Ok(child_to_validity( + array.slots_view().validity, + array.dtype().nullability(), + )) } } @@ -284,12 +316,17 @@ 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(), - 1, - array.dtype().clone(), + validity, + dtype, ctx, )?; diff --git a/vortex-tensor/src/encodings/normalized/compress.rs b/vortex-tensor/src/encodings/normalized/compress.rs index 727483296e4..08b72847357 100644 --- a/vortex-tensor/src/encodings/normalized/compress.rs +++ b/vortex-tensor/src/encodings/normalized/compress.rs @@ -17,6 +17,7 @@ use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; @@ -40,6 +41,7 @@ use crate::encodings::normalized::Normalized; use crate::encodings::normalized::NormalizedArray; 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( @@ -106,38 +116,42 @@ impl Scheme for NormalizedScheme { exec_ctx, )?; - // SAFETY: Cascading preserves the split's child lengths and dtypes. - Ok(unsafe { Normalized::new_unchecked(normalized, norms) }.into_array()) + let validity = normalized_array.validity()?; + + // SAFETY: Cascading preserves the split's child lengths and dtypes, and the validity is + // carried over from the split unchanged. + Ok(unsafe { Normalized::new_unchecked(normalized, norms, 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,31 +159,49 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult(); + let norm_values = norms.as_slice::(); let total_elements = row_count * tensor_flat_size; let mut elements = BufferMut::::with_capacity(total_elements); for i in 0..row_count { - let is_valid = norms_valid.value(i); let norm = norm_values[i]; + // A null row arrives here with a filled zero norm, so its coordinates are zeroed + // alongside the genuine zero vectors rather than carrying whatever the masked-out + // storage happened to hold. + // // 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,8 +210,6 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult 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 +287,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..c958f8fbeab 100644 --- a/vortex-tensor/src/encodings/normalized/execute.rs +++ b/vortex-tensor/src/encodings/normalized/execute.rs @@ -25,34 +25,33 @@ use vortex_error::VortexResult; use crate::matcher::AnyTensor; use crate::utils::extract_flat_elements; -use crate::utils::unit_norm_tolerance; +use crate::utils::reattach_validity; /// 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. Both children are +/// non-nullable, so the reconstructed column carries the parent's nullability rather than either +/// child's. /// /// [`NormalizedArray`]: crate::encodings::normalized::NormalizedArray pub(super) fn denormalize( normalized: &ArrayRef, norms: &ArrayRef, - row_count: usize, + validity: Validity, 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); } + // Both children are validated to have the array's length, so either one gives the row count. + let row_count = normalized.len(); + let normalized: ExtensionArray = normalized.clone().execute(ctx)?; let norms: PrimitiveArray = norms.clone().execute(ctx)?; @@ -76,9 +75,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 +86,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 reattach_validity(normalized.clone(), validity); } let normalized: ExtensionArray = normalized.clone().execute(ctx)?; @@ -106,8 +107,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()) }; diff --git a/vortex-tensor/src/encodings/normalized/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs index 545236bba7d..08a96815075 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. //! @@ -25,7 +31,6 @@ mod array; pub use array::Normalized; pub use array::NormalizedArray; pub use array::NormalizedArraySlotsExt; -pub use array::NormalizedMetadata; pub use array::NormalizedSlots; mod compress; @@ -41,7 +46,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..47e470abc48 100644 --- a/vortex-tensor/src/encodings/normalized/rules.rs +++ b/vortex-tensor/src/encodings/normalized/rules.rs @@ -37,12 +37,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.validity()?.slice(range.clone())?, ) } .into_array(), @@ -69,12 +70,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.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..0f3f26bc1b2 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; @@ -9,6 +8,7 @@ use vortex_array::ArrayVTable; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Extension; @@ -17,6 +17,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; @@ -33,10 +34,10 @@ use vortex_mask::Mask; use crate::encodings::normalized::Normalized; 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,60 @@ 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, // row 0, decodes to [3.0, 4.0] + 0.0, 0.0, // row 1, null and zeroed + 1.0, 0.0, // row 2, decodes to [2.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<()> { - 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(); +fn validity_comes_from_the_stored_null_map() -> VortexResult<()> { + let normalized = vector_array( + 2, + &[ + 1.0, 0.0, // row 0 + 1.0, 0.0, // row 1 + 1.0, 0.0, // row 2 + ], + )?; + 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 +217,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 actual = eval_normalized(normalized.clone(), norms)?; + 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)?; - assert_tensor_arrays_eq(actual, normalized) + assert_eq!(elements.as_slice::()[0], near_unit); + + Ok(()) } #[test] @@ -196,7 +247,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 +260,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 +320,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 +371,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 +393,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 +405,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 +413,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 +429,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,15 +488,25 @@ fn normalize_zeroes_rows_with_zero_norms() -> VortexResult<()> { } #[test] -fn normalize_preserves_nulls_through_the_norms_child() -> VortexResult<()> { - let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 1.0])?; +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, // row 0, norm 5.0 + 1.0, 0.0, // row 1, masked out despite being unit-norm + 0.0, 1.0, // row 2, norm 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() @@ -419,6 +516,16 @@ fn normalize_preserves_nulls_through_the_norms_child() -> VortexResult<()> { 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 +572,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 +631,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 +708,72 @@ 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])?; +/// 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 serialization_carries_no_metadata() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + 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"); + } - Ok(MaskedArray::try_new(vectors, Validity::from_iter([true, false, true]))?.into_array()) + assert_eq!(nullable.nchildren(), NormalizedSlots::COUNT); + assert_eq!(non_nullable.nchildren(), NormalizedSlots::COUNT - 1); + + Ok(()) } -/// The parent dtype supplies the tensor shape and element ptype, while metadata records the two -/// independently nullable children. +/// `validity_to_child` writes no child for `AllValid`, so a nullable column with no null rows +/// serializes with only its two data children and the parent dtype is all that is left to recover +/// the nullability from. The constant fast path reaches this state whenever the input dtype is +/// nullable but its stored row is not null. #[test] -fn serialized_metadata_pins_child_nullabilities() -> VortexResult<()> { +fn serde_round_trip_of_a_nullable_column_with_no_null_rows() -> 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 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 original = + Normalized::try_new(normalized, norms, Validity::AllValid, &mut ctx)?.into_array(); + let children: Vec = original.children(); + + assert!(original.dtype().is_nullable()); + assert_eq!(children.len(), NormalizedSlots::COUNT - 1); - let bytes = SESSION - .array_serialize(&normalized_array.clone().into_array())? + let metadata = SESSION + .array_serialize(&original)? .expect("Normalized must serialize"); - let metadata = NormalizedMetadata::decode(bytes.as_slice())?; + let recovered = ArrayPlugin::deserialize( + &Normalized, + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; - 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!(recovered.dtype(), original.dtype()); + assert!(matches!(recovered.validity()?, Validity::AllValid)); 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 +789,38 @@ 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.as_ref().validity()?.execute_mask(2, &mut ctx)?; + assert!(mask.value(0)); + assert!(!mask.value(1)); + + Ok(()) +} + +#[test] +fn deserialize_rejects_validity_child_for_non_nullable_dtype() -> VortexResult<()> { + let normalized = vector_array(2, &[0.6f64, 0.8, 1.0, 0.0])?; + let dtype = normalized.dtype().clone(); + let children = vec![ + normalized, + PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(), + BoolArray::from_iter([true, false]).into_array(), + ]; + + let error = ArrayPlugin::deserialize(&Normalized, &dtype, 2, &[], &[], &children, &SESSION) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("Normalized validity child requires a nullable dtype") + ); + Ok(()) } @@ -645,9 +868,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 +905,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..9ef3b86645e 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 only carry a validity slot when its dtype is nullable, got {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, which is 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..b4984068bf2 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -47,6 +47,7 @@ use crate::encodings::normalized::Normalized; use crate::matcher::AnyTensor; use crate::utils::extract_flat_elements; use crate::utils::extract_normalized_children; +use crate::utils::reattach_validity; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -131,8 +132,12 @@ 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. + // + // The stored norms are non-nullable, because nulls live on the `Normalized` array itself, so + // a nullable input needs its null map reattached to reach `norm_dtype`. if input_ref.is::() { let (_, norms) = extract_normalized_children(&input_ref); + let norms = reattach_validity(norms, input_ref.validity()?)?; vortex_ensure_eq!(norms.dtype(), &norm_dtype); return Ok(norms); } @@ -275,11 +280,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 +415,44 @@ mod tests { Ok(()) } + /// The read-through returns the stored norms child, which is always non-nullable because 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. + #[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..dd861cd50b8 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -9,6 +9,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::MaskedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn; use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; @@ -20,6 +21,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::dtype::proto::dtype as pb; use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -78,6 +80,21 @@ pub fn extract_normalized_children(array: &ArrayRef) -> (ArrayRef, ArrayRef) { ) } +/// Puts `validity` back onto a non-nullable array, leaving a non-nullable column untouched. +/// +/// Every [`Normalized`] child is non-nullable, so a path that hands a child back in place of the +/// decoded column has to restore the parent's null map first. The [`Validity::NonNullable`] arm +/// **must** stay: [`MaskedArray::try_new`] widens its child's dtype to nullable unconditionally, so +/// wrapping a non-nullable column would report a nullable dtype the parent never had. +/// +/// [`Normalized`]: crate::encodings::normalized::Normalized +pub(crate) fn reattach_validity(array: ArrayRef, validity: Validity) -> VortexResult { + match validity { + Validity::NonNullable => Ok(array), + validity => Ok(MaskedArray::try_new(array, validity)?.into_array()), + } +} + /// Validates that `input_dtype` is a float-valued tensor-like extension dtype. pub fn validate_tensor_float_input(input_dtype: &DType) -> VortexResult> { let ext = input_dtype @@ -370,9 +387,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 +399,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", ], };