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::page_index::PageIndex;
24use crate::file::metadata::parser::{MetadataParser, parse_page_index};
25use crate::file::metadata::{
26    ColumnChunkMask, FooterTail, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions,
27};
28use crate::file::reader::ChunkReader;
29use bytes::Bytes;
30use std::ops::Range;
31use std::sync::Arc;
32
33/// A push decoder for [`ParquetMetaData`].
34///
35/// This structure implements a push API for decoding Parquet metadata, which
36/// decouples IO from the metadata decoding logic (sometimes referred to as
37/// [Sans-IO]).
38///
39/// See [`ParquetMetaDataReader`] for a pull-based API that incorporates IO and
40/// is simpler to use for basic use cases. This decoder is best for customizing
41/// your IO operations to minimize bytes read, prefetch data, or use async IO.
42///
43/// [Sans-IO]: https://sans-io.readthedocs.io
44/// [`ParquetMetaDataReader`]: crate::file::metadata::ParquetMetaDataReader
45///
46/// # Example
47///
48/// The most basic usage is to feed the decoder with the necessary byte ranges
49/// as requested as shown below. This minimizes the number of bytes read, but
50/// requires the most IO operations - one to read the footer and then one
51/// to read the metadata, and possibly more if page indexes are requested.
52///
53#[cfg_attr(
54    feature = "arrow",
55    doc = r#"
56```rust
57# use std::ops::Range;
58# use bytes::Bytes;
59# use arrow_array::record_batch;
60# use parquet::DecodeResult;
61# use parquet::arrow::ArrowWriter;
62# use parquet::errors::ParquetError;
63# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder};
64#
65# fn decode_metadata() -> Result<ParquetMetaData, ParquetError> {
66# let file_bytes = {
67#   let mut buffer = vec![0];
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#
81# let file_len = file_bytes.len() as u64;
82// The `ParquetMetaDataPushDecoder` needs to know the file length.
83let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
84// try to decode the metadata. If more data is needed, the decoder will tell you what ranges
85loop {
86    match decoder.try_decode() {
87       Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful
88       Ok(DecodeResult::NeedsData(ranges)) => {
89          // The decoder needs more data
90          //
91          // In this example, we call a function that returns the bytes for each given range.
92          // In a real application, you would likely read the data from a file or network.
93          let data = ranges.iter().map(|range| get_range(range)).collect();
94          // Push the data into the decoder and try to decode again on the next iteration.
95          decoder.push_ranges(ranges, data).unwrap();
96       }
97       Ok(DecodeResult::Finished) => { unreachable!("returned metadata in previous match arm") }
98       Err(e) => return Err(e),
99    }
100}
101# }
102```
103"#
104)]
105///
106/// # Example with "prefetching"
107///
108/// By default, the [`ParquetMetaDataPushDecoder`] will request only the exact byte
109/// ranges it needs. This minimizes the number of bytes read, however it
110/// requires at least two IO operations to read the metadata - one to read the
111/// footer and then one to read the metadata.
112///
113/// If the file has a "Page Index" (see [Self::with_page_index_policy]), three
114/// IO operations are required to read the metadata, as the page index is
115/// not part of the normal metadata footer.
116///
117/// To reduce the number of IO operations in systems with high per operation
118/// overhead (e.g. cloud storage), you can "prefetch" the data and then push
119/// the data into the decoder before calling [`Self::try_decode`]. If you do
120/// not push enough bytes, the decoder will return the ranges that are still
121/// needed.
122///
123/// This approach can also be used when you have the entire file already in memory
124/// for other reasons.
125#[cfg_attr(
126    feature = "arrow",
127    doc = r#"
128```rust
129# use std::ops::Range;
130# use bytes::Bytes;
131# use arrow_array::record_batch;
132# use parquet::DecodeResult;
133# use parquet::arrow::ArrowWriter;
134# use parquet::errors::ParquetError;
135# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder};
136#
137# fn decode_metadata() -> Result<ParquetMetaData, ParquetError> {
138# let file_bytes = {
139#   let mut buffer = vec![0];
140#   let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
141#   let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
142#   writer.write(&batch).unwrap();
143#   writer.close().unwrap();
144#   Bytes::from(buffer)
145# };
146#
147let file_len = file_bytes.len() as u64;
148// For this example, we "prefetch" all the bytes which we have in memory,
149// but in a real application, you would likely read a chunk from the end
150// for example 1MB.
151let prefetched_bytes = file_bytes.clone();
152let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
153// push the prefetched bytes into the decoder
154decoder.push_ranges(vec![0..file_len], vec![prefetched_bytes]).unwrap();
155// The decoder will now be able to decode the metadata. Note in a real application,
156// unless you can guarantee that the pushed data is enough to decode the metadata,
157// you still need to call `try_decode` in a loop until it returns `DecodeResult::Data`
158// as shown in  the previous example
159    match decoder.try_decode() {
160        Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful
161        other => { panic!("expected DecodeResult::Data, got: {other:?}") }
162    }
163# }
164```
165"#
166)]
167///
168/// # Example using [`AsyncRead`]
169///
170/// [`ParquetMetaDataPushDecoder`] is designed to work with any data source that can
171/// provide byte ranges, including async IO sources. However, it does not
172/// implement async IO itself. To use async IO, you simply write an async
173/// wrapper around it that reads the required byte ranges and pushes them into the
174/// decoder.
175#[cfg_attr(
176    feature = "arrow",
177    doc = r#"
178```rust
179# use std::ops::Range;
180# use bytes::Bytes;
181use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeek, AsyncSeekExt};
182# use arrow_array::record_batch;
183# use parquet::DecodeResult;
184# use parquet::arrow::ArrowWriter;
185# use parquet::errors::ParquetError;
186# use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataPushDecoder};
187#
188// This function decodes Parquet Metadata from anything that implements
189// [`AsyncRead`] and [`AsyncSeek`] such as a tokio::fs::File
190async fn decode_metadata(
191  file_len: u64,
192  mut async_source: impl AsyncRead + AsyncSeek + Unpin
193) -> Result<ParquetMetaData, ParquetError> {
194  // We need a ParquetMetaDataPushDecoder to decode the metadata.
195  let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
196  loop {
197    match decoder.try_decode() {
198       Ok(DecodeResult::Data(metadata)) => { return Ok(metadata); } // decode successful
199       Ok(DecodeResult::NeedsData(ranges)) => {
200          // The decoder needs more data
201          //
202          // In this example we use the AsyncRead and AsyncSeek traits to read the
203          // required ranges from the async source.
204          let mut data = Vec::with_capacity(ranges.len());
205          for range in &ranges {
206            let mut buffer = vec![0; (range.end - range.start) as usize];
207            async_source.seek(std::io::SeekFrom::Start(range.start)).await?;
208            async_source.read_exact(&mut buffer).await?;
209            data.push(Bytes::from(buffer));
210          }
211          // Push the data into the decoder and try to decode again on the next iteration.
212          decoder.push_ranges(ranges, data).unwrap();
213       }
214       Ok(DecodeResult::Finished) => { unreachable!("returned metadata in previous match arm") }
215       Err(e) => return Err(e),
216    }
217  }
218}
219```
220"#
221)]
222/// [`AsyncRead`]: tokio::io::AsyncRead
223#[derive(Debug)]
224pub struct ParquetMetaDataPushDecoder {
225    /// Decoding state
226    state: DecodeState,
227    /// policy for loading ColumnIndex (part of the PageIndex)
228    column_index_policy: PageIndexPolicy,
229    /// policy for loading OffsetIndex (part of the PageIndex)
230    offset_index_policy: PageIndexPolicy,
231    /// which rows and columns of the column index should be read
232    column_index_mask: ColumnChunkMask,
233    /// which rows and columns of the offset index should be read
234    offset_index_mask: ColumnChunkMask,
235    /// Underlying buffers
236    buffers: crate::util::push_buffers::PushBuffers,
237    /// Encryption API
238    metadata_parser: MetadataParser,
239}
240
241impl ParquetMetaDataPushDecoder {
242    /// Create a new `ParquetMetaDataPushDecoder` with the given file length.
243    ///
244    /// By default, this will read page indexes and column indexes. See
245    /// [`ParquetMetaDataPushDecoder::with_page_index_policy`] for more detail.
246    ///
247    /// See examples on [`ParquetMetaDataPushDecoder`].
248    pub fn try_new(file_len: u64) -> Result<Self> {
249        if file_len < 8 {
250            return Err(ParquetError::General(format!(
251                "Parquet files are at least 8 bytes long, but file length is {file_len}"
252            )));
253        }
254
255        Ok(Self {
256            state: DecodeState::ReadingFooter,
257            column_index_policy: PageIndexPolicy::Optional,
258            offset_index_policy: PageIndexPolicy::Optional,
259            column_index_mask: ColumnChunkMask::all(),
260            offset_index_mask: ColumnChunkMask::all(),
261            buffers: crate::util::push_buffers::PushBuffers::new(file_len),
262            metadata_parser: MetadataParser::new(),
263        })
264    }
265
266    /// Begin decoding from the given footer tail.
267    pub(crate) fn try_new_with_footer_tail(file_len: u64, footer_tail: FooterTail) -> Result<Self> {
268        let mut new_self = Self::try_new(file_len)?;
269        new_self.state = DecodeState::ReadingMetadata(footer_tail);
270        Ok(new_self)
271    }
272
273    /// Create a decoder with the given `ParquetMetaData` already known.
274    ///
275    /// This can be used to parse and populate the page index structures
276    /// after the metadata has already been decoded.
277    ///
278    /// The page indexes can then be decoded in one of two ways:
279    ///
280    /// * [`Self::try_decode`] attaches the decoded page indexes to the returned
281    ///   metadata, replacing (or clearing) any existing page indexes
282    ///   See [`ParquetMetaDataReader::read_page_indexes`] for details.
283    /// * [`Self::try_decode_page_index`] returns the decoded page indexes
284    ///   separately and leaves any page indexes already attached to `metadata`
285    ///   unchanged.
286    ///
287    /// [`ParquetMetaDataReader::read_page_indexes`]: [`crate::file::metadata::ParquetMetaDataReader::read_page_indexes`]
288    pub fn try_new_with_metadata(file_len: u64, metadata: ParquetMetaData) -> Result<Self> {
289        let mut new_self = Self::try_new(file_len)?;
290        new_self.state = DecodeState::ReadingPageIndex(Box::new(metadata));
291        Ok(new_self)
292    }
293
294    /// Enable or disable reading the page index structures described in
295    /// "[Parquet page index] Layout to Support Page Skipping".
296    ///
297    /// Defaults to [`PageIndexPolicy::Optional`]
298    ///
299    /// This requires
300    /// 1. The Parquet file to have been written with page indexes
301    /// 2. Additional data to be pushed into the decoder (as the page indexes are not part of the thrift footer)
302    ///
303    /// [Parquet page index]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
304    pub fn with_page_index_policy(mut self, page_index_policy: PageIndexPolicy) -> Self {
305        self.column_index_policy = page_index_policy;
306        self.offset_index_policy = page_index_policy;
307        self
308    }
309
310    /// Set the policy for reading the ColumnIndex (part of the PageIndex)
311    pub fn with_column_index_policy(mut self, column_index_policy: PageIndexPolicy) -> Self {
312        self.column_index_policy = column_index_policy;
313        self
314    }
315
316    /// Set the policy for reading the OffsetIndex (part of the PageIndex)
317    pub fn with_offset_index_policy(mut self, offset_index_policy: PageIndexPolicy) -> Self {
318        self.offset_index_policy = offset_index_policy;
319        self
320    }
321
322    // TODO: deprecate the following as part of https://github.com/apache/arrow-rs/issues/11395
323
324    /// Select the row groups and columns for which both page index structures are read.
325    pub fn with_page_index_mask(mut self, mask: ColumnChunkMask) -> Self {
326        self.column_index_mask = mask.clone();
327        self.offset_index_mask = mask;
328        self
329    }
330
331    /// Select the row groups and columns for which column indexes are read.
332    pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
333        self.column_index_mask = mask;
334        self
335    }
336
337    /// Select the row groups and columns for which offset indexes are read.
338    pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
339        self.offset_index_mask = mask;
340        self
341    }
342
343    /// Set the options to use when decoding the Parquet metadata.
344    pub fn with_metadata_options(mut self, options: Option<Arc<ParquetMetaDataOptions>>) -> Self {
345        self.metadata_parser = self.metadata_parser.with_metadata_options(options);
346        self
347    }
348
349    #[cfg(feature = "encryption")]
350    /// Provide decryption properties for decoding encrypted Parquet files
351    pub fn with_file_decryption_properties(
352        mut self,
353        file_decryption_properties: Option<std::sync::Arc<FileDecryptionProperties>>,
354    ) -> Self {
355        self.metadata_parser = self
356            .metadata_parser
357            .with_file_decryption_properties(file_decryption_properties);
358        self
359    }
360
361    /// Push the data into the decoder's buffer.
362    ///
363    /// The decoder does not immediately attempt to decode the metadata
364    /// after pushing data. Instead, it accumulates the pushed data until you
365    /// call [`Self::try_decode`].
366    ///
367    /// # Determining required data:
368    ///
369    /// To determine what ranges are required to decode the metadata, you can
370    /// either:
371    ///
372    /// 1. Call [`Self::try_decode`] first to get the exact ranges required (see
373    ///    example on [`Self`])
374    ///
375    /// 2. Speculatively push any data that you have available, which may
376    ///    include more than the footer data or requested bytes.
377    ///
378    /// Speculatively pushing data can be used when  "prefetching" data. See
379    /// example on [`Self`]
380    pub fn push_ranges(&mut self, ranges: Vec<Range<u64>>, buffers: Vec<Bytes>) -> Result<()> {
381        if matches!(&self.state, DecodeState::Finished) {
382            return Err(general_err!(
383                "ParquetMetaDataPushDecoder: cannot push data after decoding is finished"
384            ));
385        }
386        self.buffers.push_ranges(ranges, buffers)?;
387        Ok(())
388    }
389
390    /// Pushes a single range of data into the decoder's buffer.
391    pub fn push_range(&mut self, range: Range<u64>, buffer: Bytes) -> Result<()> {
392        if matches!(&self.state, DecodeState::Finished) {
393            return Err(general_err!(
394                "ParquetMetaDataPushDecoder: cannot push data after decoding is finished"
395            ));
396        }
397        self.buffers.push_range(range, buffer)?;
398        Ok(())
399    }
400
401    /// Clear any staged byte ranges currently buffered for future decode work.
402    pub fn clear_all_ranges(&mut self) {
403        self.buffers.clear_all_ranges();
404    }
405
406    /// Try to decode the metadata from the pushed data, returning the
407    /// decoded metadata or an error if not enough data is available.
408    pub fn try_decode(&mut self) -> Result<DecodeResult<ParquetMetaData>> {
409        let file_len = self.buffers.file_len();
410        let footer_len = FOOTER_SIZE as u64;
411        loop {
412            match std::mem::replace(&mut self.state, DecodeState::Intermediate) {
413                DecodeState::ReadingFooter => {
414                    // need to have the last 8 bytes of the file to decode the metadata
415                    let footer_start = file_len.saturating_sub(footer_len);
416                    let footer_range = footer_start..file_len;
417
418                    if !self.buffers.has_range(&footer_range) {
419                        self.state = DecodeState::ReadingFooter;
420                        return Ok(needs_range(footer_range));
421                    }
422                    let footer_bytes = self.get_bytes(&footer_range)?;
423                    let footer_tail = FooterTail::try_from(footer_bytes.as_ref())?;
424
425                    self.state = DecodeState::ReadingMetadata(footer_tail);
426                }
427
428                DecodeState::ReadingMetadata(footer_tail) => {
429                    let metadata_len: u64 = footer_tail.metadata_length() as u64;
430                    // A corrupt footer can claim more metadata than the file holds
431                    let needed = footer_len + metadata_len;
432                    let Some(metadata_start) = file_len.checked_sub(needed) else {
433                        return Err(eof_err!(
434                            "Parquet file too small. Size is {file_len} but need {needed}"
435                        ));
436                    };
437                    let metadata_end = metadata_start + metadata_len;
438                    let metadata_range = metadata_start..metadata_end;
439
440                    if !self.buffers.has_range(&metadata_range) {
441                        self.state = DecodeState::ReadingMetadata(footer_tail);
442                        return Ok(needs_range(metadata_range));
443                    }
444
445                    let metadata = self.metadata_parser.decode_metadata(
446                        &self.get_bytes(&metadata_range)?,
447                        footer_tail.is_encrypted_footer(),
448                    )?;
449                    // Note: ReadingPageIndex first checks if page indexes are needed
450                    // and is a no-op if not
451                    self.state = DecodeState::ReadingPageIndex(Box::new(metadata));
452                }
453
454                DecodeState::ReadingPageIndex(metadata) => {
455                    match self.try_decode_page_index_for_metadata(metadata)? {
456                        DecodeResult::NeedsData(ranges) => {
457                            return Ok(DecodeResult::NeedsData(ranges));
458                        }
459                        DecodeResult::Data((mut metadata, page_index)) => {
460                            // Install the newly decoded indexes, or clear any existing indexes if
461                            // none were requested or present in the file.
462                            if let Some(page_index) = page_index {
463                                metadata.set_page_index(Some(Arc::new(page_index)));
464                            } else {
465                                metadata.set_page_index(None);
466                            }
467                            return Ok(DecodeResult::Data(*metadata));
468                        }
469                        DecodeResult::Finished => return Ok(DecodeResult::Finished),
470                    }
471                }
472
473                DecodeState::Finished => return Ok(DecodeResult::Finished),
474                DecodeState::Intermediate => {
475                    return Err(general_err!(
476                        "ParquetMetaDataPushDecoder: internal error, invalid state"
477                    ));
478                }
479            }
480        }
481    }
482
483    /// Try to decode the page indexes from the pushed data, returning the stored
484    /// [`ParquetMetaData`] together with the decoded [`PageIndex`].
485    ///
486    /// This method can only be called once the footer metadata is known: either
487    /// after [`Self::try_new_with_metadata`], or after [`Self::try_decode`] has
488    /// decoded the footer metadata and requested page index data. Calling it
489    /// earlier returns an error and leaves the decoder unchanged.
490    ///
491    /// This method is useful when you want to:
492    /// 1. First decode the footer metadata without reading the page indexes
493    ///    (using [`PageIndexPolicy::Skip`])
494    /// 2. Later selectively decode a subset of the page indexes using
495    ///    [`Self::with_column_index_mask`] or [`Self::with_offset_index_mask`]
496    ///
497    /// When the data can be accessed through a [`ChunkReader`], prefer
498    /// [`ParquetMetaDataReader::read_page_index_and_finish`] (or
499    /// [`ParquetMetaDataReader::load_page_index_and_finish`] for asynchronous
500    /// I/O), which handles the decoder state machine and I/O requests. See the
501    /// [`read_page_index_and_finish` example][`ParquetMetaDataReader::read_page_index_and_finish`]
502    /// for the complete two-stage workflow.
503    ///
504    /// # Returns
505    ///
506    /// * `Ok(DecodeResult::Data((metadata, Some(page_index))))` - Successfully decoded the
507    ///   requested page indexes based on the configured policies and masks. The newly decoded page
508    ///   index is returned separately; any page index already attached to `metadata` is unchanged.
509    /// * `Ok(DecodeResult::Data((metadata, None)))` - No page indexes were requested or found
510    ///   according to the configured policies and masks. The stored metadata is still returned.
511    /// * `Ok(DecodeResult::NeedsData(ranges))` - More data is needed to decode
512    ///   the page indexes. Push the requested ranges and call this method again.
513    /// * `Ok(DecodeResult::Finished)` - The decoder has finished and no more
514    ///   data can be decoded
515    /// * `Err(_)` - An error occurred during decoding
516    ///
517    /// [`ParquetMetaDataReader::read_page_index_and_finish`]: crate::file::metadata::ParquetMetaDataReader::read_page_index_and_finish
518    /// [`ParquetMetaDataReader::load_page_index_and_finish`]: crate::file::metadata::ParquetMetaDataReader::load_page_index_and_finish
519    pub fn try_decode_page_index(
520        &mut self,
521    ) -> Result<DecodeResult<(ParquetMetaData, Option<PageIndex>)>> {
522        match std::mem::replace(&mut self.state, DecodeState::Intermediate) {
523            DecodeState::ReadingPageIndex(metadata) => {
524                Ok(match self.try_decode_page_index_for_metadata(metadata)? {
525                    DecodeResult::NeedsData(ranges) => DecodeResult::NeedsData(ranges),
526                    DecodeResult::Data((metadata, page_index)) => {
527                        DecodeResult::Data((*metadata, page_index))
528                    }
529                    DecodeResult::Finished => DecodeResult::Finished,
530                })
531            }
532            DecodeState::Finished => {
533                self.state = DecodeState::Finished;
534                Ok(DecodeResult::Finished)
535            }
536            other => {
537                // Restore the original state so the decoder remains usable
538                self.state = other;
539                Err(general_err!(
540                    "ParquetMetaDataPushDecoder: try_decode_page_index requires a decoder created with try_new_with_metadata"
541                ))
542            }
543        }
544    }
545
546    fn try_decode_page_index_for_metadata(
547        &mut self,
548        metadata: Box<ParquetMetaData>,
549    ) -> Result<DecodeResult<(Box<ParquetMetaData>, Option<PageIndex>)>> {
550        let range = range_for_page_index(
551            &metadata,
552            self.column_index_policy,
553            self.offset_index_policy,
554            &self.column_index_mask,
555            &self.offset_index_mask,
556        );
557
558        let Some(page_index_range) = range else {
559            self.state = DecodeState::Finished;
560            return Ok(DecodeResult::Data((metadata, None)));
561        };
562
563        if !self.buffers.has_range(&page_index_range) {
564            self.state = DecodeState::ReadingPageIndex(metadata);
565            return Ok(DecodeResult::NeedsData(vec![page_index_range]));
566        }
567
568        let buffer = match self.get_bytes(&page_index_range) {
569            Ok(buffer) => buffer,
570            Err(error) => {
571                self.state = DecodeState::ReadingPageIndex(metadata);
572                return Err(error);
573            }
574        };
575        let page_index = match parse_page_index(
576            &metadata,
577            self.column_index_policy,
578            self.offset_index_policy,
579            &self.column_index_mask,
580            &self.offset_index_mask,
581            &buffer,
582            page_index_range.start,
583        ) {
584            Ok(page_index) => page_index,
585            Err(error) => {
586                self.state = DecodeState::ReadingPageIndex(metadata);
587                return Err(error);
588            }
589        };
590
591        self.state = DecodeState::Finished;
592        Ok(DecodeResult::Data((metadata, page_index)))
593    }
594
595    /// Returns the bytes for the given range from the internal buffer
596    fn get_bytes(&self, range: &Range<u64>) -> Result<Bytes> {
597        let start = range.start;
598        let raw_len = range.end - range.start;
599        let len: usize = raw_len.try_into().map_err(|_| {
600            ParquetError::General(format!(
601                "ParquetMetaDataPushDecoder: Range length too large to fit in usize: {raw_len}",
602            ))
603        })?;
604        self.buffers.get_bytes(start, len)
605    }
606}
607
608/// returns a DecodeResults that describes needing the given range
609fn needs_range<T: std::fmt::Debug>(range: Range<u64>) -> DecodeResult<T> {
610    DecodeResult::NeedsData(vec![range])
611}
612
613/// Decoding state machine
614#[derive(Debug)]
615enum DecodeState {
616    /// Reading the last 8 bytes of the file
617    ReadingFooter,
618    /// Reading the metadata thrift structure
619    ReadingMetadata(FooterTail),
620    // Actively reading the page index
621    ReadingPageIndex(Box<ParquetMetaData>),
622    // Decoding is complete
623    Finished,
624    /// State left during the `try_decode` method so something valid is present.
625    /// This state should never be observed.
626    Intermediate,
627}
628
629/// Returns the byte range needed to read the offset/page indexes, based on the
630/// specified policies and masks
631///
632/// Returns None if no page indexes are needed
633pub(crate) fn range_for_page_index(
634    metadata: &ParquetMetaData,
635    column_index_policy: PageIndexPolicy,
636    offset_index_policy: PageIndexPolicy,
637    column_index_mask: &ColumnChunkMask,
638    offset_index_mask: &ColumnChunkMask,
639) -> Option<Range<u64>> {
640    let mut range: Option<Range<u64>> = None;
641    for (policy, mask, column_index) in [
642        (column_index_policy, column_index_mask, true),
643        (offset_index_policy, offset_index_mask, false),
644    ] {
645        if policy == PageIndexPolicy::Skip {
646            continue;
647        }
648        for row_group in mask.row_group_indices(metadata.num_row_groups()) {
649            for column in mask.column_indices(metadata.row_group(row_group).num_columns()) {
650                let column = metadata.row_group(row_group).column(column);
651                let index = if column_index {
652                    column.column_index_range()
653                } else {
654                    column.offset_index_range()
655                };
656                if let Some(index) = index {
657                    range = Some(match range {
658                        Some(range) => range.start.min(index.start)..range.end.max(index.end),
659                        None => index,
660                    });
661                }
662            }
663        }
664    }
665    range
666}
667
668// These tests use the arrow writer to create a parquet file in memory
669// so they need the arrow feature and the test feature
670#[cfg(all(test, feature = "arrow"))]
671mod tests {
672    use super::*;
673    use crate::arrow::ArrowWriter;
674    use crate::file::metadata::page_index::PageIndexProvider;
675    use crate::file::properties::WriterProperties;
676    use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringViewArray};
677    use bytes::Bytes;
678    use std::fmt::Debug;
679    use std::ops::Range;
680    use std::sync::{Arc, LazyLock};
681
682    /// It is possible to decode the metadata from the entire file at once before being asked
683    #[test]
684    fn test_metadata_decoder_all_data() {
685        let file_len = test_file_len();
686        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
687        // Push the entire file data into the metadata decoder
688        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![test_file_range()]);
689
690        // should be able to decode the metadata without needing more data
691        let metadata = expect_data(metadata_decoder.try_decode());
692
693        assert_eq!(metadata.num_row_groups(), 2);
694        assert_eq!(metadata.row_group(0).num_rows(), 200);
695        assert_eq!(metadata.row_group(1).num_rows(), 200);
696        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
697    }
698
699    /// A footer whose metadata length exceeds the file length is an error, not a panic
700    #[test]
701    fn test_metadata_decoder_metadata_length_exceeds_file() {
702        // 4 bytes of data, then a footer claiming 1000 bytes of metadata
703        let mut file = vec![0u8; 4];
704        file.extend_from_slice(&1000u32.to_le_bytes());
705        file.extend_from_slice(b"PAR1");
706        let file_len = file.len() as u64;
707
708        let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
709        decoder.push_range(0..file_len, Bytes::from(file)).unwrap();
710        let err = decoder.try_decode().unwrap_err();
711        assert_eq!(
712            err.to_string(),
713            "EOF: Parquet file too small. Size is 12 but need 1008"
714        );
715    }
716
717    #[test]
718    fn test_range_for_page_index_respects_column_mask() {
719        let mut decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap();
720        push_ranges_to_metadata_decoder(&mut decoder, vec![test_file_range()]);
721        let metadata = expect_data(decoder.try_decode());
722
723        let all = range_for_page_index(
724            &metadata,
725            PageIndexPolicy::Required,
726            PageIndexPolicy::Skip,
727            &ColumnChunkMask::all(),
728            &ColumnChunkMask::all(),
729        )
730        .unwrap();
731        let masked = range_for_page_index(
732            &metadata,
733            PageIndexPolicy::Required,
734            PageIndexPolicy::Skip,
735            &ColumnChunkMask::columns([0]),
736            &ColumnChunkMask::all(),
737        )
738        .unwrap();
739
740        assert!(all.start <= masked.start && masked.end <= all.end);
741        assert!(masked.end - masked.start < all.end - all.start);
742    }
743
744    #[test]
745    fn test_decode_page_index_after_metadata() {
746        let metadata = test_metadata_without_page_index();
747        assert!(metadata.page_index().is_none());
748
749        let mut decoder =
750            ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(), metadata)
751                .unwrap()
752                .with_column_index_policy(PageIndexPolicy::Required)
753                .with_offset_index_policy(PageIndexPolicy::Required)
754                .with_column_index_mask(ColumnChunkMask::row_groups_and_columns([1], [0, 2]))
755                .with_offset_index_mask(ColumnChunkMask::row_groups_and_columns([0], [1]));
756
757        let ranges = expect_needs_data(decoder.try_decode_page_index());
758        assert_eq!(ranges.len(), 1);
759        push_ranges_to_metadata_decoder(&mut decoder, ranges);
760
761        let (metadata, page_index) = expect_data(decoder.try_decode_page_index());
762        assert!(metadata.page_index().is_none());
763        let page_index = page_index.unwrap();
764        expect_finished(decoder.try_decode_page_index());
765
766        for row_group_idx in 0..2 {
767            for column_idx in 0..3 {
768                assert_eq!(
769                    page_index.column_index(row_group_idx, column_idx).is_some(),
770                    row_group_idx == 1 && matches!(column_idx, 0 | 2)
771                );
772                assert_eq!(
773                    page_index.offset_index(row_group_idx, column_idx).is_some(),
774                    row_group_idx == 0 && column_idx == 1
775                );
776            }
777        }
778    }
779
780    #[test]
781    fn test_page_index_decoder_no_indexes_requested() {
782        let metadata = test_metadata_without_page_index();
783        let mut decoder =
784            ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(), metadata)
785                .unwrap()
786                .with_page_index_policy(PageIndexPolicy::Skip);
787
788        assert!(expect_data(decoder.try_decode_page_index()).1.is_none());
789        expect_finished(decoder.try_decode_page_index());
790
791        let metadata = test_metadata_without_page_index();
792        let mut decoder =
793            ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(), metadata)
794                .unwrap()
795                .with_page_index_mask(ColumnChunkMask::none());
796
797        assert!(expect_data(decoder.try_decode_page_index()).1.is_none());
798        expect_finished(decoder.try_decode_page_index());
799    }
800
801    #[test]
802    fn test_page_index_decoder_requires_metadata() {
803        let mut decoder = ParquetMetaDataPushDecoder::try_new(test_file_len()).unwrap();
804        let err = decoder.try_decode_page_index().unwrap_err();
805        assert!(
806            err.to_string()
807                .contains("requires a decoder created with try_new_with_metadata"),
808            "unexpected error: {err}"
809        );
810    }
811
812    #[test]
813    fn test_page_index_decoder_restores_state_after_decode_error() {
814        let metadata = test_metadata_without_page_index();
815        let mut decoder =
816            ParquetMetaDataPushDecoder::try_new_with_metadata(test_file_len(), metadata)
817                .unwrap()
818                .with_page_index_policy(PageIndexPolicy::Required);
819
820        let ranges = expect_needs_data(decoder.try_decode_page_index());
821        let invalid_data = ranges
822            .iter()
823            .map(|range| Bytes::from(vec![0; (range.end - range.start) as usize]))
824            .collect();
825        decoder.push_ranges(ranges.clone(), invalid_data).unwrap();
826        decoder.try_decode_page_index().unwrap_err();
827
828        decoder.clear_all_ranges();
829        push_ranges_to_metadata_decoder(&mut decoder, ranges);
830        let (metadata, page_index) = expect_data(decoder.try_decode_page_index());
831        assert!(metadata.page_index().is_none());
832        assert!(page_index.is_some_and(|index| index.is_complete()));
833    }
834
835    /// A failed call to `try_decode_page_index` must not leave the decoder in
836    /// an unusable state: decoding the metadata via `try_decode` should still
837    /// work afterwards, at every stage of the footer decoding state machine.
838    #[test]
839    fn test_decode_after_failed_page_index_decode() {
840        let file_len = test_file_len();
841        let mut decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
842
843        // Before any data is pushed, the decoder is waiting for the footer
844        decoder.try_decode_page_index().unwrap_err();
845        let footer_ranges = expect_needs_data(decoder.try_decode());
846        assert_eq!(footer_ranges.len(), 1);
847        push_ranges_to_metadata_decoder(&mut decoder, footer_ranges);
848
849        // The footer is available but not yet decoded
850        decoder.try_decode_page_index().unwrap_err();
851        let metadata_ranges = expect_needs_data(decoder.try_decode());
852        assert_eq!(metadata_ranges.len(), 1);
853        push_ranges_to_metadata_decoder(&mut decoder, metadata_ranges);
854
855        // The metadata is available but not yet decoded
856        decoder.try_decode_page_index().unwrap_err();
857        let index_ranges = expect_needs_data(decoder.try_decode());
858        assert_eq!(index_ranges.len(), 1);
859
860        // The metadata is decoded and the decoder is waiting for page index
861        // data, so try_decode_page_index is now valid and asks for the same range
862        let index_ranges_again = expect_needs_data(decoder.try_decode_page_index());
863        assert_eq!(index_ranges, index_ranges_again);
864        push_ranges_to_metadata_decoder(&mut decoder, index_ranges);
865
866        // The regular decode path still works and attaches the page index
867        let metadata = expect_data(decoder.try_decode());
868        assert_eq!(metadata.num_row_groups(), 2);
869        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
870
871        // Once finished, both entry points keep reporting Finished
872        expect_finished(decoder.try_decode_page_index());
873        expect_finished(decoder.try_decode_page_index());
874        expect_finished(decoder.try_decode());
875    }
876
877    /// It is possible to feed some, but not all, of the footer into the metadata decoder
878    /// before asked. This avoids multiple IO requests
879    #[test]
880    fn test_metadata_decoder_prefetch_success() {
881        let file_len = test_file_len();
882        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
883        // simulate pre-fetching the last 2k bytes of the file without asking the decoder
884        let prefetch_range = (file_len - 2 * 1024)..file_len;
885        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![prefetch_range]);
886
887        // expect the decoder has enough data to decode the metadata
888        let metadata = expect_data(metadata_decoder.try_decode());
889        expect_finished(metadata_decoder.try_decode());
890        assert_eq!(metadata.num_row_groups(), 2);
891        assert_eq!(metadata.row_group(0).num_rows(), 200);
892        assert_eq!(metadata.row_group(1).num_rows(), 200);
893        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
894    }
895
896    /// It is possible to pre-fetch some, but not all, of the necessary data
897    /// data
898    #[test]
899    fn test_metadata_decoder_prefetch_retry() {
900        let file_len = test_file_len();
901        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
902        // simulate pre-fetching the last 1500 bytes of the file.
903        // this is enough to read the footer thrift metadata, but not the offset indexes
904        let prefetch_range = (file_len - 1500)..file_len;
905        push_ranges_to_metadata_decoder(&mut metadata_decoder, vec![prefetch_range]);
906
907        // expect another request is needed to read the offset indexes (note
908        // try_decode only returns NeedsData once, whereas without any prefetching it would
909        // return NeedsData three times)
910        let ranges = expect_needs_data(metadata_decoder.try_decode());
911        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
912
913        // expect the decoder has enough data to decode the metadata
914        let metadata = expect_data(metadata_decoder.try_decode());
915        expect_finished(metadata_decoder.try_decode());
916
917        assert_eq!(metadata.num_row_groups(), 2);
918        assert_eq!(metadata.row_group(0).num_rows(), 200);
919        assert_eq!(metadata.row_group(1).num_rows(), 200);
920        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
921    }
922
923    #[test]
924    fn test_metadata_decoder_clear_all_ranges() {
925        let file_len = test_file_len();
926        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
927
928        metadata_decoder
929            .push_range(test_file_range(), TEST_FILE_DATA.clone())
930            .unwrap();
931        assert_eq!(metadata_decoder.buffers.buffered_bytes(), test_file_len());
932
933        metadata_decoder.clear_all_ranges();
934        assert_eq!(metadata_decoder.buffers.buffered_bytes(), 0);
935
936        let ranges = expect_needs_data(metadata_decoder.try_decode());
937        assert_eq!(ranges, vec![test_file_len() - 8..test_file_len()]);
938    }
939
940    /// Decode the metadata incrementally, simulating a scenario where exactly the data needed
941    /// is read in each step
942    #[test]
943    fn test_metadata_decoder_incremental() {
944        let file_len = TEST_FILE_DATA.len() as u64;
945        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len).unwrap();
946        let ranges = expect_needs_data(metadata_decoder.try_decode());
947        assert_eq!(ranges.len(), 1);
948        assert_eq!(ranges[0], test_file_len() - 8..test_file_len());
949        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
950
951        // expect the first request to read the footer
952        let ranges = expect_needs_data(metadata_decoder.try_decode());
953        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
954
955        // expect the second request to read the offset indexes
956        let ranges = expect_needs_data(metadata_decoder.try_decode());
957        // since we're reading all indexes, there should be a single range needed
958        assert_eq!(ranges.len(), 1);
959        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
960
961        // expect the third request to read the actual data
962        let metadata = expect_data(metadata_decoder.try_decode());
963        expect_finished(metadata_decoder.try_decode());
964
965        assert_eq!(metadata.num_row_groups(), 2);
966        assert_eq!(metadata.row_group(0).num_rows(), 200);
967        assert_eq!(metadata.row_group(1).num_rows(), 200);
968        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
969    }
970
971    /// Decode the metadata incrementally, but without reading the page indexes
972    /// (so only two requests)
973    #[test]
974    fn test_metadata_decoder_incremental_no_page_index() {
975        let file_len = TEST_FILE_DATA.len() as u64;
976        let mut metadata_decoder = ParquetMetaDataPushDecoder::try_new(file_len)
977            .unwrap()
978            .with_page_index_policy(PageIndexPolicy::Skip);
979        let ranges = expect_needs_data(metadata_decoder.try_decode());
980        assert_eq!(ranges.len(), 1);
981        assert_eq!(ranges[0], test_file_len() - 8..test_file_len());
982        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
983
984        // expect the first request to read the footer
985        let ranges = expect_needs_data(metadata_decoder.try_decode());
986        push_ranges_to_metadata_decoder(&mut metadata_decoder, ranges);
987
988        // expect NO second request to read the offset indexes, should just cough up the metadata
989        let metadata = expect_data(metadata_decoder.try_decode());
990        expect_finished(metadata_decoder.try_decode());
991
992        assert_eq!(metadata.num_row_groups(), 2);
993        assert_eq!(metadata.row_group(0).num_rows(), 200);
994        assert_eq!(metadata.row_group(1).num_rows(), 200);
995        assert!(metadata.page_index().is_none()); // of course, we did not read the page index
996    }
997
998    static TEST_BATCH: LazyLock<RecordBatch> = LazyLock::new(|| {
999        // Input batch has 400 rows, with 3 columns: "a", "b", "c"
1000        // Note c is a different types (so the data page sizes will be different)
1001        let a: ArrayRef = Arc::new(Int64Array::from_iter_values(0..400));
1002        let b: ArrayRef = Arc::new(Int64Array::from_iter_values(400..800));
1003        let c: ArrayRef = Arc::new(StringViewArray::from_iter_values((0..400).map(|i| {
1004            if i % 2 == 0 {
1005                format!("string_{i}")
1006            } else {
1007                format!("A string larger than 12 bytes and thus not inlined {i}")
1008            }
1009        })));
1010
1011        RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
1012    });
1013
1014    /// Create a parquet file in memory for testing. See [`test_file_range`] for details.
1015    static TEST_FILE_DATA: LazyLock<Bytes> = LazyLock::new(|| {
1016        let input_batch = &TEST_BATCH;
1017        let mut output = Vec::new();
1018
1019        let writer_options = WriterProperties::builder()
1020            .set_max_row_group_row_count(Some(200))
1021            .set_data_page_row_count_limit(100)
1022            .build();
1023        let mut writer =
1024            ArrowWriter::try_new(&mut output, input_batch.schema(), Some(writer_options)).unwrap();
1025
1026        // since the limits are only enforced on batch boundaries, write the input
1027        // batch in chunks of 50
1028        let mut row_remain = input_batch.num_rows();
1029        while row_remain > 0 {
1030            let chunk_size = row_remain.min(50);
1031            let chunk = input_batch.slice(input_batch.num_rows() - row_remain, chunk_size);
1032            writer.write(&chunk).unwrap();
1033            row_remain -= chunk_size;
1034        }
1035        writer.close().unwrap();
1036        Bytes::from(output)
1037    });
1038
1039    /// Return the length of the test file in bytes
1040    fn test_file_len() -> u64 {
1041        TEST_FILE_DATA.len() as u64
1042    }
1043
1044    /// Return the range of the entire test file
1045    fn test_file_range() -> Range<u64> {
1046        0..test_file_len()
1047    }
1048
1049    /// Return a slice of the test file data from the given range
1050    pub fn test_file_slice(range: Range<u64>) -> Bytes {
1051        let start: usize = range.start.try_into().unwrap();
1052        let end: usize = range.end.try_into().unwrap();
1053        TEST_FILE_DATA.slice(start..end)
1054    }
1055
1056    /// Push the given ranges to the metadata decoder, simulating reading from a file
1057    fn push_ranges_to_metadata_decoder(
1058        metadata_decoder: &mut ParquetMetaDataPushDecoder,
1059        ranges: Vec<Range<u64>>,
1060    ) {
1061        let data = ranges
1062            .iter()
1063            .map(|range| test_file_slice(range.clone()))
1064            .collect::<Vec<_>>();
1065        metadata_decoder.push_ranges(ranges, data).unwrap();
1066    }
1067
1068    fn test_metadata_without_page_index() -> ParquetMetaData {
1069        let mut decoder = ParquetMetaDataPushDecoder::try_new(test_file_len())
1070            .unwrap()
1071            .with_page_index_policy(PageIndexPolicy::Skip);
1072        push_ranges_to_metadata_decoder(&mut decoder, vec![test_file_range()]);
1073        let metadata = expect_data(decoder.try_decode());
1074        assert!(metadata.page_index().is_none());
1075        metadata
1076    }
1077
1078    /// Expect that the [`DecodeResult`] is a [`DecodeResult::Data`] and return the corresponding element
1079    fn expect_data<T: Debug>(result: Result<DecodeResult<T>>) -> T {
1080        match result.expect("Expected Ok(DecodeResult::Data(T))") {
1081            DecodeResult::Data(data) => data,
1082            result => panic!("Expected DecodeResult::Data, got {result:?}"),
1083        }
1084    }
1085
1086    /// Expect that the [`DecodeResult`] is a [`DecodeResult::NeedsData`] and return the corresponding ranges
1087    fn expect_needs_data<T: Debug>(result: Result<DecodeResult<T>>) -> Vec<Range<u64>> {
1088        match result.expect("Expected Ok(DecodeResult::NeedsData(ranges))") {
1089            DecodeResult::NeedsData(ranges) => ranges,
1090            result => panic!("Expected DecodeResult::NeedsData, got {result:?}"),
1091        }
1092    }
1093
1094    fn expect_finished<T: Debug>(result: Result<DecodeResult<T>>) {
1095        match result.expect("Expected Ok(DecodeResult::Finished)") {
1096            DecodeResult::Finished => {}
1097            result => panic!("Expected DecodeResult::Finished, got {result:?}"),
1098        }
1099    }
1100}
parquet/file/metadata/push_decoder.rs

parquet/file/metadata/
push_decoder.rs
