feat: make gRPC max message size consistently configurable

ballista/core/src/config.rs

@@ -212,7 +212,7 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
         ConfigEntry::new(BALLISTA_CLIENT_GRPC_MAX_MESSAGE_SIZE.to_string(),
                          "Configuration for max message size in gRPC clients".to_string(),
                          DataType::UInt64,
-                         Some((16 * 1024 * 1024).to_string())),
+                         Some((128 * 1024 * 1024).to_string())),
         ConfigEntry::new(BALLISTA_CLIENT_GRPC_CONNECT_TIMEOUT_SECONDS.to_string(),
                          "Connection timeout for gRPC client in seconds".to_string(),
                          DataType::UInt64,
@@ -934,10 +934,37 @@ mod tests {
     #[test]
     fn default_config() -> Result<()> {
         let config = BallistaConfig::default();
-        assert_eq!(16777216, config.grpc_client_max_message_size());
+        assert_eq!(134217728, config.grpc_client_max_message_size());
         Ok(())
     }
 
+    #[test]
+    fn cluster_config_plumbs_grpc_max_message_size() {
+        // Sanity check that setting `ballista.client.grpc_max_message_size`
+        // via `SessionConfig::options_mut().set(...)` — the path Python
+        // clients' `cluster_config` overrides use — actually reaches the
+        // reader consulted by `distributed_query`.
+        use crate::extension::SessionConfigExt;
+        use datafusion::prelude::SessionConfig;
+
+        let mut config = SessionConfig::new_with_ballista();
+        assert_eq!(
+            config.ballista_config().grpc_client_max_message_size(),
+            128 * 1024 * 1024,
+            "unexpected default (did the bump land?)",
+        );
+
+        let r = config
+            .options_mut()
+            .set("ballista.client.grpc_max_message_size", "67108864");
+        assert!(r.is_ok(), "set failed: {r:?}");
+
+        assert_eq!(
+            config.ballista_config().grpc_client_max_message_size(),
+            64 * 1024 * 1024,
+        );
+    }
+
     // The default must stay non-zero: `0` disables the fit check, and since AQE
     // no longer consults `prefer_hash_join`, a disabled check would leave the
     // Partitioned arm unconditionally on hash join and make SortMergeJoin — the

ballista/core/src/utils.rs

@@ -107,7 +107,7 @@ impl Default for GrpcClientConfig {
             tcp_keepalive_seconds: 3600,
             http2_keepalive_interval_seconds: 300,
             use_tls: false,
-            max_message_size: 16 * 1024 * 1024,
+            max_message_size: 128 * 1024 * 1024,
             io_retries_times: 3,
             io_retry_wait_time_ms: 3000,
             initial_connection_window_size: 67108864,

ballista/executor/src/config.rs

@@ -169,17 +169,17 @@ pub struct Config {
         help = "Tracing log rotation policy."
     )]
     pub log_rotation_policy: ballista_core::config::LogRotationPolicy,
-    /// Maximum size of incoming gRPC messages in bytes (default: 16MB).
+    /// Maximum size of incoming gRPC messages in bytes (default: 128MB).
     #[arg(
         long,
-        default_value_t = 16777216,
+        default_value_t = 134217728,
         help = "The maximum size of a decoded message at the grpc server side."
     )]
     pub grpc_server_max_decoding_message_size: u32,
-    /// Maximum size of outgoing gRPC messages in bytes (default: 16MB).
+    /// Maximum size of outgoing gRPC messages in bytes (default: 128MB).
     #[arg(
         long,
-        default_value_t = 16777216,
+        default_value_t = 134217728,
         help = "The maximum size of an encoded message at the grpc server side."
     )]
     pub grpc_server_max_encoding_message_size: u32,

ballista/executor/src/executor_process.rs

@@ -375,8 +375,8 @@ impl Default for ExecutorProcessConfig {
             log_rotation_policy: Default::default(),
             job_data_ttl_seconds: 604800,
             job_data_clean_up_interval_seconds: 0,
-            grpc_max_decoding_message_size: 16777216,
-            grpc_max_encoding_message_size: 16777216,
+            grpc_max_decoding_message_size: 134217728,
+            grpc_max_encoding_message_size: 134217728,
             grpc_server_config: Default::default(),
             executor_heartbeat_interval_seconds: 60,
             metric_collection_policy: ExecutorMetricCollectionPolicy::default(),

ballista/scheduler/src/config.rs

@@ -148,17 +148,27 @@ pub struct Config {
     /// Maximum size of decoded gRPC messages.
     #[arg(
         long,
-        default_value_t = 16777216,
+        default_value_t = 134217728,
         help = "The maximum size of a decoded message at the grpc server side."
     )]
     pub grpc_server_max_decoding_message_size: u32,
     /// Maximum size of encoded gRPC messages.
     #[arg(
         long,
-        default_value_t = 16777216,
+        default_value_t = 134217728,
         help = "The maximum size of an encoded message at the grpc server side."
     )]
     pub grpc_server_max_encoding_message_size: u32,
+    /// Maximum size of messages sent by the scheduler's outbound gRPC clients
+    /// (e.g. task assignment to executors). Should be at least as large as the
+    /// executor's `--grpc-server-max-decoding-message-size` for those RPCs to
+    /// succeed with big encoded plans.
+    #[arg(
+        long,
+        default_value_t = 134217728,
+        help = "The maximum size of a message sent by the scheduler's outbound gRPC clients (in bytes)."
+    )]
+    pub grpc_client_max_message_size: u32,
     /// Timeout in seconds before marking an executor as dead.
     #[arg(
         long,
@@ -256,6 +266,8 @@ pub struct SchedulerConfig {
     pub grpc_server_max_decoding_message_size: u32,
     /// The maximum size of an encoded message at the grpc server side.
     pub grpc_server_max_encoding_message_size: u32,
+    /// The maximum size of a message sent by the scheduler's outbound gRPC clients.
+    pub grpc_client_max_message_size: u32,
     /// The executor timeout in seconds. It should be longer than executor's heartbeat intervals.
     pub executor_timeout_seconds: u64,
     /// The interval to check expired or dead executors
@@ -305,8 +317,9 @@ impl Default for SchedulerConfig {
             job_resubmit_interval_ms: None,
             executor_termination_grace_period: 0,
             scheduler_event_expected_processing_duration: 0,
-            grpc_server_max_decoding_message_size: 16777216,
-            grpc_server_max_encoding_message_size: 16777216,
+            grpc_server_max_decoding_message_size: 134217728,
+            grpc_server_max_encoding_message_size: 134217728,
+            grpc_client_max_message_size: 134217728,
             executor_timeout_seconds: 180,
             expire_dead_executor_interval_seconds: 15,
             override_config_producer: None,
@@ -426,6 +439,12 @@ impl SchedulerConfig {
         self
     }
 
+    /// Sets the maximum message size for the scheduler's outbound gRPC clients.
+    pub fn with_grpc_client_max_message_size(mut self, value: u32) -> Self {
+        self.grpc_client_max_message_size = value;
+        self
+    }
+
     /// Sets a custom config producer.
     pub fn with_override_config_producer(
         mut self,
@@ -539,6 +558,7 @@ impl TryFrom<Config> for SchedulerConfig {
                 .grpc_server_max_decoding_message_size,
             grpc_server_max_encoding_message_size: opt
                 .grpc_server_max_encoding_message_size,
+            grpc_client_max_message_size: opt.grpc_client_max_message_size,
             executor_timeout_seconds: opt.executor_timeout_seconds,
             expire_dead_executor_interval_seconds: opt
                 .expire_dead_executor_interval_seconds,

ballista/scheduler/src/state/executor_manager.rs

@@ -79,13 +79,22 @@ impl ExecutorManager {
         cluster_state: Arc<dyn ClusterState>,
         config: Arc<SchedulerConfig>,
     ) -> Self {
+        // Prefer an explicit override_config_producer if the embedder wired one,
+        // so a full BallistaConfig (with all its grpc-client knobs) still takes
+        // precedence. Otherwise, use `default()` but override
+        // `max_message_size` from the scheduler's `grpc_client_max_message_size`
+        // CLI flag so users can raise the ceiling for outbound task-assignment
+        // RPCs without having to write a config-producer in Rust.
         let grpc_client_config =
             if let Some(config_producer) = &config.override_config_producer {
                 let session_config = config_producer();
                 let ballista_config = session_config.ballista_config();
                 GrpcClientConfig::from(&ballista_config)
             } else {
-                GrpcClientConfig::default()
+                GrpcClientConfig {
+                    max_message_size: config.grpc_client_max_message_size as usize,
+                    ..GrpcClientConfig::default()
+                }
             };
         Self {
             cluster_state,

docs/source/user-guide/configs.md

@@ -109,7 +109,7 @@ standard DataFusion settings.
 | ballista.cache.noop                                    | Boolean | true                          | Disable default cache node extension                                                                                                                                                                                                                                                                                                                                                                                                               |
 | ballista.client.grpc_connect_timeout_seconds           | UInt64  | 20                            | Connection timeout for gRPC client in seconds                                                                                                                                                                                                                                                                                                                                                                                                      |
 | ballista.client.grpc_http2_keepalive_interval_seconds  | UInt64  | 300                           | HTTP/2 keep-alive interval for gRPC client in seconds                                                                                                                                                                                                                                                                                                                                                                                              |
-| ballista.client.grpc_max_message_size                  | UInt64  | 16777216                      | Configuration for max message size in gRPC clients                                                                                                                                                                                                                                                                                                                                                                                                 |
+| ballista.client.grpc_max_message_size                  | UInt64  | 134217728                     | Configuration for max message size in gRPC clients                                                                                                                                                                                                                                                                                                                                                                                                 |
 | ballista.client.grpc_tcp_keepalive_seconds             | UInt64  | 3600                          | TCP keep-alive interval for gRPC client in seconds                                                                                                                                                                                                                                                                                                                                                                                                 |
 | ballista.client.grpc_timeout_seconds                   | UInt64  | 20                            | Request timeout for gRPC client in seconds                                                                                                                                                                                                                                                                                                                                                                                                         |
 | ballista.client.initial_connection_window_size         | UInt64  | 67108864                      | HTTP/2 initial connection-level flow-control window for gRPC data-plane clients, in bytes. Should be >= the shuffle governor byte budget so the governor, not the transport window, is the binding backpressure. 0 leaves the tonic default.                                                                                                                                                                                                       |