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