alamb commented on code in PR #10968:
URL: https://github.com/apache/arrow-rs/pull/10968#discussion_r4105011601
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -976,6 +1034,51 @@ mod tests {
}
}
+ #[test]
+ fn test_expand() {
+ fn reference(values: u64, mask: u64) -> u64 {
+ let mut expected = 0;
+ let mut input_idx = 0;
+ for output_idx in 0..64 {
+ if mask & (1 << output_idx) != 0 {
+ expected |= ((values >> input_idx) & 1) << output_idx;
+ input_idx += 1;
+ }
+ }
+ expected
+ }
+
+ assert_eq!(expand(0b1010, 0b1111), 0b1010);
+ assert_eq!(expand(0b11, 0b1010), 0b1010);
+ assert_eq!(expand(u64::MAX, 0), 0);
+ assert_eq!(expand(0, u64::MAX), 0);
+ assert_eq!(expand(0, 0x5555_5555_5555_5555), 0);
+ assert_eq!(expand(u64::MAX, u64::MAX), u64::MAX);
+
+ let mut rng = StdRng::seed_from_u64(0x2b7e_1516_28ae_d2a6);
+
+ // Masks with at most one unset bit or at most one set bit
+ for bit in 0..64 {
+ for mask in [u64::MAX, !(1 << bit), 1 << bit, 0] {
+ for _ in 0..16 {
+ let values = rng.random::<u64>();
+ assert_eq!(expand(values, mask), reference(values, mask),
"{mask:#x}");
+ }
+ }
+ }
+
+ // Masks across the full density range, exercising both loops
+ for _ in 0..20_000 {
Review Comment:
in case anyone is curious, this takes 300ms on my laptop
```
PASS [ 0.302s] arrow-buffer util::bit_util::tests::test_expand
```
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -81,6 +81,64 @@ pub fn compress(value: u64, mask: u64) -> u64 {
}
}
+/// Parallel bit deposit: scatter the lowest `mask.count_ones()` bits of
+/// `value` into the set positions of `mask`, preserving their order.
+/// All other bits in the result are zero; excess input bits are ignored.
+///
+/// This is the inverse of [`compress`] on the selected bits:
+/// `expand(compress(value, mask), mask) == value & mask`.
+///
+/// Equivalent to the x86 BMI2 `PDEP` instruction, implemented with a portable
+/// scalar loop that visits whichever is fewer: unset or set bits in `mask`.
+///
+/// # Functional Example
+///
+/// Using 8 bits for brevity (the function operates on all 64). The low bits
+/// of `value` are scattered into the set positions of `mask`:
+///
+/// ```text
+/// bit: 7 6 5 4 3 2 1 0
+/// value: 0 0 0 b c e f h
+/// mask: 0 1 1 0 1 1 0 1
+/// result: 0 b c 0 e f 0 h
+/// ```
+///
+/// # Code Example
+///
+/// ```
+/// # use arrow_buffer::bit_util::{compress, expand};
+/// assert_eq!(expand(0b0000_1010, 0b0110_1101), 0b0010_0100);
+/// let value = 0b1011_0100;
+/// let mask = 0b0110_1101;
+/// assert_eq!(expand(compress(value, mask), mask), value & mask);
+/// ```
+#[inline]
+pub fn expand(mut value: u64, mask: u64) -> u64 {
Review Comment:
❤️
##########
parquet/src/arrow/arrow_reader/selection/algebra.rs:
##########
@@ -430,11 +456,49 @@ fn and_then_masks(mask: &BooleanBuffer, other:
&BooleanBuffer) -> BooleanBuffer
builder.finish()
}
+/// Scatters the next `mask_word.count_ones()` bits of `other` into each mask
word.
+/// Requires `other.len() == mask.count_set_bits()`.
+#[inline(never)]
+fn and_then_dense_masks(mask: &BooleanBuffer, other: &BooleanBuffer) ->
BooleanBuffer {
+ let mut other_chunks = other.bit_chunks().iter_padded();
Review Comment:
can we please add a debug_assert for this:
> /// Requires `other.len() == mask.count_set_bits()`.
```rust
debug_assert!(other.len(), mask.count_set_bits());
```
To ensure it is good
##########
parquet/src/arrow/arrow_reader/selection/algebra.rs:
##########
@@ -728,6 +792,96 @@ mod tests {
);
}
+ #[test]
+ fn test_dense_mask_and_then_mask_with_offsets() {
+ for len in [8192, 8193, 10_000] {
+ for outer_stride in [100, 20, 4] {
+ let outer_offset = 3;
+ let outer = BooleanBuffer::from_iter(
+ (0..len + outer_offset).map(|i| (i + 1) % outer_stride !=
0),
+ )
+ .slice(outer_offset, len);
+
+ let inner_offset = 5;
+ let inner_len = outer.count_set_bits();
+ let inner = BooleanBuffer::from_iter(
+ (0..inner_len + inner_offset).map(|i| (i + 2) % 97 != 0),
+ )
+ .slice(inner_offset, inner_len);
+
+ assert!(should_use_dense_mask(
+ outer.len(),
+ inner.len(),
+ inner.count_set_bits()
+ ));
+
+ let mut inner_idx = 0;
Review Comment:
this validation code seems to be copy/pasted in each of the following tests.
Can we refactor it into a fixture? Something like this maybe?
```rust
AndThenTest { inner, outer, dense_mask_expected: true }.run()
```
##########
parquet/src/arrow/arrow_reader/selection/algebra.rs:
##########
@@ -27,6 +27,14 @@ use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder,
MutableBuffer, bit_util}
use std::cmp::Ordering;
use std::iter::Peekable;
+// Use word-at-a-time expansion for masks with at least 8192 rows, roughly
+// 75% outer selectivity and 5% inner selectivity. Smaller or sparser inputs
+// use set indices to avoid scanning every output word.
+// The selectivity constants are divisors: at most 1/4 dropped, at least 1/20
kept.
+const AND_THEN_DENSE_MASK_MIN_LEN: usize = 8192;
Review Comment:
how did you arrive at these constants? Do you measure a sweep on some
machine?
It would also help to explain what you mean here by "outer" and "inner" if
possible -- I assume you mean `outer.and_then(inner)` but that would be good to
explicitly say
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -976,6 +1034,51 @@ mod tests {
}
}
+ #[test]
+ fn test_expand() {
+ fn reference(values: u64, mask: u64) -> u64 {
+ let mut expected = 0;
+ let mut input_idx = 0;
+ for output_idx in 0..64 {
+ if mask & (1 << output_idx) != 0 {
+ expected |= ((values >> input_idx) & 1) << output_idx;
+ input_idx += 1;
+ }
+ }
+ expected
+ }
+
+ assert_eq!(expand(0b1010, 0b1111), 0b1010);
+ assert_eq!(expand(0b11, 0b1010), 0b1010);
+ assert_eq!(expand(u64::MAX, 0), 0);
+ assert_eq!(expand(0, u64::MAX), 0);
+ assert_eq!(expand(0, 0x5555_5555_5555_5555), 0);
+ assert_eq!(expand(u64::MAX, u64::MAX), u64::MAX);
+
+ let mut rng = StdRng::seed_from_u64(0x2b7e_1516_28ae_d2a6);
+
+ // Masks with at most one unset bit or at most one set bit
+ for bit in 0..64 {
+ for mask in [u64::MAX, !(1 << bit), 1 << bit, 0] {
+ for _ in 0..16 {
+ let values = rng.random::<u64>();
+ assert_eq!(expand(values, mask), reference(values, mask),
"{mask:#x}");
+ }
+ }
+ }
+
+ // Masks across the full density range, exercising both loops
+ for _ in 0..20_000 {
Review Comment:
I also verified coverage with
```shell
cargo llvm-cov test --html -p arrow-buffer -- test_expand
```
And all the branches are covered:
<img width="1164" height="928" alt="Image"
src="https://github.com/user-attachments/assets/8c9286c2-2e2d-4f1d-bb5b-4babe193bc02";
/>
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -81,6 +81,64 @@ pub fn compress(value: u64, mask: u64) -> u64 {
}
}
+/// Parallel bit deposit: scatter the lowest `mask.count_ones()` bits of
+/// `value` into the set positions of `mask`, preserving their order.
+/// All other bits in the result are zero; excess input bits are ignored.
+///
+/// This is the inverse of [`compress`] on the selected bits:
+/// `expand(compress(value, mask), mask) == value & mask`.
+///
+/// Equivalent to the x86 BMI2 `PDEP` instruction, implemented with a portable
+/// scalar loop that visits whichever is fewer: unset or set bits in `mask`.
+///
+/// # Functional Example
+///
+/// Using 8 bits for brevity (the function operates on all 64). The low bits
+/// of `value` are scattered into the set positions of `mask`:
+///
+/// ```text
+/// bit: 7 6 5 4 3 2 1 0
+/// value: 0 0 0 b c e f h
+/// mask: 0 1 1 0 1 1 0 1
+/// result: 0 b c 0 e f 0 h
+/// ```
+///
+/// # Code Example
+///
+/// ```
+/// # use arrow_buffer::bit_util::{compress, expand};
+/// assert_eq!(expand(0b0000_1010, 0b0110_1101), 0b0010_0100);
+/// let value = 0b1011_0100;
+/// let mask = 0b0110_1101;
+/// assert_eq!(expand(compress(value, mask), mask), value & mask);
+/// ```
+#[inline]
+pub fn expand(mut value: u64, mask: u64) -> u64 {
Review Comment:
I wonder if we should put the x86 pdep call here while we are at it (it is
like 4 extra lines...)
--
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]
