This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new bad29dde10 feat(parquet): Make ParquetMetaData more easily clonable
(#11297)
bad29dde10 is described below
commit bad29dde10217915e808f364d5aee92759e33ea4
Author: Ed Seidl <[email protected]>
AuthorDate: Fri Oct 30 23:46:09 2026 +0000
feat(parquet): Make ParquetMetaData more easily clonable (#11297)
# Which issue does this PR close?
- Closes #11328
- Related to discussion in
https://github.com/apache/datafusion/issues/24288
# Rationale for this change
One source of pain with the `ParquetMetaData` is the fact that cloning
it is an expensive operation. This makes operations like cloning a
cached footer metadata and adding a custom page index provider for the
purposes of a single query inefficient.
# What changes are included in this PR?
Wrap the `file_metadata` and `row_groups` fields of `ParquetMetaData`
with `Arc` so they can be easily shared. Modifies the
`ParquetMetaDataBuilder` in the same way. Now simply replacing the page
index in a shared `ParquetMetaData` can be done cheaply. Modifications
to the `row_groups` will be copy-on-write operations that might incur a
penalty.
Also updates the custom page index provider example to show how a
caching `PageIndexProvider` could be implemented.
# Are these changes tested?
Should be covered by existing tests
# Are there any user-facing changes?
No API changes, but memory for the metadata increases due to the `Arc`
wrappers. Certain operations on `row_groups` may now trigger a `clone`.
---------
Co-authored-by: Andrew Lamb <[email protected]>
---
parquet/examples/custom_page_index.rs | 216 ++++++++++++++++---------
parquet/src/file/metadata/mod.rs | 168 +++++++++++++++----
parquet/src/file/metadata/thrift/encryption.rs | 14 +-
parquet/src/file/metadata/thrift/mod.rs | 22 ++-
parquet/src/file/metadata/writer.rs | 4 +-
parquet/src/file/serialized_reader.rs | 82 ++++++++--
parquet/tests/arrow_reader/page_index.rs | 20 +--
7 files changed, 378 insertions(+), 148 deletions(-)
diff --git a/parquet/examples/custom_page_index.rs
b/parquet/examples/custom_page_index.rs
index dfc2a73104..9f9ee9cead 100644
--- a/parquet/examples/custom_page_index.rs
+++ b/parquet/examples/custom_page_index.rs
@@ -18,8 +18,9 @@
//! Example of implementing a custom PageIndexProvider
//!
//! This example demonstrates how to create a custom page index provider that:
-//! - Only stores page indexes for specified columns (selective storage)
-//! - Uses nested HashMaps for efficient storage and lookup
+//! - Loads page indexes only for specified row groups and columns
+//! - Caches decoded page indexes for reuse across queries
+//! - Builds query-specific providers containing only the requested indexes
//! - Implements all required PageIndexProvider trait methods
//!
//! This approach can significantly reduce memory usage and reduce load time
@@ -36,86 +37,124 @@ use
parquet::file::page_index::column_index::ColumnIndexMetaData;
use parquet::file::page_index::index_reader::{decode_column_index,
decode_offset_index};
use parquet::file::page_index::offset_index::OffsetIndexMetaData;
use std::collections::HashMap;
-use std::collections::hash_map::Entry;
use std::fs::File;
use std::path::PathBuf;
use std::sync::Arc;
use tempfile::TempDir;
-// A custom PageIndexProvider that only stores indexes for a subset of columns
-//
-// This provider contains the parsed footer metadata and the entire contents
of a
-// Parquet file. Indexes are lazily populated as they are requested.
-#[derive(Debug, Clone)]
-pub struct OnDemandPageIndexProvider {
- metadata: ParquetMetaData,
- file_bytes: Bytes,
- // indexes are accessed first by row_group index and then by column index
- column_indexes: Option<HashMap<usize, HashMap<usize,
ColumnIndexMetaData>>>,
- offset_indexes: Option<HashMap<usize, HashMap<usize,
OffsetIndexMetaData>>>,
+/// A custom [`PageIndexProvider`] containing indexes for a subset of row
groups
+/// and columns.
+///
+/// Loading is performed separately by [`SelectivePageIndexLoader`], so this
+/// provider retains neither the Parquet footer metadata nor the file contents.
+#[derive(Debug)]
+struct SparsePageIndexProvider {
+ column_indexes: HashMap<(usize, usize), Arc<ColumnIndexMetaData>>,
+ offset_indexes: HashMap<(usize, usize), Arc<OffsetIndexMetaData>>,
+}
+
+/// A cache of decoded page indexes for a single Parquet file.
+///
+/// A cache shared by multiple files would also need to include a stable file
+/// identity and freshness information in each key.
+#[derive(Default)]
+struct PageIndexCache {
+ column_indexes: HashMap<(usize, usize), Arc<ColumnIndexMetaData>>,
+ offset_indexes: HashMap<(usize, usize), Arc<OffsetIndexMetaData>>,
}
-impl OnDemandPageIndexProvider {
- fn new(metadata: ParquetMetaData, file_bytes: Bytes) -> Self {
+/// Loads selected page indexes and constructs a [`SparsePageIndexProvider`].
+struct SelectivePageIndexLoader<'a> {
+ metadata: &'a ParquetMetaData,
+ file_bytes: &'a Bytes,
+ cache: &'a mut PageIndexCache,
+ provider: SparsePageIndexProvider,
+}
+
+impl<'a> SelectivePageIndexLoader<'a> {
+ fn new(
+ metadata: &'a ParquetMetaData,
+ file_bytes: &'a Bytes,
+ cache: &'a mut PageIndexCache,
+ ) -> Self {
Self {
metadata,
file_bytes,
- column_indexes: None,
- offset_indexes: None,
+ cache,
+ provider: SparsePageIndexProvider {
+ column_indexes: HashMap::new(),
+ offset_indexes: HashMap::new(),
+ },
}
}
- // fetches the bytes for and parses the column index for the given row
group and
- // column. If already fetched, this does nothing.
- fn fetch_column_index(&mut self, row_group_idx: usize, column_idx: usize)
-> Result<()> {
- let map = self.column_indexes.get_or_insert_with(HashMap::new);
- let rg = map.entry(row_group_idx).or_default();
- if let Entry::Vacant(e) = rg.entry(column_idx) {
- // get the `ColumnIndexMetaData` for the requested row group and
column...
- let column =
self.metadata.row_group(row_group_idx).column(column_idx);
- // and get the location of the column index for this chunk
- let range = column.column_index_range();
- if let Some(range) = range {
- let idx_bytes = self
- .file_bytes
- .slice(range.start as usize..range.end as usize);
- let idx = decode_column_index(&idx_bytes,
column.column_type())?;
- e.insert(idx);
- }
+ /// Loads and decodes the column index for the specified column chunk.
+ /// Repeated requests for the same entry do nothing.
+ fn load_column_index(&mut self, row_group_idx: usize, column_idx: usize)
-> Result<()> {
+ let key = (row_group_idx, column_idx);
+ if self.provider.column_indexes.contains_key(&key) {
+ return Ok(());
+ }
+ if let Some(index) = self.cache.column_indexes.get(&key) {
+ println!(" Column index {key:?}: cache hit");
+ self.provider.column_indexes.insert(key, Arc::clone(index));
+ return Ok(());
+ }
+
+ let column = self.metadata.row_group(row_group_idx).column(column_idx);
+ if let Some(range) = column.column_index_range() {
+ println!(" Column index {key:?}: cache miss, decoding");
+ let idx_bytes = self
+ .file_bytes
+ .slice(range.start as usize..range.end as usize);
+ let index = Arc::new(decode_column_index(&idx_bytes,
column.column_type())?);
+ self.cache.column_indexes.insert(key, Arc::clone(&index));
+ self.provider.column_indexes.insert(key, index);
}
Ok(())
}
- // fetches the bytes for and parses the offset index for the given row
group and
- // column. If already fetched, this does nothing.
- fn fetch_offset_index(&mut self, row_group_idx: usize, column_idx: usize)
-> Result<()> {
- let map = self.offset_indexes.get_or_insert_with(HashMap::new);
- let rg = map.entry(row_group_idx).or_default();
- if let Entry::Vacant(e) = rg.entry(column_idx) {
- // get the `ColumnIndexMetaData` for the requested row group and
column...
- let column =
self.metadata.row_group(row_group_idx).column(column_idx);
- // and get the location of the offset index for this chunk
- let range = column.offset_index_range();
- if let Some(range) = range {
- let idx_bytes = self
- .file_bytes
- .slice(range.start as usize..range.end as usize);
- let idx = decode_offset_index(&idx_bytes)?;
- e.insert(idx);
- }
+ /// Loads and decodes the offset index for the specified column chunk.
+ /// Repeated requests for the same entry do nothing.
+ fn load_offset_index(&mut self, row_group_idx: usize, column_idx: usize)
-> Result<()> {
+ let key = (row_group_idx, column_idx);
+ if self.provider.offset_indexes.contains_key(&key) {
+ return Ok(());
+ }
+ if let Some(index) = self.cache.offset_indexes.get(&key) {
+ println!(" Offset index {key:?}: cache hit");
+ self.provider.offset_indexes.insert(key, Arc::clone(index));
+ return Ok(());
+ }
+
+ let column = self.metadata.row_group(row_group_idx).column(column_idx);
+ if let Some(range) = column.offset_index_range() {
+ println!(" Offset index {key:?}: cache miss, decoding");
+ let idx_bytes = self
+ .file_bytes
+ .slice(range.start as usize..range.end as usize);
+ let index = Arc::new(decode_offset_index(&idx_bytes)?);
+ self.cache.offset_indexes.insert(key, Arc::clone(&index));
+ self.provider.offset_indexes.insert(key, index);
}
Ok(())
}
+
+ /// Finishes loading and returns an immutable provider containing only the
+ /// selected indexes.
+ fn finish(self) -> SparsePageIndexProvider {
+ self.provider
+ }
}
// Custom providers must implement the `PageIndexProvider` trait.
-impl PageIndexProvider for OnDemandPageIndexProvider {
+impl PageIndexProvider for SparsePageIndexProvider {
fn has_offset_indexes(&self) -> bool {
- self.offset_indexes.is_some()
+ !self.offset_indexes.is_empty()
}
fn has_column_indexes(&self) -> bool {
- self.column_indexes.is_some()
+ !self.column_indexes.is_empty()
}
fn column_index(
@@ -124,9 +163,8 @@ impl PageIndexProvider for OnDemandPageIndexProvider {
column_idx: usize,
) -> Option<&ColumnIndexMetaData> {
self.column_indexes
- .as_ref()?
- .get(&row_group_idx)?
- .get(&column_idx)
+ .get(&(row_group_idx, column_idx))
+ .map(Arc::as_ref)
}
fn offset_index(
@@ -135,9 +173,8 @@ impl PageIndexProvider for OnDemandPageIndexProvider {
column_idx: usize,
) -> Option<&OffsetIndexMetaData> {
self.offset_indexes
- .as_ref()?
- .get(&row_group_idx)?
- .get(&column_idx)
+ .get(&(row_group_idx, column_idx))
+ .map(Arc::as_ref)
}
fn as_any(&self) -> &dyn std::any::Any {
@@ -192,6 +229,9 @@ fn main() -> Result<()> {
}
};
+ // Store the footer-only metadata in a shared metadata cache.
+ let cached_metadata = Arc::new(metadata);
+
// In this scenario, a query is performed which projects columns 0, 1, and
4
// after applying a predicate to column 0. Chunk level statistics on
column 0
// have already indicated only row group 0 satisfies the predicate. We now
wish
@@ -199,25 +239,53 @@ fn main() -> Result<()> {
// for the predicate column 0, and only fetch the Offset Index (page
locations)
// for the projected columns 0, 1, and 4. Both indexes are only fetched for
// row group 0.
- let mut provider = OnDemandPageIndexProvider::new(metadata.clone(),
file_bytes);
+ let mut page_index_cache = PageIndexCache::default();
+ println!("Query 1 page index cache lookups:");
+ let mut loader =
+ SelectivePageIndexLoader::new(cached_metadata.as_ref(), &file_bytes,
&mut page_index_cache);
let row_group_idx = 0;
- provider.fetch_column_index(row_group_idx, 0)?;
- provider.fetch_offset_index(row_group_idx, 0)?;
- provider.fetch_offset_index(row_group_idx, 1)?;
- provider.fetch_offset_index(row_group_idx, 4)?;
-
- // convert the retrieved metadata into a `ParquetMetaDataBuilder`
- let mut builder = metadata.into_builder();
-
- // and add the provider to the metadata
- builder = builder.set_page_index(Some(Arc::new(provider)));
-
- let metadata = builder.build();
+ loader.load_column_index(row_group_idx, 0)?;
+ loader.load_offset_index(row_group_idx, 0)?;
+ loader.load_offset_index(row_group_idx, 1)?;
+ loader.load_offset_index(row_group_idx, 4)?;
+ let provider = loader.finish();
+
+ // Create query-specific metadata by cheaply cloning the cached footer and
+ // attaching the sparse provider. The provider retains only the selected,
+ // decoded indexes; it does not retain `file_bytes` or another metadata
copy.
+ let metadata = cached_metadata
+ .as_ref()
+ .clone()
+ .into_builder()
+ .set_page_index(Some(Arc::new(provider)))
+ .build();
println!("== Indexes should be partially populated");
print_page_index(&metadata, 0)?;
println!("== Indexes should be unpopulated");
print_page_index(&metadata, 1)?;
+ // A second query uses the same predicate column but projects columns 0, 2,
+ // and 4. It reuses three indexes loaded by the first query and loads only
+ // the newly requested offset index for column 2.
+ println!("=== Query 2: Selective PageIndex (columns 0, 2, 4 only) ===");
+ println!("Query 2 page index cache lookups:");
+ let mut loader =
+ SelectivePageIndexLoader::new(cached_metadata.as_ref(), &file_bytes,
&mut page_index_cache);
+ loader.load_column_index(row_group_idx, 0)?;
+ loader.load_offset_index(row_group_idx, 0)?;
+ loader.load_offset_index(row_group_idx, 2)?;
+ loader.load_offset_index(row_group_idx, 4)?;
+ let provider = loader.finish();
+
+ let metadata = cached_metadata
+ .as_ref()
+ .clone()
+ .into_builder()
+ .set_page_index(Some(Arc::new(provider)))
+ .build();
+ println!("== Indexes should be partially populated");
+ print_page_index(&metadata, 0)?;
+
Ok(())
}
diff --git a/parquet/src/file/metadata/mod.rs b/parquet/src/file/metadata/mod.rs
index f5ee3b20a2..99b6bd6acd 100644
--- a/parquet/src/file/metadata/mod.rs
+++ b/parquet/src/file/metadata/mod.rs
@@ -116,20 +116,23 @@ pub(crate) use writer::ThriftMetadataWriter;
/// This structure is read by the various readers in this crate or can be read
/// directly from a file using the [`ParquetMetaDataReader`] struct.
///
+/// The individual fields of this structure are stored in a way that make
cloning
+/// this structure low-cost.
+///
/// See the [`ParquetMetaDataBuilder`] to create and modify this structure.
///
/// [`parquet.thrift`]:
https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift
#[derive(Debug, Clone)]
pub struct ParquetMetaData {
/// File level metadata
- file_metadata: FileMetaData,
+ file_metadata: Arc<FileMetaData>,
/// Row group metadata
- row_groups: Vec<RowGroupMetaData>,
+ row_groups: Arc<[RowGroupMetaData]>,
/// Page level index for each page in each column chunk
page_index: Option<Arc<dyn PageIndexProvider>>,
/// Optional file decryptor
#[cfg(feature = "encryption")]
- file_decryptor: Option<Box<FileDecryptor>>,
+ file_decryptor: Option<Arc<FileDecryptor>>,
}
impl ParquetMetaData {
@@ -137,22 +140,18 @@ impl ParquetMetaData {
/// group metadata
pub fn new(file_metadata: FileMetaData, row_groups: Vec<RowGroupMetaData>)
-> Self {
ParquetMetaData {
- file_metadata,
- row_groups,
+ file_metadata: Arc::new(file_metadata),
+ row_groups: row_groups.into(),
page_index: None,
#[cfg(feature = "encryption")]
file_decryptor: None,
}
}
- /// Adds [`FileDecryptor`] to this metadata instance to enable decryption
of
- /// encrypted data.
- #[cfg(feature = "encryption")]
- pub(crate) fn with_file_decryptor(&mut self, file_decryptor:
Option<FileDecryptor>) {
- self.file_decryptor = file_decryptor.map(Box::new);
- }
-
- /// Convert this ParquetMetaData into a [`ParquetMetaDataBuilder`]
+ /// Convert this ParquetMetaData into a [`ParquetMetaDataBuilder`].
+ ///
+ /// The builder initially shares any reference-counted metadata. Mutating
shared
+ /// fields may clone their contents using copy-on-write.
pub fn into_builder(self) -> ParquetMetaDataBuilder {
self.into()
}
@@ -296,6 +295,49 @@ impl PartialEq for ParquetMetaData {
}
}
+// helper enum for ParquetMetaDataBuilder
+enum RowGroups {
+ /// Owned vector - cheap to mutate
+ Owned(Vec<RowGroupMetaData>),
+ /// Shared Arc slice - cheap to clone
+ Shared(Arc<[RowGroupMetaData]>),
+}
+
+impl RowGroups {
+ /// Get a reference to the row groups
+ fn as_slice(&self) -> &[RowGroupMetaData] {
+ match self {
+ RowGroups::Owned(v) => v.as_slice(),
+ RowGroups::Shared(arc) => arc.as_ref(),
+ }
+ }
+
+ /// Take ownership of the row groups, converting to Vec
+ fn into_vec(self) -> Vec<RowGroupMetaData> {
+ match self {
+ RowGroups::Owned(v) => v,
+ RowGroups::Shared(arc) => arc.to_vec(),
+ }
+ }
+
+ /// Convert to Owned variant, cloning if necessary
+ fn convert_to_owned(&mut self) {
+ if matches!(self, RowGroups::Shared(_)) {
+ let old = std::mem::replace(self, RowGroups::Owned(Vec::new()));
+ *self = RowGroups::Owned(old.into_vec());
+ }
+ }
+
+ /// Get a mutable reference to the Vec, converting from Shared if necessary
+ fn as_mut_vec(&mut self) -> &mut Vec<RowGroupMetaData> {
+ self.convert_to_owned();
+ match self {
+ RowGroups::Owned(v) => v,
+ RowGroups::Shared(_) => unreachable!(),
+ }
+ }
+}
+
/// A builder for creating / manipulating [`ParquetMetaData`]
///
/// # Example creating a new [`ParquetMetaData`]
@@ -333,28 +375,49 @@ impl PartialEq for ParquetMetaData {
/// .add_row_group(last_row_group)
/// .build();
/// ```
-pub struct ParquetMetaDataBuilder(ParquetMetaData);
+pub struct ParquetMetaDataBuilder {
+ file_metadata: Arc<FileMetaData>,
+ row_groups: RowGroups,
+ page_index: Option<Arc<dyn PageIndexProvider>>,
+ #[cfg(feature = "encryption")]
+ file_decryptor: Option<Arc<FileDecryptor>>,
+}
impl ParquetMetaDataBuilder {
/// Create a new builder from a file metadata, with no row groups
pub fn new(file_meta_data: FileMetaData) -> Self {
- Self(ParquetMetaData::new(file_meta_data, vec![]))
+ Self {
+ file_metadata: Arc::new(file_meta_data),
+ row_groups: RowGroups::Owned(Vec::new()),
+ page_index: None,
+ #[cfg(feature = "encryption")]
+ file_decryptor: None,
+ }
}
/// Create a new builder from an existing ParquetMetaData
pub fn new_from_metadata(metadata: ParquetMetaData) -> Self {
- Self(metadata)
+ Self {
+ file_metadata: metadata.file_metadata,
+ row_groups: RowGroups::Shared(metadata.row_groups),
+ page_index: metadata.page_index,
+ #[cfg(feature = "encryption")]
+ file_decryptor: metadata.file_decryptor,
+ }
}
/// Adds a row group to the metadata
+ ///
+ /// If the row groups are shared with another `ParquetMetaData`, this
clones
+ /// the existing row groups before appending.
pub fn add_row_group(mut self, row_group: RowGroupMetaData) -> Self {
- self.0.row_groups.push(row_group);
+ self.row_groups.as_mut_vec().push(row_group);
self
}
/// Sets all the row groups to the specified list
pub fn set_row_groups(mut self, row_groups: Vec<RowGroupMetaData>) -> Self
{
- self.0.row_groups = row_groups;
+ self.row_groups = RowGroups::Owned(row_groups);
self
}
@@ -363,13 +426,16 @@ impl ParquetMetaDataBuilder {
///
/// This can be used for more efficient creation of a new ParquetMetaData
/// from an existing one.
+ ///
+ /// Returns the existing allocation when uniquely owned; otherwise clones
the
+ /// shared row groups.
pub fn take_row_groups(&mut self) -> Vec<RowGroupMetaData> {
- std::mem::take(&mut self.0.row_groups)
+ std::mem::replace(&mut self.row_groups,
RowGroups::Owned(Vec::new())).into_vec()
}
/// Return a reference to the current row groups
pub fn row_groups(&self) -> &[RowGroupMetaData] {
- &self.0.row_groups
+ self.row_groups.as_slice()
}
/// Sets the [`PageIndexProvider`]
@@ -381,37 +447,45 @@ impl ParquetMetaDataBuilder {
///
/// [`custom_page_index.rs`]:
https://github.com/apache/arrow-rs/blob/main/parquet/examples/custom_page_index.rs
pub fn set_page_index(mut self, page_index: Option<Arc<dyn
PageIndexProvider>>) -> Self {
- self.0.page_index = page_index;
+ self.page_index = page_index;
self
}
/// Returns the current [`PageIndexProvider`] from the builder, replacing
it with `None`
pub fn take_page_index(&mut self) -> Option<Arc<dyn PageIndexProvider>> {
- std::mem::take(&mut self.0.page_index)
+ std::mem::take(&mut self.page_index)
}
/// Return a reference to the current [`PageIndexProvider`], if any
pub fn page_index(&self) -> Option<&dyn PageIndexProvider> {
- self.0.page_index.as_ref().map(|arc| arc.as_ref())
+ self.page_index.as_ref().map(|arc| arc.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.file_decryptor = file_decryptor.map(Arc::new);
self
}
/// Creates a new ParquetMetaData from the builder
pub fn build(self) -> ParquetMetaData {
- let Self(metadata) = self;
- metadata
+ ParquetMetaData {
+ file_metadata: self.file_metadata,
+ row_groups: match self.row_groups {
+ RowGroups::Owned(v) => v.into(),
+ RowGroups::Shared(arc) => arc,
+ },
+ page_index: self.page_index,
+ #[cfg(feature = "encryption")]
+ file_decryptor: self.file_decryptor,
+ }
}
}
impl From<ParquetMetaData> for ParquetMetaDataBuilder {
fn from(meta_data: ParquetMetaData) -> Self {
- Self(meta_data)
+ Self::new_from_metadata(meta_data)
}
}
@@ -694,6 +768,16 @@ impl RowGroupMetaData {
}
}
+impl HeapSize for Arc<[RowGroupMetaData]> {
+ fn heap_size(&self) -> usize {
+ // account for the contents of the slice
+ let rg_size = self.iter().map(|t| t.heap_size()).sum::<usize>();
+ // The heap block contains the strong and weak counts followed by the
slice,
+ // The fat pointer itself is stored inline.
+ 2 * std::mem::size_of::<usize>() +
std::mem::size_of_val(self.as_ref()) + rg_size
+ }
+}
+
/// Builder for row group metadata.
pub struct RowGroupMetaDataBuilder(RowGroupMetaData);
@@ -2085,9 +2169,9 @@ mod tests {
// Base size without page index
#[cfg(not(feature = "encryption"))]
- let base_expected_size = 2766;
+ let base_expected_size = 2798;
#[cfg(feature = "encryption")]
- let base_expected_size = 2934;
+ let base_expected_size = 2966;
assert_eq!(parquet_meta.memory_size(), base_expected_size);
@@ -2122,9 +2206,9 @@ mod tests {
// Size with page index (includes Arc overhead plus PageIndex heap
size)
#[cfg(not(feature = "encryption"))]
- let bigger_expected_size = 3280;
+ let bigger_expected_size = 3312;
#[cfg(feature = "encryption")]
- let bigger_expected_size = 3448;
+ let bigger_expected_size = 3480;
// more set fields means more memory usage
assert!(bigger_expected_size > base_expected_size);
@@ -2171,7 +2255,7 @@ mod tests {
.set_row_groups(row_group_meta.clone())
.build();
- let base_expected_size = 2042;
+ let base_expected_size = 2074;
assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
let footer_key = b"0123456789012345";
@@ -2197,7 +2281,7 @@ mod tests {
.set_file_decryptor(Some(decryptor))
.build();
- let expected_size_with_decryptor = 3056;
+ let expected_size_with_decryptor = 3104;
assert!(expected_size_with_decryptor > base_expected_size);
assert_eq!(
@@ -2258,6 +2342,24 @@ mod tests {
assert!(!page_index.has_offset_indexes());
}
+ #[test]
+ fn test_replace_page_index_preserves_shared_metadata() {
+ let schema_descr = get_test_schema_descr();
+ let file_metadata = FileMetaData::new(1, 0, None, None, schema_descr,
None);
+ let cached = Arc::new(ParquetMetaData::new(file_metadata, vec![]));
+
+ let page_index = Arc::new(PageIndexBuilder::new(0, 0).build());
+ let rebuilt = cached
+ .as_ref()
+ .clone()
+ .into_builder()
+ .set_page_index(Some(page_index))
+ .build();
+
+ assert!(Arc::ptr_eq(&cached.file_metadata, &rebuilt.file_metadata));
+ assert!(Arc::ptr_eq(&cached.row_groups, &rebuilt.row_groups));
+ }
+
#[test]
fn test_rebuild_page_index() {
// Create builder with one rowgroup and two columns
diff --git a/parquet/src/file/metadata/thrift/encryption.rs
b/parquet/src/file/metadata/thrift/encryption.rs
index 5540820726..6c233f7f32 100644
--- a/parquet/src/file/metadata/thrift/encryption.rs
+++ b/parquet/src/file/metadata/thrift/encryption.rs
@@ -24,7 +24,10 @@ use crate::{
column_crypto_metadata::ColumnCryptoMetaData,
metadata::{
HeapSize, ParquetMetaData, ParquetMetaDataOptions,
RowGroupMetaData,
- thrift::{parquet_metadata_from_bytes, read_column_metadata,
validate_column_metadata},
+ thrift::{
+ file_metadata_and_row_groups_from_bytes, read_column_metadata,
+ validate_column_metadata,
+ },
},
},
parquet_thrift::{
@@ -292,16 +295,9 @@ pub(crate) fn parquet_metadata_with_encryption(
}
}
- let parquet_meta = parquet_metadata_from_bytes(buf, options)
+ let (mut file_metadata, row_groups) =
file_metadata_and_row_groups_from_bytes(buf, options)
.map_err(|e| general_err!("Could not parse metadata: {}", e))?;
- let ParquetMetaData {
- mut file_metadata,
- row_groups,
- page_index: _,
- file_decryptor: _,
- } = parquet_meta;
-
// Take the encryption algorithm and footer signing key metadata as they
are no longer
// needed after this.
if let (Some(algo), Some(file_decryption_properties)) = (
diff --git a/parquet/src/file/metadata/thrift/mod.rs
b/parquet/src/file/metadata/thrift/mod.rs
index dffaacf440..c4508ba6f2 100644
--- a/parquet/src/file/metadata/thrift/mod.rs
+++ b/parquet/src/file/metadata/thrift/mod.rs
@@ -42,9 +42,9 @@ use crate::{
errors::{ParquetError, Result},
file::{
metadata::{
- ColumnChunkMetaData, ColumnChunkMetaDataBuilder, KeyValue,
LevelHistogram,
- PageEncodingStats, ParquetMetaData, ParquetMetaDataOptions,
ParquetPageEncodingStats,
- RowGroupMetaData, RowGroupMetaDataBuilder, SortingColumn,
+ ColumnChunkMetaData, ColumnChunkMetaDataBuilder, FileMetaData,
KeyValue,
+ LevelHistogram, PageEncodingStats, ParquetMetaData,
ParquetMetaDataOptions,
+ ParquetPageEncodingStats, RowGroupMetaData,
RowGroupMetaDataBuilder, SortingColumn,
},
statistics::ValueStatistics,
},
@@ -757,6 +757,14 @@ pub(crate) fn parquet_metadata_from_bytes(
buf: &[u8],
options: Option<&ParquetMetaDataOptions>,
) -> Result<ParquetMetaData> {
+ let (file_metadata, row_groups) =
file_metadata_and_row_groups_from_bytes(buf, options)?;
+ Ok(ParquetMetaData::new(file_metadata, row_groups))
+}
+
+pub(crate) fn file_metadata_and_row_groups_from_bytes(
+ buf: &[u8],
+ options: Option<&ParquetMetaDataOptions>,
+) -> Result<(FileMetaData, Vec<RowGroupMetaData>)> {
let mut prot = ThriftSliceInputProtocol::new(buf);
// begin reading the file metadata
@@ -899,7 +907,7 @@ pub(crate) fn parquet_metadata_from_bytes(
});
#[cfg(not(feature = "encryption"))]
- let fmd = crate::file::metadata::FileMetaData::new(
+ let fmd = FileMetaData::new(
version,
num_rows,
created_by,
@@ -908,7 +916,7 @@ pub(crate) fn parquet_metadata_from_bytes(
column_orders,
);
#[cfg(feature = "encryption")]
- let fmd = crate::file::metadata::FileMetaData::new(
+ let fmd = FileMetaData::new(
version,
num_rows,
created_by,
@@ -919,7 +927,7 @@ pub(crate) fn parquet_metadata_from_bytes(
.with_encryption_algorithm(encryption_algorithm)
.with_footer_signing_key_metadata(footer_signing_key_metadata.map(|v|
v.to_vec()));
- Ok(ParquetMetaData::new(fmd, row_groups))
+ Ok((fmd, row_groups))
}
/// Ensure [`RowGroupMetaData::ordinal`] is usable after decode without
@@ -1402,7 +1410,7 @@ pub(super) fn serialize_column_meta_data<W: Write>(
// temp struct used for writing
pub(super) struct FileMeta<'a> {
- pub(super) file_metadata: &'a crate::file::metadata::FileMetaData,
+ pub(super) file_metadata: &'a FileMetaData,
pub(super) row_groups: &'a Vec<RowGroupMetaData>,
// If true, then write the `path_in_schema` field in the ColumnMetaData
struct.
pub(super) write_path_in_schema: bool,
diff --git a/parquet/src/file/metadata/writer.rs
b/parquet/src/file/metadata/writer.rs
index fcbc63a9c5..ce2e425f4e 100644
--- a/parquet/src/file/metadata/writer.rs
+++ b/parquet/src/file/metadata/writer.rs
@@ -439,7 +439,9 @@ impl<'a, W: Write> ParquetMetaDataWriter<'a, W> {
let schema_descr = Arc::new(SchemaDescriptor::new(schema.clone()));
let created_by = file_metadata.created_by().map(str::to_string);
- let row_groups = self.metadata.row_groups.clone();
+ // need to clone because ThriftMetadataWriter wants to update page
+ // index pointers. we throw this clone away upon calling
encoder.finish()
+ let row_groups = self.metadata.row_groups.as_ref().to_vec();
let key_value_metadata = file_metadata.key_value_metadata().cloned();
diff --git a/parquet/src/file/serialized_reader.rs
b/parquet/src/file/serialized_reader.rs
index 1f6924c4b0..f8021d7b77 100644
--- a/parquet/src/file/serialized_reader.rs
+++ b/parquet/src/file/serialized_reader.rs
@@ -246,28 +246,30 @@ impl<R: 'static + ChunkReader> SerializedFileReader<R> {
/// Creates file reader from a Parquet file with read options.
/// Returns an error if the Parquet file does not exist or is corrupt.
pub fn new_with_options(chunk_reader: R, options: ReadOptions) ->
Result<Self> {
- let mut metadata_builder = ParquetMetaDataReader::new()
+ let mut metadata = ParquetMetaDataReader::new()
.with_metadata_options(Some(options.metadata_options.clone()))
- .parse_and_finish(&chunk_reader)?
- .into_builder();
+ .parse_and_finish(&chunk_reader)?;
let mut predicates = options.predicates;
- // Filter row groups based on the predicates
- for (i, rg_meta) in
metadata_builder.take_row_groups().into_iter().enumerate() {
- let mut keep = true;
- for predicate in &mut predicates {
- if !predicate(&rg_meta, i) {
- keep = false;
- break;
+ if !predicates.is_empty() {
+ let row_groups = metadata.row_groups();
+ let mut filtered = None;
+
+ // Only clone retained row groups if a predicate rejects a row
group.
+ for (i, rg_meta) in row_groups.iter().enumerate() {
+ let keep = predicates.iter_mut().all(|predicate|
predicate(rg_meta, i));
+ if !keep {
+ filtered.get_or_insert_with(|| row_groups[..i].to_vec());
+ } else if let Some(kept) = &mut filtered {
+ kept.push(rg_meta.clone());
}
}
- if keep {
- metadata_builder = metadata_builder.add_row_group(rg_meta);
+
+ if let Some(row_groups) = filtered {
+ metadata =
metadata.into_builder().set_row_groups(row_groups).build();
}
}
- let mut metadata = metadata_builder.build();
-
// If page indexes are desired, build them with the filtered set of
row groups
if options.enable_page_index {
let mut reader = ParquetMetaDataReader::new_with_metadata(metadata)
@@ -1298,7 +1300,9 @@ impl<R: ChunkReader> PageReader for
SerializedPageReader<R> {
#[cfg(test)]
mod tests {
+ use std::cell::RefCell;
use std::collections::HashSet;
+ use std::rc::Rc;
use bytes::Buf;
@@ -2891,6 +2895,56 @@ mod tests {
assert_eq!(file_metadata.file_metadata().num_rows(), 25);
assert_eq!(file_metadata.num_row_groups(), 5);
+ // Stateful predicates must run once per group in order, with
short-circuiting
+ // and original indexes, whether or not any groups are rejected.
+ for filter in [false, true] {
+ let calls = Rc::new(RefCell::new(Vec::new()));
+ let first_calls = Rc::clone(&calls);
+ let second_calls = Rc::clone(&calls);
+ let mut next_index = 0;
+ let read_options = ReadOptionsBuilder::new()
+ .with_page_index()
+ .with_predicate(Box::new(move |rg_meta, i| {
+ assert_eq!(i, next_index);
+ assert_eq!(rg_meta.ordinal(), Some(i as i32));
+ next_index += 1;
+ first_calls.borrow_mut().push((i, 0));
+ !filter || i != 1
+ }))
+ .with_predicate(Box::new(move |_, i| {
+ second_calls.borrow_mut().push((i, 1));
+ !filter || i != 3
+ }))
+ .build();
+ let reader =
+
SerializedFileReader::new_with_options(file.try_clone().unwrap(), read_options)
+ .unwrap();
+ let metadata = reader.metadata();
+ let expected_ordinals = if filter {
+ vec![0, 2, 4]
+ } else {
+ vec![0, 1, 2, 3, 4]
+ };
+ let ordinals: Vec<_> = metadata
+ .row_groups()
+ .iter()
+ .map(|rg| rg.ordinal().unwrap())
+ .collect();
+ assert_eq!(ordinals, expected_ordinals);
+ assert!(metadata.page_index().is_some_and(|idx|
idx.is_complete()));
+
+ let expected_calls: Vec<_> = (0..5)
+ .flat_map(|i| {
+ if filter && i == 1 {
+ vec![(i, 0)]
+ } else {
+ vec![(i, 0), (i, 1)]
+ }
+ })
+ .collect();
+ assert_eq!(*calls.borrow(), expected_calls);
+ }
+
// read only the 3rd row group
let read_options = ReadOptionsBuilder::new()
.with_page_index()
diff --git a/parquet/tests/arrow_reader/page_index.rs
b/parquet/tests/arrow_reader/page_index.rs
index e1a8300286..de3bae70c1 100644
--- a/parquet/tests/arrow_reader/page_index.rs
+++ b/parquet/tests/arrow_reader/page_index.rs
@@ -370,9 +370,9 @@ fn test_page_index_sizes() {
let metadata = reader.finish().unwrap();
assert!(metadata.page_index().is_none());
#[cfg(not(feature = "encryption"))]
- assert_eq!(metadata.memory_size(), 7393);
+ assert_eq!(metadata.memory_size(), 7425);
#[cfg(feature = "encryption")]
- assert_eq!(metadata.memory_size(), 7817);
+ assert_eq!(metadata.memory_size(), 7849);
// full index
let mut reader =
ParquetMetaDataReader::new().with_page_index_policy(PageIndexPolicy::Required);
@@ -382,9 +382,9 @@ fn test_page_index_sizes() {
let metadata = reader.finish().unwrap();
assert!(metadata.page_index().is_some());
#[cfg(not(feature = "encryption"))]
- assert_eq!(metadata.memory_size(), 13897);
+ assert_eq!(metadata.memory_size(), 13929);
#[cfg(feature = "encryption")]
- assert_eq!(metadata.memory_size(), 14321);
+ assert_eq!(metadata.memory_size(), 14353);
// populate column 0 for column index and all columns for offset index
let mut reader = ParquetMetaDataReader::new()
@@ -397,9 +397,9 @@ fn test_page_index_sizes() {
let metadata = reader.finish().unwrap();
assert!(metadata.page_index().is_some());
#[cfg(not(feature = "encryption"))]
- assert_eq!(metadata.memory_size(), 12776);
+ assert_eq!(metadata.memory_size(), 12808);
#[cfg(feature = "encryption")]
- assert_eq!(metadata.memory_size(), 13200);
+ assert_eq!(metadata.memory_size(), 13232);
// populate column 0 for column index and columns 0 & 2 for the offset
index
let mut reader = ParquetMetaDataReader::new()
@@ -412,9 +412,9 @@ fn test_page_index_sizes() {
let metadata = reader.finish().unwrap();
assert!(metadata.page_index().is_some());
#[cfg(not(feature = "encryption"))]
- assert_eq!(metadata.memory_size(), 12056);
+ assert_eq!(metadata.memory_size(), 12088);
#[cfg(feature = "encryption")]
- assert_eq!(metadata.memory_size(), 12480);
+ assert_eq!(metadata.memory_size(), 12512);
// populate only row group 1, column index gets column 0, offset index gets
// columns 0 and 2.
@@ -428,7 +428,7 @@ fn test_page_index_sizes() {
let metadata = reader.finish().unwrap();
assert!(metadata.page_index().is_some());
#[cfg(not(feature = "encryption"))]
- assert_eq!(metadata.memory_size(), 11326);
+ assert_eq!(metadata.memory_size(), 11358);
#[cfg(feature = "encryption")]
- assert_eq!(metadata.memory_size(), 11750);
+ assert_eq!(metadata.memory_size(), 11782);
}
