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 3fe1f0e  [SPARK-51621] Support `sparkSession` for `DataFrame`
3fe1f0e is described below

commit 3fe1f0e8407606b3f6db06d518b55f1f8d62010d
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Wed Mar 26 20:25:24 2025 -0700

    [SPARK-51621] Support `sparkSession` for `DataFrame`
    
    ### What changes were proposed in this pull request?
    
    This PR aims to support `sparkSession` API for `DataFrame`.
    
    ### Why are the changes needed?
    
    For feature parity.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No, this is a new addition.
    
    ### How was this patch tested?
    
    Pass the CIs.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    No.
    
    Closes #28 from dongjoon-hyun/SPARK-51621.
    
    Authored-by: Dongjoon Hyun <[email protected]>
    Signed-off-by: Dongjoon Hyun <[email protected]>
---
 Sources/SparkConnect/DataFrame.swift         | 6 ++++++
 Sources/SparkConnect/Extension.swift         | 6 ++++++
 Sources/SparkConnect/SparkSession.swift      | 2 +-
 Tests/SparkConnectTests/DataFrameTests.swift | 7 +++++++
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/Sources/SparkConnect/DataFrame.swift 
b/Sources/SparkConnect/DataFrame.swift
index 237c08c..4669943 100644
--- a/Sources/SparkConnect/DataFrame.swift
+++ b/Sources/SparkConnect/DataFrame.swift
@@ -62,6 +62,12 @@ public actor DataFrame: Sendable {
     self.batches.append(contentsOf: batches)
   }
 
+  /// Return the `SparkSession` of this `DataFrame`.
+  /// - Returns: A `SparkSession`
+  public func sparkSession() -> SparkSession {
+    return self.spark
+  }
+
   /// A method to access the underlying Spark's `RDD`.
   /// In `Spark Connect`, this feature is not allowed by design.
   public func rdd() throws {
diff --git a/Sources/SparkConnect/Extension.swift 
b/Sources/SparkConnect/Extension.swift
index 848a96e..da330cc 100644
--- a/Sources/SparkConnect/Extension.swift
+++ b/Sources/SparkConnect/Extension.swift
@@ -77,3 +77,9 @@ extension Data {
   /// Get an `Int32` value from unsafe 4 bytes.
   var int32: Int32 { withUnsafeBytes({ $0.load(as: Int32.self) }) }
 }
+
+extension SparkSession: Equatable {
+  public static func == (lhs: SparkSession, rhs: SparkSession) -> Bool {
+    return lhs.sessionID == rhs.sessionID
+  }
+}
diff --git a/Sources/SparkConnect/SparkSession.swift 
b/Sources/SparkConnect/SparkSession.swift
index e68deef..524f46a 100644
--- a/Sources/SparkConnect/SparkSession.swift
+++ b/Sources/SparkConnect/SparkSession.swift
@@ -62,7 +62,7 @@ public actor SparkSession {
   }
 
   /// A unique session ID for this session from client.
-  var sessionID: String = UUID().uuidString
+  nonisolated let sessionID: String = UUID().uuidString
 
   /// Get the current session ID
   /// - Returns: the current session ID
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift 
b/Tests/SparkConnectTests/DataFrameTests.swift
index c49fa15..7cb1e1f 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -23,6 +23,13 @@ import Testing
 
 /// A test suite for `DataFrame`
 struct DataFrameTests {
+  @Test
+  func sparkSession() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    #expect(try await spark.range(1).sparkSession() == spark)
+    await spark.stop()
+  }
+
   @Test
   func rdd() async throws {
     let spark = try await SparkSession.builder.getOrCreate()


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to