You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "xi-db (via GitHub)" <gi...@apache.org> on 2025/09/08 13:27:27 UTC
[PR] [SPARK-53525] Spark Connect ArrowBatch Result Chunking [spark]
xi-db opened a new pull request, #52271:
URL: https://github.com/apache/spark/pull/52271
### 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 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 followup PR.)
To reproduce the issue:
```
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 will fails with `StatusCode.RESOURCE_EXHAUSTED` error with message `Received message larger than max (314570608 vs. 134217728)`.
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 server side and client side.
### Was this patch authored or co-authored using generative AI tooling?
No.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2336727323
##########
sql/connect/common/src/main/protobuf/spark/connect/base.proto:
##########
@@ -819,6 +828,19 @@ message ReattachOptions {
bool reattachable = 1;
}
+message ResultChunkingOptions {
Review Comment:
Yes, I introduced a client-side option `preferred_arrow_chunk_size` ([this commit](https://github.com/apache/spark/pull/52271/commits/b68461f384d476a199bf46bb7ee584288ac4a5ba)).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on PR #52271:
URL: https://github.com/apache/spark/pull/52271#issuecomment-3269698566
> @xi-db are you adding scala support in a follow-up?
Yes, I'm implementing the scala support as a follow-up.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2340463706
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala:
##########
@@ -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
Review Comment:
Nit: `option.getResultChunkingOptions.getAllowArrowBatchChunking` should be sufficient, since the default value is false. `option.getResultChunkingOptions` will return a default message even if it is not set.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330564421
##########
sql/connect/common/src/main/protobuf/spark/connect/base.proto:
##########
@@ -819,6 +828,19 @@ message ReattachOptions {
bool reattachable = 1;
}
+message ResultChunkingOptions {
Review Comment:
Also, can we name this ResultOptions? I can also see us setting a max arrow batch size here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2333327019
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala:
##########
@@ -374,6 +374,132 @@ class SparkConnectServiceSuite
}
}
+ test("Arrow batch chunking") {
+ withEvents { verifyEvents =>
+ val overriddenMaxChunkSize = 100
+ withSparkConf(
+ Connect.CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE.key ->
+ overriddenMaxChunkSize.toString) {
+ 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, exp(id) as eid from range(0, 100, 1, 4)"))
+ .build()
+ val request = proto.ExecutePlanRequest
+ .newBuilder()
+ .setPlan(plan)
+ .addRequestOptions(
+ proto.ExecutePlanRequest.RequestOption
+ .newBuilder()
+ .setResultChunkingOptions(proto.ResultChunkingOptions
+ .newBuilder()
+ .setAllowArrowBatchChunking(true)
+ .build())
+ .build())
+ .setUserContext(context)
+ .setSessionId(UUID.randomUUID.toString())
+ .build()
+
+ // Execute plan.
+ @volatile 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(100))
+ // 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() <= overriddenMaxChunkSize)
+ }
+
+ // Reassemble the chunks into a single Arrow batch and validate its content.
+ val batchData: ByteString =
Review Comment:
Yeah, it's a good point, updated the test as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2340476581
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala:
##########
@@ -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 =>
Review Comment:
nit: `option.getResultChunkingOptions.hasPreferredArrowChunkSize` is sufficient
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala:
##########
@@ -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 =>
Review Comment:
nit: `option.getResultChunkingOptions.hasPreferredArrowChunkSize` should be sufficient
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330557376
##########
sql/connect/common/src/main/protobuf/spark/connect/base.proto:
##########
@@ -819,6 +828,19 @@ message ReattachOptions {
bool reattachable = 1;
}
+message ResultChunkingOptions {
Review Comment:
Should the client be able to set chunk size?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on PR #52271:
URL: https://github.com/apache/spark/pull/52271#issuecomment-3281464057
Merging to master. You can fix the NITS in a follow-up.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330855138
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala:
##########
@@ -374,6 +374,132 @@ class SparkConnectServiceSuite
}
}
+ test("Arrow batch chunking") {
+ withEvents { verifyEvents =>
+ val overriddenMaxChunkSize = 100
+ withSparkConf(
+ Connect.CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE.key ->
+ overriddenMaxChunkSize.toString) {
+ 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, exp(id) as eid from range(0, 100, 1, 4)"))
+ .build()
+ val request = proto.ExecutePlanRequest
+ .newBuilder()
+ .setPlan(plan)
+ .addRequestOptions(
+ proto.ExecutePlanRequest.RequestOption
+ .newBuilder()
+ .setResultChunkingOptions(proto.ResultChunkingOptions
+ .newBuilder()
+ .setAllowArrowBatchChunking(true)
+ .build())
+ .build())
+ .setUserContext(context)
+ .setSessionId(UUID.randomUUID.toString())
+ .build()
+
+ // Execute plan.
+ @volatile var done = false
+ val responses = mutable.Buffer.empty[proto.ExecutePlanResponse]
Review Comment:
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)...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330841927
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala:
##########
@@ -141,19 +143,48 @@ 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)
+ if (isResultChunkingEnabled) {
+ // Result chunking is enabled. Split the result batch into multiple chunks if needed.
+ val maxChunkSize = 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())
+ for (i <- 0 until numChunks) {
Review Comment:
Please don't use for comprehensions... Either use an actual loop, or the functional equivalent this translates into.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2340463706
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala:
##########
@@ -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
Review Comment:
nit: `option.getResultChunkingOptions.getAllowArrowBatchChunking` should be sufficient, since the default value is false. `option.getResultChunkingOptions` will return a default message even if it is not set. In proto3, you won’t get a null pointer when accessing an unset field.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330564421
##########
sql/connect/common/src/main/protobuf/spark/connect/base.proto:
##########
@@ -819,6 +828,19 @@ message ReattachOptions {
bool reattachable = 1;
}
+message ResultChunkingOptions {
Review Comment:
Also, can we name this ResultOptions?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2333371840
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala:
##########
@@ -141,19 +143,48 @@ 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)
+ if (isResultChunkingEnabled) {
+ // Result chunking is enabled. Split the result batch into multiple chunks if needed.
+ val maxChunkSize = 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())
+ for (i <- 0 until numChunks) {
Review Comment:
Yeah, right, updated to an explicit loop with foreach.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "asf-gitbox-commits (via GitHub)" <gi...@apache.org>.
asf-gitbox-commits closed pull request #52271: [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking
URL: https://github.com/apache/spark/pull/52271
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330843792
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala:
##########
@@ -141,19 +143,48 @@ 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)
+ if (isResultChunkingEnabled) {
+ // Result chunking is enabled. Split the result batch into multiple chunks if needed.
+ val maxChunkSize = 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())
+ for (i <- 0 until numChunks) {
+ val from = i * maxChunkSize
+ val to = math.min(from + maxChunkSize, bytes.length)
+ val length = to - from
+
+ val response = proto.ExecutePlanResponse
Review Comment:
You can reuse the builder.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on PR #52271:
URL: https://github.com/apache/spark/pull/52271#issuecomment-3267299417
@xi-db are you adding scala support in a follow-up?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2332677082
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -392,4 +392,20 @@ 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." +
+ " However, excessively large plans remain unsupported due to Spark internals and JVM" +
Review Comment:
Sounds good, removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2340463706
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/ExecuteHolder.scala:
##########
@@ -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
Review Comment:
Nit: `option.getResultChunkingOptions.getAllowArrowBatchChunking` should be sufficient, since the default value is false. `option.getResultChunkingOptions` will return a default message even if it is not set. In proto3, you won’t get a null pointer when accessing an unset field.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330919792
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala:
##########
@@ -374,6 +374,132 @@ class SparkConnectServiceSuite
}
}
+ test("Arrow batch chunking") {
+ withEvents { verifyEvents =>
+ val overriddenMaxChunkSize = 100
+ withSparkConf(
+ Connect.CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE.key ->
+ overriddenMaxChunkSize.toString) {
+ 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, exp(id) as eid from range(0, 100, 1, 4)"))
+ .build()
+ val request = proto.ExecutePlanRequest
+ .newBuilder()
+ .setPlan(plan)
+ .addRequestOptions(
+ proto.ExecutePlanRequest.RequestOption
+ .newBuilder()
+ .setResultChunkingOptions(proto.ResultChunkingOptions
+ .newBuilder()
+ .setAllowArrowBatchChunking(true)
+ .build())
+ .build())
+ .setUserContext(context)
+ .setSessionId(UUID.randomUUID.toString())
+ .build()
+
+ // Execute plan.
+ @volatile 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(100))
+ // 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() <= overriddenMaxChunkSize)
+ }
+
+ // Reassemble the chunks into a single Arrow batch and validate its content.
+ val batchData: ByteString =
Review Comment:
When you build this for scala, please use a Concatenating InputStream or something like that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2330846917
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -392,4 +392,20 @@ 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." +
+ " However, excessively large plans remain unsupported due to Spark internals and JVM" +
Review Comment:
Remove these two lines. They are not related to the conf.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2332669292
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/execution/SparkConnectPlanExecution.scala:
##########
@@ -141,19 +143,48 @@ 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)
+ if (isResultChunkingEnabled) {
+ // Result chunking is enabled. Split the result batch into multiple chunks if needed.
+ val maxChunkSize = 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())
+ for (i <- 0 until numChunks) {
+ val from = i * maxChunkSize
+ val to = math.min(from + maxChunkSize, bytes.length)
+ val length = to - from
+
+ val response = proto.ExecutePlanResponse
Review Comment:
Yes, updated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-53525][CONNECT] Spark Connect ArrowBatch Result Chunking [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #52271:
URL: https://github.com/apache/spark/pull/52271#discussion_r2332818320
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectServiceSuite.scala:
##########
@@ -374,6 +374,132 @@ class SparkConnectServiceSuite
}
}
+ test("Arrow batch chunking") {
+ withEvents { verifyEvents =>
+ val overriddenMaxChunkSize = 100
+ withSparkConf(
+ Connect.CONNECT_SESSION_RESULT_CHUNKING_MAX_CHUNK_SIZE.key ->
+ overriddenMaxChunkSize.toString) {
+ 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, exp(id) as eid from range(0, 100, 1, 4)"))
+ .build()
+ val request = proto.ExecutePlanRequest
+ .newBuilder()
+ .setPlan(plan)
+ .addRequestOptions(
+ proto.ExecutePlanRequest.RequestOption
+ .newBuilder()
+ .setResultChunkingOptions(proto.ResultChunkingOptions
+ .newBuilder()
+ .setAllowArrowBatchChunking(true)
+ .build())
+ .build())
+ .setUserContext(context)
+ .setSessionId(UUID.randomUUID.toString())
+ .build()
+
+ // Execute plan.
+ @volatile var done = false
+ val responses = mutable.Buffer.empty[proto.ExecutePlanResponse]
Review Comment:
`done` here doesn't need volatile. I added volatile because other test cases have it. Removed as it is not needed in our case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org