haohuaijin commented on code in PR #10446:
URL: https://github.com/apache/arrow-rs/pull/10446#discussion_r3730182266
##########
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:
apply in
[fae46dd](https://github.com/apache/arrow-rs/pull/10446/commits/fae46dde05e81859d0b40623d54623ce95107217)
##########
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:
apply in
[fae46dd](https://github.com/apache/arrow-rs/pull/10446/commits/fae46dde05e81859d0b40623d54623ce95107217)
--
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]
