Shuffle-fetch failures lose their type, so the map-stage resubmit never fires and the job fails

[andygrove]

Describe the bug

When an executor dies after producing shuffle output, Ballista fails the query instead of re-running the map stage. The FetchPartitionError recovery path — arguably the most important HA path in the engine — never fires, because the typed error that triggers it is destroyed before the scheduler sees it.

Root cause

ballista/core/src/execution_plans/shuffle_reader.rs (lines ~1756 and ~1773) converts a Ballista error into a DataFusion error by Debug-formatting it into a string:

.map_err(|e| DataFusionError::Execution(format!("{e:?}")))

A typed BallistaError::FetchFailed(executor_id, map_stage_id, map_partition_id, desc) is thereby flattened into an opaque DataFusionError::Execution("FetchFailed(\"...\", 1, 2, \"...\")").

By the time the executor builds its task status, the error is BallistaError::DataFusionError(Execution(..)), not BallistaError::FetchFailed(..). So this arm in ballista/core/src/error.rs (~line 208):

BallistaError::FetchFailed(executor_id, map_stage_id, map_partition_id, desc) => FailedTask {
    failed_reason: Some(FailedReason::FetchPartitionError(FetchPartitionError { .. })),
    ..
}

never matches. The error falls through to the catch-all (~line 248) and is reported as:

retryable: false,
failed_reason: Some(FailedReason::ExecutionError(..)),

The scheduler therefore never receives a FetchPartitionError, so the logic in ballista/scheduler/src/state/execution_graph.rs (~line 826) that removes the lost input partitions and resubmits the map stage is unreachable via this path. The stage is failed instead, and the job fails.

To Reproduce

Reproduced end to end on a real multi-process cluster by the chaos harness in #2026 (scenario executor_killed_mid_stage_is_recovered, which fails under both AQE on and AQE off):

1. Start a scheduler and 2 executors.
2. Run a query whose plan has a shuffle (join + grouped aggregate).
3. SIGKILL one executor while stage 1 has running tasks.
4. The downstream stage tries to fetch shuffle partitions from the dead executor.

Observed job failure:

Job failed due to stage 2 failed: Task failed due to runtime execution error:
DataFusionError(Execution("FetchFailed(\"e40d5c70-cf84-466a-b390-91cc861401ff\", 1, 2,
  \"Error connecting to Ballista scheduler or executor at http://127.0.0.1:49643:
   tonic::transport::Error(Transport, ConnectError(... Connection refused ...))\")"))

Note Task failed due to runtime execution error — that is the catch-all message from error.rs, confirming the FetchFailed arm did not match.

Expected behavior

The scheduler should classify this as a FetchPartitionError, remove the input partitions produced by the dead executor, resubmit the map stage, and complete the query with the correct result. Ballista already implements all of that; it simply never gets the chance.

Additional context

This is one of two bugs with the same underlying shape: structured Ballista errors lose their type when they cross the DataFusion error boundary, and retryability is then decided by shallow pattern matches that cannot see through the wrapping. See the sibling issue on DataFusionError::Shared and retryable IO errors.

A fix likely needs BallistaError to survive the round-trip through DataFusionError (e.g. carried as a boxed external error rather than a formatted string) so that error.rs can recover the typed variant — plus a downcast/unwrap at the classification site rather than a shallow matches!.

[villebro]

I'm happy to take this on if it's still up for grabs.