dongjoon-hyun commented on code in PR #597:
URL:
https://github.com/apache/spark-connect-swift/pull/597#discussion_r4298485569
##########
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:
No problem, thanks for catching that. Switched to `reduce(into:)` with the
capacity reserved up front in 90da317.
##########
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:
Good point. Replaced the two checks with `executeAndCount() == 10000` and
switched `DataFrameTests` to `@testable import SparkConnect` in 90da317.
--
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]
