sdf-jkl commented on code in PR #11055:
URL: https://github.com/apache/arrow-rs/pull/11055#discussion_r4008588170
##########
arrow-select/src/filter.rs:
##########
@@ -676,6 +677,74 @@ where
RunArray::try_new(&run_ends, &values)
}
+/// Extract bits from `src` at positions where `filter` has a 1, packed
densely,
+/// processing 64 filter bits at a time
+fn gather_bits(src: &BooleanBuffer, filter: &BooleanBuffer, count: usize) ->
Buffer {
+ let filter_chunks = filter.bit_chunks();
+ let src_chunks = BitChunks::new(src.values(), src.offset(), filter.len());
+
+ let out_u64s = bit_util::ceil(count, 64);
+ let mut out: Vec<u64> = Vec::with_capacity(out_u64s);
+ let mut current_word = 0u64;
+ let mut bits_filled = 0usize;
+
+ // Appends `bits_selected` densely-packed bits from `selected_bits` into
the output word stream.
+ let mut push_chunk = |selected_bits: u64, bits_selected: usize| {
+ let bits_remaining = 64 - bits_filled;
+ current_word |= selected_bits << bits_filled;
+ if bits_selected < bits_remaining {
+ bits_filled += bits_selected;
+ } else {
+ // Current word is full; carry the overflow into the next word.
+ out.push(current_word);
+ current_word = if bits_selected == bits_remaining {
+ 0
+ } else {
+ selected_bits >> bits_remaining
+ };
+ bits_filled = bits_selected - bits_remaining;
+ }
+ };
+
+ for (filter_word, src_word) in filter_chunks.iter().zip(src_chunks.iter())
{
+ if filter_word == 0 {
+ continue;
+ }
+ let n_set = filter_word.count_ones() as usize;
+ push_chunk(pext64(src_word, filter_word), n_set);
+ }
+
+ let rem_filter = filter_chunks.remainder_bits();
+ if rem_filter != 0 {
+ let n_set = rem_filter.count_ones() as usize;
+ push_chunk(pext64(src_chunks.remainder_bits(), rem_filter), n_set);
+ }
+
+ if bits_filled > 0 {
+ out.push(current_word);
+ }
+
+ let mut buf: MutableBuffer = out.into();
+ buf.truncate(bit_util::ceil(count, 8));
+ buf.into()
+}
+
+/// Collects the bits of `val` wherever `mask` is 1, packed into the low bits
of the result.
+#[inline(always)]
+fn pext64(val: u64, mut mask: u64) -> u64 {
Review Comment:
There already is an implementation of PEXT in the codebase --
https://github.com/apache/arrow-rs/blob/4727e6a654d28f6371a9fd352a34d8a77af9c0e8/parquet/src/util/bit_util.rs#L959-L984
Yours seems to perform better though (on my machine 🤓 ) If anything we can
drop drop the parquet one and make it import your implementation.
Rust has a nightly impl--
doc.rust-lang.org/std/primitive.u64.html#method.extract_bits
The issue here - https://github.com/rust-lang/rust/issues/149069 - explains
that the implementation is waiting on supporting the new LLVM intrinsics that
support automatically using the PEXT/PDEP machine instructions if machine
supports them. When the instructions are available the perf is pretty epic.
##########
arrow-select/src/filter.rs:
##########
@@ -676,6 +677,74 @@ where
RunArray::try_new(&run_ends, &values)
}
+/// Extract bits from `src` at positions where `filter` has a 1, packed
densely,
+/// processing 64 filter bits at a time
+fn gather_bits(src: &BooleanBuffer, filter: &BooleanBuffer, count: usize) ->
Buffer {
+ let filter_chunks = filter.bit_chunks();
+ let src_chunks = BitChunks::new(src.values(), src.offset(), filter.len());
+
+ let out_u64s = bit_util::ceil(count, 64);
+ let mut out: Vec<u64> = Vec::with_capacity(out_u64s);
+ let mut current_word = 0u64;
+ let mut bits_filled = 0usize;
+
+ // Appends `bits_selected` densely-packed bits from `selected_bits` into
the output word stream.
+ let mut push_chunk = |selected_bits: u64, bits_selected: usize| {
+ let bits_remaining = 64 - bits_filled;
+ current_word |= selected_bits << bits_filled;
+ if bits_selected < bits_remaining {
+ bits_filled += bits_selected;
+ } else {
+ // Current word is full; carry the overflow into the next word.
+ out.push(current_word);
+ current_word = if bits_selected == bits_remaining {
+ 0
+ } else {
+ selected_bits >> bits_remaining
+ };
+ bits_filled = bits_selected - bits_remaining;
+ }
+ };
+
+ for (filter_word, src_word) in filter_chunks.iter().zip(src_chunks.iter())
{
Review Comment:
Processing each word is independent from each other so this _could_ be
paralellizable.
I killed some time on it today. You can take a look here --
https://github.com/apache/arrow-rs/pull/11086
--
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]
