This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark-connect-rust.git
The following commit(s) were added to refs/heads/master by this push:
new 6088999 [SPARK-59037][CORE] Apply the gRPC max message size (raise
tonic's 4 MiB default)
6088999 is described below
commit 60889992d9155d7c2d3af6291624aae51bd4478a
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Fri Aug 28 16:48:27 2026 +0900
[SPARK-59037][CORE] Apply the gRPC max message size (raise tonic's 4 MiB
default)
### What changes were proposed in this pull request?
`GRPC_MAX_MESSAGE_LENGTH_DEFAULT` (128 MiB) was declared in `channel.rs`
but referenced nowhere, so tonic's **4 MiB** default receive cap stood. This
applies that 128 MiB cap **unconditionally** — matching the reference client —
via `max_decoding_message_size` / `max_encoding_message_size` on:
- the `SparkConnectServiceClient` stub, and
- every raw `tonic::client::Grpc` call (`execute_plan_raw` /
`reattach_execute_raw` / `analyze_plan_raw`).
The reference client (`pyspark/sql/connect/client/core.py`) sets
`grpc.max_{send,receive}_message_length` from `GRPC_MAX_MESSAGE_LENGTH_DEFAULT`
unconditionally via `GRPC_DEFAULT_OPTIONS`; there is no `grpc_max_message_size`
connection parameter. So this does **not** add a connection-string option or
any public API — it just applies the already-declared constant. `channel.rs`
has no net change.
(An earlier revision of this PR added a
`ChannelBuilder::max_message_size()` method parsing a `grpc_max_message_size`
connection param; that was dropped in `13ad5e1` after review, because the
reference client has no such parameter and adding one would both diverge from
it and leak the key into gRPC metadata.)
### Why are the changes needed?
Any `collect()` returning a >4 MiB Arrow batch fails with a decode-size
error that the reference (grpcio) client handles fine, because it raises the
limit. This is the same "declared value never applied" bug-class as the
keepalive fix in SPARK-59032; found during that PR's review by viirya.
### Does this PR introduce _any_ user-facing change?
No API or option change. The client now raises the gRPC send/receive cap to
128 MiB (from tonic's 4 MiB default), matching the reference client, so large
results no longer hit a spurious decode error.
### How was this patch tested?
The cap is a constant applied at the stub and the three raw `Grpc` call
sites, so its correctness rests on the tonic API contract
(`max_decoding_message_size` / `max_encoding_message_size`, whose 4 MiB default
this overrides); the earlier param-parsing unit test was removed together with
the parameter. All 93 `apache-spark-connect-core` unit tests pass.
A new server-gated e2e test,
`e2e_integration::test_collect_batch_larger_than_grpc_default_cap`, collects a
single ~5 MiB string cell — one Arrow batch, hence one gRPC response message,
above tonic's 4 MiB default. It fails with a decode-size error without this fix
and passes with it, and asserts the full value round-trips. It runs in the
existing e2e CI job (`SPARK_REMOTE` against a live Spark 4.2.0 Connect server),
closing the end-to-end gap.
This pull request and its description were written by Isaac.
Closes #64 from HyukjinKwon/fix-grpc-max-message-size.
Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
crates/spark-connect-core/src/client.rs | 22 +++++++++++---
crates/spark-connect/tests/e2e_integration.rs | 44 +++++++++++++++++++++++++++
2 files changed, 61 insertions(+), 5 deletions(-)
diff --git a/crates/spark-connect-core/src/client.rs
b/crates/spark-connect-core/src/client.rs
index 5bbaf9e..4ec73cb 100644
--- a/crates/spark-connect-core/src/client.rs
+++ b/crates/spark-connect-core/src/client.rs
@@ -18,7 +18,7 @@ use spark_connect_proto::{
};
use crate::artifact::{build_artifact_request_stream, FileArtifact,
NamedArtifact};
-use crate::channel::ChannelBuilder;
+use crate::channel::{ChannelBuilder, GRPC_MAX_MESSAGE_LENGTH_DEFAULT};
use crate::error::{Result, SparkError};
use crate::reattach::ExecutePlanResponseReattachableIterator;
use crate::retries::{RetryPolicy, RetryPolicyState};
@@ -103,7 +103,13 @@ impl SparkConnectClient {
.connect_lazy()
};
- let stub = SparkConnectServiceClient::new(channel.clone());
+ // Raise the message-size cap from tonic's 4 MiB default (too small
for real Arrow
+ // `collect()` results) to 128 MiB, applied to both the stub and the
raw Grpc calls
+ // below. The reference client sets this unconditionally via
GRPC_DEFAULT_OPTIONS
+ // (grpc.max_{send,receive}_message_length); it is not a
connection-string option.
+ let stub = SparkConnectServiceClient::new(channel.clone())
+ .max_decoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT)
+ .max_encoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT);
let metadata = builder.metadata().into_iter().collect();
@@ -173,7 +179,9 @@ impl SparkConnectClient {
/// and deep plans don't hit the recursion limit - the server sees exactly
the bytes
/// the reference client built.
pub async fn execute_plan_raw(&self, request: Vec<u8>) ->
Result<Streaming<Vec<u8>>> {
- let mut grpc = tonic::client::Grpc::new(self.channel.clone());
+ let mut grpc = tonic::client::Grpc::new(self.channel.clone())
+ .max_decoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT)
+ .max_encoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT);
wait_ready(&mut grpc).await?;
let mut req = Request::new(request);
self._attach_metadata(&mut req);
@@ -189,7 +197,9 @@ impl SparkConnectClient {
/// ReattachExecute, byte-exact passthrough returning a raw response
stream.
pub async fn reattach_execute_raw(&self, request: Vec<u8>) ->
Result<Streaming<Vec<u8>>> {
- let mut grpc = tonic::client::Grpc::new(self.channel.clone());
+ let mut grpc = tonic::client::Grpc::new(self.channel.clone())
+ .max_decoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT)
+ .max_encoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT);
wait_ready(&mut grpc).await?;
let mut req = Request::new(request);
self._attach_metadata(&mut req);
@@ -239,7 +249,9 @@ impl SparkConnectClient {
/// returns the raw response bytes, avoiding the client-side limit
(matching the
/// reference client, which never re-decodes the request it built).
pub async fn analyze_plan_raw(&self, request: Vec<u8>) -> Result<Vec<u8>> {
- let mut grpc = tonic::client::Grpc::new(self.channel.clone());
+ let mut grpc = tonic::client::Grpc::new(self.channel.clone())
+ .max_decoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT)
+ .max_encoding_message_size(GRPC_MAX_MESSAGE_LENGTH_DEFAULT);
wait_ready(&mut grpc).await?;
let mut req = Request::new(request);
self._attach_metadata(&mut req);
diff --git a/crates/spark-connect/tests/e2e_integration.rs
b/crates/spark-connect/tests/e2e_integration.rs
index f72bfa8..f6a14d5 100644
--- a/crates/spark-connect/tests/e2e_integration.rs
+++ b/crates/spark-connect/tests/e2e_integration.rs
@@ -434,3 +434,47 @@ fn test_to_arrow_full_path() {
let rows: usize = reader.map(|b| b.expect("batch").num_rows()).sum();
assert_eq!(rows, 5, "to_arrow must round-trip all 5 rows of range(5)");
}
+
+/// End-to-end guard for SPARK-59037: a single ~5 MiB string cell yields one
Arrow
+/// batch, and thus one gRPC response message, larger than tonic's 4 MiB
default
+/// decode cap. Before the client raised `max_decoding_message_size` to 128
MiB this
+/// `collect()` failed with a gRPC "message length too large" decode error
(the same
+/// batch the reference grpcio client returns fine). It must now succeed and
return
+/// the whole value intact.
+#[test]
+fn test_collect_batch_larger_than_grpc_default_cap() {
+ if !should_run() {
+ println!(
+ "Skipping test_collect_batch_larger_than_grpc_default_cap - set
SPARK_REMOTE to run"
+ );
+ return;
+ }
+
+ let remote_url =
+ std::env::var("SPARK_REMOTE").unwrap_or_else(|_|
"sc://localhost:15002".to_string());
+ let session = SparkSession::builder()
+ .remote(&remote_url)
+ .get_or_create()
+ .expect("Failed to create session");
+
+ // 5 MiB, comfortably above tonic's 4 MiB default receive cap. A single
row cannot be
+ // split across Arrow batches, so this forces one oversized gRPC response
message.
+ let n: usize = 5 * 1024 * 1024;
+ let df = session
+ .sql(&format!("SELECT repeat('x', {n}) AS s"))
+ .expect("Failed to build repeat() DataFrame");
+ let rows = df
+ .collect()
+ .expect("collect() of a >4 MiB batch must succeed with the raised gRPC
cap");
+
+ assert_eq!(rows.len(), 1, "Expected exactly one row");
+ let s = rows[0]
+ .get(0)
+ .and_then(|v| v.as_str())
+ .expect("Expected a string cell");
+ assert_eq!(
+ s.len(),
+ n,
+ "Expected the full {n}-byte string to round-trip through collect()"
+ );
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
