Add ExecutionEngine abstraction

ballista/core/src/execution_plans/mod.rs

@@ -25,5 +25,5 @@ mod unresolved_shuffle;
 
 pub use distributed_query::DistributedQueryExec;
 pub use shuffle_reader::ShuffleReaderExec;
-pub use shuffle_writer::ShuffleWriterExec;
+pub use shuffle_writer::{ShuffleWriter, ShuffleWriterExec};
 pub use unresolved_shuffle::UnresolvedShuffleExec;

ballista/core/src/execution_plans/shuffle_writer.rs

@@ -22,7 +22,9 @@
 
 use datafusion::physical_plan::expressions::PhysicalSortExpr;
 
+use async_trait::async_trait;
 use std::any::Any;
+use std::fmt::Debug;
 use std::future::Future;
 use std::iter::Iterator;
 use std::path::PathBuf;
@@ -57,6 +59,19 @@ use datafusion::physical_plan::repartition::BatchPartitioner;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use log::{debug, info};
 
+/// ShuffleWriter represents a section of a query plan that has consistent partitioning and
+/// can be executed as one unit with each partition being executed in parallel. The output of each
+/// partition is re-partitioned and streamed to disk in Arrow IPC format. Future stages of the query
+/// will use the ShuffleReaderExec to read these results.
+#[async_trait]
+pub trait ShuffleWriter: Sync + Send + Debug {
+    async fn execute_shuffle_write(
+        &self,
+        input_partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<Vec<ShuffleWritePartition>>;
+}
+
 /// ShuffleWriterExec represents a section of a query plan that has consistent partitioning and
 /// can be executed as one unit with each partition being executed in parallel. The output of each
 /// partition is re-partitioned and streamed to disk in Arrow IPC format. Future stages of the query
@@ -139,7 +154,7 @@ impl ShuffleWriterExec {
         self.shuffle_output_partitioning.as_ref()
     }
 
-    pub fn execute_shuffle_write(
+    fn execute_shuffle_write_internal(
         &self,
         input_partition: usize,
         context: Arc<TaskContext>,
@@ -288,6 +303,18 @@ impl ShuffleWriterExec {
     }
 }
 
+#[async_trait]
+impl ShuffleWriter for ShuffleWriterExec {
+    async fn execute_shuffle_write(
+        &self,
+        input_partition: usize,
+        context: Arc<TaskContext>,
+    ) -> Result<Vec<ShuffleWritePartition>> {
+        self.execute_shuffle_write_internal(input_partition, context)
+            .await
+    }
+}
+
 impl ExecutionPlan for ShuffleWriterExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -335,7 +362,7 @@ impl ExecutionPlan for ShuffleWriterExec {
 
         let schema_captured = schema.clone();
         let fut_stream = self
-            .execute_shuffle_write(partition, context)
+            .execute_shuffle_write_internal(partition, context)
             .and_then(|part_loc| async move {
                 // build metadata result batch
                 let num_writers = part_loc.len();

ballista/executor/src/executor.rs

@@ -23,13 +23,13 @@ use std::sync::Arc;
 
 use crate::metrics::ExecutorMetricsCollector;
 use ballista_core::error::BallistaError;
+use ballista_core::execution_plans::ShuffleWriter;
 use ballista_core::execution_plans::ShuffleWriterExec;
 use ballista_core::serde::protobuf;
 use ballista_core::serde::protobuf::ExecutorRegistration;
 use datafusion::error::DataFusionError;
 use datafusion::execution::context::TaskContext;
 use datafusion::execution::runtime_env::RuntimeEnv;
-
 use datafusion::physical_plan::udaf::AggregateUDF;
 use datafusion::physical_plan::udf::ScalarUDF;
 use datafusion::physical_plan::{ExecutionPlan, Partitioning};
@@ -98,7 +98,7 @@ impl Executor {
         &self,
         task_id: usize,
         partition: PartitionId,
-        shuffle_writer: Arc<ShuffleWriterExec>,
+        shuffle_writer: Arc<dyn ShuffleWriter>,
         task_ctx: Arc<TaskContext>,
         _shuffle_output_partitioning: Option<Partitioning>,
     ) -> Result<Vec<protobuf::ShuffleWritePartition>, BallistaError> {

ballista/executor/src/metrics/mod.rs

@@ -15,8 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use ballista_core::execution_plans::ShuffleWriterExec;
-use datafusion::physical_plan::display::DisplayableExecutionPlan;
+use ballista_core::execution_plans::ShuffleWriter;
 use log::info;
 use std::sync::Arc;
 
@@ -32,7 +31,7 @@ pub trait ExecutorMetricsCollector: Send + Sync {
         job_id: &str,
         stage_id: usize,
         partition: usize,
-        plan: Arc<ShuffleWriterExec>,
+        plan: Arc<dyn ShuffleWriter>,
     );
 }
 
@@ -47,14 +46,11 @@ impl ExecutorMetricsCollector for LoggingMetricsCollector {
         job_id: &str,
         stage_id: usize,
         partition: usize,
-        plan: Arc<ShuffleWriterExec>,
+        plan: Arc<dyn ShuffleWriter>,
     ) {
         info!(
-            "=== [{}/{}/{}] Physical plan with metrics ===\n{}\n",
-            job_id,
-            stage_id,
-            partition,
-            DisplayableExecutionPlan::with_metrics(plan.as_ref()).indent()
+            "=== [{}/{}/{}] Physical plan with metrics ===\n{:?}\n",
+            job_id, stage_id, partition, plan
         );
     }
 }