You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by hv...@apache.org on 2025/09/11 15:35:09 UTC
(spark) branch master updated: [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking
This is an automated email from the ASF dual-hosted git repository.
hvanhovell pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new bb41e19ae9ee [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking
bb41e19ae9ee is described below
commit bb41e19ae9ee0f7c228ac395800e740944ee355f
Author: Xi Lyu <xi...@databricks.com>
AuthorDate: Thu Sep 11 11:34:54 2025 -0400
[SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking
### 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.
Closes #52271 from xi-db/arrow-batch-chunking.
Authored-by: Xi Lyu <xi...@databricks.com>
Signed-off-by: Herman van Hovell <he...@databricks.com>
---
python/pyspark/sql/connect/client/core.py | 100 +++++++--
python/pyspark/sql/connect/proto/base_pb2.py | 224 +++++++++++----------
python/pyspark/sql/connect/proto/base_pb2.pyi | 111 +++++++++-
.../sql/tests/connect/test_connect_basic.py | 71 ++++++-
.../src/main/protobuf/spark/connect/base.proto | 27 +++
.../apache/spark/sql/connect/config/Connect.scala | 15 ++
.../execution/SparkConnectPlanExecution.scala | 65 ++++--
.../spark/sql/connect/service/ExecuteHolder.scala | 30 ++-
.../connect/planner/SparkConnectServiceSuite.scala | 148 ++++++++++++++
9 files changed, 644 insertions(+), 147 deletions(-)
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index 9d2e18ebb760..468e496d229c 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -607,6 +607,8 @@ class SparkConnectClient(object):
retry_policy: Optional[Dict[str, Any]] = None,
use_reattachable_execute: bool = True,
session_hooks: Optional[list["SparkSession.Hook"]] = None,
+ allow_arrow_batch_chunking: bool = True,
+ preferred_arrow_chunk_size: Optional[int] = None,
):
"""
Creates a new SparkSession for the Spark Connect interface.
@@ -639,6 +641,21 @@ class SparkConnectClient(object):
Enable reattachable execution.
session_hooks: list[SparkSession.Hook], optional
List of session hooks to call.
+ allow_arrow_batch_chunking: bool
+ Whether to allow the server to split large Arrow batches into smaller chunks.
+ Although Arrow results are split into batches with a size limit according to estimation,
+ the size of the batches is not guaranteed to be less than the limit, especially when a
+ single row is larger than the limit, in which case the server will fail to split it
+ further into smaller batches. As a result, the client may encounter a gRPC error stating
+ "Received message larger than max" when a batch is too large.
+ If true, the server will split large Arrow batches into smaller chunks, and the client
+ is expected to handle the chunked Arrow batches.
+ If false, the server will not chunk large Arrow batches.
+ preferred_arrow_chunk_size: Optional[int]
+ Optional preferred Arrow batch size in bytes for the server to use when sending Arrow
+ results.
+ The server will attempt to use this size if it is set and within the valid range
+ ([1KB, max batch size on server]). Otherwise, the server's maximum batch size is used.
"""
self.thread_local = threading.local()
@@ -678,6 +695,8 @@ class SparkConnectClient(object):
self._user_id, self._session_id, self._channel, self._builder.metadata()
)
self._use_reattachable_execute = use_reattachable_execute
+ self._allow_arrow_batch_chunking = allow_arrow_batch_chunking
+ self._preferred_arrow_chunk_size = preferred_arrow_chunk_size
self._session_hooks = session_hooks or []
# Configure logging for the SparkConnect client.
@@ -1235,6 +1254,15 @@ class SparkConnectClient(object):
req.client_observed_server_side_session_id = self._server_session_id
if self._user_id:
req.user_context.user_id = self._user_id
+ # Add request option to allow result chunking.
+ req.request_options.append(
+ pb2.ExecutePlanRequest.RequestOption(
+ result_chunking_options=pb2.ResultChunkingOptions(
+ allow_arrow_batch_chunking=self._allow_arrow_batch_chunking,
+ preferred_arrow_chunk_size=self._preferred_arrow_chunk_size,
+ )
+ )
+ )
if operation_id is not None:
try:
uuid.UUID(operation_id, version=4)
@@ -1408,6 +1436,7 @@ class SparkConnectClient(object):
req = hook.on_execute_plan(req)
num_records = 0
+ arrow_batch_chunks_to_assemble: List[bytes] = []
def handle_response(
b: pb2.ExecutePlanResponse,
@@ -1495,32 +1524,65 @@ class SparkConnectClient(object):
if b.HasField("arrow_batch"):
logger.debug(
f"Received arrow batch rows={b.arrow_batch.row_count} "
+ f"Number of chunks in batch={b.arrow_batch.num_chunks_in_batch} "
+ f"Chunk index={b.arrow_batch.chunk_index} "
f"size={len(b.arrow_batch.data)}"
)
+ if arrow_batch_chunks_to_assemble:
+ # Expect next chunk of the same batch
+ if b.arrow_batch.chunk_index != len(arrow_batch_chunks_to_assemble):
+ raise SparkConnectException(
+ f"Expected chunk index {len(arrow_batch_chunks_to_assemble)} of the "
+ f"arrow batch but got {b.arrow_batch.chunk_index}."
+ )
+ else:
+ # Expect next batch
+ if (
+ b.arrow_batch.HasField("start_offset")
+ and num_records != b.arrow_batch.start_offset
+ ):
+ # Expect next batch
+ raise SparkConnectException(
+ f"Expected arrow batch to start at row offset {num_records} in "
+ + "results, but received arrow batch starting at offset "
+ + f"{b.arrow_batch.start_offset}."
+ )
+ if b.arrow_batch.chunk_index != 0:
+ raise SparkConnectException(
+ f"Expected chunk index 0 of the next arrow batch "
+ f"but got {b.arrow_batch.chunk_index}."
+ )
+
+ arrow_batch_chunks_to_assemble.append(b.arrow_batch.data)
+ # Assemble the chunks to an arrow batch to process if
+ # (a) chunking is not enabled (num_chunks_in_batch is not set or is 0,
+ # in this case, it is the single chunk in the batch)
+ # (b) or the client has received all chunks of the batch.
if (
- b.arrow_batch.HasField("start_offset")
- and num_records != b.arrow_batch.start_offset
+ not b.arrow_batch.HasField("num_chunks_in_batch")
+ or b.arrow_batch.num_chunks_in_batch == 0
+ or len(arrow_batch_chunks_to_assemble) == b.arrow_batch.num_chunks_in_batch
):
- raise SparkConnectException(
- f"Expected arrow batch to start at row offset {num_records} in results, "
- + "but received arrow batch starting at offset "
- + f"{b.arrow_batch.start_offset}."
+ arrow_batch_data = b"".join(arrow_batch_chunks_to_assemble)
+ arrow_batch_chunks_to_assemble.clear()
+ logger.debug(
+ f"Assembling arrow batch of size {len(arrow_batch_data)} from "
+ f"{b.arrow_batch.num_chunks_in_batch} chunks."
)
- num_records_in_batch = 0
- with pa.ipc.open_stream(b.arrow_batch.data) as reader:
- for batch in reader:
- assert isinstance(batch, pa.RecordBatch)
- num_records_in_batch += batch.num_rows
- yield batch
-
- if num_records_in_batch != b.arrow_batch.row_count:
- raise SparkConnectException(
- f"Expected {b.arrow_batch.row_count} rows in arrow batch but got "
- + f"{num_records_in_batch}."
- )
- num_records += num_records_in_batch
+ num_records_in_batch = 0
+ with pa.ipc.open_stream(arrow_batch_data) as reader:
+ for batch in reader:
+ assert isinstance(batch, pa.RecordBatch)
+ num_records_in_batch += batch.num_rows
+ if num_records_in_batch != b.arrow_batch.row_count:
+ raise SparkConnectException(
+ f"Expected {b.arrow_batch.row_count} rows in arrow batch but "
+ + f"got {num_records_in_batch}."
+ )
+ num_records += num_records_in_batch
+ yield batch
if b.HasField("create_resource_profile_command_result"):
profile_id = b.create_resource_profile_command_result.profile_id
yield {"create_resource_profile_command_result": profile_id}
diff --git a/python/pyspark/sql/connect/proto/base_pb2.py b/python/pyspark/sql/connect/proto/base_pb2.py
index 7d0c681576b9..f690adf40c6b 100644
--- a/python/pyspark/sql/connect/proto/base_pb2.py
+++ b/python/pyspark/sql/connect/proto/base_pb2.py
@@ -45,7 +45,7 @@ from pyspark.sql.connect.proto import pipelines_pb2 as spark_dot_connect_dot_pip
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
- b'\n\x18spark/connect/base.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1cspark/connect/commands.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\x1a\x19spark/connect/types.proto\x1a\x16spark/connect/ml.proto\x1a\x1dspark/connect/pipelines.proto"t\n\x04Plan\x12-\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x04root\x12\x32\n\x07\x63ommand\x18\x02 \x01(\x0b\x32\x16.spark.connect.Com [...]
+ b'\n\x18spark/connect/base.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1cspark/connect/commands.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/relations.proto\x1a\x19spark/connect/types.proto\x1a\x16spark/connect/ml.proto\x1a\x1dspark/connect/pipelines.proto"t\n\x04Plan\x12-\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x04root\x12\x32\n\x07\x63ommand\x18\x02 \x01(\x0b\x32\x16.spark.connect.Com [...]
)
_globals = globals()
@@ -137,123 +137,125 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals["_ANALYZEPLANRESPONSE_JSONTODDL"]._serialized_start = 5011
_globals["_ANALYZEPLANRESPONSE_JSONTODDL"]._serialized_end = 5053
_globals["_EXECUTEPLANREQUEST"]._serialized_start = 5066
- _globals["_EXECUTEPLANREQUEST"]._serialized_end = 5741
+ _globals["_EXECUTEPLANREQUEST"]._serialized_end = 5837
_globals["_EXECUTEPLANREQUEST_REQUESTOPTION"]._serialized_start = 5500
- _globals["_EXECUTEPLANREQUEST_REQUESTOPTION"]._serialized_end = 5665
- _globals["_EXECUTEPLANRESPONSE"]._serialized_start = 5744
- _globals["_EXECUTEPLANRESPONSE"]._serialized_end = 8926
- _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_start = 7700
- _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_end = 7771
- _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_start = 7773
- _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_end = 7891
- _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_start = 7894
- _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_end = 8411
- _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_start = 7989
- _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_end = 8321
+ _globals["_EXECUTEPLANREQUEST_REQUESTOPTION"]._serialized_end = 5761
+ _globals["_EXECUTEPLANRESPONSE"]._serialized_start = 5840
+ _globals["_EXECUTEPLANRESPONSE"]._serialized_end = 9153
+ _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_start = 7796
+ _globals["_EXECUTEPLANRESPONSE_SQLCOMMANDRESULT"]._serialized_end = 7867
+ _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_start = 7870
+ _globals["_EXECUTEPLANRESPONSE_ARROWBATCH"]._serialized_end = 8118
+ _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_start = 8121
+ _globals["_EXECUTEPLANRESPONSE_METRICS"]._serialized_end = 8638
+ _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_start = 8216
+ _globals["_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT"]._serialized_end = 8548
_globals[
"_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY"
- ]._serialized_start = 8198
+ ]._serialized_start = 8425
_globals[
"_EXECUTEPLANRESPONSE_METRICS_METRICOBJECT_EXECUTIONMETRICSENTRY"
- ]._serialized_end = 8321
- _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_start = 8323
- _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_end = 8411
- _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_start = 8414
- _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_end = 8555
- _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_start = 8557
- _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_end = 8573
- _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_start = 8576
- _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_end = 8909
- _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_start = 8732
- _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_end = 8909
- _globals["_KEYVALUE"]._serialized_start = 8928
- _globals["_KEYVALUE"]._serialized_end = 8993
- _globals["_CONFIGREQUEST"]._serialized_start = 8996
- _globals["_CONFIGREQUEST"]._serialized_end = 10195
- _globals["_CONFIGREQUEST_OPERATION"]._serialized_start = 9304
- _globals["_CONFIGREQUEST_OPERATION"]._serialized_end = 9802
- _globals["_CONFIGREQUEST_SET"]._serialized_start = 9804
- _globals["_CONFIGREQUEST_SET"]._serialized_end = 9896
- _globals["_CONFIGREQUEST_GET"]._serialized_start = 9898
- _globals["_CONFIGREQUEST_GET"]._serialized_end = 9923
- _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_start = 9925
- _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_end = 9988
- _globals["_CONFIGREQUEST_GETOPTION"]._serialized_start = 9990
- _globals["_CONFIGREQUEST_GETOPTION"]._serialized_end = 10021
- _globals["_CONFIGREQUEST_GETALL"]._serialized_start = 10023
- _globals["_CONFIGREQUEST_GETALL"]._serialized_end = 10071
- _globals["_CONFIGREQUEST_UNSET"]._serialized_start = 10073
- _globals["_CONFIGREQUEST_UNSET"]._serialized_end = 10100
- _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_start = 10102
- _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_end = 10136
- _globals["_CONFIGRESPONSE"]._serialized_start = 10198
- _globals["_CONFIGRESPONSE"]._serialized_end = 10373
- _globals["_ADDARTIFACTSREQUEST"]._serialized_start = 10376
- _globals["_ADDARTIFACTSREQUEST"]._serialized_end = 11378
- _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_start = 10851
- _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_end = 10904
- _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_start = 10906
- _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_end = 11017
- _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_start = 11019
- _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_end = 11112
- _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_start = 11115
- _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_end = 11308
- _globals["_ADDARTIFACTSRESPONSE"]._serialized_start = 11381
- _globals["_ADDARTIFACTSRESPONSE"]._serialized_end = 11653
- _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_start = 11572
- _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_end = 11653
- _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_start = 11656
- _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_end = 11982
- _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_start = 11985
- _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_end = 12337
- _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_start = 12180
- _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_end = 12295
- _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_start = 12297
- _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_end = 12337
- _globals["_INTERRUPTREQUEST"]._serialized_start = 12340
- _globals["_INTERRUPTREQUEST"]._serialized_end = 12943
- _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_start = 12743
- _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_end = 12871
- _globals["_INTERRUPTRESPONSE"]._serialized_start = 12946
- _globals["_INTERRUPTRESPONSE"]._serialized_end = 13090
- _globals["_REATTACHOPTIONS"]._serialized_start = 13092
- _globals["_REATTACHOPTIONS"]._serialized_end = 13145
- _globals["_REATTACHEXECUTEREQUEST"]._serialized_start = 13148
- _globals["_REATTACHEXECUTEREQUEST"]._serialized_end = 13554
- _globals["_RELEASEEXECUTEREQUEST"]._serialized_start = 13557
- _globals["_RELEASEEXECUTEREQUEST"]._serialized_end = 14142
- _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_start = 14011
- _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_end = 14023
- _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_start = 14025
- _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_end = 14072
- _globals["_RELEASEEXECUTERESPONSE"]._serialized_start = 14145
- _globals["_RELEASEEXECUTERESPONSE"]._serialized_end = 14310
- _globals["_RELEASESESSIONREQUEST"]._serialized_start = 14313
- _globals["_RELEASESESSIONREQUEST"]._serialized_end = 14525
- _globals["_RELEASESESSIONRESPONSE"]._serialized_start = 14527
- _globals["_RELEASESESSIONRESPONSE"]._serialized_end = 14635
- _globals["_FETCHERRORDETAILSREQUEST"]._serialized_start = 14638
- _globals["_FETCHERRORDETAILSREQUEST"]._serialized_end = 14970
- _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_start = 14973
- _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_end = 16528
- _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_start = 15202
- _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_end = 15376
- _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_start = 15379
- _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_end = 15747
- _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_start = 15710
- _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_end = 15747
- _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_start = 15750
- _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_end = 16159
+ ]._serialized_end = 8548
+ _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_start = 8550
+ _globals["_EXECUTEPLANRESPONSE_METRICS_METRICVALUE"]._serialized_end = 8638
+ _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_start = 8641
+ _globals["_EXECUTEPLANRESPONSE_OBSERVEDMETRICS"]._serialized_end = 8782
+ _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_start = 8784
+ _globals["_EXECUTEPLANRESPONSE_RESULTCOMPLETE"]._serialized_end = 8800
+ _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_start = 8803
+ _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS"]._serialized_end = 9136
+ _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_start = 8959
+ _globals["_EXECUTEPLANRESPONSE_EXECUTIONPROGRESS_STAGEINFO"]._serialized_end = 9136
+ _globals["_KEYVALUE"]._serialized_start = 9155
+ _globals["_KEYVALUE"]._serialized_end = 9220
+ _globals["_CONFIGREQUEST"]._serialized_start = 9223
+ _globals["_CONFIGREQUEST"]._serialized_end = 10422
+ _globals["_CONFIGREQUEST_OPERATION"]._serialized_start = 9531
+ _globals["_CONFIGREQUEST_OPERATION"]._serialized_end = 10029
+ _globals["_CONFIGREQUEST_SET"]._serialized_start = 10031
+ _globals["_CONFIGREQUEST_SET"]._serialized_end = 10123
+ _globals["_CONFIGREQUEST_GET"]._serialized_start = 10125
+ _globals["_CONFIGREQUEST_GET"]._serialized_end = 10150
+ _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_start = 10152
+ _globals["_CONFIGREQUEST_GETWITHDEFAULT"]._serialized_end = 10215
+ _globals["_CONFIGREQUEST_GETOPTION"]._serialized_start = 10217
+ _globals["_CONFIGREQUEST_GETOPTION"]._serialized_end = 10248
+ _globals["_CONFIGREQUEST_GETALL"]._serialized_start = 10250
+ _globals["_CONFIGREQUEST_GETALL"]._serialized_end = 10298
+ _globals["_CONFIGREQUEST_UNSET"]._serialized_start = 10300
+ _globals["_CONFIGREQUEST_UNSET"]._serialized_end = 10327
+ _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_start = 10329
+ _globals["_CONFIGREQUEST_ISMODIFIABLE"]._serialized_end = 10363
+ _globals["_CONFIGRESPONSE"]._serialized_start = 10425
+ _globals["_CONFIGRESPONSE"]._serialized_end = 10600
+ _globals["_ADDARTIFACTSREQUEST"]._serialized_start = 10603
+ _globals["_ADDARTIFACTSREQUEST"]._serialized_end = 11605
+ _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_start = 11078
+ _globals["_ADDARTIFACTSREQUEST_ARTIFACTCHUNK"]._serialized_end = 11131
+ _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_start = 11133
+ _globals["_ADDARTIFACTSREQUEST_SINGLECHUNKARTIFACT"]._serialized_end = 11244
+ _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_start = 11246
+ _globals["_ADDARTIFACTSREQUEST_BATCH"]._serialized_end = 11339
+ _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_start = 11342
+ _globals["_ADDARTIFACTSREQUEST_BEGINCHUNKEDARTIFACT"]._serialized_end = 11535
+ _globals["_ADDARTIFACTSRESPONSE"]._serialized_start = 11608
+ _globals["_ADDARTIFACTSRESPONSE"]._serialized_end = 11880
+ _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_start = 11799
+ _globals["_ADDARTIFACTSRESPONSE_ARTIFACTSUMMARY"]._serialized_end = 11880
+ _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_start = 11883
+ _globals["_ARTIFACTSTATUSESREQUEST"]._serialized_end = 12209
+ _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_start = 12212
+ _globals["_ARTIFACTSTATUSESRESPONSE"]._serialized_end = 12564
+ _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_start = 12407
+ _globals["_ARTIFACTSTATUSESRESPONSE_STATUSESENTRY"]._serialized_end = 12522
+ _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_start = 12524
+ _globals["_ARTIFACTSTATUSESRESPONSE_ARTIFACTSTATUS"]._serialized_end = 12564
+ _globals["_INTERRUPTREQUEST"]._serialized_start = 12567
+ _globals["_INTERRUPTREQUEST"]._serialized_end = 13170
+ _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_start = 12970
+ _globals["_INTERRUPTREQUEST_INTERRUPTTYPE"]._serialized_end = 13098
+ _globals["_INTERRUPTRESPONSE"]._serialized_start = 13173
+ _globals["_INTERRUPTRESPONSE"]._serialized_end = 13317
+ _globals["_REATTACHOPTIONS"]._serialized_start = 13319
+ _globals["_REATTACHOPTIONS"]._serialized_end = 13372
+ _globals["_RESULTCHUNKINGOPTIONS"]._serialized_start = 13375
+ _globals["_RESULTCHUNKINGOPTIONS"]._serialized_end = 13556
+ _globals["_REATTACHEXECUTEREQUEST"]._serialized_start = 13559
+ _globals["_REATTACHEXECUTEREQUEST"]._serialized_end = 13965
+ _globals["_RELEASEEXECUTEREQUEST"]._serialized_start = 13968
+ _globals["_RELEASEEXECUTEREQUEST"]._serialized_end = 14553
+ _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_start = 14422
+ _globals["_RELEASEEXECUTEREQUEST_RELEASEALL"]._serialized_end = 14434
+ _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_start = 14436
+ _globals["_RELEASEEXECUTEREQUEST_RELEASEUNTIL"]._serialized_end = 14483
+ _globals["_RELEASEEXECUTERESPONSE"]._serialized_start = 14556
+ _globals["_RELEASEEXECUTERESPONSE"]._serialized_end = 14721
+ _globals["_RELEASESESSIONREQUEST"]._serialized_start = 14724
+ _globals["_RELEASESESSIONREQUEST"]._serialized_end = 14936
+ _globals["_RELEASESESSIONRESPONSE"]._serialized_start = 14938
+ _globals["_RELEASESESSIONRESPONSE"]._serialized_end = 15046
+ _globals["_FETCHERRORDETAILSREQUEST"]._serialized_start = 15049
+ _globals["_FETCHERRORDETAILSREQUEST"]._serialized_end = 15381
+ _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_start = 15384
+ _globals["_FETCHERRORDETAILSRESPONSE"]._serialized_end = 16939
+ _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_start = 15613
+ _globals["_FETCHERRORDETAILSRESPONSE_STACKTRACEELEMENT"]._serialized_end = 15787
+ _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_start = 15790
+ _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT"]._serialized_end = 16158
+ _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_start = 16121
+ _globals["_FETCHERRORDETAILSRESPONSE_QUERYCONTEXT_CONTEXTTYPE"]._serialized_end = 16158
+ _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_start = 16161
+ _globals["_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE"]._serialized_end = 16570
_globals[
"_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE_MESSAGEPARAMETERSENTRY"
- ]._serialized_start = 16061
+ ]._serialized_start = 16472
_globals[
"_FETCHERRORDETAILSRESPONSE_SPARKTHROWABLE_MESSAGEPARAMETERSENTRY"
- ]._serialized_end = 16129
- _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_start = 16162
- _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_end = 16509
- _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_start = 16530
- _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_end = 16620
- _globals["_SPARKCONNECTSERVICE"]._serialized_start = 16623
- _globals["_SPARKCONNECTSERVICE"]._serialized_end = 17569
+ ]._serialized_end = 16540
+ _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_start = 16573
+ _globals["_FETCHERRORDETAILSRESPONSE_ERROR"]._serialized_end = 16920
+ _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_start = 16941
+ _globals["_CHECKPOINTCOMMANDRESULT"]._serialized_end = 17031
+ _globals["_SPARKCONNECTSERVICE"]._serialized_start = 17034
+ _globals["_SPARKCONNECTSERVICE"]._serialized_end = 17980
# @@protoc_insertion_point(module_scope)
diff --git a/python/pyspark/sql/connect/proto/base_pb2.pyi b/python/pyspark/sql/connect/proto/base_pb2.pyi
index a4bb0d2fcfd7..93dff6639665 100644
--- a/python/pyspark/sql/connect/proto/base_pb2.pyi
+++ b/python/pyspark/sql/connect/proto/base_pb2.pyi
@@ -1093,16 +1093,20 @@ class ExecutePlanRequest(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
REATTACH_OPTIONS_FIELD_NUMBER: builtins.int
+ RESULT_CHUNKING_OPTIONS_FIELD_NUMBER: builtins.int
EXTENSION_FIELD_NUMBER: builtins.int
@property
def reattach_options(self) -> global___ReattachOptions: ...
@property
+ def result_chunking_options(self) -> global___ResultChunkingOptions: ...
+ @property
def extension(self) -> google.protobuf.any_pb2.Any:
"""Extension type for request options"""
def __init__(
self,
*,
reattach_options: global___ReattachOptions | None = ...,
+ result_chunking_options: global___ResultChunkingOptions | None = ...,
extension: google.protobuf.any_pb2.Any | None = ...,
) -> None: ...
def HasField(
@@ -1114,6 +1118,8 @@ class ExecutePlanRequest(google.protobuf.message.Message):
b"reattach_options",
"request_option",
b"request_option",
+ "result_chunking_options",
+ b"result_chunking_options",
],
) -> builtins.bool: ...
def ClearField(
@@ -1125,11 +1131,16 @@ class ExecutePlanRequest(google.protobuf.message.Message):
b"reattach_options",
"request_option",
b"request_option",
+ "result_chunking_options",
+ b"result_chunking_options",
],
) -> None: ...
def WhichOneof(
self, oneof_group: typing_extensions.Literal["request_option", b"request_option"]
- ) -> typing_extensions.Literal["reattach_options", "extension"] | None: ...
+ ) -> (
+ typing_extensions.Literal["reattach_options", "result_chunking_options", "extension"]
+ | None
+ ): ...
SESSION_ID_FIELD_NUMBER: builtins.int
CLIENT_OBSERVED_SERVER_SIDE_SESSION_ID_FIELD_NUMBER: builtins.int
@@ -1308,38 +1319,78 @@ class ExecutePlanResponse(google.protobuf.message.Message):
ROW_COUNT_FIELD_NUMBER: builtins.int
DATA_FIELD_NUMBER: builtins.int
START_OFFSET_FIELD_NUMBER: builtins.int
+ CHUNK_INDEX_FIELD_NUMBER: builtins.int
+ NUM_CHUNKS_IN_BATCH_FIELD_NUMBER: builtins.int
row_count: builtins.int
"""Count rows in `data`. Must match the number of rows inside `data`."""
data: builtins.bytes
"""Serialized Arrow data."""
start_offset: builtins.int
"""If set, row offset of the start of this ArrowBatch in execution results."""
+ chunk_index: builtins.int
+ """Index of this chunk in the batch if chunking is enabled. The index starts from 0."""
+ num_chunks_in_batch: builtins.int
+ """Total number of chunks in this batch if chunking is enabled.
+ It is missing when chunking is disabled - the batch is returned whole
+ and client will treat this response as the batch.
+ """
def __init__(
self,
*,
row_count: builtins.int = ...,
data: builtins.bytes = ...,
start_offset: builtins.int | None = ...,
+ chunk_index: builtins.int | None = ...,
+ num_chunks_in_batch: builtins.int | None = ...,
) -> None: ...
def HasField(
self,
field_name: typing_extensions.Literal[
- "_start_offset", b"_start_offset", "start_offset", b"start_offset"
+ "_chunk_index",
+ b"_chunk_index",
+ "_num_chunks_in_batch",
+ b"_num_chunks_in_batch",
+ "_start_offset",
+ b"_start_offset",
+ "chunk_index",
+ b"chunk_index",
+ "num_chunks_in_batch",
+ b"num_chunks_in_batch",
+ "start_offset",
+ b"start_offset",
],
) -> builtins.bool: ...
def ClearField(
self,
field_name: typing_extensions.Literal[
+ "_chunk_index",
+ b"_chunk_index",
+ "_num_chunks_in_batch",
+ b"_num_chunks_in_batch",
"_start_offset",
b"_start_offset",
+ "chunk_index",
+ b"chunk_index",
"data",
b"data",
+ "num_chunks_in_batch",
+ b"num_chunks_in_batch",
"row_count",
b"row_count",
"start_offset",
b"start_offset",
],
) -> None: ...
+ @typing.overload
+ def WhichOneof(
+ self, oneof_group: typing_extensions.Literal["_chunk_index", b"_chunk_index"]
+ ) -> typing_extensions.Literal["chunk_index"] | None: ...
+ @typing.overload
+ def WhichOneof(
+ self,
+ oneof_group: typing_extensions.Literal["_num_chunks_in_batch", b"_num_chunks_in_batch"],
+ ) -> typing_extensions.Literal["num_chunks_in_batch"] | None: ...
+ @typing.overload
def WhichOneof(
self, oneof_group: typing_extensions.Literal["_start_offset", b"_start_offset"]
) -> typing_extensions.Literal["start_offset"] | None: ...
@@ -2942,6 +2993,62 @@ class ReattachOptions(google.protobuf.message.Message):
global___ReattachOptions = ReattachOptions
+class ResultChunkingOptions(google.protobuf.message.Message):
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+ ALLOW_ARROW_BATCH_CHUNKING_FIELD_NUMBER: builtins.int
+ PREFERRED_ARROW_CHUNK_SIZE_FIELD_NUMBER: builtins.int
+ allow_arrow_batch_chunking: builtins.bool
+ """Although Arrow results are split into batches with a size limit according to estimation, the
+ size of the batches is not guaranteed to be less than the limit, especially when a single row
+ is larger than the limit, in which case the server will fail to split it further into smaller
+ batches. As a result, the client may encounter a gRPC error stating “Received message larger
+ than max” when a batch is too large.
+ If allow_arrow_batch_chunking=true, the server will split large Arrow batches into smaller chunks,
+ and the client is expected to handle the chunked Arrow batches.
+
+ If false, the server will not chunk large Arrow batches.
+ """
+ preferred_arrow_chunk_size: builtins.int
+ """Optional preferred Arrow batch size in bytes for the server to use when sending Arrow results.
+ The server will attempt to use this size if it is set and within the valid range
+ ([1KB, max batch size on server]). Otherwise, the server's maximum batch size is used.
+ """
+ def __init__(
+ self,
+ *,
+ allow_arrow_batch_chunking: builtins.bool = ...,
+ preferred_arrow_chunk_size: builtins.int | None = ...,
+ ) -> None: ...
+ def HasField(
+ self,
+ field_name: typing_extensions.Literal[
+ "_preferred_arrow_chunk_size",
+ b"_preferred_arrow_chunk_size",
+ "preferred_arrow_chunk_size",
+ b"preferred_arrow_chunk_size",
+ ],
+ ) -> builtins.bool: ...
+ def ClearField(
+ self,
+ field_name: typing_extensions.Literal[
+ "_preferred_arrow_chunk_size",
+ b"_preferred_arrow_chunk_size",
+ "allow_arrow_batch_chunking",
+ b"allow_arrow_batch_chunking",
+ "preferred_arrow_chunk_size",
+ b"preferred_arrow_chunk_size",
+ ],
+ ) -> None: ...
+ def WhichOneof(
+ self,
+ oneof_group: typing_extensions.Literal[
+ "_preferred_arrow_chunk_size", b"_preferred_arrow_chunk_size"
+ ],
+ ) -> typing_extensions.Literal["preferred_arrow_chunk_size"] | None: ...
+
+global___ResultChunkingOptions = ResultChunkingOptions
+
class ReattachExecuteRequest(google.protobuf.message.Message):
DESCRIPTOR: google.protobuf.descriptor.Descriptor
diff --git a/python/pyspark/sql/tests/connect/test_connect_basic.py b/python/pyspark/sql/tests/connect/test_connect_basic.py
index 2aa383f39937..473bd328d5ad 100755
--- a/python/pyspark/sql/tests/connect/test_connect_basic.py
+++ b/python/pyspark/sql/tests/connect/test_connect_basic.py
@@ -46,7 +46,7 @@ from pyspark.testing.pandasutils import PandasOnSparkTestUtils
if should_test_connect:
- from pyspark.sql.connect.proto import Expression as ProtoExpression
+ from pyspark.sql.connect.proto import ExecutePlanResponse, Expression as ProtoExpression
from pyspark.sql.connect.column import Column
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.connect.dataframe import DataFrame as CDataFrame
@@ -1508,6 +1508,75 @@ class SparkConnectGCTests(SparkConnectSQLTestCase):
eventually(catch_assertions=True)(condition)()
+ def test_arrow_batch_result_chunking(self):
+ # Two cases are tested here:
+ # (a) client preferred chunk size is set: the server should respect it
+ # (b) client preferred chunk size is not set: the server should use its own max chunk size
+ for preferred_chunk_size_optional, max_chunk_size_optional in ((1024, None), (None, 1024)):
+ sql_query = "select id, CAST(id + 0.5 AS DOUBLE) from range(0, 2000, 1, 4)"
+ cdf = self.connect.sql(sql_query)
+ sdf = self.spark.sql(sql_query)
+
+ original_verify_response_integrity = self.connect._client._verify_response_integrity
+ captured_chunks = []
+
+ def patched_verify_response_integrity(response):
+ original_verify_response_integrity(response)
+ if isinstance(response, ExecutePlanResponse) and response.HasField("arrow_batch"):
+ captured_chunks.append(response.arrow_batch)
+
+ try:
+ # Patch the response verifier for testing to access the chunked arrow batch
+ # responses.
+ self.connect._client._verify_response_integrity = patched_verify_response_integrity
+ # Override the chunk size to 1024 bytes for testing
+ if preferred_chunk_size_optional:
+ self.connect._client._preferred_arrow_chunk_size = preferred_chunk_size_optional
+ if max_chunk_size_optional:
+ self.connect.conf.set(
+ "spark.connect.session.resultChunking.maxChunkSize",
+ max_chunk_size_optional,
+ )
+
+ # Execute the query, and assert the results are correct.
+ self.assertEqual(cdf.collect(), sdf.collect())
+
+ # Verify the metadata of arrow batch chunks.
+ def split_into_batches(chunks):
+ batches = []
+ i = 0
+ n = len(chunks)
+ while i < n:
+ num_chunks = chunks[i].num_chunks_in_batch
+ batch = chunks[i : i + num_chunks]
+ batches.append(batch)
+ i += num_chunks
+ return batches
+
+ batches = split_into_batches(captured_chunks)
+ # There are 4 batches (partitions) in total.
+ self.assertEqual(len(batches), 4)
+ for batch in batches:
+ # In this example, the max chunk size is set to a small value, so each Arrow
+ # batch should be split into multiple chunks.
+ self.assertTrue(len(batch) > 5)
+ row_count = batch[0].row_count
+ row_start_offset = batch[0].start_offset
+ for i, chunk in enumerate(batch):
+ self.assertEqual(chunk.chunk_index, i)
+ self.assertEqual(chunk.num_chunks_in_batch, len(batch))
+ self.assertEqual(chunk.row_count, row_count)
+ self.assertEqual(chunk.start_offset, row_start_offset)
+ self.assertTrue(len(chunk.data) > 0)
+ self.assertTrue(
+ len(chunk.data)
+ <= (preferred_chunk_size_optional or max_chunk_size_optional)
+ )
+ finally:
+ self.connect._client._verify_response_integrity = original_verify_response_integrity
+ self.connect._client._preferred_arrow_chunk_size = None
+ self.connect.conf.unset("spark.connect.session.resultChunking.maxChunkSize")
+
if __name__ == "__main__":
from pyspark.sql.tests.connect.test_connect_basic import * # noqa: F401
diff --git a/sql/connect/common/src/main/protobuf/spark/connect/base.proto b/sql/connect/common/src/main/protobuf/spark/connect/base.proto
index 7f317defd47b..89d22f7455db 100644
--- a/sql/connect/common/src/main/protobuf/spark/connect/base.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/base.proto
@@ -333,6 +333,7 @@ message ExecutePlanRequest {
message RequestOption {
oneof request_option {
ReattachOptions reattach_options = 1;
+ ResultChunkingOptions result_chunking_options = 2;
// Extension type for request options
google.protobuf.Any extension = 999;
}
@@ -435,6 +436,14 @@ message ExecutePlanResponse {
// If set, row offset of the start of this ArrowBatch in execution results.
optional int64 start_offset = 3;
+
+ // Index of this chunk in the batch if chunking is enabled. The index starts from 0.
+ optional int64 chunk_index = 4;
+
+ // Total number of chunks in this batch if chunking is enabled.
+ // It is missing when chunking is disabled - the batch is returned whole
+ // and client will treat this response as the batch.
+ optional int64 num_chunks_in_batch = 5;
}
message Metrics {
@@ -819,6 +828,24 @@ message ReattachOptions {
bool reattachable = 1;
}
+message ResultChunkingOptions {
+ // Although Arrow results are split into batches with a size limit according to estimation, the
+ // size of the batches is not guaranteed to be less than the limit, especially when a single row
+ // is larger than the limit, in which case the server will fail to split it further into smaller
+ // batches. As a result, the client may encounter a gRPC error stating “Received message larger
+ // than max” when a batch is too large.
+ // If allow_arrow_batch_chunking=true, the server will split large Arrow batches into smaller chunks,
+ // and the client is expected to handle the chunked Arrow batches.
+ //
+ // If false, the server will not chunk large Arrow batches.
+ bool allow_arrow_batch_chunking = 1;
+
+ // Optional preferred Arrow batch size in bytes for the server to use when sending Arrow results.
+ // The server will attempt to use this size if it is set and within the valid range
+ // ([1KB, max batch size on server]). Otherwise, the server's maximum batch size is used.
+ optional int64 preferred_arrow_chunk_size = 2;
+}
+
message ReattachExecuteRequest {
// (Required)
//
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
index 1887e4ede04d..128823f2bb9c 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
@@ -392,4 +392,19 @@ object Connect {
.internal()
.bytesConf(ByteUnit.BYTE)
.createWithDefaultString("10g")
+
+ val CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE =
+ buildConf("spark.connect.session.resultChunking.maxChunkSize")
+ .doc("The max size of a chunk in responses for a result batch. Result chunking is enabled" +
+ " if this config is set to a value greater than 0 and if the client allows it in" +
+ " ResultChunkingOptions. Otherwise, for example if set to -1, this feature is disabled." +
+ " 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." +
+ " The size is in bytes.")
+ .version("4.1.0")
+ .internal()
+ .bytesConf(ByteUnit.BYTE)
+ // 90% of the max message size by default to allow for some overhead.
+ .createWithDefault((ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE * 0.9).toInt)
}
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala
index 0e89681440ea..f5cb2696d849 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala
@@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.classic.{DataFrame, Dataset}
import org.apache.spark.sql.connect.common.DataTypeProtoConverter
import org.apache.spark.sql.connect.common.LiteralValueProtoConverter.toLiteralProto
-import org.apache.spark.sql.connect.config.Connect.CONNECT_GRPC_ARROW_MAX_BATCH_SIZE
+import org.apache.spark.sql.connect.config.Connect.{CONNECT_GRPC_ARROW_MAX_BATCH_SIZE, CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE}
import org.apache.spark.sql.connect.planner.{InvalidInputErrors, SparkConnectPlanner}
import org.apache.spark.sql.connect.service.ExecuteHolder
import org.apache.spark.sql.connect.utils.MetricGenerator
@@ -131,6 +131,8 @@ private[execution] class SparkConnectPlanExecution(executeHolder: ExecuteHolder)
val largeVarTypes = spark.sessionState.conf.arrowUseLargeVarTypes
// Conservatively sets it 70% because the size is not accurate but estimated.
val maxBatchSize = (SparkEnv.get.conf.get(CONNECT_GRPC_ARROW_MAX_BATCH_SIZE) * 0.7).toLong
+ // Whether to enable arrow batch chunking for large result batches.
+ val isResultChunkingEnabled = executePlan.resultChunkingEnabled
val converter = rowToArrowConverter(
schema,
@@ -142,19 +144,56 @@ private[execution] class SparkConnectPlanExecution(executeHolder: ExecuteHolder)
var numSent = 0
def sendBatch(bytes: Array[Byte], count: Long, startOffset: Long): Unit = {
- val response = proto.ExecutePlanResponse
- .newBuilder()
- .setSessionId(sessionId)
- .setServerSideSessionId(sessionHolder.serverSessionId)
+ def newExecutePlanResponseBuilder = {
+ proto.ExecutePlanResponse
+ .newBuilder()
+ .setSessionId(sessionId)
+ .setServerSideSessionId(sessionHolder.serverSessionId)
+ }
+ if (isResultChunkingEnabled) {
+ // Result chunking is enabled. Split the result batch into multiple chunks if needed.
+ // The server will attempt to use the preferred chunk size if it is set and within the
+ // valid range ([1KB, max batch size on server]). Otherwise, the server's max batch size
+ // is used.
+ val maxChunkSize = executePlan.preferredArrowChunkSize match {
+ case Some(preferred)
+ if preferred >= 1024 &&
+ preferred <= spark.conf.get(CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE) =>
+ preferred
+ case _ =>
+ spark.conf.get(CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE).toInt
+ }
+ val numChunks = (bytes.length + maxChunkSize - 1) / maxChunkSize
- val batch = proto.ExecutePlanResponse.ArrowBatch
- .newBuilder()
- .setRowCount(count)
- .setData(ByteString.copyFrom(bytes))
- .setStartOffset(startOffset)
- .build()
- response.setArrowBatch(batch)
- responseObserver.onNext(response.build())
+ (0 until numChunks).foreach { i =>
+ val from = i * maxChunkSize
+ val to = math.min(from + maxChunkSize, bytes.length)
+ val length = to - from
+
+ val response = newExecutePlanResponseBuilder
+ val batch = proto.ExecutePlanResponse.ArrowBatch
+ .newBuilder()
+ .setRowCount(count) // same for all chunks - total rows
+ .setData(ByteString.copyFrom(bytes, from, length))
+ .setStartOffset(startOffset)
+ .setChunkIndex(i)
+ .setNumChunksInBatch(numChunks) // same for all chunks - num chunks
+ .build()
+ response.setArrowBatch(batch)
+ responseObserver.onNext(response.build())
+ }
+ } else {
+ // Result chunking is not enabled. Return the entire batch without chunking.
+ val response = newExecutePlanResponseBuilder
+ val batch = proto.ExecutePlanResponse.ArrowBatch
+ .newBuilder()
+ .setRowCount(count)
+ .setData(ByteString.copyFrom(bytes))
+ .setStartOffset(startOffset)
+ .build()
+ response.setArrowBatch(batch)
+ responseObserver.onNext(response.build())
+ }
numSent += 1
}
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala
index 9d8603d95c65..42574b1f8d43 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala
@@ -29,7 +29,7 @@ import org.apache.spark.connect.proto
import org.apache.spark.internal.Logging
import org.apache.spark.sql.Observation
import org.apache.spark.sql.connect.common.ProtoUtils
-import org.apache.spark.sql.connect.config.Connect.CONNECT_EXECUTE_REATTACHABLE_ENABLED
+import org.apache.spark.sql.connect.config.Connect.{CONNECT_EXECUTE_REATTACHABLE_ENABLED, CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE}
import org.apache.spark.sql.connect.execution.{ExecuteGrpcResponseSender, ExecuteResponseObserver, ExecuteThreadRunner}
import org.apache.spark.util.SystemClock
@@ -76,6 +76,34 @@ private[connect] class ExecuteHolder(
}
}
+ /**
+ * If result chunking is enabled, it will split large arrow batches into smaller chunks in
+ * responses.
+ */
+ lazy val resultChunkingEnabled: Boolean = {
+ sessionHolder.session.conf.get(CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE) > 0 &&
+ request.getRequestOptionsList.asScala.exists { option =>
+ option.hasResultChunkingOptions &&
+ option.getResultChunkingOptions.getAllowArrowBatchChunking == true
+ }
+ }
+
+ /**
+ * The optional preferred arrow chunk size set by the client.
+ */
+ lazy val preferredArrowChunkSize: Option[Int] = {
+ if (resultChunkingEnabled) {
+ request.getRequestOptionsList.asScala.iterator.collectFirst {
+ case option
+ if option.hasResultChunkingOptions &&
+ option.getResultChunkingOptions.hasPreferredArrowChunkSize =>
+ option.getResultChunkingOptions.getPreferredArrowChunkSize.toInt
+ }
+ } else {
+ None
+ }
+ }
+
val responseObserver: ExecuteResponseObserver[proto.ExecutePlanResponse] =
new ExecuteResponseObserver[proto.ExecutePlanResponse](this)
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala
index b02507f79648..2989471d36a0 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala
@@ -16,6 +16,7 @@
*/
package org.apache.spark.sql.connect.planner
+import java.io.SequenceInputStream
import java.util.UUID
import java.util.concurrent.Semaphore
@@ -376,6 +377,153 @@ class SparkConnectServiceSuite
}
}
+ Seq((Some(1024), None), (None, Some(1024))).foreach {
+ // Two cases are tested here:
+ // (a) client preferred chunk size is set: the server should respect it
+ // (b) client preferred chunk size is not set: the server should use its own max chunk size
+ case (preferredChunkSizeOption, maxChunkSizeOption) =>
+ test(
+ s"Arrow batch chunking with preferredMaxChunkSize=$preferredChunkSizeOption, " +
+ s"maxChunkSize=$maxChunkSizeOption") {
+ withEvents { verifyEvents =>
+ val maxChunkSizeConfigValue = maxChunkSizeOption match {
+ case Some(maxChunkSize) => maxChunkSize.toString
+ case None => Connect.CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE.defaultValueString
+ }
+ withSparkConf(
+ Connect.CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE.key ->
+ maxChunkSizeConfigValue) {
+ val instance = new SparkConnectService(false)
+ val connect = new MockRemoteSession()
+ val context = proto.UserContext
+ .newBuilder()
+ .setUserId("c1")
+ .build()
+ val plan = proto.Plan
+ .newBuilder()
+ .setRoot(
+ connect.sql("select id, CAST(id + 0.5 AS DOUBLE) from range(0, 2000, 1, 4)"))
+ .build()
+ val request = proto.ExecutePlanRequest
+ .newBuilder()
+ .setPlan(plan)
+ .addRequestOptions(
+ proto.ExecutePlanRequest.RequestOption
+ .newBuilder()
+ .setResultChunkingOptions({
+ val builder = proto.ResultChunkingOptions
+ .newBuilder()
+ .setAllowArrowBatchChunking(true)
+ preferredChunkSizeOption.foreach { size =>
+ builder.setPreferredArrowChunkSize(size.toLong)
+ }
+ builder.build()
+ })
+ .build())
+ .setUserContext(context)
+ .setSessionId(UUID.randomUUID.toString())
+ .build()
+
+ // Execute plan.
+ var done = false
+ val responses = mutable.Buffer.empty[proto.ExecutePlanResponse]
+ instance.executePlan(
+ request,
+ new StreamObserver[proto.ExecutePlanResponse] {
+ override def onNext(v: proto.ExecutePlanResponse): Unit = {
+ responses += v
+ verifyEvents.onNext(v)
+ }
+
+ override def onError(throwable: Throwable): Unit = {
+ verifyEvents.onError(throwable)
+ throw throwable
+ }
+
+ override def onCompleted(): Unit = {
+ done = true
+ }
+ })
+ verifyEvents.onCompleted(Some(2000))
+ // The current implementation is expected to be blocking. This is here to make sure it
+ // is.
+ assert(done)
+
+ // ArrowBatch chunks of 4 partitions + Metrics + optional progress messages
+ val filteredResponses = responses.filter(!_.hasExecutionProgress)
+
+ // Make sure the first response is schema only
+ val head = filteredResponses.head
+ assert(head.hasSchema && !head.hasArrowBatch && !head.hasMetrics)
+
+ // Make sure the last response is metrics only
+ val last = filteredResponses.last
+ assert(last.hasMetrics && !last.hasSchema && !last.hasArrowBatch)
+
+ val allocator = new RootAllocator()
+
+ // Check the 'data' batches
+ val arrowResponseQueue = mutable.Queue(filteredResponses.tail.dropRight(1).toSeq: _*)
+
+ def dequeueNextBatch(): List[proto.ExecutePlanResponse] = {
+ if (arrowResponseQueue.isEmpty) return Nil
+ val n = arrowResponseQueue.front.getArrowBatch.getNumChunksInBatch.toInt
+ List.fill(n)(arrowResponseQueue.dequeue())
+ }
+
+ var batchCount = 0
+ var expectedId = 0L
+ var previousEId = 0.0d
+ while (arrowResponseQueue.nonEmpty) {
+ val batch = dequeueNextBatch()
+ // In this example, the max chunk size is set to a small value, so each Arrow batch
+ // should be split into multiple chunks.
+ assert(batch.size > 5)
+ batchCount += 1
+
+ val rowCount = batch.head.getArrowBatch.getRowCount
+ val rowStartOffset = batch.head.getArrowBatch.getStartOffset
+ batch.zipWithIndex.foreach { case (chunk, chunkIndex) =>
+ assert(chunk.getArrowBatch.getChunkIndex == chunkIndex)
+ assert(chunk.getArrowBatch.getNumChunksInBatch == batch.size)
+ assert(chunk.getArrowBatch.getRowCount == rowCount)
+ assert(chunk.getArrowBatch.getStartOffset == rowStartOffset)
+ assert(chunk.getArrowBatch.getData != null)
+ assert(chunk.getArrowBatch.getData.size() > 0)
+ assert(
+ chunk.getArrowBatch.getData.size() <= preferredChunkSizeOption.getOrElse(
+ maxChunkSizeOption.get))
+ }
+
+ // Reassemble the chunks into a single Arrow batch and validate its content.
+ val input = new SequenceInputStream(
+ batch.iterator.map(_.getArrowBatch.getData.newInput()).asJavaEnumeration)
+ val reader = new ArrowStreamReader(input, allocator)
+ while (reader.loadNextBatch()) {
+ val root = reader.getVectorSchemaRoot
+ val idVector = root.getVector(0).asInstanceOf[BigIntVector]
+ val eidVector = root.getVector(1).asInstanceOf[Float8Vector]
+ val numRows = root.getRowCount
+ var i = 0
+ while (i < numRows) {
+ assert(idVector.get(i) == expectedId)
+ expectedId += 1
+ val eid = eidVector.get(i)
+ assert(eid > previousEId)
+ previousEId = eid
+ i += 1
+ }
+ }
+ reader.close()
+ }
+ // There are 4 batches (partitions) in total.
+ assert(batchCount == 4)
+ allocator.close()
+ }
+ }
+ }
+ }
+
gridTest("SPARK-43923: commands send events")(
Seq(
(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org