viirya commented on code in PR #597:
URL:
https://github.com/apache/spark-connect-swift/pull/597#discussion_r4298259156
##########
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:
This will break as soon as another request option is added. How about
checking for the option itself?
```swift
#expect(request.requestOptions.contains {
$0.resultChunkingOptions.allowArrowBatchChunking })
```
##########
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:
nit: "the same with" → "the same as".
##########
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:
A question about smaller limits. If a user sets `grpc_max_message_size`
below about 115 MiB (for example 64 MiB), the server still uses its default
~115 MiB chunks, so chunking won't help for batches between the two sizes.
Scala and Python don't handle this either, so it's fine to leave as is. Would
it be worth sending `preferred_arrow_chunk_size` (for example `maxMessageSize *
9 / 10`) when `maxMessageSize` is below the default? A follow-up would be fine
too.
##########
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:
It's worth noting here why both limits must be set. In
grpc-swift-nio-transport 2.11.0, the client's inbound deframer is sized from
`maxRequestMessageBytes ?? 4 MiB` (`Connection.swift`, `maxPayloadSize:
maxRequestSize`). `maxResponseMessageBytes` isn't used by the transport for
responses. So the 128 MiB response limit takes effect only because the request
limit is set too. If someone later drops the request limit, responses would
silently fall back to 4 MiB. Something like:
```swift
// Note that grpc-swift-nio-transport limits inbound response messages by
// `maxRequestMessageBytes`, so both limits are required.
```
##########
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:
nit: For the multi-chunk branch, `Data(chunks.joined())` is shorter and
avoids growing the buffer repeatedly:
```swift
return chunks.count == 1 ? chunks[0] : Data(chunks.joined())
```
--
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]
