diff --git a/ballista/core/src/error.rs b/ballista/core/src/error.rs index 33504fdedf..6d5cd868b5 100644 --- a/ballista/core/src/error.rs +++ b/ballista/core/src/error.rs @@ -76,6 +76,14 @@ impl Into> for BallistaError { } } +impl BallistaError { + /// Keeps a Ballista error structural across DataFusion's stream boundary, so + /// task-failure classification can recover it instead of parsing a string. + pub fn into_datafusion(self) -> DataFusionError { + DataFusionError::External(Box::new(self)) + } +} + /// Creates a general Ballista error from a string message. pub fn ballista_error(message: &str) -> BallistaError { BallistaError::General(message.to_owned()) @@ -117,6 +125,14 @@ impl From for BallistaError { fn from(e: DataFusionError) -> Self { match e { DataFusionError::ArrowError(e, _) => Self::from(*e), + // A Ballista error carried across DataFusion arrives as + // `External(Box)`; recover the original. + DataFusionError::External(inner) => match inner.downcast::() { + Ok(b) => *b, + Err(other) => BallistaError::DataFusionError(Box::new( + DataFusionError::External(other), + )), + }, _ => BallistaError::DataFusionError(Box::new(e)), } } @@ -211,9 +227,8 @@ struct FetchFailedDetails { desc: String, } -/// Recovers a shuffle fetch failure that crossed DataFusion as -/// `ArrowError::ExternalError(FetchFailed)` and may now be wrapped in -/// `DataFusionError` layers. +/// Recovers a shuffle fetch failure carried across DataFusion, seeing through +/// any wrapper layers. fn find_fetch_failed(e: &BallistaError) -> Option { match e { BallistaError::FetchFailed(executor_id, map_stage_id, map_partition_id, desc) => { @@ -224,30 +239,30 @@ fn find_fetch_failed(e: &BallistaError) -> Option { desc: desc.clone(), }) } - BallistaError::ArrowError(e) => fetch_failed_in_arrow(e), - BallistaError::DataFusionError(e) => fetch_failed_in_datafusion(e), + BallistaError::DataFusionError(e) => match e.find_root() { + DataFusionError::External(inner) => inner + .downcast_ref::() + .and_then(find_fetch_failed), + _ => None, + }, _ => None, } } -fn fetch_failed_in_datafusion(e: &DataFusionError) -> Option { - match e.find_root() { - DataFusionError::ArrowError(e, _) => fetch_failed_in_arrow(e), - _ => None, - } -} - -fn fetch_failed_in_arrow(e: &ArrowError) -> Option { - let ArrowError::ExternalError(inner) = e else { - return None; - }; - if let Some(e) = inner.downcast_ref::() { - return find_fetch_failed(e); - } - if let Some(e) = inner.downcast_ref::() { - return fetch_failed_in_datafusion(e); +/// Whether the error is a retryable IO failure, native or carried across +/// DataFusion under any wrapper layer. +fn is_retryable_io(e: &BallistaError) -> bool { + match e { + BallistaError::IoError(_) => true, + BallistaError::DataFusionError(e) => match e.find_root() { + DataFusionError::IoError(_) => true, + DataFusionError::External(inner) => inner + .downcast_ref::() + .is_some_and(is_retryable_io), + _ => false, + }, + _ => false, } - None } impl From for FailedTask { @@ -273,20 +288,9 @@ impl From for FailedTask { count_to_failures: false, failed_reason: Some(FailedReason::TaskKilled(TaskKilled {})), }, - BallistaError::IoError(io) => { + ref e if is_retryable_io(e) => { FailedTask { - error: format!("Task failed due to Ballista IO error: {io:?}"), - // IO error is considered to be temporary and retryable - retryable: true, - count_to_failures: true, - failed_reason: Some(FailedReason::IoError(IoError {})), - } - } - BallistaError::DataFusionError(e) - if matches!(e.find_root(), DataFusionError::IoError(_)) => - { - FailedTask { - error: format!("Task failed due to DataFusion IO error: {e:?}"), + error: format!("Task failed due to IO error: {e:?}"), // IO error is considered to be temporary and retryable retryable: true, count_to_failures: true, @@ -329,40 +333,48 @@ mod tests { } #[test] - fn bare_datafusion_io_error_is_retryable() { - let e = BallistaError::DataFusionError(Box::new(DataFusionError::IoError( - io::Error::new(io::ErrorKind::ConnectionReset, "connection reset"), - ))); - let task = io_failed_task(e); - assert!(task.retryable); - assert!(matches!(task.failed_reason, Some(FailedReason::IoError(_)))); - } - - #[test] - fn shared_wrapped_io_error_is_retryable() { - // Errors from a join's shared build side arrive as Shared(Arc); - // the classifier must see through the wrapper or it will not retry a - // transient IO failure. - let inner = DataFusionError::IoError(io::Error::new( - io::ErrorKind::ConnectionReset, - "connection reset", - )); - let shared = DataFusionError::Shared(Arc::new(inner)); - let e = BallistaError::DataFusionError(Box::new(shared)); - let task = io_failed_task(e); - assert!(task.retryable); - assert!(matches!(task.failed_reason, Some(FailedReason::IoError(_)))); - } - - #[test] - fn context_wrapped_shared_io_error_is_retryable() { - let inner = DataFusionError::IoError(io::Error::other("s3 timeout")); - let shared = DataFusionError::Shared(Arc::new(inner)); - let ctx = shared.context("reading join build side"); - let e = BallistaError::DataFusionError(Box::new(ctx)); - let task = io_failed_task(e); - assert!(task.retryable); - assert!(matches!(task.failed_reason, Some(FailedReason::IoError(_)))); + fn io_error_is_retryable_through_any_wrapper() { + // Both a native DataFusion IoError and a BallistaError::IoError carried + // across DataFusion as External stay retryable, including under the + // Shared/Context layers DataFusion adds (e.g. a join's shared build side). + let df_io = || { + DataFusionError::IoError(io::Error::new( + io::ErrorKind::ConnectionReset, + "connection reset", + )) + }; + let cases: Vec<(&str, BallistaError)> = vec![ + ("native", BallistaError::DataFusionError(Box::new(df_io()))), + ( + "shared", + BallistaError::DataFusionError(Box::new(DataFusionError::Shared( + Arc::new(df_io()), + ))), + ), + ( + "context+shared", + BallistaError::DataFusionError(Box::new( + DataFusionError::Shared(Arc::new(df_io())) + .context("reading join build side"), + )), + ), + ( + "external", + wrap_in_external(BallistaError::IoError(io::Error::new( + io::ErrorKind::ConnectionReset, + "connection reset", + ))), + ), + ]; + for (label, e) in cases { + let task = io_failed_task(e); + assert!(task.retryable, "{label} should be retryable"); + assert!(task.count_to_failures, "{label} should count to failures"); + assert!( + matches!(task.failed_reason, Some(FailedReason::IoError(_))), + "{label} should classify as IoError", + ); + } } #[test] @@ -402,25 +414,18 @@ mod tests { )); } - /// 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))) + /// Builds the single-wrap shape that reaches task failure classification. + fn wrap_in_external(inner: BallistaError) -> BallistaError { + BallistaError::DataFusionError(Box::new(inner.into_datafusion())) } - 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))); + fn wrap_in_shared_external(inner: BallistaError) -> BallistaError { + let df = DataFusionError::Shared(Arc::new(inner.into_datafusion())); BallistaError::DataFusionError(Box::new(df)) } - fn wrap_in_context_arrow_external(inner: BallistaError) -> BallistaError { - let df = datafusion_arrow_external(inner).context("reading shuffle partition"); + fn wrap_in_context_external(inner: BallistaError) -> BallistaError { + let df = inner.into_datafusion().context("reading shuffle partition"); BallistaError::DataFusionError(Box::new(df)) } @@ -451,13 +456,10 @@ mod tests { } #[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", - ))); + fn datafusion_external_fetch_failed_converts_to_bare_fetch_failed() { + let e = BallistaError::from( + fetch_failed("exec-1", 3, 7, "connection reset").into_datafusion(), + ); match e { BallistaError::FetchFailed( @@ -477,29 +479,28 @@ mod tests { #[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 e = wrap_in_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 e = wrap_in_shared_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 e = wrap_in_context_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 e = wrap_in_external(BallistaError::General("boom".to_string())); let task = FailedTask::from(e); assert!(!task.retryable); assert!(matches!( diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index a27abd0a17..256ce37306 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -18,6 +18,7 @@ use crate::JobId; use crate::client::BallistaClient; use crate::config::BallistaConfig; +use crate::error::BallistaError; use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt}; use crate::serde::protobuf::get_job_status_result::FlightProxy; use crate::serde::protobuf::{ @@ -28,7 +29,6 @@ use crate::serde::protobuf::{ use crate::serde::protobuf::{ExecutorMetadata, SuccessfulJob}; use crate::utils::{GrpcClientConfig, create_grpc_client_endpoint}; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::arrow::error::ArrowError; use datafusion::arrow::record_batch::RecordBatch; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::context::TaskContext; @@ -47,7 +47,7 @@ use datafusion_proto::logical_plan::{ AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec, }; use datafusion_proto::physical_plan::{AsExecutionPlan, PhysicalExtensionCodec}; -use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt}; +use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, error, info}; use parking_lot::Mutex; use std::fmt::Debug; @@ -298,20 +298,17 @@ impl ExecutionPlan for DistributedQueryExec { let session_config = context.session_config().clone(); if session_config.ballista_config().client_pull() { - let stream = futures::stream::once( - execute_query_pull( - self.scheduler_url.clone(), - self.session_id.clone(), - query, - self.config.grpc_client_max_message_size(), - GrpcClientConfig::from(&self.config), - Arc::new(self.metrics.clone()), - Arc::clone(&self.job_id), - partition, - session_config, - ) - .map_err(|e| ArrowError::ExternalError(Box::new(e))), - ) + let stream = futures::stream::once(execute_query_pull( + self.scheduler_url.clone(), + self.session_id.clone(), + query, + self.config.grpc_client_max_message_size(), + GrpcClientConfig::from(&self.config), + Arc::new(self.metrics.clone()), + Arc::clone(&self.job_id), + partition, + session_config, + )) .try_flatten() .inspect(move |batch| { metric_total_bytes.add( @@ -327,19 +324,16 @@ impl ExecutionPlan for DistributedQueryExec { let schema = self.schema(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } else { - let stream = futures::stream::once( - execute_query_push( - self.scheduler_url.clone(), - query, - self.config.grpc_client_max_message_size(), - GrpcClientConfig::from(&self.config), - Arc::new(self.metrics.clone()), - Arc::clone(&self.job_id), - partition, - session_config, - ) - .map_err(|e| ArrowError::ExternalError(Box::new(e))), - ) + let stream = futures::stream::once(execute_query_push( + self.scheduler_url.clone(), + query, + self.config.grpc_client_max_message_size(), + GrpcClientConfig::from(&self.config), + Arc::new(self.metrics.clone()), + Arc::clone(&self.job_id), + partition, + session_config, + )) .try_flatten() .inspect(move |batch| { metric_total_bytes.add( @@ -610,8 +604,7 @@ async fn execute_query_pull( use_tls, io_retries_times, io_retry_wait_time_ms, - ) - .map_err(|e| ArrowError::ExternalError(Box::new(e))); + ); futures::stream::once(f).try_flatten() }); @@ -776,8 +769,7 @@ async fn execute_query_push( use_tls, io_retries_times, io_retry_wait_time_ms, - ) - .map_err(|e| ArrowError::ExternalError(Box::new(e))); + ); futures::stream::once(f).try_flatten() }); @@ -881,7 +873,7 @@ async fn fetch_partition( flight_transport, ) .await - .map_err(|e| DataFusionError::External(Box::new(e))) + .map_err(BallistaError::into_datafusion) } #[cfg(test)] diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index 6f4ec1a6a0..37af2903bb 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/ballista/core/src/execution_plans/shuffle_reader.rs @@ -25,7 +25,6 @@ use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt}; use crate::serde::scheduler::{PartitionLocation, PartitionStats}; use crate::utils::GrpcClientConfig; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::arrow::error::ArrowError; use datafusion::arrow::ipc::reader::StreamReader; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::runtime::SpawnedTask; @@ -606,7 +605,7 @@ impl AbortableReceiverStream { } impl Stream for AbortableReceiverStream { - type Item = result::Result; + type Item = result::Result; fn poll_next( mut self: std::pin::Pin<&mut Self>, @@ -614,7 +613,7 @@ impl Stream for AbortableReceiverStream { ) -> std::task::Poll> { self.inner .poll_next_unpin(cx) - .map_err(|e| ArrowError::ExternalError(Box::new(e))) + .map_err(BallistaError::into_datafusion) } } @@ -1618,7 +1617,7 @@ mod tests { assert!(batches.is_err()); - // BallistaError::FetchFailed -> ArrowError::ExternalError -> ballistaError::FetchFailed + // BallistaError::FetchFailed -> DataFusionError::External -> BallistaError::FetchFailed let ballista_error = batches.unwrap_err(); assert!(matches!( ballista_error, diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 23257e633a..76948984e9 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -27,6 +27,7 @@ use std::sync::Arc; use std::time::Instant; use crate::JobId; +use crate::error::BallistaError; use crate::execution_plans::{ OrderedRangeRepartitionExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, create_shuffle_path, @@ -55,7 +56,6 @@ use datafusion::physical_plan::{ }; use futures::TryStreamExt; -use datafusion::arrow::error::ArrowError; use datafusion::execution::context::TaskContext; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; @@ -575,12 +575,7 @@ impl ShuffleWriterExec { compression_type, ) .await - .map_err(|e| { - DataFusionError::ArrowError( - Box::new(ArrowError::ExternalError(Box::new(e))), - None, - ) - })?; + .map_err(BallistaError::into_datafusion)?; let rows = stats.num_rows.unwrap_or(0) as usize; write_metrics.input_rows.add(rows); write_metrics.output_rows.add(rows); @@ -756,7 +751,6 @@ impl ExecutionPlan for ShuffleWriterExec { &job_id, stage_id, ) - .map_err(|e| ArrowError::ExternalError(Box::new(e))) }) .try_flatten(), ))) @@ -1009,15 +1003,12 @@ mod tests { matches!(e.as_ref(), DataFusionError::Shared(_)), "expected shared DataFusionError, got {e:?}" ); - let DataFusionError::ArrowError(arrow, _) = e.find_root() else { - panic!("expected ArrowError root, got {:?}", e.find_root()); - }; - let ArrowError::ExternalError(inner) = arrow.as_ref() else { - panic!("expected Arrow external error, got {arrow:?}"); + let DataFusionError::External(inner) = e.find_root() else { + panic!("expected External root, got {:?}", e.find_root()); }; assert!( inner.downcast_ref::().is_some(), - "expected Arrow external BallistaError, got {inner:?}" + "expected External BallistaError, got {inner:?}" ); } diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 4d8405d196..9cb9e73f51 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -42,7 +42,6 @@ use crate::serde::protobuf::ShuffleWritePartition; use crate::utils::create_write_options; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::arrow::error::ArrowError; use datafusion::arrow::ipc::CompressionType; use datafusion::arrow::ipc::writer::StreamWriter; use datafusion::arrow::record_batch::RecordBatch; @@ -971,7 +970,6 @@ impl ExecutionPlan for SortShuffleWriterExec { &job_id, stage_id, ) - .map_err(|e| ArrowError::ExternalError(Box::new(e))) }) .try_flatten(), ))) diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 99b7c38d5d..ce0e0d073d 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -1824,7 +1824,6 @@ mod test { self, ExecutionError, FailedTask, FetchPartitionError, IoError, JobStatus, TaskKilled, failed_task, job_status, task_status, }; - use datafusion::arrow::error::ArrowError; use datafusion::common::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::TaskContext; use datafusion::physical_plan::{ @@ -3223,17 +3222,15 @@ mod test { map_stage_id: usize, map_partition_id: usize, ) -> FailedTask { - let err = BallistaError::DataFusionError(Box::new(DataFusionError::ArrowError( - Box::new(ArrowError::ExternalError(Box::new( - BallistaError::FetchFailed( - executor_id.to_owned(), - map_stage_id, - map_partition_id, - "FetchPartitionError".to_owned(), - ), - ))), - None, - ))); + let err = BallistaError::DataFusionError(Box::new( + BallistaError::FetchFailed( + executor_id.to_owned(), + map_stage_id, + map_partition_id, + "FetchPartitionError".to_owned(), + ) + .into_datafusion(), + )); FailedTask::from(err) }