fix: classify wrapped FetchFailed as fetch partition error by villebro · Pull Request #2202 · apache/datafusion-ballista

villebro · 2026-07-28T21:13:16Z

Closes #2027.

Rationale for this change

Issue #2027 describes scheduler recovery failing after an executor dies after producing shuffle output. The shuffle fetch path starts with a structured FetchFailed { executor_id, map_stage_id, map_partition_id, message }, but that error could be flattened into a generic DataFusionError::Execution("...") string while crossing the shuffle-writer handoff. Once flattened, the scheduler saw a non-retryable ExecutionError instead of the FetchPartitionError signal it needs to invalidate and re-run the lost map output.

The exact location hints in the issue reflect the codebase when the issue was written. The production type loss fixed here is the wrapping/stringification path above.

What changes are included in this PR?

The core fix keeps shuffle-fetch failures structural until task failure classification. The regular and sort-shuffle writers now share the original structural error with every output-partition handoff instead of formatting it into a string, and the executor uses the existing BallistaError conversion at the query-stage boundary.

Error behavior before and after:

Before	After
A shuffle fetch failure could reach the scheduler as a generic `ExecutionError` string.	Bare or wrapped `FetchFailed` errors become `FetchPartitionError`, including when wrapped through `ArrowError::ExternalError`, `DataFusionError::ArrowError`, `DataFusionError::Shared`, or `DataFusionError::Context`.
Sort-shuffle retryable IO errors crossing the writer handoff could be flattened into generic `ExecutionError`s.	The sort-shuffle coordinator shares the original `DataFusionError`, preserving the `IoError` root so it stays retryable.
Executor-reported `Cancelled` from stale task cleanup could be treated as a fatal execution failure.	It is classified as retryable, non-counting `TaskKilled`; explicit job cancellation still marks the job terminal before executor tasks are cancelled.

Enabling the chaos regression also exposed adjacent recovery gaps needed for that end-to-end scenario:

Executor-loss handling now wakes push scheduling after reset work.
AQE now preserves stage input location metadata so lost map output can be invalidated and resubmitted.
AQE now drops stale active shuffle locations for a lost executor before constructing downstream stages.

The chaos harness changes are mostly regression coverage for this fix. Two harness-stability changes are slightly unrelated to the production fix but were needed to avoid CI flakiness while testing it: the chaos executor now enables the existing scheduler-connect retry with a 10s timeout, and the startup canary uses a longer executor timeout.

The ArrowError::ExternalError wrapping is intentional: it matches the shape produced by the shuffle fetch stream boundary, and Ballista's conversion/classification logic now recovers that shape through Shared and Context wrappers. This PR does not introduce a separate DataFusionError::External(Box<BallistaError>) fetch path.

Are there any user-facing changes?

No API or documentation changes.

Behavior change: shuffle fetch failures that were previously reported as non-retryable execution errors are now classified as fetch-partition errors, so the scheduler can resubmit lost map output instead of failing the job. Sort-shuffle IO errors crossing the same handoff keep their retryable IO classification. Executor-side cancellation of stale task attempts during recovery is treated as retryable cleanup, while explicit job cancellation still marks the job terminal before executor tasks are cancelled. AQE also drops stale active shuffle locations for a lost executor before constructing downstream stages.

No breaking changes to public APIs.

villebro · 2026-07-29T03:13:58Z

I'm marking this as draft until #2026 is merged so we can enable the disabled chaos harness tests related to shuffle fetching as part of this PR.

andygrove · 2026-07-29T04:39:52Z

I'm marking this as draft until #2026 is merged so we can enable the disabled chaos harness tests related to shuffle fetching as part of this PR.

#2026 is now merged

villebro · 2026-07-29T16:50:43Z

I rebased and I'm now enabling the new tests, will push and remove draft when I have everything working.

When an executor dies after producing shuffle output, the shuffle fetch fails with a typed BallistaError::FetchFailed. The scheduler has a recovery path for this via FailedReason::FetchPartitionError, but the job failed instead of recovering because the typed error was hidden before task failure classification. As the fetch error propagates up the stream it can be carried as ArrowError::ExternalError(FetchFailed), then as DataFusionError::ArrowError(..), and then as BallistaError::DataFusionError before building the task status. By the time it reaches From<BallistaError> for FailedTask, it is no longer a bare FetchFailed, so the classifier falls through to FailedReason::ExecutionError. Fix this in two places: use result.map_err(BallistaError::from) at the executor boundary so the immediate DataFusion/Arrow wrapper can be unwrapped, and add classifier-side recovery for FetchFailed values that still arrive wrapped in DataFusionError, ArrowError::ExternalError, Shared, or Context. Add regression tests for executor-boundary conversion, bare and wrapped FetchFailed classification, and non-fetch classification. Update the existing scheduler fetch-failure recovery test to build its failed task through the real classifier, covering the classification -> recovery -> success path without duplicating the graph scenario.

milenkovicm · 2026-08-07T16:08:09Z

One maybe unreleased question, what do you think if we remove BallistaError from interfaces and return DataFusionError. We can keep simplified BallistaError but it would be packed in DFE:: External?

We have a mix of DF and Ballista errors across interfaces and sometimes it gets confusing

villebro · 2026-08-07T16:46:58Z

One maybe unreleased question, what do you think if we remove BallistaError from interfaces and return DataFusionError. We can keep simplified BallistaError but it would be packed in DFE:: External?

We have a mix of DF and Ballista errors across interfaces and sometimes it gets confusing

@milenkovicm I agree we should simplify/consolidate as much as possible. However, I think it's a good idea to do this as a follow-up, as the consolidation work will probably branch out into many unrelated areas (this PR is already getting quite loaded). If this sounds good I can take this on after this PR is merged?

milenkovicm · 2026-08-07T16:50:18Z

I agree, we can take it as a follow up it you believe it makes sense.

Had a quick look at this pr, it makes sense, will do another pass later

villebro · 2026-08-07T16:59:05Z

Had a quick look at this pr, it makes sense, will do another pass later

Thanks for reviewing. Btw, I removed Draft from this PR yesterday as I thought it was blocking CI. Turns out GH was just borked all day. But let me make sure CI is green and do one last review pass and then ping you for a final review when I feel it's completely done.

milenkovicm

thanks @villebro great addition, i have just one question/comment regarding errors used im not sure i understand

This stays targeted to the known shuffle fetch path, where FetchFailed travels through ArrowError::ExternalError. A direct DataFusionError::External(Box) fetch path is not known to be produced here and is outside this change.

milenkovicm · 2026-08-07T18:30:11Z

                    .await
-                    .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
+                    .map_err(|e| {
+                        DataFusionError::ArrowError(


why dont we use DataFusionError::External ?

Good point. I agree it would be a cleaner/shorter wrapper shape, but switching to it here would require adding an explicit downcast/unwrapping path and corresponding tests, which would probably increase the complexity of this PR. I think it could be an excellent follow-up PR focused specifically on harmonizing these error paths, and one I can take on after this is merged.

milenkovicm · 2026-08-07T18:36:51Z

 /// Wraps stage plan with addition of references to previous stages
 pub(crate) struct AdaptiveStageInfo {
    pub(crate) plan: Arc<dyn ShuffleWriter>,
-    #[allow(dead_code)] // TODO: still not sure if this is needed


i like to call myself API visionist I had a hunch it'll be needed in the future 😀

milenkovicm · 2026-08-07T18:38:22Z

+        map_stage_id: usize,
+        map_partition_id: usize,
+    ) -> FailedTask {
+        let err = BallistaError::DataFusionError(Box::new(DataFusionError::ArrowError(


same question here, cant we use DFE::External

villebro · 2026-08-07T21:11:10Z

@milenkovicm I left a comment: let me know if you're ok with following up with a targeted follow-up PR for consolidating/simplifying the error paths or if you prefer to have those changes in this PR.

milenkovicm · 2026-08-07T21:24:22Z

Thanks @villebro we can split them if you prefer it, let's merge this

andygrove requested a review from milenkovicm July 28, 2026 22:58

villebro marked this pull request as draft July 29, 2026 03:11

villebro force-pushed the fix/shuffle-fetch-error-type branch 2 times, most recently from 04578f8 to cfd9ad0 Compare July 30, 2026 20:32

villebro force-pushed the fix/shuffle-fetch-error-type branch from cfd9ad0 to 1a243dd Compare August 6, 2026 21:31

villebro marked this pull request as ready for review August 6, 2026 21:32

villebro added 6 commits August 7, 2026 08:20

enable chaos harness tests

b537884

fix failing test

112e66f

widen scope

f01f213

Fix AQE shuffle recovery after executor loss

fde2b82

Share shuffle writer errors across output handoffs

6a2e64c

villebro force-pushed the fix/shuffle-fetch-error-type branch from 1a243dd to 6a2e64c Compare August 7, 2026 15:25

increase timeout to decrease flakiness

1ed91e9

villebro added 2 commits August 7, 2026 10:30

Document shared shuffle writer error handoff

9337428

Clean up chaos harness README wording

6960df1

milenkovicm approved these changes Aug 7, 2026

View reviewed changes

milenkovicm merged commit a9d9a72 into apache:main Aug 7, 2026
25 checks passed

villebro deleted the fix/shuffle-fetch-error-type branch August 7, 2026 21:24

villebro mentioned this pull request Aug 10, 2026

refactor(core): use DataFusionError::External for structured task errors #2276

Merged

Uh oh!

Conversation

villebro commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

villebro commented Jul 29, 2026

Uh oh!

andygrove commented Jul 29, 2026

Uh oh!

villebro commented Jul 29, 2026

Uh oh!

milenkovicm commented Aug 7, 2026

Uh oh!

villebro commented Aug 7, 2026

Uh oh!

milenkovicm commented Aug 7, 2026

Uh oh!

villebro commented Aug 7, 2026

Uh oh!

milenkovicm left a comment

Choose a reason for hiding this comment

Uh oh!

milenkovicm Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

villebro Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

milenkovicm Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

milenkovicm Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

villebro commented Aug 7, 2026

Uh oh!

milenkovicm commented Aug 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

villebro commented Jul 28, 2026 •

edited

Loading