Jefffrey commented on code in PR #11918:
URL: https://github.com/apache/arrow-rs/pull/11918#discussion_r4490311291
##########
arrow-select/src/filter.rs:
##########
@@ -800,81 +810,88 @@ fn filter_bits_strategy(buffer: &BooleanBuffer,
predicate: &FilterPredicate) ->
}
}
+/// Packs bits selected from successive words into complete output words.
+#[derive(Default)]
+pub(crate) struct Packer {
+ /// Bits waiting to be written.
+ current: u64,
+ /// Number of populated low bits in `current`.
+ filled: u32,
+}
+
+impl Packer {
+ #[inline(always)]
+ pub(crate) fn push(&mut self, values: u64, mask: u64) -> Option<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;
+ return None;
+ }
+
+ let complete = self.current;
+ // `bits >> (64 - filled)`, written so that `filled == 0`
+ // shifts everything out
+ self.current = (bits >> 1) >> (63 - self.filled);
+ self.filled = total - 64;
+ Some(complete)
+ }
+
+ #[inline]
+ pub(crate) fn finish(self) -> Option<(u64, usize)> {
+ (self.filled != 0).then_some((self.current, self.filled as usize))
+ }
+}
+
/// Filter the packed bitmask `buffer` with `predicate` by extracting the kept
-/// bits of each 64-bit word with [`bit_util::compress`] (`pext`)
+/// bits of each 64-bit word with [`bit_util::compress`] (`pext`).
///
-/// Not inlined: within `filter_array` the packing state spills to the stack
+/// 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
Review Comment:
any reason we're removing this comment?
##########
arrow-select/src/filter.rs:
##########
@@ -800,81 +810,88 @@ fn filter_bits_strategy(buffer: &BooleanBuffer,
predicate: &FilterPredicate) ->
}
}
+/// Packs bits selected from successive words into complete output words.
+#[derive(Default)]
+pub(crate) struct Packer {
+ /// Bits waiting to be written.
+ current: u64,
+ /// Number of populated low bits in `current`.
+ filled: u32,
+}
+
+impl Packer {
+ #[inline(always)]
+ pub(crate) fn push(&mut self, values: u64, mask: u64) -> Option<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;
+ return None;
+ }
+
+ let complete = self.current;
+ // `bits >> (64 - filled)`, written so that `filled == 0`
+ // shifts everything out
+ self.current = (bits >> 1) >> (63 - self.filled);
+ self.filled = total - 64;
+ Some(complete)
+ }
+
+ #[inline]
+ pub(crate) fn finish(self) -> Option<(u64, usize)> {
+ (self.filled != 0).then_some((self.current, self.filled as usize))
+ }
+}
+
/// Filter the packed bitmask `buffer` with `predicate` by extracting the kept
-/// bits of each 64-bit word with [`bit_util::compress`] (`pext`)
+/// bits of each 64-bit word with [`bit_util::compress`] (`pext`).
///
-/// Not inlined: within `filter_array` the packing state spills to the stack
+/// 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,
- };
+ let ptr = out.as_mut_ptr();
+ let mut idx = 0;
+ let mut packer = Packer::default();
for (index, mask) in mask_chunks.iter().enumerate() {
// Words with no kept bits are skipped before the corresponding values
// are read, so only the mask is touched for them
if mask == 0 {
continue;
}
- packer.push(value_chunks.chunk(index), mask);
+ if let Some(complete) = packer.push(value_chunks.chunk(index), mask) {
+ debug_assert!(idx < out.capacity());
+ // SAFETY: `count` is the number of set bits in the filter, so at
+ // most `count / 64` words are completed and the output has room
Review Comment:
this safety comment is a bit vague, would be good to tighten it up; for
example it refers to `count` but its not clear this is referring to
`predicate.count` (and even then its hard to see the relation of that to the
safety here)
--
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]
