Skip to main content

arrow_select/
filter.rs

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