perf: batch fragmented CASE branch execution

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/filter/rules.rs

@@ -67,7 +67,14 @@ impl ArrayReduceRule<Filter> for TrivialFilterRule {
         match array.filter_mask() {
             Mask::AllTrue(_) => Ok(Some(array.child().clone())),
             Mask::AllFalse(_) => Ok(Some(Canonical::empty(array.dtype()).into_array())),
-            Mask::Values(_) => Ok(None),
+            mask @ Mask::Values(_) => {
+                if let (Some(first), Some(last)) = (mask.first(), mask.last())
+                    && last - first + 1 == mask.true_count()
+                {
+                    return Ok(Some(array.child().slice(first..last + 1)?));
+                }
+                Ok(None)
+            }
         }
     }
 }

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

@@ -27,7 +27,7 @@ pub(super) const RULES: ReduceRuleSet<ScalarFn> =
     ReduceRuleSet::new(&[&ScalarFnPackToStructRule, &ScalarFnAbstractReduceRule]);
 
 pub(super) const PARENT_RULES: ParentRuleSet<ScalarFn> = ParentRuleSet::new(&[
-    ParentRuleSet::lift(&ScalarFnUnaryFilterPushDownRule),
+    ParentRuleSet::lift(&ScalarFnFilterPushDownRule),
     ParentRuleSet::lift(&ScalarFnSliceReduceRule),
 ]);
 
@@ -95,9 +95,9 @@ impl ArrayReduceRule<ScalarFn> for ScalarFnAbstractReduceRule {
 }
 
 #[derive(Debug)]
-struct ScalarFnUnaryFilterPushDownRule;
+struct ScalarFnFilterPushDownRule;
 
-impl ArrayParentReduceRule<ScalarFn> for ScalarFnUnaryFilterPushDownRule {
+impl ArrayParentReduceRule<ScalarFn> for ScalarFnFilterPushDownRule {
     type Parent = Filter;
 
     fn reduce_parent(
@@ -106,31 +106,21 @@ impl ArrayParentReduceRule<ScalarFn> for ScalarFnUnaryFilterPushDownRule {
         parent: ArrayView<'_, Filter>,
         _child_idx: usize,
     ) -> VortexResult<Option<ArrayRef>> {
-        // If we only have one non-constant child, then it is _always_ cheaper to push down the
-        // filter over the children of the scalar function array.
-        if child
+        // Selection precedes value evaluation, including errors in unselected rows.
+        let new_children: Vec<_> = child
             .iter_children()
-            .filter(|c| !c.is::<Constant>())
-            .count()
-            == 1
-        {
-            let new_children: Vec<_> = child
-                .iter_children()
-                .map(|c| match c.as_opt::<Constant>() {
-                    Some(array) => {
-                        Ok(ConstantArray::new(array.scalar().clone(), parent.len()).into_array())
-                    }
-                    None => c.filter(parent.filter_mask().clone()),
-                })
-                .try_collect()?;
-
-            let new_array =
-                ScalarFnArray::try_new(child.scalar_fn().clone(), new_children)?.into_array();
-
-            return Ok(Some(new_array));
-        }
+            .map(|c| match c.as_opt::<Constant>() {
+                Some(array) => {
+                    Ok(ConstantArray::new(array.scalar().clone(), parent.len()).into_array())
+                }
+                None => c.filter(parent.filter_mask().clone()),
+            })
+            .try_collect()?;
 
-        Ok(None)
+        Ok(Some(
+            ScalarFnArray::try_new_with_len(child.scalar_fn().clone(), new_children, parent.len())?
+                .into_array(),
+        ))
     }
 }