Push LIMIT / OFFSET into the last RowFilter predicate and skip unused row groups

parquet/src/arrow/array_reader/mod.rs

@@ -49,7 +49,7 @@ mod row_number;
 mod struct_array;
 
 #[cfg(test)]
-mod test_util;
+pub(crate) mod test_util;
 
 // Note that this crate is public under the `experimental` feature flag.
 use crate::file::metadata::RowGroupMetaData;

parquet/src/arrow/arrow_reader/mod.rs

@@ -49,7 +49,7 @@ use crate::schema::types::SchemaDescriptor;
 
 use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics;
 // Exposed so integration tests and benchmarks can temporarily override the threshold.
-pub use read_plan::{ReadPlan, ReadPlanBuilder};
+pub use read_plan::{PredicateOptions, ReadPlan, ReadPlanBuilder};
 
 mod filter;
 pub mod metrics;

parquet/src/arrow/arrow_reader/read_plan.rs

@@ -25,10 +25,57 @@ use crate::arrow::arrow_reader::{
     ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelectionCursor, RowSelector,
 };
 use crate::errors::{ParquetError, Result};
-use arrow_array::Array;
+use arrow_array::{Array, BooleanArray};
+use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
 use arrow_select::filter::prep_null_mask_filter;
 use std::collections::VecDeque;
 
+/// 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 {
@@ -144,12 +191,42 @@ impl ReadPlanBuilder {
     /// or if the [`ParquetRecordBatchReader`] specified an explicit
     /// [`RowSelection`] in addition to one or more predicates.
     pub fn with_predicate(
-        mut self,
+        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();
@@ -161,10 +238,37 @@ impl ReadPlanBuilder {
                     filter.len()
                 ));
             }
-            match filter.null_count() {
-                0 => filters.push(filter),
-                _ => filters.push(prep_null_mask_filter(&filter)),
+            let filter = match filter.null_count() {
+                0 => filter,
+                _ => prep_null_mask_filter(&filter),
             };
+
+            processed_rows += input_rows;
+
+            match limit {
+                Some(limit) if matched_rows + filter.true_count() >= limit => {
+                    let needed = limit - matched_rows;
+                    let truncated = truncate_filter_after_n_trues(filter, needed);
+                    filters.push(truncated);
+                    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 {
+            if 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 and there is no prior selection,
@@ -305,6 +409,35 @@ impl LimitedReadPlanBuilder {
     }
 }
 
+/// Produce a new `BooleanArray` of the same length as `filter` in which only
+/// the first `n` `true` positions from `filter` remain `true`; any `true`
+/// positions beyond the first `n` are replaced with `false`.
+///
+/// `filter` must not contain nulls (callers apply [`prep_null_mask_filter`]
+/// first). If `filter` has at most `n` `true` values, a clone is returned.
+fn truncate_filter_after_n_trues(filter: BooleanArray, n: usize) -> BooleanArray {
+    if filter.true_count() <= n {
+        return filter;
+    }
+    let len = filter.len();
+    if n == 0 {
+        return BooleanArray::new(BooleanBuffer::new_unset(len), None);
+    }
+    // `set_indices` scans 64 bits at a time via `trailing_zeros`, so locating
+    // the `n`-th set bit is cheaper than visiting every bit. Everything up to
+    // and including that position is copied verbatim; the rest is zeroed.
+    let values = filter.values();
+    let last_kept = values
+        .set_indices()
+        .nth(n - 1)
+        .expect("n - 1 < true_count, checked above");
+
+    let mut builder = BooleanBufferBuilder::new(len);
+    builder.append_buffer(&values.slice(0, last_kept + 1));
+    builder.append_n(len - last_kept - 1, false);
+    BooleanArray::new(builder.finish(), None)
+}
+
 /// A plan reading specific rows from a Parquet Row Group.
 ///
 /// See [`ReadPlanBuilder`] to create `ReadPlan`s
@@ -367,4 +500,89 @@ mod tests {
             RowSelectionStrategy::Selectors
         );
     }
+
+    #[test]
+    fn truncate_filter_after_n_trues_keeps_first_n_matches() {
+        let f = BooleanArray::from(vec![true, false, true, true, false, true, true]);
+        // true positions: 0, 2, 3, 5, 6
+        let t = truncate_filter_after_n_trues(f.clone(), 3);
+        assert_eq!(t.len(), f.len());
+        assert_eq!(t.true_count(), 3);
+        let out: Vec<bool> = (0..t.len()).map(|i| t.value(i)).collect();
+        assert_eq!(
+            out,
+            vec![true, false, true, true, false, false, false],
+            "first three trues should survive, the rest become false"
+        );
+    }
+
+    #[test]
+    fn truncate_filter_after_n_trues_passes_through_when_already_small_enough() {
+        let f = BooleanArray::from(vec![true, false, true, false]);
+        let t = truncate_filter_after_n_trues(f.clone(), 5);
+        assert_eq!(t.len(), f.len());
+        assert_eq!(t.true_count(), 2);
+    }
+
+    #[test]
+    fn truncate_filter_after_n_trues_zero_returns_all_false() {
+        let f = BooleanArray::from(vec![true, true, true]);
+        let t = truncate_filter_after_n_trues(f, 0);
+        assert_eq!(t.len(), 3);
+        assert_eq!(t.true_count(), 0);
+    }
+
+    #[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::InMemoryArrayReader;
+        use crate::arrow::arrow_reader::ArrowPredicateFn;
+        use arrow_array::Int32Array;
+        use arrow_schema::{DataType as ArrowType, Field, Fields};
+        use std::sync::Arc;
+
+        // 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 array = Arc::new(Int32Array::from(data));
+        let leaf = InMemoryArrayReader::new(ArrowType::Int32, array.clone(), None, None);
+        let struct_type = ArrowType::Struct(Fields::from(vec![Field::new(
+            "c0",
+            ArrowType::Int32,
+            false,
+        )]));
+        let struct_reader = StructArrayReader::new(struct_type, vec![Box::new(leaf)], 0, 0, false);
+
+        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.
+        let total: usize = selection.iter().map(|s| s.row_count).sum();
+        assert_eq!(
+            total, TOTAL_ROWS,
+            "selection must span the full row group, not only the prefix evaluated before the limit"
+        );
+    }
 }