[SPARK-51621] Support sparkSession for DataFrame

Author: dongjoon-hyun

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 {

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
+  }
+}

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

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()