Execute tensor L2 norm with RowFn

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

@@ -55,6 +55,11 @@ impl<'a> BorrowedExecutionArgs<'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/execution.rs

@@ -3,11 +3,12 @@
 
 //! Applies columnar semantics around one typed row kernel invocation.
 //!
-//! [`Batch`] owns strict null propagation, constant broadcasting, execution strategy selection, and
-//! output validation. The row kernel therefore handles only decoded values and its selected output
-//! capability.
+//! [`Batch`] owns strict null propagation, encoded reductions, constant broadcasting, execution
+//! strategy selection, and output validation. The row kernel therefore handles only decoded values
+//! and its selected output capability.
 
 use smallvec::SmallVec;
+use vortex_error::VortexError;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
@@ -121,12 +122,17 @@ impl Batch {
         })
     }
 
-    /// Apply constant folding and null handling around `kernel`.
+    /// Apply encoded reductions, constant folding, and null handling around `kernel`.
     ///
-    /// For a mixed validity mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects
-    /// filter-and-scatter. Every kernel result is checked against the planned shape and dtype.
+    /// `reduce` receives the original inputs before constant broadcasting. For a mixed validity
+    /// mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects filter-and-scatter. Every
+    /// kernel result is checked against the planned shape and dtype.
     pub fn execute(
         &self,
+        reduce: impl FnOnce(
+            BorrowedExecutionArgs<'_>,
+            &mut ExecutionCtx,
+        ) -> VortexResult<Option<RowExecution>>,
         kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult<RowExecution>,
         try_unfiltered: impl FnOnce(
             BorrowedExecutionArgs<'_>,
@@ -147,6 +153,20 @@ impl Batch {
             return Ok(self.all_null());
         }
 
+        // An empty mask is both all-true and all-false, so deferred encoded evidence cannot be
+        // attributed to an observable row. Let the ordinary policy construct the typed empty
+        // output instead.
+        if self.row_count > 0
+            && let Some(execution) = reduce(self.execution_args(&self.inputs, self.row_count), ctx)?
+        {
+            match execution {
+                RowExecution::Output(values) => return self.finalize_reduced(values, ctx),
+                RowExecution::DeferredError(error) => {
+                    return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx);
+                }
+            }
+        }
+
         // All inputs constant, and their conjoined validity proves every row non-null. This sees
         // through extension and masked wrappers just like argument decoding does.
         if self.row_count > 0
@@ -331,6 +351,56 @@ impl Batch {
         ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array()
     }
 
+    /// Reconcile an encoding-aware result and apply the batch's strict input validity.
+    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)
+    }
+
+    /// Resolve deferred evidence from the encoded path by executing only observable rows.
+    fn resolve_reduced_error(
+        &self,
+        error: VortexError,
+        kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult<RowExecution>,
+        try_unfiltered: impl FnOnce(
+            BorrowedExecutionArgs<'_>,
+            &Mask,
+            &mut ExecutionCtx,
+        ) -> VortexResult<Option<RowExecution>>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let valid = self.validity.clone().execute_mask(self.row_count, ctx)?;
+
+        if valid.all_true() {
+            return Err(error);
+        }
+        if valid.all_false() {
+            return Ok(self.all_null());
+        }
+
+        if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? {
+            return Ok(result);
+        }
+
+        self.filter_and_scatter(kernel, &valid, ctx)
+    }
+
     /// Pair an input view with this batch's planning metadata.
     fn execution_args<'b>(
         &'b self,

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

@@ -31,7 +31,6 @@ use crate::dtype::DType;
 use crate::dtype::NativePType;
 use crate::dtype::Nullability;
 use crate::scalar_fn::EmptyOptions;
-use crate::scalar_fn::ExecutionArgs;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::VecExecutionArgs;
 use crate::scalar_fn::unstable::row::InputElement;
@@ -59,7 +58,13 @@ struct AddShort;
 struct ShortDecodeI64;
 
 #[derive(Clone)]
-struct Identity;
+struct OriginalInputReducer;
+
+#[derive(Clone)]
+struct InvalidEncodedReduction;
+
+#[derive(Clone)]
+struct DeferredOriginalReducer;
 
 #[derive(Clone)]
 struct SinkOptions;
@@ -308,16 +313,103 @@ impl RowFn for RetryConstantAdd {
             },
         )
     }
+
+    fn reduce_encoded(
+        &self,
+        _options: &Self::Options,
+        args: &[ArrayRef],
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<RowExecution>> {
+        if args[0].len() == 1 {
+            return Ok(Some(RowExecution::Output(
+                ConstantArray::new(0u8, args[0].len()).into_array(),
+            )));
+        }
+
+        Ok(None)
+    }
 }
 
-impl RowFn for Identity {
+impl RowFn for OriginalInputReducer {
     type Options = EmptyOptions;
 
     const ARG_NAMES: &'static [&'static str] = &["value"];
     const FALLIBLE: bool = false;
 
     fn id(&self) -> ScalarFnId {
-        static ID: CachedId = CachedId::new("test.identity");
+        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<RowExecution>> {
+        if args[0].len() == 3 {
+            return Ok(Some(RowExecution::Output(
+                ConstantArray::new(42_i64, 3).into_array(),
+            )));
+        }
+
+        Ok(None)
+    }
+}
+
+impl RowFn for InvalidEncodedReduction {
+    type Options = usize;
+
+    const ARG_NAMES: &'static [&'static str] = &["value"];
+    const FALLIBLE: bool = false;
+
+    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<RowExecution>> {
+        Ok(Some(RowExecution::Output(
+            PrimitiveArray::new(
+                vec![10_i64, 20],
+                Validity::from_iter((0..2).map(|index| index != *null_index)),
+            )
+            .into_array(),
+        )))
+    }
+}
+
+impl RowFn for DeferredOriginalReducer {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["value"];
+    const FALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.deferred_original_reducer");
         *ID
     }
 
@@ -329,6 +421,17 @@ impl RowFn for Identity {
     ) -> VortexResult<V::VisitResult> {
         visitor.visit::<(i64,), i64>(|(value,)| value)
     }
+
+    fn reduce_encoded(
+        &self,
+        _options: &Self::Options,
+        _args: &[ArrayRef],
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<RowExecution>> {
+        Ok(Some(RowExecution::DeferredError(vortex_err!(
+            InvalidArgument: "encoded payload failed"
+        ))))
+    }
 }
 
 impl RowFn for SinkOptions {
@@ -461,7 +564,7 @@ fn test_short_decode_beside_constant_is_rejected() -> VortexResult<()> {
 }
 
 #[test]
-fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> {
+fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> {
     let lhs =
         PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array();
     let rhs = ConstantArray::new(1u8, 2).into_array();
@@ -489,13 +592,89 @@ fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn test_reduce_encoded_defers_errors_behind_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(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?;
+
+    assert_arrays_eq!(&actual, &input, &mut ctx);
+    Ok(())
+}
+
+#[test]
+fn test_empty_batch_skips_deferred_encoded_error() -> VortexResult<()> {
+    let input = PrimitiveArray::from_iter(Vec::<i64>::new()).into_array();
+    let args = VecExecutionArgs::new(vec![input.clone()], 0);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?;
+
+    assert_arrays_eq!(&actual, &input, &mut ctx);
+    Ok(())
+}
+
+#[rstest]
+#[case::all_valid(Validity::AllValid)]
+#[case::mixed(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_constant_input_broadcasts_one_row() -> VortexResult<()> {
     let input = ConstantArray::new(7_i64, 2).into_array();
     let args = VecExecutionArgs::new(vec![input.clone()], 2);
     let mut ctx = array_session().create_execution_ctx();
 
-    let actual = execute_rows(&Identity, &EmptyOptions, &args, &mut ctx)?;
+    let actual = execute_rows(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?;
 
     assert_arrays_eq!(&actual, &input, &mut ctx);
     Ok(())
@@ -519,7 +698,8 @@ fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResul
     let mut ctx = array_session().create_execution_ctx();
 
     let actual = batch.execute(
-        |args, _ctx| Ok(RowExecution::Output(args.get(0)?)),
+        |_args, _ctx| Ok(None),
+        |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())),
         |_args, _valid, _ctx| Ok(None),
         &mut ctx,
     )?;
@@ -547,7 +727,8 @@ fn test_valid_only_filters_and_scatters() -> VortexResult<()> {
     let mut ctx = array_session().create_execution_ctx();
 
     let actual = batch.execute(
-        |args, _ctx| Ok(RowExecution::Output(args.get(0)?)),
+        |_args, _ctx| Ok(None),
+        |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())),
         |_args, _valid, _ctx| Ok(None),
         &mut ctx,
     )?;