diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index af97a192fa7ce..3b08e7dc49b9c 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1147,16 +1147,28 @@ impl FiltersPreparedParquetOpen { impl RowGroupsPrunedParquetOpen { /// Load the page index if pruning requires it and metadata did not include it. async fn load_page_index(mut self) -> Result { - self.prepared.loaded.reader_metadata = load_page_index( + let options = self + .prepared + .loaded + .options + .clone() + .with_page_index_policy(PageIndexPolicy::Optional); + let reader_metadata = load_page_index( self.prepared.loaded.reader_metadata.clone(), &mut self.prepared.loaded.prepared.async_file_reader, - self.prepared - .loaded - .options - .clone() - .with_page_index_policy(PageIndexPolicy::Optional), + options.clone(), ) .await?; + self.prepared + .loaded + .prepared + .parquet_file_reader_factory + .cache_metadata( + &self.prepared.loaded.prepared.partitioned_file, + Arc::clone(reader_metadata.metadata()), + &options, + ); + self.prepared.loaded.reader_metadata = reader_metadata; Ok(self) } @@ -3162,6 +3174,92 @@ mod test { ); } + #[tokio::test] + async fn test_deferred_page_index_is_cached() { + use parquet::file::properties::WriterProperties; + + let store = Arc::new(InMemory::new()) as Arc; + let metadata_cache: Arc = + Arc::new(DefaultCache::::new( + 64 * 1024 * 1024, + )); + let values: Vec = (1..=100).collect(); + let batch = record_batch!(( + "a", + Int32, + values.iter().map(|v| Some(*v)).collect::>() + )) + .unwrap(); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(10) + .set_write_batch_size(10) + .build(); + let schema = batch.schema(); + let data_len = write_parquet_batches( + Arc::clone(&store), + "test.parquet", + vec![batch], + Some(props), + ) + .await; + let file = PartitionedFile::new( + "test.parquet".to_string(), + u64::try_from(data_len).unwrap(), + ); + let predicate = logical2physical(&col("a").gt(lit(90i32)), &schema); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate)) + .with_enable_page_index(true) + .with_pushdown_filters(false) + .with_row_group_stats_pruning(false) + .with_parquet_file_reader_factory(Arc::new( + CachedParquetFileReaderFactory::new( + Arc::clone(&store), + Arc::clone(&metadata_cache), + ), + )) + .build(); + + let (_, rows) = + count_batches_and_rows(open_file(&morselizer, file.clone()).await.unwrap()) + .await; + assert_eq!(rows, 10); + + let cached = metadata_cache + .get(&Path::from("test.parquet")) + .expect("metadata cache should contain the file"); + let page_index_cached = cached.file_metadata.extra_info(); + assert_eq!( + page_index_cached.get("page_index").map(String::as_str), + Some("true"), + "deferred page index should be written back to the metadata cache" + ); + + let second_metrics = ExecutionPlanMetricsSet::new(); + let second_morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(schema) + .with_predicate(predicate) + .with_enable_page_index(true) + .with_pushdown_filters(false) + .with_row_group_stats_pruning(false) + .with_metrics(second_metrics.clone()) + .with_parquet_file_reader_factory(Arc::new( + CachedParquetFileReaderFactory::new(store, metadata_cache), + )) + .build(); + + let _stream = open_file(&second_morselizer, file).await.unwrap(); + assert_eq!( + counter_metric_value(&second_metrics, "bytes_scanned"), + 0, + "reopening the file should reuse the cached page index" + ); + } + #[tokio::test] async fn test_page_index_loaded_when_not_fully_matched() { use parquet::file::properties::WriterProperties; diff --git a/datafusion/datasource-parquet/src/reader.rs b/datafusion/datasource-parquet/src/reader.rs index 4df636b894940..64681e7016424 100644 --- a/datafusion/datasource-parquet/src/reader.rs +++ b/datafusion/datasource-parquet/src/reader.rs @@ -23,6 +23,7 @@ use crate::metadata::DFParquetMetadata; use bytes::Bytes; use datafusion_common::HashMap; use datafusion_datasource::PartitionedFile; +use datafusion_execution::cache::cache_manager::CachedFileMetadataEntry; use datafusion_execution::cache::cache_manager::FileMetadata; use datafusion_execution::cache::cache_manager::FileMetadataCache; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; @@ -66,6 +67,20 @@ pub trait ParquetFileReaderFactory: Debug + Send + Sync + 'static { metadata_size_hint: Option, metrics: &ExecutionPlanMetricsSet, ) -> datafusion_common::Result>; + + /// Stores metadata loaded after the initial metadata request, if this + /// factory has a metadata cache. + /// + /// The default implementation does nothing because most factories do not + /// cache metadata. Implementations that do cache metadata can override + /// this hook so deferred page-index loads are available to later readers. + fn cache_metadata( + &self, + _partitioned_file: &PartitionedFile, + _metadata: Arc, + _options: &ArrowReaderOptions, + ) { + } } /// Default implementation of [`ParquetFileReaderFactory`] @@ -232,6 +247,28 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { metadata_size_hint, ))) } + + fn cache_metadata( + &self, + partitioned_file: &PartitionedFile, + metadata: Arc, + options: &ArrowReaderOptions, + ) { + #[cfg(feature = "parquet_encryption")] + if options.file_decryption_properties().is_some() { + return; + } + #[cfg(not(feature = "parquet_encryption"))] + let _ = options; + + self.metadata_cache.put( + &partitioned_file.object_meta.location, + CachedFileMetadataEntry::new( + partitioned_file.object_meta.clone(), + Arc::new(crate::metadata::CachedParquetMetaData::new(metadata)), + ), + ); + } } /// Implements [`AsyncFileReader`] for a Parquet file in object storage. Reads the file metadata