alamb commented on code in PR #10136:
URL: https://github.com/apache/arrow-rs/pull/10136#discussion_r4087289204
##########
arrow-buffer/src/util/bit_chunk_iterator.rs:
##########
@@ -360,6 +374,38 @@ impl<'a> IntoIterator for &BitChunks<'a> {
}
}
+/// Reads the `index`th complete 64-bit chunk of `buffer`, whose bits start
+/// at `bit_offset` (in `0..8`)
+///
+/// # Safety
+///
+/// `index` must be less than the number of complete chunks, so that the
+/// buffer holds at least `index * 8 + 8` bytes, plus one more byte when
+/// `bit_offset != 0` (the remainder byte the constructor guarantees)
+#[inline]
+unsafe fn read_chunk(buffer: &[u8], bit_offset: usize, index: usize) -> u64 {
+ // cast to *const u64 should be fine since we are using read_unaligned
below
Review Comment:
I recommend slapping some `debug_assert` in here to make the safety claims
precise
```rust
debug_assert!(bit_offset != 0);
debug_assert!(index < self.chunk_len, "chunk index out of bounds");
```
##########
arrow-select/src/filter.rs:
##########
@@ -676,8 +677,27 @@ where
RunArray::try_new(&run_ends, &values)
}
-/// Filter the packed bitmask `buffer`, with `predicate` starting at bit
offset `offset`
+/// Filter the packed bitmask `buffer` with `predicate`, choosing between the
+/// strategy-based and compress-based kernels by filter density
fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
+ // Compressing scans the whole mask a word at a time, so it loses to the
+ // slices strategies once fewer than one bit per word is dropped, and to
+ // precomputed `Indices` once fewer than one bit per word is kept. The lazy
+ // `IndexIterator` scans the mask anyway, so it never beats compressing
+ let len = predicate.filter.len();
+ let count = predicate.count;
+ let dense = count >= len - len / 64;
+ let sparse_indices =
+ count <= len / 64 && matches!(predicate.strategy,
IterationStrategy::Indices(_));
+ if !dense && !sparse_indices {
+ return filter_bits_compress(buffer, predicate);
Review Comment:
Would this be more naturally expressed as a method on `IterationStrategy`
(and maybe a new enum variant?)
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -19,6 +19,45 @@
use crate::bit_chunk_iterator::BitChunks;
+/// Parallel bit extract: for each set bit in `mask`, extract the
+/// corresponding bit from `value` and pack them contiguously into the low
+/// bits of the return value.
+///
+/// Equivalent to the x86 BMI2 `PEXT` instruction. When compiled with the
+/// `bmi2` target feature enabled (for example `-C target-cpu=x86-64-v3`)
+/// this lowers to the hardware `pext` instruction; otherwise it falls back
+/// to a portable scalar loop.
Review Comment:
I always like examples; here is a propoal:
```suggestion
/// to a portable scalar loop.
///
/// # Functional Example
///
/// Using 8 bits for brevity (the function operates on all 64). Each
/// set bit in `mask` selects the bit at the same position in `value`; the
/// selected bits are then shifted down so they are contiguous in the low
/// bits of the result, in their original order:
///
/// ```text
/// bit: 7 6 5 4 3 2 1 0
/// value: a b c d e f g h
/// mask: 0 1 1 0 1 1 0 1 set bits select b, c, e, f and h
/// | | | | |
/// v v v v v copy the relevant bits into result
/// result: 0 0 0 b c e f h
/// ```
///
/// # Code Example
///
/// ```
/// # use arrow_buffer::bit_util::compress;
/// assert_eq!(compress(0b1011_0100, 0b0110_1101), 0b0000_1010);
/// ```
```
##########
arrow-select/src/filter.rs:
##########
@@ -676,8 +677,27 @@ where
RunArray::try_new(&run_ends, &values)
}
-/// Filter the packed bitmask `buffer`, with `predicate` starting at bit
offset `offset`
+/// Filter the packed bitmask `buffer` with `predicate`, choosing between the
+/// strategy-based and compress-based kernels by filter density
fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
+ // Compressing scans the whole mask a word at a time, so it loses to the
+ // slices strategies once fewer than one bit per word is dropped, and to
+ // precomputed `Indices` once fewer than one bit per word is kept. The lazy
+ // `IndexIterator` scans the mask anyway, so it never beats compressing
+ let len = predicate.filter.len();
+ let count = predicate.count;
+ let dense = count >= len - len / 64;
+ let sparse_indices =
+ count <= len / 64 && matches!(predicate.strategy,
IterationStrategy::Indices(_));
+ if !dense && !sparse_indices {
+ return filter_bits_compress(buffer, predicate);
Review Comment:
I kind of like how it looked in @Rich-T-kid 's pR here
- https://github.com/apache/arrow-rs/pull/11055
where it was added to the relevant IterationStrategy branches
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -19,6 +19,45 @@
use crate::bit_chunk_iterator::BitChunks;
+/// Parallel bit extract: for each set bit in `mask`, extract the
+/// corresponding bit from `value` and pack them contiguously into the low
+/// bits of the return value.
+///
+/// Equivalent to the x86 BMI2 `PEXT` instruction. When compiled with the
+/// `bmi2` target feature enabled (for example `-C target-cpu=x86-64-v3`)
+/// this lowers to the hardware `pext` instruction; otherwise it falls back
+/// to a portable scalar loop.
+//
+// Replace with `value.compress(mask)` when `uint_gather_scatter_bits` is
+// stabilised: <https://github.com/rust-lang/rust/issues/149069>
+#[inline]
+pub fn compress(value: u64, mask: u64) -> u64 {
+ #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
+ {
+ // SAFETY: the `bmi2` target feature is statically enabled for this
+ // build, so the `pext` instruction is guaranteed to be available.
+ unsafe { std::arch::x86_64::_pext_u64(value, mask) }
+ }
+
+ #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
+ {
+ let mut mask = mask;
+ let mut result = 0_u64;
+ let mut dest_bit = 1_u64;
+ while mask != 0 {
Review Comment:
Claude flagged this as not very efficient as it is still doing bit at a time
Here is some version it wrote that operates using a precomputed table
(citing the [Hacker's
Delight](https://www.amazon.com/dp/0321842685?lv=shuf&channelId=500&plpRedirect=mhFallback),
one of my favorite books -- I think @XiangpengHao has a copy!)
Perhaps we can flag this as a potential improvement for follow on work in
another ticket
```rust
{
// Nibble at a time through a 256 byte table: 16 steps per word
// regardless of how many bits `mask` selects
let mut result = 0_u64;
let mut shift = 0_u32;
let mut i = 0;
while i < 16 {
let m = ((mask >> (i * 4)) & 0xF) as usize;
let v = ((value >> (i * 4)) & 0xF) as usize;
result |= (COMPRESS_NIBBLE[m][v] as u64) << shift;
shift += m.count_ones();
i += 1;
}
result
}
/// `COMPRESS_NIBBLE[mask][value]` is `compress(value, mask)` for 4-bit
inputs
#[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
static COMPRESS_NIBBLE: [[u8; 16]; 16] = {
let mut t = [[0u8; 16]; 16];
let mut m = 0;
while m < 16 {
let mut v = 0;
while v < 16 {
let mut r = 0u8;
let mut d = 0;
let mut b = 0;
while b < 4 {
if m >> b & 1 == 1 {
r |= ((v >> b & 1) as u8) << d;
d += 1;
}
b += 1;
}
t[m][v] = r;
v += 1;
}
m += 1;
}
t
};
```
--
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]
