From 2ef06534568a9f93efcdde0c6fe0ee933fe24935 Mon Sep 17 00:00:00 2001 From: Jake Wallin Date: Tue, 20 Oct 2026 11:27:53 +0000 Subject: [PATCH] fix: Validate file cache identity and isolate object stores Key file caches by object store URL and path. Validate cached metadata using strong ETags, versions, or size and modification time, while checking the schema independently. Carry store identity through file-format inference and Parquet readers. Expose store URLs separately from relative paths in CLI cache functions, with regression coverage and upgrade guidance. --- datafusion-cli/src/functions.rs | 22 +- datafusion-cli/src/main.rs | 75 ++++ .../custom_data_source/custom_file_format.rs | 8 +- datafusion/catalog-listing/src/options.rs | 5 +- datafusion/catalog-listing/src/table.rs | 51 ++- .../core/src/datasource/file_format/csv.rs | 13 +- .../core/src/datasource/file_format/json.rs | 15 +- .../core/src/datasource/file_format/mod.rs | 15 +- .../src/datasource/file_format/parquet.rs | 134 +++++-- .../core/src/datasource/physical_plan/avro.rs | 21 +- .../core/src/datasource/physical_plan/json.rs | 2 +- .../src/datasource/physical_plan/parquet.rs | 14 +- .../tests/datasource/cache_object_identity.rs | 329 +++++++++++++++++ datafusion/core/tests/datasource/mod.rs | 1 + .../tests/datasource/object_store_access.rs | 41 ++- datafusion/core/tests/parquet/page_pruning.rs | 7 +- .../datasource-arrow/src/file_format.rs | 5 + datafusion/datasource-avro/src/file_format.rs | 3 + datafusion/datasource-csv/src/file_format.rs | 3 + datafusion/datasource-json/src/file_format.rs | 3 + .../datasource-parquet/src/file_format.rs | 23 +- datafusion/datasource-parquet/src/metadata.rs | 135 ++++++- .../datasource-parquet/src/opener/mod.rs | 39 +- datafusion/datasource-parquet/src/reader.rs | 27 +- datafusion/datasource-parquet/src/source.rs | 9 +- datafusion/datasource/src/file_format.rs | 23 +- datafusion/datasource/src/url.rs | 87 ++++- .../execution/src/cache/cache_manager.rs | 224 +++++++++++- .../execution/src/cache/default_cache.rs | 336 ++++++++++++------ datafusion/execution/src/cache/mod.rs | 93 ++++- .../library-user-guide/upgrading/56.0.0.md | 33 ++ docs/source/user-guide/cli/functions.md | 81 +++-- 32 files changed, 1591 insertions(+), 286 deletions(-) create mode 100644 datafusion/core/tests/datasource/cache_object_identity.rs diff --git a/datafusion-cli/src/functions.rs b/datafusion-cli/src/functions.rs index d5889bd2a50ba..4389940307bdb 100644 --- a/datafusion-cli/src/functions.rs +++ b/datafusion-cli/src/functions.rs @@ -701,6 +701,7 @@ impl TableFunctionImpl for MetadataCacheFunc { let schema = Arc::new(Schema::new(vec![ Field::new("path", DataType::Utf8, false), + Field::new("object_store_url", DataType::Utf8, false), Field::new( "file_modified", DataType::Timestamp(TimeUnit::Millisecond, None), @@ -718,6 +719,7 @@ impl TableFunctionImpl for MetadataCacheFunc { // construct record batch from metadata let mut path_arr = Vec::with_capacity(cached_entries.len()); + let mut object_store_url_arr = Vec::with_capacity(cached_entries.len()); let mut file_modified_arr = Vec::with_capacity(cached_entries.len()); let mut file_size_bytes_arr = Vec::with_capacity(cached_entries.len()); let mut e_tag_arr = Vec::with_capacity(cached_entries.len()); @@ -728,7 +730,8 @@ impl TableFunctionImpl for MetadataCacheFunc { let mut extra = Vec::new(); for (path, entry) in cached_entries { - path_arr.push(path.to_string()); + path_arr.push(path.path.to_string()); + object_store_url_arr.push(path.object_store_url.to_string()); file_modified_arr .push(Some(entry.value.meta.last_modified.timestamp_millis())); file_size_bytes_arr.push(entry.value.meta.size); @@ -754,6 +757,7 @@ impl TableFunctionImpl for MetadataCacheFunc { schema.clone(), vec![ Arc::new(StringArray::from(path_arr)), + Arc::new(StringArray::from(object_store_url_arr)), Arc::new(TimestampMillisecondArray::from(file_modified_arr)), Arc::new(UInt64Array::from(file_size_bytes_arr)), Arc::new(StringArray::from(e_tag_arr)), @@ -821,6 +825,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { let schema = Arc::new(Schema::new(vec![ Field::new("path", DataType::Utf8, false), + Field::new("object_store_url", DataType::Utf8, false), Field::new("table", DataType::Utf8, false), Field::new( "file_modified", @@ -839,6 +844,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { // construct record batch from metadata let mut path_arr = vec![]; + let mut object_store_url_arr = vec![]; let mut table_arr = vec![]; let mut file_modified_arr = vec![]; let mut file_size_bytes_arr = vec![]; @@ -853,7 +859,8 @@ impl TableFunctionImpl for StatisticsCacheFunc { if let Some(file_statistics_cache) = self.cache_manager.get_file_statistic_cache() { for (path, entry) in file_statistics_cache.list_entries() { - path_arr.push(path.path.to_string()); + path_arr.push(path.store_path.path.to_string()); + object_store_url_arr.push(path.store_path.object_store_url.to_string()); table_arr .push(path.table.map_or_else(|| "".to_string(), |t| t.to_string())); file_modified_arr @@ -881,6 +888,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { schema.clone(), vec![ Arc::new(StringArray::from(path_arr)), + Arc::new(StringArray::from(object_store_url_arr)), Arc::new(StringArray::from(table_arr)), Arc::new(TimestampMillisecondArray::from(file_modified_arr)), Arc::new(UInt64Array::from(file_size_bytes_arr)), @@ -912,11 +920,13 @@ impl TableFunctionImpl for StatisticsCacheFunc { /// +---------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------+ /// | column_name | data_type | is_nullable | /// +---------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------+ -/// | table | Utf8 | NO | +/// | table | Utf8 | YES | /// | path | Utf8 | NO | +/// | object_store_url | Utf8 | NO | /// | metadata_size_bytes | UInt64 | NO | /// | expires_in | Duration(ms) | YES | /// | metadata_list | List(Struct("file_path": non-null Utf8, "file_modified": non-null Timestamp(ms), "file_size_bytes": non-null UInt64, "e_tag": Utf8, "version": Utf8), field: 'metadata') | YES | +/// | hits | UInt64 | NO | /// +---------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------+ /// ``` #[derive(Debug)] @@ -986,6 +996,7 @@ impl TableFunctionImpl for ListFilesCacheFunc { let schema = Arc::new(Schema::new(vec![ Field::new("table", DataType::Utf8, true), Field::new("path", DataType::Utf8, false), + Field::new("object_store_url", DataType::Utf8, false), Field::new("metadata_size_bytes", DataType::UInt64, false), // expires field in ListFilesEntry has type Instant when set, from which we cannot get "the number of seconds", hence using Duration instead of Timestamp as data type. Field::new( @@ -1003,6 +1014,7 @@ impl TableFunctionImpl for ListFilesCacheFunc { let mut table_arr = vec![]; let mut path_arr = vec![]; + let mut object_store_url_arr = vec![]; let mut metadata_size_bytes_arr = vec![]; let mut expires_arr = vec![]; @@ -1020,7 +1032,8 @@ impl TableFunctionImpl for ListFilesCacheFunc { for (path, entry) in list_files_cache.list_entries() { table_arr.push(path.table.map(|t| t.to_string())); - path_arr.push(path.path.to_string()); + path_arr.push(path.store_path.path.to_string()); + object_store_url_arr.push(path.store_path.object_store_url.to_string()); metadata_size_bytes_arr.push(entry.size_bytes as u64); // calculates time left before entry expires expires_arr.push( @@ -1062,6 +1075,7 @@ impl TableFunctionImpl for ListFilesCacheFunc { vec![ Arc::new(StringArray::from(table_arr)), Arc::new(StringArray::from(path_arr)), + Arc::new(StringArray::from(object_store_url_arr)), Arc::new(UInt64Array::from(metadata_size_bytes_arr)), Arc::new(DurationMillisecondArray::from(expires_arr)), Arc::new(GenericListArray::new( diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index d24f924bfb5bd..1a43c77aa5246 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -499,6 +499,8 @@ mod tests { use std::time::Duration; use super::*; + use arrow::array::{ArrayRef, RecordBatch, StringArray}; + use arrow::datatypes::DataType; use datafusion::execution::cache::default_cache::DefaultCache; use datafusion::{ common::test_util::batches_to_string, @@ -685,6 +687,79 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_cache_functions_object_store_identity() -> Result<(), DataFusionError> { + let ctx = SessionContext::new(); + let cache_manager = &ctx.runtime_env().cache_manager; + ctx.register_udtf( + "metadata_cache", + Arc::new(MetadataCacheFunc::new(Arc::clone(cache_manager))), + ); + ctx.register_udtf( + "statistics_cache", + Arc::new(StatisticsCacheFunc::new(Arc::clone(cache_manager))), + ); + ctx.register_udtf( + "list_files_cache", + Arc::new(ListFilesCacheFunc::new(Arc::clone(cache_manager))), + ); + + for store in ["first", "second"] { + ctx.register_object_store( + &Url::parse(&format!("mem://{store}")).unwrap(), + Arc::new(InMemory::new()), + ); + ctx.sql(&format!( + "COPY (SELECT 1 AS value) TO 'mem://{store}/data/file.parquet' STORED AS PARQUET" + )) + .await? + .collect() + .await?; + ctx.read_parquet( + format!("mem://{store}/data/"), + ParquetReadOptions::default(), + ) + .await? + .collect() + .await?; + } + + for (function, expected_path) in [ + ("metadata_cache", "data/file.parquet"), + ("statistics_cache", "data/file.parquet"), + ("list_files_cache", "data"), + ] { + let df = ctx + .sql(&format!( + "SELECT object_store_url, path FROM {function}() ORDER BY object_store_url" + )) + .await?; + let field = df + .schema() + .field_with_unqualified_name("object_store_url")?; + assert_eq!(field.data_type(), &DataType::Utf8, "{function}"); + assert!(!field.is_nullable(), "{function}"); + let results = df.collect().await?; + let expected = RecordBatch::try_from_iter([ + ( + "object_store_url", + Arc::new(StringArray::from(vec!["mem://first/", "mem://second/"])) + as ArrayRef, + ), + ( + "path", + Arc::new(StringArray::from(vec![expected_path, expected_path])), + ), + ])?; + assert_eq!( + batches_to_string(&results), + batches_to_string(&[expected]), + "{function}" + ); + } + Ok(()) + } + #[tokio::test] async fn test_metadata_cache() -> Result<(), DataFusionError> { let ctx = SessionContext::new(); diff --git a/datafusion-examples/examples/custom_data_source/custom_file_format.rs b/datafusion-examples/examples/custom_data_source/custom_file_format.rs index 0cfbe11877e4d..4f62c004913b9 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_format.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_format.rs @@ -36,7 +36,7 @@ use datafusion::{ table_schema::TableSchema, }, error::Result, - execution::session_state::SessionStateBuilder, + execution::{object_store::ObjectStoreUrl, session_state::SessionStateBuilder}, physical_expr_common::sort_expr::LexRequirement, physical_plan::ExecutionPlan, prelude::SessionContext, @@ -124,10 +124,11 @@ impl FileFormat for TSVFileFormat { &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, objects: &[ObjectMeta], ) -> Result { self.csv_file_format - .infer_schema(state, store, objects) + .infer_schema(state, store, object_store_url, objects) .await } @@ -135,11 +136,12 @@ impl FileFormat for TSVFileFormat { &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, object: &ObjectMeta, ) -> Result { self.csv_file_format - .infer_stats(state, store, table_schema, object) + .infer_stats(state, store, object_store_url, table_schema, object) .await } diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 44337e52a1e05..48cba85f92f35 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -300,7 +300,10 @@ impl ListingOptions { .filter(|object_meta| object_meta.size > 0) .collect(); - let schema = self.format.infer_schema(state, &store, &files).await?; + let schema = self + .format + .infer_schema(state, &store, &table_path.object_store(), &files) + .await?; Ok(schema) } diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index a72d2c2bc7d40..d5c2f97b2dac3 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -40,8 +40,10 @@ use datafusion_datasource::{ ListingTableUrl, PartitionedFile, TableSchemaBuilder, compute_all_files_statistics, }; use datafusion_execution::cache::cache_manager::{ - CachedFileMetadata, FileStatisticsCache, SchemaFingerprint, TableScopedPath, + CachedFileMetadata, FileStatisticsCache, SchemaFingerprint, StoreScopedPath, + TableScopedPath, }; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; @@ -867,13 +869,15 @@ impl ListingTable { lfc.drop_table_entries(table_ref)?; } else { let table_prefix = table_path.prefix(); + let object_store_url = table_path.object_store(); let keys: Vec<_> = lfc .list_entries() .into_keys() .filter(|key| { key.table.is_none() - && (key.path.prefix_matches(table_prefix) - || table_prefix.prefix_matches(&key.path)) + && key.store_path.object_store_url == object_store_url + && (key.store_path.path.prefix_matches(table_prefix) + || table_prefix.prefix_matches(&key.store_path.path)) }) .collect(); for key in keys { @@ -937,6 +941,7 @@ impl ListingTable { &'a self, ctx: &'a dyn Session, store: &'a Arc, + object_store_url: &'a ObjectStoreUrl, listing_time_filters: &'a [Expr], file_limit: Option, ) -> datafusion_common::Result<(FileGroup, bool)> { @@ -968,8 +973,13 @@ impl ListingTable { .map(|part_file| async { let part_file = part_file?; let (statistics, ordering) = if ctx.config().collect_statistics() { - self.do_collect_statistics_and_ordering(ctx, store, &part_file) - .await? + self.do_collect_statistics_and_ordering( + ctx, + store, + object_store_url, + &part_file, + ) + .await? } else { (Arc::new(Statistics::new_unknown(&self.file_schema)), None) }; @@ -998,8 +1008,8 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? + let (store, object_store_url) = if let Some(url) = self.table_paths.first() { + (ctx.runtime_env().object_store(url)?, url.object_store()) } else { return Ok(ListFilesResult { file_groups: vec![], @@ -1008,7 +1018,7 @@ impl ListingTable { }); }; let (file_group, inexact_stats) = self - .collect_files_for_scan(ctx, &store, filters, limit) + .collect_files_for_scan(ctx, &store, &object_store_url, filters, limit) .await?; // Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N @@ -1060,8 +1070,8 @@ impl ListingTable { ); } - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? + let (store, object_store_url) = if let Some(url) = self.table_paths.first() { + (ctx.runtime_env().object_store(url)?, url.object_store()) } else { return Ok(ListFilesResult { file_groups: vec![], @@ -1069,8 +1079,9 @@ impl ListingTable { grouped_by_partition: false, }); }; - let (file_group, inexact_stats) = - self.collect_files_for_scan(ctx, &store, &[], None).await?; + let (file_group, inexact_stats) = self + .collect_files_for_scan(ctx, &store, &object_store_url, &[], None) + .await?; let mut file_groups = file_group.split_files(file_group_count); if !file_groups.is_empty() { file_groups.resize_with(file_group_count, || FileGroup::new(vec![])); @@ -1141,15 +1152,19 @@ impl ListingTable { &self, ctx: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, part_file: &PartitionedFile, ) -> datafusion_common::Result<(Arc, Option)> { let path = TableScopedPath { table: part_file.table_reference.clone(), - path: part_file.object_meta.location.clone(), + store_path: StoreScopedPath::new( + object_store_url.clone(), + part_file.object_meta.location.clone(), + ), }; let meta = &part_file.object_meta; - // Check cache first. The key stays `{table, path}` for cheap lookups; + // Check cache first. The key identifies the store, table and path; // the cached value carries the schema fingerprint to prevent reusing // stats computed under a different file schema. if let Some(cache) = &self.collected_statistics @@ -1164,7 +1179,13 @@ impl ListingTable { let file_meta = self .options .format - .infer_stats_and_ordering(ctx, store, Arc::clone(&self.file_schema), meta) + .infer_stats_and_ordering( + ctx, + store, + object_store_url, + Arc::clone(&self.file_schema), + meta, + ) .await?; let statistics = Arc::new(file_meta.statistics); diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs index 8442b00900002..797e323bd2f81 100644 --- a/datafusion/core/src/datasource/file_format/csv.rs +++ b/datafusion/core/src/datasource/file_format/csv.rs @@ -44,6 +44,7 @@ mod tests { use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::FileFormat; use datafusion_datasource::write::BatchSerializer; + use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::{col, lit}; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::{ExecutionPlan, collect}; @@ -397,6 +398,7 @@ mod tests { .infer_schema( &state, &(variable_object_store.clone() as Arc), + &ObjectStoreUrl::parse("memory://test")?, &[object_meta], ) .await?; @@ -458,6 +460,7 @@ mod tests { .infer_schema( &state, &(variable_object_store.clone() as Arc), + &ObjectStoreUrl::parse("memory://test")?, &[object_meta], ) .await?; @@ -513,6 +516,7 @@ mod tests { .infer_schema( &state, &(chunked_object_store as Arc), + &ObjectStoreUrl::parse("memory://test")?, &[object_meta], ) .await?; @@ -1394,6 +1398,7 @@ mod tests { .infer_schema( &state, &(variable_object_store.clone() as Arc), + &ObjectStoreUrl::parse("memory://test")?, &[object_meta], ) .await?; @@ -1469,6 +1474,7 @@ mod tests { .infer_schema( &state, &(variable_object_store.clone() as Arc), + &ObjectStoreUrl::parse("memory://test")?, &[object_meta], ) .await; @@ -1634,7 +1640,12 @@ mod tests { let ctx = SessionContext::new().state(); let error = CsvFormat::default() .with_has_header(true) - .infer_schema(&ctx, &store, std::slice::from_ref(&meta)) + .infer_schema( + &ctx, + &store, + &ObjectStoreUrl::local_filesystem(), + std::slice::from_ref(&meta), + ) .await .expect_err("duplicate header names must not infer a schema") .to_string(); diff --git a/datafusion/core/src/datasource/file_format/json.rs b/datafusion/core/src/datasource/file_format/json.rs index 1f6f27242e723..c379c8e87d3dc 100644 --- a/datafusion/core/src/datasource/file_format/json.rs +++ b/datafusion/core/src/datasource/file_format/json.rs @@ -36,6 +36,7 @@ mod tests { BatchDeserializer, DecoderDeserializer, DeserializerOutput, }; use datafusion_datasource::file_format::FileFormat; + use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::{ExecutionPlan, collect}; @@ -75,7 +76,12 @@ mod tests { let store = Arc::new(LocalFileSystem::new()) as _; let format = JsonFormat::default().with_newline_delimited(false); format - .infer_schema(&ctx, &store, &[local_unpartitioned_file(&path)]) + .infer_schema( + &ctx, + &store, + &ObjectStoreUrl::local_filesystem(), + &[local_unpartitioned_file(&path)], + ) .await } @@ -214,7 +220,12 @@ mod tests { let format = JsonFormat::default().with_schema_infer_max_rec(3); let file_schema = format - .infer_schema(&ctx, &store, &[local_unpartitioned_file(filename)]) + .infer_schema( + &ctx, + &store, + &ObjectStoreUrl::local_filesystem(), + &[local_unpartitioned_file(filename)], + ) .await .expect("Schema inference"); diff --git a/datafusion/core/src/datasource/file_format/mod.rs b/datafusion/core/src/datasource/file_format/mod.rs index c46b472bd6404..b74148cfeedef 100644 --- a/datafusion/core/src/datasource/file_format/mod.rs +++ b/datafusion/core/src/datasource/file_format/mod.rs @@ -63,14 +63,25 @@ pub(crate) mod test_util { file_schema } else { format - .infer_schema(state, &store, std::slice::from_ref(&meta)) + .infer_schema( + state, + &store, + &ObjectStoreUrl::local_filesystem(), + std::slice::from_ref(&meta), + ) .await? }; let table_schema = TableSchema::from(&file_schema); let statistics = format - .infer_stats(state, &store, file_schema.clone(), &meta) + .infer_stats( + state, + &store, + &ObjectStoreUrl::local_filesystem(), + file_schema.clone(), + &meta, + ) .await?; let file_groups = vec![vec![PartitionedFile::new_from_meta(meta)].into()]; diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index c5685e1d2dfcb..530aae20ef5f9 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -194,12 +194,17 @@ mod tests { ForceViews::No => false, }; let format = ParquetFormat::default().with_force_view_types(force_views); - let schema = format.infer_schema(&ctx, &store, &meta).await?; + let schema = format + .infer_schema(&ctx, &store, &ObjectStoreUrl::local_filesystem(), &meta) + .await?; let file_metadata_cache = ctx.runtime_env().cache_manager.get_file_metadata_cache(); let stats = DFParquetMetadata::new(&store, &meta[0]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ) .fetch_statistics(&schema) .await?; @@ -210,7 +215,10 @@ mod tests { assert_eq!(c2_stats.null_count, Precision::Exact(3)); let stats = DFParquetMetadata::new(&store, &meta[1]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ) .fetch_statistics(&schema) .await?; @@ -257,7 +265,9 @@ mod tests { let session = SessionContext::new(); let ctx = session.state(); let format = ParquetFormat::default(); - let schema = format.infer_schema(&ctx, &store, &meta).await?; + let schema = format + .infer_schema(&ctx, &store, &ObjectStoreUrl::local_filesystem(), &meta) + .await?; let order: Vec<_> = ["a", "b", "c", "d"] .into_iter() @@ -387,8 +397,10 @@ mod tests { df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 2); - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ); // Increases by 3 because cache has no entries yet df_meta.fetch_metadata().await?; @@ -399,7 +411,8 @@ mod tests { assert_eq!(store.request_count(), 5); // Increase by 2 because `get_file_metadata_cache()` is None - let df_meta = df_meta.with_file_metadata_cache(None); + let df_meta = + df_meta.with_file_metadata_cache(None, ObjectStoreUrl::local_filesystem()); df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 7); @@ -411,15 +424,31 @@ mod tests { .with_metadata_size_hint(Some(9)) .with_force_view_types(force_views); // Increase by 3, partial cache being used. - let _schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let _schema = format + .infer_schema( + &ctx, + &store.upcast(), + &ObjectStoreUrl::local_filesystem(), + &meta, + ) + .await?; assert_eq!(store.request_count(), 10); // No increase, full cache being used. - let schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let schema = format + .infer_schema( + &ctx, + &store.upcast(), + &ObjectStoreUrl::local_filesystem(), + &meta, + ) + .await?; assert_eq!(store.request_count(), 10); // No increase, cache being used - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ); let stats = df_meta.fetch_statistics(&schema).await?; assert_eq!(store.request_count(), 10); @@ -446,8 +475,10 @@ mod tests { let ctx = session.state(); let file_metadata_cache = ctx.runtime_env().cache_manager.get_file_metadata_cache(); - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ); // Increases by 1 because cache has no entries yet and new session context df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 2); @@ -457,7 +488,8 @@ mod tests { assert_eq!(store.request_count(), 2); // Increase by 1 because `get_file_metadata_cache` is None - let df_meta = df_meta.with_file_metadata_cache(None); + let df_meta = + df_meta.with_file_metadata_cache(None, ObjectStoreUrl::local_filesystem()); df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 3); @@ -465,14 +497,30 @@ mod tests { .with_metadata_size_hint(Some(size_hint)) .with_force_view_types(force_views); // Increase by 1, partial cache being used. - let _schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let _schema = format + .infer_schema( + &ctx, + &store.upcast(), + &ObjectStoreUrl::local_filesystem(), + &meta, + ) + .await?; assert_eq!(store.request_count(), 4); // No increase, full cache being used. - let schema = format.infer_schema(&ctx, &store.upcast(), &meta).await?; + let schema = format + .infer_schema( + &ctx, + &store.upcast(), + &ObjectStoreUrl::local_filesystem(), + &meta, + ) + .await?; assert_eq!(store.request_count(), 4); // No increase, cache being used - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ); let stats = df_meta.fetch_statistics(&schema).await?; assert_eq!(store.request_count(), 4); @@ -495,8 +543,10 @@ mod tests { assert_eq!(store.request_count(), 1); // No increase because cache has an entry - let df_meta = - df_meta.with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))); + let df_meta = df_meta.with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ); df_meta.fetch_metadata().await?; assert_eq!(store.request_count(), 1); @@ -544,17 +594,34 @@ mod tests { let state = SessionContext::new().state(); // Make metadata size hint None to keep original behavior let format = ParquetFormat::default().with_metadata_size_hint(None); - let _schema = format.infer_schema(&state, &store.upcast(), &files).await?; + let _schema = format + .infer_schema( + &state, + &store.upcast(), + &ObjectStoreUrl::local_filesystem(), + &files, + ) + .await?; assert_eq!(store.request_count(), 3); // No increase, cache being used. - let schema = format.infer_schema(&state, &store.upcast(), &files).await?; + let schema = format + .infer_schema( + &state, + &store.upcast(), + &ObjectStoreUrl::local_filesystem(), + &files, + ) + .await?; assert_eq!(store.request_count(), 3); // No increase in request count because cache is not empty let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); let stats = DFParquetMetadata::new(store.as_ref(), &files[0]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ) .fetch_statistics(&schema) .await?; assert_eq!(stats.num_rows, Precision::Exact(4)); @@ -620,7 +687,14 @@ mod tests { let format = ParquetFormat::default() .with_force_view_types(force_views) .with_metadata_size_hint(None); - let schema = format.infer_schema(&state, &store.upcast(), &files).await?; + let schema = format + .infer_schema( + &state, + &store.upcast(), + &ObjectStoreUrl::local_filesystem(), + &files, + ) + .await?; assert_eq!(store.request_count(), 6); let null_i64 = ScalarValue::Int64(None); @@ -634,7 +708,10 @@ mod tests { let file_metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); let stats = DFParquetMetadata::new(store.as_ref(), &files[0]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ) .fetch_statistics(&schema) .await?; assert_eq!(store.request_count(), 6); @@ -663,7 +740,10 @@ mod tests { // No increase in request count because cache is not empty let stats = DFParquetMetadata::new(store.as_ref(), &files[1]) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + ObjectStoreUrl::local_filesystem(), + ) .fetch_statistics(&schema) .await?; assert_eq!(store.request_count(), 6); @@ -1836,7 +1916,7 @@ mod tests { let ctx = SessionContext::new().state(); let error = ParquetFormat::default() - .infer_schema(&ctx, &store, &meta) + .infer_schema(&ctx, &store, &ObjectStoreUrl::local_filesystem(), &meta) .await .expect_err("duplicate field names must not infer a schema") .to_string(); diff --git a/datafusion/core/src/datasource/physical_plan/avro.rs b/datafusion/core/src/datasource/physical_plan/avro.rs index c9ee2cc407783..0c730d04f4f94 100644 --- a/datafusion/core/src/datasource/physical_plan/avro.rs +++ b/datafusion/core/src/datasource/physical_plan/avro.rs @@ -78,7 +78,12 @@ mod tests { let meta = local_unpartitioned_file(filename); let file_schema = AvroFormat {} - .infer_schema(&state, &store, std::slice::from_ref(&meta)) + .infer_schema( + &state, + &store, + &ObjectStoreUrl::local_filesystem(), + std::slice::from_ref(&meta), + ) .await?; let source = Arc::new(AvroSource::new(Arc::clone(&file_schema))); @@ -143,7 +148,12 @@ mod tests { let object_store_url = ObjectStoreUrl::local_filesystem(); let meta = local_unpartitioned_file(filename); let actual_schema = AvroFormat {} - .infer_schema(&state, &object_store, std::slice::from_ref(&meta)) + .infer_schema( + &state, + &object_store, + &object_store_url, + std::slice::from_ref(&meta), + ) .await?; let mut builder = SchemaBuilder::from(actual_schema.fields()); @@ -216,7 +226,12 @@ mod tests { let object_store_url = ObjectStoreUrl::local_filesystem(); let meta = local_unpartitioned_file(filename); let file_schema = AvroFormat {} - .infer_schema(&state, &object_store, std::slice::from_ref(&meta)) + .infer_schema( + &state, + &object_store, + &object_store_url, + std::slice::from_ref(&meta), + ) .await?; let mut partitioned_file = PartitionedFile::from(meta); diff --git a/datafusion/core/src/datasource/physical_plan/json.rs b/datafusion/core/src/datasource/physical_plan/json.rs index 6b4361e0c4d07..055e9c66a0b58 100644 --- a/datafusion/core/src/datasource/physical_plan/json.rs +++ b/datafusion/core/src/datasource/physical_plan/json.rs @@ -90,7 +90,7 @@ mod tests { .object_meta; let schema = JsonFormat::default() .with_file_compression_type(file_compression_type.to_owned()) - .infer_schema(state, &store, std::slice::from_ref(&meta)) + .infer_schema(state, &store, &store_url, std::slice::from_ref(&meta)) .await .unwrap(); diff --git a/datafusion/core/src/datasource/physical_plan/parquet.rs b/datafusion/core/src/datasource/physical_plan/parquet.rs index e0b186406b50d..5659cef92ccdf 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet.rs +++ b/datafusion/core/src/datasource/physical_plan/parquet.rs @@ -1583,7 +1583,12 @@ mod tests { let store = Arc::new(LocalFileSystem::new()) as _; let file_schema = ParquetFormat::default() - .infer_schema(&state, &store, std::slice::from_ref(&meta)) + .infer_schema( + &state, + &store, + &ObjectStoreUrl::local_filesystem(), + std::slice::from_ref(&meta), + ) .await?; let group_empty = vec![FileGroup::new(vec![file_range(&meta, 0, 2)])]; @@ -1615,7 +1620,12 @@ mod tests { let meta = local_unpartitioned_file(filename); let schema = ParquetFormat::default() - .infer_schema(&state, &store, std::slice::from_ref(&meta)) + .infer_schema( + &state, + &store, + &object_store_url, + std::slice::from_ref(&meta), + ) .await .unwrap(); diff --git a/datafusion/core/tests/datasource/cache_object_identity.rs b/datafusion/core/tests/datasource/cache_object_identity.rs new file mode 100644 index 0000000000000..74fd5c684e22d --- /dev/null +++ b/datafusion/core/tests/datasource/cache_object_identity.rs @@ -0,0 +1,329 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; +use async_trait::async_trait; +use bytes::Bytes; +use datafusion::prelude::{ParquetReadOptions, SessionContext, col, lit}; +use datafusion_common::assert_batches_eq; +use datafusion_execution::cache::{StoreScopedPath, TableScopedPath}; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_functions_aggregate::expr_fn::{count, max, min}; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; +use object_store::memory::InMemory; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, +}; +use parquet::arrow::ArrowWriter; +use parquet::basic::{Compression, Encoding}; +use parquet::file::properties::WriterProperties; +use std::fmt::{Display, Formatter}; +use std::sync::Arc; +use url::Url; + +fn parquet_bytes(values: Vec) -> Bytes { + let values: ArrayRef = Arc::new(Int32Array::from(values)); + let batch = RecordBatch::try_from_iter([("value", values)]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .set_compression(Compression::UNCOMPRESSED) + .build(); + let mut bytes = vec![]; + let mut writer = + ArrowWriter::try_new(&mut bytes, batch.schema(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + bytes.into() +} + +async fn aggregate_file(ctx: &SessionContext, url: &str) -> Vec { + ctx.read_parquet(url, ParquetReadOptions::default()) + .await + .unwrap() + .aggregate( + vec![], + vec![ + min(col("value")).alias("min"), + max(col("value")).alias("max"), + count(col("value")).alias("count"), + ], + ) + .unwrap() + .collect() + .await + .unwrap() +} + +#[tokio::test] +async fn parquet_caches_separate_object_stores() { + let first = Arc::new(FixedTimeStore::default()); + let second = Arc::new(FixedTimeStore::default()); + let path = Path::from("data/file.parquet"); + let first_bytes = parquet_bytes(vec![1, 2, 3]); + let second_bytes = parquet_bytes(vec![7, 8, 9]); + assert_eq!(first_bytes.len(), second_bytes.len()); + first.put(&path, first_bytes.into()).await.unwrap(); + second.put(&path, second_bytes.into()).await.unwrap(); + // Every ObjectMeta field, including ETag, is identical across the stores. + assert_eq!( + first.head(&path).await.unwrap(), + second.head(&path).await.unwrap() + ); + + let ctx = SessionContext::new(); + ctx.runtime_env() + .register_object_store(&Url::parse("mem://first").unwrap(), first); + ctx.runtime_env() + .register_object_store(&Url::parse("mem://second").unwrap(), second); + let first_result = aggregate_file(&ctx, "mem://first/data/").await; + assert_batches_eq!( + [ + "+-----+-----+-------+", + "| min | max | count |", + "+-----+-----+-------+", + "| 1 | 3 | 3 |", + "+-----+-----+-------+", + ], + &first_result + ); + let second_result = aggregate_file(&ctx, "mem://second/data/").await; + assert_batches_eq!( + [ + "+-----+-----+-------+", + "| min | max | count |", + "+-----+-----+-------+", + "| 7 | 9 | 3 |", + "+-----+-----+-------+", + ], + &second_result + ); + + // Execute scans as well as metadata-backed aggregates, including repeated + // reads through the Parquet reader's metadata cache. + for (url, value) in [("mem://first/data/", 2), ("mem://second/data/", 8)] { + for _ in 0..2 { + let results = ctx + .read_parquet(url, ParquetReadOptions::default()) + .await + .unwrap() + .filter(col("value").eq(lit(value))) + .unwrap() + .collect() + .await + .unwrap(); + let values = results + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect::>(); + assert_eq!(values, vec![Some(value)], "{url}"); + } + } + + let caches = &ctx.runtime_env().cache_manager; + let metadata = caches.get_file_metadata_cache(); + let statistics = caches.get_file_statistic_cache().unwrap(); + let listings = caches.get_list_files_cache().unwrap(); + for url in ["mem://first", "mem://second"] { + let object_store_url = ObjectStoreUrl::parse(url).unwrap(); + assert!( + metadata + .get(&StoreScopedPath::new( + object_store_url.clone(), + path.clone() + )) + .is_some() + ); + assert!( + statistics + .get(&TableScopedPath { + table: None, + store_path: StoreScopedPath::new( + object_store_url.clone(), + path.clone() + ), + }) + .is_some() + ); + assert!( + listings + .get(&TableScopedPath { + table: None, + store_path: StoreScopedPath::new( + object_store_url, + Path::from("data") + ), + }) + .is_some() + ); + } + assert_eq!(metadata.len(), 2); + assert_eq!(statistics.len(), 2); + assert_eq!(listings.len(), 2); +} + +#[tokio::test] +async fn parquet_caches_refresh_changed_etag() { + let store = Arc::new(FixedTimeStore::default()); + let path = Path::from("data.parquet"); + let original = parquet_bytes(vec![1, 2, 3]); + let replacement = parquet_bytes(vec![7, 8, 9]); + assert_eq!(original.len(), replacement.len()); + store.put(&path, original.into()).await.unwrap(); + let before = store.head(&path).await.unwrap(); + let ctx = SessionContext::new(); + ctx.runtime_env() + .register_object_store(&Url::parse("mem://overwrite").unwrap(), store.clone()); + // A single-file URL obtains fresh metadata without a cached directory listing. + let original_result = aggregate_file(&ctx, "mem://overwrite/data.parquet").await; + assert_batches_eq!( + [ + "+-----+-----+-------+", + "| min | max | count |", + "+-----+-----+-------+", + "| 1 | 3 | 3 |", + "+-----+-----+-------+", + ], + &original_result + ); + + store.put(&path, replacement.into()).await.unwrap(); + let after = store.head(&path).await.unwrap(); + assert_eq!(before.size, after.size); + assert_eq!(before.last_modified, after.last_modified); + assert_ne!(before.e_tag, after.e_tag); + let replacement_result = aggregate_file(&ctx, "mem://overwrite/data.parquet").await; + assert_batches_eq!( + [ + "+-----+-----+-------+", + "| min | max | count |", + "+-----+-----+-------+", + "| 7 | 9 | 3 |", + "+-----+-----+-------+", + ], + &replacement_result + ); + let object_store_url = ObjectStoreUrl::parse("mem://overwrite").unwrap(); + let caches = &ctx.runtime_env().cache_manager; + let metadata = caches + .get_file_metadata_cache() + .get(&StoreScopedPath::new( + object_store_url.clone(), + path.clone(), + )) + .unwrap(); + let statistics = caches + .get_file_statistic_cache() + .unwrap() + .get(&TableScopedPath { + table: None, + store_path: StoreScopedPath::new(object_store_url, path), + }) + .unwrap(); + assert_eq!(metadata.meta.e_tag, after.e_tag); + assert_eq!(statistics.meta.e_tag, after.e_tag); +} + +/// Keeps modification times identical so the tests isolate store identity and ETag. +#[derive(Debug, Default)] +struct FixedTimeStore(InMemory); + +impl Display for FixedTimeStore { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "FixedTimeStore") + } +} + +#[async_trait] +impl ObjectStore for FixedTimeStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> object_store::Result { + self.0.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> object_store::Result> { + self.0.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + let mut result = self.0.get_opts(location, options).await?; + result.meta.last_modified = chrono::DateTime::UNIX_EPOCH; + Ok(result) + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> BoxStream<'static, object_store::Result> { + self.0 + .list(prefix) + .map_ok(|mut meta| { + meta.last_modified = chrono::DateTime::UNIX_EPOCH; + meta + }) + .boxed() + } + + async fn list_with_delimiter( + &self, + prefix: Option<&Path>, + ) -> object_store::Result { + let mut result = self.0.list_with_delimiter(prefix).await?; + for meta in &mut result.objects { + meta.last_modified = chrono::DateTime::UNIX_EPOCH; + } + Ok(result) + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + self.0.delete_stream(locations) + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> object_store::Result<()> { + self.0.copy_opts(from, to, options).await + } +} diff --git a/datafusion/core/tests/datasource/mod.rs b/datafusion/core/tests/datasource/mod.rs index 3785aa0766182..109f19aaa02d1 100644 --- a/datafusion/core/tests/datasource/mod.rs +++ b/datafusion/core/tests/datasource/mod.rs @@ -19,6 +19,7 @@ //! //! Note tests for the Parquet format are in `parquet_integration` binary +mod cache_object_identity; // Include tests in csv module mod csv; mod object_store_access; diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 215725cf04961..9dd56d4614635 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -36,7 +36,8 @@ use datafusion_common::assert_batches_eq; use datafusion_datasource::ListingTableUrl; use datafusion_datasource_csv::CsvFormat; use datafusion_datasource_json::JsonFormat; -use datafusion_execution::cache::TableScopedPath; +use datafusion_execution::cache::{StoreScopedPath, TableScopedPath}; +use datafusion_execution::object_store::ObjectStoreUrl; use futures::stream::BoxStream; use insta::assert_snapshot; use object_store::memory::InMemory; @@ -219,6 +220,13 @@ async fn insert_invalidates_overlapping_unscoped_listings() { let ctx = SessionContext::new(); ctx.runtime_env() .register_object_store(&Url::parse("mem://").unwrap(), store); + let other_store = Arc::new(InMemory::new()); + other_store + .put(&Path::from("table/data.csv"), "4\n".into()) + .await + .unwrap(); + ctx.runtime_env() + .register_object_store(&Url::parse("mem://other").unwrap(), other_store); ctx.sql("SET datafusion.execution.listing_table_ignore_subdirectory = false") .await .unwrap() @@ -236,15 +244,16 @@ async fn insert_invalidates_overlapping_unscoped_listings() { .get_list_files_cache() .unwrap(); let tables = [ - ("root_table", "table"), - ("child_table", "table/region=US"), - ("descendant_table", "table/region=US/q1"), - ("sibling_table", "table/region=EU"), + ("root_table", "mem://", "table"), + ("child_table", "mem://", "table/region=US"), + ("descendant_table", "mem://", "table/region=US/q1"), + ("sibling_table", "mem://", "table/region=EU"), + ("other_store_table", "mem://other", "table"), ]; - for (name, path) in tables { + for (name, store_url, path) in tables { // Keep the URLs unscoped so overlapping tables share path-based cache entries. - let url = ListingTableUrl::parse(format!("mem:///{path}/")).unwrap(); + let url = ListingTableUrl::parse(format!("{store_url}/{path}/")).unwrap(); let options = ListingOptions::new(Arc::new(CsvFormat::default().with_has_header(false))) .with_file_extension(".csv"); @@ -261,7 +270,10 @@ async fn insert_invalidates_overlapping_unscoped_listings() { .unwrap(); let key = TableScopedPath { table: None, - path: Path::from(path), + store_path: StoreScopedPath::new( + ObjectStoreUrl::parse(store_url).unwrap(), + Path::from(path), + ), }; assert!(cache.get(&key).is_some()); } @@ -273,13 +285,20 @@ async fn insert_invalidates_overlapping_unscoped_listings() { .await .unwrap(); - for (name, path) in tables { + for (name, store_url, path) in tables { let key = TableScopedPath { table: None, - path: Path::from(path), + store_path: StoreScopedPath::new( + ObjectStoreUrl::parse(store_url).unwrap(), + Path::from(path), + ), }; let cached = cache.get(&key); - assert_eq!(cached.is_some(), name == "sibling_table", "{name}"); + assert_eq!( + cached.is_some(), + matches!(name, "sibling_table" | "other_store_table"), + "{name}" + ); } let batches = ctx diff --git a/datafusion/core/tests/parquet/page_pruning.rs b/datafusion/core/tests/parquet/page_pruning.rs index 70052ec6b7d60..48b03ac0046b8 100644 --- a/datafusion/core/tests/parquet/page_pruning.rs +++ b/datafusion/core/tests/parquet/page_pruning.rs @@ -67,7 +67,12 @@ async fn get_parquet_exec( }; let schema = ParquetFormat::default() - .infer_schema(state, &store, std::slice::from_ref(&meta)) + .infer_schema( + state, + &store, + &object_store_url, + std::slice::from_ref(&meta), + ) .await .unwrap(); diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index f31527b3caedb..3c8a8a6e046a1 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -58,6 +58,7 @@ use datafusion_datasource::file_format::{FileFormat, FileFormatFactory}; use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::write::demux::DemuxedStreamReceiver; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; use futures::StreamExt; @@ -179,6 +180,7 @@ impl FileFormat for ArrowFormat { &self, _state: &dyn Session, store: &Arc, + _object_store_url: &ObjectStoreUrl, objects: &[ObjectMeta], ) -> Result { let mut schemas = Vec::with_capacity(objects.len()); @@ -217,6 +219,7 @@ impl FileFormat for ArrowFormat { &self, _state: &dyn Session, _store: &Arc, + _object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, _object: &ObjectMeta, ) -> Result { @@ -719,6 +722,7 @@ mod tests { .infer_schema( &state, &(store.clone() as Arc), + &ObjectStoreUrl::parse("memory://test")?, std::slice::from_ref(&object_meta), ) .await?; @@ -758,6 +762,7 @@ mod tests { .infer_schema( &state, &(store.clone() as Arc), + &ObjectStoreUrl::parse("memory://test")?, std::slice::from_ref(&object_meta), ) .await; diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs index 6aa28667f48d0..0d5d83be01f11 100644 --- a/datafusion/datasource-avro/src/file_format.rs +++ b/datafusion/datasource-avro/src/file_format.rs @@ -48,6 +48,7 @@ use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::{ ObjectWriterBuilder, SharedBuffer, get_writer_schema, }; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; @@ -132,6 +133,7 @@ impl FileFormat for AvroFormat { &self, _state: &dyn Session, store: &Arc, + _object_store_url: &ObjectStoreUrl, objects: &[ObjectMeta], ) -> Result { let mut schemas = Vec::with_capacity(objects.len()); @@ -157,6 +159,7 @@ impl FileFormat for AvroFormat { &self, _state: &dyn Session, _store: &Arc, + _object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, _object: &ObjectMeta, ) -> Result { diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 15e24df4bf7c7..a1498c8b2cf4a 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -49,6 +49,7 @@ use datafusion_datasource::sink::{DataSink, DataSinkExec}; use datafusion_datasource::write::BatchSerializer; use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; @@ -377,6 +378,7 @@ impl FileFormat for CsvFormat { &self, state: &dyn Session, store: &Arc, + _object_store_url: &ObjectStoreUrl, objects: &[ObjectMeta], ) -> Result { let mut schemas = vec![]; @@ -424,6 +426,7 @@ impl FileFormat for CsvFormat { &self, _state: &dyn Session, _store: &Arc, + _object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, _object: &ObjectMeta, ) -> Result { diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index aef89560a3b5f..3a0badb3b159c 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -52,6 +52,7 @@ use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::write::BatchSerializer; use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::orchestration::spawn_writer_tasks_and_join; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; @@ -256,6 +257,7 @@ impl FileFormat for JsonFormat { &self, _state: &dyn Session, store: &Arc, + _object_store_url: &ObjectStoreUrl, objects: &[ObjectMeta], ) -> Result { let mut schemas = Vec::new(); @@ -335,6 +337,7 @@ impl FileFormat for JsonFormat { &self, _state: &dyn Session, _store: &Arc, + _object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, _object: &ObjectMeta, ) -> Result { diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 76d0e83134793..2aff020842f83 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -55,6 +55,7 @@ use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::write::get_writer_schema; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement}; use datafusion_physical_plan::ExecutionPlan; @@ -358,6 +359,7 @@ impl FileFormat for ParquetFormat { &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, objects: &[ObjectMeta], ) -> Result { let coerce_int96 = self.coerce_int96().map(arrow::datatypes::TimeUnit::from); @@ -382,7 +384,10 @@ impl FileFormat for ParquetFormat { let meta = DFParquetMetadata::new(store.as_ref(), object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_file_metadata_cache( + Some(Arc::clone(&file_metadata_cache)), + object_store_url.clone(), + ) .with_coerce_int96(coerce_int96) .with_coerce_int96_tz(coerce_int96_tz.clone()) .with_enable_rle_to_dictionary( @@ -471,6 +476,7 @@ impl FileFormat for ParquetFormat { &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, object: &ObjectMeta, ) -> Result { @@ -482,7 +488,7 @@ impl FileFormat for ParquetFormat { DFParquetMetadata::new(store, object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) + .with_file_metadata_cache(Some(file_metadata_cache), object_store_url.clone()) .fetch_statistics(&table_schema) .await } @@ -491,6 +497,7 @@ impl FileFormat for ParquetFormat { &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, object: &ObjectMeta, ) -> Result> { @@ -502,7 +509,7 @@ impl FileFormat for ParquetFormat { let metadata = DFParquetMetadata::new(store, object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) + .with_file_metadata_cache(Some(file_metadata_cache), object_store_url.clone()) .fetch_metadata() .await?; crate::metadata::ordering_from_parquet_metadata(&metadata, &table_schema) @@ -512,6 +519,7 @@ impl FileFormat for ParquetFormat { &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, object: &ObjectMeta, ) -> Result { @@ -523,7 +531,7 @@ impl FileFormat for ParquetFormat { let metadata = DFParquetMetadata::new(store, object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) + .with_file_metadata_cache(Some(file_metadata_cache), object_store_url.clone()) .fetch_metadata() .await?; let statistics = DFParquetMetadata::statistics_from_parquet_metadata( @@ -561,8 +569,11 @@ impl FileFormat for ParquetFormat { let store = state .runtime_env() .object_store(conf.object_store_url.clone())?; - let cached_parquet_read_factory = - Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); + let cached_parquet_read_factory = Arc::new(CachedParquetFileReaderFactory::new( + store, + conf.object_store_url.clone(), + metadata_cache, + )); source = source.with_parquet_file_reader_factory(cached_parquet_read_factory); if let Some(metadata_size_hint) = metadata_size_hint { diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index 0cfcb980a788f..e1019ca80d064 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -32,8 +32,9 @@ use datafusion_common::{ internal_datafusion_err, }; use datafusion_execution::cache::cache_manager::{ - CachedFileMetadataEntry, FileMetadata, FileMetadataCache, + CachedFileMetadataEntry, FileMetadata, FileMetadataCache, StoreScopedPath, }; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_functions_aggregate_common::min_max::{MaxAccumulator, MinAccumulator}; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; @@ -140,8 +141,8 @@ pub struct DFParquetMetadata<'a> { /// [`ParquetMetaDataReader::with_decryption_properties`]. decryption_properties: Option>, /// Optional cache of previously fetched [`ParquetMetaData`], keyed by - /// file location. - file_metadata_cache: Option>, + /// object store URL and file location. + file_metadata_cache: Option<(Arc, ObjectStoreUrl)>, /// Policy controlling whether the Parquet page index (column and offset /// indexes) is fetched, mirroring /// [`ParquetMetaDataReader::with_page_index_policy`]. @@ -209,11 +210,15 @@ impl<'a> DFParquetMetadata<'a> { /// Set an optional [`FileMetadataCache`] used to avoid re-fetching /// [`ParquetMetaData`] for files that have already been read. + /// `object_store_url` identifies the store supplying this file, so equal + /// paths in different stores cannot reuse each other's metadata. pub fn with_file_metadata_cache( mut self, file_metadata_cache: Option>, + object_store_url: ObjectStoreUrl, ) -> Self { - self.file_metadata_cache = file_metadata_cache; + self.file_metadata_cache = + file_metadata_cache.map(|cache| (cache, object_store_url)); self } @@ -280,8 +285,12 @@ impl<'a> DFParquetMetadata<'a> { let page_index_policy = self.effective_page_index_policy(cache_metadata); if cache_metadata - && let Some(file_metadata_cache) = self.file_metadata_cache.as_ref() - && let Some(cached) = file_metadata_cache.get(&self.object_meta.location) + && let Some((file_metadata_cache, object_store_url)) = + self.file_metadata_cache.as_ref() + && let Some(cached) = file_metadata_cache.get(&StoreScopedPath::new( + object_store_url.clone(), + self.object_meta.location.clone(), + )) && cached.is_valid_for(self.object_meta) && let Some(cached_parquet) = cached .file_metadata @@ -339,14 +348,17 @@ impl<'a> DFParquetMetadata<'a> { } /// Store `metadata` in the configured [`FileMetadataCache`], keyed by - /// the file's location. + /// the object store URL and file's location. /// /// This is a no-op unless a cache has been configured via /// [`Self::with_file_metadata_cache`]. fn cache_metadata(&self, metadata: Arc) -> Result<()> { - if let Some(file_metadata_cache) = &self.file_metadata_cache { + if let Some((file_metadata_cache, object_store_url)) = &self.file_metadata_cache { file_metadata_cache.put( - &self.object_meta.location, + &StoreScopedPath::new( + object_store_url.clone(), + self.object_meta.location.clone(), + ), CachedFileMetadataEntry::new( self.object_meta.clone(), Arc::new(CachedParquetMetaData::new(metadata)), @@ -1267,6 +1279,111 @@ mod tests { use arrow::array::Int32Array; use arrow::compute::SortOptions; use arrow::datatypes::Field; + use arrow::record_batch::RecordBatch; + use datafusion_execution::cache::default_cache::DefaultCache; + use object_store::ObjectStoreExt; + use object_store::memory::InMemory; + use parquet::arrow::ArrowWriter; + + async fn write_cached_test_file(store: &InMemory, value: i32) -> ObjectMeta { + let batch = RecordBatch::try_from_iter(vec![( + "value", + Arc::new(Int32Array::from(vec![value])) as ArrayRef, + )]) + .unwrap(); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let path = Path::from("data.parquet"); + store.put(&path, bytes.into()).await.unwrap(); + store.head(&path).await.unwrap() + } + + #[tokio::test] + async fn metadata_cache_invalidates_changed_object_identity() { + for use_version in [false, true] { + let store = InMemory::new(); + let cache: Arc = Arc::new(DefaultCache::new(1024 * 1024)); + let url = ObjectStoreUrl::parse("s3://bucket").unwrap(); + let mut original = write_cached_test_file(&store, 100).await; + if use_version { + original.e_tag = None; + original.version = Some("v1".into()); + } + let original_reader = DFParquetMetadata::new(&store, &original) + .with_file_metadata_cache(Some(Arc::clone(&cache)), url.clone()) + .with_page_index_policy(Some(PageIndexPolicy::Skip)); + let original_metadata = original_reader.fetch_metadata().await.unwrap(); + assert!(Arc::ptr_eq( + &original_metadata, + &original_reader.fetch_metadata().await.unwrap() + )); + + let mut overwritten = write_cached_test_file(&store, 200).await; + assert_eq!(original.size, overwritten.size); + overwritten.last_modified = original.last_modified; + if use_version { + overwritten.e_tag = None; + overwritten.version = Some("v2".into()); + } else { + assert_ne!(original.e_tag, overwritten.e_tag); + } + let reader = DFParquetMetadata::new(&store, &overwritten) + .with_file_metadata_cache(Some(Arc::clone(&cache)), url) + .with_page_index_policy(Some(PageIndexPolicy::Skip)); + let metadata = reader.fetch_metadata().await.unwrap(); + assert!(!Arc::ptr_eq(&original_metadata, &metadata)); + let ParquetStatistics::Int32(stats) = + metadata.row_group(0).column(0).statistics().unwrap() + else { + panic!("expected Int32 statistics"); + }; + assert_eq!(stats.min_opt(), Some(&200)); + assert_eq!(stats.max_opt(), Some(&200)); + assert!(Arc::ptr_eq( + &metadata, + &reader.fetch_metadata().await.unwrap() + )); + assert_eq!(cache.len(), 1); + } + } + + #[tokio::test] + async fn metadata_cache_separates_object_stores() { + let first_store = InMemory::new(); + let second_store = InMemory::new(); + let meta = write_cached_test_file(&first_store, 100).await; + let second_meta = write_cached_test_file(&second_store, 200).await; + assert_eq!(meta.size, second_meta.size); + let cache: Arc = Arc::new(DefaultCache::new(1024 * 1024)); + // Use identical metadata deliberately: store identity must distinguish + // the files even when their paths and all validation fields coincide. + for _ in 0..2 { + for (store, url, expected) in [ + (&first_store, "s3://first", 100), + (&second_store, "s3://second", 200), + ] { + let metadata = DFParquetMetadata::new(store, &meta) + .with_file_metadata_cache( + Some(Arc::clone(&cache)), + ObjectStoreUrl::parse(url).unwrap(), + ) + .with_page_index_policy(Some(PageIndexPolicy::Skip)) + .fetch_metadata() + .await + .unwrap(); + let ParquetStatistics::Int32(stats) = + metadata.row_group(0).column(0).statistics().unwrap() + else { + panic!("expected Int32 statistics"); + }; + assert_eq!(stats.min_opt(), Some(&expected)); + assert_eq!(stats.max_opt(), Some(&expected)); + } + } + assert_eq!(cache.len(), 2); + } #[test] fn test_lex_ordering_to_sorting_columns_uses_writer_schema() -> Result<()> { diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index b048a23eea896..21aca47c43764 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -2156,6 +2156,7 @@ mod test { use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::morsel::{Morsel, Morselizer}; use datafusion_datasource::{PartitionedFile, TableSchema, TableSchemaBuilder}; + use datafusion_execution::cache::StoreScopedPath; use datafusion_execution::cache::cache_manager::{ CachedFileMetadataEntry, FileMetadataCache, }; @@ -4042,6 +4043,7 @@ mod test { .with_parquet_file_reader_factory(Arc::new( CachedParquetFileReaderFactory::new( Arc::clone(&store), + ObjectStoreUrl::local_filesystem(), Arc::clone(&cache), ), )) @@ -4094,7 +4096,10 @@ mod test { .await; let cache: Arc = Arc::new(DefaultCache::new(1024 * 1024)); let original = DFParquetMetadata::new(store.as_ref(), &file.object_meta) - .with_file_metadata_cache(Some(Arc::clone(&cache))) + .with_file_metadata_cache( + Some(Arc::clone(&cache)), + ObjectStoreUrl::local_filesystem(), + ) .with_page_index_policy(Some(PageIndexPolicy::Optional)) .fetch_metadata() .await @@ -4110,6 +4115,7 @@ mod test { .with_parquet_file_reader_factory(Arc::new( CachedParquetFileReaderFactory::new( Arc::clone(&store), + ObjectStoreUrl::local_filesystem(), Arc::clone(&cache), ), )) @@ -4139,7 +4145,12 @@ mod test { prepared.loaded.reader_metadata.metadata() )); } - let cached = cache.get(&file.object_meta.location).unwrap(); + let cached = cache + .get(&StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + file.object_meta.location.clone(), + )) + .unwrap(); let cached = cached .file_metadata .as_any() @@ -4237,6 +4248,7 @@ mod test { .with_parquet_file_reader_factory(Arc::new( CachedParquetFileReaderFactory::new( Arc::clone(&store), + ObjectStoreUrl::local_filesystem(), Arc::clone(&cache), ), )) @@ -4266,7 +4278,12 @@ mod test { }; assert_eq!(pruning_metrics.pruned(), 2); assert!( - cache.get(&file.object_meta.location).is_none(), + cache + .get(&StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + file.object_meta.location.clone(), + )) + .is_none(), "encrypted metadata must not enter the shared cache" ); } @@ -5232,10 +5249,12 @@ mod test { #[tokio::test] async fn test_page_index_skipped_with_cached_reader_factory() { let store = Arc::new(InMemory::new()) as Arc; - let metadata_cache: Arc = - Arc::new(DefaultCache::::new( - 64 * 1024 * 1024, - )); + let metadata_cache: Arc = Arc::new(DefaultCache::< + StoreScopedPath, + CachedFileMetadataEntry, + >::new( + 64 * 1024 * 1024 + )); let values: Vec = (1..=100).collect(); let batch = record_batch!(( "a", @@ -5273,6 +5292,7 @@ mod test { .with_parquet_file_reader_factory(Arc::new( CachedParquetFileReaderFactory::new( Arc::clone(&store), + ObjectStoreUrl::local_filesystem(), Arc::clone(&metadata_cache), ), )) @@ -5284,7 +5304,10 @@ mod test { assert_eq!(counter_metric_value(&metrics, "page_index_load_skipped"), 1); let cached = metadata_cache - .get(&Path::from("test.parquet")) + .get(&StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + Path::from("test.parquet"), + )) .expect("metadata cache should contain the file"); let extra_info = cached.file_metadata.extra_info(); let page_index_cached = extra_info.get("page_index").map(String::as_str); diff --git a/datafusion/datasource-parquet/src/reader.rs b/datafusion/datasource-parquet/src/reader.rs index 06ab40fdd27cd..3902fbe21b6a2 100644 --- a/datafusion/datasource-parquet/src/reader.rs +++ b/datafusion/datasource-parquet/src/reader.rs @@ -26,6 +26,7 @@ use datafusion_common::HashMap; use datafusion_datasource::PartitionedFile; use datafusion_execution::cache::cache_manager::FileMetadata; use datafusion_execution::cache::cache_manager::FileMetadataCache; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use futures::FutureExt; use futures::TryFutureExt; @@ -143,16 +144,20 @@ impl ParquetFileReaderFactory for DefaultParquetFileReaderFactory { #[derive(Debug)] pub struct CachedParquetFileReaderFactory { store: Arc, + object_store_url: ObjectStoreUrl, metadata_cache: Arc, } impl CachedParquetFileReaderFactory { + /// Create a factory whose metadata cache entries are scoped to `object_store_url`. pub fn new( store: Arc, + object_store_url: ObjectStoreUrl, metadata_cache: Arc, ) -> Self { Self { store, + object_store_url, metadata_cache, } } @@ -178,7 +183,10 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { partitioned_file, ) .with_metadata_hint(metadata_size_hint) - .with_metadata_cache(Some(Arc::clone(&self.metadata_cache))); + .with_metadata_cache( + Some(Arc::clone(&self.metadata_cache)), + self.object_store_url.clone(), + ); Ok(Box::new(reader)) } @@ -202,7 +210,7 @@ pub struct ParquetFileReader { file_metrics: ParquetFileMetrics, store: Arc, partitioned_file: PartitionedFile, - metadata_cache: Option>, + metadata_cache: Option<(Arc, ObjectStoreUrl)>, metadata_size_hint: Option, } @@ -239,12 +247,13 @@ impl ParquetFileReader { &self.partitioned_file } - /// Set the [`FileMetadataCache`] for this reader + /// Set the [`FileMetadataCache`] and the URL identifying this reader's store. pub fn with_metadata_cache( mut self, metadata_cache: Option>, + object_store_url: ObjectStoreUrl, ) -> Self { - self.metadata_cache = metadata_cache; + self.metadata_cache = metadata_cache.map(|cache| (cache, object_store_url)); self } @@ -306,9 +315,13 @@ impl AsyncFileReader for ParquetFileReader { let page_index_policy = options.map(|o| o.column_index_policy()); - DFParquetMetadata::new(&self.store, &object_meta) - .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(metadata_cache) + let mut metadata = DFParquetMetadata::new(&self.store, &object_meta) + .with_decryption_properties(file_decryption_properties); + if let Some((cache, object_store_url)) = metadata_cache { + metadata = + metadata.with_file_metadata_cache(Some(cache), object_store_url); + } + metadata .with_metadata_size_hint(self.metadata_size_hint) .with_page_index_policy(page_index_policy) .fetch_metadata() diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index fbb6265847afc..eb913fc92e100 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -1253,14 +1253,17 @@ impl ParquetSource { let store = ctx .task_ctx() .runtime_env() - .object_store(object_store_url)?; + .object_store(&object_store_url)?; let metadata_cache = ctx .task_ctx() .runtime_env() .cache_manager .get_file_metadata_cache(); - let reader_factory = - Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); + let reader_factory = Arc::new(CachedParquetFileReaderFactory::new( + store, + object_store_url.clone(), + metadata_cache, + )); let mut source = ParquetSource::new(table_schema) .with_parquet_file_reader_factory(reader_factory) diff --git a/datafusion/datasource/src/file_format.rs b/datafusion/datasource/src/file_format.rs index 2c2a94e8cc4b4..55dc6196118dc 100644 --- a/datafusion/datasource/src/file_format.rs +++ b/datafusion/datasource/src/file_format.rs @@ -33,6 +33,7 @@ use datafusion_common::file_options::file_type::FileType; use datafusion_common::{ GetExt, Result, SchemaError, Statistics, internal_err, not_impl_err, schema_err, }; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::LexRequirement; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::ExecutionPlan; @@ -117,10 +118,13 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// be analysed up to a given number of records or files (as specified in the /// format config) then give the estimated common schema. This might fail if /// the files have schemas that cannot be merged. + /// + /// `object_store_url` identifies `store` and scopes any cached file metadata. async fn infer_schema( &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, objects: &[ObjectMeta], ) -> Result; @@ -130,11 +134,14 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// `table_schema` is the (combined) schema of the overall table /// and may be a superset of the schema contained in this file. /// + /// `object_store_url` identifies `store` and scopes any cached file metadata. + /// /// TODO: should the file source return statistics for only columns referred to in the table schema? async fn infer_stats( &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, object: &ObjectMeta, ) -> Result; @@ -148,10 +155,13 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// and may be a superset of the schema contained in this file. /// /// The default implementation returns `Ok(None)`. + /// + /// `object_store_url` identifies `store` and scopes any cached file metadata. async fn infer_ordering( &self, _state: &dyn Session, _store: &Arc, + _object_store_url: &ObjectStoreUrl, _table_schema: SchemaRef, _object: &ObjectMeta, ) -> Result> { @@ -166,18 +176,27 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// /// The default implementation calls both methods separately. File formats /// that can extract both from a single read should override this method. + /// + /// `object_store_url` identifies `store` and scopes any cached file metadata. async fn infer_stats_and_ordering( &self, state: &dyn Session, store: &Arc, + object_store_url: &ObjectStoreUrl, table_schema: SchemaRef, object: &ObjectMeta, ) -> Result { let statistics = self - .infer_stats(state, store, Arc::clone(&table_schema), object) + .infer_stats( + state, + store, + object_store_url, + Arc::clone(&table_schema), + object, + ) .await?; let ordering = self - .infer_ordering(state, store, table_schema, object) + .infer_ordering(state, store, object_store_url, table_schema, object) .await?; Ok(FileMeta { statistics, diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 50d277803e38e..f2fe4c969c315 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use datafusion_common::{DataFusionError, Result, TableReference}; use datafusion_execution::cache::cache_manager::CachedFileList; -use datafusion_execution::cache::cache_manager::TableScopedPath; +use datafusion_execution::cache::cache_manager::{StoreScopedPath, TableScopedPath}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_session::Session; @@ -267,6 +267,7 @@ impl ListingTableUrl { list_with_cache( ctx, store, + &self.object_store(), self.table_ref.as_ref(), &self.prefix, prefix.as_ref(), @@ -283,6 +284,7 @@ impl ListingTableUrl { list_with_cache( ctx, store, + &self.object_store(), self.table_ref.as_ref(), &self.prefix, prefix.as_ref(), @@ -364,6 +366,8 @@ impl ListingTableUrl { /// # Arguments /// * `ctx` - The session context /// * `store` - The object store to list from +/// * `object_store_url` - The registered store URL used to scope cache entries +/// * `table_ref` - Optional table reference used for table-level cache invalidation /// * `table_base_path` - The table's base path (the stable cache key) /// * `prefix` - Optional prefix relative to table base for filtering results /// @@ -373,6 +377,7 @@ impl ListingTableUrl { async fn list_with_cache<'b>( ctx: &'b dyn Session, store: &'b dyn ObjectStore, + object_store_url: &ObjectStoreUrl, table_ref: Option<&TableReference>, table_base_path: &Path, prefix: Option<&Path>, @@ -397,11 +402,17 @@ async fn list_with_cache<'b>( let table_scoped_base_path = TableScopedPath { table: table_ref.cloned(), - path: table_base_path.clone(), + store_path: StoreScopedPath::new( + object_store_url.clone(), + table_base_path.clone(), + ), }; let table_scoped_list_path = TableScopedPath { table: table_ref.cloned(), - path: full_prefix.clone(), + store_path: StoreScopedPath::new( + object_store_url.clone(), + full_prefix.clone(), + ), }; let vec = if let Some(cached) = cache.get(&table_scoped_base_path) { @@ -854,7 +865,10 @@ mod tests { create_file(&store, "/table/year=2024/month=12/data5.parquet").await; // Session WITHOUT cache - let session_no_cache = MockSession::new(); + let runtime_no_cache = RuntimeEnvBuilder::new() + .with_object_list_cache_limit(0) + .build_arc()?; + let session_no_cache = MockSession::with_runtime_env(runtime_no_cache); // Session WITH cache - use RuntimeEnvBuilder with cache limit (no TTL needed for this test) let runtime_with_cache = RuntimeEnvBuilder::new() @@ -862,33 +876,43 @@ mod tests { .build_arc()?; let session_with_cache = MockSession::with_runtime_env(runtime_with_cache); - // Test cases: (url, prefix, description) + // Test cases: (url, prefix, description, expected file count) let test_cases = vec![ - ("/table/", None, "full table listing"), + ("mem://bucket/table/", None, "full table listing", 5), ( - "/table/", + "mem://bucket/table/", Some(Path::from("year=2023")), "single partition filter", + 2, ), ( - "/table/", + "mem://bucket/table/", Some(Path::from("year=2024")), "different partition filter", + 3, ), ( - "/table/", + "mem://bucket/table/", Some(Path::from("year=2024/month=06")), "nested partition filter", + 1, ), ( - "/table/", + "mem://bucket/table/", Some(Path::from("year=2025")), "non-existent partition", + 0, ), ]; - for (url_str, prefix, description) in test_cases { + for (url_str, prefix, description, expected_count) in test_cases { let url = ListingTableUrl::parse(url_str)?; + session_with_cache + .runtime_env + .cache_manager + .get_list_files_cache() + .unwrap() + .clear(); // Get results WITHOUT cache (sorted for comparison) let mut results_no_cache: Vec = url @@ -900,6 +924,7 @@ mod tests { .map(|m| m.location.to_string()) .collect(); results_no_cache.sort(); + assert_eq!(results_no_cache.len(), expected_count, "{description}"); // Get results WITH cache (first call - cache miss, sorted for comparison) let mut results_with_cache_miss: Vec = url @@ -942,6 +967,40 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_listing_cache_separates_object_stores() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let first = MockObjectStore::new(); + let second = MockObjectStore::new(); + create_file(&first, "table/first.parquet").await; + create_file(&second, "table/second.parquet").await; + let runtime = RuntimeEnvBuilder::new() + .with_object_list_cache_limit(1024 * 1024) + .build_arc()?; + let session = MockSession::with_runtime_env(runtime); + + // Reuse both listings to verify that each store has its own cache entry. + for _ in 0..2 { + for (bucket, store, expected) in [ + ("first", &first, "table/first.parquet"), + ("second", &second, "table/second.parquet"), + ] { + let url = ListingTableUrl::parse(format!("s3://{bucket}/table/"))?; + let files = url + .list_prefixed_files(&session, store, None, "parquet") + .await? + .try_collect::>() + .await?; + assert_eq!(files.len(), 1); + assert_eq!(files[0].location.as_ref(), expected); + } + } + assert_eq!(first.list_prefixes(), vec![Some(Path::from("table"))]); + assert_eq!(second.list_prefixes(), vec![Some(Path::from("table"))]); + Ok(()) + } + #[tokio::test] async fn test_cache_miss_with_prefix_lists_prefixed_path() -> Result<()> { use datafusion_execution::runtime_env::RuntimeEnvBuilder; @@ -955,7 +1014,7 @@ mod tests { .with_object_list_cache_limit(1024 * 1024) .build_arc()?; let session = MockSession::with_runtime_env(runtime); - let url = ListingTableUrl::parse("/table/")?; + let url = ListingTableUrl::parse("mem://bucket/table/")?; let prefix = Path::from("year=2024/month=06"); let results: Vec = url @@ -1009,7 +1068,7 @@ mod tests { .options_mut() .execution .listing_table_ignore_subdirectory = false; - let url = ListingTableUrl::parse("/sales/")?; + let url = ListingTableUrl::parse("mem://bucket/sales/")?; let q1_results: Vec = url .list_prefixed_files( @@ -1086,7 +1145,7 @@ mod tests { .build_arc()?; let session = MockSession::with_runtime_env(runtime); - let url = ListingTableUrl::parse("/sales/")?; + let url = ListingTableUrl::parse("mem://bucket/sales/")?; // First: query full table (populates cache) let full_results: Vec = url diff --git a/datafusion/execution/src/cache/cache_manager.rs b/datafusion/execution/src/cache/cache_manager.rs index bbbcaa4c77559..1ba880c5259e8 100644 --- a/datafusion/execution/src/cache/cache_manager.rs +++ b/datafusion/execution/src/cache/cache_manager.rs @@ -16,7 +16,9 @@ // under the License. use crate::cache::default_cache::DefaultCache; -pub use crate::cache::{Cache, CacheValue, SchemaFingerprint, TableScopedPath}; +pub use crate::cache::{ + Cache, CacheValue, SchemaFingerprint, StoreScopedPath, TableScopedPath, +}; use datafusion_common::HashMap; use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use datafusion_common::{Result, Statistics}; @@ -40,7 +42,7 @@ pub const DEFAULT_METADATA_CACHE_LIMIT: usize = 50 * 1024 * 1024; // 50M /// A cache for file statistics and orderings. /// /// This cache stores [`CachedFileMetadata`] which includes: -/// - File metadata for validation (size, last_modified) +/// - File metadata for validation (size, last_modified, e_tag, version) /// - Statistics for the file /// - Ordering information for the file /// @@ -75,6 +77,7 @@ pub type ListFilesCache = dyn Cache; /// /// This cache stores per-file metadata in the form of [`CachedFileMetadataEntry`], /// which includes the [`ObjectMeta`] for validation. +/// Entries are keyed by the object store URL and the path within that store. /// /// For example, the built in [`ListingTable`] uses this cache to avoid parsing /// Parquet footers multiple times for the same file. @@ -87,7 +90,7 @@ pub type ListFilesCache = dyn Cache; /// See [`crate::runtime_env::RuntimeEnv`] for more details. /// /// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html -pub type FileMetadataCache = dyn Cache; +pub type FileMetadataCache = dyn Cache; /// Cached metadata for a file, including statistics and ordering. /// @@ -95,7 +98,7 @@ pub type FileMetadataCache = dyn Cache; /// the `file_schema` fingerprint, cached statistics, and ordering information. #[derive(Debug, Clone, PartialEq, Eq)] pub struct CachedFileMetadata { - /// File metadata used for cache validation (size, last_modified). + /// File metadata used for cache validation (size, last_modified, e_tag, version). pub meta: ObjectMeta, /// Fingerprint of the `file_schema` used to compute `statistics`. pub schema_fingerprint: Arc, @@ -123,20 +126,42 @@ impl CachedFileMetadata { /// Check if this cached entry is still valid for the given metadata. /// - /// Returns true if the file size, last modified time, and schema match. + /// Validates the file using a shared strong ETag, otherwise a shared version, + /// otherwise size and last modified time. The schema must always match. pub fn is_valid_for( &self, current_meta: &ObjectMeta, current_schema_fingerprint: &Arc, ) -> bool { - self.meta.size == current_meta.size - && self.meta.last_modified == current_meta.last_modified + file_metadata_matches(&self.meta, current_meta) && (Arc::ptr_eq(&self.schema_fingerprint, current_schema_fingerprint) || self.schema_fingerprint.as_ref() == current_schema_fingerprint.as_ref()) } } +/// Prefer content identity, then version identity, falling back to size and time +/// only when neither identifier can be compared. +fn file_metadata_matches(cached: &ObjectMeta, current: &ObjectMeta) -> bool { + // The HTTP object store passes through weak ETags, which do not guarantee + // identical bytes. Treat them as missing and try the remaining validators. + if let (Some(cached), Some(current)) = + (cached.e_tag.as_deref(), current.e_tag.as_deref()) + && !cached.starts_with("W/") + && !current.starts_with("W/") + { + // Compare opaque values verbatim. Some stores (e.g. Azure) return + // unquoted ETags in listings and quoted ETags in HEAD responses. This + // can cause conservative cache misses, but cannot reuse stale metadata. + return cached == current; + } + if let Some((cached, current)) = cached.version.as_ref().zip(current.version.as_ref()) + { + return cached == current; + } + cached.size == current.size && cached.last_modified == current.last_modified +} + impl CacheValue for CachedFileMetadata { fn size(&self) -> usize { DFHeapSize::heap_size(self, &mut DFHeapSizeCtx::default()) @@ -259,7 +284,7 @@ pub trait FileMetadata: Any + Send + Sync { /// Cached file metadata entry with validation information. #[derive(Clone)] pub struct CachedFileMetadataEntry { - /// File metadata used for cache validation (size, last_modified). + /// File metadata used for cache validation (size, last_modified, e_tag, version). pub meta: ObjectMeta, /// The cached file metadata. pub file_metadata: Arc, @@ -281,9 +306,11 @@ impl CachedFileMetadataEntry { } /// Check if this cached entry is still valid for the given metadata. + /// + /// Uses a shared strong ETag, otherwise a shared version, otherwise the file + /// size and last modified time. Weak ETags are treated as missing. pub fn is_valid_for(&self, current_meta: &ObjectMeta) -> bool { - self.meta.size == current_meta.size - && self.meta.last_modified == current_meta.last_modified + file_metadata_matches(&self.meta, current_meta) } } @@ -512,6 +539,183 @@ impl CacheManagerConfig { #[cfg(test)] mod tests { use super::*; + use arrow::datatypes::Schema; + + struct TestFileMetadata; + + impl FileMetadata for TestFileMetadata { + fn as_any(&self) -> &dyn Any { + self + } + + fn memory_size(&self) -> usize { + 0 + } + + fn extra_info(&self) -> HashMap { + HashMap::new() + } + } + + fn test_object_meta() -> ObjectMeta { + ObjectMeta { + location: Path::from("data.parquet"), + last_modified: chrono::DateTime::UNIX_EPOCH, + size: 100, + e_tag: None, + version: None, + } + } + + fn assert_metadata_validity( + cached_meta: &ObjectMeta, + current_meta: &ObjectMeta, + expected: bool, + ) { + let schema = Schema::empty(); + let fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); + let cached_statistics = CachedFileMetadata::new( + cached_meta.clone(), + Arc::clone(&fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + let cached_metadata = + CachedFileMetadataEntry::new(cached_meta.clone(), Arc::new(TestFileMetadata)); + + assert_eq!( + cached_statistics.is_valid_for(current_meta, &fingerprint), + expected, + "statistics: cached={cached_meta:?}, current={current_meta:?}" + ); + assert_eq!( + cached_metadata.is_valid_for(current_meta), + expected, + "metadata: cached={cached_meta:?}, current={current_meta:?}" + ); + } + + #[test] + fn test_metadata_validation_etag_and_version() { + let etags = [ + (Some("same"), Some("same"), Some(true)), + (Some("old"), Some("new"), Some(false)), + (Some("old"), None, None), + (None, Some("new"), None), + (None, None, None), + (Some("W/\"same\""), Some("W/\"same\""), None), + (Some("W/\"old\""), Some("W/\"new\""), None), + (Some("W/\"same\""), Some("\"same\""), None), + (Some("\"same\""), Some("W/\"same\""), None), + // Azure listing and HEAD values remain opaque, not normalized. + ( + Some("0x8D93C7D4629C227"), + Some("\"0x8D93C7D4629C227\""), + Some(false), + ), + ]; + let versions = [ + (Some("same"), Some("same"), Some(true)), + (Some("old"), Some("new"), Some(false)), + (Some("old"), None, None), + (None, Some("new"), None), + (None, None, None), + ]; + + for (cached_etag, current_etag, etag_matches) in etags { + for (cached_version, current_version, version_matches) in versions { + let cached_meta = ObjectMeta { + e_tag: cached_etag.map(str::to_owned), + version: cached_version.map(str::to_owned), + ..test_object_meta() + }; + let current_meta = ObjectMeta { + e_tag: current_etag.map(str::to_owned), + version: current_version.map(str::to_owned), + ..test_object_meta() + }; + assert_metadata_validity( + &cached_meta, + ¤t_meta, + etag_matches.or(version_matches).unwrap_or(true), + ); + } + } + } + + #[test] + fn test_metadata_validation_size_and_last_modified() { + // Missing, weak, and one-sided identifiers all use the fallback. + for (cached_etag, current_etag) in [ + (None, None), + (Some("W/\"same\""), Some("W/\"same\"")), + (Some("same"), None), + (None, Some("same")), + (Some("W/\"same\""), Some("\"same\"")), + (Some("\"same\""), Some("W/\"same\"")), + ] { + let cached_meta = ObjectMeta { + e_tag: cached_etag.map(str::to_owned), + ..test_object_meta() + }; + let current_meta = ObjectMeta { + e_tag: current_etag.map(str::to_owned), + ..test_object_meta() + }; + assert_metadata_validity(&cached_meta, ¤t_meta, true); + let mut current_meta = current_meta; + current_meta.size += 1; + assert_metadata_validity(&cached_meta, ¤t_meta, false); + + current_meta.size = cached_meta.size; + current_meta.last_modified += chrono::Duration::seconds(1); + assert_metadata_validity(&cached_meta, ¤t_meta, false); + } + } + + #[test] + fn test_metadata_validation_identifiers_override_size_and_time() { + for (etag, cached_version, current_version) in [ + (Some("same"), Some("old"), Some("new")), + (None, Some("same"), Some("same")), + (Some("W/\"same\""), Some("same"), Some("same")), + ] { + let cached_meta = ObjectMeta { + e_tag: etag.map(str::to_owned), + version: cached_version.map(str::to_owned), + ..test_object_meta() + }; + let current_meta = ObjectMeta { + size: cached_meta.size + 1, + last_modified: cached_meta.last_modified + chrono::Duration::seconds(1), + version: current_version.map(str::to_owned), + ..cached_meta.clone() + }; + assert_metadata_validity(&cached_meta, ¤t_meta, true); + } + } + + #[test] + fn test_matching_etag_does_not_override_schema_validation() { + use arrow::datatypes::{DataType, Field}; + + let meta = ObjectMeta { + e_tag: Some("same".into()), + ..test_object_meta() + }; + let original = Schema::new(vec![Field::new("value", DataType::Int32, false)]); + let changed = Schema::new(vec![Field::new("value", DataType::Int64, false)]); + let cached = CachedFileMetadata::new( + meta.clone(), + Arc::new(SchemaFingerprint::from_schema(&original)), + Arc::new(Statistics::new_unknown(&original)), + None, + ); + assert!( + !cached + .is_valid_for(&meta, &Arc::new(SchemaFingerprint::from_schema(&changed))) + ); + } /// Test to verify that TTL is preserved when not explicitly set in config. /// This fixes issue #19396 where TTL was being unset from DefaultListFilesCache diff --git a/datafusion/execution/src/cache/default_cache.rs b/datafusion/execution/src/cache/default_cache.rs index bfe326f3a47e1..25799ebe1782b 100644 --- a/datafusion/execution/src/cache/default_cache.rs +++ b/datafusion/execution/src/cache/default_cache.rs @@ -310,7 +310,8 @@ mod tests { use crate::cache::default_cache::TimeProvider; use crate::cache::{Cache, CacheEntryInfo}; use crate::cache::{CacheKey, CacheValue}; - use crate::cache::{SchemaFingerprint, TableScopedPath}; + use crate::cache::{SchemaFingerprint, StoreScopedPath, TableScopedPath}; + use crate::object_store::ObjectStoreUrl; use arrow::array::{Int32Array, ListArray, RecordBatch}; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; @@ -366,6 +367,56 @@ mod tests { } } + fn metadata_cache_key(path: &Path) -> StoreScopedPath { + StoreScopedPath::new(ObjectStoreUrl::local_filesystem(), path.clone()) + } + + #[test] + fn test_file_metadata_cache_separates_object_stores() { + let cache = DefaultCache::new(1024); + let meta = create_test_object_meta("data.parquet", 100); + let first_key = StoreScopedPath::new( + ObjectStoreUrl::parse("s3://first").unwrap(), + meta.location.clone(), + ); + let second_key = StoreScopedPath::new( + ObjectStoreUrl::parse("s3://second").unwrap(), + meta.location.clone(), + ); + let first: Arc = Arc::new(TestFileMetadata { + metadata: "first".to_owned(), + }); + let second: Arc = Arc::new(TestFileMetadata { + metadata: "second".to_owned(), + }); + + cache.put( + &first_key, + CachedFileMetadataEntry::new(meta.clone(), Arc::clone(&first)), + ); + assert!(cache.get(&second_key).is_none()); + cache.put( + &second_key, + CachedFileMetadataEntry::new(meta, Arc::clone(&second)), + ); + assert_eq!(cache.len(), 2); + assert!(Arc::ptr_eq( + &cache.get(&first_key).unwrap().file_metadata, + &first + )); + assert!(Arc::ptr_eq( + &cache.get(&second_key).unwrap().file_metadata, + &second + )); + assert_eq!( + cache.memory_used(), + first_key.size() + + second_key.size() + + first.memory_size() + + second.memory_size() + ); + } + #[test] fn test_default_file_metadata_cache() { let object_meta = create_test_object_meta("test", 1024); @@ -377,50 +428,62 @@ mod tests { let cache = DefaultCache::new(1024 * 1024); // Cache miss - assert!(cache.get(&object_meta.location).is_none()); + assert!( + cache + .get(&metadata_cache_key(&object_meta.location)) + .is_none() + ); // Put a value let cached_entry = CachedFileMetadataEntry::new(object_meta.clone(), Arc::clone(&metadata)); - cache.put(&object_meta.location, cached_entry); + cache.put(&metadata_cache_key(&object_meta.location), cached_entry); // Verify the cached value - assert!(cache.contains_key(&object_meta.location)); - let result = cache.get(&object_meta.location).unwrap(); + assert!(cache.contains_key(&metadata_cache_key(&object_meta.location))); + let result = cache + .get(&metadata_cache_key(&object_meta.location)) + .unwrap(); let test_file_metadata = Arc::downcast::(result.file_metadata); assert!(test_file_metadata.is_ok()); assert_eq!(test_file_metadata.unwrap().metadata, "retrieved_metadata"); // Cache hit - check validation - let result2 = cache.get(&object_meta.location).unwrap(); + let result2 = cache + .get(&metadata_cache_key(&object_meta.location)) + .unwrap(); assert!(result2.is_valid_for(&object_meta)); // File size changed - closure should detect invalidity let object_meta2 = create_test_object_meta("test", 2048); - let result3 = cache.get(&object_meta2.location).unwrap(); + let result3 = cache + .get(&metadata_cache_key(&object_meta2.location)) + .unwrap(); // Cached entry should NOT be valid for new meta assert!(!result3.is_valid_for(&object_meta2)); // Return new entry let new_entry = CachedFileMetadataEntry::new(object_meta2.clone(), Arc::clone(&metadata)); - cache.put(&object_meta2.location, new_entry); + cache.put(&metadata_cache_key(&object_meta2.location), new_entry); - let result4 = cache.get(&object_meta2.location).unwrap(); + let result4 = cache + .get(&metadata_cache_key(&object_meta2.location)) + .unwrap(); assert_eq!(result4.meta.size, 2048); // remove - cache.remove(&object_meta.location); - assert!(!cache.contains_key(&object_meta.location)); + cache.remove(&metadata_cache_key(&object_meta.location)); + assert!(!cache.contains_key(&metadata_cache_key(&object_meta.location))); // len and clear let object_meta3 = create_test_object_meta("test3", 100); cache.put( - &object_meta.location, + &metadata_cache_key(&object_meta.location), CachedFileMetadataEntry::new(object_meta.clone(), Arc::clone(&metadata)), ); cache.put( - &object_meta3.location, + &metadata_cache_key(&object_meta3.location), CachedFileMetadataEntry::new(object_meta3.clone(), Arc::clone(&metadata)), ); assert_eq!(cache.len(), 2); @@ -447,86 +510,87 @@ mod tests { #[test] fn test_default_file_metadata_cache_with_limit() { - // Create a cache with 1000 bytes capacity + 4 keys each key 2 bytes - let cache = DefaultCache::new(1000 + 4 * 2); + // Allow 1000 bytes of metadata plus four store-qualified keys. + let key_size = metadata_cache_key(&Path::from("01")).size(); + let cache = DefaultCache::new(1000 + 4 * key_size); let (object_meta1, metadata1) = generate_test_metadata_with_size("01", 100); let (object_meta2, metadata2) = generate_test_metadata_with_size("02", 500); let (object_meta3, metadata3) = generate_test_metadata_with_size("03", 300); cache.put( - &object_meta1.location, + &metadata_cache_key(&object_meta1.location), CachedFileMetadataEntry::new(object_meta1.clone(), metadata1), ); cache.put( - &object_meta2.location, + &metadata_cache_key(&object_meta2.location), CachedFileMetadataEntry::new(object_meta2.clone(), metadata2), ); cache.put( - &object_meta3.location, + &metadata_cache_key(&object_meta3.location), CachedFileMetadataEntry::new(object_meta3.clone(), metadata3), ); // all entries will fit assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 906); - assert!(cache.contains_key(&object_meta1.location)); - assert!(cache.contains_key(&object_meta2.location)); - assert!(cache.contains_key(&object_meta3.location)); + assert_eq!(cache.memory_used(), 900 + 3 * key_size); + assert!(cache.contains_key(&metadata_cache_key(&object_meta1.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta2.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta3.location))); // add a new entry which will remove the least recently used ("1") let (object_meta4, metadata4) = generate_test_metadata_with_size("04", 200); cache.put( - &object_meta4.location, + &metadata_cache_key(&object_meta4.location), CachedFileMetadataEntry::new(object_meta4.clone(), metadata4), ); assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 1006); - assert!(!cache.contains_key(&object_meta1.location)); - assert!(cache.contains_key(&object_meta4.location)); + assert_eq!(cache.memory_used(), 1000 + 3 * key_size); + assert!(!cache.contains_key(&metadata_cache_key(&object_meta1.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta4.location))); // get entry "2", which will move it to the top of the queue, and add a new one which will // remove the new least recently used ("3") - let _ = cache.get(&object_meta2.location); + let _ = cache.get(&metadata_cache_key(&object_meta2.location)); let (object_meta5, metadata5) = generate_test_metadata_with_size("05", 100); cache.put( - &object_meta5.location, + &metadata_cache_key(&object_meta5.location), CachedFileMetadataEntry::new(object_meta5.clone(), metadata5), ); assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 806); - assert!(!cache.contains_key(&object_meta3.location)); - assert!(cache.contains_key(&object_meta5.location)); + assert_eq!(cache.memory_used(), 800 + 3 * key_size); + assert!(!cache.contains_key(&metadata_cache_key(&object_meta3.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta5.location))); // new entry which will not be able to fit in the 1000 bytes allocated let (object_meta6, metadata6) = generate_test_metadata_with_size("06", 1200); cache.put( - &object_meta6.location, + &metadata_cache_key(&object_meta6.location), CachedFileMetadataEntry::new(object_meta6.clone(), metadata6), ); assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 806); - assert!(!cache.contains_key(&object_meta6.location)); + assert_eq!(cache.memory_used(), 800 + 3 * key_size); + assert!(!cache.contains_key(&metadata_cache_key(&object_meta6.location))); // new entry which is able to fit without removing any entry let (object_meta7, metadata7) = generate_test_metadata_with_size("07", 200); cache.put( - &object_meta7.location, + &metadata_cache_key(&object_meta7.location), CachedFileMetadataEntry::new(object_meta7.clone(), metadata7), ); assert_eq!(cache.len(), 4); - assert_eq!(cache.memory_used(), 1008); - assert!(cache.contains_key(&object_meta7.location)); + assert_eq!(cache.memory_used(), 1000 + 4 * key_size); + assert!(cache.contains_key(&metadata_cache_key(&object_meta7.location))); // new entry which will remove all other entries let (object_meta8, metadata8) = generate_test_metadata_with_size("08", 999); cache.put( - &object_meta8.location, + &metadata_cache_key(&object_meta8.location), CachedFileMetadataEntry::new(object_meta8.clone(), metadata8), ); assert_eq!(cache.len(), 1); - assert_eq!(cache.memory_used(), 1001); - assert!(cache.contains_key(&object_meta8.location)); + assert_eq!(cache.memory_used(), 999 + key_size); + assert!(cache.contains_key(&metadata_cache_key(&object_meta8.location))); // when updating an entry, the previous ones are not unnecessarily removed let (object_meta9, metadata9) = generate_test_metadata_with_size("09", 300); @@ -534,49 +598,49 @@ mod tests { let (object_meta11_v1, metadata11_v1) = generate_test_metadata_with_size("11", 400); cache.put( - &object_meta9.location, + &metadata_cache_key(&object_meta9.location), CachedFileMetadataEntry::new(object_meta9.clone(), metadata9), ); cache.put( - &object_meta10.location, + &metadata_cache_key(&object_meta10.location), CachedFileMetadataEntry::new(object_meta10.clone(), metadata10), ); cache.put( - &object_meta11_v1.location, + &metadata_cache_key(&object_meta11_v1.location), CachedFileMetadataEntry::new(object_meta11_v1.clone(), metadata11_v1), ); - assert_eq!(cache.memory_used(), 906); + assert_eq!(cache.memory_used(), 900 + 3 * key_size); assert_eq!(cache.len(), 3); let (object_meta11_v2, metadata11_v2) = generate_test_metadata_with_size("11", 500); cache.put( - &object_meta11_v2.location, + &metadata_cache_key(&object_meta11_v2.location), CachedFileMetadataEntry::new(object_meta11_v2.clone(), metadata11_v2), ); - assert_eq!(cache.memory_used(), 1006); + assert_eq!(cache.memory_used(), 1000 + 3 * key_size); assert_eq!(cache.len(), 3); - assert!(cache.contains_key(&object_meta9.location)); - assert!(cache.contains_key(&object_meta10.location)); - assert!(cache.contains_key(&object_meta11_v2.location)); + assert!(cache.contains_key(&metadata_cache_key(&object_meta9.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta10.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta11_v2.location))); // when updating an entry that now exceeds the limit, the LRU ("09") needs to be removed let (object_meta11_v3, metadata11_v3) = - generate_test_metadata_with_size("11", 510); + generate_test_metadata_with_size("11", 510 + key_size); cache.put( - &object_meta11_v3.location, + &metadata_cache_key(&object_meta11_v3.location), CachedFileMetadataEntry::new(object_meta11_v3.clone(), metadata11_v3), ); - assert_eq!(cache.memory_used(), 714); + assert_eq!(cache.memory_used(), 710 + 3 * key_size); assert_eq!(cache.len(), 2); - assert!(cache.contains_key(&object_meta10.location)); - assert!(cache.contains_key(&object_meta11_v3.location)); + assert!(cache.contains_key(&metadata_cache_key(&object_meta10.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta11_v3.location))); // manually removing an entry that is not the LRU - cache.remove(&object_meta11_v3.location); + cache.remove(&metadata_cache_key(&object_meta11_v3.location)); assert_eq!(cache.len(), 1); - assert_eq!(cache.memory_used(), 202); - assert!(cache.contains_key(&object_meta10.location)); - assert!(!cache.contains_key(&object_meta11_v3.location)); + assert_eq!(cache.memory_used(), 200 + key_size); + assert!(cache.contains_key(&metadata_cache_key(&object_meta10.location))); + assert!(!cache.contains_key(&metadata_cache_key(&object_meta11_v3.location))); // clear cache.clear(); @@ -588,31 +652,32 @@ mod tests { let (object_meta13, metadata13) = generate_test_metadata_with_size("13", 200); let (object_meta14, metadata14) = generate_test_metadata_with_size("14", 500); cache.put( - &object_meta12.location, + &metadata_cache_key(&object_meta12.location), CachedFileMetadataEntry::new(object_meta12.clone(), metadata12), ); cache.put( - &object_meta13.location, + &metadata_cache_key(&object_meta13.location), CachedFileMetadataEntry::new(object_meta13.clone(), metadata13), ); cache.put( - &object_meta14.location, + &metadata_cache_key(&object_meta14.location), CachedFileMetadataEntry::new(object_meta14.clone(), metadata14), ); assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 1006); + assert_eq!(cache.memory_used(), 1000 + 3 * key_size); cache.update_cache_limit(600); assert_eq!(cache.len(), 1); - assert_eq!(cache.memory_used(), 502); - assert!(!cache.contains_key(&object_meta12.location)); - assert!(!cache.contains_key(&object_meta13.location)); - assert!(cache.contains_key(&object_meta14.location)); + assert_eq!(cache.memory_used(), 500 + key_size); + assert!(!cache.contains_key(&metadata_cache_key(&object_meta12.location))); + assert!(!cache.contains_key(&metadata_cache_key(&object_meta13.location))); + assert!(cache.contains_key(&metadata_cache_key(&object_meta14.location))); } #[test] fn test_default_file_metadata_cache_entries_info() { - // Create a cache with 1000 bytes + 4 bytes for 4 keys each key 1 byte - let cache = DefaultCache::new(1000 + 4); + // Allow 1000 bytes of metadata plus four store-qualified keys. + let key_size = metadata_cache_key(&Path::from("1")).size(); + let cache = DefaultCache::new(1000 + 4 * key_size); let (object_meta1, metadata1) = generate_test_metadata_with_size("1", 100); let (object_meta2, metadata2) = generate_test_metadata_with_size("2", 200); @@ -625,16 +690,16 @@ mod tests { // Build a cache which fits exactly these 3 entries - cache.put(&object_meta1.location, entry_1.clone()); - cache.put(&object_meta2.location, entry_2.clone()); - cache.put(&object_meta3.location, entry_3.clone()); + cache.put(&metadata_cache_key(&object_meta1.location), entry_1.clone()); + cache.put(&metadata_cache_key(&object_meta2.location), entry_2.clone()); + cache.put(&metadata_cache_key(&object_meta3.location), entry_3.clone()); let entries = cache.list_entries(); assert_eq!( entries, HashMap::from([ ( - Path::from("1"), + metadata_cache_key(&Path::from("1")), CacheEntryInfo { value: entry_1.clone(), size_bytes: 100, @@ -643,7 +708,7 @@ mod tests { } ), ( - Path::from("2"), + metadata_cache_key(&Path::from("2")), CacheEntryInfo { value: entry_2.clone(), size_bytes: 200, @@ -652,7 +717,7 @@ mod tests { } ), ( - Path::from("3"), + metadata_cache_key(&Path::from("3")), CacheEntryInfo { value: entry_3.clone(), size_bytes: 300, @@ -664,12 +729,12 @@ mod tests { ); // new hit on "1" - let _ = cache.get(&object_meta1.location); + let _ = cache.get(&metadata_cache_key(&object_meta1.location)); assert_eq!( cache.list_entries(), HashMap::from([ ( - Path::from("1"), + metadata_cache_key(&Path::from("1")), CacheEntryInfo { value: entry_1.clone(), size_bytes: 100, @@ -678,7 +743,7 @@ mod tests { } ), ( - Path::from("2"), + metadata_cache_key(&Path::from("2")), CacheEntryInfo { value: entry_2.clone(), size_bytes: 200, @@ -687,7 +752,7 @@ mod tests { } ), ( - Path::from("3"), + metadata_cache_key(&Path::from("3")), CacheEntryInfo { value: entry_3.clone(), size_bytes: 300, @@ -701,12 +766,12 @@ mod tests { // new entry, will evict "2" let (object_meta4, metadata4) = generate_test_metadata_with_size("4", 600); let entry_4 = CachedFileMetadataEntry::new(object_meta4.clone(), metadata4); - cache.put(&object_meta4.location, entry_4.clone()); + cache.put(&metadata_cache_key(&object_meta4.location), entry_4.clone()); assert_eq!( cache.list_entries(), HashMap::from([ ( - Path::from("1"), + metadata_cache_key(&Path::from("1")), CacheEntryInfo { value: entry_1.clone(), size_bytes: 100, @@ -715,7 +780,7 @@ mod tests { } ), ( - Path::from("3"), + metadata_cache_key(&Path::from("3")), CacheEntryInfo { value: entry_3.clone(), size_bytes: 300, @@ -724,7 +789,7 @@ mod tests { } ), ( - Path::from("4"), + metadata_cache_key(&Path::from("4")), CacheEntryInfo { value: entry_4.clone(), size_bytes: 600, @@ -739,12 +804,15 @@ mod tests { let (object_meta1_new, metadata1_new) = generate_test_metadata_with_size("1", 50); let entry_1 = CachedFileMetadataEntry::new(object_meta1_new.clone(), metadata1_new); - cache.put(&object_meta1_new.location, entry_1.clone()); + cache.put( + &metadata_cache_key(&object_meta1_new.location), + entry_1.clone(), + ); assert_eq!( cache.list_entries(), HashMap::from([ ( - Path::from("1"), + metadata_cache_key(&Path::from("1")), CacheEntryInfo { value: entry_1.clone(), size_bytes: 50, @@ -753,7 +821,7 @@ mod tests { } ), ( - Path::from("3"), + metadata_cache_key(&Path::from("3")), CacheEntryInfo { value: entry_3.clone(), size_bytes: 300, @@ -762,7 +830,7 @@ mod tests { } ), ( - Path::from("4"), + metadata_cache_key(&Path::from("4")), CacheEntryInfo { value: entry_4.clone(), size_bytes: 600, @@ -774,12 +842,12 @@ mod tests { ); // remove entry "4" - cache.remove(&object_meta4.location); + cache.remove(&metadata_cache_key(&object_meta4.location)); assert_eq!( cache.list_entries(), HashMap::from([ ( - Path::from("1"), + metadata_cache_key(&Path::from("1")), CacheEntryInfo { value: entry_1.clone(), size_bytes: 50, @@ -788,7 +856,7 @@ mod tests { } ), ( - Path::from("3"), + metadata_cache_key(&Path::from("3")), CacheEntryInfo { value: entry_3.clone(), size_bytes: 300, @@ -828,8 +896,11 @@ mod tests { )]); let path = TableScopedPath { - path: meta.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta.location.clone(), + ), }; let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); @@ -889,8 +960,11 @@ mod tests { assert_eq!(entries.len(), 1); let path_3 = TableScopedPath { - path: Path::from("test"), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + Path::from("test"), + ), }; let entry = entries.get(&path_3).unwrap(); @@ -964,8 +1038,11 @@ mod tests { ); let path = TableScopedPath { - path: meta.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta.location.clone(), + ), }; cache.put(&path, cached_value); @@ -993,8 +1070,11 @@ mod tests { fn test_cache_invalidation_on_file_modification() { let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); let path = TableScopedPath { - path: Path::from("test.parquet"), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + Path::from("test.parquet"), + ), }; let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); @@ -1035,15 +1115,18 @@ mod tests { fn test_ordering_cache_invalidation_on_file_modification() { let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); let path = TableScopedPath { - path: Path::from("test.parquet"), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + Path::from("test.parquet"), + ), }; let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); // Cache with original metadata and ordering let meta_v1 = ObjectMeta { - location: path.path.clone(), + location: path.store_path.path.clone(), last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") .unwrap() .into(), @@ -1067,7 +1150,7 @@ mod tests { // File modified (size changed) let meta_v2 = ObjectMeta { - location: path.path.clone(), + location: path.store_path.path.clone(), last_modified: DateTime::parse_from_rfc3339("2022-09-28T10:00:00+02:00") .unwrap() .into(), @@ -1115,8 +1198,11 @@ mod tests { ); let path_1 = TableScopedPath { - path: meta1.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta1.location.clone(), + ), }; cache.put(&path_1, cached_value_1.clone()); @@ -1129,8 +1215,11 @@ mod tests { ); let path_2 = TableScopedPath { - path: meta2.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta2.location.clone(), + ), }; cache.put(&path_2, cached_value_2.clone()); @@ -1172,7 +1261,10 @@ mod tests { let mut ctx = DFHeapSizeCtx::default(); - let limit_for_2_entries = meta_1.location.as_ref().heap_size(&mut ctx) + let limit_for_2_entries = 2 * ObjectStoreUrl::local_filesystem() + .as_str() + .heap_size(&mut ctx) + + meta_1.location.as_ref().heap_size(&mut ctx) + value_1.heap_size(&mut ctx) + meta_2.location.as_ref().heap_size(&mut ctx) + value_2.heap_size(&mut ctx); @@ -1180,13 +1272,19 @@ mod tests { // create a cache with a limit which fits exactly 2 entries let cache = DefaultCache::new(limit_for_2_entries); let path_1 = TableScopedPath { - path: meta_1.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta_1.location.clone(), + ), }; let path_2 = TableScopedPath { - path: meta_2.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta_2.location.clone(), + ), }; cache.put(&path_1, value_1.clone()); @@ -1201,8 +1299,11 @@ mod tests { assert_eq!(result_2.unwrap(), value_2); let path_3 = TableScopedPath { - path: meta_3.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta_3.location.clone(), + ), }; // adding the third entry evicts the first entry @@ -1230,7 +1331,7 @@ mod tests { assert_eq!(cache.len(), 1); assert_eq!( cache.memory_used(), - meta_3.location.as_ref().heap_size(&mut ctx) + value_3.heap_size(&mut ctx) + path_3.size() + value_3.heap_size(&mut ctx) ); cache.clear(); @@ -1249,8 +1350,11 @@ mod tests { let cache = DefaultCache::new(limit_less_than_the_entry); let path_1 = TableScopedPath { - path: meta.location.clone(), table: None, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + meta.location.clone(), + ), }; cache.put(&path_1, value_too_large.clone()); @@ -1263,7 +1367,7 @@ mod tests { cache.put(&path_1, value_fits.clone()); assert_eq!(cache.len(), 1); - assert_eq!(cache.memory_used(), 1514); + assert_eq!(cache.memory_used(), path_1.size() + value_fits.size()); // now add an entry which is over the limit and make sure the old stale entry is removed let stale_entry = cache.put(&path_1, value_too_large.clone()); @@ -1364,7 +1468,10 @@ mod tests { ) -> (TableScopedPath, CachedFileList) { let key = TableScopedPath { table, - path: Path::from(path), + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + Path::from(path), + ), }; let metas: Vec = (0..count) .map(|i| create_object_meta(&format!("file{i}"), meta_size)) @@ -1380,7 +1487,7 @@ mod tests { let path = Path::from("test_path"); let key = TableScopedPath { table: table_ref.clone(), - path, + store_path: StoreScopedPath::new(ObjectStoreUrl::local_filesystem(), path), }; // Initially cache is empty @@ -1895,7 +2002,10 @@ mod tests { let table_ref = Some(TableReference::from("table")); let key = TableScopedPath { table: table_ref, - path: table_base, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + table_base, + ), }; cache.put(&key, CachedFileList::new(files)); @@ -1939,7 +2049,10 @@ mod tests { let table_ref = Some(TableReference::from("table")); let key = TableScopedPath { table: table_ref, - path: table_base, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + table_base, + ), }; cache.put(&key, CachedFileList::new(files)); let result = cache.get(&key).unwrap(); @@ -1973,7 +2086,10 @@ mod tests { let table_ref = Some(TableReference::from("table")); let key = TableScopedPath { table: table_ref, - path: table_base, + store_path: StoreScopedPath::new( + ObjectStoreUrl::local_filesystem(), + table_base, + ), }; cache.put(&key, CachedFileList::new(files)); let result = cache.get(&key).unwrap(); diff --git a/datafusion/execution/src/cache/mod.rs b/datafusion/execution/src/cache/mod.rs index 2eba533e04b67..30521d6a01451 100644 --- a/datafusion/execution/src/cache/mod.rs +++ b/datafusion/execution/src/cache/mod.rs @@ -30,6 +30,8 @@ use std::fmt::{Debug, Display, Formatter}; use std::hash::{Hash, Hasher}; use std::time::Duration; +use crate::object_store::ObjectStoreUrl; + /// Base trait for cache implementations with common operations. /// /// This trait provides the fundamental cache operations (`get`, `put`, `remove`, etc.) @@ -143,27 +145,68 @@ impl CacheKey for TableScopedPath { } } -/// Each entry is scoped to its use within a specific table so that the cache -/// can differentiate between identical paths in different tables, and +/// Identifies an object within its registered object store. +#[derive(PartialEq, Eq, Hash, Clone, Debug)] +pub struct StoreScopedPath { + /// URL identifying the registered object store. + pub object_store_url: ObjectStoreUrl, + /// Location relative to that store. + pub path: Path, +} + +impl StoreScopedPath { + /// Create a cache key for a path in the given object store. + pub fn new(object_store_url: ObjectStoreUrl, path: Path) -> Self { + Self { + object_store_url, + path, + } + } +} + +impl CacheKey for StoreScopedPath { + fn size(&self) -> usize { + self.heap_size(&mut DFHeapSizeCtx::default()) + } + + fn table_ref(&self) -> Option<&TableReference> { + None + } +} + +impl DFHeapSize for StoreScopedPath { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.object_store_url.as_str().heap_size(ctx) + self.path.as_ref().heap_size(ctx) + } +} + +impl Display for StoreScopedPath { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}{}", self.object_store_url, self.path) + } +} + +/// Scopes a path to its object store and optional table, supporting /// table-level cache invalidation. #[derive(PartialEq, Eq, Hash, Clone, Debug)] pub struct TableScopedPath { pub table: Option, - pub path: Path, + /// Object location, including its registered store. + pub store_path: StoreScopedPath, } impl DFHeapSize for TableScopedPath { fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { - self.path.as_ref().heap_size(ctx) + self.table.heap_size(ctx) + self.store_path.heap_size(ctx) + self.table.heap_size(ctx) } } impl Display for TableScopedPath { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { if let Some(table) = &self.table { - write!(f, "{}, {}", self.path, table) + write!(f, "{}, {}", self.store_path, table) } else { - write!(f, "{}", self.path) + Display::fmt(&self.store_path, f) } } } @@ -274,3 +317,41 @@ mod schema_fingerprint_tests { assert_eq!(plain, schema_md, "schema metadata must be ignored"); } } + +#[cfg(test)] +mod path_tests { + use super::*; + + #[test] + fn scoped_paths_preserve_store_identity_and_table_scope() { + let store_path = StoreScopedPath::new( + ObjectStoreUrl::parse("s3://bucket").unwrap(), + Path::from("data.parquet"), + ); + assert_eq!(store_path.to_string(), "s3://bucket/data.parquet"); + assert_eq!(store_path.size(), "s3://bucket/data.parquet".len()); + assert_eq!(store_path.table_ref(), None); + + let unscoped = TableScopedPath { + table: None, + store_path: store_path.clone(), + }; + let table = TableReference::bare("example"); + let scoped = TableScopedPath { + table: Some(table.clone()), + store_path, + }; + assert_eq!(unscoped.to_string(), "s3://bucket/data.parquet"); + assert_eq!(scoped.to_string(), "s3://bucket/data.parquet, example"); + assert_eq!(scoped.table_ref(), Some(&table)); + assert_eq!( + scoped.size(), + unscoped.size() + 2 * size_of::() + "example".len() + ); + + // Entries for the same object can coexist inside and outside a table. + let entries = HashMap::from([(unscoped.clone(), 1), (scoped.clone(), 2)]); + assert_eq!(entries.get(&unscoped), Some(&1)); + assert_eq!(entries.get(&scoped), Some(&2)); + } +} diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 8f5604a5cd3d2..bb26ddcb6a257 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1247,3 +1247,36 @@ without declaring one, use `FileScanConfigBuilder::with_preserve_order(true)`. Code that reads `FileScanConfig::preserve_order` to decide whether files may be reordered should call `FileScanConfig::preserves_file_order()` instead, which also accounts for the output ordering. + +### File caches include the object store URL + +File listing, statistics, and Parquet metadata caches now include the registered +object store URL in their keys. This prevents identical paths in different buckets +from sharing cached results. Statistics and Parquet metadata validation compares +strong ETags when both metadata values provide them, otherwise versions when both +provide them, otherwise size and modification time. Weak HTTP ETags are treated +as missing. ETags remain opaque: differences such as quoted versus unquoted values +cause a cache miss. + +**Migration guide:** + +- `FileFormat::infer_schema`, `infer_stats`, `infer_ordering`, and + `infer_stats_and_ordering` take a new `object_store_url: &ObjectStoreUrl` + argument immediately after `store`. Update custom implementations and callers + to pass the URL identifying that store. +- `TableScopedPath` contains a `store_path: StoreScopedPath` field alongside its + optional `table`, replacing the previous `path` field. +- `FileMetadataCache` uses `StoreScopedPath` instead of `object_store::path::Path` + as its key. Construct keys with + `datafusion_execution::cache::StoreScopedPath::new(object_store_url, path)`. +- `CachedParquetFileReaderFactory::new` now takes + `(store, object_store_url, metadata_cache)`. +- `DFParquetMetadata::with_file_metadata_cache` and + `ParquetFileReader::with_metadata_cache` take an `ObjectStoreUrl` argument + immediately after the cache argument. + +The CLI's `metadata_cache()`, `statistics_cache()`, and `list_files_cache()` +functions expose an `object_store_url` column. Their `path` columns remain relative +to that store. + +See [issue #25841](https://github.com/apache/datafusion/issues/25841) for details. diff --git a/docs/source/user-guide/cli/functions.md b/docs/source/user-guide/cli/functions.md index eb7f70e66a1a5..874f976ff9fb3 100644 --- a/docs/source/user-guide/cli/functions.md +++ b/docs/source/user-guide/cli/functions.md @@ -150,6 +150,8 @@ The columns of the returned table are: The `metadata_cache` function shows information about the default File Metadata Cache that is used by the [`ListingTable`] implementation in DataFusion. This cache is used to speed up reading metadata from files when scanning directories with many files. +The `path` column is relative to the object store; `object_store_url` identifies +the store. Together they distinguish files with the same path in different stores. For example, after creating a table with the [CREATE EXTERNAL TABLE](../sql/ddl.md#create-external-table) command: @@ -164,15 +166,14 @@ You can inspect the metadata cache by querying the `metadata_cache` function: ```sql > select * from metadata_cache(); -+----------------------------------------------------+---------------------+-----------------+---------------------------------------+---------+---------------------+------+------------------+ -| path | file_modified | file_size_bytes | e_tag | version | metadata_size_bytes | hits | extra | -+----------------------------------------------------+---------------------+-----------------+---------------------------------------+---------+---------------------+------+------------------+ -| hits_compatible/athena_partitioned/hits_61.parquet | 2022-07-03T15:40:34 | 117270944 | "5db11cad1ca0d80d748fc92c914b010a-6" | NULL | 212949 | 0 | page_index=false | -| hits_compatible/athena_partitioned/hits_32.parquet | 2022-07-03T15:37:17 | 94506004 | "2f7db49a9fe242179590b615b94a39d2-5" | NULL | 278157 | 0 | page_index=false | -| hits_compatible/athena_partitioned/hits_40.parquet | 2022-07-03T15:38:07 | 142508647 | "9e5852b45a469d5a05bf270a286eab8a-8" | NULL | 212917 | 0 | page_index=false | -| hits_compatible/athena_partitioned/hits_93.parquet | 2022-07-03T15:44:07 | 127987774 | "751100bf0dac7d489b9836abf3108b99-7" | NULL | 278318 | 0 | page_index=false | -| . | -+----------------------------------------------------+---------------------+-----------------+---------------------------------------+---------+---------------------+------+------------------+ ++----------------------------------------------------+----------------------------------+---------------------+-----------------+--------------------------------------+---------+---------------------+------+------------------+ +| path | object_store_url | file_modified | file_size_bytes | e_tag | version | metadata_size_bytes | hits | extra | ++----------------------------------------------------+----------------------------------+---------------------+-----------------+--------------------------------------+---------+---------------------+------+------------------+ +| hits_compatible/athena_partitioned/hits_61.parquet | s3://clickhouse-public-datasets/ | 2022-07-03T15:40:34 | 117270944 | "5db11cad1ca0d80d748fc92c914b010a-6" | NULL | 212949 | 0 | page_index=false | +| hits_compatible/athena_partitioned/hits_32.parquet | s3://clickhouse-public-datasets/ | 2022-07-03T15:37:17 | 94506004 | "2f7db49a9fe242179590b615b94a39d2-5" | NULL | 278157 | 0 | page_index=false | +| hits_compatible/athena_partitioned/hits_40.parquet | s3://clickhouse-public-datasets/ | 2022-07-03T15:38:07 | 142508647 | "9e5852b45a469d5a05bf270a286eab8a-8" | NULL | 212917 | 0 | page_index=false | +| hits_compatible/athena_partitioned/hits_93.parquet | s3://clickhouse-public-datasets/ | 2022-07-03T15:44:07 | 127987774 | "751100bf0dac7d489b9836abf3108b99-7" | NULL | 278318 | 0 | page_index=false | ++----------------------------------------------------+----------------------------------+---------------------+-----------------+--------------------------------------+---------+---------------------+------+------------------+ ``` Since `metadata_cache` is a normal table function, you can use it in most places you can use @@ -194,6 +195,7 @@ The columns of the returned table are: | column_name | data_type | Description | | ------------------- | --------- | ----------------------------------------------------------------------------------------- | | path | Utf8 | File path relative to the object store / filesystem root | +| object_store_url | Utf8 | URL identifying the object store, such as `s3://bucket/` or `file:///` | | file_modified | Timestamp | Last modified time of the file | | file_size_bytes | UInt64 | Size of the file in bytes | | e_tag | Utf8 | [Entity Tag] (ETag) of the file if available | @@ -213,11 +215,11 @@ You can inspect the statistics cache by querying the `statistics_cache` function ```sql > select * from statistics_cache(); -+------------------+---------------------+-----------------+------------------------+---------+-----------------+-------------+--------------------+-----------------------+ -| path | file_modified | file_size_bytes | e_tag | version | num_rows | num_columns | table_size_bytes | statistics_size_bytes | -+------------------+---------------------+-----------------+------------------------+---------+-----------------+-------------+--------------------+-----------------------+ -| .../hits.parquet | 2022-06-25T22:22:22 | 14779976446 | 0-5e24d1ee16380-370f48 | NULL | Exact(99997497) | 105 | Exact(36445943240) | 0 | -+------------------+---------------------+-----------------+------------------------+---------+-----------------+-------------+--------------------+-----------------------+ ++--------------+----------------------------------+-------+---------------------+-----------------+------------------------+---------+-----------------+-------------+--------------------+-----------------------+------+ +| path | object_store_url | table | file_modified | file_size_bytes | e_tag | version | num_rows | num_columns | table_size_bytes | statistics_size_bytes | hits | ++--------------+----------------------------------+-------+---------------------+-----------------+------------------------+---------+-----------------+-------------+--------------------+-----------------------+------+ +| hits.parquet | s3://clickhouse-public-datasets/ | hits | 2022-06-25T22:22:22 | 14779976446 | 0-5e24d1ee16380-370f48 | NULL | Exact(99997497) | 105 | Exact(36445943240) | 0 | 0 | ++--------------+----------------------------------+-------+---------------------+-----------------+------------------------+---------+-----------------+-------------+--------------------+-----------------------+------+ ``` The columns of the returned table are: @@ -225,6 +227,8 @@ The columns of the returned table are: | column_name | data_type | Description | | --------------------- | --------- | ---------------------------------------------------------------------------- | | path | Utf8 | File path relative to the object store / filesystem root | +| object_store_url | Utf8 | URL identifying the object store, such as `s3://bucket/` or `file:///` | +| table | Utf8 | Name of the table, or an empty string for an anonymous read | | file_modified | Timestamp | Last modified time of the file | | file_size_bytes | UInt64 | Size of the file in bytes | | e_tag | Utf8 | [Entity Tag] (ETag) of the file if available | @@ -237,7 +241,7 @@ The columns of the returned table are: ## `list_files_cache` -The `list_files_cache` function shows information about the `ListFilesCache` that is used by the [`ListingTable`] implementation in DataFusion. When creating a [`ListingTable`], DataFusion lists the files in the table's location and caches results in the `ListFilesCache`. Subsequent queries against the same table can reuse this cached information instead of re-listing the files. Cache entries are scoped to tables. +The `list_files_cache` function shows information about the `ListFilesCache` that is used by the [`ListingTable`] implementation in DataFusion. When creating a [`ListingTable`], DataFusion lists the files in the table's location and caches results in the `ListFilesCache`. Subsequent queries against the same table can reuse this cached information instead of re-listing the files. Cache entries are scoped to the object store and, when available, the table. You can inspect the cache by querying the `list_files_cache` function. For example, @@ -247,33 +251,34 @@ You can inspect the cache by querying the `list_files_cache` function. For examp stored as parquet location 's3://overturemaps-us-west-2/release/2025-12-17.0/theme=base/type=infrastructure'; 0 row(s) fetched. -> select table, path, metadata_size_bytes, expires_in, unnest(metadata_list)['file_size_bytes'] as file_size_bytes, unnest(metadata_list)['e_tag'] as e_tag from list_files_cache() limit 10; -+--------------+-----------------------------------------------------+---------------------+-----------------------------------+-----------------+---------------------------------------+ -| table | path | metadata_size_bytes | expires_in | file_size_bytes | e_tag | -+--------------+-----------------------------------------------------+---------------------+-----------------------------------+-----------------+---------------------------------------+ -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 999055952 | "35fc8fbe8400960b54c66fbb408c48e8-60" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 975592768 | "8a16e10b722681cdc00242564b502965-59" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1082925747 | "24cd13ddb5e0e438952d2499f5dabe06-65" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1008425557 | "37663e31c7c64d4ef355882bcd47e361-61" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1065561905 | "4e7c50d2d1b3c5ed7b82b4898f5ac332-64" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1045655427 | "8fff7e6a72d375eba668727c55d4f103-63" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1086822683 | "b67167d8022d778936c330a52a5f1922-65" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1016732378 | "6d70857a0473ed9ed3fc6e149814168b-61" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 991363784 | "c9cafb42fcbb413f851691c895dd7c2b-60" | -| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1032469715 | "7540252d0d67158297a67038a3365e0f-62" | -+--------------+-----------------------------------------------------+---------------------+-----------------------------------+-----------------+---------------------------------------+ +> select table, path, object_store_url, metadata_size_bytes, expires_in, unnest(metadata_list)['file_size_bytes'] as file_size_bytes, unnest(metadata_list)['e_tag'] as e_tag from list_files_cache() limit 10; ++--------------+-----------------------------------------------------+------------------------------+---------------------+-----------------------------------+-----------------+---------------------------------------+ +| table | path | object_store_url | metadata_size_bytes | expires_in | file_size_bytes | e_tag | ++--------------+-----------------------------------------------------+------------------------------+---------------------+-----------------------------------+-----------------+---------------------------------------+ +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 999055952 | "35fc8fbe8400960b54c66fbb408c48e8-60" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 975592768 | "8a16e10b722681cdc00242564b502965-59" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1082925747 | "24cd13ddb5e0e438952d2499f5dabe06-65" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1008425557 | "37663e31c7c64d4ef355882bcd47e361-61" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1065561905 | "4e7c50d2d1b3c5ed7b82b4898f5ac332-64" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1045655427 | "8fff7e6a72d375eba668727c55d4f103-63" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1086822683 | "b67167d8022d778936c330a52a5f1922-65" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1016732378 | "6d70857a0473ed9ed3fc6e149814168b-61" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 991363784 | "c9cafb42fcbb413f851691c895dd7c2b-60" | +| overturemaps | release/2025-12-17.0/theme=base/type=infrastructure | s3://overturemaps-us-west-2/ | 2750 | 0 days 0 hours 0 mins 25.264 secs | 1032469715 | "7540252d0d67158297a67038a3365e0f-62" | ++--------------+-----------------------------------------------------+------------------------------+---------------------+-----------------------------------+-----------------+---------------------------------------+ ``` The columns of the returned table are: -| column_name | data_type | Description | -| ------------------- | ------------ | ------------------------------------------------------------------- | -| table | Utf8 | Name of the table | -| path | Utf8 | File path relative to the object store / filesystem root | -| metadata_size_bytes | UInt64 | Size of the cached metadata in memory (not its thrift encoded form) | -| expires_in | Duration(ms) | Last modified time of the file | -| hits | UInt64 | Number of times the cached metadata has been accessed | -| metadata_list | List(Struct) | List of metadatas, one for each file under the path. | +| column_name | data_type | Description | +| ------------------- | ------------ | ---------------------------------------------------------------------- | +| table | Utf8 | Name of the table | +| path | Utf8 | File path relative to the object store / filesystem root | +| object_store_url | Utf8 | URL identifying the object store, such as `s3://bucket/` or `file:///` | +| metadata_size_bytes | UInt64 | Size of the cached metadata in memory (not its thrift encoded form) | +| expires_in | Duration(ms) | Last modified time of the file | +| hits | UInt64 | Number of times the cached metadata has been accessed | +| metadata_list | List(Struct) | List of metadatas, one for each file under the path. | A metadata struct in the metadata_list contains the following fields: