mightsleep commented on code in PR #11554:
URL: https://github.com/apache/arrow-rs/pull/11554#discussion_r4286639711
##########
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:
Agreed, they belong in expand, and the merge kernels from #11437 will call
it too (there the taken validity bits are usually all set, so they hit often).
My next PR replaces the fallback of expand with a byte network (the inverse of
compress_bytes from #11322, 1.7 to 3x on mixed words) and needs a small
dispatch by bit count anyway, so I'd move these shortcuts into that dispatch
there and drop them from deposit, rather than change expand twice. is it ok to
keep them here until then?
--
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]
