parquet/file/metadata/
options.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//! Options used to control metadata parsing
6
7use std::collections::HashSet;
8use std::sync::Arc;
9
10use crate::schema::types::SchemaDescPtr;
11
12/// Enum to control decoding of some Parquet statistics fields.
13///
14/// # Example
15/// ```rust
16/// use parquet::file::metadata::ParquetStatisticsPolicy;
17/// use parquet::file::serialized_reader::ReadOptionsBuilder;
18/// use parquet::arrow::arrow_reader::ArrowReaderOptions;
19///
20/// // Set arrow options to skip encoding statistics for all columns.
21/// let options =
22///     ArrowReaderOptions::new().with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll);
23///
24/// // Set serialized reader options to decode encoding statistics for all columns.
25/// let options =
26///     ReadOptionsBuilder::new().with_encoding_stats_policy(ParquetStatisticsPolicy::KeepAll)
27///     .build();
28///
29/// // Set arrow options to skip encoding statistics for all columns, but to decode statistics
30/// // for columns 0 and 1.
31/// let options = ArrowReaderOptions::new()
32///     .with_encoding_stats_policy(ParquetStatisticsPolicy::skip_except(&[0, 1]));
33/// ```
34#[derive(Default, Debug, Clone)]
35pub enum ParquetStatisticsPolicy {
36    /// Decode the relevant statistics for all columns.
37    #[default]
38    KeepAll,
39    /// Skip decoding the relevant statistics for all columns.
40    SkipAll,
41    /// Skip decoding the relevant statistics for all columns not in the provided set
42    /// of column indices.
43    SkipExcept(Arc<HashSet<usize>>),
44}
45
46impl ParquetStatisticsPolicy {
47    /// Create a `ParquetStatisticsPolicy` to skip all columns except those in `keep`.
48    ///
49    /// If `keep` is empty, then this returns [`Self::SkipAll`]
50    pub fn skip_except(keep: &[usize]) -> Self {
51        if keep.is_empty() {
52            Self::SkipAll
53        } else {
54            let mut keep_set = HashSet::<usize>::with_capacity(keep.len());
55            keep_set.extend(keep.iter());
56            Self::SkipExcept(Arc::new(keep_set))
57        }
58    }
59
60    /// Returns whether the policy for the given column index is to skip the statistics.
61    pub(crate) fn is_skip(&self, col_index: usize) -> bool {
62        match self {
63            Self::KeepAll => false,
64            Self::SkipAll => true,
65            Self::SkipExcept(keep) => !keep.contains(&col_index),
66        }
67    }
68}
69
70/// Options that can be set to control what parts of the Parquet file footer
71/// metadata will be decoded and made present in the [`ParquetMetaData`] returned
72/// by [`ParquetMetaDataReader`] and [`ParquetMetaDataPushDecoder`].
73///
74/// [`ParquetMetaData`]: crate::file::metadata::ParquetMetaData
75/// [`ParquetMetaDataReader`]: crate::file::metadata::ParquetMetaDataReader
76/// [`ParquetMetaDataPushDecoder`]: crate::file::metadata::ParquetMetaDataPushDecoder
77#[derive(Debug, Clone)]
78pub struct ParquetMetaDataOptions {
79    schema_descr: Option<SchemaDescPtr>,
80    encoding_stats_as_mask: bool,
81    encoding_stats_policy: ParquetStatisticsPolicy,
82    column_stats_policy: ParquetStatisticsPolicy,
83    size_stats_policy: ParquetStatisticsPolicy,
84}
85
86impl Default for ParquetMetaDataOptions {
87    fn default() -> Self {
88        Self {
89            schema_descr: None,
90            encoding_stats_as_mask: true,
91            encoding_stats_policy: ParquetStatisticsPolicy::KeepAll,
92            column_stats_policy: ParquetStatisticsPolicy::KeepAll,
93            size_stats_policy: ParquetStatisticsPolicy::KeepAll,
94        }
95    }
96}
97
98impl ParquetMetaDataOptions {
99    /// Return a new default [`ParquetMetaDataOptions`].
100    pub fn new() -> Self {
101        Default::default()
102    }
103
104    /// Returns an optional [`SchemaDescPtr`] to use when decoding. If this is not `None` then
105    /// the schema in the footer will be skipped.
106    pub fn schema(&self) -> Option<&SchemaDescPtr> {
107        self.schema_descr.as_ref()
108    }
109
110    /// Provide a schema to use when decoding the metadata.
111    pub fn set_schema(&mut self, val: SchemaDescPtr) {
112        self.schema_descr = Some(val);
113    }
114
115    /// Call [`Self::set_schema`] and return `Self` for chaining.
116    pub fn with_schema(mut self, val: SchemaDescPtr) -> Self {
117        self.set_schema(val);
118        self
119    }
120
121    /// Returns whether to present the [`encoding_stats`] field of the Parquet `ColumnMetaData`
122    /// as a bitmask (defaults to `true`).
123    ///
124    /// See [`ColumnChunkMetaData::page_encoding_stats_mask`] for an explanation of why this
125    /// might be desirable.
126    ///
127    /// [`ColumnChunkMetaData::page_encoding_stats_mask`]:
128    /// crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask
129    /// [`encoding_stats`]:
130    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917
131    pub fn encoding_stats_as_mask(&self) -> bool {
132        self.encoding_stats_as_mask
133    }
134
135    /// Convert [`encoding_stats`] from a vector of [`PageEncodingStats`] to a bitmask. This can
136    /// speed up metadata decoding while still enabling some use cases served by the full stats.
137    ///
138    /// Note that if for a given column both this option and `skip_encoding_stats` are `true`, the
139    /// stats will be skipped and not be returned as a mask.
140    ///
141    /// See [`ColumnChunkMetaData::page_encoding_stats_mask`] for more information.
142    ///
143    /// [`PageEncodingStats`]: crate::file::metadata::PageEncodingStats
144    /// [`ColumnChunkMetaData::page_encoding_stats_mask`]:
145    /// crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask
146    /// [`encoding_stats`]:
147    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917
148    pub fn set_encoding_stats_as_mask(&mut self, val: bool) {
149        self.encoding_stats_as_mask = val;
150    }
151
152    /// Call [`Self::set_encoding_stats_as_mask`] and return `Self` for chaining.
153    pub fn with_encoding_stats_as_mask(mut self, val: bool) -> Self {
154        self.set_encoding_stats_as_mask(val);
155        self
156    }
157
158    /// Returns whether to skip decoding the [`encoding_stats`] in the Parquet `ColumnMetaData`
159    /// for the column indexed by `col_index`.
160    ///
161    /// [`encoding_stats`]:
162    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917
163    pub fn skip_encoding_stats(&self, col_index: usize) -> bool {
164        self.encoding_stats_policy.is_skip(col_index)
165    }
166
167    /// Sets the decoding policy for [`encoding_stats`] in the Parquet `ColumnMetaData`.
168    ///
169    /// The default policy is to decode all `encoding_stats`.
170    ///
171    /// This option takes precedence over [`Self::encoding_stats_as_mask`].
172    ///
173    /// [`encoding_stats`]:
174    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917
175    pub fn set_encoding_stats_policy(&mut self, policy: ParquetStatisticsPolicy) {
176        self.encoding_stats_policy = policy;
177    }
178
179    /// Call [`Self::set_encoding_stats_policy`] and return `Self` for chaining.
180    pub fn with_encoding_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
181        self.set_encoding_stats_policy(policy);
182        self
183    }
184
185    /// Returns whether to skip decoding the [`statistics`] in the Parquet `ColumnMetaData`
186    /// for the column indexed by `col_index`.
187    ///
188    /// [`statistics`]:
189    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L912
190    pub fn skip_column_stats(&self, col_index: usize) -> bool {
191        self.column_stats_policy.is_skip(col_index)
192    }
193
194    /// Sets the decoding policy for [`statistics`] in the Parquet `ColumnMetaData`.
195    ///
196    /// The default policy is to decode all `statistics`.
197    ///
198    /// [`statistics`]:
199    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L912
200    pub fn set_column_stats_policy(&mut self, policy: ParquetStatisticsPolicy) {
201        self.column_stats_policy = policy;
202    }
203
204    /// Call [`Self::set_column_stats_policy`] and return `Self` for chaining.
205    pub fn with_column_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
206        self.set_column_stats_policy(policy);
207        self
208    }
209
210    /// Returns whether to skip decoding the [`size_statistics`] in the Parquet `ColumnMetaData`
211    /// for the column indexed by `col_index`.
212    ///
213    /// [`size_statistics`]:
214    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L936
215    pub fn skip_size_stats(&self, col_index: usize) -> bool {
216        self.size_stats_policy.is_skip(col_index)
217    }
218
219    /// Sets the decoding policy for [`size_statistics`] in the Parquet `ColumnMetaData`.
220    ///
221    /// The default policy is to decode all `size_statistics`.
222    ///
223    /// [`size_statistics`]:
224    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L936
225    pub fn set_size_stats_policy(&mut self, policy: ParquetStatisticsPolicy) {
226        self.size_stats_policy = policy;
227    }
228
229    /// Call [`Self::set_size_stats_policy`] and return `Self` for chaining.
230    pub fn with_size_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
231        self.set_size_stats_policy(policy);
232        self
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use bytes::Bytes;
239
240    use crate::{
241        DecodeResult,
242        file::metadata::{ParquetMetaDataOptions, ParquetMetaDataPushDecoder},
243        util::test_common::file_util::get_test_file,
244    };
245    use std::{io::Read, sync::Arc};
246
247    #[test]
248    fn test_options_default() {
249        let options = ParquetMetaDataOptions::default();
250        assert!(options.encoding_stats_as_mask());
251    }
252
253    #[test]
254    fn test_provide_schema() {
255        let mut buf: Vec<u8> = Vec::new();
256        get_test_file("alltypes_plain.parquet")
257            .read_to_end(&mut buf)
258            .unwrap();
259
260        let data = Bytes::from(buf);
261        let mut decoder = ParquetMetaDataPushDecoder::try_new(data.len() as u64).unwrap();
262        decoder
263            .push_range(0..data.len() as u64, data.clone())
264            .unwrap();
265
266        let DecodeResult::Data(expected) = decoder.try_decode().unwrap() else {
267            panic!("could not parse metadata")
268        };
269        let expected_schema = expected.file_metadata().schema_descr_ptr();
270
271        let mut options = ParquetMetaDataOptions::new();
272        options.set_schema(expected_schema);
273        let options = Arc::new(options);
274
275        let mut decoder = ParquetMetaDataPushDecoder::try_new(data.len() as u64)
276            .unwrap()
277            .with_metadata_options(Some(options));
278        decoder.push_range(0..data.len() as u64, data).unwrap();
279        let DecodeResult::Data(metadata) = decoder.try_decode().unwrap() else {
280            panic!("could not parse metadata")
281        };
282
283        assert_eq!(expected, metadata);
284        // the schema pointers should be the same
285        assert!(Arc::ptr_eq(
286            &expected.file_metadata().schema_descr_ptr(),
287            &metadata.file_metadata().schema_descr_ptr()
288        ));
289    }
290}
parquet/file/metadata/options.rs

parquet/file/metadata/
options.rs
