Skip to content

Commit a9cfe2f

Browse files
committed
Support nullary RowFn execution
Signed-off-by: Connor Tsui <[email protected]>
1 parent 46a8d39 commit a9cfe2f

3 files changed

Lines changed: 64 additions & 7 deletions

File tree

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ mod dense;
2323
mod valid_only;
2424

2525
mod output;
26-
#[cfg(test)]
2726
pub(crate) use output::finalize_kernel_output;
2827

2928
impl RowFnExecutionArgs {

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ mod args;
2424
pub(super) use args::BorrowedRowFnArgs;
2525

2626
mod execute;
27-
#[cfg(test)]
2827
pub(super) use execute::finalize_kernel_output;
2928

3029
mod planning;

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

Lines changed: 64 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,13 @@
99
//! that delegate to a private row kernel.
1010
1111
use vortex_error::VortexResult;
12-
use vortex_error::vortex_ensure;
1312
use vortex_error::vortex_ensure_eq;
1413
use vortex_mask::Mask;
1514
use vortex_session::VortexSession;
1615

1716
use super::batch::BorrowedRowFnArgs;
1817
use super::batch::RowFnExecutionArgs;
18+
use super::batch::finalize_kernel_output;
1919
use super::row_fn::RowFn;
2020
use super::visitor::BatchPlanner;
2121
use super::visitor::ExecuteRows;
@@ -104,17 +104,20 @@ pub fn row_fn_return_dtype<F: RowFn>(
104104
/// A type cannot implement both [`RowFn`] and [`ScalarFnVTable`] because every `RowFn` receives the
105105
/// standard vtable automatically. Existing vtables can keep their custom hooks on one type and
106106
/// delegate row execution to a private `RowFn` kernel through this function.
107+
///
108+
/// Nullary functions execute for `args.row_count()` rows without batch validity handling because
109+
/// they have no input validity to propagate.
107110
pub fn execute_rows<F: RowFn>(
108111
function: &F,
109112
options: &F::Options,
110113
args: &dyn ExecutionArgs,
111114
ctx: &mut ExecutionCtx,
112115
) -> VortexResult<ArrayRef> {
113116
ensure_arity(function, args.num_inputs())?;
114-
vortex_ensure!(
115-
args.num_inputs() != 0,
116-
"row-function execution does not support nullary kernels"
117-
);
117+
118+
if args.num_inputs() == 0 {
119+
return execute_nullary_rows(function, options, args.row_count(), ctx);
120+
}
118121

119122
let batch = prepare_batch(function, options, args)?;
120123
batch.execute(
@@ -124,6 +127,22 @@ pub fn execute_rows<F: RowFn>(
124127
)
125128
}
126129

130+
/// Execute a nullary kernel without batch validity or constant handling.
131+
fn execute_nullary_rows<F: RowFn>(
132+
function: &F,
133+
options: &F::Options,
134+
row_count: usize,
135+
ctx: &mut ExecutionCtx,
136+
) -> VortexResult<ArrayRef> {
137+
let plan = function.dispatch(options, &[], BatchPlanner::<F>::new(&[], options))?;
138+
let result_dtype = plan.result_dtype(&[]);
139+
let args = BorrowedRowFnArgs::new(&[], row_count, &[], &plan.output_dtype, plan.policy);
140+
141+
let values = execute_row_kernel(function, options, args, ctx)?;
142+
143+
finalize_kernel_output(RowFn::id(function), &result_dtype, row_count, values, ctx)
144+
}
145+
127146
fn ensure_arity<F: RowFn>(function: &F, actual: usize) -> VortexResult<()> {
128147
let expected = F::ARG_NAMES.len();
129148
vortex_ensure_eq!(
@@ -198,6 +217,7 @@ mod tests {
198217
use std::sync::atomic::AtomicUsize;
199218
use std::sync::atomic::Ordering;
200219

220+
use rstest::rstest;
201221
use vortex_error::VortexError;
202222
use vortex_error::VortexResult;
203223
use vortex_session::registry::CachedId;
@@ -208,6 +228,7 @@ mod tests {
208228
use crate::VortexSessionExecute;
209229
use crate::array_session;
210230
use crate::arrays::PrimitiveArray;
231+
use crate::assert_arrays_eq;
211232
use crate::dtype::DType;
212233
use crate::scalar_fn::EmptyOptions;
213234
use crate::scalar_fn::ScalarFnId;
@@ -219,6 +240,9 @@ mod tests {
219240
#[derive(Clone)]
220241
struct IndexingRowFn;
221242

243+
#[derive(Clone)]
244+
struct NullarySeven;
245+
222246
#[derive(Clone)]
223247
struct ChangingDispatchRowFn {
224248
dispatches: Arc<AtomicUsize>,
@@ -231,6 +255,27 @@ mod tests {
231255
Element,
232256
}
233257

258+
impl RowFn for NullarySeven {
259+
type Options = EmptyOptions;
260+
261+
const ARG_NAMES: &'static [&'static str] = &[];
262+
const INFALLIBLE: bool = true;
263+
264+
fn id(&self) -> ScalarFnId {
265+
static ID: CachedId = CachedId::new("test.nullary_seven");
266+
*ID
267+
}
268+
269+
fn dispatch<V: RowVisitor<Self::Options>>(
270+
&self,
271+
_options: &Self::Options,
272+
_args: &[DType],
273+
visitor: V,
274+
) -> VortexResult<V::VisitResult> {
275+
visitor.visit::<(), i64>(|()| 7)
276+
}
277+
}
278+
234279
impl RowFn for IndexingRowFn {
235280
type Options = EmptyOptions;
236281

@@ -302,6 +347,20 @@ mod tests {
302347
assert_arity_error(error);
303348
}
304349

350+
#[rstest]
351+
#[case::empty(0)]
352+
#[case::nonempty(3)]
353+
fn test_execute_nullary_rows(#[case] row_count: usize) -> VortexResult<()> {
354+
let args = VecExecutionArgs::new(vec![], row_count);
355+
let mut ctx = array_session().create_execution_ctx();
356+
357+
let actual = execute_rows(&NullarySeven, &EmptyOptions, &args, &mut ctx)?;
358+
let expected = PrimitiveArray::from_iter(vec![7_i64; row_count]).into_array();
359+
360+
assert_arrays_eq!(&actual, &expected, &mut ctx);
361+
Ok(())
362+
}
363+
305364
#[test]
306365
fn test_execute_rejects_dispatch_that_changes_after_planning() -> VortexResult<()> {
307366
let function = ChangingDispatchRowFn {

0 commit comments

Comments
 (0)