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