[arrow-ballista] branch main updated: Implement the with_new_children for ShuffleReaderExec (#821)

[nj...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

nju_yaho pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-ballista.git


The following commit(s) were added to refs/heads/main by this push:
     new b65464e4 Implement the with_new_children for ShuffleReaderExec (#821)
b65464e4 is described below

commit b65464e4b73590470fa69aad5b6954300ad243a0
Author: yahoNanJing <90...@users.noreply.github.com>
AuthorDate: Wed Jun 28 14:44:41 2023 +0800

    Implement the with_new_children for ShuffleReaderExec (#821)
    
    Co-authored-by: yangzhong <ya...@ebay.com>
---
 ballista/core/src/execution_plans/shuffle_reader.rs | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs
index 5acde9ec..f80002cb 100644
--- a/ballista/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/core/src/execution_plans/shuffle_reader.rs
@@ -33,7 +33,7 @@ use datafusion::arrow::error::ArrowError;
 use datafusion::arrow::ipc::reader::FileReader;
 use datafusion::arrow::record_batch::RecordBatch;
 
-use datafusion::error::{DataFusionError, Result};
+use datafusion::error::Result;
 use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use datafusion::physical_plan::{
@@ -110,9 +110,11 @@ impl ExecutionPlan for ShuffleReaderExec {
         self: Arc<Self>,
         _children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        Err(DataFusionError::Plan(
-            "Ballista ShuffleReaderExec does not support with_new_children()".to_owned(),
-        ))
+        Ok(Arc::new(ShuffleReaderExec::try_new(
+            self.stage_id,
+            self.partition.clone(),
+            self.schema.clone(),
+        )?))
     }
 
     fn execute(
@@ -432,6 +434,7 @@ mod tests {
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use datafusion::arrow::ipc::writer::FileWriter;
     use datafusion::arrow::record_batch::RecordBatch;
+    use datafusion::common::DataFusionError;
     use datafusion::physical_expr::expressions::Column;
     use datafusion::physical_plan::common;
     use datafusion::physical_plan::memory::MemoryExec;