diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index 60a23a1154d..d80091844c7 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -23,7 +23,6 @@ mod dense; mod valid_only; mod output; -#[cfg(test)] pub(crate) use output::finalize_kernel_output; impl RowFnExecutionArgs { diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs index dd6cfa9ba0f..b2784009011 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -24,7 +24,6 @@ mod args; pub(super) use args::BorrowedRowFnArgs; mod execute; -#[cfg(test)] pub(super) use execute::finalize_kernel_output; mod planning; diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index c2b8f44f6c3..bb0d7815e22 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -9,13 +9,13 @@ //! that delegate to a private row kernel. use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; use vortex_session::VortexSession; use super::batch::BorrowedRowFnArgs; use super::batch::RowFnExecutionArgs; +use super::batch::finalize_kernel_output; use super::row_fn::RowFn; use super::visitor::BatchPlanner; use super::visitor::ExecuteRows; @@ -104,6 +104,9 @@ pub fn row_fn_return_dtype( /// A type cannot implement both [`RowFn`] and [`ScalarFnVTable`] because every `RowFn` receives the /// standard vtable automatically. Existing vtables can keep their custom hooks on one type and /// delegate row execution to a private `RowFn` kernel through this function. +/// +/// Nullary functions execute for `args.row_count()` rows without batch validity handling because +/// they have no input validity to propagate. pub fn execute_rows( function: &F, options: &F::Options, @@ -111,10 +114,10 @@ pub fn execute_rows( ctx: &mut ExecutionCtx, ) -> VortexResult { ensure_arity(function, args.num_inputs())?; - vortex_ensure!( - args.num_inputs() != 0, - "row-function execution does not support nullary kernels" - ); + + if args.num_inputs() == 0 { + return execute_nullary_rows(function, options, args.row_count(), ctx); + } let batch = prepare_batch(function, options, args)?; batch.execute( @@ -124,6 +127,22 @@ pub fn execute_rows( ) } +/// Execute a nullary kernel without batch validity or constant handling. +fn execute_nullary_rows( + function: &F, + options: &F::Options, + row_count: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let plan = function.dispatch(options, &[], BatchPlanner::::new(&[], options))?; + let result_dtype = plan.result_dtype(&[]); + let args = BorrowedRowFnArgs::new(&[], row_count, &[], &plan.output_dtype, plan.policy); + + let values = execute_row_kernel(function, options, args, ctx)?; + + finalize_kernel_output(RowFn::id(function), &result_dtype, row_count, values, ctx) +} + fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { let expected = F::ARG_NAMES.len(); vortex_ensure_eq!( @@ -198,6 +217,7 @@ mod tests { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use rstest::rstest; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_session::registry::CachedId; @@ -208,6 +228,7 @@ mod tests { use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; use crate::dtype::DType; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; @@ -219,6 +240,9 @@ mod tests { #[derive(Clone)] struct IndexingRowFn; + #[derive(Clone)] + struct NullarySeven; + #[derive(Clone)] struct ChangingDispatchRowFn { dispatches: Arc, @@ -231,6 +255,27 @@ mod tests { Element, } + impl RowFn for NullarySeven { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_seven"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(), i64>(|()| 7) + } + } + impl RowFn for IndexingRowFn { type Options = EmptyOptions; @@ -302,6 +347,20 @@ mod tests { assert_arity_error(error); } + #[rstest] + #[case::empty(0)] + #[case::nonempty(3)] + fn test_execute_nullary_rows(#[case] row_count: usize) -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], row_count); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter(vec![7_i64; row_count]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + #[test] fn test_execute_rejects_dispatch_that_changes_after_planning() -> VortexResult<()> { let function = ChangingDispatchRowFn {