Skip to content

Execute tensor L2 norm with RowFn #9768

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’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
connortsui20 wants to merge 1 commit into from
Draft

Execute tensor L2 norm with RowFn #9768

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
214 changes: 96 additions & 118 deletions vortex-tensor/src/scalar_fns/l2_norm.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,45 +9,31 @@

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::ExtensionArray;
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::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::dtype::DType;
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_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::ScalarFnVTableExt;
use vortex_array::scalar_fn::unstable::row::InitializedElement;
use vortex_array::scalar_fn::unstable::row::RowFn;
use vortex_array::scalar_fn::unstable::row::RowVisitor;
use vortex_array::scalar_fn::unstable::row::UninitElementSink;
use vortex_array::serde::ArrayChildren;
use vortex_buffer::Buffer;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_err;
use vortex_session::VortexSession;
use vortex_session::registry::CachedId;

use crate::matcher::AnyTensor;
use crate::scalar_fns::arithmetic::l2_norm_row;
use crate::utils::extract_flat_elements;
use crate::utils::validate_tensor_float_input;
use crate::scalar_fns::row::TensorRow;
use crate::scalar_fns::row::tensor_element_ptype;

/// L2 norm (Euclidean norm) of a tensor or vector column.
///
Expand Down Expand Up @@ -75,118 +61,49 @@ impl L2Norm {
}
}

impl ScalarFnVTable for L2Norm {
impl RowFn for L2Norm {
type Options = EmptyOptions;

const ARG_NAMES: &'static [&'static str] = &["input"];
const INFALLIBLE: bool = true;

fn id(&self) -> ScalarFnId {
static ID: CachedId = CachedId::new("vortex.tensor.l2_norm");
*ID
}

fn arity(&self, _options: &Self::Options) -> Arity {
Arity::Exact(1)
}

fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
match child_idx {
0 => ChildName::from("input"),
_ => unreachable!("L2Norm must have exactly one child"),
}
}

fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
let input_dtype = &arg_dtypes[0];
let tensor_match = validate_tensor_float_input(input_dtype)?;
let ptype = tensor_match.element_ptype();

let nullability = Nullability::from(input_dtype.is_nullable());
Ok(DType::Primitive(ptype, nullability))
fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
Ok(Some(vec![]))
}

fn execute(
fn deserialize(
&self,
_options: &Self::Options,
args: &dyn ExecutionArgs,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
let input_ref = args.get(0)?;
let row_count = args.row_count();

let ext = input_ref.dtype().as_extension();
let tensor_match = ext
.metadata_opt::<AnyTensor>()
.vortex_expect("L2Norm::return_dtype validated the input tensor metadata");
let tensor_flat_size = tensor_match.list_size() as usize;
let element_ptype = tensor_match.element_ptype();

let norm_dtype = DType::Primitive(element_ptype, ext.nullability());

// Optimize for the constant array case.
if let Some(array) = input_ref.as_opt::<Constant>() {
let scalar = array.scalar().as_extension().to_storage_scalar();

let Some(elements) = scalar.as_list().elements() else {
return Ok(ConstantArray::new(Scalar::null(norm_dtype), row_count).into_array());
};

let norm_scalar = match_each_float_ptype!(element_ptype, |T| {
let values: Vec<T> = elements
.iter()
.map(|element| {
element
.as_primitive()
.as_::<T>()
.vortex_expect("L2Norm::return_dtype validated the float element type")
})
.collect();
let norm = l2_norm_row::<T>(&values);

Scalar::try_new(norm_dtype, Some(norm.into()))
})?;

let norms = ConstantArray::new(norm_scalar, row_count).into_array();
return Ok(norms);
}

let input: ExtensionArray = input_ref.execute(ctx)?;
let validity = input.as_ref().validity()?;

let storage = input.storage_array();
let flat = extract_flat_elements(storage, tensor_flat_size, ctx)?;

match_each_float_ptype!(flat.ptype(), |T| {
let buffer: Buffer<T> = (0..row_count)
.map(|row_index| l2_norm_row(flat.row::<T>(row_index)))
.collect();

// SAFETY: The buffer length equals `row_count`, which matches the source validity
// length.
Ok(unsafe { PrimitiveArray::new_unchecked(buffer, validity) }.into_array())
})
_metadata: &[u8],
_session: &VortexSession,
) -> VortexResult<Self::Options> {
Ok(EmptyOptions)
}

fn validity(
fn dispatch<V: RowVisitor>(
&self,
_options: &Self::Options,
expression: &Expression,
) -> VortexResult<Option<Expression>> {
union_child_validities(expression)
}

fn is_strict(&self, _options: &Self::Options) -> bool {
true
}

fn is_infallible(&self, _options: &Self::Options) -> bool {
true
args: &[DType],
visitor: V,
) -> VortexResult<V::VisitResult> {
match_each_float_ptype!(tensor_element_ptype(args)?, |T| {
visitor.visit_into::<(TensorRow<T>,), UninitElementSink<T>, _>((), |(row,), output| {
// SAFETY: `output` is the `UninitElementSink` row supplied for this callback.
unsafe { InitializedElement::write(output, l2_norm_row(row)) }
})
})
}
}

/// Metadata for a serialized [`L2Norm`] array: the single `input` child's [`DType`], which carries
/// the extension type (`FixedShapeTensor` vs `Vector`), dimension, and nullability that are not
/// recoverable from the parent's primitive-float output.
#[derive(Clone, prost::Message)]
pub(super) struct L2NormMetadata {
struct L2NormMetadata {
#[prost(message, optional, tag = "1")]
input_dtype: Option<pb::DType>,
}
Expand Down Expand Up @@ -240,6 +157,7 @@ mod tests {
use vortex_array::arrays::PrimitiveArray;
use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin;
use vortex_array::dtype::DType;
use vortex_array::dtype::NativePType;
use vortex_array::dtype::Nullability;
use vortex_array::dtype::PType;
use vortex_array::dtype::extension::ExtDType;
Expand All @@ -254,13 +172,29 @@ mod tests {
use crate::utils::test_helpers::literal_vector_array;
use crate::utils::test_helpers::tensor_array;
use crate::utils::test_helpers::vector_array;
use crate::utils::test_helpers::zero_width_vector_array;

/// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec<f64>`.
fn eval_l2_norm(input: ArrayRef) -> VortexResult<Vec<f64>> {
fn evaluate_l2_norm<T: NativePType>(input: ArrayRef) -> VortexResult<Vec<T>> {
let result = L2Norm::try_new(input)?;
let mut ctx = SESSION.create_execution_ctx();
let prim: PrimitiveArray = result.into_array().execute(&mut ctx)?;
Ok(prim.as_slice::<f64>().to_vec())
let output: PrimitiveArray = result.into_array().execute(&mut ctx)?;

Ok(output.as_slice::<T>().to_vec())
}

#[test]
fn test_zero_width_and_empty_inputs() -> VortexResult<()> {
assert_close(
&evaluate_l2_norm(zero_width_vector_array::<f64>(3)?)?,
&[0.0, 0.0, 0.0],
);
assert!(evaluate_l2_norm::<f64>(vector_array::<f64>(2, &[])?)?.is_empty());
assert_close(
&evaluate_l2_norm(Vector::constant_array::<f64>(&[], 3)?)?,
&[0.0, 0.0, 0.0],
);

Ok(())
}

#[rstest]
Expand All @@ -274,7 +208,7 @@ mod tests {
#[case] expected: &[f64],
) -> VortexResult<()> {
let arr = tensor_array(shape, elements)?;
assert_close(&eval_l2_norm(arr)?, expected);
assert_close(&evaluate_l2_norm(arr)?, expected);
Ok(())
}

Expand All @@ -288,7 +222,7 @@ mod tests {
1.0, 1.0, 1.0, // norm = sqrt(3)
],
)?;
assert_close(&eval_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]);
assert_close(&evaluate_l2_norm(arr)?, &[5.0, 0.0, 3.0_f64.sqrt()]);
Ok(())
}

Expand All @@ -301,7 +235,7 @@ mod tests {
3.0, 4.0, 0.0, // norm = 5.0
],
)?;
assert_close(&eval_l2_norm(arr)?, &[1.0, 5.0]);
assert_close(&evaluate_l2_norm(arr)?, &[1.0, 5.0]);
Ok(())
}

Expand Down Expand Up @@ -347,6 +281,50 @@ mod tests {
Ok(())
}

#[test]
fn test_extension_backed_constant_yields_constant_output() -> VortexResult<()> {
let input = Vector::constant_array(&[3.0f64, 4.0], 4)?;

let result = L2Norm::try_new(input)?.into_array();
let mut ctx = SESSION.create_execution_ctx();
let output = result.execute_until::<Constant>(&mut ctx)?;

let constant = output
.as_opt::<Constant>()
.expect("L2Norm over constant-backed extension storage must be constant");
assert_eq!(constant.len(), 4);
let norm = constant
.scalar()
.as_primitive()
.as_::<f64>()
.expect("norm scalar must be a non-null primitive");
assert_eq!(norm.to_bits(), 5.0f64.to_bits());

Ok(())
}

#[test]
fn test_encoded_constants_match_materialized_rows_bitwise() -> VortexResult<()> {
let row = [f32::MAX, 1.0, -1.0];
let materialized = vector_array(3, &[row, row, row].concat())?;
let storage_constant = Vector::constant_array(&row, 3)?;
let literal_constant = literal_vector_array(&row, 3);

let expected: Vec<_> = evaluate_l2_norm::<f32>(materialized)?
.into_iter()
.map(f32::to_bits)
.collect();
for encoded in [storage_constant, literal_constant] {
let actual: Vec<_> = evaluate_l2_norm::<f32>(encoded)?
.into_iter()
.map(f32::to_bits)
.collect();
assert_eq!(actual, expected);
}

Ok(())
}

/// A constant input whose scalar is null should short-circuit to a null [`ConstantArray`] of
/// the correct primitive dtype and length.
#[test]
Expand Down
12 changes: 12 additions & 0 deletions vortex-tensor/src/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,18 @@ pub mod test_helpers {
Vector::try_new_vector_array(flat_fsl(elements, dim))
}

/// Builds `row_count` zero-width vectors over an empty typed element buffer.
pub fn zero_width_vector_array<T: NativePType>(row_count: usize) -> VortexResult<ArrayRef> {
let storage = FixedSizeListArray::new(
Buffer::<T>::empty().into_array(),
0,
Validity::NonNullable,
row_count,
)
.into_array();
Vector::try_new_vector_array(storage)
}

/// Builds a [`FixedShapeTensor`] extension array whose storage is a [`ConstantArray`],
/// representing a single query tensor broadcast to `len` rows.
pub fn constant_tensor_array<T: NativePType + Into<PValue>>(
Expand Down
Loading