From 3a64b9eb294ec5c6defcb4d1f4e5a3ec977f9743 Mon Sep 17 00:00:00 2001 From: Devan Date: Fri, 12 Jun 2026 07:46:54 -0500 Subject: [PATCH 1/7] feat: Working on BMI impl for filter_bits --- Cargo.lock | 1 + arrow-buffer/src/util/bit_util.rs | 65 +++++++++++++++++++++ arrow-select/Cargo.toml | 5 ++ arrow-select/benches/filter_bits.rs | 90 +++++++++++++++++++++++++++++ arrow-select/src/filter.rs | 46 +++++++++++++++ 5 files changed, 207 insertions(+) create mode 100644 arrow-select/benches/filter_bits.rs diff --git a/Cargo.lock b/Cargo.lock index af1c980d45da..a1fd86831b31 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -511,6 +511,7 @@ dependencies = [ "arrow-buffer", "arrow-data", "arrow-schema", + "criterion", "num-traits", "rand 0.9.4", ] diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index 920ecba6d60c..e8ac1de05909 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -19,6 +19,71 @@ use crate::bit_chunk_iterator::BitChunks; +/// Extracts the bits of `value` selected by `mask` and packs them into the +/// least significant bits of the result (parallel bits extract, `pext`) +#[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(all(target_arch = "x86_64", not(target_feature = "bmi2")))] + { + if std::arch::is_x86_feature_detected!("bmi2") { + // SAFETY: `bmi2` support was detected at runtime + unsafe { compress_bmi2(value, mask) } + } else { + compress_fallback(value, mask) + } + } + + #[cfg(not(target_arch = "x86_64"))] + { + compress_fallback(value, mask) + } +} + +#[cfg(all(target_arch = "x86_64", not(target_feature = "bmi2")))] +#[target_feature(enable = "bmi2")] +#[inline] +unsafe fn compress_bmi2(value: u64, mask: u64) -> u64 { + // SAFETY: the caller guarantees the `bmi2` target feature is available + unsafe { std::arch::x86_64::_pext_u64(value, mask) } +} + +#[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] +#[inline] +fn compress_fallback(value: u64, mask: u64) -> u64 { + 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`] is hardware accelerated (`pext`) +#[inline] +pub fn compress_available() -> bool { + #[cfg(target_arch = "x86_64")] + { + std::arch::is_x86_feature_detected!("bmi2") + } + #[cfg(not(target_arch = "x86_64"))] + { + false + } +} + /// 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 { diff --git a/arrow-select/Cargo.toml b/arrow-select/Cargo.toml index 443094e6c986..3e124d0728e9 100644 --- a/arrow-select/Cargo.toml +++ b/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 diff --git a/arrow-select/benches/filter_bits.rs b/arrow-select/benches/filter_bits.rs new file mode 100644 index 000000000000..380d0de6bb40 --- /dev/null +++ b/arrow-select/benches/filter_bits.rs @@ -0,0 +1,90 @@ +// 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(); + + c.bench_function(&format!("filter_bits ({label})"), |b| { + b.iter(|| bench_filter_bits(&lazy, &data)) + }); + + c.bench_function(&format!("filter_bits optimized ({label})"), |b| { + b.iter(|| bench_filter_bits(&optimized, &data)) + }); + + c.bench_function(&format!("filter_bits sliced ({label})"), |b| { + b.iter(|| bench_filter_bits(&lazy, &sliced)) + }); + } +} + +criterion_group!(benches, add_benchmark); +criterion_main!(benches); diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index b1f3a21a3ed3..c116aefad90f 100644 --- a/arrow-select/src/filter.rs +++ b/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, }; @@ -679,6 +680,14 @@ fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { let offset = buffer.offset(); assert!(buffer.len() >= predicate.filter.len()); + // 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 + // keeps more than one bit per word on average; otherwise visiting each + // kept bit individually is faster + if bit_util::compress_available() && predicate.count > predicate.filter.len() / 64 { + return filter_bits_compress(buffer, predicate); + } + match &predicate.strategy { IterationStrategy::IndexIterator => { let bits = @@ -716,6 +725,43 @@ 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 { + 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().zip(mask_chunks.iter()).chain([( + value_chunks.remainder_bits(), + mask_chunks.remainder_bits(), + )]); + + 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(¤t.to_le_bytes()); + filled = filled + count - 64; + // The bits of `bits` that did not fit in `current`, if any + current = if filled == 0 { 0 } else { bits >> (count - filled) }; + } else { + filled += count; + } + } + + if filled > 0 { + out.extend_from_slice(¤t.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); From 903f9c3d2c1f6cab45a06d292c04636de2acbd7d Mon Sep 17 00:00:00 2001 From: Devan Date: Fri, 12 Jun 2026 07:55:45 -0500 Subject: [PATCH 2/7] chore: use code from https://github.com/apache/arrow-rs/pull/9848/changes#diff-f2fe5b601534125d0ba033762b62db93ae2f81c71a543b6694df33773ebfa308R939-R976 --- arrow-buffer/src/util/bit_util.rs | 72 +++++++++++-------------------- 1 file changed, 25 insertions(+), 47 deletions(-) diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index e8ac1de05909..1c0dbc8024b3 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -19,8 +19,17 @@ use crate::bit_chunk_iterator::BitChunks; -/// Extracts the bits of `value` selected by `mask` and packs them into the -/// least significant bits of the result (parallel bits extract, `pext`) +/// 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: #[inline] pub fn compress(value: u64, mask: u64) -> u64 { #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))] @@ -30,58 +39,27 @@ pub fn compress(value: u64, mask: u64) -> u64 { unsafe { std::arch::x86_64::_pext_u64(value, mask) } } - #[cfg(all(target_arch = "x86_64", not(target_feature = "bmi2")))] + #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] { - if std::arch::is_x86_feature_detected!("bmi2") { - // SAFETY: `bmi2` support was detected at runtime - unsafe { compress_bmi2(value, mask) } - } else { - compress_fallback(value, mask) - } - } - - #[cfg(not(target_arch = "x86_64"))] - { - compress_fallback(value, mask) - } -} - -#[cfg(all(target_arch = "x86_64", not(target_feature = "bmi2")))] -#[target_feature(enable = "bmi2")] -#[inline] -unsafe fn compress_bmi2(value: u64, mask: u64) -> u64 { - // SAFETY: the caller guarantees the `bmi2` target feature is available - unsafe { std::arch::x86_64::_pext_u64(value, mask) } -} - -#[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] -#[inline] -fn compress_fallback(value: u64, mask: u64) -> u64 { - 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; + 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; } - dest_bit <<= 1; - mask ^= lowest; + result } - result } -/// Returns true if [`compress`] is hardware accelerated (`pext`) +/// Returns true if [`compress`] lowers to the hardware `pext` instruction #[inline] pub fn compress_available() -> bool { - #[cfg(target_arch = "x86_64")] - { - std::arch::is_x86_feature_detected!("bmi2") - } - #[cfg(not(target_arch = "x86_64"))] - { - false - } + cfg!(all(target_arch = "x86_64", target_feature = "bmi2")) } /// Returns the nearest number that is `>=` than `num` and is a multiple of 64 From 76c7efc6540ee8efc694fab80ad63569de800475 Mon Sep 17 00:00:00 2001 From: devanbenz Date: Fri, 12 Jun 2026 11:51:57 -0500 Subject: [PATCH 3/7] fix: attempt to resolve regression for dense bitmaps --- arrow-select/src/filter.rs | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index c116aefad90f..e8c181db9a86 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -682,9 +682,15 @@ 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 - // keeps more than one bit per word on average; otherwise visiting each - // kept bit individually is faster - if bit_util::compress_available() && predicate.count > predicate.filter.len() / 64 { + // 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(); + if bit_util::compress_available() + && predicate.count > len / 64 + && predicate.count < len - len / 64 + { return filter_bits_compress(buffer, predicate); } @@ -737,10 +743,10 @@ fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> let mut current = 0_u64; let mut filled = 0_u32; - let chunks = value_chunks.iter().zip(mask_chunks.iter()).chain([( - value_chunks.remainder_bits(), - mask_chunks.remainder_bits(), - )]); + let chunks = value_chunks + .iter() + .zip(mask_chunks.iter()) + .chain([(value_chunks.remainder_bits(), mask_chunks.remainder_bits())]); for (values, mask) in chunks { let bits = bit_util::compress(values, mask); @@ -750,7 +756,11 @@ fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> out.extend_from_slice(¤t.to_le_bytes()); filled = filled + count - 64; // The bits of `bits` that did not fit in `current`, if any - current = if filled == 0 { 0 } else { bits >> (count - filled) }; + current = if filled == 0 { + 0 + } else { + bits >> (count - filled) + }; } else { filled += count; } From 8e326b4dd2fd5ef0522515c9ba2ec12fa0306ec8 Mon Sep 17 00:00:00 2001 From: devanbenz Date: Sat, 27 Jun 2026 15:22:55 -0500 Subject: [PATCH 4/7] feat: simplify code a bit --- arrow-select/benches/filter_bits.rs | 20 +++++++++----------- arrow-select/src/filter.rs | 12 +++--------- 2 files changed, 12 insertions(+), 20 deletions(-) diff --git a/arrow-select/benches/filter_bits.rs b/arrow-select/benches/filter_bits.rs index 380d0de6bb40..2643ecf7fbdd 100644 --- a/arrow-select/benches/filter_bits.rs +++ b/arrow-select/benches/filter_bits.rs @@ -72,17 +72,15 @@ fn add_benchmark(c: &mut Criterion) { // Precomputed strategies: Slices / Indices let optimized = FilterBuilder::new(&filter_array).optimize().build(); - c.bench_function(&format!("filter_bits ({label})"), |b| { - b.iter(|| bench_filter_bits(&lazy, &data)) - }); - - c.bench_function(&format!("filter_bits optimized ({label})"), |b| { - b.iter(|| bench_filter_bits(&optimized, &data)) - }); - - c.bench_function(&format!("filter_bits sliced ({label})"), |b| { - b.iter(|| bench_filter_bits(&lazy, &sliced)) - }); + 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)) + }); + } } } diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 871b449e9c95..0d7d6cbbb42e 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -755,10 +755,7 @@ fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> let mut current = 0_u64; let mut filled = 0_u32; - let chunks = value_chunks - .iter() - .zip(mask_chunks.iter()) - .chain([(value_chunks.remainder_bits(), mask_chunks.remainder_bits())]); + let chunks = value_chunks.iter_padded().zip(mask_chunks.iter_padded()); for (values, mask) in chunks { let bits = bit_util::compress(values, mask); @@ -768,11 +765,8 @@ fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> out.extend_from_slice(¤t.to_le_bytes()); filled = filled + count - 64; // The bits of `bits` that did not fit in `current`, if any - current = if filled == 0 { - 0 - } else { - bits >> (count - filled) - }; + // (`checked_shr` yields 0 when the carry is a full word) + current = bits.checked_shr(count - filled).unwrap_or(0); } else { filled += count; } From 5bbbdd80b98c20425b49bb03f0009d12e6da1a70 Mon Sep 17 00:00:00 2001 From: devanbenz Date: Thu, 2 Jul 2026 10:18:47 -0500 Subject: [PATCH 5/7] feat: Moves compress from parquet workspace to arrow-buffer --- arrow-buffer/src/util/bit_util.rs | 49 ++++++++++++- .../arrow/record_reader/definition_levels.rs | 12 ++- parquet/src/util/bit_util.rs | 73 ------------------- 3 files changed, 53 insertions(+), 81 deletions(-) diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index 1c0dbc8024b3..32130396b856 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -868,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(n: usize) -> Vec + where + StandardUniform: Distribution, + { + 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::(1024); + let masks = random_numbers::(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() { diff --git a/parquet/src/arrow/record_reader/definition_levels.rs b/parquet/src/arrow/record_reader/definition_levels.rs index 0720c6cdbe0d..da3c6d088756 100644 --- a/parquet/src/arrow/record_reader/definition_levels.rs +++ b/parquet/src/arrow/record_reader/definition_levels.rs @@ -15,11 +15,6 @@ // specific language governing permissions and limitations // under the License. -use arrow_array::builder::BooleanBufferBuilder; -use arrow_buffer::Buffer; -use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk; -use bytes::Bytes; - use crate::arrow::buffer::bit_util::count_set_bits; use crate::basic::Encoding; use crate::column::reader::decoder::{ @@ -27,6 +22,11 @@ use crate::column::reader::decoder::{ }; use crate::errors::{ParquetError, Result}; use crate::schema::types::ColumnDescPtr; +use arrow_array::builder::BooleanBufferBuilder; +use arrow_buffer::Buffer; +use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk; +use arrow_buffer::bit_util::compress; +use bytes::Bytes; enum BufferInner { /// Compute levels and null mask @@ -214,8 +214,6 @@ pub(crate) fn build_filtered_validity_bitmap( item_count } -use crate::util::bit_util::compress; - enum MaybePacked { Packed(PackedDecoder), Fallback(DefinitionLevelDecoderImpl), diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs index d5f5945a12a9..298efeadbad9 100644 --- a/parquet/src/util/bit_util.rs +++ b/parquet/src/util/bit_util.rs @@ -921,43 +921,7 @@ impl From> for BitReader { } } -/// 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: -#[allow(dead_code)] -#[inline] -pub(crate) 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 - } -} #[cfg(test)] mod tests { @@ -967,43 +931,6 @@ mod tests { use rand::distr::{Distribution, StandardUniform}; use std::fmt::Debug; - #[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); - - // Randomised 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::(1024); - let masks = random_numbers::(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_ceil() { From 9b545436729bca28f4ff627dcc2b26dc2e09dbd8 Mon Sep 17 00:00:00 2001 From: devanbenz Date: Thu, 2 Jul 2026 12:10:47 -0500 Subject: [PATCH 6/7] chore: parquet fmt --- parquet/src/util/bit_util.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/parquet/src/util/bit_util.rs b/parquet/src/util/bit_util.rs index 298efeadbad9..326f4a61e214 100644 --- a/parquet/src/util/bit_util.rs +++ b/parquet/src/util/bit_util.rs @@ -921,8 +921,6 @@ impl From> for BitReader { } } - - #[cfg(test)] mod tests { use super::*; @@ -931,7 +929,6 @@ mod tests { use rand::distr::{Distribution, StandardUniform}; use std::fmt::Debug; - #[test] fn test_ceil() { assert_eq!(ceil(0, 1), 0); From a9e1587f46cc7d06f0a3788800675dbcf845fcf4 Mon Sep 17 00:00:00 2001 From: devanbenz Date: Thu, 2 Jul 2026 13:15:49 -0500 Subject: [PATCH 7/7] feat: Adds test for bmi2 path --- arrow-select/src/filter.rs | 44 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index f56a80fba91b..49b53e80e05d 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -1699,6 +1699,50 @@ mod tests { test_case_filter_sliced_list_view::(); } + /// [`filter_bits_compress`] is only dispatched to when `pext` is + /// available in hardware, so exercise it directly to get coverage of the + /// word packing logic on every platform + #[test] + fn test_filter_bits_compress() { + 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 actual = filter_bits_compress(&values, &predicate); + let actual = BooleanBuffer::new(actual, 0, predicate.count); + + assert_eq!( + actual, expected, + "len={len} density={density} offset={offset}" + ); + } + } + } + } + #[test] fn test_slice_iterator_bits() { let filter_values = (0..64).map(|i| i == 1).collect::>();