parquet/file/metadata/
writer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::file::metadata::page_index::PageIndexProvider;
19use crate::file::metadata::thrift::FileMeta;
20use crate::file::metadata::{ColumnChunkMetaData, RowGroupMetaData};
21use crate::schema::types::{SchemaDescPtr, SchemaDescriptor};
22use crate::{
23    basic::ColumnOrder,
24    file::metadata::{FileMetaData, ParquetMetaDataBuilder},
25};
26#[cfg(feature = "encryption")]
27use crate::{
28    encryption::{
29        encrypt::{FileEncryptor, encrypt_thrift_object, write_signed_plaintext_thrift_object},
30        modules::{ModuleType, create_footer_aad, create_module_aad},
31    },
32    file::column_crypto_metadata::ColumnCryptoMetaData,
33    file::metadata::thrift::encryption::{AesGcmV1, EncryptionAlgorithm, FileCryptoMetaData},
34};
35use crate::{errors::Result, file::page_index::column_index::ColumnIndexMetaData};
36
37use crate::{
38    file::writer::{TrackedWrite, get_file_magic},
39    parquet_thrift::WriteThrift,
40};
41use crate::{
42    file::{
43        metadata::{KeyValue, ParquetMetaData},
44        page_index::offset_index::OffsetIndexMetaData,
45    },
46    parquet_thrift::ThriftCompactOutputProtocol,
47};
48use std::io::Write;
49use std::sync::Arc;
50
51/// Writes `crate::file::metadata` structures to a thrift encoded byte stream
52///
53/// See [`ParquetMetaDataWriter`] for background and example.
54pub(crate) struct ThriftMetadataWriter<'a, W: Write> {
55    buf: &'a mut TrackedWrite<W>,
56    schema_descr: &'a SchemaDescPtr,
57    row_groups: Vec<RowGroupMetaData>,
58    page_index: Option<Arc<dyn PageIndexProvider>>,
59    key_value_metadata: Option<Vec<KeyValue>>,
60    created_by: Option<String>,
61    object_writer: MetadataObjectWriter,
62    writer_version: i32,
63    write_path_in_schema: bool,
64    preserve_page_index_locations: bool,
65}
66
67impl<'a, W: Write> ThriftMetadataWriter<'a, W> {
68    /// Serialize all the offset indexes to `self.buf`,
69    ///
70    /// Note: also updates the `ColumnChunk::offset_index_offset` and
71    /// `ColumnChunk::offset_index_length` to reflect the position and length
72    /// of the serialized offset indexes.
73    fn write_offset_indexes(&mut self, page_index: &dyn PageIndexProvider) -> Result<()> {
74        // iter row group
75        // iter each column
76        // write offset index to the file
77        for (row_group_idx, row_group) in self.row_groups.iter_mut().enumerate() {
78            for (column_idx, column_metadata) in row_group.columns.iter_mut().enumerate() {
79                if let Some(offset_index) = page_index.offset_index(row_group_idx, column_idx) {
80                    let start_offset = self.buf.bytes_written();
81                    self.object_writer.write_offset_index(
82                        offset_index,
83                        column_metadata,
84                        row_group_idx,
85                        column_idx,
86                        &mut *self.buf,
87                    )?;
88                    let end_offset = self.buf.bytes_written();
89                    // set offset and index for offset index
90                    column_metadata.offset_index_offset = Some(start_offset as i64);
91                    column_metadata.offset_index_length = Some((end_offset - start_offset) as i32);
92                } else if !self.preserve_page_index_locations {
93                    column_metadata.offset_index_offset = None;
94                    column_metadata.offset_index_length = None;
95                }
96            }
97        }
98        Ok(())
99    }
100
101    /// Serialize all the column indexes to the `self.buf`
102    ///
103    /// Note: also updates the `ColumnChunk::column_index_offset` and
104    /// `ColumnChunk::column_index_length` to reflect the position and length
105    /// of the serialized column indexes.
106    fn write_column_indexes(&mut self, page_index: &dyn PageIndexProvider) -> Result<()> {
107        // iter row group
108        // iter each column
109        // write column index to the file
110        for (row_group_idx, row_group) in self.row_groups.iter_mut().enumerate() {
111            for (column_idx, column_metadata) in row_group.columns.iter_mut().enumerate() {
112                if let Some(column_index) = page_index.column_index(row_group_idx, column_idx) {
113                    let start_offset = self.buf.bytes_written();
114                    // only update column_metadata if the write succeeds
115                    if self.object_writer.write_column_index(
116                        column_index,
117                        column_metadata,
118                        row_group_idx,
119                        column_idx,
120                        &mut *self.buf,
121                    )? {
122                        let end_offset = self.buf.bytes_written();
123                        // set offset and index for offset index
124                        column_metadata.column_index_offset = Some(start_offset as i64);
125                        column_metadata.column_index_length =
126                            Some((end_offset - start_offset) as i32);
127                    }
128                } else if !self.preserve_page_index_locations {
129                    column_metadata.column_index_offset = None;
130                    column_metadata.column_index_length = None;
131                }
132            }
133        }
134        Ok(())
135    }
136
137    fn clear_page_indexes(&mut self) {
138        for row_group in &mut self.row_groups {
139            for column_metadata in &mut row_group.columns {
140                column_metadata.column_index_offset = None;
141                column_metadata.column_index_length = None;
142                column_metadata.offset_index_offset = None;
143                column_metadata.offset_index_length = None;
144            }
145        }
146    }
147
148    /// Assembles and writes the final metadata to self.buf
149    pub(crate) fn finish(mut self) -> Result<ParquetMetaData> {
150        let num_rows = self.row_groups.iter().map(|x| x.num_rows).sum();
151
152        // serialize page indexes and update index locations
153        let page_index = self.page_index.take();
154        if let Some(pi) = page_index.as_deref() {
155            self.write_column_indexes(pi)?;
156            self.write_offset_indexes(pi)?;
157        } else if !self.preserve_page_index_locations {
158            // no page indexes at all, then clear index locations
159            self.clear_page_indexes();
160        }
161
162        // We only include ColumnOrder for leaf nodes.
163        let column_orders = self
164            .schema_descr
165            .columns()
166            .iter()
167            .map(|col| {
168                ColumnOrder::column_order_for_type(
169                    col.logical_type_ref(),
170                    col.converted_type(),
171                    col.physical_type(),
172                )
173            })
174            .collect();
175
176        // This field is optional, perhaps in cases where no min/max fields are set
177        // in any Statistics or ColumnIndex object in the whole file.
178        // But for simplicity we always set this field.
179        let column_orders = Some(column_orders);
180
181        let (row_groups, unencrypted_row_groups) = self
182            .object_writer
183            .apply_row_group_encryption(self.row_groups)?;
184
185        #[cfg(feature = "encryption")]
186        let (encryption_algorithm, footer_signing_key_metadata) =
187            self.object_writer.get_plaintext_footer_crypto_metadata();
188        #[cfg(feature = "encryption")]
189        let file_metadata = FileMetaData::new(
190            self.writer_version,
191            num_rows,
192            self.created_by,
193            self.key_value_metadata,
194            self.schema_descr.clone(),
195            column_orders,
196        )
197        .with_encryption_algorithm(encryption_algorithm)
198        .with_footer_signing_key_metadata(footer_signing_key_metadata);
199
200        #[cfg(not(feature = "encryption"))]
201        let file_metadata = FileMetaData::new(
202            self.writer_version,
203            num_rows,
204            self.created_by,
205            self.key_value_metadata,
206            self.schema_descr.clone(),
207            column_orders,
208        );
209
210        let file_meta = FileMeta {
211            file_metadata: &file_metadata,
212            row_groups: &row_groups,
213            write_path_in_schema: self.write_path_in_schema,
214        };
215
216        // Write file metadata
217        let start_pos = self.buf.bytes_written();
218        self.object_writer
219            .write_file_metadata(&file_meta, &mut *self.buf)?;
220        let end_pos = self.buf.bytes_written();
221
222        // Write footer
223        let metadata_len = (end_pos - start_pos) as u32;
224
225        self.buf.write_all(&metadata_len.to_le_bytes())?;
226        self.buf.write_all(self.object_writer.get_file_magic())?;
227
228        // Note: we simply pass the original reference counted page index provider to the
229        // builder. Either it came from ParquetMetaDataWriter where it was cloned, and this
230        // reference will go away when ParquetMetaDataWriter::finish drops the newly created
231        // metadata, or it was passed an owned version from SerializedFileWriter which then
232        // takes it back via the returned metadata.
233        let builder = ParquetMetaDataBuilder::new(file_metadata).set_page_index(page_index);
234
235        // If row group metadata was encrypted, we replace the encrypted row groups with
236        // unencrypted metadata before it is returned to users. This allows the metadata
237        // to be usable for retrieving the row group statistics for example, without users
238        // needing to decrypt the metadata.
239        Ok(match unencrypted_row_groups {
240            Some(rg) => builder.set_row_groups(rg).build(),
241            None => builder.set_row_groups(row_groups).build(),
242        })
243    }
244
245    pub(crate) fn new(
246        buf: &'a mut TrackedWrite<W>,
247        schema_descr: &'a SchemaDescPtr,
248        row_groups: Vec<RowGroupMetaData>,
249        created_by: Option<String>,
250        writer_version: i32,
251        write_path_in_schema: bool,
252        preserve_page_index_locations: bool,
253    ) -> Self {
254        Self {
255            buf,
256            schema_descr,
257            row_groups,
258            page_index: None,
259            key_value_metadata: None,
260            created_by,
261            object_writer: Default::default(),
262            writer_version,
263            write_path_in_schema,
264            preserve_page_index_locations,
265        }
266    }
267
268    pub(crate) fn with_page_index(mut self, page_index: Arc<dyn PageIndexProvider>) -> Self {
269        self.page_index = Some(page_index);
270        self
271    }
272
273    pub(crate) fn with_key_value_metadata(mut self, key_value_metadata: Vec<KeyValue>) -> Self {
274        self.key_value_metadata = Some(key_value_metadata);
275        self
276    }
277
278    #[cfg(feature = "encryption")]
279    pub(crate) fn with_file_encryptor(
280        mut self,
281        file_encryptor: Option<Arc<FileEncryptor>>,
282    ) -> Self {
283        self.object_writer = self.object_writer.with_file_encryptor(file_encryptor);
284        self
285    }
286}
287
288/// Writes [`ParquetMetaData`] to a byte stream
289///
290/// This structure handles the details of writing the various parts of Parquet
291/// metadata into a byte stream. It is used to write the metadata into a parquet
292/// file and can also write metadata into other locations (such as a store of
293/// bytes).
294///
295/// # Discussion
296///
297/// The process of writing Parquet metadata is tricky because the
298/// metadata is not stored as a single inline thrift structure. It can have
299/// several "out of band" structures such as the [`OffsetIndex`] and
300/// BloomFilters stored in separate structures whose locations are stored as
301/// offsets from the beginning of the file.
302///
303/// Note: this writer does not directly write BloomFilters. In order to write
304/// BloomFilters, write the bloom filters into the buffer before creating the
305/// metadata writer. Then set the corresponding `bloom_filter_offset` and
306/// `bloom_filter_length` on [`ColumnChunkMetaData`] passed to this writer.
307///
308/// The writer serializes the page index of any [`PageIndexProvider`] attached
309/// to the input [`ParquetMetaData`]. By default, if the provider has no index
310/// for a column chunk, the offset and length of that index in the [`ColumnChunkMetaData`]
311/// are preserved in the output metadata. This is done so one can store the footer
312/// metadata externally but keep the page indexes in the original file. If this
313/// behavior is not desired, then [`Self::with_preserve_page_index_locations`] should
314/// be set to `false`.
315///
316/// # Output Format
317///
318/// The format of the metadata is as follows:
319///
320/// 1. Optional [`ColumnIndex`] (thrift encoded)
321/// 2. Optional [`OffsetIndex`] (thrift encoded)
322/// 3. [`FileMetaData`] (thrift encoded)
323/// 4. Length of encoded `FileMetaData` (4 bytes, little endian)
324/// 5. Parquet Magic Bytes (4 bytes)
325///
326/// [`FileMetaData`]: https://github.com/apache/parquet-format/tree/master?tab=readme-ov-file#metadata
327/// [`ColumnChunkMetaData`]: crate::file::metadata::ColumnChunkMetaData
328/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
329/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
330/// [`PageIndexProvider`]: crate::file::metadata::page_index::PageIndexProvider
331///
332/// ```text
333/// ┌──────────────────────┐
334/// │                      │
335/// │         ...          │
336/// │                      │
337/// │┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ │
338/// │     ColumnIndex     ◀│─ ─ ─
339/// ││    (Optional)     │ │     │
340/// │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─  │
341/// │┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ │     │ FileMetadata
342/// │     OffsetIndex      │       contains embedded
343/// ││    (Optional)     │◀┼ ─   │ offsets to
344/// │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─  │  │    ColumnIndex and
345/// │╔═══════════════════╗ │     │ OffsetIndex
346/// │║                   ║ │  │
347/// │║                   ║ ┼ ─   │
348/// │║   FileMetadata    ║ │
349/// │║                   ║ ┼ ─ ─ ┘
350/// │║                   ║ │
351/// │╚═══════════════════╝ │
352/// │┌───────────────────┐ │
353/// ││  metadata length  │ │ length of FileMetadata  (only)
354/// │└───────────────────┘ │
355/// │┌───────────────────┐ │
356/// ││      'PAR1'       │ │ Parquet Magic Bytes
357/// │└───────────────────┘ │
358/// └──────────────────────┘
359///      Output Buffer
360/// ```
361///
362/// # Example
363/// ```no_run
364/// # use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataWriter};
365/// # fn get_metadata() -> ParquetMetaData { unimplemented!(); }
366/// // write parquet metadata to an in-memory buffer
367/// let mut buffer = vec![];
368/// let metadata: ParquetMetaData = get_metadata();
369/// let writer = ParquetMetaDataWriter::new(&mut buffer, &metadata);
370/// // write the metadata to the buffer
371/// writer.finish().unwrap();
372/// assert!(!buffer.is_empty());
373/// ```
374pub struct ParquetMetaDataWriter<'a, W: Write> {
375    buf: TrackedWrite<W>,
376    metadata: &'a ParquetMetaData,
377    write_path_in_schema: bool,
378    preserve_page_index_locations: bool,
379}
380
381impl<'a, W: Write> ParquetMetaDataWriter<'a, W> {
382    /// Create a new `ParquetMetaDataWriter` to write to `buf`
383    ///
384    /// Note any embedded offsets in the metadata will be written assuming the
385    /// metadata is at the start of the buffer. If the metadata is being written
386    /// to a location other than the start of the buffer, see [`Self::new_with_tracked`]
387    ///
388    /// See example on the struct level documentation
389    pub fn new(buf: W, metadata: &'a ParquetMetaData) -> Self {
390        Self::new_with_tracked(TrackedWrite::new(buf), metadata)
391    }
392
393    /// Create a new ParquetMetaDataWriter to write to `buf`
394    ///
395    /// This method is used when the metadata is being written to a location other
396    /// than the start of the buffer.
397    ///
398    /// See example on the struct level documentation
399    pub fn new_with_tracked(buf: TrackedWrite<W>, metadata: &'a ParquetMetaData) -> Self {
400        Self {
401            buf,
402            metadata,
403            write_path_in_schema: true,
404            preserve_page_index_locations: true,
405        }
406    }
407
408    /// Set whether or not to write the `path_in_schema` field in the Thrift `ColumnMetaData`
409    /// struct.
410    pub fn with_write_path_in_schema(self, val: bool) -> Self {
411        Self {
412            write_path_in_schema: val,
413            ..self
414        }
415    }
416
417    /// Set whether or not to preserve the page index location metadata in the Thrift
418    /// `ColumnMetaData` (defaults to `true`).
419    ///
420    /// Because this struct is often used to externalize the footer metadata, it is
421    /// usually desirable to preserve this location information, even when the
422    /// page indexes are not duplicated (for instance if the provided `ParquetMetaData`
423    /// returns no `PageIndexProvider`). As such, this defaults to `true`.
424    ///
425    /// Set this to `false` to reset the location metadata if no page indexes are
426    /// present in the provided [`ParquetMetaData`].
427    pub fn with_preserve_page_index_locations(self, val: bool) -> Self {
428        Self {
429            preserve_page_index_locations: val,
430            ..self
431        }
432    }
433
434    /// Write the metadata to the buffer
435    pub fn finish(mut self) -> Result<()> {
436        let file_metadata = self.metadata.file_metadata();
437
438        let schema = Arc::new(file_metadata.schema().clone());
439        let schema_descr = Arc::new(SchemaDescriptor::new(schema.clone()));
440        let created_by = file_metadata.created_by().map(str::to_string);
441
442        // need to clone because ThriftMetadataWriter wants to update page
443        // index pointers. we throw this clone away upon calling encoder.finish()
444        let row_groups = self.metadata.row_groups.as_ref().to_vec();
445
446        let key_value_metadata = file_metadata.key_value_metadata().cloned();
447
448        let mut encoder = ThriftMetadataWriter::new(
449            &mut self.buf,
450            &schema_descr,
451            row_groups,
452            created_by,
453            file_metadata.version(),
454            self.write_path_in_schema,
455            self.preserve_page_index_locations,
456        );
457
458        if let Some(page_index_arc) = self.metadata.page_index.clone() {
459            encoder = encoder.with_page_index(page_index_arc);
460        }
461
462        if let Some(key_value_metadata) = key_value_metadata {
463            encoder = encoder.with_key_value_metadata(key_value_metadata);
464        }
465        encoder.finish()?;
466
467        Ok(())
468    }
469}
470
471#[derive(Debug, Default)]
472struct MetadataObjectWriter {
473    #[cfg(feature = "encryption")]
474    file_encryptor: Option<Arc<FileEncryptor>>,
475}
476
477impl MetadataObjectWriter {
478    #[inline]
479    fn write_thrift_object(object: &impl WriteThrift, sink: impl Write) -> Result<()> {
480        let mut protocol = ThriftCompactOutputProtocol::new(sink);
481        object.write_thrift(&mut protocol)?;
482        Ok(())
483    }
484}
485
486/// Implementations of [`MetadataObjectWriter`] methods for when encryption is disabled
487#[cfg(not(feature = "encryption"))]
488impl MetadataObjectWriter {
489    /// Write [`FileMetaData`] in Thrift format
490    ///
491    /// [`FileMetaData`]: https://github.com/apache/parquet-format/tree/master?tab=readme-ov-file#metadata
492    fn write_file_metadata(&self, file_metadata: &FileMeta, sink: impl Write) -> Result<()> {
493        Self::write_thrift_object(file_metadata, sink)
494    }
495
496    /// Write a column [`OffsetIndex`] in Thrift format
497    ///
498    /// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
499    fn write_offset_index(
500        &self,
501        offset_index: &OffsetIndexMetaData,
502        _column_chunk: &ColumnChunkMetaData,
503        _row_group_idx: usize,
504        _column_idx: usize,
505        sink: impl Write,
506    ) -> Result<()> {
507        Self::write_thrift_object(offset_index, sink)
508    }
509
510    /// Write a column [`ColumnIndex`] in Thrift format
511    ///
512    /// Returns `true` unless there is an error.
513    ///
514    /// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
515    fn write_column_index(
516        &self,
517        column_index: &ColumnIndexMetaData,
518        _column_chunk: &ColumnChunkMetaData,
519        _row_group_idx: usize,
520        _column_idx: usize,
521        sink: impl Write,
522    ) -> Result<bool> {
523        Self::write_thrift_object(column_index, sink)?;
524        Ok(true)
525    }
526
527    /// No-op implementation of row-group metadata encryption
528    fn apply_row_group_encryption(
529        &self,
530        row_groups: Vec<RowGroupMetaData>,
531    ) -> Result<(Vec<RowGroupMetaData>, Option<Vec<RowGroupMetaData>>)> {
532        Ok((row_groups, None))
533    }
534
535    /// Get the "magic" bytes identifying the file type
536    pub fn get_file_magic(&self) -> &[u8; 4] {
537        get_file_magic()
538    }
539}
540
541/// Implementations of [`MetadataObjectWriter`] methods that rely on encryption being enabled
542#[cfg(feature = "encryption")]
543impl MetadataObjectWriter {
544    /// Set the file encryptor to use
545    fn with_file_encryptor(mut self, encryptor: Option<Arc<FileEncryptor>>) -> Self {
546        self.file_encryptor = encryptor;
547        self
548    }
549
550    /// Write [`FileMetaData`] in Thrift format, possibly encrypting it if required
551    ///
552    /// [`FileMetaData`]: https://github.com/apache/parquet-format/tree/master?tab=readme-ov-file#metadata
553    fn write_file_metadata(&self, file_metadata: &FileMeta, mut sink: impl Write) -> Result<()> {
554        match self.file_encryptor.as_ref() {
555            Some(file_encryptor) if file_encryptor.properties().encrypt_footer() => {
556                // First write FileCryptoMetadata
557                let crypto_metadata = Self::file_crypto_metadata(file_encryptor)?;
558                let mut protocol = ThriftCompactOutputProtocol::new(&mut sink);
559                crypto_metadata.write_thrift(&mut protocol)?;
560
561                // Then write encrypted footer
562                let aad = create_footer_aad(file_encryptor.file_aad())?;
563                let mut encryptor = file_encryptor.get_footer_encryptor()?;
564                encrypt_thrift_object(file_metadata, &mut encryptor, &mut sink, &aad)
565            }
566            Some(file_encryptor) if file_metadata.file_metadata.encryption_algorithm.is_some() => {
567                let aad = create_footer_aad(file_encryptor.file_aad())?;
568                let mut encryptor = file_encryptor.get_footer_encryptor()?;
569                write_signed_plaintext_thrift_object(file_metadata, &mut encryptor, &mut sink, &aad)
570            }
571            _ => Self::write_thrift_object(file_metadata, &mut sink),
572        }
573    }
574
575    /// Write a column [`OffsetIndex`] in Thrift format, possibly encrypting it if required
576    ///
577    /// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
578    fn write_offset_index(
579        &self,
580        offset_index: &OffsetIndexMetaData,
581        column_chunk: &ColumnChunkMetaData,
582        row_group_idx: usize,
583        column_idx: usize,
584        sink: impl Write,
585    ) -> Result<()> {
586        match &self.file_encryptor {
587            Some(file_encryptor) => Self::write_thrift_object_with_encryption(
588                offset_index,
589                sink,
590                file_encryptor,
591                column_chunk,
592                ModuleType::OffsetIndex,
593                row_group_idx,
594                column_idx,
595            ),
596            None => Self::write_thrift_object(offset_index, sink),
597        }
598    }
599
600    /// Write a column [`ColumnIndex`] in Thrift format, possibly encrypting it if required
601    ///
602    /// Returns `true` unless there is an error.
603    ///
604    /// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
605    fn write_column_index(
606        &self,
607        column_index: &ColumnIndexMetaData,
608        column_chunk: &ColumnChunkMetaData,
609        row_group_idx: usize,
610        column_idx: usize,
611        sink: impl Write,
612    ) -> Result<bool> {
613        match &self.file_encryptor {
614            Some(file_encryptor) => Self::write_thrift_object_with_encryption(
615                column_index,
616                sink,
617                file_encryptor,
618                column_chunk,
619                ModuleType::ColumnIndex,
620                row_group_idx,
621                column_idx,
622            )?,
623            None => Self::write_thrift_object(column_index, sink)?,
624        }
625        Ok(true)
626    }
627
628    /// If encryption is enabled and configured, encrypt row group metadata.
629    /// Returns a tuple of the row group metadata to write,
630    /// and possibly unencrypted metadata to be returned to clients if data was encrypted.
631    fn apply_row_group_encryption(
632        &self,
633        row_groups: Vec<RowGroupMetaData>,
634    ) -> Result<(Vec<RowGroupMetaData>, Option<Vec<RowGroupMetaData>>)> {
635        match &self.file_encryptor {
636            Some(file_encryptor) => {
637                let unencrypted_row_groups = row_groups.clone();
638                let encrypted_row_groups = Self::encrypt_row_groups(row_groups, file_encryptor)?;
639                Ok((encrypted_row_groups, Some(unencrypted_row_groups)))
640            }
641            None => Ok((row_groups, None)),
642        }
643    }
644
645    /// Get the "magic" bytes identifying the file type
646    fn get_file_magic(&self) -> &[u8; 4] {
647        get_file_magic(
648            self.file_encryptor
649                .as_ref()
650                .map(|encryptor| encryptor.properties()),
651        )
652    }
653
654    fn write_thrift_object_with_encryption(
655        object: &impl WriteThrift,
656        mut sink: impl Write,
657        file_encryptor: &FileEncryptor,
658        column_metadata: &ColumnChunkMetaData,
659        module_type: ModuleType,
660        row_group_index: usize,
661        column_index: usize,
662    ) -> Result<()> {
663        let column_path_vec = column_metadata.column_path().as_ref();
664
665        let joined_column_path;
666        let column_path = if column_path_vec.len() == 1 {
667            &column_path_vec[0]
668        } else {
669            joined_column_path = column_path_vec.join(".");
670            &joined_column_path
671        };
672
673        if file_encryptor.is_column_encrypted(column_path) {
674            use crate::encryption::encrypt::encrypt_thrift_object;
675
676            let aad = create_module_aad(
677                file_encryptor.file_aad(),
678                module_type,
679                row_group_index,
680                column_index,
681                None,
682            )?;
683            let mut encryptor = file_encryptor.get_column_encryptor(column_path)?;
684            encrypt_thrift_object(object, &mut encryptor, &mut sink, &aad)
685        } else {
686            Self::write_thrift_object(object, sink)
687        }
688    }
689
690    fn get_plaintext_footer_crypto_metadata(
691        &self,
692    ) -> (Option<EncryptionAlgorithm>, Option<Vec<u8>>) {
693        // Only plaintext footers may contain encryption algorithm and footer key metadata.
694        if let Some(file_encryptor) = self.file_encryptor.as_ref() {
695            let encryption_properties = file_encryptor.properties();
696            if !encryption_properties.encrypt_footer() {
697                return (
698                    Some(Self::encryption_algorithm_from_encryptor(file_encryptor)),
699                    encryption_properties.footer_key_metadata().cloned(),
700                );
701            }
702        }
703        (None, None)
704    }
705
706    fn encryption_algorithm_from_encryptor(file_encryptor: &FileEncryptor) -> EncryptionAlgorithm {
707        let supply_aad_prefix = file_encryptor
708            .properties()
709            .aad_prefix()
710            .map(|_| !file_encryptor.properties().store_aad_prefix());
711        let aad_prefix = if file_encryptor.properties().store_aad_prefix() {
712            file_encryptor.properties().aad_prefix()
713        } else {
714            None
715        };
716        EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
717            aad_prefix: aad_prefix.cloned(),
718            aad_file_unique: Some(file_encryptor.aad_file_unique().clone()),
719            supply_aad_prefix,
720        })
721    }
722
723    fn file_crypto_metadata(file_encryptor: &'_ FileEncryptor) -> Result<FileCryptoMetaData<'_>> {
724        let properties = file_encryptor.properties();
725        Ok(FileCryptoMetaData {
726            encryption_algorithm: Self::encryption_algorithm_from_encryptor(file_encryptor),
727            key_metadata: properties.footer_key_metadata().map(|v| v.as_slice()),
728        })
729    }
730
731    fn encrypt_row_groups(
732        row_groups: Vec<RowGroupMetaData>,
733        file_encryptor: &Arc<FileEncryptor>,
734    ) -> Result<Vec<RowGroupMetaData>> {
735        row_groups
736            .into_iter()
737            .enumerate()
738            .map(|(rg_idx, mut rg)| {
739                let cols: Result<Vec<ColumnChunkMetaData>> = rg
740                    .columns
741                    .into_iter()
742                    .enumerate()
743                    .map(|(col_idx, c)| {
744                        Self::encrypt_column_chunk(c, file_encryptor, rg_idx, col_idx)
745                    })
746                    .collect();
747                rg.columns = cols?;
748                Ok(rg)
749            })
750            .collect()
751    }
752
753    /// Apply column encryption to column chunk metadata
754    fn encrypt_column_chunk(
755        mut column_chunk: ColumnChunkMetaData,
756        file_encryptor: &Arc<FileEncryptor>,
757        row_group_index: usize,
758        column_index: usize,
759    ) -> Result<ColumnChunkMetaData> {
760        // Column crypto metadata should have already been set when the column was created.
761        // Here we apply the encryption by encrypting the column metadata if required.
762        let encryptor = match column_chunk.column_crypto_metadata.as_deref() {
763            None => None,
764            Some(ColumnCryptoMetaData::ENCRYPTION_WITH_FOOTER_KEY) => {
765                let is_footer_encrypted = file_encryptor.properties().encrypt_footer();
766
767                // When uniform encryption is used the footer is already encrypted,
768                // so the column chunk does not need additional encryption.
769                // Except if we're in plaintext footer mode, then we need to encrypt
770                // the column metadata here.
771                if !is_footer_encrypted {
772                    Some(file_encryptor.get_footer_encryptor()?)
773                } else {
774                    None
775                }
776            }
777            Some(ColumnCryptoMetaData::ENCRYPTION_WITH_COLUMN_KEY(col_key)) => {
778                let column_path = col_key.path_in_schema.join(".");
779                Some(file_encryptor.get_column_encryptor(&column_path)?)
780            }
781        };
782
783        if let Some(mut encryptor) = encryptor {
784            use crate::file::metadata::thrift::serialize_column_meta_data;
785
786            let aad = create_module_aad(
787                file_encryptor.file_aad(),
788                ModuleType::ColumnMetaData,
789                row_group_index,
790                column_index,
791                None,
792            )?;
793            // create temp ColumnMetaData that we can encrypt
794            let mut buffer: Vec<u8> = vec![];
795            {
796                let mut prot = ThriftCompactOutputProtocol::new(&mut buffer);
797                serialize_column_meta_data(&column_chunk, &mut prot)?;
798            }
799            let ciphertext = encryptor.encrypt(&buffer, &aad)?;
800            column_chunk.encrypted_column_metadata = Some(ciphertext);
801            // Track whether the footer is plaintext, which affects how we serialize
802            // the column metadata (we need to write stripped metadata for backward compatibility)
803            column_chunk.plaintext_footer_mode = !file_encryptor.properties().encrypt_footer();
804        }
805
806        Ok(column_chunk)
807    }
808}
parquet/file/metadata/writer.rs

parquet/file/metadata/
writer.rs
