Use interleave for fragmented zip masks

arrow-select/src/zip.rs

@@ -145,13 +145,53 @@ pub fn zip(
     zip_impl(mask, &truthy, truthy_is_scalar, &falsy, falsy_is_scalar)
 }
 
+fn count_true_runs(mask: &BooleanBuffer) -> usize {
+    let mut slices = 0;
+    let mut previous = 0;
+    for chunk in mask.bit_chunks().iter_padded() {
+        let starts = chunk & !(chunk << 1 | previous);
+        slices += starts.count_ones() as usize;
+        previous = chunk >> 63;
+    }
+    slices
+}
+
+fn should_use_interleave(mask: &BooleanBuffer) -> bool {
+    const MIN_LEN: usize = 1024;
+
+    // Interleave's fixed dispatch and index construction costs are not competitive
+    // for small arrays. For larger arrays, use the run count that determines the
+    // amount of work performed by the MutableArrayData implementation.
+    mask.len() >= MIN_LEN && count_true_runs(mask) > mask.len() / 8
+}
+
+fn interleave_arrays(
+    mask: &BooleanBuffer,
+    truthy: &ArrayData,
+    falsy: &ArrayData,
+) -> Result<ArrayRef, ArrowError> {
+    let truthy = make_array(truthy.clone());
+    let falsy = make_array(falsy.clone());
+    let indices: Vec<_> = mask
+        .iter()
+        .enumerate()
+        .map(|(idx, selected)| (usize::from(!selected), idx))
+        .collect();
+    crate::interleave::interleave(&[truthy.as_ref(), falsy.as_ref()], &indices)
+}
+
 fn zip_impl(
     mask: &BooleanArray,
     truthy: &ArrayData,
     truthy_is_scalar: bool,
     falsy: &ArrayData,
     falsy_is_scalar: bool,
 ) -> Result<ArrayRef, ArrowError> {
+    let mask_buffer = maybe_prep_null_mask_filter(mask);
+    if !truthy_is_scalar && !falsy_is_scalar && should_use_interleave(&mask_buffer) {
+        return interleave_arrays(&mask_buffer, truthy, falsy);
+    }
+
     let mut mutable = MutableArrayData::new(vec![truthy, falsy], false, truthy.len());
 
     // the SlicesIterator slices only the true values. So the gaps left by this iterator we need to
@@ -160,7 +200,6 @@ fn zip_impl(
     // keep track of how much is filled
     let mut filled = 0;
 
-    let mask_buffer = maybe_prep_null_mask_filter(mask);
     for (start, end) in SlicesIterator::from(&mask_buffer) {
         // the gap needs to be filled with falsy values
         if start > filled {
@@ -846,6 +885,96 @@ mod test {
     use super::*;
     use arrow_array::types::Int32Type;
 
+    #[test]
+    fn test_count_true_runs() {
+        let assert_runs = |values: &[bool], expected| {
+            let mask: BooleanBuffer = values.iter().copied().collect();
+            assert_eq!(count_true_runs(&mask), expected, "mask: {values:?}");
+        };
+
+        assert_runs(&[], 0);
+        assert_runs(&[false, false, false], 0);
+        assert_runs(&[true, true, true], 1);
+        assert_runs(&[true, false, true, true, false, true], 3);
+
+        // Exercise runs crossing 64-bit chunk boundaries and trailing padding.
+        let mut values = vec![false; 130];
+        values[0] = true;
+        values[63..66].fill(true);
+        values[128..].fill(true);
+        assert_runs(&values, 3);
+
+        // Exercise a non-zero bit offset, as masks may be sliced.
+        let mut offset_values = vec![false; 135];
+        offset_values[3..133].copy_from_slice(&values);
+        let offset_mask: BooleanBuffer = offset_values.into_iter().collect();
+        assert_eq!(count_true_runs(&offset_mask.slice(3, 130)), 3);
+    }
+
+    #[test]
+    fn test_should_use_interleave() {
+        let short: BooleanBuffer = (0..64).map(|i| i % 2 == 0).collect();
+        assert!(!should_use_interleave(&short));
+
+        let fragmented: BooleanBuffer = (0..8192).map(|i| i % 2 == 0).collect();
+        assert!(should_use_interleave(&fragmented));
+
+        let long_runs: BooleanBuffer = (0..8192).map(|i| i < 4096).collect();
+        assert!(!should_use_interleave(&long_runs));
+
+        let sparse: BooleanBuffer = (0..8192).map(|i| i % 10 == 0).collect();
+        assert!(!should_use_interleave(&sparse));
+
+        let dense: BooleanBuffer = (0..8192).map(|i| i % 10 != 0).collect();
+        assert!(!should_use_interleave(&dense));
+
+        let fragmented_head: BooleanBuffer = (0..8192).map(|i| i < 256 && i % 2 == 0).collect();
+        assert!(!should_use_interleave(&fragmented_head));
+
+        let fragmented_edges: BooleanBuffer = (0..8192)
+            .map(|i| !(256..7936).contains(&i) && i % 2 == 0)
+            .collect();
+        assert!(!should_use_interleave(&fragmented_edges));
+
+        // Exercise arbitrary bit offsets as masks may be sliced
+        let offset: BooleanBuffer = (0..8195).map(|i| i >= 3 && i % 2 == 1).collect();
+        assert!(should_use_interleave(&offset.slice(3, 8192)));
+    }
+
+    #[test]
+    fn test_interleave_arrays() {
+        let mask = BooleanArray::from(vec![Some(true), None, Some(true), Some(false)]);
+        let mask = maybe_prep_null_mask_filter(&mask);
+        let truthy = Int32Array::from(vec![Some(1), None, Some(3), Some(4)]).to_data();
+        let falsy = Int32Array::from(vec![Some(10), Some(20), None, Some(40)]).to_data();
+        let expected = Int32Array::from(vec![Some(1), Some(20), Some(3), Some(40)]);
+
+        let actual = interleave_arrays(&mask, &truthy, &falsy).unwrap();
+        assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
+    }
+
+    #[test]
+    fn test_zip_fragmented_array_mask() {
+        let mask: BooleanArray = (0..8192)
+            .map(|i| match i % 3 {
+                0 => Some(true),
+                1 => Some(false),
+                _ => None,
+            })
+            .collect();
+        let truthy: Int32Array = (0..8192).map(|i| (i % 7 != 0).then_some(i)).collect();
+        let falsy: Int32Array = (0..8192).map(|i| (i % 11 != 0).then_some(-i)).collect();
+        let expected: Int32Array = (0..8192)
+            .map(|i| {
+                let array = if i % 3 == 0 { &truthy } else { &falsy };
+                array.is_valid(i).then(|| array.value(i))
+            })
+            .collect();
+
+        let actual = zip(&mask, &truthy, &falsy).unwrap();
+        assert_eq!(actual.as_primitive::<Int32Type>(), &expected);
+    }
+
     #[test]
     fn test_zip_kernel_one() {
         let a = Int32Array::from(vec![Some(5), None, Some(7), None, Some(1)]);

arrow/benches/zip_kernels.rs

@@ -27,6 +27,7 @@ use std::sync::Arc;
 use arrow::array::*;
 use arrow::datatypes::*;
 use arrow::util::bench_util::*;
+use arrow_select::interleave::interleave;
 use arrow_select::zip::zip;
 
 trait InputGenerator {
@@ -169,6 +170,10 @@ fn mask_cases(len: usize) -> Vec<(&'static str, BooleanArray)> {
         ("99pct_true", create_boolean_array(len, 0.0, 0.99)),
         ("90pct_true", create_boolean_array(len, 0.0, 0.9)),
         ("50pct_true", create_boolean_array(len, 0.0, 0.5)),
+        (
+            "true_then_false",
+            BooleanArray::from_iter((0..len).map(|i| i < len / 2)),
+        ),
         ("10pct_true", create_boolean_array(len, 0.0, 0.1)),
         ("1pct_true", create_boolean_array(len, 0.0, 0.01)),
         ("all_false", create_boolean_array(len, 0.0, 0.0)),
@@ -236,6 +241,26 @@ fn bench_zip_on_input_generator(c: &mut Criterion, input_generator: &impl InputG
         &array_2_10pct_nulls,
     );
 
+    for (mask_description, mask) in &masks {
+        let id = BenchmarkId::new("array_vs_array_interleave", mask_description);
+        group.bench_with_input(id, mask, |b, mask| {
+            b.iter(|| {
+                let indices: Vec<_> = mask
+                    .iter()
+                    .enumerate()
+                    .map(|(idx, selected)| (usize::from(!selected.unwrap_or(false)), idx))
+                    .collect();
+                hint::black_box(
+                    interleave(
+                        &[array_1_10pct_nulls.as_ref(), array_2_10pct_nulls.as_ref()],
+                        &indices,
+                    )
+                    .unwrap(),
+                )
+            })
+        });
+    }
+
     group.finish();
 }