mbutrovich commented on code in PR #10136:
URL: https://github.com/apache/arrow-rs/pull/10136#discussion_r4086150523
##########
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:
Thanks for looking into it!
--
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]
