diff --git a/Cargo.lock b/Cargo.lock index c764cd8eab7..0a55082948f 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`