bug fix

Author: andygrove

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

@@ -255,8 +255,6 @@ impl ExecutionPlan for ShuffleWriterExec {
                         let output_batch =
                             RecordBatch::try_new(input_batch.schema(), columns)?;
 
-                        //TODO avoid writing empty batches!
-
                         // write batch out
                         let start = Instant::now();
                         match &mut writers[output_partition_num] {
@@ -283,6 +281,25 @@ impl ExecutionPlan for ShuffleWriterExec {
                     }
                 }
 
+                //TODO big hack to write empty partitions so that shuffle reader doesn't blow up
+                // when attempting to read empty partitions
+                // see https://github.com/apache/arrow-datafusion/issues/711 for more info
+                for output_partition_num in 0..writers.len() {
+                    if writers[output_partition_num].is_none() {
+                        let mut path = path.clone();
+                        path.push(&format!("{}", output_partition_num));
+                        std::fs::create_dir_all(&path)?;
+
+                        path.push("data.arrow");
+                        let path = path.to_str().unwrap();
+                        info!("Writing empty results to {}", path);
+
+                        let writer = ShuffleWriter::new(path, stream.schema().as_ref())?;
+
+                        writers[output_partition_num] = Some(writer);
+                    }
+                }
+
                 // build metadata result batch
                 let num_writers = writers.iter().filter(|w| w.is_some()).count();
                 let mut partition_builder = UInt32Builder::new(num_writers);

ballista/rust/executor/src/execution_loop.rs

@@ -32,6 +32,7 @@ use ballista_core::serde::protobuf::{
 use protobuf::CompletedTask;
 
 use crate::executor::Executor;
+use ballista_core::execution_plans::ShuffleWriterExec;
 
 pub async fn poll_loop(
     mut scheduler: SchedulerGrpcClient<Channel>,
@@ -108,19 +109,35 @@ async fn run_received_tasks(
     available_tasks_slots.fetch_sub(1, Ordering::SeqCst);
     let plan: Arc<dyn ExecutionPlan> = (&task.plan.unwrap()).try_into().unwrap();
 
-    tokio::spawn(async move {
-        let execution_result = executor
-            .execute_shuffle_write(plan, task_id.partition_id as usize)
-            .await;
-        info!("Done with task {}", task_id_log);
-        debug!("Statistics: {:?}", execution_result);
-        available_tasks_slots.fetch_add(1, Ordering::SeqCst);
-        let _ = task_status_sender.send(as_task_status(
-            execution_result.map(|_| ()),
-            executor_id,
-            task_id,
-        ));
-    });
+    if let Some(shuffle_writer) = plan.as_any().downcast_ref::<ShuffleWriterExec>() {
+        //TODO this is a hacky way to set the working directory for the shuffle writer
+        let exec: Arc<dyn ExecutionPlan> = Arc::new(
+            ShuffleWriterExec::try_new(
+                shuffle_writer.job_id().to_string(),
+                shuffle_writer.stage_id(),
+                shuffle_writer.children()[0].clone(),
+                executor.work_dir.clone(),
+                Some(shuffle_writer.output_partitioning()),
+            )
+            .unwrap(),
+        ); //TODO unwrap
+        tokio::spawn(async move {
+            let execution_result = executor
+                .execute_shuffle_write(exec, task_id.partition_id as usize)
+                .await;
+            info!("Done with task {}", task_id_log);
+            debug!("Statistics: {:?}", execution_result);
+            available_tasks_slots.fetch_add(1, Ordering::SeqCst);
+            let _ = task_status_sender.send(as_task_status(
+                execution_result.map(|_| ()),
+                executor_id,
+                task_id,
+            ));
+        });
+    } else {
+        //TODO error handling
+        unreachable!()
+    }
 }
 
 fn as_task_status(

ballista/rust/executor/src/executor.rs

@@ -28,7 +28,7 @@ use datafusion::physical_plan::ExecutionPlan;
 /// Ballista executor
 pub struct Executor {
     /// Directory for storing partial results
-    work_dir: String,
+    pub work_dir: String,
 }
 
 impl Executor {