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