Set message sizes to 16MB in clients

andygrove · Dec 12, 2023 · 897ca4a · 897ca4a
1 parent 4b3b35f
commit 897ca4a
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 4 deletions.
diff --git a/ballista/client/src/context.rs b/ballista/client/src/context.rs
@@ -100,7 +100,9 @@ impl BallistaContext {
         let connection = create_grpc_client_connection(scheduler_url.clone())
             .await
             .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
-        let mut scheduler = SchedulerGrpcClient::new(connection);
+        let mut scheduler = SchedulerGrpcClient::new(connection)
+            .max_encoding_message_size(16 * 1024 * 1024)
+            .max_decoding_message_size(16 * 1024 * 1024);
 
         let remote_session_id = scheduler
             .create_session(CreateSessionParams {

diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs
@@ -229,7 +229,9 @@ async fn execute_query(
         .await
         .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
 
-    let mut scheduler = SchedulerGrpcClient::new(connection);
+    let mut scheduler = SchedulerGrpcClient::new(connection)
+        .max_encoding_message_size(16 * 1024 * 1024)
+        .max_decoding_message_size(16 * 1024 * 1024);
 
     let query_result = scheduler
         .execute_query(query)

diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs
@@ -284,7 +284,9 @@ pub async fn start_executor_process(opt: Arc<ExecutorProcessConfig>) -> Result<(
         }
     }?;
 
-    let mut scheduler = SchedulerGrpcClient::new(connection);
+    let mut scheduler = SchedulerGrpcClient::new(connection)
+        .max_encoding_message_size(16 * 1024 * 1024)
+        .max_decoding_message_size(16 * 1024 * 1024);
 
     let default_codec: BallistaCodec<LogicalPlanNode, PhysicalPlanNode> =
         BallistaCodec::default();

diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs
@@ -237,7 +237,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
         } else {
             let scheduler_url = format!("http://{scheduler_id}");
             let connection = create_grpc_client_connection(scheduler_url).await?;
-            let scheduler = SchedulerGrpcClient::new(connection);
+            let scheduler = SchedulerGrpcClient::new(connection)
+                .max_encoding_message_size(16 * 1024 * 1024)
+                .max_decoding_message_size(16 * 1024 * 1024);
 
             {
                 self.schedulers