perf(scheduler): share the file statistics cache across sessions by andygrove · Pull Request #2498 · apache/datafusion-ballista

andygrove · 2026-09-27T18:53:48Z

Which issue does this PR close?

Closes #2497.

Rationale for this change

The scheduler builds a new session, with a new RuntimeEnv, for every query, so DataFusion's file statistics cache is empty for every job. Planning each job reads the footer of every Parquet file in every table it scans to collect statistics, which takes seconds per query on large tables (over 4 s at TPC-H SF1000). Details and measurements are in #2497.

What changes are included in this PR?

The new public share_file_statistics_cache wraps a SessionBuilder so every session it builds uses one file statistics cache, the one the first session was built with. The builder's configured limit applies (DataFusion's default is 20 MiB, LRU), and a builder that disables the cache also disables sharing. A session that already has the shared cache is returned unchanged.
BallistaCluster::new_memory wraps its session builder with it, which covers the scheduler binary, standalone mode and tests. InMemoryJobState::new does not, so passing one to BallistaCluster::new keeps a cache per session, and a custom JobState can opt in with one call.
The listing cache deliberately stays per session. DataFusion checks each cached entry against the size and modification time from the current listing, so a changed file is read again. Sharing the listing as well would serve stale statistics, and COUNT(*) is answered from them.
The rebuilt session keeps the session ID the builder gave it, since SessionStateBuilder::new_from_existing would otherwise assign a new one.

Four tests: sessions from new_memory share the statistics cache but not the listing cache, a table's file rewritten between sessions gives the new COUNT(*), the session ID is kept, and a builder that hands out one state gets it back unchanged. Each fails when the behavior it covers is removed, and the COUNT(*) one returns the stale count if the listing cache is shared too.

Planning time (Job [...] planning took) for TPC-H, measured locally:

	Before	After
SF100 (32 files per table), all 22 queries	1492 ms	129 ms
SF100 files linked 10 times (320 per table), Q21	1542 ms	10 ms
Same, Q1, the first query to scan `lineitem`	559 ms	542 ms

Once a table has been scanned, queries on it plan in 1 to 3 ms at SF100 and 5 to 13 ms with 320 files per table. The first scan of each table after the scheduler starts still pays the full cost. Physical plans are identical before and after for all 22 queries, and all 22 pass the benchmark's --verify check against DataFusion.

Are there any user-facing changes?

Planning is faster once a table has been scanned. The scheduler keeps one file statistics cache for its lifetime instead of one per query, so sessions reuse statistics that other sessions collected, still checked against their own file listing. That check compares size and modification time but not e-tags, so a file rewritten in place without changing either keeps its old statistics, as it would in a long-lived DataFusion SessionContext.

share_file_statistics_cache is a new public function, and BallistaCluster::new_memory now applies it. No existing API changes.

The scheduler builds a new session, and with it a new RuntimeEnv, for every query, so DataFusion's file statistics cache is empty for every job and planning re-reads the footer of every Parquet file it scans. On large tables that takes seconds for every job. Build every session the in-memory job state creates around one shared file statistics cache. Each job still lists the files itself, and cached entries are checked against that listing's size and modification time, so a changed file is read again. The listing cache stays per session.

martin-g · 2026-09-27T20:30:31Z

+pub(crate) fn share_file_statistics_cache(
+    session_builder: SessionBuilder,
+) -> SessionBuilder {
+    let shared = OnceLock::new();


Shouldn't this be a static ?

Suggested change

let shared = OnceLock::new();

static shared: FileStatitsticsCache = OnceLock::new();

I think the closure-captured OnceLock is the right scope here. A static would make the cache process-wide, so the first builder wrapped in the process would decide the limit, or whether sharing is on, for every scheduler in it (for example several standalone clusters in one test binary). It would also make assert_eq!(1, cache.len()) in test_in_memory_sessions_share_file_statistics order-dependent, because test_in_memory_sessions_reread_changed_files inserts its own entry into the same cache when it runs first in the same test binary. With the OnceLock there is one cache per wrapped builder, so one per scheduler.

I kept the OnceLock captured by the closure, for the reasons @comphead gives above. A static would mean one cache per process. If several schedulers run in one process, like the standalone clusters in our tests, the first builder wrapped would decide the cache limit, and whether sharing is on, for all of them. With the closure it's one cache per wrapped builder, so one per scheduler.

milenkovicm

thanks @andygrove this makes sense i wonder how what would be improvement on SF1000?

few observations "

share_file_statistics_cache wraps a SessionBuilder so every session it builds uses one file statistics cache, the one the first session was built with. The builder's configured limit applies (DataFusion's default is 20 MiB, LRU), and a builder that disables the cache also disables sharing.

i believe it would make more sense to follow current pattern where users can provide their session_builders (with or without cache)

InMemoryJobState::new wraps its session builder with it, which covers the scheduler binary, standalone mode and tests.

I wonder if we need to do it at this point, as users can plug in their own state they might miss adding wrapper. See code comment

The listing cache deliberately stays per session. DataFusion checks each cached entry against the size and modification time from the current listing, so a changed file is read again. Sharing the listing as well would serve stale statistics, and COUNT(*) is answered from them.

Im not sure I understand this point

The rebuilt session keeps the session ID the builder gave it, since SessionStateBuilder::new_from_existing would otherwise assign a new one.

I believe at the moment scheduler does not respect session id (but i'm not 100% sure) it might be that session id is generated on schedule for every request. so keeping session id might be no big deal

If you define new default_session_builder_with_cache and use it instead of default_session_builder could simplify design as you could define cache there as static

milenkovicm · 2026-09-27T20:32:55Z

            running_jobs: Default::default(),
            //sessions: Default::default(),
-            session_builder,
+            session_builder: share_file_statistics_cache(session_builder),


if we wrap session_builder here it would push caching all the time, users would not be able to override it / disable it. Would it make sense to wrap builder at one of following places

datafusion-ballista/ballista/scheduler/src/cluster/mod.rs

Line 116 in 32ceaf8

.unwrap_or_else(|| Arc::new(default_session_builder));

datafusion-ballista/ballista/core/src/utils.rs

Line 158 in 32ceaf8

pub fn default_session_builder(

One data point for this thread: wrapping only the default at cluster/mod.rs:116 or inside default_session_builder would miss the scheduler binary. bin/main.rs:91 always installs session_state_with_s3_support through with_override_session_builder, so new_from_config never reaches the default.

A middle ground that keeps coverage and gives an opt-out: apply the wrapper in BallistaCluster::new_memory instead of InMemoryJobState::new, and make share_file_statistics_cache pub. The binary (via new_from_config), all three standalone constructors, test_cluster_context and SchedulerTest go through new_memory, so none of them lose the speedup. InMemoryJobState::new stays a plain constructor, so BallistaCluster::new with an unwrapped state is the opt-out, and a custom JobState can opt in with one call. The two new tests in this file would then wrap the builder themselves.

Moved it in 2f8fe37, following @comphead's suggestion. BallistaCluster::new_memory now wraps the builder, InMemoryJobState::new is a plain constructor again, and share_file_statistics_cache is public.

I didn't wrap just the default in new_from_config or default_session_builder, because the scheduler binary always installs session_state_with_s3_support through with_override_session_builder. It would never get the speedup. new_memory covers the binary (via new_from_config), all three standalone constructors, test_cluster_context and SchedulerTest.

To opt out, pass an InMemoryJobState to BallistaCluster::new. A custom JobState can opt in by wrapping its builder with share_file_statistics_cache. A session builder that disables the file statistics cache also disables sharing.

comphead

Thanks @andygrove, the planning numbers are convincing. Sharing one bounded, revalidated cache per process while keeping the listing per job matches prior art (Spark's FileStatusCache, DuckDB's Parquet metadata cache, StarRocks' FE caches). I checked the details against DataFusion 55.1.0. Replies on placement and static are in the existing threads. The rest is inline.

comphead · 2026-09-28T16:00:09Z

+/// The shared cache is the one the first session was built with, so the
+/// builder's configured limit applies, and a builder that disables the cache
+/// also disables sharing.
+pub(crate) fn share_file_statistics_cache(


Worth documenting: sharing widens the staleness window from one job to the scheduler's lifetime. At DataFusion 55.1.0 an entry is keyed by TableScopedPath { table, path }, where path is store-relative (no scheme or bucket), and a hit is reused while size, last_modified and the schema fingerprint match (CachedFileMetadata::is_valid_for). e_tag and version are not checked. Two cases could now serve wrong statistics, and so wrong COUNT(*), MIN or MAX answered from them:

the same relative path in two buckets or stores with equal size, mtime and schema

an in-place overwrite that keeps the size, within the store's mtime precision (S3 Last-Modified has one-second precision)

Both are unlikely, and a long-lived DataFusion SessionContext has the same exposure, so I don't think it blocks this PR. An upstream issue to add the object store URL to the key and check e_tag/version when present would close it. Spark, DuckDB and StarRocks key their shared caches by the full URI, and DuckDB checks the ETag first.

Added a paragraph to the doc comment in 2f8fe37 that covers both cases. Agreed it doesn't need to block this. An upstream issue to add the object store URL to the key, and check e_tag or version when present, sounds right to me.

Filed apache/datafusion#25841 for this.

comphead · 2026-09-28T16:00:09Z

+        let Some(cache) = shared
+            .get_or_init(|| runtime.cache_manager.get_file_statistic_cache())
+            .clone()
+        else {
+            return Ok(state);
+        };


Minor: when the builder already returned the shared cache (the first session, every job from new_standalone_scheduler_from_state, which clones one state, or a builder that manages its own cache), this still rebuilds the runtime and the whole state for the same result. new_from_existing(..).build() re-registers every function, so the cost is small (I estimate tens of µs per job, not measured) but avoidable. An early return also leaves such builders untouched. Untested sketch:

let own = runtime.cache_manager.get_file_statistic_cache(); let Some(cache) = shared.get_or_init(|| own.clone()).clone() else { return Ok(state); }; if own.is_some_and(|own| Arc::ptr_eq(&own, &cache)) { return Ok(state); }

Done in 2f8fe37. The first session, and every session from a builder that reuses one runtime, now come back untouched. test_share_file_statistics_cache_skips_rebuild_when_shared covers the new_standalone_scheduler_from_state shape.

comphead · 2026-09-28T16:00:09Z

+        });
+        let session_builder = share_file_statistics_cache(session_builder);
+
+        let state: SessionState = session_builder(SessionConfig::new())?;


Nit: the annotation is redundant, since SessionBuilder already returns SessionState. Dropping it also lets the use datafusion::execution::SessionState; above go.

Removed, along with the import. The test now builds two sessions, because with the early return the first one comes back untouched and wouldn't go through the rebuild.

comphead · 2026-09-28T16:00:10Z

+    /// re-reading file footers if statistics outlive the session that
+    /// collected them.
+    #[tokio::test]
+    async fn test_in_memory_sessions_share_file_statistics() -> Result<()> {


This writes Parquet and plans a query only to observe one cache entry. Asserting identity, as executor/src/runtime_cache.rs does, is simpler and also pins the invariant the doc comment relies on, that the list-files cache stays per session. Untested sketch:

let caches = |c: &SessionContext| { let cm = &c.runtime_env().cache_manager; (cm.get_file_statistic_cache().unwrap(), cm.get_list_files_cache().unwrap()) }; let ((stats_0, list_0), (stats_1, list_1)) = (caches(&first), caches(&second)); assert!(Arc::ptr_eq(&stats_0, &stats_1)); assert!(!Arc::ptr_eq(&list_0, &list_1));

The test then needs no table, query or with_collect_statistics.

Switched to the identity check. It builds the sessions through BallistaCluster::new_memory, so it also covers the wiring now that the wrapper lives there.

comphead · 2026-09-28T16:00:10Z

+    /// Shared statistics must not outlive the file they describe: `COUNT(*)`
+    /// is answered from exact statistics, so stale ones give a wrong result.
+    #[tokio::test]
+    async fn test_in_memory_sessions_reread_changed_files() -> Result<()> {


A few simplifications, all checked against DataFusion 55.1.0:

register_table can be ctx.register_parquet("t", dir.path().to_str().unwrap(), ParquetReadOptions::default().schema(&schema)), which passes the schema through without inference. That also avoids reusing the name of SessionContext::register_table.

The downcast can be ctx.table("t").await?.count().await? (with rows: usize). It builds the same COUNT(*) aggregate, so exact statistics still answer it.

write_parquet_table then has one caller and can be inlined. The trailing slash, its comment and with_single_file_output(true) are not needed: an existing directory already parses as a collection, and the default Automatic mode writes a single file for a .parquet path.

Keep with_collect_statistics(true) here, since without statistics this test would pass vacuously.

That removes the AsArray, Int64Type, ListingOptions, ParquetFormat and Path imports.

Applied these in 2f8fe37 and kept with_collect_statistics(true). I also checked that the test still fails with the stale count if the listing cache is shared.

Wrap the session builder in BallistaCluster::new_memory instead of InMemoryJobState::new, and make share_file_statistics_cache public. The scheduler binary, standalone mode and the test harnesses all go through new_memory, so they keep the shared cache, while an InMemoryJobState passed to BallistaCluster::new keeps a cache per session, and a custom JobState can opt in with one call. Return a session unchanged when it already has the shared cache. That skips rebuilding the first session, and every session from a builder that reuses one runtime, such as the standalone scheduler built from a client's state. Document that sharing keeps statistics for the scheduler's lifetime and the cases the size and modification time check cannot catch. Assert cache identity in the sharing test, and simplify the COUNT(*) test.

andygrove · 2026-09-28T17:45:21Z

@milenkovicm on the rest of your review:

SF1000. I haven't measured SF1000 with the fix yet. The closest data point is the SF100 files linked 10 times (320 files per table), which approximates SF1000's footer volume. There, Q21 planning went from 1542 ms to 10 ms once lineitem had been scanned. The cost is linear in the number of files, so I'd expect SF1000 to drop from over 4 s to tens of ms after the first query on each table. That first query still pays the full cost.

The listing cache. Two DataFusion caches are involved. The listing cache maps a table's path to the files under it, including each file's size and modification time. A hit is used as is, and by default entries never expire. The statistics cache maps each file to its statistics, and a hit is only used if the size and modification time match the current listing. So each job lists the files itself, and cached statistics are checked against that fresh listing. If the listing cache were shared too, a later job would get the old listing and the old statistics would pass the check. COUNT(*) is answered from statistics, so it would return the old count. New files wouldn't show up either. test_in_memory_sessions_reread_changed_files fails with the stale count if the listing cache is shared.

Session ID. With the default builder every job does get a new random ID. But ctx.session_id() becomes the job's session ID, which is sent with each task, and executors key their runtime cache on it. A builder can set it, like the one new_standalone_scheduler_from_state uses, which hands out the client's state. So I'd rather the wrapper not change it. With the early return @comphead suggested, sessions that already have the shared cache aren't rebuilt at all.

A static cache in default_session_builder_with_cache. The binary never calls default_session_builder (see the inline thread), and a static would share one cache across every scheduler in the process. So I went with new_memory.

villebro

Great improvement! Agree on the general design, but I'd prefer to harden that comparison logic to make sure we leverage etag and version metadata when available.

villebro · 2026-09-28T17:56:56Z

+/// Sharing keeps statistics for the scheduler's lifetime instead of one job's.
+/// Entries are keyed by table and store-relative path, and the check compares
+/// neither e-tags nor versions. So a file rewritten in place at the same size,
+/// quickly enough that its modification time does not change, can still be
+/// served stale statistics, and so can a file in another store with the same
+/// path, size and modification time.


Should we consider using etag or version in the case where those are available, too? That would probably provide the best protection against staleness. So compare in this order of preference:

if both sides have etag, just compare those (minimizes cache miss if the content is byte identical)

if both sides have version, compare those (can cache miss if byte identical change was pushed)

last resort: modification timestamp + size

The order makes sense to me. Comparing only the e-tag when both sides have one also keeps a hit when a file is rewritten with the same bytes, which size plus modification time would count as a change.

The catch is that the comparison isn't in Ballista. ListingTable validates the cached entry itself by calling DataFusion's CachedFileMetadata::is_valid_for, and Cache::get only receives the table and path. So a cache we pass in never sees the current file's e-tag, and we'd have to fork the listing table to change it. I filed apache/datafusion#25841 to do this upstream. Once it lands and we upgrade, the shared cache picks it up with no change here. 43bb787 points the doc comment at that issue.

The staleness note on share_file_statistics_cache says the validity check ignores e-tags and versions. ListingTable makes that check, and a cache only sees the table and path on lookup, so it cannot be tightened from Ballista. Point at apache/datafusion#25841, which tracks it.

andygrove · 2026-09-29T17:44:00Z

There's already a PR open upstream for apache/datafusion#25841: apache/datafusion#25858. It validates e-tags and versions when both sides have them, and adds the object store URL to the statistics, listing and Parquet metadata cache keys. We should review it.

milenkovicm

Thanks @andygrove
i hope tpch1000 gets under 550.

One minor comment, would it make sense/possible to expose ballista. configuration to disable this functionality should we observe corner cases ?

andygrove · 2026-09-29T19:31:38Z

Thanks @andygrove i hope tpch1000 gets under 550.

One minor comment, would it make sense/possible to expose ballista. configuration to disable this functionality should we observe corner cases ?

Good idea. I'll follow up with a config PR.

andygrove requested review from avantgardnerio, comphead, mbutrovich and milenkovicm September 27, 2026 18:54

martin-g reviewed Sep 27, 2026

View reviewed changes

milenkovicm reviewed Sep 27, 2026

View reviewed changes

milenkovicm mentioned this pull request Sep 27, 2026

Ballista 55.0.0 Release #2369

Closed

1 task

comphead reviewed Sep 28, 2026

View reviewed changes

andygrove mentioned this pull request Sep 28, 2026

Validate cached file statistics and Parquet metadata with e_tag, not just size and mtime apache/datafusion#25841

Open

villebro reviewed Sep 28, 2026

View reviewed changes

villebro mentioned this pull request Sep 29, 2026

fix: Validate file cache identity and isolate object stores apache/datafusion#25858

Open

villebro approved these changes Sep 29, 2026

View reviewed changes

milenkovicm approved these changes Sep 29, 2026

View reviewed changes

andygrove merged commit da20520 into apache:main Sep 29, 2026
25 checks passed

andygrove deleted the tpch-sf1000-slow-planning branch September 29, 2026 19:32

This was referenced Sep 29, 2026

feat(scheduler): add a setting to disable the shared file statistics cache #2510

Merged

Discussion: release the Python client separately from the Rust crates #2511

Closed

NoahKusaba mentioned this pull request Oct 1, 2026

Performance review: shuffle, executor, scheduler and client findings (tracking) #2520

Open

	let shared = OnceLock::new();
	static shared: FileStatitsticsCache = OnceLock::new();

Conversation

andygrove commented Sep 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

milenkovicm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

comphead left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andygrove commented Sep 28, 2026

Uh oh!

villebro left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andygrove commented Sep 29, 2026

Uh oh!

milenkovicm left a comment

Choose a reason for hiding this comment

Uh oh!

andygrove commented Sep 29, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

andygrove commented Sep 27, 2026 •

edited

Loading