Skip to main content

parquet/arrow/push_decoder/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`ParquetPushDecoder`]: decodes Parquet data with data provided by the
19//! caller (rather than from an underlying reader).
20
21mod reader_builder;
22mod remaining;
23mod scan_plan;
24
25use crate::DecodeResult;
26pub use crate::arrow::arrow_reader::RowGroupSelection;
27use crate::arrow::arrow_reader::{
28    ArrowReaderBuilder, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader,
29};
30use crate::errors::ParquetError;
31use crate::file::metadata::ParquetMetaData;
32pub use crate::util::push_buffers::PushBuffers;
33use arrow_array::RecordBatch;
34use bytes::Bytes;
35use reader_builder::{RowGroupReaderBuilder, RowGroupReaderBuilderParts};
36use remaining::{RemainingRowGroups, RemainingRowGroupsParts};
37use scan_plan::RowBudget;
38pub use scan_plan::{PageKind, PlannedRange, ScanPlan};
39use std::ops::Range;
40use std::sync::Arc;
41
42/// A builder for [`ParquetPushDecoder`].
43///
44/// To create a new decoder, use [`ParquetPushDecoderBuilder::try_new_decoder`].
45///
46/// You can decode the metadata from a Parquet file using either
47/// [`ParquetMetadataReader`] or [`ParquetMetaDataPushDecoder`].
48///
49/// [`ParquetMetadataReader`]: crate::file::metadata::ParquetMetaDataReader
50/// [`ParquetMetaDataPushDecoder`]: crate::file::metadata::ParquetMetaDataPushDecoder
51///
52/// Note the "input" type is `u64` which represents the length of the Parquet file
53/// being decoded. This is needed to initialize the internal buffers that track
54/// what data has been provided to the decoder.
55///
56/// # Example
57/// ```
58/// # use std::ops::Range;
59/// # use std::sync::Arc;
60/// # use bytes::Bytes;
61/// # use arrow_array::record_batch;
62/// # use parquet::DecodeResult;
63/// # use parquet::arrow::push_decoder::ParquetPushDecoderBuilder;
64/// # use parquet::arrow::ArrowWriter;
65/// # use parquet::file::metadata::ParquetMetaDataPushDecoder;
66/// # let file_bytes = {
67/// #   let mut buffer = vec![];
68/// #   let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
69/// #   let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
70/// #   writer.write(&batch).unwrap();
71/// #   writer.close().unwrap();
72/// #   Bytes::from(buffer)
73/// # };
74/// # // mimic IO by returning a function that returns the bytes for a given range
75/// # let get_range = |range: &Range<u64>| -> Bytes {
76/// #    let start = range.start as usize;
77/// #     let end = range.end as usize;
78/// #    file_bytes.slice(start..end)
79/// # };
80/// # let file_length = file_bytes.len() as u64;
81/// # let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_length).unwrap();
82/// # metadata_decoder.push_ranges(vec![0..file_length], vec![file_bytes.clone()]).unwrap();
83/// # let DecodeResult::Data(parquet_metadata) = metadata_decoder.try_decode().unwrap() else { panic!("failed to decode metadata") };
84/// # let parquet_metadata = Arc::new(parquet_metadata);
85/// // The file length and metadata are required to create the decoder
86/// let mut decoder =
87///     ParquetPushDecoderBuilder::try_new_decoder(parquet_metadata)
88///       .unwrap()
89///       // Optionally configure the decoder, e.g. batch size
90///       .with_batch_size(1024)
91///       // Build the decoder
92///       .build()
93///       .unwrap();
94///
95///     // In a loop, ask the decoder what it needs next, and provide it with the required data
96///     loop {
97///         match decoder.try_decode().unwrap() {
98///             DecodeResult::NeedsData(ranges) => {
99///                 // The decoder needs more data. Fetch the data for the given ranges
100///                 let data = ranges.iter().map(|r| get_range(r)).collect::<Vec<_>>();
101///                 // Push the data to the decoder
102///                 decoder.push_ranges(ranges, data).unwrap();
103///                 // After pushing the data, we can try to decode again on the next iteration
104///             }
105///             DecodeResult::Data(batch) => {
106///                 // Successfully decoded a batch of data
107///                 assert!(batch.num_rows() > 0);
108///             }
109///             DecodeResult::Finished => {
110///                 // The decoder has finished decoding exit the loop
111///                 break;
112///             }
113///         }
114///     }
115/// ```
116///
117/// # Adaptive scans
118///
119/// The scan strategy is not fixed once [`build`](Self::build) is called: it
120/// can be changed *while decoding*, at row-group boundaries.
121///
122/// The important API for this is [`ParquetPushDecoder::try_next_reader`].
123/// Unlike [`try_decode`](ParquetPushDecoder::try_decode), which barrels
124/// straight through row-group boundaries, `try_next_reader` returns once per
125/// row group — leaving a clean window *between* row groups. At any such
126/// boundary, [`ParquetPushDecoder::into_builder`] hands back a
127/// `ParquetPushDecoderBuilder` for the row groups not yet decoded. Change any
128/// option on it (projection, row filter, row selection policy, …) and
129/// [`build`](Self::build) a fresh decoder that resumes from the next row
130/// group. This is how a query engine promotes or demotes filters — for
131/// example turning a row filter on or off — based on the selectivity observed
132/// in the row groups decoded so far.
133///
134/// ```
135/// # use std::ops::Range;
136/// # use std::sync::Arc;
137/// # use bytes::Bytes;
138/// # use arrow_array::record_batch;
139/// # use parquet::DecodeResult;
140/// # use parquet::arrow::ProjectionMask;
141/// # use parquet::arrow::push_decoder::ParquetPushDecoderBuilder;
142/// # use parquet::arrow::ArrowWriter;
143/// # use parquet::file::metadata::ParquetMetaDataPushDecoder;
144/// # use parquet::file::properties::WriterProperties;
145/// # let file_bytes = {
146/// #   let batch = record_batch!(
147/// #       ("a", Int32, [1, 2, 3, 4, 5, 6]),
148/// #       ("b", Int32, [6, 5, 4, 3, 2, 1])
149/// #   ).unwrap();
150/// #   // Small row groups so the test file has two of them.
151/// #   let props = WriterProperties::builder().set_max_row_group_row_count(Some(3)).build();
152/// #   let mut buffer = vec![];
153/// #   let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap();
154/// #   writer.write(&batch).unwrap();
155/// #   writer.close().unwrap();
156/// #   Bytes::from(buffer)
157/// # };
158/// # let get_range = |r: &Range<u64>| file_bytes.slice(r.start as usize..r.end as usize);
159/// # let file_length = file_bytes.len() as u64;
160/// # let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_length).unwrap();
161/// # metadata_decoder.push_ranges(vec![0..file_length], vec![file_bytes.clone()]).unwrap();
162/// # let DecodeResult::Data(parquet_metadata) = metadata_decoder.try_decode().unwrap() else { panic!() };
163/// # let parquet_metadata = Arc::new(parquet_metadata);
164/// let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(parquet_metadata)
165///     .unwrap()
166///     .build()
167///     .unwrap();
168///
169/// // Drive the decoder one row group at a time with `try_next_reader`.
170/// loop {
171///     match decoder.try_next_reader().unwrap() {
172///         DecodeResult::NeedsData(ranges) => {
173///             // Fetch and hand over the bytes the decoder asked for.
174///             let data = ranges.iter().map(|r| get_range(r)).collect();
175///             decoder.push_ranges(ranges, data).unwrap();
176///         }
177///         DecodeResult::Data(reader) => {
178///             // Decode this row group's batches.
179///             for batch in reader {
180///                 assert!(batch.unwrap().num_rows() > 0);
181///             }
182///             // We are now at a row-group boundary. Based on whatever stats
183///             // were gathered, optionally change strategy for the row groups
184///             // still to come: drop or promote a row filter, narrow or widen
185///             // the projection, etc.
186///             if decoder.is_at_row_group_boundary() && decoder.row_groups_remaining() > 0 {
187///                 let builder = decoder.into_builder().unwrap();
188///                 // e.g. column "b" turned out not to be needed.
189///                 let projection = ProjectionMask::columns(builder.parquet_schema(), ["a"]);
190///                 decoder = builder.with_projection(projection).build().unwrap();
191///             }
192///         }
193///         DecodeResult::Finished => break,
194///     }
195/// }
196/// ```
197pub type ParquetPushDecoderBuilder = ArrowReaderBuilder<PushDecoderInput>;
198
199/// The `input` of a [`ParquetPushDecoderBuilder`].
200///
201/// The shared [`ArrowReaderBuilder`] is generic over an `input`. The sync and
202/// async builders read from a file or async reader; the push decoder has no
203/// reader, so its input is the [`PushBuffers`] that caller-pushed bytes
204/// accumulate in (empty for a fresh builder).
205#[derive(Debug, Default)]
206pub struct PushDecoderInput {
207    /// Bytes pushed into the decoder, awaiting decode.
208    buffers: PushBuffers,
209}
210
211/// Methods for building a ParquetDecoder. See the base [`ArrowReaderBuilder`] for
212/// more options that can be configured.
213impl ParquetPushDecoderBuilder {
214    /// Create a new `ParquetDecoderBuilder` for configuring a Parquet decoder for the given file.
215    ///
216    /// See [`ParquetMetadataDecoder`] for a builder that can read the metadata from a Parquet file.
217    ///
218    /// [`ParquetMetadataDecoder`]: crate::file::metadata::ParquetMetaDataPushDecoder
219    ///
220    /// See example on [`ParquetPushDecoderBuilder`]
221    pub fn try_new_decoder(parquet_metadata: Arc<ParquetMetaData>) -> Result<Self, ParquetError> {
222        Self::try_new_decoder_with_options(parquet_metadata, ArrowReaderOptions::default())
223    }
224
225    /// Create a new `ParquetDecoderBuilder` for configuring a Parquet decoder for the given file
226    /// with the given reader options.
227    ///
228    /// This is similar to [`Self::try_new_decoder`] but allows configuring
229    /// options such as Arrow schema
230    pub fn try_new_decoder_with_options(
231        parquet_metadata: Arc<ParquetMetaData>,
232        arrow_reader_options: ArrowReaderOptions,
233    ) -> Result<Self, ParquetError> {
234        let arrow_reader_metadata =
235            ArrowReaderMetadata::try_new(parquet_metadata, arrow_reader_options)?;
236        Ok(Self::new_with_metadata(arrow_reader_metadata))
237    }
238
239    /// Create a new `ParquetDecoderBuilder` given [`ArrowReaderMetadata`].
240    ///
241    /// See [`ArrowReaderMetadata::try_new`] for how to create the metadata from
242    /// the Parquet metadata and reader options.
243    pub fn new_with_metadata(arrow_reader_metadata: ArrowReaderMetadata) -> Self {
244        Self::new_builder(PushDecoderInput::default(), arrow_reader_metadata)
245    }
246
247    /// Provide a preexisting [`PushBuffers`] for the built decoder to read
248    /// from, so bytes already fetched are not requested again.
249    ///
250    /// [`Self::build`] releases the bytes that the decoder does not read:
251    /// the bytes outside the read column chunks of the row groups that the
252    /// decoder reads.
253    pub fn with_buffers(self, buffers: PushBuffers) -> Self {
254        Self {
255            input: PushDecoderInput { buffers },
256            ..self
257        }
258    }
259
260    /// Select row groups and rows using row-group-local coordinates.
261    ///
262    /// Entries are decoded in the supplied order, and omitted row groups are
263    /// skipped. A row group listed more than once is decoded once per entry.
264    /// A `None` selection reads the entire row group. A selection shorter
265    /// than its row group skips the trailing rows, while a selection longer
266    /// than its row group returns an error from [`Self::build`].
267    ///
268    /// [`ArrowReaderBuilder::with_offset`] and
269    /// [`ArrowReaderBuilder::with_limit`] apply after the row-group-local
270    /// selections and any [`ArrowReaderBuilder::with_row_filter`], across the
271    /// plan as a whole and in the supplied order: the offset skips the first N
272    /// remaining rows and the limit caps the total rows emitted, regardless of
273    /// which row group they come from.
274    ///
275    /// This configuration is mutually exclusive with
276    /// [`ArrowReaderBuilder::with_row_groups`] and
277    /// [`ArrowReaderBuilder::with_row_selection`]. Combining them returns an
278    /// error from [`Self::build`]. Calling this method more than once replaces
279    /// the previous row-group-local configuration.
280    ///
281    /// This method is defined on the push decoder rather than
282    /// [`ArrowReaderBuilder`] because the synchronous reader does not support
283    /// row-group-local selections. The async stream builder exposes the same
284    /// method because it uses the push decoder internally.
285    ///
286    /// For example, if row groups 0 and 2 each contain 200 rows, the legacy
287    /// combination of `with_row_groups(vec![0, 2])` and a global selection for
288    /// rows 10..15 of row group 0 and all of row group 2 can be expressed as:
289    ///
290    /// ```no_run
291    /// # use parquet::arrow::arrow_reader::{RowSelection, RowSelector};
292    /// # use parquet::arrow::push_decoder::{ParquetPushDecoderBuilder, RowGroupSelection};
293    /// # fn configure(builder: ParquetPushDecoderBuilder) -> ParquetPushDecoderBuilder {
294    /// builder.with_row_group_selections(vec![
295    ///     RowGroupSelection::new(0, Some(RowSelection::from(vec![
296    ///         RowSelector::skip(10),
297    ///         RowSelector::select(5),
298    ///     ]))),
299    ///     RowGroupSelection::new(2, None),
300    /// ])
301    /// # }
302    /// ```
303    pub fn with_row_group_selections(
304        mut self,
305        row_group_selections: Vec<RowGroupSelection>,
306    ) -> Self {
307        self.row_group_plan
308            .set_row_group_selections(row_group_selections);
309        self
310    }
311
312    /// Create a [`ParquetPushDecoder`] with the configured options
313    pub fn build(self) -> Result<ParquetPushDecoder, ParquetError> {
314        let Self {
315            input: PushDecoderInput { buffers },
316            metadata: parquet_metadata,
317            schema,
318            fields,
319            batch_size,
320            row_group_plan,
321            projection,
322            filter,
323            limit,
324            offset,
325            metrics,
326            row_selection_policy,
327            max_predicate_cache_size,
328        } = self;
329
330        // Evaluate eligible same-projection predicates as one predicate.
331        let filter = filter.map(|filter| {
332            filter.fuse_same_projection(parquet_metadata.file_metadata().schema_descr())
333        });
334
335        let has_predicates = filter
336            .as_ref()
337            .is_some_and(|filter| !filter.predicates.is_empty());
338
339        // Prepare to build RowGroup readers. `buffers` carries any bytes the
340        // caller already pushed (preserved across `into_builder`); a fresh
341        // builder supplies an empty `PushBuffers`.
342        let row_group_reader_builder = RowGroupReaderBuilder::new(
343            batch_size,
344            projection,
345            Arc::clone(&parquet_metadata),
346            fields,
347            filter,
348            metrics,
349            max_predicate_cache_size,
350            buffers,
351            row_selection_policy,
352        );
353
354        // Initialize the decoder with the configured options
355        let mut remaining_row_groups = RemainingRowGroups::new(
356            schema,
357            parquet_metadata,
358            row_group_plan,
359            RowBudget::new(offset, limit),
360            has_predicates,
361            row_group_reader_builder,
362        )?;
363        // The buffers can hold bytes that this decoder does not read. For
364        // example, the buffers of a rebuilt decoder can hold the bytes of row
365        // groups that the new configuration skips.
366        remaining_row_groups.release_unread_bytes();
367
368        Ok(ParquetPushDecoder {
369            state: ParquetDecoderState::ReadingRowGroup {
370                remaining_row_groups: Box::new(remaining_row_groups),
371            },
372        })
373    }
374}
375
376/// Reassemble a [`ParquetPushDecoderBuilder`] from a decoder's not-yet-decoded
377/// state — the inverse of [`ParquetPushDecoderBuilder::build`]. The rebuilt
378/// builder carries the remaining row-group plan, offset/limit budget, and
379/// buffered bytes.
380fn builder_from_remaining(parts: RemainingRowGroupsParts) -> ParquetPushDecoderBuilder {
381    let RemainingRowGroupsParts {
382        metadata,
383        schema,
384        row_group_plan,
385        offset,
386        limit,
387        reader_builder,
388    } = parts;
389    let RowGroupReaderBuilderParts {
390        batch_size,
391        projection,
392        fields,
393        filter,
394        max_predicate_cache_size,
395        metrics,
396        row_selection_policy,
397        buffers,
398    } = reader_builder;
399
400    ArrowReaderBuilder {
401        input: PushDecoderInput::default(),
402        metadata,
403        schema,
404        fields,
405        batch_size,
406        row_group_plan,
407        projection,
408        filter,
409        row_selection_policy,
410        limit,
411        offset,
412        metrics,
413        max_predicate_cache_size,
414    }
415    // Carry the decoder's already-fetched bytes across the rebuild so the new
416    // decoder does not re-request them.
417    .with_buffers(buffers)
418}
419
420/// A push based Parquet Decoder
421///
422/// See [`ParquetPushDecoderBuilder`] for an example of how to build and use the decoder.
423///
424/// [`ParquetPushDecoder`] is a low level API for decoding Parquet data without an
425/// underlying reader for performing IO, and thus offers fine grained control
426/// over how data is fetched and decoded.
427///
428/// When more data is needed to make progress, instead of reading data directly
429/// from a reader, the decoder returns [`DecodeResult`] indicating what ranges
430/// are needed. Once the caller provides the requested ranges via
431/// [`Self::push_ranges`], they try to decode again by calling
432/// [`Self::try_decode`].
433///
434/// The decoder's internal state tracks what has been already decoded and what
435/// is needed next.
436#[derive(Debug)]
437pub struct ParquetPushDecoder {
438    /// The inner state.
439    ///
440    /// This state is consumed on every transition and a new state is produced
441    /// so the Rust compiler can ensure that the state is always valid and
442    /// transitions are not missed.
443    state: ParquetDecoderState,
444}
445
446impl ParquetPushDecoder {
447    /// Attempt to decode the next batch of data, or return what data is needed
448    ///
449    /// The the decoder communicates the next state with a [`DecodeResult`]
450    ///
451    /// See full example in [`ParquetPushDecoderBuilder`]
452    ///
453    /// ```no_run
454    /// # use parquet::arrow::push_decoder::ParquetPushDecoder;
455    /// use parquet::DecodeResult;
456    /// # fn get_decoder() -> ParquetPushDecoder { unimplemented!() }
457    /// # fn push_data(decoder: &mut ParquetPushDecoder, ranges: Vec<std::ops::Range<u64>>) { unimplemented!() }
458    /// let mut decoder = get_decoder();
459    /// loop {
460    ///    match decoder.try_decode().unwrap() {
461    ///       DecodeResult::NeedsData(ranges) => {
462    ///         // The decoder needs more data. Fetch the data for the given ranges
463    ///         // call decoder.push_ranges(ranges, data) and call again
464    ///         push_data(&mut decoder, ranges);
465    ///       }
466    ///       DecodeResult::Data(batch) => {
467    ///         // Successfully decoded the next batch of data
468    ///         println!("Got batch with {} rows", batch.num_rows());
469    ///       }
470    ///       DecodeResult::Finished => {
471    ///         // The decoder has finished decoding all data
472    ///         break;
473    ///       }
474    ///    }
475    /// }
476    ///```
477    pub fn try_decode(&mut self) -> Result<DecodeResult<RecordBatch>, ParquetError> {
478        let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished);
479        let (new_state, decode_result) = current_state.try_next_batch()?;
480        self.state = new_state;
481        Ok(decode_result)
482    }
483
484    /// Return a [`ParquetRecordBatchReader`] that reads the next set of rows, or
485    /// return what data is needed to produce it.
486    ///
487    /// This API can be used to get a reader for decoding the next set of
488    /// RecordBatches while proceeding to begin fetching data for the set (e.g
489    /// row group)
490    ///
491    /// Example
492    /// ```no_run
493    /// # use parquet::arrow::push_decoder::ParquetPushDecoder;
494    /// use parquet::DecodeResult;
495    /// # fn get_decoder() -> ParquetPushDecoder { unimplemented!() }
496    /// # fn push_data(decoder: &mut ParquetPushDecoder, ranges: Vec<std::ops::Range<u64>>) { unimplemented!() }
497    /// let mut decoder = get_decoder();
498    /// loop {
499    ///    match decoder.try_next_reader().unwrap() {
500    ///       DecodeResult::NeedsData(ranges) => {
501    ///         // The decoder needs more data. Fetch the data for the given ranges
502    ///         // call decoder.push_ranges(ranges, data) and call again
503    ///         push_data(&mut decoder, ranges);
504    ///       }
505    ///       DecodeResult::Data(reader) => {
506    ///          // spawn a thread to read the batches in parallel
507    ///          // with fetching the next row group / data
508    ///          std::thread::spawn(move || {
509    ///            for batch in reader {
510    ///              let batch = batch.unwrap();
511    ///              println!("Got batch with {} rows", batch.num_rows());
512    ///            }
513    ///         });
514    ///       }
515    ///       DecodeResult::Finished => {
516    ///         // The decoder has finished decoding all data
517    ///         break;
518    ///       }
519    ///    }
520    /// }
521    ///```
522    pub fn try_next_reader(
523        &mut self,
524    ) -> Result<DecodeResult<ParquetRecordBatchReader>, ParquetError> {
525        let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished);
526        let (new_state, decode_result) = current_state.try_next_reader()?;
527        self.state = new_state;
528        Ok(decode_result)
529    }
530
531    /// Push data into the decoder for processing
532    ///
533    /// This is a convenience wrapper around [`Self::push_ranges`] for pushing a
534    /// single range of data.
535    ///
536    /// Note this can be the entire file or just a part of it. If it is part of the file,
537    /// the ranges should correspond to the data ranges requested by the decoder.
538    ///
539    /// # Memory Usage
540    ///
541    /// The decoder releases the pushed bytes that it will not read again (for
542    /// example, the bytes of a row group when it is done with that row group).
543    /// However, it does not release *all* pushed bytes: it keeps the bytes of
544    /// row groups that it does not read (for example, row groups without
545    /// selected rows, or after the limit). Call [`Self::clear_all_ranges`] to
546    /// release them.
547    ///
548    /// See example in [`ParquetPushDecoderBuilder`]
549    pub fn push_range(&mut self, range: Range<u64>, data: Bytes) -> Result<(), ParquetError> {
550        self.push_ranges(vec![range], vec![data])
551    }
552
553    /// Push data into the decoder for processing
554    ///
555    /// This should correspond to the data ranges requested by the decoder.
556    ///
557    /// If you fetch data before the decoder requests it, for example with
558    /// [`Self::scan_plan`], keep it in your own cache and push exactly the
559    /// ranges in each [`DecodeResult::NeedsData`]:
560    ///
561    /// * The decoder does not use a requested range that is split over two
562    ///   pushed buffers.
563    /// * The decoder releases a pushed buffer only if its range is equal to a
564    ///   requested range. Other pushed data stays buffered.
565    pub fn push_ranges(
566        &mut self,
567        ranges: Vec<Range<u64>>,
568        data: Vec<Bytes>,
569    ) -> Result<(), ParquetError> {
570        let current_state = std::mem::replace(&mut self.state, ParquetDecoderState::Finished);
571        self.state = current_state.push_data(ranges, data)?;
572        Ok(())
573    }
574
575    /// Returns the total number of buffered bytes in the decoder
576    ///
577    /// This is the sum of the size of all [`Bytes`] that has been pushed to the
578    /// decoder but not yet consumed.
579    ///
580    /// Note that this does not include any overhead of the internal data
581    /// structures and that since [`Bytes`] are ref counted memory, this may not
582    /// reflect additional memory usage.
583    ///
584    /// This can be used to monitor memory usage of the decoder.
585    pub fn buffered_bytes(&self) -> u64 {
586        self.state.buffered_bytes()
587    }
588
589    /// Clear any staged byte ranges currently buffered for future decode work.
590    ///
591    /// This clears byte ranges still owned by the decoder's internal
592    /// `PushBuffers`. It does not affect any data that has already been handed
593    /// off to an active [`ParquetRecordBatchReader`].
594    pub fn clear_all_ranges(&mut self) {
595        self.state.clear_all_ranges();
596    }
597
598    /// True iff the decoder is at a row-group boundary, where
599    /// [`Self::into_builder`] can reconfigure the scan.
600    ///
601    /// A boundary is "between row groups": the previous row group's
602    /// [`ParquetRecordBatchReader`] has been fully extracted (via
603    /// [`Self::try_next_reader`]) or fully drained (via [`Self::try_decode`]),
604    /// and the next row group has not yet been planned. While
605    /// [`Self::try_decode`] is iterating an active row group's reader this
606    /// returns `false`; with [`Self::try_next_reader`] there is a clean
607    /// window between two consecutive returns where this is `true`.
608    pub fn is_at_row_group_boundary(&self) -> bool {
609        self.state.is_at_row_group_boundary()
610    }
611
612    /// Number of row groups left to decode after the one currently in flight.
613    /// Useful as a "should I bother reconfiguring the scan?" signal.
614    pub fn row_groups_remaining(&self) -> usize {
615        self.state.row_groups_remaining()
616    }
617
618    /// Returns the row-group index that the next call to
619    /// [`Self::try_next_reader`] will yield a reader for, after applying
620    /// any internal skipping (row selection emptiness, exhausted budget,
621    /// finished state).
622    ///
623    /// Safe to call at any time. When called mid-row-group (i.e. while a
624    /// previously-emitted [`ParquetRecordBatchReader`] is still being
625    /// drained), the returned index refers to the row group that
626    /// [`Self::try_next_reader`] will produce *after* the current one.
627    ///
628    /// Returns `Ok(None)` when:
629    /// - the decoder has no more row groups to read, or
630    /// - every remaining row group would be skipped.
631    ///
632    /// This method does not mutate decoder state. It is useful for
633    /// callers that maintain per-row-group state in lock-step with the
634    /// decoder (e.g. dynamic row-group pruners) to determine which row
635    /// group the next reader corresponds to, since
636    /// [`Self::try_next_reader`] may silently advance past row groups
637    /// based on filtering and other criteria.
638    pub fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
639        self.state.peek_next_row_group()
640    }
641
642    /// Returns the byte ranges that this decoder may still read, in the order
643    /// that decoding needs them. See [`ScanPlan`].
644    pub fn scan_plan(&self) -> ScanPlan {
645        self.state.scan_plan()
646    }
647
648    /// Decompose this decoder back into a [`ParquetPushDecoderBuilder`] for the
649    /// row groups that have *not* yet been decoded.
650    ///
651    /// This is the API for *adaptive* scans. Drive the decoder with
652    /// [`Self::try_next_reader`]; at any row-group boundary, call
653    /// `into_builder` to recover a builder, adjust it with the usual
654    /// [`ParquetPushDecoderBuilder`] setters, and
655    /// [`build`](ParquetPushDecoderBuilder::build) a fresh decoder that resumes
656    /// from the next row group:
657    ///
658    /// ```no_run
659    /// # use parquet::arrow::push_decoder::ParquetPushDecoder;
660    /// # use parquet::arrow::arrow_reader::RowFilter;
661    /// # fn get_decoder() -> ParquetPushDecoder { unimplemented!() }
662    /// # fn new_filter() -> RowFilter { unimplemented!() }
663    /// let mut decoder = get_decoder();
664    /// // ... drive `decoder.try_next_reader()` for a few row groups ...
665    /// if decoder.is_at_row_group_boundary() && decoder.row_groups_remaining() > 0 {
666    ///     decoder = decoder
667    ///         .into_builder()
668    ///         .unwrap()
669    ///         // any builder option can be changed here, e.g. promote a
670    ///         // filter into a row filter based on observed selectivity
671    ///         .with_row_filter(new_filter())
672    ///         .build()
673    ///         .unwrap();
674    /// }
675    /// ```
676    ///
677    /// The returned builder preserves the not-yet-decoded row-group plan,
678    /// including any row-group-local selections, together with the remaining
679    /// offset/limit budget. Rows from already-decoded row groups are not
680    /// produced again. Every other option — projection, row filter, row
681    /// selection policy, batch size, metrics, predicate-cache size — is left
682    /// exactly as the decoder had it and can be overridden before
683    /// [`build`](ParquetPushDecoderBuilder::build).
684    ///
685    /// Because the preserved plan pins the remaining row groups, the mutual
686    /// exclusion documented on
687    /// [`with_row_group_selections`](ParquetPushDecoderBuilder::with_row_group_selections)
688    /// applies to the rebuilt builder as well: calling it on a builder rebuilt
689    /// from a decoder configured with `with_row_groups` / `with_row_selection`
690    /// (or vice versa) returns an error from `build`.
691    ///
692    /// # Errors
693    ///
694    /// Returns `Err(ParquetError::General)` when the decoder is not at a
695    /// row-group boundary (check [`Self::is_at_row_group_boundary`] first) or
696    /// has already finished. The decoder is consumed either way.
697    ///
698    /// # Buffered bytes
699    ///
700    /// The decoder's buffered bytes are carried across the rebuild: bytes
701    /// already fetched for row groups the new configuration still reads are
702    /// not re-requested. [`build`](ParquetPushDecoderBuilder::build) releases
703    /// the bytes that the new configuration does not read. See
704    /// [`push_range`](Self::push_range) for more details.
705    pub fn into_builder(self) -> Result<ParquetPushDecoderBuilder, ParquetError> {
706        self.state.into_builder()
707    }
708}
709
710/// Internal state machine for the [`ParquetPushDecoder`]
711#[derive(Debug)]
712enum ParquetDecoderState {
713    /// Waiting for data needed to decode the next RowGroup
714    ReadingRowGroup {
715        remaining_row_groups: Box<RemainingRowGroups>,
716    },
717    /// The decoder is actively decoding a RowGroup
718    DecodingRowGroup {
719        /// Current active reader
720        record_batch_reader: Box<ParquetRecordBatchReader>,
721        remaining_row_groups: Box<RemainingRowGroups>,
722    },
723    /// The decoder has finished processing all data
724    Finished,
725}
726
727impl ParquetDecoderState {
728    /// If actively reading a RowGroup, return the currently active
729    /// ParquetRecordBatchReader and advance to the next group.
730    fn try_next_reader(
731        self,
732    ) -> Result<(Self, DecodeResult<ParquetRecordBatchReader>), ParquetError> {
733        let mut current_state = self;
734        loop {
735            let (next_state, decode_result) = current_state.transition()?;
736            // if more data is needed to transition, can't proceed further without it
737            match decode_result {
738                DecodeResult::NeedsData(ranges) => {
739                    return Ok((next_state, DecodeResult::NeedsData(ranges)));
740                }
741                // act next based on state
742                DecodeResult::Data(()) | DecodeResult::Finished => {}
743            }
744            match next_state {
745                // not ready to read yet, continue transitioning
746                Self::ReadingRowGroup { .. } => current_state = next_state,
747                // have a reader ready, so return it and set ourself to ReadingRowGroup
748                Self::DecodingRowGroup {
749                    record_batch_reader,
750                    remaining_row_groups,
751                } => {
752                    let result = DecodeResult::Data(*record_batch_reader);
753                    let next_state = Self::ReadingRowGroup {
754                        remaining_row_groups,
755                    };
756                    return Ok((next_state, result));
757                }
758                Self::Finished => {
759                    return Ok((Self::Finished, DecodeResult::Finished));
760                }
761            }
762        }
763    }
764
765    /// Current state --> next state + output
766    ///
767    /// This function is called to get the next RecordBatch
768    ///
769    /// This structure is used to reduce the indentation level of the main loop
770    /// in try_build
771    fn try_next_batch(self) -> Result<(Self, DecodeResult<RecordBatch>), ParquetError> {
772        let mut current_state = self;
773        loop {
774            let (new_state, decode_result) = current_state.transition()?;
775            // if more data is needed to transition, can't proceed further without it
776            match decode_result {
777                DecodeResult::NeedsData(ranges) => {
778                    return Ok((new_state, DecodeResult::NeedsData(ranges)));
779                }
780                // act next based on state
781                DecodeResult::Data(()) | DecodeResult::Finished => {}
782            }
783            match new_state {
784                // not ready to read yet, continue transitioning
785                Self::ReadingRowGroup { .. } => current_state = new_state,
786                // have a reader ready, so decode the next batch
787                Self::DecodingRowGroup {
788                    mut record_batch_reader,
789                    remaining_row_groups,
790                } => {
791                    match record_batch_reader.next() {
792                        // Successfully decoded a batch, return it
793                        Some(Ok(batch)) => {
794                            let result = DecodeResult::Data(batch);
795                            let next_state = Self::DecodingRowGroup {
796                                record_batch_reader,
797                                remaining_row_groups,
798                            };
799                            return Ok((next_state, result));
800                        }
801                        // No more batches in this row group, move to the next row group
802                        None => {
803                            current_state = Self::ReadingRowGroup {
804                                remaining_row_groups,
805                            }
806                        }
807                        // some error occurred while decoding, so return that
808                        Some(Err(e)) => {
809                            // TODO: preserve ArrowError in ParquetError (rather than convert to a string)
810                            return Err(ParquetError::ArrowError(e.to_string()));
811                        }
812                    }
813                }
814                Self::Finished => {
815                    return Ok((Self::Finished, DecodeResult::Finished));
816                }
817            }
818        }
819    }
820
821    /// Transition to the next state with a reader (data can be produced), if not end of stream
822    ///
823    /// This function is called in a loop until the decoder is ready to return
824    /// data (has the required pages buffered) or is finished.
825    fn transition(self) -> Result<(Self, DecodeResult<()>), ParquetError> {
826        // result returned when there is data ready
827        let data_ready = DecodeResult::Data(());
828        match self {
829            Self::ReadingRowGroup {
830                mut remaining_row_groups,
831            } => {
832                match remaining_row_groups.try_next_reader()? {
833                    // If we have a next reader, we can transition to decoding it
834                    DecodeResult::Data(record_batch_reader) => {
835                        // Transition to decoding the row group
836                        Ok((
837                            Self::DecodingRowGroup {
838                                record_batch_reader: Box::new(record_batch_reader),
839                                remaining_row_groups,
840                            },
841                            data_ready,
842                        ))
843                    }
844                    DecodeResult::NeedsData(ranges) => {
845                        // If we need more data, we return the ranges needed and stay in Reading
846                        // RowGroup state
847                        Ok((
848                            Self::ReadingRowGroup {
849                                remaining_row_groups,
850                            },
851                            DecodeResult::NeedsData(ranges),
852                        ))
853                    }
854                    // If there are no more readers, we are finished
855                    DecodeResult::Finished => {
856                        // No more row groups to read, we are finished
857                        Ok((Self::Finished, DecodeResult::Finished))
858                    }
859                }
860            }
861            // if we are already in DecodingRowGroup, just return data ready
862            Self::DecodingRowGroup { .. } => Ok((self, data_ready)),
863            // if finished, just return finished
864            Self::Finished => Ok((self, DecodeResult::Finished)),
865        }
866    }
867
868    /// Push data, and transition state if needed
869    ///
870    /// This should correspond to the data ranges requested by the decoder
871    pub fn push_data(
872        self,
873        ranges: Vec<Range<u64>>,
874        data: Vec<Bytes>,
875    ) -> Result<Self, ParquetError> {
876        match self {
877            ParquetDecoderState::ReadingRowGroup {
878                mut remaining_row_groups,
879            } => {
880                // Push data to the RowGroupReaderBuilder
881                remaining_row_groups.push_data(ranges, data)?;
882                Ok(ParquetDecoderState::ReadingRowGroup {
883                    remaining_row_groups,
884                })
885            }
886            // it is ok to get data before we asked for it
887            ParquetDecoderState::DecodingRowGroup {
888                record_batch_reader,
889                mut remaining_row_groups,
890            } => {
891                remaining_row_groups.push_data(ranges, data)?;
892                Ok(ParquetDecoderState::DecodingRowGroup {
893                    record_batch_reader,
894                    remaining_row_groups,
895                })
896            }
897            ParquetDecoderState::Finished => Err(ParquetError::General(
898                "Cannot push data to a finished decoder".to_string(),
899            )),
900        }
901    }
902
903    /// How many bytes are currently buffered in the decoder?
904    fn buffered_bytes(&self) -> u64 {
905        match self {
906            ParquetDecoderState::ReadingRowGroup {
907                remaining_row_groups,
908            } => remaining_row_groups.buffered_bytes(),
909            ParquetDecoderState::DecodingRowGroup {
910                record_batch_reader: _,
911                remaining_row_groups,
912            } => remaining_row_groups.buffered_bytes(),
913            ParquetDecoderState::Finished => 0,
914        }
915    }
916
917    /// Clear any staged ranges currently buffered in the decoder.
918    fn clear_all_ranges(&mut self) {
919        match self {
920            ParquetDecoderState::ReadingRowGroup {
921                remaining_row_groups,
922            } => remaining_row_groups.clear_all_ranges(),
923            ParquetDecoderState::DecodingRowGroup {
924                record_batch_reader: _,
925                remaining_row_groups,
926            } => remaining_row_groups.clear_all_ranges(),
927            ParquetDecoderState::Finished => {}
928        }
929    }
930
931    fn is_at_row_group_boundary(&self) -> bool {
932        match self {
933            ParquetDecoderState::ReadingRowGroup {
934                remaining_row_groups,
935            } => remaining_row_groups.is_at_row_group_boundary(),
936            // Mid-row-group: the active reader holds an `ArrayReader` and
937            // `ReadPlan` keyed to the *current* projection/filter; rebuilding
938            // would require throwing that work away.
939            ParquetDecoderState::DecodingRowGroup { .. } => false,
940            ParquetDecoderState::Finished => false,
941        }
942    }
943
944    /// See [`ParquetPushDecoder::scan_plan`].
945    fn scan_plan(&self) -> ScanPlan {
946        match self {
947            ParquetDecoderState::ReadingRowGroup {
948                remaining_row_groups,
949            }
950            | ParquetDecoderState::DecodingRowGroup {
951                remaining_row_groups,
952                ..
953            } => remaining_row_groups.scan_plan(),
954            ParquetDecoderState::Finished => ScanPlan::empty(),
955        }
956    }
957
958    fn row_groups_remaining(&self) -> usize {
959        match self {
960            ParquetDecoderState::ReadingRowGroup {
961                remaining_row_groups,
962            } => remaining_row_groups.row_groups_remaining(),
963            ParquetDecoderState::DecodingRowGroup {
964                remaining_row_groups,
965                ..
966            } => remaining_row_groups.row_groups_remaining(),
967            ParquetDecoderState::Finished => 0,
968        }
969    }
970
971    /// See [`ParquetPushDecoder::peek_next_row_group`] for the public
972    /// API contract. This inner method delegates to the underlying
973    /// [`RemainingRowGroups`] for both `ReadingRowGroup` and
974    /// `DecodingRowGroup`: mid-row-group the answer is the row group
975    /// `try_next_reader` will produce *after* the active one finishes,
976    /// which is exactly what `RemainingRowGroups::peek_next_row_group`
977    /// computes from the queued frontier.
978    fn peek_next_row_group(&self) -> Result<Option<usize>, ParquetError> {
979        match self {
980            ParquetDecoderState::ReadingRowGroup {
981                remaining_row_groups,
982            } => remaining_row_groups.peek_next_row_group(),
983            ParquetDecoderState::DecodingRowGroup {
984                remaining_row_groups,
985                ..
986            } => remaining_row_groups.peek_next_row_group(),
987            ParquetDecoderState::Finished => Ok(None),
988        }
989    }
990
991    fn into_builder(self) -> Result<ParquetPushDecoderBuilder, ParquetError> {
992        let remaining_row_groups = match self {
993            ParquetDecoderState::ReadingRowGroup {
994                remaining_row_groups,
995            } => remaining_row_groups,
996            ParquetDecoderState::DecodingRowGroup { .. } => {
997                return Err(ParquetError::General(
998                    "into_builder called while a row group is being decoded; \
999                     check is_at_row_group_boundary() first"
1000                        .to_string(),
1001                ));
1002            }
1003            ParquetDecoderState::Finished => {
1004                return Err(ParquetError::General(
1005                    "into_builder called on a finished decoder".to_string(),
1006                ));
1007            }
1008        };
1009        if !remaining_row_groups.is_at_row_group_boundary() {
1010            return Err(ParquetError::General(
1011                "into_builder called mid-row-group; check is_at_row_group_boundary() first"
1012                    .to_string(),
1013            ));
1014        }
1015        Ok(builder_from_remaining(remaining_row_groups.into_parts()))
1016    }
1017}
1018
1019#[cfg(test)]
1020mod test {
1021    use super::*;
1022    use crate::DecodeResult;
1023    use crate::arrow::arrow_reader::{ArrowPredicateFn, RowFilter, RowSelection, RowSelector};
1024    use crate::arrow::push_decoder::reader_builder::column_chunk_range;
1025    use crate::arrow::push_decoder::{
1026        ParquetPushDecoder, ParquetPushDecoderBuilder, RowGroupSelection,
1027    };
1028    use crate::arrow::{ArrowWriter, ProjectionMask};
1029    use crate::errors::ParquetError;
1030    use crate::file::metadata::page_index::PageIndexBuilder;
1031    use crate::file::metadata::{PageIndexPolicy, ParquetMetaDataPushDecoder};
1032    use crate::file::properties::WriterProperties;
1033    use arrow::compute::kernels::cmp::{gt, lt};
1034    use arrow_array::cast::AsArray;
1035    use arrow_array::types::Int64Type;
1036    use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringViewArray};
1037    use arrow_buffer::BooleanBuffer;
1038    use arrow_select::concat::concat_batches;
1039    use bytes::Bytes;
1040    use std::fmt::Debug;
1041    use std::ops::Range;
1042    use std::sync::{Arc, LazyLock};
1043
1044    /// Test decoder struct size (as they are copied around on each transition, they
1045    /// should not grow too large)
1046    #[test]
1047    fn test_decoder_size() {
1048        assert_eq!(std::mem::size_of::<ParquetDecoderState>(), 24);
1049    }
1050
1051    /// Decode the entire file at once, simulating a scenario where all data is
1052    /// available in memory
1053    #[test]
1054    fn test_decoder_all_data() {
1055        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1056            .unwrap()
1057            .build()
1058            .unwrap();
1059
1060        decoder
1061            .push_range(test_file_range(), TEST_FILE_DATA.clone())
1062            .unwrap();
1063
1064        let results = vec![
1065            // first row group should be decoded without needing more data
1066            expect_data(decoder.try_decode()),
1067            // second row group should be decoded without needing more data
1068            expect_data(decoder.try_decode()),
1069        ];
1070        expect_finished(decoder.try_decode());
1071
1072        let all_output = concat_batches(&TEST_BATCH.schema(), &results).unwrap();
1073        // Check that the output matches the input batch
1074        assert_eq!(all_output, *TEST_BATCH);
1075    }
1076
1077    /// Decode the entire file incrementally, simulating a scenario where data is
1078    /// fetched as needed
1079    #[test]
1080    fn test_decoder_incremental() {
1081        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1082            .unwrap()
1083            .build()
1084            .unwrap();
1085
1086        let mut results = vec![];
1087
1088        // First row group, expect a single request
1089        let ranges = expect_needs_data(decoder.try_decode());
1090        let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum();
1091        push_ranges_to_decoder(&mut decoder, ranges);
1092        // The decoder should currently only store the data it needs to decode the first row group
1093        assert_eq!(decoder.buffered_bytes(), num_bytes_requested);
1094        results.push(expect_data(decoder.try_decode()));
1095        // the decoder should have consumed the data for the first row group and freed it
1096        assert_eq!(decoder.buffered_bytes(), 0);
1097
1098        // Second row group,
1099        let ranges = expect_needs_data(decoder.try_decode());
1100        let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum();
1101        push_ranges_to_decoder(&mut decoder, ranges);
1102        // The decoder should currently only store the data it needs to decode the second row group
1103        assert_eq!(decoder.buffered_bytes(), num_bytes_requested);
1104        results.push(expect_data(decoder.try_decode()));
1105        // the decoder should have consumed the data for the second row group and freed it
1106        assert_eq!(decoder.buffered_bytes(), 0);
1107        expect_finished(decoder.try_decode());
1108
1109        // Check that the output matches the input batch
1110        let all_output = concat_batches(&TEST_BATCH.schema(), &results).unwrap();
1111        assert_eq!(all_output, *TEST_BATCH);
1112    }
1113
1114    /// Releasing staged ranges should free speculative buffers without affecting
1115    /// the active row group reader.
1116    #[test]
1117    fn test_decoder_clear_all_ranges() {
1118        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1119            .unwrap()
1120            .with_batch_size(100)
1121            .build()
1122            .unwrap();
1123
1124        decoder
1125            .push_range(test_file_range(), TEST_FILE_DATA.clone())
1126            .unwrap();
1127        assert_eq!(decoder.buffered_bytes(), test_file_len());
1128
1129        // The current row group reader is built from the prefetched bytes.
1130        // The decoder releases the bytes of the first row group, and the rest
1131        // of the speculative full-file range remains staged in the decoder.
1132        let batch1 = expect_data(decoder.try_decode());
1133        assert_eq!(batch1, TEST_BATCH.slice(0, 100));
1134        assert_eq!(
1135            decoder.buffered_bytes(),
1136            test_file_len() - test_file_row_group_bytes(0)
1137        );
1138
1139        // All of the buffer is released
1140        decoder.clear_all_ranges();
1141        assert_eq!(decoder.buffered_bytes(), 0);
1142
1143        // The active reader still owns the current row group's bytes, so it can
1144        // continue decoding without consulting PushBuffers.
1145        let batch2 = expect_data(decoder.try_decode());
1146        assert_eq!(batch2, TEST_BATCH.slice(100, 100));
1147        assert_eq!(decoder.buffered_bytes(), 0);
1148
1149        // Moving to the next row group now requires the decoder to ask for data
1150        // again because the staged speculative ranges were released.
1151        let ranges = expect_needs_data(decoder.try_decode());
1152        let num_bytes_requested: u64 = ranges.iter().map(|r| r.end - r.start).sum();
1153        push_ranges_to_decoder(&mut decoder, ranges);
1154        assert_eq!(decoder.buffered_bytes(), num_bytes_requested);
1155
1156        let batch3 = expect_data(decoder.try_decode());
1157        assert_eq!(batch3, TEST_BATCH.slice(200, 100));
1158        assert_eq!(decoder.buffered_bytes(), 0);
1159
1160        let batch4 = expect_data(decoder.try_decode());
1161        assert_eq!(batch4, TEST_BATCH.slice(300, 100));
1162        assert_eq!(decoder.buffered_bytes(), 0);
1163
1164        expect_finished(decoder.try_decode());
1165    }
1166
1167    /// Decode the entire file incrementally, simulating partial reads
1168    #[test]
1169    fn test_decoder_partial() {
1170        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1171            .unwrap()
1172            .build()
1173            .unwrap();
1174
1175        // First row group, expect a single request for all data needed to read "a" and "b"
1176        let ranges = expect_needs_data(decoder.try_decode());
1177        push_ranges_to_decoder(&mut decoder, ranges);
1178
1179        let batch1 = expect_data(decoder.try_decode());
1180        let expected1 = TEST_BATCH.slice(0, 200);
1181        assert_eq!(batch1, expected1);
1182
1183        // Second row group, this time provide the data in two steps
1184        let ranges = expect_needs_data(decoder.try_decode());
1185        let (ranges1, ranges2) = ranges.split_at(ranges.len() / 2);
1186        assert!(!ranges1.is_empty());
1187        assert!(!ranges2.is_empty());
1188        // push first half to simulate partial read
1189        push_ranges_to_decoder(&mut decoder, ranges1.to_vec());
1190
1191        // still expect more data
1192        let ranges = expect_needs_data(decoder.try_decode());
1193        assert_eq!(ranges, ranges2); // should be the remaining ranges
1194        // push empty ranges should be a no-op
1195        push_ranges_to_decoder(&mut decoder, vec![]);
1196        let ranges = expect_needs_data(decoder.try_decode());
1197        assert_eq!(ranges, ranges2); // should be the remaining ranges
1198        push_ranges_to_decoder(&mut decoder, ranges);
1199
1200        let batch2 = expect_data(decoder.try_decode());
1201        let expected2 = TEST_BATCH.slice(200, 200);
1202        assert_eq!(batch2, expected2);
1203
1204        expect_finished(decoder.try_decode());
1205    }
1206
1207    /// Push only the pages needed for each batch of a row group rather than
1208    /// the entire row group, and check whether the decoder can produce that
1209    /// batch before the rest of the row group's pages have been pushed.
1210    #[test]
1211    fn test_decoder_first_pages_only() {
1212        let metadata = test_file_parquet_metadata_with_offset_index();
1213        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata))
1214            .unwrap()
1215            // Each data page has 100 rows, so a batch of 100 rows needs only
1216            // the first data page of each column
1217            .with_batch_size(100)
1218            .build()
1219            .unwrap();
1220
1221        // Row group 0: the decoder asks for the entire column chunk of each
1222        // of the three columns "a", "b", and "c"
1223        let ranges = expect_needs_data(decoder.try_decode());
1224        assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]);
1225
1226        // Compute the ranges that cover only the first data page of each
1227        // column (and the dictionary page which precedes it)
1228        let page_index = metadata.page_index_for_row_group(0);
1229        let row_group = metadata.row_group(0);
1230        let mut first_page_ranges = vec![];
1231        let mut second_page_ranges = vec![];
1232        for (idx, column) in row_group.columns().iter().enumerate() {
1233            let (start, len) = column.byte_range();
1234            let locations = page_index.page_locations(idx).unwrap();
1235            assert_eq!(locations.len(), 2, "expected 2 data pages per column chunk");
1236            let second_page_start = locations[1].offset as u64;
1237            first_page_ranges.push(start..second_page_start);
1238            second_page_ranges.push(second_page_start..start + len);
1239        }
1240        // Note the first range for each column includes the dictionary page
1241        assert_eq!(first_page_ranges, vec![4..1734, 1860..3590, 3716..10936]);
1242        assert_eq!(
1243            second_page_ranges,
1244            vec![1734..1860, 3590..3716, 10936..11062]
1245        );
1246
1247        // Push only the first page of each column. This is all the data
1248        // needed to decode the first batch of 100 rows.
1249        push_ranges_to_decoder(&mut decoder, first_page_ranges);
1250
1251        // Note will likely change as part of
1252        // <https://github.com/apache/arrow-rs/issues/6946>
1253
1254        // decoder still reports it needs the (complete) ranges
1255        // it originally asked for, and does not produce a batch.
1256        let ranges = expect_needs_data(decoder.try_decode());
1257        assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]);
1258
1259        // Pushing the second pages as separate ranges does not help either:
1260        // the decoder does not coalesce adjacent pushed ranges, so a requested
1261        // range is only satisfied by a single pushed buffer that covers it.
1262        push_ranges_to_decoder(&mut decoder, second_page_ranges);
1263        let ranges = expect_needs_data(decoder.try_decode());
1264        assert_eq!(ranges, vec![4..1860, 1860..3716, 3716..11062]);
1265
1266        // Only once the exact ranges originally requested are pushed does the
1267        // decoder produce batches.
1268        push_ranges_to_decoder(&mut decoder, ranges);
1269        let batch = expect_data(decoder.try_decode());
1270        assert_eq!(batch, TEST_BATCH.slice(0, 100));
1271        let batch = expect_data(decoder.try_decode());
1272        assert_eq!(batch, TEST_BATCH.slice(100, 100));
1273    }
1274
1275    /// Push one buffer for each `NeedsData`, from the start of the first
1276    /// requested range to the end of the last one, and check that the
1277    /// decoder releases the bytes of each row group when it is done with it.
1278    ///
1279    /// <https://github.com/apache/arrow-rs/issues/11233>
1280    fn decode_with_coalesced_pushes(mut decoder: ParquetPushDecoder) -> Vec<RecordBatch> {
1281        let mut batches = vec![];
1282        loop {
1283            match decoder.try_decode().unwrap() {
1284                DecodeResult::NeedsData(ranges) => {
1285                    let start = ranges.iter().map(|r| r.start).min().unwrap();
1286                    let end = ranges.iter().map(|r| r.end).max().unwrap();
1287                    let coalesced = start..end;
1288                    push_ranges_to_decoder(&mut decoder, vec![coalesced]);
1289                }
1290                DecodeResult::Data(batch) => {
1291                    // The test file has one batch per row group.
1292                    assert_eq!(decoder.buffered_bytes(), 0);
1293                    batches.push(batch);
1294                }
1295                DecodeResult::Finished => return batches,
1296            }
1297        }
1298    }
1299
1300    /// The requested ranges are adjacent, so the pushed buffer holds only
1301    /// requested bytes.
1302    #[test]
1303    fn test_decoder_releases_coalesced_push() {
1304        let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1305            .unwrap()
1306            .build()
1307            .unwrap();
1308        let batches = decode_with_coalesced_pushes(decoder);
1309        let all_output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
1310        assert_eq!(all_output, *TEST_BATCH);
1311    }
1312
1313    /// Column "b" is not read, so the pushed buffer also holds the bytes of
1314    /// "b", which the decoder did not request.
1315    #[test]
1316    fn test_decoder_releases_unrequested_bytes_of_row_group() {
1317        let builder =
1318            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1319        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1320        let decoder = builder
1321            .with_projection(ProjectionMask::columns(&schema_descr, ["a", "c"]))
1322            .build()
1323            .unwrap();
1324        let batches = decode_with_coalesced_pushes(decoder);
1325        let all_output = concat_batches(&batches[0].schema(), &batches).unwrap();
1326        assert_eq!(all_output, TEST_BATCH.project(&[0, 2]).unwrap());
1327    }
1328
1329    /// A plan that reads a row group two times keeps its bytes until the
1330    /// second read.
1331    #[test]
1332    fn test_decoder_keeps_bytes_of_a_row_group_that_is_read_again() {
1333        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
1334            .unwrap()
1335            .with_row_groups(vec![0, 1, 0])
1336            .build()
1337            .unwrap();
1338        prefetch_test_file(&mut decoder);
1339        // Row group 0 is read again, so the decoder keeps its bytes.
1340        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 200));
1341        assert_eq!(decoder.buffered_bytes(), test_file_len());
1342
1343        assert_eq!(
1344            expect_data(decoder.try_decode()),
1345            TEST_BATCH.slice(200, 200)
1346        );
1347        assert_eq!(
1348            decoder.buffered_bytes(),
1349            test_file_len() - test_file_row_group_bytes(1)
1350        );
1351
1352        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 200));
1353        assert_eq!(
1354            decoder.buffered_bytes(),
1355            test_file_len() - test_file_row_group_bytes(0) - test_file_row_group_bytes(1)
1356        );
1357        expect_finished(decoder.try_decode());
1358    }
1359
1360    /// With a predicate, the decoder requests the bytes of the predicate
1361    /// column first, then the bytes of the other output columns.
1362    #[test]
1363    fn test_decoder_releases_coalesced_push_with_predicate() {
1364        let builder =
1365            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1366        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1367        let row_filter_a = ArrowPredicateFn::new(
1368            ProjectionMask::columns(&schema_descr, ["a"]),
1369            |batch: RecordBatch| {
1370                let scalar_250 = Int64Array::new_scalar(250);
1371                let column = batch.column(0).as_primitive::<Int64Type>();
1372                gt(column, &scalar_250)
1373            },
1374        );
1375        let decoder = builder
1376            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1377            .build()
1378            .unwrap();
1379        let batches = decode_with_coalesced_pushes(decoder);
1380        let all_output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
1381        assert_eq!(all_output, TEST_BATCH.slice(251, 149));
1382    }
1383
1384    /// Decode multiple columns "a" and "b", expect that the decoder requests
1385    /// only a single request per row group
1386    #[test]
1387    fn test_decoder_selection_does_one_request() {
1388        let builder =
1389            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1390
1391        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1392
1393        let mut decoder = builder
1394            .with_projection(
1395                ProjectionMask::columns(&schema_descr, ["a", "b"]), // read "a", "b"
1396            )
1397            .build()
1398            .unwrap();
1399
1400        // First row group, expect a single request for all data needed to read "a" and "b"
1401        let ranges = expect_needs_data(decoder.try_decode());
1402        push_ranges_to_decoder(&mut decoder, ranges);
1403
1404        let batch1 = expect_data(decoder.try_decode());
1405        let expected1 = TEST_BATCH.slice(0, 200).project(&[0, 1]).unwrap();
1406        assert_eq!(batch1, expected1);
1407
1408        // Second row group, similarly expect a single request for all data needed to read "a" and "b"
1409        let ranges = expect_needs_data(decoder.try_decode());
1410        push_ranges_to_decoder(&mut decoder, ranges);
1411
1412        let batch2 = expect_data(decoder.try_decode());
1413        let expected2 = TEST_BATCH.slice(200, 200).project(&[0, 1]).unwrap();
1414        assert_eq!(batch2, expected2);
1415
1416        expect_finished(decoder.try_decode());
1417    }
1418
1419    /// Decode with a filter that requires multiple requests, but only provide part
1420    /// of the data needed for the filter at a time simulating partial reads.
1421    #[test]
1422    fn test_decoder_single_filter_partial() {
1423        let builder =
1424            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1425
1426        // Values in column "a" range 0..399
1427        // First filter: "a" > 250  (nothing in Row Group 0, both data pages in Row Group 1)
1428        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1429
1430        // a > 250
1431        let row_filter_a = ArrowPredicateFn::new(
1432            // claim to use both a and b so we get two ranges requests for the filter pages
1433            ProjectionMask::columns(&schema_descr, ["a", "b"]),
1434            |batch: RecordBatch| {
1435                let scalar_250 = Int64Array::new_scalar(250);
1436                let column = batch.column(0).as_primitive::<Int64Type>();
1437                gt(column, &scalar_250)
1438            },
1439        );
1440
1441        let mut decoder = builder
1442            .with_projection(
1443                // read only column "a" to test that filter pages are reused
1444                ProjectionMask::columns(&schema_descr, ["a"]), // read "a"
1445            )
1446            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1447            .build()
1448            .unwrap();
1449
1450        // First row group, evaluating filters
1451        let ranges = expect_needs_data(decoder.try_decode());
1452        // only provide half the ranges
1453        let (ranges1, ranges2) = ranges.split_at(ranges.len() / 2);
1454        assert!(!ranges1.is_empty());
1455        assert!(!ranges2.is_empty());
1456        push_ranges_to_decoder(&mut decoder, ranges1.to_vec());
1457        // still expect more data
1458        let ranges = expect_needs_data(decoder.try_decode());
1459        assert_eq!(ranges, ranges2); // should be the remaining ranges
1460        let ranges = expect_needs_data(decoder.try_decode());
1461        assert_eq!(ranges, ranges2); // should be the remaining ranges
1462        push_ranges_to_decoder(&mut decoder, ranges2.to_vec());
1463
1464        // Since no rows in the first row group pass the filters, there is no
1465        // additional requests to read data pages for "b" here
1466
1467        // Second row group
1468        let ranges = expect_needs_data(decoder.try_decode());
1469        push_ranges_to_decoder(&mut decoder, ranges);
1470
1471        let batch = expect_data(decoder.try_decode());
1472        let expected = TEST_BATCH.slice(251, 149).project(&[0]).unwrap();
1473        assert_eq!(batch, expected);
1474
1475        expect_finished(decoder.try_decode());
1476    }
1477
1478    /// Decode with a filter where we also skip one of the RowGroups via a RowSelection
1479    #[test]
1480    fn test_decoder_single_filter_and_row_selection() {
1481        let builder =
1482            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1483
1484        // Values in column "a" range 0..399
1485        // First filter: "a" > 250  (nothing in Row Group 0, last data page in Row Group 1)
1486        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1487
1488        // a > 250
1489        let row_filter_a = ArrowPredicateFn::new(
1490            ProjectionMask::columns(&schema_descr, ["a"]),
1491            |batch: RecordBatch| {
1492                let scalar_250 = Int64Array::new_scalar(250);
1493                let column = batch.column(0).as_primitive::<Int64Type>();
1494                gt(column, &scalar_250)
1495            },
1496        );
1497
1498        let mut decoder = builder
1499            .with_projection(
1500                // read only column "a" to test that filter pages are reused
1501                ProjectionMask::columns(&schema_descr, ["b"]), // read "b"
1502            )
1503            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1504            .with_row_selection(RowSelection::from(vec![
1505                RowSelector::skip(200),   // skip first row group
1506                RowSelector::select(100), // first 100 rows of second row group
1507                RowSelector::skip(100),
1508            ]))
1509            .build()
1510            .unwrap();
1511
1512        // expect the first row group to be filtered out (no filter is evaluated due to row selection)
1513
1514        // First row group, first filter (a > 250)
1515        let ranges = expect_needs_data(decoder.try_decode());
1516        push_ranges_to_decoder(&mut decoder, ranges);
1517
1518        // Second row group
1519        let ranges = expect_needs_data(decoder.try_decode());
1520        push_ranges_to_decoder(&mut decoder, ranges);
1521
1522        let batch = expect_data(decoder.try_decode());
1523        let expected = TEST_BATCH.slice(251, 49).project(&[1]).unwrap();
1524        assert_eq!(batch, expected);
1525
1526        expect_finished(decoder.try_decode());
1527    }
1528
1529    /// Decode with multiple filters that require multiple requests
1530    #[test]
1531    fn test_decoder_multi_filters() {
1532        // Create a decoder for decoding parquet data (note it does not have any IO / readers)
1533        let builder =
1534            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1535
1536        // Values in column "a" range 0..399
1537        // Values in column "b" range 400..799
1538        // First filter: "a" > 175  (last data page in Row Group 0)
1539        // Second filter: "b" < 625 (last data page in Row Group 0 and first DataPage in RowGroup 1)
1540        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1541
1542        // a > 175
1543        let row_filter_a = ArrowPredicateFn::new(
1544            ProjectionMask::columns(&schema_descr, ["a"]),
1545            |batch: RecordBatch| {
1546                let scalar_175 = Int64Array::new_scalar(175);
1547                let column = batch.column(0).as_primitive::<Int64Type>();
1548                gt(column, &scalar_175)
1549            },
1550        );
1551
1552        // b < 625
1553        let row_filter_b = ArrowPredicateFn::new(
1554            ProjectionMask::columns(&schema_descr, ["b"]),
1555            |batch: RecordBatch| {
1556                let scalar_625 = Int64Array::new_scalar(625);
1557                let column = batch.column(0).as_primitive::<Int64Type>();
1558                lt(column, &scalar_625)
1559            },
1560        );
1561
1562        let mut decoder = builder
1563            .with_projection(
1564                ProjectionMask::columns(&schema_descr, ["c"]), // read "c"
1565            )
1566            .with_row_filter(RowFilter::new(vec![
1567                Box::new(row_filter_a),
1568                Box::new(row_filter_b),
1569            ]))
1570            .build()
1571            .unwrap();
1572
1573        // First row group, first filter (a > 175)
1574        let ranges = expect_needs_data(decoder.try_decode());
1575        push_ranges_to_decoder(&mut decoder, ranges);
1576
1577        // first row group, second filter (b < 625)
1578        let ranges = expect_needs_data(decoder.try_decode());
1579        push_ranges_to_decoder(&mut decoder, ranges);
1580
1581        // first row group, data pages for "c"
1582        let ranges = expect_needs_data(decoder.try_decode());
1583        push_ranges_to_decoder(&mut decoder, ranges);
1584
1585        // expect the first batch to be decoded: rows 176..199, column "c"
1586        let batch1 = expect_data(decoder.try_decode());
1587        let expected1 = TEST_BATCH.slice(176, 24).project(&[2]).unwrap();
1588        assert_eq!(batch1, expected1);
1589
1590        // Second row group, first filter (a > 175)
1591        let ranges = expect_needs_data(decoder.try_decode());
1592        push_ranges_to_decoder(&mut decoder, ranges);
1593
1594        // Second row group, second filter (b < 625)
1595        let ranges = expect_needs_data(decoder.try_decode());
1596        push_ranges_to_decoder(&mut decoder, ranges);
1597
1598        // Second row group, data pages for "c"
1599        let ranges = expect_needs_data(decoder.try_decode());
1600        push_ranges_to_decoder(&mut decoder, ranges);
1601
1602        // expect the second batch to be decoded: rows 200..224, column "c"
1603        let batch2 = expect_data(decoder.try_decode());
1604        let expected2 = TEST_BATCH.slice(200, 25).project(&[2]).unwrap();
1605        assert_eq!(batch2, expected2);
1606
1607        expect_finished(decoder.try_decode());
1608    }
1609
1610    /// Consecutive filters on the same projection share one decode stream.
1611    #[test]
1612    fn test_decoder_same_projection_filters() {
1613        let builder =
1614            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1615        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1616        let projection_a = ProjectionMask::columns(&schema_descr, ["a"]);
1617
1618        let row_filter_gt = ArrowPredicateFn::new(projection_a.clone(), |batch: RecordBatch| {
1619            let column = batch.column(0).as_primitive::<Int64Type>();
1620            gt(column, &Int64Array::new_scalar(175))
1621        });
1622        let row_filter_lt = ArrowPredicateFn::new(projection_a, |batch: RecordBatch| {
1623            let column = batch.column(0).as_primitive::<Int64Type>();
1624            assert!(column.iter().flatten().all(|value| value > 175));
1625            lt(column, &Int64Array::new_scalar(190))
1626        });
1627
1628        let mut decoder = builder
1629            .with_projection(ProjectionMask::columns(&schema_descr, ["c"]))
1630            .with_row_filter(RowFilter::new(vec![
1631                Box::new(row_filter_gt),
1632                Box::new(row_filter_lt),
1633            ]))
1634            .with_batch_size(50)
1635            .build()
1636            .unwrap();
1637
1638        // One filter-column request, followed by the selected output page.
1639        let ranges = expect_needs_data(decoder.try_decode());
1640        push_ranges_to_decoder(&mut decoder, ranges);
1641        let ranges = expect_needs_data(decoder.try_decode());
1642        push_ranges_to_decoder(&mut decoder, ranges);
1643
1644        let batch = expect_data(decoder.try_decode());
1645        let expected = TEST_BATCH.slice(176, 14).project(&[2]).unwrap();
1646        assert_eq!(batch, expected);
1647
1648        // Row group 1 is rejected by the fused predicates, so no output-column
1649        // request is made.
1650        let ranges = expect_needs_data(decoder.try_decode());
1651        push_ranges_to_decoder(&mut decoder, ranges);
1652        expect_finished(decoder.try_decode());
1653    }
1654
1655    /// Fused same-projection filters honour the output limit: evaluation
1656    /// stops once enough rows survive the whole group.
1657    #[test]
1658    fn test_decoder_same_projection_filters_with_limit() {
1659        use std::sync::atomic::{AtomicUsize, Ordering};
1660
1661        let builder =
1662            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1663        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1664        let projection_a = ProjectionMask::columns(&schema_descr, ["a"]);
1665
1666        let rows_evaluated = Arc::new(AtomicUsize::new(0));
1667        let rows_evaluated_for_predicate = Arc::clone(&rows_evaluated);
1668        let row_filter_gt =
1669            ArrowPredicateFn::new(projection_a.clone(), move |batch: RecordBatch| {
1670                rows_evaluated_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed);
1671                let column = batch.column(0).as_primitive::<Int64Type>();
1672                gt(column, &Int64Array::new_scalar(175))
1673            });
1674        let row_filter_lt = ArrowPredicateFn::new(projection_a, |batch: RecordBatch| {
1675            let column = batch.column(0).as_primitive::<Int64Type>();
1676            lt(column, &Int64Array::new_scalar(190))
1677        });
1678
1679        let mut decoder = builder
1680            .with_projection(ProjectionMask::columns(&schema_descr, ["c"]))
1681            .with_row_filter(RowFilter::new(vec![
1682                Box::new(row_filter_gt),
1683                Box::new(row_filter_lt),
1684            ]))
1685            .with_batch_size(10)
1686            .with_limit(5)
1687            .build()
1688            .unwrap();
1689
1690        // One filter-column request for the fused group, then the output page.
1691        let ranges = expect_needs_data(decoder.try_decode());
1692        push_ranges_to_decoder(&mut decoder, ranges);
1693        let ranges = expect_needs_data(decoder.try_decode());
1694        push_ranges_to_decoder(&mut decoder, ranges);
1695
1696        let batch = expect_data(decoder.try_decode());
1697        let expected = TEST_BATCH.slice(176, 5).project(&[2]).unwrap();
1698        assert_eq!(batch, expected);
1699
1700        // The limit was satisfied by row group 0.
1701        expect_finished(decoder.try_decode());
1702
1703        // Row 181 is the 6th match, so at most the batch holding it (rows
1704        // 180..189) is evaluated.
1705        let evaluated = rows_evaluated.load(Ordering::Relaxed);
1706        assert!(evaluated <= 190, "evaluated {evaluated} rows");
1707    }
1708
1709    /// Decode with a filter that uses a column that is also projected, and expect
1710    /// that the filter pages are reused (don't refetch them)
1711    #[test]
1712    fn test_decoder_reuses_filter_pages() {
1713        // Create a decoder for decoding parquet data (note it does not have any IO / readers)
1714        let builder =
1715            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1716
1717        // Values in column "a" range 0..399
1718        // First filter: "a" > 250  (nothing in Row Group 0, last data page in Row Group 1)
1719        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1720
1721        // a > 250
1722        let row_filter_a = ArrowPredicateFn::new(
1723            ProjectionMask::columns(&schema_descr, ["a"]),
1724            |batch: RecordBatch| {
1725                let scalar_250 = Int64Array::new_scalar(250);
1726                let column = batch.column(0).as_primitive::<Int64Type>();
1727                gt(column, &scalar_250)
1728            },
1729        );
1730
1731        let mut decoder = builder
1732            .with_projection(
1733                // read only column "a" to test that filter pages are reused
1734                ProjectionMask::columns(&schema_descr, ["a"]), // read "a"
1735            )
1736            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1737            .build()
1738            .unwrap();
1739
1740        // First row group, first filter (a > 175)
1741        let ranges = expect_needs_data(decoder.try_decode());
1742        push_ranges_to_decoder(&mut decoder, ranges);
1743
1744        // expect the first row group to be filtered out (no rows match)
1745
1746        // Second row group, first filter (a > 250)
1747        let ranges = expect_needs_data(decoder.try_decode());
1748        push_ranges_to_decoder(&mut decoder, ranges);
1749
1750        // expect that the second row group is decoded: rows 251..399, column "a"
1751        // Note that the filter pages for "a" should be reused and no additional data
1752        // should be requested
1753        let batch = expect_data(decoder.try_decode());
1754        let expected = TEST_BATCH.slice(251, 149).project(&[0]).unwrap();
1755        assert_eq!(batch, expected);
1756
1757        expect_finished(decoder.try_decode());
1758    }
1759
1760    #[test]
1761    fn test_decoder_empty_filters() {
1762        let builder =
1763            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1764        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1765
1766        // only read column "c", but with empty filters
1767        let mut decoder = builder
1768            .with_projection(
1769                ProjectionMask::columns(&schema_descr, ["c"]), // read "c"
1770            )
1771            .with_row_filter(RowFilter::new(vec![
1772                // empty filters should be ignored
1773            ]))
1774            .build()
1775            .unwrap();
1776
1777        // First row group
1778        let ranges = expect_needs_data(decoder.try_decode());
1779        push_ranges_to_decoder(&mut decoder, ranges);
1780
1781        // expect the first batch to be decoded: rows 0..199, column "c"
1782        let batch1 = expect_data(decoder.try_decode());
1783        let expected1 = TEST_BATCH.slice(0, 200).project(&[2]).unwrap();
1784        assert_eq!(batch1, expected1);
1785
1786        // Second row group,
1787        let ranges = expect_needs_data(decoder.try_decode());
1788        push_ranges_to_decoder(&mut decoder, ranges);
1789
1790        // expect the second batch to be decoded: rows 200..399, column "c"
1791        let batch2 = expect_data(decoder.try_decode());
1792        let expected2 = TEST_BATCH.slice(200, 200).project(&[2]).unwrap();
1793
1794        assert_eq!(batch2, expected2);
1795
1796        expect_finished(decoder.try_decode());
1797    }
1798
1799    /// When filter pushdown is combined with a `LIMIT`, the predicate must
1800    /// not be evaluated for rows beyond the `limit`-th match.
1801    ///
1802    /// Filter `a > 175` produces 24 matches in row group 0 (rows 176..199).
1803    /// With `limit = 10`, only the first 10 matches (rows 176..185) should be
1804    /// emitted, AND the predicate counter should observe that evaluation was
1805    /// short-circuited.
1806    #[test]
1807    fn test_decoder_filter_with_limit_short_circuits_within_row_group() {
1808        use std::sync::atomic::{AtomicUsize, Ordering};
1809
1810        let builder =
1811            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1812        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1813
1814        let rows_filtered = Arc::new(AtomicUsize::new(0));
1815        let rows_filtered_for_predicate = Arc::clone(&rows_filtered);
1816
1817        let row_filter_a = ArrowPredicateFn::new(
1818            ProjectionMask::columns(&schema_descr, ["a"]),
1819            move |batch: RecordBatch| {
1820                rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed);
1821                let scalar_175 = Int64Array::new_scalar(175);
1822                let column = batch.column(0).as_primitive::<Int64Type>();
1823                gt(column, &scalar_175)
1824            },
1825        );
1826
1827        // Use a small batch size so the row group is evaluated across
1828        // multiple predicate batches; that is the regime where Layer 2's
1829        // short-circuit saves predicate evaluation work. Matching rows are
1830        // 176..199 (24 rows); with batch_size = 10 those span batches 17, 18,
1831        // and 19 (rows 170..199). A limit of 10 should stop filter evaluation
1832        // in the middle of batch 18.
1833        let mut decoder = builder
1834            .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
1835            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1836            .with_batch_size(10)
1837            .with_limit(10)
1838            .build()
1839            .unwrap();
1840
1841        // First row group: filter columns fetch (predicate is evaluated here)
1842        let ranges = expect_needs_data(decoder.try_decode());
1843        push_ranges_to_decoder(&mut decoder, ranges);
1844
1845        // The first 10 matching rows come out: 176..185, column "a"
1846        let batch = expect_data(decoder.try_decode());
1847        let expected = TEST_BATCH.slice(176, 10).project(&[0]).unwrap();
1848        assert_eq!(batch, expected);
1849
1850        // no data for row group 1 should be requested — the limit
1851        // was satisfied by row group 0 and the `Start` state for row group 1
1852        // short-circuits to `Finished`.
1853        expect_finished(decoder.try_decode());
1854
1855        // Row 186 is the 11th match; the scan should stop no later than the
1856        // batch containing it (batch 18 of 10 rows = rows 180..189), so at
1857        // most 190 rows are evaluated.
1858        let evaluated = rows_filtered.load(Ordering::Relaxed);
1859        assert!(
1860            evaluated <= 190,
1861            "predicate evaluated {evaluated} rows; expected ≤ 190 (stop within batch containing 11th match)"
1862        );
1863    }
1864
1865    /// Once the limit has been satisfied by a prior row group, subsequent
1866    /// row groups should be skipped entirely — no data request for their
1867    /// filter columns.
1868    #[test]
1869    fn test_decoder_filter_with_limit_skips_later_row_groups() {
1870        let builder =
1871            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1872        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1873
1874        // `a > 175` matches rows 176..199 in row group 0 (24 matches) and
1875        // 200..399 in row group 1 (200 matches). With limit = 5, all matches
1876        // should come from row group 0.
1877        let row_filter_a = ArrowPredicateFn::new(
1878            ProjectionMask::columns(&schema_descr, ["a"]),
1879            |batch: RecordBatch| {
1880                let scalar_175 = Int64Array::new_scalar(175);
1881                let column = batch.column(0).as_primitive::<Int64Type>();
1882                gt(column, &scalar_175)
1883            },
1884        );
1885
1886        let mut decoder = builder
1887            .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
1888            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1889            .with_limit(5)
1890            .build()
1891            .unwrap();
1892
1893        // Row group 0: fetch filter pages
1894        let ranges = expect_needs_data(decoder.try_decode());
1895        push_ranges_to_decoder(&mut decoder, ranges);
1896
1897        // First 5 matches: 176..180
1898        let batch = expect_data(decoder.try_decode());
1899        let expected = TEST_BATCH.slice(176, 5).project(&[0]).unwrap();
1900        assert_eq!(batch, expected);
1901
1902        // Row group 1 must NOT request data — the limit is already satisfied
1903        // so `Start` in row group 1 short-circuits to `Finished`.
1904        expect_finished(decoder.try_decode());
1905    }
1906
1907    /// The predicate short-circuit must account for `self.offset` as well as
1908    /// `self.limit`. The post-predicate `with_offset` step skips that many
1909    /// already-selected rows before `with_limit` counts output rows — so the
1910    /// predicate must retain at least `offset + limit` matches. Without the
1911    /// fix, Layer 2 caps at just `limit` and the later `with_offset` consumes
1912    /// all of them, producing 0 rows instead of `limit`.
1913    ///
1914    /// `a > 175` matches rows 176..199 in row group 0 (24 matches). With
1915    /// `offset = 10, limit = 5`, the expected output is rows 186..190 (the
1916    /// 11th through 15th matches).
1917    #[test]
1918    fn test_decoder_filter_with_offset_and_limit() {
1919        let builder =
1920            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1921        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1922
1923        let row_filter_a = ArrowPredicateFn::new(
1924            ProjectionMask::columns(&schema_descr, ["a"]),
1925            |batch: RecordBatch| {
1926                let scalar_175 = Int64Array::new_scalar(175);
1927                let column = batch.column(0).as_primitive::<Int64Type>();
1928                gt(column, &scalar_175)
1929            },
1930        );
1931
1932        let mut decoder = builder
1933            .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
1934            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
1935            .with_offset(10)
1936            .with_limit(5)
1937            .build()
1938            .unwrap();
1939
1940        let ranges = expect_needs_data(decoder.try_decode());
1941        push_ranges_to_decoder(&mut decoder, ranges);
1942
1943        let batch = expect_data(decoder.try_decode());
1944        let expected = TEST_BATCH.slice(186, 5).project(&[0]).unwrap();
1945        assert_eq!(batch, expected);
1946
1947        expect_finished(decoder.try_decode());
1948    }
1949
1950    /// The limit short-circuit must also be correct when the limited predicate
1951    /// is the last predicate in a multi-predicate chain.
1952    ///
1953    /// `a > 175` first narrows row group 0 to rows 176..199. The final
1954    /// predicate `b < 625` is then evaluated only over those 24 rows, all of
1955    /// which match. With `limit = 10`, the final output should still be rows
1956    /// 176..185, and the second predicate should stop before consuming all 24
1957    /// selected rows.
1958    #[test]
1959    fn test_decoder_multi_filters_with_limit() {
1960        use std::sync::atomic::{AtomicUsize, Ordering};
1961
1962        let builder =
1963            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
1964        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
1965
1966        let first_predicate_rows = Arc::new(AtomicUsize::new(0));
1967        let second_predicate_rows = Arc::new(AtomicUsize::new(0));
1968
1969        let first_predicate_rows_for_filter = Arc::clone(&first_predicate_rows);
1970        let row_filter_a = ArrowPredicateFn::new(
1971            ProjectionMask::columns(&schema_descr, ["a"]),
1972            move |batch: RecordBatch| {
1973                first_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed);
1974                let scalar_175 = Int64Array::new_scalar(175);
1975                let column = batch.column(0).as_primitive::<Int64Type>();
1976                gt(column, &scalar_175)
1977            },
1978        );
1979
1980        let second_predicate_rows_for_filter = Arc::clone(&second_predicate_rows);
1981        let row_filter_b = ArrowPredicateFn::new(
1982            ProjectionMask::columns(&schema_descr, ["b"]),
1983            move |batch: RecordBatch| {
1984                second_predicate_rows_for_filter.fetch_add(batch.num_rows(), Ordering::Relaxed);
1985                let scalar_625 = Int64Array::new_scalar(625);
1986                let column = batch.column(0).as_primitive::<Int64Type>();
1987                lt(column, &scalar_625)
1988            },
1989        );
1990
1991        let mut decoder = builder
1992            .with_projection(ProjectionMask::columns(&schema_descr, ["c"]))
1993            .with_row_filter(RowFilter::new(vec![
1994                Box::new(row_filter_a),
1995                Box::new(row_filter_b),
1996            ]))
1997            .with_batch_size(10)
1998            .with_limit(10)
1999            .build()
2000            .unwrap();
2001
2002        // Row group 0, first predicate
2003        let ranges = expect_needs_data(decoder.try_decode());
2004        push_ranges_to_decoder(&mut decoder, ranges);
2005
2006        // Row group 0, second predicate
2007        let ranges = expect_needs_data(decoder.try_decode());
2008        push_ranges_to_decoder(&mut decoder, ranges);
2009
2010        // Final projected data
2011        let ranges = expect_needs_data(decoder.try_decode());
2012        push_ranges_to_decoder(&mut decoder, ranges);
2013
2014        let batch = expect_data(decoder.try_decode());
2015        let expected = TEST_BATCH.slice(176, 10).project(&[2]).unwrap();
2016        assert_eq!(batch, expected);
2017
2018        // The overall limit was satisfied by row group 0.
2019        expect_finished(decoder.try_decode());
2020
2021        assert_eq!(first_predicate_rows.load(Ordering::Relaxed), 200);
2022        assert!(
2023            second_predicate_rows.load(Ordering::Relaxed) < 24,
2024            "final predicate should short-circuit before consuming all 24 rows selected by the first predicate"
2025        );
2026    }
2027
2028    /// When a row selection already exists, limiting the predicate must still
2029    /// preserve alignment with that prior selection.
2030    ///
2031    /// The explicit selection narrows row group 0 to rows 150..199. Applying
2032    /// `a > 175` over that selection yields rows 176..199. With `limit = 10`,
2033    /// the decoder should emit rows 176..185 and stop without evaluating the
2034    /// remaining selected rows.
2035    #[test]
2036    fn test_decoder_filter_with_row_selection_and_limit() {
2037        use std::sync::atomic::{AtomicUsize, Ordering};
2038
2039        let builder =
2040            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2041        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2042
2043        let rows_filtered = Arc::new(AtomicUsize::new(0));
2044        let rows_filtered_for_predicate = Arc::clone(&rows_filtered);
2045
2046        let row_filter_a = ArrowPredicateFn::new(
2047            ProjectionMask::columns(&schema_descr, ["a"]),
2048            move |batch: RecordBatch| {
2049                rows_filtered_for_predicate.fetch_add(batch.num_rows(), Ordering::Relaxed);
2050                let scalar_175 = Int64Array::new_scalar(175);
2051                let column = batch.column(0).as_primitive::<Int64Type>();
2052                gt(column, &scalar_175)
2053            },
2054        );
2055
2056        let mut decoder = builder
2057            .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
2058            .with_row_selection(RowSelection::from(vec![
2059                RowSelector::skip(150),
2060                RowSelector::select(50),
2061            ]))
2062            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
2063            .with_batch_size(10)
2064            .with_limit(10)
2065            .build()
2066            .unwrap();
2067
2068        let ranges = expect_needs_data(decoder.try_decode());
2069        push_ranges_to_decoder(&mut decoder, ranges);
2070
2071        let batch = expect_data(decoder.try_decode());
2072        let expected = TEST_BATCH.slice(176, 10).project(&[0]).unwrap();
2073        assert_eq!(batch, expected);
2074
2075        expect_finished(decoder.try_decode());
2076
2077        assert!(
2078            rows_filtered.load(Ordering::Relaxed) < 50,
2079            "predicate should short-circuit before consuming all 50 rows from the explicit row selection"
2080        );
2081    }
2082
2083    #[test]
2084    fn test_decoder_offset_limit() {
2085        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2086            .unwrap()
2087            // skip entire first row group (200 rows) and first 25 rows of second row group
2088            .with_offset(225)
2089            // and limit to 20 rows
2090            .with_limit(20)
2091            .build()
2092            .unwrap();
2093
2094        // First row group should be skipped,
2095
2096        // Second row group
2097        let ranges = expect_needs_data(decoder.try_decode());
2098        push_ranges_to_decoder(&mut decoder, ranges);
2099
2100        // expect the first and only batch to be decoded
2101        let batch1 = expect_data(decoder.try_decode());
2102        let expected1 = TEST_BATCH.slice(225, 20);
2103        assert_eq!(batch1, expected1);
2104
2105        expect_finished(decoder.try_decode());
2106    }
2107
2108    #[test]
2109    fn test_decoder_try_next_reader_offset_limit() {
2110        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2111            .unwrap()
2112            .with_offset(225)
2113            .with_limit(20)
2114            .build()
2115            .unwrap();
2116
2117        let ranges = expect_needs_data(decoder.try_next_reader());
2118        push_ranges_to_decoder(&mut decoder, ranges);
2119
2120        let reader = expect_data(decoder.try_next_reader());
2121        let batches = reader
2122            .map(|batch| batch.expect("expected decoded batch"))
2123            .collect::<Vec<_>>();
2124        let output = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
2125        assert_eq!(output, TEST_BATCH.slice(225, 20));
2126
2127        expect_finished(decoder.try_next_reader());
2128    }
2129
2130    #[test]
2131    fn test_decoder_row_group_selection() {
2132        // take only the second row group
2133        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2134            .unwrap()
2135            .with_row_groups(vec![1])
2136            .build()
2137            .unwrap();
2138
2139        // First row group should be skipped,
2140
2141        // Second row group
2142        let ranges = expect_needs_data(decoder.try_decode());
2143        push_ranges_to_decoder(&mut decoder, ranges);
2144
2145        // expect the first and only batch to be decoded
2146        let batch1 = expect_data(decoder.try_decode());
2147        let expected1 = TEST_BATCH.slice(200, 200);
2148        assert_eq!(batch1, expected1);
2149
2150        expect_finished(decoder.try_decode());
2151    }
2152
2153    #[test]
2154    fn test_decoder_row_selection() {
2155        // take only the second row group
2156        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2157            .unwrap()
2158            .with_row_selection(RowSelection::from(vec![
2159                RowSelector::skip(225),  // skip first row group and 25 rows of second])
2160                RowSelector::select(20), // take 20 rows
2161            ]))
2162            .build()
2163            .unwrap();
2164
2165        // First row group should be skipped,
2166
2167        // Second row group
2168        let ranges = expect_needs_data(decoder.try_decode());
2169        push_ranges_to_decoder(&mut decoder, ranges);
2170
2171        // expect the first and only batch to be decoded
2172        let batch1 = expect_data(decoder.try_decode());
2173        let expected1 = TEST_BATCH.slice(225, 20);
2174        assert_eq!(batch1, expected1);
2175
2176        expect_finished(decoder.try_decode());
2177    }
2178
2179    #[test]
2180    fn test_decoder_row_group_local_selections() {
2181        let bitmap_selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(
2182            (0..45)
2183                .map(|row_idx| (25..45).contains(&row_idx))
2184                .collect::<Vec<_>>(),
2185        ));
2186        let rle_selection =
2187            RowSelection::from(vec![RowSelector::skip(190), RowSelector::select(10)]);
2188
2189        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2190            .unwrap()
2191            .with_row_group_selections(vec![
2192                RowGroupSelection::new(1, Some(bitmap_selection)),
2193                RowGroupSelection::new(0, Some(rle_selection)),
2194            ])
2195            .build()
2196            .unwrap();
2197        prefetch_test_file(&mut decoder);
2198
2199        // Row-group-local selections use local coordinates and preserve the
2200        // supplied row-group order. The bitmap is intentionally shorter than
2201        // RG1, so rows after its 45th local row are skipped.
2202        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(225, 20));
2203        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(190, 10));
2204        expect_finished(decoder.try_decode());
2205    }
2206
2207    #[test]
2208    fn test_row_group_local_selections_respect_offset_and_limit() {
2209        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2210            .unwrap()
2211            .with_row_group_selections(vec![
2212                RowGroupSelection::new(1, None),
2213                RowGroupSelection::new(0, None),
2214            ])
2215            .with_offset(195)
2216            .with_limit(10)
2217            .build()
2218            .unwrap();
2219        prefetch_test_file(&mut decoder);
2220
2221        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(395, 5));
2222        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(0, 5));
2223        expect_finished(decoder.try_decode());
2224    }
2225
2226    #[test]
2227    fn test_row_group_local_selections_allow_duplicate_row_groups() {
2228        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2229            .unwrap()
2230            .with_row_group_selections(vec![
2231                RowGroupSelection::new(0, None),
2232                RowGroupSelection::new(0, None),
2233            ])
2234            .build()
2235            .unwrap();
2236        prefetch_test_file(&mut decoder);
2237
2238        let expected = TEST_BATCH.slice(0, 200);
2239        assert_eq!(expect_data(decoder.try_decode()), expected);
2240        assert_eq!(expect_data(decoder.try_decode()), expected);
2241        expect_finished(decoder.try_decode());
2242    }
2243
2244    #[test]
2245    fn test_short_row_group_local_selection_skips_trailing_rows() {
2246        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2247            .unwrap()
2248            .with_row_group_selections(vec![RowGroupSelection::new(
2249                1,
2250                Some(RowSelection::from(vec![
2251                    RowSelector::skip(5),
2252                    RowSelector::select(3),
2253                ])),
2254            )])
2255            .build()
2256            .unwrap();
2257        prefetch_test_file(&mut decoder);
2258
2259        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(205, 3));
2260        expect_finished(decoder.try_decode());
2261    }
2262
2263    #[test]
2264    fn test_empty_row_group_local_selections_read_nothing() {
2265        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2266            .unwrap()
2267            .with_row_group_selections(vec![])
2268            .build()
2269            .unwrap();
2270
2271        expect_finished(decoder.try_decode());
2272    }
2273
2274    /// A `RowFilter` narrows each row group's local selection rather than
2275    /// replacing it: the predicate is evaluated against the locally selected
2276    /// rows, including when the selection is shorter than its row group and
2277    /// the row groups are supplied out of order.
2278    ///
2279    /// RG1 contributes local rows 10..20 ("a" 210..220), all of which pass
2280    /// `a > 195`; RG0 contributes local rows 190..200 ("a" 190..200), of which
2281    /// only 196..200 pass.
2282    fn row_group_local_selections_with_filter() -> ParquetPushDecoderBuilder {
2283        let builder =
2284            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2285        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2286
2287        // Values in column "a" range 0..399
2288        let row_filter_a = ArrowPredicateFn::new(
2289            ProjectionMask::columns(&schema_descr, ["a"]),
2290            |batch: RecordBatch| {
2291                let scalar_195 = Int64Array::new_scalar(195);
2292                let column = batch.column(0).as_primitive::<Int64Type>();
2293                gt(column, &scalar_195)
2294            },
2295        );
2296
2297        builder
2298            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
2299            .with_row_group_selections(vec![
2300                RowGroupSelection::new(
2301                    1,
2302                    Some(RowSelection::from(vec![
2303                        RowSelector::skip(10),
2304                        RowSelector::select(10),
2305                    ])),
2306                ),
2307                RowGroupSelection::new(
2308                    0,
2309                    Some(RowSelection::from(vec![
2310                        RowSelector::skip(190),
2311                        RowSelector::select(10),
2312                    ])),
2313                ),
2314            ])
2315    }
2316
2317    #[test]
2318    fn test_row_group_local_selections_with_row_filter() {
2319        let mut decoder = row_group_local_selections_with_filter().build().unwrap();
2320        prefetch_test_file(&mut decoder);
2321
2322        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(210, 10));
2323        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(196, 4));
2324        expect_finished(decoder.try_decode());
2325    }
2326
2327    /// `into_builder` mid-scan must carry both the row filter and the
2328    /// still-local selection of the not-yet-decoded row group, so the rebuilt
2329    /// decoder produces exactly what an uninterrupted scan would have.
2330    #[test]
2331    fn test_into_builder_preserves_local_selections_with_row_filter() {
2332        let mut decoder = row_group_local_selections_with_filter().build().unwrap();
2333        prefetch_test_file(&mut decoder);
2334
2335        let reader1 = expect_data(decoder.try_next_reader());
2336        let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2337        let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2338        assert_eq!(batch1, TEST_BATCH.slice(210, 10));
2339
2340        // Only RG0 and its local `skip(190) + select(10)` remain.
2341        assert!(decoder.is_at_row_group_boundary());
2342        assert_eq!(decoder.row_groups_remaining(), 1);
2343        let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2344
2345        let reader0 = expect_data(decoder.try_next_reader());
2346        let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2347        let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2348        assert_eq!(batch0, TEST_BATCH.slice(196, 4));
2349        expect_finished(decoder.try_next_reader());
2350    }
2351
2352    #[test]
2353    fn test_row_group_local_selection_replaces_previous_configuration() {
2354        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2355            .unwrap()
2356            .with_row_group_selections(vec![RowGroupSelection::new(0, None)])
2357            .with_row_group_selections(vec![RowGroupSelection::new(1, None)])
2358            .build()
2359            .unwrap();
2360        prefetch_test_file(&mut decoder);
2361
2362        // The second call replaces the first. `None` reads RG1 in full, and
2363        // the omitted RG0 is skipped.
2364        assert_eq!(
2365            expect_data(decoder.try_decode()),
2366            TEST_BATCH.slice(200, 200)
2367        );
2368        expect_finished(decoder.try_decode());
2369    }
2370
2371    #[test]
2372    fn test_legacy_row_groups_and_row_selection_remain_composable() {
2373        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2374            .unwrap()
2375            .with_row_groups(vec![1])
2376            .with_row_selection(RowSelection::from(vec![
2377                RowSelector::skip(25),
2378                RowSelector::select(20),
2379            ]))
2380            .build()
2381            .unwrap();
2382        prefetch_test_file(&mut decoder);
2383
2384        assert_eq!(expect_data(decoder.try_decode()), TEST_BATCH.slice(225, 20));
2385        expect_finished(decoder.try_decode());
2386    }
2387
2388    #[test]
2389    fn test_row_group_local_selection_is_mutually_exclusive_with_legacy_configuration() {
2390        let metadata = test_file_parquet_metadata();
2391        let new_builder =
2392            || ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata)).unwrap();
2393        let local_selection = || vec![RowGroupSelection::new(0, None)];
2394        let global_selection = || RowSelection::from(vec![RowSelector::select(1)]);
2395        let assert_conflict = |builder: ParquetPushDecoderBuilder| {
2396            let error = builder.build().unwrap_err();
2397            assert_eq!(
2398                error.to_string(),
2399                "Parquet error: with_row_group_selections cannot be combined with with_row_groups or with_row_selection"
2400            );
2401        };
2402
2403        assert_conflict(
2404            new_builder()
2405                .with_row_groups(vec![0])
2406                .with_row_group_selections(local_selection()),
2407        );
2408        assert_conflict(
2409            new_builder()
2410                .with_row_group_selections(local_selection())
2411                .with_row_groups(vec![0]),
2412        );
2413        assert_conflict(
2414            new_builder()
2415                .with_row_selection(global_selection())
2416                .with_row_group_selections(local_selection()),
2417        );
2418        assert_conflict(
2419            new_builder()
2420                .with_row_group_selections(local_selection())
2421                .with_row_selection(global_selection()),
2422        );
2423    }
2424
2425    #[test]
2426    fn test_row_group_local_selection_validates_row_group_and_length_at_build() {
2427        let metadata = test_file_parquet_metadata();
2428
2429        let error = ParquetPushDecoderBuilder::try_new_decoder(Arc::clone(&metadata))
2430            .unwrap()
2431            .with_row_group_selections(vec![RowGroupSelection::new(
2432                0,
2433                Some(RowSelection::from(vec![RowSelector::select(201)])),
2434            )])
2435            .build()
2436            .unwrap_err();
2437        assert!(
2438            error.to_string().contains(
2439                "Row selection for row group 0 contains 201 rows, but the row group has 200"
2440            ),
2441            "unexpected error: {error}"
2442        );
2443
2444        let error = ParquetPushDecoderBuilder::try_new_decoder(metadata)
2445            .unwrap()
2446            .with_row_group_selections(vec![RowGroupSelection::new(2, None)])
2447            .build()
2448            .unwrap_err();
2449        assert!(
2450            error
2451                .to_string()
2452                .contains("Row group index 2 out of bounds for file with 2 row groups"),
2453            "unexpected error: {error}"
2454        );
2455    }
2456
2457    /// `peek_next_row_group` reports the index of the row group the
2458    /// next `try_next_reader` call will hand back, matching the
2459    /// frontier's internal skip logic.
2460    #[test]
2461    fn test_peek_next_row_group_basic() {
2462        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2463            .unwrap()
2464            .build()
2465            .unwrap();
2466
2467        // Two row groups (0, 1). At boundary before any read, peek should
2468        // see RG 0.
2469        assert_eq!(decoder.peek_next_row_group().unwrap(), Some(0));
2470        assert!(decoder.is_at_row_group_boundary());
2471
2472        let ranges = expect_needs_data(decoder.try_next_reader());
2473        push_ranges_to_decoder(&mut decoder, ranges);
2474        let reader = expect_data(decoder.try_next_reader());
2475        // Once the reader for RG 0 has been handed off, the decoder is
2476        // back at a boundary waiting for RG 1 — peek must reflect that
2477        // (the active reader lives outside the decoder).
2478        assert!(decoder.is_at_row_group_boundary());
2479        assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2480
2481        // Drain RG 0's reader and consume RG 1.
2482        for batch in reader {
2483            let _ = batch.unwrap();
2484        }
2485        let ranges = expect_needs_data(decoder.try_next_reader());
2486        push_ranges_to_decoder(&mut decoder, ranges);
2487        let reader = expect_data(decoder.try_next_reader());
2488        for batch in reader {
2489            let _ = batch.unwrap();
2490        }
2491
2492        // No row groups left.
2493        assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2494    }
2495
2496    /// `peek_next_row_group` honors `with_row_groups` — restricting the
2497    /// scan to a single row group means peek reports only that one and
2498    /// then `None`.
2499    #[test]
2500    fn test_peek_next_row_group_respects_with_row_groups() {
2501        let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2502            .unwrap()
2503            .with_row_groups(vec![1])
2504            .build()
2505            .unwrap();
2506
2507        assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2508    }
2509
2510    /// When a row-selection segment leaves the next row group with zero
2511    /// selected rows, `peek_next_row_group` mirrors
2512    /// `next_readable_row_group`'s skip: it returns the *following*
2513    /// row group instead of the empty one.
2514    #[test]
2515    fn test_peek_next_row_group_skips_empty_selection() {
2516        // Each row group has 200 rows. Skip all 200 of RG 0 plus 50 of
2517        // RG 1; the next reader will be for RG 1, not RG 0.
2518        let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2519            .unwrap()
2520            .with_row_selection(RowSelection::from(vec![
2521                RowSelector::skip(250),
2522                RowSelector::select(100),
2523            ]))
2524            .build()
2525            .unwrap();
2526
2527        assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2528    }
2529
2530    /// `peek_next_row_group` returns `None` on a finished decoder.
2531    #[test]
2532    fn test_peek_next_row_group_finished() {
2533        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2534            .unwrap()
2535            .with_row_groups(vec![])
2536            .build()
2537            .unwrap();
2538
2539        // No row groups requested ⇒ already finished, no peek.
2540        expect_finished(decoder.try_next_reader());
2541        assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2542    }
2543
2544    /// `peek_next_row_group` mirrors `next_readable_row_group`'s
2545    /// budget-skipping logic: with an `OFFSET` that consumes the
2546    /// entire first row group, peek must return the *following*
2547    /// row group, not RG 0 (no predicates, so budget skips apply).
2548    #[test]
2549    fn test_peek_next_row_group_skips_for_budget() {
2550        // Each row group has 200 rows. OFFSET 200 drains RG 0 entirely
2551        // and lands the decoder's first emitted reader at RG 1.
2552        let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2553            .unwrap()
2554            .with_offset(200)
2555            .build()
2556            .unwrap();
2557
2558        assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2559    }
2560
2561    /// OFFSET larger than every remaining row group's row count
2562    /// exhausts the budget, so peek should report `None` — matching
2563    /// `next_readable_row_group`'s behavior of producing `Finished`.
2564    #[test]
2565    fn test_peek_next_row_group_budget_drains_all() {
2566        // 2 row groups × 200 rows = 400 total. OFFSET 500 cannot land
2567        // anywhere — every RG is skipped under the budget.
2568        let decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2569            .unwrap()
2570            .with_offset(500)
2571            .build()
2572            .unwrap();
2573
2574        assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2575    }
2576
2577    /// `peek_next_row_group` is safe to call while a row group's reader
2578    /// is still being drained via `try_decode`. In `DecodingRowGroup`
2579    /// state it must report the row group `try_next_reader` will yield
2580    /// a reader for *after* the active one — never `None` just because
2581    /// the decoder is mid-row-group.
2582    #[test]
2583    fn test_peek_next_row_group_during_decoding_row_group() {
2584        // Two row groups (200 rows each), small batch size so the
2585        // decoder stays in `DecodingRowGroup` across multiple
2586        // `try_decode` calls within RG 0.
2587        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2588            .unwrap()
2589            .with_batch_size(100)
2590            .build()
2591            .unwrap();
2592        decoder
2593            .push_range(test_file_range(), TEST_FILE_DATA.clone())
2594            .unwrap();
2595
2596        // First batch of RG 0 — after this, the decoder is in
2597        // `DecodingRowGroup` state with the reader retained internally
2598        // and one more 100-row batch still owed for RG 0. Peek must
2599        // therefore look past the active RG to RG 1.
2600        let _batch0 = expect_data(decoder.try_decode());
2601        assert!(!decoder.is_at_row_group_boundary());
2602        assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2603
2604        // Drain the rest of RG 0; peek still reports RG 1.
2605        let _batch1 = expect_data(decoder.try_decode());
2606        assert_eq!(decoder.peek_next_row_group().unwrap(), Some(1));
2607
2608        // Move into RG 1 — peek now sees no further row groups.
2609        let _batch2 = expect_data(decoder.try_decode());
2610        assert!(!decoder.is_at_row_group_boundary());
2611        assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2612
2613        let _batch3 = expect_data(decoder.try_decode());
2614        expect_finished(decoder.try_decode());
2615    }
2616
2617    /// Peeking is a read-only operation: calling it repeatedly between
2618    /// `try_next_reader` calls must never change which row group the
2619    /// reader path actually produces. Drives the decoder all the way
2620    /// through and asserts each peek/read pair agrees.
2621    #[test]
2622    fn test_peek_next_row_group_does_not_mutate_state() {
2623        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2624            .unwrap()
2625            .build()
2626            .unwrap();
2627
2628        // Two row groups expected — drive both, asserting peek/read agree.
2629        for expected_rg in [0usize, 1usize] {
2630            assert!(decoder.is_at_row_group_boundary());
2631
2632            // Multiple peeks before reading must all agree, and must not
2633            // disturb the upcoming read.
2634            let first_peek = decoder.peek_next_row_group().unwrap();
2635            let second_peek = decoder.peek_next_row_group().unwrap();
2636            let third_peek = decoder.peek_next_row_group().unwrap();
2637            assert_eq!(first_peek, Some(expected_rg));
2638            assert_eq!(first_peek, second_peek);
2639            assert_eq!(second_peek, third_peek);
2640
2641            // Now read for real and confirm the decoder hands back exactly
2642            // what peek promised.
2643            let ranges = expect_needs_data(decoder.try_next_reader());
2644            push_ranges_to_decoder(&mut decoder, ranges);
2645            let reader = expect_data(decoder.try_next_reader());
2646            for batch in reader {
2647                let _ = batch.unwrap();
2648            }
2649        }
2650
2651        // Decoder is drained. Peek must agree with `try_next_reader`'s
2652        // terminal state.
2653        assert_eq!(decoder.peek_next_row_group().unwrap(), None);
2654        expect_finished(decoder.try_next_reader());
2655    }
2656
2657    /// `into_builder` between row groups recovers a builder for the
2658    /// not-yet-decoded row groups; rebuilding it with a new row filter
2659    /// applies that filter to the subsequent row groups while leaving the
2660    /// already-decoded row group's results untouched.
2661    ///
2662    /// See the "Adaptive scans" section of [`ParquetPushDecoderBuilder`] for
2663    /// the high-level overview.
2664    #[test]
2665    fn test_into_builder_installs_filter_between_row_groups() {
2666        let schema_descr = test_file_parquet_metadata()
2667            .file_metadata()
2668            .schema_descr_ptr();
2669        let mut decoder = prefetched_decoder(1024);
2670
2671        // Reader for row group 0 — no filter.
2672        let reader0 = expect_data(decoder.try_next_reader());
2673        let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2674        let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2675        assert_eq!(batch0, TEST_BATCH.slice(0, 200));
2676
2677        // We're between row groups now. Rebuild with a filter on column "a".
2678        assert!(decoder.is_at_row_group_boundary());
2679        assert_eq!(decoder.row_groups_remaining(), 1);
2680        let filter =
2681            ArrowPredicateFn::new(ProjectionMask::columns(&schema_descr, ["a"]), |batch| {
2682                gt(batch.column(0), &Int64Array::new_scalar(250))
2683            });
2684        let mut decoder = decoder
2685            .into_builder()
2686            .unwrap()
2687            .with_row_filter(RowFilter::new(vec![Box::new(filter)]))
2688            .build()
2689            .unwrap();
2690
2691        // Reader for row group 1 — filter applied. The rebuilt decoder kept
2692        // the buffered bytes (see `test_into_builder_preserves_buffered_bytes`)
2693        // so no data needs to be re-supplied. Column "a" in RG1 has values
2694        // 200..399; `a > 250` keeps 251..399 = 149 rows.
2695        let reader1 = expect_data(decoder.try_next_reader());
2696        let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2697        let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2698        assert_eq!(batch1, TEST_BATCH.slice(251, 149));
2699        expect_finished(decoder.try_next_reader());
2700    }
2701
2702    /// `into_builder` is rejected while a row group's reader is being
2703    /// drained (`DecodingRowGroup`); the error points at
2704    /// `is_at_row_group_boundary`.
2705    #[test]
2706    fn test_into_builder_rejected_mid_row_group() {
2707        let mut decoder = prefetched_decoder(50);
2708
2709        // Decode one batch to land mid-row-group, inside `DecodingRowGroup`
2710        // with an active reader — not a boundary.
2711        expect_data(decoder.try_decode());
2712        assert!(!decoder.is_at_row_group_boundary());
2713
2714        let err = decoder.into_builder().unwrap_err();
2715        let err_msg = format!("{err}");
2716        assert!(
2717            err_msg.contains("is_at_row_group_boundary"),
2718            "unexpected error: {err_msg}"
2719        );
2720    }
2721
2722    /// `into_builder` is rejected once the decoder has finished.
2723    #[test]
2724    fn test_into_builder_rejected_on_finished_decoder() {
2725        let mut decoder = prefetched_decoder(1024);
2726        expect_data(decoder.try_decode());
2727        expect_data(decoder.try_decode());
2728        expect_finished(decoder.try_decode());
2729        assert!(!decoder.is_at_row_group_boundary());
2730
2731        let err = decoder.into_builder().unwrap_err();
2732        assert!(
2733            format!("{err}").contains("finished"),
2734            "unexpected error: {err}"
2735        );
2736    }
2737
2738    /// `try_next_reader` hands the active reader off to the caller and
2739    /// transitions the decoder back to `ReadingRowGroup` — so the caller
2740    /// can call `into_builder` even while still holding the returned
2741    /// reader. (The handed-off reader has no link back to the decoder's
2742    /// projection/filter; it has its own `ArrayReader` and `ReadPlan`.)
2743    #[test]
2744    fn test_into_builder_allowed_while_iterating_handed_off_reader() {
2745        let mut decoder = prefetched_decoder(1024);
2746
2747        let reader0 = expect_data(decoder.try_next_reader());
2748        // Decoder no longer owns the reader, so it considers itself
2749        // "between row groups".
2750        assert!(decoder.is_at_row_group_boundary());
2751        // Recovering the builder consumes the decoder but leaves `reader0`
2752        // valid: iterating it is independent of the decoder's state.
2753        let _builder = decoder.into_builder().unwrap();
2754        let batches: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2755        let batch0 = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
2756        assert_eq!(batch0, TEST_BATCH.slice(0, 200));
2757    }
2758
2759    /// `into_builder` recovers a builder for the *remaining* row groups and
2760    /// carries the not-yet-consumed offset/limit budget, so a rebuilt
2761    /// decoder resumes where the original left off rather than restarting.
2762    #[test]
2763    fn test_into_builder_resumes_remaining_budget() {
2764        // limit = 250 spans both 200-row row groups: all 200 rows of RG0
2765        // plus the first 50 rows of RG1.
2766        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2767            .unwrap()
2768            .with_batch_size(1024)
2769            .with_limit(250)
2770            .build()
2771            .unwrap();
2772        prefetch_test_file(&mut decoder);
2773
2774        // RG0 contributes all 200 of its rows.
2775        let reader0 = expect_data(decoder.try_next_reader());
2776        let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2777        let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2778        assert_eq!(batch0, TEST_BATCH.slice(0, 200));
2779
2780        // Rebuild without changing anything: the remaining 50-row limit and
2781        // the not-yet-decoded RG1 must carry through (as do the buffers, so
2782        // no data needs re-supplying).
2783        assert!(decoder.is_at_row_group_boundary());
2784        let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2785
2786        let reader1 = expect_data(decoder.try_next_reader());
2787        let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2788        let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2789        // Only the first 50 rows of RG1 (200..249) — the rest of the limit.
2790        assert_eq!(batch1, TEST_BATCH.slice(200, 50));
2791        expect_finished(decoder.try_next_reader());
2792    }
2793
2794    /// `into_builder` carries the decoder's buffered bytes across the
2795    /// rebuild: the rebuilt decoder keeps them and does not re-request data
2796    /// it already holds.
2797    #[test]
2798    fn test_into_builder_preserves_buffered_bytes() {
2799        let mut decoder = prefetched_decoder(1024);
2800        assert_eq!(decoder.buffered_bytes(), test_file_len());
2801
2802        // Drain RG0.
2803        let reader0 = expect_data(decoder.try_next_reader());
2804        let _: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2805        // RG1's bytes are still staged inside the decoder.
2806        let buffered = decoder.buffered_bytes();
2807        assert!(buffered > 0);
2808
2809        // Rebuilding via into_builder keeps the staged bytes that RG1 reads,
2810        // and releases the other bytes of the prefetched file.
2811        let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2812        assert!(decoder.buffered_bytes() < buffered);
2813        assert_eq!(decoder.buffered_bytes(), test_file_row_group_bytes(1));
2814
2815        // RG1's bytes are already buffered, so it decodes without a
2816        // `NeedsData` round-trip.
2817        let reader1 = expect_data(decoder.try_next_reader());
2818        let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2819        let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2820        assert_eq!(batch1, TEST_BATCH.slice(200, 200));
2821        expect_finished(decoder.try_next_reader());
2822    }
2823
2824    /// A rebuilt decoder releases the bytes of a row group that it skips,
2825    /// and keeps the bytes of the columns that it reads.
2826    #[test]
2827    fn test_into_builder_releases_bytes_of_skipped_row_groups() {
2828        let builder =
2829            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2830        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2831        let mut decoder = builder
2832            .with_row_groups(vec![1])
2833            .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b"]))
2834            .build()
2835            .unwrap();
2836        prefetch_test_file(&mut decoder);
2837        assert!(decoder.is_at_row_group_boundary());
2838
2839        // Skip RG1 and read RG0, with column "c" too.
2840        let mut decoder = decoder
2841            .into_builder()
2842            .unwrap()
2843            .with_row_groups(vec![0])
2844            .with_projection(ProjectionMask::all())
2845            .build()
2846            .unwrap();
2847        assert_eq!(decoder.buffered_bytes(), test_file_row_group_bytes(0));
2848
2849        let reader = expect_data(decoder.try_next_reader());
2850        let batches: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
2851        let batch = concat_batches(&TEST_BATCH.schema(), &batches).unwrap();
2852        assert_eq!(batch, TEST_BATCH.slice(0, 200));
2853        expect_finished(decoder.try_next_reader());
2854    }
2855
2856    /// A rebuilt decoder keeps the bytes of a column that only a predicate
2857    /// reads.
2858    #[test]
2859    fn test_into_builder_keeps_bytes_of_predicate_columns() {
2860        let builder =
2861            ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata()).unwrap();
2862        let schema_descr = builder.metadata().file_metadata().schema_descr_ptr();
2863        let row_filter_a = ArrowPredicateFn::new(
2864            ProjectionMask::columns(&schema_descr, ["a"]),
2865            |batch: RecordBatch| {
2866                let scalar_250 = Int64Array::new_scalar(250);
2867                let column = batch.column(0).as_primitive::<Int64Type>();
2868                gt(column, &scalar_250)
2869            },
2870        );
2871        let mut decoder = builder
2872            .with_row_groups(vec![1])
2873            .with_projection(ProjectionMask::columns(&schema_descr, ["b"]))
2874            .with_row_filter(RowFilter::new(vec![Box::new(row_filter_a)]))
2875            .build()
2876            .unwrap();
2877        prefetch_test_file(&mut decoder);
2878
2879        // The rebuilt decoder keeps the bytes of "a" (the predicate) and "b"
2880        // (the output) of RG1, and releases the bytes of "c".
2881        let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2882        let column_bytes = |column_idx: usize| {
2883            let range =
2884                column_chunk_range(test_file_parquet_metadata().row_group(1).column(column_idx));
2885            range.end - range.start
2886        };
2887        assert_eq!(decoder.buffered_bytes(), column_bytes(0) + column_bytes(1));
2888
2889        let reader = expect_data(decoder.try_next_reader());
2890        let batches: Vec<_> = reader.collect::<Result<_, _>>().unwrap();
2891        let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
2892        assert_eq!(batch, TEST_BATCH.slice(251, 149).project(&[1]).unwrap());
2893        expect_finished(decoder.try_next_reader());
2894    }
2895
2896    #[test]
2897    fn test_into_builder_preserves_remaining_row_group_local_selections() {
2898        let bitmap_selection = RowSelection::from_boolean_buffer(BooleanBuffer::from(
2899            (0..45)
2900                .map(|row_idx| (25..45).contains(&row_idx))
2901                .collect::<Vec<_>>(),
2902        ));
2903        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
2904            .unwrap()
2905            .with_row_group_selections(vec![
2906                RowGroupSelection::new(
2907                    0,
2908                    Some(RowSelection::from(vec![
2909                        RowSelector::skip(190),
2910                        RowSelector::select(10),
2911                    ])),
2912                ),
2913                RowGroupSelection::new(1, Some(bitmap_selection)),
2914            ])
2915            .build()
2916            .unwrap();
2917        prefetch_test_file(&mut decoder);
2918
2919        let reader0 = expect_data(decoder.try_next_reader());
2920        let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2921        let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
2922        assert_eq!(batch0, TEST_BATCH.slice(190, 10));
2923
2924        // Rebuilding must carry only RG1 and its still-local bitmap selection.
2925        assert!(decoder.is_at_row_group_boundary());
2926        assert_eq!(decoder.row_groups_remaining(), 1);
2927        let mut decoder = decoder.into_builder().unwrap().build().unwrap();
2928
2929        let reader1 = expect_data(decoder.try_next_reader());
2930        let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2931        let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2932        assert_eq!(batch1, TEST_BATCH.slice(225, 20));
2933        expect_finished(decoder.try_next_reader());
2934    }
2935
2936    /// Drive the decoder incrementally. Start with a narrow projection,
2937    /// drain RG0, then `into_builder` and widen the projection to all three
2938    /// columns. The rebuilt decoder's `NeedsData` for RG1 must request
2939    /// bytes for *all three* columns, not just the originally-projected
2940    /// "a". The expected ranges are hardcoded because `TEST_BATCH` and the
2941    /// writer settings are static; this pins the layout cleanly without a
2942    /// parallel reference decoder.
2943    #[test]
2944    fn test_into_builder_expand_projection_requests_new_bytes() {
2945        let metadata = test_file_parquet_metadata();
2946        let schema_descr = metadata.file_metadata().schema_descr_ptr();
2947
2948        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata)
2949            .unwrap()
2950            .with_batch_size(1024)
2951            .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
2952            .build()
2953            .unwrap();
2954
2955        // RG0: incrementally satisfy the narrow request — a single
2956        // contiguous range for the "a"-only projection.
2957        let ranges_rg0 = expect_needs_data(decoder.try_next_reader());
2958        assert_eq!(ranges_rg0, vec![4..1860]);
2959        push_ranges_to_decoder(&mut decoder, ranges_rg0);
2960
2961        let reader0 = expect_data(decoder.try_next_reader());
2962        let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
2963        let batch0 = concat_batches(&batches0[0].schema(), &batches0).unwrap();
2964        assert_eq!(batch0, TEST_BATCH.slice(0, 200).project(&[0]).unwrap());
2965
2966        // Widen the projection at the boundary.
2967        assert!(decoder.is_at_row_group_boundary());
2968        let mut decoder = decoder
2969            .into_builder()
2970            .unwrap()
2971            .with_projection(ProjectionMask::columns(&schema_descr, ["a", "b", "c"]))
2972            .build()
2973            .unwrap();
2974
2975        // RG1 now requests "a", "b", and "c" column chunks: ~1.8KiB each
2976        // for "a" and "b", ~7.5KiB for the StringView column "c".
2977        let ranges_rg1 = expect_needs_data(decoder.try_next_reader());
2978        assert_eq!(ranges_rg1, vec![11062..12918, 12918..14774, 14774..22230]);
2979        push_ranges_to_decoder(&mut decoder, ranges_rg1);
2980
2981        let reader1 = expect_data(decoder.try_next_reader());
2982        let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
2983        let batch1 = concat_batches(&TEST_BATCH.schema(), &batches1).unwrap();
2984        assert_eq!(batch1, TEST_BATCH.slice(200, 200));
2985        expect_finished(decoder.try_next_reader());
2986    }
2987
2988    /// Mirror of [`test_into_builder_expand_projection_requests_new_bytes`]:
2989    /// start with the full projection, drain RG0, then `into_builder` and
2990    /// narrow the projection to just column "a". RG1's `NeedsData` must
2991    /// request only the single "a" column-chunk range, not the three a
2992    /// wide projection would.
2993    #[test]
2994    fn test_into_builder_narrow_projection_requests_fewer_bytes() {
2995        let metadata = test_file_parquet_metadata();
2996        let schema_descr = metadata.file_metadata().schema_descr_ptr();
2997
2998        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(metadata)
2999            .unwrap()
3000            .with_batch_size(1024)
3001            .build()
3002            .unwrap();
3003
3004        // RG0 with the default (full) projection — three column ranges.
3005        let ranges_rg0 = expect_needs_data(decoder.try_next_reader());
3006        assert_eq!(ranges_rg0, vec![4..1860, 1860..3716, 3716..11062]);
3007        push_ranges_to_decoder(&mut decoder, ranges_rg0);
3008
3009        let reader0 = expect_data(decoder.try_next_reader());
3010        let batches0: Vec<_> = reader0.collect::<Result<_, _>>().unwrap();
3011        let batch0 = concat_batches(&TEST_BATCH.schema(), &batches0).unwrap();
3012        assert_eq!(batch0, TEST_BATCH.slice(0, 200));
3013
3014        // Narrow the projection at the boundary.
3015        assert!(decoder.is_at_row_group_boundary());
3016        let mut decoder = decoder
3017            .into_builder()
3018            .unwrap()
3019            .with_projection(ProjectionMask::columns(&schema_descr, ["a"]))
3020            .build()
3021            .unwrap();
3022
3023        // RG1 now requests column "a" only — a single 1856-byte range.
3024        let ranges_rg1 = expect_needs_data(decoder.try_next_reader());
3025        assert_eq!(ranges_rg1, vec![11062..12918]);
3026        push_ranges_to_decoder(&mut decoder, ranges_rg1);
3027
3028        let reader1 = expect_data(decoder.try_next_reader());
3029        let batches1: Vec<_> = reader1.collect::<Result<_, _>>().unwrap();
3030        let batch1 = concat_batches(&batches1[0].schema(), &batches1).unwrap();
3031        assert_eq!(batch1, TEST_BATCH.slice(200, 200).project(&[0]).unwrap());
3032        expect_finished(decoder.try_next_reader());
3033    }
3034
3035    /// Returns a batch with 400 rows, with 3 columns: "a", "b", "c"
3036    ///
3037    /// Note c is a different types (so the data page sizes will be different)
3038    static TEST_BATCH: LazyLock<RecordBatch> = LazyLock::new(|| {
3039        let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400));
3040        let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800));
3041        let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| {
3042            if i % 2 == 0 {
3043                format!("string_{i}")
3044            } else {
3045                format!("A string larger than 12 bytes and thus not inlined {i}")
3046            }
3047        })));
3048
3049        RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
3050    });
3051
3052    /// Create a parquet file in memory for testing.
3053    ///
3054    /// See [`TEST_BATCH`] for the data in the file.
3055    ///
3056    /// Each column is written in 4 data pages, each with 100 rows, across 2
3057    /// row groups. Each column in each row group has two data pages.
3058    ///
3059    /// The data is split across row groups like this
3060    ///
3061    /// Column |   Values                | Data Page | Row Group
3062    /// -------|------------------------|-----------|-----------
3063    /// a      | 0..99                  | 1         | 0
3064    /// a      | 100..199               | 2         | 0
3065    /// a      | 200..299               | 1         | 1
3066    /// a      | 300..399               | 2         | 1
3067    ///
3068    /// b      | 400..499               | 1         | 0
3069    /// b      | 500..599               | 2         | 0
3070    /// b      | 600..699               | 1         | 1
3071    /// b      | 700..799               | 2         | 1
3072    ///
3073    /// c      | "string_0".."string_99"        | 1         | 0
3074    /// c      | "string_100".."string_199"     | 2         | 0
3075    /// c      | "string_200".."string_299"     | 1         | 1
3076    /// c      | "string_300".."string_399"     | 2         | 1
3077    static TEST_FILE_DATA: LazyLock<Bytes> = LazyLock::new(|| {
3078        let input_batch = &TEST_BATCH;
3079        let mut output = Vec::new();
3080
3081        let writer_options = WriterProperties::builder()
3082            .set_max_row_group_row_count(Some(200))
3083            .set_data_page_row_count_limit(100)
3084            .build();
3085        let mut writer =
3086            ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap();
3087
3088        // since the limits are only enforced on batch boundaries, write the input
3089        // batch in chunks of 50
3090        let mut row_remain = input_batch.num_rows();
3091        while row_remain > 0 {
3092            let chunk_size = row_remain.min(50);
3093            let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size);
3094            writer.write(&chunk).unwrap();
3095            row_remain -= chunk_size;
3096        }
3097        writer.close().unwrap();
3098        Bytes::from(output)
3099    });
3100
3101    /// Return the length of [`TEST_FILE_DATA`], in bytes
3102    fn test_file_len() -> u64 {
3103        TEST_FILE_DATA.len() as u64
3104    }
3105
3106    /// Return a range that covers the entire [`TEST_FILE_DATA`]
3107    fn test_file_range() -> Range<u64> {
3108        0..test_file_len()
3109    }
3110
3111    /// Return a slice of the test file data from the given range
3112    pub fn test_file_slice(range: Range<u64>) -> Bytes {
3113        let start: usize = range.start.try_into().unwrap();
3114        let end: usize = range.end.try_into().unwrap();
3115        TEST_FILE_DATA.slice(start..end)
3116    }
3117
3118    /// return the metadata for the test file
3119    pub fn test_file_parquet_metadata() -> Arc<ParquetMetaData> {
3120        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap();
3121        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]);
3122        let metadata = metadata_decoder.try_decode().unwrap();
3123        let DecodeResult::Data(metadata) = metadata else {
3124            panic!("Expected metadata to be decoded successfully");
3125        };
3126        Arc::new(metadata)
3127    }
3128
3129    /// An offset index whose last page is longer than the column chunk gives
3130    /// an error, not a panic.
3131    #[test]
3132    fn test_decoder_offset_index_page_past_column_chunk() {
3133        let metadata = Arc::unwrap_or_clone(test_file_parquet_metadata_with_offset_index());
3134        let num_columns = metadata.row_group(0).num_columns();
3135        let mut page_index = PageIndexBuilder::new(metadata.num_row_groups(), num_columns);
3136        for row_group in 0..metadata.num_row_groups() {
3137            let row_group_index = metadata.page_index_for_row_group(row_group);
3138            for column in 0..num_columns {
3139                let mut offset_index = row_group_index.offset_index(column).unwrap().clone();
3140                if (row_group, column) == (0, 0) {
3141                    let last = offset_index.page_locations.last_mut().unwrap();
3142                    last.compressed_page_size += 1;
3143                }
3144                page_index.put_offset_index(offset_index, row_group, column);
3145            }
3146        }
3147        let metadata = metadata
3148            .into_builder()
3149            .set_page_index(Some(Arc::new(page_index.build())))
3150            .build();
3151
3152        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(Arc::new(metadata))
3153            .unwrap()
3154            .with_row_groups(vec![0])
3155            .build()
3156            .unwrap();
3157        let ranges = expect_needs_data(decoder.try_decode());
3158        push_ranges_to_decoder(&mut decoder, ranges);
3159        let err = loop {
3160            match decoder.try_decode() {
3161                Ok(DecodeResult::Data(_)) => {}
3162                Ok(other) => panic!("expected an error, got {other:?}"),
3163                Err(e) => break e.to_string(),
3164            }
3165        };
3166        assert!(err.contains("has 126 bytes, expected 127"), "{err}");
3167    }
3168
3169    /// return the metadata for the test file, including the offset index
3170    fn test_file_parquet_metadata_with_offset_index() -> Arc<ParquetMetaData> {
3171        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(test_file_len())
3172            .unwrap()
3173            .with_offset_index_policy(PageIndexPolicy::Required);
3174        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]);
3175        let metadata = metadata_decoder.try_decode().unwrap();
3176        let DecodeResult::Data(metadata) = metadata else {
3177            panic!("Expected metadata to be decoded successfully");
3178        };
3179        Arc::new(metadata)
3180    }
3181
3182    /// Push the given ranges to the metadata decoder, simulating reading from a file
3183    fn push_ranges_to_metadata_decoder(
3184        metadata_decoder: &mut ParquetMetaDataPushDecoder,
3185        ranges: Vec<Range<u64>>,
3186    ) {
3187        let data = ranges
3188            .iter()
3189            .map(|range| test_file_slice(range.clone()))
3190            .collect::<Vec<_>>();
3191        metadata_decoder.push_ranges(ranges, data).unwrap();
3192    }
3193
3194    /// Push the entire test file into `decoder`.
3195    fn prefetch_test_file(decoder: &mut ParquetPushDecoder) {
3196        decoder
3197            .push_range(test_file_range(), TEST_FILE_DATA.clone())
3198            .unwrap();
3199    }
3200
3201    /// Build a decoder over the test file with the given batch size and
3202    /// prefetch the whole file into it.
3203    fn prefetched_decoder(batch_size: usize) -> ParquetPushDecoder {
3204        let mut decoder = ParquetPushDecoderBuilder::try_new_decoder(test_file_parquet_metadata())
3205            .unwrap()
3206            .with_batch_size(batch_size)
3207            .build()
3208            .unwrap();
3209        prefetch_test_file(&mut decoder);
3210        decoder
3211    }
3212
3213    /// The number of bytes in all column chunks of a row group of the test
3214    /// file.
3215    fn test_file_row_group_bytes(row_group_idx: usize) -> u64 {
3216        test_file_parquet_metadata()
3217            .row_group(row_group_idx)
3218            .columns()
3219            .iter()
3220            .map(|column| {
3221                let range = column_chunk_range(column);
3222                range.end - range.start
3223            })
3224            .sum()
3225    }
3226
3227    fn push_ranges_to_decoder(decoder: &mut ParquetPushDecoder, ranges: Vec<Range<u64>>) {
3228        let data = ranges
3229            .iter()
3230            .map(|range| test_file_slice(range.clone()))
3231            .collect::<Vec<_>>();
3232        decoder.push_ranges(ranges, data).unwrap();
3233    }
3234
3235    /// Expect that the [`DecodeResult`] is a [`DecodeResult::Data`] and return the corresponding element
3236    fn expect_data<T: Debug>(result: Result<DecodeResult<T>, ParquetError>) -> T {
3237        match result.expect("Expected Ok(DecodeResult::Data(T))") {
3238            DecodeResult::Data(data) => data,
3239            result => panic!("Expected DecodeResult::Data, got {result:?}"),
3240        }
3241    }
3242
3243    /// Expect that the [`DecodeResult`] is a [`DecodeResult::NeedsData`] and return the corresponding ranges
3244    fn expect_needs_data<T: Debug>(
3245        result: Result<DecodeResult<T>, ParquetError>,
3246    ) -> Vec<Range<u64>> {
3247        match result.expect("Expected Ok(DecodeResult::NeedsData(ranges))") {
3248            DecodeResult::NeedsData(ranges) => ranges,
3249            result => panic!("Expected DecodeResult::NeedsData, got {result:?}"),
3250        }
3251    }
3252
3253    fn expect_finished<T: Debug>(result: Result<DecodeResult<T>, ParquetError>) {
3254        match result.expect("Expected Ok(DecodeResult::Finished)") {
3255            DecodeResult::Finished => {}
3256            result => panic!("Expected DecodeResult::Finished, got {result:?}"),
3257        }
3258    }
3259}