feat(scheduler): rebuild and resume a query that loses DataFrame::cache() data by akshaychitneni · Pull Request #2730 · apache/datafusion-ballista

akshaychitneni · 2026-11-01T15:07:53Z

Part of #1829. Follows #2633 (design doc step 5: recovery, part 2).

Rationale for this change

Since #2633, a running query that loses cache data fails cleanly and the retry
rebuilds the cache. This PR rebuilds the cache and resumes the query instead.

What changes are included in this PR?

Detection: a failed fetch is cache data when it matches a cache location
(executor, build stage, partition), with no protocol change. If the job's own
input on that executor also matches, both are recovered. On executor loss, the
cache inputs left incomplete are reported in a new
ExecutorLostOutcome::lost_inputs field.
Rebuild and resume: the reading stage rolls back and waits
(ExternalInputLost). The scheduler drops the lost build and then starts a
rebuild, joins one in flight, or reuses a newer build. ExternalInputReady
then re-resolves the stage.
Pipelined jobs: a cache lost before the job is placed is rebuilt, then the
job is placed reading the rebuild. Once placed, the job fails cleanly, because
its tasks and its results to the client are already streaming.
Guard: a job that reads the lost cache more than once, directly or through
another cache built from it, fails cleanly, so a query never sees two versions
of a cache.
Edge cases: rebuilds run off the event loop under the job's admission. A
rebuild that fails or writes a different schema fails the job. A stage out of
attempts names the lost data, so the retry rebuilds it.
Unchanged: adaptive-planner jobs still fail cleanly; that graph cannot
resume a rolled-back stage.
Tests and docs: graph, flow and pipelined-graph tests, chaos Scenario L,
and a new k8s Scenario M that kills the pod holding a cache mid-read. Updated
the design doc, tuning guide and 56.0.0 upgrade guide.

Are there any user-facing changes?

Yes. A static-planner query that loses DataFrame::cache() data now rebuilds it
and finishes, as does a pipelined query that loses it before it starts.
ExecutorLostOutcome gains lost_inputs (see the upgrade guide).

akshaychitneni · 2026-11-01T15:11:03Z

@andygrove @milenkovicm @comphead please review. This includes cache recovery. A query that loses cache data now rebuilds it and resumes, with tests for both blocking and pipelined mode (a pipelined job recovers only if the loss happens before it is placed)

andygrove · 2026-11-05T20:25:56Z

    },
+    /// A job lost `DataFrame::cache()` data it reads; the stages reading it wait
+    /// until it is provided again.
+    ExternalInputLost {


Small one on QueryStageSchedulerEvent. The two new variants will break anyone downstream doing an exhaustive match on it, since the enum is public and shipped in 55. Would you mind either adding #[non_exhaustive] to the enum while you're here, or adding a short note to the 56.0.0 upgrade guide? I'd lean toward #[non_exhaustive] so we stop paying this cost on every new event.

andygrove

Thanks for pushing this forward, the recovery story is a big improvement over failing and retrying. I went through it in detail and left comments inline.

The two I'd want fixed before merging are the double fetch failure in execution_graph.rs, which can stop a job from recovering in what I think is a pretty common case (several tasks failing on the same dead executor at once), and the nested cache hole in the read-more-than-once guard, which can still give a query two versions of a cache. I'd also like to talk through what JobReads keeps alive, since it adds per-job scheduler memory. The rest are smaller hang risks, a docs mismatch and some nits.

andygrove · 2026-11-05T21:15:26Z

+                                        map_stage_id,
+                                        map_partition_id,
+                                    );
+                                    let lost_own = lost_cache.is_empty()


I think there's a problem when two tasks fail fetching the same lost cache data in one status batch.

Say two tasks of the stage reading cache input EXT both fail fetching from dead executor E, and their statuses arrive together. The first failure removes every E location from EXT. For the second one, external_fetch_failures now finds nothing, because those locations are already gone, so lost_own = lost_cache.is_empty() comes out true.

We then call running_stage.remove_input_partitions(map_stage_id, ...), but map_stage_id here is the cache build job's stage id, not one of this graph's inputs. That returns Internal("... is not a valid child stage ID") and the ? propagates it. As far as I can tell TaskManager logs it and drops the whole batch, so there's no rollback and no ExternalInputLost, and the stage input is left half updated. The job never recovers. If the build's stage id happens to match one of the job's own input ids, we'd instead re-run the job's own producer for no reason.

Before this PR the second failure hit the "Stages was marked failed, ignore FetchPartitionError" branch, so this looks like a regression. Maybe lost_own should only be true when map_stage_id is actually one of running_stage.inputs, and a repeat failure on an already-removed cache location should just be ignored? A graph test with two fetch failures from the same executor in one update_task_status call would pin it down.

andygrove · 2026-11-05T21:15:26Z


+/// Whether a job with `reads` reads cache `key` more than once: directly, or
+/// through another cache built from it.
+fn read_more_than_once(reads: &[ResolvedRead], key: &CacheKey) -> bool {


I think this guard misses one direction of nesting. It counts reads of the lost key and caches built from it, but not the case where the lost cache was itself built from another cache the job also reads directly.

Example: a job reads c1 directly (that stage already finished) and also reads c2, which was cached from a plan over c1. E holds parts of both and dies, but only c2's input is reported lost, so this counts one read and the rebuild goes ahead. If I'm reading prepare_caches right, rebuilding c2 recurses into its subplan, claim_or_hit finds c1's entry on dead E, invalidates it, and rebuilds c1 too. Then c2 is rebuilt from the new c1 while the job already consumed the old one. With a non-reproducible subplan (random(), or LIMIT without ORDER BY) the query ends up mixing two versions of c1, which is exactly what this guard is meant to prevent.

Would it work to also fail when any cache the lost one was built from is read by the job, or to check after preparing the rebuild whether it rebuilt anything the job already read?

andygrove · 2026-11-05T21:15:27Z

+        let sender = sender.clone();
+        tokio::spawn(async move {
+            let Some(job_reads) = state.cache_linker.reads_of(&job_id) else {
+                return; // the job ended meanwhile


If reads_of comes back None here, we return silently. By this point the graph has already rolled the stage back and marked the external input incomplete, so if the job is actually still running, nothing rebuilds it, resolves it, or fails it, and the query hangs with no error to the client.

I'm not sure there's a path today where a running cache reader has no read record, so this may be theoretical. It still feels safer to check whether the job is still active and post JobRunningFailed if so, rather than assume it ended. Same thing in provide_input below.

andygrove · 2026-11-05T21:15:27Z

+            };
+            // A rebuild would give the query a second version of the cache, next
+            // to the one it already read elsewhere.
+            if read_more_than_once(&job_reads.reads, &read.key) {


Small one: the read-more-than-once check and the missing-read check run per input inside the loop, after earlier inputs have already been invalidated and their rebuilds claimed and submitted.

So with inputs [A, B], A passes, its build is dropped and a full rebuild starts under the job's admission, then B fails the guard and the job fails. The rebuild of A keeps running for a job that's gone, and A's entry was dropped even though it was healthy on the other executors. Validating every input first, before invalidating or claiming anything, would avoid that.

andygrove · 2026-11-05T21:15:27Z

+    pub(crate) ctx: Arc<SessionContext>,
+    pub(crate) queued_at: u64,
+    /// The submission's admission, for rebuilds; dropped when the job ends.
+    pub(crate) admission: Admission,


I want to check the memory and lifetime side of this. JobReads keeps the job's SessionContext, its Admission, and each cache's full LogicalPlan subplan (via ResolvedRead.cache) alive for the job's whole run. Before this, the subplan was dropped once the read was resolved. That's extra scheduler memory for every running cache-reading job, which we've been strict about.

On Admission, I see the design doc wants shutdown to wait for rebuilds, which makes sense. But Admission is documented in event.rs as held until the job is running or fails to plan, so that doc needs updating. Also, now any end path that skips on_job_terminal would leak it and the scheduler would never report idle. Could we keep only what a rebuild actually needs (the cache key and enough to re-plan it), and double check that every terminal path drops the record?

andygrove · 2026-11-05T21:15:27Z

+  read was planned with, as the read picks its columns by position; otherwise
+  the job fails.
+- **Retries and failure.** Retries count towards `stage_max_failures`, like any
+  stage's; a stage out of attempts, or a rebuild that fails, fails the job with


I don't think this matches the code for a failed rebuild. fail_cache_job reports "a cache this query reads could not be built: ..." without CACHE_DATA_LOST, so anyone matching on "lost DataFrame::cache() data" won't see it. It also means on_job_terminal doesn't take its CACHE_DATA_LOST invalidation branch for that failure. Either the message or the doc should change so they agree.

andygrove · 2026-11-05T21:15:27Z

+    /// is cancelled or fails meanwhile.
    pub fn awaits_placement(&self) -> bool {
-        !self.placed && self.is_running()
+        !self.placed && self.is_running() && self.waiting_external_inputs().is_empty()


Nit: awaits_placement is called for every unplaced pipelined job on each offer round (placement_requirements, place, and the offer loop in state/mod.rs), and waiting_external_inputs() allocates, sorts and dedups a Vec just so we can check it's empty. An .any(|(id, out)| is_external_stage_id(*id) && !out.is_complete()) over the stage inputs would short-circuit without allocating.

andygrove · 2026-11-05T21:15:27Z


+    /// Stages are resolved before placement, so a lost cache read can leave a
+    /// resolved stage waiting too.
+    fn waiting_external_inputs(&self) -> Vec<usize> {


Nit: this is a copy of the trait default with a Resolved arm added, so the two filter/sort/dedup pipelines have to be kept in sync. Since a static graph never has an incomplete external input on a Resolved stage, I think the default could just cover Resolved too and this override could go. A shared helper would also work.

andygrove · 2026-11-05T21:15:27Z

+        if reader.stage_id != input_id {
+            return Ok(Transformed::no(node));
+        }
+        let read = ShuffleReaderExec::try_new(


Not a bug today as far as I can tell, since the cache linker also builds these readers with plain try_new. But this drops whatever else the replaced reader had (broadcast, coalesced, skew join, work_dir, client pool). If a lowering ever starts touching external readers, a broadcast reader would quietly become per-partition here and give wrong join results. Maybe add an assert or a debug check that the reader being replaced is a plain one, so it fails loudly if that assumption ever changes?

andygrove · 2026-11-05T21:15:27Z

@@ -914,12 +899,6 @@ impl TestCluster {
 }

 /// Stage ids come back from the REST API as strings.


Looks like this doc comment got left behind when find_stage moved to rest.rs, so it now documents dir_has_entries. Should move with find_stage or go.

…he() data A job planned by the static planner no longer fails when it loses cache data it reads. The stage reading the cache rolls back and waits, the cache is rebuilt, and the stage resumes once the rebuild is provided. - Detection: a failed fetch is cache data when it matches a cache location's executor, build stage and partition (no fetch relabel, no protocol change). When one of the job's own inputs on that executor matches too, both are recovered. On executor loss the generic reset already drops the executor's cache locations and rolls back their readers; the cache inputs it leaves incomplete are reported. - The graph emits ExternalInputLost instead of re-running a producer, which is not in this graph. The linker drops the lost build, claims a rebuild (or waits on one in flight, or reuses a newer entry), and posts ExternalInputReady, which re-resolves the waiting stages, folding a rebuild with another partition count into the reader's. - A job that reads the lost cache more than once, directly or through another cache built from it, fails cleanly instead, so one query never reads two versions of a cache. - Rebuilds run off the event loop and carry the job's admission, released when the job ends. A rebuild that writes another schema fails the job. - The adaptive planner cannot re-resolve a rolled-back stage, so its jobs still fail cleanly. A stage out of attempts on a cache fetch now names the lost data, so a retry rebuilds it. TaskManager::executor_lost now returns the events to post for jobs that lost cache data.

… keeps each output once

…d on k8s Scenario M force-deletes the executor pod holding part of a cache while a query reads it, after some reading tasks have finished. With the 5s executor timeout this is the executor-loss path, and the Deployment reschedules a replacement. Under the static planner the query must rebuild the cache (a new build job runs) and return the baseline result; under AQE it must fail naming the lost data, and a retry must rebuild. The harness gains a known executor work dir, a lookup of the pods holding a job's data (the REST API does not say which executor ran a task), a targeted pod kill, and a shared wait for a stage's task state.

…aced A pipelined job holds until it can place all its tasks at once. If an executor holding a cache it reads is lost in that window, nothing has run or streamed yet, so failing the job is needless. Until placed, the pipelined graph now drops the lost cache locations, holds placement while a cache input is missing, and on ExternalInputReady rewrites the cache reader's locations, so the job is placed reading the rebuild. Once placed, its tasks stream to each other and its results to the client, so a restart would deliver rows twice; it still fails cleanly, and the retry rebuilds the cache. executor_lost now lets a pipelined job the loss leaves intact reach the reset, where an unplaced job reports its lost cache inputs.

- Drop the "cache lost under a colliding stage id" fetch case: a_lost_cache_is_read_again_once_rebuilt loses a cache built under a colliding id and checks the same rollback, report, and no re-run. - The static case of the executor-loss flow test is covered by the two-readers test, which now also checks the rebuilt location and that nothing stays admitted; the adaptive case stays as its own test. - k8s Scenario M runs under the static planner only; the adaptive outcome is already covered by Scenario L and the scheduler tests. - The pipelined executor-loss test waits until its job is placed, as an unplaced pipelined job now waits for a rebuild instead of failing.

github-actions Bot added the documentation Improvements or additions to documentation label Nov 1, 2026

akshaychitneni force-pushed the feat/df-cache-recovery-2 branch 2 times, most recently from 2dcc12d to 1fb2ca0 Compare November 4, 2026 00:12

andygrove reviewed Nov 5, 2026

View reviewed changes

andygrove added the api-change label Nov 5, 2026

andygrove requested changes Nov 5, 2026

View reviewed changes

akshaychitneni force-pushed the feat/df-cache-recovery-2 branch from 1fb2ca0 to bc00270 Compare November 7, 2026 17:23

Akshay Chitneni added 9 commits November 10, 2026 13:17

test(chaos): a DataFrame::cache() lost before it is read is recovered

11c8b75

docs: DataFrame::cache() recovery rebuilds and resumes a running query

673da47

test(scheduler): a stage that loses a cache after some tasks finished…

7934b3d

… keeps each output once

docs: describe the DataFrame::cache() read record kept for recovery

c0535fc

docs: describe ExecutorLostOutcome in the 56.0.0 upgrade guide

aba81f5

akshaychitneni force-pushed the feat/df-cache-recovery-2 branch from bc00270 to 28d5e05 Compare November 11, 2026 00:49

		@@ -914,12 +899,6 @@ impl TestCluster {
		}

		/// Stage ids come back from the REST API as strings.

Conversation

akshaychitneni commented Nov 1, 2026

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

akshaychitneni commented Nov 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andygrove 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!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

akshaychitneni commented Nov 1, 2026 •

edited

Loading