diff --git a/arrow-select/src/coalesce/boolean.rs b/arrow-select/src/coalesce/boolean.rs index 17c76b2b6196..63311d252ddc 100644 --- a/arrow-select/src/coalesce/boolean.rs +++ b/arrow-select/src/coalesce/boolean.rs @@ -16,7 +16,7 @@ // 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::{BooleanBuffer, BooleanBufferBuilder, NullBufferBuilder}; @@ -77,8 +77,35 @@ 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(); + 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); + }, + ); + } else { + filter_bits_compress(source.values(), filter, |values, count| { + output_values.append_word(values, count) + }); + self.nulls.append_n_non_nulls(filter.count()); + } Ok(()) } @@ -109,6 +136,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 +184,172 @@ 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 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; + } + 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..d3aac27916f7 100644 --- a/arrow-select/src/filter.rs +++ b/arrow-select/src/filter.rs @@ -751,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) } @@ -801,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)] -fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer { +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)] +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 @@ -863,18 +893,59 @@ 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, _| { + 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 out = MutableBuffer::from(out); out.truncate(bit_util::ceil(predicate.count, 8)); out.into() @@ -1955,7 +2026,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}"); 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))