From e15e3087e9ceaa88cd5da0bb5e24201e5615fb12 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 2 Aug 2026 00:45:07 +0000 Subject: [PATCH 1/2] Make L2Denorm an array encoding instead of a scalar function `L2Denorm` was registered as a `ScalarFnVTable`, but it never behaved like one. Its constructor took an `ExecutionCtx` and scanned both children to enforce a unit-norm invariant, `L2Norm` read its stored norms instead of recomputing, `CosineSimilarity` and `InnerProduct` reached into its physical children, and the compressor scheme named it as a produced encoding. Those are all properties of a physical decomposition, not of an operation over arbitrary well-typed values. Moves it to `vortex-tensor/src/encodings/l2_denorm/` as a real `VTable` with two slots (`normalized`, `norms`). Structural validation runs on construction and on deserialization, `try_new` additionally scans for the exact unit-norm invariant, and `try_new_trusted` skips that scan for lossy normalized children whose stored norms stay authoritative. Neither constructor is `unsafe`, since violating the contract produces wrong answers rather than undefined behavior. The encoding keeps the `vortex.tensor.l2_denorm` array ID and the same two-field metadata message, so the wire format is unchanged. Slice and filter now push down into both children through `reduce_parent`. The generic `ScalarFnArray` filter rule only fired when at most one child was non-constant, which for this encoding was almost never. Also makes `L2DenormScheme` cascade its two children like `TemporalScheme` does, which lets both `HACK TO SUPPORT L2 DENORMALIZATION` special cases come out of `CascadingCompressor`. The scheme now competes on measured size like every other scheme. Signed-off-by: Connor Tsui --- Cargo.lock | 2 + vortex-compressor/src/compressor/cascade.rs | 8 +- vortex-tensor/Cargo.toml | 2 + vortex-tensor/src/encodings/l2_denorm.rs | 61 - .../src/encodings/l2_denorm/array.rs | 294 +++++ .../src/encodings/l2_denorm/compress.rs | 297 +++++ .../src/encodings/l2_denorm/execute.rs | 145 +++ vortex-tensor/src/encodings/l2_denorm/mod.rs | 47 + .../src/encodings/l2_denorm/orientation.rs | 54 + .../src/encodings/l2_denorm/rules.rs | 83 ++ .../src/encodings/l2_denorm/tests.rs | 645 ++++++++++ .../src/encodings/l2_denorm/validate.rs | 182 +++ vortex-tensor/src/encodings/mod.rs | 2 +- vortex-tensor/src/lib.rs | 17 +- .../src/scalar_fns/cosine_similarity.rs | 57 +- vortex-tensor/src/scalar_fns/inner_product.rs | 16 +- vortex-tensor/src/scalar_fns/l2_denorm.rs | 1146 ----------------- vortex-tensor/src/scalar_fns/l2_norm.rs | 21 +- vortex-tensor/src/scalar_fns/mod.rs | 1 - vortex-tensor/src/utils.rs | 37 +- 20 files changed, 1836 insertions(+), 1281 deletions(-) delete mode 100644 vortex-tensor/src/encodings/l2_denorm.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/array.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/compress.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/execute.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/mod.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/orientation.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/rules.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/tests.rs create mode 100644 vortex-tensor/src/encodings/l2_denorm/validate.rs delete mode 100644 vortex-tensor/src/scalar_fns/l2_denorm.rs diff --git a/Cargo.lock b/Cargo.lock index d9f80fc40b7..db19fe686e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10493,9 +10493,11 @@ dependencies = [ "rstest", "vortex-array", "vortex-arrow", + "vortex-btrblocks", "vortex-buffer", "vortex-compressor", "vortex-error", + "vortex-mask", "vortex-session", ] diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index eb81e81241e..60a5aa3e6fc 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -23,7 +23,6 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::listview::ListViewArraySlotsExt; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::arrays::masked::MaskedArraySlotsExt; -use vortex_array::arrays::scalar_fn::AnyScalarFn; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::union::UnionArrayExt; use vortex_array::arrays::union::UnionArraySlotsExt; @@ -175,10 +174,6 @@ impl CascadingCompressor { compress_ctx, exec_ctx, )?; - // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! - if scheme_compressed.is::() { - return Ok(scheme_compressed); - } // A constant extension array (that might be masked) is already in its terminal // representation, and compressing the storage separately cannot do better. @@ -321,8 +316,7 @@ impl CascadingCompressor { let after_nbytes = compressed.nbytes(); let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); - // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!! - let accepted = after_nbytes < before_nbytes || compressed.is::(); + let accepted = after_nbytes < before_nbytes; trace::record_winner_compress_result( after_nbytes, diff --git a/vortex-tensor/Cargo.toml b/vortex-tensor/Cargo.toml index a6782eea309..abe63fde595 100644 --- a/vortex-tensor/Cargo.toml +++ b/vortex-tensor/Cargo.toml @@ -34,3 +34,5 @@ prost = { workspace = true } [dev-dependencies] rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-btrblocks = { workspace = true } +vortex-mask = { workspace = true } diff --git a/vortex-tensor/src/encodings/l2_denorm.rs b/vortex-tensor/src/encodings/l2_denorm.rs deleted file mode 100644 index 8bec6292ee0..00000000000 --- a/vortex-tensor/src/encodings/l2_denorm.rs +++ /dev/null @@ -1,61 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_array::ArrayId; -use vortex_array::ArrayRef; -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_compressor::CascadingCompressor; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::CompressorContext; -use vortex_compressor::scheme::EstimateVerdict; -use vortex_compressor::scheme::Scheme; -use vortex_compressor::stats::ArrayAndStats; -use vortex_error::VortexResult; - -use crate::matcher::AnyTensor; -use crate::scalar_fns::l2_denorm::L2Denorm; -use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm; - -#[derive(Debug)] -pub struct L2DenormScheme; - -impl Scheme for L2DenormScheme { - fn scheme_name(&self) -> &'static str { - "vortex.tensor.l2_denorm" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches!( - canonical, - Canonical::Extension(ext) if ext.ext_dtype().is::() - ) - } - - fn produced_encodings(&self) -> Vec { - vec![L2Denorm.id()] - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - _compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let l2_denorm = normalize_as_l2_denorm(data.array().clone(), exec_ctx)?; - Ok(l2_denorm.into_array()) - } -} diff --git a/vortex-tensor/src/encodings/l2_denorm/array.rs b/vortex-tensor/src/encodings/l2_denorm/array.rs new file mode 100644 index 00000000000..f3f621c3415 --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/array.rs @@ -0,0 +1,294 @@ +// 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; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::EmptyArrayData; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::array_slots; +use vortex_array::arrays::ConstantArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +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::with_empty_buffers; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::encodings::l2_denorm::execute::denormalize; +use crate::encodings::l2_denorm::rules::RULES; +use crate::encodings::l2_denorm::validate::validate_l2_denorm_children; +use crate::encodings::l2_denorm::validate::validate_l2_normalized_rows_against_norms; +use crate::utils::validate_tensor_float_input; + +/// An [`L2Denorm`]-encoded Vortex array. +pub type L2DenormArray = Array; + +/// The norm-split encoding for tensor-like columns. +/// +/// Row `i` decodes to `normalized[i] * norms[i]`, which is exactly the original tensor row when +/// `normalized[i]` is unit-norm. The encoding covers both logical tensor dtypes reachable through +/// [`AnyTensor`]: `Vector` and `FixedShapeTensor`. +/// +/// # Invariants +/// +/// Every [`L2DenormArray`] 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. +/// - both children have the array's length. +/// - the array dtype is `normalized.dtype().union_nullability(norms.nullability())`. +/// +/// 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 stored norm is non-negative. +/// - a stored norm of `0.0` is paired with an all-zero normalized row. +/// +/// # 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. +/// 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 +/// different answers than fully decoding both operands and recomputing. That difference is the +/// storage contract, not a separate lossy-compute mode. +/// +/// [`AnyTensor`]: crate::matcher::AnyTensor +/// [`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 L2Denorm; + +/// The two child arrays of an [`L2DenormArray`]. +#[array_slots(L2Denorm)] +pub struct L2DenormSlots { + /// The unit-norm (or zero) direction of each row, as a tensor-like extension array. + #[slot(0)] + pub normalized: ArrayRef, + + /// The authoritative L2 norm of each row, as a primitive float column. + #[slot(1)] + pub norms: ArrayRef, +} + +impl L2Denorm { + /// Builds an [`L2DenormArray`], 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 + /// `O(len * list_size)`. + /// + /// # Errors + /// + /// Returns an error if the children are structurally incompatible, or if they violate any of + /// the semantic invariants listed on [`L2Denorm`]. + pub fn try_new( + normalized: ArrayRef, + norms: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let len = normalized.len(); + let dtype = normalized + .dtype() + .union_nullability(norms.dtype().nullability()); + let slots = L2DenormSlots { 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 denorm = Array::try_from_parts( + ArrayParts::new(L2Denorm, dtype, len, EmptyArrayData).with_slots(slots), + )?; + validate_l2_normalized_rows_against_norms(denorm.normalized(), Some(denorm.norms()), ctx)?; + + Ok(denorm) + } + + /// Builds an [`L2DenormArray`] without validation. + /// + /// # Safety + /// + /// The caller must uphold the structural invariants listed on [`L2Denorm`]. 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. + /// + /// 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) -> L2DenormArray { + let len = normalized.len(); + let dtype = normalized + .dtype() + .union_nullability(norms.dtype().nullability()); + let slots = L2DenormSlots { normalized, norms }.into_slots(); + + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(L2Denorm, dtype, len, EmptyArrayData).with_slots(slots), + ) + } + } +} + +/// Metadata for a serialized [`L2DenormArray`]: its children's nullabilities. +/// +/// 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 L2DenormMetadata { + /// 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, +} + +impl VTable for L2Denorm { + type TypedArrayData = EmptyArrayData; + + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.tensor.l2_denorm"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let slots = L2DenormSlotsView::from_slots(slots); + + validate_l2_denorm_children(slots.normalized, slots.norms, dtype, len) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("L2DenormArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("L2DenormArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + L2DenormMetadata { + normalized_is_nullable: array.normalized().dtype().is_nullable(), + norms_is_nullable: array.norms().dtype().is_nullable(), + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + _buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = L2DenormMetadata::decode(metadata) + .map_err(|e| vortex_err!("Failed to decode L2DenormMetadata: {e}"))?; + + 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 = children.get(0, &normalized_dtype, len)?; + let norms = children.get(1, &norms_dtype, len)?; + let slots = L2DenormSlots { normalized, norms }.into_slots(); + + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + L2DenormSlots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let dtype = array.dtype().clone(); + let slots = array.slots_view(); + + denormalize(slots.normalized, slots.norms, array.len(), dtype, ctx) + .map(ExecutionResult::done) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +impl ValidityVTable for L2Denorm { + fn validity(array: ArrayView<'_, L2Denorm>) -> VortexResult { + array + .normalized() + .validity()? + .and(array.norms().validity()?) + } +} + +impl OperationsVTable for L2Denorm { + fn scalar_at( + array: ArrayView<'_, L2Denorm>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Denormalize a single row rather than the whole column: both children are collapsed to + // 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 row = denormalize( + &ConstantArray::new(normalized, 1).into_array(), + &ConstantArray::new(norms, 1).into_array(), + 1, + array.dtype().clone(), + ctx, + )?; + + row.execute_scalar(0, ctx) + } +} diff --git a/vortex-tensor/src/encodings/l2_denorm/compress.rs b/vortex-tensor/src/encodings/l2_denorm/compress.rs new file mode 100644 index 00000000000..18ce117cea6 --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/compress.rs @@ -0,0 +1,297 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::Float; +use num_traits::Zero; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Extension; +use vortex_array::arrays::ExtensionArray; +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::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_compressor::CascadingCompressor; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::encodings::l2_denorm::L2Denorm; +use crate::encodings::l2_denorm::L2DenormArray; +use crate::encodings::l2_denorm::L2DenormArraySlotsExt; +use crate::encodings::l2_denorm::L2DenormSlots; +use crate::matcher::AnyTensor; +use crate::scalar_fns::l2_norm::L2Norm; +use crate::utils::extract_constant_flat_row; +use crate::utils::extract_flat_elements; +use crate::utils::validate_tensor_float_input; + +/// The compression scheme that rewrites a tensor-like column into the [`L2Denorm`] encoding. +#[derive(Debug)] +pub struct L2DenormScheme; + +impl Scheme for L2DenormScheme { + fn scheme_name(&self) -> &'static str { + "vortex.tensor.l2_denorm" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches!( + canonical, + Canonical::Extension(ext) if ext.ext_dtype().is::() + ) + } + + fn produced_encodings(&self) -> Vec { + vec![L2Denorm.id()] + } + + /// Children: normalized=0, norms=1. + fn num_children(&self) -> usize { + L2DenormSlots::COUNT + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let denorm = normalize_as_l2_denorm(data.array().clone(), exec_ctx)?; + + // Splitting magnitude out is only worth anything if the children then compress: the + // unit-norm coordinates have a bounded range and the norms are an ordinary float column. + let normalized = compressor.compress_child( + denorm.normalized(), + &compress_ctx, + self.id(), + L2DenormSlots::NORMALIZED, + exec_ctx, + )?; + let norms = compressor.compress_child( + denorm.norms(), + &compress_ctx, + self.id(), + L2DenormSlots::NORMS, + exec_ctx, + )?; + + // SAFETY: Cascading preserves the split's child lengths and dtypes. + Ok(unsafe { L2Denorm::new_unchecked(normalized, norms) }.into_array()) + } +} + +/// Splits a tensor-like column into its exact [`L2Denorm`] representation. +/// +/// # Normalized child +/// +/// 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. +/// +/// 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. +/// +/// # Nullability +/// +/// Nullability is tracked entirely by the norms child, which inherits the input's nulls through +/// [`L2Norm`]'s validity propagation. The [`L2Denorm`] 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. +/// +/// Because this computes exact norms first and then divides by them, the returned `normalized` +/// child satisfies the strict unit-norm invariant. +pub fn normalize_as_l2_denorm( + 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 `L2Denorm` whose children are both `ConstantArray`s. + if let Some(wrapped) = try_build_constant_l2_denorm(&input, row_count, ctx)? { + return Ok(wrapped); + } + + let norms_array: ArrayRef = L2Norm + .try_new_array(row_count, EmptyOptions, [input.clone()])? + .execute(ctx)?; + let primitive_norms: PrimitiveArray = norms_array.clone().execute(ctx)?; + let norms_validity = primitive_norms.validity()?; + + let input: ExtensionArray = input.execute(ctx)?; + let normalized_dtype = input.dtype().as_nonnullable(); + let flat = extract_flat_elements(input.storage_array(), tensor_flat_size, ctx)?; + + // Resolve validity to a mask once rather than probing it per row (each `Validity::is_valid` + // executes a scalar for array-backed validity). + let norms_valid = norms_validity.execute_mask(row_count, ctx)?; + + let normalized = match_each_float_ptype!(flat.ptype(), |T| { + let norm_values = primitive_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]; + + // 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() { + unsafe { elements.push_n_unchecked(T::zero(), tensor_flat_size) }; + } else { + for &x in flat.row::(i) { + unsafe { elements.push_unchecked(x / norm) }; + } + } + } + + // Since L2Denorm's validity is the `and` of its child validities, the normalized child can + // be non-nullable. + build_normalized( + normalized_dtype, + tensor_flat_size, + row_count, + elements.freeze(), + ) + })?; + + // SAFETY: The normalized rows, norms ptype, and child lengths come directly from this split. + Ok(unsafe { L2Denorm::new_unchecked(normalized, norms_array) }) +} + +/// Attempts to build an [`L2DenormArray`] whose two children are both [`ConstantArray`]s by +/// eagerly normalizing `input`'s single stored row. +/// +/// Returns `Ok(None)` when `input` is not a tensor-like extension array whose storage is a +/// [`ConstantArray`] with a non-null fixed-size-list scalar. +/// +/// When `input` matches, the result is equivalent to [`normalize_as_l2_denorm`] but runs in +/// `O(list_size)` instead of `O(row_count * list_size)`. Keeping both children constant is what +/// lets cosine similarity and inner product short-circuit against a literal query vector. +pub(crate) fn try_build_constant_l2_denorm( + input: &ArrayRef, + len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some(ext) = input.as_opt::() else { + return Ok(None); + }; + let storage = ext.storage_array(); + let Some(const_storage) = storage.as_opt::() else { + return Ok(None); + }; + if const_storage.scalar().is_null() { + return Ok(None); + } + + // The caller is expected to have already validated that `input` is an `AnyTensor` extension + // dtype. + let tensor_match = input + .dtype() + .as_extension() + .metadata_opt::() + .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(); + + // Materialize just the single stored row; this does not expand the constant to the full column + // length. + let flat = extract_constant_flat_row(storage, ctx)?; + + let (normalized_fsl_scalar, norms_scalar) = match_each_float_ptype!(flat.ptype(), |T| { + let row = flat.as_slice::(); + + let mut sum_sq = T::zero(); + for &x in row { + sum_sq += x * x; + } + let norm_t: T = sum_sq.sqrt(); + + // Zero-norm rows must be stored as all-zeros so the unit-norm-or-zero invariant holds. + // This mirrors the per-row logic in `normalize_as_l2_denorm`. + let element_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable); + let children: Vec = if norm_t == T::zero() { + (0..list_size) + .map(|_| Scalar::zero_value(&element_dtype)) + .collect() + } else { + row.iter() + .map(|&v| Scalar::primitive(v / norm_t, Nullability::NonNullable)) + .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); + (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 norms = ConstantArray::new(norms_scalar, len).into_array(); + + // SAFETY: The constant children have matching lengths and element ptypes. + Ok(Some(unsafe { L2Denorm::new_unchecked(normalized, norms) })) +} + +/// Builds the non-nullable tensor-like extension array that becomes the `normalized` child. +fn build_normalized( + dtype: DType, + tensor_flat_size: usize, + row_count: usize, + elements: Buffer, +) -> VortexResult { + let list_size = + u32::try_from(tensor_flat_size).vortex_expect("tensor flat size must fit into `u32`"); + + // SAFETY: Tensor elements are always non-nullable, so the validity carries no length. + let elements = unsafe { PrimitiveArray::new_unchecked(elements, Validity::NonNullable) }; + + let storage = FixedSizeListArray::try_new( + elements.into_array(), + list_size, + Validity::NonNullable, + row_count, + )?; + + Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) +} diff --git a/vortex-tensor/src/encodings/l2_denorm/execute.rs b/vortex-tensor/src/encodings/l2_denorm/execute.rs new file mode 100644 index 00000000000..f1a5ee30c42 --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/execute.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::FixedSizeListArray; +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::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_float_ptype; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +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 [`L2DenormArray`]'s dtype, so the reconstructed column carries the +/// unioned nullability of both children. +/// +/// [`L2DenormArray`]: crate::encodings::l2_denorm::L2DenormArray +pub(super) fn denormalize( + normalized: &ArrayRef, + norms: &ArrayRef, + 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. + 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); + } + + let normalized: ExtensionArray = normalized.clone().execute(ctx)?; + let norms: PrimitiveArray = norms.clone().execute(ctx)?; + + let tensor_flat_size = tensor_flat_size(normalized.dtype()); + let flat = extract_flat_elements(normalized.storage_array(), tensor_flat_size, ctx)?; + + // TODO(connor): Do we want a "broadcast" expression for the List types, or is this fine? + match_each_float_ptype!(flat.ptype(), |T| { + let norms = norms.as_slice::(); + + let elements: Buffer = (0..row_count) + .flat_map(|i| { + let norm = norms[i]; + flat.row::(i).iter().map(move |&x| x * norm) + }) + .collect(); + + build_tensor_array(dtype, tensor_flat_size, row_count, validity, elements) + }) +} + +/// 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 +/// per-row loop. +fn denormalize_constant_norms( + normalized: &ArrayRef, + norm: &Scalar, + dtype: DType, + validity: Validity, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let tensor_flat_size = tensor_flat_size(normalized.dtype()); + let error = 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()); + } + + let normalized: ExtensionArray = normalized.clone().execute(ctx)?; + let storage: FixedSizeListArray = normalized.storage_array().clone().execute(ctx)?; + + 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. + let storage = unsafe { + FixedSizeListArray::new_unchecked(elements, storage.list_size(), validity, storage.len()) + }; + + Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) +} + +/// Rebuilds a tensor-like extension array from flat primitive elements. +fn build_tensor_array( + dtype: DType, + tensor_flat_size: usize, + row_count: usize, + validity: Validity, + elements: Buffer, +) -> VortexResult { + let list_size = + u32::try_from(tensor_flat_size).vortex_expect("tensor flat size must fit into `u32`"); + + // SAFETY: Tensor elements are always non-nullable, so the validity carries no length. + let elements = unsafe { PrimitiveArray::new_unchecked(elements, Validity::NonNullable) }; + + let storage = + FixedSizeListArray::try_new(elements.into_array(), list_size, validity, row_count)?; + + Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) +} + +/// Returns the flattened element count of each row of a tensor-like extension dtype. +fn tensor_flat_size(dtype: &DType) -> usize { + dtype + .as_extension() + .metadata_opt::() + .vortex_expect("the normalized child is validated to be an `AnyTensor` on construction") + .list_size() as usize +} diff --git a/vortex-tensor/src/encodings/l2_denorm/mod.rs b/vortex-tensor/src/encodings/l2_denorm/mod.rs new file mode 100644 index 00000000000..0a8e6a213ec --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/mod.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`L2Denorm`] encoding: a norm-split physical layout for tensor-like columns. +//! +//! An [`L2Denorm`] array stores a tensor or vector column as two children: +//! +//! - `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. +//! +//! 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 +//! the coordinates cheap to compress further: a unit-norm child has a bounded, well-conditioned +//! value range, and quantizing it only perturbs direction while the exact magnitude survives in +//! `norms`. +//! +//! Because the split is physical rather than logical, [`L2Norm`], [`InnerProduct`], and +//! [`CosineSimilarity`] can read straight through it instead of decoding first. +//! +//! [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm +//! [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct +//! [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity + +mod array; +pub use array::L2Denorm; +pub use array::L2DenormArray; +pub use array::L2DenormArraySlotsExt; +pub use array::L2DenormMetadata; +pub use array::L2DenormSlots; + +mod compress; +pub use compress::L2DenormScheme; +pub use compress::normalize_as_l2_denorm; +pub(crate) use compress::try_build_constant_l2_denorm; + +mod execute; + +mod orientation; +pub(crate) use orientation::DenormOrientation; + +mod rules; + +mod validate; +pub use validate::validate_l2_normalized_rows_against_norms; + +#[cfg(test)] +mod tests; diff --git a/vortex-tensor/src/encodings/l2_denorm/orientation.rs b/vortex-tensor/src/encodings/l2_denorm/orientation.rs new file mode 100644 index 00000000000..f4675410f32 --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/orientation.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; + +use crate::encodings::l2_denorm::L2Denorm; + +/// Classification of a binary operand pair by which side (if any) is [`L2Denorm`]-encoded. +/// +/// Symmetric binary tensor operators ([`CosineSimilarity`], [`InnerProduct`]) have identical fast +/// paths for "only the lhs is denormalized" and "only the rhs is denormalized", plus a separate +/// fast path for "both are denormalized". Rather than hand-rolling the commutative swap at every +/// call site, callers classify their operands with [`Self::classify`] and match on the result. +/// +/// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity +/// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct +pub(crate) enum DenormOrientation<'a> { + /// Both operands are [`L2Denorm`] arrays. + Both { + /// The left-hand operand. + lhs: &'a ArrayRef, + /// The right-hand operand. + rhs: &'a ArrayRef, + }, + + /// Exactly one operand is an [`L2Denorm`] array; the other is a plain tensor column. + One { + /// The [`L2Denorm`]-encoded operand, whichever side it came from. + denorm: &'a ArrayRef, + /// The other operand. + plain: &'a ArrayRef, + }, + + /// Neither operand is an [`L2Denorm`] array. + Neither, +} + +impl<'a> DenormOrientation<'a> { + /// Classify `(lhs, rhs)` by which side (if any) is [`L2Denorm`]-encoded. + pub(crate) fn classify(lhs: &'a ArrayRef, rhs: &'a ArrayRef) -> Self { + match (lhs.is::(), rhs.is::()) { + (true, true) => Self::Both { lhs, rhs }, + (true, false) => Self::One { + denorm: lhs, + plain: rhs, + }, + (false, true) => Self::One { + denorm: rhs, + plain: lhs, + }, + (false, false) => Self::Neither, + } + } +} diff --git a/vortex-tensor/src/encodings/l2_denorm/rules.rs b/vortex-tensor/src/encodings/l2_denorm/rules.rs new file mode 100644 index 00000000000..40ffcac98cf --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/rules.rs @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::IntoArray; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; +use vortex_array::optimizer::rules::ArrayParentReduceRule; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_error::VortexResult; + +use crate::encodings::l2_denorm::L2Denorm; +use crate::encodings::l2_denorm::array::L2DenormArraySlotsExt; + +pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&L2DenormSliceRule), + ParentRuleSet::lift(&L2DenormFilterRule), +]); + +/// Pushes a slice through the encoding into both children. +/// +/// The norm split is row-wise, so any row subset of an [`L2Denorm`] array is itself a valid +/// [`L2Denorm`] array. Rewriting the slice as two child slices keeps the column encoded instead of +/// canonicalizing it just to throw most of the rows away. +#[derive(Debug)] +struct L2DenormSliceRule; + +impl ArrayParentReduceRule for L2DenormSliceRule { + type Parent = Slice; + + fn reduce_parent( + &self, + array: ArrayView<'_, L2Denorm>, + parent: ArrayView<'_, Slice>, + _child_idx: usize, + ) -> VortexResult> { + let range = parent.slice_range(); + + // SAFETY: Slicing both children preserves their structure. + Ok(Some( + unsafe { + L2Denorm::new_unchecked( + array.normalized().slice(range.clone())?, + array.norms().slice(range.clone())?, + ) + } + .into_array(), + )) + } +} + +/// Pushes a filter through the encoding into both children. +/// +/// Same row-wise argument as [`L2DenormSliceRule`]. Unlike the generic scalar-function push-down, +/// this always fires: both children are physically per-row, so filtering them is strictly less +/// work than reconstructing the tensor column and filtering that. +#[derive(Debug)] +struct L2DenormFilterRule; + +impl ArrayParentReduceRule for L2DenormFilterRule { + type Parent = Filter; + + fn reduce_parent( + &self, + array: ArrayView<'_, L2Denorm>, + parent: ArrayView<'_, Filter>, + _child_idx: usize, + ) -> VortexResult> { + let mask = parent.filter_mask(); + + // SAFETY: Filtering both children with the same mask preserves their structure. + Ok(Some( + unsafe { + L2Denorm::new_unchecked( + array.normalized().filter(mask.clone())?, + array.norms().filter(mask.clone())?, + ) + } + .into_array(), + )) + } +} diff --git a/vortex-tensor/src/encodings/l2_denorm/tests.rs b/vortex-tensor/src/encodings/l2_denorm/tests.rs new file mode 100644 index 00000000000..5914e9c6ac7 --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/tests.rs @@ -0,0 +1,645 @@ +// 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; +use vortex_array::ArrayVTable; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Extension; +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::FixedSizeListArraySlotsExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::extension::datetime::Date; +use vortex_array::extension::datetime::TimeUnit; +use vortex_array::scalar::Scalar; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_compressor::scheme::Scheme; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::encodings::l2_denorm::L2Denorm; +use crate::encodings::l2_denorm::L2DenormArraySlotsExt; +use crate::encodings::l2_denorm::L2DenormMetadata; +use crate::encodings::l2_denorm::L2DenormScheme; +use crate::encodings::l2_denorm::normalize_as_l2_denorm; +use crate::encodings::l2_denorm::validate_l2_normalized_rows_against_norms; +use crate::tests::SESSION; +use crate::types::vector::Vector; +use crate::utils::test_helpers::assert_close; +use crate::utils::test_helpers::constant_tensor_array; +use crate::utils::test_helpers::tensor_array; +use crate::utils::test_helpers::vector_array; + +/// Builds an [`L2Denorm`] array through the checked constructor and executes it, which is the +/// end-to-end path every decode test cares about. +fn eval_l2_denorm(normalized: ArrayRef, norms: ArrayRef) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + let denorm = L2Denorm::try_new(normalized, norms, &mut ctx)?; + + denorm.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)> { + let mut ctx = SESSION.create_execution_ctx(); + let ext: ExtensionArray = array.execute(&mut ctx)?; + let validity = (0..ext.len()) + .map(|i| ext.is_valid(i, &mut ctx)) + .collect::>>()?; + 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(), + )) +} + +#[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)?; + + assert_eq!(actual_dtype, expected_dtype); + assert_eq!(actual_validity, expected_validity); + assert_close(&actual_elements, &expected_elements); + + Ok(()) +} + +fn non_tensor_extension_array() -> VortexResult { + let storage = PrimitiveArray::from_iter([1i32, 2]).into_array(); + let ext_dtype = ExtDType::::try_new(TimeUnit::Days, storage.dtype().clone())?.erased(); + + Ok(ExtensionArray::new(ext_dtype, storage).into_array()) +} + +/// Builds a non-nullable constant f64 norms array of length `len`. +fn constant_f64_norms(value: f64, len: usize) -> ArrayRef { + ConstantArray::new(Scalar::primitive(value, Nullability::NonNullable), len).into_array() +} + +// ============================================================================= +// Decoding +// ============================================================================= + +#[test] +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_l2_denorm(normalized, norms)?; + let expected = vector_array(3, &[3.0, 4.0, 0.0, 0.0, 0.0, 0.0])?; + + assert_tensor_arrays_eq(actual, expected) +} + +#[test] +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_l2_denorm(normalized, norms)?; + 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(); + + let mut ctx = SESSION.create_execution_ctx(); + let actual: ExtensionArray = eval_l2_denorm(normalized, norms)?.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_close(&elements.as_slice::()[..2], &[3.0, 4.0]); + + Ok(()) +} + +#[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(); + + let mut ctx = SESSION.create_execution_ctx(); + let denorm = L2Denorm::try_new(normalized, norms, &mut ctx)?; + + assert!(denorm.dtype().is_nullable()); + let mask = denorm.as_ref().validity()?.execute_mask(3, &mut ctx)?; + assert!(mask.value(0)); + assert!(!mask.value(1)); + assert!(!mask.value(2)); + + Ok(()) +} + +// ============================================================================= +// Constant fast paths +// ============================================================================= + +#[test] +fn constant_unit_norms_decode_to_the_normalized_child() -> VortexResult<()> { + // Every stored norm is exactly 1.0, so the fast path must short-circuit and return the + // normalized child unchanged. + 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_l2_denorm(normalized.clone(), norms)?; + + 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. + 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 actual = eval_l2_denorm(normalized.clone(), norms)?; + + assert_tensor_arrays_eq(actual, normalized) +} + +#[test] +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_l2_denorm(normalized, norms)?; + let expected = vector_array(3, &[3.0, 4.0, 0.0, 5.0, 0.0, 0.0])?; + + assert_tensor_arrays_eq(actual, expected) +} + +#[test] +fn constant_nonunit_norms_scale_fixed_shape_tensors() -> VortexResult<()> { + // The constant-scaling fast path must also cover multi-dimensional tensors, where the backing + // elements buffer spans more than one slot per row. + 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_l2_denorm(normalized, norms)?; + 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(); + + let mut ctx = SESSION.create_execution_ctx(); + let denorm = L2Denorm::try_new(normalized, norms, &mut ctx)?; + let dtype = denorm.dtype().clone(); + let decoded: ArrayRef = denorm.into_array().execute(&mut ctx)?; + + assert!(dtype.is_nullable()); + assert_eq!(decoded.dtype(), &dtype); + + Ok(()) +} + +// ============================================================================= +// Construction and validation +// ============================================================================= + +#[rstest] +#[case::non_extension_normalized( + PrimitiveArray::from_iter([1.0f64, 2.0]).into_array(), + PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(), +)] +#[case::non_tensor_extension_normalized( + non_tensor_extension_array().expect("valid date array"), + PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(), +)] +#[case::integer_tensor_normalized( + tensor_array(&[2], &[1i32, 2, 3, 4]).expect("valid tensor array"), + PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(), +)] +#[case::mismatched_norms_ptype( + vector_array(2, &[1.0f64, 0.0, 0.0, 1.0]).expect("valid vector array"), + PrimitiveArray::from_iter([1.0f32, 1.0]).into_array(), +)] +#[case::non_primitive_norms( + vector_array(2, &[1.0f64, 0.0, 0.0, 1.0]).expect("valid vector array"), + vector_array(1, &[1.0f64, 1.0]).expect("valid vector array"), +)] +#[case::mismatched_child_lengths( + vector_array(2, &[1.0f64, 0.0, 0.0, 1.0]).expect("valid vector array"), + PrimitiveArray::from_iter([1.0f64]).into_array(), +)] +fn rejects_structurally_invalid_children( + #[case] normalized: ArrayRef, + #[case] norms: ArrayRef, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + assert!(L2Denorm::try_new(normalized, norms, &mut ctx).is_err()); + + Ok(()) +} + +#[rstest] +#[case::unnormalized_child( + vector_array(2, &[3.0f64, 4.0, 1.0, 0.0]).expect("valid vector array"), + PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(), +)] +#[case::negative_norm( + vector_array(2, &[1.0f64, 0.0, 0.0, 1.0]).expect("valid vector array"), + PrimitiveArray::from_iter([1.0f64, -1.0]).into_array(), +)] +#[case::nonzero_row_with_zero_norm( + vector_array(2, &[1.0f64, 0.0, 0.0, 0.0]).expect("valid vector array"), + PrimitiveArray::from_iter([0.0f64, 0.0]).into_array(), +)] +fn checked_construction_rejects_semantic_violations( + #[case] normalized: ArrayRef, + #[case] norms: ArrayRef, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + assert!(L2Denorm::try_new(normalized, norms, &mut ctx).is_err()); + + Ok(()) +} + +#[test] +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_l2_denorm(normalized, norms)?; + let expected = vector_array(2, &[0.0, 0.0, 3.0, 0.0])?; + + assert_tensor_arrays_eq(actual, expected) +} + +#[test] +fn validate_accepts_normalized_f16_rows() -> VortexResult<()> { + let input = vector_array(2, &[3.0f32, 4.0, 0.0, 0.0].map(half::f16::from_f32))?; + let mut ctx = SESSION.create_execution_ctx(); + + let denorm = normalize_as_l2_denorm(input, &mut ctx)?; + validate_l2_normalized_rows_against_norms(&denorm.normalized().clone(), None, &mut ctx) +} + +#[test] +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()); + + Ok(()) +} + +// ============================================================================= +// Normalization +// ============================================================================= + +#[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( + tensor_array(&[2, 2], &[1.0, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]).expect("valid tensor array") +)] +#[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"))] +fn normalize_round_trips(#[case] input: ArrayRef) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input.clone(), &mut ctx)?; + let actual = denorm.into_array().execute(&mut ctx)?; + + assert_tensor_arrays_eq(actual, input) +} + +#[test] +fn normalize_keeps_constant_input_children_constant() -> VortexResult<()> { + // The constant fast path must leave both children constant, which is what lets cosine + // similarity and inner product short-circuit against a literal query vector. + let input = Vector::constant_array(&[3.0, 4.0], 16)?; + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input, &mut ctx)?; + + let normalized = denorm + .normalized() + .as_opt::() + .expect("normalized child should be an Extension array"); + assert!( + normalized.storage_array().as_opt::().is_some(), + "normalized storage should stay constant after the fast path" + ); + + let norms = denorm + .norms() + .as_opt::() + .expect("norms child should be a ConstantArray"); + assert_close( + &[norms + .scalar() + .as_primitive() + .typed_value::() + .expect("norms scalar")], + &[5.0], + ); + + Ok(()) +} + +#[test] +fn normalize_zeroes_rows_with_zero_norms() -> VortexResult<()> { + let input = vector_array(2, &[0.0, 0.0, 3.0, 4.0])?; + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input.clone(), &mut ctx)?; + + let normalized: ExtensionArray = denorm.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], &[0.0, 0.0]); + + let actual = denorm.into_array().execute(&mut ctx)?; + + assert_tensor_arrays_eq(actual, input) +} + +#[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])?; + let input = MaskedArray::try_new(input, Validity::from_iter([true, false, true]))?.into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input, &mut ctx)?; + + assert!(!denorm.normalized().dtype().is_nullable()); + assert!(denorm.norms().dtype().is_nullable()); + + let mask = denorm.as_ref().validity()?.execute_mask(3, &mut ctx)?; + assert!(mask.value(0)); + assert!(!mask.value(1)); + assert!(mask.value(2)); + + Ok(()) +} + +// ============================================================================= +// Row operations +// ============================================================================= + +#[test] +fn slice_stays_encoded_and_decodes_correctly() -> VortexResult<()> { + let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0, 5.0, 12.0])?; + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + + let sliced = denorm.slice(1..3)?.execute_until::(&mut ctx)?; + assert!( + sliced.is::(), + "slicing must push down into both children instead of decoding the column" + ); + + let expected = vector_array(2, &[1.0, 0.0, 0.0, 2.0])?; + + assert_tensor_arrays_eq(sliced, expected) +} + +#[test] +fn filter_stays_encoded_and_decodes_correctly() -> VortexResult<()> { + let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0, 5.0, 12.0])?; + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + + let mask = Mask::from_iter([true, false, true, false]); + let filtered = denorm.filter(mask)?.execute_until::(&mut ctx)?; + assert!( + filtered.is::(), + "filtering must push down into both children instead of decoding the column" + ); + + let expected = vector_array(2, &[3.0, 4.0, 0.0, 2.0])?; + + assert_tensor_arrays_eq(filtered, expected) +} + +#[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])?; + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + + let indices = PrimitiveArray::from_iter([3u64, 0, 3]).into_array(); + let taken = denorm.take(indices)?.execute::(&mut ctx)?; + let expected = vector_array(2, &[5.0, 12.0, 3.0, 4.0, 5.0, 12.0])?; + + assert_tensor_arrays_eq(taken.into_array(), expected) +} + +#[test] +fn scalar_at_reads_a_single_denormalized_row() -> VortexResult<()> { + let input = vector_array(2, &[3.0, 4.0, 5.0, 12.0])?; + let mut ctx = SESSION.create_execution_ctx(); + let denorm = normalize_as_l2_denorm(input.clone(), &mut ctx)?.into_array(); + + for i in 0..input.len() { + assert_eq!( + denorm.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_as_l2_denorm` 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( + tensor_array(&[2, 2], &[1.0, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]).expect("valid tensor array") +)] +#[case::nullable_vector(nullable_vector_input().expect("valid vector array"))] +fn serde_round_trip(#[case] input: ArrayRef) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let original = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + let children: Vec = original.children(); + + let metadata = SESSION + .array_serialize(&original)? + .expect("L2Denorm must serialize"); + let recovered = ArrayPlugin::deserialize( + &L2Denorm, + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + assert_eq!(recovered.encoding_id(), ArrayVTable::id(&L2Denorm)); + assert_eq!(recovered.dtype(), original.dtype()); + assert_eq!(recovered.len(), original.len()); + 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. +#[test] +fn serialized_metadata_pins_child_nullabilities() -> 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 denorm = normalize_as_l2_denorm(input, &mut ctx)?; + + let bytes = SESSION + .array_serialize(&denorm.clone().into_array())? + .expect("L2Denorm must serialize"); + let metadata = L2DenormMetadata::decode(bytes.as_slice())?; + + assert_eq!( + metadata.normalized_is_nullable, + denorm.normalized().dtype().is_nullable(), + ); + assert_eq!( + metadata.norms_is_nullable, + denorm.norms().dtype().is_nullable(), + ); + + 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(); + let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + let original = L2Denorm::try_new(normalized, norms, &mut ctx)?.into_array(); + let children: Vec = original.children(); + let metadata = SESSION + .array_serialize(&original)? + .expect("L2Denorm must serialize"); + + let recovered = ArrayPlugin::deserialize( + &L2Denorm, + original.dtype(), + original.len(), + &metadata, + &[], + &children, + &SESSION, + )?; + + let recovered = recovered.as_::(); + assert!(recovered.normalized().dtype().is_nullable()); + assert!(!recovered.norms().dtype().is_nullable()); + + Ok(()) +} + +#[test] +fn encoding_is_registered_under_the_original_id() { + let id = ArrayVTable::id(&L2Denorm); + + assert_eq!(id.as_ref(), "vortex.tensor.l2_denorm"); + assert!(SESSION.arrays().registry().contains_key(&id)); +} + +// ============================================================================= +// Compression +// ============================================================================= + +/// Vectors that all point the same way but vary in magnitude: the norm split turns the +/// coordinates into a repeating pattern and isolates the magnitudes into their own float column, +/// which is exactly the shape the encoding exists to exploit. +fn collinear_vectors(rows: usize) -> VortexResult { + let elements: Vec = (0..rows) + .flat_map(|i| { + let scale = 1.0 + i as f64; + [3.0 * scale, 4.0 * scale, 12.0 * scale, 0.0] + }) + .collect(); + + vector_array(4, &elements) +} + +#[rstest] +#[case::vector(vector_array(2, &[3.0, 4.0, 5.0, 12.0]).expect("valid vector array"))] +#[case::fixed_shape_tensor( + tensor_array(&[2, 2], &[1.0, 2.0, 3.0, 4.0, 0.0, 1.0, 0.0, 0.0]).expect("valid tensor array") +)] +fn scheme_matches_tensor_columns(#[case] input: ArrayRef) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let canonical: Canonical = input.execute(&mut ctx)?; + + assert!(L2DenormScheme.matches(&canonical)); + assert_eq!( + L2DenormScheme.produced_encodings(), + vec![ArrayVTable::id(&L2Denorm)] + ); + + Ok(()) +} + +#[test] +fn compressor_emits_the_dedicated_encoding() -> VortexResult<()> { + let input = collinear_vectors(1024)?; + let compressor = BtrBlocksCompressorBuilder::default() + .with_new_scheme(&L2DenormScheme) + .build(); + + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&input, &mut ctx)?; + + assert_eq!(compressed.encoding_id(), ArrayVTable::id(&L2Denorm)); + assert!(compressed.nbytes() < input.nbytes()); + assert_tensor_arrays_eq(compressed, input) +} diff --git a/vortex-tensor/src/encodings/l2_denorm/validate.rs b/vortex-tensor/src/encodings/l2_denorm/validate.rs new file mode 100644 index 00000000000..c01563944f9 --- /dev/null +++ b/vortex-tensor/src/encodings/l2_denorm/validate.rs @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::ToPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::dtype::DType; +use vortex_array::match_each_float_ptype; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +use crate::utils::extract_flat_elements; +use crate::utils::unit_norm_tolerance; +use crate::utils::validate_tensor_float_input; + +/// Validates the structural invariants of an [`L2Denorm`] array's children. +/// +/// These are the cheap, dtype-and-length checks that every [`L2DenormArray`] upholds, whichever +/// constructor built it. They run on construction and on deserialization. +/// +/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm +/// [`L2DenormArray`]: crate::encodings::l2_denorm::L2DenormArray +pub(super) fn validate_l2_denorm_children( + normalized: &ArrayRef, + norms: &ArrayRef, + dtype: &DType, + len: usize, +) -> VortexResult<()> { + vortex_ensure_eq!( + normalized.len(), + len, + "L2Denorm normalized child must have the array length ({len}), got {}", + normalized.len(), + ); + vortex_ensure_eq!( + norms.len(), + len, + "L2Denorm norms child must have the array length ({len}), got {}", + norms.len(), + ); + + 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!( + "L2Denorm norms must be a primitive float array, got {}", + norms.dtype(), + ); + }; + vortex_ensure_eq!( + *norms_ptype, + element_ptype, + "L2Denorm norms dtype must match the normalized element dtype ({element_ptype}), \ + got {norms_ptype}", + ); + + let expected = normalized + .dtype() + .union_nullability(norms.dtype().nullability()); + vortex_ensure_eq!( + *dtype, + expected, + "L2Denorm dtype must be the union of its children's nullability ({expected}), got {dtype}", + ); + + Ok(()) +} + +/// Validates that `normalized` and (when supplied) the matching `norms` jointly satisfy the +/// semantic [`L2Denorm`] 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`. +/// +/// This costs `O(len * list_size)`, which is why it is a separate step rather than part of the +/// encoding's structural validation. +/// +/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm +pub fn validate_l2_normalized_rows_against_norms( + normalized: &ArrayRef, + norms: Option<&ArrayRef>, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let row_count = normalized.len(); + if row_count == 0 { + return Ok(()); + } + + let tensor_match = validate_tensor_float_input(normalized.dtype())?; + let element_ptype = tensor_match.element_ptype(); + let tensor_flat_size = tensor_match.list_size() as usize; + let tolerance = unit_norm_tolerance(element_ptype, tensor_flat_size); + + if let Some(norms) = norms { + vortex_ensure_eq!( + norms.len(), + row_count, + "L2Denorm norms must have the same length as the normalized child ({row_count}), \ + got {}", + norms.len(), + ); + + let DType::Primitive(norms_ptype, _) = norms.dtype() else { + vortex_bail!( + "L2Denorm norms must be a primitive float array, got {}", + norms.dtype(), + ); + }; + vortex_ensure_eq!( + *norms_ptype, + element_ptype, + "L2Denorm norms ptype must match the normalized element ptype ({element_ptype}), \ + got {norms_ptype}", + ); + } + + 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() + .fold((0.0f64, true), |(sum_sq, is_zero), x| { + let value = ToPrimitive::to_f64(x).unwrap_or(f64::NAN); + (sum_sq + value * value, is_zero && value.abs() <= tolerance) + }); + let row_norm = row_norm_sq.sqrt(); + + vortex_ensure!( + row_norm == 0.0 || (row_norm - 1.0).abs() <= tolerance, + "L2Denorm normalized child must have L2 norm 1.0 or 0.0, but row {i} has \ + {row_norm:.6}", + ); + + if let Some(stored_norms) = stored_norms { + let stored_norm_f64 = ToPrimitive::to_f64(&stored_norms[i]).unwrap_or(f64::NAN); + vortex_ensure!( + stored_norm_f64 >= 0.0, + "L2Denorm norms must be non-negative, but row {i} has {stored_norm_f64:.6}", + ); + + if stored_norm_f64 == 0.0 { + vortex_ensure!( + is_zero_row, + "L2Denorm normalized child must be all zeros when norms row {i} is 0.0", + ); + } + } + } + }); + + Ok(()) +} diff --git a/vortex-tensor/src/encodings/mod.rs b/vortex-tensor/src/encodings/mod.rs index e42a8605096..627c9d7c7d0 100644 --- a/vortex-tensor/src/encodings/mod.rs +++ b/vortex-tensor/src/encodings/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Encodings for the different tensor types. +//! Array encodings for the different tensor types. // TODO(connor): // pub mod spherical; // Spherical transform on unit-normalized vectors. diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index 6826f3ad191..5c39724ef40 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -19,9 +19,9 @@ use vortex_array::session::ArraySessionExt; use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; +use crate::encodings::l2_denorm::L2Denorm; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::scalar_fns::inner_product::InnerProduct; -use crate::scalar_fns::l2_denorm::L2Denorm; use crate::scalar_fns::l2_norm::L2Norm; use crate::types::fixed_shape_tensor::FixedShapeTensor; use crate::types::vector::Vector; @@ -41,10 +41,13 @@ pub mod vector_search; mod utils; /// Environment variable that gates registration of the tensor scalar-fn array plugins (the array -/// encodings that let [`CosineSimilarity`], [`InnerProduct`], [`L2Denorm`], and [`L2Norm`] -/// persist in a Vortex file). When unset, only the scalar functions themselves -/// are registered; readers of files containing serialized tensor scalar-fn arrays will fail to -/// deserialize. Opt-in by setting the variable to any non-empty value. +/// encodings that let [`CosineSimilarity`], [`InnerProduct`], and [`L2Norm`] persist in a Vortex +/// file). When unset, only the scalar functions themselves are registered; readers of files +/// containing serialized tensor scalar-fn arrays will fail to deserialize. Opt-in by setting the +/// variable to any non-empty value. +/// +/// This does **not** gate [`L2Denorm`]. That is a real array encoding rather than a persisted +/// scalar function, and the compressor can emit it, so it always registers. pub const SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str = "VX_SCALAR_FN_ARRAY_TENSOR_PLUGIN"; /// Initialize the Vortex tensor library with a Vortex session. @@ -56,11 +59,12 @@ pub fn initialize(session: &VortexSession) { arrow_session.register_exporter(Arc::new(Vector)); arrow_session.register_importer(Arc::new(Vector)); + session.arrays().register(L2Denorm); + let session_fns = session.scalar_fns(); session_fns.register(CosineSimilarity); session_fns.register(InnerProduct); - session_fns.register(L2Denorm); session_fns.register(L2Norm); // Registering the scalar-fn array plugins lets the tensor scalar fns be serialized as array @@ -72,7 +76,6 @@ pub fn initialize(session: &VortexSession) { session_arrays.register(ScalarFnArrayPlugin::new(CosineSimilarity)); session_arrays.register(ScalarFnArrayPlugin::new(InnerProduct)); - session_arrays.register(ScalarFnArrayPlugin::new(L2Denorm)); session_arrays.register(ScalarFnArrayPlugin::new(L2Norm)); } } diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index de6f0614471..447ea973de9 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -30,9 +30,9 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::encodings::l2_denorm::DenormOrientation; +use crate::encodings::l2_denorm::try_build_constant_l2_denorm; use crate::scalar_fns::inner_product::InnerProduct; -use crate::scalar_fns::l2_denorm::DenormOrientation; -use crate::scalar_fns::l2_denorm::try_build_constant_l2_denorm; use crate::scalar_fns::l2_norm::L2Norm; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_l2_denorm_children; @@ -47,14 +47,14 @@ use crate::utils::validate_binary_tensor_float_inputs; /// Both inputs must be tensor-like extension arrays ([`FixedShapeTensor`] or [`Vector`]) with the /// same dtype and a float element type. The output is a float column of the same float type. /// -/// When either input is wrapped in [`L2Denorm`], this operator treats the stored norms and -/// normalized children as authoritative. For lossy encodings, that means the -/// optimized readthrough path may intentionally differ slightly from decoding both sides to dense +/// When either input is [`L2Denorm`]-encoded, this operator treats the stored norms and +/// normalized children as authoritative. For lossy normalized children, that means the optimized +/// read-through path may intentionally differ slightly from decoding both sides to dense /// coordinates and recomputing cosine from scratch. /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector -/// [`L2Denorm`]: crate::scalar_fns::l2_denorm::L2Denorm +/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm #[derive(Clone)] pub struct CosineSimilarity; @@ -117,16 +117,16 @@ 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-wrap it as an `L2Denorm` whose children are both `ConstantArray`s. - // The L2Denorm fast path below then picks it up. - if let Some(sfn) = try_build_constant_l2_denorm(&lhs_ref, len, ctx)? { - lhs_ref = sfn.into_array(); + // stored row and re-encode it as an `L2Denorm` whose children are both `ConstantArray`s. + // The `L2Denorm` fast path below then picks it up. + if let Some(denorm) = try_build_constant_l2_denorm(&lhs_ref, len, ctx)? { + lhs_ref = denorm.into_array(); } - if let Some(sfn) = try_build_constant_l2_denorm(&rhs_ref, len, ctx)? { - rhs_ref = sfn.into_array(); + if let Some(denorm) = try_build_constant_l2_denorm(&rhs_ref, len, ctx)? { + rhs_ref = denorm.into_array(); } - // Take any L2Denorm-wrapped fast path that applies. + // Take any L2Denorm read-through fast path that applies. match DenormOrientation::classify(&lhs_ref, &rhs_ref) { DenormOrientation::Both { lhs, rhs } => { return self.execute_both_denorm(lhs, rhs, len, ctx); @@ -219,8 +219,10 @@ impl ScalarFnArrayVTable for CosineSimilarity { } impl CosineSimilarity { - /// Both sides are `L2Denorm`: treat the normalized children as authoritative, so + /// Both sides are [`L2Denorm`]-encoded: treat the normalized children as authoritative, so /// `cosine_similarity = dot(n_l, n_r)`. + /// + /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm fn execute_both_denorm( &self, lhs_ref: &ArrayRef, @@ -261,9 +263,11 @@ impl CosineSimilarity { }) } - /// One side is `L2Denorm`: treat the normalized child as authoritative, so + /// One side is [`L2Denorm`]-encoded: treat the normalized child as authoritative, so /// `cosine_similarity = dot(n, b) / ||b||`. /// + /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm + /// /// The caller must pass the denorm array as `denorm_ref` and the plain array as `plain_ref`. fn execute_one_denorm( &self, @@ -321,8 +325,8 @@ mod tests { use vortex_array::validity::Validity; use vortex_error::VortexResult; + use crate::encodings::l2_denorm::L2Denorm; use crate::scalar_fns::cosine_similarity::CosineSimilarity; - use crate::scalar_fns::l2_denorm::L2Denorm; use crate::tests::SESSION; use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; @@ -580,7 +584,7 @@ mod tests { 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 = L2Denorm::try_new_array(normalized_r, norms_r, &mut ctx)?.into_array(); + let rhs = L2Denorm::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); let scalar_fn = CosineSimilarity::new().erased(); let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; @@ -600,14 +604,16 @@ mod tests { // children is nonzero. let normalized_l = tensor_array(&[2], &[0.6, 0.8])?; let norms_l = PrimitiveArray::from_iter([0.0f64]).into_array(); - // SAFETY: This is a focused test that intentionally violates the unit-norm invariant by - // pairing a nonzero normalized row with a stored norm of `0.0`, mimicking lossy storage. - let lhs = unsafe { L2Denorm::new_array_unchecked(normalized_l, norms_l)? }.into_array(); + // 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 { L2Denorm::new_unchecked(normalized_l, norms_l) }.into_array(); let normalized_r = tensor_array(&[2], &[0.6, 0.8])?; let norms_r = PrimitiveArray::from_iter([0.0f64]).into_array(); - // SAFETY: Same as above for the rhs operand. - let rhs = unsafe { L2Denorm::new_array_unchecked(normalized_r, norms_r)? }.into_array(); + // Same as above for the rhs operand. + // SAFETY: The children are structurally valid. + let rhs = unsafe { L2Denorm::new_unchecked(normalized_r, norms_r) }.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`. @@ -623,9 +629,10 @@ mod tests { // authoritative stored norm on the denorm side is `0.0`. let normalized = tensor_array(&[2], &[0.6, 0.8])?; let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); - // SAFETY: This is a focused test that intentionally pairs a nonzero normalized row with a - // stored norm of `0.0`, mimicking lossy storage where the stored norm is authoritative. - let denorm = unsafe { L2Denorm::new_array_unchecked(normalized, norms)? }.into_array(); + // 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 denorm = unsafe { L2Denorm::new_unchecked(normalized, norms) }.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 f46b03ac94e..33b87e971ac 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -34,8 +34,8 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::encodings::l2_denorm::DenormOrientation; use crate::matcher::AnyTensor; -use crate::scalar_fns::l2_denorm::DenormOrientation; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_flat_elements; use crate::utils::extract_l2_denorm_children; @@ -114,7 +114,7 @@ impl ScalarFnVTable for InnerProduct { let rhs_ref = args.get(1)?; let len = args.row_count(); - // Take any L2Denorm-wrapped fast path that applies. + // Take any L2Denorm read-through fast path that applies. match DenormOrientation::classify(&lhs_ref, &rhs_ref) { DenormOrientation::Both { lhs, rhs } => { return self.execute_both_denorm(lhs, rhs, len, ctx); @@ -203,7 +203,9 @@ impl ScalarFnArrayVTable for InnerProduct { } impl InnerProduct { - /// Both sides are `L2Denorm`: `inner_product = s_l * s_r * dot(n_l, n_r)`. + /// Both sides are [`L2Denorm`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. + /// + /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm fn execute_both_denorm( &self, lhs_ref: &ArrayRef, @@ -234,7 +236,9 @@ impl InnerProduct { }) } - /// One side is `L2Denorm`: `inner_product = s * dot(n, other)`. + /// One side is [`L2Denorm`]-encoded: `inner_product = s * dot(n, other)`. + /// + /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm /// /// The caller must pass the denorm array as `denorm_ref` and the plain array as `plain_ref`. fn execute_one_denorm( @@ -289,8 +293,8 @@ mod tests { use vortex_array::validity::Validity; use vortex_error::VortexResult; + use crate::encodings::l2_denorm::L2Denorm; use crate::scalar_fns::inner_product::InnerProduct; - use crate::scalar_fns::l2_denorm::L2Denorm; use crate::tests::SESSION; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::l2_denorm_array; @@ -467,7 +471,7 @@ mod tests { let norms_l = PrimitiveArray::from_option_iter([Some(5.0f64), None]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let lhs = L2Denorm::try_new_array(normalized_l, norms_l, &mut ctx)?.into_array(); + let lhs = L2Denorm::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); let rhs = l2_denorm_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_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs deleted file mode 100644 index 7195265790f..00000000000 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ /dev/null @@ -1,1146 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! L2 denormalization expression for tensor-like types. - -use num_traits::Float; -use num_traits::ToPrimitive; -use num_traits::Zero; -use prost::Message; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::Constant; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::Extension; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::FixedSizeListArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; -use vortex_array::arrays::ScalarFnArray; -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::arrays::scalar_fn::ExactScalarFn; -use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; -use vortex_array::arrays::scalar_fn::ScalarFnArrayView; -use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; -use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; -use vortex_array::dtype::NativePType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::proto::dtype as pb; -use vortex_array::expr::Expression; -use vortex_array::expr::union_child_validities; -use vortex_array::match_each_float_ptype; -use vortex_array::scalar::Scalar; -use vortex_array::scalar::ScalarValue; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; -use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; -use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::TypedScalarFnInstance; -use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::serde::ArrayChildren; -use vortex_array::validity::Validity; -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_ensure_eq; -use vortex_error::vortex_err; -use vortex_session::VortexSession; -use vortex_session::registry::CachedId; - -use crate::matcher::AnyTensor; -use crate::scalar_fns::l2_norm::L2Norm; -use crate::utils::extract_constant_flat_row; -use crate::utils::extract_flat_elements; -use crate::utils::unit_norm_tolerance; -use crate::utils::validate_tensor_float_input; - -/// Re-applies authoritative L2 norms to a normalized tensor column. -/// -/// Computes `normalized * norm` on each row over the flat backing buffer of each tensor-like type. -/// -/// The normalized input must be a tensor-like extension array with a float element type and each -/// non-null row is semantically required to already be L2-normalized. -/// -/// The norms input must be a primitive float column with the same element type as the normalized -/// tensor elements. -/// -/// [`L2Denorm`] is the norm-splitting wrapper used throughout the tensor crate. Callers that build -/// it through [`try_new_array`](Self::try_new_array) get an exact unit-norm invariant on the -/// `normalized` child. -/// -/// Advanced callers can also use [`new_array_unchecked`](Self::new_array_unchecked) to attach -/// authoritative stored norms to a lossy approximation of that child, such as quantized normalized -/// vectors. -/// -/// Downstream readthrough rules intentionally treat the stored norms and normalized child as the -/// encoding contract, even when that differs slightly from recomputing over fully decoded -/// coordinates. -#[derive(Clone)] -pub struct L2Denorm; - -impl L2Denorm { - /// Creates a new [`TypedScalarFnInstance`] wrapping the L2 denormalization operation. - /// - /// This is a low-level scalar-function descriptor constructor. To build a semantically valid - /// [`L2Denorm`] array, prefer [`try_new_array`](Self::try_new_array). - pub fn new() -> TypedScalarFnInstance { - TypedScalarFnInstance::new(L2Denorm, EmptyOptions) - } - - /// Constructs a validated [`ScalarFnArray`] that lazily re-applies `norms` to `normalized`. - /// - /// This is the correct constructor for [`L2Denorm`] arrays. In addition to the structural - /// checks performed by [`ScalarFnArray::try_new`], it validates that every valid row of the - /// `normalized` child has L2 norm `1.0` (or `0.0` for zero rows), within the tolerance implied - /// by the child element precision. It also validates that stored norms are non-negative, and - /// that any row with stored norm `0.0` has an all-zero normalized row. - /// - /// # Errors - /// - /// Returns an error if the [`ScalarFnArray`] cannot be constructed (e.g. due to dtype - /// mismatches) or if the `normalized` child is not row-wise L2-normalized. - pub fn try_new_array( - normalized: ArrayRef, - norms: ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - validate_l2_normalized_rows_against_norms(&normalized, Some(&norms), ctx)?; - - // SAFETY: We just validated that it is normalized. - unsafe { Self::new_array_unchecked(normalized, norms) } - } - - /// Constructs an [`L2Denorm`] array without validating that the `normalized` child is actually - /// row-wise L2-normalized. - /// - /// This escape hatch is intended for advanced callers that already established, or - /// intentionally relax, the normalized-child invariant. Structural validation still runs via - /// [`ScalarFnArray::try_new`]. - /// - /// # Safety - /// - /// The caller must ensure the `normalized` child is semantically suitable for L2 - /// denormalization. For exact wrappers, that means every valid row is unit-norm or zero. - /// - /// Lossy encodings may deliberately relax that invariant while still treating the stored norms - /// as authoritative. - /// - /// Violating the intended contract will not cause memory unsafety, but may produce incorrect - /// results. - pub unsafe fn new_array_unchecked( - normalized: ArrayRef, - norms: ArrayRef, - ) -> VortexResult { - ScalarFnArray::try_new(L2Denorm::new().erased(), vec![normalized, norms]) - } -} - -impl ScalarFnVTable for L2Denorm { - type Options = EmptyOptions; - - fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("vortex.tensor.l2_denorm"); - *ID - } - - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) - } - - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("normalized"), - 1 => ChildName::from("norms"), - _ => unreachable!("L2Denorm must have exactly two children"), - } - } - - fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { - let normalized = &arg_dtypes[0]; - let norms = &arg_dtypes[1]; - - let tensor_match = validate_tensor_float_input(normalized)?; - let element_ptype = tensor_match.element_ptype(); - - let DType::Primitive(norms_ptype, _) = norms else { - vortex_bail!("L2Denorm norms must be a primitive float array, got {norms}"); - }; - vortex_ensure_eq!( - *norms_ptype, - element_ptype, - "L2Denorm norms dtype must match normalized element dtype ({element_ptype}), \ - got {norms_ptype}", - ); - - Ok(normalized.union_nullability(norms.nullability())) - } - - fn execute( - &self, - _options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let normalized_ref = args.get(0)?; - let norms_ref = args.get(1)?; - let output_dtype = normalized_ref - .dtype() - .union_nullability(norms_ref.dtype().nullability()); - let validity = normalized_ref.validity()?.and(norms_ref.validity()?)?; - - if let Some(const_norms) = norms_ref.as_opt::() { - let norm_scalar = const_norms.scalar(); - vortex_ensure!( - norm_scalar.dtype().is_float(), - "L2Denorm constant norms must be a float scalar, got {}", - norm_scalar.dtype(), - ); - - if let Some(norm_value) = norm_scalar.value() { - return execute_l2_denorm_constant_norms( - normalized_ref, - norm_scalar, - norm_value, - output_dtype, - validity, - ctx, - ); - } - } - - let normalized: ExtensionArray = normalized_ref.execute(ctx)?; - let norms: PrimitiveArray = norms_ref.execute(ctx)?; - let row_count = args.row_count(); - - let tensor_match = normalized - .dtype() - .as_extension() - .metadata_opt::() - .vortex_expect("we already validated this in `return_dtype`"); - let tensor_flat_size = tensor_match.list_size() as usize; - - let flat = extract_flat_elements(normalized.storage_array(), tensor_flat_size, ctx)?; - - // TODO(connor): Do we want a "broadcast" expression for the List types, or is this fine? - match_each_float_ptype!(flat.ptype(), |T| { - let norms = norms.as_slice::(); - - let elements: Buffer = (0..row_count) - .flat_map(|i| { - let norm = norms[i]; - flat.row::(i).iter().map(move |&x| x * norm) - }) - .collect(); - - build_tensor_array( - output_dtype, - tensor_flat_size, - row_count, - validity, - elements, - ) - }) - } - - fn validity( - &self, - _options: &Self::Options, - expression: &Expression, - ) -> VortexResult> { - union_child_validities(expression) - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false - } -} - -/// Metadata for a serialized [`L2Denorm`] array: both children's full [`DType`]s. The parent's -/// dtype is `normalized.union_nullability(norms.nullability())`, which loses both children's -/// individual nullabilities, so we persist them directly. -#[derive(Clone, prost::Message)] -pub(super) struct L2DenormMetadata { - #[prost(message, optional, tag = "1")] - normalized_dtype: Option, - #[prost(message, optional, tag = "2")] - norms_dtype: Option, -} - -impl ScalarFnArrayVTable for L2Denorm { - fn serialize( - &self, - view: &ScalarFnArrayView, - _session: &VortexSession, - ) -> VortexResult>> { - let scalar_fn_array = view.as_::(); - let normalized_dtype = Some(scalar_fn_array.child_at(0).dtype().try_into()?); - let norms_dtype = Some(scalar_fn_array.child_at(1).dtype().try_into()?); - Ok(Some( - L2DenormMetadata { - normalized_dtype, - norms_dtype, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - _dtype: &DType, - len: usize, - metadata: &[u8], - children: &dyn ArrayChildren, - session: &VortexSession, - ) -> VortexResult> { - let metadata = L2DenormMetadata::decode(metadata) - .map_err(|e| vortex_err!("Failed to decode L2DenormMetadata: {e}"))?; - let normalized_pb = metadata - .normalized_dtype - .as_ref() - .ok_or_else(|| vortex_err!("L2DenormMetadata missing normalized_dtype"))?; - let norms_pb = metadata - .norms_dtype - .as_ref() - .ok_or_else(|| vortex_err!("L2DenormMetadata missing norms_dtype"))?; - let normalized_dtype = DType::from_proto(normalized_pb, session)?; - let norms_dtype = DType::from_proto(norms_pb, session)?; - let normalized = children.get(0, &normalized_dtype, len)?; - let norms = children.get(1, &norms_dtype, len)?; - Ok(ScalarFnArrayParts { - options: EmptyOptions, - children: vec![normalized, norms], - }) - } -} - -/// Optimized execution when the norms array is constant. -fn execute_l2_denorm_constant_norms( - normalized_ref: ArrayRef, - norm_scalar: &Scalar, - norm_value: &ScalarValue, - output_dtype: DType, - new_validity: Validity, - ctx: &mut ExecutionCtx, -) -> VortexResult { - // If the norms are all equal to 1 then we don't need to do anything. - let err = norm_value - .as_primitive() - .as_f64() - .vortex_expect("we know that this is a float, so it must fit in f64") - - 1.0f64; - - let tensor_match = normalized_ref - .dtype() - .as_extension_opt() - .and_then(|ext| ext.metadata_opt::()) - .ok_or_else(|| { - vortex_err!( - "L2Denorm normalized child must be a tensor-like extension, got {}", - normalized_ref.dtype(), - ) - })?; - - let tolerance = unit_norm_tolerance( - norm_scalar.dtype().as_ptype(), - tensor_match.list_size() as usize, - ); - if err.abs() < tolerance { - return Ok(normalized_ref); - } - - // Even if the norms are not all 1, if they are all the same then we can multiply - // the entire elements array by the same number. - let normalized: ExtensionArray = normalized_ref.execute(ctx)?; - let storage_fsl: FixedSizeListArray = normalized.storage_array().clone().execute(ctx)?; - - // Replace the elements array with an array that multiplies it by the constant - // norms array (with length multiplied by the dimensions of the vectors). - let const_array = - ConstantArray::new(norm_scalar.clone(), storage_fsl.elements().len()).into_array(); - let mult_elements = storage_fsl - .elements() - .clone() - .binary(const_array, Operator::Mul)?; - - // SAFETY: We just updated the elements of the array with a scalar fn, so all - // invariants still hold. - let new_fsl = unsafe { - FixedSizeListArray::new_unchecked( - mult_elements, - storage_fsl.list_size(), - new_validity, - storage_fsl.len(), - ) - }; - - Ok(ExtensionArray::new(output_dtype.as_extension().clone(), new_fsl.into_array()).into_array()) -} - -/// Builds an unexecuted [`L2Denorm`] expression by normalizing `input` and reattaching the exact -/// norms as the norms child. -/// -/// The returned array is a lazy `L2Denorm(normalized, norms)` scalar function array. -/// -/// # Normalized child -/// -/// The normalized child is always **non-nullable** with [`Validity::NonNullable`]. Every non-null -/// row with a positive L2 norm is divided by its norm to produce a unit-norm vector. -/// -/// Rows that are null in the original input are **zeroed out** in the normalized output. This is -/// necessary because null rows may have undefined (garbage) physical storage values, and we do not -/// want to let those propagate into downstream lossy encodings. -/// -/// # Nullability -/// -/// Nullability is tracked entirely by the norms child. Null input rows produce null norms via -/// [`L2Norm`]'s validity propagation. When the [`L2Denorm`] wrapper is executed, its validity is -/// `and(normalized_validity, norms_validity)`, which correctly identifies originally-null rows -/// since the normalized child is all-valid and the norms child carries the original nulls. -/// -/// Because this helper computes exact norms first and then divides by those norms, the returned -/// `normalized` child satisfies the strict unit-norm invariant required by [`L2Denorm`]. -pub fn normalize_as_l2_denorm( - 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 `L2Denorm` whose children are both `ConstantArray`s. - if let Some(wrapped) = try_build_constant_l2_denorm(&input, row_count, ctx)? { - return Ok(wrapped); - } - - // Calculate the norms of the vectors. - let norms_sfn = L2Norm::try_new_array(input.clone())?; - let norms_array: ArrayRef = norms_sfn.into_array().execute(ctx)?; - let primitive_norms: PrimitiveArray = norms_array.clone().execute(ctx)?; - let norms_validity = primitive_norms.validity()?; - - let input: ExtensionArray = input.execute(ctx)?; - let normalized_dtype = input.dtype().as_nonnullable(); - let flat = extract_flat_elements(input.storage_array(), tensor_flat_size, ctx)?; - - // Resolve validity to a mask once rather than probing it per row (each `Validity::is_valid` - // executes a scalar for array-backed validity). - let norms_valid = norms_validity.execute_mask(row_count, ctx)?; - - // Normalize all of the vectors. - let normalized = match_each_float_ptype!(flat.ptype(), |T| { - let norm_values = primitive_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]; - - // 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() { - unsafe { elements.push_n_unchecked(T::zero(), tensor_flat_size) }; - } else { - for &x in flat.row::(i) { - unsafe { elements.push_unchecked(x / norm) }; - } - } - } - - // Since L2Denorm's validity is the `and` of its child validities, we can make the - // normalized array non-nullable. - build_tensor_array( - normalized_dtype, - tensor_flat_size, - row_count, - Validity::NonNullable, - elements.freeze(), - ) - })?; - - // SAFETY: - // - `norms_array` was produced by `L2Norm(input)`, so every stored norm is non-negative and - // null rows already carry null validity through that child. - // - For every valid row, we either emit all zeros when the norm is zero or divide every - // element by the exact stored norm, so the normalized child is unit-norm (or zero) by - // construction. - // - Null rows are zeroed out above to avoid propagating arbitrary physical storage values into - // downstream lossy encodings. - unsafe { L2Denorm::new_array_unchecked(normalized, norms_array) } -} - -/// Attempts to build an [`L2Denorm`] whose two children are both [`ConstantArray`]s by eagerly -/// normalizing `input`'s single stored row. -/// -/// Returns `Ok(None)` when `input` is not a tensor-like extension array whose storage is a -/// [`ConstantArray`] with a non-null fixed-size-list scalar. -/// -/// When `input` matches, the returned [`ScalarFnArray`] is equivalent to [`normalize_as_l2_denorm`] -/// but runs in `O(list_size)` time instead of `O(row_count * list_size)`. -/// -/// This is helpful in some of the reduction steps for cosine similarity execution into inner -/// product execution. -pub(crate) fn try_build_constant_l2_denorm( - input: &ArrayRef, - len: usize, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let Some(ext) = input.as_opt::() else { - return Ok(None); - }; - let storage = ext.storage_array(); - let Some(const_storage) = storage.as_opt::() else { - return Ok(None); - }; - if const_storage.scalar().is_null() { - return Ok(None); - } - - // The caller is expected to have already validated that `input` is an `AnyTensor` - // extension dtype. - let tensor_match = input - .dtype() - .as_extension() - .metadata_opt::() - .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(); - - // Materialize just the single stored row; this does not expand the constant to the full - // column length. - let flat = extract_constant_flat_row(storage, ctx)?; - - let (normalized_fsl_scalar, norms_scalar) = match_each_float_ptype!(flat.ptype(), |T| { - let row = flat.as_slice::(); - - let mut sum_sq = T::zero(); - for &x in row { - sum_sq += x * x; - } - let norm_t: T = sum_sq.sqrt(); - - // Zero-norm rows must be stored as all-zeros so [`L2Denorm`]'s unit-norm-or-zero - // invariant holds. This mirrors the per-row logic in `normalize_as_l2_denorm`. - let element_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable); - let children: Vec = if norm_t == T::zero() { - (0..list_size) - .map(|_| Scalar::zero_value(&element_dtype)) - .collect() - } else { - row.iter() - .map(|&v| Scalar::primitive(v / norm_t, Nullability::NonNullable)) - .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); - (fsl_scalar, norms_scalar) - }); - - let normalized_storage = ConstantArray::new(normalized_fsl_scalar, len).into_array(); - let normalized_ext = ExtensionArray::new(ext_dtype, normalized_storage).into_array(); - let norms_array = ConstantArray::new(norms_scalar, len).into_array(); - - // SAFETY: Each row of `normalized_ext` is either `v / ||v||` (unit norm within floating - // point tolerance) or all zeros when `||v|| == 0`. Stored norms are non-negative by - // construction (`sqrt`). These are exactly the invariants required by - // [`L2Denorm::new_array_unchecked`]. - let wrapped = unsafe { L2Denorm::new_array_unchecked(normalized_ext, norms_array)? }; - Ok(Some(wrapped)) -} - -/// Rebuilds a tensor-like extension array from flat primitive elements. -/// -/// # Errors -/// -/// Returns an error if the elements are invalid (have incorrect lengths for the -/// `FixedSizeListArray` storage array). -fn build_tensor_array( - dtype: DType, - tensor_flat_size: usize, - row_count: usize, - validity: Validity, - elements: Buffer, -) -> VortexResult { - let list_size = - u32::try_from(tensor_flat_size).vortex_expect("tensor flat size must fit into `u32`"); - - // SAFETY: Validity has no length (because tensor elements are always non-nullable). - let elements = unsafe { PrimitiveArray::new_unchecked(elements, Validity::NonNullable) }; - - let storage = - FixedSizeListArray::try_new(elements.into_array(), list_size, validity, row_count)?; - - Ok(ExtensionArray::new(dtype.as_extension().clone(), storage.into_array()).into_array()) -} - -/// Validates that `normalized` and (when supplied) the matching `norms` jointly satisfy the -/// [`L2Denorm`] invariants: -/// -/// - Every valid row of `normalized` has L2 norm `1.0` or `0.0` (within element-precision -/// tolerance). -/// - 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`. -pub fn validate_l2_normalized_rows_against_norms( - normalized: &ArrayRef, - norms: Option<&ArrayRef>, - ctx: &mut ExecutionCtx, -) -> VortexResult<()> { - let row_count = normalized.len(); - if row_count == 0 { - return Ok(()); - } - - let tensor_match = validate_tensor_float_input(normalized.dtype())?; - let element_ptype = tensor_match.element_ptype(); - let tensor_flat_size = tensor_match.list_size() as usize; - let tolerance = unit_norm_tolerance(element_ptype, tensor_flat_size); - - if let Some(norms) = norms { - vortex_ensure_eq!( - norms.dtype().as_ptype(), - element_ptype, - "L2Denorm norms ptype must match normalized element ptype" - ); - } - - 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() - .fold((0.0f64, true), |(sum_sq, is_zero), x| { - let value = ToPrimitive::to_f64(x).unwrap_or(f64::NAN); - (sum_sq + value * value, is_zero && value.abs() <= tolerance) - }); - let row_norm = row_norm_sq.sqrt(); - - vortex_ensure!( - row_norm == 0.0 || (row_norm - 1.0).abs() <= tolerance, - "L2Denorm normalized child must have L2 norm 1.0 or 0.0, but row {i} has \ - {row_norm:.6}", - ); - - if let Some(stored_norms) = stored_norms { - let stored_norm_f64 = ToPrimitive::to_f64(&stored_norms[i]).unwrap_or(f64::NAN); - vortex_ensure!( - stored_norm_f64 >= 0.0, - "L2Denorm norms must be non-negative, but row {i} has {stored_norm_f64:.6}", - ); - - if stored_norm_f64 == 0.0 { - vortex_ensure!( - is_zero_row, - "L2Denorm normalized child must be all zeros when norms row {i} is 0.0", - ); - } - } - } - }); - - Ok(()) -} - -/// Classification of a binary operand pair by which side (if any) is wrapped in [`L2Denorm`]. -/// -/// Symmetric binary tensor operators (e.g. [`CosineSimilarity`], [`InnerProduct`]) have identical -/// fast paths for "only the lhs is denormalized" and "only the rhs is denormalized", and a separate -/// fast path for "both are denormalized". Rather than hand-rolling the commutative swap at every -/// call site, callers classify their operands with [`Self::classify`] and pattern-match on the -/// returned variant. -/// -/// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity -/// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct -pub(crate) enum DenormOrientation<'a> { - /// Both operands are [`ExactScalarFn`] arrays. - Both { - lhs: &'a ArrayRef, - rhs: &'a ArrayRef, - }, - /// Exactly one operand is an [`ExactScalarFn`]; the other is plain. - One { - denorm: &'a ArrayRef, - plain: &'a ArrayRef, - }, - /// Neither operand is an [`ExactScalarFn`]. - Neither, -} - -impl<'a> DenormOrientation<'a> { - /// Classify `(lhs, rhs)` by which side (if any) is wrapped in [`L2Denorm`]. - pub(crate) fn classify(lhs: &'a ArrayRef, rhs: &'a ArrayRef) -> Self { - let lhs_denorm = lhs.is::>(); - let rhs_denorm = rhs.is::>(); - match (lhs_denorm, rhs_denorm) { - (true, true) => Self::Both { lhs, rhs }, - (true, false) => Self::One { - denorm: lhs, - plain: rhs, - }, - (false, true) => Self::One { - denorm: rhs, - plain: lhs, - }, - (false, false) => Self::Neither, - } - } -} - -#[cfg(test)] -mod tests { - - use rstest::rstest; - use vortex_array::ArrayPlugin; - use vortex_array::ArrayRef; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Constant; - use vortex_array::arrays::ConstantArray; - use vortex_array::arrays::Extension; - 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::FixedSizeListArraySlotsExt; - use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; - use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::extension::ExtDType; - use vortex_array::extension::datetime::Date; - use vortex_array::extension::datetime::TimeUnit; - use vortex_array::scalar::Scalar; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - - use crate::scalar_fns::l2_denorm::L2Denorm; - use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm; - use crate::scalar_fns::l2_denorm::validate_l2_normalized_rows_against_norms; - use crate::tests::SESSION; - use crate::types::vector::Vector; - use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::constant_tensor_array; - use crate::utils::test_helpers::tensor_array; - use crate::utils::test_helpers::vector_array; - - /// Evaluates L2 denorm on a tensor/vector array and returns the executed array. - fn eval_l2_denorm(normalized: ArrayRef, norms: ArrayRef) -> VortexResult { - let mut ctx = SESSION.create_execution_ctx(); - let result = L2Denorm::try_new_array(normalized, norms, &mut ctx)?; - result.into_array().execute(&mut ctx) - } - - fn non_tensor_extension_array() -> VortexResult { - let storage = PrimitiveArray::from_iter([1i32, 2]).into_array(); - let ext_dtype = - ExtDType::::try_new(TimeUnit::Days, storage.dtype().clone())?.erased(); - Ok(ExtensionArray::new(ext_dtype, storage).into_array()) - } - - 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()) - .map(|i| ext.is_valid(i, &mut ctx)) - .collect::>>()?; - 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(), - )) - } - - 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)?; - - assert_eq!(actual_dtype, expected_dtype); - assert_eq!(actual_validity, expected_validity); - assert_close(&actual_elements, &expected_elements); - Ok(()) - } - - #[test] - fn l2_denorm_vectors() -> VortexResult<()> { - let lhs = vector_array(3, &[0.6, 0.8, 0.0, 0.0, 0.0, 0.0])?; - let rhs = PrimitiveArray::from_iter([5.0f64, 0.0]).into_array(); - let actual = eval_l2_denorm(lhs, rhs)?; - let expected = vector_array(3, &[3.0, 4.0, 0.0, 0.0, 0.0, 0.0])?; - - assert_tensor_arrays_eq(actual, expected)?; - Ok(()) - } - - #[test] - fn l2_denorm_fixed_shape_tensors() -> VortexResult<()> { - let lhs = tensor_array(&[2, 2], &[0.5, 0.5, 0.5, 0.5, 1.0, 0.0, 0.0, 0.0])?; - let rhs = PrimitiveArray::from_iter([4.0f64, 2.0]).into_array(); - let actual = eval_l2_denorm(lhs, rhs)?; - 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)?; - Ok(()) - } - - #[test] - fn l2_denorm_null_propagation() -> VortexResult<()> { - let lhs = vector_array(2, &[0.6, 0.8, 1.0, 0.0, 0.0, 0.0])?; - let lhs = MaskedArray::try_new(lhs, Validity::from_iter([true, false, true]))?.into_array(); - - let rhs = PrimitiveArray::from_option_iter([Some(5.0f64), Some(2.0), None]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - let actual: ExtensionArray = eval_l2_denorm(lhs, rhs)?.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_close(&elements.as_slice::()[..2], &[3.0, 4.0]); - Ok(()) - } - - #[test] - fn l2_denorm_rejects_non_extension_lhs() { - let lhs = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array(); - let rhs = PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - let result = L2Denorm::try_new_array(lhs, rhs, &mut ctx); - assert!(result.is_err()); - } - - #[test] - fn l2_denorm_rejects_non_tensor_extension_lhs() -> VortexResult<()> { - let lhs = non_tensor_extension_array()?; - let rhs = PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - let result = L2Denorm::try_new_array(lhs, rhs, &mut ctx); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn l2_denorm_rejects_integer_tensor_lhs() -> VortexResult<()> { - let lhs = tensor_array(&[2], &[1i32, 2, 3, 4])?; - let rhs = PrimitiveArray::from_iter([1.0f64, 1.0]).into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - let result = L2Denorm::try_new_array(lhs, rhs, &mut ctx); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn l2_denorm_rejects_mismatched_rhs_ptype() -> VortexResult<()> { - let lhs = vector_array(2, &[1.0, 0.0, 0.0, 1.0])?; - let rhs = PrimitiveArray::from_iter([1.0f32, 1.0]).into_array(); - - let mut ctx = SESSION.create_execution_ctx(); - let result = L2Denorm::try_new_array(lhs, rhs, &mut ctx); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn validate_l2_normalized_rows_accepts_normalized_f16_input() -> VortexResult<()> { - let input = vector_array(2, &[3.0f32, 4.0, 0.0, 0.0].map(half::f16::from_f32))?; - let mut ctx = SESSION.create_execution_ctx(); - let roundtrip = normalize_as_l2_denorm(input, &mut ctx)?; - validate_l2_normalized_rows_against_norms(&roundtrip.child_at(0).clone(), None, &mut ctx)?; - Ok(()) - } - - #[test] - fn validate_l2_normalized_rows_rejects_unnormalized_input() -> VortexResult<()> { - let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0])?; - let mut ctx = SESSION.create_execution_ctx(); - let result = validate_l2_normalized_rows_against_norms(&input, None, &mut ctx); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn l2_denorm_try_new_array_rejects_unnormalized_child() -> VortexResult<()> { - let normalized = vector_array(2, &[3.0, 4.0, 1.0, 0.0])?; - let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let result = L2Denorm::try_new_array(normalized, norms, &mut ctx); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn l2_denorm_try_new_array_rejects_nonzero_row_with_zero_norm() -> VortexResult<()> { - let normalized = vector_array(2, &[1.0, 0.0, 0.0, 0.0])?; - let norms = PrimitiveArray::from_iter([0.0f64, 0.0]).into_array(); - let mut ctx = SESSION.create_execution_ctx(); - - let result = L2Denorm::try_new_array(normalized, norms, &mut ctx); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn l2_denorm_try_new_array_rejects_negative_norms() -> VortexResult<()> { - let normalized = vector_array(2, &[1.0, 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 result = L2Denorm::try_new_array(normalized, norms, &mut ctx); - assert!(result.is_err()); - Ok(()) - } - - #[test] - fn l2_denorm_new_array_unchecked_accepts_unnormalized_child() -> VortexResult<()> { - let normalized = vector_array(2, &[3.0, 4.0, 1.0, 0.0])?; - let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); - - let result = unsafe { L2Denorm::new_array_unchecked(normalized, norms) }; - assert!(result.is_ok()); - Ok(()) - } - - #[test] - fn normalize_as_l2_denorm_roundtrips_vectors() -> VortexResult<()> { - let input = vector_array(3, &[3.0, 4.0, 0.0, 0.0, 0.0, 0.0])?; - let mut ctx = SESSION.create_execution_ctx(); - let roundtrip = normalize_as_l2_denorm(input.clone(), &mut ctx)?; - let actual = roundtrip.into_array().execute(&mut ctx)?; - - assert_tensor_arrays_eq(actual, input)?; - Ok(()) - } - - #[test] - fn normalize_as_l2_denorm_roundtrips_fixed_shape_tensors() -> VortexResult<()> { - let input = tensor_array(&[2, 2], &[1.0, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0])?; - let mut ctx = SESSION.create_execution_ctx(); - let roundtrip = normalize_as_l2_denorm(input.clone(), &mut ctx)?; - let actual = roundtrip.into_array().execute(&mut ctx)?; - - assert_tensor_arrays_eq(actual, input)?; - Ok(()) - } - - #[test] - fn normalize_as_l2_denorm_supports_constant_tensors() -> VortexResult<()> { - let input = constant_tensor_array(&[2], &[3.0, 4.0], 3)?; - let mut ctx = SESSION.create_execution_ctx(); - let roundtrip = normalize_as_l2_denorm(input.clone(), &mut ctx)?; - let actual = roundtrip.into_array().execute(&mut ctx)?; - - assert_tensor_arrays_eq(actual, input)?; - Ok(()) - } - - #[test] - fn normalize_as_l2_denorm_supports_constant_vectors() -> VortexResult<()> { - let input = Vector::constant_array(&[3.0, 4.0], 2)?; - let mut ctx = SESSION.create_execution_ctx(); - let roundtrip = normalize_as_l2_denorm(input.clone(), &mut ctx)?; - let actual = roundtrip.into_array().execute(&mut ctx)?; - - assert_tensor_arrays_eq(actual, input)?; - Ok(()) - } - - #[test] - fn normalize_as_l2_denorm_constant_input_has_constant_children() -> VortexResult<()> { - // The constant fast path in `normalize_as_l2_denorm` must produce an `L2Denorm` whose - // normalized storage and norms child are both still `ConstantArray`s. This is what - // allows downstream ops (cosine similarity, inner product) to short-circuit. - let input = Vector::constant_array(&[3.0, 4.0], 16)?; - let mut ctx = SESSION.create_execution_ctx(); - let roundtrip = normalize_as_l2_denorm(input, &mut ctx)?; - - // The normalized child must be an extension array whose storage is still constant. - let normalized = roundtrip.child_at(0).clone(); - let normalized_ext = normalized - .as_opt::() - .expect("normalized child should be an Extension array"); - assert!( - normalized_ext - .storage_array() - .as_opt::() - .is_some(), - "normalized storage should stay constant after the fast path" - ); - - // The norms child must itself be a ConstantArray with the exact precomputed norm. - let norms = roundtrip.child_at(1).clone(); - let norms_const = norms - .as_opt::() - .expect("norms child should be a ConstantArray"); - assert_close( - &[norms_const - .scalar() - .as_primitive() - .typed_value::() - .expect("norms scalar")], - &[5.0], - ); - Ok(()) - } - - #[test] - fn normalize_as_l2_denorm_uses_zero_rows_for_zero_norms() -> VortexResult<()> { - let input = vector_array(2, &[0.0, 0.0, 3.0, 4.0])?; - let mut ctx = SESSION.create_execution_ctx(); - let roundtrip = normalize_as_l2_denorm(input.clone(), &mut ctx)?; - let normalized: ExtensionArray = roundtrip.child_at(0).clone().execute(&mut ctx)?; - let storage: FixedSizeListArray = normalized.storage_array().clone().execute(&mut ctx)?; - let elements: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; - let actual = roundtrip.into_array().execute(&mut ctx)?; - - assert_close(&elements.as_slice::()[..2], &[0.0, 0.0]); - assert_tensor_arrays_eq(actual, input)?; - Ok(()) - } - - /// Builds a non-nullable constant f64 norms array of length `len`. - fn constant_f64_norms(value: f64, len: usize) -> ArrayRef { - ConstantArray::new(Scalar::primitive(value, Nullability::NonNullable), len).into_array() - } - - #[test] - fn l2_denorm_constant_unit_norms_is_noop() -> VortexResult<()> { - // Every stored norm is exactly 1.0, so the constant fast path must short-circuit and - // return the normalized child unchanged. - 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_l2_denorm(normalized.clone(), norms)?; - assert_tensor_arrays_eq(actual, normalized)?; - Ok(()) - } - - #[test] - fn l2_denorm_constant_near_unit_norms_is_noop() -> VortexResult<()> { - // A norm that differs from 1.0 by less than the f64 unit-norm tolerance must still - // hit the fast path and return the normalized child unchanged. - 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 actual = eval_l2_denorm(normalized.clone(), norms)?; - assert_tensor_arrays_eq(actual, normalized)?; - Ok(()) - } - - #[test] - fn l2_denorm_constant_nonunit_norms_scales_vectors() -> VortexResult<()> { - // A constant norm that is not 1.0 must scale every element of every row by the same - // factor via the backing elements multiplication path. - 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_l2_denorm(normalized, norms)?; - let expected = vector_array(3, &[3.0, 4.0, 0.0, 5.0, 0.0, 0.0])?; - assert_tensor_arrays_eq(actual, expected)?; - Ok(()) - } - - #[test] - fn l2_denorm_constant_nonunit_norms_scales_fixed_shape_tensors() -> VortexResult<()> { - // The same constant-scaling fast path must also cover multi-dimensional fixed-shape - // tensors, where the backing elements buffer spans more than one slot per row. - 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_l2_denorm(normalized, norms)?; - 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)?; - Ok(()) - } - - /// Build an `L2Denorm` array from a raw input (which may have nullable storage) by running - /// `normalize_as_l2_denorm`. The normalized child ends up non-nullable, and the norms child - /// inherits the input's nullability, giving us two different per-child nullabilities to - /// round-trip. - #[rstest] - #[case::vector(l2_denorm_vector_input())] - #[case::fixed_shape_tensor(l2_denorm_tensor_input())] - fn serde_round_trip(#[case] input: ArrayRef) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let original = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); - - let scalar_fn_array = original.as_::(); - let children = scalar_fn_array.children(); - - let plugin = ScalarFnArrayPlugin::new(L2Denorm); - let metadata = plugin - .serialize(&original, &SESSION)? - .expect("L2Denorm serialize must produce metadata"); - - let recovered = plugin.deserialize( - original.dtype(), - original.len(), - &metadata, - &[], - &children, - &SESSION, - )?; - - assert_eq!(recovered.dtype(), original.dtype()); - assert_eq!(recovered.len(), original.len()); - assert_eq!(recovered.encoding_id(), original.encoding_id()); - Ok(()) - } - - fn l2_denorm_vector_input() -> ArrayRef { - vector_array(3, &[3.0, 4.0, 0.0, 0.0, 0.0, 0.0]).expect("valid vector array") - } - - fn l2_denorm_tensor_input() -> ArrayRef { - tensor_array(&[2, 2], &[1.0, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]) - .expect("valid tensor array") - } -} diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 73c1538abc9..0533fa0b055 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -15,7 +15,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn as ScalarFnArrayEncoding; use vortex_array::arrays::ScalarFnArray; use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts; @@ -44,8 +43,8 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::encodings::l2_denorm::L2Denorm; use crate::matcher::AnyTensor; -use crate::scalar_fns::l2_denorm::L2Denorm; use crate::utils::extract_flat_elements; use crate::utils::extract_l2_denorm_children; use crate::utils::validate_tensor_float_input; @@ -57,10 +56,12 @@ use crate::utils::validate_tensor_float_input; /// The input must be a tensor-like extension array with a float element type. The output is a float /// column of the same float type. /// -/// When the input is wrapped in [`L2Denorm`], this operator treats the stored norms as -/// authoritative. For lossy encodings, that means `L2Norm` may intentionally -/// read the stored norms instead of re-deriving them from fully decoded coordinates. That behavior -/// is part of the lossy storage contract, not a separate lossy-compute mode. +/// When the input is [`L2Denorm`]-encoded, this operator treats the stored norms as +/// authoritative. For lossy normalized children, that means `L2Norm` intentionally reads the +/// stored norms instead of re-deriving them from fully decoded coordinates. That behavior is part +/// of the storage contract, not a separate lossy-compute mode. +/// +/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm #[derive(Clone)] pub struct L2Norm; @@ -127,10 +128,10 @@ impl ScalarFnVTable for L2Norm { let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); - // L2Norm(L2Denorm(normalized, norms)) is defined to read back the authoritative stored - // norms. Exact callers of lossy encodings opt into that storage semantics - // instead of forcing a decode-and-recompute path here. - if input_ref.is::>() { + // L2Norm over an L2Denorm-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::() { let (_, norms) = extract_l2_denorm_children(&input_ref); vortex_ensure_eq!(norms.dtype(), &norm_dtype); return Ok(norms); diff --git a/vortex-tensor/src/scalar_fns/mod.rs b/vortex-tensor/src/scalar_fns/mod.rs index 9acf59bda01..68f10ca6b01 100644 --- a/vortex-tensor/src/scalar_fns/mod.rs +++ b/vortex-tensor/src/scalar_fns/mod.rs @@ -5,5 +5,4 @@ pub mod cosine_similarity; pub mod inner_product; -pub mod l2_denorm; pub mod l2_norm; diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index e8dbf16f171..d017820aa4f 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -13,7 +13,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFn; use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::arrays::scalar_fn::ExactScalarFn; use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; use vortex_array::arrays::scalar_fn::ScalarFnArrayView; use vortex_array::dtype::DType; @@ -27,9 +26,10 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; +use crate::encodings::l2_denorm::L2Denorm; +use crate::encodings::l2_denorm::L2DenormArraySlotsExt; use crate::matcher::AnyTensor; use crate::matcher::TensorMatch; -use crate::scalar_fns::l2_denorm::L2Denorm; /// Safety factor for unit-norm tolerance. Applied as a constant multiplier on the probabilistic /// `√d · ε` bound so that legitimate round-off noise clears the check with headroom. @@ -58,18 +58,21 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root } -/// Extracts the `(normalized, norms)` children from an [`L2Denorm`] scalar function array. +/// Extracts the `(normalized, norms)` children of an [`L2Denorm`]-encoded array. /// -/// [`L2Denorm`]: crate::scalar_fns::l2_denorm::L2Denorm +/// # Panics +/// +/// Panics if `array` is not [`L2Denorm`]-encoded. Callers reach this through +/// [`DenormOrientation::classify`], which has already matched on the encoding. +/// +/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm +/// [`DenormOrientation::classify`]: crate::encodings::l2_denorm::DenormOrientation::classify pub fn extract_l2_denorm_children(array: &ArrayRef) -> (ArrayRef, ArrayRef) { - let sfn = array - .as_opt::>() - .vortex_expect("expected ScalarFnArray wrapping L2Denorm"); - ( - sfn.nth_child(0) - .vortex_expect("L2Denorm missing normalized array"), - sfn.nth_child(1).vortex_expect("L2Denorm missing norms"), - ) + let denorm = array + .as_opt::() + .vortex_expect("expected an L2Denorm-encoded array"); + + (denorm.normalized().clone(), denorm.norms().clone()) } /// Validates that `input_dtype` is a float-valued tensor-like extension dtype. @@ -296,7 +299,7 @@ pub mod test_helpers { use vortex_buffer::Buffer; use vortex_error::VortexResult; - use crate::scalar_fns::l2_denorm::L2Denorm; + use crate::encodings::l2_denorm::L2Denorm; use crate::types::fixed_shape_tensor::FixedShapeTensor; use crate::types::fixed_shape_tensor::FixedShapeTensorMetadata; use crate::types::vector::Vector; @@ -364,9 +367,9 @@ pub mod test_helpers { ConstantArray::new(ext_scalar, len).into_array() } - /// Creates an [`L2Denorm`] scalar function array from pre-normalized tensor elements and - /// matching norms. The caller must ensure every row of `normalized_elements` is unit-norm or - /// zero. + /// Creates an [`L2Denorm`] 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 l2_denorm_array( shape: &[usize], normalized_elements: &[T], @@ -376,7 +379,7 @@ pub mod test_helpers { let normalized = tensor_array(shape, normalized_elements)?; let norms = PrimitiveArray::new(Buffer::copy_from(norms), Validity::NonNullable).into_array(); - Ok(L2Denorm::try_new_array(normalized, norms, ctx)?.into_array()) + Ok(L2Denorm::try_new(normalized, norms, ctx)?.into_array()) } /// Asserts that each element in `actual` is within `1e-10` of the corresponding `expected` From c2288dc7f280949c9fbfa7b96422212da4ed1fe0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 3 Aug 2026 14:50:21 -0400 Subject: [PATCH 2/2] Rename L2Denorm to Normalized Signed-off-by: Connor Tsui --- vortex-tensor/src/encodings/mod.rs | 2 +- .../{l2_denorm => normalized}/array.rs | 88 +++++----- .../{l2_denorm => normalized}/compress.rs | 63 ++++--- .../{l2_denorm => normalized}/execute.rs | 4 +- .../{l2_denorm => normalized}/mod.rs | 22 +-- .../{l2_denorm => normalized}/orientation.rs | 32 ++-- .../{l2_denorm => normalized}/rules.rs | 32 ++-- .../{l2_denorm => normalized}/tests.rs | 162 ++++++++++-------- .../{l2_denorm => normalized}/validate.rs | 36 ++-- vortex-tensor/src/lib.rs | 6 +- .../src/scalar_fns/cosine_similarity.rs | 137 +++++++-------- vortex-tensor/src/scalar_fns/inner_product.rs | 87 +++++----- vortex-tensor/src/scalar_fns/l2_norm.rs | 14 +- vortex-tensor/src/utils.rs | 37 ++-- vortex/src/editions/unstable/v2026_04.rs | 2 +- 15 files changed, 376 insertions(+), 348 deletions(-) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/array.rs (76%) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/compress.rs (85%) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/execute.rs (97%) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/mod.rs (71%) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/orientation.rs (52%) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/rules.rs (66%) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/tests.rs (80%) rename vortex-tensor/src/encodings/{l2_denorm => normalized}/validate.rs (77%) diff --git a/vortex-tensor/src/encodings/mod.rs b/vortex-tensor/src/encodings/mod.rs index 627c9d7c7d0..8587a58c868 100644 --- a/vortex-tensor/src/encodings/mod.rs +++ b/vortex-tensor/src/encodings/mod.rs @@ -6,4 +6,4 @@ // TODO(connor): // pub mod spherical; // Spherical transform on unit-normalized vectors. -pub mod l2_denorm; +pub mod normalized; diff --git a/vortex-tensor/src/encodings/l2_denorm/array.rs b/vortex-tensor/src/encodings/normalized/array.rs similarity index 76% rename from vortex-tensor/src/encodings/l2_denorm/array.rs rename to vortex-tensor/src/encodings/normalized/array.rs index f3f621c3415..c60bf556b02 100644 --- a/vortex-tensor/src/encodings/l2_denorm/array.rs +++ b/vortex-tensor/src/encodings/normalized/array.rs @@ -28,14 +28,14 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::encodings::l2_denorm::execute::denormalize; -use crate::encodings::l2_denorm::rules::RULES; -use crate::encodings::l2_denorm::validate::validate_l2_denorm_children; -use crate::encodings::l2_denorm::validate::validate_l2_normalized_rows_against_norms; +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::utils::validate_tensor_float_input; -/// An [`L2Denorm`]-encoded Vortex array. -pub type L2DenormArray = Array; +/// An [`Normalized`]-encoded Vortex array. +pub type NormalizedArray = Array; /// The norm-split encoding for tensor-like columns. /// @@ -45,7 +45,7 @@ pub type L2DenormArray = Array; /// /// # Invariants /// -/// Every [`L2DenormArray`] structurally guarantees, via [`VTable::validate`]: +/// 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. @@ -75,11 +75,11 @@ pub type L2DenormArray = Array; /// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct /// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity #[derive(Clone, Debug)] -pub struct L2Denorm; +pub struct Normalized; -/// The two child arrays of an [`L2DenormArray`]. -#[array_slots(L2Denorm)] -pub struct L2DenormSlots { +/// The two child arrays of an [`NormalizedArray`]. +#[array_slots(Normalized)] +pub struct NormalizedSlots { /// The unit-norm (or zero) direction of each row, as a tensor-like extension array. #[slot(0)] pub normalized: ArrayRef, @@ -89,8 +89,8 @@ pub struct L2DenormSlots { pub norms: ArrayRef, } -impl L2Denorm { - /// Builds an [`L2DenormArray`], validating that `normalized` really is row-wise L2-normalized +impl Normalized { + /// Builds an [`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 @@ -99,59 +99,63 @@ impl L2Denorm { /// # Errors /// /// Returns an error if the children are structurally incompatible, or if they violate any of - /// the semantic invariants listed on [`L2Denorm`]. + /// the semantic invariants listed on [`Normalized`]. pub fn try_new( normalized: ArrayRef, norms: ArrayRef, ctx: &mut ExecutionCtx, - ) -> VortexResult { + ) -> VortexResult { let len = normalized.len(); let dtype = normalized .dtype() .union_nullability(norms.dtype().nullability()); - let slots = L2DenormSlots { normalized, norms }.into_slots(); + 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 denorm = Array::try_from_parts( - ArrayParts::new(L2Denorm, dtype, len, EmptyArrayData).with_slots(slots), + 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, )?; - validate_l2_normalized_rows_against_norms(denorm.normalized(), Some(denorm.norms()), ctx)?; - Ok(denorm) + Ok(normalized_array) } - /// Builds an [`L2DenormArray`] without validation. + /// Builds an [`NormalizedArray`] without validation. /// /// # Safety /// - /// The caller must uphold the structural invariants listed on [`L2Denorm`]. In particular, + /// 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. /// /// 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) -> L2DenormArray { + 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 = L2DenormSlots { normalized, norms }.into_slots(); + let slots = NormalizedSlots { normalized, norms }.into_slots(); unsafe { Array::from_parts_unchecked( - ArrayParts::new(L2Denorm, dtype, len, EmptyArrayData).with_slots(slots), + ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots), ) } } } -/// Metadata for a serialized [`L2DenormArray`]: its children's nullabilities. +/// Metadata for a serialized [`NormalizedArray`]: its children's nullabilities. /// /// 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 L2DenormMetadata { +pub struct NormalizedMetadata { /// Whether the `normalized` child is nullable. #[prost(bool, tag = "1")] pub normalized_is_nullable: bool, @@ -161,14 +165,14 @@ pub struct L2DenormMetadata { pub norms_is_nullable: bool, } -impl VTable for L2Denorm { +impl VTable for Normalized { type TypedArrayData = EmptyArrayData; type OperationsVTable = Self; type ValidityVTable = Self; fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.tensor.l2_denorm"); + static ID: CachedId = CachedId::new("vortex.tensor.normalized"); *ID } @@ -179,9 +183,9 @@ impl VTable for L2Denorm { len: usize, slots: &[Option], ) -> VortexResult<()> { - let slots = L2DenormSlotsView::from_slots(slots); + let slots = NormalizedSlotsView::from_slots(slots); - validate_l2_denorm_children(slots.normalized, slots.norms, dtype, len) + validate_normalized_children(slots.normalized, slots.norms, dtype, len) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -189,11 +193,11 @@ impl VTable for L2Denorm { } fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - vortex_panic!("L2DenormArray buffer index {idx} out of bounds") + vortex_panic!("NormalizedArray buffer index {idx} out of bounds") } fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - vortex_panic!("L2DenormArray buffer_name index {idx} out of bounds") + vortex_panic!("NormalizedArray buffer_name index {idx} out of bounds") } fn with_buffers( @@ -209,7 +213,7 @@ impl VTable for L2Denorm { _session: &VortexSession, ) -> VortexResult>> { Ok(Some( - L2DenormMetadata { + NormalizedMetadata { normalized_is_nullable: array.normalized().dtype().is_nullable(), norms_is_nullable: array.norms().dtype().is_nullable(), } @@ -226,8 +230,8 @@ impl VTable for L2Denorm { children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - let metadata = L2DenormMetadata::decode(metadata) - .map_err(|e| vortex_err!("Failed to decode L2DenormMetadata: {e}"))?; + let metadata = NormalizedMetadata::decode(metadata) + .map_err(|e| vortex_err!("Failed to decode NormalizedMetadata: {e}"))?; let element_ptype = validate_tensor_float_input(dtype)?.element_ptype(); let normalized_dtype = dtype.with_nullability(metadata.normalized_is_nullable.into()); @@ -235,13 +239,13 @@ impl VTable for L2Denorm { let normalized = children.get(0, &normalized_dtype, len)?; let norms = children.get(1, &norms_dtype, len)?; - let slots = L2DenormSlots { normalized, norms }.into_slots(); + let slots = NormalizedSlots { normalized, norms }.into_slots(); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - L2DenormSlots::NAMES[idx].to_string() + NormalizedSlots::NAMES[idx].to_string() } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { @@ -261,8 +265,8 @@ impl VTable for L2Denorm { } } -impl ValidityVTable for L2Denorm { - fn validity(array: ArrayView<'_, L2Denorm>) -> VortexResult { +impl ValidityVTable for Normalized { + fn validity(array: ArrayView<'_, Normalized>) -> VortexResult { array .normalized() .validity()? @@ -270,9 +274,9 @@ impl ValidityVTable for L2Denorm { } } -impl OperationsVTable for L2Denorm { +impl OperationsVTable for Normalized { fn scalar_at( - array: ArrayView<'_, L2Denorm>, + array: ArrayView<'_, Normalized>, index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { diff --git a/vortex-tensor/src/encodings/l2_denorm/compress.rs b/vortex-tensor/src/encodings/normalized/compress.rs similarity index 85% rename from vortex-tensor/src/encodings/l2_denorm/compress.rs rename to vortex-tensor/src/encodings/normalized/compress.rs index 18ce117cea6..727483296e4 100644 --- a/vortex-tensor/src/encodings/l2_denorm/compress.rs +++ b/vortex-tensor/src/encodings/normalized/compress.rs @@ -36,23 +36,23 @@ use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use crate::encodings::l2_denorm::L2Denorm; -use crate::encodings::l2_denorm::L2DenormArray; -use crate::encodings::l2_denorm::L2DenormArraySlotsExt; -use crate::encodings::l2_denorm::L2DenormSlots; +use crate::encodings::normalized::Normalized; +use crate::encodings::normalized::NormalizedArray; +use crate::encodings::normalized::NormalizedArraySlotsExt; +use crate::encodings::normalized::NormalizedSlots; use crate::matcher::AnyTensor; use crate::scalar_fns::l2_norm::L2Norm; use crate::utils::extract_constant_flat_row; use crate::utils::extract_flat_elements; use crate::utils::validate_tensor_float_input; -/// The compression scheme that rewrites a tensor-like column into the [`L2Denorm`] encoding. +/// The compression scheme that rewrites a tensor-like column into the [`Normalized`] encoding. #[derive(Debug)] -pub struct L2DenormScheme; +pub struct NormalizedScheme; -impl Scheme for L2DenormScheme { +impl Scheme for NormalizedScheme { fn scheme_name(&self) -> &'static str { - "vortex.tensor.l2_denorm" + "vortex.tensor.normalized" } fn matches(&self, canonical: &Canonical) -> bool { @@ -63,12 +63,12 @@ impl Scheme for L2DenormScheme { } fn produced_encodings(&self) -> Vec { - vec![L2Denorm.id()] + vec![Normalized.id()] } /// Children: normalized=0, norms=1. fn num_children(&self) -> usize { - L2DenormSlots::COUNT + NormalizedSlots::COUNT } fn expected_compression_ratio( @@ -87,31 +87,31 @@ impl Scheme for L2DenormScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let denorm = normalize_as_l2_denorm(data.array().clone(), exec_ctx)?; + let normalized_array = normalize(data.array().clone(), exec_ctx)?; // Splitting magnitude out is only worth anything if the children then compress: the // unit-norm coordinates have a bounded range and the norms are an ordinary float column. let normalized = compressor.compress_child( - denorm.normalized(), + normalized_array.normalized(), &compress_ctx, self.id(), - L2DenormSlots::NORMALIZED, + NormalizedSlots::NORMALIZED, exec_ctx, )?; let norms = compressor.compress_child( - denorm.norms(), + normalized_array.norms(), &compress_ctx, self.id(), - L2DenormSlots::NORMS, + NormalizedSlots::NORMS, exec_ctx, )?; // SAFETY: Cascading preserves the split's child lengths and dtypes. - Ok(unsafe { L2Denorm::new_unchecked(normalized, norms) }.into_array()) + Ok(unsafe { Normalized::new_unchecked(normalized, norms) }.into_array()) } } -/// Splits a tensor-like column into its exact [`L2Denorm`] representation. +/// Splits a tensor-like column into its exact [`Normalized`] representation. /// /// # Normalized child /// @@ -125,23 +125,20 @@ impl Scheme for L2DenormScheme { /// # Nullability /// /// Nullability is tracked entirely by the norms child, which inherits the input's nulls through -/// [`L2Norm`]'s validity propagation. The [`L2Denorm`] array's validity is the `and` of both +/// [`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. /// /// Because this computes exact norms first and then divides by them, the returned `normalized` /// child satisfies the strict unit-norm invariant. -pub fn normalize_as_l2_denorm( - input: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { +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 `L2Denorm` whose children are both `ConstantArray`s. - if let Some(wrapped) = try_build_constant_l2_denorm(&input, row_count, ctx)? { + // row once and return an `Normalized` whose children are both `ConstantArray`s. + if let Some(wrapped) = try_build_constant_normalized(&input, row_count, ctx)? { return Ok(wrapped); } @@ -181,7 +178,7 @@ pub fn normalize_as_l2_denorm( } } - // Since L2Denorm's validity is the `and` of its child validities, the normalized child can + // Since Normalized's validity is the `and` of its child validities, the normalized child can // be non-nullable. build_normalized( normalized_dtype, @@ -192,23 +189,23 @@ pub fn normalize_as_l2_denorm( })?; // SAFETY: The normalized rows, norms ptype, and child lengths come directly from this split. - Ok(unsafe { L2Denorm::new_unchecked(normalized, norms_array) }) + Ok(unsafe { Normalized::new_unchecked(normalized, norms_array) }) } -/// Attempts to build an [`L2DenormArray`] whose two children are both [`ConstantArray`]s by +/// Attempts to build an [`NormalizedArray`] whose two children are both [`ConstantArray`]s by /// eagerly normalizing `input`'s single stored row. /// /// Returns `Ok(None)` when `input` is not a tensor-like extension array whose storage is a /// [`ConstantArray`] with a non-null fixed-size-list scalar. /// -/// When `input` matches, the result is equivalent to [`normalize_as_l2_denorm`] but runs in +/// When `input` matches, the result is equivalent to [`normalize`] but runs in /// `O(list_size)` instead of `O(row_count * list_size)`. Keeping both children constant is what /// lets cosine similarity and inner product short-circuit against a literal query vector. -pub(crate) fn try_build_constant_l2_denorm( +pub(crate) fn try_build_constant_normalized( input: &ArrayRef, len: usize, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult> { let Some(ext) = input.as_opt::() else { return Ok(None); }; @@ -246,7 +243,7 @@ pub(crate) fn try_build_constant_l2_denorm( let norm_t: T = sum_sq.sqrt(); // Zero-norm rows must be stored as all-zeros so the unit-norm-or-zero invariant holds. - // This mirrors the per-row logic in `normalize_as_l2_denorm`. + // This mirrors the per-row logic in `normalize`. let element_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable); let children: Vec = if norm_t == T::zero() { (0..list_size) @@ -270,7 +267,9 @@ pub(crate) fn try_build_constant_l2_denorm( let norms = ConstantArray::new(norms_scalar, len).into_array(); // SAFETY: The constant children have matching lengths and element ptypes. - Ok(Some(unsafe { L2Denorm::new_unchecked(normalized, norms) })) + Ok(Some(unsafe { + Normalized::new_unchecked(normalized, norms) + })) } /// Builds the non-nullable tensor-like extension array that becomes the `normalized` child. diff --git a/vortex-tensor/src/encodings/l2_denorm/execute.rs b/vortex-tensor/src/encodings/normalized/execute.rs similarity index 97% rename from vortex-tensor/src/encodings/l2_denorm/execute.rs rename to vortex-tensor/src/encodings/normalized/execute.rs index f1a5ee30c42..637c8c07117 100644 --- a/vortex-tensor/src/encodings/l2_denorm/execute.rs +++ b/vortex-tensor/src/encodings/normalized/execute.rs @@ -29,10 +29,10 @@ use crate::utils::unit_norm_tolerance; /// Reconstructs the original tensor column by scaling each normalized row by its stored norm. /// -/// `dtype` is the parent [`L2DenormArray`]'s dtype, so the reconstructed column carries the +/// `dtype` is the parent [`NormalizedArray`]'s dtype, so the reconstructed column carries the /// unioned nullability of both children. /// -/// [`L2DenormArray`]: crate::encodings::l2_denorm::L2DenormArray +/// [`NormalizedArray`]: crate::encodings::normalized::NormalizedArray pub(super) fn denormalize( normalized: &ArrayRef, norms: &ArrayRef, diff --git a/vortex-tensor/src/encodings/l2_denorm/mod.rs b/vortex-tensor/src/encodings/normalized/mod.rs similarity index 71% rename from vortex-tensor/src/encodings/l2_denorm/mod.rs rename to vortex-tensor/src/encodings/normalized/mod.rs index 0a8e6a213ec..545236bba7d 100644 --- a/vortex-tensor/src/encodings/l2_denorm/mod.rs +++ b/vortex-tensor/src/encodings/normalized/mod.rs @@ -1,9 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The [`L2Denorm`] encoding: a norm-split physical layout for tensor-like columns. +//! The [`Normalized`] encoding: a norm-split physical layout for tensor-like columns. //! -//! An [`L2Denorm`] array stores a tensor or vector column as two children: +//! An [`Normalized`] array stores a tensor or vector column as two children: //! //! - `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. @@ -22,21 +22,21 @@ //! [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity mod array; -pub use array::L2Denorm; -pub use array::L2DenormArray; -pub use array::L2DenormArraySlotsExt; -pub use array::L2DenormMetadata; -pub use array::L2DenormSlots; +pub use array::Normalized; +pub use array::NormalizedArray; +pub use array::NormalizedArraySlotsExt; +pub use array::NormalizedMetadata; +pub use array::NormalizedSlots; mod compress; -pub use compress::L2DenormScheme; -pub use compress::normalize_as_l2_denorm; -pub(crate) use compress::try_build_constant_l2_denorm; +pub use compress::NormalizedScheme; +pub use compress::normalize; +pub(crate) use compress::try_build_constant_normalized; mod execute; mod orientation; -pub(crate) use orientation::DenormOrientation; +pub(crate) use orientation::NormalizedOrientation; mod rules; diff --git a/vortex-tensor/src/encodings/l2_denorm/orientation.rs b/vortex-tensor/src/encodings/normalized/orientation.rs similarity index 52% rename from vortex-tensor/src/encodings/l2_denorm/orientation.rs rename to vortex-tensor/src/encodings/normalized/orientation.rs index f4675410f32..5c1f2364c5d 100644 --- a/vortex-tensor/src/encodings/l2_denorm/orientation.rs +++ b/vortex-tensor/src/encodings/normalized/orientation.rs @@ -3,19 +3,19 @@ use vortex_array::ArrayRef; -use crate::encodings::l2_denorm::L2Denorm; +use crate::encodings::normalized::Normalized; -/// Classification of a binary operand pair by which side (if any) is [`L2Denorm`]-encoded. +/// Classification of a binary operand pair by which side (if any) is [`Normalized`]-encoded. /// /// Symmetric binary tensor operators ([`CosineSimilarity`], [`InnerProduct`]) have identical fast -/// paths for "only the lhs is denormalized" and "only the rhs is denormalized", plus a separate -/// fast path for "both are denormalized". Rather than hand-rolling the commutative swap at every -/// call site, callers classify their operands with [`Self::classify`] and match on the result. +/// paths when only one operand is [`Normalized`], plus a separate fast path when both operands are +/// [`Normalized`]. Rather than hand-rolling the commutative swap at every call site, callers +/// classify their operands with [`Self::classify`] and match on the result. /// /// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity /// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct -pub(crate) enum DenormOrientation<'a> { - /// Both operands are [`L2Denorm`] arrays. +pub(crate) enum NormalizedOrientation<'a> { + /// Both operands are [`Normalized`] arrays. Both { /// The left-hand operand. lhs: &'a ArrayRef, @@ -23,29 +23,29 @@ pub(crate) enum DenormOrientation<'a> { rhs: &'a ArrayRef, }, - /// Exactly one operand is an [`L2Denorm`] array; the other is a plain tensor column. + /// Exactly one operand is a [`Normalized`] array; the other is a plain tensor column. One { - /// The [`L2Denorm`]-encoded operand, whichever side it came from. - denorm: &'a ArrayRef, + /// The [`Normalized`]-encoded operand, whichever side it came from. + normalized_array: &'a ArrayRef, /// The other operand. plain: &'a ArrayRef, }, - /// Neither operand is an [`L2Denorm`] array. + /// Neither operand is a [`Normalized`] array. Neither, } -impl<'a> DenormOrientation<'a> { - /// Classify `(lhs, rhs)` by which side (if any) is [`L2Denorm`]-encoded. +impl<'a> NormalizedOrientation<'a> { + /// Classify `(lhs, rhs)` by which side (if any) is [`Normalized`]-encoded. pub(crate) fn classify(lhs: &'a ArrayRef, rhs: &'a ArrayRef) -> Self { - match (lhs.is::(), rhs.is::()) { + match (lhs.is::(), rhs.is::()) { (true, true) => Self::Both { lhs, rhs }, (true, false) => Self::One { - denorm: lhs, + normalized_array: lhs, plain: rhs, }, (false, true) => Self::One { - denorm: rhs, + normalized_array: rhs, plain: lhs, }, (false, false) => Self::Neither, diff --git a/vortex-tensor/src/encodings/l2_denorm/rules.rs b/vortex-tensor/src/encodings/normalized/rules.rs similarity index 66% rename from vortex-tensor/src/encodings/l2_denorm/rules.rs rename to vortex-tensor/src/encodings/normalized/rules.rs index 40ffcac98cf..7db946c25b0 100644 --- a/vortex-tensor/src/encodings/l2_denorm/rules.rs +++ b/vortex-tensor/src/encodings/normalized/rules.rs @@ -10,28 +10,28 @@ use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_error::VortexResult; -use crate::encodings::l2_denorm::L2Denorm; -use crate::encodings::l2_denorm::array::L2DenormArraySlotsExt; +use crate::encodings::normalized::Normalized; +use crate::encodings::normalized::array::NormalizedArraySlotsExt; -pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&L2DenormSliceRule), - ParentRuleSet::lift(&L2DenormFilterRule), +pub(super) const RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&NormalizedSliceRule), + ParentRuleSet::lift(&NormalizedFilterRule), ]); /// Pushes a slice through the encoding into both children. /// -/// The norm split is row-wise, so any row subset of an [`L2Denorm`] array is itself a valid -/// [`L2Denorm`] array. Rewriting the slice as two child slices keeps the column encoded instead of +/// The norm split is row-wise, so any row subset of a [`Normalized`] array is itself a valid +/// [`Normalized`] array. Rewriting the slice as two child slices keeps the column encoded instead of /// canonicalizing it just to throw most of the rows away. #[derive(Debug)] -struct L2DenormSliceRule; +struct NormalizedSliceRule; -impl ArrayParentReduceRule for L2DenormSliceRule { +impl ArrayParentReduceRule for NormalizedSliceRule { type Parent = Slice; fn reduce_parent( &self, - array: ArrayView<'_, L2Denorm>, + array: ArrayView<'_, Normalized>, parent: ArrayView<'_, Slice>, _child_idx: usize, ) -> VortexResult> { @@ -40,7 +40,7 @@ impl ArrayParentReduceRule for L2DenormSliceRule { // SAFETY: Slicing both children preserves their structure. Ok(Some( unsafe { - L2Denorm::new_unchecked( + Normalized::new_unchecked( array.normalized().slice(range.clone())?, array.norms().slice(range.clone())?, ) @@ -52,18 +52,18 @@ impl ArrayParentReduceRule for L2DenormSliceRule { /// Pushes a filter through the encoding into both children. /// -/// Same row-wise argument as [`L2DenormSliceRule`]. Unlike the generic scalar-function push-down, +/// Same row-wise argument as [`NormalizedSliceRule`]. Unlike the generic scalar-function push-down, /// this always fires: both children are physically per-row, so filtering them is strictly less /// work than reconstructing the tensor column and filtering that. #[derive(Debug)] -struct L2DenormFilterRule; +struct NormalizedFilterRule; -impl ArrayParentReduceRule for L2DenormFilterRule { +impl ArrayParentReduceRule for NormalizedFilterRule { type Parent = Filter; fn reduce_parent( &self, - array: ArrayView<'_, L2Denorm>, + array: ArrayView<'_, Normalized>, parent: ArrayView<'_, Filter>, _child_idx: usize, ) -> VortexResult> { @@ -72,7 +72,7 @@ impl ArrayParentReduceRule for L2DenormFilterRule { // SAFETY: Filtering both children with the same mask preserves their structure. Ok(Some( unsafe { - L2Denorm::new_unchecked( + Normalized::new_unchecked( array.normalized().filter(mask.clone())?, array.norms().filter(mask.clone())?, ) diff --git a/vortex-tensor/src/encodings/l2_denorm/tests.rs b/vortex-tensor/src/encodings/normalized/tests.rs similarity index 80% rename from vortex-tensor/src/encodings/l2_denorm/tests.rs rename to vortex-tensor/src/encodings/normalized/tests.rs index 5914e9c6ac7..a3dd603a7c3 100644 --- a/vortex-tensor/src/encodings/l2_denorm/tests.rs +++ b/vortex-tensor/src/encodings/normalized/tests.rs @@ -31,12 +31,12 @@ use vortex_compressor::scheme::Scheme; use vortex_error::VortexResult; use vortex_mask::Mask; -use crate::encodings::l2_denorm::L2Denorm; -use crate::encodings::l2_denorm::L2DenormArraySlotsExt; -use crate::encodings::l2_denorm::L2DenormMetadata; -use crate::encodings::l2_denorm::L2DenormScheme; -use crate::encodings::l2_denorm::normalize_as_l2_denorm; -use crate::encodings::l2_denorm::validate_l2_normalized_rows_against_norms; +use crate::encodings::normalized::Normalized; +use crate::encodings::normalized::NormalizedArraySlotsExt; +use crate::encodings::normalized::NormalizedMetadata; +use crate::encodings::normalized::NormalizedScheme; +use crate::encodings::normalized::normalize; +use crate::encodings::normalized::validate_l2_normalized_rows_against_norms; use crate::tests::SESSION; use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; @@ -44,13 +44,13 @@ use crate::utils::test_helpers::constant_tensor_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; -/// Builds an [`L2Denorm`] array through the checked constructor and executes it, which is the +/// 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_l2_denorm(normalized: ArrayRef, norms: ArrayRef) -> VortexResult { +fn eval_normalized(normalized: ArrayRef, norms: ArrayRef) -> VortexResult { let mut ctx = SESSION.create_execution_ctx(); - let denorm = L2Denorm::try_new(normalized, norms, &mut ctx)?; + let normalized_array = Normalized::try_new(normalized, norms, &mut ctx)?; - denorm.into_array().execute(&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 @@ -104,7 +104,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_l2_denorm(normalized, norms)?; + let actual = eval_normalized(normalized, norms)?; let expected = vector_array(3, &[3.0, 4.0, 0.0, 0.0, 0.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) @@ -115,7 +115,7 @@ 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_l2_denorm(normalized, norms)?; + let actual = eval_normalized(normalized, norms)?; 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) @@ -129,7 +129,7 @@ fn decodes_null_rows_from_either_child() -> VortexResult<()> { let norms = PrimitiveArray::from_option_iter([Some(5.0f64), Some(2.0), None]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let actual: ExtensionArray = eval_l2_denorm(normalized, norms)?.execute(&mut ctx)?; + let actual: ExtensionArray = eval_normalized(normalized, norms)?.execute(&mut ctx)?; let storage: FixedSizeListArray = actual.storage_array().clone().execute(&mut ctx)?; let elements: PrimitiveArray = storage.elements().clone().execute(&mut ctx)?; @@ -149,10 +149,13 @@ fn validity_is_the_intersection_of_both_children() -> VortexResult<()> { let norms = PrimitiveArray::from_option_iter([Some(1.0f64), Some(1.0), None]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let denorm = L2Denorm::try_new(normalized, norms, &mut ctx)?; + let normalized_array = Normalized::try_new(normalized, norms, &mut ctx)?; - assert!(denorm.dtype().is_nullable()); - let mask = denorm.as_ref().validity()?.execute_mask(3, &mut ctx)?; + assert!(normalized_array.dtype().is_nullable()); + let mask = normalized_array + .as_ref() + .validity()? + .execute_mask(3, &mut ctx)?; assert!(mask.value(0)); assert!(!mask.value(1)); assert!(!mask.value(2)); @@ -171,7 +174,7 @@ 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_l2_denorm(normalized.clone(), norms)?; + let actual = eval_normalized(normalized.clone(), norms)?; assert_tensor_arrays_eq(actual, normalized) } @@ -183,7 +186,7 @@ fn constant_near_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 + 1e-12, 2); - let actual = eval_l2_denorm(normalized.clone(), norms)?; + let actual = eval_normalized(normalized.clone(), norms)?; assert_tensor_arrays_eq(actual, normalized) } @@ -193,7 +196,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_l2_denorm(normalized, norms)?; + let actual = eval_normalized(normalized, norms)?; let expected = vector_array(3, &[3.0, 4.0, 0.0, 5.0, 0.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) @@ -206,7 +209,7 @@ 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_l2_denorm(normalized, norms)?; + let actual = eval_normalized(normalized, norms)?; 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) @@ -221,9 +224,9 @@ fn nullable_constant_norms_widen_the_decoded_dtype() -> VortexResult<()> { ConstantArray::new(Scalar::primitive(1.0f64, Nullability::Nullable), 2).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let denorm = L2Denorm::try_new(normalized, norms, &mut ctx)?; - let dtype = denorm.dtype().clone(); - let decoded: ArrayRef = denorm.into_array().execute(&mut ctx)?; + let normalized_array = Normalized::try_new(normalized, norms, &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); @@ -266,7 +269,7 @@ fn rejects_structurally_invalid_children( ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - assert!(L2Denorm::try_new(normalized, norms, &mut ctx).is_err()); + assert!(Normalized::try_new(normalized, norms, &mut ctx).is_err()); Ok(()) } @@ -290,7 +293,7 @@ fn checked_construction_rejects_semantic_violations( ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - assert!(L2Denorm::try_new(normalized, norms, &mut ctx).is_err()); + assert!(Normalized::try_new(normalized, norms, &mut ctx).is_err()); Ok(()) } @@ -300,7 +303,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_l2_denorm(normalized, norms)?; + let actual = eval_normalized(normalized, norms)?; let expected = vector_array(2, &[0.0, 0.0, 3.0, 0.0])?; assert_tensor_arrays_eq(actual, expected) @@ -311,8 +314,12 @@ fn validate_accepts_normalized_f16_rows() -> VortexResult<()> { let input = vector_array(2, &[3.0f32, 4.0, 0.0, 0.0].map(half::f16::from_f32))?; let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input, &mut ctx)?; - validate_l2_normalized_rows_against_norms(&denorm.normalized().clone(), None, &mut ctx) + let normalized_array = normalize(input, &mut ctx)?; + validate_l2_normalized_rows_against_norms( + &normalized_array.normalized().clone(), + None, + &mut ctx, + ) } #[test] @@ -338,8 +345,8 @@ fn validate_rejects_unnormalized_rows() -> VortexResult<()> { #[case::constant_vector(Vector::constant_array(&[3.0, 4.0], 2).expect("valid vector array"))] fn normalize_round_trips(#[case] input: ArrayRef) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input.clone(), &mut ctx)?; - let actual = denorm.into_array().execute(&mut ctx)?; + let normalized_array = normalize(input.clone(), &mut ctx)?; + let actual = normalized_array.into_array().execute(&mut ctx)?; assert_tensor_arrays_eq(actual, input) } @@ -350,9 +357,9 @@ fn normalize_keeps_constant_input_children_constant() -> VortexResult<()> { // similarity and inner product short-circuit against a literal query vector. let input = Vector::constant_array(&[3.0, 4.0], 16)?; let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input, &mut ctx)?; + let normalized_array = normalize(input, &mut ctx)?; - let normalized = denorm + let normalized = normalized_array .normalized() .as_opt::() .expect("normalized child should be an Extension array"); @@ -361,7 +368,7 @@ fn normalize_keeps_constant_input_children_constant() -> VortexResult<()> { "normalized storage should stay constant after the fast path" ); - let norms = denorm + let norms = normalized_array .norms() .as_opt::() .expect("norms child should be a ConstantArray"); @@ -381,14 +388,14 @@ fn normalize_keeps_constant_input_children_constant() -> VortexResult<()> { fn normalize_zeroes_rows_with_zero_norms() -> VortexResult<()> { let input = vector_array(2, &[0.0, 0.0, 3.0, 4.0])?; let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input.clone(), &mut ctx)?; + let normalized_array = normalize(input.clone(), &mut ctx)?; - let normalized: ExtensionArray = denorm.normalized().clone().execute(&mut ctx)?; + 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], &[0.0, 0.0]); - let actual = denorm.into_array().execute(&mut ctx)?; + let actual = normalized_array.into_array().execute(&mut ctx)?; assert_tensor_arrays_eq(actual, input) } @@ -399,12 +406,15 @@ fn normalize_preserves_nulls_through_the_norms_child() -> VortexResult<()> { let input = MaskedArray::try_new(input, Validity::from_iter([true, false, true]))?.into_array(); let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input, &mut ctx)?; + let normalized_array = normalize(input, &mut ctx)?; - assert!(!denorm.normalized().dtype().is_nullable()); - assert!(denorm.norms().dtype().is_nullable()); + assert!(!normalized_array.normalized().dtype().is_nullable()); + assert!(normalized_array.norms().dtype().is_nullable()); - let mask = denorm.as_ref().validity()?.execute_mask(3, &mut ctx)?; + let mask = normalized_array + .as_ref() + .validity()? + .execute_mask(3, &mut ctx)?; assert!(mask.value(0)); assert!(!mask.value(1)); assert!(mask.value(2)); @@ -420,11 +430,13 @@ fn normalize_preserves_nulls_through_the_norms_child() -> VortexResult<()> { fn slice_stays_encoded_and_decodes_correctly() -> VortexResult<()> { let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0, 5.0, 12.0])?; let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + let normalized_array = normalize(input, &mut ctx)?.into_array(); - let sliced = denorm.slice(1..3)?.execute_until::(&mut ctx)?; + let sliced = normalized_array + .slice(1..3)? + .execute_until::(&mut ctx)?; assert!( - sliced.is::(), + sliced.is::(), "slicing must push down into both children instead of decoding the column" ); @@ -437,12 +449,14 @@ fn slice_stays_encoded_and_decodes_correctly() -> VortexResult<()> { fn filter_stays_encoded_and_decodes_correctly() -> VortexResult<()> { let input = vector_array(2, &[3.0, 4.0, 1.0, 0.0, 0.0, 2.0, 5.0, 12.0])?; let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + let normalized_array = normalize(input, &mut ctx)?.into_array(); let mask = Mask::from_iter([true, false, true, false]); - let filtered = denorm.filter(mask)?.execute_until::(&mut ctx)?; + let filtered = normalized_array + .filter(mask)? + .execute_until::(&mut ctx)?; assert!( - filtered.is::(), + filtered.is::(), "filtering must push down into both children instead of decoding the column" ); @@ -455,10 +469,12 @@ fn filter_stays_encoded_and_decodes_correctly() -> VortexResult<()> { 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])?; let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + let normalized_array = normalize(input, &mut ctx)?.into_array(); let indices = PrimitiveArray::from_iter([3u64, 0, 3]).into_array(); - let taken = denorm.take(indices)?.execute::(&mut ctx)?; + let taken = normalized_array + .take(indices)? + .execute::(&mut ctx)?; let expected = vector_array(2, &[5.0, 12.0, 3.0, 4.0, 5.0, 12.0])?; assert_tensor_arrays_eq(taken.into_array(), expected) @@ -468,11 +484,11 @@ fn take_decodes_correctly() -> VortexResult<()> { fn scalar_at_reads_a_single_denormalized_row() -> VortexResult<()> { let input = vector_array(2, &[3.0, 4.0, 5.0, 12.0])?; let mut ctx = SESSION.create_execution_ctx(); - let denorm = normalize_as_l2_denorm(input.clone(), &mut ctx)?.into_array(); + let normalized_array = normalize(input.clone(), &mut ctx)?.into_array(); for i in 0..input.len() { assert_eq!( - denorm.execute_scalar(i, &mut ctx)?, + normalized_array.execute_scalar(i, &mut ctx)?, input.execute_scalar(i, &mut ctx)?, ); } @@ -485,7 +501,7 @@ fn scalar_at_reads_a_single_denormalized_row() -> VortexResult<()> { // ============================================================================= /// Round-trips through the array plugin registry, which is the same path a Vortex file takes. -/// `normalize_as_l2_denorm` leaves the normalized child non-nullable and the norms child nullable +/// `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"))] @@ -495,14 +511,14 @@ fn scalar_at_reads_a_single_denormalized_row() -> VortexResult<()> { #[case::nullable_vector(nullable_vector_input().expect("valid vector array"))] fn serde_round_trip(#[case] input: ArrayRef) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let original = normalize_as_l2_denorm(input, &mut ctx)?.into_array(); + let original = normalize(input, &mut ctx)?.into_array(); let children: Vec = original.children(); let metadata = SESSION .array_serialize(&original)? - .expect("L2Denorm must serialize"); + .expect("Normalized must serialize"); let recovered = ArrayPlugin::deserialize( - &L2Denorm, + &Normalized, original.dtype(), original.len(), &metadata, @@ -511,7 +527,7 @@ fn serde_round_trip(#[case] input: ArrayRef) -> VortexResult<()> { &SESSION, )?; - assert_eq!(recovered.encoding_id(), ArrayVTable::id(&L2Denorm)); + assert_eq!(recovered.encoding_id(), ArrayVTable::id(&Normalized)); assert_eq!(recovered.dtype(), original.dtype()); assert_eq!(recovered.len(), original.len()); assert_tensor_arrays_eq(recovered, original) @@ -533,20 +549,20 @@ fn serialized_metadata_pins_child_nullabilities() -> VortexResult<()> { Validity::from_iter([true, false]), )? .into_array(); - let denorm = normalize_as_l2_denorm(input, &mut ctx)?; + let normalized_array = normalize(input, &mut ctx)?; let bytes = SESSION - .array_serialize(&denorm.clone().into_array())? - .expect("L2Denorm must serialize"); - let metadata = L2DenormMetadata::decode(bytes.as_slice())?; + .array_serialize(&normalized_array.clone().into_array())? + .expect("Normalized must serialize"); + let metadata = NormalizedMetadata::decode(bytes.as_slice())?; assert_eq!( metadata.normalized_is_nullable, - denorm.normalized().dtype().is_nullable(), + normalized_array.normalized().dtype().is_nullable(), ); assert_eq!( metadata.norms_is_nullable, - denorm.norms().dtype().is_nullable(), + normalized_array.norms().dtype().is_nullable(), ); Ok(()) @@ -562,14 +578,14 @@ fn serde_round_trip_preserves_normalized_nullability() -> VortexResult<()> { let norms = PrimitiveArray::from_iter([5.0f64, 1.0]).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let original = L2Denorm::try_new(normalized, norms, &mut ctx)?.into_array(); + let original = Normalized::try_new(normalized, norms, &mut ctx)?.into_array(); let children: Vec = original.children(); let metadata = SESSION .array_serialize(&original)? - .expect("L2Denorm must serialize"); + .expect("Normalized must serialize"); let recovered = ArrayPlugin::deserialize( - &L2Denorm, + &Normalized, original.dtype(), original.len(), &metadata, @@ -578,7 +594,7 @@ fn serde_round_trip_preserves_normalized_nullability() -> VortexResult<()> { &SESSION, )?; - let recovered = recovered.as_::(); + let recovered = recovered.as_::(); assert!(recovered.normalized().dtype().is_nullable()); assert!(!recovered.norms().dtype().is_nullable()); @@ -586,10 +602,10 @@ fn serde_round_trip_preserves_normalized_nullability() -> VortexResult<()> { } #[test] -fn encoding_is_registered_under_the_original_id() { - let id = ArrayVTable::id(&L2Denorm); +fn encoding_is_registered_under_the_normalized_id() { + let id = ArrayVTable::id(&Normalized); - assert_eq!(id.as_ref(), "vortex.tensor.l2_denorm"); + assert_eq!(id.as_ref(), "vortex.tensor.normalized"); assert!(SESSION.arrays().registry().contains_key(&id)); } @@ -620,10 +636,10 @@ fn scheme_matches_tensor_columns(#[case] input: ArrayRef) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let canonical: Canonical = input.execute(&mut ctx)?; - assert!(L2DenormScheme.matches(&canonical)); + assert!(NormalizedScheme.matches(&canonical)); assert_eq!( - L2DenormScheme.produced_encodings(), - vec![ArrayVTable::id(&L2Denorm)] + NormalizedScheme.produced_encodings(), + vec![ArrayVTable::id(&Normalized)] ); Ok(()) @@ -633,13 +649,13 @@ fn scheme_matches_tensor_columns(#[case] input: ArrayRef) -> VortexResult<()> { fn compressor_emits_the_dedicated_encoding() -> VortexResult<()> { let input = collinear_vectors(1024)?; let compressor = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&L2DenormScheme) + .with_new_scheme(&NormalizedScheme) .build(); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&input, &mut ctx)?; - assert_eq!(compressed.encoding_id(), ArrayVTable::id(&L2Denorm)); + assert_eq!(compressed.encoding_id(), ArrayVTable::id(&Normalized)); assert!(compressed.nbytes() < input.nbytes()); assert_tensor_arrays_eq(compressed, input) } diff --git a/vortex-tensor/src/encodings/l2_denorm/validate.rs b/vortex-tensor/src/encodings/normalized/validate.rs similarity index 77% rename from vortex-tensor/src/encodings/l2_denorm/validate.rs rename to vortex-tensor/src/encodings/normalized/validate.rs index c01563944f9..3c27231b639 100644 --- a/vortex-tensor/src/encodings/l2_denorm/validate.rs +++ b/vortex-tensor/src/encodings/normalized/validate.rs @@ -18,14 +18,14 @@ use crate::utils::extract_flat_elements; use crate::utils::unit_norm_tolerance; use crate::utils::validate_tensor_float_input; -/// Validates the structural invariants of an [`L2Denorm`] array's children. +/// Validates the structural invariants of a [`Normalized`] array's children. /// -/// These are the cheap, dtype-and-length checks that every [`L2DenormArray`] upholds, whichever +/// These are the cheap, dtype-and-length checks that every [`NormalizedArray`] upholds, whichever /// constructor built it. They run on construction and on deserialization. /// -/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm -/// [`L2DenormArray`]: crate::encodings::l2_denorm::L2DenormArray -pub(super) fn validate_l2_denorm_children( +/// [`Normalized`]: crate::encodings::normalized::Normalized +/// [`NormalizedArray`]: crate::encodings::normalized::NormalizedArray +pub(super) fn validate_normalized_children( normalized: &ArrayRef, norms: &ArrayRef, dtype: &DType, @@ -34,13 +34,13 @@ pub(super) fn validate_l2_denorm_children( vortex_ensure_eq!( normalized.len(), len, - "L2Denorm normalized child must have the array length ({len}), got {}", + "Normalized normalized child must have the array length ({len}), got {}", normalized.len(), ); vortex_ensure_eq!( norms.len(), len, - "L2Denorm norms child must have the array length ({len}), got {}", + "Normalized norms child must have the array length ({len}), got {}", norms.len(), ); @@ -49,14 +49,14 @@ pub(super) fn validate_l2_denorm_children( let DType::Primitive(norms_ptype, _) = norms.dtype() else { vortex_bail!( - "L2Denorm norms must be a primitive float array, got {}", + "Normalized norms must be a primitive float array, got {}", norms.dtype(), ); }; vortex_ensure_eq!( *norms_ptype, element_ptype, - "L2Denorm norms dtype must match the normalized element dtype ({element_ptype}), \ + "Normalized norms dtype must match the normalized element dtype ({element_ptype}), \ got {norms_ptype}", ); @@ -66,14 +66,14 @@ pub(super) fn validate_l2_denorm_children( vortex_ensure_eq!( *dtype, expected, - "L2Denorm dtype must be the union of its children's nullability ({expected}), got {dtype}", + "Normalized dtype must be the union of its children's nullability ({expected}), got {dtype}", ); Ok(()) } /// Validates that `normalized` and (when supplied) the matching `norms` jointly satisfy the -/// semantic [`L2Denorm`] invariants: +/// semantic [`Normalized`] invariants: /// /// - Every valid row of `normalized` has L2 norm `1.0` or `0.0`, within the tolerance implied by /// the element precision. @@ -83,7 +83,7 @@ pub(super) fn validate_l2_denorm_children( /// This costs `O(len * list_size)`, which is why it is a separate step rather than part of the /// encoding's structural validation. /// -/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm +/// [`Normalized`]: crate::encodings::normalized::Normalized pub fn validate_l2_normalized_rows_against_norms( normalized: &ArrayRef, norms: Option<&ArrayRef>, @@ -103,21 +103,21 @@ pub fn validate_l2_normalized_rows_against_norms( vortex_ensure_eq!( norms.len(), row_count, - "L2Denorm norms must have the same length as the normalized child ({row_count}), \ + "Normalized norms must have the same length as the normalized child ({row_count}), \ got {}", norms.len(), ); let DType::Primitive(norms_ptype, _) = norms.dtype() else { vortex_bail!( - "L2Denorm norms must be a primitive float array, got {}", + "Normalized norms must be a primitive float array, got {}", norms.dtype(), ); }; vortex_ensure_eq!( *norms_ptype, element_ptype, - "L2Denorm norms ptype must match the normalized element ptype ({element_ptype}), \ + "Normalized norms ptype must match the normalized element ptype ({element_ptype}), \ got {norms_ptype}", ); } @@ -157,7 +157,7 @@ pub fn validate_l2_normalized_rows_against_norms( vortex_ensure!( row_norm == 0.0 || (row_norm - 1.0).abs() <= tolerance, - "L2Denorm normalized child must have L2 norm 1.0 or 0.0, but row {i} has \ + "Normalized normalized child must have L2 norm 1.0 or 0.0, but row {i} has \ {row_norm:.6}", ); @@ -165,13 +165,13 @@ pub fn validate_l2_normalized_rows_against_norms( let stored_norm_f64 = ToPrimitive::to_f64(&stored_norms[i]).unwrap_or(f64::NAN); vortex_ensure!( stored_norm_f64 >= 0.0, - "L2Denorm norms must be non-negative, but row {i} has {stored_norm_f64:.6}", + "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, - "L2Denorm normalized child must be all zeros when norms row {i} is 0.0", + "Normalized normalized child must be all zeros when norms row {i} is 0.0", ); } } diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index 5c39724ef40..fe56827c15b 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -19,7 +19,7 @@ use vortex_array::session::ArraySessionExt; use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; -use crate::encodings::l2_denorm::L2Denorm; +use crate::encodings::normalized::Normalized; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; @@ -46,7 +46,7 @@ mod utils; /// containing serialized tensor scalar-fn arrays will fail to deserialize. Opt-in by setting the /// variable to any non-empty value. /// -/// This does **not** gate [`L2Denorm`]. That is a real array encoding rather than a persisted +/// This does **not** gate [`Normalized`]. That is a real array encoding rather than a persisted /// scalar function, and the compressor can emit it, so it always registers. pub const SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str = "VX_SCALAR_FN_ARRAY_TENSOR_PLUGIN"; @@ -59,7 +59,7 @@ pub fn initialize(session: &VortexSession) { arrow_session.register_exporter(Arc::new(Vector)); arrow_session.register_importer(Arc::new(Vector)); - session.arrays().register(L2Denorm); + session.arrays().register(Normalized); let session_fns = session.scalar_fns(); diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 447ea973de9..ca8fcf0efd4 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -30,12 +30,12 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::encodings::l2_denorm::DenormOrientation; -use crate::encodings::l2_denorm::try_build_constant_l2_denorm; +use crate::encodings::normalized::NormalizedOrientation; +use crate::encodings::normalized::try_build_constant_normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_norm::L2Norm; use crate::utils::BinaryTensorOpMetadata; -use crate::utils::extract_l2_denorm_children; +use crate::utils::extract_normalized_children; use crate::utils::validate_binary_tensor_float_inputs; /// Cosine similarity between two columns. @@ -47,14 +47,14 @@ use crate::utils::validate_binary_tensor_float_inputs; /// Both inputs must be tensor-like extension arrays ([`FixedShapeTensor`] or [`Vector`]) with the /// same dtype and a float element type. The output is a float column of the same float type. /// -/// When either input is [`L2Denorm`]-encoded, this operator treats the stored norms and +/// When either input is [`Normalized`]-encoded, this operator treats the stored norms and /// normalized children as authoritative. For lossy normalized children, that means the optimized /// read-through path may intentionally differ slightly from decoding both sides to dense /// coordinates and recomputing cosine from scratch. /// /// [`FixedShapeTensor`]: crate::fixed_shape_tensor::FixedShapeTensor /// [`Vector`]: crate::vector::Vector -/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm +/// [`Normalized`]: crate::encodings::normalized::Normalized #[derive(Clone)] pub struct CosineSimilarity; @@ -117,24 +117,27 @@ 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 `L2Denorm` whose children are both `ConstantArray`s. - // The `L2Denorm` fast path below then picks it up. - if let Some(denorm) = try_build_constant_l2_denorm(&lhs_ref, len, ctx)? { - lhs_ref = denorm.into_array(); + // stored row and re-encode it as an `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(); } - if let Some(denorm) = try_build_constant_l2_denorm(&rhs_ref, len, ctx)? { - rhs_ref = denorm.into_array(); + if let Some(normalized_array) = try_build_constant_normalized(&rhs_ref, len, ctx)? { + rhs_ref = normalized_array.into_array(); } - // Take any L2Denorm read-through fast path that applies. - match DenormOrientation::classify(&lhs_ref, &rhs_ref) { - DenormOrientation::Both { lhs, rhs } => { - return self.execute_both_denorm(lhs, rhs, len, ctx); + // Take any Normalized read-through fast path that applies. + match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + NormalizedOrientation::Both { lhs, rhs } => { + return self.execute_both_normalized(lhs, rhs, len, ctx); } - DenormOrientation::One { denorm, plain } => { - return self.execute_one_denorm(denorm, plain, len, ctx); + NormalizedOrientation::One { + normalized_array, + plain, + } => { + return self.execute_one_normalized(normalized_array, plain, len, ctx); } - DenormOrientation::Neither => {} + NormalizedOrientation::Neither => {} } // Compute combined validity. @@ -219,11 +222,11 @@ impl ScalarFnArrayVTable for CosineSimilarity { } impl CosineSimilarity { - /// Both sides are [`L2Denorm`]-encoded: treat the normalized children as authoritative, so + /// Both sides are [`Normalized`]-encoded: treat the normalized children as authoritative, so /// `cosine_similarity = dot(n_l, n_r)`. /// - /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm - fn execute_both_denorm( + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn execute_both_normalized( &self, lhs_ref: &ArrayRef, rhs_ref: &ArrayRef, @@ -232,10 +235,10 @@ impl CosineSimilarity { ) -> VortexResult { let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - let (normalized_l, norms_l) = extract_l2_denorm_children(lhs_ref); - let (normalized_r, norms_r) = extract_l2_denorm_children(rhs_ref); + let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); + let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); - // `L2Denorm` makes the normalized children authoritative, so their dot product is the + // `Normalized` makes the normalized children authoritative, so their dot product is the // cosine similarity even for lossy storage wrappers, except that a zero stored norm still // represents a zero vector. let dot: PrimitiveArray = InnerProduct::try_new_array(normalized_l, normalized_r)? @@ -263,27 +266,27 @@ impl CosineSimilarity { }) } - /// One side is [`L2Denorm`]-encoded: treat the normalized child as authoritative, so + /// One side is [`Normalized`]-encoded: treat the normalized child as authoritative, so /// `cosine_similarity = dot(n, b) / ||b||`. /// - /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm + /// [`Normalized`]: crate::encodings::normalized::Normalized /// - /// The caller must pass the denorm array as `denorm_ref` and the plain array as `plain_ref`. - fn execute_one_denorm( + /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. + fn execute_one_normalized( &self, - denorm_ref: &ArrayRef, + normalized_ref: &ArrayRef, plain_ref: &ArrayRef, len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let validity = denorm_ref.validity()?.and(plain_ref.validity()?)?; + let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - let (normalized, denorm_norms) = extract_l2_denorm_children(denorm_ref); + let (normalized, normalized_norms) = extract_normalized_children(normalized_ref); let dot_arr = InnerProduct::try_new_array(normalized, plain_ref.clone())?; let dot: PrimitiveArray = dot_arr.into_array().execute(ctx)?; - let denorm_norms: PrimitiveArray = denorm_norms.execute(ctx)?; + let normalized_norms: PrimitiveArray = normalized_norms.execute(ctx)?; let norm_arr = L2Norm::try_new_array(plain_ref.clone())?; let plain_norm: PrimitiveArray = norm_arr.into_array().execute(ctx)?; @@ -292,11 +295,11 @@ impl CosineSimilarity { // TODO(connor): This can be written in a more SIMD-friendly manner. match_each_float_ptype!(dot.ptype(), |T| { let dots = dot.as_slice::(); - let denorm_norms = denorm_norms.as_slice::(); + let normalized_norms = normalized_norms.as_slice::(); let plain_norms = plain_norm.as_slice::(); let buffer: Buffer = (0..len) .map(|i| { - if denorm_norms[i] == T::zero() || plain_norms[i] == T::zero() { + if normalized_norms[i] == T::zero() || plain_norms[i] == T::zero() { T::zero() } else { dots[i] / plain_norms[i] @@ -325,13 +328,13 @@ mod tests { use vortex_array::validity::Validity; use vortex_error::VortexResult; - use crate::encodings::l2_denorm::L2Denorm; + use crate::encodings::normalized::Normalized; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::tests::SESSION; use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::constant_tensor_array; - use crate::utils::test_helpers::l2_denorm_array; + use crate::utils::test_helpers::normalized_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; @@ -515,12 +518,12 @@ mod tests { } #[test] - fn both_denorm_self_similarity() -> VortexResult<()> { + fn both_normalized_self_similarity() -> VortexResult<()> { // [3.0, 4.0] has norm 5.0, normalized [0.6, 0.8]. // [1.0, 0.0] has norm 1.0, normalized [1.0, 0.0]. let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = l2_denorm_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; // Self-similarity should always be 1.0. assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 1.0]); @@ -528,23 +531,23 @@ mod tests { } #[test] - fn both_denorm_orthogonal() -> VortexResult<()> { + fn both_normalized_orthogonal() -> VortexResult<()> { // [3.0, 0.0] normalized [1.0, 0.0], norm 3.0. // [0.0, 4.0] normalized [0.0, 1.0], norm 4.0. let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; - let rhs = l2_denorm_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; + let lhs = normalized_array(&[2], &[1.0, 0.0], &[3.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.0, 1.0], &[4.0], &mut ctx)?; assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.0]); Ok(()) } #[test] - fn both_denorm_zero_norm() -> VortexResult<()> { + fn both_normalized_zero_norm() -> VortexResult<()> { // Zero-norm row: normalized is [0.0, 0.0], norm is 0.0. let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; - let rhs = l2_denorm_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let lhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 0.0], &[5.0, 0.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; // Row 0: dot([0.6, 0.8], [0.6, 0.8]) = 1.0, row 1: dot([0,0], [1,0]) = 0.0. assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0, 0.0]); @@ -552,12 +555,12 @@ mod tests { } #[test] - fn one_side_denorm_lhs() -> VortexResult<()> { - // LHS is L2Denorm([0.6, 0.8], 5.0) representing [3.0, 4.0]. + fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. // RHS is plain [3.0, 4.0]. // cosine_similarity([3.0, 4.0], [3.0, 4.0]) = 1.0. let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; let rhs = tensor_array(&[2], &[3.0, 4.0])?; assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0]); @@ -565,26 +568,26 @@ mod tests { } #[test] - fn one_side_denorm_rhs() -> VortexResult<()> { - // LHS is plain [1.0, 0.0], RHS is L2Denorm([0.6, 0.8], 5.0) representing [3.0, 4.0]. + fn one_side_normalized_rhs() -> VortexResult<()> { + // LHS is plain [1.0, 0.0], RHS is Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. // cosine_similarity([1.0, 0.0], [3.0, 4.0]) = 3.0 / (1.0 * 5.0) = 0.6. let mut ctx = SESSION.create_execution_ctx(); let lhs = tensor_array(&[2], &[1.0, 0.0])?; - let rhs = l2_denorm_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; assert_close(&eval_cosine_similarity(lhs, rhs)?, &[0.6]); Ok(()) } #[test] - fn both_denorm_null_norms() -> VortexResult<()> { + fn both_normalized_null_norms() -> VortexResult<()> { // Row 0: valid, row 1: null (via nullable norms on rhs). let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut 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 = L2Denorm::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); + let rhs = Normalized::try_new(normalized_r, norms_r, &mut ctx)?.into_array(); let scalar_fn = CosineSimilarity::new().erased(); let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; @@ -597,7 +600,7 @@ mod tests { } #[test] - fn both_denorm_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + fn both_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { // Mimics a lossy encoding where the stored norm is authoritative but // the decoded normalized child is physically nonzero. With a stored norm of `0.0`, cosine // similarity for that row must be `0.0` even though the dot product of the normalized @@ -607,13 +610,13 @@ 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 { L2Denorm::new_unchecked(normalized_l, norms_l) }.into_array(); + let lhs = unsafe { Normalized::new_unchecked(normalized_l, norms_l) }.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 { L2Denorm::new_unchecked(normalized_r, norms_r) }.into_array(); + let rhs = unsafe { Normalized::new_unchecked(normalized_r, norms_r) }.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`. @@ -622,29 +625,29 @@ mod tests { } #[test] - fn one_side_denorm_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { + fn one_side_normalized_lossy_zero_stored_norm_returns_zero() -> VortexResult<()> { // Mimics a lossy encoding where the stored norm is authoritative but // the decoded normalized child is physically nonzero. The plain side is a normal nonzero // tensor with positive norm. cosine similarity must still be `0.0` because the - // authoritative stored norm on the denorm side is `0.0`. + // authoritative stored norm on the normalized_array side is `0.0`. let normalized = tensor_array(&[2], &[0.6, 0.8])?; let norms = PrimitiveArray::from_iter([0.0f64]).into_array(); // 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 denorm = unsafe { L2Denorm::new_unchecked(normalized, norms) }.into_array(); + let normalized_array = unsafe { Normalized::new_unchecked(normalized, norms) }.into_array(); let plain = tensor_array(&[2], &[1.0, 0.0])?; - // Denorm on the lhs: `One { denorm: lhs, plain: rhs }`. + // Normalized encoding on the lhs: `One { normalized_array: lhs, plain: rhs }`. assert_close( - &eval_cosine_similarity(denorm.clone(), plain.clone())?, + &eval_cosine_similarity(normalized_array.clone(), plain.clone())?, &[0.0], ); - // Denorm on the rhs: `One { denorm: rhs, plain: lhs }`. The same zero-norm guard must - // fire regardless of operand order. - assert_close(&eval_cosine_similarity(plain, denorm)?, &[0.0]); + // Normalized encoding on the rhs: `One { normalized_array: rhs, plain: lhs }`. The same + // zero-norm guard must fire regardless of operand order. + assert_close(&eval_cosine_similarity(plain, normalized_array)?, &[0.0]); Ok(()) } @@ -705,7 +708,7 @@ mod tests { #[test] fn constant_zero_norm_query() -> VortexResult<()> { // A zero-norm constant query must produce `0.0` for every row via the zero-norm guard in - // `execute_one_denorm` and `execute_both_denorm`. + // `execute_one_normalized` and `execute_both_normalized`. let lhs = constant_tensor_array(&[3], &[0.0, 0.0, 0.0], 3)?; let rhs = tensor_array( &[3], @@ -723,7 +726,7 @@ mod tests { fn constant_self_similarity_nonunit() -> VortexResult<()> { // A non-unit constant query compared to itself must produce `1.0`. This exercises the // helper's division: after normalization, both sides must be exactly unit so the - // L2Denorm fast path's inner product yields 1. + // Normalized fast path's inner product yields 1. let lhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; let rhs = constant_tensor_array(&[3], &[3.0, 4.0, 0.0], 5)?; assert_close(&eval_cosine_similarity(lhs, rhs)?, &[1.0; 5]); diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 33b87e971ac..53ae82eb4a2 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -34,11 +34,11 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::encodings::l2_denorm::DenormOrientation; +use crate::encodings::normalized::NormalizedOrientation; use crate::matcher::AnyTensor; use crate::utils::BinaryTensorOpMetadata; use crate::utils::extract_flat_elements; -use crate::utils::extract_l2_denorm_children; +use crate::utils::extract_normalized_children; use crate::utils::validate_binary_tensor_float_inputs; /// Inner product (dot product) between two columns. @@ -114,15 +114,18 @@ impl ScalarFnVTable for InnerProduct { let rhs_ref = args.get(1)?; let len = args.row_count(); - // Take any L2Denorm read-through fast path that applies. - match DenormOrientation::classify(&lhs_ref, &rhs_ref) { - DenormOrientation::Both { lhs, rhs } => { - return self.execute_both_denorm(lhs, rhs, len, ctx); + // Take any Normalized read-through fast path that applies. + match NormalizedOrientation::classify(&lhs_ref, &rhs_ref) { + NormalizedOrientation::Both { lhs, rhs } => { + return self.execute_both_normalized(lhs, rhs, len, ctx); } - DenormOrientation::One { denorm, plain } => { - return self.execute_one_denorm(denorm, plain, len, ctx); + NormalizedOrientation::One { + normalized_array, + plain, + } => { + return self.execute_one_normalized(normalized_array, plain, len, ctx); } - DenormOrientation::Neither => {} + NormalizedOrientation::Neither => {} } // Compute combined validity. @@ -203,10 +206,10 @@ impl ScalarFnArrayVTable for InnerProduct { } impl InnerProduct { - /// Both sides are [`L2Denorm`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. + /// Both sides are [`Normalized`]-encoded: `inner_product = s_l * s_r * dot(n_l, n_r)`. /// - /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm - fn execute_both_denorm( + /// [`Normalized`]: crate::encodings::normalized::Normalized + fn execute_both_normalized( &self, lhs_ref: &ArrayRef, rhs_ref: &ArrayRef, @@ -215,8 +218,8 @@ impl InnerProduct { ) -> VortexResult { let validity = lhs_ref.validity()?.and(rhs_ref.validity()?)?; - let (normalized_l, norms_l) = extract_l2_denorm_children(lhs_ref); - let (normalized_r, norms_r) = extract_l2_denorm_children(rhs_ref); + let (normalized_l, norms_l) = extract_normalized_children(lhs_ref); + let (normalized_r, norms_r) = extract_normalized_children(rhs_ref); let norms_l: PrimitiveArray = norms_l.execute(ctx)?; let norms_r: PrimitiveArray = norms_r.execute(ctx)?; @@ -236,22 +239,22 @@ impl InnerProduct { }) } - /// One side is [`L2Denorm`]-encoded: `inner_product = s * dot(n, other)`. + /// One side is [`Normalized`]-encoded: `inner_product = s * dot(n, other)`. /// - /// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm + /// [`Normalized`]: crate::encodings::normalized::Normalized /// - /// The caller must pass the denorm array as `denorm_ref` and the plain array as `plain_ref`. - fn execute_one_denorm( + /// The caller must pass the [`Normalized`] array as `normalized_ref` and the plain array as `plain_ref`. + fn execute_one_normalized( &self, - denorm_ref: &ArrayRef, + normalized_ref: &ArrayRef, plain_ref: &ArrayRef, len: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let validity = denorm_ref.validity()?.and(plain_ref.validity()?)?; + let validity = normalized_ref.validity()?.and(plain_ref.validity()?)?; - let (normalized, norms) = extract_l2_denorm_children(denorm_ref); - let denorm_norms: PrimitiveArray = norms.execute(ctx)?; + let (normalized, norms) = extract_normalized_children(normalized_ref); + let normalized_norms: PrimitiveArray = norms.execute(ctx)?; let dot: PrimitiveArray = InnerProduct::try_new_array(normalized, plain_ref.clone())? .into_array() @@ -259,7 +262,7 @@ impl InnerProduct { match_each_float_ptype!(dot.ptype(), |T| { let dots = dot.as_slice::(); - let ns = denorm_norms.as_slice::(); + let ns = normalized_norms.as_slice::(); let buffer: Buffer = (0..len).map(|i| ns[i] * dots[i]).collect(); // SAFETY: The buffer length equals `len`, which matches the source validity length. @@ -293,11 +296,11 @@ mod tests { use vortex_array::validity::Validity; use vortex_error::VortexResult; - use crate::encodings::l2_denorm::L2Denorm; + use crate::encodings::normalized::Normalized; use crate::scalar_fns::inner_product::InnerProduct; use crate::tests::SESSION; use crate::utils::test_helpers::assert_close; - use crate::utils::test_helpers::l2_denorm_array; + use crate::utils::test_helpers::normalized_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; @@ -413,13 +416,13 @@ mod tests { } #[test] - fn both_denorm() -> VortexResult<()> { - // LHS: [3.0, 4.0] = L2Denorm([0.6, 0.8], 5.0). - // RHS: [1.0, 0.0] = L2Denorm([1.0, 0.0], 1.0). + fn both_normalized() -> VortexResult<()> { + // LHS: [3.0, 4.0] = Normalized([0.6, 0.8], 5.0). + // RHS: [1.0, 0.0] = Normalized([1.0, 0.0], 1.0). // dot([3.0, 4.0], [1.0, 0.0]) = 3.0. let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; - let rhs = l2_denorm_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[1.0, 0.0], &[1.0], &mut ctx)?; // Expected: 5.0 * 1.0 * dot([0.6, 0.8], [1.0, 0.0]) = 5.0 * 0.6 = 3.0. assert_close(&eval_inner_product(lhs, rhs)?, &[3.0]); @@ -427,24 +430,24 @@ mod tests { } #[test] - fn both_denorm_multiple_rows() -> VortexResult<()> { + fn both_normalized_multiple_rows() -> VortexResult<()> { // Row 0: [3.0, 4.0] dot [3.0, 4.0] = 25.0. // Row 1: [1.0, 0.0] dot [0.0, 1.0] = 0.0. let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; - let rhs = l2_denorm_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; + let lhs = normalized_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8, 0.0, 1.0], &[5.0, 1.0], &mut ctx)?; assert_close(&eval_inner_product(lhs, rhs)?, &[25.0, 0.0]); Ok(()) } #[test] - fn one_side_denorm_lhs() -> VortexResult<()> { - // LHS: L2Denorm([0.6, 0.8], 5.0) representing [3.0, 4.0]. + fn one_side_normalized_lhs() -> VortexResult<()> { + // LHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. // RHS: plain [1.0, 2.0]. // dot([3.0, 4.0], [1.0, 2.0]) = 3.0 + 8.0 = 11.0. let mut ctx = SESSION.create_execution_ctx(); - let lhs = l2_denorm_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let lhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; let rhs = tensor_array(&[2], &[1.0, 2.0])?; assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); @@ -452,27 +455,27 @@ mod tests { } #[test] - fn one_side_denorm_rhs() -> VortexResult<()> { + fn one_side_normalized_rhs() -> VortexResult<()> { // LHS: plain [1.0, 2.0]. - // RHS: L2Denorm([0.6, 0.8], 5.0) representing [3.0, 4.0]. + // RHS: Normalized([0.6, 0.8], 5.0) representing [3.0, 4.0]. // dot([1.0, 2.0], [3.0, 4.0]) = 3.0 + 8.0 = 11.0. let mut ctx = SESSION.create_execution_ctx(); let lhs = tensor_array(&[2], &[1.0, 2.0])?; - let rhs = l2_denorm_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; + let rhs = normalized_array(&[2], &[0.6, 0.8], &[5.0], &mut ctx)?; assert_close(&eval_inner_product(lhs, rhs)?, &[11.0]); Ok(()) } #[test] - fn both_denorm_null_norms() -> VortexResult<()> { + fn both_normalized_null_norms() -> VortexResult<()> { // Row 0: valid, row 1: null (via nullable norms 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 mut ctx = SESSION.create_execution_ctx(); - let lhs = L2Denorm::try_new(normalized_l, norms_l, &mut ctx)?.into_array(); - let rhs = l2_denorm_array(&[2], &[0.6, 0.8, 1.0, 0.0], &[5.0, 1.0], &mut ctx)?; + let lhs = Normalized::try_new(normalized_l, norms_l, &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(); let result = ScalarFnArray::try_new(scalar_fn, vec![lhs, rhs])?; diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 0533fa0b055..b7e9060ed3f 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -43,10 +43,10 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::encodings::l2_denorm::L2Denorm; +use crate::encodings::normalized::Normalized; use crate::matcher::AnyTensor; use crate::utils::extract_flat_elements; -use crate::utils::extract_l2_denorm_children; +use crate::utils::extract_normalized_children; use crate::utils::validate_tensor_float_input; /// L2 norm (Euclidean norm) of a tensor or vector column. @@ -56,12 +56,12 @@ use crate::utils::validate_tensor_float_input; /// The input must be a tensor-like extension array with a float element type. The output is a float /// column of the same float type. /// -/// When the input is [`L2Denorm`]-encoded, this operator treats the stored norms as +/// When the input is [`Normalized`]-encoded, this operator treats the stored norms as /// authoritative. For lossy normalized children, that means `L2Norm` intentionally reads the /// stored norms instead of re-deriving them from fully decoded coordinates. That behavior is part /// of the storage contract, not a separate lossy-compute mode. /// -/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm +/// [`Normalized`]: crate::encodings::normalized::Normalized #[derive(Clone)] pub struct L2Norm; @@ -128,11 +128,11 @@ impl ScalarFnVTable for L2Norm { let norm_dtype = DType::Primitive(element_ptype, ext.nullability()); - // L2Norm over an L2Denorm-encoded column is defined to read back the authoritative stored + // 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::() { - let (_, norms) = extract_l2_denorm_children(&input_ref); + if input_ref.is::() { + let (_, norms) = extract_normalized_children(&input_ref); vortex_ensure_eq!(norms.dtype(), &norm_dtype); return Ok(norms); } diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index d017820aa4f..488694bd47f 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -26,8 +26,8 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::VortexSession; -use crate::encodings::l2_denorm::L2Denorm; -use crate::encodings::l2_denorm::L2DenormArraySlotsExt; +use crate::encodings::normalized::Normalized; +use crate::encodings::normalized::NormalizedArraySlotsExt; use crate::matcher::AnyTensor; use crate::matcher::TensorMatch; @@ -58,21 +58,24 @@ pub fn unit_norm_tolerance(element_ptype: PType, dimensions: usize) -> f64 { SAFETY_FACTOR as f64 * machine_epsilon * dimensions_root } -/// Extracts the `(normalized, norms)` children of an [`L2Denorm`]-encoded array. +/// Extracts the `(normalized, norms)` children of a [`Normalized`]-encoded array. /// /// # Panics /// -/// Panics if `array` is not [`L2Denorm`]-encoded. Callers reach this through -/// [`DenormOrientation::classify`], which has already matched on the encoding. +/// Panics if `array` is not [`Normalized`]-encoded. Callers reach this through +/// [`NormalizedOrientation::classify`], which has already matched on the encoding. /// -/// [`L2Denorm`]: crate::encodings::l2_denorm::L2Denorm -/// [`DenormOrientation::classify`]: crate::encodings::l2_denorm::DenormOrientation::classify -pub fn extract_l2_denorm_children(array: &ArrayRef) -> (ArrayRef, ArrayRef) { - let denorm = array - .as_opt::() - .vortex_expect("expected an L2Denorm-encoded array"); - - (denorm.normalized().clone(), denorm.norms().clone()) +/// [`Normalized`]: crate::encodings::normalized::Normalized +/// [`NormalizedOrientation::classify`]: crate::encodings::normalized::NormalizedOrientation::classify +pub fn extract_normalized_children(array: &ArrayRef) -> (ArrayRef, ArrayRef) { + let normalized_array = array + .as_opt::() + .vortex_expect("expected a Normalized-encoded array"); + + ( + normalized_array.normalized().clone(), + normalized_array.norms().clone(), + ) } /// Validates that `input_dtype` is a float-valued tensor-like extension dtype. @@ -299,7 +302,7 @@ pub mod test_helpers { use vortex_buffer::Buffer; use vortex_error::VortexResult; - use crate::encodings::l2_denorm::L2Denorm; + use crate::encodings::normalized::Normalized; use crate::types::fixed_shape_tensor::FixedShapeTensor; use crate::types::fixed_shape_tensor::FixedShapeTensorMetadata; use crate::types::vector::Vector; @@ -367,10 +370,10 @@ pub mod test_helpers { ConstantArray::new(ext_scalar, len).into_array() } - /// Creates an [`L2Denorm`] array from pre-normalized tensor elements and matching norms. The + /// 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. - pub fn l2_denorm_array( + pub fn normalized_array( shape: &[usize], normalized_elements: &[T], norms: &[T], @@ -379,7 +382,7 @@ pub mod test_helpers { let normalized = tensor_array(shape, normalized_elements)?; let norms = PrimitiveArray::new(Buffer::copy_from(norms), Validity::NonNullable).into_array(); - Ok(L2Denorm::try_new(normalized, norms, ctx)?.into_array()) + Ok(Normalized::try_new(normalized, norms, 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 fcc9516dea8..90955f01d04 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.l2_denorm", + &"vortex.tensor.normalized", &"vortex.tensor.l2_norm", ], };