hhhizzz commented on code in PR #10446:
URL: https://github.com/apache/arrow-rs/pull/10446#discussion_r3725499778
##########
parquet/src/arrow/arrow_reader/selection/algebra.rs:
##########
@@ -825,4 +861,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.
+ let long = BooleanBuffer::from((0..80_000).map(|i| i % 3 ==
0).collect::<Vec<bool>>());
+ assert!(long.inner().len() >= 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().len() <= 16,
Review Comment:
I don't think `len()` fully verifies the property stated by this test. A
shallow `Buffer::slice_with_length` can expose a small logical length while
still retaining the original large allocation through its `Arc`, so that
regression could pass this assertion.
Could we compare the backing capacity instead, for example:
assert!(
combined.inner().capacity() < long.inner().capacity(),
"result retained the original backing allocation"
);
The current implementation does make a compact copy; this would ensure the
test actually protects that memory-retention property.
##########
parquet/src/arrow/arrow_reader/selection/algebra.rs:
##########
@@ -269,39 +269,75 @@ 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 of the whole mask when
the
+/// operands are not byte aligned. Building the buffer directly keeps the
offset,
Review Comment:
Could we narrow this condition? “When the operands are not byte aligned” is
too broad: `(3,5)` is also non-byte-aligned, but because the operands have
different mod-64 alignments, `from_bitwise_binary_op` takes the shifting
fallback and returns an offset-zero result, so the old wrapper does not perform
an additional normalization copy.
Perhaps:
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
