You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by do...@apache.org on 2026/08/03 22:19:11 UTC
(spark-connect-swift) branch main updated: [SPARK-58528] Support cache artifact upload in `SparkConnectClient`
This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/spark-connect-swift.git
The following commit(s) were added to refs/heads/main by this push:
new dbb0e2f [SPARK-58528] Support cache artifact upload in `SparkConnectClient`
dbb0e2f is described below
commit dbb0e2f2144f2a87abddb87ef8f440d13ca446d0
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Mon Aug 3 15:19:04 2026 -0700
[SPARK-58528] Support cache artifact upload in `SparkConnectClient`
### What changes were proposed in this pull request?
This PR aims to support cache artifact upload in `SparkConnectClient` by adding a new file, `SparkConnectClient+Artifact.swift`, with the following internal APIs.
- `artifactExists`: checks whether an artifact exists in the server-side session via the `ArtifactStatus` RPC.
- `cacheArtifact`: caches the given data as a `cache/<sha256>` artifact and returns its SHA-256 hash. The upload is skipped if the artifact already exists. Data up to 32KiB is uploaded as a single-chunk batch, and larger data as a chunked artifact stream, like PySpark's `ArtifactManager.cache_artifact`.
The existing `addArtifact` is moved to the new file to group all artifact operations together.
### Why are the changes needed?
This is an infrastructure for supporting large `createDataFrame` data via `CachedLocalRelation`.
### Does this PR introduce _any_ user-facing change?
No, these are internal APIs.
### How was this patch tested?
Pass the CIs with a new test suite, `SparkConnectClientArtifactTests`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Fable 5
Closes #472 from dongjoon-hyun/SPARK-58528.
Authored-by: Dongjoon Hyun <do...@apache.org>
Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
.../SparkConnect/SparkConnectClient+Artifact.swift | 150 +++++++++++++++++++++
Sources/SparkConnect/SparkConnectClient.swift | 36 +----
.../SparkConnectClientArtifactTests.swift | 72 ++++++++++
3 files changed, 223 insertions(+), 35 deletions(-)
diff --git a/Sources/SparkConnect/SparkConnectClient+Artifact.swift b/Sources/SparkConnect/SparkConnectClient+Artifact.swift
new file mode 100644
index 0000000..cf500a0
--- /dev/null
+++ b/Sources/SparkConnect/SparkConnectClient+Artifact.swift
@@ -0,0 +1,150 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+//
+
+#if canImport(FoundationEssentials)
+import FoundationEssentials
+#else
+import Foundation
+#endif
+import GRPCCore
+import GRPCNIOTransportHTTP2
+
+/// Extension providing artifact operations on ``SparkConnectClient``.
+extension SparkConnectClient {
+ /// The chunk size used when splitting an artifact into multiple `AddArtifactsRequest` messages,
+ /// following the midpoint recommendation of 32KiB for gRPC payloads.
+ static let artifactChunkSize = 32 * 1024
+
+ func addArtifact(_ url: URL) async throws {
+ guard url.lastPathComponent.hasSuffix(".jar") else {
+ throw SparkConnectError.InvalidArgument
+ }
+
+ let JAR_PREFIX = "jars"
+ let name = "\(JAR_PREFIX)/" + url.lastPathComponent
+
+ try await withGPRC { client in
+ let service = SparkConnectService.Client(wrapping: client)
+
+ var chunk = Spark_Connect_AddArtifactsRequest.ArtifactChunk()
+ chunk.data = try Data(contentsOf: url)
+ chunk.crc = Int64(CRC32.checksum(data: chunk.data))
+
+ var singleChunk = Spark_Connect_AddArtifactsRequest.SingleChunkArtifact()
+ singleChunk.name = name
+ singleChunk.data = chunk
+ var batch = Spark_Connect_AddArtifactsRequest.Batch()
+ batch.artifacts.append(singleChunk)
+
+ var addArtifactsRequest = Spark_Connect_AddArtifactsRequest()
+ addArtifactsRequest.sessionID = self.sessionID!
+ addArtifactsRequest.userContext = self.userContext
+ addArtifactsRequest.clientType = self.clientType
+ addArtifactsRequest.batch = batch
+ let request = addArtifactsRequest
+ _ = try await service.addArtifacts(
+ request: StreamingClientRequest<Spark_Connect_AddArtifactsRequest> { x in
+ try await x.write(contentsOf: [request])
+ })
+ }
+ }
+
+ /// Check whether an artifact with the given name exists in the server-side session.
+ /// - Parameter name: An artifact name, e.g. `cache/abc123`.
+ /// - Returns: True if the artifact exists.
+ func artifactExists(_ name: String) async throws -> Bool {
+ try await withGPRC { client in
+ let service = SparkConnectService.Client(wrapping: client)
+ var request = Spark_Connect_ArtifactStatusesRequest()
+ request.sessionID = self.sessionID!
+ request.userContext = self.userContext
+ request.clientType = self.clientType
+ request.names = [name]
+ let response = try await service.artifactStatus(request)
+ return response.statuses[name]?.exists ?? false
+ }
+ }
+
+ /// Cache the given data as a `cache/` artifact in the server-side session and return its
+ /// SHA-256 hash. The upload is skipped if an artifact with the same hash already exists.
+ /// Data larger than ``artifactChunkSize`` is uploaded as a chunked artifact stream.
+ /// - Parameter data: The data to cache.
+ /// - Returns: A SHA-256 hash of the data as a lowercase hex string.
+ func cacheArtifact(_ data: Data) async throws -> String {
+ let hash = SHA256.hexString(data: data)
+ let name = "cache/\(hash)"
+ if try await artifactExists(name) {
+ return hash
+ }
+
+ try await withGPRC { client in
+ let service = SparkConnectService.Client(wrapping: client)
+ let chunkSize = Self.artifactChunkSize
+ var requests: [Spark_Connect_AddArtifactsRequest] = []
+ if data.count <= chunkSize {
+ var chunk = Spark_Connect_AddArtifactsRequest.ArtifactChunk()
+ chunk.data = data
+ chunk.crc = Int64(CRC32.checksum(data: chunk.data))
+
+ var singleChunk = Spark_Connect_AddArtifactsRequest.SingleChunkArtifact()
+ singleChunk.name = name
+ singleChunk.data = chunk
+ var batch = Spark_Connect_AddArtifactsRequest.Batch()
+ batch.artifacts.append(singleChunk)
+
+ var request = Spark_Connect_AddArtifactsRequest()
+ request.sessionID = self.sessionID!
+ request.userContext = self.userContext
+ request.clientType = self.clientType
+ request.batch = batch
+ requests.append(request)
+ } else {
+ let numChunks = (data.count + chunkSize - 1) / chunkSize
+ for i in 0..<numChunks {
+ let range = (i * chunkSize)..<min((i + 1) * chunkSize, data.count)
+ var chunk = Spark_Connect_AddArtifactsRequest.ArtifactChunk()
+ chunk.data = data.subdata(in: range)
+ chunk.crc = Int64(CRC32.checksum(data: chunk.data))
+
+ var request = Spark_Connect_AddArtifactsRequest()
+ request.sessionID = self.sessionID!
+ request.userContext = self.userContext
+ request.clientType = self.clientType
+ if i == 0 {
+ var beginChunk = Spark_Connect_AddArtifactsRequest.BeginChunkedArtifact()
+ beginChunk.name = name
+ beginChunk.totalBytes = Int64(data.count)
+ beginChunk.numChunks = Int64(numChunks)
+ beginChunk.initialChunk = chunk
+ request.beginChunk = beginChunk
+ } else {
+ request.chunk = chunk
+ }
+ requests.append(request)
+ }
+ }
+ let allRequests = requests
+ _ = try await service.addArtifacts(
+ request: StreamingClientRequest<Spark_Connect_AddArtifactsRequest> { x in
+ try await x.write(contentsOf: allRequests)
+ })
+ }
+ return hash
+ }
+}
diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index aa762ea..6195125 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -132,7 +132,7 @@ public actor SparkConnectClient {
}
}
- private func withGPRC<Result: Sendable>(
+ func withGPRC<Result: Sendable>(
_ f: (GRPCClient<GRPCNIOTransportHTTP2.HTTP2ClientTransport.Posix>) async throws -> Result
) async throws -> Result {
try await withGRPCClient(
@@ -893,40 +893,6 @@ public actor SparkConnectClient {
return plan
}
- func addArtifact(_ url: URL) async throws {
- guard url.lastPathComponent.hasSuffix(".jar") else {
- throw SparkConnectError.InvalidArgument
- }
-
- let JAR_PREFIX = "jars"
- let name = "\(JAR_PREFIX)/" + url.lastPathComponent
-
- try await withGPRC { client in
- let service = SparkConnectService.Client(wrapping: client)
-
- var chunk = Spark_Connect_AddArtifactsRequest.ArtifactChunk()
- chunk.data = try Data(contentsOf: url)
- chunk.crc = Int64(CRC32.checksum(data: chunk.data))
-
- var singleChunk = Spark_Connect_AddArtifactsRequest.SingleChunkArtifact()
- singleChunk.name = name
- singleChunk.data = chunk
- var batch = Spark_Connect_AddArtifactsRequest.Batch()
- batch.artifacts.append(singleChunk)
-
- var addArtifactsRequest = Spark_Connect_AddArtifactsRequest()
- addArtifactsRequest.sessionID = self.sessionID!
- addArtifactsRequest.userContext = self.userContext
- addArtifactsRequest.clientType = self.clientType
- addArtifactsRequest.batch = batch
- let request = addArtifactsRequest
- _ = try await service.addArtifacts(
- request: StreamingClientRequest<Spark_Connect_AddArtifactsRequest> { x in
- try await x.write(contentsOf: [request])
- })
- }
- }
-
/// Add a tag to be assigned to all the operations started by this thread in this session.
/// - Parameter tag: The tag to be added. Cannot contain ',' (comma) character or be an empty string.
public func addTag(tag: String) throws {
diff --git a/Tests/SparkConnectTests/SparkConnectClientArtifactTests.swift b/Tests/SparkConnectTests/SparkConnectClientArtifactTests.swift
new file mode 100644
index 0000000..9955b95
--- /dev/null
+++ b/Tests/SparkConnectTests/SparkConnectClientArtifactTests.swift
@@ -0,0 +1,72 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+//
+
+#if canImport(FoundationEssentials)
+import FoundationEssentials
+#else
+import Foundation
+#endif
+import Testing
+
+@testable import SparkConnect
+
+/// A test suite for the artifact operations of `SparkConnectClient`
+@Suite(.serialized)
+struct SparkConnectClientArtifactTests {
+ let TEST_REMOTE = ProcessInfo.processInfo.environment["SPARK_REMOTE"] ?? "sc://localhost"
+
+ @Test
+ func cacheArtifact() async throws {
+ let client = try SparkConnectClient(remote: TEST_REMOTE)
+ try await client.connect(UUID().uuidString)
+
+ let data = "Apache Spark Connect Client for Swift".data(using: .utf8)!
+ let hash = try await client.cacheArtifact(data)
+ #expect(hash == SHA256.hexString(data: data))
+ #expect(try await client.artifactExists("cache/\(hash)"))
+
+ // The second call skips the upload and returns the same hash.
+ #expect(try await client.cacheArtifact(data) == hash)
+ await client.stop()
+ }
+
+ @Test
+ func cacheArtifactChunked() async throws {
+ let client = try SparkConnectClient(remote: TEST_REMOTE)
+ try await client.connect(UUID().uuidString)
+
+ // Larger than the 32KiB chunk size in order to exercise the chunked upload path.
+ var data = Data(capacity: 100 * 1024)
+ for i in 0..<(100 * 1024) {
+ data.append(UInt8(i % 251))
+ }
+ let hash = try await client.cacheArtifact(data)
+ #expect(hash == SHA256.hexString(data: data))
+ #expect(try await client.artifactExists("cache/\(hash)"))
+ await client.stop()
+ }
+
+ @Test
+ func artifactExists() async throws {
+ let client = try SparkConnectClient(remote: TEST_REMOTE)
+ try await client.connect(UUID().uuidString)
+ #expect(try await client.artifactExists("cache/nonexistent") == false)
+ await client.stop()
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org