diff --git a/src/common/query/src/error.rs b/src/common/query/src/error.rs index 538c18b704..4c53be86b6 100644 --- a/src/common/query/src/error.rs +++ b/src/common/query/src/error.rs @@ -270,46 +270,79 @@ pub fn datafusion_status_code( e: &DataFusionError, default_status: Option, ) -> StatusCode { - match e { + let mut error = e; + loop { + error = match error { + DataFusionError::Shared(inner) => inner, + DataFusionError::Context(_, inner) | DataFusionError::Diagnostic(_, inner) => inner, + _ => break, + }; + } + + match error { DataFusionError::Internal(_) => StatusCode::Internal, DataFusionError::NotImplemented(_) => StatusCode::Unsupported, DataFusionError::Plan(_) => StatusCode::PlanQuery, - DataFusionError::External(e) => { - if let Some(ext) = (*e).downcast_ref::() { + DataFusionError::External(error) => { + if let Some(ext) = (*error).downcast_ref::() { ext.status_code() - } else if let Some(ext) = (*e).downcast_ref::() { + } else if let Some(ext) = (*error).downcast_ref::() { ext.status_code() } else { default_status.unwrap_or(StatusCode::EngineExecuteQuery) } } - DataFusionError::Diagnostic(_, e) => datafusion_status_code::(e, default_status), _ => default_status.unwrap_or(StatusCode::EngineExecuteQuery), } } #[cfg(test)] mod tests { + use std::sync::Arc; + use common_error::ext::PlainError; use super::*; #[test] - fn test_datafusion_status_code_external_errors() { - let boxed_error = || { + fn test_datafusion_status_code_preserves_external_errors_through_wrappers() { + let boxed_error = |status| { DataFusionError::External(Box::new(BoxedError::new(PlainError::new( "neutral error".to_string(), - StatusCode::RequestOutdated, + status, )))) }; - assert_eq!( - datafusion_status_code::(&boxed_error(), None), - StatusCode::RequestOutdated - ); - assert_eq!( - datafusion_status_code::(&boxed_error(), Some(StatusCode::PlanQuery)), - StatusCode::RequestOutdated - ); + + for status in [ + StatusCode::RequestOutdated, + StatusCode::Unknown, + StatusCode::Unsupported, + ] { + let errors = [ + boxed_error(status), + DataFusionError::Shared(Arc::new(boxed_error(status))), + DataFusionError::Context("context".to_string(), Box::new(boxed_error(status))), + DataFusionError::Diagnostic( + Box::new(datafusion_common::Diagnostic::new_error("diagnostic", None)), + Box::new(boxed_error(status)), + ), + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion_common::Diagnostic::new_error("diagnostic", None)), + Box::new(boxed_error(status)), + )), + ))), + ]; + + for error in errors { + assert_eq!(datafusion_status_code::(&error, None), status); + assert_eq!( + datafusion_status_code::(&error, Some(StatusCode::PlanQuery)), + status + ); + } + } let direct_error = DataFusionError::External(Box::new(Error::DynFilterPayloadTooLarge { payload_size_bytes: 2, @@ -317,19 +350,78 @@ mod tests { location: Location::default(), })); assert_eq!( - datafusion_status_code::(&direct_error, None), + datafusion_status_code::(&direct_error, Some(StatusCode::Internal)), StatusCode::PlanQuery ); - let unknown_error = - || DataFusionError::External(Box::new(std::io::Error::other("neutral error"))); + let boundary_error = + DataFusionError::Shared(Arc::new(DataFusionError::External(Box::new( + BoxedError::new(common_recordbatch::error::Error::PhysicalExpr { + error: DataFusionError::NotImplemented("inner error".to_string()), + location: Location::default(), + }), + )))); assert_eq!( - datafusion_status_code::(&unknown_error(), None), - StatusCode::EngineExecuteQuery - ); - assert_eq!( - datafusion_status_code::(&unknown_error(), Some(StatusCode::PlanQuery)), - StatusCode::PlanQuery + datafusion_status_code::(&boundary_error, Some(StatusCode::PlanQuery)), + StatusCode::Internal ); } + + #[test] + fn test_datafusion_status_code_uses_default_for_untyped_errors() { + let wrap = |error| { + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion_common::Diagnostic::new_error("diagnostic", None)), + Box::new(error), + )), + ))) + }; + let errors = || { + [ + ( + DataFusionError::External(Box::new(std::io::Error::other("neutral io error"))), + StatusCode::EngineExecuteQuery, + StatusCode::Unknown, + ), + ( + DataFusionError::Internal("neutral internal error".to_string()), + StatusCode::Internal, + StatusCode::Internal, + ), + ( + DataFusionError::NotImplemented("neutral not implemented error".to_string()), + StatusCode::Unsupported, + StatusCode::Unsupported, + ), + ( + DataFusionError::Plan("neutral plan error".to_string()), + StatusCode::PlanQuery, + StatusCode::PlanQuery, + ), + ( + DataFusionError::External(Box::new(DataFusionError::Internal( + "inner error".to_string(), + ))), + StatusCode::EngineExecuteQuery, + StatusCode::Unknown, + ), + ] + }; + + for (error, none_expected, default_expected) in + errors() + .into_iter() + .chain(errors().map(|(error, none_expected, default_expected)| { + (wrap(error), none_expected, default_expected) + })) + { + assert_eq!(datafusion_status_code::(&error, None), none_expected); + assert_eq!( + datafusion_status_code::(&error, Some(StatusCode::Unknown)), + default_expected + ); + } + } } diff --git a/src/common/recordbatch/src/error.rs b/src/common/recordbatch/src/error.rs index f1425492fb..bddb8ae803 100644 --- a/src/common/recordbatch/src/error.rs +++ b/src/common/recordbatch/src/error.rs @@ -205,16 +205,26 @@ impl ErrorExt for Error { | Error::PhysicalExpr { .. } | Error::RecordBatchSliceIndexOverflow { .. } => StatusCode::Internal, - Error::PollStream { - error: datafusion::error::DataFusionError::External(source), - .. - } => source - .downcast_ref::() - .map_or(StatusCode::EngineExecuteQuery, |source| { - source.status_code() - }), + Error::PollStream { error, .. } => { + let mut error = error; + loop { + error = match error { + datafusion::error::DataFusionError::Shared(inner) => inner, + datafusion::error::DataFusionError::Context(_, inner) + | datafusion::error::DataFusionError::Diagnostic(_, inner) => inner, + _ => break, + }; + } - Error::PollStream { .. } => StatusCode::EngineExecuteQuery, + match error { + datafusion::error::DataFusionError::External(source) => source + .downcast_ref::() + .map_or(StatusCode::EngineExecuteQuery, |source| { + source.status_code() + }), + _ => StatusCode::EngineExecuteQuery, + } + } Error::ArrowCompute { .. } => StatusCode::IllegalState, @@ -254,38 +264,98 @@ impl ErrorExt for Error { #[cfg(test)] mod tests { + use std::sync::Arc; + use common_error::ext::PlainError; + use datafusion::error::DataFusionError; use super::*; #[test] - fn poll_stream_status_code_preserves_direct_external_boxed_error() { - let cases = [ - (StatusCode::RequestOutdated, StatusCode::RequestOutdated), - (StatusCode::Unknown, StatusCode::Unknown), - ]; + fn poll_stream_status_code_preserves_boxed_error_through_wrappers() { + let boxed_error = |status| { + DataFusionError::External(Box::new(BoxedError::new(PlainError::new( + "neutral error".to_string(), + status, + )))) + }; - for (source_status, expected_status) in cases { - let error = Error::PollStream { - error: datafusion::error::DataFusionError::External(Box::new(BoxedError::new( - PlainError::new("neutral error".to_string(), source_status), + for status in [ + StatusCode::RequestOutdated, + StatusCode::Unknown, + StatusCode::Unsupported, + ] { + let errors = [ + boxed_error(status), + DataFusionError::Shared(Arc::new(boxed_error(status))), + DataFusionError::Context("context".to_string(), Box::new(boxed_error(status))), + DataFusionError::Diagnostic( + Box::new(datafusion::common::Diagnostic::new_error( + "diagnostic", + None, + )), + Box::new(boxed_error(status)), + ), + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion::common::Diagnostic::new_error( + "diagnostic", + None, + )), + Box::new(boxed_error(status)), + )), ))), - location: Location::default(), - }; - assert_eq!(error.status_code(), expected_status); + ]; + + for error in errors { + let error = Error::PollStream { + error, + location: Location::default(), + }; + assert_eq!(error.status_code(), status); + } } + + let error = Error::PollStream { + error: DataFusionError::Shared(Arc::new(DataFusionError::External(Box::new( + BoxedError::new(Error::PhysicalExpr { + error: DataFusionError::NotImplemented("inner error".to_string()), + location: Location::default(), + }), + )))), + location: Location::default(), + }; + assert_eq!(error.status_code(), StatusCode::Internal); } #[test] - fn poll_stream_status_code_defaults_for_unrecognized_datafusion_errors() { - let errors = [ - datafusion::error::DataFusionError::External(Box::new(std::io::Error::other( - "neutral io error", - ))), - datafusion::error::DataFusionError::Internal("neutral internal error".to_string()), - ]; + fn poll_stream_status_code_defaults_for_other_datafusion_errors() { + let wrap = |error| { + DataFusionError::Shared(Arc::new(DataFusionError::Context( + "context".to_string(), + Box::new(DataFusionError::Diagnostic( + Box::new(datafusion::common::Diagnostic::new_error( + "diagnostic", + None, + )), + Box::new(error), + )), + ))) + }; + let errors = || { + [ + DataFusionError::External(Box::new(std::io::Error::other("neutral io error"))), + DataFusionError::Internal("neutral internal error".to_string()), + DataFusionError::NotImplemented("neutral not implemented error".to_string()), + DataFusionError::Plan("neutral plan error".to_string()), + DataFusionError::External(Box::new(DataFusionError::Internal( + "inner error".to_string(), + ))), + ] + }; - for error in errors { + for error in errors().into_iter().chain(errors().map(wrap)) { let error = Error::PollStream { error, location: Location::default(), diff --git a/src/flow/src/batching_mode/task/test.rs b/src/flow/src/batching_mode/task/test.rs index e2c417aa0d..e9f8e3e89d 100644 --- a/src/flow/src/batching_mode/task/test.rs +++ b/src/flow/src/batching_mode/task/test.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::Arc; use catalog::memory::MemoryCatalogManager; use catalog::{DeregisterTableRequest, RegisterTableRequest}; @@ -1611,11 +1612,28 @@ async fn test_fenced_repair_stale_fence_next_plan_is_scoped_base_repair() { { let mut state = task.state.write().unwrap(); + let error = Err::<(), _>(BoxedError::new( + common_recordbatch::error::Error::PollStream { + error: datafusion::error::DataFusionError::Shared(Arc::new( + datafusion::error::DataFusionError::External(Box::new(BoxedError::new( + MockError::new(StatusCode::RequestOutdated), + ))), + )), + location: snafu::Location::default(), + }, + )) + .context(crate::error::ExternalSnafu) + .unwrap_err(); + let reason = BatchingTask::query_failure_reason( + &error, + &QueryCoverage::FencedRepairChunk { high: high.clone() }, + ); + assert_eq!(reason, FlowQueryFallbackReason::SnapshotFenceExpired); let decision = BatchingTask::apply_query_failure_to_state( &mut state, std::time::Duration::from_millis(1), &QueryCoverage::FencedRepairChunk { high }, - FlowQueryFallbackReason::SnapshotFenceExpired, + reason, ); assert_eq!( decision, @@ -1625,9 +1643,11 @@ async fn test_fenced_repair_stale_fence_next_plan_is_scoped_base_repair() { }) ); assert!(state.pending_fenced_repair().is_none()); + assert_eq!(state.dirty_time_windows.len(), 1); // Simulate the outer execution failure restore for the in-flight chunk. state.restore_scoped_windows(&filter); + assert_eq!(state.dirty_time_windows.len(), 2); } let plan = task diff --git a/src/query/src/dist_plan/merge_scan.rs b/src/query/src/dist_plan/merge_scan.rs index ffe154ade0..6311d1ff3b 100644 --- a/src/query/src/dist_plan/merge_scan.rs +++ b/src/query/src/dist_plan/merge_scan.rs @@ -1426,6 +1426,7 @@ mod tests { use datafusion::config::ConfigOptions; use datafusion::execution::SessionStateBuilder; use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult; + use datafusion::physical_plan::repartition::RepartitionExec; use datafusion_common::TableReference; use datafusion_expr::{LogicalPlanBuilder, col, lit}; use datafusion_physical_expr::Distribution; @@ -1890,7 +1891,7 @@ mod tests { } #[tokio::test] - async fn merge_scan_later_stream_error_preserves_status_code() { + async fn repartitioned_merge_scan_later_stream_error_preserves_status_code() { let region_id = RegionId::new(1024, 1); let handler = Arc::new(TestRegionQueryHandler::with_responses(vec![( region_id, @@ -1902,10 +1903,26 @@ mod tests { )), ))], )])); - let exec = - merge_scan_exec_with_handler(vec![region_id], expected_int64_schema(), handler, 1); + let merge_scan = Arc::new(merge_scan_exec_with_handler( + vec![region_id], + expected_int64_schema(), + handler, + 1, + )); + let repartition = + RepartitionExec::try_new(merge_scan, Partitioning::RoundRobinBatch(2)).unwrap(); + assert_eq!( + repartition + .properties() + .output_partitioning() + .partition_count(), + 2 + ); + let mut stream = common_recordbatch::adapter::RecordBatchStreamAdapter::try_new( - exec.to_stream(Arc::new(TaskContext::default()), 0).unwrap(), + repartition + .execute(0, Arc::new(TaskContext::default())) + .unwrap(), ) .unwrap();