This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-20047-fed096cd516efd6bc2f50188a908c6b0db0fd798 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit d3792410d9444afe9eb6d67d383210f81dad3cdf Author: Michael Kleen <[email protected]> AuthorDate: Mon May 18 22:52:25 2026 +0200 Add a memory bound FileStatisticsCache for the Listing Table (#20047) ## Which issue does this PR close? This change introduces a default `FileStatisticsCache` implementation for the Listing-Table with a size limit, implementing the following steps following https://github.com/apache/datafusion/issues/19052#issuecomment-3603796097 : - Add heap size estimation for file statistics and the relevant data types used in caching (This is temporary until https://github.com/apache/arrow-rs/pull/9138 is resolved) - Redesign `DefaultFileStatisticsCache` to use a `LruQueue` to make it memory-bound following https://github.com/apache/datafusion/pull/18855 - Introduce a size limit and use it together with the heap-size to limit the memory usage of the cache - Move `FileStatisticsCache` creation into `CacheManager`, making it session-scoped and shared across statements and tables - Closes https://github.com/apache/datafusion/issues/19217 - Closes https://github.com/apache/datafusion/issues/19052 ## Rationale for this change See above. ## What changes are included in this PR? See above. ## Are these changes tested? Yes. ## Are there any user-facing changes? A new runtime setting `datafusion.runtime.file_statistics.cache_limit` --- benchmarks/src/bin/external_aggr.rs | 4 +- benchmarks/src/imdb/run.rs | 4 +- benchmarks/src/sort_pushdown.rs | 4 +- benchmarks/src/sort_tpch.rs | 4 +- benchmarks/src/tpcds/run.rs | 4 +- benchmarks/src/tpch/run.rs | 4 +- datafusion-cli/src/functions.rs | 7 +- datafusion-cli/src/main.rs | 74 +- datafusion/catalog-listing/src/helpers.rs | 34 +- datafusion/catalog-listing/src/table.rs | 35 +- datafusion/common/src/heap_size.rs | 561 ++++++++++++++++ datafusion/common/src/lib.rs | 2 + .../core/src/datasource/listing_table_factory.rs | 6 +- datafusion/core/src/execution/context/mod.rs | 59 +- datafusion/core/tests/parquet/file_statistics.rs | 19 +- datafusion/core/tests/sql/runtime_config.rs | 46 ++ datafusion/datasource/src/mod.rs | 16 +- datafusion/execution/src/cache/cache_manager.rs | 85 ++- datafusion/execution/src/cache/cache_unit.rs | 410 ------------ .../execution/src/cache/file_statistics_cache.rs | 744 +++++++++++++++++++++ datafusion/execution/src/cache/list_files_cache.rs | 30 +- datafusion/execution/src/cache/mod.rs | 2 +- datafusion/execution/src/runtime_env.rs | 26 +- .../sqllogictest/test_files/encrypted_parquet.slt | 2 +- .../sqllogictest/test_files/information_schema.slt | 2 + .../sqllogictest/test_files/set_variable.slt | 19 + docs/source/library-user-guide/upgrading/54.0.0.md | 39 ++ docs/source/user-guide/configs.md | 17 +- 28 files changed, 1695 insertions(+), 564 deletions(-) diff --git a/benchmarks/src/bin/external_aggr.rs b/benchmarks/src/bin/external_aggr.rs index ee604ec736..a6e322c7fa 100644 --- a/benchmarks/src/bin/external_aggr.rs +++ b/benchmarks/src/bin/external_aggr.rs @@ -326,7 +326,9 @@ impl ExternalAggrConfig { let config = ListingTableConfig::new(table_path).with_listing_options(options); let config = config.infer_schema(&state).await?; - Ok(Arc::new(ListingTable::try_new(config)?)) + Ok(Arc::new(ListingTable::try_new(config)?.with_cache( + ctx.runtime_env().cache_manager.get_file_statistic_cache(), + ))) } fn iterations(&self) -> usize { diff --git a/benchmarks/src/imdb/run.rs b/benchmarks/src/imdb/run.rs index ca9710a920..6d3b5c6baf 100644 --- a/benchmarks/src/imdb/run.rs +++ b/benchmarks/src/imdb/run.rs @@ -470,7 +470,9 @@ impl RunOpt { _ => unreachable!(), }; - Ok(Arc::new(ListingTable::try_new(config)?)) + Ok(Arc::new(ListingTable::try_new(config)?.with_cache( + ctx.runtime_env().cache_manager.get_file_statistic_cache(), + ))) } fn iterations(&self) -> usize { diff --git a/benchmarks/src/sort_pushdown.rs b/benchmarks/src/sort_pushdown.rs index e7fce1921e..8e34706ac1 100644 --- a/benchmarks/src/sort_pushdown.rs +++ b/benchmarks/src/sort_pushdown.rs @@ -273,7 +273,9 @@ impl RunOpt { .with_listing_options(options) .with_schema(schema); - Ok(Arc::new(ListingTable::try_new(config)?)) + Ok(Arc::new(ListingTable::try_new(config)?.with_cache( + ctx.runtime_env().cache_manager.get_file_statistic_cache(), + ))) } fn iterations(&self) -> usize { diff --git a/benchmarks/src/sort_tpch.rs b/benchmarks/src/sort_tpch.rs index 95c90d826d..206911c45a 100644 --- a/benchmarks/src/sort_tpch.rs +++ b/benchmarks/src/sort_tpch.rs @@ -351,7 +351,9 @@ impl RunOpt { .with_listing_options(options) .with_schema(schema); - Ok(Arc::new(ListingTable::try_new(config)?)) + Ok(Arc::new(ListingTable::try_new(config)?.with_cache( + ctx.runtime_env().cache_manager.get_file_statistic_cache(), + ))) } fn iterations(&self) -> usize { diff --git a/benchmarks/src/tpcds/run.rs b/benchmarks/src/tpcds/run.rs index f7ef699151..5882134003 100644 --- a/benchmarks/src/tpcds/run.rs +++ b/benchmarks/src/tpcds/run.rs @@ -347,7 +347,9 @@ impl RunOpt { .with_listing_options(options) .with_schema(schema); - Ok(Arc::new(ListingTable::try_new(config)?)) + Ok(Arc::new(ListingTable::try_new(config)?.with_cache( + ctx.runtime_env().cache_manager.get_file_statistic_cache(), + ))) } fn iterations(&self) -> usize { diff --git a/benchmarks/src/tpch/run.rs b/benchmarks/src/tpch/run.rs index ec7aa8c554..75983ee141 100644 --- a/benchmarks/src/tpch/run.rs +++ b/benchmarks/src/tpch/run.rs @@ -342,7 +342,9 @@ impl RunOpt { .with_listing_options(options) .with_schema(schema); - Ok(Arc::new(ListingTable::try_new(config)?)) + Ok(Arc::new(ListingTable::try_new(config)?.with_cache( + ctx.runtime_env().cache_manager.get_file_statistic_cache(), + ))) } fn iterations(&self) -> usize { diff --git a/datafusion-cli/src/functions.rs b/datafusion-cli/src/functions.rs index 26f007cdd3..df066992fb 100644 --- a/datafusion-cli/src/functions.rs +++ b/datafusion-cli/src/functions.rs @@ -633,6 +633,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { let schema = Arc::new(Schema::new(vec![ Field::new("path", DataType::Utf8, false), + Field::new("table", DataType::Utf8, false), Field::new( "file_modified", DataType::Timestamp(TimeUnit::Millisecond, None), @@ -649,6 +650,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { // construct record batch from metadata let mut path_arr = vec![]; + let mut table_arr = vec![]; let mut file_modified_arr = vec![]; let mut file_size_bytes_arr = vec![]; let mut e_tag_arr = vec![]; @@ -661,7 +663,9 @@ 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.to_string()); + path_arr.push(path.path.to_string()); + table_arr + .push(path.table.map_or_else(|| "".to_string(), |t| t.to_string())); file_modified_arr .push(Some(entry.object_meta.last_modified.timestamp_millis())); file_size_bytes_arr.push(entry.object_meta.size); @@ -678,6 +682,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { schema.clone(), vec![ Arc::new(StringArray::from(path_arr)), + Arc::new(StringArray::from(table_arr)), Arc::new(TimestampMillisecondArray::from(file_modified_arr)), Arc::new(UInt64Array::from(file_size_bytes_arr)), Arc::new(StringArray::from(e_tag_arr)), diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index 6bfe1160ec..935bf0a974 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -443,10 +443,7 @@ mod tests { use super::*; use datafusion::{ common::test_util::batches_to_string, - execution::cache::{ - DefaultListFilesCache, cache_manager::CacheManagerConfig, - cache_unit::DefaultFileStatisticsCache, - }, + execution::cache::{DefaultListFilesCache, cache_manager::CacheManagerConfig}, prelude::{ParquetReadOptions, col, lit, split_part}, }; use insta::assert_snapshot; @@ -656,8 +653,6 @@ mod tests { Ok(()) } - /// Shows that the statistics cache is not enabled by default yet - /// See https://github.com/apache/datafusion/issues/19217 #[tokio::test] async fn test_statistics_cache_default() -> Result<(), DataFusionError> { let ctx = SessionContext::new(); @@ -687,68 +682,17 @@ mod tests { .await?; } - // When the cache manager creates a StatisticsCache by default, - // the contents will show up here - let sql = "SELECT split_part(path, '/', -1) as filename, file_size_bytes, num_rows, num_columns, table_size_bytes from statistics_cache() order by filename"; + let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, table_size_bytes from statistics_cache() order by filename"; let df = ctx.sql(sql).await?; let rbs = df.collect().await?; assert_snapshot!(batches_to_string(&rbs),@r" - ++ - ++ - "); - - Ok(()) - } - - // Can be removed when https://github.com/apache/datafusion/issues/19217 is resolved - #[tokio::test] - async fn test_statistics_cache_override() -> Result<(), DataFusionError> { - // Install a specific StatisticsCache implementation - let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default()); - let cache_config = CacheManagerConfig::default() - .with_files_statistics_cache(Some(file_statistics_cache.clone())); - let runtime = RuntimeEnvBuilder::new() - .with_cache_manager(cache_config) - .build()?; - let config = SessionConfig::new().with_collect_statistics(true); - let ctx = SessionContext::new_with_config_rt(config, Arc::new(runtime)); - - ctx.register_udtf( - "statistics_cache", - Arc::new(StatisticsCacheFunc::new( - ctx.task_ctx().runtime_env().cache_manager.clone(), - )), - ); - - for filename in [ - "alltypes_plain", - "alltypes_tiny_pages", - "lz4_raw_compressed_larger", - ] { - ctx.sql( - format!( - "create external table {filename} - stored as parquet - location '../parquet-testing/data/{filename}.parquet'", - ) - .as_str(), - ) - .await? - .collect() - .await?; - } - - let sql = "SELECT split_part(path, '/', -1) as filename, file_size_bytes, num_rows, num_columns, table_size_bytes from statistics_cache() order by filename"; - let df = ctx.sql(sql).await?; - let rbs = df.collect().await?; - assert_snapshot!(batches_to_string(&rbs),@r" - +-----------------------------------+-----------------+--------------+-------------+------------------+ - | filename | file_size_bytes | num_rows | num_columns | table_size_bytes | - +-----------------------------------+-----------------+--------------+-------------+------------------+ - | alltypes_plain.parquet | 1851 | Exact(8) | 11 | Absent | - | alltypes_tiny_pages.parquet | 454233 | Exact(7300) | 13 | Absent | - | lz4_raw_compressed_larger.parquet | 380836 | Exact(10000) | 1 | Absent | - +-----------------------------------+-----------------+--------------+-------------+------------------+ + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+ + | filename | table | file_size_bytes | num_rows | num_columns | table_size_bytes | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+ + | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | Absent | + | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | Absent | + | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | Absent | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+ "); Ok(()) diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index 2abe2a49a9..0389b3cb17 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -21,9 +21,11 @@ use std::mem; use std::sync::Arc; use datafusion_catalog::Session; -use datafusion_common::{HashMap, Result, ScalarValue, assert_or_internal_err}; -use datafusion_datasource::ListingTableUrl; +use datafusion_common::{ + HashMap, Result, ScalarValue, TableReference, assert_or_internal_err, +}; use datafusion_datasource::PartitionedFile; +use datafusion_datasource::{FileExtensions, ListingTableUrl}; use datafusion_expr::{BinaryExpr, Operator, lit, utils}; use arrow::{ @@ -383,6 +385,7 @@ fn try_into_partitioned_file( let mut pf: PartitionedFile = object_meta.into(); pf.partition_values = partition_values; + pf.table_reference.clone_from(table_path.get_table_ref()); Ok(Some(pf)) } @@ -417,8 +420,15 @@ pub async fn pruned_partition_list<'a>( table_path ); - // if no partition col => simply list all the files - Ok(objects.map_ok(|object_meta| object_meta.into()).boxed()) + // if no partition col => list all the files + Ok(objects + .try_filter_map(|object_meta| { + futures::future::ready(object_meta_to_partitioned_file( + object_meta, + table_path.get_table_ref(), + )) + }) + .boxed()) } else { let df_schema = DFSchema::from_unqualified_fields( partition_cols @@ -443,6 +453,22 @@ pub async fn pruned_partition_list<'a>( } } +fn object_meta_to_partitioned_file( + object_meta: ObjectMeta, + table_ref: &Option<TableReference>, +) -> Result<Option<PartitionedFile>> { + Ok(Some(PartitionedFile { + object_meta, + partition_values: vec![], + range: None, + statistics: None, + ordering: None, + extensions: FileExtensions::new(), + metadata_size_hint: None, + table_reference: table_ref.clone(), + })) +} + /// Extract the partition values for the given `file_path` (in the given `table_path`) /// associated to the partitions defined by `table_partition_cols` pub fn parse_partitions_for_path<'a, I>( diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 06ba8c8113..7ee743a6ab 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -36,7 +36,6 @@ use datafusion_datasource::{ }; use datafusion_execution::cache::TableScopedPath; use datafusion_execution::cache::cache_manager::FileStatisticsCache; -use datafusion_execution::cache::cache_unit::DefaultFileStatisticsCache; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; @@ -187,7 +186,7 @@ pub struct ListingTable { /// The SQL definition for this table, if any definition: Option<String>, /// Cache for collected file statistics - collected_statistics: Arc<dyn FileStatisticsCache>, + collected_statistics: Option<Arc<dyn FileStatisticsCache>>, /// Constraints applied to this table constraints: Constraints, /// Column default expressions for columns that are not physically present in the data files @@ -231,7 +230,7 @@ impl ListingTable { schema_source, options, definition: None, - collected_statistics: Arc::new(DefaultFileStatisticsCache::default()), + collected_statistics: None, constraints: Constraints::default(), column_defaults: HashMap::new(), expr_adapter_factory: config.expr_adapter_factory, @@ -260,10 +259,8 @@ impl ListingTable { /// Setting a statistics cache on the `SessionContext` can avoid refetching statistics /// multiple times in the same session. /// - /// If `None`, creates a new [`DefaultFileStatisticsCache`] scoped to this query. pub fn with_cache(mut self, cache: Option<Arc<dyn FileStatisticsCache>>) -> Self { - self.collected_statistics = - cache.unwrap_or_else(|| Arc::new(DefaultFileStatisticsCache::default())); + self.collected_statistics = cache; self } @@ -802,11 +799,15 @@ impl ListingTable { ) -> datafusion_common::Result<(Arc<Statistics>, Option<LexOrdering>)> { use datafusion_execution::cache::cache_manager::CachedFileMetadata; - let path = &part_file.object_meta.location; + let path = TableScopedPath { + table: part_file.table_reference.clone(), + path: part_file.object_meta.location.clone(), + }; let meta = &part_file.object_meta; // Check cache first - if we have valid cached statistics and ordering - if let Some(cached) = self.collected_statistics.get(path) + if let Some(cache) = &self.collected_statistics + && let Some(cached) = cache.get(&path) && cached.is_valid_for(meta) { // Return cached statistics and ordering @@ -823,14 +824,16 @@ impl ListingTable { let statistics = Arc::new(file_meta.statistics); // Store in cache - self.collected_statistics.put( - path, - CachedFileMetadata::new( - meta.clone(), - Arc::clone(&statistics), - file_meta.ordering.clone(), - ), - ); + if let Some(cache) = &self.collected_statistics { + cache.put( + &path, + CachedFileMetadata::new( + meta.clone(), + Arc::clone(&statistics), + file_meta.ordering.clone(), + ), + ); + } Ok((statistics, file_meta.ordering)) } diff --git a/datafusion/common/src/heap_size.rs b/datafusion/common/src/heap_size.rs new file mode 100644 index 0000000000..edb64709d5 --- /dev/null +++ b/datafusion/common/src/heap_size.rs @@ -0,0 +1,561 @@ +// 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 crate::stats::Precision; +use crate::{ColumnStatistics, ScalarValue, Statistics, TableReference}; +use arrow::array::{ + Array, FixedSizeListArray, LargeListArray, LargeListViewArray, ListArray, + ListViewArray, MapArray, StructArray, +}; +use arrow::datatypes::{ + DataType, Field, Fields, IntervalDayTime, IntervalMonthDayNano, IntervalUnit, + TimeUnit, UnionFields, UnionMode, i256, +}; +use chrono::{DateTime, Utc}; +use half::f16; +use hashbrown::HashSet; +use std::collections::HashMap; +use std::fmt::Debug; +use std::sync::Arc; + +/// This is a temporary solution until <https://github.com/apache/datafusion/pull/19599> and +/// <https://github.com/apache/arrow-rs/pull/9138> are resolved. +/// Trait for calculating the size of various containers +pub trait DFHeapSize { + /// Return the size of any bytes allocated on the heap by this object, + /// including heap memory in those structures + /// + /// Note that the size of the type itself is not included in the result -- + /// instead, that size is added by the caller (e.g. container). + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize; +} + +#[derive(Default)] +pub struct DFHeapSizeCtx { + seen: HashSet<usize>, +} + +impl DFHeapSize for Statistics { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.num_rows.heap_size(ctx) + + self.total_byte_size.heap_size(ctx) + + self.column_statistics.heap_size(ctx) + } +} + +impl DFHeapSize for TableReference { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + match self { + TableReference::Bare { table } => table.heap_size(ctx), + TableReference::Partial { schema, table } => { + schema.heap_size(ctx) + table.heap_size(ctx) + } + TableReference::Full { + catalog, + schema, + table, + } => catalog.heap_size(ctx) + schema.heap_size(ctx) + table.heap_size(ctx), + } + } +} + +impl<T: Debug + Clone + PartialEq + Eq + PartialOrd + DFHeapSize> DFHeapSize + for Precision<T> +{ + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.get_value().map_or_else(|| 0, |v| v.heap_size(ctx)) + } +} + +impl DFHeapSize for ColumnStatistics { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.null_count.heap_size(ctx) + + self.max_value.heap_size(ctx) + + self.min_value.heap_size(ctx) + + self.sum_value.heap_size(ctx) + + self.distinct_count.heap_size(ctx) + + self.byte_size.heap_size(ctx) + } +} + +impl DFHeapSize for ScalarValue { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + use crate::scalar::ScalarValue::*; + match self { + Null => 0, + Boolean(b) => b.heap_size(ctx), + Float16(f) => f.heap_size(ctx), + Float32(f) => f.heap_size(ctx), + Float64(f) => f.heap_size(ctx), + Decimal32(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx), + Decimal64(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx), + Decimal128(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx), + Decimal256(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx), + Int8(i) => i.heap_size(ctx), + Int16(i) => i.heap_size(ctx), + Int32(i) => i.heap_size(ctx), + Int64(i) => i.heap_size(ctx), + UInt8(u) => u.heap_size(ctx), + UInt16(u) => u.heap_size(ctx), + UInt32(u) => u.heap_size(ctx), + UInt64(u) => u.heap_size(ctx), + Utf8(u) => u.heap_size(ctx), + Utf8View(u) => u.heap_size(ctx), + LargeUtf8(l) => l.heap_size(ctx), + Binary(b) => b.heap_size(ctx), + BinaryView(b) => b.heap_size(ctx), + FixedSizeBinary(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + LargeBinary(l) => l.heap_size(ctx), + FixedSizeList(f) => f.heap_size(ctx), + List(l) => l.heap_size(ctx), + LargeList(l) => l.heap_size(ctx), + Struct(s) => s.heap_size(ctx), + Map(m) => m.heap_size(ctx), + Date32(d) => d.heap_size(ctx), + Date64(d) => d.heap_size(ctx), + Time32Second(t) => t.heap_size(ctx), + Time32Millisecond(t) => t.heap_size(ctx), + Time64Microsecond(t) => t.heap_size(ctx), + Time64Nanosecond(t) => t.heap_size(ctx), + TimestampSecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + TimestampMillisecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + TimestampMicrosecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + TimestampNanosecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + IntervalYearMonth(i) => i.heap_size(ctx), + IntervalDayTime(i) => i.heap_size(ctx), + IntervalMonthDayNano(i) => i.heap_size(ctx), + DurationSecond(d) => d.heap_size(ctx), + DurationMillisecond(d) => d.heap_size(ctx), + DurationMicrosecond(d) => d.heap_size(ctx), + DurationNanosecond(d) => d.heap_size(ctx), + Union(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx), + Dictionary(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + RunEndEncoded(a, b, c) => { + a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx) + } + ListView(a) => a.heap_size(ctx), + LargeListView(a) => a.heap_size(ctx), + } + } +} + +impl DFHeapSize for DataType { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + use DataType::*; + match self { + Null => 0, + Boolean => 0, + Int8 => 0, + Int16 => 0, + Int32 => 0, + Int64 => 0, + UInt8 => 0, + UInt16 => 0, + UInt32 => 0, + UInt64 => 0, + Float16 => 0, + Float32 => 0, + Float64 => 0, + Timestamp(t, s) => t.heap_size(ctx) + s.heap_size(ctx), + Date32 => 0, + Date64 => 0, + Time32(t) => t.heap_size(ctx), + Time64(t) => t.heap_size(ctx), + Duration(t) => t.heap_size(ctx), + Interval(i) => i.heap_size(ctx), + Binary => 0, + FixedSizeBinary(i) => i.heap_size(ctx), + LargeBinary => 0, + BinaryView => 0, + Utf8 => 0, + LargeUtf8 => 0, + Utf8View => 0, + List(v) => v.heap_size(ctx), + ListView(v) => v.heap_size(ctx), + FixedSizeList(f, i) => f.heap_size(ctx) + i.heap_size(ctx), + LargeList(l) => l.heap_size(ctx), + LargeListView(l) => l.heap_size(ctx), + Struct(s) => s.heap_size(ctx), + Union(u, m) => u.heap_size(ctx) + m.heap_size(ctx), + Dictionary(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + Decimal32(p, s) => p.heap_size(ctx) + s.heap_size(ctx), + Decimal64(p, s) => p.heap_size(ctx) + s.heap_size(ctx), + Decimal128(p, s) => p.heap_size(ctx) + s.heap_size(ctx), + Decimal256(p, s) => p.heap_size(ctx) + s.heap_size(ctx), + Map(m, b) => m.heap_size(ctx) + b.heap_size(ctx), + RunEndEncoded(a, b) => a.heap_size(ctx) + b.heap_size(ctx), + } + } +} + +impl<T: DFHeapSize> DFHeapSize for Vec<T> { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + let item_size = size_of::<T>(); + // account for the contents of the Vec + (self.capacity() * item_size) + + // add any heap allocations by contents + self.iter().map(|t| t.heap_size(ctx)).sum::<usize>() + } +} + +impl<K: DFHeapSize, V: DFHeapSize> DFHeapSize for HashMap<K, V> { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + let capacity = self.capacity(); + if capacity == 0 { + return 0; + } + + // HashMap doesn't provide a way to get its heap size, so this is an approximation based on + // the behavior of hashbrown::HashMap as at version 0.16.0, and may become inaccurate + // if the implementation changes. + let key_val_size = size_of::<(K, V)>(); + // Overhead for the control tags group, which may be smaller depending on architecture + let group_size = 16; + // 1 byte of metadata stored per bucket. + let metadata_size = 1; + + // Compute the number of buckets for the capacity. Based on hashbrown's capacity_to_buckets + let buckets = if capacity < 15 { + let min_cap = match key_val_size { + 0..=1 => 14, + 2..=3 => 7, + _ => 3, + }; + let cap = min_cap.max(capacity); + if cap < 4 { + 4 + } else if cap < 8 { + 8 + } else { + 16 + } + } else { + (capacity.saturating_mul(8) / 7).next_power_of_two() + }; + + group_size + + (buckets * (key_val_size + metadata_size)) + + self.keys().map(|k| k.heap_size(ctx)).sum::<usize>() + + self.values().map(|v| v.heap_size(ctx)).sum::<usize>() + } +} + +impl<T: DFHeapSize> DFHeapSize for Arc<T> { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + let ptr = Arc::as_ptr(self) as usize; + + if !ctx.seen.insert(ptr) { + return 0; + } + + // Arc stores weak and strong counts on the heap alongside an instance of T + 2 * size_of::<usize>() + size_of::<T>() + self.as_ref().heap_size(ctx) + } +} + +impl DFHeapSize for Arc<str> { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + let ptr = Arc::as_ptr(self) as *const i32 as usize; + + if !ctx.seen.insert(ptr) { + return 0; + } + + // Arc stores weak and strong counts on the heap alongside an instance of T + 2 * size_of::<usize>() + self.as_ref().heap_size(ctx) + } +} + +impl DFHeapSize for Arc<dyn DFHeapSize> { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + let ptr = Arc::as_ptr(self) as *const i32 as usize; + + if !ctx.seen.insert(ptr) { + return 0; + } + + // Arc stores weak and strong counts on the heap alongside an instance of T + 2 * size_of::<usize>() + size_of_val(self.as_ref()) + self.as_ref().heap_size(ctx) + } +} + +impl DFHeapSize for Fields { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.into_iter().map(|f| f.heap_size(ctx)).sum::<usize>() + } +} + +impl DFHeapSize for StructArray { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } +} + +impl DFHeapSize for LargeListArray { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } +} + +impl DFHeapSize for LargeListViewArray { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } +} + +impl DFHeapSize for ListArray { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } +} + +impl DFHeapSize for ListViewArray { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } +} + +impl DFHeapSize for FixedSizeListArray { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } +} +impl DFHeapSize for MapArray { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } +} + +impl<T: DFHeapSize> DFHeapSize for Box<T> { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + size_of::<T>() + self.as_ref().heap_size(ctx) + } +} + +impl<T: DFHeapSize> DFHeapSize for Option<T> { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.as_ref().map(|inner| inner.heap_size(ctx)).unwrap_or(0) + } +} + +impl<A, B> DFHeapSize for (A, B) +where + A: DFHeapSize, + B: DFHeapSize, +{ + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.0.heap_size(ctx) + self.1.heap_size(ctx) + } +} + +impl DFHeapSize for String { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.capacity() + } +} + +impl DFHeapSize for str { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.len() + } +} + +impl DFHeapSize for UnionFields { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.iter() + .map(|f| f.0.heap_size(ctx) + f.1.heap_size(ctx)) + .sum() + } +} + +impl DFHeapSize for UnionMode { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for TimeUnit { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for IntervalUnit { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for Field { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.name().heap_size(ctx) + + self.data_type().heap_size(ctx) + + self.is_nullable().heap_size(ctx) + + self.dict_is_ordered().heap_size(ctx) + + self.metadata().heap_size(ctx) + } +} + +impl DFHeapSize for IntervalMonthDayNano { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.days.heap_size(ctx) + + self.months.heap_size(ctx) + + self.nanoseconds.heap_size(ctx) + } +} + +impl DFHeapSize for IntervalDayTime { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.days.heap_size(ctx) + self.milliseconds.heap_size(ctx) + } +} + +impl DFHeapSize for DateTime<Utc> { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for bool { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} +impl DFHeapSize for u8 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for u16 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for u32 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for u64 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for i8 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for i16 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for i32 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} +impl DFHeapSize for i64 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for i128 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for i256 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for f16 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for f32 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} +impl DFHeapSize for f64 { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +impl DFHeapSize for usize { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_heap_size_arc_avoid_double_accounting() { + let a1 = Arc::new(vec![1, 2, 3]); + let mut ctx = DFHeapSizeCtx::default(); + let heap_size = a1.heap_size(&mut ctx); + + let a2 = Arc::clone(&a1); + let a3 = Arc::clone(&a1); + let a4 = Arc::clone(&a3); + + let mut ctx = DFHeapSizeCtx::default(); + let heap_size_with_clones = a1.heap_size(&mut ctx) + + a2.heap_size(&mut ctx) + + a3.heap_size(&mut ctx) + + a4.heap_size(&mut ctx); + + assert_eq!(heap_size, heap_size_with_clones); + } + + #[test] + fn test_heap_size_arc_str_avoid_double_accounting() { + let a1: Arc<str> = Arc::from("Hello"); + let mut ctx = DFHeapSizeCtx::default(); + let heap_size = a1.heap_size(&mut ctx); + + let a2 = Arc::clone(&a1); + let a3 = Arc::clone(&a1); + let a4 = Arc::clone(&a3); + + let mut ctx = DFHeapSizeCtx::default(); + let heap_size_with_clones = a1.heap_size(&mut ctx) + + a2.heap_size(&mut ctx) + + a3.heap_size(&mut ctx) + + a4.heap_size(&mut ctx); + + assert_eq!(heap_size, heap_size_with_clones); + } +} diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index 06e567cb12..e865c548bb 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -47,6 +47,7 @@ pub mod extensions; pub mod file_options; pub mod format; pub mod hash_utils; +pub mod heap_size; pub mod instant; pub mod metadata; pub mod nested_struct; @@ -62,6 +63,7 @@ pub mod test_util; pub mod tree_node; pub mod types; pub mod utils; + /// Reexport arrow crate pub use arrow; pub use column::Column; diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs index 7fbf5696de..a1eb7ffb64 100644 --- a/datafusion/core/src/datasource/listing_table_factory.rs +++ b/datafusion/core/src/datasource/listing_table_factory.rs @@ -231,7 +231,7 @@ mod tests { }; use datafusion_execution::cache::CacheAccessor; use datafusion_execution::cache::cache_manager::CacheManagerConfig; - use datafusion_execution::cache::cache_unit::DefaultFileStatisticsCache; + use datafusion_execution::cache::file_statistics_cache::DefaultFileStatisticsCache; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use glob::Pattern; @@ -486,7 +486,7 @@ mod tests { // Test with collect_statistics enabled let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default()); let cache_config = CacheManagerConfig::default() - .with_files_statistics_cache(Some(file_statistics_cache.clone())); + .with_file_statistics_cache(Some(file_statistics_cache.clone())); let runtime = RuntimeEnvBuilder::new() .with_cache_manager(cache_config) .build_arc() @@ -516,7 +516,7 @@ mod tests { // Test with collect_statistics disabled let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default()); let cache_config = CacheManagerConfig::default() - .with_files_statistics_cache(Some(file_statistics_cache.clone())); + .with_file_statistics_cache(Some(file_statistics_cache.clone())); let runtime = RuntimeEnvBuilder::new() .with_cache_manager(cache_config) .build_arc() diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index d84ef0c898..72d6c4c3a6 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -102,6 +102,7 @@ use datafusion_session::SessionStore; use async_trait::async_trait; use chrono::{DateTime, Utc}; +use datafusion_execution::cache::file_statistics_cache::DEFAULT_FILE_STATISTICS_MEMORY_LIMIT; use object_store::ObjectStore; use parking_lot::RwLock; use url::Url; @@ -1183,6 +1184,10 @@ impl SessionContext { let duration = Self::parse_duration(variable, value)?; builder.with_object_list_cache_ttl(Some(duration)) } + "file_statistics_cache_limit" => { + let limit = Self::parse_capacity_limit(variable, value)?; + builder.with_file_statistics_cache_limit(limit) + } _ => return plan_err!("Unknown runtime configuration: {variable}"), // Remember to update `reset_runtime_variable()` when adding new options }; @@ -1222,9 +1227,13 @@ impl SessionContext { builder = builder.with_object_list_cache_ttl(DEFAULT_LIST_FILES_CACHE_TTL); } + "file_statistics_cache_limit" => { + builder = builder.with_file_statistics_cache_limit( + DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, + ); + } _ => return plan_err!("Unknown runtime configuration: {variable}"), }; - *state = SessionStateBuilder::from(state.clone()) .with_runtime_env(Arc::new(builder.build()?)) .build(); @@ -1420,17 +1429,29 @@ impl SessionContext { && table_provider.table_type() == table_type { schema.deregister_table(&table)?; - if table_type == TableType::Base - && let Some(lfc) = self.runtime_env().cache_manager.get_list_files_cache() - { - lfc.drop_table_entries(&Some(table_ref))?; - } + self.invalidate_caches(&Some(table_ref.clone()), table_type)?; return Ok(true); } - Ok(false) } + fn invalidate_caches( + &self, + table_ref: &Option<TableReference>, + table_type: TableType, + ) -> Result<()> { + if table_type == TableType::Base { + if let Some(lfc) = self.runtime_env().cache_manager.get_list_files_cache() { + lfc.drop_table_entries(table_ref)?; + } + if let Some(fsc) = self.runtime_env().cache_manager.get_file_statistic_cache() + { + fsc.drop_table_entries(table_ref)?; + } + } + Ok(()) + } + async fn create_function(&self, stmt: CreateFunction) -> Result<DataFrame> { let function = { let state = self.state.read().clone(); @@ -1661,7 +1682,8 @@ impl SessionContext { let config = ListingTableConfig::new_with_multi_paths(table_paths) .with_listing_options(listing_options) .with_schema(resolved_schema); - let provider = ListingTable::try_new(config)?; + let provider = ListingTable::try_new(config)? + .with_cache(self.runtime_env().cache_manager.get_file_statistic_cache()); self.read_table(Arc::new(provider)) } @@ -1748,7 +1770,9 @@ impl SessionContext { provided_schema: Option<SchemaRef>, sql_definition: Option<String>, ) -> Result<()> { - let table_path = ListingTableUrl::parse(table_path)?; + let table_ref = table_ref.into(); + let table_path = + ListingTableUrl::parse(table_path)?.with_table_ref(table_ref.clone()); let resolved_schema = match provided_schema { Some(s) => s, None => options.infer_schema(&self.state(), &table_path).await?, @@ -1756,7 +1780,9 @@ impl SessionContext { let config = ListingTableConfig::new(table_path) .with_listing_options(options) .with_schema(resolved_schema); - let table = ListingTable::try_new(config)?.with_definition(sql_definition); + let table = ListingTable::try_new(config)? + .with_definition(sql_definition) + .with_cache(self.runtime_env().cache_manager.get_file_statistic_cache()); self.register_table(table_ref, Arc::new(table))?; Ok(()) } @@ -1860,10 +1886,17 @@ impl SessionContext { ) -> Result<Option<Arc<dyn TableProvider>>> { let table_ref = table_ref.into(); let table = table_ref.table().to_owned(); - self.state + let result = self + .state .read() - .schema_for_ref(table_ref)? - .deregister_table(&table) + .schema_for_ref(table_ref.clone())? + .deregister_table(&table); + + if let Ok(Some(ref table_provider)) = result { + self.invalidate_caches(&Some(table_ref), table_provider.table_type())?; + } + + result } /// Return `true` if the specified table exists in the schema provider. diff --git a/datafusion/core/tests/parquet/file_statistics.rs b/datafusion/core/tests/parquet/file_statistics.rs index 84396be8a6..3e3b90a348 100644 --- a/datafusion/core/tests/parquet/file_statistics.rs +++ b/datafusion/core/tests/parquet/file_statistics.rs @@ -30,8 +30,10 @@ use datafusion::prelude::SessionContext; use datafusion_common::DFSchema; use datafusion_common::stats::Precision; use datafusion_execution::cache::DefaultListFilesCache; -use datafusion_execution::cache::cache_manager::CacheManagerConfig; -use datafusion_execution::cache::cache_unit::DefaultFileStatisticsCache; +use datafusion_execution::cache::cache_manager::{ + CacheManagerConfig, FileStatisticsCache, +}; +use datafusion_execution::cache::file_statistics_cache::DefaultFileStatisticsCache; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::{Expr, col, lit}; @@ -236,7 +238,7 @@ async fn list_files_with_session_level_cache() { async fn get_listing_table( table_path: &ListingTableUrl, - static_cache: Option<Arc<DefaultFileStatisticsCache>>, + static_cache: Option<Arc<dyn FileStatisticsCache>>, opt: &ListingOptions, ) -> ListingTable { let schema = opt @@ -249,12 +251,9 @@ async fn get_listing_table( let config1 = ListingTableConfig::new(table_path.clone()) .with_listing_options(opt.clone()) .with_schema(schema); - let table = ListingTable::try_new(config1).unwrap(); - if let Some(c) = static_cache { - table.with_cache(Some(c)) - } else { - table - } + ListingTable::try_new(config1) + .unwrap() + .with_cache(static_cache) } fn get_cache_runtime_state() -> ( @@ -267,7 +266,7 @@ fn get_cache_runtime_state() -> ( let list_file_cache = Arc::new(DefaultListFilesCache::default()); let cache_config = cache_config - .with_files_statistics_cache(Some(file_static_cache.clone())) + .with_file_statistics_cache(Some(file_static_cache.clone())) .with_list_files_cache(Some(list_file_cache.clone())); let rt = RuntimeEnvBuilder::new() diff --git a/datafusion/core/tests/sql/runtime_config.rs b/datafusion/core/tests/sql/runtime_config.rs index ccc11afbff..2db1e1ce12 100644 --- a/datafusion/core/tests/sql/runtime_config.rs +++ b/datafusion/core/tests/sql/runtime_config.rs @@ -25,6 +25,7 @@ use datafusion::execution::context::TaskContext; use datafusion::prelude::SessionConfig; use datafusion_execution::cache::DefaultListFilesCache; use datafusion_execution::cache::cache_manager::CacheManagerConfig; +use datafusion_execution::cache::file_statistics_cache::DefaultFileStatisticsCache; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_plan::common::collect; @@ -344,6 +345,51 @@ async fn test_list_files_cache_ttl() { assert_eq!(get_limit(&ctx), Duration::from_secs(90)); } +#[tokio::test] +async fn test_file_statistics_cache_limit() { + let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default()); + + let rt = RuntimeEnvBuilder::new() + .with_cache_manager( + CacheManagerConfig::default() + .with_file_statistics_cache(Some(file_statistics_cache)), + ) + .build_arc() + .unwrap(); + + let ctx = SessionContext::new_with_config_rt(SessionConfig::default(), rt); + + let update_limit = async |ctx: &SessionContext, limit: &str| { + ctx.sql( + format!("SET datafusion.runtime.file_statistics_cache_limit = '{limit}'") + .as_str(), + ) + .await + .unwrap() + .collect() + .await + .unwrap(); + }; + + let get_limit = |ctx: &SessionContext| -> usize { + ctx.task_ctx() + .runtime_env() + .cache_manager + .get_file_statistic_cache() + .unwrap() + .cache_limit() + }; + + update_limit(&ctx, "1M").await; + assert_eq!(get_limit(&ctx), 1024 * 1024); + + update_limit(&ctx, "42G").await; + assert_eq!(get_limit(&ctx), 42 * 1024 * 1024 * 1024); + + update_limit(&ctx, "23K").await; + assert_eq!(get_limit(&ctx), 23 * 1024); +} + #[tokio::test] async fn test_unknown_runtime_config() { let ctx = SessionContext::new(); diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index c5f9d72200..84daf608b5 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -56,7 +56,7 @@ pub use self::url::ListingTableUrl; use crate::file_groups::FileGroup; use chrono::TimeZone; use datafusion_common::stats::Precision; -use datafusion_common::{ColumnStatistics, Result, exec_datafusion_err}; +use datafusion_common::{ColumnStatistics, Result, TableReference, exec_datafusion_err}; use datafusion_common::{ScalarValue, Statistics}; use datafusion_physical_expr::LexOrdering; use futures::{Stream, StreamExt}; @@ -162,6 +162,7 @@ pub struct PartitionedFile { pub extensions: FileExtensions, /// The estimated size of the parquet metadata, in bytes pub metadata_size_hint: Option<usize>, + pub table_reference: Option<TableReference>, } impl PartitionedFile { @@ -181,6 +182,7 @@ impl PartitionedFile { ordering: None, extensions: FileExtensions::new(), metadata_size_hint: None, + table_reference: None, } } @@ -194,6 +196,7 @@ impl PartitionedFile { ordering: None, extensions: FileExtensions::new(), metadata_size_hint: None, + table_reference: None, } } @@ -213,6 +216,7 @@ impl PartitionedFile { ordering: None, extensions: FileExtensions::new(), metadata_size_hint: None, + table_reference: None, } .with_range(start, end) } @@ -224,6 +228,14 @@ impl PartitionedFile { self } + pub fn with_table_reference( + mut self, + table_reference: Option<TableReference>, + ) -> Self { + self.table_reference = table_reference; + self + } + /// Size of the file to be scanned (taking into account the range, if present). pub fn effective_size(&self) -> u64 { if let Some(range) = &self.range { @@ -370,6 +382,7 @@ impl From<ObjectMeta> for PartitionedFile { ordering: None, extensions: FileExtensions::new(), metadata_size_hint: None, + table_reference: None, } } } @@ -567,6 +580,7 @@ pub fn generate_test_files(num_files: usize, overlap_factor: f64) -> Vec<FileGro ordering: None, extensions: FileExtensions::new(), metadata_size_hint: None, + table_reference: None, }; files.push(file); } diff --git a/datafusion/execution/src/cache/cache_manager.rs b/datafusion/execution/src/cache/cache_manager.rs index 0868c968c3..08a8dc9fd9 100644 --- a/datafusion/execution/src/cache/cache_manager.rs +++ b/datafusion/execution/src/cache/cache_manager.rs @@ -17,10 +17,14 @@ use crate::cache::CacheAccessor; use crate::cache::DefaultListFilesCache; -use crate::cache::cache_unit::DefaultFilesMetadataCache; +use crate::cache::file_statistics_cache::{ + DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, DefaultFileStatisticsCache, + DefaultFilesMetadataCache, +}; use crate::cache::list_files_cache::ListFilesEntry; use crate::cache::list_files_cache::TableScopedPath; use datafusion_common::TableReference; +use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use datafusion_common::stats::Precision; use datafusion_common::{Result, Statistics}; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -41,7 +45,7 @@ pub use super::list_files_cache::{ /// /// This struct embeds the [`ObjectMeta`] used for cache validation, /// along with the cached statistics and ordering information. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct CachedFileMetadata { /// File metadata used for cache validation (size, last_modified). pub meta: ObjectMeta, @@ -81,7 +85,7 @@ impl CachedFileMetadata { /// - Statistics for the file /// - Ordering information for the file /// -/// If enabled via [`CacheManagerConfig::with_files_statistics_cache`] this +/// If enabled via [`CacheManagerConfig::with_file_statistics_cache`] this /// cache avoids inferring the same file statistics repeatedly during the /// session lifetime. /// @@ -91,9 +95,31 @@ impl CachedFileMetadata { /// 3. If invalid or missing, compute new value and call `put(path, new_value)` /// /// See [`crate::runtime_env::RuntimeEnv`] for more details -pub trait FileStatisticsCache: CacheAccessor<Path, CachedFileMetadata> { +pub trait FileStatisticsCache: + CacheAccessor<TableScopedPath, CachedFileMetadata> +{ + /// Cache memory limit in bytes. + fn cache_limit(&self) -> usize; + + /// Updates the cache with a new memory limit in bytes. + fn update_cache_limit(&self, limit: usize); + /// Retrieves the information about the entries currently cached. - fn list_entries(&self) -> HashMap<Path, FileStatisticsCacheEntry>; + fn list_entries(&self) -> HashMap<TableScopedPath, FileStatisticsCacheEntry>; + + fn drop_table_entries(&self, table_ref: &Option<TableReference>) -> Result<()>; +} + +impl DFHeapSize for CachedFileMetadata { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.meta.size.heap_size(ctx) + + self.meta.last_modified.heap_size(ctx) + + self.meta.version.heap_size(ctx) + + self.meta.e_tag.heap_size(ctx) + + self.meta.location.as_ref().heap_size(ctx) + + self.statistics.heap_size(ctx) + //TODO add ordering once LexOrdering/PhysicalExpr implements DFHeapSize + } } /// Represents information about a cached statistics entry. @@ -330,8 +356,19 @@ pub struct CacheManager { impl CacheManager { pub fn try_new(config: &CacheManagerConfig) -> Result<Arc<Self>> { - let file_statistic_cache = - config.table_files_statistics_cache.as_ref().map(Arc::clone); + let file_statistic_cache = match &config.file_statistics_cache { + Some(fsc) if config.file_statistics_cache_limit > 0 => { + fsc.update_cache_limit(config.file_statistics_cache_limit); + Some(Arc::clone(fsc)) + } + None if config.file_statistics_cache_limit > 0 => { + let fsc: Arc<dyn FileStatisticsCache> = Arc::new( + DefaultFileStatisticsCache::new(config.file_statistics_cache_limit), + ); + Some(fsc) + } + _ => None, + }; let list_files_cache = match &config.list_files_cache { Some(lfc) if config.list_files_cache_limit > 0 => { @@ -371,11 +408,18 @@ impl CacheManager { })) } - /// Get the cache of listing files statistics. + /// Get the file statistics cache. pub fn get_file_statistic_cache(&self) -> Option<Arc<dyn FileStatisticsCache>> { self.file_statistic_cache.clone() } + /// Get the memory limit of the file statistics cache. + pub fn get_file_statistic_cache_limit(&self) -> usize { + self.file_statistic_cache + .as_ref() + .map_or(0, |c| c.cache_limit()) + } + /// Get the cache for storing the result of listing [`ObjectMeta`]s under the same path. pub fn get_list_files_cache(&self) -> Option<Arc<dyn ListFilesCache>> { self.list_files_cache.clone() @@ -410,15 +454,17 @@ pub const DEFAULT_METADATA_CACHE_LIMIT: usize = 50 * 1024 * 1024; // 50M pub struct CacheManagerConfig { /// Enable caching of file statistics when listing files. /// Enabling the cache avoids repeatedly reading file statistics in a DataFusion session. - /// Default is disabled. Currently only Parquet files are supported. - pub table_files_statistics_cache: Option<Arc<dyn FileStatisticsCache>>, + /// Default is enabled. Currently only Parquet files are supported. + pub file_statistics_cache: Option<Arc<dyn FileStatisticsCache>>, + /// Limit of the file statistics cache, in bytes. Default: 20MiB. + pub file_statistics_cache_limit: usize, /// Enable caching of file metadata when listing files. /// Enabling the cache avoids repeat list and object metadata fetch operations, which may be /// expensive in certain situations (e.g. remote object storage), for objects under paths that /// are cached. /// Note that if this option is enabled, DataFusion will not see any updates to the underlying /// storage for at least `list_files_cache_ttl` duration. - /// Default is disabled. + /// Default is enabled. pub list_files_cache: Option<Arc<dyn ListFilesCache>>, /// Limit of the `list_files_cache`, in bytes. Default: 1MiB. pub list_files_cache_limit: usize, @@ -437,7 +483,8 @@ pub struct CacheManagerConfig { impl Default for CacheManagerConfig { fn default() -> Self { Self { - table_files_statistics_cache: Default::default(), + file_statistics_cache: Default::default(), + file_statistics_cache_limit: DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, list_files_cache: Default::default(), list_files_cache_limit: DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, list_files_cache_ttl: DEFAULT_LIST_FILES_CACHE_TTL, @@ -448,14 +495,18 @@ impl Default for CacheManagerConfig { } impl CacheManagerConfig { - /// Set the cache for files statistics. - /// - /// Default is `None` (disabled). - pub fn with_files_statistics_cache( + /// Set the cache for file statistics. + pub fn with_file_statistics_cache( mut self, cache: Option<Arc<dyn FileStatisticsCache>>, ) -> Self { - self.table_files_statistics_cache = cache; + self.file_statistics_cache = cache; + self + } + + /// Specifies the memory limit for the file statistics cache, in bytes. + pub fn with_file_statistics_cache_limit(mut self, limit: usize) -> Self { + self.file_statistics_cache_limit = limit; self } diff --git a/datafusion/execution/src/cache/cache_unit.rs b/datafusion/execution/src/cache/cache_unit.rs deleted file mode 100644 index 49e16ca4b6..0000000000 --- a/datafusion/execution/src/cache/cache_unit.rs +++ /dev/null @@ -1,410 +0,0 @@ -// 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 std::collections::HashMap; - -use crate::cache::CacheAccessor; -use crate::cache::cache_manager::{ - CachedFileMetadata, FileStatisticsCache, FileStatisticsCacheEntry, -}; - -use dashmap::DashMap; -use object_store::path::Path; - -pub use crate::cache::DefaultFilesMetadataCache; - -/// Default implementation of [`FileStatisticsCache`] -/// -/// Stores cached file metadata (statistics and orderings) for files. -/// -/// The typical usage pattern is: -/// 1. Call `get(path)` to check for cached value -/// 2. If `Some(cached)`, validate with `cached.is_valid_for(¤t_meta)` -/// 3. If invalid or missing, compute new value and call `put(path, new_value)` -/// -/// Uses DashMap for lock-free concurrent access. -/// -/// [`FileStatisticsCache`]: crate::cache::cache_manager::FileStatisticsCache -#[derive(Default)] -pub struct DefaultFileStatisticsCache { - cache: DashMap<Path, CachedFileMetadata>, -} - -impl CacheAccessor<Path, CachedFileMetadata> for DefaultFileStatisticsCache { - fn get(&self, key: &Path) -> Option<CachedFileMetadata> { - self.cache.get(key).map(|entry| entry.value().clone()) - } - - fn put(&self, key: &Path, value: CachedFileMetadata) -> Option<CachedFileMetadata> { - self.cache.insert(key.clone(), value) - } - - fn remove(&self, k: &Path) -> Option<CachedFileMetadata> { - self.cache.remove(k).map(|(_, entry)| entry) - } - - fn contains_key(&self, k: &Path) -> bool { - self.cache.contains_key(k) - } - - fn len(&self) -> usize { - self.cache.len() - } - - fn clear(&self) { - self.cache.clear(); - } - - fn name(&self) -> String { - "DefaultFileStatisticsCache".to_string() - } -} - -impl FileStatisticsCache for DefaultFileStatisticsCache { - fn list_entries(&self) -> HashMap<Path, FileStatisticsCacheEntry> { - let mut entries = HashMap::<Path, FileStatisticsCacheEntry>::new(); - - for entry in self.cache.iter() { - let path = entry.key(); - let cached = entry.value(); - entries.insert( - path.clone(), - FileStatisticsCacheEntry { - object_meta: cached.meta.clone(), - num_rows: cached.statistics.num_rows, - num_columns: cached.statistics.column_statistics.len(), - table_size_bytes: cached.statistics.total_byte_size, - statistics_size_bytes: 0, // TODO: set to the real size in the future - has_ordering: cached.ordering.is_some(), - }, - ); - } - - entries - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::cache::cache_manager::{ - CachedFileMetadata, FileStatisticsCache, FileStatisticsCacheEntry, - }; - use arrow::array::RecordBatch; - use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; - use chrono::DateTime; - use datafusion_common::Statistics; - use datafusion_common::stats::Precision; - use datafusion_expr::ColumnarValue; - use datafusion_physical_expr_common::physical_expr::PhysicalExpr; - use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; - use object_store::ObjectMeta; - use std::sync::Arc; - - fn create_test_meta(path: &str, size: u64) -> ObjectMeta { - ObjectMeta { - location: Path::from(path), - last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") - .unwrap() - .into(), - size, - e_tag: None, - version: None, - } - } - - #[test] - fn test_statistics_cache() { - let meta = create_test_meta("test", 1024); - let cache = DefaultFileStatisticsCache::default(); - - let schema = Schema::new(vec![Field::new( - "test_column", - DataType::Timestamp(TimeUnit::Second, None), - false, - )]); - - // Cache miss - assert!(cache.get(&meta.location).is_none()); - - // Put a value - let cached_value = CachedFileMetadata::new( - meta.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&meta.location, cached_value); - - // Cache hit - let result = cache.get(&meta.location); - assert!(result.is_some()); - let cached = result.unwrap(); - assert!(cached.is_valid_for(&meta)); - - // File size changed - validation should fail - let meta2 = create_test_meta("test", 2048); - let cached = cache.get(&meta2.location).unwrap(); - assert!(!cached.is_valid_for(&meta2)); - - // Update with new value - let cached_value2 = CachedFileMetadata::new( - meta2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&meta2.location, cached_value2); - - // Test list_entries - let entries = cache.list_entries(); - assert_eq!(entries.len(), 1); - let entry = entries.get(&Path::from("test")).unwrap(); - assert_eq!(entry.object_meta.size, 2048); // Should be updated value - } - - #[derive(Clone, Debug, PartialEq, Eq, Hash)] - struct MockExpr {} - - impl std::fmt::Display for MockExpr { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "MockExpr") - } - } - - impl PhysicalExpr for MockExpr { - fn data_type( - &self, - _input_schema: &Schema, - ) -> datafusion_common::Result<DataType> { - Ok(DataType::Int32) - } - - fn nullable(&self, _input_schema: &Schema) -> datafusion_common::Result<bool> { - Ok(false) - } - - fn evaluate( - &self, - _batch: &RecordBatch, - ) -> datafusion_common::Result<ColumnarValue> { - unimplemented!() - } - - fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { - vec![] - } - - fn with_new_children( - self: Arc<Self>, - children: Vec<Arc<dyn PhysicalExpr>>, - ) -> datafusion_common::Result<Arc<dyn PhysicalExpr>> { - assert!(children.is_empty()); - Ok(self) - } - - fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "MockExpr") - } - } - - fn ordering() -> LexOrdering { - let expr = Arc::new(MockExpr {}) as Arc<dyn PhysicalExpr>; - LexOrdering::new(vec![PhysicalSortExpr::new_default(expr)]).unwrap() - } - - #[test] - fn test_ordering_cache() { - let meta = create_test_meta("test.parquet", 100); - let cache = DefaultFileStatisticsCache::default(); - - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - // Cache statistics with no ordering - let cached_value = CachedFileMetadata::new( - meta.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, // No ordering yet - ); - cache.put(&meta.location, cached_value); - - let result = cache.get(&meta.location).unwrap(); - assert!(result.ordering.is_none()); - - // Update to add ordering - let mut cached = cache.get(&meta.location).unwrap(); - if cached.is_valid_for(&meta) && cached.ordering.is_none() { - cached.ordering = Some(ordering()); - } - cache.put(&meta.location, cached); - - let result2 = cache.get(&meta.location).unwrap(); - assert!(result2.ordering.is_some()); - - // Verify list_entries shows has_ordering = true - let entries = cache.list_entries(); - assert_eq!(entries.len(), 1); - assert!(entries.get(&meta.location).unwrap().has_ordering); - } - - #[test] - fn test_cache_invalidation_on_file_modification() { - let cache = DefaultFileStatisticsCache::default(); - let path = Path::from("test.parquet"); - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - let meta_v1 = create_test_meta("test.parquet", 100); - - // Cache initial value - let cached_value = CachedFileMetadata::new( - meta_v1.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&path, cached_value); - - // File modified (size changed) - let meta_v2 = create_test_meta("test.parquet", 200); - - let cached = cache.get(&path).unwrap(); - // Should not be valid for new meta - assert!(!cached.is_valid_for(&meta_v2)); - - // Compute new value and update - let new_cached = CachedFileMetadata::new( - meta_v2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&path, new_cached); - - // Should have new metadata - let result = cache.get(&path).unwrap(); - assert_eq!(result.meta.size, 200); - } - - #[test] - fn test_ordering_cache_invalidation_on_file_modification() { - let cache = DefaultFileStatisticsCache::default(); - let path = Path::from("test.parquet"); - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - // Cache with original metadata and ordering - let meta_v1 = ObjectMeta { - location: path.clone(), - last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") - .unwrap() - .into(), - size: 100, - e_tag: None, - version: None, - }; - let ordering_v1 = ordering(); - let cached_v1 = CachedFileMetadata::new( - meta_v1.clone(), - Arc::new(Statistics::new_unknown(&schema)), - Some(ordering_v1), - ); - cache.put(&path, cached_v1); - - // Verify cached ordering is valid - let cached = cache.get(&path).unwrap(); - assert!(cached.is_valid_for(&meta_v1)); - assert!(cached.ordering.is_some()); - - // File modified (size changed) - let meta_v2 = ObjectMeta { - location: path.clone(), - last_modified: DateTime::parse_from_rfc3339("2022-09-28T10:00:00+02:00") - .unwrap() - .into(), - size: 200, // Changed - e_tag: None, - version: None, - }; - - // Cache entry exists but should be invalid for new metadata - let cached = cache.get(&path).unwrap(); - assert!(!cached.is_valid_for(&meta_v2)); - - // Cache new version with different ordering - let ordering_v2 = ordering(); // New ordering instance - let cached_v2 = CachedFileMetadata::new( - meta_v2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - Some(ordering_v2), - ); - cache.put(&path, cached_v2); - - // Old metadata should be invalid - let cached = cache.get(&path).unwrap(); - assert!(!cached.is_valid_for(&meta_v1)); - - // New metadata should be valid - assert!(cached.is_valid_for(&meta_v2)); - assert!(cached.ordering.is_some()); - } - - #[test] - fn test_list_entries() { - let cache = DefaultFileStatisticsCache::default(); - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - let meta1 = create_test_meta("test1.parquet", 100); - - let cached_value = CachedFileMetadata::new( - meta1.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&meta1.location, cached_value); - let meta2 = create_test_meta("test2.parquet", 200); - let cached_value = CachedFileMetadata::new( - meta2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - Some(ordering()), - ); - cache.put(&meta2.location, cached_value); - - let entries = cache.list_entries(); - assert_eq!( - entries, - HashMap::from([ - ( - Path::from("test1.parquet"), - FileStatisticsCacheEntry { - object_meta: meta1, - num_rows: Precision::Absent, - num_columns: 1, - table_size_bytes: Precision::Absent, - statistics_size_bytes: 0, - has_ordering: false, - } - ), - ( - Path::from("test2.parquet"), - FileStatisticsCacheEntry { - object_meta: meta2, - num_rows: Precision::Absent, - num_columns: 1, - table_size_bytes: Precision::Absent, - statistics_size_bytes: 0, - has_ordering: true, - } - ), - ]) - ); - } -} diff --git a/datafusion/execution/src/cache/file_statistics_cache.rs b/datafusion/execution/src/cache/file_statistics_cache.rs new file mode 100644 index 0000000000..12f0bb1b8a --- /dev/null +++ b/datafusion/execution/src/cache/file_statistics_cache.rs @@ -0,0 +1,744 @@ +// 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 crate::cache::cache_manager::{ + CachedFileMetadata, FileStatisticsCache, FileStatisticsCacheEntry, +}; +use crate::cache::{CacheAccessor, TableScopedPath}; +use std::collections::HashMap; +use std::sync::Mutex; + +pub use crate::cache::DefaultFilesMetadataCache; +use crate::cache::lru_queue::LruQueue; +use datafusion_common::TableReference; +use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; + +/// Default implementation of [`FileStatisticsCache`] +/// +/// Stores cached file metadata (statistics and orderings) for files. +/// +/// The typical usage pattern is: +/// 1. Call `get(path)` to check for cached value +/// 2. If `Some(cached)`, validate with `cached.is_valid_for(¤t_meta)` +/// 3. If invalid or missing, compute new value and call `put(path, new_value)` +/// +/// # Internal details +/// +/// The `memory_limit` controls the maximum size of the cache, which uses a +/// Least Recently Used eviction algorithm. When adding a new entry, if the total +/// size of the cached entries exceeds `memory_limit`, the least recently used entries +/// are evicted until the total size is lower than `memory_limit`. +/// +/// +/// [`FileStatisticsCache`]: crate::cache::cache_manager::FileStatisticsCache +#[derive(Default)] +pub struct DefaultFileStatisticsCache { + state: Mutex<DefaultFileStatisticsCacheState>, +} + +impl DefaultFileStatisticsCache { + pub fn new(memory_limit: usize) -> Self { + Self { + state: Mutex::new(DefaultFileStatisticsCacheState::new(memory_limit)), + } + } + + /// Returns the size of the cached memory, in bytes. + pub fn memory_used(&self) -> usize { + let state = self.state.lock().unwrap(); + state.memory_used + } +} + +struct DefaultFileStatisticsCacheState { + lru_queue: LruQueue<TableScopedPath, CachedFileMetadata>, + memory_limit: usize, + memory_used: usize, +} + +pub const DEFAULT_FILE_STATISTICS_MEMORY_LIMIT: usize = 20 * 1024 * 1024; // 20MiB + +impl Default for DefaultFileStatisticsCacheState { + fn default() -> Self { + Self { + lru_queue: LruQueue::new(), + memory_limit: DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, + memory_used: 0, + } + } +} + +impl DefaultFileStatisticsCacheState { + fn new(memory_limit: usize) -> Self { + Self { + lru_queue: LruQueue::new(), + memory_limit, + memory_used: 0, + } + } + fn get(&mut self, key: &TableScopedPath) -> Option<CachedFileMetadata> { + self.lru_queue.get(key).cloned() + } + + fn put( + &mut self, + key: &TableScopedPath, + value: CachedFileMetadata, + ) -> Option<CachedFileMetadata> { + let mut ctx = DFHeapSizeCtx::default(); + let key_size = key.heap_size(&mut ctx); + let entry_size = value.heap_size(&mut ctx); + + if entry_size + key_size > self.memory_limit { + // Remove potential stale entry + return self.remove(key); + } + + self.memory_used += entry_size; + self.memory_used += key_size; + + let old_value = self.lru_queue.put(key.clone(), value); + if let Some(old_entry) = &old_value { + let mut ctx = DFHeapSizeCtx::default(); + self.memory_used -= old_entry.heap_size(&mut ctx); + self.memory_used -= key_size; + } + + self.evict_entries(); + + old_value + } + + fn remove(&mut self, k: &TableScopedPath) -> Option<CachedFileMetadata> { + if let Some(old_entry) = self.lru_queue.remove(k) { + let mut ctx = DFHeapSizeCtx::default(); + self.memory_used -= k.heap_size(&mut ctx); + self.memory_used -= old_entry.heap_size(&mut ctx); + Some(old_entry) + } else { + None + } + } + + fn contains_key(&self, k: &TableScopedPath) -> bool { + self.lru_queue.contains_key(k) + } + + fn len(&self) -> usize { + self.lru_queue.len() + } + + fn clear(&mut self) { + self.lru_queue.clear(); + self.memory_used = 0; + } + + fn evict_entries(&mut self) { + while self.memory_used > self.memory_limit { + if let Some(removed) = self.lru_queue.pop() { + let mut ctx = DFHeapSizeCtx::default(); + self.memory_used -= removed.0.heap_size(&mut ctx); + self.memory_used -= removed.1.heap_size(&mut ctx); + } else { + // cache is empty while memory_used > memory_limit, cannot happen + log::error!( + "File statistics cache memory accounting bug: memory_used={} but cache is empty. \ + Please report this to the Apache DataFusion developers.", + self.memory_used + ); + debug_assert!( + false, + "memory_used={} but cache is empty", + self.memory_used + ); + self.memory_used = 0; + return; + } + } + } +} +impl CacheAccessor<TableScopedPath, CachedFileMetadata> for DefaultFileStatisticsCache { + fn get(&self, key: &TableScopedPath) -> Option<CachedFileMetadata> { + let mut state = self.state.lock().unwrap(); + state.get(key) + } + + fn put( + &self, + key: &TableScopedPath, + value: CachedFileMetadata, + ) -> Option<CachedFileMetadata> { + let mut state = self.state.lock().unwrap(); + state.put(key, value) + } + + fn remove(&self, key: &TableScopedPath) -> Option<CachedFileMetadata> { + let mut state = self.state.lock().unwrap(); + state.remove(key) + } + + fn contains_key(&self, k: &TableScopedPath) -> bool { + let state = self.state.lock().unwrap(); + state.contains_key(k) + } + + fn len(&self) -> usize { + let state = self.state.lock().unwrap(); + state.len() + } + + fn clear(&self) { + let mut state = self.state.lock().unwrap(); + state.clear(); + } + + fn name(&self) -> String { + "DefaultFileStatisticsCache".to_string() + } +} + +impl FileStatisticsCache for DefaultFileStatisticsCache { + fn cache_limit(&self) -> usize { + let state = self.state.lock().unwrap(); + state.memory_limit + } + + fn update_cache_limit(&self, limit: usize) { + let mut state = self.state.lock().unwrap(); + state.memory_limit = limit; + state.evict_entries(); + } + + fn list_entries(&self) -> HashMap<TableScopedPath, FileStatisticsCacheEntry> { + let mut entries = HashMap::<TableScopedPath, FileStatisticsCacheEntry>::new(); + let mut ctx = DFHeapSizeCtx::default(); + for entry in self.state.lock().unwrap().lru_queue.list_entries() { + let path = entry.0.clone(); + let cached = entry.1; + entries.insert( + path, + FileStatisticsCacheEntry { + object_meta: cached.meta.clone(), + num_rows: cached.statistics.num_rows, + num_columns: cached.statistics.column_statistics.len(), + table_size_bytes: cached.statistics.total_byte_size, + statistics_size_bytes: cached.statistics.heap_size(&mut ctx), + has_ordering: cached.ordering.is_some(), + }, + ); + } + + entries + } + + fn drop_table_entries( + &self, + table_ref: &Option<TableReference>, + ) -> datafusion_common::Result<()> { + let mut state = self.state.lock().unwrap(); + let mut table_paths = vec![]; + for (path, _) in state.lru_queue.list_entries() { + if path.table == *table_ref { + table_paths.push(path.clone()); + } + } + for path in table_paths { + state.remove(&path); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cache::cache_manager::{ + CachedFileMetadata, FileStatisticsCache, FileStatisticsCacheEntry, + }; + use arrow::array::{Int32Array, ListArray, RecordBatch}; + use arrow::buffer::{OffsetBuffer, ScalarBuffer}; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use chrono::DateTime; + use datafusion_common::heap_size::DFHeapSizeCtx; + use datafusion_common::stats::Precision; + use datafusion_common::{ColumnStatistics, ScalarValue, Statistics}; + use datafusion_expr::ColumnarValue; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; + use object_store::ObjectMeta; + use object_store::path::Path; + use std::sync::Arc; + + fn create_test_meta(path: &str, size: u64) -> ObjectMeta { + ObjectMeta { + location: Path::from(path), + last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") + .unwrap() + .into(), + size, + e_tag: None, + version: None, + } + } + + #[test] + fn test_statistics_cache() { + let meta = create_test_meta("test", 1024); + let cache = DefaultFileStatisticsCache::default(); + + let schema = Schema::new(vec![Field::new( + "test_column", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]); + + let path = TableScopedPath { + path: meta.location.clone(), + table: None, + }; + + // Cache miss + assert!(cache.get(&path).is_none()); + + // Put a value + let cached_value = CachedFileMetadata::new( + meta.clone(), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path, cached_value); + + // Cache hit + let result = cache.get(&path); + assert!(result.is_some()); + + let cached = result.unwrap(); + assert!(cached.is_valid_for(&meta)); + + // File size changed - validation should fail + let meta2 = create_test_meta("test", 2048); + + let path_2 = TableScopedPath { + path: meta2.location.clone(), + table: None, + }; + + let cached = cache.get(&path_2).unwrap(); + assert!(!cached.is_valid_for(&meta2)); + + // Update with new value + let cached_value2 = CachedFileMetadata::new( + meta2.clone(), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path_2, cached_value2); + + // Test list_entries + let entries = cache.list_entries(); + assert_eq!(entries.len(), 1); + + let path_3 = TableScopedPath { + path: Path::from("test"), + table: None, + }; + + let entry = entries.get(&path_3).unwrap(); + assert_eq!(entry.object_meta.size, 2048); // Should be updated value + } + + #[derive(Clone, Debug, PartialEq, Eq, Hash)] + struct MockExpr {} + + impl std::fmt::Display for MockExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "MockExpr") + } + } + + impl PhysicalExpr for MockExpr { + fn data_type( + &self, + _input_schema: &Schema, + ) -> datafusion_common::Result<DataType> { + Ok(DataType::Int32) + } + + fn nullable(&self, _input_schema: &Schema) -> datafusion_common::Result<bool> { + Ok(false) + } + + fn evaluate( + &self, + _batch: &RecordBatch, + ) -> datafusion_common::Result<ColumnarValue> { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { + vec![] + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn PhysicalExpr>>, + ) -> datafusion_common::Result<Arc<dyn PhysicalExpr>> { + assert!(children.is_empty()); + Ok(self) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "MockExpr") + } + } + + fn ordering() -> LexOrdering { + let expr = Arc::new(MockExpr {}) as Arc<dyn PhysicalExpr>; + LexOrdering::new(vec![PhysicalSortExpr::new_default(expr)]).unwrap() + } + + #[test] + fn test_ordering_cache() { + let meta = create_test_meta("test.parquet", 100); + let cache = DefaultFileStatisticsCache::default(); + + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + + // Cache statistics with no ordering + let cached_value = CachedFileMetadata::new( + meta.clone(), + Arc::new(Statistics::new_unknown(&schema)), + None, // No ordering yet + ); + + let path = TableScopedPath { + path: meta.location.clone(), + table: None, + }; + + cache.put(&path, cached_value); + + let result = cache.get(&path).unwrap(); + assert!(result.ordering.is_none()); + + // Update to add ordering + let mut cached = cache.get(&path).unwrap(); + if cached.is_valid_for(&meta) && cached.ordering.is_none() { + cached.ordering = Some(ordering()); + } + cache.put(&path, cached); + + let result2 = cache.get(&path).unwrap(); + assert!(result2.ordering.is_some()); + + // Verify list_entries shows has_ordering = true + let entries = cache.list_entries(); + assert_eq!(entries.len(), 1); + assert!(entries.get(&path).unwrap().has_ordering); + } + + #[test] + fn test_cache_invalidation_on_file_modification() { + let cache = DefaultFileStatisticsCache::default(); + let path = TableScopedPath { + path: Path::from("test.parquet"), + table: None, + }; + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + + let meta_v1 = create_test_meta("test.parquet", 100); + + // Cache initial value + let cached_value = CachedFileMetadata::new( + meta_v1.clone(), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path, cached_value); + + // File modified (size changed) + let meta_v2 = create_test_meta("test.parquet", 200); + + let cached = cache.get(&path).unwrap(); + // Should not be valid for new meta + assert!(!cached.is_valid_for(&meta_v2)); + + // Compute new value and update + let new_cached = CachedFileMetadata::new( + meta_v2.clone(), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path, new_cached); + + // Should have new metadata + let result = cache.get(&path).unwrap(); + assert_eq!(result.meta.size, 200); + } + + #[test] + fn test_ordering_cache_invalidation_on_file_modification() { + let cache = DefaultFileStatisticsCache::default(); + let path = TableScopedPath { + path: Path::from("test.parquet"), + table: None, + }; + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + + // Cache with original metadata and ordering + let meta_v1 = ObjectMeta { + location: path.path.clone(), + last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") + .unwrap() + .into(), + size: 100, + e_tag: None, + version: None, + }; + let ordering_v1 = ordering(); + let cached_v1 = CachedFileMetadata::new( + meta_v1.clone(), + Arc::new(Statistics::new_unknown(&schema)), + Some(ordering_v1), + ); + cache.put(&path, cached_v1); + + // Verify cached ordering is valid + let cached = cache.get(&path).unwrap(); + assert!(cached.is_valid_for(&meta_v1)); + assert!(cached.ordering.is_some()); + + // File modified (size changed) + let meta_v2 = ObjectMeta { + location: path.path.clone(), + last_modified: DateTime::parse_from_rfc3339("2022-09-28T10:00:00+02:00") + .unwrap() + .into(), + size: 200, // Changed + e_tag: None, + version: None, + }; + + // Cache entry exists but should be invalid for new metadata + let cached = cache.get(&path).unwrap(); + assert!(!cached.is_valid_for(&meta_v2)); + + // Cache new version with different ordering + let ordering_v2 = ordering(); // New ordering instance + let cached_v2 = CachedFileMetadata::new( + meta_v2.clone(), + Arc::new(Statistics::new_unknown(&schema)), + Some(ordering_v2), + ); + cache.put(&path, cached_v2); + + // Old metadata should be invalid + let cached = cache.get(&path).unwrap(); + assert!(!cached.is_valid_for(&meta_v1)); + + // New metadata should be valid + assert!(cached.is_valid_for(&meta_v2)); + assert!(cached.ordering.is_some()); + } + + #[test] + fn test_list_entries() { + let cache = DefaultFileStatisticsCache::default(); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + + let meta1 = create_test_meta("test1.parquet", 100); + + let cached_value = CachedFileMetadata::new( + meta1.clone(), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + + let path_1 = TableScopedPath { + path: meta1.location.clone(), + table: None, + }; + + cache.put(&path_1, cached_value); + let meta2 = create_test_meta("test2.parquet", 200); + let cached_value = CachedFileMetadata::new( + meta2.clone(), + Arc::new(Statistics::new_unknown(&schema)), + Some(ordering()), + ); + + let path_2 = TableScopedPath { + path: meta2.location.clone(), + table: None, + }; + + cache.put(&path_2, cached_value); + + let entries = cache.list_entries(); + assert_eq!( + entries, + HashMap::from([ + ( + path_1, + FileStatisticsCacheEntry { + object_meta: meta1, + num_rows: Precision::Absent, + num_columns: 1, + table_size_bytes: Precision::Absent, + statistics_size_bytes: 360, + has_ordering: false, + } + ), + ( + path_2, + FileStatisticsCacheEntry { + object_meta: meta2, + num_rows: Precision::Absent, + num_columns: 1, + table_size_bytes: Precision::Absent, + statistics_size_bytes: 360, + has_ordering: true, + } + ), + ]) + ); + } + + #[test] + fn test_cache_entry_added_when_entries_are_within_cache_limit() { + let (meta_1, value_1) = create_cached_file_metadata_with_stats("test1.parquet"); + let (meta_2, value_2) = create_cached_file_metadata_with_stats("test2.parquet"); + let (meta_3, value_3) = create_cached_file_metadata_with_stats("test3.parquet"); + + let mut ctx = DFHeapSizeCtx::default(); + + let limit_for_2_entries = 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); + + // create a cache with a limit which fits exactly 2 entries + let cache = DefaultFileStatisticsCache::new(limit_for_2_entries); + let path_1 = TableScopedPath { + path: meta_1.location.clone(), + table: None, + }; + + let path_2 = TableScopedPath { + path: meta_2.location.clone(), + table: None, + }; + + cache.put(&path_1, value_1.clone()); + cache.put(&path_2, value_2.clone()); + + assert_eq!(cache.len(), 2); + assert_eq!(cache.memory_used(), limit_for_2_entries); + + let result_1 = cache.get(&path_1); + let result_2 = cache.get(&path_2); + assert_eq!(result_1.unwrap(), value_1); + assert_eq!(result_2.unwrap(), value_2); + + let path_3 = TableScopedPath { + path: meta_3.location.clone(), + table: None, + }; + + // adding the third entry evicts the first entry + cache.put(&path_3, value_3.clone()); + assert_eq!(cache.len(), 2); + assert_eq!(cache.memory_used(), limit_for_2_entries); + + let result_1 = cache.get(&path_1); + assert!(result_1.is_none()); + + let result_2 = cache.get(&path_2); + let result_3 = cache.get(&path_3); + + assert_eq!(result_2.unwrap(), value_2); + assert_eq!(result_3.unwrap(), value_3); + + // add the third entry again, making sure memory usage remains the same + cache.put(&path_3, value_3.clone()); + assert_eq!(cache.memory_used(), limit_for_2_entries); + cache.put(&path_3, value_3.clone()); + assert_eq!(cache.memory_used(), limit_for_2_entries); + + let mut ctx = DFHeapSizeCtx::default(); + cache.remove(&path_2); + 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) + ); + + cache.clear(); + assert_eq!(cache.len(), 0); + assert_eq!(cache.memory_used(), 0); + } + + #[test] + fn test_cache_rejects_entry_which_is_too_large() { + let (meta, value) = create_cached_file_metadata_with_stats("test1.parquet"); + let mut ctx = DFHeapSizeCtx::default(); + let limit_less_than_the_entry = value.heap_size(&mut ctx) - 1; + + // create a cache with a size less than the entry + let cache = DefaultFileStatisticsCache::new(limit_less_than_the_entry); + + let path_1 = TableScopedPath { + path: meta.location.clone(), + table: None, + }; + + cache.put(&path_1, value); + + assert_eq!(cache.len(), 0); + assert_eq!(cache.memory_used(), 0); + } + + fn create_cached_file_metadata_with_stats( + file_name: &str, + ) -> (ObjectMeta, CachedFileMetadata) { + let series: Vec<i32> = (0..=10).collect(); + let values = Int32Array::from(series); + let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 11])); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let list_array = ListArray::new(field, offsets, Arc::new(values), None); + + let column_statistics = ColumnStatistics { + null_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), + min_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), + sum_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), + distinct_count: Precision::Exact(10), + byte_size: Precision::Absent, + }; + + let stats = Statistics { + num_rows: Precision::Exact(100), + total_byte_size: Precision::Exact(100), + column_statistics: vec![column_statistics.clone()], + }; + let mut ctx = DFHeapSizeCtx::default(); + let object_meta = create_test_meta(file_name, stats.heap_size(&mut ctx) as u64); + let value = + CachedFileMetadata::new(object_meta.clone(), Arc::new(stats.clone()), None); + (object_meta, value) + } +} diff --git a/datafusion/execution/src/cache/list_files_cache.rs b/datafusion/execution/src/cache/list_files_cache.rs index b1b8e6b500..a3cdf7c5e9 100644 --- a/datafusion/execution/src/cache/list_files_cache.rs +++ b/datafusion/execution/src/cache/list_files_cache.rs @@ -15,6 +15,13 @@ // specific language governing permissions and limitations // under the License. +use crate::cache::{ + CacheAccessor, + cache_manager::{CachedFileList, ListFilesCache}, + lru_queue::LruQueue, +}; + +use std::fmt::{Debug, Display, Formatter}; use std::mem::size_of; use std::{ collections::HashMap, @@ -23,15 +30,10 @@ use std::{ }; use datafusion_common::TableReference; +use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use datafusion_common::instant::Instant; use object_store::{ObjectMeta, path::Path}; -use crate::cache::{ - CacheAccessor, - cache_manager::{CachedFileList, ListFilesCache}, - lru_queue::LruQueue, -}; - pub trait TimeProvider: Send + Sync + 'static { fn now(&self) -> Instant; } @@ -169,6 +171,22 @@ impl Default for DefaultListFilesCacheState { } } +impl DFHeapSize for TableScopedPath { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.path.as_ref().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) + } else { + write!(f, "{}", self.path) + } + } +} + impl DefaultListFilesCacheState { fn new(memory_limit: usize, ttl: Option<Duration>) -> Self { Self { diff --git a/datafusion/execution/src/cache/mod.rs b/datafusion/execution/src/cache/mod.rs index 0380e50c09..76bd660e6c 100644 --- a/datafusion/execution/src/cache/mod.rs +++ b/datafusion/execution/src/cache/mod.rs @@ -16,7 +16,7 @@ // under the License. pub mod cache_manager; -pub mod cache_unit; +pub mod file_statistics_cache; pub mod lru_queue; mod file_metadata_cache; diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs index 08f5339e7b..5b90f28a14 100644 --- a/datafusion/execution/src/runtime_env.rs +++ b/datafusion/execution/src/runtime_env.rs @@ -103,6 +103,7 @@ fn create_runtime_config_entries( metadata_cache_limit: Option<String>, list_files_cache_limit: Option<String>, list_files_cache_ttl: Option<String>, + file_statistics_cache_limit: Option<String>, ) -> Vec<ConfigEntry> { vec![ ConfigEntry { @@ -135,6 +136,11 @@ fn create_runtime_config_entries( value: list_files_cache_ttl, description: "TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes.", }, + ConfigEntry { + key: "datafusion.runtime.file_statistics_cache_limit".to_string(), + value: file_statistics_cache_limit, + description: "Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", + }, ] } @@ -296,6 +302,14 @@ impl RuntimeEnv { .get_list_files_cache_ttl() .map(format_duration); + let file_statistics_cache_limit = + self.cache_manager.get_file_statistic_cache_limit(); + let file_statistics_cache_value = format_byte_size( + file_statistics_cache_limit + .try_into() + .expect("File statistics cache size conversion failed"), + ); + create_runtime_config_entries( memory_limit_value, Some(max_temp_dir_value), @@ -303,6 +317,7 @@ impl RuntimeEnv { Some(metadata_cache_value), Some(list_files_cache_value), list_files_cache_ttl, + Some(file_statistics_cache_value), ) } } @@ -438,6 +453,11 @@ impl RuntimeEnvBuilder { self } + pub fn with_file_statistics_cache_limit(mut self, limit: usize) -> Self { + self.cache_manager = self.cache_manager.with_file_statistics_cache_limit(limit); + self + } + /// Build a RuntimeEnv pub fn build(self) -> Result<RuntimeEnv> { let Self { @@ -475,9 +495,10 @@ impl RuntimeEnvBuilder { /// Create a new RuntimeEnvBuilder from an existing RuntimeEnv pub fn from_runtime_env(runtime_env: &RuntimeEnv) -> Self { let cache_config = CacheManagerConfig { - table_files_statistics_cache: runtime_env + file_statistics_cache: runtime_env.cache_manager.get_file_statistic_cache(), + file_statistics_cache_limit: runtime_env .cache_manager - .get_file_statistic_cache(), + .get_file_statistic_cache_limit(), list_files_cache: runtime_env.cache_manager.get_list_files_cache(), list_files_cache_limit: runtime_env .cache_manager @@ -514,6 +535,7 @@ impl RuntimeEnvBuilder { Some("50M".to_owned()), Some("1M".to_owned()), None, + Some("20M".to_owned()), ) } diff --git a/datafusion/sqllogictest/test_files/encrypted_parquet.slt b/datafusion/sqllogictest/test_files/encrypted_parquet.slt index d580b7d1ad..f51d84f0c8 100644 --- a/datafusion/sqllogictest/test_files/encrypted_parquet.slt +++ b/datafusion/sqllogictest/test_files/encrypted_parquet.slt @@ -85,5 +85,5 @@ float_field float ) STORED AS PARQUET LOCATION 'test_files/scratch/encrypted_parquet/' -query error DataFusion error: Parquet error: Parquet error: Parquet file has an encrypted footer but decryption properties were not provided +query error Parquet error: Parquet error: Parquet file has an encrypted footer but decryption properties were not provided SELECT * FROM parquet_table diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 8396a60137..c940782d67 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -332,6 +332,7 @@ datafusion.optimizer.skip_failed_rules false datafusion.optimizer.subset_repartition_threshold 4 datafusion.optimizer.top_down_join_key_reordering true datafusion.optimizer.use_statistics_registry false +datafusion.runtime.file_statistics_cache_limit 20M datafusion.runtime.list_files_cache_limit 1M datafusion.runtime.list_files_cache_ttl NULL datafusion.runtime.max_temp_directory_size 100G @@ -479,6 +480,7 @@ datafusion.optimizer.skip_failed_rules false When set to true, the logical plan datafusion.optimizer.subset_repartition_threshold 4 Partition count threshold for subset satisfaction optimization. When the current partition count is >= this threshold, DataFusion will skip repartitioning if the required partitioning expression is a subset of the current partition expression such as Hash(a) satisfies Hash(a, b). When the current partition count is < this threshold, DataFusion will repartition to increase parallelism even when subset satisfaction applies. Set to 0 to al [...] datafusion.optimizer.top_down_join_key_reordering true When set to true, the physical plan optimizer will run a top down process to reorder the join keys datafusion.optimizer.use_statistics_registry false When set to true, the physical plan optimizer uses the pluggable `StatisticsRegistry` for statistics propagation across operators. This enables more accurate cardinality estimates compared to each operator's built-in `partition_statistics`. +datafusion.runtime.file_statistics_cache_limit 20M Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_limit 1M Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_ttl NULL TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes. datafusion.runtime.max_temp_directory_size 100G Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. diff --git a/datafusion/sqllogictest/test_files/set_variable.slt b/datafusion/sqllogictest/test_files/set_variable.slt index 508b301fe3..6f58e5fb31 100644 --- a/datafusion/sqllogictest/test_files/set_variable.slt +++ b/datafusion/sqllogictest/test_files/set_variable.slt @@ -351,6 +351,12 @@ RESET datafusion.runtime.memory_limit statement ok EXPLAIN ANALYZE SELECT * FROM generate_series(1, 1000) AS t1(v1) ORDER BY v1 DESC +statement ok +SET datafusion.runtime.file_statistics_cache_limit = '1K' + +statement ok +RESET datafusion.runtime.file_statistics_cache_limit + statement ok SET datafusion.runtime.list_files_cache_limit = '1K' @@ -605,6 +611,18 @@ SHOW datafusion.runtime.max_temp_directory_size ---- datafusion.runtime.max_temp_directory_size 10G +# Test SET and SHOW runtime.file_statistics_cache_limit +statement ok +SET datafusion.runtime.file_statistics_cache_limit = '42M' + +query TT +SHOW datafusion.runtime.file_statistics_cache_limit +---- +datafusion.runtime.file_statistics_cache_limit 42M + +statement ok +RESET datafusion.runtime.file_statistics_cache_limit + # Test SET and SHOW runtime.metadata_cache_limit statement ok SET datafusion.runtime.metadata_cache_limit = '200M' @@ -648,6 +666,7 @@ datafusion.runtime.list_files_cache_limit 0 query T SELECT name FROM information_schema.df_settings WHERE name LIKE 'datafusion.runtime.%' ORDER BY name ---- +datafusion.runtime.file_statistics_cache_limit datafusion.runtime.list_files_cache_limit datafusion.runtime.list_files_cache_ttl datafusion.runtime.max_temp_directory_size diff --git a/docs/source/library-user-guide/upgrading/54.0.0.md b/docs/source/library-user-guide/upgrading/54.0.0.md index 0117a776b2..076b65444f 100644 --- a/docs/source/library-user-guide/upgrading/54.0.0.md +++ b/docs/source/library-user-guide/upgrading/54.0.0.md @@ -716,3 +716,42 @@ move the borrow behind an interior-mutability primitive such as `Arc<Mutex<_>>` or `Arc<OnceLock<_>>`. See [PR #21803](https://github.com/apache/datafusion/pull/21803) for details. + +[20047]: https://github.com/apache/datafusion/pull/20047 + +### File statistics cache is now memory-limited and managed by the `CacheManager` + +The file statistics cache used by `ListingTable` is now memory-limited and +centrally managed through the `CacheManager`. + +To configure the cache size use the `file_statistics_cache_limit` setting: + +```sql +SET datafusion.runtime.file_statistics_cache_limit = '10M' +``` + +To disable the file statistics cache, set the limit to 0. + +The file statistics cache is no longer created inside the `ListingTable`. +Instead, it is created within the `CacheManager` and must be passed to the `ListingTable`. + +**Who is affected:** + +- Users who want to limit the memory usage of the file statistics cache. +- Users who want to disable the file statistics cache. +- Users who want to create a `ListingTable` programmatically with a file statistics cache. + +**Migration guide:** + +Disable the cache by setting the configuration value to 0: + +```sql +SET datafusion.runtime.file_statistics_cache_limit = '0' +``` + +Use the file statistics cache provided by the `CacheManager` when initializing a new `ListingTable`: + +```rust,ignore +ListingTable::try_new(config)? + .with_cache(ctx.runtime_env().cache_manager.get_file_statistic_cache()) +``` diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index f7f9426e2b..f0bab1c4db 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -229,14 +229,15 @@ SET datafusion.runtime.memory_limit = '2G'; The following runtime configuration settings are available: -| key | default | description | -| ------------------------------------------ | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| datafusion.runtime.list_files_cache_limit | 1M | Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | -| datafusion.runtime.list_files_cache_ttl | NULL | TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes. | -| datafusion.runtime.max_temp_directory_size | 100G | Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | -| datafusion.runtime.memory_limit | NULL | Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | -| datafusion.runtime.metadata_cache_limit | 50M | Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | -| datafusion.runtime.temp_directory | NULL | The path to the temporary file directory. | +| key | default | description | +| ---------------------------------------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| datafusion.runtime.file_statistics_cache_limit | 20M | Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | +| datafusion.runtime.list_files_cache_limit | 1M | Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | +| datafusion.runtime.list_files_cache_ttl | NULL | TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes. | +| datafusion.runtime.max_temp_directory_size | 100G | Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | +| datafusion.runtime.memory_limit | NULL | Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | +| datafusion.runtime.metadata_cache_limit | 50M | Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | +| datafusion.runtime.temp_directory | NULL | The path to the temporary file directory. | # Tuning Guide --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
