[GitHub] [arrow-datafusion] alamb commented on a diff in pull request #6308: Parallel merge sort (#6162)

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

alamb commented on code in PR #6308:
URL: https://github.com/apache/arrow-datafusion/pull/6308#discussion_r1189868392


##########
datafusion/core/src/physical_plan/common.rs:
##########
@@ -180,6 +181,31 @@ pub(crate) fn spawn_execution(
     })
 }
 
+/// If running in a tokio context spawns the execution of `stream` to a separate task
+/// allowing it to execute in parallel with an intermediate buffer of size `buffer`
+pub(crate) fn spawn_buffered(
+    mut input: SendableRecordBatchStream,
+    buffer: usize,
+) -> SendableRecordBatchStream {
+    // Use tokio only if running from a tokio context (#2201)

Review Comment:
   Nice



##########
datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs:
##########
@@ -181,29 +180,12 @@ impl ExecutionPlan for SortPreservingMergeExec {
                 result
             }
             _ => {
-                // Use tokio only if running from a tokio context (#2201)
-                let receivers = match tokio::runtime::Handle::try_current() {
-                    Ok(_) => (0..input_partitions)
-                        .map(|part_i| {
-                            let (sender, receiver) = mpsc::channel(1);
-                            let join_handle = spawn_execution(
-                                self.input.clone(),
-                                sender,
-                                part_i,
-                                context.clone(),
-                            );
-
-                            RecordBatchReceiverStream::create(
-                                &schema,
-                                receiver,
-                                join_handle,
-                            )
-                        })
-                        .collect(),
-                    Err(_) => (0..input_partitions)
-                        .map(|partition| self.input.execute(partition, context.clone()))
-                        .collect::<Result<_>>()?,
-                };
+                let receivers = (0..input_partitions)

Review Comment:
   This looks beautiful



-- 
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