Make ExecutionPlan sync (apache#2307)

ballista/rust/core/src/execution_plans/distributed_query.rs

@@ -41,7 +41,6 @@ use datafusion::physical_plan::{
 
 use crate::serde::protobuf::execute_query_params::OptionalSessionId;
 use crate::serde::{AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec};
-use async_trait::async_trait;
 use datafusion::arrow::error::{ArrowError, Result as ArrowResult};
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::execution::context::TaskContext;
@@ -122,7 +121,6 @@ impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
     }
 }
 
-#[async_trait]
 impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
     fn as_any(&self) -> &dyn Any {
         self
@@ -162,7 +160,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
         }))
     }
 
-    async fn execute(
+    fn execute(
         &self,
         partition: usize,
         _context: Arc<TaskContext>,

ballista/rust/core/src/execution_plans/shuffle_reader.rs

@@ -21,7 +21,6 @@ use std::sync::Arc;
 use crate::client::BallistaClient;
 use crate::serde::scheduler::{PartitionLocation, PartitionStats};
 
-use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
 
 use datafusion::error::{DataFusionError, Result};
@@ -64,7 +63,6 @@ impl ShuffleReaderExec {
     }
 }
 
-#[async_trait]
 impl ExecutionPlan for ShuffleReaderExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -101,7 +99,7 @@ impl ExecutionPlan for ShuffleReaderExec {
         ))
     }
 
-    async fn execute(
+    fn execute(
         &self,
         partition: usize,
         _context: Arc<TaskContext>,

ballista/rust/core/src/execution_plans/shuffle_writer.rs

@@ -33,7 +33,6 @@ use crate::utils;
 
 use crate::serde::protobuf::ShuffleWritePartition;
 use crate::serde::scheduler::PartitionStats;
-use async_trait::async_trait;
 use datafusion::arrow::array::{
     ArrayBuilder, ArrayRef, StringBuilder, StructBuilder, UInt32Builder, UInt64Builder,
 };
@@ -155,7 +154,7 @@ impl ShuffleWriterExec {
 
         async move {
             let now = Instant::now();
-            let mut stream = plan.execute(input_partition, context).await?;
+            let mut stream = plan.execute(input_partition, context)?;
 
             match output_partitioning {
                 None => {
@@ -293,7 +292,6 @@ impl ShuffleWriterExec {
     }
 }
 
-#[async_trait]
 impl ExecutionPlan for ShuffleWriterExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -336,7 +334,7 @@ impl ExecutionPlan for ShuffleWriterExec {
         )?))
     }
 
-    async fn execute(
+    fn execute(
         &self,
         partition: usize,
         context: Arc<TaskContext>,
@@ -459,7 +457,7 @@ mod tests {
             work_dir.into_path().to_str().unwrap().to_owned(),
             Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)),
         )?;
-        let mut stream = query_stage.execute(0, task_ctx).await?;
+        let mut stream = query_stage.execute(0, task_ctx)?;
         let batches = utils::collect_stream(&mut stream)
             .await
             .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?;
@@ -516,7 +514,7 @@ mod tests {
             work_dir.into_path().to_str().unwrap().to_owned(),
             Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)),
         )?;
-        let mut stream = query_stage.execute(0, task_ctx).await?;
+        let mut stream = query_stage.execute(0, task_ctx)?;
         let batches = utils::collect_stream(&mut stream)
             .await
             .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?;

ballista/rust/core/src/execution_plans/unresolved_shuffle.rs

@@ -18,7 +18,6 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::context::TaskContext;
@@ -63,7 +62,6 @@ impl UnresolvedShuffleExec {
     }
 }
 
-#[async_trait]
 impl ExecutionPlan for UnresolvedShuffleExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -101,7 +99,7 @@ impl ExecutionPlan for UnresolvedShuffleExec {
         ))
     }
 
-    async fn execute(
+    fn execute(
         &self,
         _partition: usize,
         _context: Arc<TaskContext>,

ballista/rust/core/src/serde/mod.rs

@@ -477,7 +477,6 @@ mod tests {
         }
     }
 
-    #[async_trait]
     impl ExecutionPlan for TopKExec {
         /// Return a reference to Any that can be used for downcasting
         fn as_any(&self) -> &dyn Any {
@@ -515,7 +514,7 @@ mod tests {
         }
 
         /// Execute one partition and return an iterator over RecordBatch
-        async fn execute(
+        fn execute(
             &self,
             _partition: usize,
             _context: Arc<TaskContext>,

ballista/rust/executor/src/collect.rs

@@ -22,7 +22,6 @@ use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::{any::Any, pin::Pin};
 
-use async_trait::async_trait;
 use datafusion::arrow::{
     datatypes::SchemaRef, error::Result as ArrowResult, record_batch::RecordBatch,
 };
@@ -49,7 +48,6 @@ impl CollectExec {
     }
 }
 
-#[async_trait]
 impl ExecutionPlan for CollectExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -78,18 +76,16 @@ impl ExecutionPlan for CollectExec {
         unimplemented!()
     }
 
-    async fn execute(
+    fn execute(
         &self,
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
         assert_eq!(0, partition);
         let num_partitions = self.plan.output_partitioning().partition_count();
 
-        let futures = (0..num_partitions).map(|i| self.plan.execute(i, context.clone()));
-        let streams = futures::future::join_all(futures)
-            .await
-            .into_iter()
+        let streams = (0..num_partitions)
+            .map(|i| self.plan.execute(i, context.clone()))
             .collect::<Result<Vec<_>>>()
             .map_err(|e| DataFusionError::Execution(format!("BallistaError: {:?}", e)))?;
 

datafusion-examples/examples/custom_datasource.rs

@@ -196,7 +196,6 @@ impl CustomExec {
     }
 }
 
-#[async_trait]
 impl ExecutionPlan for CustomExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -225,7 +224,7 @@ impl ExecutionPlan for CustomExec {
         Ok(self)
     }
 
-    async fn execute(
+    fn execute(
         &self,
         _partition: usize,
         _context: Arc<TaskContext>,
@@ -243,7 +242,7 @@ impl ExecutionPlan for CustomExec {
             account_array.append_value(user.bank_account)?;
         }
 
-        return Ok(Box::pin(MemoryStream::try_new(
+        Ok(Box::pin(MemoryStream::try_new(
             vec![RecordBatch::try_new(
                 self.projected_schema.clone(),
                 vec![
@@ -253,7 +252,7 @@ impl ExecutionPlan for CustomExec {
             )?],
             self.schema(),
             None,
-        )?));
+        )?))
     }
 
     fn statistics(&self) -> Statistics {

datafusion/core/src/datasource/file_format/csv.rs

@@ -161,7 +161,7 @@ mod tests {
         let projection = Some(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12]);
         let exec = get_exec("aggregate_test_100.csv", &projection, None).await?;
         let task_ctx = ctx.task_ctx();
-        let stream = exec.execute(0, task_ctx).await?;
+        let stream = exec.execute(0, task_ctx)?;
 
         let tt_batches: i32 = stream
             .map(|batch| {

datafusion/core/src/datasource/file_format/json.rs

@@ -131,7 +131,7 @@ mod tests {
         let projection = None;
         let exec = get_exec(&projection, None).await?;
         let task_ctx = ctx.task_ctx();
-        let stream = exec.execute(0, task_ctx).await?;
+        let stream = exec.execute(0, task_ctx)?;
 
         let tt_batches: i32 = stream
             .map(|batch| {

datafusion/core/src/datasource/file_format/parquet.rs

@@ -512,7 +512,7 @@ mod tests {
         let projection = None;
         let exec = get_exec("alltypes_plain.parquet", &projection, None).await?;
         let task_ctx = ctx.task_ctx();
-        let stream = exec.execute(0, task_ctx).await?;
+        let stream = exec.execute(0, task_ctx)?;
 
         let tt_batches = stream
             .map(|batch| {

datafusion/core/src/datasource/memory.rs

@@ -76,7 +76,7 @@ impl MemTable {
                 let context1 = context.clone();
                 let exec = exec.clone();
                 tokio::spawn(async move {
-                    let stream = exec.execute(part_i, context1.clone()).await?;
+                    let stream = exec.execute(part_i, context1.clone())?;
                     common::collect(stream).await
                 })
             })
@@ -103,7 +103,7 @@ impl MemTable {
             let mut output_partitions = vec![];
             for i in 0..exec.output_partitioning().partition_count() {
                 // execute this *output* partition and collect all batches
-                let mut stream = exec.execute(i, context.clone()).await?;
+                let mut stream = exec.execute(i, context.clone())?;
                 let mut batches = vec![];
                 while let Some(result) = stream.next().await {
                     batches.push(result?);
@@ -177,7 +177,7 @@ mod tests {
 
         // scan with projection
         let exec = provider.scan(&Some(vec![2, 1]), &[], None).await?;
-        let mut it = exec.execute(0, task_ctx).await?;
+        let mut it = exec.execute(0, task_ctx)?;
         let batch2 = it.next().await.unwrap()?;
         assert_eq!(2, batch2.schema().fields().len());
         assert_eq!("c", batch2.schema().field(0).name());
@@ -209,7 +209,7 @@ mod tests {
         let provider = MemTable::try_new(schema, vec![vec![batch]])?;
 
         let exec = provider.scan(&None, &[], None).await?;
-        let mut it = exec.execute(0, task_ctx).await?;
+        let mut it = exec.execute(0, task_ctx)?;
         let batch1 = it.next().await.unwrap()?;
         assert_eq!(3, batch1.schema().fields().len());
         assert_eq!(3, batch1.num_columns());
@@ -365,7 +365,7 @@ mod tests {
             MemTable::try_new(Arc::new(merged_schema), vec![vec![batch1, batch2]])?;
 
         let exec = provider.scan(&None, &[], None).await?;
-        let mut it = exec.execute(0, task_ctx).await?;
+        let mut it = exec.execute(0, task_ctx)?;
         let batch1 = it.next().await.unwrap()?;
         assert_eq!(3, batch1.schema().fields().len());
         assert_eq!(3, batch1.num_columns());

datafusion/core/src/physical_optimizer/aggregate_statistics.rs

@@ -304,7 +304,7 @@ mod tests {
 
         // A ProjectionExec is a sign that the count optimization was applied
         assert!(optimized.as_any().is::<ProjectionExec>());
-        let result = common::collect(optimized.execute(0, task_ctx).await?).await?;
+        let result = common::collect(optimized.execute(0, task_ctx)?).await?;
         assert_eq!(result[0].schema(), Arc::new(Schema::new(vec![col])));
         assert_eq!(
             result[0]

datafusion/core/src/physical_plan/aggregates/mod.rs

@@ -29,7 +29,6 @@ use crate::physical_plan::{
 };
 use arrow::array::ArrayRef;
 use arrow::datatypes::{Field, Schema, SchemaRef};
-use async_trait::async_trait;
 use datafusion_common::Result;
 use datafusion_expr::Accumulator;
 use datafusion_physical_expr::expressions::Column;
@@ -145,7 +144,6 @@ impl AggregateExec {
     }
 }
 
-#[async_trait]
 impl ExecutionPlan for AggregateExec {
     /// Return a reference to Any that can be used for down-casting
     fn as_any(&self) -> &dyn Any {
@@ -196,12 +194,12 @@ impl ExecutionPlan for AggregateExec {
         )?))
     }
 
-    async fn execute(
+    fn execute(
         &self,
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
-        let input = self.input.execute(partition, context).await?;
+        let input = self.input.execute(partition, context)?;
         let group_expr = self.group_expr.iter().map(|x| x.0.clone()).collect();
 
         let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
@@ -417,7 +415,6 @@ mod tests {
     use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
     use arrow::error::Result as ArrowResult;
     use arrow::record_batch::RecordBatch;
-    use async_trait::async_trait;
     use datafusion_common::{DataFusionError, Result};
     use datafusion_physical_expr::{AggregateExpr, PhysicalExpr, PhysicalSortExpr};
     use futures::{FutureExt, Stream};
@@ -489,8 +486,7 @@ mod tests {
         )?);
 
         let result =
-            common::collect(partial_aggregate.execute(0, task_ctx.clone()).await?)
-                .await?;
+            common::collect(partial_aggregate.execute(0, task_ctx.clone())?).await?;
 
         let expected = vec![
             "+---+---------------+-------------+",
@@ -522,7 +518,7 @@ mod tests {
         )?);
 
         let result =
-            common::collect(merged_aggregate.execute(0, task_ctx.clone()).await?).await?;
+            common::collect(merged_aggregate.execute(0, task_ctx.clone())?).await?;
         assert_eq!(result.len(), 1);
 
         let batch = &result[0];
@@ -556,7 +552,6 @@ mod tests {
         pub yield_first: bool,
     }
 
-    #[async_trait]
     impl ExecutionPlan for TestYieldingExec {
         fn as_any(&self) -> &dyn Any {
             self
@@ -587,7 +582,7 @@ mod tests {
             )))
         }
 
-        async fn execute(
+        fn execute(
             &self,
             _partition: usize,
             _context: Arc<TaskContext>,