Skip to main content

parquet/file/metadata/
mod.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//! Parquet metadata API
19//!
20//! Users should use these structures to interact with Parquet metadata.
21//!
22//! * [`ParquetMetaData`]: Top level metadata container, read from the Parquet
23//!   file footer.
24//!
25//! * [`FileMetaData`]: File level metadata such as schema, row counts and
26//!   version.
27//!
28//! * [`RowGroupMetaData`]: Metadata for each Row Group with a File, such as
29//!   location and number of rows, and column chunks.
30//!
31//! * [`ColumnChunkMetaData`]: Metadata for each column chunk (primitive leaf)
32//!   within a Row Group including encoding and compression information,
33//!   number of values, statistics, etc.
34//!
35//! # APIs for working with Parquet Metadata
36//!
37//! The Parquet readers and writers in this crate handle reading and writing
38//! metadata into parquet files. To work with metadata directly,
39//! the following APIs are available:
40//!
41//! * [`ParquetMetaDataReader`] for reading metadata from an I/O source (sync and async)
42//! * [`ParquetMetaDataPushDecoder`] for decoding from bytes without I/O
43//! * [`ParquetMetaDataWriter`] for writing.
44//!
45//! # Examples
46//!
47//! Please see [`external_metadata.rs`]
48//!
49//! [`external_metadata.rs`]: https://github.com/apache/arrow-rs/tree/master/parquet/examples/external_metadata.rs
50//!
51//! # Metadata Encodings and Structures
52//!
53//! There are three different encodings of Parquet Metadata in this crate:
54//!
55//! 1. `bytes`:encoded with the Thrift `TCompactProtocol` as defined in
56//!    [parquet.thrift]
57//!
58//! 2. [`format`]: Rust structures automatically generated by the thrift compiler
59//!    from [parquet.thrift]. These structures are low level and mirror
60//!    the thrift definitions.
61//!
62//! 3. [`file::metadata`] (this module): Easier to use Rust structures
63//!    with a more idiomatic API. Note that, confusingly, some but not all
64//!    of these structures have the same name as the [`format`] structures.
65//!
66//! [`file::metadata`]: crate::file::metadata
67//! [parquet.thrift]:  https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift
68//!
69//! Graphically, this is how the different structures relate to each other:
70//!
71//! ```text
72//!                          ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─         ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
73//!                            ┌──────────────┐     │         ┌───────────────────────┐ │
74//!                          │ │ ColumnIndex  │              ││    ParquetMetaData    │
75//!                            └──────────────┘     │         └───────────────────────┘ │
76//! ┌──────────────┐         │ ┌────────────────┐            │┌───────────────────────┐
77//! │   ..0x24..   │ ◀────▶    │  OffsetIndex   │   │ ◀────▶  │    ParquetMetaData    │ │
78//! └──────────────┘         │ └────────────────┘            │└───────────────────────┘
79//!                                     ...         │                   ...             │
80//!                          │ ┌──────────────────┐          │ ┌──────────────────┐
81//! bytes                      │  FileMetaData*   │ │          │  FileMetaData*   │     │
82//! (thrift encoded)         │ └──────────────────┘          │ └──────────────────┘
83//!                           ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘         ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘
84//!
85//!                          format::meta structures          file::metadata structures
86//!
87//!                         * Same name, different struct
88//! ```
89mod footer_tail;
90mod memory;
91mod options;
92mod parser;
93mod push_decoder;
94pub(crate) mod reader;
95pub(crate) mod thrift;
96mod writer;
97
98use crate::basic::{EncodingMask, PageType};
99#[cfg(feature = "encryption")]
100use crate::encryption::decrypt::FileDecryptor;
101#[cfg(feature = "encryption")]
102use crate::file::column_crypto_metadata::ColumnCryptoMetaData;
103pub(crate) use crate::file::metadata::memory::HeapSize;
104#[cfg(feature = "encryption")]
105use crate::file::metadata::thrift::encryption::EncryptionAlgorithm;
106use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex};
107use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation};
108use crate::file::statistics::Statistics;
109use crate::geospatial::statistics as geo_statistics;
110use crate::schema::types::{
111    ColumnDescPtr, ColumnDescriptor, ColumnPath, SchemaDescPtr, SchemaDescriptor,
112    Type as SchemaType,
113};
114use crate::thrift_struct;
115use crate::{
116    basic::BoundaryOrder,
117    errors::{ParquetError, Result},
118};
119use crate::{
120    basic::{ColumnOrder, Compression, Encoding, Type},
121    parquet_thrift::{
122        ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
123        ThriftCompactOutputProtocol, WriteThrift, WriteThriftField,
124    },
125};
126use crate::{
127    data_type::private::ParquetValueType, file::page_index::offset_index::OffsetIndexMetaData,
128};
129
130pub use footer_tail::FooterTail;
131pub use options::{ParquetMetaDataOptions, ParquetStatisticsPolicy};
132pub use push_decoder::ParquetMetaDataPushDecoder;
133pub use reader::{PageIndexPolicy, ParquetMetaDataReader};
134use std::io::Write;
135use std::ops::Range;
136use std::sync::Arc;
137pub use writer::ParquetMetaDataWriter;
138pub(crate) use writer::ThriftMetadataWriter;
139
140/// Page level statistics for each column chunk of each row group.
141///
142/// This structure is an in-memory representation of multiple [`ColumnIndex`]
143/// structures in a parquet file footer, as described in the Parquet [PageIndex
144/// documentation]. Each [`ColumnIndex`] holds statistics about all the pages in a
145/// particular column chunk.
146///
147/// `column_index[row_group_number][column_number]` holds the
148/// [`ColumnIndex`] corresponding to column `column_number` of row group
149/// `row_group_number`.
150///
151/// For example `column_index[2][3]` holds the [`ColumnIndex`] for the fourth
152/// column in the third row group of the parquet file.
153///
154/// [PageIndex documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
155/// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
156pub type ParquetColumnIndex = Vec<Vec<ColumnIndexMetaData>>;
157
158/// [`OffsetIndexMetaData`] for each data page of each row group of each column
159///
160/// This structure is the parsed representation of the [`OffsetIndex`] from the
161/// Parquet file footer, as described in the Parquet [PageIndex documentation].
162///
163/// `offset_index[row_group_number][column_number]` holds
164/// the [`OffsetIndexMetaData`] corresponding to column
165/// `column_number`of row group `row_group_number`.
166///
167/// [PageIndex documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
168/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
169pub type ParquetOffsetIndex = Vec<Vec<OffsetIndexMetaData>>;
170
171/// Parsed metadata for a single Parquet file
172///
173/// This structure is stored in the footer of Parquet files, in the format
174/// defined by [`parquet.thrift`].
175///
176/// # Overview
177/// The fields of this structure are:
178/// * [`FileMetaData`]: Information about the overall file (such as the schema) (See [`Self::file_metadata`])
179/// * [`RowGroupMetaData`]: Information about each Row Group (see [`Self::row_groups`])
180/// * [`ParquetColumnIndex`] and [`ParquetOffsetIndex`]: Optional "Page Index" structures (see [`Self::column_index`] and [`Self::offset_index`])
181///
182/// This structure is read by the various readers in this crate or can be read
183/// directly from a file using the [`ParquetMetaDataReader`] struct.
184///
185/// See the [`ParquetMetaDataBuilder`] to create and modify this structure.
186///
187/// [`parquet.thrift`]: https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift
188#[derive(Debug, Clone, PartialEq)]
189pub struct ParquetMetaData {
190    /// File level metadata
191    file_metadata: FileMetaData,
192    /// Row group metadata
193    row_groups: Vec<RowGroupMetaData>,
194    /// Page level index for each page in each column chunk
195    column_index: Option<ParquetColumnIndex>,
196    /// Offset index for each page in each column chunk
197    offset_index: Option<ParquetOffsetIndex>,
198    /// Optional file decryptor
199    #[cfg(feature = "encryption")]
200    file_decryptor: Option<Box<FileDecryptor>>,
201}
202
203impl ParquetMetaData {
204    /// Creates Parquet metadata from file metadata and a list of row
205    /// group metadata
206    pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>) -> Self {
207        ParquetMetaData {
208            file_metadata,
209            row_groups,
210            column_index: None,
211            offset_index: None,
212            #[cfg(feature = "encryption")]
213            file_decryptor: None,
214        }
215    }
216
217    /// Adds [`FileDecryptor`] to this metadata instance to enable decryption of
218    /// encrypted data.
219    #[cfg(feature = "encryption")]
220    pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option<FileDecryptor>) {
221        self.file_decryptor = file_decryptor.map(Box::new);
222    }
223
224    /// Convert this ParquetMetaData into a [`ParquetMetaDataBuilder`]
225    pub fn into_builder(self) -> ParquetMetaDataBuilder {
226        self.into()
227    }
228
229    /// Returns file metadata as reference.
230    pub fn file_metadata(&self) -> &FileMetaData {
231        &self.file_metadata
232    }
233
234    /// Returns file decryptor as reference.
235    #[cfg(feature = "encryption")]
236    pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
237        self.file_decryptor.as_deref()
238    }
239
240    /// Returns number of row groups in this file.
241    pub fn num_row_groups(&self) -> usize {
242        self.row_groups.len()
243    }
244
245    /// Returns row group metadata for `i`th position.
246    /// Position should be less than number of row groups `num_row_groups`.
247    pub fn row_group(&self, i: usize) -> &RowGroupMetaData {
248        &self.row_groups[i]
249    }
250
251    /// Returns slice of row groups in this file.
252    pub fn row_groups(&self) -> &[RowGroupMetaData] {
253        &self.row_groups
254    }
255
256    /// Returns the column index for this file if loaded
257    ///
258    /// Returns `None` if the parquet file does not have a `ColumnIndex` or
259    /// [ArrowReaderOptions::with_page_index] was set to false.
260    ///
261    /// [ArrowReaderOptions::with_page_index]: https://docs.rs/parquet/latest/parquet/arrow/arrow_reader/struct.ArrowReaderOptions.html#method.with_page_index
262    pub fn column_index(&self) -> Option<&ParquetColumnIndex> {
263        self.column_index.as_ref()
264    }
265
266    /// Returns offset indexes in this file, if loaded
267    ///
268    /// Returns `None` if the parquet file does not have a `OffsetIndex` or
269    /// [ArrowReaderOptions::with_page_index] was set to false.
270    ///
271    /// [ArrowReaderOptions::with_page_index]: https://docs.rs/parquet/latest/parquet/arrow/arrow_reader/struct.ArrowReaderOptions.html#method.with_page_index
272    pub fn offset_index(&self) -> Option<&ParquetOffsetIndex> {
273        self.offset_index.as_ref()
274    }
275
276    /// Estimate of the bytes allocated to store `ParquetMetadata`
277    ///
278    /// # Notes:
279    ///
280    /// 1. Includes size of self
281    ///
282    /// 2. Includes heap memory for sub fields such as [`FileMetaData`] and
283    ///    [`RowGroupMetaData`].
284    ///
285    /// 3. Includes memory from shared pointers (e.g. [`SchemaDescPtr`]). This
286    ///    means `memory_size` will over estimate the memory size if such pointers
287    ///    are shared.
288    ///
289    /// 4. Does not include any allocator overheads
290    pub fn memory_size(&self) -> usize {
291        #[cfg(feature = "encryption")]
292        let encryption_size = self.file_decryptor.heap_size();
293        #[cfg(not(feature = "encryption"))]
294        let encryption_size = 0usize;
295
296        std::mem::size_of::<Self>()
297            + self.file_metadata.heap_size()
298            + self.row_groups.heap_size()
299            + self.column_index.heap_size()
300            + self.offset_index.heap_size()
301            + encryption_size
302    }
303
304    /// Override the column index
305    pub(crate) fn set_column_index(&mut self, index: Option<ParquetColumnIndex>) {
306        self.column_index = index;
307    }
308
309    /// Override the offset index
310    pub(crate) fn set_offset_index(&mut self, index: Option<ParquetOffsetIndex>) {
311        self.offset_index = index;
312    }
313}
314
315/// A builder for creating / manipulating [`ParquetMetaData`]
316///
317/// # Example creating a new [`ParquetMetaData`]
318///
319///```no_run
320/// # use parquet::file::metadata::{FileMetaData, ParquetMetaData, ParquetMetaDataBuilder, RowGroupMetaData, RowGroupMetaDataBuilder};
321/// # fn get_file_metadata() -> FileMetaData { unimplemented!(); }
322/// // Create a new builder given the file metadata
323/// let file_metadata = get_file_metadata();
324/// // Create a row group
325/// let row_group = RowGroupMetaData::builder(file_metadata.schema_descr_ptr())
326///    .set_num_rows(100)
327///    // ... (A real row group needs more than just the number of rows)
328///    .build()
329///    .unwrap();
330/// // Create the final metadata
331/// let metadata: ParquetMetaData = ParquetMetaDataBuilder::new(file_metadata)
332///   .add_row_group(row_group)
333///   .build();
334/// ```
335///
336/// # Example modifying an existing [`ParquetMetaData`]
337/// ```no_run
338/// # use parquet::file::metadata::ParquetMetaData;
339/// # fn load_metadata() -> ParquetMetaData { unimplemented!(); }
340/// // Modify the metadata so only the last RowGroup remains
341/// let metadata: ParquetMetaData = load_metadata();
342/// let mut builder = metadata.into_builder();
343///
344/// // Take existing row groups to modify
345/// let mut row_groups = builder.take_row_groups();
346/// let last_row_group = row_groups.pop().unwrap();
347///
348/// let metadata = builder
349///   .add_row_group(last_row_group)
350///   .build();
351/// ```
352pub struct ParquetMetaDataBuilder(ParquetMetaData);
353
354impl ParquetMetaDataBuilder {
355    /// Create a new builder from a file metadata, with no row groups
356    pub fn new(file_meta_data: FileMetaData) -> Self {
357        Self(ParquetMetaData::new(file_meta_data, vec![]))
358    }
359
360    /// Create a new builder from an existing ParquetMetaData
361    pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
362        Self(metadata)
363    }
364
365    /// Adds a row group to the metadata
366    pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
367        self.0.row_groups.push(row_group);
368        self
369    }
370
371    /// Sets all the row groups to the specified list
372    pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self {
373        self.0.row_groups = row_groups;
374        self
375    }
376
377    /// Takes ownership of the row groups in this builder, and clears the list
378    /// of row groups.
379    ///
380    /// This can be used for more efficient creation of a new ParquetMetaData
381    /// from an existing one.
382    pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
383        std::mem::take(&mut self.0.row_groups)
384    }
385
386    /// Return a reference to the current row groups
387    pub fn row_groups(&self) -> &[RowGroupMetaData] {
388        &self.0.row_groups
389    }
390
391    /// Sets the column index
392    pub fn set_column_index(mut self, column_index: Option<ParquetColumnIndex>) -> Self {
393        self.0.column_index = column_index;
394        self
395    }
396
397    /// Returns the current column index from the builder, replacing it with `None`
398    pub fn take_column_index(&mut self) -> Option<ParquetColumnIndex> {
399        std::mem::take(&mut self.0.column_index)
400    }
401
402    /// Return a reference to the current column index, if any
403    pub fn column_index(&self) -> Option<&ParquetColumnIndex> {
404        self.0.column_index.as_ref()
405    }
406
407    /// Sets the offset index
408    pub fn set_offset_index(mut self, offset_index: Option<ParquetOffsetIndex>) -> Self {
409        self.0.offset_index = offset_index;
410        self
411    }
412
413    /// Returns the current offset index from the builder, replacing it with `None`
414    pub fn take_offset_index(&mut self) -> Option<ParquetOffsetIndex> {
415        std::mem::take(&mut self.0.offset_index)
416    }
417
418    /// Return a reference to the current offset index, if any
419    pub fn offset_index(&self) -> Option<&ParquetOffsetIndex> {
420        self.0.offset_index.as_ref()
421    }
422
423    /// Sets the file decryptor needed to decrypt this metadata.
424    #[cfg(feature = "encryption")]
425    pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
426        self.0.with_file_decryptor(file_decryptor);
427        self
428    }
429
430    /// Creates a new ParquetMetaData from the builder
431    pub fn build(self) -> ParquetMetaData {
432        let Self(metadata) = self;
433        metadata
434    }
435}
436
437impl From<ParquetMetaData> for ParquetMetaDataBuilder {
438    fn from(meta_data: ParquetMetaData) -> Self {
439        Self(meta_data)
440    }
441}
442
443thrift_struct!(
444/// A key-value pair for [`FileMetaData`].
445pub struct KeyValue {
446  1: required string key
447  2: optional string value
448}
449);
450
451impl KeyValue {
452    /// Create a new key value pair
453    pub fn new<F2>(key: String, value: F2) -> KeyValue
454    where
455        F2: Into<Option<String>>,
456    {
457        KeyValue {
458            key,
459            value: value.into(),
460        }
461    }
462}
463
464thrift_struct!(
465/// PageEncodingStats for a column chunk and data page.
466pub struct PageEncodingStats {
467  1: required PageType page_type;
468  2: required Encoding encoding;
469  3: required i32 count;
470}
471);
472
473/// Internal representation of the page encoding stats in the [`ColumnChunkMetaData`].
474/// This is not publicly exposed, with different getters defined for each variant.
475#[derive(Debug, Clone, PartialEq)]
476enum ParquetPageEncodingStats {
477    /// The full array of stats as defined in the Parquet spec.
478    Full(Vec<PageEncodingStats>),
479    /// A condensed version of only page encodings seen.
480    Mask(EncodingMask),
481}
482
483/// Reference counted pointer for [`FileMetaData`].
484pub type FileMetaDataPtr = Arc<FileMetaData>;
485
486/// File level metadata for a Parquet file.
487///
488/// Includes the version of the file, metadata, number of rows, schema, and column orders
489#[derive(Debug, Clone, PartialEq)]
490pub struct FileMetaData {
491    version: i32,
492    num_rows: i64,
493    created_by: Option<String>,
494    key_value_metadata: Option<Vec<KeyValue>>,
495    schema_descr: SchemaDescPtr,
496    column_orders: Option<Vec<ColumnOrder>>,
497    #[cfg(feature = "encryption")]
498    encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
499    #[cfg(feature = "encryption")]
500    footer_signing_key_metadata: Option<Vec<u8>>,
501}
502
503impl FileMetaData {
504    /// Creates new file metadata.
505    pub fn new(
506        version: i32,
507        num_rows: i64,
508        created_by: Option<String>,
509        key_value_metadata: Option<Vec<KeyValue>>,
510        schema_descr: SchemaDescPtr,
511        column_orders: Option<Vec<ColumnOrder>>,
512    ) -> Self {
513        FileMetaData {
514            version,
515            num_rows,
516            created_by,
517            key_value_metadata,
518            schema_descr,
519            column_orders,
520            #[cfg(feature = "encryption")]
521            encryption_algorithm: None,
522            #[cfg(feature = "encryption")]
523            footer_signing_key_metadata: None,
524        }
525    }
526
527    #[cfg(feature = "encryption")]
528    pub(crate) fn with_encryption_algorithm(
529        mut self,
530        encryption_algorithm: Option<EncryptionAlgorithm>,
531    ) -> Self {
532        self.encryption_algorithm = encryption_algorithm.map(Box::new);
533        self
534    }
535
536    #[cfg(feature = "encryption")]
537    pub(crate) fn with_footer_signing_key_metadata(
538        mut self,
539        footer_signing_key_metadata: Option<Vec<u8>>,
540    ) -> Self {
541        self.footer_signing_key_metadata = footer_signing_key_metadata;
542        self
543    }
544
545    /// Returns version of this file.
546    pub fn version(&self) -> i32 {
547        self.version
548    }
549
550    /// Returns number of rows in the file.
551    pub fn num_rows(&self) -> i64 {
552        self.num_rows
553    }
554
555    /// String message for application that wrote this file.
556    ///
557    /// This should have the following format:
558    /// `<application> version <application version> (build <application build hash>)`.
559    ///
560    /// ```shell
561    /// parquet-mr version 1.8.0 (build 0fda28af84b9746396014ad6a415b90592a98b3b)
562    /// ```
563    pub fn created_by(&self) -> Option<&str> {
564        self.created_by.as_deref()
565    }
566
567    /// Returns key_value_metadata of this file.
568    pub fn key_value_metadata(&self) -> Option<&Vec<KeyValue>> {
569        self.key_value_metadata.as_ref()
570    }
571
572    /// Returns Parquet [`Type`] that describes schema in this file.
573    ///
574    /// [`Type`]: crate::schema::types::Type
575    pub fn schema(&self) -> &SchemaType {
576        self.schema_descr.root_schema()
577    }
578
579    /// Returns a reference to schema descriptor.
580    pub fn schema_descr(&self) -> &SchemaDescriptor {
581        &self.schema_descr
582    }
583
584    /// Returns reference counted clone for schema descriptor.
585    pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
586        self.schema_descr.clone()
587    }
588
589    /// Column (sort) order used for `min` and `max` values of each column in this file.
590    ///
591    /// Each column order corresponds to one column, determined by its position in the
592    /// list, matching the position of the column in the schema.
593    ///
594    /// When `None` is returned, there are no column orders available, and each column
595    /// should be assumed to have undefined (legacy) column order.
596    pub fn column_orders(&self) -> Option<&Vec<ColumnOrder>> {
597        self.column_orders.as_ref()
598    }
599
600    /// Returns column order for `i`th column in this file.
601    /// If column orders are not available, returns undefined (legacy) column order.
602    pub fn column_order(&self, i: usize) -> ColumnOrder {
603        self.column_orders
604            .as_ref()
605            .map(|data| data[i])
606            .unwrap_or(ColumnOrder::UNDEFINED)
607    }
608}
609
610thrift_struct!(
611/// Sort order within a RowGroup of a leaf column
612pub struct SortingColumn {
613  /// The ordinal position of the column (in this row group)
614  1: required i32 column_idx
615
616  /// If true, indicates this column is sorted in descending order.
617  2: required bool descending
618
619  /// If true, nulls will come before non-null values, otherwise,
620  /// nulls go at the end. */
621  3: required bool nulls_first
622}
623);
624
625/// Reference counted pointer for [`RowGroupMetaData`].
626pub type RowGroupMetaDataPtr = Arc<RowGroupMetaData>;
627
628/// Metadata for a row group
629///
630/// Includes [`ColumnChunkMetaData`] for each column in the row group, the number of rows
631/// the total byte size of the row group, and the [`SchemaDescriptor`] for the row group.
632#[derive(Debug, Clone, PartialEq)]
633pub struct RowGroupMetaData {
634    columns: Vec<ColumnChunkMetaData>,
635    num_rows: i64,
636    sorting_columns: Option<Vec<SortingColumn>>,
637    total_byte_size: i64,
638    schema_descr: SchemaDescPtr,
639    /// We can't infer from file offset of first column since there may empty columns in row group.
640    file_offset: Option<i64>,
641    /// Ordinal position of this row group in file
642    ordinal: Option<i16>,
643}
644
645impl RowGroupMetaData {
646    /// Returns builder for row group metadata.
647    pub fn builder(schema_descr: SchemaDescPtr) -> RowGroupMetaDataBuilder {
648        RowGroupMetaDataBuilder::new(schema_descr)
649    }
650
651    /// Number of columns in this row group.
652    pub fn num_columns(&self) -> usize {
653        self.columns.len()
654    }
655
656    /// Returns column chunk metadata for `i`th column.
657    pub fn column(&self, i: usize) -> &ColumnChunkMetaData {
658        &self.columns[i]
659    }
660
661    /// Returns slice of column chunk metadata.
662    pub fn columns(&self) -> &[ColumnChunkMetaData] {
663        &self.columns
664    }
665
666    /// Returns mutable slice of column chunk metadata.
667    pub fn columns_mut(&mut self) -> &mut [ColumnChunkMetaData] {
668        &mut self.columns
669    }
670
671    /// Number of rows in this row group.
672    pub fn num_rows(&self) -> i64 {
673        self.num_rows
674    }
675
676    /// Returns the sort ordering of the rows in this RowGroup if any
677    pub fn sorting_columns(&self) -> Option<&Vec<SortingColumn>> {
678        self.sorting_columns.as_ref()
679    }
680
681    /// Total byte size of all uncompressed column data in this row group.
682    pub fn total_byte_size(&self) -> i64 {
683        self.total_byte_size
684    }
685
686    /// Total size of all compressed column data in this row group.
687    pub fn compressed_size(&self) -> i64 {
688        self.columns.iter().map(|c| c.total_compressed_size).sum()
689    }
690
691    /// Returns reference to a schema descriptor.
692    pub fn schema_descr(&self) -> &SchemaDescriptor {
693        self.schema_descr.as_ref()
694    }
695
696    /// Returns reference counted clone of schema descriptor.
697    pub fn schema_descr_ptr(&self) -> SchemaDescPtr {
698        self.schema_descr.clone()
699    }
700
701    /// Returns ordinal position of this row group in file.
702    ///
703    /// For example if this is the first row group in the file, this will return 0.
704    /// If this is the second row group in the file, this will return 1.
705    #[inline(always)]
706    pub fn ordinal(&self) -> Option<i16> {
707        self.ordinal
708    }
709
710    /// Returns file offset of this row group in file.
711    #[inline(always)]
712    pub fn file_offset(&self) -> Option<i64> {
713        self.file_offset
714    }
715
716    /// Converts this [`RowGroupMetaData`] into a [`RowGroupMetaDataBuilder`]
717    pub fn into_builder(self) -> RowGroupMetaDataBuilder {
718        RowGroupMetaDataBuilder(self)
719    }
720}
721
722/// Builder for row group metadata.
723pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
724
725impl RowGroupMetaDataBuilder {
726    /// Creates new builder from schema descriptor.
727    fn new(schema_descr: SchemaDescPtr) -> Self {
728        Self(RowGroupMetaData {
729            columns: Vec::with_capacity(schema_descr.num_columns()),
730            schema_descr,
731            file_offset: None,
732            num_rows: 0,
733            sorting_columns: None,
734            total_byte_size: 0,
735            ordinal: None,
736        })
737    }
738
739    /// Sets number of rows in this row group.
740    pub fn set_num_rows(mut self, value: i64) -> Self {
741        self.0.num_rows = value;
742        self
743    }
744
745    /// Sets the sorting order for columns
746    pub fn set_sorting_columns(mut self, value: Option<Vec<SortingColumn>>) -> Self {
747        self.0.sorting_columns = value;
748        self
749    }
750
751    /// Sets total size in bytes for this row group.
752    pub fn set_total_byte_size(mut self, value: i64) -> Self {
753        self.0.total_byte_size = value;
754        self
755    }
756
757    /// Takes ownership of the the column metadata in this builder, and clears
758    /// the list of columns.
759    ///
760    /// This can be used for more efficient creation of a new RowGroupMetaData
761    /// from an existing one.
762    pub fn take_columns(&mut self) -> Vec<ColumnChunkMetaData> {
763        std::mem::take(&mut self.0.columns)
764    }
765
766    /// Sets column metadata for this row group.
767    pub fn set_column_metadata(mut self, value: Vec<ColumnChunkMetaData>) -> Self {
768        self.0.columns = value;
769        self
770    }
771
772    /// Adds a column metadata to this row group
773    pub fn add_column_metadata(mut self, value: ColumnChunkMetaData) -> Self {
774        self.0.columns.push(value);
775        self
776    }
777
778    /// Sets ordinal for this row group.
779    pub fn set_ordinal(mut self, value: i16) -> Self {
780        self.0.ordinal = Some(value);
781        self
782    }
783
784    /// Sets file offset for this row group.
785    pub fn set_file_offset(mut self, value: i64) -> Self {
786        self.0.file_offset = Some(value);
787        self
788    }
789
790    /// Builds row group metadata.
791    pub fn build(self) -> Result<RowGroupMetaData> {
792        if self.0.schema_descr.num_columns() != self.0.columns.len() {
793            return Err(general_err!(
794                "Column length mismatch: {} != {}",
795                self.0.schema_descr.num_columns(),
796                self.0.columns.len()
797            ));
798        }
799
800        Ok(self.0)
801    }
802
803    /// Build row group metadata without validation.
804    pub(super) fn build_unchecked(self) -> RowGroupMetaData {
805        self.0
806    }
807}
808
809/// Metadata for a column chunk.
810#[derive(Debug, Clone, PartialEq)]
811pub struct ColumnChunkMetaData {
812    column_descr: ColumnDescPtr,
813    encodings: EncodingMask,
814    file_path: Option<String>,
815    file_offset: i64,
816    num_values: i64,
817    compression: Compression,
818    total_compressed_size: i64,
819    total_uncompressed_size: i64,
820    data_page_offset: i64,
821    index_page_offset: Option<i64>,
822    dictionary_page_offset: Option<i64>,
823    statistics: Option<Statistics>,
824    geo_statistics: Option<Box<geo_statistics::GeospatialStatistics>>,
825    encoding_stats: Option<ParquetPageEncodingStats>,
826    bloom_filter_offset: Option<i64>,
827    bloom_filter_length: Option<i32>,
828    offset_index_offset: Option<i64>,
829    offset_index_length: Option<i32>,
830    column_index_offset: Option<i64>,
831    column_index_length: Option<i32>,
832    unencoded_byte_array_data_bytes: Option<i64>,
833    repetition_level_histogram: Option<LevelHistogram>,
834    definition_level_histogram: Option<LevelHistogram>,
835    #[cfg(feature = "encryption")]
836    column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
837    #[cfg(feature = "encryption")]
838    encrypted_column_metadata: Option<Vec<u8>>,
839    /// When true, indicates the footer is plaintext (not encrypted).
840    /// This affects how column metadata is serialized when `encrypted_column_metadata` is present.
841    /// This field is only used at write time and is not needed when reading metadata.
842    #[cfg(feature = "encryption")]
843    plaintext_footer_mode: bool,
844}
845
846/// Histograms for repetition and definition levels.
847///
848/// Each histogram is a vector of length `max_level + 1`. The value at index `i` is the number of
849/// values at level `i`.
850///
851/// For example, `vec[0]` is the number of rows with level 0, `vec[1]` is the
852/// number of rows with level 1, and so on.
853///
854#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
855pub struct LevelHistogram {
856    inner: Vec<i64>,
857}
858
859impl LevelHistogram {
860    /// Creates a new level histogram data.
861    ///
862    /// Length will be `max_level + 1`.
863    ///
864    /// Returns `None` when `max_level == 0` (because histograms are not necessary in this case)
865    pub fn try_new(max_level: i16) -> Option<Self> {
866        if max_level > 0 {
867            Some(Self {
868                inner: vec![0; max_level as usize + 1],
869            })
870        } else {
871            None
872        }
873    }
874    /// Returns a reference to the the histogram's values.
875    pub fn values(&self) -> &[i64] {
876        &self.inner
877    }
878
879    /// Return the inner vector, consuming self
880    pub fn into_inner(self) -> Vec<i64> {
881        self.inner
882    }
883
884    /// Returns the histogram value at the given index.
885    ///
886    /// The value of `i` is the number of values with level `i`. For example,
887    /// `get(1)` returns the number of values with level 1.
888    ///
889    /// Returns `None` if the index is out of bounds.
890    pub fn get(&self, index: usize) -> Option<i64> {
891        self.inner.get(index).copied()
892    }
893
894    /// Adds the values from the other histogram to this histogram
895    ///
896    /// # Panics
897    /// If the histograms have different lengths
898    pub fn add(&mut self, other: &Self) {
899        assert_eq!(self.len(), other.len());
900        for (dst, src) in self.inner.iter_mut().zip(other.inner.iter()) {
901            *dst += src;
902        }
903    }
904
905    /// return the length of the histogram
906    pub fn len(&self) -> usize {
907        self.inner.len()
908    }
909
910    /// returns if the histogram is empty
911    pub fn is_empty(&self) -> bool {
912        self.inner.is_empty()
913    }
914
915    /// Sets the values of all histogram levels to 0.
916    pub fn reset(&mut self) {
917        for value in self.inner.iter_mut() {
918            *value = 0;
919        }
920    }
921
922    /// Increments the count for a level value by `count`.
923    #[inline]
924    pub fn increment_by(&mut self, level: i16, count: i64) {
925        self.inner[level as usize] += count;
926    }
927
928    /// Updates histogram values using provided repetition levels
929    ///
930    /// # Panics
931    /// if any of the levels is greater than the length of the histogram (
932    /// the argument supplied to [`Self::try_new`])
933    #[deprecated(since = "58.2.0", note = "Use `increment_by` instead")]
934    pub fn update_from_levels(&mut self, levels: &[i16]) {
935        for &level in levels {
936            self.increment_by(level, 1);
937        }
938    }
939}
940
941impl From<Vec<i64>> for LevelHistogram {
942    fn from(inner: Vec<i64>) -> Self {
943        Self { inner }
944    }
945}
946
947impl From<LevelHistogram> for Vec<i64> {
948    fn from(value: LevelHistogram) -> Self {
949        value.into_inner()
950    }
951}
952
953impl HeapSize for LevelHistogram {
954    fn heap_size(&self) -> usize {
955        self.inner.heap_size()
956    }
957}
958
959/// Represents common operations for a column chunk.
960impl ColumnChunkMetaData {
961    /// Returns builder for column chunk metadata.
962    pub fn builder(column_descr: ColumnDescPtr) -> ColumnChunkMetaDataBuilder {
963        ColumnChunkMetaDataBuilder::new(column_descr)
964    }
965
966    /// File where the column chunk is stored.
967    ///
968    /// If not set, assumed to belong to the same file as the metadata.
969    /// This path is relative to the current file.
970    pub fn file_path(&self) -> Option<&str> {
971        self.file_path.as_deref()
972    }
973
974    /// Byte offset of `ColumnMetaData` in `file_path()`.
975    ///
976    /// Note that the meaning of this field has been inconsistent between implementations
977    /// so its use has since been deprecated in the Parquet specification. Modern implementations
978    /// will set this to `0` to indicate that the `ColumnMetaData` is solely contained in the
979    /// `ColumnChunk` struct.
980    pub fn file_offset(&self) -> i64 {
981        self.file_offset
982    }
983
984    /// Type of this column. Must be primitive.
985    pub fn column_type(&self) -> Type {
986        self.column_descr.physical_type()
987    }
988
989    /// Path (or identifier) of this column.
990    pub fn column_path(&self) -> &ColumnPath {
991        self.column_descr.path()
992    }
993
994    /// Descriptor for this column.
995    pub fn column_descr(&self) -> &ColumnDescriptor {
996        self.column_descr.as_ref()
997    }
998
999    /// Reference counted clone of descriptor for this column.
1000    pub fn column_descr_ptr(&self) -> ColumnDescPtr {
1001        self.column_descr.clone()
1002    }
1003
1004    /// All encodings used for this column.
1005    pub fn encodings(&self) -> impl Iterator<Item = Encoding> {
1006        self.encodings.encodings()
1007    }
1008
1009    /// All encodings used for this column, returned as a bitmask.
1010    pub fn encodings_mask(&self) -> &EncodingMask {
1011        &self.encodings
1012    }
1013
1014    /// Total number of values in this column chunk.
1015    pub fn num_values(&self) -> i64 {
1016        self.num_values
1017    }
1018
1019    /// Compression for this column.
1020    pub fn compression(&self) -> Compression {
1021        self.compression
1022    }
1023
1024    /// Returns the total compressed data size of this column chunk.
1025    pub fn compressed_size(&self) -> i64 {
1026        self.total_compressed_size
1027    }
1028
1029    /// Returns the total uncompressed data size of this column chunk.
1030    pub fn uncompressed_size(&self) -> i64 {
1031        self.total_uncompressed_size
1032    }
1033
1034    /// Returns the offset for the column data.
1035    pub fn data_page_offset(&self) -> i64 {
1036        self.data_page_offset
1037    }
1038
1039    /// Returns the offset for the index page.
1040    pub fn index_page_offset(&self) -> Option<i64> {
1041        self.index_page_offset
1042    }
1043
1044    /// Returns the offset for the dictionary page, if any.
1045    pub fn dictionary_page_offset(&self) -> Option<i64> {
1046        self.dictionary_page_offset
1047    }
1048
1049    /// Returns the offset and length in bytes of the column chunk within the file
1050    pub fn byte_range(&self) -> (u64, u64) {
1051        let col_start = match self.dictionary_page_offset() {
1052            Some(dictionary_page_offset) => dictionary_page_offset,
1053            None => self.data_page_offset(),
1054        };
1055        let col_len = self.compressed_size();
1056        assert!(
1057            col_start >= 0 && col_len >= 0,
1058            "column start and length should not be negative"
1059        );
1060        (col_start as u64, col_len as u64)
1061    }
1062
1063    /// Returns statistics that are set for this column chunk,
1064    /// or `None` if no statistics are available.
1065    pub fn statistics(&self) -> Option<&Statistics> {
1066        self.statistics.as_ref()
1067    }
1068
1069    /// Returns geospatial statistics that are set for this column chunk,
1070    /// or `None` if no geospatial statistics are available.
1071    pub fn geo_statistics(&self) -> Option<&geo_statistics::GeospatialStatistics> {
1072        self.geo_statistics.as_deref()
1073    }
1074
1075    /// Returns the page encoding statistics, or `None` if no page encoding statistics
1076    /// are available (or they were converted to a mask).
1077    ///
1078    /// Note: By default, this crate converts page encoding statistics to a mask for performance
1079    /// reasons. To get the full statistics, you must set [`ParquetMetaDataOptions::with_encoding_stats_as_mask`]
1080    /// to `false`.
1081    pub fn page_encoding_stats(&self) -> Option<&Vec<PageEncodingStats>> {
1082        match self.encoding_stats.as_ref() {
1083            Some(ParquetPageEncodingStats::Full(stats)) => Some(stats),
1084            _ => None,
1085        }
1086    }
1087
1088    /// Returns the page encoding statistics reduced to a bitmask, or `None` if statistics are
1089    /// not available (or they were left in their original form).
1090    ///
1091    /// Note: This is the default behavior for this crate.
1092    ///
1093    /// The [`PageEncodingStats`] struct was added to the Parquet specification specifically to
1094    /// enable fast determination of whether all pages in a column chunk are dictionary encoded
1095    /// (see <https://github.com/apache/parquet-format/pull/16>).
1096    /// Decoding the full page encoding statistics, however, can be very costly, and is not
1097    /// necessary to support the aforementioned use case. As an alternative, this crate can
1098    /// instead distill the list of `PageEncodingStats` down to a bitmask of just the encodings
1099    /// used for data pages
1100    /// (see [`ParquetMetaDataOptions::set_encoding_stats_as_mask`]).
1101    /// To test for an all-dictionary-encoded chunk one could use this bitmask in the following way:
1102    ///
1103    /// ```rust
1104    /// use parquet::basic::Encoding;
1105    /// use parquet::file::metadata::ColumnChunkMetaData;
1106    /// // test if all data pages in the column chunk are dictionary encoded
1107    /// fn is_all_dictionary_encoded(col_meta: &ColumnChunkMetaData) -> bool {
1108    ///     // check that dictionary encoding was used
1109    ///     col_meta.dictionary_page_offset().is_some()
1110    ///         && col_meta.page_encoding_stats_mask().is_some_and(|mask| {
1111    ///             // mask should only have one bit set, either for PLAIN_DICTIONARY or
1112    ///             // RLE_DICTIONARY
1113    ///             mask.is_only(Encoding::PLAIN_DICTIONARY) || mask.is_only(Encoding::RLE_DICTIONARY)
1114    ///         })
1115    /// }
1116    /// ```
1117    pub fn page_encoding_stats_mask(&self) -> Option<&EncodingMask> {
1118        match self.encoding_stats.as_ref() {
1119            Some(ParquetPageEncodingStats::Mask(stats)) => Some(stats),
1120            _ => None,
1121        }
1122    }
1123
1124    /// Returns the offset for the bloom filter.
1125    pub fn bloom_filter_offset(&self) -> Option<i64> {
1126        self.bloom_filter_offset
1127    }
1128
1129    /// Returns the offset for the bloom filter.
1130    pub fn bloom_filter_length(&self) -> Option<i32> {
1131        self.bloom_filter_length
1132    }
1133
1134    /// Returns the offset for the column index.
1135    pub fn column_index_offset(&self) -> Option<i64> {
1136        self.column_index_offset
1137    }
1138
1139    /// Returns the offset for the column index length.
1140    pub fn column_index_length(&self) -> Option<i32> {
1141        self.column_index_length
1142    }
1143
1144    /// Returns the range for the offset index if any
1145    pub(crate) fn column_index_range(&self) -> Option<Range<u64>> {
1146        let offset = u64::try_from(self.column_index_offset?).ok()?;
1147        let length = u64::try_from(self.column_index_length?).ok()?;
1148        Some(offset..(offset + length))
1149    }
1150
1151    /// Returns the offset for the offset index.
1152    pub fn offset_index_offset(&self) -> Option<i64> {
1153        self.offset_index_offset
1154    }
1155
1156    /// Returns the offset for the offset index length.
1157    pub fn offset_index_length(&self) -> Option<i32> {
1158        self.offset_index_length
1159    }
1160
1161    /// Returns the range for the offset index if any
1162    pub(crate) fn offset_index_range(&self) -> Option<Range<u64>> {
1163        let offset = u64::try_from(self.offset_index_offset?).ok()?;
1164        let length = u64::try_from(self.offset_index_length?).ok()?;
1165        Some(offset..(offset + length))
1166    }
1167
1168    /// Returns the number of bytes of variable length data after decoding.
1169    ///
1170    /// Only set for BYTE_ARRAY columns. This field may not be set by older
1171    /// writers.
1172    pub fn unencoded_byte_array_data_bytes(&self) -> Option<i64> {
1173        self.unencoded_byte_array_data_bytes
1174    }
1175
1176    /// Returns the repetition level histogram.
1177    ///
1178    /// The returned value `vec[i]` is how many values are at repetition level `i`. For example,
1179    /// `vec[0]` indicates how many rows the page contains.
1180    /// This field may not be set by older writers.
1181    pub fn repetition_level_histogram(&self) -> Option<&LevelHistogram> {
1182        self.repetition_level_histogram.as_ref()
1183    }
1184
1185    /// Returns the definition level histogram.
1186    ///
1187    /// The returned value `vec[i]` is how many values are at definition level `i`. For example,
1188    /// `vec[max_definition_level]` indicates how many non-null values are present in the page.
1189    /// This field may not be set by older writers.
1190    pub fn definition_level_histogram(&self) -> Option<&LevelHistogram> {
1191        self.definition_level_histogram.as_ref()
1192    }
1193
1194    /// Returns the encryption metadata for this column chunk.
1195    #[cfg(feature = "encryption")]
1196    pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
1197        self.column_crypto_metadata.as_deref()
1198    }
1199
1200    /// Converts this [`ColumnChunkMetaData`] into a [`ColumnChunkMetaDataBuilder`]
1201    pub fn into_builder(self) -> ColumnChunkMetaDataBuilder {
1202        ColumnChunkMetaDataBuilder::from(self)
1203    }
1204}
1205
1206/// Builder for [`ColumnChunkMetaData`]
1207///
1208/// This builder is used to create a new column chunk metadata or modify an
1209/// existing one.
1210///
1211/// # Example
1212/// ```no_run
1213/// # use parquet::file::metadata::{ColumnChunkMetaData, ColumnChunkMetaDataBuilder};
1214/// # fn get_column_chunk_metadata() -> ColumnChunkMetaData { unimplemented!(); }
1215/// let column_chunk_metadata = get_column_chunk_metadata();
1216/// // create a new builder from existing column chunk metadata
1217/// let builder = ColumnChunkMetaDataBuilder::from(column_chunk_metadata);
1218/// // clear the statistics:
1219/// let column_chunk_metadata: ColumnChunkMetaData = builder
1220///   .clear_statistics()
1221///   .build()
1222///   .unwrap();
1223/// ```
1224pub struct ColumnChunkMetaDataBuilder(ColumnChunkMetaData);
1225
1226impl ColumnChunkMetaDataBuilder {
1227    /// Creates new column chunk metadata builder.
1228    ///
1229    /// See also [`ColumnChunkMetaData::builder`]
1230    fn new(column_descr: ColumnDescPtr) -> Self {
1231        Self(ColumnChunkMetaData {
1232            column_descr,
1233            encodings: Default::default(),
1234            file_path: None,
1235            file_offset: 0,
1236            num_values: 0,
1237            compression: Compression::UNCOMPRESSED,
1238            total_compressed_size: 0,
1239            total_uncompressed_size: 0,
1240            data_page_offset: 0,
1241            index_page_offset: None,
1242            dictionary_page_offset: None,
1243            statistics: None,
1244            geo_statistics: None,
1245            encoding_stats: None,
1246            bloom_filter_offset: None,
1247            bloom_filter_length: None,
1248            offset_index_offset: None,
1249            offset_index_length: None,
1250            column_index_offset: None,
1251            column_index_length: None,
1252            unencoded_byte_array_data_bytes: None,
1253            repetition_level_histogram: None,
1254            definition_level_histogram: None,
1255            #[cfg(feature = "encryption")]
1256            column_crypto_metadata: None,
1257            #[cfg(feature = "encryption")]
1258            encrypted_column_metadata: None,
1259            #[cfg(feature = "encryption")]
1260            plaintext_footer_mode: false,
1261        })
1262    }
1263
1264    /// Sets list of encodings for this column chunk.
1265    pub fn set_encodings(mut self, encodings: Vec<Encoding>) -> Self {
1266        self.0.encodings = EncodingMask::new_from_encodings(encodings.iter());
1267        self
1268    }
1269
1270    /// Sets the encodings mask for this column chunk.
1271    pub fn set_encodings_mask(mut self, encodings: EncodingMask) -> Self {
1272        self.0.encodings = encodings;
1273        self
1274    }
1275
1276    /// Sets optional file path for this column chunk.
1277    pub fn set_file_path(mut self, value: String) -> Self {
1278        self.0.file_path = Some(value);
1279        self
1280    }
1281
1282    /// Sets number of values.
1283    pub fn set_num_values(mut self, value: i64) -> Self {
1284        self.0.num_values = value;
1285        self
1286    }
1287
1288    /// Sets compression.
1289    pub fn set_compression(mut self, value: Compression) -> Self {
1290        self.0.compression = value;
1291        self
1292    }
1293
1294    /// Sets total compressed size in bytes.
1295    pub fn set_total_compressed_size(mut self, value: i64) -> Self {
1296        self.0.total_compressed_size = value;
1297        self
1298    }
1299
1300    /// Sets total uncompressed size in bytes.
1301    pub fn set_total_uncompressed_size(mut self, value: i64) -> Self {
1302        self.0.total_uncompressed_size = value;
1303        self
1304    }
1305
1306    /// Sets data page offset in bytes.
1307    pub fn set_data_page_offset(mut self, value: i64) -> Self {
1308        self.0.data_page_offset = value;
1309        self
1310    }
1311
1312    /// Sets optional dictionary page offset in bytes.
1313    pub fn set_dictionary_page_offset(mut self, value: Option<i64>) -> Self {
1314        self.0.dictionary_page_offset = value;
1315        self
1316    }
1317
1318    /// Sets optional index page offset in bytes.
1319    pub fn set_index_page_offset(mut self, value: Option<i64>) -> Self {
1320        self.0.index_page_offset = value;
1321        self
1322    }
1323
1324    /// Sets statistics for this column chunk.
1325    pub fn set_statistics(mut self, value: Statistics) -> Self {
1326        self.0.statistics = Some(value);
1327        self
1328    }
1329
1330    /// Sets geospatial statistics for this column chunk.
1331    pub fn set_geo_statistics(mut self, value: Box<geo_statistics::GeospatialStatistics>) -> Self {
1332        self.0.geo_statistics = Some(value);
1333        self
1334    }
1335
1336    /// Clears the statistics for this column chunk.
1337    pub fn clear_statistics(mut self) -> Self {
1338        self.0.statistics = None;
1339        self
1340    }
1341
1342    /// Sets page encoding stats for this column chunk.
1343    ///
1344    /// This will overwrite any existing stats, either `Vec` based or bitmask.
1345    pub fn set_page_encoding_stats(mut self, value: Vec<PageEncodingStats>) -> Self {
1346        self.0.encoding_stats = Some(ParquetPageEncodingStats::Full(value));
1347        self
1348    }
1349
1350    /// Sets page encoding stats mask for this column chunk.
1351    ///
1352    /// This will overwrite any existing stats, either `Vec` based or bitmask.
1353    pub fn set_page_encoding_stats_mask(mut self, value: EncodingMask) -> Self {
1354        self.0.encoding_stats = Some(ParquetPageEncodingStats::Mask(value));
1355        self
1356    }
1357
1358    /// Clears the page encoding stats for this column chunk.
1359    pub fn clear_page_encoding_stats(mut self) -> Self {
1360        self.0.encoding_stats = None;
1361        self
1362    }
1363
1364    /// Sets optional bloom filter offset in bytes.
1365    pub fn set_bloom_filter_offset(mut self, value: Option<i64>) -> Self {
1366        self.0.bloom_filter_offset = value;
1367        self
1368    }
1369
1370    /// Sets optional bloom filter length in bytes.
1371    pub fn set_bloom_filter_length(mut self, value: Option<i32>) -> Self {
1372        self.0.bloom_filter_length = value;
1373        self
1374    }
1375
1376    /// Sets optional offset index offset in bytes.
1377    pub fn set_offset_index_offset(mut self, value: Option<i64>) -> Self {
1378        self.0.offset_index_offset = value;
1379        self
1380    }
1381
1382    /// Sets optional offset index length in bytes.
1383    pub fn set_offset_index_length(mut self, value: Option<i32>) -> Self {
1384        self.0.offset_index_length = value;
1385        self
1386    }
1387
1388    /// Sets optional column index offset in bytes.
1389    pub fn set_column_index_offset(mut self, value: Option<i64>) -> Self {
1390        self.0.column_index_offset = value;
1391        self
1392    }
1393
1394    /// Sets optional column index length in bytes.
1395    pub fn set_column_index_length(mut self, value: Option<i32>) -> Self {
1396        self.0.column_index_length = value;
1397        self
1398    }
1399
1400    /// Sets optional length of variable length data in bytes.
1401    pub fn set_unencoded_byte_array_data_bytes(mut self, value: Option<i64>) -> Self {
1402        self.0.unencoded_byte_array_data_bytes = value;
1403        self
1404    }
1405
1406    /// Sets optional repetition level histogram
1407    pub fn set_repetition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1408        self.0.repetition_level_histogram = value;
1409        self
1410    }
1411
1412    /// Sets optional repetition level histogram
1413    pub fn set_definition_level_histogram(mut self, value: Option<LevelHistogram>) -> Self {
1414        self.0.definition_level_histogram = value;
1415        self
1416    }
1417
1418    #[cfg(feature = "encryption")]
1419    /// Set the encryption metadata for an encrypted column
1420    pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
1421        self.0.column_crypto_metadata = value.map(Box::new);
1422        self
1423    }
1424
1425    #[cfg(feature = "encryption")]
1426    /// Set the encryption metadata for an encrypted column
1427    pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
1428        self.0.encrypted_column_metadata = value;
1429        self
1430    }
1431
1432    /// Builds column chunk metadata.
1433    pub fn build(self) -> Result<ColumnChunkMetaData> {
1434        Ok(self.0)
1435    }
1436}
1437
1438/// Builder for Parquet [`ColumnIndex`], part of the Parquet [PageIndex]
1439///
1440/// [PageIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
1441/// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
1442pub struct ColumnIndexBuilder {
1443    column_type: Type,
1444    null_pages: Vec<bool>,
1445    min_values: Vec<Vec<u8>>,
1446    max_values: Vec<Vec<u8>>,
1447    null_counts: Vec<i64>,
1448    boundary_order: BoundaryOrder,
1449    /// contains the concatenation of the histograms of all pages
1450    repetition_level_histograms: Option<Vec<i64>>,
1451    /// contains the concatenation of the histograms of all pages
1452    definition_level_histograms: Option<Vec<i64>>,
1453    /// Is the information in the builder valid?
1454    ///
1455    /// Set to `false` if any entry in the page doesn't have statistics for
1456    /// some reason, so statistics for that page won't be written to the file.
1457    /// This might happen if the page is entirely null, or
1458    /// is a floating point column without any non-nan values
1459    /// e.g. <https://github.com/apache/parquet-format/pull/196>
1460    valid: bool,
1461}
1462
1463impl ColumnIndexBuilder {
1464    /// Creates a new column index builder.
1465    pub fn new(column_type: Type) -> Self {
1466        ColumnIndexBuilder {
1467            column_type,
1468            null_pages: Vec::new(),
1469            min_values: Vec::new(),
1470            max_values: Vec::new(),
1471            null_counts: Vec::new(),
1472            boundary_order: BoundaryOrder::UNORDERED,
1473            repetition_level_histograms: None,
1474            definition_level_histograms: None,
1475            valid: true,
1476        }
1477    }
1478
1479    /// Append statistics for the next page
1480    pub fn append(
1481        &mut self,
1482        null_page: bool,
1483        min_value: Vec<u8>,
1484        max_value: Vec<u8>,
1485        null_count: i64,
1486    ) {
1487        self.null_pages.push(null_page);
1488        self.min_values.push(min_value);
1489        self.max_values.push(max_value);
1490        self.null_counts.push(null_count);
1491    }
1492
1493    /// Append the given page-level histograms to the [`ColumnIndex`] histograms.
1494    /// Does nothing if the `ColumnIndexBuilder` is not in the `valid` state.
1495    ///
1496    /// [`ColumnIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
1497    pub fn append_histograms(
1498        &mut self,
1499        repetition_level_histogram: &Option<LevelHistogram>,
1500        definition_level_histogram: &Option<LevelHistogram>,
1501    ) {
1502        if !self.valid {
1503            return;
1504        }
1505        if let Some(rep_lvl_hist) = repetition_level_histogram {
1506            let hist = self.repetition_level_histograms.get_or_insert(Vec::new());
1507            hist.reserve(rep_lvl_hist.len());
1508            hist.extend(rep_lvl_hist.values());
1509        }
1510        if let Some(def_lvl_hist) = definition_level_histogram {
1511            let hist = self.definition_level_histograms.get_or_insert(Vec::new());
1512            hist.reserve(def_lvl_hist.len());
1513            hist.extend(def_lvl_hist.values());
1514        }
1515    }
1516
1517    /// Set the boundary order of the column index
1518    pub fn set_boundary_order(&mut self, boundary_order: BoundaryOrder) {
1519        self.boundary_order = boundary_order;
1520    }
1521
1522    /// Mark this column index as invalid
1523    pub fn to_invalid(&mut self) {
1524        self.valid = false;
1525    }
1526
1527    /// Is the information in the builder valid?
1528    pub fn valid(&self) -> bool {
1529        self.valid
1530    }
1531
1532    /// Build and get the column index
1533    ///
1534    /// Note: callers should check [`Self::valid`] before calling this method
1535    pub fn build(self) -> Result<ColumnIndexMetaData> {
1536        Ok(match self.column_type {
1537            Type::BOOLEAN => {
1538                let index = self.build_page_index()?;
1539                ColumnIndexMetaData::BOOLEAN(index)
1540            }
1541            Type::INT32 => {
1542                let index = self.build_page_index()?;
1543                ColumnIndexMetaData::INT32(index)
1544            }
1545            Type::INT64 => {
1546                let index = self.build_page_index()?;
1547                ColumnIndexMetaData::INT64(index)
1548            }
1549            Type::INT96 => {
1550                let index = self.build_page_index()?;
1551                ColumnIndexMetaData::INT96(index)
1552            }
1553            Type::FLOAT => {
1554                let index = self.build_page_index()?;
1555                ColumnIndexMetaData::FLOAT(index)
1556            }
1557            Type::DOUBLE => {
1558                let index = self.build_page_index()?;
1559                ColumnIndexMetaData::DOUBLE(index)
1560            }
1561            Type::BYTE_ARRAY => {
1562                let index = self.build_byte_array_index()?;
1563                ColumnIndexMetaData::BYTE_ARRAY(index)
1564            }
1565            Type::FIXED_LEN_BYTE_ARRAY => {
1566                let index = self.build_byte_array_index()?;
1567                ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(index)
1568            }
1569        })
1570    }
1571
1572    fn build_page_index<T>(self) -> Result<PrimitiveColumnIndex<T>>
1573    where
1574        T: ParquetValueType,
1575    {
1576        let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1577        let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1578
1579        PrimitiveColumnIndex::try_new(
1580            self.null_pages,
1581            self.boundary_order,
1582            Some(self.null_counts),
1583            self.repetition_level_histograms,
1584            self.definition_level_histograms,
1585            min_values,
1586            max_values,
1587        )
1588    }
1589
1590    fn build_byte_array_index(self) -> Result<ByteArrayColumnIndex> {
1591        let min_values: Vec<&[u8]> = self.min_values.iter().map(|v| v.as_slice()).collect();
1592        let max_values: Vec<&[u8]> = self.max_values.iter().map(|v| v.as_slice()).collect();
1593
1594        ByteArrayColumnIndex::try_new(
1595            self.null_pages,
1596            self.boundary_order,
1597            Some(self.null_counts),
1598            self.repetition_level_histograms,
1599            self.definition_level_histograms,
1600            min_values,
1601            max_values,
1602        )
1603    }
1604}
1605
1606impl From<ColumnChunkMetaData> for ColumnChunkMetaDataBuilder {
1607    fn from(value: ColumnChunkMetaData) -> Self {
1608        ColumnChunkMetaDataBuilder(value)
1609    }
1610}
1611
1612/// Builder for offset index, part of the Parquet [PageIndex].
1613///
1614/// [PageIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
1615pub struct OffsetIndexBuilder {
1616    offset_array: Vec<i64>,
1617    compressed_page_size_array: Vec<i32>,
1618    first_row_index_array: Vec<i64>,
1619    unencoded_byte_array_data_bytes_array: Option<Vec<i64>>,
1620    current_first_row_index: i64,
1621}
1622
1623impl Default for OffsetIndexBuilder {
1624    fn default() -> Self {
1625        Self::new()
1626    }
1627}
1628
1629impl OffsetIndexBuilder {
1630    /// Creates a new offset index builder.
1631    pub fn new() -> Self {
1632        OffsetIndexBuilder {
1633            offset_array: Vec::new(),
1634            compressed_page_size_array: Vec::new(),
1635            first_row_index_array: Vec::new(),
1636            unencoded_byte_array_data_bytes_array: None,
1637            current_first_row_index: 0,
1638        }
1639    }
1640
1641    /// Append the row count of the next page.
1642    pub fn append_row_count(&mut self, row_count: i64) {
1643        let current_page_row_index = self.current_first_row_index;
1644        self.first_row_index_array.push(current_page_row_index);
1645        self.current_first_row_index += row_count;
1646    }
1647
1648    /// Append the offset and size of the next page.
1649    pub fn append_offset_and_size(&mut self, offset: i64, compressed_page_size: i32) {
1650        self.offset_array.push(offset);
1651        self.compressed_page_size_array.push(compressed_page_size);
1652    }
1653
1654    /// Append the unencoded byte array data bytes of the next page.
1655    pub fn append_unencoded_byte_array_data_bytes(
1656        &mut self,
1657        unencoded_byte_array_data_bytes: Option<i64>,
1658    ) {
1659        if let Some(val) = unencoded_byte_array_data_bytes {
1660            self.unencoded_byte_array_data_bytes_array
1661                .get_or_insert(Vec::new())
1662                .push(val);
1663        }
1664    }
1665
1666    /// Build and get the thrift metadata of offset index
1667    pub fn build(self) -> OffsetIndexMetaData {
1668        let locations = self
1669            .offset_array
1670            .iter()
1671            .zip(self.compressed_page_size_array.iter())
1672            .zip(self.first_row_index_array.iter())
1673            .map(|((offset, size), row_index)| PageLocation {
1674                offset: *offset,
1675                compressed_page_size: *size,
1676                first_row_index: *row_index,
1677            })
1678            .collect::<Vec<_>>();
1679        OffsetIndexMetaData {
1680            page_locations: locations,
1681            unencoded_byte_array_data_bytes: self.unencoded_byte_array_data_bytes_array,
1682        }
1683    }
1684}
1685
1686#[cfg(test)]
1687mod tests {
1688    use super::*;
1689    use crate::basic::{PageType, SortOrder};
1690    use crate::file::metadata::thrift::tests::{
1691        read_column_chunk, read_column_chunk_with_options, read_row_group,
1692    };
1693
1694    #[test]
1695    #[allow(deprecated)]
1696    fn test_level_histogram_update_from_levels_compat() {
1697        let mut histogram = LevelHistogram::try_new(2).unwrap();
1698        histogram.update_from_levels(&[0, 2, 1, 2, 2]);
1699        assert_eq!(histogram.values(), &[1, 1, 3]);
1700    }
1701
1702    #[test]
1703    fn test_row_group_metadata_thrift_conversion() {
1704        let schema_descr = get_test_schema_descr();
1705
1706        let mut columns = vec![];
1707        for ptr in schema_descr.columns() {
1708            let column = ColumnChunkMetaData::builder(ptr.clone()).build().unwrap();
1709            columns.push(column);
1710        }
1711        let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1712            .set_num_rows(1000)
1713            .set_total_byte_size(2000)
1714            .set_column_metadata(columns)
1715            .set_ordinal(1)
1716            .build()
1717            .unwrap();
1718
1719        let mut buf = Vec::new();
1720        let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1721        row_group_meta.write_thrift(&mut writer).unwrap();
1722
1723        let row_group_res = read_row_group(&mut buf, schema_descr).unwrap();
1724
1725        assert_eq!(row_group_res, row_group_meta);
1726    }
1727
1728    #[test]
1729    fn test_row_group_metadata_thrift_conversion_empty() {
1730        let schema_descr = get_test_schema_descr();
1731
1732        let row_group_meta = RowGroupMetaData::builder(schema_descr).build();
1733
1734        assert!(row_group_meta.is_err());
1735        if let Err(e) = row_group_meta {
1736            assert_eq!(
1737                format!("{e}"),
1738                "Parquet error: Column length mismatch: 2 != 0"
1739            );
1740        }
1741    }
1742
1743    /// Test reading a corrupted Parquet file with 3 columns in its schema but only 2 in its row group
1744    #[test]
1745    fn test_row_group_metadata_thrift_corrupted() {
1746        let schema_descr_2cols = Arc::new(SchemaDescriptor::new(Arc::new(
1747            SchemaType::group_type_builder("schema")
1748                .with_fields(vec![
1749                    Arc::new(
1750                        SchemaType::primitive_type_builder("a", Type::INT32)
1751                            .build()
1752                            .unwrap(),
1753                    ),
1754                    Arc::new(
1755                        SchemaType::primitive_type_builder("b", Type::INT32)
1756                            .build()
1757                            .unwrap(),
1758                    ),
1759                ])
1760                .build()
1761                .unwrap(),
1762        )));
1763
1764        let schema_descr_3cols = Arc::new(SchemaDescriptor::new(Arc::new(
1765            SchemaType::group_type_builder("schema")
1766                .with_fields(vec![
1767                    Arc::new(
1768                        SchemaType::primitive_type_builder("a", Type::INT32)
1769                            .build()
1770                            .unwrap(),
1771                    ),
1772                    Arc::new(
1773                        SchemaType::primitive_type_builder("b", Type::INT32)
1774                            .build()
1775                            .unwrap(),
1776                    ),
1777                    Arc::new(
1778                        SchemaType::primitive_type_builder("c", Type::INT32)
1779                            .build()
1780                            .unwrap(),
1781                    ),
1782                ])
1783                .build()
1784                .unwrap(),
1785        )));
1786
1787        let row_group_meta_2cols = RowGroupMetaData::builder(schema_descr_2cols.clone())
1788            .set_num_rows(1000)
1789            .set_total_byte_size(2000)
1790            .set_column_metadata(vec![
1791                ColumnChunkMetaData::builder(schema_descr_2cols.column(0))
1792                    .build()
1793                    .unwrap(),
1794                ColumnChunkMetaData::builder(schema_descr_2cols.column(1))
1795                    .build()
1796                    .unwrap(),
1797            ])
1798            .set_ordinal(1)
1799            .build()
1800            .unwrap();
1801        let mut buf = Vec::new();
1802        let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1803        row_group_meta_2cols.write_thrift(&mut writer).unwrap();
1804
1805        let err = read_row_group(&mut buf, schema_descr_3cols)
1806            .unwrap_err()
1807            .to_string();
1808        assert_eq!(
1809            err,
1810            "Parquet error: Column count mismatch. Schema has 3 columns while Row Group has 2"
1811        );
1812    }
1813
1814    #[test]
1815    fn test_column_chunk_metadata_thrift_conversion() {
1816        let column_descr = get_test_schema_descr().column(0);
1817        let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1818            .set_encodings_mask(EncodingMask::new_from_encodings(
1819                [Encoding::PLAIN, Encoding::RLE].iter(),
1820            ))
1821            .set_file_path("file_path".to_owned())
1822            .set_num_values(1000)
1823            .set_compression(Compression::SNAPPY)
1824            .set_total_compressed_size(2000)
1825            .set_total_uncompressed_size(3000)
1826            .set_data_page_offset(4000)
1827            .set_dictionary_page_offset(Some(5000))
1828            .set_page_encoding_stats(vec![
1829                PageEncodingStats {
1830                    page_type: PageType::DATA_PAGE,
1831                    encoding: Encoding::PLAIN,
1832                    count: 3,
1833                },
1834                PageEncodingStats {
1835                    page_type: PageType::DATA_PAGE,
1836                    encoding: Encoding::RLE,
1837                    count: 5,
1838                },
1839            ])
1840            .set_bloom_filter_offset(Some(6000))
1841            .set_bloom_filter_length(Some(25))
1842            .set_offset_index_offset(Some(7000))
1843            .set_offset_index_length(Some(25))
1844            .set_column_index_offset(Some(8000))
1845            .set_column_index_length(Some(25))
1846            .set_unencoded_byte_array_data_bytes(Some(2000))
1847            .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1848            .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1849            .build()
1850            .unwrap();
1851
1852        let mut buf = Vec::new();
1853        let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1854        col_metadata.write_thrift(&mut writer).unwrap();
1855        let col_chunk_res = read_column_chunk(&mut buf, column_descr.clone()).unwrap();
1856
1857        let expected_metadata = ColumnChunkMetaData::builder(column_descr)
1858            .set_encodings_mask(EncodingMask::new_from_encodings(
1859                [Encoding::PLAIN, Encoding::RLE].iter(),
1860            ))
1861            .set_file_path("file_path".to_owned())
1862            .set_num_values(1000)
1863            .set_compression(Compression::SNAPPY)
1864            .set_total_compressed_size(2000)
1865            .set_total_uncompressed_size(3000)
1866            .set_data_page_offset(4000)
1867            .set_dictionary_page_offset(Some(5000))
1868            .set_page_encoding_stats_mask(EncodingMask::new_from_encodings(
1869                [Encoding::PLAIN, Encoding::RLE].iter(),
1870            ))
1871            .set_bloom_filter_offset(Some(6000))
1872            .set_bloom_filter_length(Some(25))
1873            .set_offset_index_offset(Some(7000))
1874            .set_offset_index_length(Some(25))
1875            .set_column_index_offset(Some(8000))
1876            .set_column_index_length(Some(25))
1877            .set_unencoded_byte_array_data_bytes(Some(2000))
1878            .set_repetition_level_histogram(Some(LevelHistogram::from(vec![100, 100])))
1879            .set_definition_level_histogram(Some(LevelHistogram::from(vec![0, 200])))
1880            .build()
1881            .unwrap();
1882
1883        assert_eq!(col_chunk_res, expected_metadata);
1884    }
1885
1886    #[test]
1887    fn test_column_chunk_metadata_thrift_conversion_full_stats() {
1888        let column_descr = get_test_schema_descr().column(0);
1889        let stats = vec![
1890            PageEncodingStats {
1891                page_type: PageType::DATA_PAGE,
1892                encoding: Encoding::PLAIN,
1893                count: 3,
1894            },
1895            PageEncodingStats {
1896                page_type: PageType::DATA_PAGE,
1897                encoding: Encoding::RLE,
1898                count: 5,
1899            },
1900        ];
1901        let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1902            .set_encodings_mask(EncodingMask::new_from_encodings(
1903                [Encoding::PLAIN, Encoding::RLE].iter(),
1904            ))
1905            .set_num_values(1000)
1906            .set_compression(Compression::SNAPPY)
1907            .set_total_compressed_size(2000)
1908            .set_total_uncompressed_size(3000)
1909            .set_data_page_offset(4000)
1910            .set_page_encoding_stats(stats)
1911            .build()
1912            .unwrap();
1913
1914        let mut buf = Vec::new();
1915        let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1916        col_metadata.write_thrift(&mut writer).unwrap();
1917
1918        let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
1919        let col_chunk_res =
1920            read_column_chunk_with_options(&mut buf, column_descr, Some(&options)).unwrap();
1921
1922        assert_eq!(col_chunk_res, col_metadata);
1923    }
1924
1925    #[test]
1926    fn test_column_chunk_metadata_thrift_conversion_empty() {
1927        let column_descr = get_test_schema_descr().column(0);
1928
1929        let col_metadata = ColumnChunkMetaData::builder(column_descr.clone())
1930            .build()
1931            .unwrap();
1932
1933        let mut buf = Vec::new();
1934        let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1935        col_metadata.write_thrift(&mut writer).unwrap();
1936        let col_chunk_res = read_column_chunk(&mut buf, column_descr).unwrap();
1937
1938        assert_eq!(col_chunk_res, col_metadata);
1939    }
1940
1941    #[test]
1942    fn test_compressed_size() {
1943        let schema_descr = get_test_schema_descr();
1944
1945        let mut columns = vec![];
1946        for column_descr in schema_descr.columns() {
1947            let column = ColumnChunkMetaData::builder(column_descr.clone())
1948                .set_total_compressed_size(500)
1949                .set_total_uncompressed_size(700)
1950                .build()
1951                .unwrap();
1952            columns.push(column);
1953        }
1954        let row_group_meta = RowGroupMetaData::builder(schema_descr)
1955            .set_num_rows(1000)
1956            .set_column_metadata(columns)
1957            .build()
1958            .unwrap();
1959
1960        let compressed_size_res: i64 = row_group_meta.compressed_size();
1961        let compressed_size_exp: i64 = 1000;
1962
1963        assert_eq!(compressed_size_res, compressed_size_exp);
1964    }
1965
1966    #[test]
1967    fn test_memory_size() {
1968        let schema_descr = get_test_schema_descr();
1969
1970        let columns = schema_descr
1971            .columns()
1972            .iter()
1973            .map(|column_descr| {
1974                ColumnChunkMetaData::builder(column_descr.clone())
1975                    .set_statistics(Statistics::new::<i32>(None, None, None, None, false))
1976                    .build()
1977            })
1978            .collect::<Result<Vec<_>>>()
1979            .unwrap();
1980        let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
1981            .set_num_rows(1000)
1982            .set_column_metadata(columns)
1983            .build()
1984            .unwrap();
1985        let row_group_meta = vec![row_group_meta];
1986
1987        let version = 2;
1988        let num_rows = 1000;
1989        let created_by = Some(String::from("test harness"));
1990        let key_value_metadata = Some(vec![KeyValue::new(
1991            String::from("Foo"),
1992            Some(String::from("bar")),
1993        )]);
1994        let column_orders = Some(vec![
1995            ColumnOrder::UNDEFINED,
1996            ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED),
1997        ]);
1998        let file_metadata = FileMetaData::new(
1999            version,
2000            num_rows,
2001            created_by,
2002            key_value_metadata,
2003            schema_descr.clone(),
2004            column_orders,
2005        );
2006
2007        // Now, add in Exact Statistics
2008        let columns_with_stats = schema_descr
2009            .columns()
2010            .iter()
2011            .map(|column_descr| {
2012                ColumnChunkMetaData::builder(column_descr.clone())
2013                    .set_statistics(Statistics::new::<i32>(
2014                        Some(0),
2015                        Some(100),
2016                        None,
2017                        None,
2018                        false,
2019                    ))
2020                    .build()
2021            })
2022            .collect::<Result<Vec<_>>>()
2023            .unwrap();
2024
2025        let row_group_meta_with_stats = RowGroupMetaData::builder(schema_descr)
2026            .set_num_rows(1000)
2027            .set_column_metadata(columns_with_stats)
2028            .build()
2029            .unwrap();
2030        let row_group_meta_with_stats = vec![row_group_meta_with_stats];
2031
2032        let parquet_meta = ParquetMetaDataBuilder::new(file_metadata.clone())
2033            .set_row_groups(row_group_meta_with_stats)
2034            .build();
2035
2036        #[cfg(not(feature = "encryption"))]
2037        let base_expected_size = 2766;
2038        #[cfg(feature = "encryption")]
2039        let base_expected_size = 2934;
2040
2041        assert_eq!(parquet_meta.memory_size(), base_expected_size);
2042
2043        let mut column_index = ColumnIndexBuilder::new(Type::BOOLEAN);
2044        column_index.append(false, vec![1u8], vec![2u8, 3u8], 4);
2045        let column_index = column_index.build().unwrap();
2046        let native_index = match column_index {
2047            ColumnIndexMetaData::BOOLEAN(index) => index,
2048            _ => panic!("wrong type of column index"),
2049        };
2050
2051        // Now, add in OffsetIndex
2052        let mut offset_index = OffsetIndexBuilder::new();
2053        offset_index.append_row_count(1);
2054        offset_index.append_offset_and_size(2, 3);
2055        offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2056        offset_index.append_row_count(1);
2057        offset_index.append_offset_and_size(2, 3);
2058        offset_index.append_unencoded_byte_array_data_bytes(Some(10));
2059        let offset_index = offset_index.build();
2060
2061        let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
2062            .set_row_groups(row_group_meta)
2063            .set_column_index(Some(vec![vec![ColumnIndexMetaData::BOOLEAN(native_index)]]))
2064            .set_offset_index(Some(vec![vec![offset_index]]))
2065            .build();
2066
2067        #[cfg(not(feature = "encryption"))]
2068        let bigger_expected_size = 3192;
2069        #[cfg(feature = "encryption")]
2070        let bigger_expected_size = 3360;
2071
2072        // more set fields means more memory usage
2073        assert!(bigger_expected_size > base_expected_size);
2074        assert_eq!(parquet_meta.memory_size(), bigger_expected_size);
2075    }
2076
2077    #[test]
2078    #[cfg(feature = "encryption")]
2079    fn test_memory_size_with_decryptor() {
2080        use crate::encryption::decrypt::FileDecryptionProperties;
2081        use crate::file::metadata::thrift::encryption::AesGcmV1;
2082
2083        let schema_descr = get_test_schema_descr();
2084
2085        let columns = schema_descr
2086            .columns()
2087            .iter()
2088            .map(|column_descr| ColumnChunkMetaData::builder(column_descr.clone()).build())
2089            .collect::<Result<Vec<_>>>()
2090            .unwrap();
2091        let row_group_meta = RowGroupMetaData::builder(schema_descr.clone())
2092            .set_num_rows(1000)
2093            .set_column_metadata(columns)
2094            .build()
2095            .unwrap();
2096        let row_group_meta = vec![row_group_meta];
2097
2098        let version = 2;
2099        let num_rows = 1000;
2100        let aad_file_unique = vec![1u8; 8];
2101        let aad_prefix = vec![2u8; 8];
2102        let encryption_algorithm = EncryptionAlgorithm::AES_GCM_V1(AesGcmV1 {
2103            aad_prefix: Some(aad_prefix.clone()),
2104            aad_file_unique: Some(aad_file_unique.clone()),
2105            supply_aad_prefix: Some(true),
2106        });
2107        let footer_key_metadata = Some(vec![3u8; 8]);
2108        let file_metadata =
2109            FileMetaData::new(version, num_rows, None, None, schema_descr.clone(), None)
2110                .with_encryption_algorithm(Some(encryption_algorithm))
2111                .with_footer_signing_key_metadata(footer_key_metadata.clone());
2112
2113        let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2114            .set_row_groups(row_group_meta.clone())
2115            .build();
2116
2117        let base_expected_size = 2058;
2118        assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
2119
2120        let footer_key = "0123456789012345".as_bytes();
2121        let column_key = "1234567890123450".as_bytes();
2122        let mut decryption_properties_builder =
2123            FileDecryptionProperties::builder(footer_key.to_vec())
2124                .with_aad_prefix(aad_prefix.clone());
2125        for column in schema_descr.columns() {
2126            decryption_properties_builder = decryption_properties_builder
2127                .with_column_key(&column.path().string(), column_key.to_vec());
2128        }
2129        let decryption_properties = decryption_properties_builder.build().unwrap();
2130        let decryptor = FileDecryptor::new(
2131            &decryption_properties,
2132            footer_key_metadata.as_deref(),
2133            aad_file_unique,
2134            aad_prefix,
2135        )
2136        .unwrap();
2137
2138        let parquet_meta_data = ParquetMetaDataBuilder::new(file_metadata.clone())
2139            .set_row_groups(row_group_meta.clone())
2140            .set_file_decryptor(Some(decryptor))
2141            .build();
2142
2143        let expected_size_with_decryptor = 3072;
2144        assert!(expected_size_with_decryptor > base_expected_size);
2145
2146        assert_eq!(
2147            parquet_meta_data.memory_size(),
2148            expected_size_with_decryptor
2149        );
2150    }
2151
2152    /// Returns sample schema descriptor so we can create column metadata.
2153    fn get_test_schema_descr() -> SchemaDescPtr {
2154        let schema = SchemaType::group_type_builder("schema")
2155            .with_fields(vec![
2156                Arc::new(
2157                    SchemaType::primitive_type_builder("a", Type::INT32)
2158                        .build()
2159                        .unwrap(),
2160                ),
2161                Arc::new(
2162                    SchemaType::primitive_type_builder("b", Type::INT32)
2163                        .build()
2164                        .unwrap(),
2165                ),
2166            ])
2167            .build()
2168            .unwrap();
2169
2170        Arc::new(SchemaDescriptor::new(Arc::new(schema)))
2171    }
2172}