|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! Normalized vector extension type: a refinement of [`Vector`](crate::vector::Vector) whose |
| 5 | +//! rows are guaranteed (or asserted, for lossy encodings) to have unit L2 norm. |
| 6 | +
|
| 7 | +use num_traits::ToPrimitive; |
| 8 | +use vortex_array::ArrayRef; |
| 9 | +use vortex_array::ExecutionCtx; |
| 10 | +use vortex_array::IntoArray; |
| 11 | +use vortex_array::arrays::ExtensionArray; |
| 12 | +use vortex_array::arrays::extension::ExtensionArrayExt; |
| 13 | +use vortex_array::dtype::PType; |
| 14 | +use vortex_array::extension::EmptyMetadata; |
| 15 | +use vortex_array::match_each_float_ptype; |
| 16 | +use vortex_error::VortexResult; |
| 17 | +use vortex_error::vortex_ensure; |
| 18 | + |
| 19 | +use crate::utils::extract_flat_elements; |
| 20 | +use crate::utils::validate_tensor_float_input; |
| 21 | + |
| 22 | +/// Refinement of [`Vector`](crate::vector::Vector) that asserts every valid row is L2-normalized |
| 23 | +/// (unit-norm) or the zero vector. |
| 24 | +/// |
| 25 | +/// The storage shape is identical to [`Vector`](crate::vector::Vector): a `FixedSizeList<float, |
| 26 | +/// dim, nullability>` with non-nullable float elements. Downstream operators such as |
| 27 | +/// [`L2Denorm`](crate::scalar_fns::l2_denorm::L2Denorm), |
| 28 | +/// [`L2Norm`](crate::scalar_fns::l2_norm::L2Norm), |
| 29 | +/// [`InnerProduct`](crate::scalar_fns::inner_product::InnerProduct), and |
| 30 | +/// [`CosineSimilarity`](crate::scalar_fns::cosine_similarity::CosineSimilarity) short-circuit |
| 31 | +/// arithmetic when they see this refinement. |
| 32 | +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] |
| 33 | +pub struct NormalizedVector; |
| 34 | + |
| 35 | +impl NormalizedVector { |
| 36 | + /// Wraps `storage` as a [`NormalizedVector`] extension array after checking that every valid |
| 37 | + /// row is unit-norm or the zero vector. |
| 38 | + /// |
| 39 | + /// # Errors |
| 40 | + /// |
| 41 | + /// Returns an error if the extension dtype rejects `storage`, if `storage` is not a tensor |
| 42 | + /// with float elements, or if any valid row's L2 norm is not `1.0` (or `0.0`) within the |
| 43 | + /// tolerance implied by the element precision. |
| 44 | + pub fn try_new(storage: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> { |
| 45 | + let ext = ExtensionArray::try_new_from_vtable(NormalizedVector, EmptyMetadata, storage)? |
| 46 | + .into_array(); |
| 47 | + validate_unit_norm_rows(&ext, ctx)?; |
| 48 | + Ok(ext) |
| 49 | + } |
| 50 | + |
| 51 | + /// Wraps `storage` as a [`NormalizedVector`] extension array **without** validating that |
| 52 | + /// rows are unit-norm. |
| 53 | + /// |
| 54 | + /// # Safety |
| 55 | + /// |
| 56 | + /// Every valid row must be unit-norm or the zero vector. Lossy approximations (e.g. |
| 57 | + /// TurboQuant) deliberately relax this, but still treat the claim as authoritative |
| 58 | + /// downstream. Violating this does not cause memory unsafety but will produce silently |
| 59 | + /// incorrect results. |
| 60 | + /// |
| 61 | + /// # Errors |
| 62 | + /// |
| 63 | + /// Returns an error if the extension dtype rejects `storage` (e.g. non-FSL storage, wrong |
| 64 | + /// element dtype, or nullable elements). |
| 65 | + pub unsafe fn new_unchecked(storage: ArrayRef) -> VortexResult<ArrayRef> { |
| 66 | + Ok( |
| 67 | + ExtensionArray::try_new_from_vtable(NormalizedVector, EmptyMetadata, storage)? |
| 68 | + .into_array(), |
| 69 | + ) |
| 70 | + } |
| 71 | +} |
| 72 | + |
| 73 | +/// Returns the acceptable unit-norm drift for the given element precision. |
| 74 | +pub(crate) fn unit_norm_tolerance(element_ptype: PType) -> f64 { |
| 75 | + match element_ptype { |
| 76 | + PType::F16 => 2e-3, |
| 77 | + PType::F32 => 2e-6, |
| 78 | + PType::F64 => 1e-10, |
| 79 | + _ => unreachable!("NormalizedVector requires float elements, got {element_ptype:?}"), |
| 80 | + } |
| 81 | +} |
| 82 | + |
| 83 | +/// Validates that every valid row of a [`NormalizedVector`] extension array has L2 norm `1.0` |
| 84 | +/// or `0.0` within the element-precision tolerance. |
| 85 | +fn validate_unit_norm_rows(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { |
| 86 | + let row_count = array.len(); |
| 87 | + if row_count == 0 { |
| 88 | + return Ok(()); |
| 89 | + } |
| 90 | + |
| 91 | + let tensor_match = validate_tensor_float_input(array.dtype())?; |
| 92 | + let element_ptype = tensor_match.element_ptype(); |
| 93 | + let tolerance = unit_norm_tolerance(element_ptype); |
| 94 | + let tensor_flat_size = tensor_match.list_size() as usize; |
| 95 | + |
| 96 | + let ext: ExtensionArray = array.clone().execute(ctx)?; |
| 97 | + let validity = ext.as_ref().validity()?; |
| 98 | + let flat = extract_flat_elements(ext.storage_array(), tensor_flat_size, ctx)?; |
| 99 | + |
| 100 | + match_each_float_ptype!(element_ptype, |T| { |
| 101 | + for i in 0..row_count { |
| 102 | + if !validity.is_valid(i)? { |
| 103 | + continue; |
| 104 | + } |
| 105 | + |
| 106 | + let row_norm_sq = flat.row::<T>(i).iter().fold(0.0f64, |sum_sq, x| { |
| 107 | + let value = ToPrimitive::to_f64(x).unwrap_or(f64::NAN); |
| 108 | + sum_sq + value * value |
| 109 | + }); |
| 110 | + let row_norm = row_norm_sq.sqrt(); |
| 111 | + |
| 112 | + vortex_ensure!( |
| 113 | + row_norm == 0.0 || (row_norm - 1.0).abs() <= tolerance, |
| 114 | + "NormalizedVector row {i} has L2 norm {row_norm:.6}, expected 1.0 or 0.0", |
| 115 | + ); |
| 116 | + } |
| 117 | + }); |
| 118 | + |
| 119 | + Ok(()) |
| 120 | +} |
| 121 | + |
| 122 | +mod matcher; |
| 123 | +mod vtable; |
| 124 | + |
| 125 | +pub use matcher::AnyNormalizedVector; |
0 commit comments