Skip to main content

arrow_select/
filter.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Defines filter kernels
19
20use std::ops::AddAssign;
21use std::sync::Arc;
22
23use arrow_array::builder::BooleanBufferBuilder;
24use arrow_array::cast::AsArray;
25use arrow_array::types::{
26    ArrowDictionaryKeyType, ArrowPrimitiveType, ByteArrayType, ByteViewType, RunEndIndexType,
27};
28use arrow_array::*;
29use arrow_buffer::{
30    ArrowNativeType, BooleanBuffer, NullBuffer, OffsetBuffer, RunEndBuffer, ScalarBuffer, bit_util,
31};
32use arrow_buffer::{Buffer, MutableBuffer};
33use arrow_data::bit_iterator::{BitIndexIterator, BitSliceIterator};
34use arrow_data::transform::MutableArrayData;
35use arrow_schema::*;
36
37/// If the filter selects more than this fraction of rows, use
38/// [`SlicesIterator`] to copy ranges of values. Otherwise iterate
39/// over individual rows using [`IndexIterator`]
40///
41/// Threshold of 0.8 chosen based on <https://dl.acm.org/doi/abs/10.1145/3465998.3466009>
42///
43const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8;
44
45/// An iterator of `(usize, usize)` each representing an interval
46/// `[start, end)` whose slots of a bitmap [Buffer] are true.
47///
48/// Each interval corresponds to a contiguous region of memory to be
49/// "taken" from an array to be filtered.
50///
51/// ## Notes:
52///
53/// 1. Ignores the validity bitmap (ignores nulls)
54///
55/// 2. Only performant for filters that copy across long contiguous runs
56#[derive(Debug)]
57pub struct SlicesIterator<'a>(BitSliceIterator<'a>);
58
59impl<'a> SlicesIterator<'a> {
60    /// Creates a new iterator from a [BooleanArray]
61    pub fn new(filter: &'a BooleanArray) -> Self {
62        filter.values().into()
63    }
64}
65
66impl<'a> From<&'a BooleanBuffer> for SlicesIterator<'a> {
67    fn from(filter: &'a BooleanBuffer) -> Self {
68        Self(filter.set_slices())
69    }
70}
71
72impl Iterator for SlicesIterator<'_> {
73    type Item = (usize, usize);
74
75    fn next(&mut self) -> Option<Self::Item> {
76        self.0.next()
77    }
78}
79
80/// An iterator of `usize` whose index in [`BooleanArray`] is true
81///
82/// This provides the best performance on most predicates, apart from those which keep
83/// large runs and therefore favour [`SlicesIterator`]
84pub(crate) struct IndexIterator<'a> {
85    remaining: usize,
86    iter: BitIndexIterator<'a>,
87}
88
89impl<'a> IndexIterator<'a> {
90    pub(crate) fn new(filter: &'a BooleanArray, remaining: usize) -> Self {
91        assert_eq!(filter.null_count(), 0);
92        let iter = filter.values().set_indices();
93        Self { remaining, iter }
94    }
95
96    /// Collect this iterator as a [`Vec`]
97    /// This is more efficient than the standard `collect` as we can
98    /// pre-allocate the entire uninitialized buffer and then fill it (roughly 1.6x faster)
99    pub fn collect(mut self) -> Vec<usize> {
100        let len = self.remaining;
101        let mut result = Vec::with_capacity(len);
102        let ptr: *mut usize = result.as_mut_ptr();
103        for i in 0..len {
104            // SAFETY: we have allocated enough space in `result` and remaining
105            // correctly tracks the number of elements
106            let next = self.iter.next();
107            debug_assert!(next.is_some(), "IndexIterator exhausted early");
108            unsafe {
109                *ptr.add(i) = next.unwrap_unchecked();
110            }
111        }
112        // SAFETY: we have initialized `len` elements
113        unsafe {
114            result.set_len(len);
115        }
116        result
117    }
118}
119
120impl Iterator for IndexIterator<'_> {
121    type Item = usize;
122
123    fn next(&mut self) -> Option<Self::Item> {
124        if self.remaining != 0 {
125            // Fascinatingly swapping these two lines around results in a 50%
126            // performance regression for some benchmarks
127            let next = self.iter.next().expect("IndexIterator exhausted early");
128            self.remaining -= 1;
129            // Must panic if exhausted early as trusted length iterator
130            return Some(next);
131        }
132        None
133    }
134
135    fn size_hint(&self) -> (usize, Option<usize>) {
136        (self.remaining, Some(self.remaining))
137    }
138}
139
140/// Convert all null values in `BooleanArray` to `false`
141///
142/// This is useful for filter-like operations which select only `true`
143/// values, but not `false` or `NULL` values
144///
145/// Internally this is implemented as a bitwise `AND` operation with null bits
146/// and the boolean bits.
147///
148/// # Example
149/// ```
150/// # use arrow_array::{Array, BooleanArray};
151/// # use arrow_select::filter::prep_null_mask_filter;
152/// let filter = BooleanArray::from(vec![
153///   Some(true),
154///   Some(false),
155///   None
156/// ]);
157/// // convert Boolean array to a filter mask
158/// let null_mask = prep_null_mask_filter(&filter);
159/// // there are no nulls in the output mask
160/// assert!(null_mask.nulls().is_none());
161/// assert_eq!(null_mask, BooleanArray::from(vec![
162///  true,
163///  false,
164///  false, // Null is converted to false
165/// ]));
166/// ```
167pub fn prep_null_mask_filter(filter: &BooleanArray) -> BooleanArray {
168    let nulls = filter.nulls().unwrap();
169    let mask = filter.values() & nulls.inner();
170    BooleanArray::new(mask, None)
171}
172
173/// Returns a filtered `values` [`Array`] where the corresponding elements of
174/// `predicate` are `true`.
175///
176/// If multiple arrays (or record batches) need to be filtered using the same predicate array,
177/// consider using [FilterBuilder] to create a single [FilterPredicate] and then
178/// calling [FilterPredicate::filter_record_batch].
179///
180/// In contrast to this function, it is then the responsibility of the caller
181/// to use [FilterBuilder::optimize] if appropriate.
182///
183/// # See also
184/// * [`FilterBuilder`] for more control over the filtering process.
185/// * [`filter_record_batch`] to filter a [`RecordBatch`]
186/// * [`BatchCoalescer`]: to filter multiple [`RecordBatch`] and coalesce
187///   the results into a single array.
188///
189/// [`BatchCoalescer`]: crate::coalesce::BatchCoalescer
190///
191/// # Example
192/// ```rust
193/// # use arrow_array::{Int32Array, BooleanArray};
194/// # use arrow_select::filter::filter;
195/// let array = Int32Array::from(vec![5, 6, 7, 8, 9]);
196/// let filter_array = BooleanArray::from(vec![true, false, false, true, false]);
197/// let c = filter(&array, &filter_array).unwrap();
198/// let c = c.as_any().downcast_ref::<Int32Array>().unwrap();
199/// assert_eq!(c, &Int32Array::from(vec![5, 8]));
200/// ```
201pub fn filter(values: &dyn Array, predicate: &BooleanArray) -> Result<ArrayRef, ArrowError> {
202    let mut filter_builder = FilterBuilder::new(predicate);
203
204    if FilterBuilder::is_optimize_beneficial(values.data_type()) {
205        // Only optimize if filtering more than one array
206        // Otherwise, the overhead of optimization can be more than the benefit
207        filter_builder = filter_builder.optimize();
208    }
209
210    let predicate = filter_builder.build();
211
212    filter_array(values, &predicate)
213}
214
215/// Returns a filtered [RecordBatch] where the corresponding elements of
216/// `predicate` are true.
217///
218/// This is the equivalent of calling [filter] on each column of the [RecordBatch].
219///
220/// If multiple record batches (or arrays) need to be filtered using the same predicate array,
221/// consider using [FilterBuilder] to create a single [FilterPredicate] and then
222/// calling [FilterPredicate::filter_record_batch].
223/// In contrast to this function, it is then the responsibility of the caller
224/// to use [FilterBuilder::optimize] if appropriate.
225pub fn filter_record_batch(
226    record_batch: &RecordBatch,
227    predicate: &BooleanArray,
228) -> Result<RecordBatch, ArrowError> {
229    let mut filter_builder = FilterBuilder::new(predicate);
230    let num_cols = record_batch.num_columns();
231    if num_cols > 1
232        || (num_cols > 0
233            && FilterBuilder::is_optimize_beneficial(
234                record_batch.schema_ref().field(0).data_type(),
235            ))
236    {
237        // Only optimize if filtering more than one column or if the column contains multiple internal arrays
238        // Otherwise, the overhead of optimization can be more than the benefit
239        filter_builder = filter_builder.optimize();
240    }
241    let filter = filter_builder.build();
242
243    filter.filter_record_batch(record_batch)
244}
245
246/// A builder to construct [`FilterPredicate`]
247#[derive(Debug)]
248pub struct FilterBuilder {
249    filter: BooleanArray,
250    count: usize,
251    strategy: IterationStrategy,
252}
253
254impl FilterBuilder {
255    /// Create a new [`FilterBuilder`] that can be used to construct a [`FilterPredicate`]
256    pub fn new(filter: &BooleanArray) -> Self {
257        Self::new_with_count(filter, filter.true_count())
258    }
259
260    pub(crate) fn new_with_count(filter: &BooleanArray, count: usize) -> Self {
261        let filter = match filter.null_count() {
262            0 => filter.clone(),
263            _ => prep_null_mask_filter(filter),
264        };
265
266        let strategy = IterationStrategy::default_strategy(filter.len(), count);
267
268        Self {
269            filter,
270            count,
271            strategy,
272        }
273    }
274
275    /// Compute an optimized representation of the provided `filter` mask that can be
276    /// applied to an array more quickly.
277    ///
278    /// When filtering multiple arrays (e.g. a [`RecordBatch`] or a
279    /// [`StructArray`] with multiple fields), optimizing the filter can provide
280    /// significant performance benefits.
281    ///
282    /// However, optimization takes time and can have a larger memory footprint
283    /// than the original mask, so it is often faster to filter a single array,
284    /// without filter optimization.
285    pub fn optimize(mut self) -> Self {
286        match self.strategy {
287            IterationStrategy::SlicesIterator => {
288                let slices = SlicesIterator::new(&self.filter).collect();
289                self.strategy = IterationStrategy::Slices(slices)
290            }
291            IterationStrategy::IndexIterator => {
292                let indices = IndexIterator::new(&self.filter, self.count).collect();
293                self.strategy = IterationStrategy::Indices(indices)
294            }
295            _ => {}
296        }
297        self
298    }
299
300    /// Determines if calling [FilterBuilder::optimize] is beneficial for the
301    /// given type even when filtering just a single array.
302    ///
303    /// See [`FilterBuilder::optimize`] for more details.
304    pub fn is_optimize_beneficial(data_type: &DataType) -> bool {
305        match data_type {
306            DataType::Struct(fields) => {
307                fields.len() > 1
308                    || fields.len() == 1
309                        && FilterBuilder::is_optimize_beneficial(fields[0].data_type())
310            }
311            DataType::Union(fields, UnionMode::Sparse) => !fields.is_empty(),
312            _ => false,
313        }
314    }
315
316    /// Construct the final `FilterPredicate`
317    pub fn build(self) -> FilterPredicate {
318        FilterPredicate {
319            filter: self.filter,
320            count: self.count,
321            strategy: self.strategy,
322        }
323    }
324}
325
326/// The iteration strategy used to evaluate [`FilterPredicate`]
327#[derive(Debug)]
328enum IterationStrategy {
329    /// A lazily evaluated iterator of ranges
330    SlicesIterator,
331    /// A lazily evaluated iterator of indices
332    IndexIterator,
333    /// A precomputed list of indices
334    Indices(Vec<usize>),
335    /// A precomputed array of ranges
336    Slices(Vec<(usize, usize)>),
337    /// Select all rows
338    All,
339    /// Select no rows
340    None,
341}
342
343impl IterationStrategy {
344    /// The default [`IterationStrategy`] for a filter of length `filter_length`
345    /// and selecting `filter_count` rows
346    fn default_strategy(filter_length: usize, filter_count: usize) -> Self {
347        if filter_length == 0 || filter_count == 0 {
348            return IterationStrategy::None;
349        }
350
351        if filter_count == filter_length {
352            return IterationStrategy::All;
353        }
354
355        // Compute the selectivity of the predicate by dividing the number of true
356        // bits in the predicate by the predicate's total length
357        //
358        // This can then be used as a heuristic for the optimal iteration strategy
359        let selectivity_frac = filter_count as f64 / filter_length as f64;
360        if selectivity_frac > FILTER_SLICES_SELECTIVITY_THRESHOLD {
361            return IterationStrategy::SlicesIterator;
362        }
363        IterationStrategy::IndexIterator
364    }
365}
366
367/// Borrowed description of which rows a [`FilterPredicate`] selects.
368pub(crate) enum FilterSelection<'a> {
369    None,
370    All { len: usize },
371    Slices(FilterSlices<'a>),
372    Indices(FilterIndices<'a>),
373}
374
375pub(crate) type FilterSlices<'a> =
376    FilterIterator<std::iter::Copied<std::slice::Iter<'a, (usize, usize)>>, SlicesIterator<'a>>;
377
378pub(crate) type FilterIndices<'a> =
379    FilterIterator<std::iter::Copied<std::slice::Iter<'a, usize>>, IndexIterator<'a>>;
380
381/// Holds either materialized rows or a lazy iterator.
382///
383/// This does not implement [`Iterator`] on purpose. Callers use
384/// [`Self::for_each`] or [`Self::try_for_each`] so the enum is matched once
385/// before the loop, not once per row in `next`.
386pub(crate) enum FilterIterator<M, I> {
387    Materialized(M),
388    Lazy(I),
389}
390
391impl<M, I> FilterIterator<M, I>
392where
393    M: Iterator,
394    I: Iterator<Item = M::Item>,
395{
396    pub(crate) fn for_each<F>(self, f: F)
397    where
398        F: FnMut(M::Item),
399    {
400        match self {
401            Self::Materialized(iter) => iter.for_each(f),
402            Self::Lazy(iter) => iter.for_each(f),
403        }
404    }
405
406    pub(crate) fn try_for_each<F, E>(self, mut f: F) -> Result<(), E>
407    where
408        F: FnMut(M::Item) -> Result<(), E>,
409    {
410        match self {
411            Self::Materialized(iter) => {
412                for item in iter {
413                    f(item)?;
414                }
415            }
416            Self::Lazy(iter) => {
417                for item in iter {
418                    f(item)?;
419                }
420            }
421        }
422
423        Ok(())
424    }
425}
426
427/// A filtering predicate that can be applied to an [`Array`]
428#[derive(Debug)]
429pub struct FilterPredicate {
430    filter: BooleanArray,
431    count: usize,
432    strategy: IterationStrategy,
433}
434
435impl FilterPredicate {
436    /// Selects rows from `values` based on this [`FilterPredicate`]
437    pub fn filter(&self, values: &dyn Array) -> Result<ArrayRef, ArrowError> {
438        filter_array(values, self)
439    }
440
441    /// Returns a filtered [`RecordBatch`] containing only the rows that are selected by this
442    /// [`FilterPredicate`].
443    ///
444    /// This is the equivalent of calling [filter] on each column of the [`RecordBatch`].
445    pub fn filter_record_batch(
446        &self,
447        record_batch: &RecordBatch,
448    ) -> Result<RecordBatch, ArrowError> {
449        let filtered_arrays = record_batch
450            .columns()
451            .iter()
452            .map(|a| filter_array(a, self))
453            .collect::<Result<Vec<_>, _>>()?;
454
455        // SAFETY: we know that the set of filtered arrays will match the schema of the original
456        // record batch
457        unsafe {
458            Ok(RecordBatch::new_unchecked(
459                record_batch.schema(),
460                filtered_arrays,
461                self.count,
462            ))
463        }
464    }
465
466    /// Number of rows being selected based on this [`FilterPredicate`]
467    pub fn count(&self) -> usize {
468        self.count
469    }
470
471    pub(crate) fn selection(&self) -> FilterSelection<'_> {
472        match &self.strategy {
473            IterationStrategy::None => FilterSelection::None,
474            IterationStrategy::All => FilterSelection::All { len: self.count },
475            IterationStrategy::Slices(slices) => {
476                FilterSelection::Slices(FilterIterator::Materialized(slices.iter().copied()))
477            }
478            IterationStrategy::SlicesIterator => {
479                FilterSelection::Slices(FilterIterator::Lazy(SlicesIterator::new(&self.filter)))
480            }
481            IterationStrategy::Indices(indices) => {
482                FilterSelection::Indices(FilterIterator::Materialized(indices.iter().copied()))
483            }
484            IterationStrategy::IndexIterator => FilterSelection::Indices(FilterIterator::Lazy(
485                IndexIterator::new(&self.filter, self.count),
486            )),
487        }
488    }
489
490    /// Filters the given `nulls` buffer using this predicate.
491    ///
492    /// Returns `None` when there is nothing to track in the output, either
493    /// because the input `nulls` was `None`, the input had no nulls, or the
494    /// filtered result has no nulls. Otherwise returns the filtered
495    /// [`NullBuffer`] with its precomputed null count.
496    ///
497    /// # Panics
498    ///
499    /// Panics if the input contains nulls and is shorter than this predicate.
500    pub fn filter_nulls(&self, nulls: Option<&NullBuffer>) -> Option<NullBuffer> {
501        let nulls = nulls?;
502        if nulls.null_count() == 0 {
503            return None;
504        }
505
506        assert!(nulls.len() >= self.filter.len());
507        match self.strategy {
508            IterationStrategy::None => return None,
509            IterationStrategy::All => {
510                let nulls = nulls.slice(0, self.count);
511                return (nulls.null_count() != 0).then_some(nulls);
512            }
513            _ => {}
514        }
515
516        let nulls = filter_bits(nulls.inner(), self);
517        // The filtered `nulls` has a length of `self.count` bits and therefore
518        // the null count is this minus the number of valid bits
519        let null_count = self.count - nulls.count_set_bits_offset(0, self.count);
520
521        if null_count == 0 {
522            return None;
523        }
524
525        let buffer = BooleanBuffer::new(nulls, 0, self.count);
526        debug_assert_eq!(null_count, buffer.len() - buffer.count_set_bits());
527        // SAFETY: `null_count` was derived from `buffer` above, so it matches
528        // the number of unset bits as required by `new_unchecked`.
529        Some(unsafe { NullBuffer::new_unchecked(buffer, null_count) })
530    }
531}
532
533fn filter_array(values: &dyn Array, predicate: &FilterPredicate) -> Result<ArrayRef, ArrowError> {
534    if predicate.filter.len() > values.len() {
535        return Err(ArrowError::InvalidArgumentError(format!(
536            "Filter predicate of length {} is larger than target array of length {}",
537            predicate.filter.len(),
538            values.len()
539        )));
540    }
541
542    match predicate.strategy {
543        IterationStrategy::None => Ok(new_empty_array(values.data_type())),
544        IterationStrategy::All => Ok(values.slice(0, predicate.count)),
545        // actually filter
546        _ => downcast_primitive_array! {
547            values => Ok(Arc::new(filter_primitive(values, predicate))),
548            DataType::Boolean => {
549                let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
550                Ok(Arc::new(filter_boolean(values, predicate)))
551            }
552            DataType::Utf8 => {
553                Ok(Arc::new(filter_bytes(values.as_string::<i32>(), predicate)))
554            }
555            DataType::LargeUtf8 => {
556                Ok(Arc::new(filter_bytes(values.as_string::<i64>(), predicate)))
557            }
558            DataType::Utf8View => {
559                Ok(Arc::new(filter_byte_view(values.as_string_view(), predicate)))
560            }
561            DataType::Binary => {
562                Ok(Arc::new(filter_bytes(values.as_binary::<i32>(), predicate)))
563            }
564            DataType::LargeBinary => {
565                Ok(Arc::new(filter_bytes(values.as_binary::<i64>(), predicate)))
566            }
567            DataType::BinaryView => {
568                Ok(Arc::new(filter_byte_view(values.as_binary_view(), predicate)))
569            }
570            DataType::FixedSizeBinary(_) => {
571                Ok(Arc::new(filter_fixed_size_binary(values.as_fixed_size_binary(), predicate)))
572            }
573            DataType::ListView(_) => {
574                Ok(Arc::new(filter_list_view::<i32>(values.as_list_view(), predicate)))
575            }
576            DataType::LargeListView(_) => {
577                Ok(Arc::new(filter_list_view::<i64>(values.as_list_view(), predicate)))
578            }
579            DataType::RunEndEncoded(_, _) => {
580                downcast_run_array!{
581                    values => Ok(Arc::new(filter_run_end_array(values, predicate)?)),
582                    t => unimplemented!("Filter not supported for RunEndEncoded type {:?}", t)
583                }
584            }
585            DataType::Dictionary(_, _) => downcast_dictionary_array! {
586                values => Ok(Arc::new(filter_dict(values, predicate))),
587                t => unimplemented!("Filter not supported for dictionary type {:?}", t)
588            }
589            DataType::Struct(_) => {
590                Ok(Arc::new(filter_struct(values.as_struct(), predicate)?))
591            }
592            DataType::Union(_, UnionMode::Sparse) => {
593                Ok(Arc::new(filter_sparse_union(values.as_union(), predicate)?))
594            }
595            _ => {
596                let data = values.to_data();
597                // fallback to using MutableArrayData
598                let mut mutable = MutableArrayData::new(
599                    vec![&data],
600                    false,
601                    predicate.count,
602                );
603
604                match &predicate.strategy {
605                    IterationStrategy::Slices(slices) => {
606                        for (start, end) in slices {
607                            mutable.try_extend(0, *start, *end)?;
608                        }
609                    }
610                    _ => {
611                        let iter = SlicesIterator::new(&predicate.filter);
612                        for (start, end) in iter {
613                            mutable.try_extend(0, start, end)?;
614                        }
615                    }
616                }
617
618                let data = mutable.freeze();
619                Ok(make_array(data))
620            }
621        },
622    }
623}
624
625/// Filter any supported [`RunArray`] based on a [`FilterPredicate`]
626fn filter_run_end_array<R: RunEndIndexType>(
627    array: &RunArray<R>,
628    predicate: &FilterPredicate,
629) -> Result<RunArray<R>, ArrowError>
630where
631    R::Native: Into<i64> + From<bool>,
632    R::Native: AddAssign,
633{
634    let run_ends: &RunEndBuffer<R::Native> = array.run_ends();
635    let start_physical = run_ends.get_start_physical_index();
636    let end_physical = run_ends.get_end_physical_index();
637    let physical_len = end_physical - start_physical + 1;
638
639    let mut new_run_ends = vec![R::default_value(); physical_len];
640    let offset = run_ends.offset() as u64;
641
642    let mut start = 0u64;
643    let mut j = 0;
644    let mut count = R::default_value();
645    let filter_values = predicate.filter.values();
646    let run_ends = run_ends.inner();
647
648    let pred: BooleanArray = BooleanBuffer::collect_bool(physical_len, |i| {
649        let mut keep = false;
650        let mut end = (run_ends[i + start_physical].into() as u64).saturating_sub(offset);
651        let difference = end.saturating_sub(filter_values.len() as u64);
652        end -= difference;
653
654        // Safety: we subtract the difference off `end` so we are always within bounds
655        for pred in (start..end).map(|i| unsafe { filter_values.value_unchecked(i as usize) }) {
656            count += R::Native::from(pred);
657            keep |= pred
658        }
659        // this is to avoid branching
660        new_run_ends[j] = count;
661        j += keep as usize;
662
663        start = end;
664        keep
665    })
666    .into();
667
668    new_run_ends.truncate(j);
669
670    let values = array.values_slice();
671    let values = filter(values.as_ref(), &pred)?;
672
673    let run_ends = PrimitiveArray::<R>::try_new(new_run_ends.into(), None)?;
674    RunArray::try_new(&run_ends, &values)
675}
676
677/// Filter the packed bitmask `buffer`, with `predicate` starting at bit offset `offset`
678fn filter_bits(buffer: &BooleanBuffer, predicate: &FilterPredicate) -> Buffer {
679    let src = buffer.values();
680    let offset = buffer.offset();
681    assert!(buffer.len() >= predicate.filter.len());
682
683    match &predicate.strategy {
684        IterationStrategy::IndexIterator => {
685            let bits =
686                // SAFETY: IndexIterator uses the filter predicate to derive indices
687                IndexIterator::new(&predicate.filter, predicate.count).map(|src_idx| unsafe {
688                    bit_util::get_bit_raw(buffer.values().as_ptr(), src_idx + offset)
689                });
690
691            // SAFETY: `IndexIterator` reports its size correctly
692            unsafe { MutableBuffer::from_trusted_len_iter_bool(bits).into() }
693        }
694        IterationStrategy::Indices(indices) => {
695            // SAFETY: indices were derived from the filter predicate
696            let bits = indices.iter().map(|src_idx| unsafe {
697                bit_util::get_bit_raw(buffer.values().as_ptr(), *src_idx + offset)
698            });
699            // SAFETY: `Vec::iter()` reports its size correctly
700            unsafe { MutableBuffer::from_trusted_len_iter_bool(bits).into() }
701        }
702        IterationStrategy::SlicesIterator => {
703            let mut builder = BooleanBufferBuilder::new(predicate.count);
704            for (start, end) in SlicesIterator::new(&predicate.filter) {
705                builder.append_packed_range(start + offset..end + offset, src)
706            }
707            builder.into()
708        }
709        IterationStrategy::Slices(slices) => {
710            let mut builder = BooleanBufferBuilder::new(predicate.count);
711            for (start, end) in slices {
712                builder.append_packed_range(*start + offset..*end + offset, src)
713            }
714            builder.into()
715        }
716        IterationStrategy::All | IterationStrategy::None => unreachable!(),
717    }
718}
719
720/// `filter` implementation for boolean buffers
721fn filter_boolean(array: &BooleanArray, predicate: &FilterPredicate) -> BooleanArray {
722    let buffer = filter_bits(array.values(), predicate);
723    let values = BooleanBuffer::new(buffer, 0, predicate.count);
724    let nulls = predicate.filter_nulls(array.nulls());
725
726    BooleanArray::new(values, nulls)
727}
728
729#[inline(never)]
730pub(crate) fn filter_native<T: ArrowNativeType>(
731    values: &[T],
732    predicate: &FilterPredicate,
733) -> Buffer {
734    assert!(values.len() >= predicate.filter.len());
735
736    match &predicate.strategy {
737        IterationStrategy::SlicesIterator => {
738            let mut buffer = Vec::with_capacity(predicate.count);
739            for (start, end) in SlicesIterator::new(&predicate.filter) {
740                // SAFETY: indices were derived from the filter predicate
741                buffer.extend_from_slice(unsafe { values.get_unchecked(start..end) });
742            }
743            buffer.into()
744        }
745        IterationStrategy::Slices(slices) => {
746            let mut buffer = Vec::with_capacity(predicate.count);
747            for (start, end) in slices {
748                // SAFETY: indices were derived from the filter predicate
749                buffer.extend_from_slice(unsafe { values.get_unchecked(*start..*end) });
750            }
751            buffer.into()
752        }
753        IterationStrategy::IndexIterator => {
754            // SAFETY: indices were derived from the filter predicate
755            let iter = IndexIterator::new(&predicate.filter, predicate.count)
756                .map(|x| unsafe { *values.get_unchecked(x) });
757
758            // SAFETY: IndexIterator is trusted length
759            unsafe { MutableBuffer::from_trusted_len_iter(iter) }.into()
760        }
761        IterationStrategy::Indices(indices) => {
762            // SAFETY: indices were derived from the filter predicate
763            let iter = indices.iter().map(|x| unsafe { *values.get_unchecked(*x) });
764            iter.collect::<Vec<_>>().into()
765        }
766        IterationStrategy::All | IterationStrategy::None => unreachable!(),
767    }
768}
769
770/// `filter` implementation for primitive arrays
771fn filter_primitive<T>(array: &PrimitiveArray<T>, predicate: &FilterPredicate) -> PrimitiveArray<T>
772where
773    T: ArrowPrimitiveType,
774{
775    let buffer = filter_native(array.values(), predicate);
776    let values = ScalarBuffer::new(buffer, 0, predicate.count);
777    let nulls = predicate.filter_nulls(array.nulls());
778    let filtered = PrimitiveArray::new(values, nulls);
779
780    // Avoid the compatibility check when the physical type already matches.
781    if array.data_type() == &T::DATA_TYPE {
782        filtered
783    } else {
784        filtered.with_data_type(array.data_type().clone())
785    }
786}
787
788/// [`FilterBytes`] is created from a source [`GenericByteArray`] and can be
789/// used to build a new [`GenericByteArray`] by copying values from the source
790///
791/// TODO(raphael): Could this be used for the take kernel as well?
792struct FilterBytes<'a, OffsetSize> {
793    src_offsets: &'a [OffsetSize],
794    src_values: &'a [u8],
795    dst_offsets: Vec<OffsetSize>,
796    dst_values: Vec<u8>,
797    cur_offset: OffsetSize,
798}
799
800impl<'a, OffsetSize> FilterBytes<'a, OffsetSize>
801where
802    OffsetSize: OffsetSizeTrait,
803{
804    fn new<T>(capacity: usize, array: &'a GenericByteArray<T>) -> Self
805    where
806        T: ByteArrayType<Offset = OffsetSize>,
807    {
808        let dst_values = Vec::new();
809        let mut dst_offsets: Vec<OffsetSize> = Vec::with_capacity(capacity + 1);
810        let cur_offset = OffsetSize::from_usize(0).unwrap();
811
812        dst_offsets.push(cur_offset);
813
814        Self {
815            src_offsets: array.value_offsets(),
816            src_values: array.value_data(),
817            dst_offsets,
818            dst_values,
819            cur_offset,
820        }
821    }
822
823    /// Returns the byte offset at `idx`
824    #[inline]
825    fn get_value_offset(&self, idx: usize) -> usize {
826        self.src_offsets[idx].as_usize()
827    }
828
829    /// Returns the start and end of the value at index `idx` along with its length
830    #[inline]
831    fn get_value_range(&self, idx: usize) -> (usize, usize, OffsetSize) {
832        // These can only fail if `array` contains invalid data
833        let start = self.get_value_offset(idx);
834        let end = self.get_value_offset(idx + 1);
835        let len = OffsetSize::from_usize(end - start).expect("illegal offset range");
836        (start, end, len)
837    }
838
839    fn extend_offsets_idx(&mut self, iter: impl Iterator<Item = usize>) {
840        self.dst_offsets.extend(iter.map(|idx| {
841            let start = self.src_offsets[idx].as_usize();
842            let end = self.src_offsets[idx + 1].as_usize();
843            let len = OffsetSize::from_usize(end - start).expect("illegal offset range");
844            self.cur_offset += len;
845
846            self.cur_offset
847        }));
848    }
849
850    /// Extends the in-progress array by the indexes in the provided iterator
851    fn extend_idx(&mut self, iter: impl Iterator<Item = usize>) {
852        self.dst_values.reserve_exact(self.cur_offset.as_usize());
853
854        for idx in iter {
855            let start = self.src_offsets[idx].as_usize();
856            let end = self.src_offsets[idx + 1].as_usize();
857            self.dst_values
858                .extend_from_slice(&self.src_values[start..end]);
859        }
860    }
861
862    fn extend_offsets_slices(&mut self, iter: impl Iterator<Item = (usize, usize)>, count: usize) {
863        self.dst_offsets.reserve_exact(count);
864        for (start, end) in iter {
865            // These can only fail if `array` contains invalid data
866            for idx in start..end {
867                let (_, _, len) = self.get_value_range(idx);
868                self.cur_offset += len;
869                self.dst_offsets.push(self.cur_offset);
870            }
871        }
872    }
873
874    /// Extends the in-progress array by the ranges in the provided iterator
875    fn extend_slices(&mut self, iter: impl Iterator<Item = (usize, usize)>) {
876        self.dst_values.reserve_exact(self.cur_offset.as_usize());
877
878        for (start, end) in iter {
879            let value_start = self.get_value_offset(start);
880            let value_end = self.get_value_offset(end);
881            self.dst_values
882                .extend_from_slice(&self.src_values[value_start..value_end]);
883        }
884    }
885}
886
887/// `filter` implementation for byte arrays
888///
889/// Note: NULLs with a non-zero slot length in `array` will have the corresponding
890/// data copied across. This allows handling the null mask separately from the data
891fn filter_bytes<T>(array: &GenericByteArray<T>, predicate: &FilterPredicate) -> GenericByteArray<T>
892where
893    T: ByteArrayType,
894{
895    let mut filter = FilterBytes::new(predicate.count, array);
896
897    match &predicate.strategy {
898        IterationStrategy::SlicesIterator => {
899            filter.extend_offsets_slices(SlicesIterator::new(&predicate.filter), predicate.count);
900            filter.extend_slices(SlicesIterator::new(&predicate.filter))
901        }
902        IterationStrategy::Slices(slices) => {
903            filter.extend_offsets_slices(slices.iter().cloned(), predicate.count);
904            filter.extend_slices(slices.iter().cloned())
905        }
906        IterationStrategy::IndexIterator => {
907            filter.extend_offsets_idx(IndexIterator::new(&predicate.filter, predicate.count));
908            filter.extend_idx(IndexIterator::new(&predicate.filter, predicate.count))
909        }
910        IterationStrategy::Indices(indices) => {
911            filter.extend_offsets_idx(indices.iter().cloned());
912            filter.extend_idx(indices.iter().cloned())
913        }
914        IterationStrategy::All | IterationStrategy::None => unreachable!(),
915    }
916
917    // SAFETY: `dst_offsets` starts at `[0]` and only grows by the running
918    // `cur_offset`, so it is monotonically non-decreasing.
919    let offsets = unsafe { OffsetBuffer::new_unchecked(filter.dst_offsets.into()) };
920    let nulls = predicate.filter_nulls(array.nulls());
921
922    // SAFETY: `offsets` index into `dst_values` by construction, and each slot
923    // is a byte-for-byte copy from `array`, so UTF-8 validity (if any) is preserved.
924    // Length invariant: `offsets.len() - 1 == predicate.count == nulls.len()`.
925    unsafe { GenericByteArray::new_unchecked(offsets, filter.dst_values.into(), nulls) }
926}
927
928/// `filter` implementation for byte view arrays.
929fn filter_byte_view<T: ByteViewType>(
930    array: &GenericByteViewArray<T>,
931    predicate: &FilterPredicate,
932) -> GenericByteViewArray<T> {
933    let new_view_buffer = filter_native(array.views(), predicate);
934    let views = ScalarBuffer::new(new_view_buffer, 0, predicate.count);
935    let buffers = array.data_buffers().to_vec();
936    let nulls = predicate.filter_nulls(array.nulls());
937
938    // SAFETY: each view is copied unchanged from `array.views()` and `buffers`
939    // is the same buffer list, so every view still points to an in-bounds
940    // (and, for strings, UTF-8 valid) range.
941    unsafe { GenericByteViewArray::new_unchecked(views, buffers, nulls) }
942}
943
944fn filter_fixed_size_binary(
945    array: &FixedSizeBinaryArray,
946    predicate: &FilterPredicate,
947) -> FixedSizeBinaryArray {
948    let values: &[u8] = array.values();
949    let value_length = array.value_length() as usize;
950    let calculate_offset_from_index = |index: usize| index * value_length;
951    let buffer = match &predicate.strategy {
952        IterationStrategy::SlicesIterator => {
953            let mut buffer = MutableBuffer::with_capacity(predicate.count * value_length);
954            for (start, end) in SlicesIterator::new(&predicate.filter) {
955                buffer.extend_from_slice(
956                    &values[calculate_offset_from_index(start)..calculate_offset_from_index(end)],
957                );
958            }
959            buffer
960        }
961        IterationStrategy::Slices(slices) => {
962            let mut buffer = MutableBuffer::with_capacity(predicate.count * value_length);
963            for (start, end) in slices {
964                buffer.extend_from_slice(
965                    &values[calculate_offset_from_index(*start)..calculate_offset_from_index(*end)],
966                );
967            }
968            buffer
969        }
970        IterationStrategy::IndexIterator => {
971            let iter = IndexIterator::new(&predicate.filter, predicate.count).map(|x| {
972                &values[calculate_offset_from_index(x)..calculate_offset_from_index(x + 1)]
973            });
974
975            let mut buffer = MutableBuffer::new(predicate.count * value_length);
976            iter.for_each(|item| buffer.extend_from_slice(item));
977            buffer
978        }
979        IterationStrategy::Indices(indices) => {
980            let iter = indices.iter().map(|x| {
981                &values[calculate_offset_from_index(*x)..calculate_offset_from_index(*x + 1)]
982            });
983
984            let mut buffer = MutableBuffer::new(predicate.count * value_length);
985            iter.for_each(|item| buffer.extend_from_slice(item));
986            buffer
987        }
988        IterationStrategy::All | IterationStrategy::None => unreachable!(),
989    };
990
991    let nulls = predicate.filter_nulls(array.nulls());
992
993    FixedSizeBinaryArray::new(array.value_length(), buffer.into(), nulls)
994}
995
996/// `filter` implementation for dictionaries
997fn filter_dict<K: ArrowDictionaryKeyType>(
998    array: &DictionaryArray<K>,
999    predicate: &FilterPredicate,
1000) -> DictionaryArray<K> {
1001    // SAFETY:
1002    // Keys were valid before, filtered subset is therefore still valid
1003    let new_keys = filter_primitive(array.keys(), predicate);
1004    unsafe { DictionaryArray::new_unchecked(new_keys, array.values().clone()) }
1005}
1006
1007/// `filter` implementation for structs
1008fn filter_struct(
1009    array: &StructArray,
1010    predicate: &FilterPredicate,
1011) -> Result<StructArray, ArrowError> {
1012    let columns = array
1013        .columns()
1014        .iter()
1015        .map(|column| filter_array(column, predicate))
1016        .collect::<Result<_, _>>()?;
1017
1018    let nulls = predicate.filter_nulls(array.nulls());
1019
1020    Ok(unsafe {
1021        StructArray::new_unchecked_with_length(
1022            array.fields().clone(),
1023            columns,
1024            nulls,
1025            predicate.count(),
1026        )
1027    })
1028}
1029
1030/// `filter` implementation for sparse unions
1031fn filter_sparse_union(
1032    array: &UnionArray,
1033    predicate: &FilterPredicate,
1034) -> Result<UnionArray, ArrowError> {
1035    let DataType::Union(fields, UnionMode::Sparse) = array.data_type() else {
1036        unreachable!()
1037    };
1038
1039    let type_ids = filter_primitive(
1040        &Int8Array::try_new(array.type_ids().clone(), None)?,
1041        predicate,
1042    );
1043
1044    let children = fields
1045        .iter()
1046        .map(|(child_type_id, _)| filter_array(array.child(child_type_id), predicate))
1047        .collect::<Result<_, _>>()?;
1048
1049    Ok(unsafe {
1050        UnionArray::new_unchecked(fields.clone(), type_ids.into_parts().1, None, children)
1051    })
1052}
1053
1054/// `filter` implementation for list views
1055fn filter_list_view<OffsetType: OffsetSizeTrait>(
1056    array: &GenericListViewArray<OffsetType>,
1057    predicate: &FilterPredicate,
1058) -> GenericListViewArray<OffsetType> {
1059    let filtered_offsets = filter_native::<OffsetType>(array.offsets(), predicate);
1060    let filtered_sizes = filter_native::<OffsetType>(array.sizes(), predicate);
1061
1062    let field = match array.data_type() {
1063        DataType::ListView(field) | DataType::LargeListView(field) => field.clone(),
1064        _ => unreachable!(),
1065    };
1066    let offsets = ScalarBuffer::new(filtered_offsets, 0, predicate.count);
1067    let sizes = ScalarBuffer::new(filtered_sizes, 0, predicate.count);
1068    let values = array.values().clone();
1069    let nulls = predicate.filter_nulls(array.nulls());
1070
1071    // SAFETY: each `(offset, size)` pair is copied unchanged from `array` and
1072    // indexes into the same `values` child, so every range stays in-bounds.
1073    // `field` and `values`' data type are unchanged from `array`.
1074    unsafe { GenericListViewArray::new_unchecked(field, offsets, sizes, values, nulls) }
1075}
1076
1077#[cfg(test)]
1078mod tests {
1079    use super::*;
1080    use arrow_array::builder::*;
1081    use arrow_array::cast::as_run_array;
1082    use arrow_array::types::*;
1083    use rand::distr::uniform::{UniformSampler, UniformUsize};
1084    use rand::distr::{Alphanumeric, StandardUniform};
1085    use rand::prelude::*;
1086    use rand::rng;
1087
1088    #[test]
1089    fn test_filter_nulls_all() {
1090        let predicate = FilterBuilder::new(&BooleanArray::from(vec![true, true])).build();
1091        let nulls = NullBuffer::from(vec![true, false]);
1092        assert_eq!(predicate.filter_nulls(Some(&nulls)), Some(nulls));
1093    }
1094
1095    #[test]
1096    fn test_filter_nulls_none() {
1097        let predicate = FilterBuilder::new(&BooleanArray::from(vec![false, false])).build();
1098        let nulls = NullBuffer::from(vec![true, false]);
1099        assert_eq!(predicate.filter_nulls(Some(&nulls)), None);
1100    }
1101
1102    #[test]
1103    fn test_filter_nulls_selection() {
1104        // Cover full and partial prefixes, empty selections, and both general
1105        // iteration strategies, with and without materializing the selection.
1106        let filters = [
1107            vec![],
1108            vec![true],
1109            vec![true, true],
1110            vec![true; 6],
1111            vec![false; 6],
1112            vec![true, false, false, false, false, false],
1113            vec![true, true, true, true, true, false],
1114        ];
1115        for offset in [0, 3, 9] {
1116            let mut validity = vec![false; offset];
1117            validity.extend([true, false, true, false, true, false]);
1118            let nulls = NullBuffer::from(validity).slice(offset, 6);
1119            for filter in &filters {
1120                let expected: NullBuffer = filter
1121                    .iter()
1122                    .enumerate()
1123                    .filter(|(_, selected)| **selected)
1124                    .map(|(i, _)| nulls.is_valid(i))
1125                    .collect();
1126                let expected = (expected.null_count() != 0).then_some(expected);
1127                let filter = BooleanArray::from(filter.clone());
1128                for optimize in [false, true] {
1129                    let builder = FilterBuilder::new(&filter);
1130                    let predicate = if optimize {
1131                        builder.optimize()
1132                    } else {
1133                        builder
1134                    }
1135                    .build();
1136                    assert_eq!(predicate.filter_nulls(Some(&nulls)), expected);
1137                    assert_eq!(predicate.filter_nulls(None), None);
1138                    assert_eq!(
1139                        predicate.filter_nulls(Some(&NullBuffer::new_valid(6))),
1140                        None
1141                    );
1142                }
1143            }
1144        }
1145    }
1146
1147    #[test]
1148    fn test_filter_nulls_length_check() {
1149        for filter in [vec![true, true], vec![false, false], vec![true, false]] {
1150            let predicate = FilterBuilder::new(&BooleanArray::from(filter)).build();
1151            let nulls = NullBuffer::new_null(1);
1152            assert!(std::panic::catch_unwind(|| predicate.filter_nulls(Some(&nulls))).is_err());
1153            // Inputs without nulls continue to short-circuit before checking length.
1154            assert_eq!(
1155                predicate.filter_nulls(Some(&NullBuffer::new_valid(1))),
1156                None
1157            );
1158            assert_eq!(predicate.filter_nulls(None), None);
1159        }
1160    }
1161
1162    macro_rules! def_temporal_test {
1163        ($test:ident, $array_type: ident, $data: expr) => {
1164            #[test]
1165            fn $test() {
1166                let a = $data;
1167                let b = BooleanArray::from(vec![true, false, true, false]);
1168                let c = filter(&a, &b).unwrap();
1169                let d = c.as_ref().as_any().downcast_ref::<$array_type>().unwrap();
1170                assert_eq!(2, d.len());
1171                assert_eq!(1, d.value(0));
1172                assert_eq!(3, d.value(1));
1173            }
1174        };
1175    }
1176
1177    def_temporal_test!(
1178        test_filter_date32,
1179        Date32Array,
1180        Date32Array::from(vec![1, 2, 3, 4])
1181    );
1182    def_temporal_test!(
1183        test_filter_date64,
1184        Date64Array,
1185        Date64Array::from(vec![1, 2, 3, 4])
1186    );
1187    def_temporal_test!(
1188        test_filter_time32_second,
1189        Time32SecondArray,
1190        Time32SecondArray::from(vec![1, 2, 3, 4])
1191    );
1192    def_temporal_test!(
1193        test_filter_time32_millisecond,
1194        Time32MillisecondArray,
1195        Time32MillisecondArray::from(vec![1, 2, 3, 4])
1196    );
1197    def_temporal_test!(
1198        test_filter_time64_microsecond,
1199        Time64MicrosecondArray,
1200        Time64MicrosecondArray::from(vec![1, 2, 3, 4])
1201    );
1202    def_temporal_test!(
1203        test_filter_time64_nanosecond,
1204        Time64NanosecondArray,
1205        Time64NanosecondArray::from(vec![1, 2, 3, 4])
1206    );
1207    def_temporal_test!(
1208        test_filter_duration_second,
1209        DurationSecondArray,
1210        DurationSecondArray::from(vec![1, 2, 3, 4])
1211    );
1212    def_temporal_test!(
1213        test_filter_duration_millisecond,
1214        DurationMillisecondArray,
1215        DurationMillisecondArray::from(vec![1, 2, 3, 4])
1216    );
1217    def_temporal_test!(
1218        test_filter_duration_microsecond,
1219        DurationMicrosecondArray,
1220        DurationMicrosecondArray::from(vec![1, 2, 3, 4])
1221    );
1222    def_temporal_test!(
1223        test_filter_duration_nanosecond,
1224        DurationNanosecondArray,
1225        DurationNanosecondArray::from(vec![1, 2, 3, 4])
1226    );
1227    def_temporal_test!(
1228        test_filter_timestamp_second,
1229        TimestampSecondArray,
1230        TimestampSecondArray::from(vec![1, 2, 3, 4])
1231    );
1232    def_temporal_test!(
1233        test_filter_timestamp_millisecond,
1234        TimestampMillisecondArray,
1235        TimestampMillisecondArray::from(vec![1, 2, 3, 4])
1236    );
1237    def_temporal_test!(
1238        test_filter_timestamp_microsecond,
1239        TimestampMicrosecondArray,
1240        TimestampMicrosecondArray::from(vec![1, 2, 3, 4])
1241    );
1242    def_temporal_test!(
1243        test_filter_timestamp_nanosecond,
1244        TimestampNanosecondArray,
1245        TimestampNanosecondArray::from(vec![1, 2, 3, 4])
1246    );
1247
1248    #[test]
1249    fn test_filter_array_slice() {
1250        let a = Int32Array::from(vec![5, 6, 7, 8, 9]).slice(1, 4);
1251        let b = BooleanArray::from(vec![true, false, false, true]);
1252        // filtering with sliced filter array is not currently supported
1253        // let b_slice = BooleanArray::from(vec![true, false, false, true, false]).slice(1, 4);
1254        // let b = b_slice.as_any().downcast_ref().unwrap();
1255        let c = filter(&a, &b).unwrap();
1256        let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1257        assert_eq!(2, d.len());
1258        assert_eq!(6, d.value(0));
1259        assert_eq!(9, d.value(1));
1260    }
1261
1262    #[test]
1263    fn test_filter_array_low_density() {
1264        // this test exercises the all 0's branch of the filter algorithm
1265        let mut data_values = (1..=65).collect::<Vec<i32>>();
1266        let mut filter_values = (1..=65).map(|i| matches!(i % 65, 0)).collect::<Vec<bool>>();
1267        // set up two more values after the batch
1268        data_values.extend_from_slice(&[66, 67]);
1269        filter_values.extend_from_slice(&[false, true]);
1270        let a = Int32Array::from(data_values);
1271        let b = BooleanArray::from(filter_values);
1272        let c = filter(&a, &b).unwrap();
1273        let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1274        assert_eq!(2, d.len());
1275        assert_eq!(65, d.value(0));
1276        assert_eq!(67, d.value(1));
1277    }
1278
1279    #[test]
1280    fn test_filter_array_high_density() {
1281        // this test exercises the all 1's branch of the filter algorithm
1282        let mut data_values = (1..=65).map(Some).collect::<Vec<_>>();
1283        let mut filter_values = (1..=65)
1284            .map(|i| !matches!(i % 65, 0))
1285            .collect::<Vec<bool>>();
1286        // set second data value to null
1287        data_values[1] = None;
1288        // set up two more values after the batch
1289        data_values.extend_from_slice(&[Some(66), None, Some(67), None]);
1290        filter_values.extend_from_slice(&[false, true, true, true]);
1291        let a = Int32Array::from(data_values);
1292        let b = BooleanArray::from(filter_values);
1293        let c = filter(&a, &b).unwrap();
1294        let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1295        assert_eq!(67, d.len());
1296        assert_eq!(3, d.null_count());
1297        assert_eq!(1, d.value(0));
1298        assert!(d.is_null(1));
1299        assert_eq!(64, d.value(63));
1300        assert!(d.is_null(64));
1301        assert_eq!(67, d.value(65));
1302    }
1303
1304    #[test]
1305    fn test_filter_string_array_simple() {
1306        let a = StringArray::from(vec!["hello", " ", "world", "!"]);
1307        let b = BooleanArray::from(vec![true, false, true, false]);
1308        let c = filter(&a, &b).unwrap();
1309        let d = c.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
1310        assert_eq!(2, d.len());
1311        assert_eq!("hello", d.value(0));
1312        assert_eq!("world", d.value(1));
1313    }
1314
1315    #[test]
1316    fn test_filter_primitive_array_with_null() {
1317        let a = Int32Array::from(vec![Some(5), None]);
1318        let b = BooleanArray::from(vec![false, true]);
1319        let c = filter(&a, &b).unwrap();
1320        let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1321        assert_eq!(1, d.len());
1322        assert!(d.is_null(0));
1323    }
1324
1325    #[test]
1326    fn test_filter_string_array_with_null() {
1327        let a = StringArray::from(vec![Some("hello"), None, Some("world"), None]);
1328        let b = BooleanArray::from(vec![true, false, false, true]);
1329        let c = filter(&a, &b).unwrap();
1330        let d = c.as_ref().as_any().downcast_ref::<StringArray>().unwrap();
1331        assert_eq!(2, d.len());
1332        assert_eq!("hello", d.value(0));
1333        assert!(!d.is_null(0));
1334        assert!(d.is_null(1));
1335    }
1336
1337    #[test]
1338    fn test_filter_binary_array_with_null() {
1339        let data: Vec<Option<&[u8]>> = vec![Some(b"hello"), None, Some(b"world"), None];
1340        let a = BinaryArray::from(data);
1341        let b = BooleanArray::from(vec![true, false, false, true]);
1342        let c = filter(&a, &b).unwrap();
1343        let d = c.as_ref().as_any().downcast_ref::<BinaryArray>().unwrap();
1344        assert_eq!(2, d.len());
1345        assert_eq!(b"hello", d.value(0));
1346        assert!(!d.is_null(0));
1347        assert!(d.is_null(1));
1348    }
1349
1350    fn _test_filter_byte_view<T>()
1351    where
1352        T: ByteViewType,
1353        str: AsRef<T::Native>,
1354        T::Native: PartialEq,
1355    {
1356        let array = {
1357            // ["hello", "world", null, "large payload over 12 bytes", "lulu"]
1358            let mut builder = GenericByteViewBuilder::<T>::new();
1359            builder.append_value("hello");
1360            builder.append_value("world");
1361            builder.append_null();
1362            builder.append_value("large payload over 12 bytes");
1363            builder.append_value("lulu");
1364            builder.finish()
1365        };
1366
1367        {
1368            let predicate = BooleanArray::from(vec![true, false, true, true, false]);
1369            let actual = filter(&array, &predicate).unwrap();
1370
1371            assert_eq!(actual.len(), 3);
1372
1373            let expected = {
1374                // ["hello", null, "large payload over 12 bytes"]
1375                let mut builder = GenericByteViewBuilder::<T>::new();
1376                builder.append_value("hello");
1377                builder.append_null();
1378                builder.append_value("large payload over 12 bytes");
1379                builder.finish()
1380            };
1381
1382            assert_eq!(actual.as_ref(), &expected);
1383        }
1384
1385        {
1386            let predicate = BooleanArray::from(vec![true, false, false, false, true]);
1387            let actual = filter(&array, &predicate).unwrap();
1388
1389            assert_eq!(actual.len(), 2);
1390
1391            let expected = {
1392                // ["hello", "lulu"]
1393                let mut builder = GenericByteViewBuilder::<T>::new();
1394                builder.append_value("hello");
1395                builder.append_value("lulu");
1396                builder.finish()
1397            };
1398
1399            assert_eq!(actual.as_ref(), &expected);
1400        }
1401    }
1402
1403    #[test]
1404    fn test_filter_string_view() {
1405        _test_filter_byte_view::<StringViewType>()
1406    }
1407
1408    #[test]
1409    fn test_filter_binary_view() {
1410        _test_filter_byte_view::<BinaryViewType>()
1411    }
1412
1413    #[test]
1414    fn test_filter_fixed_binary() {
1415        let v1 = [1_u8, 2];
1416        let v2 = [3_u8, 4];
1417        let v3 = [5_u8, 6];
1418        let v = vec![&v1, &v2, &v3];
1419        let a = FixedSizeBinaryArray::try_from(v).unwrap();
1420        let b = BooleanArray::from(vec![true, false, true]);
1421        let c = filter(&a, &b).unwrap();
1422        let d = c
1423            .as_ref()
1424            .as_any()
1425            .downcast_ref::<FixedSizeBinaryArray>()
1426            .unwrap();
1427        assert_eq!(d.len(), 2);
1428        assert_eq!(d.value(0), &v1);
1429        assert_eq!(d.value(1), &v3);
1430        let c2 = FilterBuilder::new(&b)
1431            .optimize()
1432            .build()
1433            .filter(&a)
1434            .unwrap();
1435        let d2 = c2
1436            .as_ref()
1437            .as_any()
1438            .downcast_ref::<FixedSizeBinaryArray>()
1439            .unwrap();
1440        assert_eq!(d, d2);
1441
1442        let b = BooleanArray::from(vec![false, false, false]);
1443        let c = filter(&a, &b).unwrap();
1444        let d = c
1445            .as_ref()
1446            .as_any()
1447            .downcast_ref::<FixedSizeBinaryArray>()
1448            .unwrap();
1449        assert_eq!(d.len(), 0);
1450
1451        let b = BooleanArray::from(vec![true, true, true]);
1452        let c = filter(&a, &b).unwrap();
1453        let d = c
1454            .as_ref()
1455            .as_any()
1456            .downcast_ref::<FixedSizeBinaryArray>()
1457            .unwrap();
1458        assert_eq!(d.len(), 3);
1459        assert_eq!(d.value(0), &v1);
1460        assert_eq!(d.value(1), &v2);
1461        assert_eq!(d.value(2), &v3);
1462
1463        let b = BooleanArray::from(vec![false, false, true]);
1464        let c = filter(&a, &b).unwrap();
1465        let d = c
1466            .as_ref()
1467            .as_any()
1468            .downcast_ref::<FixedSizeBinaryArray>()
1469            .unwrap();
1470        assert_eq!(d.len(), 1);
1471        assert_eq!(d.value(0), &v3);
1472        let c2 = FilterBuilder::new(&b)
1473            .optimize()
1474            .build()
1475            .filter(&a)
1476            .unwrap();
1477        let d2 = c2
1478            .as_ref()
1479            .as_any()
1480            .downcast_ref::<FixedSizeBinaryArray>()
1481            .unwrap();
1482        assert_eq!(d, d2);
1483    }
1484
1485    #[test]
1486    fn test_filter_array_slice_with_null() {
1487        let a = Int32Array::from(vec![Some(5), None, Some(7), Some(8), Some(9)]).slice(1, 4);
1488        let b = BooleanArray::from(vec![true, false, false, true]);
1489        // filtering with sliced filter array is not currently supported
1490        // let b_slice = BooleanArray::from(vec![true, false, false, true, false]).slice(1, 4);
1491        // let b = b_slice.as_any().downcast_ref().unwrap();
1492        let c = filter(&a, &b).unwrap();
1493        let d = c.as_ref().as_any().downcast_ref::<Int32Array>().unwrap();
1494        assert_eq!(2, d.len());
1495        assert!(d.is_null(0));
1496        assert!(!d.is_null(1));
1497        assert_eq!(9, d.value(1));
1498    }
1499
1500    #[test]
1501    fn test_filter_run_end_encoding_array() {
1502        let run_ends = Int64Array::from(vec![2, 3, 8]);
1503        let values = Int64Array::from(vec![7, -2, 9]);
1504        let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1505        let b = BooleanArray::from(vec![true, false, true, false, true, false, true, false]);
1506        let c = filter(&a, &b).unwrap();
1507        let actual: &RunArray<Int64Type> = as_run_array(&c);
1508        assert_eq!(4, actual.len());
1509
1510        let expected = RunArray::try_new(
1511            &Int64Array::from(vec![1, 2, 4]),
1512            &Int64Array::from(vec![7, -2, 9]),
1513        )
1514        .expect("Failed to make expected RunArray test is broken");
1515
1516        assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1517        assert_eq!(actual.values(), expected.values())
1518    }
1519
1520    #[test]
1521    fn test_filter_run_end_encoding_array_sliced() {
1522        let run_ends = Int64Array::from(vec![2, 3, 8]);
1523        let values = Int64Array::from(vec![7, -2, 9]);
1524        let a = RunArray::try_new(&run_ends, &values).unwrap(); // [7, 7, -2, 9, 9, 9, 9, 9]
1525        let a = a.slice(2, 3); // [-2, 9, 9]
1526        let b = BooleanArray::from(vec![true, false, true]);
1527        let result = filter(&a, &b).unwrap();
1528
1529        let result = result.as_run::<Int64Type>();
1530        let result = result.downcast::<Int64Array>().unwrap();
1531
1532        let expected = vec![-2, 9];
1533        let actual = result.into_iter().flatten().collect::<Vec<_>>();
1534        assert_eq!(expected, actual);
1535    }
1536
1537    #[test]
1538    fn test_filter_run_end_encoding_array_remove_value() {
1539        let run_ends = Int32Array::from(vec![2, 3, 8, 10]);
1540        let values = Int32Array::from(vec![7, -2, 9, -8]);
1541        let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1542        let b = BooleanArray::from(vec![
1543            false, true, false, false, true, false, true, false, false, false,
1544        ]);
1545        let c = filter(&a, &b).unwrap();
1546        let actual: &RunArray<Int32Type> = as_run_array(&c);
1547        assert_eq!(3, actual.len());
1548
1549        let expected =
1550            RunArray::try_new(&Int32Array::from(vec![1, 3]), &Int32Array::from(vec![7, 9]))
1551                .expect("Failed to make expected RunArray test is broken");
1552
1553        assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1554        assert_eq!(actual.values(), expected.values())
1555    }
1556
1557    #[test]
1558    fn test_filter_run_end_encoding_array_remove_all_but_one() {
1559        let run_ends = Int16Array::from(vec![2, 3, 8, 10]);
1560        let values = Int16Array::from(vec![7, -2, 9, -8]);
1561        let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1562        let b = BooleanArray::from(vec![
1563            false, false, false, false, false, false, true, false, false, false,
1564        ]);
1565        let c = filter(&a, &b).unwrap();
1566        let actual: &RunArray<Int16Type> = as_run_array(&c);
1567        assert_eq!(1, actual.len());
1568
1569        let expected = RunArray::try_new(&Int16Array::from(vec![1]), &Int16Array::from(vec![9]))
1570            .expect("Failed to make expected RunArray test is broken");
1571
1572        assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1573        assert_eq!(actual.values(), expected.values())
1574    }
1575
1576    #[test]
1577    fn test_filter_run_end_encoding_array_empty() {
1578        let run_ends = Int64Array::from(vec![2, 3, 8, 10]);
1579        let values = Int64Array::from(vec![7, -2, 9, -8]);
1580        let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1581        let b = BooleanArray::from(vec![
1582            false, false, false, false, false, false, false, false, false, false,
1583        ]);
1584        let c = filter(&a, &b).unwrap();
1585        let actual: &RunArray<Int64Type> = as_run_array(&c);
1586        assert_eq!(0, actual.len());
1587    }
1588
1589    #[test]
1590    fn test_filter_run_end_encoding_array_max_value_gt_predicate_len() {
1591        let run_ends = Int64Array::from(vec![2, 3, 8, 10]);
1592        let values = Int64Array::from(vec![7, -2, 9, -8]);
1593        let a = RunArray::try_new(&run_ends, &values).expect("Failed to create RunArray");
1594        let b = BooleanArray::from(vec![false, true, true]);
1595        let c = filter(&a, &b).unwrap();
1596        let actual: &RunArray<Int64Type> = as_run_array(&c);
1597        assert_eq!(2, actual.len());
1598
1599        let expected = RunArray::try_new(
1600            &Int64Array::from(vec![1, 2]),
1601            &Int64Array::from(vec![7, -2]),
1602        )
1603        .expect("Failed to make expected RunArray test is broken");
1604
1605        assert_eq!(&actual.run_ends().values(), &expected.run_ends().values());
1606        assert_eq!(actual.values(), expected.values())
1607    }
1608
1609    #[test]
1610    fn test_filter_dictionary_array() {
1611        let values = [Some("hello"), None, Some("world"), Some("!")];
1612        let a: Int8DictionaryArray = values.iter().copied().collect();
1613        let b = BooleanArray::from(vec![false, true, true, false]);
1614        let c = filter(&a, &b).unwrap();
1615        let d = c
1616            .as_ref()
1617            .as_any()
1618            .downcast_ref::<Int8DictionaryArray>()
1619            .unwrap();
1620        let value_array = d.values();
1621        let values = value_array.as_any().downcast_ref::<StringArray>().unwrap();
1622        // values are cloned in the filtered dictionary array
1623        assert_eq!(3, values.len());
1624        // but keys are filtered
1625        assert_eq!(2, d.len());
1626        assert!(d.is_null(0));
1627        assert_eq!("world", values.value(d.keys().value(1) as usize));
1628    }
1629
1630    #[test]
1631    fn test_filter_list_array() {
1632        let field = Arc::new(Field::new_list_field(DataType::Int32, false));
1633        let offsets = OffsetBuffer::new(vec![0i64, 3, 6, 8, 8].into());
1634        let value_array = Arc::new(Int32Array::from_iter_values(0..8));
1635        let nulls = Some(NullBuffer::from(vec![true, true, true, false]));
1636        //  a = [[0, 1, 2], [3, 4, 5], [6, 7], null]
1637        let a = LargeListArray::new(field.clone(), offsets, value_array, nulls);
1638        let b = BooleanArray::from(vec![false, true, false, true]);
1639        let result = filter(&a, &b).unwrap();
1640
1641        // expected: [[3, 4, 5], null]
1642        let offsets = OffsetBuffer::new(vec![0i64, 3, 3].into());
1643        let value_array = Arc::new(Int32Array::from_iter_values([3, 4, 5]));
1644        let nulls = Some(NullBuffer::from(vec![true, false]));
1645        let expected: ArrayRef = Arc::new(LargeListArray::new(field, offsets, value_array, nulls));
1646
1647        assert_eq!(&expected, &result);
1648    }
1649
1650    fn test_case_filter_list_view<T: OffsetSizeTrait>() {
1651        // [[1, 2], null, [], [3,4]]
1652        let mut list_array = GenericListViewBuilder::<T, _>::new(Int32Builder::new());
1653        list_array.append_value([Some(1), Some(2)]);
1654        list_array.append_null();
1655        list_array.append_value([]);
1656        list_array.append_value([Some(3), Some(4)]);
1657
1658        let list_array = list_array.finish();
1659        let predicate = BooleanArray::from_iter([true, false, true, false]);
1660
1661        // Filter result: [[1, 2], []]
1662        let filtered = filter(&list_array, &predicate)
1663            .unwrap()
1664            .as_list_view::<T>()
1665            .clone();
1666
1667        let mut expected =
1668            GenericListViewBuilder::<T, _>::with_capacity(Int32Builder::with_capacity(5), 3);
1669        expected.append_value([Some(1), Some(2)]);
1670        expected.append_value([]);
1671        let expected = expected.finish();
1672
1673        assert_eq!(&filtered, &expected);
1674    }
1675
1676    fn test_case_filter_sliced_list_view<T: OffsetSizeTrait>() {
1677        // [[1, 2], null, [], [3,4]]
1678        let mut list_array =
1679            GenericListViewBuilder::<T, _>::with_capacity(Int32Builder::with_capacity(6), 4);
1680        list_array.append_value([Some(1), Some(2)]);
1681        list_array.append_null();
1682        list_array.append_value([]);
1683        list_array.append_value([Some(3), Some(4)]);
1684
1685        let list_array = list_array.finish();
1686
1687        // Sliced: [null, [], [3, 4]]
1688        let sliced = list_array.slice(1, 3);
1689        let predicate = BooleanArray::from_iter([false, false, true]);
1690
1691        // Filter result: [[1, 2], []]
1692        let filtered = filter(&sliced, &predicate)
1693            .unwrap()
1694            .as_list_view::<T>()
1695            .clone();
1696
1697        let mut expected = GenericListViewBuilder::<T, _>::new(Int32Builder::new());
1698        expected.append_value([Some(3), Some(4)]);
1699        let expected = expected.finish();
1700
1701        assert_eq!(&filtered, &expected);
1702    }
1703
1704    #[test]
1705    fn test_filter_list_view_array() {
1706        test_case_filter_list_view::<i32>();
1707        test_case_filter_list_view::<i64>();
1708
1709        test_case_filter_sliced_list_view::<i32>();
1710        test_case_filter_sliced_list_view::<i64>();
1711    }
1712
1713    #[test]
1714    fn test_slice_iterator_bits() {
1715        let filter_values = (0..64).map(|i| i == 1).collect::<Vec<bool>>();
1716        let filter = BooleanArray::from(filter_values);
1717        let filter_count = filter.true_count();
1718
1719        let iter = SlicesIterator::new(&filter);
1720        let chunks = iter.collect::<Vec<_>>();
1721
1722        assert_eq!(chunks, vec![(1, 2)]);
1723        assert_eq!(filter_count, 1);
1724    }
1725
1726    #[test]
1727    fn test_slice_iterator_bits1() {
1728        let filter_values = (0..64).map(|i| i != 1).collect::<Vec<bool>>();
1729        let filter = BooleanArray::from(filter_values);
1730        let filter_count = filter.true_count();
1731
1732        let iter = SlicesIterator::new(&filter);
1733        let chunks = iter.collect::<Vec<_>>();
1734
1735        assert_eq!(chunks, vec![(0, 1), (2, 64)]);
1736        assert_eq!(filter_count, 64 - 1);
1737    }
1738
1739    #[test]
1740    fn test_slice_iterator_chunk_and_bits() {
1741        let filter_values = (0..130).map(|i| i % 62 != 0).collect::<Vec<bool>>();
1742        let filter = BooleanArray::from(filter_values);
1743        let filter_count = filter.true_count();
1744
1745        let iter = SlicesIterator::new(&filter);
1746        let chunks = iter.collect::<Vec<_>>();
1747
1748        assert_eq!(chunks, vec![(1, 62), (63, 124), (125, 130)]);
1749        assert_eq!(filter_count, 61 + 61 + 5);
1750    }
1751
1752    #[test]
1753    fn test_filter_selection_iterators() {
1754        let slices = [(0, 2), (4, 5)];
1755        let mut ranges = Vec::new();
1756        let selection: FilterSlices<'_> = FilterIterator::Materialized(slices.iter().copied());
1757        selection.for_each(|range| ranges.push(range));
1758        assert_eq!(ranges, slices);
1759
1760        let filter = BooleanArray::from(vec![true, true, false, false, true]);
1761        let mut ranges = Vec::new();
1762        let selection: FilterSlices<'_> = FilterIterator::Lazy(SlicesIterator::new(&filter));
1763        selection
1764            .try_for_each(|range| {
1765                ranges.push(range);
1766                Ok::<(), ArrowError>(())
1767            })
1768            .unwrap();
1769        assert_eq!(ranges, vec![(0, 2), (4, 5)]);
1770
1771        let indices = [1, 3, 5];
1772        let mut selected = Vec::new();
1773        let selection: FilterIndices<'_> = FilterIterator::Materialized(indices.iter().copied());
1774        selection.for_each(|idx| selected.push(idx));
1775        assert_eq!(selected, indices);
1776
1777        let filter = BooleanArray::from(vec![false, true, false, true]);
1778        let mut selected = Vec::new();
1779        let selection: FilterIndices<'_> = FilterIterator::Lazy(IndexIterator::new(&filter, 2));
1780        selection
1781            .try_for_each(|idx| {
1782                selected.push(idx);
1783                Ok::<(), ArrowError>(())
1784            })
1785            .unwrap();
1786        assert_eq!(selected, vec![1, 3]);
1787    }
1788
1789    #[test]
1790    fn test_null_mask() {
1791        let a = Int64Array::from(vec![Some(1), Some(2), None]);
1792
1793        let mask1 = BooleanArray::from(vec![Some(true), Some(true), None]);
1794        let out = filter(&a, &mask1).unwrap();
1795        assert_eq!(out.as_ref(), &a.slice(0, 2));
1796    }
1797
1798    #[test]
1799    fn test_filter_record_batch_no_columns() {
1800        let pred = BooleanArray::from(vec![Some(true), Some(true), None]);
1801        let options = RecordBatchOptions::default().with_row_count(Some(100));
1802        let record_batch =
1803            RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options).unwrap();
1804        let out = filter_record_batch(&record_batch, &pred).unwrap();
1805
1806        assert_eq!(out.num_rows(), 2);
1807    }
1808
1809    #[test]
1810    fn test_fast_path() {
1811        let a: PrimitiveArray<Int64Type> = PrimitiveArray::from(vec![Some(1), Some(2), None]);
1812
1813        // all true
1814        let mask = BooleanArray::from(vec![true, true, true]);
1815        let out = filter(&a, &mask).unwrap();
1816        let b = out
1817            .as_any()
1818            .downcast_ref::<PrimitiveArray<Int64Type>>()
1819            .unwrap();
1820        assert_eq!(&a, b);
1821
1822        // all false
1823        let mask = BooleanArray::from(vec![false, false, false]);
1824        let out = filter(&a, &mask).unwrap();
1825        assert_eq!(out.len(), 0);
1826        assert_eq!(out.data_type(), &DataType::Int64);
1827    }
1828
1829    #[test]
1830    fn test_slices() {
1831        // takes up 2 u64s
1832        let bools = std::iter::repeat_n(true, 10)
1833            .chain(std::iter::repeat_n(false, 30))
1834            .chain(std::iter::repeat_n(true, 20))
1835            .chain(std::iter::repeat_n(false, 17))
1836            .chain(std::iter::repeat_n(true, 4));
1837
1838        let bool_array: BooleanArray = bools.map(Some).collect();
1839
1840        let slices: Vec<_> = SlicesIterator::new(&bool_array).collect();
1841        let expected = vec![(0, 10), (40, 60), (77, 81)];
1842        assert_eq!(slices, expected);
1843
1844        // slice with offset and truncated len
1845        let len = bool_array.len();
1846        let sliced_array = bool_array.slice(7, len - 10);
1847        let sliced_array = sliced_array
1848            .as_any()
1849            .downcast_ref::<BooleanArray>()
1850            .unwrap();
1851        let slices: Vec<_> = SlicesIterator::new(sliced_array).collect();
1852        let expected = vec![(0, 3), (33, 53), (70, 71)];
1853        assert_eq!(slices, expected);
1854    }
1855
1856    fn test_slices_fuzz(mask_len: usize, offset: usize, truncate: usize) {
1857        let mut rng = rng();
1858
1859        let bools: Vec<bool> = std::iter::from_fn(|| Some(rng.random()))
1860            .take(mask_len)
1861            .collect();
1862
1863        let buffer = Buffer::from_iter(bools.iter().cloned());
1864
1865        let truncated_length = mask_len - offset - truncate;
1866
1867        let filter = BooleanArray::new(BooleanBuffer::new(buffer, offset, truncated_length), None);
1868
1869        let slice_bits: Vec<_> = SlicesIterator::new(&filter)
1870            .flat_map(|(start, end)| start..end)
1871            .collect();
1872
1873        let count = filter.true_count();
1874        let index_bits: Vec<_> = IndexIterator::new(&filter, count).collect();
1875
1876        let expected_bits: Vec<_> = bools
1877            .iter()
1878            .skip(offset)
1879            .take(truncated_length)
1880            .enumerate()
1881            .flat_map(|(idx, v)| v.then(|| idx))
1882            .collect();
1883
1884        assert_eq!(slice_bits, expected_bits);
1885        assert_eq!(index_bits, expected_bits);
1886    }
1887
1888    #[test]
1889    #[cfg_attr(miri, ignore)]
1890    fn fuzz_test_slices_iterator() {
1891        let mut rng = rng();
1892
1893        let uusize = UniformUsize::new(usize::MIN, usize::MAX).unwrap();
1894        for _ in 0..100 {
1895            let mask_len = rng.random_range(0..1024);
1896            let max_offset = 64.min(mask_len);
1897            let offset = uusize.sample(&mut rng).checked_rem(max_offset).unwrap_or(0);
1898
1899            let max_truncate = 128.min(mask_len - offset);
1900            let truncate = uusize
1901                .sample(&mut rng)
1902                .checked_rem(max_truncate)
1903                .unwrap_or(0);
1904
1905            test_slices_fuzz(mask_len, offset, truncate);
1906        }
1907
1908        test_slices_fuzz(64, 0, 0);
1909        test_slices_fuzz(64, 8, 0);
1910        test_slices_fuzz(64, 8, 8);
1911        test_slices_fuzz(32, 8, 8);
1912        test_slices_fuzz(32, 5, 9);
1913    }
1914
1915    /// Filters `values` by `predicate` using standard rust iterators
1916    fn filter_rust<T>(values: impl IntoIterator<Item = T>, predicate: &[bool]) -> Vec<T> {
1917        values
1918            .into_iter()
1919            .zip(predicate)
1920            .filter(|(_, x)| **x)
1921            .map(|(a, _)| a)
1922            .collect()
1923    }
1924
1925    /// Generates an array of length `len` with `valid_percent` non-null values
1926    fn gen_primitive<T>(len: usize, valid_percent: f64) -> Vec<Option<T>>
1927    where
1928        StandardUniform: Distribution<T>,
1929    {
1930        let mut rng = rng();
1931        (0..len)
1932            .map(|_| rng.random_bool(valid_percent).then(|| rng.random()))
1933            .collect()
1934    }
1935
1936    /// Generates an array of length `len` with `valid_percent` non-null values
1937    fn gen_strings(
1938        len: usize,
1939        valid_percent: f64,
1940        str_len_range: std::ops::Range<usize>,
1941    ) -> Vec<Option<String>> {
1942        let mut rng = rng();
1943        (0..len)
1944            .map(|_| {
1945                rng.random_bool(valid_percent).then(|| {
1946                    let len = rng.random_range(str_len_range.clone());
1947                    (0..len)
1948                        .map(|_| char::from(rng.sample(Alphanumeric)))
1949                        .collect()
1950                })
1951            })
1952            .collect()
1953    }
1954
1955    /// Returns an iterator that calls `Option::as_deref` on each item
1956    fn as_deref<T: std::ops::Deref>(src: &[Option<T>]) -> impl Iterator<Item = Option<&T::Target>> {
1957        src.iter().map(|x| x.as_deref())
1958    }
1959
1960    #[test]
1961    #[cfg_attr(miri, ignore)]
1962    fn fuzz_filter() {
1963        let mut rng = rng();
1964
1965        for i in 0..100 {
1966            let filter_percent = match i {
1967                0..=4 => 1.,
1968                5..=10 => 0.,
1969                _ => rng.random_range(0.0..1.0),
1970            };
1971
1972            let valid_percent = rng.random_range(0.0..1.0);
1973
1974            let array_len = rng.random_range(32..256);
1975            let array_offset = rng.random_range(0..10);
1976
1977            // Construct a predicate
1978            let filter_offset = rng.random_range(0..10);
1979            let filter_truncate = rng.random_range(0..10);
1980            let bools: Vec<_> = std::iter::from_fn(|| Some(rng.random_bool(filter_percent)))
1981                .take(array_len + filter_offset - filter_truncate)
1982                .collect();
1983
1984            let predicate = BooleanArray::from_iter(bools.iter().cloned().map(Some));
1985
1986            // Offset predicate
1987            let predicate = predicate.slice(filter_offset, array_len - filter_truncate);
1988            let predicate = predicate.as_any().downcast_ref::<BooleanArray>().unwrap();
1989            let bools = &bools[filter_offset..];
1990
1991            // Test i32
1992            let values = gen_primitive(array_len + array_offset, valid_percent);
1993            let src = Int32Array::from_iter(values.iter().cloned());
1994
1995            let src = src.slice(array_offset, array_len);
1996            let src = src.as_any().downcast_ref::<Int32Array>().unwrap();
1997            let values = &values[array_offset..];
1998
1999            let filtered = filter(src, predicate).unwrap();
2000            let array = filtered.as_any().downcast_ref::<Int32Array>().unwrap();
2001            let actual: Vec<_> = array.iter().collect();
2002
2003            assert_eq!(actual, filter_rust(values.iter().cloned(), bools));
2004
2005            // Test string
2006            let strings = gen_strings(array_len + array_offset, valid_percent, 0..20);
2007            let src = StringArray::from_iter(as_deref(&strings));
2008
2009            let src = src.slice(array_offset, array_len);
2010            let src = src.as_any().downcast_ref::<StringArray>().unwrap();
2011
2012            let filtered = filter(src, predicate).unwrap();
2013            let array = filtered.as_any().downcast_ref::<StringArray>().unwrap();
2014            let actual: Vec<_> = array.iter().collect();
2015
2016            let expected_strings = filter_rust(as_deref(&strings[array_offset..]), bools);
2017            assert_eq!(actual, expected_strings);
2018
2019            // Test string dictionary
2020            let src = DictionaryArray::<Int32Type>::from_iter(as_deref(&strings));
2021
2022            let src = src.slice(array_offset, array_len);
2023            let src = src
2024                .as_any()
2025                .downcast_ref::<DictionaryArray<Int32Type>>()
2026                .unwrap();
2027
2028            let filtered = filter(src, predicate).unwrap();
2029
2030            let array = filtered
2031                .as_any()
2032                .downcast_ref::<DictionaryArray<Int32Type>>()
2033                .unwrap();
2034
2035            let values = array
2036                .values()
2037                .as_any()
2038                .downcast_ref::<StringArray>()
2039                .unwrap();
2040
2041            let actual: Vec<_> = array
2042                .keys()
2043                .iter()
2044                .map(|key| key.map(|key| values.value(key as usize)))
2045                .collect();
2046
2047            assert_eq!(actual, expected_strings);
2048        }
2049    }
2050
2051    #[test]
2052    fn test_filter_map() {
2053        let mut builder =
2054            MapBuilder::new(None, StringBuilder::new(), Int64Builder::with_capacity(4));
2055        // [{"key1": 1}, {"key2": 2, "key3": 3}, null, {"key1": 1}
2056        builder.keys().append_value("key1");
2057        builder.values().append_value(1);
2058        builder.append(true).unwrap();
2059        builder.keys().append_value("key2");
2060        builder.keys().append_value("key3");
2061        builder.values().append_value(2);
2062        builder.values().append_value(3);
2063        builder.append(true).unwrap();
2064        builder.append(false).unwrap();
2065        builder.keys().append_value("key1");
2066        builder.values().append_value(1);
2067        builder.append(true).unwrap();
2068        let maparray = Arc::new(builder.finish()) as ArrayRef;
2069
2070        let indices = vec![Some(true), Some(false), Some(false), Some(true)]
2071            .into_iter()
2072            .collect::<BooleanArray>();
2073        let got = filter(&maparray, &indices).unwrap();
2074
2075        let mut builder =
2076            MapBuilder::new(None, StringBuilder::new(), Int64Builder::with_capacity(2));
2077        builder.keys().append_value("key1");
2078        builder.values().append_value(1);
2079        builder.append(true).unwrap();
2080        builder.keys().append_value("key1");
2081        builder.values().append_value(1);
2082        builder.append(true).unwrap();
2083        let expected = Arc::new(builder.finish()) as ArrayRef;
2084
2085        assert_eq!(&expected, &got);
2086    }
2087
2088    #[test]
2089    fn test_filter_fixed_size_list_arrays() {
2090        let field = Arc::new(Field::new_list_field(DataType::Int32, false));
2091        let value_array = Arc::new(Int32Array::from_iter_values(0..9));
2092        let array = FixedSizeListArray::new(field, 3, value_array, None);
2093
2094        let filter_array = BooleanArray::from(vec![true, false, false]);
2095
2096        let c = filter(&array, &filter_array).unwrap();
2097        let filtered = c.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
2098
2099        assert_eq!(filtered.len(), 1);
2100
2101        let list = filtered.value(0);
2102        assert_eq!(
2103            &[0, 1, 2],
2104            list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2105        );
2106
2107        let filter_array = BooleanArray::from(vec![true, false, true]);
2108
2109        let c = filter(&array, &filter_array).unwrap();
2110        let filtered = c.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
2111
2112        assert_eq!(filtered.len(), 2);
2113
2114        let list = filtered.value(0);
2115        assert_eq!(
2116            &[0, 1, 2],
2117            list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2118        );
2119        let list = filtered.value(1);
2120        assert_eq!(
2121            &[6, 7, 8],
2122            list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2123        );
2124    }
2125
2126    #[test]
2127    fn test_filter_fixed_size_list_arrays_with_null() {
2128        let field = Arc::new(Field::new_list_field(DataType::Int32, false));
2129        let value_array = Arc::new(Int32Array::from_iter_values(0..10));
2130        let nulls = Some(NullBuffer::from(vec![true, false, false, true, true]));
2131        let array = FixedSizeListArray::new(field, 2, value_array, nulls);
2132
2133        let filter_array = BooleanArray::from(vec![true, true, false, true, false]);
2134
2135        let c = filter(&array, &filter_array).unwrap();
2136        let filtered = c.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
2137
2138        assert_eq!(filtered.len(), 3);
2139
2140        let list = filtered.value(0);
2141        assert_eq!(
2142            &[0, 1],
2143            list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2144        );
2145        assert!(filtered.is_null(1));
2146        let list = filtered.value(2);
2147        assert_eq!(
2148            &[6, 7],
2149            list.as_any().downcast_ref::<Int32Array>().unwrap().values()
2150        );
2151    }
2152
2153    fn test_filter_union_array(array: UnionArray) {
2154        let filter_array = BooleanArray::from(vec![true, false, false]);
2155        let c = filter(&array, &filter_array).unwrap();
2156        let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2157
2158        let mut builder = UnionBuilder::new_dense();
2159        builder.append::<Int32Type>("A", 1).unwrap();
2160        let expected_array = builder.build().unwrap();
2161
2162        compare_union_arrays(filtered, &expected_array);
2163
2164        let filter_array = BooleanArray::from(vec![true, false, true]);
2165        let c = filter(&array, &filter_array).unwrap();
2166        let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2167
2168        let mut builder = UnionBuilder::new_dense();
2169        builder.append::<Int32Type>("A", 1).unwrap();
2170        builder.append::<Int32Type>("A", 34).unwrap();
2171        let expected_array = builder.build().unwrap();
2172
2173        compare_union_arrays(filtered, &expected_array);
2174
2175        let filter_array = BooleanArray::from(vec![true, true, false]);
2176        let c = filter(&array, &filter_array).unwrap();
2177        let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2178
2179        let mut builder = UnionBuilder::new_dense();
2180        builder.append::<Int32Type>("A", 1).unwrap();
2181        builder.append::<Float64Type>("B", 3.2).unwrap();
2182        let expected_array = builder.build().unwrap();
2183
2184        compare_union_arrays(filtered, &expected_array);
2185    }
2186
2187    #[test]
2188    fn test_filter_union_array_dense() {
2189        let mut builder = UnionBuilder::new_dense();
2190        builder.append::<Int32Type>("A", 1).unwrap();
2191        builder.append::<Float64Type>("B", 3.2).unwrap();
2192        builder.append::<Int32Type>("A", 34).unwrap();
2193        let array = builder.build().unwrap();
2194
2195        test_filter_union_array(array);
2196    }
2197
2198    #[test]
2199    fn test_filter_run_union_array_dense() {
2200        let mut builder = UnionBuilder::new_dense();
2201        builder.append::<Int32Type>("A", 1).unwrap();
2202        builder.append::<Int32Type>("A", 3).unwrap();
2203        builder.append::<Int32Type>("A", 34).unwrap();
2204        let array = builder.build().unwrap();
2205
2206        let filter_array = BooleanArray::from(vec![true, true, false]);
2207        let c = filter(&array, &filter_array).unwrap();
2208        let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2209
2210        let mut builder = UnionBuilder::new_dense();
2211        builder.append::<Int32Type>("A", 1).unwrap();
2212        builder.append::<Int32Type>("A", 3).unwrap();
2213        let expected = builder.build().unwrap();
2214
2215        assert_eq!(filtered.to_data(), expected.to_data());
2216    }
2217
2218    #[test]
2219    fn test_filter_union_array_dense_with_nulls() {
2220        let mut builder = UnionBuilder::new_dense();
2221        builder.append::<Int32Type>("A", 1).unwrap();
2222        builder.append::<Float64Type>("B", 3.2).unwrap();
2223        builder.append_null::<Float64Type>("B").unwrap();
2224        builder.append::<Int32Type>("A", 34).unwrap();
2225        let array = builder.build().unwrap();
2226
2227        let filter_array = BooleanArray::from(vec![true, true, false, false]);
2228        let c = filter(&array, &filter_array).unwrap();
2229        let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2230
2231        let mut builder = UnionBuilder::new_dense();
2232        builder.append::<Int32Type>("A", 1).unwrap();
2233        builder.append::<Float64Type>("B", 3.2).unwrap();
2234        let expected_array = builder.build().unwrap();
2235
2236        compare_union_arrays(filtered, &expected_array);
2237
2238        let filter_array = BooleanArray::from(vec![true, false, true, false]);
2239        let c = filter(&array, &filter_array).unwrap();
2240        let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2241
2242        let mut builder = UnionBuilder::new_dense();
2243        builder.append::<Int32Type>("A", 1).unwrap();
2244        builder.append_null::<Float64Type>("B").unwrap();
2245        let expected_array = builder.build().unwrap();
2246
2247        compare_union_arrays(filtered, &expected_array);
2248    }
2249
2250    #[test]
2251    fn test_filter_union_array_sparse() {
2252        let mut builder = UnionBuilder::new_sparse();
2253        builder.append::<Int32Type>("A", 1).unwrap();
2254        builder.append::<Float64Type>("B", 3.2).unwrap();
2255        builder.append::<Int32Type>("A", 34).unwrap();
2256        let array = builder.build().unwrap();
2257
2258        test_filter_union_array(array);
2259    }
2260
2261    #[test]
2262    fn test_filter_union_array_sparse_with_nulls() {
2263        let mut builder = UnionBuilder::new_sparse();
2264        builder.append::<Int32Type>("A", 1).unwrap();
2265        builder.append::<Float64Type>("B", 3.2).unwrap();
2266        builder.append_null::<Float64Type>("B").unwrap();
2267        builder.append::<Int32Type>("A", 34).unwrap();
2268        let array = builder.build().unwrap();
2269
2270        let filter_array = BooleanArray::from(vec![true, false, true, false]);
2271        let c = filter(&array, &filter_array).unwrap();
2272        let filtered = c.as_any().downcast_ref::<UnionArray>().unwrap();
2273
2274        let mut builder = UnionBuilder::new_sparse();
2275        builder.append::<Int32Type>("A", 1).unwrap();
2276        builder.append_null::<Float64Type>("B").unwrap();
2277        let expected_array = builder.build().unwrap();
2278
2279        compare_union_arrays(filtered, &expected_array);
2280    }
2281
2282    fn compare_union_arrays(union1: &UnionArray, union2: &UnionArray) {
2283        assert_eq!(union1.len(), union2.len());
2284
2285        for i in 0..union1.len() {
2286            let type_id = union1.type_id(i);
2287
2288            let slot1 = union1.value(i);
2289            let slot2 = union2.value(i);
2290
2291            assert_eq!(slot1.is_null(0), slot2.is_null(0));
2292
2293            if !slot1.is_null(0) && !slot2.is_null(0) {
2294                match type_id {
2295                    0 => {
2296                        let slot1 = slot1.as_any().downcast_ref::<Int32Array>().unwrap();
2297                        assert_eq!(slot1.len(), 1);
2298                        let value1 = slot1.value(0);
2299
2300                        let slot2 = slot2.as_any().downcast_ref::<Int32Array>().unwrap();
2301                        assert_eq!(slot2.len(), 1);
2302                        let value2 = slot2.value(0);
2303                        assert_eq!(value1, value2);
2304                    }
2305                    1 => {
2306                        let slot1 = slot1.as_any().downcast_ref::<Float64Array>().unwrap();
2307                        assert_eq!(slot1.len(), 1);
2308                        let value1 = slot1.value(0);
2309
2310                        let slot2 = slot2.as_any().downcast_ref::<Float64Array>().unwrap();
2311                        assert_eq!(slot2.len(), 1);
2312                        let value2 = slot2.value(0);
2313                        assert_eq!(value1, value2);
2314                    }
2315                    _ => unreachable!(),
2316                }
2317            }
2318        }
2319    }
2320
2321    #[test]
2322    fn test_filter_struct() {
2323        let predicate = BooleanArray::from(vec![true, false, true, false]);
2324
2325        let a = Arc::new(StringArray::from(vec!["hello", " ", "world", "!"]));
2326        let a_filtered = Arc::new(StringArray::from(vec!["hello", "world"]));
2327
2328        let b = Arc::new(Int32Array::from(vec![5, 6, 7, 8]));
2329        let b_filtered = Arc::new(Int32Array::from(vec![5, 7]));
2330
2331        let null_mask = NullBuffer::from(vec![true, false, false, true]);
2332        let null_mask_filtered = NullBuffer::from(vec![true, false]);
2333
2334        let a_field = Field::new("a", DataType::Utf8, false);
2335        let b_field = Field::new("b", DataType::Int32, false);
2336
2337        let array = StructArray::new(vec![a_field.clone()].into(), vec![a.clone()], None);
2338        let expected =
2339            StructArray::new(vec![a_field.clone()].into(), vec![a_filtered.clone()], None);
2340
2341        let result = filter(&array, &predicate).unwrap();
2342
2343        assert_eq!(result.to_data(), expected.to_data());
2344
2345        let array = StructArray::new(
2346            vec![a_field.clone()].into(),
2347            vec![a.clone()],
2348            Some(null_mask.clone()),
2349        );
2350        let expected = StructArray::new(
2351            vec![a_field.clone()].into(),
2352            vec![a_filtered.clone()],
2353            Some(null_mask_filtered.clone()),
2354        );
2355
2356        let result = filter(&array, &predicate).unwrap();
2357
2358        assert_eq!(result.to_data(), expected.to_data());
2359
2360        let array = StructArray::new(
2361            vec![a_field.clone(), b_field.clone()].into(),
2362            vec![a.clone(), b.clone()],
2363            None,
2364        );
2365        let expected = StructArray::new(
2366            vec![a_field.clone(), b_field.clone()].into(),
2367            vec![a_filtered.clone(), b_filtered.clone()],
2368            None,
2369        );
2370
2371        let result = filter(&array, &predicate).unwrap();
2372
2373        assert_eq!(result.to_data(), expected.to_data());
2374
2375        let array = StructArray::new(
2376            vec![a_field.clone(), b_field.clone()].into(),
2377            vec![a.clone(), b.clone()],
2378            Some(null_mask.clone()),
2379        );
2380
2381        let expected = StructArray::new(
2382            vec![a_field.clone(), b_field.clone()].into(),
2383            vec![a_filtered.clone(), b_filtered.clone()],
2384            Some(null_mask_filtered.clone()),
2385        );
2386
2387        let result = filter(&array, &predicate).unwrap();
2388
2389        assert_eq!(result.to_data(), expected.to_data());
2390    }
2391
2392    #[test]
2393    fn test_filter_empty_struct() {
2394        /*
2395            "a": {
2396                "b": int64,
2397                "c": {}
2398            },
2399        */
2400        let fields = arrow_schema::Field::new(
2401            "a",
2402            arrow_schema::DataType::Struct(arrow_schema::Fields::from(vec![
2403                arrow_schema::Field::new("b", arrow_schema::DataType::Int64, true),
2404                arrow_schema::Field::new(
2405                    "c",
2406                    arrow_schema::DataType::Struct(arrow_schema::Fields::empty()),
2407                    true,
2408                ),
2409            ])),
2410            true,
2411        );
2412
2413        /* Test record
2414            {"a":{"c": {}}}
2415            {"a":{"c": {}}}
2416            {"a":{"c": {}}}
2417        */
2418
2419        // Create the record batch with the nested struct array
2420        let schema = Arc::new(Schema::new(vec![fields]));
2421
2422        let b = Arc::new(Int64Array::from(vec![None, None, None]));
2423        let c = Arc::new(StructArray::new_empty_fields(
2424            3,
2425            Some(NullBuffer::from(vec![true, true, true])),
2426        ));
2427        let a = StructArray::new(
2428            vec![
2429                Field::new("b", DataType::Int64, true),
2430                Field::new("c", DataType::Struct(Fields::empty()), true),
2431            ]
2432            .into(),
2433            vec![b.clone(), c.clone()],
2434            Some(NullBuffer::from(vec![true, true, true])),
2435        );
2436        let record_batch = RecordBatch::try_new(schema, vec![Arc::new(a)]).unwrap();
2437        println!("{record_batch:?}");
2438
2439        // Apply the filter
2440        let predicate = BooleanArray::from(vec![true, false, true]);
2441        let filtered_batch = filter_record_batch(&record_batch, &predicate).unwrap();
2442
2443        // The filtered batch should have 2 rows (the 1st and 3rd)
2444        assert_eq!(filtered_batch.num_rows(), 2);
2445    }
2446
2447    #[test]
2448    #[should_panic]
2449    fn test_filter_bits_too_large() {
2450        let buffer = BooleanBuffer::from(vec![false; 8]);
2451        let predicate = BooleanArray::from(vec![true; 9]);
2452        let filter = FilterBuilder::new(&predicate).build();
2453        filter_bits(&buffer, &filter);
2454    }
2455
2456    #[test]
2457    #[should_panic]
2458    fn test_filter_native_too_large() {
2459        let values = vec![1; 8];
2460        let predicate = BooleanArray::from(vec![false; 9]);
2461        let filter = FilterBuilder::new(&predicate).build();
2462        filter_native(&values, &filter);
2463    }
2464}