devanbenz commented on code in PR #10136:
URL: https://github.com/apache/arrow-rs/pull/10136#discussion_r4086054678
##########
arrow-select/src/filter.rs:
##########
@@ -719,6 +739,86 @@ fn filter_bits(buffer: &BooleanBuffer, predicate:
&FilterPredicate) -> Buffer {
}
}
+/// Filter the packed bitmask `buffer` with `predicate` by extracting the kept
+/// bits of each 64-bit word with [`bit_util::compress`] (`pext`)
+///
+/// Not inlined: within `filter_array` the packing state spills to the stack
+#[inline(never)]
+fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate)
-> Buffer {
+ /// Packs the bits extracted from successive words into the low `filled`
+ /// bits of `current`; once complete it is written at `idx` and restarts
+ /// from the bits that did not fit
+ struct Packer {
+ ptr: *mut u64,
+ idx: usize,
+ current: u64,
+ filled: u32,
+ }
+
+ impl Packer {
+ #[inline(always)]
+ fn push(&mut self, values: u64, mask: u64) {
+ let bits = bit_util::compress(values, mask);
+ self.current |= bits << self.filled;
+ let total = self.filled + mask.count_ones();
+ if total < 64 {
+ self.filled = total;
+ } else {
+ // SAFETY: `count` is the number of set bits in the filter, so
+ // at most `count / 64` words are ever completed and the
+ // buffer holds `count / 64 + 1`
+ unsafe { self.ptr.add(self.idx).write(self.current) };
+ self.idx += 1;
+ // `bits >> (64 - filled)`, written so that `filled == 0`
+ // shifts everything out
+ self.current = (bits >> 1) >> (63 - self.filled);
+ self.filled = total - 64;
+ }
+ }
+ }
+
+ assert!(buffer.len() >= predicate.filter.len());
+ let mask_chunks = predicate.filter.values().bit_chunks();
+ let value_chunks = BitChunks::new(buffer.values(), buffer.offset(),
predicate.filter.len());
+ // `count` is the filter's set bit count, which the buffer size and the
+ // raw writes below rely on, and both chunk views cover
+ // `predicate.filter.len()` bits, so indexing `value_chunks` by the
+ // position in `mask_chunks` stays in bounds
+ debug_assert_eq!(predicate.count, predicate.filter.true_count());
+ debug_assert_eq!(mask_chunks.chunk_len(), value_chunks.chunk_len());
+
+ // One word beyond the complete ones for the trailing partial word
+ let mut out: Vec<u64> = Vec::with_capacity(predicate.count / 64 + 1);
+ let mut packer = Packer {
+ ptr: out.as_mut_ptr(),
+ idx: 0,
+ current: 0,
+ filled: 0,
+ };
+
+ for (index, mask) in mask_chunks.iter().enumerate() {
+ // Skipping words with no kept bits before touching the values makes
+ // sparse filters cost a load and a test per word; at moderate
+ // densities the branch is never taken
+ if mask == 0 {
+ continue;
+ }
+ packer.push(value_chunks.chunk(index), mask);
Review Comment:
| Rows kept | scalar chunk(index) | scalar zip | Δ | bmi2 chunk(index) |
bmi2 zip | Δ |
|---|---|---|---|---|---|---|
| 9/10 | 31.8 µs | 33.7 µs | +6% | 1.68 µs | 1.90 µs | +13% |
| 1/2 | 19.1 µs | 21.9 µs | +15% | 1.59 µs | 1.74 µs | +9% |
| 1/10 | 4.23 µs | 4.51 µs | +7% | 1.41 µs | 1.46 µs | +4% |
| 1/64 | 1.50 µs | 1.58 µs | +5% | 968 ns | 1025 ns | +6% |
| 1/256 | 641 ns | 761 ns | +19% | 558 ns | 599 ns | +7% |
| 1/1024 | 348 ns | 465 ns | +34% | 326 ns | 432 ns | +33% |
| sliced 1/2 | 20.3 µs | 21.6 µs | +6% | 1.84 µs | 2.00 µs | +9% |
| sliced 1/1024 | 360 ns | 546 ns | +52% | 353 ns | 502 ns | +42% |
It appears that x86 codegen does better using the chunk api rather than the
zip. I think we should keep the `chunk` implementation and I'll add some
unit-tests for coverage + adjust the comment.
--
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]
