Skip to main content

parquet/file/metadata/
reader.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5#[cfg(feature = "arrow")]
6use crate::arrow::ProjectionMask;
7#[cfg(feature = "encryption")]
8use crate::encryption::decrypt::FileDecryptionProperties;
9use crate::errors::{ParquetError, Result};
10use crate::file::FOOTER_SIZE;
11#[cfg(feature = "arrow")]
12use crate::file::metadata::dictionary::decode_dictionary_page;
13use crate::file::metadata::page_index::{PageIndex, PageIndexProvider};
14use crate::file::metadata::parser::decode_metadata;
15use crate::file::metadata::thrift::parquet_schema_from_bytes;
16use crate::file::metadata::{
17    FooterTail, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataPushDecoder,
18};
19use crate::file::reader::ChunkReader;
20use crate::schema::types::SchemaDescriptor;
21#[cfg(feature = "arrow")]
22use arrow_array::ArrayRef;
23use bytes::Bytes;
24use std::collections::BTreeSet;
25use std::iter::FusedIterator;
26use std::slice::Iter;
27use std::sync::Arc;
28use std::{io::Read, ops::Range};
29
30use crate::DecodeResult;
31#[cfg(all(feature = "async", feature = "arrow"))]
32use crate::arrow::async_reader::{MetadataFetch, MetadataSuffixFetch};
33
34/// Reads [`ParquetMetaData`] from a byte stream, with either synchronous or
35/// asynchronous I/O.
36///
37/// There are two flavors of APIs:
38/// * Synchronous: [`Self::try_parse()`], [`Self::try_parse_sized()`], [`Self::parse_and_finish()`], etc.
39/// * Asynchronous (requires `async` and `arrow` features): [`Self::try_load()`], etc
40///
41///  See the [`ParquetMetaDataPushDecoder`] for an API that does not require I/O.
42///
43/// # Format Notes
44///
45/// Parquet metadata is not necessarily contiguous in a Parquet file: a portion is stored
46/// in the footer (the last bytes of the file), but other portions (such as the
47/// PageIndex) can be stored elsewhere.
48/// See [`crate::file::metadata::ParquetMetaDataWriter#output-format`] for more details of
49/// Parquet metadata.
50///
51/// This reader handles reading the footer as well as the non contiguous parts
52/// of the metadata (`PageIndex` and `ColumnIndex`). It does not handle reading Bloom Filters.
53///
54/// # Example
55/// ```no_run
56/// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
57/// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); }
58/// // read parquet metadata including page indexes from a file
59/// let file = open_parquet_file("some_path.parquet");
60/// let mut reader = ParquetMetaDataReader::new()
61///     .with_page_index_policy(PageIndexPolicy::Required);
62/// reader.try_parse(&file).unwrap();
63/// let metadata = reader.finish().unwrap();
64/// assert!(metadata.page_index().is_some());
65/// ```
66#[derive(Default, Debug)]
67pub struct ParquetMetaDataReader {
68    metadata: Option<ParquetMetaData>,
69    column_index: PageIndexPolicy,
70    offset_index: PageIndexPolicy,
71    column_index_mask: ColumnChunkMask,
72    offset_index_mask: ColumnChunkMask,
73    prefetch_hint: Option<usize>,
74    metadata_options: Option<Arc<ParquetMetaDataOptions>>,
75    // Size of the serialized thrift metadata plus the 8 byte footer. Only set if
76    // `self.parse_metadata` is called.
77    metadata_size: Option<usize>,
78    #[cfg(feature = "encryption")]
79    file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
80}
81
82// TODO: https://github.com/apache/arrow-rs/issues/11395 is tracking unifying PageIndexPolicy
83// with the ColumnChunkMask to reduce the dizzying number of APIs that exist to control reading
84// the page indexes.
85
86/// Describes the policy for reading page indexes
87#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
88pub enum PageIndexPolicy {
89    /// Do not read the page index.
90    #[default]
91    Skip,
92    /// Read the page index if it exists, otherwise do not error.
93    Optional,
94    /// Require offset indexes for every selected column chunk when any selected
95    /// page index ranges exist.
96    ///
97    /// For compatibility, a file with no page index ranges is accepted. This is only
98    /// enforced for the offset index; this is the same as [`Self::Optional`] for the
99    /// column index.
100    Required,
101}
102
103impl From<bool> for PageIndexPolicy {
104    fn from(value: bool) -> Self {
105        match value {
106            true => Self::Required,
107            false => Self::Skip,
108        }
109    }
110}
111
112/// A cheaply cloneable struct to specify column chunks for which metadata is required.
113///
114/// Column chunks are identified by row group index and leaf column index (the index of the
115/// column in [`SchemaDescriptor::columns`], not the index of a root or Arrow field). This struct
116/// allows for specifying vertical slices of column chunk data (via [`Self::columns`]),
117/// horizontal slices (via [`Self::row_groups`]), or the intersection of the two
118/// (via [`Self::row_groups_and_columns`]).
119///
120/// Because Parquet collection cannot contain more than [`i32::MAX`] values, user provided indices
121/// outside that range will be discarded.
122///
123/// At present this is only used to select elements of the [Page Index] for decoding.
124///
125/// # Examples
126///
127/// To select columns 0 and 1 from all row groups:
128/// ```rust
129/// # use parquet::file::metadata::ColumnChunkMask;
130/// let mask = ColumnChunkMask::columns([0, 1]);
131/// ```
132///
133/// To select all columns from row group 2:
134/// ```rust
135/// # use parquet::file::metadata::ColumnChunkMask;
136/// let mask = ColumnChunkMask::row_groups([2]);
137/// ```
138///
139/// To select columns 1 and 3 from row group 0:
140/// ```rust
141/// # use parquet::file::metadata::ColumnChunkMask;
142/// let mask = ColumnChunkMask::row_groups_and_columns([0], [1, 3]);
143/// ```
144///
145/// [Page Index]: https://parquet.apache.org/docs/file-format/pageindex/
146#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
147pub struct ColumnChunkMask {
148    // `row_groups` and `columns` are optional `u32` slices, with elements
149    // sorted in ascending order (to permit binary searching for member indices).
150    // `None` means all indices are included, while an empty slice means none are.
151    // Indices are stored as `u32` because Parquet/Thrift collections cannot contain
152    // more than i32::MAX entries.
153    row_groups: Option<Arc<[u32]>>,
154    columns: Option<Arc<[u32]>>,
155}
156
157impl ColumnChunkMask {
158    /// Select all row groups and columns.
159    pub fn all() -> Self {
160        Self::default()
161    }
162
163    /// Select no row groups or columns.
164    pub fn none() -> Self {
165        Self {
166            row_groups: Some(Arc::from([])),
167            columns: Some(Arc::from([])),
168        }
169    }
170
171    /// Select only the listed columns.
172    ///
173    /// Passing an empty iterator selects no columns.
174    pub fn columns(columns: impl IntoIterator<Item = usize>) -> Self {
175        Self {
176            row_groups: None,
177            columns: Self::iter_to_set(columns),
178        }
179    }
180
181    /// Select only the listed row groups.
182    ///
183    /// Passing an empty iterator selects no row groups.
184    pub fn row_groups(row_groups: impl IntoIterator<Item = usize>) -> Self {
185        Self {
186            row_groups: Self::iter_to_set(row_groups),
187            columns: None,
188        }
189    }
190
191    /// Select only the listed row groups and columns.
192    ///
193    /// An empty iterator for either dimension selects no column chunks.
194    pub fn row_groups_and_columns(
195        row_groups: impl IntoIterator<Item = usize>,
196        columns: impl IntoIterator<Item = usize>,
197    ) -> Self {
198        Self {
199            row_groups: Self::iter_to_set(row_groups),
200            columns: Self::iter_to_set(columns),
201        }
202    }
203
204    /// Returns whether row group `idx` is selected and within the bounds `[0, num_row_groups)`.
205    pub fn includes_row_group(&self, idx: usize, num_row_groups: usize) -> bool {
206        idx < num_row_groups && Self::includes_index(self.row_groups.as_ref(), idx)
207    }
208
209    /// Returns whether column `idx` is selected and within the bounds `[0, num_columns)`.
210    pub fn includes_column(&self, idx: usize, num_columns: usize) -> bool {
211        idx < num_columns && Self::includes_index(self.columns.as_ref(), idx)
212    }
213
214    fn includes_index(keep: Option<&Arc<[u32]>>, idx: usize) -> bool {
215        // return false for out-of-bounds index
216        let Ok(idx) = i32::try_from(idx) else {
217            return false;
218        };
219        let idx = idx as u32;
220        keep.is_none_or(|keep| keep.binary_search(&idx).is_ok())
221    }
222
223    /// Returns `true` when this mask selects every column chunk.
224    pub fn is_all(&self) -> bool {
225        self.row_groups.is_none() && self.columns.is_none()
226    }
227
228    /// Returns `true` when either configured axis is empty (and thus would select no column chunks).
229    ///
230    /// This does not consider the dimensions of a particular Parquet file, so may return `false`
231    /// when [`Self::column_indices`] or [`Self::row_group_indices`] returns an empty iterator due
232    /// to all elements exceeding the provided upper bound.
233    pub fn is_none(&self) -> bool {
234        self.row_groups.as_ref().is_some_and(|i| i.is_empty())
235            || self.columns.as_ref().is_some_and(|i| i.is_empty())
236    }
237
238    /// Creates a mask selecting the leaf columns in an Arrow projection.
239    #[cfg(feature = "arrow")]
240    pub fn from_projection(projection: &ProjectionMask, schema: &SchemaDescriptor) -> Self {
241        Self::columns((0..schema.num_columns()).filter(|&i| projection.leaf_included(i)))
242    }
243
244    /// Returns an iterator over the row group indices selected by this mask. If the mask is `all`
245    /// then all indices in the range `[0, num_row_groups)` are returned.
246    pub fn row_group_indices(&self, num_row_groups: usize) -> impl Iterator<Item = usize> + '_ {
247        Self::axis_indices(self.row_groups.as_deref(), num_row_groups)
248    }
249
250    /// Returns an iterator over the column indices selected by this mask. If the mask is `all`
251    /// then all indices in the range `[0, num_columns)` are returned.
252    pub fn column_indices(&self, num_columns: usize) -> impl Iterator<Item = usize> + '_ {
253        Self::axis_indices(self.columns.as_deref(), num_columns)
254    }
255
256    fn axis_indices(axis: Option<&[u32]>, len: usize) -> AxisIndices<'_> {
257        match axis {
258            None => AxisIndices::All(0..len),
259            Some(indices) => {
260                let end = indices.partition_point(|&i| (i as usize) < len);
261                AxisIndices::Selected(indices[..end].iter())
262            }
263        }
264    }
265
266    // Used during initialization, this helper takes an iterator over `usize` elements,
267    // discards out-of-range values and then sorts.
268    fn iter_to_set(indices: impl IntoIterator<Item = usize>) -> Option<Arc<[u32]>> {
269        Some(
270            indices
271                .into_iter()
272                .filter_map(|i| i32::try_from(i).ok()) // remove out-of-bounds values
273                .map(|i| i as u32) // store as u32. safe because 0 <= i <= i32::MAX
274                .collect::<BTreeSet<_>>()
275                .into_iter()
276                .collect::<Vec<_>>()
277                .into(),
278        )
279    }
280}
281
282// helper for ColumnChunkMask selected index traversal
283enum AxisIndices<'a> {
284    All(Range<usize>),
285    Selected(Iter<'a, u32>),
286}
287
288impl Iterator for AxisIndices<'_> {
289    type Item = usize;
290
291    fn next(&mut self) -> Option<Self::Item> {
292        match self {
293            Self::All(indices) => indices.next(),
294            Self::Selected(indices) => indices.next().map(|&i| i as usize),
295        }
296    }
297
298    fn size_hint(&self) -> (usize, Option<usize>) {
299        match self {
300            Self::All(indices) => indices.size_hint(),
301            Self::Selected(indices) => indices.size_hint(),
302        }
303    }
304}
305
306impl ExactSizeIterator for AxisIndices<'_> {}
307impl FusedIterator for AxisIndices<'_> {}
308
309impl ParquetMetaDataReader {
310    /// Create a new [`ParquetMetaDataReader`]
311    pub fn new() -> Self {
312        Default::default()
313    }
314
315    /// Create a new [`ParquetMetaDataReader`] populated with a [`ParquetMetaData`] struct
316    /// obtained via other means.
317    pub fn new_with_metadata(metadata: ParquetMetaData) -> Self {
318        Self {
319            metadata: Some(metadata),
320            ..Default::default()
321        }
322    }
323
324    /// Sets the [`PageIndexPolicy`] for the column and offset indexes
325    pub fn with_page_index_policy(self, policy: PageIndexPolicy) -> Self {
326        self.with_column_index_policy(policy)
327            .with_offset_index_policy(policy)
328    }
329
330    /// Sets the [`PageIndexPolicy`] for the column index
331    pub fn with_column_index_policy(mut self, policy: PageIndexPolicy) -> Self {
332        self.column_index = policy;
333        self
334    }
335
336    /// Sets the [`PageIndexPolicy`] for the offset index
337    pub fn with_offset_index_policy(mut self, policy: PageIndexPolicy) -> Self {
338        self.offset_index = policy;
339        self
340    }
341
342    // TODO: deprecate the following as part of https://github.com/apache/arrow-rs/issues/11395
343
344    /// Selects the row groups and leaf columns for which both page index structures are read.
345    ///
346    /// The mask applies only to index structures whose policy is not [`PageIndexPolicy::Skip`].
347    /// This reader fetches one byte range covering the selected indexes, so a sparse mask
348    /// primarily reduces decoding and memory rather than necessarily reducing I/O.
349    pub fn with_page_index_mask(self, mask: ColumnChunkMask) -> Self {
350        self.with_column_index_mask(mask.clone())
351            .with_offset_index_mask(mask)
352    }
353
354    /// Selects the row groups and leaf columns for which column indexes are read.
355    ///
356    /// The mask applies only when the column index policy is not [`PageIndexPolicy::Skip`].
357    /// [`PageIndexPolicy::Required`] validates only selected column chunks.
358    pub fn with_column_index_mask(mut self, mask: ColumnChunkMask) -> Self {
359        self.column_index_mask = mask;
360        self
361    }
362
363    /// Selects the row groups and leaf columns for which offset indexes are read.
364    ///
365    /// The mask applies only when the offset index policy is not [`PageIndexPolicy::Skip`].
366    /// [`PageIndexPolicy::Required`] validates only selected column chunks.
367    pub fn with_offset_index_mask(mut self, mask: ColumnChunkMask) -> Self {
368        self.offset_index_mask = mask;
369        self
370    }
371
372    /// Sets the [`ParquetMetaDataOptions`] to use when decoding
373    pub fn with_metadata_options(mut self, options: Option<ParquetMetaDataOptions>) -> Self {
374        self.metadata_options = options.map(Arc::new);
375        self
376    }
377
378    /// Provide a hint as to the number of bytes needed to fully parse the [`ParquetMetaData`].
379    /// Only used for the asynchronous [`Self::try_load()`] method.
380    ///
381    /// By default, the reader will first fetch the last 8 bytes of the input file to obtain the
382    /// size of the footer metadata. A second fetch will be performed to obtain the needed bytes.
383    /// After parsing the footer metadata, a third fetch will be performed to obtain the bytes
384    /// needed to decode the page index structures, if they have been requested. To avoid
385    /// unnecessary fetches, `prefetch` can be set to an estimate of the number of bytes needed
386    /// to fully decode the [`ParquetMetaData`], which can reduce the number of fetch requests and
387    /// reduce latency. Setting `prefetch` too small will not trigger an error, but will result
388    /// in extra fetches being performed.
389    pub fn with_prefetch_hint(mut self, prefetch: Option<usize>) -> Self {
390        self.prefetch_hint = prefetch;
391        self
392    }
393
394    /// Provide the FileDecryptionProperties to use when decrypting the file.
395    ///
396    /// This is only necessary when the file is encrypted.
397    #[cfg(feature = "encryption")]
398    pub fn with_decryption_properties(
399        mut self,
400        properties: Option<std::sync::Arc<FileDecryptionProperties>>,
401    ) -> Self {
402        self.file_decryption_properties = properties;
403        self
404    }
405
406    /// Indicates whether this reader has a [`ParquetMetaData`] internally.
407    pub fn has_metadata(&self) -> bool {
408        self.metadata.is_some()
409    }
410
411    /// Return the parsed [`ParquetMetaData`] struct, leaving `None` in its place.
412    pub fn finish(&mut self) -> Result<ParquetMetaData> {
413        self.metadata
414            .take()
415            .ok_or_else(|| general_err!("could not parse parquet metadata"))
416    }
417
418    /// Given a [`ChunkReader`], parse and return the [`ParquetMetaData`] in a single pass.
419    ///
420    /// If `reader` is [`Bytes`] based, then the buffer must contain sufficient bytes to complete
421    /// the request, and must include the Parquet footer. If page indexes are desired, the buffer
422    /// must contain the entire file, or [`Self::try_parse_sized()`] should be used.
423    ///
424    /// This call will consume `self`.
425    ///
426    /// # Example
427    /// ```no_run
428    /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
429    /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); }
430    /// // read parquet metadata including page indexes
431    /// let file = open_parquet_file("some_path.parquet");
432    /// let metadata = ParquetMetaDataReader::new()
433    ///     .with_page_index_policy(PageIndexPolicy::Required)
434    ///     .parse_and_finish(&file).unwrap();
435    /// ```
436    pub fn parse_and_finish<R: ChunkReader>(mut self, reader: &R) -> Result<ParquetMetaData> {
437        self.try_parse(reader)?;
438        self.finish()
439    }
440
441    /// Attempts to parse the footer metadata (and optionally page indexes) given a [`ChunkReader`].
442    ///
443    /// If `reader` is [`Bytes`] based, then the buffer must contain sufficient bytes to complete
444    /// the request, and must include the Parquet footer. If page indexes are desired, the buffer
445    /// must contain the entire file, or [`Self::try_parse_sized()`] should be used.
446    pub fn try_parse<R: ChunkReader>(&mut self, reader: &R) -> Result<()> {
447        self.try_parse_sized(reader, reader.len())
448    }
449
450    /// Same as [`Self::try_parse()`], but provide the original file size in the case that `reader`
451    /// is a [`Bytes`] struct that does not contain the entire file. This information is necessary
452    /// when the page indexes are desired. `reader` must have access to the Parquet footer.
453    ///
454    /// Using this function also allows for retrying with a larger buffer.
455    ///
456    /// # Errors
457    ///
458    /// This function will return [`ParquetError::NeedMoreData`] in the event `reader` does not
459    /// provide enough data to fully parse the metadata (see example below). The returned error
460    /// will be populated with a `usize` field indicating the number of bytes required from the
461    /// tail of the file to completely parse the requested metadata.
462    ///
463    /// Other errors returned include [`ParquetError::General`] and [`ParquetError::EOF`].
464    ///
465    /// # Example
466    /// ```no_run
467    /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
468    /// # use parquet::errors::ParquetError;
469    /// # use crate::parquet::file::reader::Length;
470    /// # fn get_bytes(file: &std::fs::File, range: std::ops::Range<u64>) -> bytes::Bytes { unimplemented!(); }
471    /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); }
472    /// let file = open_parquet_file("some_path.parquet");
473    /// let len = file.len();
474    /// // Speculatively read 1 kilobyte from the end of the file
475    /// let bytes = get_bytes(&file, len - 1024..len);
476    /// let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
477    /// match reader.try_parse_sized(&bytes, len) {
478    ///     Ok(_) => (),
479    ///     Err(ParquetError::NeedMoreData(needed)) => {
480    ///         // Read the needed number of bytes from the end of the file
481    ///         let bytes = get_bytes(&file, len - needed as u64..len);
482    ///         reader.try_parse_sized(&bytes, len).unwrap();
483    ///     }
484    ///     _ => panic!("unexpected error")
485    /// }
486    /// let metadata = reader.finish().unwrap();
487    /// ```
488    ///
489    /// Note that it is possible for the file metadata to be completely read, but there are
490    /// insufficient bytes available to read the page indexes. [`Self::has_metadata()`] can be used
491    /// to test for this. In the event the file metadata is present, re-parsing of the file
492    /// metadata can be skipped by using [`Self::read_page_indexes_sized()`], as shown below.
493    /// ```no_run
494    /// # use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
495    /// # use parquet::errors::ParquetError;
496    /// # use crate::parquet::file::reader::Length;
497    /// # fn get_bytes(file: &std::fs::File, range: std::ops::Range<u64>) -> bytes::Bytes { unimplemented!(); }
498    /// # fn open_parquet_file(path: &str) -> std::fs::File { unimplemented!(); }
499    /// let file = open_parquet_file("some_path.parquet");
500    /// let len = file.len();
501    /// // Speculatively read 1 kilobyte from the end of the file
502    /// let mut bytes = get_bytes(&file, len - 1024..len);
503    /// let mut reader = ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
504    /// // Loop until `bytes` is large enough
505    /// loop {
506    ///     match reader.try_parse_sized(&bytes, len) {
507    ///         Ok(_) => break,
508    ///         Err(ParquetError::NeedMoreData(needed)) => {
509    ///             // Read the needed number of bytes from the end of the file
510    ///             bytes = get_bytes(&file, len - needed as u64..len);
511    ///             // If file metadata was read only read page indexes, otherwise continue loop
512    ///             if reader.has_metadata() {
513    ///                 reader.read_page_indexes_sized(&bytes, len).unwrap();
514    ///                 break;
515    ///             }
516    ///         }
517    ///         _ => panic!("unexpected error")
518    ///     }
519    /// }
520    /// let metadata = reader.finish().unwrap();
521    /// ```
522    pub fn try_parse_sized<R: ChunkReader>(&mut self, reader: &R, file_size: u64) -> Result<()> {
523        self.metadata = match self.parse_metadata(reader) {
524            Ok(metadata) => Some(metadata),
525            Err(ParquetError::NeedMoreData(needed)) => {
526                // If reader is the same length as `file_size` then presumably there is no more to
527                // read, so return an EOF error.
528                return if file_size == reader.len() || needed as u64 > file_size {
529                    Err(eof_err!(
530                        "Parquet file too small. Size is {} but need {}",
531                        file_size,
532                        needed
533                    ))
534                } else {
535                    // Ask for a larger buffer
536                    Err(ParquetError::NeedMoreData(needed))
537                };
538            }
539            Err(e) => return Err(e),
540        };
541
542        // we can return if page indexes aren't requested
543        if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip
544        {
545            return Ok(());
546        }
547
548        self.read_page_indexes_sized(reader, file_size)
549    }
550
551    /// Read the page index structures when a [`ParquetMetaData`] has already been obtained.
552    /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`].
553    ///
554    /// # Behavior on multiple calls
555    ///
556    /// On success, this method replaces any existing page index in the metadata. Index types
557    /// configured as [`PageIndexPolicy::Skip`] are cleared; other index types contain only the
558    /// indexes selected by their configured masks and available according to their policies. If
559    /// neither index type is populated, the page index is cleared.
560    pub fn read_page_indexes<R: ChunkReader>(&mut self, reader: &R) -> Result<()> {
561        self.read_page_indexes_sized(reader, reader.len())
562    }
563
564    /// Reads the page index structures described by the stored [`ParquetMetaData`], returning the
565    /// metadata together with the decoded page index.
566    ///
567    /// Unlike [`Self::read_page_indexes`], this returns the newly decoded [`PageIndex`] separately;
568    /// any page index already attached to the returned metadata is unchanged. This is useful for
569    /// populating an external page index cache. The configured page index policies and masks
570    /// determine which indexes are returned. The page index is `None` when no indexes were
571    /// requested or found.
572    ///
573    /// This operation consumes `self`, but returns the stored [`ParquetMetaData`] as the first
574    /// tuple element. Use [`Self::new_with_metadata`] to construct the reader.
575    ///
576    /// # Example
577    ///
578    /// Read the footer metadata first, then load selected page indexes separately when they are
579    /// needed:
580    ///
581    /// ```no_run
582    /// # use parquet::errors::Result;
583    /// use parquet::file::metadata::page_index::PageIndexProvider;
584    /// use parquet::file::metadata::{
585    ///     ColumnChunkMask, PageIndexPolicy, ParquetMetaDataReader,
586    /// };
587    /// # fn open_parquet_file() -> std::fs::File { unimplemented!() }
588    /// # fn example() -> Result<()> {
589    /// let file = open_parquet_file();
590    /// let metadata = ParquetMetaDataReader::new()
591    ///     .with_page_index_policy(PageIndexPolicy::Skip)
592    ///     .parse_and_finish(&file)?;
593    /// assert!(metadata.page_index().is_none());
594    ///
595    /// let mask = ColumnChunkMask::columns([0]);
596    /// let (metadata, page_index) = ParquetMetaDataReader::new_with_metadata(metadata)
597    ///     .with_column_index_policy(PageIndexPolicy::Required)
598    ///     .with_offset_index_policy(PageIndexPolicy::Required)
599    ///     .with_column_index_mask(mask.clone())
600    ///     .with_offset_index_mask(mask)
601    ///     .read_page_index_and_finish(&file)?;
602    ///
603    /// // The metadata is returned unchanged and the selected indexes can be cached separately.
604    /// assert!(metadata.page_index().is_none());
605    /// let page_index = page_index.expect("page indexes are required");
606    /// assert!(page_index.column_index(0, 0).is_some());
607    /// assert!(page_index.offset_index(0, 0).is_some());
608    /// # Ok(())
609    /// # }
610    /// ```
611    pub fn read_page_index_and_finish<R: ChunkReader>(
612        self,
613        reader: &R,
614    ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
615        self.read_page_index_sized_and_finish(reader, reader.len())
616    }
617
618    /// Reads the page index structures described by the stored [`ParquetMetaData`], returning the
619    /// metadata and decoded page index, and using `file_size` as the size of the original Parquet
620    /// file.
621    ///
622    /// This variant supports a [`Bytes`] reader containing only a suffix of the file. It returns
623    /// [`ParquetError::NeedMoreData`] if that suffix does not contain the requested page indexes.
624    /// See [`Self::read_page_index_and_finish`] for ownership and configuration details.
625    pub fn read_page_index_sized_and_finish<R: ChunkReader>(
626        mut self,
627        reader: &R,
628        file_size: u64,
629    ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
630        self.read_page_index_sized_impl(reader, file_size)
631    }
632
633    fn read_page_index_sized_impl<R: ChunkReader>(
634        &mut self,
635        reader: &R,
636        file_size: u64,
637    ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
638        let Some(metadata) = self.metadata.take() else {
639            return Err(general_err!(
640                "Tried to read page indexes without ParquetMetaData metadata"
641            ));
642        };
643
644        let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)?
645            .with_offset_index_policy(self.offset_index)
646            .with_column_index_policy(self.column_index)
647            .with_offset_index_mask(self.offset_index_mask.clone())
648            .with_column_index_mask(self.column_index_mask.clone())
649            .with_metadata_options(self.metadata_options.clone());
650        let mut push_decoder = self.prepare_push_decoder(push_decoder);
651
652        let range = match needs_page_index_data(&mut push_decoder)? {
653            NeedsIndexData::No(page_index) => return Ok(page_index),
654            NeedsIndexData::Yes(range) => range,
655        };
656
657        let file_range = file_size.saturating_sub(reader.len())..file_size;
658        if !(file_range.contains(&range.start) && file_range.contains(&range.end)) {
659            return if range.end > file_size {
660                Err(eof_err!(
661                    "Parquet file too small. Range {range:?} is beyond file bounds {file_size}",
662                ))
663            } else {
664                Err(ParquetError::NeedMoreData(
665                    (file_size - range.start).try_into()?,
666                ))
667            };
668        }
669
670        if let Some(metadata_size) = self.metadata_size {
671            let metadata_range = file_size.saturating_sub(metadata_size as u64)..file_size;
672            if range.end > metadata_range.start {
673                return Err(eof_err!(
674                    "Parquet file too small. Page index range {range:?} overlaps with file metadata {metadata_range:?}",
675                ));
676            }
677        }
678
679        let bytes_needed = usize::try_from(range.end - range.start)?;
680        let bytes = reader.get_bytes(range.start - file_range.start, bytes_needed)?;
681        push_decoder.push_range(range, bytes)?;
682        parse_page_index_data(&mut push_decoder)
683    }
684
685    /// Read the page index structures when a [`ParquetMetaData`] has already been obtained.
686    /// This variant is used when `reader` cannot access the entire Parquet file (e.g. it is
687    /// a [`Bytes`] struct containing the tail of the file).
688    /// See [`Self::new_with_metadata()`] and [`Self::has_metadata()`]. Like
689    /// [`Self::try_parse_sized()`] this function may return [`ParquetError::NeedMoreData`].
690    ///
691    /// See [`Self::read_page_indexes()`] for a description of the replacement behavior on
692    /// multiple calls.
693    pub fn read_page_indexes_sized<R: ChunkReader>(
694        &mut self,
695        reader: &R,
696        file_size: u64,
697    ) -> Result<()> {
698        let (mut metadata, page_index) = self.read_page_index_sized_impl(reader, file_size)?;
699        metadata
700            .set_page_index(page_index.map(|index| Arc::new(index) as Arc<dyn PageIndexProvider>));
701        self.metadata = Some(metadata);
702        Ok(())
703    }
704
705    /// Given a [`MetadataFetch`], parse and return the [`ParquetMetaData`] in a single pass.
706    ///
707    /// This call will consume `self`.
708    ///
709    /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches
710    /// performed by this function.
711    #[cfg(all(feature = "async", feature = "arrow"))]
712    pub async fn load_and_finish<F: MetadataFetch>(
713        mut self,
714        fetch: F,
715        file_size: u64,
716    ) -> Result<ParquetMetaData> {
717        self.try_load(fetch, file_size).await?;
718        self.finish()
719    }
720
721    /// Given a [`MetadataSuffixFetch`], parse and return the [`ParquetMetaData`] in a single pass.
722    ///
723    /// This call will consume `self`.
724    ///
725    /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches
726    /// performed by this function.
727    #[cfg(all(feature = "async", feature = "arrow"))]
728    pub async fn load_via_suffix_and_finish<F: MetadataSuffixFetch>(
729        mut self,
730        fetch: F,
731    ) -> Result<ParquetMetaData> {
732        self.try_load_via_suffix(fetch).await?;
733        self.finish()
734    }
735    /// Attempts to (asynchronously) parse the footer metadata (and optionally page indexes)
736    /// given a [`MetadataFetch`].
737    ///
738    /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches
739    /// performed by this function.
740    #[cfg(all(feature = "async", feature = "arrow"))]
741    pub async fn try_load<F: MetadataFetch>(&mut self, mut fetch: F, file_size: u64) -> Result<()> {
742        let (metadata, remainder) = self.load_metadata(&mut fetch, file_size).await?;
743
744        self.metadata = Some(metadata);
745
746        // we can return if page indexes aren't requested
747        if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip
748        {
749            return Ok(());
750        }
751
752        self.load_page_index_with_remainder(fetch, remainder).await
753    }
754
755    /// Attempts to (asynchronously) parse the footer metadata (and optionally page indexes)
756    /// given a [`MetadataSuffixFetch`].
757    ///
758    /// See [`Self::with_prefetch_hint`] for a discussion of how to reduce the number of fetches
759    /// performed by this function.
760    #[cfg(all(feature = "async", feature = "arrow"))]
761    pub async fn try_load_via_suffix<F: MetadataSuffixFetch>(
762        &mut self,
763        mut fetch: F,
764    ) -> Result<()> {
765        let metadata = self.load_metadata_via_suffix(&mut fetch).await?;
766
767        self.metadata = Some(metadata);
768
769        // we can return if page indexes aren't requested
770        if self.column_index == PageIndexPolicy::Skip && self.offset_index == PageIndexPolicy::Skip
771        {
772            return Ok(());
773        }
774
775        self.load_page_index_with_remainder(fetch, None).await
776    }
777
778    /// Asynchronously fetch the page index structures when a [`ParquetMetaData`] has already
779    /// been obtained. See [`Self::new_with_metadata()`].
780    ///
781    /// See [`Self::read_page_indexes()`] for a description of the replacement behavior on
782    /// multiple calls.
783    #[cfg(all(feature = "async", feature = "arrow"))]
784    pub async fn load_page_index<F: MetadataFetch>(&mut self, fetch: F) -> Result<()> {
785        self.load_page_index_with_remainder(fetch, None).await
786    }
787
788    /// Asynchronously loads the page index structures described by the stored [`ParquetMetaData`],
789    /// returning the metadata together with the decoded page index.
790    ///
791    /// Unlike [`Self::load_page_index`], this returns the newly decoded [`PageIndex`] separately;
792    /// any page index already attached to the returned metadata is unchanged. This is useful for
793    /// populating an external page index cache. The page index is `None` when no indexes were
794    /// requested or found. This operation consumes `self`, but returns the stored metadata as the
795    /// first tuple element.
796    #[cfg(all(feature = "async", feature = "arrow"))]
797    pub async fn load_page_index_and_finish<F: MetadataFetch>(
798        mut self,
799        fetch: F,
800    ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
801        self.load_page_index_with_remainder_impl(fetch, None).await
802    }
803
804    /// Reads and decodes the dictionary page of a column chunk into an Arrow array.
805    ///
806    /// Returns `Ok(None)` if the column chunk has no dictionary page, or if
807    /// its physical type is not `BYTE_ARRAY` (the only physical type
808    /// currently supported).
809    ///
810    /// The returned array contains raw `Binary` values, even for columns
811    /// annotated as strings. Callers can compare byte slices directly or
812    /// convert values to UTF-8 explicitly.
813    ///
814    /// This can be used to inspect dictionary values when selecting or pruning
815    /// row groups before reading their data pages.
816    ///
817    /// Note this does not verify that the *entire* column chunk is
818    /// dictionary-encoded (i.e. that the dictionary contains every value in
819    /// the chunk) -- callers that need that guarantee should check
820    /// [`crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask`].
821    #[cfg(feature = "arrow")]
822    pub fn read_column_dictionary<R: ChunkReader>(
823        reader: &R,
824        metadata: &ParquetMetaData,
825        row_group_idx: usize,
826        column_idx: usize,
827    ) -> Result<Option<ArrayRef>> {
828        let Some(range) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? else {
829            return Ok(None);
830        };
831        let length = usize::try_from(range.end - range.start)?;
832        let buffer = reader.get_bytes(range.start, length)?;
833        decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some)
834    }
835
836    /// Asynchronous version of [`Self::read_column_dictionary`].
837    #[cfg(all(feature = "async", feature = "arrow"))]
838    pub async fn read_column_dictionary_async<F: MetadataFetch>(
839        mut fetch: F,
840        metadata: &ParquetMetaData,
841        row_group_idx: usize,
842        column_idx: usize,
843    ) -> Result<Option<ArrayRef>> {
844        let Some(range) = dictionary_page_byte_range(metadata, row_group_idx, column_idx)? else {
845            return Ok(None);
846        };
847        let buffer = fetch.fetch(range).await?;
848        decode_dictionary_page(buffer, metadata, row_group_idx, column_idx).map(Some)
849    }
850
851    #[cfg(all(feature = "async", feature = "arrow"))]
852    async fn load_page_index_with_remainder<F: MetadataFetch>(
853        &mut self,
854        fetch: F,
855        remainder: Option<(usize, Bytes)>,
856    ) -> Result<()> {
857        let (mut metadata, page_index) = self
858            .load_page_index_with_remainder_impl(fetch, remainder)
859            .await?;
860        metadata
861            .set_page_index(page_index.map(|index| Arc::new(index) as Arc<dyn PageIndexProvider>));
862        self.metadata = Some(metadata);
863        Ok(())
864    }
865
866    #[cfg(all(feature = "async", feature = "arrow"))]
867    async fn load_page_index_with_remainder_impl<F: MetadataFetch>(
868        &mut self,
869        mut fetch: F,
870        remainder: Option<(usize, Bytes)>,
871    ) -> Result<(ParquetMetaData, Option<PageIndex>)> {
872        let Some(metadata) = self.metadata.take() else {
873            return Err(general_err!("Footer metadata is not present"));
874        };
875
876        // in this case we don't actually know what the file size is, so just use u64::MAX
877        // this is ok since the offsets in the metadata are always valid
878        let file_size = u64::MAX;
879        let push_decoder = ParquetMetaDataPushDecoder::try_new_with_metadata(file_size, metadata)?
880            .with_offset_index_policy(self.offset_index)
881            .with_column_index_policy(self.column_index)
882            .with_offset_index_mask(self.offset_index_mask.clone())
883            .with_column_index_mask(self.column_index_mask.clone())
884            .with_metadata_options(self.metadata_options.clone());
885        let mut push_decoder = self.prepare_push_decoder(push_decoder);
886
887        // Get bounds needed for page indexes (if any are present in the file).
888        let range = match needs_page_index_data(&mut push_decoder)? {
889            NeedsIndexData::No(page_index) => return Ok(page_index),
890            NeedsIndexData::Yes(range) => range,
891        };
892
893        let bytes = match &remainder {
894            Some((remainder_start, remainder)) if *remainder_start as u64 <= range.start => {
895                let remainder_start = *remainder_start as u64;
896                let offset = usize::try_from(range.start - remainder_start)?;
897                let end = usize::try_from(range.end - remainder_start)?;
898                if end > remainder.len() {
899                    return Err(general_err!(
900                        "Corrupted parquet file: index data range ({:?}) exceeds remainder length ({})",
901                        range,
902                        remainder.len()
903                    ));
904                }
905                remainder.slice(offset..end)
906            }
907            // Note: this will potentially fetch data already in remainder, this keeps things simple
908            _ => fetch.fetch(range.clone()).await?,
909        };
910
911        // Sanity check
912        if bytes.len() as u64 != range.end - range.start {
913            return Err(general_err!(
914                "Corrupted parquet file: index data length mismatch, expected {}, got {}",
915                range.end - range.start,
916                bytes.len()
917            ));
918        }
919        push_decoder.push_range(range.clone(), bytes)?;
920        parse_page_index_data(&mut push_decoder)
921    }
922
923    // One-shot parse of footer.
924    // Side effect: this will set `self.metadata_size`
925    fn parse_metadata<R: ChunkReader>(&mut self, chunk_reader: &R) -> Result<ParquetMetaData> {
926        // check file is large enough to hold footer
927        let file_size = chunk_reader.len();
928        if file_size < (FOOTER_SIZE as u64) {
929            return Err(ParquetError::NeedMoreData(FOOTER_SIZE));
930        }
931
932        let mut footer = [0_u8; FOOTER_SIZE];
933        chunk_reader
934            .get_read(file_size - FOOTER_SIZE as u64)?
935            .read_exact(&mut footer)?;
936
937        let footer = FooterTail::try_new(&footer)?;
938        let metadata_len = footer.metadata_length();
939        let footer_metadata_len = FOOTER_SIZE + metadata_len;
940        self.metadata_size = Some(footer_metadata_len);
941
942        if footer_metadata_len as u64 > file_size {
943            return Err(ParquetError::NeedMoreData(footer_metadata_len));
944        }
945
946        let start = file_size - footer_metadata_len as u64;
947        let bytes = chunk_reader.get_bytes(start, metadata_len)?;
948        self.decode_footer_metadata(bytes, file_size, footer)
949    }
950
951    /// Size of the serialized thrift metadata plus the 8 byte footer. Only set if
952    /// `self.parse_metadata` is called.
953    pub fn metadata_size(&self) -> Option<usize> {
954        self.metadata_size
955    }
956
957    /// Return the number of bytes to read in the initial pass. If `prefetch_size` has
958    /// been provided, then return that value if it is larger than the size of the Parquet
959    /// file footer (8 bytes). Otherwise returns `8`.
960    #[cfg(all(feature = "async", feature = "arrow"))]
961    fn get_prefetch_size(&self) -> usize {
962        if let Some(prefetch) = self.prefetch_hint
963            && prefetch > FOOTER_SIZE
964        {
965            return prefetch;
966        }
967        FOOTER_SIZE
968    }
969
970    #[cfg(all(feature = "async", feature = "arrow"))]
971    async fn load_metadata<F: MetadataFetch>(
972        &self,
973        fetch: &mut F,
974        file_size: u64,
975    ) -> Result<(ParquetMetaData, Option<(usize, Bytes)>)> {
976        let prefetch = self.get_prefetch_size() as u64;
977
978        if file_size < FOOTER_SIZE as u64 {
979            return Err(eof_err!("file size of {} is less than footer", file_size));
980        }
981
982        // If a size hint is provided, read more than the minimum size
983        // to try and avoid a second fetch.
984        // Note: prefetch > file_size is ok since we're using saturating_sub.
985        let footer_start = file_size.saturating_sub(prefetch);
986
987        let suffix = fetch.fetch(footer_start..file_size).await?;
988        let suffix_len = suffix.len();
989        let fetch_len = (file_size - footer_start)
990            .try_into()
991            .expect("footer size should never be larger than u32");
992        if suffix_len < fetch_len {
993            return Err(eof_err!(
994                "metadata requires {} bytes, but could only read {}",
995                fetch_len,
996                suffix_len
997            ));
998        }
999
1000        let mut footer = [0; FOOTER_SIZE];
1001        footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]);
1002
1003        let footer = FooterTail::try_new(&footer)?;
1004        let length = footer.metadata_length();
1005
1006        if file_size < (length + FOOTER_SIZE) as u64 {
1007            return Err(eof_err!(
1008                "file size of {} is less than footer + metadata {}",
1009                file_size,
1010                length + FOOTER_SIZE
1011            ));
1012        }
1013
1014        // Did not fetch the entire file metadata in the initial read, need to make a second request
1015        if length > suffix_len - FOOTER_SIZE {
1016            let metadata_start = file_size - (length + FOOTER_SIZE) as u64;
1017            let meta = fetch
1018                .fetch(metadata_start..(file_size - FOOTER_SIZE as u64))
1019                .await?;
1020            Ok((self.decode_footer_metadata(meta, file_size, footer)?, None))
1021        } else {
1022            let metadata_start = (file_size - (length + FOOTER_SIZE) as u64 - footer_start)
1023                .try_into()
1024                .expect("metadata length should never be larger than u32");
1025            let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE);
1026            Ok((
1027                self.decode_footer_metadata(slice, file_size, footer)?,
1028                Some((footer_start as usize, suffix.slice(..metadata_start))),
1029            ))
1030        }
1031    }
1032
1033    #[cfg(all(feature = "async", feature = "arrow"))]
1034    // Unlike load_metadata, the file size is not known so it is not safe
1035    // to use any leftover bytes that may have been pre-fetched. Thus this
1036    // returns only the metadata.
1037    async fn load_metadata_via_suffix<F: MetadataSuffixFetch>(
1038        &self,
1039        fetch: &mut F,
1040    ) -> Result<ParquetMetaData> {
1041        let prefetch = self.get_prefetch_size();
1042
1043        let suffix = fetch.fetch_suffix(prefetch).await?;
1044        let suffix_len = suffix.len();
1045
1046        if suffix_len < FOOTER_SIZE {
1047            return Err(eof_err!(
1048                "footer metadata requires {} bytes, but could only read {}",
1049                FOOTER_SIZE,
1050                suffix_len
1051            ));
1052        }
1053
1054        let mut footer = [0; FOOTER_SIZE];
1055        footer.copy_from_slice(&suffix[suffix_len - FOOTER_SIZE..suffix_len]);
1056
1057        let footer = FooterTail::try_new(&footer)?;
1058        let length = footer.metadata_length();
1059        // fake file size as we are only parsing the footer metadata here
1060        // (cant be parsing page indexes without the full file size)
1061        let file_size = (length + FOOTER_SIZE) as u64;
1062
1063        // Did not fetch the entire file metadata in the initial read, need to make a second request
1064        let metadata_offset = length + FOOTER_SIZE;
1065        if length > suffix_len - FOOTER_SIZE {
1066            let meta = fetch.fetch_suffix(metadata_offset).await?;
1067
1068            if meta.len() < metadata_offset {
1069                return Err(eof_err!(
1070                    "metadata requires {} bytes, but could only read {}",
1071                    metadata_offset,
1072                    meta.len()
1073                ));
1074            }
1075
1076            // need to slice off the footer or decryption fails
1077            let meta = meta.slice(0..length);
1078            Ok(self.decode_footer_metadata(meta, file_size, footer)?)
1079        } else {
1080            let metadata_start = suffix_len - metadata_offset;
1081            let slice = suffix.slice(metadata_start..suffix_len - FOOTER_SIZE);
1082            Ok(self.decode_footer_metadata(slice, file_size, footer)?)
1083        }
1084    }
1085
1086    /// Decodes [`ParquetMetaData`] from the provided bytes.
1087    ///
1088    /// Typically, this is used to decode the metadata from the end of a parquet
1089    /// file. The format of `buf` is the Thrift compact binary protocol, as specified
1090    /// by the [Parquet Spec].
1091    ///
1092    /// It does **NOT** include the 8-byte footer.
1093    ///
1094    /// This method handles using either `decode_metadata` or
1095    /// `decode_metadata_with_encryption` depending on whether the encryption
1096    /// feature is enabled.
1097    ///
1098    /// [Parquet Spec]: https://github.com/apache/parquet-format#metadata
1099    pub(crate) fn decode_footer_metadata(
1100        &self,
1101        buf: Bytes,
1102        file_size: u64,
1103        footer_tail: FooterTail,
1104    ) -> Result<ParquetMetaData> {
1105        // The push decoder expects the metadata to be at the end of the file
1106        // (... data ...) + (metadata) + (footer)
1107        // so we need to provide the starting offset of the metadata
1108        // within the file.
1109        let ending_offset = file_size.checked_sub(FOOTER_SIZE as u64).ok_or_else(|| {
1110            general_err!(
1111                "file size {file_size} is smaller than footer size {}",
1112                FOOTER_SIZE
1113            )
1114        })?;
1115
1116        let starting_offset = ending_offset.checked_sub(buf.len() as u64).ok_or_else(|| {
1117            general_err!(
1118                "file size {file_size} is smaller than buffer size {} + footer size {}",
1119                buf.len(),
1120                FOOTER_SIZE
1121            )
1122        })?;
1123
1124        let range = starting_offset..ending_offset;
1125
1126        let push_decoder =
1127            ParquetMetaDataPushDecoder::try_new_with_footer_tail(file_size, footer_tail)?
1128                // NOTE: DO NOT enable page indexes here, they are handled separately
1129                .with_page_index_policy(PageIndexPolicy::Skip)
1130                .with_metadata_options(self.metadata_options.clone());
1131
1132        let mut push_decoder = self.prepare_push_decoder(push_decoder);
1133        push_decoder.push_range(range, buf)?;
1134        match push_decoder.try_decode()? {
1135            DecodeResult::Data(metadata) => Ok(metadata),
1136            DecodeResult::Finished => Err(general_err!(
1137                "could not parse parquet metadata -- previously finished"
1138            )),
1139            DecodeResult::NeedsData(ranges) => Err(general_err!(
1140                "could not parse parquet metadata, needs ranges {:?}",
1141                ranges
1142            )),
1143        }
1144    }
1145
1146    /// Prepares a push decoder and runs it to decode the metadata.
1147    #[cfg(feature = "encryption")]
1148    fn prepare_push_decoder(
1149        &self,
1150        push_decoder: ParquetMetaDataPushDecoder,
1151    ) -> ParquetMetaDataPushDecoder {
1152        push_decoder.with_file_decryption_properties(
1153            self.file_decryption_properties
1154                .as_ref()
1155                .map(std::sync::Arc::clone),
1156        )
1157    }
1158    #[cfg(not(feature = "encryption"))]
1159    fn prepare_push_decoder(
1160        &self,
1161        push_decoder: ParquetMetaDataPushDecoder,
1162    ) -> ParquetMetaDataPushDecoder {
1163        push_decoder
1164    }
1165
1166    /// Decodes [`ParquetMetaData`] from the provided bytes.
1167    ///
1168    /// Typically this is used to decode the metadata from the end of a parquet
1169    /// file. The format of `buf` is the Thrift compact binary protocol, as specified
1170    /// by the [Parquet Spec].
1171    ///
1172    /// [Parquet Spec]: https://github.com/apache/parquet-format#metadata
1173    pub fn decode_metadata(buf: &[u8]) -> Result<ParquetMetaData> {
1174        decode_metadata(buf, None)
1175    }
1176
1177    /// Decodes [`ParquetMetaData`] from the provided bytes.
1178    ///
1179    /// Like [`Self::decode_metadata`] but this also accepts
1180    /// metadata parsing options.
1181    pub fn decode_metadata_with_options(
1182        buf: &[u8],
1183        options: Option<&ParquetMetaDataOptions>,
1184    ) -> Result<ParquetMetaData> {
1185        decode_metadata(buf, options)
1186    }
1187
1188    /// Decodes the schema from the Parquet footer in `buf`. Returned as
1189    /// a [`SchemaDescriptor`].
1190    pub fn decode_schema(buf: &[u8]) -> Result<Arc<SchemaDescriptor>> {
1191        Ok(Arc::new(parquet_schema_from_bytes(buf)?))
1192    }
1193}
1194
1195/// The bounds needed to read page indexes
1196// this is an internal enum, so it is ok to allow differences in enum size
1197enum NeedsIndexData<T> {
1198    /// no additional data is needed (e.g. the indexes weren't requested)
1199    No(T),
1200    /// Additional data is needed, with the range that are required
1201    Yes(Range<u64>),
1202}
1203
1204/// Determines a single combined range of bytes needed to read the page indexes, or returns the
1205/// metadata and decoded page index if no additional data is needed.
1206fn needs_page_index_data(
1207    push_decoder: &mut ParquetMetaDataPushDecoder,
1208) -> Result<NeedsIndexData<(ParquetMetaData, Option<PageIndex>)>> {
1209    match push_decoder.try_decode_page_index()? {
1210        DecodeResult::NeedsData(ranges) => {
1211            let range = ranges
1212                .into_iter()
1213                .reduce(|a, b| a.start.min(b.start)..a.end.max(b.end))
1214                .ok_or_else(|| general_err!("Internal error: no ranges provided"))?;
1215            Ok(NeedsIndexData::Yes(range))
1216        }
1217        DecodeResult::Data(page_index) => Ok(NeedsIndexData::No(page_index)),
1218        DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")),
1219    }
1220}
1221
1222/// Given a push decoder that has had the needed ranges pushed to it, decode and return the metadata
1223/// and page indexes.
1224fn parse_page_index_data(
1225    push_decoder: &mut ParquetMetaDataPushDecoder,
1226) -> Result<(ParquetMetaData, Option<PageIndex>)> {
1227    match push_decoder.try_decode_page_index()? {
1228        DecodeResult::NeedsData(_) => Err(general_err!(
1229            "Internal error: decoder still needs data after reading required range"
1230        )),
1231        DecodeResult::Data(page_index) => Ok(page_index),
1232        DecodeResult::Finished => Err(general_err!("Internal error: decoder was finished")),
1233    }
1234}
1235
1236/// Returns the `[start, end)` byte range of a column chunk's dictionary
1237/// page, or `None` if the column chunk has no dictionary page or is not a
1238/// `BYTE_ARRAY` column (the only physical type [`decode_dictionary_page`]
1239/// currently supports).
1240#[cfg(feature = "arrow")]
1241fn dictionary_page_byte_range(
1242    metadata: &ParquetMetaData,
1243    row_group_idx: usize,
1244    column_idx: usize,
1245) -> Result<Option<Range<u64>>> {
1246    let column_metadata = metadata.row_group(row_group_idx).column(column_idx);
1247    let column_descriptor = column_metadata.column_descr();
1248
1249    if column_descriptor.physical_type() != crate::basic::Type::BYTE_ARRAY {
1250        return Ok(None);
1251    }
1252    let Some(start) = column_metadata.dictionary_page_offset() else {
1253        return Ok(None);
1254    };
1255    let start: u64 = start
1256        .try_into()
1257        .map_err(|_| ParquetError::General("Dictionary page offset is invalid".to_string()))?;
1258    let end: u64 = column_metadata
1259        .data_page_offset()
1260        .try_into()
1261        .map_err(|_| ParquetError::General("Data page offset is invalid".to_string()))?;
1262    if end < start {
1263        return Err(ParquetError::General(
1264            "Data page offset precedes dictionary page offset".to_string(),
1265        ));
1266    }
1267
1268    Ok(Some(start..end))
1269}
1270
1271#[cfg(test)]
1272mod tests {
1273    use super::*;
1274    use crate::file::reader::Length;
1275    use crate::util::test_common::file_util::get_test_file;
1276    use std::ops::Range;
1277
1278    #[test]
1279    fn test_parse_metadata_size_smaller_than_footer() {
1280        let test_file = tempfile::tempfile().unwrap();
1281        let err = ParquetMetaDataReader::new()
1282            .parse_metadata(&test_file)
1283            .unwrap_err();
1284        assert!(matches!(err, ParquetError::NeedMoreData(FOOTER_SIZE)));
1285    }
1286
1287    #[test]
1288    fn test_parse_metadata_corrupt_footer() {
1289        let data = Bytes::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
1290        let reader_result = ParquetMetaDataReader::new().parse_metadata(&data);
1291        assert_eq!(
1292            reader_result.unwrap_err().to_string(),
1293            "Parquet error: Invalid Parquet file. Corrupt footer"
1294        );
1295    }
1296
1297    #[test]
1298    fn test_parse_metadata_invalid_start() {
1299        let test_file = Bytes::from(vec![255, 0, 0, 0, b'P', b'A', b'R', b'1']);
1300        let err = ParquetMetaDataReader::new()
1301            .parse_metadata(&test_file)
1302            .unwrap_err();
1303        assert!(matches!(err, ParquetError::NeedMoreData(263)));
1304    }
1305
1306    #[test]
1307    #[cfg_attr(miri, ignore)] // Takes too long
1308    fn test_try_parse() {
1309        let file = get_test_file("alltypes_tiny_pages.parquet");
1310        let len = file.len();
1311
1312        let mut reader =
1313            ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
1314
1315        let bytes_for_range = |range: Range<u64>| {
1316            file.get_bytes(range.start, (range.end - range.start).try_into().unwrap())
1317                .unwrap()
1318        };
1319
1320        // read entire file
1321        let bytes = bytes_for_range(0..len);
1322        reader.try_parse(&bytes).unwrap();
1323        let metadata = reader.finish().unwrap();
1324        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1325
1326        // read more than enough of file
1327        let bytes = bytes_for_range(320000..len);
1328        reader.try_parse_sized(&bytes, len).unwrap();
1329        let metadata = reader.finish().unwrap();
1330        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1331
1332        // exactly enough
1333        let bytes = bytes_for_range(323583..len);
1334        reader.try_parse_sized(&bytes, len).unwrap();
1335        let metadata = reader.finish().unwrap();
1336        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1337
1338        // not enough for page index
1339        let bytes = bytes_for_range(323584..len);
1340        // should fail
1341        match reader.try_parse_sized(&bytes, len).unwrap_err() {
1342            // expected error, try again with provided bounds
1343            ParquetError::NeedMoreData(needed) => {
1344                let bytes = bytes_for_range(len - needed as u64..len);
1345                reader.try_parse_sized(&bytes, len).unwrap();
1346                let metadata = reader.finish().unwrap();
1347                assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1348            }
1349            _ => panic!("unexpected error"),
1350        }
1351
1352        // not enough for file metadata, but keep trying until page indexes are read
1353        let mut reader =
1354            ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
1355        let mut bytes = bytes_for_range(452505..len);
1356        loop {
1357            match reader.try_parse_sized(&bytes, len) {
1358                Ok(()) => break,
1359                Err(ParquetError::NeedMoreData(needed)) => {
1360                    bytes = bytes_for_range(len - needed as u64..len);
1361                    if reader.has_metadata() {
1362                        reader.read_page_indexes_sized(&bytes, len).unwrap();
1363                        break;
1364                    }
1365                }
1366                _ => panic!("unexpected error"),
1367            }
1368        }
1369        let metadata = reader.finish().unwrap();
1370        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1371
1372        // not enough for page index but lie about file size
1373        let bytes = bytes_for_range(323584..len);
1374        let reader_result = reader.try_parse_sized(&bytes, len - 323584).unwrap_err();
1375        assert_eq!(
1376            reader_result.to_string(),
1377            "EOF: Parquet file too small. Range 323583..452504 is beyond file bounds 130649"
1378        );
1379
1380        // not enough for file metadata
1381        let mut reader = ParquetMetaDataReader::new();
1382        let bytes = bytes_for_range(452505..len);
1383        // should fail
1384        match reader.try_parse_sized(&bytes, len).unwrap_err() {
1385            // expected error, try again with provided bounds
1386            ParquetError::NeedMoreData(needed) => {
1387                let bytes = bytes_for_range(len - needed as u64..len);
1388                reader.try_parse_sized(&bytes, len).unwrap();
1389                reader.finish().unwrap();
1390            }
1391            _ => panic!("unexpected error"),
1392        }
1393
1394        // not enough for file metadata but use try_parse()
1395        let reader_result = reader.try_parse(&bytes).unwrap_err();
1396        assert_eq!(
1397            reader_result.to_string(),
1398            "EOF: Parquet file too small. Size is 1728 but need 1729"
1399        );
1400
1401        // read head of file rather than tail
1402        let bytes = bytes_for_range(0..1000);
1403        let reader_result = reader.try_parse_sized(&bytes, len).unwrap_err();
1404        assert_eq!(
1405            reader_result.to_string(),
1406            "Parquet error: Invalid Parquet file. Corrupt footer"
1407        );
1408
1409        // lie about file size
1410        let bytes = bytes_for_range(452510..len);
1411        let reader_result = reader.try_parse_sized(&bytes, len - 452505).unwrap_err();
1412        assert_eq!(
1413            reader_result.to_string(),
1414            "EOF: Parquet file too small. Size is 1728 but need 1729"
1415        );
1416    }
1417
1418    #[test]
1419    #[cfg_attr(miri, ignore)] // Takes too long
1420    fn test_read_page_index_and_finish() {
1421        use crate::file::metadata::page_index::PageIndexProvider;
1422
1423        let file = get_test_file("alltypes_tiny_pages.parquet");
1424        let metadata = ParquetMetaDataReader::new()
1425            .with_page_index_policy(PageIndexPolicy::Required)
1426            .parse_and_finish(&file)
1427            .unwrap();
1428        let original_page_index = Arc::clone(metadata.page_index().unwrap());
1429
1430        let reader = ParquetMetaDataReader::new_with_metadata(metadata);
1431        let (metadata, page_index) = reader.read_page_index_and_finish(&file).unwrap();
1432        assert!(page_index.is_none());
1433        assert!(Arc::ptr_eq(
1434            metadata.page_index().unwrap(),
1435            &original_page_index
1436        ));
1437
1438        let reader = ParquetMetaDataReader::new_with_metadata(metadata)
1439            .with_page_index_policy(PageIndexPolicy::Required)
1440            .with_column_index_mask(ColumnChunkMask::columns([0]))
1441            .with_offset_index_mask(ColumnChunkMask::columns([0]));
1442        let (metadata, page_index) = reader.read_page_index_and_finish(&file).unwrap();
1443        let page_index = page_index.unwrap();
1444        assert!(Arc::ptr_eq(
1445            metadata.page_index().unwrap(),
1446            &original_page_index
1447        ));
1448
1449        assert!(page_index.column_index(0, 0).is_some());
1450        assert!(page_index.offset_index(0, 0).is_some());
1451        assert!(page_index.column_index(0, 1).is_none());
1452        assert!(page_index.offset_index(0, 1).is_none());
1453    }
1454
1455    #[test]
1456    fn test_chunk_mask() {
1457        let mask = ColumnChunkMask::row_groups_and_columns([0], [1]);
1458        assert!(mask.includes_row_group(0, 1));
1459        assert!(!mask.includes_row_group(1, 1));
1460        assert!(!mask.includes_column(0, 2));
1461        assert!(mask.includes_column(1, 2));
1462
1463        let mask = ColumnChunkMask::columns([]);
1464        assert!(!mask.includes_column(0, 3));
1465        assert_eq!(mask.column_indices(3).count(), 0);
1466
1467        let all = ColumnChunkMask::all();
1468        assert!(all.is_all());
1469        assert!(all.includes_column(2, 3));
1470        assert!(!all.includes_column(3, 3));
1471
1472        assert_eq!(all.row_group_indices(3).collect::<Vec<_>>(), [0, 1, 2]);
1473        assert_eq!(all.column_indices(2).collect::<Vec<_>>(), [0, 1]);
1474
1475        let none = ColumnChunkMask::none();
1476        assert_eq!(none.row_group_indices(3).count(), 0);
1477        assert_eq!(none.column_indices(3).count(), 0);
1478
1479        let sparse = ColumnChunkMask::row_groups_and_columns([0, 2, 4], [1, 3, 5]);
1480        assert_eq!(sparse.row_group_indices(3).collect::<Vec<_>>(), [0, 2]);
1481        assert_eq!(sparse.column_indices(4).collect::<Vec<_>>(), [1, 3]);
1482    }
1483}
1484
1485#[cfg(all(feature = "async", feature = "arrow", test))]
1486mod async_tests {
1487    use super::*;
1488
1489    use arrow::{array::Int32Array, datatypes::DataType};
1490    use arrow_array::RecordBatch;
1491    use arrow_schema::{Field, Schema};
1492    use bytes::Bytes;
1493    use futures::FutureExt;
1494    use futures::future::BoxFuture;
1495    use std::fs::File;
1496    use std::future::Future;
1497    use std::io::{Read, Seek, SeekFrom};
1498    use std::ops::Range;
1499    use std::sync::Arc;
1500    use std::sync::atomic::{AtomicUsize, Ordering};
1501    use tempfile::NamedTempFile;
1502
1503    use crate::arrow::ArrowWriter;
1504    use crate::file::properties::WriterProperties;
1505    use crate::file::reader::Length;
1506    use crate::util::test_common::file_util::get_test_file;
1507
1508    struct MetadataFetchFn<F>(F);
1509
1510    impl<F, Fut> MetadataFetch for MetadataFetchFn<F>
1511    where
1512        F: FnMut(Range<u64>) -> Fut + Send,
1513        Fut: Future<Output = Result<Bytes>> + Send,
1514    {
1515        fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
1516            async move { self.0(range).await }.boxed()
1517        }
1518    }
1519
1520    struct MetadataSuffixFetchFn<F1, F2>(F1, F2);
1521
1522    impl<F1, Fut, F2> MetadataFetch for MetadataSuffixFetchFn<F1, F2>
1523    where
1524        F1: FnMut(Range<u64>) -> Fut + Send,
1525        Fut: Future<Output = Result<Bytes>> + Send,
1526        F2: Send,
1527    {
1528        fn fetch(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes>> {
1529            async move { self.0(range).await }.boxed()
1530        }
1531    }
1532
1533    impl<F1, Fut, F2> MetadataSuffixFetch for MetadataSuffixFetchFn<F1, F2>
1534    where
1535        F1: FnMut(Range<u64>) -> Fut + Send,
1536        F2: FnMut(usize) -> Fut + Send,
1537        Fut: Future<Output = Result<Bytes>> + Send,
1538    {
1539        fn fetch_suffix(&mut self, suffix: usize) -> BoxFuture<'_, Result<Bytes>> {
1540            async move { self.1(suffix).await }.boxed()
1541        }
1542    }
1543
1544    fn read_range(file: &mut File, range: Range<u64>) -> Result<Bytes> {
1545        file.seek(SeekFrom::Start(range.start))?;
1546        let len = range.end - range.start;
1547        let mut buf = Vec::with_capacity(len.try_into().unwrap());
1548        file.take(len).read_to_end(&mut buf)?;
1549        Ok(buf.into())
1550    }
1551
1552    fn read_suffix(file: &mut File, suffix: usize) -> Result<Bytes> {
1553        let file_len = file.len();
1554        // Don't seek before beginning of file
1555        file.seek(SeekFrom::End(0 - suffix.min(file_len as _) as i64))?;
1556        let mut buf = Vec::with_capacity(suffix);
1557        file.take(suffix as _).read_to_end(&mut buf)?;
1558        Ok(buf.into())
1559    }
1560
1561    #[tokio::test]
1562    async fn test_load_page_index_and_finish() {
1563        use crate::file::metadata::page_index::PageIndexProvider;
1564
1565        let mut file = get_test_file("alltypes_tiny_pages.parquet");
1566        let metadata = ParquetMetaDataReader::new()
1567            .parse_and_finish(&file)
1568            .unwrap();
1569        let reader = ParquetMetaDataReader::new_with_metadata(metadata)
1570            .with_page_index_policy(PageIndexPolicy::Required)
1571            .with_column_index_mask(ColumnChunkMask::columns([0]))
1572            .with_offset_index_mask(ColumnChunkMask::columns([0]));
1573        let fetch = MetadataFetchFn(|range| futures::future::ready(read_range(&mut file, range)));
1574
1575        let page_index = reader
1576            .load_page_index_and_finish(fetch)
1577            .await
1578            .unwrap()
1579            .1
1580            .unwrap();
1581
1582        assert!(page_index.column_index(0, 0).is_some());
1583        assert!(page_index.offset_index(0, 0).is_some());
1584        assert!(page_index.column_index(0, 1).is_none());
1585        assert!(page_index.offset_index(0, 1).is_none());
1586    }
1587
1588    #[tokio::test]
1589    async fn test_simple() {
1590        let mut file = get_test_file("nulls.snappy.parquet");
1591        let len = file.len();
1592
1593        let expected = ParquetMetaDataReader::new()
1594            .parse_and_finish(&file)
1595            .unwrap();
1596        let expected = expected.file_metadata().schema();
1597        let fetch_count = AtomicUsize::new(0);
1598
1599        let mut fetch = |range| {
1600            fetch_count.fetch_add(1, Ordering::SeqCst);
1601            futures::future::ready(read_range(&mut file, range))
1602        };
1603
1604        let input = MetadataFetchFn(&mut fetch);
1605        let actual = ParquetMetaDataReader::new()
1606            .load_and_finish(input, len)
1607            .await
1608            .unwrap();
1609        assert_eq!(actual.file_metadata().schema(), expected);
1610        assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1611
1612        // Metadata hint too small - below footer size
1613        fetch_count.store(0, Ordering::SeqCst);
1614        let input = MetadataFetchFn(&mut fetch);
1615        let actual = ParquetMetaDataReader::new()
1616            .with_prefetch_hint(Some(7))
1617            .load_and_finish(input, len)
1618            .await
1619            .unwrap();
1620        assert_eq!(actual.file_metadata().schema(), expected);
1621        assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1622
1623        // Metadata hint too small
1624        fetch_count.store(0, Ordering::SeqCst);
1625        let input = MetadataFetchFn(&mut fetch);
1626        let actual = ParquetMetaDataReader::new()
1627            .with_prefetch_hint(Some(10))
1628            .load_and_finish(input, len)
1629            .await
1630            .unwrap();
1631        assert_eq!(actual.file_metadata().schema(), expected);
1632        assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1633
1634        // Metadata hint too large
1635        fetch_count.store(0, Ordering::SeqCst);
1636        let input = MetadataFetchFn(&mut fetch);
1637        let actual = ParquetMetaDataReader::new()
1638            .with_prefetch_hint(Some(500))
1639            .load_and_finish(input, len)
1640            .await
1641            .unwrap();
1642        assert_eq!(actual.file_metadata().schema(), expected);
1643        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1644
1645        // Metadata hint exactly correct
1646        fetch_count.store(0, Ordering::SeqCst);
1647        let input = MetadataFetchFn(&mut fetch);
1648        let actual = ParquetMetaDataReader::new()
1649            .with_prefetch_hint(Some(428))
1650            .load_and_finish(input, len)
1651            .await
1652            .unwrap();
1653        assert_eq!(actual.file_metadata().schema(), expected);
1654        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1655
1656        let input = MetadataFetchFn(&mut fetch);
1657        let err = ParquetMetaDataReader::new()
1658            .load_and_finish(input, 4)
1659            .await
1660            .unwrap_err()
1661            .to_string();
1662        assert_eq!(err, "EOF: file size of 4 is less than footer");
1663
1664        let input = MetadataFetchFn(&mut fetch);
1665        let err = ParquetMetaDataReader::new()
1666            .load_and_finish(input, 20)
1667            .await
1668            .unwrap_err()
1669            .to_string();
1670        assert_eq!(err, "Parquet error: Invalid Parquet file. Corrupt footer");
1671    }
1672
1673    #[tokio::test]
1674    async fn test_suffix() {
1675        let mut file = get_test_file("nulls.snappy.parquet");
1676        let mut file2 = file.try_clone().unwrap();
1677
1678        let expected = ParquetMetaDataReader::new()
1679            .parse_and_finish(&file)
1680            .unwrap();
1681        let expected = expected.file_metadata().schema();
1682        let fetch_count = AtomicUsize::new(0);
1683        let suffix_fetch_count = AtomicUsize::new(0);
1684
1685        let mut fetch = |range| {
1686            fetch_count.fetch_add(1, Ordering::SeqCst);
1687            futures::future::ready(read_range(&mut file, range))
1688        };
1689        let mut suffix_fetch = |suffix| {
1690            suffix_fetch_count.fetch_add(1, Ordering::SeqCst);
1691            futures::future::ready(read_suffix(&mut file2, suffix))
1692        };
1693
1694        let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1695        let actual = ParquetMetaDataReader::new()
1696            .load_via_suffix_and_finish(input)
1697            .await
1698            .unwrap();
1699        assert_eq!(actual.file_metadata().schema(), expected);
1700        assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1701        assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2);
1702
1703        // Metadata hint too small - below footer size
1704        fetch_count.store(0, Ordering::SeqCst);
1705        suffix_fetch_count.store(0, Ordering::SeqCst);
1706        let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1707        let actual = ParquetMetaDataReader::new()
1708            .with_prefetch_hint(Some(7))
1709            .load_via_suffix_and_finish(input)
1710            .await
1711            .unwrap();
1712        assert_eq!(actual.file_metadata().schema(), expected);
1713        assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1714        assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2);
1715
1716        // Metadata hint too small
1717        fetch_count.store(0, Ordering::SeqCst);
1718        suffix_fetch_count.store(0, Ordering::SeqCst);
1719        let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1720        let actual = ParquetMetaDataReader::new()
1721            .with_prefetch_hint(Some(10))
1722            .load_via_suffix_and_finish(input)
1723            .await
1724            .unwrap();
1725        assert_eq!(actual.file_metadata().schema(), expected);
1726        assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1727        assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 2);
1728
1729        // Metadata hint too large
1730        fetch_count.store(0, Ordering::SeqCst);
1731        suffix_fetch_count.store(0, Ordering::SeqCst);
1732        let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1733        let actual = ParquetMetaDataReader::new()
1734            .with_prefetch_hint(Some(500))
1735            .load_via_suffix_and_finish(input)
1736            .await
1737            .unwrap();
1738        assert_eq!(actual.file_metadata().schema(), expected);
1739        assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1740        assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1);
1741
1742        // Metadata hint exactly correct
1743        fetch_count.store(0, Ordering::SeqCst);
1744        suffix_fetch_count.store(0, Ordering::SeqCst);
1745        let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1746        let actual = ParquetMetaDataReader::new()
1747            .with_prefetch_hint(Some(428))
1748            .load_via_suffix_and_finish(input)
1749            .await
1750            .unwrap();
1751        assert_eq!(actual.file_metadata().schema(), expected);
1752        assert_eq!(fetch_count.load(Ordering::SeqCst), 0);
1753        assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1);
1754    }
1755
1756    #[cfg(feature = "encryption")]
1757    #[tokio::test]
1758    async fn test_suffix_with_encryption() {
1759        let mut file = get_test_file("uniform_encryption.parquet.encrypted");
1760        let mut file2 = file.try_clone().unwrap();
1761
1762        let mut fetch = |range| futures::future::ready(read_range(&mut file, range));
1763        let mut suffix_fetch = |suffix| futures::future::ready(read_suffix(&mut file2, suffix));
1764
1765        let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1766
1767        let key_code: &[u8] = b"0123456789012345";
1768        let decryption_properties = FileDecryptionProperties::builder(key_code.to_vec())
1769            .build()
1770            .unwrap();
1771
1772        // just make sure the metadata is properly decrypted and read
1773        let expected = ParquetMetaDataReader::new()
1774            .with_decryption_properties(Some(decryption_properties))
1775            .load_via_suffix_and_finish(input)
1776            .await
1777            .unwrap();
1778        assert_eq!(expected.num_row_groups(), 1);
1779    }
1780
1781    #[tokio::test]
1782    async fn test_page_index() {
1783        let mut file = get_test_file("alltypes_tiny_pages.parquet");
1784        let len = file.len();
1785        let fetch_count = AtomicUsize::new(0);
1786        let mut fetch = |range| {
1787            fetch_count.fetch_add(1, Ordering::SeqCst);
1788            futures::future::ready(read_range(&mut file, range))
1789        };
1790
1791        let f = MetadataFetchFn(&mut fetch);
1792        let mut loader =
1793            ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
1794        loader.try_load(f, len).await.unwrap();
1795        assert_eq!(fetch_count.load(Ordering::SeqCst), 3);
1796        let metadata = loader.finish().unwrap();
1797        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1798
1799        // Prefetch just footer exactly
1800        fetch_count.store(0, Ordering::SeqCst);
1801        let f = MetadataFetchFn(&mut fetch);
1802        let mut loader = ParquetMetaDataReader::new()
1803            .with_page_index_policy(PageIndexPolicy::Required)
1804            .with_prefetch_hint(Some(1729));
1805        loader.try_load(f, len).await.unwrap();
1806        assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1807        let metadata = loader.finish().unwrap();
1808        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1809
1810        // Prefetch more than footer but not enough
1811        fetch_count.store(0, Ordering::SeqCst);
1812        let f = MetadataFetchFn(&mut fetch);
1813        let mut loader = ParquetMetaDataReader::new()
1814            .with_page_index_policy(PageIndexPolicy::Required)
1815            .with_prefetch_hint(Some(130649));
1816        loader.try_load(f, len).await.unwrap();
1817        assert_eq!(fetch_count.load(Ordering::SeqCst), 2);
1818        let metadata = loader.finish().unwrap();
1819        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1820
1821        // Prefetch exactly enough
1822        fetch_count.store(0, Ordering::SeqCst);
1823        let f = MetadataFetchFn(&mut fetch);
1824        let metadata = ParquetMetaDataReader::new()
1825            .with_page_index_policy(PageIndexPolicy::Required)
1826            .with_prefetch_hint(Some(130650))
1827            .load_and_finish(f, len)
1828            .await
1829            .unwrap();
1830        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1831        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1832
1833        // Prefetch more than enough but less than the entire file
1834        fetch_count.store(0, Ordering::SeqCst);
1835        let f = MetadataFetchFn(&mut fetch);
1836        let metadata = ParquetMetaDataReader::new()
1837            .with_page_index_policy(PageIndexPolicy::Required)
1838            .with_prefetch_hint(Some((len - 1000) as usize)) // prefetch entire file
1839            .load_and_finish(f, len)
1840            .await
1841            .unwrap();
1842        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1843        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1844
1845        // Prefetch the entire file
1846        fetch_count.store(0, Ordering::SeqCst);
1847        let f = MetadataFetchFn(&mut fetch);
1848        let metadata = ParquetMetaDataReader::new()
1849            .with_page_index_policy(PageIndexPolicy::Required)
1850            .with_prefetch_hint(Some(len as usize)) // prefetch entire file
1851            .load_and_finish(f, len)
1852            .await
1853            .unwrap();
1854        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1855        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1856
1857        // Prefetch more than the entire file
1858        fetch_count.store(0, Ordering::SeqCst);
1859        let f = MetadataFetchFn(&mut fetch);
1860        let metadata = ParquetMetaDataReader::new()
1861            .with_page_index_policy(PageIndexPolicy::Required)
1862            .with_prefetch_hint(Some((len + 1000) as usize)) // prefetch entire file
1863            .load_and_finish(f, len)
1864            .await
1865            .unwrap();
1866        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1867        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1868    }
1869
1870    #[tokio::test]
1871    async fn test_page_index_via_suffix_with_prefetch() {
1872        let mut file = get_test_file("alltypes_tiny_pages.parquet");
1873        let mut suffix_file = file.try_clone().unwrap();
1874        let len = file.len();
1875        let fetch_count = AtomicUsize::new(0);
1876        let suffix_fetch_count = AtomicUsize::new(0);
1877
1878        let mut fetch = |range| {
1879            fetch_count.fetch_add(1, Ordering::SeqCst);
1880            futures::future::ready(read_range(&mut file, range))
1881        };
1882        let mut suffix_fetch = |suffix| {
1883            suffix_fetch_count.fetch_add(1, Ordering::SeqCst);
1884            futures::future::ready(read_suffix(&mut suffix_file, suffix))
1885        };
1886
1887        let input = MetadataSuffixFetchFn(&mut fetch, &mut suffix_fetch);
1888        let metadata = ParquetMetaDataReader::new()
1889            .with_page_index_policy(PageIndexPolicy::Required)
1890            .with_prefetch_hint(Some((len - 1000) as usize))
1891            .load_via_suffix_and_finish(input)
1892            .await
1893            .unwrap();
1894
1895        assert_eq!(suffix_fetch_count.load(Ordering::SeqCst), 1);
1896        assert_eq!(fetch_count.load(Ordering::SeqCst), 1);
1897        assert!(metadata.page_index().is_some_and(|idx| idx.is_complete()));
1898    }
1899
1900    fn write_parquet_file(offset_index_disabled: bool) -> Result<NamedTempFile> {
1901        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1902        let batch = RecordBatch::try_new(
1903            schema.clone(),
1904            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
1905        )?;
1906
1907        let file = NamedTempFile::new().unwrap();
1908
1909        // Write properties with page index disabled
1910        let props = WriterProperties::builder()
1911            .set_offset_index_disabled(offset_index_disabled)
1912            .build();
1913
1914        let mut writer = ArrowWriter::try_new(file.reopen()?, schema, Some(props))?;
1915        writer.write(&batch)?;
1916        writer.close()?;
1917
1918        Ok(file)
1919    }
1920
1921    fn read_and_check(file: &File, policy: PageIndexPolicy) -> Result<ParquetMetaData> {
1922        let mut reader = ParquetMetaDataReader::new().with_page_index_policy(policy);
1923        reader.try_parse(file)?;
1924        reader.finish()
1925    }
1926
1927    #[test]
1928    fn test_page_index_policy() {
1929        // With page index
1930        let f = write_parquet_file(false).unwrap();
1931        read_and_check(f.as_file(), PageIndexPolicy::Required).unwrap();
1932        read_and_check(f.as_file(), PageIndexPolicy::Optional).unwrap();
1933        read_and_check(f.as_file(), PageIndexPolicy::Skip).unwrap();
1934
1935        // Without page index
1936        let f = write_parquet_file(true).unwrap();
1937        let res = read_and_check(f.as_file(), PageIndexPolicy::Required);
1938        assert!(matches!(
1939            res,
1940            Err(ParquetError::General(e)) if e == "missing offset index"
1941        ));
1942        read_and_check(f.as_file(), PageIndexPolicy::Optional).unwrap();
1943        read_and_check(f.as_file(), PageIndexPolicy::Skip).unwrap();
1944    }
1945}