Skip to main content

arrow_select/
filter.rs

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