Skip to main content

parquet/arrow/arrow_reader/
filter.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::arrow::ProjectionMask;
6use crate::arrow::arrow_reader::{RowSelection, RowSelectionPolicy};
7use crate::schema::types::SchemaDescriptor;
8use arrow_array::{Array, BooleanArray, RecordBatch};
9use arrow_buffer::BooleanBuffer;
10use arrow_schema::ArrowError;
11use arrow_select::filter::{SlicesIterator, filter_record_batch, prep_null_mask_filter};
12use std::fmt::{Debug, Formatter};
13
14/// A predicate operating on [`RecordBatch`]
15///
16/// See also:
17/// * [`RowFilter`] for more information  on applying filters during the
18///   Parquet decoding process.
19/// * [`ArrowPredicateFn`] for a concrete implementation based on a function
20pub trait ArrowPredicate: Send + 'static {
21    /// Returns the [`ProjectionMask`] that describes the columns required
22    /// to evaluate this predicate.
23    ///
24    /// All projected columns will be provided in the `batch` passed to
25    /// [`evaluate`](Self::evaluate). The projection mask should be as small as
26    /// possible because any columns needed for the overall projection mask are
27    /// decoded again after a predicate is applied.
28    fn projection(&self) -> &ProjectionMask;
29
30    /// Evaluate this predicate for the given [`RecordBatch`] containing the columns
31    /// identified by [`Self::projection`]
32    ///
33    /// Must return a [`BooleanArray`] that has the same length as the input
34    /// `batch` where each row indicates whether the row should be returned:
35    /// * `true`:the row should be returned
36    /// * `false` or `null`: the row should not be returned
37    fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray, ArrowError>;
38}
39
40/// An [`ArrowPredicate`] created from an [`FnMut`] and a [`ProjectionMask`]
41///
42/// See [`RowFilter`] for more information on applying filters during the
43/// Parquet decoding process.
44///
45/// The function is passed `RecordBatch`es with only the columns specified in
46/// the [`ProjectionMask`].
47///
48/// The function must return a [`BooleanArray`] that has the same length as the
49/// input `batch` where each row indicates whether the row should be returned:
50/// * `true`: the row should be returned
51/// * `false` or `null`: the row should not be returned
52///
53/// # Example:
54///
55/// Given an input schema: `"a:int64", "b:int64"`, you can create a predicate that
56/// evaluates `b > 0` like this:
57///
58/// ```
59/// # use std::sync::Arc;
60/// # use arrow::compute::kernels::cmp::gt;
61/// # use arrow_array::{BooleanArray, Int64Array, RecordBatch};
62/// # use arrow_array::cast::AsArray;
63/// # use arrow_array::types::Int64Type;
64/// # use parquet::arrow::arrow_reader::ArrowPredicateFn;
65/// # use parquet::arrow::ProjectionMask;
66/// # use parquet::schema::types::{SchemaDescriptor, Type};
67/// # use parquet::basic; // note there are two `Type`s that are different
68/// # // Schema for a table with one columns: "a" (int64) and "b" (int64)
69/// # let descriptor = SchemaDescriptor::new(
70/// #  Arc::new(
71/// #    Type::group_type_builder("my_schema")
72/// #      .with_fields(vec![
73/// #        Arc::new(
74/// #         Type::primitive_type_builder("a", basic::Type::INT64)
75/// #          .build().unwrap()
76/// #        ),
77/// #        Arc::new(
78/// #         Type::primitive_type_builder("b", basic::Type::INT64)
79/// #          .build().unwrap()
80/// #        ),
81/// #     ])
82/// #     .build().unwrap()
83/// #  )
84/// # );
85/// // Create a mask for selecting only the second column "b" (index 1)
86/// let projection_mask = ProjectionMask::leaves(&descriptor, [1]);
87/// // Closure that evaluates "b > 0"
88/// let predicate = |batch: RecordBatch| {
89///    let scalar_0 = Int64Array::new_scalar(0);
90///    let column = batch.column(0).as_primitive::<Int64Type>();
91///    // call the gt kernel to compute `>` which returns a BooleanArray
92///    gt(column, &scalar_0)
93///  };
94/// // Create ArrowPredicateFn that can be passed to RowFilter
95/// let arrow_predicate = ArrowPredicateFn::new(projection_mask, predicate);
96/// ```
97pub struct ArrowPredicateFn<F> {
98    f: F,
99    projection: ProjectionMask,
100}
101
102impl<F> ArrowPredicateFn<F>
103where
104    F: FnMut(RecordBatch) -> Result<BooleanArray, ArrowError> + Send + 'static,
105{
106    /// Create a new [`ArrowPredicateFn`] that invokes `f` on the columns
107    /// specified in `projection`.
108    pub fn new(projection: ProjectionMask, f: F) -> Self {
109        Self { f, projection }
110    }
111}
112
113impl<F> ArrowPredicate for ArrowPredicateFn<F>
114where
115    F: FnMut(RecordBatch) -> Result<BooleanArray, ArrowError> + Send + 'static,
116{
117    fn projection(&self) -> &ProjectionMask {
118        &self.projection
119    }
120
121    fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray, ArrowError> {
122        (self.f)(batch)
123    }
124}
125
126/// Filter applied *during* the parquet read process
127///
128/// See example on [`ArrowReaderBuilder::with_row_filter`]
129///
130/// [`RowFilter`] applies predicates in order, after decoding only the columns
131/// required. As predicates eliminate rows, fewer rows from subsequent columns
132/// may be required, thus potentially reducing IO and decode. This process is
133/// also known as *push down* filtering and  *late materialization*.
134///
135/// A `RowFilter` consists of a list of [`ArrowPredicate`]s. Only the rows for which
136/// all the predicates evaluate to `true` will be returned.
137/// Any [`RowSelection`] provided to the reader will be applied prior
138/// to the first predicate, and each predicate in turn will then be used to compute
139/// a more refined [`RowSelection`] used when evaluating the subsequent predicates.
140///
141/// Once all predicates have been evaluated, the final [`RowSelection`] is applied
142/// to the top-level [`ProjectionMask`] to produce the final output [`RecordBatch`].
143///
144/// This design has a couple of implications:
145///
146/// * [`RowFilter`] can be used to skip entire pages, and thus IO, in addition to CPU decode overheads
147/// * Columns may be decoded multiple times if they appear in multiple [`ProjectionMask`].
148///   Consecutive predicates whose projection is the same single top-level, non-repeated
149///   column are evaluated together on one decoded stream, so ordering such predicates
150///   next to each other avoids decoding that column again
151/// * IO will be deferred until needed by a [`ProjectionMask`]
152///
153/// As such there is a trade-off between a single large predicate, or multiple predicates,
154/// that will depend on the shape of the data. Whilst multiple smaller predicates may
155/// minimise the amount of data scanned/decoded, it may not be faster overall.
156///
157/// For example, if a predicate that needs a single column of data filters out all but
158/// 1% of the rows, applying it as one of the early `ArrowPredicateFn` will likely significantly
159/// improve performance.
160///
161/// As a counter example, if a predicate needs several columns of data to evaluate but
162/// leaves 99% of the rows, it may be better to not filter the data from parquet and
163/// apply the filter after the RecordBatch has been fully decoded.
164///
165/// Additionally, even if a predicate eliminates a moderate number of rows, it may still be faster
166/// to filter the data after the RecordBatch has been fully decoded, if the eliminated rows are
167/// not contiguous.
168///
169/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection
170/// [`ArrowReaderBuilder::with_row_filter`]: crate::arrow::arrow_reader::ArrowReaderBuilder::with_row_filter
171pub struct RowFilter {
172    /// A list of [`ArrowPredicate`]
173    pub(crate) predicates: Vec<Box<dyn ArrowPredicate>>,
174}
175
176impl Debug for RowFilter {
177    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
178        write!(f, "RowFilter {{ {} predicates: }}", self.predicates.len())
179    }
180}
181
182impl RowFilter {
183    /// Create a new [`RowFilter`] from an array of [`ArrowPredicate`]
184    pub fn new(predicates: Vec<Box<dyn ArrowPredicate>>) -> Self {
185        Self { predicates }
186    }
187    /// Returns the inner predicates
188    pub fn predicates(&self) -> &Vec<Box<dyn ArrowPredicate>> {
189        &self.predicates
190    }
191    /// Returns the inner predicates, consuming self
192    pub fn into_predicates(self) -> Vec<Box<dyn ArrowPredicate>> {
193        self.predicates
194    }
195
196    /// Fuse consecutive predicates on the same single top-level, non-repeated leaf.
197    /// This avoids repeated decoding or predicate-cache replay of that column.
198    pub(crate) fn fuse_same_projection(self, parquet_schema: &SchemaDescriptor) -> Self {
199        let mut predicates: Vec<Box<dyn ArrowPredicate>> =
200            Vec::with_capacity(self.predicates.len());
201        let mut group: Vec<Box<dyn ArrowPredicate>> = Vec::new();
202        let mut flush_group = |group: &mut Vec<Box<dyn ArrowPredicate>>| {
203            if group.len() > 1 && can_fuse_projection(group[0].projection(), parquet_schema) {
204                let group = std::mem::take(group);
205                predicates.push(Box::new(FusedPredicate::new(group)));
206            } else {
207                predicates.append(group);
208            }
209        };
210
211        for predicate in self.predicates {
212            if group
213                .last()
214                .is_some_and(|last| last.projection() != predicate.projection())
215            {
216                flush_group(&mut group);
217            }
218            group.push(predicate);
219        }
220        flush_group(&mut group);
221
222        Self { predicates }
223    }
224}
225
226/// Restrict fusion to one top-level, non-repeated leaf to limit compaction costs.
227fn can_fuse_projection(projection: &ProjectionMask, parquet_schema: &SchemaDescriptor) -> bool {
228    let mut leaf_indices =
229        (0..parquet_schema.num_columns()).filter(|idx| projection.leaf_included(*idx));
230    let Some(leaf_idx) = leaf_indices.next() else {
231        return false;
232    };
233    if leaf_indices.next().is_some() {
234        return false;
235    }
236
237    let column = parquet_schema.column(leaf_idx);
238    column.path().parts().len() == 1 && column.max_rep_level() == 0
239}
240
241/// Evaluate same-projection predicates in order on one decoded batch.
242/// Later predicates see only surviving rows, sliced when contiguous or compacted
243/// otherwise.
244struct FusedPredicate {
245    /// At least two predicates, all with the same projection.
246    predicates: Vec<Box<dyn ArrowPredicate>>,
247    /// Chooses how intermediate selections are composed. The output is always
248    /// a [`BooleanArray`], so this is independent of the reader's
249    /// [`RowSelectionPolicy`], which governs decoding.
250    composition_policy: RowSelectionPolicy,
251}
252
253impl FusedPredicate {
254    /// Create a fused predicate from at least two predicates that share one
255    /// projection.
256    fn new(predicates: Vec<Box<dyn ArrowPredicate>>) -> Self {
257        debug_assert!(predicates.len() > 1);
258        debug_assert!(
259            predicates
260                .windows(2)
261                .all(|pair| pair[0].projection() == pair[1].projection())
262        );
263        Self {
264            predicates,
265            composition_policy: RowSelectionPolicy::default(),
266        }
267    }
268}
269
270impl Debug for FusedPredicate {
271    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
272        write!(
273            f,
274            "FusedPredicate {{ {} predicates }}",
275            self.predicates.len()
276        )
277    }
278}
279
280impl ArrowPredicate for FusedPredicate {
281    fn projection(&self) -> &ProjectionMask {
282        self.predicates[0].projection()
283    }
284
285    fn evaluate(&mut self, batch: RecordBatch) -> Result<BooleanArray, ArrowError> {
286        let num_rows = batch.num_rows();
287        // Positions in the original batch; None until a predicate rejects rows.
288        let mut selection: Option<RowSelection> = None;
289        let mut filtered_batch = batch;
290        let last_predicate_idx = self.predicates.len() - 1;
291
292        for (idx, predicate) in self.predicates.iter_mut().enumerate() {
293            let filter = evaluate_predicate(predicate.as_mut(), filtered_batch.clone())?;
294            // No mapping is needed if all preceding predicates accepted every row.
295            if idx == last_predicate_idx && selection.is_none() {
296                return Ok(filter);
297            }
298            let true_count = filter.true_count();
299            if true_count == 0 {
300                return Ok(BooleanArray::new(BooleanBuffer::new_unset(num_rows), None));
301            }
302            if true_count == filter.len() {
303                continue;
304            }
305
306            let predicate_selection = RowSelection::from_boolean_buffer(filter.values().clone());
307            selection = Some(match selection.take() {
308                // Only the accumulated selection drives the composition
309                // algorithm, so adapt it once, right before it is used.
310                Some(prev) => self
311                    .composition_policy
312                    .apply(prev)
313                    .and_then(&predicate_selection),
314                None => predicate_selection,
315            });
316            if idx != last_predicate_idx {
317                filtered_batch = narrow_batch(&filtered_batch, &filter, true_count)?;
318            }
319        }
320
321        let mask = match selection {
322            Some(selection) => selection.into_boolean_buffer(),
323            None => BooleanBuffer::new_set(num_rows),
324        };
325        debug_assert_eq!(mask.len(), num_rows);
326        Ok(BooleanArray::new(mask, None))
327    }
328}
329
330/// Validate the predicate result length and treat nulls as false.
331fn evaluate_predicate(
332    predicate: &mut dyn ArrowPredicate,
333    batch: RecordBatch,
334) -> Result<BooleanArray, ArrowError> {
335    let input_rows = batch.num_rows();
336    let filter = predicate.evaluate(batch)?;
337    if filter.len() != input_rows {
338        return Err(ArrowError::InvalidArgumentError(format!(
339            "ArrowPredicate predicate returned {} rows, expected {input_rows}",
340            filter.len()
341        )));
342    }
343    Ok(match filter.null_count() {
344        0 => filter,
345        _ => prep_null_mask_filter(&filter),
346    })
347}
348
349/// Restrict `batch` to the rows `filter` accepts, slicing zero-copy when they
350/// form one contiguous range and compacting otherwise.
351///
352/// `filter` must be null-free and accept some but not all rows.
353fn narrow_batch(
354    batch: &RecordBatch,
355    filter: &BooleanArray,
356    true_count: usize,
357) -> Result<RecordBatch, ArrowError> {
358    let mut slices = SlicesIterator::new(filter);
359    let (start, end) = slices
360        .next()
361        .expect("a partially selected filter has a true slice");
362    if end - start == true_count {
363        Ok(batch.slice(start, end - start))
364    } else {
365        filter_record_batch(batch, filter)
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372    use crate::arrow::array_reader::ArrayReader;
373    use crate::arrow::array_reader::StructArrayReader;
374    use crate::arrow::array_reader::test_util::make_int32_page_reader;
375    use crate::arrow::arrow_reader::ReadPlanBuilder;
376    use crate::schema::parser::parse_message_type;
377    use arrow_array::Int32Array;
378    use arrow_array::cast::AsArray;
379    use arrow_array::types::Int32Type;
380    use arrow_schema::{DataType, Field, Fields, Schema};
381    use std::sync::Arc;
382
383    fn test_schema() -> SchemaDescriptor {
384        let schema = parse_message_type(
385            "message schema {
386                REQUIRED INT32 a;
387                REQUIRED INT32 b;
388                REQUIRED INT32 c;
389                REQUIRED GROUP nested { REQUIRED INT32 d; }
390                REPEATED INT32 e;
391            }",
392        )
393        .unwrap();
394        SchemaDescriptor::new(Arc::new(schema))
395    }
396
397    #[test]
398    fn can_fuse_projection_requires_one_top_level_leaf() {
399        let schema = test_schema();
400
401        let one_leaf = ProjectionMask::leaves(&schema, [1]);
402        let two_leaves = ProjectionMask::leaves(&schema, [0, 2]);
403        let nested_leaf = ProjectionMask::leaves(&schema, [3]);
404        let repeated_leaf = ProjectionMask::leaves(&schema, [4]);
405        assert!(can_fuse_projection(&one_leaf, &schema));
406        assert!(!can_fuse_projection(&two_leaves, &schema));
407        assert!(!can_fuse_projection(&nested_leaf, &schema));
408        assert!(!can_fuse_projection(&repeated_leaf, &schema));
409        assert!(!can_fuse_projection(
410            &ProjectionMask::none(schema.num_columns()),
411            &schema
412        ));
413    }
414
415    fn accept_all(projection: ProjectionMask) -> Box<dyn ArrowPredicate> {
416        Box::new(ArrowPredicateFn::new(projection, |batch| {
417            Ok(BooleanArray::from(vec![true; batch.num_rows()]))
418        }))
419    }
420
421    #[test]
422    fn fuse_same_projection_groups_consecutive_eligible_predicates() {
423        let schema = test_schema();
424        let a = ProjectionMask::leaves(&schema, [0]);
425        let b = ProjectionMask::leaves(&schema, [1]);
426        let ac = ProjectionMask::leaves(&schema, [0, 2]);
427        let nested = ProjectionMask::leaves(&schema, [3]);
428
429        let filter = RowFilter::new(vec![
430            // fused
431            accept_all(a.clone()),
432            accept_all(a.clone()),
433            // single predicate: kept as-is
434            accept_all(b.clone()),
435            // two leaves: not eligible for fusion
436            accept_all(ac.clone()),
437            accept_all(ac.clone()),
438            // nested leaf: not eligible for fusion
439            accept_all(nested.clone()),
440            accept_all(nested.clone()),
441            // fused
442            accept_all(a.clone()),
443            accept_all(a.clone()),
444            accept_all(a.clone()),
445        ])
446        .fuse_same_projection(&schema);
447
448        let projections: Vec<_> = filter
449            .predicates()
450            .iter()
451            .map(|predicate| predicate.projection().clone())
452            .collect();
453        assert_eq!(
454            projections,
455            vec![a.clone(), b, ac.clone(), ac, nested.clone(), nested, a]
456        );
457        assert_eq!(format!("{filter:?}"), "RowFilter { 7 predicates: }");
458    }
459
460    #[test]
461    fn fuse_same_projection_keeps_single_predicates_and_empty_filters() {
462        let schema = test_schema();
463        let a = ProjectionMask::leaves(&schema, [0]);
464        let b = ProjectionMask::leaves(&schema, [1]);
465
466        let filter =
467            RowFilter::new(vec![accept_all(a), accept_all(b)]).fuse_same_projection(&schema);
468        assert_eq!(filter.predicates().len(), 2);
469
470        let filter = RowFilter::new(vec![]).fuse_same_projection(&schema);
471        assert!(filter.predicates().is_empty());
472    }
473
474    fn int_batch(values: impl IntoIterator<Item = i32>) -> RecordBatch {
475        let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)]));
476        let values = Int32Array::from_iter_values(values);
477        RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap()
478    }
479
480    /// A predicate over the `v` column that asserts `precondition` holds for
481    /// every row it sees before returning `f(v)`.
482    fn int_predicate(
483        precondition: impl Fn(i32) -> bool + Send + 'static,
484        f: impl Fn(i32) -> Option<bool> + Send + 'static,
485    ) -> Box<dyn ArrowPredicate> {
486        Box::new(ArrowPredicateFn::new(ProjectionMask::all(), move |batch| {
487            let values = batch.column(0).as_primitive::<Int32Type>();
488            assert!(
489                values.values().iter().all(|v| precondition(*v)),
490                "predicate saw rows rejected by a predecessor"
491            );
492            Ok(values.values().iter().map(|v| f(*v)).collect())
493        }))
494    }
495
496    /// Fragmented survivors: roughly every other row.
497    fn even() -> Box<dyn ArrowPredicate> {
498        int_predicate(|_| true, |v| Some(v % 2 == 0))
499    }
500
501    /// Clustered survivors: one contiguous range.
502    fn in_range() -> Box<dyn ArrowPredicate> {
503        int_predicate(|_| true, |v| Some((10..40).contains(&v)))
504    }
505
506    /// Applies after `even`: mixes nulls (rejected) into the result.
507    fn divisible_by_three_after_even() -> Box<dyn ArrowPredicate> {
508        int_predicate(
509            |v| v % 2 == 0,
510            |v| if v % 7 == 0 { None } else { Some(v % 3 == 0) },
511        )
512    }
513
514    fn all_true() -> Box<dyn ArrowPredicate> {
515        int_predicate(|_| true, |_| Some(true))
516    }
517
518    fn all_false() -> Box<dyn ArrowPredicate> {
519        int_predicate(|_| true, |_| Some(false))
520    }
521
522    fn never_called() -> Box<dyn ArrowPredicate> {
523        Box::new(ArrowPredicateFn::new(ProjectionMask::all(), |_| {
524            panic!("predicate evaluated after every row was rejected")
525        }))
526    }
527
528    /// Evaluate `predicates` one after another, each on the rows kept by its
529    /// predecessors, the way `RowFilter` applies separate predicates.
530    fn sequential(batch: &RecordBatch, predicates: &mut [Box<dyn ArrowPredicate>]) -> Vec<bool> {
531        let mut kept = vec![true; batch.num_rows()];
532        for predicate in predicates {
533            // Like `RowFilter`, stop once every row has been rejected.
534            if !kept.contains(&true) {
535                break;
536            }
537            let survivors = filter_record_batch(batch, &BooleanArray::from(kept.clone())).unwrap();
538            let filter = predicate.evaluate(survivors).unwrap();
539            for (filter_idx, keep) in kept.iter_mut().filter(|keep| **keep).enumerate() {
540                *keep = filter.is_valid(filter_idx) && filter.value(filter_idx);
541            }
542        }
543        kept
544    }
545
546    const POLICIES: [RowSelectionPolicy; 4] = [
547        RowSelectionPolicy::Selectors,
548        RowSelectionPolicy::Mask,
549        RowSelectionPolicy::Auto { threshold: 32 },
550        RowSelectionPolicy::Auto { threshold: 2 },
551    ];
552
553    type PredicateChain = fn() -> Vec<Box<dyn ArrowPredicate>>;
554
555    #[test]
556    fn fused_predicate_matches_sequential_evaluation() {
557        let cases: Vec<PredicateChain> = vec![
558            || vec![even(), divisible_by_three_after_even()],
559            || vec![even(), in_range(), divisible_by_three_after_even()],
560            || vec![in_range(), even()],
561            || vec![even(), all_true(), divisible_by_three_after_even()],
562            || vec![all_true(), even()],
563            || vec![all_true(), all_true()],
564            || vec![even(), all_false(), never_called()],
565            || vec![all_false(), never_called()],
566        ];
567
568        for num_rows in [0, 1, 7, 97, 200] {
569            let batch = int_batch(0..num_rows);
570            for make_predicates in &cases {
571                let expected = sequential(&batch, &mut make_predicates());
572                for policy in POLICIES {
573                    let mut fused = FusedPredicate {
574                        predicates: make_predicates(),
575                        composition_policy: policy,
576                    };
577                    let actual = fused.evaluate(batch.clone()).unwrap();
578                    assert_eq!(actual.null_count(), 0);
579                    assert_eq!(
580                        actual.values().iter().collect::<Vec<_>>(),
581                        expected,
582                        "{num_rows} rows, {policy:?}"
583                    );
584                }
585            }
586        }
587    }
588
589    #[test]
590    fn fused_predicate_rejects_wrong_result_length() {
591        let short = Box::new(ArrowPredicateFn::new(ProjectionMask::all(), |_| {
592            Ok(BooleanArray::from(vec![true]))
593        }));
594        let mut fused = FusedPredicate::new(vec![all_true(), short]);
595        let err = fused.evaluate(int_batch(0..4)).unwrap_err();
596        assert!(
597            err.to_string()
598                .contains("ArrowPredicate predicate returned 1 rows, expected 4"),
599            "{err}"
600        );
601    }
602
603    #[test]
604    fn narrow_batch_slices_contiguous_survivors() {
605        let batch = int_batch(0..6);
606        let filter = BooleanArray::from(vec![false, false, true, true, true, false]);
607        let narrowed = narrow_batch(&batch, &filter, 3).unwrap();
608        assert_eq!(
609            narrowed.column(0).as_primitive::<Int32Type>().values(),
610            &[2, 3, 4]
611        );
612        // Slicing shares the input buffer instead of copying it.
613        assert_eq!(narrowed.column(0).to_data().buffers()[0].as_ptr(), unsafe {
614            batch.column(0).to_data().buffers()[0].as_ptr().add(2 * 4)
615        });
616
617        let filter = BooleanArray::from(vec![false, true, false, true, false, false]);
618        let narrowed = narrow_batch(&batch, &filter, 2).unwrap();
619        assert_eq!(
620            narrowed.column(0).as_primitive::<Int32Type>().values(),
621            &[1, 3]
622        );
623    }
624
625    /// Fusing predicates into one `ReadPlanBuilder::with_predicate` call
626    /// produces the same selection as applying them one at a time.
627    #[test]
628    fn fused_predicate_matches_predicate_major_read_plan() {
629        let data: Vec<i32> = (0..97).collect();
630        let make_reader = || {
631            let levels = vec![0; data.len()];
632            let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
633            let struct_type =
634                DataType::Struct(Fields::from(vec![Field::new("c0", DataType::Int32, false)]));
635            Box::new(StructArrayReader::new(
636                struct_type,
637                vec![leaf],
638                0,
639                0,
640                false,
641                None,
642            )) as Box<dyn ArrayReader>
643        };
644        let make_predicates = || vec![even(), divisible_by_three_after_even()];
645
646        let prior = RowSelection::from_filters(&[BooleanArray::from(
647            (0..data.len()).map(|idx| idx % 5 != 0).collect::<Vec<_>>(),
648        )]);
649        for initial in [None, Some(prior)] {
650            let mut sequential = ReadPlanBuilder::new(7).with_selection(initial.clone());
651            for predicate in &mut make_predicates() {
652                sequential = sequential
653                    .with_predicate(make_reader(), predicate.as_mut())
654                    .unwrap();
655            }
656
657            for policy in POLICIES {
658                let mut fused = FusedPredicate::new(make_predicates());
659                let plan = ReadPlanBuilder::new(7)
660                    .with_selection(initial.clone())
661                    .with_row_selection_policy(policy)
662                    .with_predicate(make_reader(), &mut fused)
663                    .unwrap();
664                assert_eq!(plan.selection(), sequential.selection(), "{policy:?}");
665            }
666        }
667    }
668}