parquet/file/metadata/
push_decoder.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
18use crate::DecodeResult;
19#[cfg(feature = "encryption")]
20use crate::encryption::decrypt::FileDecryptionProperties;
21use crate::errors::{ParquetError, Result};
22use crate::file::FOOTER_SIZE;
23use crate::file::metadata::parser::{MetadataParser, parse_page_index};
24use crate::file::metadata::{
25    ColumnChunkMask, FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions,
26};
27use crate::file::reader::ChunkReader;
28use bytes::Bytes;
29use std::ops::Range;
30use std::sync::Arc;
31
32/// A push decoder for [`ParquetMetaData`].
33///
34/// This structure implements a push API for decoding Parquet metadata, which
35/// decouples IO from the metadata decoding logic (sometimes referred to as
36/// [Sans-IO]).
37///
38/// See [`ParquetMetaDataReader`] for a pull-based API that incorporates IO and
39/// is simpler to use for basic use cases. This decoder is best for customizing
40/// your IO operations to minimize bytes read, prefetch data, or use async IO.
41///
42/// [Sans-IO]: https://sans-io.readthedocs.io
43/// [`ParquetMetaDataReader`]: crate::file::metadata::ParquetMetaDataReader
44///
45/// # Example
46///
47/// The most basic usage is to feed the decoder with the necessary byte ranges
48/// as requested as shown below. This minimizes the number of bytes read, but
49/// requires the most IO operations - one to read the footer and then one
50/// to read the metadata, and possibly more if page indexes are requested.
51///
52#[cfg_attr(
53    feature = "arrow",
54    doc = r#"
55```rust
56# use std::ops::Range;
57# use bytes::Bytes;
58# use arrow_array::record_batch;
59# use parquet::DecodeResult;
60# use parquet::arrow::ArrowWriter;
61# use parquet::errors::ParquetError;
62# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder};
63#
64# fn decode_metadata() -> Result<ParquetMetaData, ParquetError> {
65# let file_bytes = {
66#   let mut buffer = vec![0];
67#   let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
68#   let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
69#   writer.write(&batch).unwrap();
70#   writer.close().unwrap();
71#   Bytes::from(buffer)
72# };
73# // mimic IO by returning a function that returns the bytes for a given range
74# let get_range = |range: &Range<u64>| -> Bytes {
75#    let start = range.start as usize;
76#     let end = range.end as usize;
77#    file_bytes.slice(start..end)
78# };
79#
80# let file_len = file_bytes.len() as u64;
81// The `ParquetMetaDataPushDecoder` needs to know the file length.
82let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
83// try to decode the metadata. If more data is needed, the decoder will tell you what ranges
84loop {
85    match decoder.try_decode() {
86       Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful
87       Ok(DecodeResult::NeedsData(ranges)) => {
88          // The decoder needs more data
89          //
90          // In this example, we call a function that returns the bytes for each given range.
91          // In a real application, you would likely read the data from a file or network.
92          let data = ranges.iter().map(|range| get_range(range)).collect();
93          // Push the data into the decoder and try to decode again on the next iteration.
94          decoder.push_ranges(ranges, data).unwrap();
95       }
96       Ok(DecodeResult::Finished) => { unreachable!("returned metadata in previous match arm") }
97       Err(e) => return Err(e),
98    }
99}
100# }
101```
102"#
103)]
104///
105/// # Example with "prefetching"
106///
107/// By default, the [`ParquetMetaDataPushDecoder`] will request only the exact byte
108/// ranges it needs. This minimizes the number of bytes read, however it
109/// requires at least two IO operations to read the metadata - one to read the
110/// footer and then one to read the metadata.
111///
112/// If the file has a "Page Index" (see [Self::with_page_index_policy]), three
113/// IO operations are required to read the metadata, as the page index is
114/// not part of the normal metadata footer.
115///
116/// To reduce the number of IO operations in systems with high per operation
117/// overhead (e.g. cloud storage), you can "prefetch" the data and then push
118/// the data into the decoder before calling [`Self::try_decode`]. If you do
119/// not push enough bytes, the decoder will return the ranges that are still
120/// needed.
121///
122/// This approach can also be used when you have the entire file already in memory
123/// for other reasons.
124#[cfg_attr(
125    feature = "arrow",
126    doc = r#"
127```rust
128# use std::ops::Range;
129# use bytes::Bytes;
130# use arrow_array::record_batch;
131# use parquet::DecodeResult;
132# use parquet::arrow::ArrowWriter;
133# use parquet::errors::ParquetError;
134# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder};
135#
136# fn decode_metadata() -> Result<ParquetMetaData, ParquetError> {
137# let file_bytes = {
138#   let mut buffer = vec![0];
139#   let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
140#   let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
141#   writer.write(&batch).unwrap();
142#   writer.close().unwrap();
143#   Bytes::from(buffer)
144# };
145#
146let file_len = file_bytes.len() as u64;
147// For this example, we "prefetch" all the bytes which we have in memory,
148// but in a real application, you would likely read a chunk from the end
149// for example 1MB.
150let prefetched_bytes = file_bytes.clone();
151let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
152// push the prefetched bytes into the decoder
153decoder.push_ranges(vec![0..file_len], vec![prefetched_bytes]).unwrap();
154// The decoder will now be able to decode the metadata. Note in a real application,
155// unless you can guarantee that the pushed data is enough to decode the metadata,
156// you still need to call `try_decode` in a loop until it returns `DecodeResult::Data`
157// as shown in  the previous example
158    match decoder.try_decode() {
159        Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful
160        other => { panic!("expected DecodeResult::Data, got: {other:?}") }
161    }
162# }
163```
164"#
165)]
166///
167/// # Example using [`AsyncRead`]
168///
169/// [`ParquetMetaDataPushDecoder`] is designed to work with any data source that can
170/// provide byte ranges, including async IO sources. However, it does not
171/// implement async IO itself. To use async IO, you simply write an async
172/// wrapper around it that reads the required byte ranges and pushes them into the
173/// decoder.
174#[cfg_attr(
175    feature = "arrow",
176    doc = r#"
177```rust
178# use std::ops::Range;
179# use bytes::Bytes;
180use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
181# use arrow_array::record_batch;
182# use parquet::DecodeResult;
183# use parquet::arrow::ArrowWriter;
184# use parquet::errors::ParquetError;
185# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder};
186#
187// This function decodes Parquet Metadata from anything that implements
188// [`AsyncRead`] and [`AsyncSeek`] such as a tokio::fs::File
189async fn decode_metadata(
190  file_len: u64,
191  mut async_source: impl AsyncRead + AsyncSeek + Unpin
192) -> Result<ParquetMetaData, ParquetError> {
193  // We need a ParquetMetaDataPushDecoder to decode the metadata.
194  let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
195  loop {
196    match decoder.try_decode() {
197       Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful
198       Ok(DecodeResult::NeedsData(ranges)) => {
199          // The decoder needs more data
200          //
201          // In this example we use the AsyncRead and AsyncSeek traits to read the
202          // required ranges from the async source.
203          let mut data = Vec::with_capacity(ranges.len());
204          for range in &ranges {
205            let mut buffer = vec![0; (range.end - range.start) as usize];
206            async_source.seek(std::io::SeekFrom::Start(range.start)).await?;
207            async_source.read_exact(&mut buffer).await?;
208            data.push(Bytes::from(buffer));
209          }
210          // Push the data into the decoder and try to decode again on the next iteration.
211          decoder.push_ranges(ranges, data).unwrap();
212       }
213       Ok(DecodeResult::Finished) => { unreachable!("returned metadata in previous match arm") }
214       Err(e) => return Err(e),
215    }
216  }
217}
218```
219"#
220)]
221/// [`AsyncRead`]: tokio::io::AsyncRead
222#[derive(Debug)]
223pub struct ParquetMetaDataPushDecoder {
224    /// Decoding state
225    state: DecodeState,
226    /// policy for loading ColumnIndex (part of the PageIndex)
227    column_index_policy: PageIndexPolicy,
228    /// policy for loading OffsetIndex (part of the PageIndex)
229    offset_index_policy: PageIndexPolicy,
230    /// which rows and columns of the column index should be read
231    column_index_mask: ColumnChunkMask,
232    /// which rows and columns of the offset index should be read
233    offset_index_mask: ColumnChunkMask,
234    /// Underlying buffers
235    buffers: crate::util::push_buffers::PushBuffers,
236    /// Encryption API
237    metadata_parser: MetadataParser,
238}
239
240impl ParquetMetaDataPushDecoder {
241    /// Create a new `ParquetMetaDataPushDecoder` with the given file length.
242    ///
243    /// By default, this will read page indexes and column indexes. See
244    /// [`ParquetMetaDataPushDecoder::with_page_index_policy`] for more detail.
245    ///
246    /// See examples on [`ParquetMetaDataPushDecoder`].
247    pub fn try_new(file_len: u64) -> Result<Self> {
248        if file_len < 8 {
249            return Err(ParquetError::General(format!(
250                "Parquet files are at least 8 bytes long, but file length is {file_len}"
251            )));
252        }
253
254        Ok(Self {
255            state: DecodeState::ReadingFooter,
256            column_index_policy: PageIndexPolicy::Optional,
257            offset_index_policy: PageIndexPolicy::Optional,
258            column_index_mask: ColumnChunkMask::all(),
259            offset_index_mask: ColumnChunkMask::all(),
260            buffers: crate::util::push_buffers::PushBuffers::new(file_len),
261            metadata_parser: MetadataParser::new(),
262        })
263    }
264
265    /// Begin decoding from the given footer tail.
266    pub(crate) fn try_new_with_footer_tail(file_len: u64, footer_tail: FooterTail) -> Result<Self> {
267        let mut new_self = Self::try_new(file_len)?;
268        new_self.state = DecodeState::ReadingMetadata(footer_tail);
269        Ok(new_self)
270    }
271
272    /// Create a decoder with the given `ParquetMetaData` already known.
273    ///
274    /// This can be used to parse and populate the page index structures
275    /// after the metadata has already been decoded.
276    ///
277    /// When decoding page indexes, any existing page indexes in `metadata` will be
278    /// replaced based on the configured policies. See
279    /// [`crate::file::metadata::ParquetMetaDataReader::read_page_indexes`] for details
280    /// on the replacement behavior.
281    pub fn try_new_with_metadata(file_len: u64, metadata: ParquetMetaData) -> Result<Self> {
282        let mut new_self = Self::try_new(file_len)?;
283        new_self.state = DecodeState::ReadingPageIndex(Box::new(metadata));
284        Ok(new_self)
285    }
286
287    /// Enable or disable reading the page index structures described in
288    /// "[Parquet page index] Layout to Support Page Skipping".
289    ///
290    /// Defaults to [`PageIndexPolicy::Optional`]
291    ///
292    /// This requires
293    /// 1. The Parquet file to have been written with page indexes
294    /// 2. Additional data to be pushed into the decoder (as the page indexes are not part of the thrift footer)
295    ///
296    /// [Parquet page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
297    pub fn with_page_index_policy(mut self, page_index_policy: PageIndexPolicy) -> Self {
298        self.column_index_policy = page_index_policy;
299        self.offset_index_policy = page_index_policy;
300        self
301    }
302
303    /// Set the policy for reading the ColumnIndex (part of the PageIndex)
304    pub fn with_column_index_policy(mut self, column_index_policy: PageIndexPolicy) -> Self {
305        self.column_index_policy = column_index_policy;
306        self
307    }
308
309    /// Set the policy for reading the OffsetIndex (part of the PageIndex)
310    pub fn with_offset_index_policy(mut self, offset_index_policy: PageIndexPolicy) -> Self {
311        self.offset_index_policy = offset_index_policy;
312        self
313    }
314
315    // TODO: deprecate the following as part of https://github.com/apache/arrow-rs/issues/11395
316
317    /// Select the row groups and columns for which both page index structures are read.
318    pub fn with_page_index_mask(mut self, mask: ColumnChunkMask) -> Self {
319        self.column_index_mask = mask.clone();
320        self.offset_index_mask = mask;
321        self
322    }
323
324    /// Select the row groups and columns for which column indexes are read.
325    pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
326        self.column_index_mask = mask;
327        self
328    }
329
330    /// Select the row groups and columns for which offset indexes are read.
331    pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
332        self.offset_index_mask = mask;
333        self
334    }
335
336    /// Set the options to use when decoding the Parquet metadata.
337    pub fn with_metadata_options(mut self, options: Option<Arc<ParquetMetaDataOptions>>) -> Self {
338        self.metadata_parser = self.metadata_parser.with_metadata_options(options);
339        self
340    }
341
342    #[cfg(feature = "encryption")]
343    /// Provide decryption properties for decoding encrypted Parquet files
344    pub fn with_file_decryption_properties(
345        mut self,
346        file_decryption_properties: Option<std::sync::Arc<FileDecryptionProperties>>,
347    ) -> Self {
348        self.metadata_parser = self
349            .metadata_parser
350            .with_file_decryption_properties(file_decryption_properties);
351        self
352    }
353
354    /// Push the data into the decoder's buffer.
355    ///
356    /// The decoder does not immediately attempt to decode the metadata
357    /// after pushing data. Instead, it accumulates the pushed data until you
358    /// call [`Self::try_decode`].
359    ///
360    /// # Determining required data:
361    ///
362    /// To determine what ranges are required to decode the metadata, you can
363    /// either:
364    ///
365    /// 1. Call [`Self::try_decode`] first to get the exact ranges required (see
366    ///    example on [`Self`])
367    ///
368    /// 2. Speculatively push any data that you have available, which may
369    ///    include more than the footer data or requested bytes.
370    ///
371    /// Speculatively pushing data can be used when  "prefetching" data. See
372    /// example on [`Self`]
373    pub fn push_ranges(&mut self, ranges: Vec<Range<u64>>, buffers: Vec<Bytes>) -> Result<()> {
374        if matches!(&self.state, DecodeState::Finished) {
375            return Err(general_err!(
376                "ParquetMetaDataPushDecoder: cannot push data after decoding is finished"
377            ));
378        }
379        self.buffers.push_ranges(ranges, buffers)?;
380        Ok(())
381    }
382
383    /// Pushes a single range of data into the decoder's buffer.
384    pub fn push_range(&mut self, range: Range<u64>, buffer: Bytes) -> Result<()> {
385        if matches!(&self.state, DecodeState::Finished) {
386            return Err(general_err!(
387                "ParquetMetaDataPushDecoder: cannot push data after decoding is finished"
388            ));
389        }
390        self.buffers.push_range(range, buffer)?;
391        Ok(())
392    }
393
394    /// Clear any staged byte ranges currently buffered for future decode work.
395    pub fn clear_all_ranges(&mut self) {
396        self.buffers.clear_all_ranges();
397    }
398
399    /// Try to decode the metadata from the pushed data, returning the
400    /// decoded metadata or an error if not enough data is available.
401    pub fn try_decode(&mut self) -> Result<DecodeResult<ParquetMetaData>> {
402        let file_len = self.buffers.file_len();
403        let footer_len = FOOTER_SIZE as u64;
404        loop {
405            match std::mem::replace(&mut self.state, DecodeState::Intermediate) {
406                DecodeState::ReadingFooter => {
407                    // need to have the last 8 bytes of the file to decode the metadata
408                    let footer_start = file_len.saturating_sub(footer_len);
409                    let footer_range = footer_start..file_len;
410
411                    if !self.buffers.has_range(&footer_range) {
412                        self.state = DecodeState::ReadingFooter;
413                        return Ok(needs_range(footer_range));
414                    }
415                    let footer_bytes = self.get_bytes(&footer_range)?;
416                    let footer_tail = FooterTail::try_from(footer_bytes.as_ref())?;
417
418                    self.state = DecodeState::ReadingMetadata(footer_tail);
419                }
420
421                DecodeState::ReadingMetadata(footer_tail) => {
422                    let metadata_len: u64 = footer_tail.metadata_length() as u64;
423                    let metadata_start = file_len - footer_len - metadata_len;
424                    let metadata_end = metadata_start + metadata_len;
425                    let metadata_range = metadata_start..metadata_end;
426
427                    if !self.buffers.has_range(&metadata_range) {
428                        self.state = DecodeState::ReadingMetadata(footer_tail);
429                        return Ok(needs_range(metadata_range));
430                    }
431
432                    let metadata = self.metadata_parser.decode_metadata(
433                        &self.get_bytes(&metadata_range)?,
434                        footer_tail.is_encrypted_footer(),
435                    )?;
436                    // Note: ReadingPageIndex first checks if page indexes are needed
437                    // and is a no-op if not
438                    self.state = DecodeState::ReadingPageIndex(Box::new(metadata));
439                }
440
441                DecodeState::ReadingPageIndex(mut metadata) => {
442                    // First determine if any page indexes are needed based on
443                    // the specified policies
444                    let range = range_for_page_index(
445                        &metadata,
446                        self.column_index_policy,
447                        self.offset_index_policy,
448                        &self.column_index_mask,
449                        &self.offset_index_mask,
450                    );
451
452                    let Some(page_index_range) = range else {
453                        // There is nothing to decode because both policies are Skip, the masks
454                        // select nothing, or the file has no index ranges. Clear any previously
455                        // loaded indexes while preserving the compatibility behavior that
456                        // PageIndexPolicy::Required accepts files with no page indexes at all.
457                        metadata.set_page_index(None);
458                        self.state = DecodeState::Finished;
459                        return Ok(DecodeResult::Data(*metadata));
460                    };
461
462                    if !self.buffers.has_range(&page_index_range) {
463                        self.state = DecodeState::ReadingPageIndex(metadata);
464                        return Ok(needs_range(page_index_range));
465                    }
466
467                    let buffer = self.get_bytes(&page_index_range)?;
468                    let offset = page_index_range.start;
469                    parse_page_index(
470                        &mut metadata,
471                        self.column_index_policy,
472                        self.offset_index_policy,
473                        &self.column_index_mask,
474                        &self.offset_index_mask,
475                        &buffer,
476                        offset,
477                    )?;
478                    self.state = DecodeState::Finished;
479                    return Ok(DecodeResult::Data(*metadata));
480                }
481
482                DecodeState::Finished => return Ok(DecodeResult::Finished),
483                DecodeState::Intermediate => {
484                    return Err(general_err!(
485                        "ParquetMetaDataPushDecoder: internal error, invalid state"
486                    ));
487                }
488            }
489        }
490    }
491
492    /// Returns the bytes for the given range from the internal buffer
493    fn get_bytes(&self, range: &Range<u64>) -> Result<Bytes> {
494        let start = range.start;
495        let raw_len = range.end - range.start;
496        let len: usize = raw_len.try_into().map_err(|_| {
497            ParquetError::General(format!(
498                "ParquetMetaDataPushDecoder: Range length too large to fit in usize: {raw_len}",
499            ))
500        })?;
501        self.buffers.get_bytes(start, len)
502    }
503}
504
505/// returns a DecodeResults that describes needing the given range
506fn needs_range(range: Range<u64>) -> DecodeResult<ParquetMetaData> {
507    DecodeResult::NeedsData(vec![range])
508}
509
510/// Decoding state machine
511#[derive(Debug)]
512enum DecodeState {
513    /// Reading the last 8 bytes of the file
514    ReadingFooter,
515    /// Reading the metadata thrift structure
516    ReadingMetadata(FooterTail),
517    // Actively reading the page index
518    ReadingPageIndex(Box<ParquetMetaData>),
519    // Decoding is complete
520    Finished,
521    /// State left during the `try_decode` method so something valid is present.
522    /// This state should never be observed.
523    Intermediate,
524}
525
526/// Returns the byte range needed to read the offset/page indexes, based on the
527/// specified policies and masks
528///
529/// Returns None if no page indexes are needed
530pub(crate) fn range_for_page_index(
531    metadata: &ParquetMetaData,
532    column_index_policy: PageIndexPolicy,
533    offset_index_policy: PageIndexPolicy,
534    column_index_mask: &ColumnChunkMask,
535    offset_index_mask: &ColumnChunkMask,
536) -> Option<Range<u64>> {
537    let mut range: Option<Range<u64>> = None;
538    for (policy, mask, column_index) in [
539        (column_index_policy, column_index_mask, true),
540        (offset_index_policy, offset_index_mask, false),
541    ] {
542        if policy == PageIndexPolicy::Skip {
543            continue;
544        }
545        for row_group in mask.row_group_indices(metadata.num_row_groups()) {
546            for column in mask.column_indices(metadata.row_group(row_group).num_columns()) {
547                let column = metadata.row_group(row_group).column(column);
548                let index = if column_index {
549                    column.column_index_range()
550                } else {
551                    column.offset_index_range()
552                };
553                if let Some(index) = index {
554                    range = Some(match range {
555                        Some(range) => range.start.min(index.start)..range.end.max(index.end),
556                        None => index,
557                    });
558                }
559            }
560        }
561    }
562    range
563}
564
565// These tests use the arrow writer to create a parquet file in memory
566// so they need the arrow feature and the test feature
567#[cfg(all(test, feature = "arrow"))]
568mod tests {
569    use super::*;
570    use crate::arrow::ArrowWriter;
571    use crate::file::properties::WriterProperties;
572    use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringViewArray};
573    use bytes::Bytes;
574    use std::fmt::Debug;
575    use std::ops::Range;
576    use std::sync::{Arc, LazyLock};
577
578    /// It is possible to decode the metadata from the entire file at once before being asked
579    #[test]
580    fn test_metadata_decoder_all_data() {
581        let file_len = test_file_len();
582        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
583        // Push the entire file data into the metadata decoder
584        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]);
585
586        // should be able to decode the metadata without needing more data
587        let metadata = expect_data(metadata_decoder.try_decode());
588
589        assert_eq!(metadata.num_row_groups(), 2);
590        assert_eq!(metadata.row_group(0).num_rows(), 200);
591        assert_eq!(metadata.row_group(1).num_rows(), 200);
592        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
593    }
594
595    #[test]
596    fn test_range_for_page_index_respects_column_mask() {
597        let mut decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap();
598        push_ranges_to_metadata_decoder(&mut decoder, vec![test_file_range()]);
599        let metadata = expect_data(decoder.try_decode());
600
601        let all = range_for_page_index(
602            &metadata,
603            PageIndexPolicy::Required,
604            PageIndexPolicy::Skip,
605            &ColumnChunkMask::all(),
606            &ColumnChunkMask::all(),
607        )
608        .unwrap();
609        let masked = range_for_page_index(
610            &metadata,
611            PageIndexPolicy::Required,
612            PageIndexPolicy::Skip,
613            &ColumnChunkMask::columns([0]),
614            &ColumnChunkMask::all(),
615        )
616        .unwrap();
617
618        assert!(all.start <= masked.start && masked.end <= all.end);
619        assert!(masked.end - masked.start < all.end - all.start);
620    }
621
622    /// It is possible to feed some, but not all, of the footer into the metadata decoder
623    /// before asked. This avoids multiple IO requests
624    #[test]
625    fn test_metadata_decoder_prefetch_success() {
626        let file_len = test_file_len();
627        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
628        // simulate pre-fetching the last 2k bytes of the file without asking the decoder
629        let prefetch_range = (file_len - 2 * 1024)..file_len;
630        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![prefetch_range]);
631
632        // expect the decoder has enough data to decode the metadata
633        let metadata = expect_data(metadata_decoder.try_decode());
634        expect_finished(metadata_decoder.try_decode());
635        assert_eq!(metadata.num_row_groups(), 2);
636        assert_eq!(metadata.row_group(0).num_rows(), 200);
637        assert_eq!(metadata.row_group(1).num_rows(), 200);
638        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
639    }
640
641    /// It is possible to pre-fetch some, but not all, of the necessary data
642    /// data
643    #[test]
644    fn test_metadata_decoder_prefetch_retry() {
645        let file_len = test_file_len();
646        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
647        // simulate pre-fetching the last 1500 bytes of the file.
648        // this is enough to read the footer thrift metadata, but not the offset indexes
649        let prefetch_range = (file_len - 1500)..file_len;
650        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![prefetch_range]);
651
652        // expect another request is needed to read the offset indexes (note
653        // try_decode only returns NeedsData once, whereas without any prefetching it would
654        // return NeedsData three times)
655        let ranges = expect_needs_data(metadata_decoder.try_decode());
656        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
657
658        // expect the decoder has enough data to decode the metadata
659        let metadata = expect_data(metadata_decoder.try_decode());
660        expect_finished(metadata_decoder.try_decode());
661
662        assert_eq!(metadata.num_row_groups(), 2);
663        assert_eq!(metadata.row_group(0).num_rows(), 200);
664        assert_eq!(metadata.row_group(1).num_rows(), 200);
665        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
666    }
667
668    #[test]
669    fn test_metadata_decoder_clear_all_ranges() {
670        let file_len = test_file_len();
671        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
672
673        metadata_decoder
674            .push_range(test_file_range(), TEST_FILE_DATA.clone())
675            .unwrap();
676        assert_eq!(metadata_decoder.buffers.buffered_bytes(), test_file_len());
677
678        metadata_decoder.clear_all_ranges();
679        assert_eq!(metadata_decoder.buffers.buffered_bytes(), 0);
680
681        let ranges = expect_needs_data(metadata_decoder.try_decode());
682        assert_eq!(ranges, vec![test_file_len() - 8..test_file_len()]);
683    }
684
685    /// Decode the metadata incrementally, simulating a scenario where exactly the data needed
686    /// is read in each step
687    #[test]
688    fn test_metadata_decoder_incremental() {
689        let file_len = TEST_FILE_DATA.len() as u64;
690        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
691        let ranges = expect_needs_data(metadata_decoder.try_decode());
692        assert_eq!(ranges.len(), 1);
693        assert_eq!(ranges[0], test_file_len() - 8..test_file_len());
694        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
695
696        // expect the first request to read the footer
697        let ranges = expect_needs_data(metadata_decoder.try_decode());
698        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
699
700        // expect the second request to read the offset indexes
701        let ranges = expect_needs_data(metadata_decoder.try_decode());
702        // since we're reading all indexes, there should be a single range needed
703        assert_eq!(ranges.len(), 1);
704        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
705
706        // expect the third request to read the actual data
707        let metadata = expect_data(metadata_decoder.try_decode());
708        expect_finished(metadata_decoder.try_decode());
709
710        assert_eq!(metadata.num_row_groups(), 2);
711        assert_eq!(metadata.row_group(0).num_rows(), 200);
712        assert_eq!(metadata.row_group(1).num_rows(), 200);
713        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
714    }
715
716    /// Decode the metadata incrementally, but without reading the page indexes
717    /// (so only two requests)
718    #[test]
719    fn test_metadata_decoder_incremental_no_page_index() {
720        let file_len = TEST_FILE_DATA.len() as u64;
721        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len)
722            .unwrap()
723            .with_page_index_policy(PageIndexPolicy::Skip);
724        let ranges = expect_needs_data(metadata_decoder.try_decode());
725        assert_eq!(ranges.len(), 1);
726        assert_eq!(ranges[0], test_file_len() - 8..test_file_len());
727        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
728
729        // expect the first request to read the footer
730        let ranges = expect_needs_data(metadata_decoder.try_decode());
731        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
732
733        // expect NO second request to read the offset indexes, should just cough up the metadata
734        let metadata = expect_data(metadata_decoder.try_decode());
735        expect_finished(metadata_decoder.try_decode());
736
737        assert_eq!(metadata.num_row_groups(), 2);
738        assert_eq!(metadata.row_group(0).num_rows(), 200);
739        assert_eq!(metadata.row_group(1).num_rows(), 200);
740        assert!(metadata.page_index().is_none()); // of course, we did not read the page index
741    }
742
743    static TEST_BATCH: LazyLock<RecordBatch> = LazyLock::new(|| {
744        // Input batch has 400 rows, with 3 columns: "a", "b", "c"
745        // Note c is a different types (so the data page sizes will be different)
746        let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400));
747        let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800));
748        let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| {
749            if i % 2 == 0 {
750                format!("string_{i}")
751            } else {
752                format!("A string larger than 12 bytes and thus not inlined {i}")
753            }
754        })));
755
756        RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
757    });
758
759    /// Create a parquet file in memory for testing. See [`test_file_range`] for details.
760    static TEST_FILE_DATA: LazyLock<Bytes> = LazyLock::new(|| {
761        let input_batch = &TEST_BATCH;
762        let mut output = Vec::new();
763
764        let writer_options = WriterProperties::builder()
765            .set_max_row_group_row_count(Some(200))
766            .set_data_page_row_count_limit(100)
767            .build();
768        let mut writer =
769            ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap();
770
771        // since the limits are only enforced on batch boundaries, write the input
772        // batch in chunks of 50
773        let mut row_remain = input_batch.num_rows();
774        while row_remain > 0 {
775            let chunk_size = row_remain.min(50);
776            let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size);
777            writer.write(&chunk).unwrap();
778            row_remain -= chunk_size;
779        }
780        writer.close().unwrap();
781        Bytes::from(output)
782    });
783
784    /// Return the length of the test file in bytes
785    fn test_file_len() -> u64 {
786        TEST_FILE_DATA.len() as u64
787    }
788
789    /// Return the range of the entire test file
790    fn test_file_range() -> Range<u64> {
791        0..test_file_len()
792    }
793
794    /// Return a slice of the test file data from the given range
795    pub fn test_file_slice(range: Range<u64>) -> Bytes {
796        let start: usize = range.start.try_into().unwrap();
797        let end: usize = range.end.try_into().unwrap();
798        TEST_FILE_DATA.slice(start..end)
799    }
800
801    /// Push the given ranges to the metadata decoder, simulating reading from a file
802    fn push_ranges_to_metadata_decoder(
803        metadata_decoder: &mut ParquetMetaDataPushDecoder,
804        ranges: Vec<Range<u64>>,
805    ) {
806        let data = ranges
807            .iter()
808            .map(|range| test_file_slice(range.clone()))
809            .collect::<Vec<_>>();
810        metadata_decoder.push_ranges(ranges, data).unwrap();
811    }
812
813    /// Expect that the [`DecodeResult`] is a [`DecodeResult::Data`] and return the corresponding element
814    fn expect_data<T: Debug>(result: Result<DecodeResult<T>>) -> T {
815        match result.expect("Expected Ok(DecodeResult::Data(T))") {
816            DecodeResult::Data(data) => data,
817            result => panic!("Expected DecodeResult::Data, got {result:?}"),
818        }
819    }
820
821    /// Expect that the [`DecodeResult`] is a [`DecodeResult::NeedsData`] and return the corresponding ranges
822    fn expect_needs_data<T: Debug>(result: Result<DecodeResult<T>>) -> Vec<Range<u64>> {
823        match result.expect("Expected Ok(DecodeResult::NeedsData(ranges))") {
824            DecodeResult::NeedsData(ranges) => ranges,
825            result => panic!("Expected DecodeResult::NeedsData, got {result:?}"),
826        }
827    }
828
829    fn expect_finished<T: Debug>(result: Result<DecodeResult<T>>) {
830        match result.expect("Expected Ok(DecodeResult::Finished)") {
831            DecodeResult::Finished => {}
832            result => panic!("Expected DecodeResult::Finished, got {result:?}"),
833        }
834    }
835}
parquet/file/metadata/push_decoder.rs

parquet/file/metadata/
push_decoder.rs
