(spark-connect-swift) branch main updated: [SPARK-52756] Support `defineFlow`

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon 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 a7e7feb  [SPARK-52756] Support `defineFlow`
a7e7feb is described below

commit a7e7feb1c12eeeb282af14be793eada971b02be3
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Thu Jul 10 07:48:48 2025 -0700

    [SPARK-52756] Support `defineFlow`
    
    ### What changes were proposed in this pull request?
    
    This PR aims to support `defineFlow ` API in order to support `Declarative Pipelines` (SPARK-51727) of Apache Spark `4.1.0-preview1`.
    
    ### Why are the changes needed?
    
    To support the new feature incrementally.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No, this is a new feature.
    
    ### How was this patch tested?
    
    Pass the CIs with `4.1.0-preview1` test pipeline.
    - #210
    
    <img width="1000" height="373" alt="Screenshot 2025-07-10 at 07 25 37" src="https://github.com/user-attachments/assets/b4e214f6-de6c-4c31-8482-58e8de1dd4ff" />
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    No.
    
    Closes #212 from dongjoon-hyun/SPARK-52756.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 Sources/SparkConnect/SparkConnectClient.swift      | 31 ++++++++++++++++++++++
 .../SparkConnectClientTests.swift                  | 18 +++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index 109ee0d..8ea6250 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -1272,6 +1272,37 @@ public actor SparkConnectClient {
     }
   }
 
+  @discardableResult
+  func defineFlow(
+    _ dataflowGraphID: String,
+    _ flowName: String,
+    _ targetDatasetName: String,
+    _ relation: Relation
+  ) async throws -> Bool {
+    try await withGPRC { client in
+      if UUID(uuidString: dataflowGraphID) == nil {
+        throw SparkConnectError.InvalidArgument
+      }
+
+      var defineFlow = Spark_Connect_PipelineCommand.DefineFlow()
+      defineFlow.dataflowGraphID = dataflowGraphID
+      defineFlow.flowName = flowName
+      defineFlow.targetDatasetName = targetDatasetName
+      defineFlow.plan = relation
+
+      var pipelineCommand = Spark_Connect_PipelineCommand()
+      pipelineCommand.commandType = .defineFlow(defineFlow)
+
+      var command = Spark_Connect_Command()
+      command.commandType = .pipelineCommand(pipelineCommand)
+
+      let responses = try await execute(self.sessionID!, command)
+      return responses.contains {
+        $0.responseType == .pipelineCommandResult(Spark_Connect_PipelineCommandResult())
+      }
+    }
+  }
+
   private enum URIParams {
     static let PARAM_GRPC_MAX_MESSAGE_SIZE = "grpc_max_message_size"
     static let PARAM_SESSION_ID = "session_id"
diff --git a/Tests/SparkConnectTests/SparkConnectClientTests.swift b/Tests/SparkConnectTests/SparkConnectClientTests.swift
index 1983be2..72e31ba 100644
--- a/Tests/SparkConnectTests/SparkConnectClientTests.swift
+++ b/Tests/SparkConnectTests/SparkConnectClientTests.swift
@@ -146,4 +146,22 @@ struct SparkConnectClientTests {
     }
     await client.stop()
   }
+
+  @Test
+  func defineFlow() async throws {
+    let client = SparkConnectClient(remote: TEST_REMOTE)
+    let response = try await client.connect(UUID().uuidString)
+
+    try await #require(throws: SparkConnectError.InvalidArgument) {
+      try await client.defineFlow("not-a-uuid-format", "f1", "ds1", Relation())
+    }
+
+    if response.sparkVersion.version.starts(with: "4.1") {
+      let dataflowGraphID = try await client.createDataflowGraph()
+      #expect(UUID(uuidString: dataflowGraphID) != nil)
+      let relation = await client.getLocalRelation().root
+      #expect(try await client.defineFlow(dataflowGraphID, "f1", "ds1", relation))
+    }
+    await client.stop()
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org