refactor(12089): handle join error when using SpawnedTask::join_unwind

influxdata · Aug 21, 2024 · 52f8595 · 52f8595
1 parent 780f1b6
commit 52f8595
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 8 deletions.
diff --git a/datafusion/core/src/datasource/file_format/arrow.rs b/datafusion/core/src/datasource/file_format/arrow.rs
@@ -341,7 +341,10 @@ impl DataSink for ArrowFileSink {
             }
         }
 
-        demux_task.join_unwind().await?;
+        demux_task
+            .join_unwind()
+            .await
+            .map_err(DataFusionError::ExecutionJoin)??;
         Ok(row_count as u64)
     }
 }

diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs
@@ -836,7 +836,10 @@ impl DataSink for ParquetSink {
             }
         }
 
-        demux_task.join_unwind().await?;
+        demux_task
+            .join_unwind()
+            .await
+            .map_err(DataFusionError::ExecutionJoin)??;
 
         Ok(row_count as u64)
     }
@@ -942,7 +945,10 @@ fn spawn_rg_join_and_finalize_task(
         let num_cols = column_writer_tasks.len();
         let mut finalized_rg = Vec::with_capacity(num_cols);
         for task in column_writer_tasks.into_iter() {
-            let (writer, _col_reservation) = task.join_unwind().await?;
+            let (writer, _col_reservation) = task
+                .join_unwind()
+                .await
+                .map_err(DataFusionError::ExecutionJoin)??;
             let encoded_size = writer.get_estimated_total_bytes();
             rg_reservation.grow(encoded_size);
             finalized_rg.push(writer.close()?);
@@ -1070,7 +1076,8 @@ async fn concatenate_parallel_row_groups(
     while let Some(task) = serialize_rx.recv().await {
         let result = task.join_unwind().await;
         let mut rg_out = parquet_writer.next_row_group()?;
-        let (serialized_columns, mut rg_reservation, _cnt) = result?;
+        let (serialized_columns, mut rg_reservation, _cnt) =
+            result.map_err(DataFusionError::ExecutionJoin)??;
         for chunk in serialized_columns {
             chunk.append_to_row_group(&mut rg_out)?;
             rg_reservation.free();
@@ -1134,7 +1141,10 @@ async fn output_single_parquet_file_parallelized(
     )
     .await?;
 
-    launch_serialization_task.join_unwind().await?;
+    launch_serialization_task
+        .join_unwind()
+        .await
+        .map_err(DataFusionError::ExecutionJoin)??;
     Ok(file_metadata)
 }
 

diff --git a/datafusion/core/src/datasource/file_format/write/orchestration.rs b/datafusion/core/src/datasource/file_format/write/orchestration.rs
@@ -298,8 +298,8 @@ pub(crate) async fn stateless_multipart_put(
         write_coordinator_task.join_unwind(),
         demux_task.join_unwind()
     );
-    r1?;
-    r2?;
+    r1.map_err(DataFusionError::ExecutionJoin)??;
+    r2.map_err(DataFusionError::ExecutionJoin)??;
 
     let total_count = rx_row_cnt.await.map_err(|_| {
         internal_datafusion_err!("Did not receive row count from write coordinator")

diff --git a/datafusion/core/src/datasource/stream.rs b/datafusion/core/src/datasource/stream.rs
@@ -438,6 +438,9 @@ impl DataSink for StreamWrite {
             }
         }
         drop(sender);
-        write_task.join_unwind().await
+        write_task
+            .join_unwind()
+            .await
+            .map_err(DataFusionError::ExecutionJoin)?
     }
 }