alamb commented on code in PR #11554:
URL: https://github.com/apache/arrow-rs/pull/11554#discussion_r4286533054
##########
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:
This sounds to me like an argument for improving the fallback implementation
of `bit_util::expand` in general
What is the rationale for using a different fallback (only) 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]
