Skip to content

Store Normalized validity on the array instead of in its children #9202

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
connortsui20 wants to merge 1 commit into
base: develop
Choose a base branch
from
Draft

Store Normalized validity on the array instead of in its children #9202

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
211 changes: 133 additions & 78 deletions vortex-tensor/src/encodings/normalized/array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use prost::Message;
use vortex_array::Array;
use vortex_array::ArrayId;
use vortex_array::ArrayParts;
Expand All @@ -15,26 +14,30 @@ use vortex_array::array_slots;
use vortex_array::arrays::ConstantArray;
use vortex_array::buffer::BufferHandle;
use vortex_array::dtype::DType;
use vortex_array::dtype::Nullability;
use vortex_array::scalar::Scalar;
use vortex_array::serde::ArrayChildren;
use vortex_array::validity::Validity;
use vortex_array::vtable::OperationsVTable;
use vortex_array::vtable::VTable;
use vortex_array::vtable::ValidityVTable;
use vortex_array::vtable::child_to_validity;
use vortex_array::vtable::validity_to_child;
use vortex_array::vtable::with_empty_buffers;
use vortex_error::VortexResult;
use vortex_error::vortex_err;
use vortex_error::vortex_ensure;
use vortex_error::vortex_ensure_eq;
use vortex_error::vortex_panic;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::encodings::normalized::execute::denormalize;
use crate::encodings::normalized::rules::RULES;
use crate::encodings::normalized::validate::validate_l2_normalized_rows_against_norms;
use crate::encodings::normalized::validate::validate_normalized_children;
use crate::encodings::normalized::validate::validate_normalized_rows;
use crate::utils::validate_tensor_float_input;

/// An [`Normalized`]-encoded Vortex array.
/// A [`Normalized`]-encoded Vortex array.
pub type NormalizedArray = Array<Normalized>;

/// The norm-split encoding for tensor-like columns.
Expand All @@ -47,18 +50,29 @@ pub type NormalizedArray = Array<Normalized>;
///
/// Every [`NormalizedArray`] structurally guarantees, via [`VTable::validate`]:
///
/// - `normalized` is a tensor-like extension array with a float element type.
/// - `norms` is a primitive column whose ptype equals the tensor element ptype.
/// - `normalized` is a non-nullable tensor-like extension array with a float element type, whose
/// dtype is the array's own dtype with nullability stripped.
/// - `norms` is a non-nullable primitive column whose ptype equals the tensor element ptype.
/// - both children have the array's length.
/// - the array dtype is `normalized.dtype().union_nullability(norms.nullability())`.
/// - the `validity` slot is present only when the array's dtype is nullable, in which case it is a
/// non-nullable boolean column of the array's length.
///
/// Nulls therefore live on the array itself rather than in either child, which is what keeps the
/// two children free to be reshaped independently: neither the decode path nor the read-through
/// operators ever have to widen a child's dtype to match the parent's.
///
/// On top of that, [`try_new`](Self::try_new) enforces the semantic invariants that make the split
/// lossless:
///
/// - every valid row of `normalized` has L2 norm `1.0` or `0.0`, within the tolerance implied by
/// the element precision.
/// - every row of `normalized` has L2 norm `1.0` or `0.0`, within the tolerance implied by the
/// element precision.
/// - every stored norm is non-negative.
/// - a stored norm of `0.0` is paired with an all-zero normalized row.
/// - a stored norm of `0.0` is paired with an all-zero normalized row, and an all-zero normalized
/// row is paired with a stored norm of `0.0`.
///
/// Those checks run over every row, including rows the `validity` marks null. [`normalize`] zeroes
/// both children at null positions, which satisfies them, so callers building a nullable column
/// should go through `normalize` rather than pairing raw children with a mask.
///
/// # Lossy normalized children
///
Expand All @@ -71,26 +85,38 @@ pub type NormalizedArray = Array<Normalized>;
/// storage contract, not a separate lossy-compute mode.
///
/// [`AnyTensor`]: crate::matcher::AnyTensor
/// [`normalize`]: crate::encodings::normalized::normalize
/// [`L2Norm`]: crate::scalar_fns::l2_norm::L2Norm
/// [`InnerProduct`]: crate::scalar_fns::inner_product::InnerProduct
/// [`CosineSimilarity`]: crate::scalar_fns::cosine_similarity::CosineSimilarity
#[derive(Clone, Debug)]
pub struct Normalized;

/// The two child arrays of an [`NormalizedArray`].
/// The slots of a [`NormalizedArray`]: its two children plus its validity.
#[array_slots(Normalized)]
pub struct NormalizedSlots {
/// The unit-norm (or zero) direction of each row, as a tensor-like extension array.
/// The unit-norm (or zero) direction of each row, as a non-nullable tensor-like extension
/// array.
#[slot(0)]
pub normalized: ArrayRef,

/// The authoritative L2 norm of each row, as a primitive float column.
/// The authoritative L2 norm of each row, as a non-nullable primitive float column.
#[slot(1)]
pub norms: ArrayRef,

/// The validity / null map of the array.
///
/// Both children are non-nullable, so this is the column's only record of which rows are null.
#[slot(2)]
pub validity: Option<ArrayRef>,
}

/// The number of required slots: `normalized` and `norms`. The `validity` slot is optional, so a
/// serialized array has either this many children or one more.
pub(super) const DATA_CHILDREN: usize = NormalizedSlots::COUNT - 1;

impl Normalized {
/// Builds an [`NormalizedArray`], validating that `normalized` really is row-wise L2-normalized
/// Builds a [`NormalizedArray`], validating that `normalized` really is row-wise L2-normalized
/// against `norms`.
///
/// This is the constructor for exact norm splits. It scans both children, so it costs
Expand All @@ -103,68 +129,72 @@ impl Normalized {
pub fn try_new(
normalized: ArrayRef,
norms: ArrayRef,
validity: Validity,
ctx: &mut ExecutionCtx,
) -> VortexResult<NormalizedArray> {
let len = normalized.len();
let dtype = normalized
.dtype()
.union_nullability(norms.dtype().nullability());
let slots = NormalizedSlots { normalized, norms }.into_slots();

// Structural validation has to come first: the row scan walks both children in lockstep
// and assumes they are a matching-length tensor/float pair.
let normalized_array = Array::try_from_parts(
ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots),
)?;
validate_l2_normalized_rows_against_norms(
normalized_array.normalized(),
Some(normalized_array.norms()),
ctx,
)?;
let array = Array::try_from_parts(normalized_parts(normalized, norms, validity))?;
validate_normalized_rows(array.normalized(), Some(array.norms()), ctx)?;

Ok(normalized_array)
Ok(array)
}

/// Builds an [`NormalizedArray`] without validation.
/// Builds a [`NormalizedArray`] without validation.
///
/// # Safety
///
/// The caller must uphold the structural invariants listed on [`Normalized`]. In particular,
/// both children must have the same length, `normalized` must be a float tensor, and `norms`
/// must be a primitive column with the same element ptype.
/// both children must be non-nullable and have the same length, `normalized` must be a float
/// tensor, and `norms` must be a primitive column with the same element ptype.
///
/// This does not check the unit-norm relationship. Violating it can produce wrong answers but
/// not memory unsafety.
pub unsafe fn new_unchecked(normalized: ArrayRef, norms: ArrayRef) -> NormalizedArray {
let len = normalized.len();
let dtype = normalized
.dtype()
.union_nullability(norms.dtype().nullability());
let slots = NormalizedSlots { normalized, norms }.into_slots();

unsafe {
Array::from_parts_unchecked(
ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots),
)
}
pub unsafe fn new_unchecked(
normalized: ArrayRef,
norms: ArrayRef,
validity: Validity,
) -> NormalizedArray {
unsafe { Array::from_parts_unchecked(normalized_parts(normalized, norms, validity)) }
}
}

/// Metadata for a serialized [`NormalizedArray`]: its children's nullabilities.
/// Assembles the [`ArrayParts`] shared by both constructors and by deserialization.
///
/// The parent dtype supplies the tensor shape and element ptype. Its nullability is the union of
/// the children, so it cannot identify which child is nullable.
#[derive(Clone, prost::Message)]
pub struct NormalizedMetadata {
/// Whether the `normalized` child is nullable.
#[prost(bool, tag = "1")]
pub normalized_is_nullable: bool,

/// Whether the `norms` child is nullable.
#[prost(bool, tag = "2")]
pub norms_is_nullable: bool,
/// The array's dtype is the `normalized` child's dtype widened by whatever nullability `validity`
/// implies, which is the only place the parent's nullability comes from.
fn normalized_parts(
normalized: ArrayRef,
norms: ArrayRef,
validity: Validity,
) -> ArrayParts<Normalized> {
let len = normalized.len();
let dtype = normalized.dtype().union_nullability(validity.nullability());
let slots = NormalizedSlots {
normalized,
norms,
validity: validity_to_child(&validity, len),
}
.into_slots();

ArrayParts::new(Normalized, dtype, len, EmptyArrayData).with_slots(slots)
}

/// Accessors for [`NormalizedArray`] that are derived from its slots rather than being one.
pub trait NormalizedArrayExt: NormalizedArraySlotsExt {
/// The column's validity.
///
/// Both children are non-nullable, so this is the array's complete null information.
fn normalized_validity(&self) -> Validity {
child_to_validity(
self.as_ref().slots()[NormalizedSlots::VALIDITY].as_ref(),
self.as_ref().dtype().nullability(),
)
}
}

impl<T: NormalizedArraySlotsExt> NormalizedArrayExt for T {}

impl VTable for Normalized {
type TypedArrayData = EmptyArrayData;

Expand All @@ -185,7 +215,7 @@ impl VTable for Normalized {
) -> VortexResult<()> {
let slots = NormalizedSlotsView::from_slots(slots);

validate_normalized_children(slots.normalized, slots.norms, dtype, len)
validate_normalized_children(slots.normalized, slots.norms, slots.validity, dtype, len)
}

fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
Expand All @@ -208,17 +238,13 @@ impl VTable for Normalized {
with_empty_buffers(self, array, buffers)
}

/// The array carries no metadata: the parent dtype supplies the tensor shape, element ptype,
/// and nullability, and both children's dtypes follow from it.
fn serialize(
array: ArrayView<'_, Self>,
_array: ArrayView<'_, Self>,
_session: &VortexSession,
) -> VortexResult<Option<Vec<u8>>> {
Ok(Some(
NormalizedMetadata {
normalized_is_nullable: array.normalized().dtype().is_nullable(),
norms_is_nullable: array.norms().dtype().is_nullable(),
}
.encode_to_vec(),
))
Ok(Some(vec![]))
}

fn deserialize(
Expand All @@ -230,18 +256,35 @@ impl VTable for Normalized {
children: &dyn ArrayChildren,
_session: &VortexSession,
) -> VortexResult<ArrayParts<Self>> {
let metadata = NormalizedMetadata::decode(metadata)
.map_err(|e| vortex_err!("Failed to decode NormalizedMetadata: {e}"))?;
vortex_ensure!(
metadata.is_empty(),
"NormalizedArray expects empty metadata, got {} bytes",
metadata.len(),
);

let element_ptype = validate_tensor_float_input(dtype)?.element_ptype();
let normalized_dtype = dtype.with_nullability(metadata.normalized_is_nullable.into());
let norms_dtype = DType::Primitive(element_ptype, metadata.norms_is_nullable.into());
let normalized_dtype = dtype.as_nonnullable();
let norms_dtype = DType::Primitive(element_ptype, Nullability::NonNullable);

let normalized = children.get(0, &normalized_dtype, len)?;
let norms = children.get(1, &norms_dtype, len)?;
let slots = NormalizedSlots { normalized, norms }.into_slots();

Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots))
// An absent validity child means "no nulls"; the parent's nullability is what distinguishes
// `NonNullable` from `AllValid`.
let validity = if children.len() == NormalizedSlots::COUNT {
Validity::Array(children.get(NormalizedSlots::VALIDITY, &Validity::DTYPE, len)?)
} else {
vortex_ensure_eq!(
children.len(),
DATA_CHILDREN,
"NormalizedArray expects {DATA_CHILDREN} or {} children, got {}",
NormalizedSlots::COUNT,
children.len(),
);
Validity::from(dtype.nullability())
};

Ok(normalized_parts(normalized, norms, validity))
}

fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
Expand All @@ -250,10 +293,19 @@ impl VTable for Normalized {

fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
let dtype = array.dtype().clone();
let validity = array.normalized_validity();
let row_count = array.len();
let slots = array.slots_view();

denormalize(slots.normalized, slots.norms, array.len(), dtype, ctx)
.map(ExecutionResult::done)
denormalize(
slots.normalized,
slots.norms,
validity,
row_count,
dtype,
ctx,
)
.map(ExecutionResult::done)
}

fn reduce_parent(
Expand All @@ -267,10 +319,7 @@ impl VTable for Normalized {

impl ValidityVTable<Normalized> for Normalized {
fn validity(array: ArrayView<'_, Normalized>) -> VortexResult<Validity> {
array
.normalized()
.validity()?
.and(array.norms().validity()?)
Ok(array.normalized_validity())
}
}

Expand All @@ -284,12 +333,18 @@ impl OperationsVTable<Normalized> for Normalized {
// one-row constants, which also lets the constant-norms fast path do the multiply.
let normalized = array.normalized().execute_scalar(index, ctx)?;
let norms = array.norms().execute_scalar(index, ctx)?;
let dtype = array.dtype().clone();

// `Array::execute_scalar` resolves null rows before dispatching here, so this row is valid
// and only the parent's nullability has to be reproduced.
let validity = Validity::from(dtype.nullability());

let row = denormalize(
&ConstantArray::new(normalized, 1).into_array(),
&ConstantArray::new(norms, 1).into_array(),
validity,
1,
array.dtype().clone(),
dtype,
ctx,
)?;

Expand Down
Loading
Loading