Dandandan commented on a change in pull request #9043:
URL: https://github.com/apache/arrow/pull/9043#discussion_r549861337
##########
File path: rust/datafusion/src/physical_plan/mod.rs
##########
@@ -104,6 +104,26 @@ pub async fn collect(plan: Arc<dyn ExecutionPlan>) ->
Result<Vec<RecordBatch>> {
}
}
+/// Execute the [ExecutionPlan] and collect the results in memory
+pub async fn collect_partitioned(
+ plan: Arc<dyn ExecutionPlan>,
+) -> Result<Vec<Vec<RecordBatch>>> {
+ match plan.output_partitioning().partition_count() {
+ 0 => Ok(vec![]),
+ 1 => {
+ let it = plan.execute(0).await?;
+ Ok(vec![common::collect(it).await?])
+ }
+ _ => {
+ let mut partitions = vec![];
+ for i in 0..plan.output_partitioning().partition_count() {
Review comment:
Could bind this to a variable and reuse it in the code?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
