Batch fragmented CASE branches and push selections before scalar execution

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-array/benches/expr/case_when_bench.rs

@@ -14,7 +14,13 @@ use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
 use vortex_array::array_session;
 use vortex_array::arrays::BoolArray;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::ScalarFnArray;
 use vortex_array::arrays::StructArray;
+use vortex_array::arrays::VarBinViewArray;
+use vortex_array::builtins::ArrayBuiltins;
+use vortex_array::dtype::DecimalDType;
 use vortex_array::expr::case_when;
 use vortex_array::expr::case_when_no_else;
 use vortex_array::expr::eq;
@@ -24,6 +30,12 @@ use vortex_array::expr::lit;
 use vortex_array::expr::lt;
 use vortex_array::expr::nested_case_when;
 use vortex_array::expr::root;
+use vortex_array::scalar::Scalar;
+use vortex_array::scalar_fn::ScalarFnVTableExt;
+use vortex_array::scalar_fn::fns::case_when::CaseWhen;
+use vortex_array::scalar_fn::fns::case_when::CaseWhenOptions;
+use vortex_array::scalar_fn::fns::operators::Operator;
+use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_session::VortexSession;
 
@@ -37,6 +49,53 @@ fn main() {
     divan::main();
 }
 
+#[divan::bench(args = [1, 4, 16, 64, 256, 4096])]
+fn case_when_decimal_product(bencher: Bencher, run_length: usize) {
+    let len = 65_536;
+    let lhs = DecimalArray::new(
+        Buffer::from_iter((0..len).map(|i| (i % 10_000) as i64)),
+        DecimalDType::new(15, 2),
+        Validity::NonNullable,
+    )
+    .into_array();
+    let rhs = DecimalArray::new(
+        Buffer::from_iter((0..len).map(|i| (i % 100) as i64)),
+        DecimalDType::new(15, 2),
+        Validity::NonNullable,
+    )
+    .into_array();
+    let product = lhs.binary(rhs, Operator::Mul).unwrap();
+    let otherwise = ConstantArray::new(Scalar::zero_value(product.dtype()), len).into_array();
+    bench_branches(bencher, run_length, product, otherwise);
+}
+
+#[divan::bench(args = [1, 4, 16, 64, 256, 4096])]
+fn case_when_string_values(bencher: Bencher, run_length: usize) {
+    let len = 65_536;
+    let lhs = VarBinViewArray::from_iter_str((0..len).map(|i| format!("long-left-value-{i}")))
+        .into_array();
+    let rhs = VarBinViewArray::from_iter_str((0..len).map(|i| format!("long-right-value-{i}")))
+        .into_array();
+    bench_branches(bencher, run_length, lhs, rhs);
+}
+
+fn bench_branches(bencher: Bencher, run_length: usize, lhs: ArrayRef, rhs: ArrayRef) {
+    let mask = BoolArray::from_iter((0..lhs.len()).map(|i| (i / run_length).is_multiple_of(2)))
+        .into_array();
+    let array = ScalarFnArray::try_new(
+        CaseWhen.bind(CaseWhenOptions {
+            num_when_then_pairs: 1,
+            has_else: true,
+        }),
+        vec![mask, lhs, rhs],
+    )
+    .unwrap()
+    .into_array();
+    bencher
+        .with_inputs(|| SESSION.create_execution_ctx())
+        .bench_refs(|ctx| array.clone().execute::<Canonical>(ctx).unwrap());
+}
+
 fn make_struct_array(size: usize) -> ArrayRef {
     let data: Buffer<i32> = (0..size as i32).collect();
     let field = data.into_array();
@@ -264,7 +323,7 @@ fn case_when_all_false(bencher: Bencher, size: usize) {
         });
 }
 
-/// Benchmark CASE WHEN cycling through 3 branches per row (triggers merge_row_by_row).
+/// Benchmark compact CASE WHEN assembly with three interleaved branches.
 /// Run length = 1; exercises branch 0, branch 1, and the else fallback at every 3rd row.
 #[divan::bench(args = [100, 400])]
 fn case_when_fragmented(bencher: Bencher, size: usize) {

vortex-array/src/arrays/scalar_fn/rules.rs

@@ -117,6 +117,7 @@ impl ArrayParentReduceRule<ScalarFn> for ScalarFilterPushdownRule {
             prepare_mask_for_reuse(values, nchildren);
         }
 
+        // Selection precedes value evaluation, including errors in unselected rows.
         let new_children: Vec<_> = child
             .iter_children()
             .map(|c| match c.as_opt::<Constant>() {