fix: Unwind panic in spawned threads

Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>

Author: MazterQyou

datafusion/core/src/physical_plan/hash_aggregate.rs

@@ -25,10 +25,10 @@ use std::vec;
 use ahash::RandomState;
 use futures::{
     stream::{Stream, StreamExt},
-    Future,
+    Future, FutureExt,
 };
 
-use crate::error::Result;
+use crate::error::{DataFusionError, Result};
 use crate::physical_plan::hash_utils::create_hashes;
 use crate::physical_plan::{
     Accumulator, AggregateExpr, DisplayFormatType, Distribution, ExecutionPlan,
@@ -576,16 +576,32 @@ impl GroupedHashAggregateStream {
         let elapsed_compute = baseline_metrics.elapsed_compute().clone();
 
         let join_handle = tokio::spawn(async move {
-            let result = compute_grouped_hash_aggregate(
-                mode,
-                schema_clone,
-                group_expr,
-                aggr_expr,
-                input,
-                elapsed_compute,
-            )
+            let result = std::panic::AssertUnwindSafe(async move {
+                compute_grouped_hash_aggregate(
+                    mode,
+                    schema_clone,
+                    group_expr,
+                    aggr_expr,
+                    input,
+                    elapsed_compute,
+                )
+                .await
+                .record_output(&baseline_metrics)
+            })
+            .catch_unwind()
             .await
-            .record_output(&baseline_metrics);
+            .unwrap_or_else(|panic_payload| {
+                let msg = if let Some(s) = panic_payload.downcast_ref::<&str>() {
+                    s
+                } else if let Some(s) = panic_payload.downcast_ref::<String>() {
+                    s.as_str()
+                } else {
+                    "unknown panic"
+                };
+                Err(ArrowError::ExternalError(Box::new(
+                    DataFusionError::Execution(format!("Panic: {}", msg)),
+                )))
+            });
 
             // failing here is OK, the receiver is gone and does not care about the result
             tx.send(result).ok();
@@ -804,15 +820,31 @@ impl HashAggregateStream {
         let schema_clone = schema.clone();
         let elapsed_compute = baseline_metrics.elapsed_compute().clone();
         let join_handle = tokio::spawn(async move {
-            let result = compute_hash_aggregate(
-                mode,
-                schema_clone,
-                aggr_expr,
-                input,
-                elapsed_compute,
-            )
+            let result = std::panic::AssertUnwindSafe(async move {
+                compute_hash_aggregate(
+                    mode,
+                    schema_clone,
+                    aggr_expr,
+                    input,
+                    elapsed_compute,
+                )
+                .await
+                .record_output(&baseline_metrics)
+            })
+            .catch_unwind()
             .await
-            .record_output(&baseline_metrics);
+            .unwrap_or_else(|panic_payload| {
+                let msg = if let Some(s) = panic_payload.downcast_ref::<&str>() {
+                    s
+                } else if let Some(s) = panic_payload.downcast_ref::<String>() {
+                    s.as_str()
+                } else {
+                    "unknown panic"
+                };
+                Err(ArrowError::ExternalError(Box::new(
+                    DataFusionError::Execution(format!("Panic: {}", msg)),
+                )))
+            });
 
             // failing here is OK, the receiver is gone and does not care about the result
             tx.send(result).ok();

datafusion/core/src/physical_plan/windows/window_agg_exec.rs

@@ -17,7 +17,7 @@
 
 //! Stream and channel implementations for window function expressions.
 
-use crate::error::Result;
+use crate::error::{DataFusionError, Result};
 use crate::execution::context::TaskContext;
 use crate::physical_plan::common::AbortOnDropSingle;
 use crate::physical_plan::expressions::PhysicalSortExpr;
@@ -256,10 +256,25 @@ impl WindowAggStream {
         let schema_clone = schema.clone();
         let elapsed_compute = baseline_metrics.elapsed_compute().clone();
         let join_handle = tokio::spawn(async move {
-            let schema = schema_clone.clone();
-            let result =
+            let result = std::panic::AssertUnwindSafe(async move {
+                let schema = schema_clone.clone();
                 WindowAggStream::process(input, window_expr, schema, elapsed_compute)
-                    .await;
+                    .await
+            })
+            .catch_unwind()
+            .await
+            .unwrap_or_else(|panic_payload| {
+                let msg = if let Some(s) = panic_payload.downcast_ref::<&str>() {
+                    s
+                } else if let Some(s) = panic_payload.downcast_ref::<String>() {
+                    s.as_str()
+                } else {
+                    "unknown panic"
+                };
+                Err(ArrowError::ExternalError(Box::new(
+                    DataFusionError::Execution(format!("Panic: {}", msg)),
+                )))
+            });
 
             // failing here is OK, the receiver is gone and does not care about the result
             tx.send(result).ok();