fix(query): preserve typed errors through DataFusion wrappers

Signed-off-by: discord9 <discord9@163.com>

src/common/query/src/error.rs

@@ -270,46 +270,79 @@ pub fn datafusion_status_code<T: ErrorExt + 'static>(
     e: &DataFusionError,
     default_status: Option<StatusCode>,
 ) -> 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::<T>() {
+        DataFusionError::External(error) => {
+            if let Some(ext) = (*error).downcast_ref::<T>() {
                 ext.status_code()
-            } else if let Some(ext) = (*e).downcast_ref::<BoxedError>() {
+            } else if let Some(ext) = (*error).downcast_ref::<BoxedError>() {
                 ext.status_code()
             } else {
                 default_status.unwrap_or(StatusCode::EngineExecuteQuery)
             }
         }
-        DataFusionError::Diagnostic(_, e) => datafusion_status_code::<T>(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::<Error>(&boxed_error(), None),
-            StatusCode::RequestOutdated
-        );
-        assert_eq!(
-            datafusion_status_code::<Error>(&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>(&error, None), status);
+                assert_eq!(
+                    datafusion_status_code::<Error>(&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::<Error>(&direct_error, None),
+            datafusion_status_code::<Error>(&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::<Error>(&unknown_error(), None),
-            StatusCode::EngineExecuteQuery
-        );
-        assert_eq!(
-            datafusion_status_code::<Error>(&unknown_error(), Some(StatusCode::PlanQuery)),
-            StatusCode::PlanQuery
+            datafusion_status_code::<Error>(&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>(&error, None), none_expected);
+            assert_eq!(
+                datafusion_status_code::<Error>(&error, Some(StatusCode::Unknown)),
+                default_expected
+            );
+        }
+    }
 }

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::<BoxedError>()
-                .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::<BoxedError>()
+                        .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(),

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

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();