Skip to content

Commit 68cdf00

Browse files
connortsui20connortsui20
authored
Carry partial RowFn validity with MaskValues (#9534)
## Summary - Tracking Issue: #9128 Row batch execution resolves validity once, but the valid-row executors repeat the check that the materialized mask is partial. Array-backed validity stays lazy on optimistic paths because resolving it can execute and linearly scan the full mask for the uncommon all-valid and all-null cases. ## Changes Matches the materialized `Mask` once and passes `MaskValues` through valid-row visitors, owned execution, and sink execution. Marks the all-null dense path as unreachable and documents the cost boundary on `Validity::execute_mask`. --------- Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent acd8bb4 commit 68cdf00

9 files changed

Lines changed: 78 additions & 99 deletions

File tree

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

Lines changed: 16 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use vortex_error::VortexError;
55
use vortex_error::VortexResult;
66
use vortex_error::vortex_panic;
77
use vortex_mask::Mask;
8+
use vortex_mask::MaskValuesRef;
89

910
use super::super::RowFnExecutionArgs;
1011
use super::super::args::BorrowedRowFnArgs;
@@ -39,7 +40,7 @@ impl RowFnExecutionArgs {
3940
) -> VortexResult<DenseAttempt>,
4041
try_valid_rows: impl FnOnce(
4142
BorrowedRowFnArgs<'_>,
42-
&Mask,
43+
MaskValuesRef,
4344
&mut ExecutionCtx,
4445
) -> VortexResult<Option<ArrayRef>>,
4546
ctx: &mut ExecutionCtx,
@@ -60,30 +61,26 @@ impl RowFnExecutionArgs {
6061
deferred_error: VortexError,
6162
try_valid_rows: impl FnOnce(
6263
BorrowedRowFnArgs<'_>,
63-
&Mask,
64+
MaskValuesRef,
6465
&mut ExecutionCtx,
6566
) -> VortexResult<Option<ArrayRef>>,
6667
ctx: &mut ExecutionCtx,
6768
) -> VortexResult<ArrayRef> {
68-
let valid_rows = self.validity.execute_mask(self.row_count, ctx)?;
69-
70-
// An array-backed validity can materialize to all valid even though the cheap checks in
71-
// `RowFnExecutionArgs::execute` could not prove that. The deferred error therefore came
72-
// from an observable row and remains terminal. Check all-true before all-false because an
73-
// empty mask is both.
74-
if valid_rows.all_true() {
75-
return Err(deferred_error);
76-
}
77-
78-
if valid_rows.all_false() {
79-
return Ok(self.all_null());
80-
}
69+
let valid_rows = match self.validity.execute_mask(self.row_count, ctx)? {
70+
// An array-backed validity can materialize to all valid even though the cheap checks
71+
// in `RowFnExecutionArgs::execute` could not prove that. The deferred error therefore
72+
// came from an observable row and remains terminal. An empty mask is both all-valid
73+
// and all-null, so preserve the all-valid behavior.
74+
Mask::AllTrue(_) | Mask::AllFalse(0) => return Err(deferred_error),
75+
Mask::AllFalse(_) => return Ok(self.all_null()),
76+
Mask::Values(valid_rows) => valid_rows,
77+
};
8178

8279
// Reduced evidence does not identify which rows failed. Discard the dense error before
8380
// retrying only observable rows.
8481
drop(deferred_error);
8582

86-
if let Some(result) = self.try_execute_valid_rows(try_valid_rows, &valid_rows, ctx)? {
83+
if let Some(result) = self.try_execute_valid_rows(try_valid_rows, valid_rows, ctx)? {
8784
return Ok(result);
8885
}
8986

@@ -105,8 +102,9 @@ impl RowFnExecutionArgs {
105102
self.finalize_output(values, self.row_count)
106103
}
107104
Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count),
108-
// Handled by the guard in `RowFnExecutionArgs::execute`, before the kernel ran.
109-
Validity::AllInvalid => Ok(self.all_null()),
105+
Validity::AllInvalid => {
106+
unreachable!("all-invalid validity is handled before dense row execution")
107+
}
110108
}
111109
}
112110
}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
//! valid-only execution.
88
99
use vortex_error::VortexResult;
10-
use vortex_mask::Mask;
10+
use vortex_mask::MaskValuesRef;
1111

1212
use super::RowFnExecutionArgs;
1313
use super::RowPolicy;
@@ -40,7 +40,7 @@ impl RowFnExecutionArgs {
4040
) -> VortexResult<DenseAttempt>,
4141
try_valid_rows: impl FnOnce(
4242
BorrowedRowFnArgs<'_>,
43-
&Mask,
43+
MaskValuesRef,
4444
&mut ExecutionCtx,
4545
) -> VortexResult<Option<ArrayRef>>,
4646
ctx: &mut ExecutionCtx,
@@ -67,8 +67,8 @@ impl RowFnExecutionArgs {
6767
return self.execute_all_constant(kernel, ctx);
6868
}
6969

70-
// A known all-valid batch does not need to materialize validity, even when its row policy
71-
// only permits valid rows.
70+
// Do not resolve array-backed validity for the uncommon all-valid or all-null cases here.
71+
// That can execute and scan the full mask; each policy resolves it only when necessary.
7272
if self.validity.definitely_no_nulls() {
7373
return self.execute_dense(kernel, ctx);
7474
}

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

Lines changed: 19 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -4,24 +4,14 @@
44
use vortex_error::VortexResult;
55
use vortex_error::vortex_panic;
66
use vortex_mask::Mask;
7+
use vortex_mask::MaskValuesRef;
78

89
use super::super::RowFnExecutionArgs;
910
use super::super::args::BorrowedRowFnArgs;
1011
use crate::ArrayRef;
1112
use crate::ExecutionCtx;
1213
use crate::IntoArray;
13-
use crate::arrays::BoolArray;
1414
use crate::builtins::ArrayBuiltins;
15-
use crate::validity::Validity;
16-
17-
/// The result of resolving batch validity.
18-
enum ResolvedValidity {
19-
/// The output for an all-valid or all-null batch.
20-
Output(ArrayRef),
21-
22-
/// A mask with both valid and invalid rows.
23-
PartiallyValid(Mask),
24-
}
2515

2616
impl RowFnExecutionArgs {
2717
/// Resolve validity, then execute valid rows over the original inputs.
@@ -35,17 +25,26 @@ impl RowFnExecutionArgs {
3525
kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
3626
try_valid_rows: impl FnOnce(
3727
BorrowedRowFnArgs<'_>,
38-
&Mask,
28+
MaskValuesRef,
3929
&mut ExecutionCtx,
4030
) -> VortexResult<Option<ArrayRef>>,
4131
ctx: &mut ExecutionCtx,
4232
) -> VortexResult<ArrayRef> {
43-
let valid = match self.resolve_validity(&kernel, ctx)? {
44-
ResolvedValidity::Output(output) => return Ok(output),
45-
ResolvedValidity::PartiallyValid(valid) => valid,
33+
let validity = self.validity.clone().execute_mask(self.row_count, ctx)?;
34+
35+
let valid_rows = match validity {
36+
// An empty mask is both all-valid and all-null. Preserve the all-valid behavior.
37+
Mask::AllTrue(_) | Mask::AllFalse(0) => {
38+
let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?;
39+
let values = self.validate_kernel_output(values, self.row_count, ctx)?;
40+
41+
return self.finalize_output(values, self.row_count);
42+
}
43+
Mask::AllFalse(_) => return Ok(self.all_null()),
44+
Mask::Values(valid_rows) => valid_rows,
4645
};
4746

48-
if let Some(result) = self.try_execute_valid_rows(try_valid_rows, &valid, ctx)? {
47+
if let Some(result) = self.try_execute_valid_rows(try_valid_rows, valid_rows, ctx)? {
4948
return Ok(result);
5049
}
5150

@@ -55,54 +54,28 @@ impl RowFnExecutionArgs {
5554
)
5655
}
5756

58-
/// Materialize validity and handle all-valid or all-null batches.
59-
fn resolve_validity(
60-
&self,
61-
kernel: &impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult<ArrayRef>,
62-
ctx: &mut ExecutionCtx,
63-
) -> VortexResult<ResolvedValidity> {
64-
let valid = self.validity.clone().execute_mask(self.row_count, ctx)?;
65-
66-
// An array-backed validity can materialize to all valid even though the cheap checks in
67-
// `RowFnExecutionArgs::execute` could not prove that. Run the full-row kernel in that
68-
// case. Check all-true before all-false because an empty mask is both.
69-
if valid.all_true() {
70-
let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?;
71-
let values = self.validate_kernel_output(values, self.row_count, ctx)?;
72-
let values = self.finalize_output(values, self.row_count)?;
73-
74-
return Ok(ResolvedValidity::Output(values));
75-
}
76-
77-
if valid.all_false() {
78-
return Ok(ResolvedValidity::Output(self.all_null()));
79-
}
80-
81-
Ok(ResolvedValidity::PartiallyValid(valid))
82-
}
83-
8457
/// Try execution against the original inputs, then mask a returned full-length result.
8558
pub(super) fn try_execute_valid_rows(
8659
&self,
8760
try_valid_rows: impl FnOnce(
8861
BorrowedRowFnArgs<'_>,
89-
&Mask,
62+
MaskValuesRef,
9063
&mut ExecutionCtx,
9164
) -> VortexResult<Option<ArrayRef>>,
92-
valid: &Mask,
65+
valid: MaskValuesRef,
9366
ctx: &mut ExecutionCtx,
9467
) -> VortexResult<Option<ArrayRef>> {
9568
let Some(values) = try_valid_rows(
9669
self.execution_args(&self.inputs, self.row_count),
97-
valid,
70+
MaskValuesRef::clone(&valid),
9871
ctx,
9972
)?
10073
else {
10174
return Ok(None);
10275
};
10376
let values = self.validate_kernel_output(values, valid.len(), ctx)?;
10477

105-
let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array();
78+
let mask = valid.as_ref().into_array();
10679
self.finalize_output(values.mask(mask)?, valid.len())
10780
.map(Some)
10881
}

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

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,9 @@ use std::ops::BitOrAssign;
1111

1212
use vortex_compute::lane_kernels::IndexedSourceExt;
1313
use vortex_error::VortexResult;
14-
use vortex_error::vortex_bail;
1514
use vortex_error::vortex_ensure;
1615
use vortex_error::vortex_ensure_eq;
17-
use vortex_mask::AllOr;
18-
use vortex_mask::Mask;
16+
use vortex_mask::MaskValuesRef;
1917

2018
use crate::ArrayRef;
2119
use crate::ExecutionCtx;
@@ -56,7 +54,7 @@ where
5654
/// Decode nullable inputs, then store one output for each valid row from an infallible kernel.
5755
pub(crate) fn execute_owned_infallible_valid_rows<Args, Out, Prepared>(
5856
args: &dyn ExecutionArgs,
59-
valid: &Mask,
57+
valid: &MaskValuesRef,
6058
ctx: &mut ExecutionCtx,
6159
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
6260
apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out,
@@ -78,7 +76,7 @@ where
7876
/// Decode nullable inputs, then store outputs and combine failure evidence for valid rows.
7977
pub(crate) fn execute_owned_valid_rows<Args, Out, Prepared, Fail>(
8078
args: &dyn ExecutionArgs,
81-
valid: &Mask,
79+
valid: &MaskValuesRef,
8280
ctx: &mut ExecutionCtx,
8381
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
8482
apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail),
@@ -96,11 +94,7 @@ where
9694
};
9795

9896
let row_count = args.row_count();
99-
let AllOr::Some(valid_rows) = valid.bit_buffer() else {
100-
vortex_bail!(
101-
"execute_owned_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask"
102-
);
103-
};
97+
let valid_rows = valid.bit_buffer();
10498
vortex_ensure_eq!(
10599
valid_rows.len(),
106100
row_count,

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

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,7 @@ use vortex_buffer::BitBuffer;
1111
use vortex_error::VortexResult;
1212
use vortex_error::vortex_bail;
1313
use vortex_error::vortex_ensure_eq;
14-
use vortex_mask::AllOr;
15-
use vortex_mask::Mask;
14+
use vortex_mask::MaskValuesRef;
1615

1716
use crate::ArrayRef;
1817
use crate::ExecutionCtx;
@@ -102,7 +101,7 @@ where
102101
/// how to handle the decline.
103102
pub(crate) fn execute_sink_valid_rows<Args, Prepared, Sink, ApplyResult, Options>(
104103
args: &dyn ExecutionArgs,
105-
valid: &Mask,
104+
valid: &MaskValuesRef,
106105
ctx: &mut ExecutionCtx,
107106
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
108107
apply: impl Fn(&Prepared, Args::Elems<'_>, <Sink as OutputSink<Options>>::Row<'_>) -> ApplyResult,
@@ -209,7 +208,7 @@ where
209208
/// Resolve the capabilities, inputs, sink, and validity mask for skip-invalid execution.
210209
fn setup_sink_valid_rows<'valid, Args, Sink, Options>(
211210
args: &dyn ExecutionArgs,
212-
valid: &'valid Mask,
211+
valid: &'valid MaskValuesRef,
213212
ctx: &mut ExecutionCtx,
214213
) -> VortexResult<Option<ValidRowsSetup<'valid, Args, Sink, Options>>>
215214
where
@@ -235,12 +234,7 @@ where
235234
// checks.
236235
let sink = <Sink as OutputSink<Options>>::with_capacity(row_count)?;
237236

238-
// Batch execution resolves all-valid and all-null inputs before selecting this path.
239-
let AllOr::Some(valid_rows) = valid.bit_buffer() else {
240-
vortex_bail!(
241-
"execute_sink_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask"
242-
);
243-
};
237+
let valid_rows = valid.bit_buffer();
244238
vortex_ensure_eq!(
245239
valid_rows.len(),
246240
row_count,
@@ -347,7 +341,9 @@ mod tests {
347341
fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> {
348342
let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array();
349343
let args = VecExecutionArgs::new(vec![input], 2);
350-
let valid = Mask::from_iter([true, false]);
344+
let Mask::Values(valid) = Mask::from_iter([true, false]) else {
345+
vortex_bail!("the test validity must be partially valid");
346+
};
351347
let mut ctx = array_session().create_execution_ctx();
352348

353349
let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>(
@@ -367,7 +363,9 @@ mod tests {
367363
fn test_skip_invalid_sink_rechecks_rows_after_initialization() -> VortexResult<()> {
368364
let input = PrimitiveArray::from_iter([10_i64, 20]).into_array();
369365
let args = VecExecutionArgs::new(vec![input], 2);
370-
let valid = Mask::from_iter([false, true]);
366+
let Mask::Values(valid) = Mask::from_iter([false, true]) else {
367+
vortex_bail!("the test validity must be partially valid");
368+
};
371369
let mut ctx = array_session().create_execution_ctx();
372370

373371
let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>(

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

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
//! signature cannot execute over the original inputs.
99
1010
use vortex_error::VortexResult;
11-
use vortex_mask::Mask;
11+
use vortex_mask::MaskValuesRef;
1212

1313
use super::RowPolicy;
1414
use super::RowVisitor;
@@ -180,7 +180,7 @@ pub(crate) struct ExecuteValidRows<'args, 'ctx, F: RowFn> {
180180
policy: RowPolicy,
181181

182182
/// The conjoined validity, containing both valid and invalid rows.
183-
valid: &'args Mask,
183+
valid: MaskValuesRef,
184184

185185
/// The execution context used to decode the input columns.
186186
ctx: &'ctx mut ExecutionCtx,
@@ -193,7 +193,7 @@ impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> {
193193
options: &'args F::Options,
194194
output_dtype: &'args DType,
195195
policy: RowPolicy,
196-
valid: &'args Mask,
196+
valid: MaskValuesRef,
197197
ctx: &'ctx mut ExecutionCtx,
198198
) -> Self {
199199
Self {
@@ -231,7 +231,11 @@ impl<F: RowFn> RowVisitor<F::Options> for ExecuteValidRows<'_, '_, F> {
231231
)?;
232232

233233
execute_owned_infallible_valid_rows::<Args, Out, Prepared>(
234-
self.args, self.valid, self.ctx, prepare, apply,
234+
self.args,
235+
&self.valid,
236+
self.ctx,
237+
prepare,
238+
apply,
235239
)
236240
}
237241

@@ -258,7 +262,11 @@ impl<F: RowFn> RowVisitor<F::Options> for ExecuteValidRows<'_, '_, F> {
258262
)?;
259263

260264
execute_sink_valid_rows::<Args, Prepared, Sink, ApplyResult, F::Options>(
261-
self.args, self.valid, self.ctx, prepare, apply,
265+
self.args,
266+
&self.valid,
267+
self.ctx,
268+
prepare,
269+
apply,
262270
)
263271
}
264272

@@ -283,7 +291,7 @@ impl<F: RowFn> RowVisitor<F::Options> for ExecuteValidRows<'_, '_, F> {
283291

284292
execute_owned_valid_rows::<Args, Out, Prepared, Fail>(
285293
self.args,
286-
self.valid,
294+
&self.valid,
287295
self.ctx,
288296
prepare,
289297
apply,

0 commit comments

Comments
 (0)