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,21 @@ 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>>;
+
+    fn collect_plan_metrics(&self) -> Vec<MetricsSet>;
+}
+
 /// 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 +156,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 +305,22 @@ 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
+    }
+
+    fn collect_plan_metrics(&self) -> Vec<MetricsSet> {
+        utils::collect_plan_metrics(self.plan.as_ref())
+    }
+}
+
 impl ExecutionPlan for ShuffleWriterExec {
     fn as_any(&self) -> &dyn Any {
         self
@@ -335,7 +368,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/execution_loop.rs

@@ -29,7 +29,6 @@ use crate::{as_task_status, TaskExecutionTimes};
 use ballista_core::error::BallistaError;
 use ballista_core::serde::scheduler::{ExecutorSpecification, PartitionId};
 use ballista_core::serde::BallistaCodec;
-use ballista_core::utils::collect_plan_metrics;
 use datafusion::execution::context::TaskContext;
 use datafusion_proto::logical_plan::AsLogicalPlan;
 use datafusion_proto::physical_plan::from_proto::parse_protobuf_hash_partitioning;
@@ -209,7 +208,7 @@ async fn run_received_task<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
         plan.schema().as_ref(),
     )?;
 
-    let shuffle_writer_plan =
+    let shuffle_writer =
         executor.new_shuffle_writer(job_id.clone(), stage_id as usize, plan)?;
     dedicated_executor.spawn(async move {
         use std::panic::AssertUnwindSafe;
@@ -222,7 +221,7 @@ async fn run_received_task<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
         let execution_result = match AssertUnwindSafe(executor.execute_shuffle_write(
             task_id as usize,
             part.clone(),
-            shuffle_writer_plan.clone(),
+            shuffle_writer.clone(),
             task_context,
             shuffle_output_partitioning,
         ))
@@ -240,7 +239,7 @@ async fn run_received_task<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
         info!("Done with task {}", task_identity);
         debug!("Statistics: {:?}", execution_result);
 
-        let plan_metrics = collect_plan_metrics(shuffle_writer_plan.as_ref());
+        let plan_metrics = shuffle_writer.collect_plan_metrics();
         let operator_metrics = plan_metrics
             .into_iter()
             .map(|m| m.try_into())

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> {
@@ -129,7 +129,7 @@ impl Executor {
         job_id: String,
         stage_id: usize,
         plan: Arc<dyn ExecutionPlan>,
-    ) -> Result<Arc<ShuffleWriterExec>, BallistaError> {
+    ) -> Result<Arc<dyn ShuffleWriter>, BallistaError> {
         let exec = if let Some(shuffle_writer) =
             plan.as_any().downcast_ref::<ShuffleWriterExec>()
         {

ballista/executor/src/executor_server.rs

@@ -41,9 +41,7 @@ use ballista_core::serde::protobuf::{
 use ballista_core::serde::scheduler::PartitionId;
 use ballista_core::serde::scheduler::TaskDefinition;
 use ballista_core::serde::BallistaCodec;
-use ballista_core::utils::{
-    collect_plan_metrics, create_grpc_client_connection, create_grpc_server,
-};
+use ballista_core::utils::{create_grpc_client_connection, create_grpc_server};
 use dashmap::DashMap;
 use datafusion::execution::context::TaskContext;
 use datafusion::physical_plan::ExecutionPlan;
@@ -341,7 +339,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
         let stage_id = task.stage_id;
         let stage_attempt_num = task.stage_attempt_num;
         let partition_id = task.partition_id;
-        let shuffle_writer_plan =
+        let shuffle_writer =
             self.executor
                 .new_shuffle_writer(job_id.clone(), stage_id, plan)?;
 
@@ -358,15 +356,15 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
             .execute_shuffle_write(
                 task_id,
                 part.clone(),
-                shuffle_writer_plan.clone(),
+                shuffle_writer.clone(),
                 task_context,
                 shuffle_output_partitioning,
             )
             .await;
         info!("Done with task {}", task_identity);
         debug!("Statistics: {:?}", execution_result);
 
-        let plan_metrics = collect_plan_metrics(shuffle_writer_plan.as_ref());
+        let plan_metrics = shuffle_writer.collect_plan_metrics();
         let operator_metrics = plan_metrics
             .into_iter()
             .map(|m| m.try_into())

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
         );
     }
 }