Read channel capacity from BallistaConfig in create_query_stage_exec

Move channel_capacity extraction from serde deserialization to the
executor's create_query_stage_exec, which now has access to SessionConfig
(upstream #1542). This is the clean place for executor-side late binding.

Author: hcrosse

ballista/core/Cargo.toml

@@ -64,7 +64,7 @@ prost = { workspace = true }
 prost-types = { workspace = true }
 rand = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
-tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
+tokio = { workspace = true, features = ["rt-multi-thread"] }
 tokio-stream = { workspace = true, features = ["net"] }
 tonic = { workspace = true }
 tonic-prost = { workspace = true }

ballista/core/src/serde/mod.rs

@@ -18,7 +18,6 @@
 //! 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};
 
@@ -54,7 +53,7 @@ use std::{convert::TryInto, io::Cursor};
 
 use crate::execution_plans::sort_shuffle::SortShuffleConfig;
 use crate::execution_plans::{
-    DEFAULT_SHUFFLE_CHANNEL_CAPACITY, ShuffleReaderExec, ShuffleWriterExec,
+    ShuffleReaderExec, ShuffleWriterExec,
     SortShuffleWriterExec, UnresolvedShuffleExec,
 };
 use crate::serde::protobuf::{
@@ -353,24 +352,13 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                     &converter,
                 )?;
 
-                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(
-                        shuffle_writer.job_id.clone(),
-                        shuffle_writer.stage_id as usize,
-                        input,
-                        "".to_string(), // this is intentional but hacky - the executor will fill this in
-                        shuffle_output_partitioning,
-                    )?
-                    .with_channel_capacity(channel_capacity),
-                ))
+                Ok(Arc::new(ShuffleWriterExec::try_new(
+                    shuffle_writer.job_id.clone(),
+                    shuffle_writer.stage_id as usize,
+                    input,
+                    "".to_string(), // this is intentional but hacky - the executor will fill this in
+                    shuffle_output_partitioning,
+                )?))
             }
             PhysicalPlanType::SortShuffleWriter(sort_shuffle_writer) => {
                 let input = inputs[0].clone();

ballista/executor/src/execution_engine.rs

@@ -22,8 +22,11 @@
 //! for creating query stage executors from physical plans.
 
 use async_trait::async_trait;
+use ballista_core::config::BallistaConfig;
 use ballista_core::execution_plans::sort_shuffle::SortShuffleWriterExec;
-use ballista_core::execution_plans::{ShuffleReaderExec, ShuffleWriterExec};
+use ballista_core::execution_plans::{
+    ShuffleReaderExec, ShuffleWriterExec, DEFAULT_SHUFFLE_CHANNEL_CAPACITY,
+};
 use ballista_core::serde::protobuf::ShuffleWritePartition;
 use ballista_core::utils;
 use datafusion::common::tree_node::{Transformed, TreeNode};
@@ -91,7 +94,7 @@ impl ExecutionEngine for DefaultExecutionEngine {
         stage_id: usize,
         plan: Arc<dyn ExecutionPlan>,
         work_dir: &str,
-        _config: &SessionConfig,
+        config: &SessionConfig,
     ) -> Result<Arc<dyn QueryStageExecutor>> {
         let plan = plan
             .transform(|p| {
@@ -104,6 +107,14 @@ impl ExecutionEngine for DefaultExecutionEngine {
             })?
             .data;
 
+        let channel_capacity = config
+            .options()
+            .extensions
+            .get::<BallistaConfig>()
+            .map(|c| c.shuffle_writer_channel_capacity())
+            .unwrap_or(DEFAULT_SHUFFLE_CHANNEL_CAPACITY)
+            .max(1);
+
         // the query plan created by the scheduler always starts with a shuffle writer
         // (either ShuffleWriterExec or SortShuffleWriterExec)
         if let Some(shuffle_writer) = plan.as_any().downcast_ref::<ShuffleWriterExec>() {
@@ -115,7 +126,7 @@ impl ExecutionEngine for DefaultExecutionEngine {
                 work_dir.to_string(),
                 shuffle_writer.shuffle_output_partitioning().cloned(),
             )?
-            .with_channel_capacity(shuffle_writer.channel_capacity());
+            .with_channel_capacity(channel_capacity);
             Ok(Arc::new(DefaultQueryStageExec::new(
                 ShuffleWriterVariant::Hash(exec),
             )))