parquet/file/metadata/
parser.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Internal metadata parsing routines
19//!
20//! These functions parse thrift-encoded metadata from a byte slice
21//! into the corresponding Rust structures
22
23use std::sync::Arc;
24
25use crate::errors::ParquetError;
26use crate::file::metadata::page_index::{PageIndexBuilder, PageIndexProvider};
27use crate::file::metadata::thrift::parquet_metadata_from_bytes;
28use crate::file::metadata::{
29    ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions,
30};
31
32use crate::file::page_index::column_index::ColumnIndexMetaData;
33use crate::file::page_index::index_reader::{decode_column_index, decode_offset_index};
34use crate::file::page_index::offset_index::OffsetIndexMetaData;
35use bytes::Bytes;
36
37/// Helper struct for metadata parsing
38///
39/// This structure parses thrift-encoded bytes into the correct Rust structs,
40/// such as [`ParquetMetaData`], handling decryption if necessary.
41//
42// Note this structure is used to minimize the number of
43// places to add `#[cfg(feature = "encryption")]` checks.
44pub(crate) use inner::MetadataParser;
45
46#[cfg(feature = "encryption")]
47mod inner {
48    use std::sync::Arc;
49
50    use super::*;
51    use crate::encryption::decrypt::FileDecryptionProperties;
52    use crate::errors::Result;
53
54    /// API for decoding metadata that may be encrypted
55    #[derive(Debug, Default)]
56    pub(crate) struct MetadataParser {
57        // the credentials and keys needed to decrypt metadata
58        file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
59        // metadata parsing options
60        metadata_options: Option<Arc<ParquetMetaDataOptions>>,
61    }
62
63    impl MetadataParser {
64        pub(crate) fn new() -> Self {
65            MetadataParser::default()
66        }
67
68        pub(crate) fn with_file_decryption_properties(
69            mut self,
70            file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
71        ) -> Self {
72            self.file_decryption_properties = file_decryption_properties;
73            self
74        }
75
76        pub(crate) fn with_metadata_options(
77            self,
78            options: Option<Arc<ParquetMetaDataOptions>>,
79        ) -> Self {
80            Self {
81                metadata_options: options,
82                ..self
83            }
84        }
85
86        pub(crate) fn decode_metadata(
87            &self,
88            buf: &[u8],
89            encrypted_footer: bool,
90        ) -> Result<ParquetMetaData> {
91            if encrypted_footer || self.file_decryption_properties.is_some() {
92                crate::file::metadata::thrift::encryption::parquet_metadata_with_encryption(
93                    self.file_decryption_properties.as_ref(),
94                    encrypted_footer,
95                    buf,
96                    self.metadata_options.as_deref(),
97                )
98            } else {
99                decode_metadata(buf, self.metadata_options.as_deref())
100            }
101        }
102    }
103
104    pub(super) fn parse_single_column_index(
105        bytes: &[u8],
106        metadata: &ParquetMetaData,
107        column: &ColumnChunkMetaData,
108        row_group_index: usize,
109        col_index: usize,
110    ) -> crate::errors::Result<ColumnIndexMetaData> {
111        use crate::encryption::decrypt::CryptoContext;
112        match &column.column_crypto_metadata {
113            Some(crypto_metadata) => {
114                let file_decryptor = metadata.file_decryptor.as_ref().ok_or_else(|| {
115                    general_err!("Cannot decrypt column index, no file decryptor set")
116                })?;
117                let crypto_context = CryptoContext::for_column(
118                    file_decryptor,
119                    crypto_metadata,
120                    row_group_index,
121                    col_index,
122                )?;
123                let column_decryptor = crypto_context.metadata_decryptor();
124                let aad = crypto_context.create_column_index_aad()?;
125                let plaintext = column_decryptor.decrypt(bytes, &aad)?;
126                decode_column_index(&plaintext, column.column_type())
127            }
128            None => decode_column_index(bytes, column.column_type()),
129        }
130    }
131
132    pub(super) fn parse_single_offset_index(
133        bytes: &[u8],
134        metadata: &ParquetMetaData,
135        column: &ColumnChunkMetaData,
136        row_group_index: usize,
137        col_index: usize,
138    ) -> crate::errors::Result<OffsetIndexMetaData> {
139        use crate::encryption::decrypt::CryptoContext;
140        match &column.column_crypto_metadata {
141            Some(crypto_metadata) => {
142                let file_decryptor = metadata.file_decryptor.as_ref().ok_or_else(|| {
143                    general_err!("Cannot decrypt offset index, no file decryptor set")
144                })?;
145                let crypto_context = CryptoContext::for_column(
146                    file_decryptor,
147                    crypto_metadata,
148                    row_group_index,
149                    col_index,
150                )?;
151                let column_decryptor = crypto_context.metadata_decryptor();
152                let aad = crypto_context.create_offset_index_aad()?;
153                let plaintext = column_decryptor.decrypt(bytes, &aad)?;
154                decode_offset_index(&plaintext)
155            }
156            None => decode_offset_index(bytes),
157        }
158    }
159}
160
161#[cfg(not(feature = "encryption"))]
162mod inner {
163    use super::*;
164    use crate::errors::Result;
165    use std::sync::Arc;
166    /// parallel implementation when encryption feature is not enabled
167    ///
168    /// This has the same API as the encryption-enabled version
169    #[derive(Debug, Default)]
170    pub(crate) struct MetadataParser {
171        // metadata parsing options
172        metadata_options: Option<Arc<ParquetMetaDataOptions>>,
173    }
174
175    impl MetadataParser {
176        pub(crate) fn new() -> Self {
177            MetadataParser::default()
178        }
179
180        pub(crate) fn with_metadata_options(
181            self,
182            options: Option<Arc<ParquetMetaDataOptions>>,
183        ) -> Self {
184            Self {
185                metadata_options: options,
186            }
187        }
188
189        pub(crate) fn decode_metadata(
190            &self,
191            buf: &[u8],
192            encrypted_footer: bool,
193        ) -> Result<ParquetMetaData> {
194            if encrypted_footer {
195                Err(general_err!(
196                    "Parquet file has an encrypted footer but the encryption feature is disabled"
197                ))
198            } else {
199                decode_metadata(buf, self.metadata_options.as_deref())
200            }
201        }
202    }
203
204    pub(super) fn parse_single_column_index(
205        bytes: &[u8],
206        _metadata: &ParquetMetaData,
207        column: &ColumnChunkMetaData,
208        _row_group_index: usize,
209        _col_index: usize,
210    ) -> crate::errors::Result<ColumnIndexMetaData> {
211        decode_column_index(bytes, column.column_type())
212    }
213
214    pub(super) fn parse_single_offset_index(
215        bytes: &[u8],
216        _metadata: &ParquetMetaData,
217        _column: &ColumnChunkMetaData,
218        _row_group_index: usize,
219        _col_index: usize,
220    ) -> crate::errors::Result<OffsetIndexMetaData> {
221        decode_offset_index(bytes)
222    }
223}
224
225/// Decodes [`ParquetMetaData`] from the provided bytes.
226///
227/// Typically this is used to decode the metadata from the end of a parquet
228/// file. The format of `buf` is the Thrift compact binary protocol, as specified
229/// by the [Parquet Spec].
230///
231/// [Parquet Spec]: https://github.com/apache/parquet-format#metadata
232pub(crate) fn decode_metadata(
233    buf: &[u8],
234    options: Option<&ParquetMetaDataOptions>,
235) -> crate::errors::Result<ParquetMetaData> {
236    parquet_metadata_from_bytes(buf, options)
237}
238
239/// Parses page index from the provided bytes and adds it to the metadata.
240///
241/// Arguments
242/// * `metadata` - The ParquetMetaData to which the parsed column index will be added.
243/// * `column_index_policy` - The policy for handling column index parsing (e.g.,
244///   Required, Optional, Skip).
245/// * `offset_index_policy` - The policy for handling offset index parsing (e.g.,
246///   Required, Optional, Skip).
247/// * `bytes` - The byte slice containing the page index data.
248/// * `start_offset` - The offset where `bytes` begin in the file.
249pub(crate) fn parse_page_index(
250    metadata: &mut ParquetMetaData,
251    column_index_policy: PageIndexPolicy,
252    offset_index_policy: PageIndexPolicy,
253    bytes: &Bytes,
254    start_offset: u64,
255) -> crate::errors::Result<()> {
256    if column_index_policy == PageIndexPolicy::Skip && offset_index_policy == PageIndexPolicy::Skip
257    {
258        return Ok(());
259    }
260    let num_row_groups = metadata.num_row_groups();
261    let num_columns = metadata.file_metadata().schema_descr().num_columns();
262    let mut builder = PageIndexBuilder::default();
263    if column_index_policy != PageIndexPolicy::Skip {
264        builder.allocate_column_indexes(num_row_groups, num_columns);
265        parse_column_index(
266            metadata,
267            column_index_policy,
268            &mut builder,
269            bytes,
270            start_offset,
271        )?;
272    }
273    if offset_index_policy != PageIndexPolicy::Skip {
274        builder.allocate_offset_indexes(num_row_groups, num_columns);
275        parse_offset_index(
276            metadata,
277            offset_index_policy,
278            &mut builder,
279            bytes,
280            start_offset,
281        )?;
282    }
283
284    let page_index = builder.build();
285    // if both indexes are missing from the file, return without modifying `metadata`
286    if !page_index.has_column_indexes() && !page_index.has_offset_indexes() {
287        return Ok(());
288    }
289    metadata.set_page_index(Some(Arc::new(page_index)));
290    Ok(())
291}
292
293fn parse_column_index(
294    metadata: &ParquetMetaData,
295    column_index_policy: PageIndexPolicy,
296    page_index_builder: &mut PageIndexBuilder,
297    bytes: &Bytes,
298    start_offset: u64,
299) -> crate::errors::Result<()> {
300    if column_index_policy == PageIndexPolicy::Skip {
301        return Ok(());
302    }
303    for rg_idx in 0..metadata.num_row_groups() {
304        let rg = metadata.row_group(rg_idx);
305        for col_idx in 0..rg.num_columns() {
306            let col = rg.column(col_idx);
307            if let Some(r) = col.column_index_range() {
308                let r_start = usize::try_from(r.start - start_offset)?;
309                let r_end = usize::try_from(r.end - start_offset)?;
310                let idx = inner::parse_single_column_index(
311                    &bytes[r_start..r_end],
312                    metadata,
313                    col,
314                    rg_idx,
315                    col_idx,
316                )?;
317                page_index_builder.put_column_index(idx, rg_idx, col_idx);
318            }
319        }
320    }
321
322    Ok(())
323}
324
325fn parse_offset_index(
326    metadata: &ParquetMetaData,
327    offset_index_policy: PageIndexPolicy,
328    page_index_builder: &mut PageIndexBuilder,
329    bytes: &Bytes,
330    start_offset: u64,
331) -> crate::errors::Result<()> {
332    if offset_index_policy == PageIndexPolicy::Skip {
333        return Ok(());
334    }
335    for rg_idx in 0..metadata.num_row_groups() {
336        let rg = metadata.row_group(rg_idx);
337        for col_idx in 0..rg.num_columns() {
338            let col = rg.column(col_idx);
339            if let Some(r) = col.offset_index_range() {
340                let r_start = usize::try_from(r.start - start_offset)?;
341                let r_end = usize::try_from(r.end - start_offset)?;
342                let idx = inner::parse_single_offset_index(
343                    &bytes[r_start..r_end],
344                    metadata,
345                    col,
346                    rg_idx,
347                    col_idx,
348                )?;
349                page_index_builder.put_offset_index(idx, rg_idx, col_idx);
350            } else if offset_index_policy == PageIndexPolicy::Required {
351                return Err(general_err!("missing offset index"));
352            }
353        }
354    }
355
356    Ok(())
357}
parquet/file/metadata/parser.rs

parquet/file/metadata/
parser.rs
