This is an automated email from the ASF dual-hosted git repository.
alamb 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 b5794fbda0 perf(parquet): avoid redundant copies in mask-backed
intersection/union (#10446)
b5794fbda0 is described below
commit b5794fbda02c45675e6ce83c916dae70c5df477a
Author: Huaijin <[email protected]>
AuthorDate: Sat Aug 15 00:16:21 2026 +0800
perf(parquet): avoid redundant copies in mask-backed intersection/union
(#10446)
# Which issue does this PR close?
- Closes #10425.
# Rationale for this change
The bitwise `intersection`/`union` path, taken when both operands are
mask-backed, does avoidable copying in two places:
- **Unequal lengths** allocate twice — once for the result over the
common prefix, then again for a `BooleanBufferBuilder` that appends that
prefix and the longer side's tail.
- **Equal lengths** go through `BitAnd`/`BitOr` on `&BooleanBuffer`,
which normalise the result to a zero bit offset. That normalisation is a
second allocation plus a shifting copy of the whole mask, and it happens
whenever `from_bitwise_binary_op` returns a non-zero offset — i.e. when
both operands share a non-zero sub-64-bit alignment.
Masks here come from `BooleanBuffer::slice`, so a non-zero offset is the
ordinary case.
# What changes are included in this PR?
Unequal lengths now copy the longer mask once and apply `&=`/`|=` in
place over the common prefix, leaving the tail where it already is.
Equal lengths build the result directly with
`BooleanBuffer::from_bitwise_binary_op`, keeping whatever offset it
produces. Both keep the sub-byte offset rather than re-aligning, and
consumers are already offset-aware.
Only the longer mask's own byte range is copied, so a result derived
from a small slice of a large buffer does not retain the large
allocation.
The uneven path is also what motivated #10444 — it reads the bits past
the common prefix back out, so it relies on the in-place op leaving them
alone.
# Are these changes tested?
`test_mask_algebra_with_offsets` sweeps a 7×7 offset grid against twelve
length pairs, equal and unequal; `test_mask_algebra_fuzz` runs 200
randomized rounds biased towards equal lengths;
`test_mask_algebra_does_not_retain_backing_buffer` covers the
allocation-retention case above. All compare against a bit-by-bit
reference. `cargo test -p parquet --all-features` passes (1305 tests);
clippy and fmt clean.
This PR also adds `mask_intersection`/`mask_union` to
`parquet/benches/row_selector.rs`, since the existing
`intersection`/`union` benchmarks build operands with `from_filters` and
never reach this path.
```
CARGO_PROFILE_BENCH_CODEGEN_UNITS=1 cargo bench -p parquet --bench
row_selector -- 'mask_intersection|mask_union'
```
3M rows, ~1/3 density. Rows are the operand length ratio, columns are
the `(left, right)` bit offsets the operands carry. Whether the two
offsets agree mod 64 decides which path `from_bitwise_binary_op` takes,
so both are covered.
Speedup (`main` / this PR), as `intersection` / `union`:
| | `(0, 0)` | `(3, 3)` | `(3, 67)` | `(3, 5)` |
| --- | --- | --- | --- | --- |
| equal | 1.00 / 1.09 | **2.23 / 2.12** | **2.09 / 2.24** | **3.55 /
3.26** |
| tail1 | 1.18 / 1.17 | **5.90 / 5.96** | **6.03 / 5.98** | **5.15 /
4.55** |
| tail1of3 | 1.22 / 1.21 | **5.68 / 5.69** | **5.74 / 5.38** | **5.03 /
4.30** |
| tail_most | 1.49 / 1.51 | **2.75 / 2.60** | **2.77 / 2.71** | **2.63 /
2.69** |
`equal` at `(0, 0)` is the one case with nothing to save: the operands
are already aligned, so the old code never paid for normalisation, and
old and new produce identical buffers. Its confidence intervals overlap
(11.3±0.21µs vs 11.3±0.28µs, 10.0±0.92µs vs 10.9±0.86µs).
`codegen-units=1` because the default bench profile is noisy here: the
shifting path is sensitive to codegen-unit partitioning, enough to move
these numbers by ±10% in either direction.
# Are there any user-facing changes?
No public signature changes and no change to which rows a selection
selects. The underlying mask layout can differ, though:
`RowSelection::as_mask()` is public, and where equal-length
`intersection`/`union` previously always returned a zero-offset
`BooleanBuffer`, the result may now carry a non-zero offset, which is
observable through `values()`, `inner()` and `ptr_eq()`. `as_mask()`
does not promise a normalised layout and callers must already honour
`BooleanBuffer::offset()`, so this should not be a breaking change.
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
.../src/arrow/arrow_reader/selection/algebra.rs | 220 ++++++++++++++++++---
1 file changed, 191 insertions(+), 29 deletions(-)
diff --git a/parquet/src/arrow/arrow_reader/selection/algebra.rs
b/parquet/src/arrow/arrow_reader/selection/algebra.rs
index ce44882858..3d11a13540 100644
--- a/parquet/src/arrow/arrow_reader/selection/algebra.rs
+++ b/parquet/src/arrow/arrow_reader/selection/algebra.rs
@@ -23,7 +23,7 @@
//! [`BooleanBuffer`] masks.
use super::{MaskRunIter, RowSelection, RowSelectionInner, RowSelector};
-use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
+use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer,
bit_util};
use std::cmp::Ordering;
use std::iter::Peekable;
@@ -269,39 +269,76 @@ pub(super) fn union_row_selections(left: &[RowSelector],
right: &[RowSelector])
/// Bitwise AND of two mask-backed selections. Longer side's tail passes
through.
pub(super) fn intersect_masks(l: &BooleanBuffer, r: &BooleanBuffer) ->
BooleanBuffer {
if l.len() == r.len() {
- return l & r;
- }
- let common = l.len().min(r.len());
- let head = &l.slice(0, common) & &r.slice(0, common);
- let (longer, longer_len) = if l.len() > r.len() {
- (l, l.len())
- } else {
- (r, r.len())
- };
- let tail = longer.slice(common, longer_len - common);
- let mut builder = BooleanBufferBuilder::new(longer_len);
- builder.append_buffer(&head);
- builder.append_buffer(&tail);
- builder.finish()
+ return combine_equal_length_masks(l, r, |a, b| a & b);
+ }
+ combine_unequal_length_masks(l, r, |a, b| a & b)
}
/// Bitwise OR of two mask-backed selections. Longer side's tail passes
through.
pub(super) fn union_masks(l: &BooleanBuffer, r: &BooleanBuffer) ->
BooleanBuffer {
if l.len() == r.len() {
- return l | r;
- }
- let common = l.len().min(r.len());
- let head = &l.slice(0, common) | &r.slice(0, common);
- let (longer, longer_len) = if l.len() > r.len() {
- (l, l.len())
- } else {
- (r, r.len())
- };
- let tail = longer.slice(common, longer_len - common);
- let mut builder = BooleanBufferBuilder::new(longer_len);
- builder.append_buffer(&head);
- builder.append_buffer(&tail);
- builder.finish()
+ return combine_equal_length_masks(l, r, |a, b| a | b);
+ }
+ combine_unequal_length_masks(l, r, |a, b| a | b)
+}
+
+/// Combines two masks of equal length with the bitwise operation `op`.
+///
+/// `BitAnd`/`BitOr` on `&BooleanBuffer` normalise the result to a zero bit
offset,
+/// which costs a second allocation and a shifting copy when both operands
+/// share the same non-zero sub-64-bit alignment, causing
+/// `from_bitwise_binary_op` to return a non-zero-offset result. Building the
+/// buffer directly keeps the offset, as the unequal-length path does.
+fn combine_equal_length_masks<F>(l: &BooleanBuffer, r: &BooleanBuffer, op: F)
-> BooleanBuffer
+where
+ F: FnMut(u64, u64) -> u64,
+{
+ BooleanBuffer::from_bitwise_binary_op(
+ l.values(),
+ l.offset(),
+ r.values(),
+ r.offset(),
+ l.len(),
+ op,
+ )
+}
+
+/// Combines two masks of unequal length with the bitwise operation `op`,
+/// passing the longer side's tail through unchanged.
+///
+/// The longer mask is copied once into a [`MutableBuffer`] and `op` is then
+/// applied in place over the common prefix. This avoids materialising the
+/// prefix into its own buffer and copying both prefix and tail again through a
+/// [`BooleanBufferBuilder`].
+///
+/// Neither the mask offsets nor the prefix length are assumed to be byte
+/// aligned: the copy keeps the longer mask's offset within its first byte so
it
+/// stays a plain byte copy, and that offset is carried over to the result.
Only
+/// the longer mask's own byte range is copied, so the result does not retain
the
+/// backing allocation it was sliced from.
+fn combine_unequal_length_masks<F>(l: &BooleanBuffer, r: &BooleanBuffer, op:
F) -> BooleanBuffer
+where
+ F: FnMut(u64, u64) -> u64,
+{
+ let (longer, shorter) = if l.len() > r.len() { (l, r) } else { (r, l) };
+
+ let sub_byte_offset = longer.offset() % 8;
+ let start_byte = longer.offset() / 8;
+ let end_byte = bit_util::ceil(longer.offset() + longer.len(), 8);
+ let bytes = &longer.values()[start_byte..end_byte];
+ let mut buffer = MutableBuffer::new(bytes.len());
+ buffer.extend_from_slice(bytes);
+
+ bit_util::apply_bitwise_binary_op(
+ buffer.as_slice_mut(),
+ sub_byte_offset,
+ shorter.values(),
+ shorter.offset(),
+ shorter.len(),
+ op,
+ );
+
+ BooleanBuffer::new(buffer.into(), sub_byte_offset, longer.len())
}
/// Applies `other` to the selected rows of `mask`, preserving the original
row domain.
@@ -825,4 +862,129 @@ mod tests {
let bits: Vec<bool> = (0..5).map(|i| r_mask.value(i)).collect();
assert_eq!(bits, vec![true, true, false, false, true]);
}
+
+ /// Expected result of combining two masks of possibly differing lengths:
+ /// `op` over the common prefix, then the longer side's tail unchanged.
+ fn expected_combined(l: &[bool], r: &[bool], op: fn(bool, bool) -> bool)
-> Vec<bool> {
+ let common = l.len().min(r.len());
+ let longer = if l.len() > r.len() { l } else { r };
+ (0..common)
+ .map(|i| op(l[i], r[i]))
+ .chain(longer[common..].iter().copied())
+ .collect()
+ }
+
+ fn assert_mask_eq(actual: &BooleanBuffer, expected: &[bool], context:
&str) {
+ assert_eq!(actual.len(), expected.len(), "{context}: length");
+ let actual: Vec<bool> = actual.iter().collect();
+ assert_eq!(actual, expected, "{context}");
+ }
+
+ #[test]
+ fn test_mask_algebra_with_offsets() {
+ // Offsets and lengths that are not byte (or word) aligned on either
side,
+ // so the common prefix can start and end mid byte. Covers both the
equal
+ // and uneven length paths.
+ let base: Vec<bool> = (0..600).map(|i| i % 7 == 0 || i % 3 ==
1).collect();
+ let other: Vec<bool> = (0..600).map(|i| i % 5 == 2 || i % 11 ==
4).collect();
+ let base = BooleanBuffer::from(base);
+ let other = BooleanBuffer::from(other);
+
+ for l_offset in [0, 1, 5, 8, 13, 64, 67] {
+ for r_offset in [0, 1, 3, 8, 60, 64, 70] {
+ for (l_len, r_len) in [
+ (0, 9),
+ (9, 0),
+ (1, 200),
+ (200, 1),
+ (63, 130),
+ (321, 65),
+ (0, 0),
+ (1, 1),
+ (63, 63),
+ (64, 64),
+ (200, 200),
+ (321, 321),
+ ] {
+ let l = base.slice(l_offset, l_len);
+ let r = other.slice(r_offset, r_len);
+ let l_bits: Vec<bool> = l.iter().collect();
+ let r_bits: Vec<bool> = r.iter().collect();
+ let context =
+ format!("l_offset={l_offset} r_offset={r_offset}
lens=({l_len},{r_len})");
+
+ assert_mask_eq(
+ &intersect_masks(&l, &r),
+ &expected_combined(&l_bits, &r_bits, |a, b| a && b),
+ &format!("intersect {context}"),
+ );
+ assert_mask_eq(
+ &union_masks(&l, &r),
+ &expected_combined(&l_bits, &r_bits, |a, b| a || b),
+ &format!("union {context}"),
+ );
+ }
+ }
+ }
+ }
+
+ #[test]
+ fn test_mask_algebra_does_not_retain_backing_buffer() {
+ // A short slice of a long mask must not keep the long allocation
alive,
+ // including when the other operand is empty and contributes nothing.
+ // Comparing capacities rather than lengths, since a shallow slice
reports a
+ // short length while still holding the original allocation through
its `Arc`.
+ let long = BooleanBuffer::from((0..80_000).map(|i| i % 3 ==
0).collect::<Vec<bool>>());
+ assert!(long.inner().capacity() >= 10_000);
+
+ for (l, r) in [
+ (long.slice(5, 40), BooleanBuffer::new_unset(0)),
+ (long.slice(5, 40), BooleanBuffer::new_set(7)),
+ (BooleanBuffer::new_set(7), long.slice(5, 40)),
+ ] {
+ for combined in [intersect_masks(&l, &r), union_masks(&l, &r)] {
+ assert!(
+ combined.inner().capacity() < long.inner().capacity(),
+ "result retained the original backing allocation"
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn test_mask_algebra_fuzz() {
+ let mut rng = rng();
+ for _ in 0..200 {
+ let l_offset = rng.random_range(0..70);
+ let r_offset = rng.random_range(0..70);
+ let l_len = rng.random_range(0..300);
+ // Bias towards equal lengths so that path is hit often
+ let r_len = match rng.random_bool(0.25) {
+ true => l_len,
+ false => rng.random_range(0..300),
+ };
+
+ let l_bits: Vec<bool> = (0..l_offset + l_len)
+ .map(|_| rng.random_bool(0.5))
+ .collect();
+ let r_bits: Vec<bool> = (0..r_offset + r_len)
+ .map(|_| rng.random_bool(0.5))
+ .collect();
+ let l = BooleanBuffer::from(l_bits).slice(l_offset, l_len);
+ let r = BooleanBuffer::from(r_bits).slice(r_offset, r_len);
+ let l_bits: Vec<bool> = l.iter().collect();
+ let r_bits: Vec<bool> = r.iter().collect();
+
+ assert_mask_eq(
+ &intersect_masks(&l, &r),
+ &expected_combined(&l_bits, &r_bits, |a, b| a && b),
+ "intersect",
+ );
+ assert_mask_eq(
+ &union_masks(&l, &r),
+ &expected_combined(&l_bits, &r_bits, |a, b| a || b),
+ "union",
+ );
+ }
+ }
}
