dongjoon-hyun commented on code in PR #597:
URL:
https://github.com/apache/spark-connect-swift/pull/597#discussion_r4298338982
##########
Sources/SparkConnect/SparkConnectClient.swift:
##########
@@ -109,17 +112,13 @@ public actor SparkConnectClient {
self.transportSecurity = .plaintext
}
self.maxMessageSize = maxMessageSize
- if let maxMessageSize {
- // An empty service name is the default configuration for all services
and methods.
- self.serviceConfig = ServiceConfig(methodConfig: [
- MethodConfig(
- names: [MethodConfig.Name(service: "")],
- maxRequestMessageBytes: maxMessageSize,
- maxResponseMessageBytes: maxMessageSize)
- ])
- } else {
- self.serviceConfig = ServiceConfig()
- }
+ // An empty service name is the default configuration for all services and
methods.
Review Comment:
Good catch, thanks. Added the note in fefb833.
##########
Sources/SparkConnect/SparkConnectClient.swift:
##########
@@ -486,6 +485,12 @@ public actor SparkConnectClient {
request.operationID = UUID().uuidString
request.tags = Array(tags)
request.plan = plan
+ // Allow the server to split large Arrow batches into chunks which
`DataFrame` assembles.
+ var chunkingOptions = ResultChunkingOptions()
+ chunkingOptions.allowArrowBatchChunking = true
Review Comment:
Thanks for the suggestion. I'd like to keep this PR aligned with Scala and
Python, which don't handle this either. I'll do it in a follow-up PR.
##########
Sources/SparkConnect/DataFrame+Actions.swift:
##########
@@ -38,15 +38,30 @@ extension DataFrame {
/// Batches are local to each call since concurrent actions can interleave
on this actor.
private func execute() async throws -> [RecordBatch] {
let batches = Mutex([RecordBatch]())
+ // The chunks of the Arrow batch being received. The server splits a large
batch into chunks.
+ let chunks = Mutex([Data]())
try await
spark.client.executePlanWithReattach(spark.client.getExecutePlanRequest(plan)) {
m in
if m.hasSchema {
// The original schema should arrive before ArrowBatches
try await self.setSchema(DataType(m.schema))
}
- let ipcStreamBytes = m.arrowBatch.data
- if !ipcStreamBytes.isEmpty && m.arrowBatch.rowCount > 0 {
+ if !m.arrowBatch.data.isEmpty && m.arrowBatch.rowCount > 0,
+ let ipcStreamBytes = try chunks.withLock({ chunks -> Data? in
+ guard m.arrowBatch.chunkIndex == chunks.count else {
+ throw SparkConnectError.InvalidArrowData
+ }
+ chunks.append(m.arrowBatch.data)
+ // `numChunksInBatch` is missing when the batch is not chunked.
+ guard m.arrowBatch.numChunksInBatch == 0 || chunks.count ==
m.arrowBatch.numChunksInBatch
+ else {
+ return nil
+ }
+ defer { chunks.removeAll() }
+ return chunks.count == 1 ? chunks[0] : chunks.reduce(into: Data()) {
$0.append($1) }
Review Comment:
Done in fefb833.
##########
Tests/SparkConnectTests/SparkConnectClientTests.swift:
##########
@@ -123,6 +125,16 @@ struct SparkConnectClientTests {
await client.stop()
}
+ @Test
+ func resultChunking() async throws {
+ let client = try SparkConnectClient(remote: TEST_REMOTE)
+ try await client.connect(UUID().uuidString)
+ let request = await client.getExecutePlanRequest(await
client.getPlanRange(0, 1, 1))
+ #expect(request.requestOptions.count == 1)
Review Comment:
Done in fefb833.
##########
Sources/SparkConnect/SparkConnectClient.swift:
##########
@@ -28,14 +28,17 @@ import Synchronization
/// Conceptually the remote spark session that communicates with the server
public actor SparkConnectClient {
+ /// The default `grpc_max_message_size`, 128 MiB, which is the same with
Scala and Python clients.
Review Comment:
Done in fefb833.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
