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::{ops::Range, sync::Arc};
24
25use bytes::Bytes;
26
27use crate::errors::{ParquetError, Result};
28use crate::file::metadata::page_index::{PageIndexBuilder, PageIndexProvider};
29use crate::file::metadata::thrift::parquet_metadata_from_bytes;
30use crate::file::metadata::{
31    ColumnChunkMask, ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions,
32};
33
34use crate::file::page_index::column_index::ColumnIndexMetaData;
35use crate::file::page_index::index_reader::{decode_column_index, decode_offset_index};
36use crate::file::page_index::offset_index::OffsetIndexMetaData;
37
38/// Helper struct for metadata parsing
39///
40/// This structure parses thrift-encoded bytes into the correct Rust structs,
41/// such as [`ParquetMetaData`], handling decryption if necessary.
42//
43// Note this structure is used to minimize the number of
44// places to add `#[cfg(feature = "encryption")]` checks.
45pub(crate) use inner::MetadataParser;
46
47#[cfg(feature = "encryption")]
48mod inner {
49    use std::sync::Arc;
50
51    use super::*;
52    use crate::encryption::decrypt::FileDecryptionProperties;
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    ) -> 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    ) -> 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 std::sync::Arc;
165    /// parallel implementation when encryption feature is not enabled
166    ///
167    /// This has the same API as the encryption-enabled version
168    #[derive(Debug, Default)]
169    pub(crate) struct MetadataParser {
170        // metadata parsing options
171        metadata_options: Option<Arc<ParquetMetaDataOptions>>,
172    }
173
174    impl MetadataParser {
175        pub(crate) fn new() -> Self {
176            MetadataParser::default()
177        }
178
179        pub(crate) fn with_metadata_options(
180            self,
181            options: Option<Arc<ParquetMetaDataOptions>>,
182        ) -> Self {
183            Self {
184                metadata_options: options,
185            }
186        }
187
188        pub(crate) fn decode_metadata(
189            &self,
190            buf: &[u8],
191            encrypted_footer: bool,
192        ) -> Result<ParquetMetaData> {
193            if encrypted_footer {
194                Err(general_err!(
195                    "Parquet file has an encrypted footer but the encryption feature is disabled"
196                ))
197            } else {
198                decode_metadata(buf, self.metadata_options.as_deref())
199            }
200        }
201    }
202
203    pub(super) fn parse_single_column_index(
204        bytes: &[u8],
205        _metadata: &ParquetMetaData,
206        column: &ColumnChunkMetaData,
207        _row_group_index: usize,
208        _col_index: usize,
209    ) -> Result<ColumnIndexMetaData> {
210        decode_column_index(bytes, column.column_type())
211    }
212
213    pub(super) fn parse_single_offset_index(
214        bytes: &[u8],
215        _metadata: &ParquetMetaData,
216        _column: &ColumnChunkMetaData,
217        _row_group_index: usize,
218        _col_index: usize,
219    ) -> Result<OffsetIndexMetaData> {
220        decode_offset_index(bytes)
221    }
222}
223
224/// Decodes [`ParquetMetaData`] from the provided bytes.
225///
226/// Typically this is used to decode the metadata from the end of a parquet
227/// file. The format of `buf` is the Thrift compact binary protocol, as specified
228/// by the [Parquet Spec].
229///
230/// [Parquet Spec]: https://github.com/apache/parquet-format#metadata
231pub(crate) fn decode_metadata(
232    buf: &[u8],
233    options: Option<&ParquetMetaDataOptions>,
234) -> Result<ParquetMetaData> {
235    parquet_metadata_from_bytes(buf, options)
236}
237
238/// Parses page indexes from the provided bytes and replaces those in the metadata.
239///
240/// Arguments
241/// * `metadata` - The ParquetMetaData whose page index will be replaced.
242/// * `column_index_policy` - The policy for handling column index parsing (e.g.,
243///   Required, Optional, Skip).
244/// * `offset_index_policy` - The policy for handling offset index parsing (e.g.,
245///   Required, Optional, Skip).
246/// * `column_index_mask` - The row groups and leaf columns whose column indexes are parsed.
247///   Ignored when `column_index_policy` is [`PageIndexPolicy::Skip`].
248/// * `offset_index_mask` - The row groups and leaf columns whose offset indexes are parsed.
249///   Ignored when `offset_index_policy` is [`PageIndexPolicy::Skip`].
250/// * `bytes` - The byte slice containing the page index data.
251/// * `start_offset` - The offset where `bytes` begin in the file.
252pub(crate) fn parse_page_index(
253    metadata: &mut ParquetMetaData,
254    column_index_policy: PageIndexPolicy,
255    offset_index_policy: PageIndexPolicy,
256    column_index_mask: &ColumnChunkMask,
257    offset_index_mask: &ColumnChunkMask,
258    bytes: &Bytes,
259    start_offset: u64,
260) -> Result<()> {
261    let num_row_groups = metadata.num_row_groups();
262    let num_columns = metadata.file_metadata().schema_descr().num_columns();
263    let mut builder = PageIndexBuilder::default();
264
265    if column_index_policy != PageIndexPolicy::Skip {
266        builder.allocate_column_indexes(num_row_groups, num_columns);
267        parse_column_index(
268            metadata,
269            column_index_policy,
270            column_index_mask,
271            &mut builder,
272            bytes,
273            start_offset,
274        )?;
275    }
276    if offset_index_policy != PageIndexPolicy::Skip {
277        builder.allocate_offset_indexes(num_row_groups, num_columns);
278        parse_offset_index(
279            metadata,
280            offset_index_policy,
281            offset_index_mask,
282            &mut builder,
283            bytes,
284            start_offset,
285        )?;
286    }
287
288    // Always replace the page index, even if both policies are Skip.
289    // This ensures that repeated calls to read_page_indexes replace the existing index
290    // rather than preserving it.
291    let page_index = builder.build();
292    if page_index.has_column_indexes() || page_index.has_offset_indexes() {
293        metadata.set_page_index(Some(Arc::new(page_index)));
294    } else {
295        // If no indexes were read, clear the page index
296        metadata.set_page_index(None);
297    }
298    Ok(())
299}
300
301fn parse_column_index(
302    metadata: &ParquetMetaData,
303    column_index_policy: PageIndexPolicy,
304    mask: &ColumnChunkMask,
305    page_index_builder: &mut PageIndexBuilder,
306    bytes: &Bytes,
307    start_offset: u64,
308) -> Result<()> {
309    if column_index_policy == PageIndexPolicy::Skip {
310        return Ok(());
311    }
312    for rg_idx in mask.row_group_indices(metadata.num_row_groups()) {
313        let rg = metadata.row_group(rg_idx);
314        for col_idx in mask.column_indices(rg.num_columns()) {
315            let col = rg.column(col_idx);
316            if let Some(r) = col.column_index_range() {
317                let idx_bytes = get_index_bytes(bytes, start_offset, r)?;
318                let idx =
319                    inner::parse_single_column_index(idx_bytes, metadata, col, rg_idx, col_idx)?;
320                page_index_builder.put_column_index(idx, rg_idx, col_idx);
321            }
322        }
323    }
324
325    Ok(())
326}
327
328fn parse_offset_index(
329    metadata: &ParquetMetaData,
330    offset_index_policy: PageIndexPolicy,
331    mask: &ColumnChunkMask,
332    page_index_builder: &mut PageIndexBuilder,
333    bytes: &Bytes,
334    start_offset: u64,
335) -> Result<()> {
336    if offset_index_policy == PageIndexPolicy::Skip {
337        return Ok(());
338    }
339    for rg_idx in mask.row_group_indices(metadata.num_row_groups()) {
340        let rg = metadata.row_group(rg_idx);
341        for col_idx in mask.column_indices(rg.num_columns()) {
342            let col = rg.column(col_idx);
343            if let Some(r) = col.offset_index_range() {
344                let idx_bytes = get_index_bytes(bytes, start_offset, r)?;
345                let idx =
346                    inner::parse_single_offset_index(idx_bytes, metadata, col, rg_idx, col_idx)?;
347                page_index_builder.put_offset_index(idx, rg_idx, col_idx);
348            } else if offset_index_policy == PageIndexPolicy::Required {
349                return Err(general_err!("missing offset index"));
350            }
351        }
352    }
353
354    Ok(())
355}
356
357fn get_index_bytes(bytes: &[u8], start_offset: u64, range: Range<u64>) -> Result<&[u8]> {
358    let start = usize::try_from(range.start - start_offset)?;
359    let end = usize::try_from(range.end - start_offset)?;
360    Ok(&bytes[start..end])
361}
parquet/file/metadata/parser.rs

parquet/file/metadata/
parser.rs
