viirya commented on code in PR #597:
URL:
https://github.com/apache/spark-connect-swift/pull/597#discussion_r4298463584
##########
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:
Sorry, my earlier suggestion was wrong. `FlattenSequence<[Data]>` doesn't
provide contiguous storage, and its `underestimatedCount` is 0. So
`Data.init(_: some Sequence<UInt8>)` falls back to copying byte by byte (the
`slowElements` path in swift-foundation). This branch only runs for batches
over about 115 MiB, so that path matters. The original `reduce(into:)` with
`append` copied whole chunks. Reserving capacity up front keeps that and avoids
regrowth:
```swift
return chunks.count == 1
? chunks[0]
: chunks.reduce(into: Data(capacity: chunks.reduce(0) { $0 + $1.count }))
{ $0.append($1) }
```
Reverting to the original `reduce(into: Data())` is also fine.
##########
Tests/SparkConnectTests/DataFrameTests.swift:
##########
@@ -704,6 +704,22 @@ struct DataFrameTests {
await spark.stop()
}
+ @Test
+ func resultChunking() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ if await isSparkVersionAtLeast(spark.version, "4.1") {
+ let df = try await spark.range(10000).selectExpr("id", "repeat('a', 100)
AS s")
+ let expected = try await df.collect()
+ // Force the server to split each Arrow batch into chunks.
+ try await
spark.conf.set("spark.connect.session.resultChunking.maxChunkSize", "1024")
+ #expect(try await df.collect() == expected)
+ #expect(try await df.count() == 10000)
Review Comment:
nit: These two checks don't exercise the `chunkIndex == 0` counting in
`executeAndCount()`. `count()` is computed on the server and returns a single
row, and `isEmpty()` uses `limit(1)`, where over-counting can't turn a non-zero
result into zero. To cover it, `#expect(try await df.executeAndCount() ==
10000)` would need `@testable import SparkConnect`. Otherwise it's fine to drop
these two lines, since `collect()` already covers chunk assembly.
--
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]
