[SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking #52271

xi-db · 2025-09-08T13:27:24Z

What changes were proposed in this pull request?

Currently, we enforce gRPC message limits on both the client and the server. These limits are largely meant to protect both sides from potential OOMs by rejecting abnormally large messages. However, there are cases in which the server incorrectly sends oversized messages that exceed these limits and cause execution failures.

Specifically, the large message issue from the server to the client we’re solving here, comes from the Arrow batch data in ExecutePlanResponse being too large. It’s caused by a single arrow row exceeding the 128MB message limit, and Arrow cannot partition further and it has to return the single large row in one gRPC message.

To improve Spark Connect stability, this PR implements chunking large Arrow batches when returning query results from the server to the client, ensuring each ExecutePlanResponse chunk remains within the size limit, and the chunks from a batch will be reassembled on the client when parsing as an arrow batch.

(Scala client changes are being implemented in a follow-up PR.)

To reproduce the existing issue we are solving here, run this code on Spark Connect:

repeat_num_per_mb = 1024 * 1024 // len('Apache Spark ')
res = spark.sql(f"select repeat('Apache Spark ', {repeat_num_per_mb * 300}) as huge_col from range(1)").collect()
print(len(res))

It fails with StatusCode.RESOURCE_EXHAUSTED error with message Received message larger than max (314570608 vs. 134217728), because the server is trying to send an ExecutePlanResponse of ~300MB to the client.

With the improvement introduced by the PR, the above code runs successfully and prints the expected result.

Why are the changes needed?

It improves Spark Connect stability when returning large rows.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

New tests on both the server side and the client side.

Was this patch authored or co-authored using generative AI tooling?

No.

hvanhovell · 2025-09-08T15:11:53Z

sql/connect/common/src/main/protobuf/spark/connect/base.proto

@@ -819,6 +828,19 @@ message ReattachOptions {
  bool reattachable = 1;
 }

+message ResultChunkingOptions {


Should the client be able to set chunk size?

Also, can we name this ResultOptions? I can also see us setting a max arrow batch size here.

Yes, I introduced a client-side option preferred_arrow_chunk_size (this commit).

hvanhovell · 2025-09-08T17:06:06Z

...server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala

-        .build()
-      response.setArrowBatch(batch)
-      responseObserver.onNext(response.build())
+        for (i <- 0 until numChunks) {


Please don't use for comprehensions... Either use an actual loop, or the functional equivalent this translates into.

Yeah, right, updated to an explicit loop with foreach.

hvanhovell · 2025-09-08T17:07:01Z

...server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala

+          val to = math.min(from + maxChunkSize, bytes.length)
+          val length = to - from
+
+          val response = proto.ExecutePlanResponse


You can reuse the builder.

Yes, updated.

hvanhovell · 2025-09-08T17:08:10Z

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala

+        " While spark.connect.grpc.arrow.maxBatchSize determines the max size of a result batch," +
+        " maxChunkSize defines the max size of each individual chunk that is part of the batch" +
+        " that will be sent in a response. This allows the server to send large rows to clients." +
+        " However, excessively large plans remain unsupported due to Spark internals and JVM" +


Remove these two lines. They are not related to the conf.

Sounds good, removed.

hvanhovell · 2025-09-08T17:11:20Z

...ct/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala

+
+        // Execute plan.
+        @volatile var done = false
+        val responses = mutable.Buffer.empty[proto.ExecutePlanResponse]


If done need volatile, then this should be synchronized... Unless you are relying on some happens-before cleverness with the volatile variable (if so, then you need to document this)...

done here doesn't need volatile. I added volatile because other test cases have it. Removed as it is not needed in our case.

hvanhovell · 2025-09-08T17:44:17Z

...ct/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala

+          }
+
+          // Reassemble the chunks into a single Arrow batch and validate its content.
+          val batchData: ByteString =


When you build this for scala, please use a Concatenating InputStream or something like that.

Yeah, it's a good point, updated the test as well.

hvanhovell · 2025-09-08T17:44:50Z

@xi-db are you adding scala support in a follow-up?

xi-db · 2025-09-09T09:18:52Z

@xi-db are you adding scala support in a follow-up?

Yes, I'm implementing the scala support as a follow-up.

Spark Connect ArrowBatch Result Chunking

e3196c3

github-actions bot added SQL PYTHON CONNECT labels Sep 8, 2025

xi-db changed the title ~~[SPARK-53525] Spark Connect ArrowBatch Result Chunking~~ [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking Sep 8, 2025

xi-db added 2 commits September 8, 2025 13:46

Reformat code

a3fdc99

Correct the place to add request option

a81ad62

hvanhovell reviewed Sep 8, 2025

View reviewed changes

xi-db added 3 commits September 9, 2025 09:12

Extract response builder logic

0b3f071

Update config doc

66fb544

Rename the helper method to newExecutePlanResponseBuilder

a96e823

xi-db added 8 commits September 9, 2025 09:48

Remove unnecessary volatile annotation

20f81d3

Use SequenceInputStream in tests to avoid copying bytes

332a2db

Add missing variable definition

21235b7

Use explicit foreach loop instead of for-comprehension

9453896

Merge branch 'master' into arrow-batch-chunking

83c1b50

Linting

d7d1c4b

Merge branch 'master' into arrow-batch-chunking

c3e325f

Introduce client option preferred_arrow_chunk_size

b68461f

xi-db force-pushed the arrow-batch-chunking branch from fb5689a to b68461f Compare September 10, 2025 12:09

xi-db added 2 commits September 10, 2025 13:11

Fix failing tests

801b796

Fix failing tests

43117d2

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking #52271

[SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking #52271

xi-db commented Sep 8, 2025 •

edited

Loading

Uh oh!

hvanhovell Sep 8, 2025

Uh oh!

hvanhovell Sep 8, 2025 •

edited

Loading

Uh oh!

xi-db Sep 10, 2025

Uh oh!

hvanhovell Sep 8, 2025

Uh oh!

xi-db Sep 9, 2025

Uh oh!

hvanhovell Sep 8, 2025

Uh oh!

xi-db Sep 9, 2025

Uh oh!

hvanhovell Sep 8, 2025

Uh oh!

xi-db Sep 9, 2025

Uh oh!

hvanhovell Sep 8, 2025

Uh oh!

xi-db Sep 9, 2025

Uh oh!

hvanhovell Sep 8, 2025

Uh oh!

xi-db Sep 9, 2025

Uh oh!

hvanhovell commented Sep 8, 2025

Uh oh!

xi-db commented Sep 9, 2025

Uh oh!

Uh oh!

[SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking #52271

Are you sure you want to change the base?

[SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking #52271

Conversation

xi-db commented Sep 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hvanhovell Sep 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hvanhovell commented Sep 8, 2025

Uh oh!

xi-db commented Sep 9, 2025

Uh oh!

Uh oh!

xi-db commented Sep 8, 2025 •

edited

Loading

hvanhovell Sep 8, 2025 •

edited

Loading