Skip to content

feat: Use Cached Metadata for ListingTable Statistics #17022

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alamb merged 24 commits into from
Aug 11, 2025
Merged

feat: Use Cached Metadata for ListingTable Statistics #17022

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
c51e16e
make CachedParquetMetaData pub
shehabgamin Aug 3, 2025
74359e0
make parquet reader pub
shehabgamin Aug 3, 2025
c67ba4a
undo
shehabgamin Aug 3, 2025
84f91b0
add impl
shehabgamin Aug 3, 2025
e3efe06
cached meta in reader
shehabgamin Aug 3, 2025
961fc24
Merge branch 'main' of github.com:lakehq/datafusion into parquet-meta…
shehabgamin Aug 3, 2025
f95ea46
update function docs
shehabgamin Aug 3, 2025
0d6b5ff
cache after fetch
shehabgamin Aug 3, 2025
c12bc46
cache meta
shehabgamin Aug 3, 2025
c20b142
return arc
shehabgamin Aug 3, 2025
65fa8f7
Merge branch 'main' into parquet-metadata
shehabgamin Aug 4, 2025
2479a9f
Merge branch 'apache:main' into parquet-metadata
shehabgamin Aug 6, 2025
230256c
add cache metadata tests
shehabgamin Aug 6, 2025
17366f5
add parquet cache tests
shehabgamin Aug 6, 2025
489763c
Merge branch 'main' of github.com:lakehq/datafusion into parquet-meta…
shehabgamin Aug 8, 2025
15914eb
remove cache_metadata option
shehabgamin Aug 8, 2025
a6d0bc0
parquet metadata cache impacts test
shehabgamin Aug 8, 2025
fd761e7
Merge branch 'main' of github.com:lakehq/datafusion into parquet-meta…
shehabgamin Aug 9, 2025
f3ad6f1
add comment for reader
shehabgamin Aug 9, 2025
5b4129a
refactor fetch parquet metadata
shehabgamin Aug 9, 2025
c39b29f
parquet metadata cleanup
shehabgamin Aug 9, 2025
2313aed
parquet metadata cleanup
shehabgamin Aug 9, 2025
3bb321b
refactor parquet metadata cache
shehabgamin Aug 9, 2025
7d9128b
feature flag
shehabgamin Aug 9, 2025
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
90 changes: 73 additions & 17 deletions datafusion/core/src/datasource/file_format/parquet.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,18 +195,33 @@ mod tests {
let format = ParquetFormat::default().with_force_view_types(force_views);
let schema = format.infer_schema(&ctx, &store, &meta).await.unwrap();

let stats =
fetch_statistics(store.as_ref(), schema.clone(), &meta[0], None, None)
.await?;
let stats = fetch_statistics(
store.as_ref(),
schema.clone(),
&meta[0],
None,
None,
format.options().global.cache_metadata,
ctx.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;

assert_eq!(stats.num_rows, Precision::Exact(3));
let c1_stats = &stats.column_statistics[0];
let c2_stats = &stats.column_statistics[1];
assert_eq!(c1_stats.null_count, Precision::Exact(1));
assert_eq!(c2_stats.null_count, Precision::Exact(3));

let stats =
fetch_statistics(store.as_ref(), schema, &meta[1], None, None).await?;
let stats = fetch_statistics(
store.as_ref(),
schema,
&meta[1],
None,
None,
format.options().global.cache_metadata,
ctx.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;
assert_eq!(stats.num_rows, Precision::Exact(3));
let c1_stats = &stats.column_statistics[0];
let c2_stats = &stats.column_statistics[1];
Expand Down Expand Up @@ -383,6 +398,8 @@ mod tests {
&meta[0],
Some(9),
None,
false,
None,
)
.await
.expect("error reading metadata with hint");
Expand All @@ -409,6 +426,8 @@ mod tests {
&meta[0],
Some(9),
None,
format.options().global.cache_metadata,
ctx.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;

Expand All @@ -425,9 +444,16 @@ mod tests {
// Use the file size as the hint so we can get the full metadata from the first fetch
let size_hint = meta[0].size as usize;

fetch_parquet_metadata(store.upcast().as_ref(), &meta[0], Some(size_hint), None)
.await
.expect("error reading metadata with hint");
fetch_parquet_metadata(
store.upcast().as_ref(),
&meta[0],
Some(size_hint),
None,
format.options().global.cache_metadata,
ctx.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await
.expect("error reading metadata with hint");

// ensure the requests were coalesced into a single request
assert_eq!(store.request_count(), 1);
Expand All @@ -445,6 +471,8 @@ mod tests {
&meta[0],
Some(size_hint),
None,
format.options().global.cache_metadata,
ctx.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;

Expand All @@ -461,9 +489,16 @@ mod tests {
// Use the a size hint larger than the file size to make sure we don't panic
let size_hint = (meta[0].size + 100) as usize;

fetch_parquet_metadata(store.upcast().as_ref(), &meta[0], Some(size_hint), None)
.await
.expect("error reading metadata with hint");
fetch_parquet_metadata(
store.upcast().as_ref(),
&meta[0],
Some(size_hint),
None,
format.options().global.cache_metadata,
ctx.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await
.expect("error reading metadata with hint");

assert_eq!(store.request_count(), 1);

Expand Down Expand Up @@ -500,8 +535,15 @@ mod tests {
let schema = format.infer_schema(&state, &store, &files).await.unwrap();

// Fetch statistics for first file
let pq_meta =
fetch_parquet_metadata(store.as_ref(), &files[0], None, None).await?;
let pq_meta = fetch_parquet_metadata(
store.as_ref(),
&files[0],
None,
None,
format.options().global.cache_metadata,
state.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;
let stats = statistics_from_parquet_meta_calc(&pq_meta, schema.clone())?;
assert_eq!(stats.num_rows, Precision::Exact(4));

Expand Down Expand Up @@ -559,8 +601,15 @@ mod tests {
};

// Fetch statistics for first file
let pq_meta =
fetch_parquet_metadata(store.as_ref(), &files[0], None, None).await?;
let pq_meta = fetch_parquet_metadata(
store.as_ref(),
&files[0],
None,
None,
format.options().global.cache_metadata,
state.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;
let stats = statistics_from_parquet_meta_calc(&pq_meta, schema.clone())?;
assert_eq!(stats.num_rows, Precision::Exact(3));
// column c1
Expand All @@ -586,8 +635,15 @@ mod tests {
assert_eq!(c2_stats.min_value, Precision::Exact(null_i64.clone()));

// Fetch statistics for second file
let pq_meta =
fetch_parquet_metadata(store.as_ref(), &files[1], None, None).await?;
let pq_meta = fetch_parquet_metadata(
store.as_ref(),
&files[1],
None,
None,
format.options().global.cache_metadata,
state.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;
let stats = statistics_from_parquet_meta_calc(&pq_meta, schema.clone())?;
assert_eq!(stats.num_rows, Precision::Exact(3));
// column c1: missing from the file so the table treats all 3 rows as null
Expand Down
4 changes: 3 additions & 1 deletion datafusion/core/tests/parquet/custom_reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,14 +242,16 @@ impl AsyncFileReader for ParquetFileReader {
&self.meta,
self.metadata_size_hint,
None,
false,
None,
)
.await
.map_err(|e| {
ParquetError::General(format!(
"AsyncChunkReader::get_metadata error: {e}"
))
})?;
Ok(Arc::new(metadata))
Ok(metadata)
})
}
}
71 changes: 60 additions & 11 deletions datafusion/datasource-parquet/src/file_format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ use datafusion_physical_plan::Accumulator;
use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
use datafusion_session::Session;

use crate::reader::CachedParquetFileReaderFactory;
use crate::reader::{CachedParquetFileReaderFactory, CachedParquetMetaData};
use crate::source::{parse_coerce_int96_string, ParquetSource};
use async_trait::async_trait;
use bytes::Bytes;
Expand All @@ -83,6 +83,7 @@ use parquet::arrow::async_reader::MetadataFetch;
use parquet::arrow::{parquet_to_arrow_schema, ArrowSchemaConverter, AsyncArrowWriter};
use parquet::basic::Type;

use datafusion_execution::cache::cache_manager::FileMetadataCache;
use parquet::errors::ParquetError;
use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader, RowGroupMetaData};
use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder};
Expand Down Expand Up @@ -310,6 +311,8 @@ async fn fetch_schema_with_location(
metadata_size_hint: Option<usize>,
file_decryption_properties: Option<&FileDecryptionProperties>,
coerce_int96: Option<TimeUnit>,
cache_metadata: bool,
file_metadata_cache: Option<Arc<dyn FileMetadataCache>>,

@alamb alamb Aug 11, 2025

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.

I re-reviewed these changes and since I think we always now have a file_metadata_cache we could probably make this be non optional. However, there are many tests that need to change

) -> Result<(Path, Schema)> {
let loc_path = file.location.clone();
let schema = fetch_schema(
Expand All @@ -318,6 +321,8 @@ async fn fetch_schema_with_location(
metadata_size_hint,
file_decryption_properties,
coerce_int96,
cache_metadata,
file_metadata_cache,
)
.await?;
Ok((loc_path, schema))
Expand Down Expand Up @@ -371,6 +376,8 @@ impl FileFormat for ParquetFormat {
self.metadata_size_hint(),
file_decryption_properties.as_ref(),
coerce_int96,
self.options.global.cache_metadata,
state.runtime_env().cache_manager.get_file_metadata_cache(),
)
})
.boxed() // Workaround https://github.com/rust-lang/rust/issues/64552
Expand Down Expand Up @@ -414,7 +421,7 @@ impl FileFormat for ParquetFormat {

async fn infer_stats(
&self,
_state: &dyn Session,
state: &dyn Session,
store: &Arc<dyn ObjectStore>,
table_schema: SchemaRef,
object: &ObjectMeta,
Expand All @@ -429,6 +436,8 @@ impl FileFormat for ParquetFormat {
object,
self.metadata_size_hint(),
file_decryption_properties.as_ref(),
self.options.global.cache_metadata,
state.runtime_env().cache_manager.get_file_metadata_cache(),
)
.await?;
Ok(stats)
Expand Down Expand Up @@ -974,19 +983,47 @@ pub async fn fetch_parquet_metadata(
meta: &ObjectMeta,
size_hint: Option<usize>,
#[allow(unused)] decryption_properties: Option<&FileDecryptionProperties>,
) -> Result<ParquetMetaData> {
cache_metadata: bool,
file_metadata_cache: Option<Arc<dyn FileMetadataCache>>,
) -> Result<Arc<ParquetMetaData>> {
// Check cache first if caching is enabled
if cache_metadata {
if let Some(cache) = &file_metadata_cache {
if let Some(cached_metadata) = cache.get(meta) {
if let Some(parquet_metadata) = cached_metadata
.as_any()
.downcast_ref::<CachedParquetMetaData>()
{
return Ok(Arc::clone(parquet_metadata.parquet_metadata()));
}
}
}
}

let file_size = meta.size;
let fetch = ObjectStoreFetch::new(store, meta);

let reader = ParquetMetaDataReader::new().with_prefetch_hint(size_hint);

#[cfg(feature = "parquet_encryption")]
let reader = reader.with_decryption_properties(decryption_properties);

reader
.load_and_finish(fetch, file_size)
.await
.map_err(DataFusionError::from)
let metadata = Arc::new(
reader
.load_and_finish(fetch, file_size)
.await
.map_err(DataFusionError::from)?,
);

if cache_metadata {
if let Some(cache) = file_metadata_cache {
cache.put(
meta,
Arc::new(CachedParquetMetaData::new(Arc::clone(&metadata))),
);
}
}
Comment on lines +996 to +1094

@nuno-faria nuno-faria Aug 8, 2025

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.

I think there is an issue with the fetch_parquet_metadata function. When this function is initially called to retrieve the schema (in fetch_schema), it will read the metadata and update the cache. When the CachedParquetFileReader tries to get the metadata, it checks that it is present in the cache. However, the cached metadata does not contain the page index, as it is not retrieved in the fetch_parquet_metadata, meaning it will have to be read in every query.

So this fetch_parquet_metadata needs to retrieve the entire metadata for the caching to be effective. On the other hand, after this we will have two different places where the entire metadata is read and cached (CachedParquetFileReader and fetch_parquet_metadata), so creating an utility function retrieve_full_parquet_metadata -> Arc<ParquetMetaData> might be useful to avoid duplicate modifications in the future.

@shehabgamin shehabgamin Aug 9, 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.

Ahh nice catch, will do!


Ok(metadata)
}

/// Read and parse the schema of the Parquet file at location `path`
Expand All @@ -996,12 +1033,16 @@ async fn fetch_schema(
metadata_size_hint: Option<usize>,
file_decryption_properties: Option<&FileDecryptionProperties>,
coerce_int96: Option<TimeUnit>,
cache_metadata: bool,
file_metadata_cache: Option<Arc<dyn FileMetadataCache>>,
) -> Result<Schema> {
let metadata = fetch_parquet_metadata(
store,
file,
metadata_size_hint,
file_decryption_properties,
cache_metadata,
file_metadata_cache,
)
.await?;
let file_metadata = metadata.file_metadata();
Expand All @@ -1026,10 +1067,18 @@ pub async fn fetch_statistics(
file: &ObjectMeta,
metadata_size_hint: Option<usize>,
decryption_properties: Option<&FileDecryptionProperties>,
cache_metadata: bool,
file_metadata_cache: Option<Arc<dyn FileMetadataCache>>,
) -> Result<Statistics> {
let metadata =
fetch_parquet_metadata(store, file, metadata_size_hint, decryption_properties)
.await?;
let metadata = fetch_parquet_metadata(
store,
file,
metadata_size_hint,
decryption_properties,
cache_metadata,
file_metadata_cache,
)
.await?;
statistics_from_parquet_meta_calc(&metadata, table_schema)
}

Expand Down
2 changes: 1 addition & 1 deletion datafusion/datasource-parquet/src/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ pub mod file_format;
mod metrics;
mod opener;
mod page_filter;
mod reader;
pub mod reader;

@alamb alamb Aug 7, 2025

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.

why does this need to be made pub? I reverted the change and things still seem to compile just fine

@shehabgamin shehabgamin Aug 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.

So downstream crates can use it. The reader file looked like it could generally be useful for downstream crates which is why I made the entire reader public.

Currently in Sail we build our own cache for each cache type in CacheManagerConfig. I was unable to access CachedParquetMetaData to do something like the following unless I made CachedParquetMetaData public (full code):

if let Some(parquet_metadata) =
    value.1.as_any().downcast_ref::<CachedParquetMetaData>()
{
    parquet_metadata
        .parquet_metadata()
        .memory_size()
        .min(u32::MAX as usize)
}

@alamb alamb Aug 8, 2025

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.

I see -- my concern is that this change might be easy to accidentally undo / break in the future. Maybe to make it more deliberate, you could leave mod reader and then pub use both CachedParquetMetaData and CachedParquetMetaData 🤔

@shehabgamin shehabgamin Aug 9, 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.

Works with me!

mod row_filter;
mod row_group_filter;
pub mod source;
Expand Down
20 changes: 15 additions & 5 deletions datafusion/datasource-parquet/src/reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ impl DefaultParquetFileReaderFactory {
/// This implementation does not coalesce I/O operations or cache bytes. Such
/// optimizations can be done either at the object store level or by providing a
/// custom implementation of [`ParquetFileReaderFactory`].
pub(crate) struct ParquetFileReader {
pub struct ParquetFileReader {
pub file_metrics: ParquetFileMetrics,
pub inner: ParquetObjectReader,
}
Expand Down Expand Up @@ -212,7 +212,7 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory {
/// Implements [`AsyncFileReader`] for a Parquet file in object storage. Reads the file metadata
/// from the [`FileMetadataCache`], if available, otherwise reads it directly from the file and then
/// updates the cache.
pub(crate) struct CachedParquetFileReader {
pub struct CachedParquetFileReader {

@alamb alamb Aug 7, 2025

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.

Likewise, I locally reverted these changes to visibility and everything seems to have compiled just fine

@shehabgamin shehabgamin Aug 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.

Same reason as here:
https://github.com/apache/datafusion/pull/17022/files#r2262208618

pub file_metrics: ParquetFileMetrics,
pub inner: ParquetObjectReader,
file_meta: FileMeta,
Expand Down Expand Up @@ -253,10 +253,10 @@ impl AsyncFileReader for CachedParquetFileReader {

// lookup if the metadata is already cached
if let Some(metadata) = metadata_cache.get(object_meta) {
if let Some(parquet_metadata) =
if let Some(cached_parquet_metadata) =
metadata.as_any().downcast_ref::<CachedParquetMetaData>()
{
return Ok(Arc::clone(&parquet_metadata.0));
return Ok(Arc::clone(cached_parquet_metadata.parquet_metadata()));
}
}

Expand All @@ -282,7 +282,17 @@ impl AsyncFileReader for CachedParquetFileReader {
}

/// Wrapper to implement [`FileMetadata`] for [`ParquetMetaData`].
struct CachedParquetMetaData(Arc<ParquetMetaData>);
pub struct CachedParquetMetaData(Arc<ParquetMetaData>);

impl CachedParquetMetaData {
pub fn new(metadata: Arc<ParquetMetaData>) -> Self {
Self(metadata)
}

pub fn parquet_metadata(&self) -> &Arc<ParquetMetaData> {
&self.0
}
}

impl FileMetadata for CachedParquetMetaData {
fn as_any(&self) -> &dyn Any {
Expand Down
Loading