Execute tensor L2 norm with RowFn

vortex-array/src/scalar_fn/unstable/row/batch/args.rs

@@ -55,6 +55,11 @@ impl<'a> BorrowedRowFnArgs<'a> {
         }
     }
 
+    /// Return the concrete arrays used by encoding-aware execution.
+    pub(crate) fn arrays(&self) -> &'a [ArrayRef] {
+        self.arrays
+    }
+
     /// Return the original input dtypes used to select the row implementation.
     pub(crate) fn dtypes(&self) -> &'a [DType] {
         self.dtypes

vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs

@@ -3,8 +3,8 @@
 
 //! Selects a batch execution strategy.
 //!
-//! [`RowFnExecutionArgs::execute`] handles universal fast paths, then delegates to dense or
-//! valid-only execution.
+//! [`RowFnExecutionArgs::execute`] handles universal fast paths and encoded reductions, then
+//! delegates to dense or valid-only execution.
 
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
@@ -27,12 +27,14 @@ mod output;
 pub(crate) use output::finalize_kernel_output;
 
 impl RowFnExecutionArgs {
-    /// Apply constant folding and null handling around `kernel`.
+    /// Apply encoded reductions, constant folding, and null handling around `kernel`.
     ///
-    /// For a partially valid batch, `try_valid_rows` executes only valid rows over the original
-    /// inputs. Every result is checked against the planned shape and dtype.
+    /// `reduce` receives the original inputs before constant broadcasting. For a partially valid
+    /// batch, `try_valid_rows` executes only valid rows over the original inputs. Every result is
+    /// checked against the planned shape and dtype.
     pub(crate) fn execute(
         &self,
+        reduce: impl FnOnce(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<Option<ArrayRef>>,
         kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
         try_valid_rows: impl FnOnce(
             BorrowedRowFnArgs<'_>,
@@ -53,6 +55,13 @@ impl RowFnExecutionArgs {
             return Ok(self.all_null());
         }
 
+        // Let the ordinary policy construct the typed output for an empty batch.
+        if self.row_count > 0
+            && let Some(values) = reduce(self.execution_args(&self.inputs, self.row_count), ctx)?
+        {
+            return self.finalize_reduced(values, ctx);
+        }
+
         // All inputs are constant, and their conjoined validity proves that every row is non-null.
         // The constant check sees through extension and masked wrappers, just like argument
         // decoding.

vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs

@@ -14,6 +14,7 @@ use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::scalar::Scalar;
 use crate::scalar_fn::ScalarFnId;
+use crate::validity::Validity;
 
 impl RowFnExecutionArgs {
     pub(super) fn all_null(&self) -> ArrayRef {
@@ -31,6 +32,32 @@ impl RowFnExecutionArgs {
         cast_output_nullability(&self.result_dtype, values)
     }
 
+    /// Reconcile an encoding-aware result and apply the batch's strict input validity.
+    pub(super) fn finalize_reduced(
+        &self,
+        values: ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        validate_output(self.id, &self.result_dtype, self.row_count, &values)?;
+
+        let input_valid = self.validity.execute_mask(self.row_count, ctx)?;
+        let output_valid = values.validity()?.execute_mask(self.row_count, ctx)?;
+        vortex_ensure!(
+            input_valid.bitand_not(&output_valid).all_false(),
+            "the {} encoded reduction produced nulls for valid rows",
+            self.id,
+        );
+
+        let values = match self.validity.clone() {
+            Validity::NonNullable | Validity::AllValid => values,
+            Validity::Array(valid) => values.mask(valid)?,
+            // Handled before the encoding-aware hook runs.
+            Validity::AllInvalid => return Ok(self.all_null()),
+        };
+
+        cast_output_nullability(&self.result_dtype, values)
+    }
+
     /// Validate the output from a row function before batch validity is attached.
     pub(super) fn validate_kernel_output(
         &self,

vortex-array/src/scalar_fn/unstable/row/batch/tests.rs

@@ -3,13 +3,15 @@
 
 use std::sync::Arc;
 
+use rstest::rstest;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_session::registry::CachedId;
 
 use super::finalize_kernel_output;
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
 use crate::array_session;
@@ -30,6 +32,12 @@ use crate::scalar_fn::unstable::row::RowVisitor;
 use crate::scalar_fn::unstable::row::execute_rows;
 use crate::validity::Validity;
 
+#[derive(Clone)]
+struct OriginalInputReducer;
+
+#[derive(Clone)]
+struct InvalidEncodedReduction;
+
 #[derive(Clone)]
 struct DeferredAdd;
 
@@ -87,6 +95,76 @@ unsafe impl<Options> OutputSink<Options> for I64Sink {
     }
 }
 
+impl RowFn for OriginalInputReducer {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["value"];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.original_input_reducer");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor<Self::Options>>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit::<(i64,), i64>(|(value,)| value)
+    }
+
+    fn reduce_encoded(
+        &self,
+        _options: &Self::Options,
+        args: &[ArrayRef],
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        if args[0].len() == 3 {
+            return Ok(Some(ConstantArray::new(42_i64, 3).into_array()));
+        }
+
+        Ok(None)
+    }
+}
+
+impl RowFn for InvalidEncodedReduction {
+    type Options = usize;
+
+    const ARG_NAMES: &'static [&'static str] = &["value"];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.invalid_encoded_reduction");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor<Self::Options>>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit::<(i64,), i64>(|(value,)| value)
+    }
+
+    fn reduce_encoded(
+        &self,
+        null_index: &Self::Options,
+        _args: &[ArrayRef],
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        Ok(Some(
+            PrimitiveArray::new(
+                vec![10_i64, 20],
+                Validity::from_iter((0..2).map(|index| index != *null_index)),
+            )
+            .into_array(),
+        ))
+    }
+}
+
 impl RowFn for DeferredAdd {
     type Options = EmptyOptions;
 
@@ -182,6 +260,57 @@ fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<(
     Ok(())
 }
 
+#[rstest]
+#[case::all_valid(Validity::AllValid)]
+#[case::partially_valid(Validity::from_iter([true, false]))]
+fn test_reduce_encoded_rejects_nulls_on_valid_rows(#[case] validity: Validity) -> VortexResult<()> {
+    let input = PrimitiveArray::new(vec![10_i64, 20], validity).into_array();
+    let args = VecExecutionArgs::new(vec![input], 2);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let error = match execute_rows(&InvalidEncodedReduction, &0, &args, &mut ctx) {
+        Err(error) => error,
+        Ok(_) => vortex_bail!("an encoded reduction introduced a null on a valid row"),
+    };
+    let error = error.to_string();
+
+    assert!(
+        error.contains("test.invalid_encoded_reduction"),
+        "the boundary error must name the function, got {error}",
+    );
+    assert!(
+        error.contains("encoded reduction produced nulls for valid rows"),
+        "the boundary error must identify invalid reduced output, got {error}",
+    );
+    Ok(())
+}
+
+#[test]
+fn test_reduce_encoded_preserves_input_nulls() -> VortexResult<()> {
+    let input =
+        PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array();
+    let args = VecExecutionArgs::new(vec![input.clone()], 2);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&InvalidEncodedReduction, &1, &args, &mut ctx)?;
+
+    assert_arrays_eq!(&actual, &input, &mut ctx);
+    Ok(())
+}
+
+#[test]
+fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> {
+    let input = ConstantArray::new(7_i64, 3).into_array();
+    let args = VecExecutionArgs::new(vec![input], 3);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?;
+    let expected = ConstantArray::new(42_i64, 3).into_array();
+
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+    Ok(())
+}
+
 #[test]
 fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> {
     let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array();

vortex-array/src/scalar_fn/unstable/row/row_fn.rs

@@ -4,8 +4,8 @@
 //! The [`RowFn`] contract for scalar functions whose natural kernel computes one row at a time.
 //!
 //! Implementations declare their arity and fallibility, then use [`RowFn::dispatch`] to select the
-//! typed row signature for each supported dtype combination. Optional methods provide
-//! serialization without putting persistence plumbing in the row kernel.
+//! typed row signature for each supported dtype combination. Optional hooks provide serialization
+//! and encoding-aware execution without putting columnar plumbing in the row kernel.
 
 use std::fmt::Debug;
 use std::fmt::Display;
@@ -16,6 +16,8 @@ use vortex_error::vortex_bail;
 use vortex_session::VortexSession;
 
 use super::visitor::RowVisitor;
+use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::dtype::DType;
 use crate::scalar_fn::ScalarFnId;
 
@@ -47,12 +49,14 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync {
     /// The arguments in display order. Its length is the function's exact arity.
     const ARG_NAMES: &'static [&'static str];
 
-    /// Whether every dispatch is infallible.
+    /// Whether every dispatch and encoded reduction is infallible.
     ///
     /// See [`ScalarFnVTable::is_infallible`](crate::scalar_fn::ScalarFnVTable::is_infallible) for
     /// a more detailed explanation of semantic errors.
     ///
-    /// The framework checks dispatched element and result types. A conservative `false` is allowed.
+    /// The framework checks dispatched element and result types, but cannot inspect
+    /// [`reduce_encoded`](Self::reduce_encoded). Set this to `false` when that hook can return a
+    /// semantic error. A conservative `false` is allowed.
     const INFALLIBLE: bool;
 
     /// Returns the ID of the scalar function.
@@ -83,4 +87,29 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync {
         args: &[DType],
         visitor: V,
     ) -> VortexResult<V::VisitResult>;
+
+    /// Try an encoding-aware implementation before decoding the inputs into row elements.
+    ///
+    /// `None` continues to the row loop. The returned array can remain encoded or lazy. Batch
+    /// execution calls this hook at most once with the original nonempty inputs. Nullary functions,
+    /// empty batches, slices, and compacted retries skip it.
+    ///
+    /// Like a dense row closure, this hook must be total over every stored payload, including
+    /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or
+    /// retried through the row layer.
+    ///
+    /// # Requirements
+    ///
+    /// - `output.len()` **must** equal `args[0].len()`.
+    /// - The output dtype **must** match the planned dtype when ignoring outer nullability.
+    /// - The output **must not** introduce a null where every input is valid.
+    fn reduce_encoded(
+        &self,
+        options: &Self::Options,
+        args: &[ArrayRef],
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        _ = (options, args, ctx);
+        Ok(None)
+    }
 }

vortex-array/src/scalar_fn/unstable/row/vtable.rs

@@ -118,6 +118,7 @@ pub fn execute_rows<F: RowFn>(
 
     let batch = prepare_batch(function, options, args)?;
     batch.execute(
+        |args, ctx| function.reduce_encoded(options, args.arrays(), ctx),
         |args, ctx| execute_row_kernel(function, options, args, ctx),
         |args, valid, ctx| try_execute_valid_rows(function, options, args, valid, ctx),
         ctx,

vortex-tensor/Cargo.toml

@@ -17,7 +17,7 @@ version = { workspace = true }
 workspace = true
 
 [dependencies]
-vortex-array = { workspace = true }
+vortex-array = { workspace = true, features = ["unstable_row_fns"] }
 vortex-arrow = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-compressor = { workspace = true }

vortex-tensor/benches/l2_norm.rs

@@ -15,11 +15,17 @@ use divan::Bencher;
 use divan::counter::ItemsCount;
 use mimalloc::MiMalloc;
 use vortex_array::ArrayRef;
+use vortex_array::EmptyMetadata;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::MaskedArray;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
+use vortex_array::scalar::Scalar;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_tensor::scalar_fns::l2_norm::L2Norm;
@@ -54,6 +60,16 @@ fn vectors(width: usize) -> ArrayRef {
     Vector::try_new_vector_array(storage).unwrap()
 }
 
+fn constant_vector(width: usize) -> ArrayRef {
+    let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
+    let children = (0..width)
+        .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable))
+        .collect();
+    let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
+    let vector = Scalar::extension::<Vector>(EmptyMetadata, storage);
+    ConstantArray::new(vector, ELEMENTS / width).into_array()
+}
+
 fn bench_l2_norm(bencher: Bencher, input: ArrayRef) {
     let session = vortex_array::array_session();
     bencher
@@ -80,3 +96,17 @@ fn nullable(bencher: Bencher, width: usize) {
         .into_array();
     bench_l2_norm(bencher, input);
 }
+
+#[divan::bench(args = WIDTHS)]
+fn constant(bencher: Bencher, width: usize) {
+    bench_l2_norm(bencher, constant_vector(width));
+}
+
+#[divan::bench(args = WIDTHS)]
+fn nullable_constant(bencher: Bencher, width: usize) {
+    let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0));
+    let input = MaskedArray::try_new(constant_vector(width), validity)
+        .unwrap()
+        .into_array();
+    bench_l2_norm(bencher, input);
+}