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