Skip to content

Commit 1312f0a

Browse files
connortsui20connortsui20
committed
Use packed Boolean collection for RowFn dense retry
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent a01de34 commit 1312f0a

4 files changed

Lines changed: 173 additions & 3 deletions

File tree

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@ use crate::arrays::ConstantArray;
2525
use crate::arrays::ExtensionArray;
2626
use crate::arrays::FixedSizeListArray;
2727
use crate::arrays::PrimitiveArray;
28+
use crate::arrays::VarBinViewArray;
29+
use crate::arrays::varbinview::BinaryView;
2830
use crate::assert_arrays_eq;
31+
use crate::buffer::BufferHandle;
2932
use crate::dtype::DType;
3033
use crate::dtype::NativePType;
3134
use crate::dtype::Nullability;
@@ -43,6 +46,8 @@ use crate::scalar_fn::unstable::row::OutputElement;
4346
use crate::scalar_fn::unstable::row::OutputSink;
4447
use crate::scalar_fn::unstable::row::RowFn;
4548
use crate::scalar_fn::unstable::row::RowVisitor;
49+
use crate::scalar_fn::unstable::row::Utf8Column;
50+
use crate::scalar_fn::unstable::row::execute::execute_bool_dense_attempt;
4651
use crate::scalar_fn::unstable::row::execute_rows;
4752
use crate::scalar_fn::unstable::row::row_fn_return_dtype;
4853
use crate::validity::Validity;
@@ -83,6 +88,20 @@ struct PackedGreaterThan<const MULTIVERSIONED: bool>;
8388
#[derive(Clone)]
8489
struct DeferredGreaterThan<const MULTIVERSIONED: bool>;
8590

91+
#[derive(Clone, Copy)]
92+
enum BooleanInputShape {
93+
Columns,
94+
ConstantLhs,
95+
ConstantRhs,
96+
}
97+
98+
#[derive(Clone, Copy)]
99+
enum BooleanValidity {
100+
AllValid,
101+
AllNull,
102+
Partial,
103+
}
104+
86105
#[derive(Clone)]
87106
struct ValidOnlyPositive;
88107

@@ -866,6 +885,100 @@ fn test_deferred_bool_output_handles_partial_constants() -> VortexResult<()> {
866885
Ok(())
867886
}
868887

888+
#[rstest]
889+
#[case::empty(0)]
890+
#[case::one(1)]
891+
#[case::word_tail(63)]
892+
#[case::word(64)]
893+
#[case::word_remainder(65)]
894+
#[case::multiword(130)]
895+
fn test_deferred_bool_sliced_validity_and_constants(
896+
#[case] len: usize,
897+
#[values(false, true)] multiversioned: bool,
898+
#[values(
899+
BooleanInputShape::Columns,
900+
BooleanInputShape::ConstantLhs,
901+
BooleanInputShape::ConstantRhs
902+
)]
903+
input_shape: BooleanInputShape,
904+
#[values(
905+
BooleanValidity::AllValid,
906+
BooleanValidity::AllNull,
907+
BooleanValidity::Partial
908+
)]
909+
validity_pattern: BooleanValidity,
910+
) -> VortexResult<()> {
911+
let valid = |index: usize| match validity_pattern {
912+
BooleanValidity::AllValid => true,
913+
BooleanValidity::AllNull => false,
914+
BooleanValidity::Partial => !index.is_multiple_of(3),
915+
};
916+
let validity = Validity::Array(BoolArray::from_iter((0..len + 2).map(valid)).into_array());
917+
let varying = PrimitiveArray::new(
918+
(0..len + 2)
919+
.map(|index| if valid(index) { index as i64 } else { i64::MIN })
920+
.collect::<Vec<_>>(),
921+
validity,
922+
)
923+
.into_array()
924+
.slice(1..len + 1)?;
925+
let constant = ConstantArray::new(32_i64, len).into_array();
926+
let other = PrimitiveArray::from_iter(std::iter::repeat_n(32_i64, len)).into_array();
927+
let (lhs, rhs) = match input_shape {
928+
BooleanInputShape::Columns => (varying, other),
929+
BooleanInputShape::ConstantLhs => (constant, varying),
930+
BooleanInputShape::ConstantRhs => (varying, constant),
931+
};
932+
let args = VecExecutionArgs::new(vec![lhs, rhs], len);
933+
let mut ctx = array_session().create_execution_ctx();
934+
let actual = if multiversioned {
935+
execute_rows(&DeferredGreaterThan::<true>, &EmptyOptions, &args, &mut ctx)?
936+
} else {
937+
execute_rows(
938+
&DeferredGreaterThan::<false>,
939+
&EmptyOptions,
940+
&args,
941+
&mut ctx,
942+
)?
943+
};
944+
let expected = BoolArray::from_iter((1..len + 1).map(|index| {
945+
valid(index).then_some(if matches!(input_shape, BooleanInputShape::ConstantLhs) {
946+
32 > index as i64
947+
} else {
948+
index as i64 > 32
949+
})
950+
}))
951+
.into_array();
952+
assert_arrays_eq!(&actual, &expected, &mut ctx);
953+
Ok(())
954+
}
955+
956+
#[test]
957+
fn test_deferred_bool_decode_error_is_terminal() {
958+
let views = Buffer::from(vec![BinaryView::make_view(&[0xff], 0, 0)]);
959+
let input = VarBinViewArray::new_handle(
960+
BufferHandle::new_host(views.into_byte_buffer()),
961+
Default::default(),
962+
DType::Utf8(Nullability::NonNullable),
963+
Validity::NonNullable,
964+
)
965+
.into_array();
966+
let args = VecExecutionArgs::new(vec![input], 1);
967+
let mut ctx = array_session().create_execution_ctx();
968+
969+
let result = execute_bool_dense_attempt::<(Utf8Column,), (), bool, true>(
970+
&args,
971+
&mut ctx,
972+
|_| (),
973+
|&(), (value,)| (!value.is_empty(), false),
974+
|_| Ok(()),
975+
);
976+
assert!(
977+
result.is_err(),
978+
"a decode error must not become a deferred row error"
979+
);
980+
}
981+
869982
#[test]
870983
fn test_deferred_bool_output_reports_valid_row_failure() -> VortexResult<()> {
871984
let lhs = PrimitiveArray::from_iter([1_i64, i64::MIN, -1]).into_array();

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ pub(super) use owned::execute_owned_infallible_valid_rows;
1515
pub(super) use owned::execute_owned_valid_rows;
1616

1717
mod packed_bool;
18+
pub(super) use packed_bool::execute_bool_dense_attempt;
1819
pub(super) use packed_bool::execute_owned_bool;
1920
pub(super) use packed_bool::execute_owned_infallible_bool;
2021

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

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use vortex_compute::lane_kernels::IndexedSource;
1313
use vortex_error::VortexResult;
1414
use vortex_error::vortex_bail;
1515

16+
use super::DenseAttempt;
1617
use crate::ArrayRef;
1718
use crate::ExecutionCtx;
1819
use crate::IntoArray;
@@ -61,6 +62,30 @@ pub(crate) fn execute_owned_bool<Args, Prepared, Fail, const MULTIVERSIONED: boo
6162
apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail),
6263
finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
6364
) -> VortexResult<ArrayRef>
65+
where
66+
Args: IndexedElementTuple,
67+
Fail: FailureEvidence,
68+
{
69+
match execute_bool_dense_attempt::<Args, Prepared, Fail, MULTIVERSIONED>(
70+
args,
71+
ctx,
72+
prepare,
73+
apply,
74+
finish_failure,
75+
)? {
76+
DenseAttempt::Values(values) => Ok(values),
77+
DenseAttempt::DeferredError(error) => Err(error),
78+
}
79+
}
80+
81+
/// Pack a dense attempt, keeping row failure separate from terminal decoding errors.
82+
pub(crate) fn execute_bool_dense_attempt<Args, Prepared, Fail, const MULTIVERSIONED: bool>(
83+
args: &dyn ExecutionArgs,
84+
ctx: &mut ExecutionCtx,
85+
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
86+
apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail),
87+
finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
88+
) -> VortexResult<DenseAttempt>
6489
where
6590
Args: IndexedElementTuple,
6691
Fail: FailureEvidence,
@@ -96,7 +121,10 @@ where
96121
BitBuffer::collect_bool(row_count, collect)
97122
};
98123

99-
finish_failure(failure)?;
100-
101-
Ok(BoolArray::new(values, Validity::NonNullable).into_array())
124+
match finish_failure(failure) {
125+
Ok(()) => Ok(DenseAttempt::Values(
126+
BoolArray::new(values, Validity::NonNullable).into_array(),
127+
)),
128+
Err(error) => Ok(DenseAttempt::DeferredError(error)),
129+
}
102130
}

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

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ use crate::scalar_fn::unstable::row::RowFn;
3131
use crate::scalar_fn::unstable::row::SinkResult;
3232
use crate::scalar_fn::unstable::row::batch::BorrowedRowFnArgs;
3333
use crate::scalar_fn::unstable::row::execute::DenseAttempt;
34+
use crate::scalar_fn::unstable::row::execute::execute_bool_dense_attempt;
3435
use crate::scalar_fn::unstable::row::execute::execute_owned_dense_attempt;
3536
use crate::scalar_fn::unstable::row::execute::execute_owned_infallible;
3637
use crate::scalar_fn::unstable::row::execute::execute_sink;
@@ -152,4 +153,31 @@ impl<F: RowFn> RowVisitor for ExecuteDenseWithRetry<'_, '_, '_, F> {
152153
finish_failure,
153154
)
154155
}
156+
157+
fn visit_prepared_deferred_bool<Args, Prepared, Fail, const MULTIVERSIONED: bool>(
158+
self,
159+
prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
160+
apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail),
161+
finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
162+
) -> VortexResult<Self::VisitResult>
163+
where
164+
Args: IndexedElementTuple,
165+
Fail: FailureEvidence,
166+
{
167+
const { assert_deferred_visit_contract::<F, Args, bool, Fail>() };
168+
let visited = BatchPlan::new(
169+
validate_owned_visit::<Args, bool>(self.args.dtypes())?,
170+
self.output_dtype,
171+
RowPolicy::for_deferred_output::<Args>(),
172+
)?;
173+
self.args.plan().ensure_reproduced_by(&visited)?;
174+
175+
execute_bool_dense_attempt::<Args, Prepared, Fail, MULTIVERSIONED>(
176+
self.args,
177+
self.ctx,
178+
prepare,
179+
apply,
180+
finish_failure,
181+
)
182+
}
155183
}

0 commit comments

Comments
 (0)