From 16e65c357d13959f34ea544041ed0cb2a9565865 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Mon, 26 Oct 2026 12:00:29 -0400 Subject: [PATCH 01/10] perf(coalesce): fuse Boolean filter copying Pack selected Boolean values and validity bits directly into the in-progress builders, avoiding temporary filtered arrays. Check filter masks before reading source words so sparse inputs skip empty chunks. Use separate loops for sources with and without actual nulls, and carry compressed fragments across chunk boundaries while preserving lazy validity allocation. --- arrow-select/src/coalesce/boolean.rs | 340 ++++++++++++++++++++++++++- arrow-select/src/filter.rs | 5 + 2 files changed, 342 insertions(+), 3 deletions(-) diff --git a/arrow-select/src/coalesce/boolean.rs b/arrow-select/src/coalesce/boolean.rs index 17c76b2b6196..732e0dcc57cb 100644 --- a/arrow-select/src/coalesce/boolean.rs +++ b/arrow-select/src/coalesce/boolean.rs @@ -19,10 +19,136 @@ use super::InProgressArray; use crate::filter::FilterPredicate; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, BooleanArray}; -use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder}; +use arrow_buffer::bit_chunk_iterator::BitChunks; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder, bit_util}; use arrow_schema::ArrowError; use std::sync::Arc; +/// Joins compressed bit fragments into complete output words. +#[derive(Default)] +struct WordPacker { + /// Bits waiting to be appended. + pending_bits: u64, + /// Number of populated low bits in `pending_bits`. + pending_len: u32, +} + +impl WordPacker { + fn push(&mut self, compressed_bits: u64, bit_count: u32) -> Option { + self.pending_bits |= compressed_bits << self.pending_len; + let total_len = self.pending_len + bit_count; + if total_len < 64 { + self.pending_len = total_len; + return None; + } + + let complete_word = self.pending_bits; + // Keep bits that crossed the word boundary. Split the shift to avoid + // shifting by 64 when `pending_len == 0`. + self.pending_bits = (compressed_bits >> 1) >> (63 - self.pending_len); + self.pending_len = total_len - 64; + Some(complete_word) + } + + fn finish(self) -> Option<(u64, usize)> { + (self.pending_len != 0).then_some((self.pending_bits, self.pending_len as usize)) + } +} + +fn copy_filter_words_without_nulls( + output_values: &mut BooleanBufferBuilder, + source_value_chunks: &BitChunks<'_>, + filter_chunks: &BitChunks<'_>, +) -> usize { + let mut value_packer = WordPacker::default(); + let mut selected_count = 0; + + // Check the mask first so sparse filters skip reading most source words. + for (index, filter_mask) in filter_chunks.iter().enumerate() { + let selected_in_word = filter_mask.count_ones(); + if selected_in_word == 0 { + continue; + } + + let compressed_values = bit_util::compress(source_value_chunks.chunk(index), filter_mask); + if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { + output_values.append_word(complete_word, 64); + } + selected_count += selected_in_word as usize; + } + + let filter_mask = filter_chunks.remainder_bits(); + let selected_in_word = filter_mask.count_ones(); + if selected_in_word != 0 { + let compressed_values = + bit_util::compress(source_value_chunks.remainder_bits(), filter_mask); + if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { + output_values.append_word(complete_word, 64); + } + selected_count += selected_in_word as usize; + } + + if let Some((pending_bits, pending_len)) = value_packer.finish() { + output_values.append_word(pending_bits, pending_len); + } + selected_count +} + +fn copy_filter_words_with_nulls( + output_values: &mut BooleanBufferBuilder, + output_validity: &mut NullBufferBuilder, + source_value_chunks: &BitChunks<'_>, + source_validity_chunks: &BitChunks<'_>, + filter_chunks: &BitChunks<'_>, +) -> usize { + let mut value_packer = WordPacker::default(); + let mut validity_packer = WordPacker::default(); + let mut selected_count = 0; + + // Check the mask first so sparse filters skip reading most source words. + for (index, filter_mask) in filter_chunks.iter().enumerate() { + let selected_in_word = filter_mask.count_ones(); + if selected_in_word == 0 { + continue; + } + + let compressed_values = bit_util::compress(source_value_chunks.chunk(index), filter_mask); + if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { + output_values.append_word(complete_word, 64); + } + let compressed_validity = + bit_util::compress(source_validity_chunks.chunk(index), filter_mask); + if let Some(complete_word) = validity_packer.push(compressed_validity, selected_in_word) { + output_validity.append_word(complete_word, 64); + } + selected_count += selected_in_word as usize; + } + + let filter_mask = filter_chunks.remainder_bits(); + let selected_in_word = filter_mask.count_ones(); + if selected_in_word != 0 { + let compressed_values = + bit_util::compress(source_value_chunks.remainder_bits(), filter_mask); + if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { + output_values.append_word(complete_word, 64); + } + let compressed_validity = + bit_util::compress(source_validity_chunks.remainder_bits(), filter_mask); + if let Some(complete_word) = validity_packer.push(compressed_validity, selected_in_word) { + output_validity.append_word(complete_word, 64); + } + selected_count += selected_in_word as usize; + } + + if let Some((pending_bits, pending_len)) = value_packer.finish() { + output_values.append_word(pending_bits, pending_len); + } + if let Some((pending_bits, pending_len)) = validity_packer.finish() { + output_validity.append_word(pending_bits, pending_len); + } + selected_count +} + #[derive(Debug)] pub(crate) struct InProgressBooleanArray { source: Option, @@ -77,8 +203,57 @@ impl InProgressArray for InProgressBooleanArray { source: ArrayRef, filter: &FilterPredicate, ) -> Result<(), ArrowError> { - let filtered = filter.filter(source.as_ref())?; - self.append_array(filtered.as_boolean()); + if filter.filter_len() > source.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Filter predicate of length {} is larger than target array of length {}", + filter.filter_len(), + source.len() + ))); + } + + let source = source.as_boolean(); + // Read values and mask in aligned words even when the source is sliced. + let filter_chunks = filter.filter_mask().bit_chunks(); + let source_values = source.values(); + let source_value_chunks = BitChunks::new( + source_values.values(), + source_values.offset(), + filter.filter_len(), + ); + // Ignore present-but-all-valid buffers to avoid allocating output validity + // for data with no actual nulls. + let source_validity_chunks = + source + .nulls() + .filter(|nulls| nulls.null_count() != 0) + .map(|nulls| { + let nulls = nulls.inner(); + BitChunks::new(nulls.values(), nulls.offset(), filter.filter_len()) + }); + + let output_values = self + .values + .get_or_insert_with(|| BooleanBufferBuilder::new(self.batch_size)); + + // Branch once here so non-null inputs avoid validity work in the chunk loop. + let selected_count = if let Some(source_validity_chunks) = &source_validity_chunks { + copy_filter_words_with_nulls( + output_values, + &mut self.nulls, + &source_value_chunks, + source_validity_chunks, + &filter_chunks, + ) + } else { + let selected_count = copy_filter_words_without_nulls( + output_values, + &source_value_chunks, + &filter_chunks, + ); + self.nulls.append_n_non_nulls(selected_count); + selected_count + }; + debug_assert_eq!(selected_count, filter.count()); Ok(()) } @@ -109,6 +284,7 @@ impl InProgressArray for InProgressBooleanArray { mod tests { use super::*; use crate::filter::FilterBuilder; + use arrow_buffer::NullBuffer; #[test] fn copies_slices_and_sparse_filters() { @@ -156,4 +332,162 @@ mod tests { assert_eq!(second.as_boolean(), &source.slice(2, 2)); assert_eq!(in_progress.finish().unwrap().len(), 0); } + + const FILTER_LEN: usize = 269; + const SOURCE_OFFSET: usize = 3; + + #[test] + fn filtered_words_without_nulls_match_filter_kernel() { + let len = SOURCE_OFFSET + FILTER_LEN; + let values = (0..len).map(source_value).collect(); + let all_valid = BooleanArray::new(values, Some(NullBuffer::new_valid(len))) + .slice(SOURCE_OFFSET, FILTER_LEN); + let without_validity = non_null_source(); + assert!(without_validity.nulls().is_none()); + assert!(all_valid.nulls().is_some()); + + for (case, source) in [ + ("without validity", without_validity), + ("all-valid validity", all_valid), + ] { + assert_eq!(source.null_count(), 0); + assert_filtered_words_match_filter_kernel(case, source); + } + } + + #[test] + fn filtered_words_with_nulls_match_filter_kernel() { + let source = BooleanArray::from_iter( + (0..SOURCE_OFFSET + FILTER_LEN) + .map(|index| (index % 11 != 0).then_some(index % 3 == 0 || index % 7 == 0)), + ) + .slice(SOURCE_OFFSET, FILTER_LEN); + assert!(source.null_count() > 0); + + assert_filtered_words_match_filter_kernel("with nulls", source); + } + + fn assert_filtered_words_match_filter_kernel(source_case: &str, source: BooleanArray) { + for (filter_case, filter) in filter_cases() { + assert_filter_copy_matches_kernel(source_case, filter_case, &source, &filter); + } + } + + fn filter_cases() -> Vec<(&'static str, BooleanArray)> { + let selections = [ + ("none", vec![]), + ("all", (0..FILTER_LEN).collect()), + ( + "word boundaries", + vec![0, 1, 7, 63, 64, 65, 127, 128, 191, 192, 255, 256, 268], + ), + ( + "carry across source words", + // The first two chunks leave 24 carried bits; the third completes a word. + (0..40).chain(64..112).chain(128..168).collect(), + ), + ("alternating", (0..FILTER_LEN).step_by(2).collect()), + ]; + let mut filters: Vec<(&str, BooleanArray)> = selections + .into_iter() + .map(|(case, selected_indices)| { + let filter = (0..FILTER_LEN) + .map(|index| selected_indices.contains(&index)) + .collect(); + (case, filter) + }) + .collect(); + + const FILTER_OFFSET: usize = 5; + let sliced_filter: BooleanArray = std::iter::repeat_n(false, FILTER_OFFSET) + .chain((0..FILTER_LEN).map(|index| index.is_multiple_of(20))) + .collect(); + let sliced_filter = sliced_filter.slice(FILTER_OFFSET, FILTER_LEN); + assert_eq!(sliced_filter.values().offset(), FILTER_OFFSET); + filters.push(("sliced filter", sliced_filter)); + + let mut filter_values = vec![false; FILTER_LEN]; + for index in [1, 65, 129, 268] { + filter_values[index] = true; + } + let mut filter_validity = vec![true; FILTER_LEN]; + filter_validity[65] = false; + let filter_with_null = BooleanArray::new( + BooleanBuffer::from(filter_values), + Some(NullBuffer::new(BooleanBuffer::from(filter_validity))), + ); + assert!(filter_with_null.value(65)); + assert!(filter_with_null.is_null(65)); + filters.push(("null filter entry", filter_with_null)); + filters + } + + fn assert_filter_copy_matches_kernel( + source_case: &str, + filter_case: &str, + source: &BooleanArray, + filter: &BooleanArray, + ) { + const PREFIX_LEN: usize = 5; + let expect_lazy_validity = source.null_count() == 0; + + for optimize in [false, true] { + let builder = FilterBuilder::new(filter); + let predicate = if optimize { + builder.optimize().build() + } else { + builder.build() + }; + let filtered = predicate.filter(source).unwrap(); + let prefix = source.slice(0, PREFIX_LEN); + let expected = crate::concat::concat(&[&prefix, filtered.as_ref()]).unwrap(); + + let source = Arc::new(source.clone()) as ArrayRef; + let mut in_progress = InProgressBooleanArray::new(FILTER_LEN); + in_progress.set_source(Some(Arc::clone(&source))); + in_progress.copy_rows(0, PREFIX_LEN).unwrap(); + in_progress.set_source(None); + in_progress + .copy_rows_by_filter_from(source, &predicate) + .unwrap(); + + let output = in_progress.finish().unwrap(); + assert_eq!( + output.as_boolean(), + expected.as_boolean(), + "{source_case}, {filter_case}, optimize={optimize}" + ); + if expect_lazy_validity { + assert!( + output.nulls().is_none(), + "{source_case}, {filter_case}, optimize={optimize} should preserve lazy validity" + ); + } + } + } + + fn non_null_source() -> BooleanArray { + let source: BooleanArray = (0..SOURCE_OFFSET + FILTER_LEN).map(source_value).collect(); + source.slice(SOURCE_OFFSET, FILTER_LEN) + } + + fn source_value(index: usize) -> bool { + index.is_multiple_of(3) || index.is_multiple_of(7) + } + + #[test] + fn filter_longer_than_source_returns_error() { + let source = Arc::new(BooleanArray::from(vec![true, false])) as ArrayRef; + let filter_mask = BooleanArray::from(vec![true, false, true]); + let predicate = FilterBuilder::new(&filter_mask).build(); + let mut in_progress = InProgressBooleanArray::new(8); + + let error = in_progress + .copy_rows_by_filter_from(source, &predicate) + .unwrap_err(); + assert_eq!( + error.to_string(), + "Invalid argument error: Filter predicate of length 3 is larger than target array of length 2" + ); + } } diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 79966c933708..fb96ceaf7864 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -530,6 +530,11 @@ impl FilterPredicate { self.filter.len() } + /// Returns the normalized filter mask, with nulls replaced by `false`. + pub(crate) fn filter_mask(&self) -> &BooleanBuffer { + self.filter.values() + } + /// Return a [`FilterSelection`] for iterating over the rows selected by /// this [`FilterPredicate`]. pub(crate) fn selection(&self) -> FilterSelection<'_> { From 4d091cd4980c284256784f98b41eb264de1fb55a Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Mon, 26 Oct 2026 15:47:46 -0400 Subject: [PATCH 02/10] Reuse Boolean filter compression in coalescer Let the Boolean coalescer append words produced by the existing compression kernel instead of maintaining a second copy of its packing loop. Support values and validity in one mask traversal, and cover shorter filters and exact word boundaries. --- arrow-select/src/coalesce/boolean.rs | 194 ++++----------------------- arrow-select/src/filter.rs | 171 +++++++++++++++-------- 2 files changed, 146 insertions(+), 219 deletions(-) diff --git a/arrow-select/src/coalesce/boolean.rs b/arrow-select/src/coalesce/boolean.rs index 732e0dcc57cb..63311d252ddc 100644 --- a/arrow-select/src/coalesce/boolean.rs +++ b/arrow-select/src/coalesce/boolean.rs @@ -16,139 +16,13 @@ // under the License. use super::InProgressArray; -use crate::filter::FilterPredicate; +use crate::filter::{FilterPredicate, filter_bits_compress, filter_bits_compress_pair}; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, BooleanArray}; -use arrow_buffer::bit_chunk_iterator::BitChunks; -use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder, bit_util}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder}; use arrow_schema::ArrowError; use std::sync::Arc; -/// Joins compressed bit fragments into complete output words. -#[derive(Default)] -struct WordPacker { - /// Bits waiting to be appended. - pending_bits: u64, - /// Number of populated low bits in `pending_bits`. - pending_len: u32, -} - -impl WordPacker { - fn push(&mut self, compressed_bits: u64, bit_count: u32) -> Option { - self.pending_bits |= compressed_bits << self.pending_len; - let total_len = self.pending_len + bit_count; - if total_len < 64 { - self.pending_len = total_len; - return None; - } - - let complete_word = self.pending_bits; - // Keep bits that crossed the word boundary. Split the shift to avoid - // shifting by 64 when `pending_len == 0`. - self.pending_bits = (compressed_bits >> 1) >> (63 - self.pending_len); - self.pending_len = total_len - 64; - Some(complete_word) - } - - fn finish(self) -> Option<(u64, usize)> { - (self.pending_len != 0).then_some((self.pending_bits, self.pending_len as usize)) - } -} - -fn copy_filter_words_without_nulls( - output_values: &mut BooleanBufferBuilder, - source_value_chunks: &BitChunks<'_>, - filter_chunks: &BitChunks<'_>, -) -> usize { - let mut value_packer = WordPacker::default(); - let mut selected_count = 0; - - // Check the mask first so sparse filters skip reading most source words. - for (index, filter_mask) in filter_chunks.iter().enumerate() { - let selected_in_word = filter_mask.count_ones(); - if selected_in_word == 0 { - continue; - } - - let compressed_values = bit_util::compress(source_value_chunks.chunk(index), filter_mask); - if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { - output_values.append_word(complete_word, 64); - } - selected_count += selected_in_word as usize; - } - - let filter_mask = filter_chunks.remainder_bits(); - let selected_in_word = filter_mask.count_ones(); - if selected_in_word != 0 { - let compressed_values = - bit_util::compress(source_value_chunks.remainder_bits(), filter_mask); - if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { - output_values.append_word(complete_word, 64); - } - selected_count += selected_in_word as usize; - } - - if let Some((pending_bits, pending_len)) = value_packer.finish() { - output_values.append_word(pending_bits, pending_len); - } - selected_count -} - -fn copy_filter_words_with_nulls( - output_values: &mut BooleanBufferBuilder, - output_validity: &mut NullBufferBuilder, - source_value_chunks: &BitChunks<'_>, - source_validity_chunks: &BitChunks<'_>, - filter_chunks: &BitChunks<'_>, -) -> usize { - let mut value_packer = WordPacker::default(); - let mut validity_packer = WordPacker::default(); - let mut selected_count = 0; - - // Check the mask first so sparse filters skip reading most source words. - for (index, filter_mask) in filter_chunks.iter().enumerate() { - let selected_in_word = filter_mask.count_ones(); - if selected_in_word == 0 { - continue; - } - - let compressed_values = bit_util::compress(source_value_chunks.chunk(index), filter_mask); - if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { - output_values.append_word(complete_word, 64); - } - let compressed_validity = - bit_util::compress(source_validity_chunks.chunk(index), filter_mask); - if let Some(complete_word) = validity_packer.push(compressed_validity, selected_in_word) { - output_validity.append_word(complete_word, 64); - } - selected_count += selected_in_word as usize; - } - - let filter_mask = filter_chunks.remainder_bits(); - let selected_in_word = filter_mask.count_ones(); - if selected_in_word != 0 { - let compressed_values = - bit_util::compress(source_value_chunks.remainder_bits(), filter_mask); - if let Some(complete_word) = value_packer.push(compressed_values, selected_in_word) { - output_values.append_word(complete_word, 64); - } - let compressed_validity = - bit_util::compress(source_validity_chunks.remainder_bits(), filter_mask); - if let Some(complete_word) = validity_packer.push(compressed_validity, selected_in_word) { - output_validity.append_word(complete_word, 64); - } - selected_count += selected_in_word as usize; - } - - if let Some((pending_bits, pending_len)) = value_packer.finish() { - output_values.append_word(pending_bits, pending_len); - } - if let Some((pending_bits, pending_len)) = validity_packer.finish() { - output_validity.append_word(pending_bits, pending_len); - } - selected_count -} - #[derive(Debug)] pub(crate) struct InProgressBooleanArray { source: Option, @@ -212,48 +86,26 @@ impl InProgressArray for InProgressBooleanArray { } let source = source.as_boolean(); - // Read values and mask in aligned words even when the source is sliced. - let filter_chunks = filter.filter_mask().bit_chunks(); - let source_values = source.values(); - let source_value_chunks = BitChunks::new( - source_values.values(), - source_values.offset(), - filter.filter_len(), - ); - // Ignore present-but-all-valid buffers to avoid allocating output validity - // for data with no actual nulls. - let source_validity_chunks = - source - .nulls() - .filter(|nulls| nulls.null_count() != 0) - .map(|nulls| { - let nulls = nulls.inner(); - BitChunks::new(nulls.values(), nulls.offset(), filter.filter_len()) - }); - let output_values = self .values .get_or_insert_with(|| BooleanBufferBuilder::new(self.batch_size)); - - // Branch once here so non-null inputs avoid validity work in the chunk loop. - let selected_count = if let Some(source_validity_chunks) = &source_validity_chunks { - copy_filter_words_with_nulls( - output_values, - &mut self.nulls, - &source_value_chunks, - source_validity_chunks, - &filter_chunks, - ) - } else { - let selected_count = copy_filter_words_without_nulls( - output_values, - &source_value_chunks, - &filter_chunks, + // Treat all-valid sources as non-null to keep output validity lazy. + if let Some(source_nulls) = source.nulls().filter(|nulls| nulls.null_count() != 0) { + filter_bits_compress_pair( + source.values(), + source_nulls.inner(), + filter, + |values, validity, count| { + output_values.append_word(values, count); + self.nulls.append_word(validity, count); + }, ); - self.nulls.append_n_non_nulls(selected_count); - selected_count - }; - debug_assert_eq!(selected_count, filter.count()); + } else { + filter_bits_compress(source.values(), filter, |values, count| { + output_values.append_word(values, count) + }); + self.nulls.append_n_non_nulls(filter.count()); + } Ok(()) } @@ -406,6 +258,16 @@ mod tests { assert_eq!(sliced_filter.values().offset(), FILTER_OFFSET); filters.push(("sliced filter", sliced_filter)); + let shorter_filter = (0_usize..197) + .map(|index| index.is_multiple_of(17)) + .collect(); + filters.push(("shorter than source", shorter_filter)); + + let aligned_filter = (0_usize..256) + .map(|index| index.is_multiple_of(19)) + .collect(); + filters.push(("empty remainder", aligned_filter)); + let mut filter_values = vec![false; FILTER_LEN]; for index in [1, 65, 129, 268] { filter_values[index] = true; diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index fb96ceaf7864..c481ca9abd12 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -530,11 +530,6 @@ impl FilterPredicate { self.filter.len() } - /// Returns the normalized filter mask, with nulls replaced by `false`. - pub(crate) fn filter_mask(&self) -> &BooleanBuffer { - self.filter.values() - } - /// Return a [`FilterSelection`] for iterating over the rows selected by /// this [`FilterPredicate`]. pub(crate) fn selection(&self) -> FilterSelection<'_> { @@ -756,7 +751,7 @@ fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { let sparse_indices = count <= len / 64 && matches!(predicate.strategy, IterationStrategy::Indices(_)); if !dense && !sparse_indices { - return filter_bits_compress(buffer, predicate); + return collect_filter_bits_compress(buffer, predicate); } filter_bits_strategy(buffer, predicate) } @@ -806,61 +801,91 @@ fn filter_bits_strategy(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> } /// Filter the packed bitmask `buffer` with `predicate` by extracting the kept -/// bits of each 64-bit word with [`bit_util::compress`] (`pext`) -/// -/// Not inlined: within `filter_array` the packing state spills to the stack +/// bits of each 64-bit word with [`bit_util::compress`] (`pext`). +#[inline(never)] +pub(crate) fn filter_bits_compress( + buffer: &BooleanBuffer, + predicate: &FilterPredicate, + mut append: impl FnMut(u64, usize), +) { + filter_bits_compress_impl::(buffer, None, predicate, |word, _, count| { + append(word, count) + }) +} + +/// Filters two packed bitmaps using the same predicate and mask traversal. #[inline(never)] -fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { +pub(crate) fn filter_bits_compress_pair( + first: &BooleanBuffer, + second: &BooleanBuffer, + predicate: &FilterPredicate, + mut append: impl FnMut(u64, u64, usize), +) { + filter_bits_compress_impl::(first, Some(second), predicate, |first, second, count| { + append(first, second.unwrap(), count) + }) +} + +#[inline(always)] +fn filter_bits_compress_impl( + first: &BooleanBuffer, + second: Option<&BooleanBuffer>, + predicate: &FilterPredicate, + mut append: impl FnMut(u64, Option, usize), +) { /// Packs the bits extracted from successive words into the low `filled` - /// bits of `current`; once complete it is written at `idx` and restarts - /// from the bits that did not fit + /// bits of `current`. + #[derive(Default)] struct Packer { - ptr: *mut u64, - idx: usize, + /// Bits waiting to be written. current: u64, + /// Number of populated low bits in `current`. filled: u32, } impl Packer { #[inline(always)] - fn push(&mut self, values: u64, mask: u64) { - let bits = bit_util::compress(values, mask); + fn push(&mut self, bits: u64, count: u32) -> Option { self.current |= bits << self.filled; - let total = self.filled + mask.count_ones(); + let total = self.filled + count; if total < 64 { self.filled = total; - } else { - // SAFETY: `count` is the number of set bits in the filter, so - // at most `count / 64` words are ever completed and the - // buffer holds `count / 64 + 1` - unsafe { self.ptr.add(self.idx).write(self.current) }; - self.idx += 1; - // `bits >> (64 - filled)`, written so that `filled == 0` - // shifts everything out - self.current = (bits >> 1) >> (63 - self.filled); - self.filled = total - 64; + return None; } + + let complete = self.current; + // `bits >> (64 - filled)`, written so that `filled == 0` + // shifts everything out + self.current = (bits >> 1) >> (63 - self.filled); + self.filled = total - 64; + Some(complete) + } + + #[inline] + fn finish(self) -> Option<(u64, usize)> { + (self.filled != 0).then_some((self.current, self.filled as usize)) } } - assert!(buffer.len() >= predicate.filter.len()); + assert!(first.len() >= predicate.filter.len()); + if PAIR { + assert!(second.unwrap().len() >= predicate.filter.len()); + } let mask_chunks = predicate.filter.values().bit_chunks(); - let value_chunks = BitChunks::new(buffer.values(), buffer.offset(), predicate.filter.len()); - // `count` is the filter's set bit count, which the buffer size and the - // raw writes below rely on, and both chunk views cover - // `predicate.filter.len()` bits, so indexing `value_chunks` by the - // position in `mask_chunks` stays in bounds + let first_chunks = BitChunks::new(first.values(), first.offset(), predicate.filter.len()); + let second_chunks = second + .map(|buffer| BitChunks::new(buffer.values(), buffer.offset(), predicate.filter.len())); debug_assert_eq!(predicate.count, predicate.filter.true_count()); - debug_assert_eq!(mask_chunks.chunk_len(), value_chunks.chunk_len()); + debug_assert_eq!(mask_chunks.chunk_len(), first_chunks.chunk_len()); + if PAIR { + debug_assert_eq!( + mask_chunks.chunk_len(), + second_chunks.as_ref().unwrap().chunk_len() + ); + } - // One word beyond the complete ones for the trailing partial word - let mut out: Vec = Vec::with_capacity(predicate.count / 64 + 1); - let mut packer = Packer { - ptr: out.as_mut_ptr(), - idx: 0, - current: 0, - filled: 0, - }; + let mut first_packer = Packer::default(); + let mut second_packer = Packer::default(); for (index, mask) in mask_chunks.iter().enumerate() { // Words with no kept bits are skipped before the corresponding values @@ -868,18 +893,58 @@ fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> if mask == 0 { continue; } - packer.push(value_chunks.chunk(index), mask); + let count = mask.count_ones(); + let first = first_packer.push(bit_util::compress(first_chunks.chunk(index), mask), count); + let second = PAIR.then(|| { + second_packer.push( + bit_util::compress(second_chunks.as_ref().unwrap().chunk(index), mask), + count, + ) + }); + if let Some(first) = first { + append(first, second.flatten(), 64); + } } - packer.push(value_chunks.remainder_bits(), mask_chunks.remainder_bits()); - - // The trailing partial word; its bits above `filled` are zero - // SAFETY: `idx <= count / 64`, so this and every word below it is - // within the buffer and written - debug_assert!(packer.idx < out.capacity()); - unsafe { - packer.ptr.add(packer.idx).write(packer.current); - out.set_len(packer.idx + 1); + let mask = mask_chunks.remainder_bits(); + if mask != 0 { + let count = mask.count_ones(); + let first = first_packer.push( + bit_util::compress(first_chunks.remainder_bits(), mask), + count, + ); + let second = PAIR.then(|| { + second_packer.push( + bit_util::compress(second_chunks.as_ref().unwrap().remainder_bits(), mask), + count, + ) + }); + if let Some(first) = first { + append(first, second.flatten(), 64); + } + } + if let Some((first, count)) = first_packer.finish() { + let second = PAIR.then(|| second_packer.finish().unwrap().0); + append(first, second, count); } +} + +/// Collects compressed bits into a new buffer. +/// +/// Not inlined: within `filter_array` the packing state spills to the stack. +#[inline(never)] +fn collect_filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { + // One word beyond the complete ones for the trailing partial word + let mut out: Vec = Vec::with_capacity(predicate.count / 64 + 1); + let ptr = out.as_mut_ptr(); + let mut idx = 0; + filter_bits_compress(buffer, predicate, |word, _| { + // SAFETY: `count` is the number of set bits in the filter, so at most + // `count / 64 + 1` words are emitted and the output has room + unsafe { ptr.add(idx).write(word) }; + idx += 1; + }); + // SAFETY: every word below `idx` was initialized by the callback above + unsafe { out.set_len(idx) }; let mut out = MutableBuffer::from(out); out.truncate(bit_util::ceil(predicate.count, 8)); out.into() @@ -1960,7 +2025,7 @@ mod tests { predicate.strategy ); - let compressed = filter_bits_compress(&values, predicate); + let compressed = collect_filter_bits_compress(&values, predicate); let compressed = BooleanBuffer::new(compressed, 0, predicate.count); assert_eq!(compressed, expected, "compress {case}"); From 5bcede663d5e69b89423e6d7f3ef39cafd950ef7 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 27 Oct 2026 23:31:31 -0400 Subject: [PATCH 03/10] fix(filter): assert compressed output bounds --- arrow-select/src/filter.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index c481ca9abd12..d3aac27916f7 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -938,6 +938,7 @@ fn collect_filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredic let ptr = out.as_mut_ptr(); let mut idx = 0; filter_bits_compress(buffer, predicate, |word, _| { + debug_assert!(idx <= predicate.count / 64); // SAFETY: `count` is the number of set bits in the filter, so at most // `count / 64 + 1` words are emitted and the output has room unsafe { ptr.add(idx).write(word) }; From f059b1ecbb148409f813fe79189252c138b4bac5 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Thu, 29 Oct 2026 12:25:01 -0400 Subject: [PATCH 04/10] bench(filter): add Boolean data cases --- arrow/benches/filter_kernels.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/arrow/benches/filter_kernels.rs b/arrow/benches/filter_kernels.rs index 8d7403ae0456..2428cb14f64e 100644 --- a/arrow/benches/filter_kernels.rs +++ b/arrow/benches/filter_kernels.rs @@ -81,6 +81,19 @@ fn add_benchmark(c: &mut Criterion) { b.iter(|| bench_built_filter(&data_array, &sparse_filter)) }); + let data_array = create_boolean_array(size, 0.0, 0.5); + c.bench_function("filter context boolean (kept 1/2)", |b| { + b.iter(|| bench_built_filter(&data_array, &filter)) + }); + c.bench_function( + "filter context boolean high selectivity (kept 1023/1024)", + |b| b.iter(|| bench_built_filter(&data_array, &dense_filter)), + ); + c.bench_function( + "filter context boolean low selectivity (kept 1/1024)", + |b| b.iter(|| bench_built_filter(&data_array, &sparse_filter)), + ); + let data_array = create_primitive_array::(size, 0.0); c.bench_function("filter i32 (kept 1/2)", |b| { b.iter(|| bench_filter(&data_array, &filter_array)) From da0882c21d892192527363db2ccc0743b20660b6 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 3 Nov 2026 11:05:57 -0500 Subject: [PATCH 05/10] test(coalesce): cover differing Boolean buffer offsets --- arrow-select/src/coalesce/boolean.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/arrow-select/src/coalesce/boolean.rs b/arrow-select/src/coalesce/boolean.rs index 63311d252ddc..f68fcc2f7b00 100644 --- a/arrow-select/src/coalesce/boolean.rs +++ b/arrow-select/src/coalesce/boolean.rs @@ -219,6 +219,19 @@ mod tests { assert_filtered_words_match_filter_kernel("with nulls", source); } + #[test] + fn filtered_words_with_unaligned_validity_match_filter_kernel() { + let values: BooleanBuffer = (0..SOURCE_OFFSET + FILTER_LEN).map(source_value).collect(); + let validity: BooleanBuffer = (0..FILTER_LEN + 11).map(|index| index % 11 != 0).collect(); + let source = BooleanArray::new( + values.slice(SOURCE_OFFSET, FILTER_LEN), + Some(NullBuffer::new(validity.slice(11, FILTER_LEN))), + ); + assert_ne!(source.values().offset(), source.nulls().unwrap().offset()); + + assert_filtered_words_match_filter_kernel("unaligned validity", source); + } + fn assert_filtered_words_match_filter_kernel(source_case: &str, source: BooleanArray) { for (filter_case, filter) in filter_cases() { assert_filter_copy_matches_kernel(source_case, filter_case, &source, &filter); From d306443d1ff0b82d20ad5136320989e4f1a0f4f5 Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 3 Nov 2026 11:24:54 -0500 Subject: [PATCH 06/10] fix(coalesce): preserve sparse Boolean indices path --- arrow-select/src/coalesce/boolean.rs | 6 ++++++ arrow-select/src/filter.rs | 10 +++++++--- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/arrow-select/src/coalesce/boolean.rs b/arrow-select/src/coalesce/boolean.rs index f68fcc2f7b00..308f6bb69fdd 100644 --- a/arrow-select/src/coalesce/boolean.rs +++ b/arrow-select/src/coalesce/boolean.rs @@ -85,6 +85,12 @@ impl InProgressArray for InProgressBooleanArray { ))); } + if filter.uses_sparse_indices() { + let filtered = filter.filter(source.as_ref())?; + self.append_array(filtered.as_boolean()); + return Ok(()); + } + let source = source.as_boolean(); let output_values = self .values diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index d3aac27916f7..6ac691b2c104 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -530,6 +530,12 @@ impl FilterPredicate { self.filter.len() } + /// Returns whether bit filtering should gather precomputed indices instead of compressing. + pub(crate) fn uses_sparse_indices(&self) -> bool { + self.count <= self.filter.len() / 64 + && matches!(self.strategy, IterationStrategy::Indices(_)) + } + /// Return a [`FilterSelection`] for iterating over the rows selected by /// this [`FilterPredicate`]. pub(crate) fn selection(&self) -> FilterSelection<'_> { @@ -748,9 +754,7 @@ fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { 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 { + if !dense && !predicate.uses_sparse_indices() { return collect_filter_bits_compress(buffer, predicate); } filter_bits_strategy(buffer, predicate) From bee8cdc0f5401a239b39cf3d978443ff9145b03c Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 3 Nov 2026 14:18:20 -0500 Subject: [PATCH 07/10] perf(filter): preserve direct Boolean collection --- arrow-select/src/filter.rs | 71 ++++++++++++++++++++++++++++++++------ 1 file changed, 60 insertions(+), 11 deletions(-) diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 6ac691b2c104..06091295bcc7 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -937,19 +937,68 @@ fn filter_bits_compress_impl( /// Not inlined: within `filter_array` the packing state spills to the stack. #[inline(never)] fn collect_filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { + /// Packs the bits extracted from successive words into the low `filled` + /// bits of `current`; once complete it is written at `idx` and restarts + /// from the bits that did not fit. + struct Packer { + ptr: *mut u64, + idx: usize, + current: u64, + filled: u32, + } + + impl Packer { + #[inline(always)] + fn push(&mut self, values: u64, mask: u64) { + let bits = bit_util::compress(values, mask); + self.current |= bits << self.filled; + let total = self.filled + mask.count_ones(); + if total < 64 { + self.filled = total; + } else { + // SAFETY: `count` is the number of set bits in the filter, so + // at most `count / 64` words are ever completed and the + // buffer holds `count / 64 + 1`. + unsafe { self.ptr.add(self.idx).write(self.current) }; + self.idx += 1; + // `bits >> (64 - filled)`, written so that `filled == 0` + // shifts everything out. + self.current = (bits >> 1) >> (63 - self.filled); + self.filled = total - 64; + } + } + } + + 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()); + debug_assert_eq!(predicate.count, predicate.filter.true_count()); + debug_assert_eq!(mask_chunks.chunk_len(), value_chunks.chunk_len()); + // One word beyond the complete ones for the trailing partial word let mut out: Vec = Vec::with_capacity(predicate.count / 64 + 1); - let ptr = out.as_mut_ptr(); - let mut idx = 0; - filter_bits_compress(buffer, predicate, |word, _| { - debug_assert!(idx <= predicate.count / 64); - // SAFETY: `count` is the number of set bits in the filter, so at most - // `count / 64 + 1` words are emitted and the output has room - unsafe { ptr.add(idx).write(word) }; - idx += 1; - }); - // SAFETY: every word below `idx` was initialized by the callback above - unsafe { out.set_len(idx) }; + let mut packer = Packer { + ptr: out.as_mut_ptr(), + idx: 0, + current: 0, + filled: 0, + }; + + for (index, mask) in mask_chunks.iter().enumerate() { + if mask == 0 { + continue; + } + packer.push(value_chunks.chunk(index), mask); + } + packer.push(value_chunks.remainder_bits(), mask_chunks.remainder_bits()); + + // SAFETY: `idx <= count / 64`, so this and every word below it is + // within the buffer and written. + debug_assert!(packer.idx < out.capacity()); + unsafe { + packer.ptr.add(packer.idx).write(packer.current); + out.set_len(packer.idx + 1); + } let mut out = MutableBuffer::from(out); out.truncate(bit_util::ceil(predicate.count, 8)); out.into() From 85477318f90fa8fcdcc521f1c92cb0cc0acfd2eb Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 3 Nov 2026 15:04:42 -0500 Subject: [PATCH 08/10] perf(filter): preserve Boolean compression dispatch --- arrow-select/src/filter.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 06091295bcc7..b02b61464f84 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -754,7 +754,8 @@ fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { let len = predicate.filter.len(); let count = predicate.count; let dense = count >= len - len / 64; - if !dense && !predicate.uses_sparse_indices() { + let sparse_indices = predicate.uses_sparse_indices(); + if !dense && !sparse_indices { return collect_filter_bits_compress(buffer, predicate); } filter_bits_strategy(buffer, predicate) From 3ddd98a4d61f851c606c59bdafd61f4e0afd4d6d Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Tue, 3 Nov 2026 15:15:23 -0500 Subject: [PATCH 09/10] perf(filter): keep Boolean compression output direct --- arrow-select/src/coalesce/boolean.rs | 131 ++++++++++++++-- arrow-select/src/filter.rs | 227 +++++++-------------------- 2 files changed, 174 insertions(+), 184 deletions(-) diff --git a/arrow-select/src/coalesce/boolean.rs b/arrow-select/src/coalesce/boolean.rs index 308f6bb69fdd..0b778563acf9 100644 --- a/arrow-select/src/coalesce/boolean.rs +++ b/arrow-select/src/coalesce/boolean.rs @@ -16,13 +16,97 @@ // under the License. use super::InProgressArray; -use crate::filter::{FilterPredicate, filter_bits_compress, filter_bits_compress_pair}; +use crate::filter::{FilterPredicate, Packer}; use arrow_array::cast::AsArray; use arrow_array::{Array, ArrayRef, BooleanArray}; +use arrow_buffer::bit_chunk_iterator::BitChunks; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder}; use arrow_schema::ArrowError; use std::sync::Arc; +fn copy_filter_words_without_nulls( + output_values: &mut BooleanBufferBuilder, + source_value_chunks: &BitChunks<'_>, + filter_chunks: &BitChunks<'_>, +) { + let mut value_packer = Packer::default(); + + // Check the mask first so sparse filters skip reading most source words. + for (index, filter_mask) in filter_chunks.iter().enumerate() { + if filter_mask == 0 { + continue; + } + + if let Some(complete_word) = + value_packer.push(source_value_chunks.chunk(index), filter_mask) + { + output_values.append_word(complete_word, 64); + } + } + + let filter_mask = filter_chunks.remainder_bits(); + if filter_mask != 0 + && let Some(complete_word) = + value_packer.push(source_value_chunks.remainder_bits(), filter_mask) + { + output_values.append_word(complete_word, 64); + } + + if let Some((pending_bits, pending_len)) = value_packer.finish() { + output_values.append_word(pending_bits, pending_len); + } +} + +fn copy_filter_words_with_nulls( + output_values: &mut BooleanBufferBuilder, + output_validity: &mut NullBufferBuilder, + source_value_chunks: &BitChunks<'_>, + source_validity_chunks: &BitChunks<'_>, + filter_chunks: &BitChunks<'_>, +) { + let mut value_packer = Packer::default(); + let mut validity_packer = Packer::default(); + + // Check the mask first so sparse filters skip reading most source words. + for (index, filter_mask) in filter_chunks.iter().enumerate() { + if filter_mask == 0 { + continue; + } + + if let Some(complete_word) = + value_packer.push(source_value_chunks.chunk(index), filter_mask) + { + output_values.append_word(complete_word, 64); + } + if let Some(complete_word) = + validity_packer.push(source_validity_chunks.chunk(index), filter_mask) + { + output_validity.append_word(complete_word, 64); + } + } + + let filter_mask = filter_chunks.remainder_bits(); + if filter_mask != 0 { + if let Some(complete_word) = + value_packer.push(source_value_chunks.remainder_bits(), filter_mask) + { + output_values.append_word(complete_word, 64); + } + if let Some(complete_word) = + validity_packer.push(source_validity_chunks.remainder_bits(), filter_mask) + { + output_validity.append_word(complete_word, 64); + } + } + + if let Some((pending_bits, pending_len)) = value_packer.finish() { + output_values.append_word(pending_bits, pending_len); + } + if let Some((pending_bits, pending_len)) = validity_packer.finish() { + output_validity.append_word(pending_bits, pending_len); + } +} + #[derive(Debug)] pub(crate) struct InProgressBooleanArray { source: Option, @@ -92,25 +176,42 @@ impl InProgressArray for InProgressBooleanArray { } let source = source.as_boolean(); + // Read values and mask in aligned words even when the source is sliced. + let filter_chunks = filter.filter_mask().bit_chunks(); + let source_values = source.values(); + let source_value_chunks = BitChunks::new( + source_values.values(), + source_values.offset(), + filter.filter_len(), + ); + // Ignore present-but-all-valid buffers to avoid allocating output validity + // for data with no actual nulls. + let source_validity_chunks = + source + .nulls() + .filter(|nulls| nulls.null_count() != 0) + .map(|nulls| { + let nulls = nulls.inner(); + BitChunks::new(nulls.values(), nulls.offset(), filter.filter_len()) + }); + let output_values = self .values .get_or_insert_with(|| BooleanBufferBuilder::new(self.batch_size)); - // Treat all-valid sources as non-null to keep output validity lazy. - if let Some(source_nulls) = source.nulls().filter(|nulls| nulls.null_count() != 0) { - filter_bits_compress_pair( - source.values(), - source_nulls.inner(), - filter, - |values, validity, count| { - output_values.append_word(values, count); - self.nulls.append_word(validity, count); - }, + + // Branch once here so non-null inputs avoid validity work in the chunk loop. + let selected_count = filter.count(); + if let Some(source_validity_chunks) = &source_validity_chunks { + copy_filter_words_with_nulls( + output_values, + &mut self.nulls, + &source_value_chunks, + source_validity_chunks, + &filter_chunks, ); } else { - filter_bits_compress(source.values(), filter, |values, count| { - output_values.append_word(values, count) - }); - self.nulls.append_n_non_nulls(filter.count()); + copy_filter_words_without_nulls(output_values, &source_value_chunks, &filter_chunks); + self.nulls.append_n_non_nulls(selected_count); } Ok(()) } diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index b02b61464f84..64fbc7e30baf 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -530,6 +530,11 @@ impl FilterPredicate { self.filter.len() } + /// Returns the normalized filter mask, with nulls replaced by `false`. + pub(crate) fn filter_mask(&self) -> &BooleanBuffer { + self.filter.values() + } + /// Returns whether bit filtering should gather precomputed indices instead of compressing. pub(crate) fn uses_sparse_indices(&self) -> bool { self.count <= self.filter.len() / 64 @@ -756,7 +761,7 @@ fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { let dense = count >= len - len / 64; let sparse_indices = predicate.uses_sparse_indices(); if !dense && !sparse_indices { - return collect_filter_bits_compress(buffer, predicate); + return filter_bits_compress(buffer, predicate); } filter_bits_strategy(buffer, predicate) } @@ -805,171 +810,46 @@ fn filter_bits_strategy(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> } } -/// Filter the packed bitmask `buffer` with `predicate` by extracting the kept -/// bits of each 64-bit word with [`bit_util::compress`] (`pext`). -#[inline(never)] -pub(crate) fn filter_bits_compress( - buffer: &BooleanBuffer, - predicate: &FilterPredicate, - mut append: impl FnMut(u64, usize), -) { - filter_bits_compress_impl::(buffer, None, predicate, |word, _, count| { - append(word, count) - }) +/// Packs bits selected from successive words into complete output words. +#[derive(Default)] +pub(crate) struct Packer { + /// Bits waiting to be written. + current: u64, + /// Number of populated low bits in `current`. + filled: u32, } -/// Filters two packed bitmaps using the same predicate and mask traversal. -#[inline(never)] -pub(crate) fn filter_bits_compress_pair( - first: &BooleanBuffer, - second: &BooleanBuffer, - predicate: &FilterPredicate, - mut append: impl FnMut(u64, u64, usize), -) { - filter_bits_compress_impl::(first, Some(second), predicate, |first, second, count| { - append(first, second.unwrap(), count) - }) -} - -#[inline(always)] -fn filter_bits_compress_impl( - first: &BooleanBuffer, - second: Option<&BooleanBuffer>, - predicate: &FilterPredicate, - mut append: impl FnMut(u64, Option, usize), -) { - /// Packs the bits extracted from successive words into the low `filled` - /// bits of `current`. - #[derive(Default)] - struct Packer { - /// Bits waiting to be written. - current: u64, - /// Number of populated low bits in `current`. - filled: u32, - } - - impl Packer { - #[inline(always)] - fn push(&mut self, bits: u64, count: u32) -> Option { - self.current |= bits << self.filled; - let total = self.filled + count; - if total < 64 { - self.filled = total; - return None; - } - - let complete = self.current; - // `bits >> (64 - filled)`, written so that `filled == 0` - // shifts everything out - self.current = (bits >> 1) >> (63 - self.filled); - self.filled = total - 64; - Some(complete) - } - - #[inline] - fn finish(self) -> Option<(u64, usize)> { - (self.filled != 0).then_some((self.current, self.filled as usize)) +impl Packer { + #[inline(always)] + pub(crate) fn push(&mut self, values: u64, mask: u64) -> Option { + let bits = bit_util::compress(values, mask); + self.current |= bits << self.filled; + let total = self.filled + mask.count_ones(); + if total < 64 { + self.filled = total; + return None; } - } - assert!(first.len() >= predicate.filter.len()); - if PAIR { - assert!(second.unwrap().len() >= predicate.filter.len()); + let complete = self.current; + // `bits >> (64 - filled)`, written so that `filled == 0` + // shifts everything out + self.current = (bits >> 1) >> (63 - self.filled); + self.filled = total - 64; + Some(complete) } - let mask_chunks = predicate.filter.values().bit_chunks(); - let first_chunks = BitChunks::new(first.values(), first.offset(), predicate.filter.len()); - let second_chunks = second - .map(|buffer| BitChunks::new(buffer.values(), buffer.offset(), predicate.filter.len())); - debug_assert_eq!(predicate.count, predicate.filter.true_count()); - debug_assert_eq!(mask_chunks.chunk_len(), first_chunks.chunk_len()); - if PAIR { - debug_assert_eq!( - mask_chunks.chunk_len(), - second_chunks.as_ref().unwrap().chunk_len() - ); - } - - let mut first_packer = Packer::default(); - let mut second_packer = Packer::default(); - for (index, mask) in mask_chunks.iter().enumerate() { - // Words with no kept bits are skipped before the corresponding values - // are read, so only the mask is touched for them - if mask == 0 { - continue; - } - let count = mask.count_ones(); - let first = first_packer.push(bit_util::compress(first_chunks.chunk(index), mask), count); - let second = PAIR.then(|| { - second_packer.push( - bit_util::compress(second_chunks.as_ref().unwrap().chunk(index), mask), - count, - ) - }); - if let Some(first) = first { - append(first, second.flatten(), 64); - } - } - let mask = mask_chunks.remainder_bits(); - if mask != 0 { - let count = mask.count_ones(); - let first = first_packer.push( - bit_util::compress(first_chunks.remainder_bits(), mask), - count, - ); - let second = PAIR.then(|| { - second_packer.push( - bit_util::compress(second_chunks.as_ref().unwrap().remainder_bits(), mask), - count, - ) - }); - if let Some(first) = first { - append(first, second.flatten(), 64); - } - } - if let Some((first, count)) = first_packer.finish() { - let second = PAIR.then(|| second_packer.finish().unwrap().0); - append(first, second, count); + #[inline] + pub(crate) fn finish(self) -> Option<(u64, usize)> { + (self.filled != 0).then_some((self.current, self.filled as usize)) } } -/// Collects compressed bits into a new buffer. +/// Filter the packed bitmask `buffer` with `predicate` by extracting the kept +/// bits of each 64-bit word with [`bit_util::compress`] (`pext`). /// /// Not inlined: within `filter_array` the packing state spills to the stack. #[inline(never)] -fn collect_filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { - /// Packs the bits extracted from successive words into the low `filled` - /// bits of `current`; once complete it is written at `idx` and restarts - /// from the bits that did not fit. - struct Packer { - ptr: *mut u64, - idx: usize, - current: u64, - filled: u32, - } - - impl Packer { - #[inline(always)] - fn push(&mut self, values: u64, mask: u64) { - let bits = bit_util::compress(values, mask); - self.current |= bits << self.filled; - let total = self.filled + mask.count_ones(); - if total < 64 { - self.filled = total; - } else { - // SAFETY: `count` is the number of set bits in the filter, so - // at most `count / 64` words are ever completed and the - // buffer holds `count / 64 + 1`. - unsafe { self.ptr.add(self.idx).write(self.current) }; - self.idx += 1; - // `bits >> (64 - filled)`, written so that `filled == 0` - // shifts everything out. - self.current = (bits >> 1) >> (63 - self.filled); - self.filled = total - 64; - } - } - } - +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()); @@ -978,28 +858,37 @@ fn collect_filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredic // One word beyond the complete ones for the trailing partial word let mut out: Vec = Vec::with_capacity(predicate.count / 64 + 1); - let mut packer = Packer { - ptr: out.as_mut_ptr(), - idx: 0, - current: 0, - filled: 0, - }; + let ptr = out.as_mut_ptr(); + let mut idx = 0; + let mut packer = Packer::default(); for (index, mask) in mask_chunks.iter().enumerate() { + // Words with no kept bits are skipped before the corresponding values + // are read, so only the mask is touched for them if mask == 0 { continue; } - packer.push(value_chunks.chunk(index), mask); + if let Some(complete) = packer.push(value_chunks.chunk(index), mask) { + // SAFETY: `count` is the number of set bits in the filter, so at + // most `count / 64` words are completed and the output has room + unsafe { ptr.add(idx).write(complete) }; + idx += 1; + } + } + if let Some(complete) = packer.push(value_chunks.remainder_bits(), mask_chunks.remainder_bits()) + { + // SAFETY: as above, the completed word fits in the output allocation + unsafe { ptr.add(idx).write(complete) }; + idx += 1; } - packer.push(value_chunks.remainder_bits(), mask_chunks.remainder_bits()); - // SAFETY: `idx <= count / 64`, so this and every word below it is - // within the buffer and written. - debug_assert!(packer.idx < out.capacity()); - unsafe { - packer.ptr.add(packer.idx).write(packer.current); - out.set_len(packer.idx + 1); + if let Some((pending, _)) = packer.finish() { + // SAFETY: `idx <= count / 64`, so the trailing partial word fits + unsafe { ptr.add(idx).write(pending) }; + idx += 1; } + // SAFETY: every word below `idx` was initialized above + unsafe { out.set_len(idx) }; let mut out = MutableBuffer::from(out); out.truncate(bit_util::ceil(predicate.count, 8)); out.into() @@ -2080,7 +1969,7 @@ mod tests { predicate.strategy ); - let compressed = collect_filter_bits_compress(&values, predicate); + let compressed = filter_bits_compress(&values, predicate); let compressed = BooleanBuffer::new(compressed, 0, predicate.count); assert_eq!(compressed, expected, "compress {case}"); From c3f0045a8811f48fe2a9adb1221703c31618dc3c Mon Sep 17 00:00:00 2001 From: Richard Baah Date: Wed, 4 Nov 2026 09:19:13 -0500 Subject: [PATCH 10/10] fix(filter): restore compressed write assertions --- arrow-select/src/filter.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs index 64fbc7e30baf..cd00af00075a 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -869,6 +869,7 @@ fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> continue; } if let Some(complete) = packer.push(value_chunks.chunk(index), mask) { + debug_assert!(idx < out.capacity()); // SAFETY: `count` is the number of set bits in the filter, so at // most `count / 64` words are completed and the output has room unsafe { ptr.add(idx).write(complete) }; @@ -877,12 +878,14 @@ fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> } if let Some(complete) = packer.push(value_chunks.remainder_bits(), mask_chunks.remainder_bits()) { + debug_assert!(idx < out.capacity()); // SAFETY: as above, the completed word fits in the output allocation unsafe { ptr.add(idx).write(complete) }; idx += 1; } if let Some((pending, _)) = packer.finish() { + debug_assert!(idx < out.capacity()); // SAFETY: `idx <= count / 64`, so the trailing partial word fits unsafe { ptr.add(idx).write(pending) }; idx += 1;