This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11297-bf0cde345b14e5abdf3e8abb57f55d43da48b49f in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
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); }
