parquet/arrow/arrow_reader/

read_plan.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! [`ReadPlan`] and [`ReadPlanBuilder`] for determining which rows to read
//! from a Parquet file

use crate::arrow::array_reader::ArrayReader;
use crate::arrow::arrow_reader::selection::{
    LoadedRowRanges, RowSelectionInner, RowSelectionPolicy, RowSelectionStrategy,
};
use crate::arrow::arrow_reader::{
    ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector,
};
use crate::errors::{ParquetError, Result};
use arrow_array::{Array, BooleanArray};
use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
use arrow_select::filter::prep_null_mask_filter;
use std::sync::Arc;

/// Options for [`ReadPlanBuilder::with_predicate_options`].
pub struct PredicateOptions<'a> {
    array_reader: Box<dyn ArrayReader>,
    predicate: &'a mut dyn ArrowPredicate,
    limit: Option<usize>,
    total_rows: usize,
}

impl<'a> PredicateOptions<'a> {
    /// Create options for evaluating `predicate` against rows produced by
    /// `array_reader`.
    ///
    /// By default there is no match-count limit; the predicate is evaluated
    /// over every row the reader yields. Use [`Self::with_limit`] to enable
    /// early termination.
    pub fn new(array_reader: Box<dyn ArrayReader>, predicate: &'a mut dyn ArrowPredicate) -> Self {
        Self {
            array_reader,
            predicate,
            limit: None,
            total_rows: 0,
        }
    }

    /// Stop scanning `array_reader` once `limit` matches have accumulated.
    ///
    /// Performance optimization for `LIMIT` / TopK: when the cumulative
    /// `true_count` reaches `limit`, the current filter batch is truncated
    /// at the `limit`-th match and remaining batches are never decoded.
    ///
    /// `limit` counts predicate matches, not output rows — callers applying
    /// an offset must pass `offset + limit`.
    ///
    /// `total_rows` is the row count `array_reader` would yield if iterated
    /// to completion. It is used to pad un-evaluated trailing rows as "not
    /// selected" so the returned [`RowSelection`] covers the full row group.
    ///
    /// Only valid for the *last* predicate in a filter chain: intermediate
    /// predicates' match counts do not map 1:1 to output rows.
    pub fn with_limit(mut self, limit: usize, total_rows: usize) -> Self {
        self.limit = Some(limit);
        self.total_rows = total_rows;
        self
    }
}

/// A builder for [`ReadPlan`]
#[derive(Clone, Debug)]
pub struct ReadPlanBuilder {
    batch_size: usize,
    /// Which rows to select. Includes the result of all filters applied so far
    selection: Option<RowSelection>,
    /// Policy to use when materializing the row selection
    row_selection_policy: RowSelectionPolicy,
    /// Row ranges with page data loaded for the current projection.
    loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
}

impl ReadPlanBuilder {
    /// Create a `ReadPlanBuilder` with the given batch size
    pub fn new(batch_size: usize) -> Self {
        Self {
            batch_size,
            selection: None,
            row_selection_policy: RowSelectionPolicy::default(),
            loaded_row_ranges: None,
        }
    }

    /// Set the current selection to the given value
    pub fn with_selection(mut self, selection: Option<RowSelection>) -> Self {
        self.selection = selection;
        self
    }

    /// Configure the policy to use when materialising the [`RowSelection`]
    ///
    /// Defaults to [`RowSelectionPolicy::Auto`]
    pub fn with_row_selection_policy(mut self, policy: RowSelectionPolicy) -> Self {
        self.row_selection_policy = policy;
        self
    }

    pub(crate) fn with_loaded_row_ranges(mut self, ranges: Option<LoadedRowRanges>) -> Self {
        self.loaded_row_ranges = ranges.map(Arc::new);
        self
    }

    /// Returns the current row selection policy
    pub fn row_selection_policy(&self) -> &RowSelectionPolicy {
        &self.row_selection_policy
    }

    /// Returns the current selection, if any
    pub fn selection(&self) -> Option<&RowSelection> {
        self.selection.as_ref()
    }

    /// Specifies the number of rows in the row group, before filtering is applied.
    ///
    /// Returns a [`LimitedReadPlanBuilder`] that can apply
    /// offset and limit.
    ///
    /// Call [`LimitedReadPlanBuilder::build_limited`] to apply the limits to this
    /// selection.
    pub(crate) fn limited(self, row_count: usize) -> LimitedReadPlanBuilder {
        LimitedReadPlanBuilder::new(self, row_count)
    }

    /// Returns true if the current plan selects any rows
    pub fn selects_any(&self) -> bool {
        self.selection
            .as_ref()
            .map(|s| s.selects_any())
            .unwrap_or(true)
    }

    /// Returns the number of rows selected, or `None` if all rows are selected.
    pub fn num_rows_selected(&self) -> Option<usize> {
        self.selection.as_ref().map(|s| s.row_count())
    }

    /// Returns the [`RowSelectionStrategy`] for this plan.
    ///
    /// Guarantees to return either `Selectors` or `Mask`, never `Auto`.
    pub(crate) fn resolve_selection_strategy(&self) -> RowSelectionStrategy {
        match self.row_selection_policy {
            RowSelectionPolicy::Selectors => RowSelectionStrategy::Selectors,
            RowSelectionPolicy::Mask => RowSelectionStrategy::Mask,
            RowSelectionPolicy::Auto { threshold, .. } => {
                let Some(selection) = self.selection.as_ref() else {
                    return RowSelectionStrategy::Selectors;
                };

                selection.auto_selection_strategy(threshold)
            }
        }
    }

    /// Evaluates an [`ArrowPredicate`], updating this plan's `selection`
    ///
    /// If the current `selection` is `Some`, the resulting [`RowSelection`]
    /// will be the conjunction of the existing selection and the rows selected
    /// by `predicate`.
    ///
    /// Note: pre-existing selections may come from evaluating a previous predicate
    /// or if the [`ParquetRecordBatchReader`] specified an explicit
    /// [`RowSelection`] in addition to one or more predicates.
    pub fn with_predicate(
        self,
        array_reader: Box<dyn ArrayReader>,
        predicate: &mut dyn ArrowPredicate,
    ) -> Result<Self> {
        self.with_predicate_options(PredicateOptions::new(array_reader, predicate))
    }

    /// Evaluates an [`ArrowPredicate`] with the given [`PredicateOptions`],
    /// updating this plan's `selection`.
    ///
    /// Like [`Self::with_predicate`], but allows additional options such as a
    /// match-count limit for early termination (see
    /// [`PredicateOptions::with_limit`]).
    pub fn with_predicate_options(mut self, options: PredicateOptions<'_>) -> Result<Self> {
        let PredicateOptions {
            array_reader,
            predicate,
            limit,
            total_rows,
        } = options;

        // Target length for the concatenated filter output:
        // - Prior selection ⇒ the reader yields that many rows; `and_then`
        //   below requires the filter output to match.
        // - No prior selection ⇒ the reader yields `total_rows`. We only
        //   need to pad when `limit` may short-circuit the loop; otherwise
        //   iteration naturally exhausts.
        let expected_rows = match self.selection.as_ref() {
            Some(s) => Some(s.row_count()),
            None => limit.map(|_| total_rows),
        };

        let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build());
        let mut filters = vec![];
        let mut processed_rows: usize = 0;
        let mut matched_rows: usize = 0;
        for maybe_batch in reader {
            let maybe_batch = maybe_batch?;
            let input_rows = maybe_batch.num_rows();
            let filter = predicate.evaluate(maybe_batch)?;
            // Since user supplied predicate, check error here to catch bugs quickly
            if filter.len() != input_rows {
                return Err(arrow_err!(
                    "ArrowPredicate predicate returned {} rows, expected {input_rows}",
                    filter.len()
                ));
            }
            let filter = match filter.null_count() {
                0 => filter,
                // RowSelection::from_filters expects non-null filters. Convert
                // NULL predicate results to false so they are not selected.
                _ => prep_null_mask_filter(&filter),
            };

            processed_rows += input_rows;

            match limit {
                Some(limit) if limit - matched_rows <= filter.len() => {
                    let truncated = filter.take_n_true(limit - matched_rows);
                    matched_rows += truncated.true_count();
                    filters.push(truncated);
                    if matched_rows >= limit {
                        break;
                    }
                }
                _ => {
                    matched_rows += filter.true_count();
                    filters.push(filter);
                }
            }
        }

        // Pad the tail so the filters cover `expected_rows` total. This keeps
        // the invariant that the resulting `RowSelection` spans every row the
        // reader would have produced — rows past the early break are marked
        // "not selected". When no limit is set the loop always exhausts and
        // no padding is needed.
        if let Some(expected) = expected_rows
            && processed_rows < expected
        {
            let pad_len = expected - processed_rows;
            filters.push(BooleanArray::new(BooleanBuffer::new_unset(pad_len), None));
        }

        // If the predicate selected all rows, applying it is a no-op. With no
        // prior selection this keeps selection as None, enabling coalesced page
        // fetches; with a prior selection it avoids rebuilding the same
        // selection.
        let all_selected = filters.iter().all(|f| f.true_count() == f.len());
        if all_selected {
            return Ok(self);
        }
        let raw = if self
            .selection
            .as_ref()
            .is_some_and(|s| s.as_mask().is_some())
        {
            RowSelection::from_boolean_buffer(filters_to_boolean_buffer(&filters))
        } else {
            RowSelection::from_filters(&filters)
        };
        self.selection = match self.selection.take() {
            Some(selection) => Some(selection.and_then(&raw)),
            None => Some(raw),
        };
        Ok(self)
    }

    /// Create a final `ReadPlan` the read plan for the scan
    pub fn build(mut self) -> ReadPlan {
        // If selection is empty, truncate
        if !self.selects_any() {
            self.selection = Some(RowSelection::from(vec![]));
        }

        // Preferred strategy must not be Auto
        let selection_strategy = self.resolve_selection_strategy();

        let Self {
            batch_size,
            selection,
            row_selection_policy: _,
            loaded_row_ranges,
        } = self;

        let row_selection_cursor = selection
            .map(|s| build_cursor(s.trim(), selection_strategy, loaded_row_ranges))
            .unwrap_or_else(RowSelectionCursor::new_all);

        ReadPlan {
            batch_size,
            row_selection_cursor,
        }
    }
}

/// Lower a [`RowSelection`] to the cursor form requested by the resolved strategy.
fn build_cursor(
    selection: RowSelection,
    strategy: RowSelectionStrategy,
    loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
) -> RowSelectionCursor {
    match (strategy, selection.into_inner()) {
        (RowSelectionStrategy::Mask, RowSelectionInner::Mask(mask)) => {
            RowSelectionCursor::new_mask_from_buffer((*mask).into_mask(), loaded_row_ranges)
        }
        (RowSelectionStrategy::Mask, RowSelectionInner::Selectors(selectors)) => {
            RowSelectionCursor::new_mask_from_selectors(selectors, loaded_row_ranges)
        }
        (RowSelectionStrategy::Selectors, RowSelectionInner::Selectors(selectors)) => {
            RowSelectionCursor::new_selectors(selectors)
        }
        (RowSelectionStrategy::Selectors, RowSelectionInner::Mask(mask)) => {
            RowSelectionCursor::new_selectors((*mask).into_selectors())
        }
    }
}

/// Builder for [`ReadPlan`] that applies a limit and offset to the read plan
///
/// See [`ReadPlanBuilder::limited`] to create this builder.
pub(crate) struct LimitedReadPlanBuilder {
    /// The underlying builder
    inner: ReadPlanBuilder,
    /// Total number of rows in the row group before the selection, limit or
    /// offset are applied
    row_count: usize,
    /// The offset to apply, if any
    offset: Option<usize>,
    /// The limit to apply, if any
    limit: Option<usize>,
}

impl LimitedReadPlanBuilder {
    /// Create a new `LimitedReadPlanBuilder` from the existing builder and number of rows
    fn new(inner: ReadPlanBuilder, row_count: usize) -> Self {
        Self {
            inner,
            row_count,
            offset: None,
            limit: None,
        }
    }

    /// Set the offset to apply to the read plan
    pub(crate) fn with_offset(mut self, offset: Option<usize>) -> Self {
        self.offset = offset;
        self
    }

    /// Set the limit to apply to the read plan
    pub(crate) fn with_limit(mut self, limit: Option<usize>) -> Self {
        self.limit = limit;
        self
    }

    /// Apply offset and limit, updating the selection on the underlying builder
    /// and returning it.
    pub(crate) fn build_limited(self) -> ReadPlanBuilder {
        let Self {
            mut inner,
            row_count,
            offset,
            limit,
        } = self;

        // If the selection is empty, truncate
        if !inner.selects_any() {
            inner.selection = Some(RowSelection::from(vec![]));
        }

        // If an offset is defined, apply it to the `selection`
        if let Some(offset) = offset {
            inner.selection = Some(match row_count.checked_sub(offset) {
                None => RowSelection::from(vec![]),
                Some(remaining) => inner
                    .selection
                    .map(|selection| selection.offset(offset))
                    .unwrap_or_else(|| {
                        RowSelection::from(vec![
                            RowSelector::skip(offset),
                            RowSelector::select(remaining),
                        ])
                    }),
            });
        }

        // If a limit is defined, apply it to the final `selection`
        if let Some(limit) = limit {
            inner.selection = Some(
                inner
                    .selection
                    .map(|selection| selection.limit(limit))
                    .unwrap_or_else(|| {
                        RowSelection::from(vec![RowSelector::select(limit.min(row_count))])
                    }),
            );
        }

        inner
    }
}

fn filters_to_boolean_buffer(filters: &[BooleanArray]) -> BooleanBuffer {
    let total_rows = filters.iter().map(|f| f.len()).sum();
    let mut builder = BooleanBufferBuilder::new(total_rows);
    for filter in filters {
        assert_eq!(filter.null_count(), 0);
        builder.append_buffer(filter.values());
    }
    builder.finish()
}

/// A plan reading specific rows from a Parquet Row Group.
///
/// See [`ReadPlanBuilder`] to create `ReadPlan`s
#[derive(Debug)]
pub struct ReadPlan {
    /// The number of rows to read in each batch
    batch_size: usize,
    /// Row ranges to be selected from the data source
    row_selection_cursor: RowSelectionCursor,
}

impl ReadPlan {
    /// Returns a mutable reference to the row selection cursor
    pub fn row_selection_cursor_mut(&mut self) -> &mut RowSelectionCursor {
        &mut self.row_selection_cursor
    }

    /// Return the number of rows to read in each output batch
    #[inline(always)]
    pub fn batch_size(&self) -> usize {
        self.batch_size
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn builder_with_selection(selection: RowSelection) -> ReadPlanBuilder {
        ReadPlanBuilder::new(1024).with_selection(Some(selection))
    }

    #[test]
    fn preferred_selection_strategy_prefers_mask_by_default() {
        let selection = RowSelection::from(vec![RowSelector::select(8)]);
        let builder = builder_with_selection(selection);
        assert_eq!(
            builder.resolve_selection_strategy(),
            RowSelectionStrategy::Mask
        );
    }

    #[test]
    fn preferred_selection_strategy_prefers_selectors_when_threshold_small() {
        let selection = RowSelection::from(vec![RowSelector::select(8)]);
        let builder = builder_with_selection(selection)
            .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 1 });
        assert_eq!(
            builder.resolve_selection_strategy(),
            RowSelectionStrategy::Selectors
        );
    }

    #[test]
    fn preferred_selection_strategy_handles_dense_mask_backing() {
        let bits: Vec<_> = (0..16).map(|i| i % 2 == 0).collect();
        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits));
        let builder = builder_with_selection(selection)
            .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 4 });
        assert_eq!(
            builder.resolve_selection_strategy(),
            RowSelectionStrategy::Mask
        );
    }

    #[test]
    fn preferred_selection_strategy_handles_sparse_mask_backing() {
        let bits: Vec<_> = (0..128).map(|i| i < 64).collect();
        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits));
        let builder = builder_with_selection(selection)
            .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 32 });
        assert_eq!(
            builder.resolve_selection_strategy(),
            RowSelectionStrategy::Selectors
        );
    }

    #[test]
    fn preferred_selection_strategy_preserves_mask_threshold_boundaries() {
        let mask = BooleanBuffer::from(vec![true; 8]);

        let empty = builder_with_selection(RowSelection::from_boolean_buffer(
            BooleanBuffer::new_unset(0),
        ));
        assert_eq!(
            empty.resolve_selection_strategy(),
            RowSelectionStrategy::Mask
        );

        let disabled = builder_with_selection(RowSelection::from_boolean_buffer(mask.clone()))
            .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 0 });
        assert_eq!(
            disabled.resolve_selection_strategy(),
            RowSelectionStrategy::Selectors
        );

        // Equality does not satisfy the policy's strict inequality: 8 < 1 * 8.
        let equal = builder_with_selection(RowSelection::from_boolean_buffer(mask.clone()))
            .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 8 });
        assert_eq!(
            equal.resolve_selection_strategy(),
            RowSelectionStrategy::Selectors
        );

        let above = builder_with_selection(RowSelection::from_boolean_buffer(mask))
            .with_row_selection_policy(RowSelectionPolicy::Auto { threshold: 9 });
        assert_eq!(
            above.resolve_selection_strategy(),
            RowSelectionStrategy::Mask
        );
    }

    #[test]
    fn preferred_selection_strategy_mask_matches_selector_backing() {
        use rand::{RngExt, rng};

        let mut rand = rng();
        for _ in 0..200 {
            let len = rand.random_range(0..256);
            let bits: Vec<_> = (0..len).map(|_| rand.random_bool(0.5)).collect();
            let mask_backed = RowSelection::from_boolean_buffer(BooleanBuffer::from(bits.clone()));
            let selector_backed = RowSelection::from_filters(&[BooleanArray::from(bits)]);

            for threshold in [0, 1, 2, 8, 32, 64, usize::MAX] {
                assert_eq!(
                    mask_backed.auto_selection_strategy(threshold),
                    selector_backed.auto_selection_strategy(threshold),
                    "strategy differs for len {len} and threshold {threshold}"
                );
            }
        }
    }

    #[test]
    fn mask_plan_trims_trailing_skips_before_chunking() {
        let mut plan = ReadPlanBuilder::new(8)
            .with_selection(Some(RowSelection::from(vec![
                RowSelector::select(1),
                RowSelector::skip(7),
            ])))
            .with_row_selection_policy(RowSelectionPolicy::Mask)
            .build();
        let RowSelectionCursor::Mask(cursor) = plan.row_selection_cursor_mut() else {
            panic!("expected a Mask cursor");
        };

        let chunk = cursor.next_chunk(8).unwrap();
        assert_eq!(chunk.chunk_rows, 1);
        assert_eq!(chunk.selected_rows, 1);
        assert!(cursor.is_empty());
    }

    #[test]
    fn selectors_policy_lowers_mask_backed_selection() {
        let selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
            true, false, false, true, true,
        ]));
        let mut plan = ReadPlanBuilder::new(4)
            .with_selection(Some(selection))
            .with_row_selection_policy(RowSelectionPolicy::Selectors)
            .build();
        let RowSelectionCursor::Selectors(cursor) = plan.row_selection_cursor_mut() else {
            panic!("expected a Selectors cursor");
        };

        assert_eq!(cursor.next_selector(), RowSelector::select(1));
        assert_eq!(cursor.next_selector(), RowSelector::skip(2));
        assert_eq!(cursor.next_selector(), RowSelector::select(2));
        assert!(cursor.is_empty());
    }

    #[test]
    fn mask_backed_plan_respects_loaded_row_ranges() {
        // Start from a non-byte-aligned BooleanBuffer::slice so the mask-backed
        // cursor path (new_mask_from_buffer) is exercised with a bit offset.
        // Rows 0 and 11 of 12 are selected; pages for rows [4, 10) are unloaded.
        let mut bits = vec![false; 5];
        bits.push(true); // row 0
        bits.extend(std::iter::repeat_n(false, 10)); // rows 1..=10
        bits.push(true); // row 11
        bits.extend([true, false, true]); // trailing padding outside the slice
        let mask = BooleanBuffer::from(bits).slice(5, 12);
        let selection = RowSelection::from_boolean_buffer(mask);
        assert!(selection.as_mask().is_some());

        let loaded = LoadedRowRanges::from_selection(RowSelection::from(vec![
            RowSelector::select(4),
            RowSelector::skip(6),
            RowSelector::select(2),
        ]));

        let mut plan = ReadPlanBuilder::new(12)
            .with_selection(Some(selection))
            .with_row_selection_policy(RowSelectionPolicy::Mask)
            .with_loaded_row_ranges(Some(loaded))
            .build();
        let RowSelectionCursor::Mask(cursor) = plan.row_selection_cursor_mut() else {
            panic!("expected a Mask cursor");
        };

        // The first chunk stops at its final selected row instead of carrying
        // trailing skipped rows to the loaded range boundary.
        let first = cursor.next_chunk(12).unwrap();
        assert_eq!(first.initial_skip, 0);
        assert_eq!(first.chunk_rows, 1);
        assert_eq!(first.selected_rows, 1);

        // The second chunk skips directly to the next selected row.
        let second = cursor.next_chunk(12).unwrap();
        assert_eq!(second.initial_skip, 10);
        assert_eq!(second.chunk_rows, 1);
        assert_eq!(second.selected_rows, 1);
        assert!(cursor.is_empty());
    }

    #[test]
    fn with_predicate_options_limit_pads_tail_when_no_prior_selection() {
        use crate::arrow::ProjectionMask;
        use crate::arrow::array_reader::StructArrayReader;
        use crate::arrow::array_reader::test_util::make_int32_page_reader;
        use crate::arrow::arrow_reader::ArrowPredicateFn;
        use arrow_schema::{DataType as ArrowType, Field, Fields};

        // 100 rows, all match the predicate. Limit stops the loop after 10
        // matches — but the resulting RowSelection must still describe the
        // full 100-row row group (90 trailing rows as "not selected"), not
        // only the 10 rows we happened to evaluate before breaking.
        const TOTAL_ROWS: usize = 100;
        const LIMIT: usize = 10;

        let data: Vec<i32> = (0..TOTAL_ROWS as i32).collect();
        let levels = vec![0; TOTAL_ROWS];
        let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
        let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
            "c0",
            ArrowType::Int32,
            false,
        )]));
        let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None);

        let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| {
            Ok(BooleanArray::from(vec![true; batch.num_rows()]))
        });

        let builder = ReadPlanBuilder::new(16)
            .with_predicate_options(
                PredicateOptions::new(Box::new(struct_reader), &mut predicate)
                    .with_limit(LIMIT, TOTAL_ROWS),
            )
            .unwrap();

        let selection = builder
            .selection()
            .expect("limit-driven early break must produce a selection");

        // `row_count` counts selected rows — must equal the limit.
        assert_eq!(selection.row_count(), LIMIT);

        // Total rows covered (selects + skips) must equal the full row group
        // so downstream offset/limit math stays in absolute-row space.
        assert_eq!(
            selection.total_row_count(),
            TOTAL_ROWS,
            "selection must span the full row group, not only the prefix evaluated before the limit"
        );
    }

    #[test]
    fn with_predicate_options_preserves_mask_selection() {
        use crate::arrow::ProjectionMask;
        use crate::arrow::array_reader::StructArrayReader;
        use crate::arrow::array_reader::test_util::make_int32_page_reader;
        use crate::arrow::arrow_reader::ArrowPredicateFn;
        use arrow_schema::{DataType as ArrowType, Field, Fields};

        let data: Vec<i32> = (0..6).collect();
        let levels = vec![0; data.len()];
        let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
        let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
            "c0",
            ArrowType::Int32,
            false,
        )]));
        let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None);

        let prior = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
            true, false, true, true, false, true,
        ]));
        let mut filters = vec![BooleanArray::from(vec![true, false, true, false])];
        let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), move |batch| {
            assert_eq!(batch.num_rows(), 4);
            Ok(filters.remove(0))
        });

        let builder = ReadPlanBuilder::new(16)
            .with_selection(Some(prior))
            .with_predicate_options(PredicateOptions::new(
                Box::new(struct_reader),
                &mut predicate,
            ))
            .unwrap();

        let selection = builder.selection().unwrap();
        assert!(selection.as_mask().is_some());

        let expected = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
            true, false, false, true, false, false,
        ]));
        assert_eq!(selection, &expected);
    }

    #[test]
    fn with_predicate_options_limit_handles_null_filters() {
        use crate::arrow::ProjectionMask;
        use crate::arrow::array_reader::StructArrayReader;
        use crate::arrow::array_reader::test_util::make_int32_page_reader;
        use crate::arrow::arrow_reader::ArrowPredicateFn;
        use arrow_schema::{DataType as ArrowType, Field, Fields};

        const TOTAL_ROWS: usize = 100;
        const LIMIT: usize = 10;

        let data: Vec<i32> = (0..TOTAL_ROWS as i32).collect();
        let levels = vec![0; TOTAL_ROWS];
        let leaf = make_int32_page_reader(&data, &levels, &levels, 0, 0, None);
        let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
            "c0",
            ArrowType::Int32,
            false,
        )]));
        let struct_reader = StructArrayReader::new(struct_type, vec![leaf], 0, 0, false, None);

        let mut predicate = ArrowPredicateFn::new(ProjectionMask::all(), |batch| {
            Ok((0..batch.num_rows())
                .map(|i| match i % 4 {
                    0 | 2 => Some(true),
                    1 => None,
                    _ => Some(false),
                })
                .collect::<BooleanArray>())
        });

        let builder = ReadPlanBuilder::new(16)
            .with_predicate_options(
                PredicateOptions::new(Box::new(struct_reader), &mut predicate)
                    .with_limit(LIMIT, TOTAL_ROWS),
            )
            .unwrap();

        let selection = builder
            .selection()
            .expect("limit-driven early break must produce a selection");

        assert_eq!(selection.row_count(), LIMIT);

        assert_eq!(selection.total_row_count(), TOTAL_ROWS);
    }
}