diff --git a/Cargo.lock b/Cargo.lock index 1e5dc95ff383..8030e5405902 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -512,6 +512,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..32130396b856 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/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: +#[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(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/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..2643ecf7fbdd --- /dev/null +++ b/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); diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 8b69811d4a8b..49b53e80e05d 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, }; @@ -693,6 +694,20 @@ 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 + // 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); + } + match &predicate.strategy { IterationStrategy::IndexIterator => { let bits = @@ -730,6 +745,41 @@ 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_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(¤t.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(¤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); @@ -1649,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::>(); 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..326f4a61e214 100644 --- a/parquet/src/util/bit_util.rs +++ b/parquet/src/util/bit_util.rs @@ -921,44 +921,6 @@ 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 { use super::*; @@ -967,44 +929,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() { assert_eq!(ceil(0, 1), 0);