Add physical parameters to OutputSink

vortex-array/src/scalar_fn/fns/binary/numeric/row.rs

@@ -116,7 +116,7 @@ where
     // vectorization. Check each divide immediately and stop at the first failure.
     // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when
     // they need to skip invalid rows.
-    visitor.visit_into::<(T, T), UninitElementSink<T>, _>(|(lhs, rhs), output| {
+    visitor.visit_into::<(T, T), UninitElementSink<T>, _>((), |(lhs, rhs), output| {
         let (value, failed) = CheckedDiv::apply(lhs, rhs);
         if failed {
             return Err(numeric_error(<CheckedDiv as CheckedPrimitiveOp<T>>::ERROR));

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

@@ -9,6 +9,8 @@ use rstest::rstest;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
 use vortex_session::registry::CachedId;
 
 use super::finalize_kernel_output;
@@ -18,6 +20,7 @@ use crate::VortexSessionExecute;
 use crate::array_session;
 use crate::arrays::ConstantArray;
 use crate::arrays::ExtensionArray;
+use crate::arrays::FixedSizeListArray;
 use crate::arrays::PrimitiveArray;
 use crate::assert_arrays_eq;
 use crate::dtype::DType;
@@ -30,6 +33,8 @@ use crate::scalar::Scalar;
 use crate::scalar_fn::EmptyOptions;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::VecExecutionArgs;
+use crate::scalar_fn::unstable::row::FixedSizeListSink;
+use crate::scalar_fn::unstable::row::InitializedRow;
 use crate::scalar_fn::unstable::row::OutputElement;
 use crate::scalar_fn::unstable::row::OutputSink;
 use crate::scalar_fn::unstable::row::RowFn;
@@ -78,15 +83,16 @@ struct I64Sink(BufferMut<i64>);
 // SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that
 // initialized slice. The `()` write token therefore proves no additional invariant.
 unsafe impl OutputSink for I64Sink {
+    type Params = ();
     type Rows<'a> = &'a mut [i64];
     type Row<'a> = &'a mut i64;
     type WriteToken = ();
 
-    fn storage_dtype() -> DType {
+    fn storage_dtype(_params: &Self::Params) -> DType {
         DType::from(i64::PTYPE)
     }
 
-    fn with_capacity(rows: usize) -> VortexResult<Self> {
+    fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult<Self> {
         Ok(Self(BufferMut::zeroed(rows)))
     }
 
@@ -104,6 +110,38 @@ unsafe impl OutputSink for I64Sink {
     }
 }
 
+#[derive(Clone)]
+struct RepeatValue;
+
+impl RowFn for RepeatValue {
+    type Options = usize;
+
+    const ARG_NAMES: &'static [&'static str] = &["value"];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.repeat_value");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        width: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        vortex_ensure!(
+            u32::try_from(*width).is_ok(),
+            InvalidArgument:
+            "test.repeat_value width must fit in the fixed-size-list u32 list size, got {width}",
+        );
+
+        visitor.visit_into::<(i64,), FixedSizeListSink<i64>, _>(*width, |(value,), row| {
+            InitializedRow::fill(row, |_| value)
+        })
+    }
+}
+
 impl RowFn for DeferredAdd {
     type Options = EmptyOptions;
 
@@ -156,7 +194,7 @@ impl RowFn for ValidOnlyIdentity {
         _args: &[DType],
         visitor: V,
     ) -> VortexResult<V::VisitResult> {
-        visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| {
+        visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>((), |(value,), output| {
             *output = value;
             Ok(())
         })
@@ -184,6 +222,58 @@ impl RowFn for InvalidKernelOutput {
     }
 }
 
+#[rstest]
+#[case::dense_width_two(
+    vec![1_i64, 2],
+    Validity::NonNullable,
+    2,
+    vec![1_i64, 1, 2, 2],
+)]
+#[case::dense_width_four(
+    vec![3_i64, 4],
+    Validity::NonNullable,
+    4,
+    vec![3_i64, 3, 3, 3, 4, 4, 4, 4],
+)]
+#[case::empty(vec![], Validity::NonNullable, 3, vec![])]
+#[case::zero_width(vec![3_i64, 4], Validity::NonNullable, 0, vec![])]
+#[case::all_null(
+    vec![5_i64, 6],
+    Validity::AllInvalid,
+    2,
+    vec![0_i64, 0, 0, 0],
+)]
+#[case::partially_valid(
+    vec![7_i64, 8, 9],
+    Validity::from_iter([true, false, true]),
+    3,
+    vec![7_i64, 7, 7, 0, 0, 0, 9, 9, 9],
+)]
+fn test_fixed_size_list_sink_uses_runtime_width(
+    #[case] input_values: Vec<i64>,
+    #[case] validity: Validity,
+    #[case] width: usize,
+    #[case] expected_elements: Vec<i64>,
+) -> VortexResult<()> {
+    let row_count = input_values.len();
+    let input = PrimitiveArray::new(input_values, validity.clone()).into_array();
+    let args = VecExecutionArgs::new(vec![input], row_count);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&RepeatValue, &width, &args, &mut ctx)?;
+    let expected = FixedSizeListArray::new(
+        PrimitiveArray::from_iter(expected_elements).into_array(),
+        u32::try_from(width)
+            .map_err(|_| vortex_err!(InvalidArgument: "test width must fit in u32, got {width}"))?,
+        validity,
+        row_count,
+    )
+    .into_array();
+
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+    Ok(())
+}
+
 #[test]
 fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<()> {
     static ID: CachedId = CachedId::new("test.finalize_kernel_output");
@@ -403,7 +493,7 @@ impl RowFn for DeclaredSinkOutput {
     ) -> VortexResult<V::VisitResult> {
         visitor
             .with_output_dtype(timestamp_dtype(Nullability::NonNullable))
-            .visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| {
+            .visit_into::<(i64,), I64Sink, VortexResult<()>>((), |(value,), output| {
                 *output = value;
                 Ok(())
             })

vortex-array/src/scalar_fn/unstable/row/execute/sink.rs

@@ -28,6 +28,7 @@ use crate::scalar_fn::unstable::row::ViewLen;
 /// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant.
 pub(crate) fn execute_sink<Args, Prepared, Sink, ApplyResult>(
     args: &dyn ExecutionArgs,
+    params: &Sink::Params,
     ctx: &mut ExecutionCtx,
     prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
     apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult,
@@ -43,7 +44,7 @@ where
     let const_values = Args::const_values(&columns);
     let prepared = prepare(const_values);
 
-    let mut sink = Sink::with_capacity(row_count)?;
+    let mut sink = Sink::with_capacity(row_count, params)?;
 
     // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink.
     {
@@ -100,6 +101,7 @@ where
 pub(crate) fn execute_sink_valid_rows<Args, Prepared, Sink, ApplyResult>(
     args: &dyn ExecutionArgs,
     valid: &MaskValuesRef,
+    params: &Sink::Params,
     ctx: &mut ExecutionCtx,
     prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
     apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult,
@@ -115,7 +117,7 @@ where
         valid_rows,
         row_count,
         mut sink,
-    }) = setup_sink_valid_rows::<Args, Sink>(args, valid, ctx)?
+    }) = setup_sink_valid_rows::<Args, Sink>(args, valid, params, ctx)?
     else {
         return Ok(None);
     };
@@ -205,6 +207,7 @@ where
 fn setup_sink_valid_rows<'valid, Args, Sink>(
     args: &dyn ExecutionArgs,
     valid: &'valid MaskValuesRef,
+    params: &Sink::Params,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<Option<ValidRowsSetup<'valid, Args, Sink>>>
 where
@@ -227,7 +230,7 @@ where
     // Keep allocation before the validity and length checks. With multiple CGUs and no LTO,
     // moving it later inlines `Args::get` into every sparse callback, duplicating its bounds
     // checks.
-    let sink = Sink::with_capacity(row_count)?;
+    let sink = Sink::with_capacity(row_count, params)?;
 
     let valid_rows = valid.bit_buffer();
     vortex_ensure_eq!(
@@ -272,15 +275,16 @@ mod tests {
     // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or
     // `finish` through the executor. The row-initialization requirements are therefore vacuous.
     unsafe impl OutputSink for NonSkippingSink {
+        type Params = ();
         type Rows<'a> = ();
         type Row<'a> = ();
         type WriteToken = ();
 
-        fn storage_dtype() -> DType {
+        fn storage_dtype(_params: &Self::Params) -> DType {
             DType::from(i64::PTYPE)
         }
 
-        fn with_capacity(_rows: usize) -> VortexResult<Self> {
+        fn with_capacity(_rows: usize, _params: &Self::Params) -> VortexResult<Self> {
             Err(vortex_err!(
                 "a non-skipping sink must decline before allocation"
             ))
@@ -300,6 +304,7 @@ mod tests {
     // post-initialization length check. If execution incorrectly continues, safe indexing in
     // `row_unchecked` panics instead of accessing invalid memory.
     unsafe impl OutputSink for ShrinkingSink {
+        type Params = ();
         type Rows<'a> = &'a mut Vec<i64>;
         type Row<'a> = &'a mut i64;
         type WriteToken = ();
@@ -310,11 +315,11 @@ mod tests {
             })
         }
 
-        fn storage_dtype() -> DType {
+        fn storage_dtype(_params: &Self::Params) -> DType {
             DType::from(i64::PTYPE)
         }
 
-        fn with_capacity(rows: usize) -> VortexResult<Self> {
+        fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult<Self> {
             Ok(Self(vec![0; rows]))
         }
 
@@ -343,6 +348,7 @@ mod tests {
         let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>(
             &args,
             &valid,
+            &(),
             &mut ctx,
             |_| (),
             |_, _, _| (),
@@ -365,6 +371,7 @@ mod tests {
         let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, ()>(
             &args,
             &valid,
+            &(),
             &mut ctx,
             |_| (),
             |_, (value,), output| {

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

@@ -31,8 +31,10 @@ pub use row_fn::RowFn;
 mod types;
 pub use types::ElementTuple;
 pub use types::FailureEvidence;
+pub use types::FixedSizeListSink;
 pub use types::IndexedElementTuple;
 pub use types::InitializedElement;
+pub use types::InitializedRow;
 pub use types::InputElement;
 pub use types::OutputElement;
 pub use types::OutputSink;

vortex-array/src/scalar_fn/unstable/row/types/mod.rs

@@ -20,7 +20,9 @@ pub use result::FailureEvidence;
 pub use result::SinkResult;
 
 mod sink;
+pub use sink::FixedSizeListSink;
 pub use sink::InitializedElement;
+pub use sink::InitializedRow;
 pub use sink::OutputSink;
 pub use sink::UninitElementSink;
 

vortex-array/src/scalar_fn/unstable/row/types/result.rs

@@ -4,14 +4,15 @@
 //! Return types for row closures.
 //!
 //! [`FailureEvidence`] represents deferred failures from owned row closures. [`SinkResult`] lets
-//! the executor handle initialized sinks and sinks that require an [`InitializedElement`] token,
-//! with either infallible or immediate-error callbacks.
+//! the executor handle initialized sinks and sinks that require an [`InitializedElement`] or
+//! [`InitializedRow`] token, with either infallible or immediate-error callbacks.
 
 use std::ops::BitOrAssign;
 
 use vortex_error::VortexResult;
 
 use super::InitializedElement;
+use super::InitializedRow;
 
 /// Compact failure evidence that can be OR-reduced across rows.
 ///
@@ -57,6 +58,17 @@ impl SinkResult for InitializedElement {
     }
 }
 
+impl private::Sealed for InitializedRow {}
+
+impl SinkResult for InitializedRow {
+    type WriteToken = InitializedRow;
+    const INFALLIBLE: bool = true;
+
+    fn into_result(self) -> VortexResult<()> {
+        Ok(())
+    }
+}
+
 impl private::Sealed for VortexResult<()> {}
 
 impl SinkResult for VortexResult<()> {