Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion parquet/src/arrow/arrow_writer/byte_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::encodings::encoding::{DeltaBitPackEncoder, Encoder};
use crate::encodings::rle::RleEncoder;
use crate::errors::{ParquetError, Result};
use crate::file::properties::{EnabledStatistics, WriterProperties, WriterVersion};
use crate::geospatial::accumulator::{try_new_geo_stats_accumulator, GeoStatsAccumulator};
use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator};
use crate::geospatial::statistics::GeospatialStatistics;
use crate::schema::types::ColumnDescPtr;
use crate::util::bit_util::num_required_bits;
Expand Down
13 changes: 6 additions & 7 deletions parquet/src/arrow/arrow_writer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4161,13 +4161,12 @@ mod tests {
true,
)]));
let parquet_schema = Type::group_type_builder("root")
.with_fields(vec![Type::primitive_type_builder(
"integers",
crate::basic::Type::INT64,
)
.build()
.unwrap()
.into()])
.with_fields(vec![
Type::primitive_type_builder("integers", crate::basic::Type::INT64)
.build()
.unwrap()
.into(),
])
.build()
.unwrap();
let parquet_schema_descr = SchemaDescriptor::new(parquet_schema.into());
Expand Down
2 changes: 1 addition & 1 deletion parquet/src/column/writer/encoder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ use crate::data_type::private::ParquetValueType;
use crate::encodings::encoding::{DictEncoder, Encoder, get_encoder};
use crate::errors::{ParquetError, Result};
use crate::file::properties::{EnabledStatistics, WriterProperties};
use crate::geospatial::accumulator::{try_new_geo_stats_accumulator, GeoStatsAccumulator};
use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator};
use crate::geospatial::statistics::GeospatialStatistics;
use crate::schema::types::{ColumnDescPtr, ColumnDescriptor};

Expand Down
15 changes: 15 additions & 0 deletions parquet/src/file/metadata/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,12 @@ impl<T: HeapSize> HeapSize for Arc<T> {
}
}

impl<T: HeapSize> HeapSize for Box<T> {
fn heap_size(&self) -> usize {
self.as_ref().heap_size()
}
}
Comment on lines +59 to +63

@etseidl etseidl Oct 8, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if this is correct. Should this also include the size of T?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Per the description of

/// Return the size of any bytes allocated on the heap by this object,
/// including heap memory in those structures
///
/// Note that the size of the type itself is not included in the result --
/// instead, that size is added by the caller (e.g. container).

So in that case I do think the size of T should be included. The size of self (the pointer/Box) should not be included, but the memory it points to should be

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I'll fix.


impl<T: HeapSize> HeapSize for Option<T> {
fn heap_size(&self) -> usize {
self.as_ref().map(|inner| inner.heap_size()).unwrap_or(0)
Expand All @@ -70,10 +76,17 @@ impl HeapSize for String {

impl HeapSize for FileMetaData {
fn heap_size(&self) -> usize {
#[cfg(feature = "encryption")]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, we still need to implement HeapSize for FileDecryptor.

let encryption_heap_size =
self.encryption_algorithm.heap_size() + self.footer_signing_key_metadata.heap_size();
#[cfg(not(feature = "encryption"))]
let encryption_heap_size = 0;

self.created_by.heap_size()
+ self.key_value_metadata.heap_size()
+ self.schema_descr.heap_size()
+ self.column_orders.heap_size()
+ encryption_heap_size
}
}

Expand Down Expand Up @@ -109,6 +122,7 @@ impl HeapSize for ColumnChunkMetaData {
+ self.unencoded_byte_array_data_bytes.heap_size()
+ self.repetition_level_histogram.heap_size()
+ self.definition_level_histogram.heap_size()
+ self.geo_statistics.heap_size()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💯

+ encryption_heap_size
}
}
Expand Down Expand Up @@ -141,6 +155,7 @@ impl HeapSize for PageType {
0 // no heap allocations
}
}

impl HeapSize for Statistics {
fn heap_size(&self) -> usize {
match self {
Expand Down
58 changes: 50 additions & 8 deletions parquet/src/file/metadata/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,8 @@ use crate::encryption::decrypt::FileDecryptor;
#[cfg(feature = "encryption")]
use crate::file::column_crypto_metadata::ColumnCryptoMetaData;
pub(crate) use crate::file::metadata::memory::HeapSize;
#[cfg(feature = "encryption")]
use crate::file::metadata::thrift_gen::EncryptionAlgorithm;
use crate::file::page_index::column_index::{ByteArrayColumnIndex, PrimitiveColumnIndex};
use crate::file::page_index::{column_index::ColumnIndexMetaData, offset_index::PageLocation};
use crate::file::statistics::Statistics;
Expand Down Expand Up @@ -194,7 +196,7 @@ pub struct ParquetMetaData {
offset_index: Option<ParquetOffsetIndex>,
/// Optional file decryptor
#[cfg(feature = "encryption")]
file_decryptor: Option<FileDecryptor>,
file_decryptor: Option<Box<FileDecryptor>>,
}

impl ParquetMetaData {
Expand All @@ -215,7 +217,7 @@ impl ParquetMetaData {
/// encrypted data.
#[cfg(feature = "encryption")]
pub(crate) fn with_file_decryptor(&mut self, file_decryptor: Option<FileDecryptor>) {
self.file_decryptor = file_decryptor;
self.file_decryptor = file_decryptor.map(Box::new);
}

/// Convert this ParquetMetaData into a [`ParquetMetaDataBuilder`]
Expand All @@ -231,7 +233,7 @@ impl ParquetMetaData {
/// Returns file decryptor as reference.
#[cfg(feature = "encryption")]
pub(crate) fn file_decryptor(&self) -> Option<&FileDecryptor> {
self.file_decryptor.as_ref()
self.file_decryptor.as_deref()
}

/// Returns number of row groups in this file.
Expand Down Expand Up @@ -411,6 +413,13 @@ impl ParquetMetaDataBuilder {
self.0.offset_index.as_ref()
}

/// Sets the file decryptor needed to decrypt this metadata.
#[cfg(feature = "encryption")]
pub(crate) fn set_file_decryptor(mut self, file_decryptor: Option<FileDecryptor>) -> Self {
self.0.with_file_decryptor(file_decryptor);
self
}

/// Creates a new ParquetMetaData from the builder
pub fn build(self) -> ParquetMetaData {
let Self(metadata) = self;
Expand Down Expand Up @@ -468,6 +477,10 @@ pub struct FileMetaData {
key_value_metadata: Option<Vec<KeyValue>>,
schema_descr: SchemaDescPtr,
column_orders: Option<Vec<ColumnOrder>>,
#[cfg(feature = "encryption")]
encryption_algorithm: Option<Box<EncryptionAlgorithm>>,
#[cfg(feature = "encryption")]
footer_signing_key_metadata: Option<Vec<u8>>,
}

impl FileMetaData {
Expand All @@ -487,9 +500,31 @@ impl FileMetaData {
key_value_metadata,
schema_descr,
column_orders,
#[cfg(feature = "encryption")]
encryption_algorithm: None,
#[cfg(feature = "encryption")]
footer_signing_key_metadata: None,
}
}

#[cfg(feature = "encryption")]
pub(crate) fn with_encryption_algorithm(
mut self,
encryption_algorithm: Option<EncryptionAlgorithm>,
) -> Self {
self.encryption_algorithm = encryption_algorithm.map(Box::new);
self
}

#[cfg(feature = "encryption")]
pub(crate) fn with_footer_signing_key_metadata(
mut self,
footer_signing_key_metadata: Option<Vec<u8>>,
) -> Self {
self.footer_signing_key_metadata = footer_signing_key_metadata;
self
}

/// Returns version of this file.
pub fn version(&self) -> i32 {
self.version
Expand Down Expand Up @@ -776,7 +811,7 @@ pub struct ColumnChunkMetaData {
repetition_level_histogram: Option<LevelHistogram>,
definition_level_histogram: Option<LevelHistogram>,
#[cfg(feature = "encryption")]
column_crypto_metadata: Option<ColumnCryptoMetaData>,
column_crypto_metadata: Option<Box<ColumnCryptoMetaData>>,
#[cfg(feature = "encryption")]
encrypted_column_metadata: Option<Vec<u8>>,
}
Expand Down Expand Up @@ -1077,7 +1112,7 @@ impl ColumnChunkMetaData {
/// Returns the encryption metadata for this column chunk.
#[cfg(feature = "encryption")]
pub fn crypto_metadata(&self) -> Option<&ColumnCryptoMetaData> {
self.column_crypto_metadata.as_ref()
self.column_crypto_metadata.as_deref()
}

/// Converts this [`ColumnChunkMetaData`] into a [`ColumnChunkMetaDataBuilder`]
Expand Down Expand Up @@ -1283,7 +1318,14 @@ impl ColumnChunkMetaDataBuilder {
#[cfg(feature = "encryption")]
/// Set the encryption metadata for an encrypted column
pub fn set_column_crypto_metadata(mut self, value: Option<ColumnCryptoMetaData>) -> Self {
self.0.column_crypto_metadata = value;
self.0.column_crypto_metadata = value.map(Box::new);
self
}

#[cfg(feature = "encryption")]
/// Set the encryption metadata for an encrypted column
pub fn set_encrypted_column_metadata(mut self, value: Option<Vec<u8>>) -> Self {
self.0.encrypted_column_metadata = value;
self
}

Expand Down Expand Up @@ -1818,7 +1860,7 @@ mod tests {
#[cfg(not(feature = "encryption"))]
let base_expected_size = 2312;
#[cfg(feature = "encryption")]
let base_expected_size = 2744;
let base_expected_size = 2480;

assert_eq!(parquet_meta.memory_size(), base_expected_size);

Expand Down Expand Up @@ -1849,7 +1891,7 @@ mod tests {
#[cfg(not(feature = "encryption"))]
let bigger_expected_size = 2738;
#[cfg(feature = "encryption")]
let bigger_expected_size = 3170;
let bigger_expected_size = 2906;

// more set fields means more memory usage
assert!(bigger_expected_size > base_expected_size);
Expand Down
19 changes: 11 additions & 8 deletions parquet/src/file/metadata/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,12 @@
//! into the corresponding Rust structures

use crate::errors::ParquetError;
use crate::file::metadata::thrift_gen::parquet_metadata_from_bytes;
use crate::file::metadata::{ColumnChunkMetaData, PageIndexPolicy, ParquetMetaData};

use crate::file::page_index::column_index::ColumnIndexMetaData;
use crate::file::page_index::index_reader::{decode_column_index, decode_offset_index};
use crate::file::page_index::offset_index::OffsetIndexMetaData;
use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
use bytes::Bytes;

/// Helper struct for metadata parsing
Expand Down Expand Up @@ -71,11 +71,15 @@ mod inner {
buf: &[u8],
encrypted_footer: bool,
) -> Result<ParquetMetaData> {
crate::file::metadata::thrift_gen::parquet_metadata_with_encryption(
self.file_decryption_properties.as_deref(),
encrypted_footer,
buf,
)
if encrypted_footer || self.file_decryption_properties.is_some() {
crate::file::metadata::thrift_gen::parquet_metadata_with_encryption(
self.file_decryption_properties.as_deref(),
encrypted_footer,
buf,
)
} else {
decode_metadata(buf)
}
}
}

Expand Down Expand Up @@ -195,8 +199,7 @@ mod inner {
///
/// [Parquet Spec]: https://github.com/apache/parquet-format#metadata
pub(crate) fn decode_metadata(buf: &[u8]) -> crate::errors::Result<ParquetMetaData> {
let mut prot = ThriftSliceInputProtocol::new(buf);
ParquetMetaData::read_thrift(&mut prot)
parquet_metadata_from_bytes(buf)
}

/// Parses column index from the provided bytes and adds it to the metadata.
Expand Down
Loading
Loading