fix: classify wrapped FetchFailed as fetch partition error

ballista/core/src/error.rs

@@ -24,7 +24,9 @@ use std::{
 };
 
 use crate::serde::protobuf::failed_task::FailedReason;
-use crate::serde::protobuf::{ExecutionError, FailedTask, FetchPartitionError, IoError};
+use crate::serde::protobuf::{
+    ExecutionError, FailedTask, FetchPartitionError, IoError, TaskKilled,
+};
 use datafusion::error::DataFusionError;
 use datafusion::{arrow::error::ArrowError, sql::sqlparser::parser};
 use futures::future::Aborted;
@@ -202,29 +204,75 @@ impl Display for BallistaError {
     }
 }
 
+struct FetchFailedDetails {
+    executor_id: String,
+    map_stage_id: usize,
+    map_partition_id: usize,
+    desc: String,
+}
+
+/// Recovers a shuffle fetch failure that crossed DataFusion as
+/// `ArrowError::ExternalError(FetchFailed)` and may now be wrapped in
+/// `DataFusionError` layers.
+fn find_fetch_failed(e: &BallistaError) -> Option<FetchFailedDetails> {
+    match e {
+        BallistaError::FetchFailed(executor_id, map_stage_id, map_partition_id, desc) => {
+            Some(FetchFailedDetails {
+                executor_id: executor_id.clone(),
+                map_stage_id: *map_stage_id,
+                map_partition_id: *map_partition_id,
+                desc: desc.clone(),
+            })
+        }
+        BallistaError::ArrowError(e) => fetch_failed_in_arrow(e),
+        BallistaError::DataFusionError(e) => fetch_failed_in_datafusion(e),
+        _ => None,
+    }
+}
+
+fn fetch_failed_in_datafusion(e: &DataFusionError) -> Option<FetchFailedDetails> {
+    match e.find_root() {
+        DataFusionError::ArrowError(e, _) => fetch_failed_in_arrow(e),
+        _ => None,
+    }
+}
+
+fn fetch_failed_in_arrow(e: &ArrowError) -> Option<FetchFailedDetails> {
+    let ArrowError::ExternalError(inner) = e else {
+        return None;
+    };
+    if let Some(e) = inner.downcast_ref::<BallistaError>() {
+        return find_fetch_failed(e);
+    }
+    if let Some(e) = inner.downcast_ref::<DataFusionError>() {
+        return fetch_failed_in_datafusion(e);
+    }
+    None
+}
+
 impl From<BallistaError> for FailedTask {
     fn from(e: BallistaError) -> Self {
+        if let Some(fetch_failed) = find_fetch_failed(&e) {
+            return FailedTask {
+                error: fetch_failed.desc,
+                retryable: false,
+                count_to_failures: false,
+                failed_reason: Some(FailedReason::FetchPartitionError(
+                    FetchPartitionError {
+                        executor_id: fetch_failed.executor_id,
+                        map_stage_id: fetch_failed.map_stage_id as u32,
+                        map_partition_id: fetch_failed.map_partition_id as u32,
+                    },
+                )),
+            };
+        }
         match e {
-            BallistaError::FetchFailed(
-                executor_id,
-                map_stage_id,
-                map_partition_id,
-                desc,
-            ) => {
-                FailedTask {
-                    error: desc,
-                    // fetch partition error is considered to be non-retryable
-                    retryable: false,
-                    count_to_failures: false,
-                    failed_reason: Some(FailedReason::FetchPartitionError(
-                        FetchPartitionError {
-                            executor_id,
-                            map_stage_id: map_stage_id as u32,
-                            map_partition_id: map_partition_id as u32,
-                        },
-                    )),
-                }
-            }
+            BallistaError::Cancelled => FailedTask {
+                error: "Task cancelled".to_string(),
+                retryable: true,
+                count_to_failures: false,
+                failed_reason: Some(FailedReason::TaskKilled(TaskKilled {})),
+            },
             BallistaError::IoError(io) => {
                 FailedTask {
                     error: format!("Task failed due to Ballista IO error: {io:?}"),
@@ -266,6 +314,20 @@ mod tests {
         FailedTask::from(e)
     }
 
+    fn fetch_failed(
+        executor_id: &str,
+        map_stage_id: usize,
+        map_partition_id: usize,
+        desc: &str,
+    ) -> BallistaError {
+        BallistaError::FetchFailed(
+            executor_id.to_owned(),
+            map_stage_id,
+            map_partition_id,
+            desc.to_owned(),
+        )
+    }
+
     #[test]
     fn bare_datafusion_io_error_is_retryable() {
         let e = BallistaError::DataFusionError(Box::new(DataFusionError::IoError(
@@ -316,6 +378,17 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn cancelled_task_is_retryable_without_counting_to_failures() {
+        let task = FailedTask::from(BallistaError::Cancelled);
+        assert!(task.retryable);
+        assert!(!task.count_to_failures);
+        assert!(matches!(
+            task.failed_reason,
+            Some(FailedReason::TaskKilled(_))
+        ));
+    }
+
     #[test]
     fn shared_wrapped_non_io_error_stays_non_retryable() {
         let inner = DataFusionError::Plan("bad plan".to_string());
@@ -328,4 +401,110 @@ mod tests {
             Some(FailedReason::ExecutionError(_))
         ));
     }
+
+    /// Builds the wrapped shape that can reach task failure classification.
+    fn wrap_in_arrow_external(inner: BallistaError) -> BallistaError {
+        BallistaError::DataFusionError(Box::new(datafusion_arrow_external(inner)))
+    }
+
+    fn datafusion_arrow_external(inner: BallistaError) -> DataFusionError {
+        DataFusionError::ArrowError(
+            Box::new(ArrowError::ExternalError(Box::new(inner))),
+            None,
+        )
+    }
+
+    fn wrap_in_shared_arrow_external(inner: BallistaError) -> BallistaError {
+        let df = DataFusionError::Shared(Arc::new(datafusion_arrow_external(inner)));
+        BallistaError::DataFusionError(Box::new(df))
+    }
+
+    fn wrap_in_context_arrow_external(inner: BallistaError) -> BallistaError {
+        let df = datafusion_arrow_external(inner).context("reading shuffle partition");
+        BallistaError::DataFusionError(Box::new(df))
+    }
+
+    fn assert_fetch_partition_error(
+        task: FailedTask,
+        executor_id: &str,
+        map_stage_id: u32,
+        map_partition_id: u32,
+        error: &str,
+    ) {
+        assert!(!task.retryable);
+        assert!(!task.count_to_failures);
+        assert_eq!(task.error, error);
+        match task.failed_reason {
+            Some(FailedReason::FetchPartitionError(fp)) => {
+                assert_eq!(fp.executor_id, executor_id);
+                assert_eq!(fp.map_stage_id, map_stage_id);
+                assert_eq!(fp.map_partition_id, map_partition_id);
+            }
+            other => panic!("expected FetchPartitionError, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn bare_fetch_failed_maps_to_fetch_partition_error() {
+        let task = FailedTask::from(fetch_failed("exec-1", 3, 7, "boom"));
+        assert_fetch_partition_error(task, "exec-1", 3, 7, "boom");
+    }
+
+    #[test]
+    fn datafusion_arrow_external_fetch_failed_converts_to_bare_fetch_failed() {
+        let e = BallistaError::from(datafusion_arrow_external(fetch_failed(
+            "exec-1",
+            3,
+            7,
+            "connection reset",
+        )));
+
+        match e {
+            BallistaError::FetchFailed(
+                executor_id,
+                map_stage_id,
+                map_partition_id,
+                desc,
+            ) => {
+                assert_eq!(executor_id, "exec-1");
+                assert_eq!(map_stage_id, 3);
+                assert_eq!(map_partition_id, 7);
+                assert_eq!(desc, "connection reset");
+            }
+            other => panic!("expected bare FetchFailed, got {other:?}"),
+        }
+    }
+
+    #[test]
+    fn wrapped_fetch_failed_is_recovered_as_fetch_partition_error() {
+        let e = wrap_in_arrow_external(fetch_failed("exec-1", 3, 7, "connection reset"));
+        let task = FailedTask::from(e);
+        assert_fetch_partition_error(task, "exec-1", 3, 7, "connection reset");
+    }
+
+    #[test]
+    fn shared_wrapped_fetch_failed_is_recovered() {
+        let e =
+            wrap_in_shared_arrow_external(fetch_failed("exec-2", 1, 2, "peer closed"));
+        let task = FailedTask::from(e);
+        assert_fetch_partition_error(task, "exec-2", 1, 2, "peer closed");
+    }
+
+    #[test]
+    fn context_wrapped_fetch_failed_is_recovered() {
+        let e = wrap_in_context_arrow_external(fetch_failed("exec-3", 5, 9, "timeout"));
+        let task = FailedTask::from(e);
+        assert_fetch_partition_error(task, "exec-3", 5, 9, "timeout");
+    }
+
+    #[test]
+    fn wrapped_non_fetch_error_stays_execution_error() {
+        let e = wrap_in_arrow_external(BallistaError::General("boom".to_string()));
+        let task = FailedTask::from(e);
+        assert!(!task.retryable);
+        assert!(matches!(
+            task.failed_reason,
+            Some(FailedReason::ExecutionError(_))
+        ));
+    }
 }