diff --git a/parquet/src/arrow/array_reader/mod.rs b/parquet/src/arrow/array_reader/mod.rs index 726eae1f51c5..989c93480993 100644 --- a/parquet/src/arrow/array_reader/mod.rs +++ b/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; diff --git a/parquet/src/arrow/arrow_reader/mod.rs b/parquet/src/arrow/arrow_reader/mod.rs index 8a7e602618c2..f9ef4eac651f 100644 --- a/parquet/src/arrow/arrow_reader/mod.rs +++ b/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; diff --git a/parquet/src/arrow/arrow_reader/read_plan.rs b/parquet/src/arrow/arrow_reader/read_plan.rs index 99ffe0febc95..ac2e105ecf4f 100644 --- a/parquet/src/arrow/arrow_reader/read_plan.rs +++ b/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, + predicate: &'a mut dyn ArrowPredicate, + limit: Option, + 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, 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, predicate: &mut dyn ArrowPredicate, ) -> Result { + 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 { + 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 = (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 = (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" + ); + } } diff --git a/parquet/src/arrow/push_decoder/mod.rs b/parquet/src/arrow/push_decoder/mod.rs index 24384471a4ec..4c667e534366 100644 --- a/parquet/src/arrow/push_decoder/mod.rs +++ b/parquet/src/arrow/push_decoder/mod.rs @@ -1093,6 +1093,290 @@ mod test { expect_finished(decoder.try_decode()); } + /// When filter pushdown is combined with a `LIMIT`, the predicate must + /// not be evaluated for rows beyond the `limit`-th match. + /// + /// Filter `a > 175` produces 24 matches in row group 0 (rows 176..199). + /// With `limit = 10`, only the first 10 matches (rows 176..185) should be + /// emitted, AND the predicate counter should observe that evaluation was + /// short-circuited. + #[test] + fn test_decoder_filter_with_limit_short_circuits_within_row_group() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let rows_filtered = Arc::new(AtomicUsize::new(0)); + let rows_filtered_for_predicate = Arc::clone(&rows_filtered); + + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + move |batch: RecordBatch| { + rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + // Use a small batch size so the row group is evaluated across + // multiple predicate batches; that is the regime where Layer 2's + // short-circuit saves predicate evaluation work. Matching rows are + // 176..199 (24 rows); with batch_size = 10 those span batches 17, 18, + // and 19 (rows 170..199). A limit of 10 should stop filter evaluation + // in the middle of batch 18. + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_batch_size(10) + .with_limit(10) + .build() + .unwrap(); + + // First row group: filter columns fetch (predicate is evaluated here) + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // The first 10 matching rows come out: 176..185, column "a" + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 10).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + // no data for row group 1 should be requested — the limit + // was satisfied by row group 0 and the `Start` state for row group 1 + // short-circuits to `Finished`. + expect_finished(decoder.try_decode()); + + // Row 186 is the 11th match; the scan should stop no later than the + // batch containing it (batch 18 of 10 rows = rows 180..189), so at + // most 190 rows are evaluated. + let evaluated = rows_filtered.load(Ordering::Relaxed); + assert!( + evaluated <= 190, + "predicate evaluated {evaluated} rows; expected ≤ 190 (stop within batch containing 11th match)" + ); + } + + /// Once the limit has been satisfied by a prior row group, subsequent + /// row groups should be skipped entirely — no data request for their + /// filter columns. + #[test] + fn test_decoder_filter_with_limit_skips_later_row_groups() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + // `a > 175` matches rows 176..199 in row group 0 (24 matches) and + // 200..399 in row group 1 (200 matches). With limit = 5, all matches + // should come from row group 0. + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_limit(5) + .build() + .unwrap(); + + // Row group 0: fetch filter pages + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // First 5 matches: 176..180 + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 5).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + // Row group 1 must NOT request data — the limit is already satisfied + // so `Start` in row group 1 short-circuits to `Finished`. + expect_finished(decoder.try_decode()); + } + + /// The predicate short-circuit must account for `self.offset` as well as + /// `self.limit`. The post-predicate `with_offset` step skips that many + /// already-selected rows before `with_limit` counts output rows — so the + /// predicate must retain at least `offset + limit` matches. Without the + /// fix, Layer 2 caps at just `limit` and the later `with_offset` consumes + /// all of them, producing 0 rows instead of `limit`. + /// + /// `a > 175` matches rows 176..199 in row group 0 (24 matches). With + /// `offset = 10, limit = 5`, the expected output is rows 186..190 (the + /// 11th through 15th matches). + #[test] + fn test_decoder_filter_with_offset_and_limit() { + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + |batch: RecordBatch| { + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_offset(10) + .with_limit(5) + .build() + .unwrap(); + + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(186, 5).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + expect_finished(decoder.try_decode()); + } + + /// The limit short-circuit must also be correct when the limited predicate + /// is the last predicate in a multi-predicate chain. + /// + /// `a > 175` first narrows row group 0 to rows 176..199. The final + /// predicate `b < 625` is then evaluated only over those 24 rows, all of + /// which match. With `limit = 10`, the final output should still be rows + /// 176..185, and the second predicate should stop before consuming all 24 + /// selected rows. + #[test] + fn test_decoder_multi_filters_with_limit() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let first_predicate_rows = Arc::new(AtomicUsize::new(0)); + let second_predicate_rows = Arc::new(AtomicUsize::new(0)); + + let first_predicate_rows_for_filter = Arc::clone(&first_predicate_rows); + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + move |batch: RecordBatch| { + first_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let second_predicate_rows_for_filter = Arc::clone(&second_predicate_rows); + let row_filter_b = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["b"]), + move |batch: RecordBatch| { + second_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_625 = Int64Array::new_scalar(625); + let column = batch.column(0).as_primitive::(); + lt(column, &scalar_625) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["c"])) + .with_row_filter(RowFilter::new(vec![ + Box::new(row_filter_a), + Box::new(row_filter_b), + ])) + .with_batch_size(10) + .with_limit(10) + .build() + .unwrap(); + + // Row group 0, first predicate + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // Row group 0, second predicate + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + // Final projected data + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 10).project(&[2]).unwrap(); + assert_eq!(batch, expected); + + // The overall limit was satisfied by row group 0. + expect_finished(decoder.try_decode()); + + assert_eq!(first_predicate_rows.load(Ordering::Relaxed), 200); + assert!( + second_predicate_rows.load(Ordering::Relaxed) < 24, + "final predicate should short-circuit before consuming all 24 rows selected by the first predicate" + ); + } + + /// When a row selection already exists, limiting the predicate must still + /// preserve alignment with that prior selection. + /// + /// The explicit selection narrows row group 0 to rows 150..199. Applying + /// `a > 175` over that selection yields rows 176..199. With `limit = 10`, + /// the decoder should emit rows 176..185 and stop without evaluating the + /// remaining selected rows. + #[test] + fn test_decoder_filter_with_row_selection_and_limit() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let builder = + ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap(); + let schema_descr = builder.metadata().file_metadata().schema_descr_ptr(); + + let rows_filtered = Arc::new(AtomicUsize::new(0)); + let rows_filtered_for_predicate = Arc::clone(&rows_filtered); + + let row_filter_a = ArrowPredicateFn::new( + ProjectionMask::columns(&schema_descr, ["a"]), + move |batch: RecordBatch| { + rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed); + let scalar_175 = Int64Array::new_scalar(175); + let column = batch.column(0).as_primitive::(); + gt(column, &scalar_175) + }, + ); + + let mut decoder = builder + .with_projection(ProjectionMask::columns(&schema_descr, ["a"])) + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(150), + RowSelector::select(50), + ])) + .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)])) + .with_batch_size(10) + .with_limit(10) + .build() + .unwrap(); + + let ranges = expect_needs_data(decoder.try_decode()); + push_ranges_to_decoder(&mut decoder, ranges); + + let batch = expect_data(decoder.try_decode()); + let expected = TEST_BATCH.slice(176, 10).project(&[0]).unwrap(); + assert_eq!(batch, expected); + + expect_finished(decoder.try_decode()); + + assert!( + rows_filtered.load(Ordering::Relaxed) < 50, + "predicate should short-circuit before consuming all 50 rows from the explicit row selection" + ); + } + #[test] fn test_decoder_offset_limit() { let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()) diff --git a/parquet/src/arrow/push_decoder/reader_builder/filter.rs b/parquet/src/arrow/push_decoder/reader_builder/filter.rs index 380211cca66e..219ab06da9b3 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/filter.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/filter.rs @@ -126,6 +126,13 @@ impl FilterInfo { .as_ref() } + /// Returns `true` if the current predicate is the last one in the chain + /// (i.e. the next call to [`Self::advance`] will return + /// [`AdvanceResult::Done`]). + pub(super) fn is_last(&self) -> bool { + self.next_predicate.get() == self.filter.predicates.len() + } + /// Return a reference to the cache projection pub(super) fn cache_projection(&self) -> &ProjectionMask { &self.cache_info.cache_projection diff --git a/parquet/src/arrow/push_decoder/reader_builder/mod.rs b/parquet/src/arrow/push_decoder/reader_builder/mod.rs index a95be9d87d37..60e50d29524e 100644 --- a/parquet/src/arrow/push_decoder/reader_builder/mod.rs +++ b/parquet/src/arrow/push_decoder/reader_builder/mod.rs @@ -24,7 +24,8 @@ use crate::arrow::array_reader::{ArrayReaderBuilder, CacheOptions, RowGroupCache use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics; use crate::arrow::arrow_reader::selection::RowSelectionStrategy; use crate::arrow::arrow_reader::{ - ParquetRecordBatchReader, ReadPlanBuilder, RowFilter, RowSelection, RowSelectionPolicy, + ParquetRecordBatchReader, PredicateOptions, ReadPlanBuilder, RowFilter, RowSelection, + RowSelectionPolicy, }; use crate::arrow::in_memory_row_group::ColumnChunkData; use crate::arrow::push_decoder::reader_builder::data::DataRequestBuilder; @@ -309,6 +310,19 @@ impl RowGroupReaderBuilder { ) -> Result { let result = match current_state { RowGroupDecoderState::Start { row_group_info } => { + // Short-circuit once the overall output limit is exhausted. + // + // `self.limit` tracks how many more rows the reader is still + // allowed to emit and is decremented as each row group is + // planned in `StartData`, so `Some(0)` means earlier row + // groups have already produced the full requested output. + if matches!(self.limit, Some(0)) { + return Ok(NextState::result( + RowGroupDecoderState::Finished, + DecodeResult::Finished, + )); + } + let column_chunks = None; // no prior column chunks let Some(filter) = self.filter.take() else { @@ -458,10 +472,24 @@ impl RowGroupReaderBuilder { predicate.projection(), self.row_group_offset_index(row_group_idx), ); - // `with_predicate` actually evaluates the filter - plan_builder = - plan_builder.with_predicate(array_reader, filter_info.current_mut())?; + // When this is the final predicate in the chain and an output + // limit is set, tell the filter evaluation to stop once enough + // matching rows have been accumulated. + let predicate_limit = self + .limit + .filter(|_| filter_info.is_last()) + .map(|l| l.saturating_add(self.offset.unwrap_or(0))); + + // Evaluate the filter via `with_predicate_options`, opting into + // early termination when this is the final predicate and an + // output limit was set. + let mut predicate_options = + PredicateOptions::new(array_reader, filter_info.current_mut()); + if let Some(limit) = predicate_limit { + predicate_options = predicate_options.with_limit(limit, row_count); + } + plan_builder = plan_builder.with_predicate_options(predicate_options)?; let row_group_info = RowGroupInfo { row_group_idx,