mightsleep commented on code in PR #11554:
URL: https://github.com/apache/arrow-rs/pull/11554#discussion_r4351488583
##########
parquet/src/arrow/arrow_reader/selection/algebra.rs:
##########
@@ -484,38 +484,109 @@ fn and_then_sparse_masks(mask: &BooleanBuffer, other:
&BooleanBuffer) -> Boolean
fn and_then_dense_masks(mask: &BooleanBuffer, other: &BooleanBuffer) ->
BooleanBuffer {
debug_assert_eq!(other.len(), mask.count_set_bits());
- let mut other_chunks = other.bit_chunks().iter_padded();
- let mut other_remaining = other.len();
- let mut pending = 0_u128;
- let mut pending_len = 0;
- let mut output = MutableBuffer::with_capacity(mask.len().div_ceil(8));
-
- for mask_word in mask
- .bit_chunks()
- .iter_padded()
- .take(mask.len().div_ceil(64))
+ let other_bits = BitReader::new(other);
+ let mut output = MutableBuffer::from_len_zeroed(mask.len().div_ceil(64) *
8);
+ let mut words = output
+ .typed_data_mut::<u64>()
+ .iter_mut()
+ .zip(mask.bit_chunks().iter_padded());
+ let mut position = 0;
+
+ // A word's read depends only on `position`, not on the previous deposit,
so consecutive
+ // words overlap. Bits of `source` past the `count` taken ones are ignored
by `deposit`.
+ // Reads that would run past the backing bytes go to the loop below.
+ while let Some(source) = other_bits.read(position) {
+ let Some((slot, mask_word)) = words.next() else {
+ break;
+ };
+ let (word, count) = deposit(source, mask_word);
+ *slot = word;
+ position += count;
+ }
+ for (slot, mask_word) in words {
+ let (word, count) = deposit(other_bits.read_padded(position),
mask_word);
+ *slot = word;
+ position += count;
+ }
+
+ debug_assert_eq!(position, other.len());
+ output.truncate(mask.len().div_ceil(8));
+ BooleanBuffer::new(output.into(), 0, mask.len())
+}
+
+/// Returns the low `mask_word.count_ones()` bits of `source` scattered into
the set bits of
+/// `mask_word`, and that count.
+#[inline]
+fn deposit(source: u64, mask_word: u64) -> (u64, usize) {
+ // Without `pdep`, `expand` loops over bits. Filters on row groups and
pages leave whole
+ // words set or clear, in `mask_word` or in the taken bits of `source`,
and those skip it.
+ // With `pdep` these branches cost more than they save.
+ #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
Review Comment:
Yes hello. after your comment I tried the network, and it turned into
something more interesting than inverting compress_bytes. It's not done yet. I
still need to try it on more shapes and on AArch64 and let it settle a bit, so
I'd rather not hold this PR on it. In the meantime I did roughly what you
suggested: the shortcuts now live in the expand fallback, so every caller gets
them, and deposit keeps only the full-word check. and_then stays where it was
(geomean 0.99 vs the previous push). If that's ok with you, I'll update the PR
description.
--
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]
