Skip to content

Commit 95ff440

Browse files
connortsui20connortsui20
committed
Add physical parameters to OutputSink
Signed-off-by: Connor Tsui <[email protected]>
1 parent 3715f42 commit 95ff440

14 files changed

Lines changed: 449 additions & 129 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ where
116116
// vectorization. Check each divide immediately and stop at the first failure.
117117
// Dense execution leaves output uninitialized. Nullable branches fill placeholders only when
118118
// they need to skip invalid rows.
119-
visitor.visit_into::<(T, T), UninitElementSink<T>, _>(|(lhs, rhs), output| {
119+
visitor.visit_into::<(T, T), UninitElementSink<T>, _>((), |(lhs, rhs), output| {
120120
let (value, failed) = CheckedDiv::apply(lhs, rhs);
121121
if failed {
122122
return Err(numeric_error(<CheckedDiv as CheckedPrimitiveOp<T>>::ERROR));

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

Lines changed: 94 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ use rstest::rstest;
99
use vortex_buffer::BufferMut;
1010
use vortex_error::VortexResult;
1111
use vortex_error::vortex_bail;
12+
use vortex_error::vortex_ensure;
13+
use vortex_error::vortex_err;
1214
use vortex_session::registry::CachedId;
1315

1416
use super::finalize_kernel_output;
@@ -18,6 +20,7 @@ use crate::VortexSessionExecute;
1820
use crate::array_session;
1921
use crate::arrays::ConstantArray;
2022
use crate::arrays::ExtensionArray;
23+
use crate::arrays::FixedSizeListArray;
2124
use crate::arrays::PrimitiveArray;
2225
use crate::assert_arrays_eq;
2326
use crate::dtype::DType;
@@ -30,6 +33,8 @@ use crate::scalar::Scalar;
3033
use crate::scalar_fn::EmptyOptions;
3134
use crate::scalar_fn::ScalarFnId;
3235
use crate::scalar_fn::VecExecutionArgs;
36+
use crate::scalar_fn::unstable::row::FixedSizeListSink;
37+
use crate::scalar_fn::unstable::row::InitializedRow;
3338
use crate::scalar_fn::unstable::row::OutputElement;
3439
use crate::scalar_fn::unstable::row::OutputSink;
3540
use crate::scalar_fn::unstable::row::RowFn;
@@ -78,15 +83,16 @@ struct I64Sink(BufferMut<i64>);
7883
// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that
7984
// initialized slice. The `()` write token therefore proves no additional invariant.
8085
unsafe impl OutputSink for I64Sink {
86+
type Params = ();
8187
type Rows<'a> = &'a mut [i64];
8288
type Row<'a> = &'a mut i64;
8389
type WriteToken = ();
8490

85-
fn storage_dtype() -> DType {
91+
fn storage_dtype(_params: &Self::Params) -> DType {
8692
DType::from(i64::PTYPE)
8793
}
8894

89-
fn with_capacity(rows: usize) -> VortexResult<Self> {
95+
fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult<Self> {
9096
Ok(Self(BufferMut::zeroed(rows)))
9197
}
9298

@@ -104,6 +110,38 @@ unsafe impl OutputSink for I64Sink {
104110
}
105111
}
106112

113+
#[derive(Clone)]
114+
struct RepeatValue;
115+
116+
impl RowFn for RepeatValue {
117+
type Options = usize;
118+
119+
const ARG_NAMES: &'static [&'static str] = &["value"];
120+
const INFALLIBLE: bool = true;
121+
122+
fn id(&self) -> ScalarFnId {
123+
static ID: CachedId = CachedId::new("test.repeat_value");
124+
*ID
125+
}
126+
127+
fn dispatch<V: RowVisitor>(
128+
&self,
129+
width: &Self::Options,
130+
_args: &[DType],
131+
visitor: V,
132+
) -> VortexResult<V::VisitResult> {
133+
vortex_ensure!(
134+
u32::try_from(*width).is_ok(),
135+
InvalidArgument:
136+
"test.repeat_value width must fit in the fixed-size-list u32 list size, got {width}",
137+
);
138+
139+
visitor.visit_into::<(i64,), FixedSizeListSink<i64>, _>(*width, |(value,), row| {
140+
InitializedRow::fill(row, |_| value)
141+
})
142+
}
143+
}
144+
107145
impl RowFn for DeferredAdd {
108146
type Options = EmptyOptions;
109147

@@ -156,7 +194,7 @@ impl RowFn for ValidOnlyIdentity {
156194
_args: &[DType],
157195
visitor: V,
158196
) -> VortexResult<V::VisitResult> {
159-
visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| {
197+
visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>((), |(value,), output| {
160198
*output = value;
161199
Ok(())
162200
})
@@ -184,6 +222,58 @@ impl RowFn for InvalidKernelOutput {
184222
}
185223
}
186224

225+
#[rstest]
226+
#[case::dense_width_two(
227+
vec![1_i64, 2],
228+
Validity::NonNullable,
229+
2,
230+
vec![1_i64, 1, 2, 2],
231+
)]
232+
#[case::dense_width_four(
233+
vec![3_i64, 4],
234+
Validity::NonNullable,
235+
4,
236+
vec![3_i64, 3, 3, 3, 4, 4, 4, 4],
237+
)]
238+
#[case::empty(vec![], Validity::NonNullable, 3, vec![])]
239+
#[case::zero_width(vec![3_i64, 4], Validity::NonNullable, 0, vec![])]
240+
#[case::all_null(
241+
vec![5_i64, 6],
242+
Validity::AllInvalid,
243+
2,
244+
vec![0_i64, 0, 0, 0],
245+
)]
246+
#[case::partially_valid(
247+
vec![7_i64, 8, 9],
248+
Validity::from_iter([true, false, true]),
249+
3,
250+
vec![7_i64, 7, 7, 0, 0, 0, 9, 9, 9],
251+
)]
252+
fn test_fixed_size_list_sink_uses_runtime_width(
253+
#[case] input_values: Vec<i64>,
254+
#[case] validity: Validity,
255+
#[case] width: usize,
256+
#[case] expected_elements: Vec<i64>,
257+
) -> VortexResult<()> {
258+
let row_count = input_values.len();
259+
let input = PrimitiveArray::new(input_values, validity.clone()).into_array();
260+
let args = VecExecutionArgs::new(vec![input], row_count);
261+
let mut ctx = array_session().create_execution_ctx();
262+
263+
let actual = execute_rows(&RepeatValue, &width, &args, &mut ctx)?;
264+
let expected = FixedSizeListArray::new(
265+
PrimitiveArray::from_iter(expected_elements).into_array(),
266+
u32::try_from(width)
267+
.map_err(|_| vortex_err!(InvalidArgument: "test width must fit in u32, got {width}"))?,
268+
validity,
269+
row_count,
270+
)
271+
.into_array();
272+
273+
assert_arrays_eq!(&actual, &expected, &mut ctx);
274+
Ok(())
275+
}
276+
187277
#[test]
188278
fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<()> {
189279
static ID: CachedId = CachedId::new("test.finalize_kernel_output");
@@ -403,7 +493,7 @@ impl RowFn for DeclaredSinkOutput {
403493
) -> VortexResult<V::VisitResult> {
404494
visitor
405495
.with_output_dtype(timestamp_dtype(Nullability::NonNullable))
406-
.visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| {
496+
.visit_into::<(i64,), I64Sink, VortexResult<()>>((), |(value,), output| {
407497
*output = value;
408498
Ok(())
409499
})

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

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use crate::scalar_fn::unstable::row::ViewLen;
2828
/// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant.
2929
pub(crate) fn execute_sink<Args, Prepared, Sink, ApplyResult>(
3030
args: &dyn ExecutionArgs,
31+
params: &Sink::Params,
3132
ctx: &mut ExecutionCtx,
3233
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
3334
apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult,
@@ -43,7 +44,7 @@ where
4344
let const_values = Args::const_values(&columns);
4445
let prepared = prepare(const_values);
4546

46-
let mut sink = Sink::with_capacity(row_count)?;
47+
let mut sink = Sink::with_capacity(row_count, params)?;
4748

4849
// Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink.
4950
{
@@ -100,6 +101,7 @@ where
100101
pub(crate) fn execute_sink_valid_rows<Args, Prepared, Sink, ApplyResult>(
101102
args: &dyn ExecutionArgs,
102103
valid: &MaskValuesRef,
104+
params: &Sink::Params,
103105
ctx: &mut ExecutionCtx,
104106
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
105107
apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult,
@@ -115,7 +117,7 @@ where
115117
valid_rows,
116118
row_count,
117119
mut sink,
118-
}) = setup_sink_valid_rows::<Args, Sink>(args, valid, ctx)?
120+
}) = setup_sink_valid_rows::<Args, Sink>(args, valid, params, ctx)?
119121
else {
120122
return Ok(None);
121123
};
@@ -205,6 +207,7 @@ where
205207
fn setup_sink_valid_rows<'valid, Args, Sink>(
206208
args: &dyn ExecutionArgs,
207209
valid: &'valid MaskValuesRef,
210+
params: &Sink::Params,
208211
ctx: &mut ExecutionCtx,
209212
) -> VortexResult<Option<ValidRowsSetup<'valid, Args, Sink>>>
210213
where
@@ -227,7 +230,7 @@ where
227230
// Keep allocation before the validity and length checks. With multiple CGUs and no LTO,
228231
// moving it later inlines `Args::get` into every sparse callback, duplicating its bounds
229232
// checks.
230-
let sink = Sink::with_capacity(row_count)?;
233+
let sink = Sink::with_capacity(row_count, params)?;
231234

232235
let valid_rows = valid.bit_buffer();
233236
vortex_ensure_eq!(
@@ -272,15 +275,16 @@ mod tests {
272275
// SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or
273276
// `finish` through the executor. The row-initialization requirements are therefore vacuous.
274277
unsafe impl OutputSink for NonSkippingSink {
278+
type Params = ();
275279
type Rows<'a> = ();
276280
type Row<'a> = ();
277281
type WriteToken = ();
278282

279-
fn storage_dtype() -> DType {
283+
fn storage_dtype(_params: &Self::Params) -> DType {
280284
DType::from(i64::PTYPE)
281285
}
282286

283-
fn with_capacity(_rows: usize) -> VortexResult<Self> {
287+
fn with_capacity(_rows: usize, _params: &Self::Params) -> VortexResult<Self> {
284288
Err(vortex_err!(
285289
"a non-skipping sink must decline before allocation"
286290
))
@@ -300,6 +304,7 @@ mod tests {
300304
// post-initialization length check. If execution incorrectly continues, safe indexing in
301305
// `row_unchecked` panics instead of accessing invalid memory.
302306
unsafe impl OutputSink for ShrinkingSink {
307+
type Params = ();
303308
type Rows<'a> = &'a mut Vec<i64>;
304309
type Row<'a> = &'a mut i64;
305310
type WriteToken = ();
@@ -310,11 +315,11 @@ mod tests {
310315
})
311316
}
312317

313-
fn storage_dtype() -> DType {
318+
fn storage_dtype(_params: &Self::Params) -> DType {
314319
DType::from(i64::PTYPE)
315320
}
316321

317-
fn with_capacity(rows: usize) -> VortexResult<Self> {
322+
fn with_capacity(rows: usize, _params: &Self::Params) -> VortexResult<Self> {
318323
Ok(Self(vec![0; rows]))
319324
}
320325

@@ -343,6 +348,7 @@ mod tests {
343348
let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>(
344349
&args,
345350
&valid,
351+
&(),
346352
&mut ctx,
347353
|_| (),
348354
|_, _, _| (),
@@ -365,6 +371,7 @@ mod tests {
365371
let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, ()>(
366372
&args,
367373
&valid,
374+
&(),
368375
&mut ctx,
369376
|_| (),
370377
|_, (value,), output| {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,10 @@ pub use row_fn::RowFn;
3131
mod types;
3232
pub use types::ElementTuple;
3333
pub use types::FailureEvidence;
34+
pub use types::FixedSizeListSink;
3435
pub use types::IndexedElementTuple;
3536
pub use types::InitializedElement;
37+
pub use types::InitializedRow;
3638
pub use types::InputElement;
3739
pub use types::OutputElement;
3840
pub use types::OutputSink;

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,9 @@ pub use result::FailureEvidence;
2020
pub use result::SinkResult;
2121

2222
mod sink;
23+
pub use sink::FixedSizeListSink;
2324
pub use sink::InitializedElement;
25+
pub use sink::InitializedRow;
2426
pub use sink::OutputSink;
2527
pub use sink::UninitElementSink;
2628

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,15 @@
44
//! Return types for row closures.
55
//!
66
//! [`FailureEvidence`] represents deferred failures from owned row closures. [`SinkResult`] lets
7-
//! the executor handle initialized sinks and sinks that require an [`InitializedElement`] token,
8-
//! with either infallible or immediate-error callbacks.
7+
//! the executor handle initialized sinks and sinks that require an [`InitializedElement`] or
8+
//! [`InitializedRow`] token, with either infallible or immediate-error callbacks.
99
1010
use std::ops::BitOrAssign;
1111

1212
use vortex_error::VortexResult;
1313

1414
use super::InitializedElement;
15+
use super::InitializedRow;
1516

1617
/// Compact failure evidence that can be OR-reduced across rows.
1718
///
@@ -57,6 +58,17 @@ impl SinkResult for InitializedElement {
5758
}
5859
}
5960

61+
impl private::Sealed for InitializedRow {}
62+
63+
impl SinkResult for InitializedRow {
64+
type WriteToken = InitializedRow;
65+
const INFALLIBLE: bool = true;
66+
67+
fn into_result(self) -> VortexResult<()> {
68+
Ok(())
69+
}
70+
}
71+
6072
impl private::Sealed for VortexResult<()> {}
6173

6274
impl SinkResult for VortexResult<()> {

0 commit comments

Comments
 (0)