This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 34b49efced perf: Use interleave for fragmented zip masks (#10368)
34b49efced is described below
commit 34b49efced33a5c6fce086b69fdc9fa5f0d29620
Author: Johan Vaz <[email protected]>
AuthorDate: Sat Aug 29 14:49:22 2026 +0800
perf: Use interleave for fragmented zip masks (#10368)
# Which issue does this PR close?
- Closes #5097.
# Rationale for this change
`zip` currently uses `MutableArrayData`, which is efficient when a mask
contains long runs but performs one extension per run. Fragmented
array/array masks therefore pay substantial per-run overhead, while the
specialized interleave kernels are faster for this shape.
# What changes are included in this PR?
- Count true runs in large masks and route fragmented array/array inputs
to `interleave`.
- Keep the existing `MutableArrayData` path for inputs shorter than
1,024 rows, scalar inputs, and masks with fewer runs.
- Treat null mask entries as false before dispatch, preserving existing
`zip` semantics.
- Add a contiguous `true_then_false` case to the existing zip benchmark
matrix. The broader direct-interleave benchmark expansion was removed
after review.
The run-count threshold is conservative: interleave is selected only
above one true run per eight rows. Counting the complete bitmap also
avoids sampling errors for masks whose fragmentation is unevenly
distributed.
# Are these changes tested?
Yes. The tests cover dispatch decisions for short, fragmented,
sparse/dense, contiguous, unevenly fragmented, and offset masks, plus
end-to-end fragmented array selection with null mask and input values.
Validation on final head `709a774a`:
- `CARGO_BUILD_JOBS=2 cargo test -p arrow-select --all-features` — 394
unit tests and 17 doctests passed.
- `CARGO_BUILD_JOBS=2 cargo clippy -p arrow-select --all-targets
--all-features -- -D warnings` — passed.
- `CARGO_BUILD_JOBS=2 cargo clippy -p arrow --bench zip_kernels
--features=async,test_utils -- -D warnings` — passed.
- `cargo fmt --all -- --check` — passed.
- `git diff --check` — passed.
Criterion used 30 samples, a 2-second warm-up, and a 5-second
measurement. On 8,192-row array/array inputs with random 50%-true masks,
the existing run path was forced by temporarily raising the local
dispatch threshold; the benchmark-only control was then restored:
| Input | Run path | Adaptive | Improvement |
|---|---:|---:|---:|
| i32 | 76.889 us | 33.676 us | 56.20% |
| short strings | 101.600 us | 69.580 us | 31.52% |
| long strings | 214.880 us | 125.560 us | 41.57% |
| short bytes | 103.680 us | 70.169 us | 32.32% |
| long bytes | 218.090 us | 123.040 us | 43.58% |
| short string views | 85.114 us | 38.201 us | 55.12% |
| longer string views | 84.633 us | 49.308 us | 41.74% |
At the 1,024-row boundary, i32 measured 14.350 us on the forced run path
versus 4.6005 us on the adaptive path. The contiguous `true_then_false`
case remains on the run-oriented path.
# Are there any user-facing changes?
No API or semantic changes. Fragmented array/array `zip` calls use a
faster internal implementation; scalar, short, and contiguous-mask calls
retain the existing path.
# AI assistance disclosure
AI assistance was used to draft parts of the adaptive dispatch, tests,
benchmark changes, and review follow-up. The resulting diff was reviewed
against the `zip`, `MutableArrayData`, `BooleanBuffer`, and `interleave`
implementations; an edge-sampling flaw found during review was replaced
with a complete run count and covered by a regression test. The
validation commands above were run on the final branch head.
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
arrow-select/src/zip.rs | 131 ++++++++++++++++++++++++++++++++++++++++++-
arrow/benches/zip_kernels.rs | 4 ++
2 files changed, 134 insertions(+), 1 deletion(-)
diff --git a/arrow-select/src/zip.rs b/arrow-select/src/zip.rs
index 14a3dd60a4..c0dcf6fbd9 100644
--- a/arrow-select/src/zip.rs
+++ b/arrow-select/src/zip.rs
@@ -145,6 +145,41 @@ 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,
@@ -152,6 +187,11 @@ fn zip_impl(
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 {
@@ -866,6 +905,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)]);
diff --git a/arrow/benches/zip_kernels.rs b/arrow/benches/zip_kernels.rs
index c798bee3b6..03f20214a8 100644
--- a/arrow/benches/zip_kernels.rs
+++ b/arrow/benches/zip_kernels.rs
@@ -169,6 +169,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)),
