Propagate ExecutionCtx through CASE WHEN binary merges (#8040)

Signed-off-by: "Dimitar Dimitrov" <dimitar@spiraldb.com>
Signed-off-by: Dimitar Dimitrov <dimitar@spiraldb.com>

Author: dimitarvdimitrov

vortex-array/src/scalar_fn/fns/case_when.rs

@@ -258,7 +258,7 @@ impl ScalarFnVTable for CaseWhen {
     }
 }
 
-/// Average run length at which slicing + `extend_from_array` becomes cheaper than `scalar_at`.
+/// Average run length at which slicing + context-aware builder appends become cheaper than `scalar_at`.
 /// Measured empirically via benchmarks.
 const SLICE_CROSSOVER_RUN_LEN: usize = 4;
 
@@ -272,7 +272,7 @@ fn merge_case_branches(
 ) -> VortexResult<ArrayRef> {
     if branches.len() == 1 {
         let (mask, then_value) = &branches[0];
-        return zip_impl(then_value, &else_value, mask);
+        return zip_impl(then_value, &else_value, mask, ctx);
     }
 
     let output_nullability = branches
@@ -314,7 +314,14 @@ fn merge_case_branches(
             ctx,
         )
     } else {
-        merge_run_by_run(&branch_arrays, &else_value, &spans, &output_dtype, builder)
+        merge_run_by_run(
+            &branch_arrays,
+            &else_value,
+            &spans,
+            &output_dtype,
+            builder,
+            ctx,
+        )
     }
 }
 
@@ -348,7 +355,7 @@ fn merge_row_by_row(
     Ok(builder.finish())
 }
 
-/// Bulk-copies each span via `slice()` + `extend_from_array`.
+/// Bulk-copies each span via `slice()` and context-aware builder appends.
 /// Preferred when runs are long enough that memcpy dominates over per-slice allocation cost.
 /// Lazy cast via `arr.cast(output_dtype)` is executed once per span as a block.
 fn merge_run_by_run(
@@ -357,21 +364,25 @@ fn merge_run_by_run(
     spans: &[(usize, usize, usize)],
     output_dtype: &DType,
     mut builder: Box<dyn ArrayBuilder>,
+    ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
     let else_value = else_value.cast(output_dtype.clone())?;
     let len = else_value.len();
     for (start, end, branch_idx) in spans {
         if builder.len() < *start {
-            builder.extend_from_array(&else_value.slice(builder.len()..*start)?);
+            else_value
+                .slice(builder.len()..*start)?
+                .append_to_builder(builder.as_mut(), ctx)?;
         }
-        builder.extend_from_array(
-            &branch_arrays[*branch_idx]
-                .cast(output_dtype.clone())?
-                .slice(*start..*end)?,
-        );
+        branch_arrays[*branch_idx]
+            .cast(output_dtype.clone())?
+            .slice(*start..*end)?
+            .append_to_builder(builder.as_mut(), ctx)?;
     }
     if builder.len() < len {
-        builder.extend_from_array(&else_value.slice(builder.len()..len)?);
+        else_value
+            .slice(builder.len()..len)?
+            .append_to_builder(builder.as_mut(), ctx)?;
     }
 
     Ok(builder.finish())

vortex-array/src/scalar_fn/fns/zip/mod.rs

@@ -139,7 +139,7 @@ impl ScalarFnVTable for Zip {
             return mask.into_array().zip(if_true, if_false);
         }
 
-        zip_impl(&if_true, &if_false, &mask)
+        zip_impl(&if_true, &if_false, &mask, ctx)
     }
 
     fn simplify(
@@ -176,6 +176,7 @@ pub(crate) fn zip_impl(
     if_true: &ArrayRef,
     if_false: &ArrayRef,
     mask: &Mask,
+    ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
     assert_eq!(
         if_true.len(),
@@ -195,12 +196,18 @@ pub(crate) fn zip_impl(
         return if_false.cast(return_type);
     }
 
+    // `append_to_builder` requires exact dtype equality, so normalize branch
+    // nullability to the output dtype before appending slices into the builder.
+    let if_true = if_true.cast(return_type.clone())?;
+    let if_false = if_false.cast(return_type.clone())?;
+
     zip_impl_with_builder(
-        if_true,
-        if_false,
+        &if_true,
+        &if_false,
         mask.values()
             .vortex_expect("zip_impl_with_builder: mask is not all-true or all-false"),
         builder_with_capacity(&return_type, if_true.len()),
+        ctx,
     )
 }
 
@@ -209,13 +216,22 @@ fn zip_impl_with_builder(
     if_false: &ArrayRef,
     mask: &MaskValues,
     mut builder: Box<dyn ArrayBuilder>,
+    ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
     for (start, end) in mask.slices() {
-        builder.extend_from_array(&if_false.slice(builder.len()..*start)?);
-        builder.extend_from_array(&if_true.slice(*start..*end)?);
+        if builder.len() < *start {
+            if_false
+                .slice(builder.len()..*start)?
+                .append_to_builder(builder.as_mut(), ctx)?;
+        }
+        if_true
+            .slice(*start..*end)?
+            .append_to_builder(builder.as_mut(), ctx)?;
     }
     if builder.len() < if_false.len() {
-        builder.extend_from_array(&if_false.slice(builder.len()..if_false.len())?);
+        if_false
+            .slice(builder.len()..if_false.len())?
+            .append_to_builder(builder.as_mut(), ctx)?;
     }
     Ok(builder.finish())
 }
@@ -319,7 +335,8 @@ mod tests {
         let if_false =
             PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array();
 
-        let result = zip_impl(&if_true, &if_false, &mask)?;
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?;
         assert_arrays_eq!(
             result,
             PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30), Some(40)])
@@ -336,7 +353,8 @@ mod tests {
             PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array();
         let if_false = buffer![1i32, 2, 3, 4].into_array();
 
-        let result = zip_impl(&if_true, &if_false, &mask)?;
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?;
         assert_arrays_eq!(
             result,
             PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array()