[GitHub] [arrow-datafusion] DDtKey commented on a diff in pull request #5196: fix: make `write_csv/json/parquet` cancel-safe

["DDtKey (via GitHub)" <gi...@apache.org>]

DDtKey commented on code in PR #5196:
URL: https://github.com/apache/arrow-datafusion/pull/5196#discussion_r1097222724


##########
datafusion/core/src/physical_plan/file_format/csv.rs:
##########
@@ -286,38 +288,44 @@ pub async fn plan_to_csv(
     let path = path.as_ref();
     // create directory to contain the CSV files (one per partition)
     let fs_path = Path::new(path);
-    match fs::create_dir(fs_path) {
-        Ok(()) => {
-            let mut tasks = vec![];
-            for i in 0..plan.output_partitioning().partition_count() {
-                let plan = plan.clone();
-                let filename = format!("part-{i}.csv");
-                let path = fs_path.join(filename);
-                let file = fs::File::create(path)?;
-                let mut writer = csv::Writer::new(file);
-                let task_ctx = Arc::new(TaskContext::from(state));
-                let stream = plan.execute(i, task_ctx)?;
-                let handle: JoinHandle<Result<()>> = task::spawn(async move {
-                    stream
-                        .map(|batch| writer.write(&batch?))
-                        .try_collect()
-                        .await
-                        .map_err(DataFusionError::from)
-                });
-                tasks.push(handle);
-            }
-            futures::future::join_all(tasks)
-                .await
-                .into_iter()
-                .try_for_each(|result| {
-                    result.map_err(|e| DataFusionError::Execution(format!("{e}")))?
-                })?;
-            Ok(())
-        }
-        Err(e) => Err(DataFusionError::Execution(format!(
+    if let Err(e) = fs::create_dir(fs_path) {
+        return Err(DataFusionError::Execution(format!(
             "Could not create directory {path}: {e:?}"
-        ))),
+        )));
     }
+
+    let mut tasks = vec![];
+    // Create cancellation-token to interrupt background execution on drop.
+    let cancellation_token = CancellationToken::new();

Review Comment:
   I've decided not to introduce a new pattern here and changed it to `AbortOnDropSingle` per task. 
   You still can check how it looked like before: https://github.com/apache/arrow-datafusion/pull/5196/commits/3918d49bf9f130989a61391a3bc45fbe5318a14c



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org