Read shuffle writer channel capacity from BallistaConfig instead of protobuf

Author: hcrosse

ballista/core/proto/ballista.proto

@@ -60,7 +60,6 @@ message ShuffleWriterExecNode {
   uint32 stage_id = 2;
   datafusion.PhysicalPlanNode input = 3;
   datafusion.PhysicalHashRepartition output_partitioning = 4;
-  uint32 channel_capacity = 5;
 }
 
 // Sort-based shuffle writer that produces consolidated files with index

ballista/core/src/serde/generated/ballista.rs

@@ -65,8 +65,6 @@ pub struct ShuffleWriterExecNode {
     pub output_partitioning: ::core::option::Option<
         ::datafusion_proto::protobuf::PhysicalHashRepartition,
     >,
-    #[prost(uint32, tag = "5")]
-    pub channel_capacity: u32,
 }
 /// Sort-based shuffle writer that produces consolidated files with index
 #[derive(Clone, PartialEq, ::prost::Message)]

ballista/core/src/serde/mod.rs

@@ -18,6 +18,7 @@
 //! This crate contains code generated from the Ballista Protocol Buffer Definition as well
 //! as convenience code for interacting with the generated code.
 
+use crate::config::BallistaConfig;
 use crate::extension::BallistaCacheNode;
 use crate::{error::BallistaError, serde::scheduler::Action as BallistaAction};
 
@@ -352,12 +353,13 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                     &converter,
                 )?;
 
-                let channel_capacity = shuffle_writer.channel_capacity as usize;
-                let channel_capacity = if channel_capacity == 0 {
-                    DEFAULT_SHUFFLE_CHANNEL_CAPACITY
-                } else {
-                    channel_capacity
-                };
+                let channel_capacity = ctx
+                    .session_config()
+                    .options()
+                    .extensions
+                    .get::<BallistaConfig>()
+                    .map(|c| c.shuffle_writer_channel_capacity())
+                    .unwrap_or(DEFAULT_SHUFFLE_CHANNEL_CAPACITY);
 
                 Ok(Arc::new(
                     ShuffleWriterExec::try_new(
@@ -501,7 +503,6 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                         stage_id: exec.stage_id() as u32,
                         input: None,
                         output_partitioning,
-                        channel_capacity: exec.channel_capacity() as u32,
                     },
                 )),
             };

ballista/scheduler/src/planner.rs

@@ -374,16 +374,13 @@ fn create_shuffle_writer_with_config(
     }
 
     // Fall back to standard shuffle writer
-    Ok(Arc::new(
-        ShuffleWriterExec::try_new(
-            job_id.to_owned(),
-            stage_id,
-            plan,
-            "".to_owned(),
-            partitioning,
-        )?
-        .with_channel_capacity(ballista_config.shuffle_writer_channel_capacity()),
-    ))
+    Ok(Arc::new(ShuffleWriterExec::try_new(
+        job_id.to_owned(),
+        stage_id,
+        plan,
+        "".to_owned(),
+        partitioning,
+    )?))
 }
 
 #[cfg(test)]