feat: BMI for bit filtering

arrow-buffer/src/util/bit_util.rs

@@ -19,6 +19,49 @@
 
 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: u64 = 0;
+        let mut dest_bit: u64 = 1;
+        while mask != 0 {
+            let lowest = mask & mask.wrapping_neg();
+            if value & lowest != 0 {
+                result |= dest_bit;
+            }
+            dest_bit <<= 1;
+            mask ^= lowest;
+        }
+        result
+    }
+}
+
+/// Returns true if [`compress`] lowers to the hardware `pext` instruction
+#[inline]
+pub fn compress_available() -> bool {
+    cfg!(all(target_arch = "x86_64", target_feature = "bmi2"))
+}
+
 /// Returns the nearest number that is `>=` than `num` and is a multiple of 64
 #[inline]
 pub fn round_upto_multiple_of_64(num: usize) -> usize {
@@ -825,8 +868,55 @@ mod tests {
     use super::*;
     use crate::bit_iterator::BitIterator;
     use crate::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer};
+    use rand::distr::{Distribution, StandardUniform};
     use rand::rngs::StdRng;
-    use rand::{Rng, SeedableRng};
+    use rand::{Rng, SeedableRng, rng};
+
+    fn random_numbers<T>(n: usize) -> Vec<T>
+    where
+        StandardUniform: Distribution<T>,
+    {
+        let mut rng = rng();
+        StandardUniform.sample_iter(&mut rng).take(n).collect()
+    }
+
+    #[test]
+    fn test_compress() {
+        // Reference: gather the `mask`-selected bits of `value` into
+        // contiguous low bits, least-significant first.
+        fn reference(value: u64, mut mask: u64) -> u64 {
+            let mut result = 0u64;
+            let mut dest = 0u32;
+            while mask != 0 {
+                let lowest = mask & mask.wrapping_neg();
+                result |= (((value & lowest) != 0) as u64) << dest;
+                dest += 1;
+                mask ^= lowest;
+            }
+            result
+        }
+
+        // Hand-picked edge cases.
+        assert_eq!(compress(0b1010, 0b1111), 0b1010);
+        assert_eq!(compress(0b1010, 0b1010), 0b11);
+        assert_eq!(compress(0b1010, 0b0101), 0);
+        assert_eq!(compress(u64::MAX, 0), 0);
+        assert_eq!(compress(0, u64::MAX), 0);
+        assert_eq!(compress(u64::MAX, u64::MAX), u64::MAX);
+
+        // Randomized cross-check against the reference. On a `bmi2` build
+        // this validates the hardware `pext` path; otherwise it exercises
+        // the portable fallback.
+        let values = random_numbers::<u64>(1024);
+        let masks = random_numbers::<u64>(1024);
+        for (&value, &mask) in values.iter().zip(masks.iter()) {
+            assert_eq!(
+                compress(value, mask),
+                reference(value, mask),
+                "value={value:#x} mask={mask:#x}"
+            );
+        }
+    }
 
     #[test]
     fn test_round_upto_multiple_of_64() {

arrow-select/Cargo.toml

@@ -44,4 +44,9 @@ num-traits = { version = "0.2.19", default-features = false, features = ["std"]
 ahash = { version = "0.8", default-features = false}
 
 [dev-dependencies]
+criterion = { workspace = true, default-features = false }
 rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"] }
+
+[[bench]]
+name = "filter_bits"
+harness = false

arrow-select/benches/filter_bits.rs

@@ -0,0 +1,88 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Benchmarks for the internal `filter_bits` kernel.
+//!
+//! `filter_bits` is private, so it is exercised through
+//! [`FilterPredicate::filter`] on a [`BooleanArray`] without nulls, which
+//! dispatches directly to `filter_bits` on the array's value buffer.
+//!
+//! The filter selectivity determines which `IterationStrategy` is used:
+//! selectivity above 0.8 selects `SlicesIterator`, below selects
+//! `IndexIterator`, and [`FilterBuilder::optimize`] converts these into their
+//! precomputed `Slices` / `Indices` counterparts.
+
+use arrow_array::BooleanArray;
+use arrow_select::filter::{FilterBuilder, FilterPredicate};
+use criterion::{Criterion, criterion_group, criterion_main};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use std::hint;
+
+fn create_boolean_array(size: usize, true_density: f64, rng: &mut StdRng) -> BooleanArray {
+    (0..size)
+        .map(|_| Some(rng.random_bool(true_density)))
+        .collect()
+}
+
+fn bench_filter_bits(predicate: &FilterPredicate, array: &BooleanArray) {
+    hint::black_box(predicate.filter(array).unwrap());
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    const SIZE: usize = 65536;
+    let mut rng = StdRng::seed_from_u64(42);
+
+    let data = create_boolean_array(SIZE, 0.5, &mut rng);
+
+    // Slice off a non-byte-aligned prefix to exercise the bit offset handling
+    // in `filter_bits`
+    let padded = create_boolean_array(SIZE + 3, 0.5, &mut rng);
+    let sliced = padded.slice(3, SIZE);
+
+    // (label, true_density): densities above the 0.8 selectivity threshold use
+    // the slices strategies, those below use the index strategies
+    let cases = [
+        ("slices, kept 1023/1024", 1.0 - 1.0 / 1024.0),
+        ("slices, kept 9/10", 0.9),
+        ("indices, kept 1/2", 0.5),
+        ("indices, kept 1/10", 0.1),
+        ("indices, kept 1/1024", 1.0 / 1024.0),
+    ];
+
+    for (label, true_density) in cases {
+        let filter_array = create_boolean_array(SIZE, true_density, &mut rng);
+
+        // Lazy strategies: SlicesIterator / IndexIterator
+        let lazy = FilterBuilder::new(&filter_array).build();
+        // Precomputed strategies: Slices / Indices
+        let optimized = FilterBuilder::new(&filter_array).optimize().build();
+
+        for (suffix, predicate, array) in [
+            ("", &lazy, &data),
+            (" optimized", &optimized, &data),
+            (" sliced", &lazy, &sliced),
+        ] {
+            c.bench_function(&format!("filter_bits{suffix} ({label})"), |b| {
+                b.iter(|| bench_filter_bits(predicate, array))
+            });
+        }
+    }
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

arrow-select/src/filter.rs

@@ -26,6 +26,7 @@ use arrow_array::types::{
     ArrowDictionaryKeyType, ArrowPrimitiveType, ByteArrayType, ByteViewType, RunEndIndexType,
 };
 use arrow_array::*;
+use arrow_buffer::bit_chunk_iterator::BitChunks;
 use arrow_buffer::{
     ArrowNativeType, BooleanBuffer, NullBuffer, OffsetBuffer, RunEndBuffer, ScalarBuffer, bit_util,
 };
@@ -689,6 +690,26 @@ pub(crate) fn filter_null_mask(
 
 /// Filter the packed bitmask `buffer`, with `predicate` starting at bit offset `offset`
 fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
+    // The compress path scans the entire filter mask a word at a time, so it
+    // is only worthwhile when `pext` is available in hardware and the filter
+    // is neither very sparse nor very dense: it must keep more than one bit
+    // per word on average (otherwise visiting each kept bit individually is
+    // faster) and drop more than one bit per word on average (otherwise
+    // copying whole ranges via the slices strategies is faster)
+    let len = predicate.filter.len();
+    let predicate_count = predicate.count;
+    let upper = predicate_count < len - (len / 64);
+    let lower = predicate_count > len / 64;
+    if bit_util::compress_available() && upper && lower {
+        return filter_bits_compress(buffer, predicate);
+    }
+
+    filter_bits_strategy(buffer, predicate)
+}
+
+/// Filter the packed bitmask `buffer` with `predicate` using its
+/// [`IterationStrategy`]
+fn filter_bits_strategy(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
     let src = buffer.values();
     let offset = buffer.offset();
     assert!(buffer.len() >= predicate.filter.len());
@@ -730,6 +751,42 @@ fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
     }
 }
 
+/// Filter the packed bitmask `buffer` with `predicate` by extracting the kept
+/// bits of each 64-bit word with [`bit_util::compress`] (`pext`)
+fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
+    assert!(buffer.len() >= predicate.filter.len());
+    let mask_chunks = predicate.filter.values().bit_chunks();
+    let value_chunks = BitChunks::new(buffer.values(), buffer.offset(), predicate.filter.len());
+
+    let mut out = MutableBuffer::new(bit_util::ceil(predicate.count, 8));
+    // Bits extracted from each chunk are packed into `current` until it
+    // contains a complete word, which is then flushed to `out`
+    let mut current = 0_u64;
+    let mut filled = 0_u32;
+
+    let chunks = value_chunks.iter_padded().zip(mask_chunks.iter_padded());
+
+    for (values, mask) in chunks {
+        let bits = bit_util::compress(values, mask);
+        let count = mask.count_ones();
+        current |= bits << filled;
+        if filled + count >= 64 {
+            out.extend_from_slice(&current.to_le_bytes());
+            filled = filled + count - 64;
+            // The bits of `bits` that did not fit in `current`, if any
+            // (`checked_shr` yields 0 when the carry is a full word)
+            current = bits.checked_shr(count - filled).unwrap_or(0);
+        } else {
+            filled += count;
+        }
+    }
+
+    if filled > 0 {
+        out.extend_from_slice(&current.to_le_bytes()[..bit_util::ceil(filled as usize, 8)]);
+    }
+    out.into()
+}
+
 /// `filter` implementation for boolean buffers
 fn filter_boolean(array: &BooleanArray, predicate: &FilterPredicate) -> BooleanArray {
     let buffer = filter_bits(array.values(), predicate);
@@ -1649,6 +1706,77 @@ mod tests {
         test_case_filter_sliced_list_view::<i64>();
     }
 
+    /// Tests [`filter_bits_compress`] and [`filter_bits_strategy`] on the
+    /// same inputs against a naive bit-by-bit filter, verifying both pathways
+    /// produce the same output. Both are called directly rather than through
+    /// [`filter_bits`], whose dispatch depends on whether the build has
+    /// hardware `pext`, so both get coverage on every platform
+    #[test]
+    fn test_filter_bits() {
+        let mut rng = StdRng::seed_from_u64(42);
+
+        // Lengths exercising partial words, exact word multiples, and the
+        // carry logic across flushed words
+        let lens = [0, 1, 7, 63, 64, 65, 127, 128, 200, 1024, 4099];
+        // Densities covering empty, sparse, balanced, dense and full masks
+        let densities = [0.0, 0.01, 0.5, 0.9, 1.0];
+        // Bit offsets of the value buffer, including non byte-aligned ones
+        let offsets = [0, 3, 8, 67];
+
+        for len in lens {
+            for density in densities {
+                for offset in offsets {
+                    let values: BooleanBuffer =
+                        (0..len + offset).map(|_| rng.random_bool(0.5)).collect();
+                    let values = values.slice(offset, len);
+                    let filter: BooleanArray =
+                        (0..len).map(|_| Some(rng.random_bool(density))).collect();
+
+                    let predicate = FilterBuilder::new(&filter).build();
+
+                    let expected: BooleanBuffer = values
+                        .iter()
+                        .zip(filter.values().iter())
+                        .filter_map(|(value, keep)| keep.then_some(value))
+                        .collect();
+
+                    let compressed = filter_bits_compress(&values, &predicate);
+                    let compressed = BooleanBuffer::new(compressed, 0, predicate.count);
+
+                    assert_eq!(
+                        compressed, expected,
+                        "compress: len={len} density={density} offset={offset}"
+                    );
+
+                    // `filter_bits` is never reached with the `All` / `None`
+                    // strategies, they are short-circuited by the callers
+                    match predicate.strategy {
+                        IterationStrategy::All | IterationStrategy::None => continue,
+                        _ => {}
+                    }
+
+                    let strategy = filter_bits_strategy(&values, &predicate);
+                    let strategy = BooleanBuffer::new(strategy, 0, predicate.count);
+
+                    assert_eq!(
+                        strategy, expected,
+                        "{:?}: len={len} density={density} offset={offset}",
+                        predicate.strategy
+                    );
+
+                    // Also cover the dispatch between the two pathways
+                    let dispatched = filter_bits(&values, &predicate);
+                    let dispatched = BooleanBuffer::new(dispatched, 0, predicate.count);
+
+                    assert_eq!(
+                        dispatched, expected,
+                        "dispatch: len={len} density={density} offset={offset}"
+                    );
+                }
+            }
+        }
+    }
+
     #[test]
     fn test_slice_iterator_bits() {
         let filter_values = (0..64).map(|i| i == 1).collect::<Vec<bool>>();