From 70ef41c610dc71a4e7702fc9ada92774ab55f35b Mon Sep 17 00:00:00 2001
From: Dongjoon Hyun <dongjoon@apache.org>
Date: Thu, 13 Mar 2025 19:04:16 -0700
Subject: [PATCH 1/2] [SPARK-51508] Support `collect(): [[String?]]` for
 `DataFrame`

---
 Sources/SparkConnect/DataFrame.swift          | 29 +++++++++++++++----
 Sources/SparkConnect/SparkConnectClient.swift |  2 ++
 Sources/SparkConnect/SparkSession.swift       |  6 ++--
 Tests/SparkConnectTests/DataFrameTests.swift  | 21 ++++++++++++--
 .../SparkConnectTests/SparkSessionTests.swift |  4 +++
 5 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/Sources/SparkConnect/DataFrame.swift b/Sources/SparkConnect/DataFrame.swift
index 7ab4fc6..81b74b1 100644
--- a/Sources/SparkConnect/DataFrame.swift
+++ b/Sources/SparkConnect/DataFrame.swift
@@ -58,7 +58,7 @@ public actor DataFrame: Sendable {
 
   /// Add `Apache Arrow`'s `RecordBatch`s to the internal array.
   /// - Parameter batches: An array of ``RecordBatch``.
-  private func addBathes(_ batches: [RecordBatch]) {
+  private func addBatches(_ batches: [RecordBatch]) {
     self.batches.append(contentsOf: batches)
   }
 
@@ -153,16 +153,35 @@ public actor DataFrame: Sendable {
             let arrowResult = ArrowReader.makeArrowReaderResult()
             _ = reader.fromMessage(schema, dataBody: Data(), result: arrowResult)
             _ = reader.fromMessage(dataHeader, dataBody: dataBody, result: arrowResult)
-            await self.addBathes(arrowResult.batches)
+            await self.addBatches(arrowResult.batches)
           }
         }
       }
     }
   }
 
-  /// This is designed not to support this feature in order to simplify the Swift client.
-  public func collect() async throws {
-    throw SparkConnectError.UnsupportedOperationException
+  /// Execute the plan and return the result as ``[[String?]]``.
+  /// - Returns: ``[[String?]]``
+  public func collect() async throws -> [[String?]] {
+    try await execute()
+
+    var result: [[String?]] = []
+    for batch in self.batches {
+      for i in 0..<batch.length {
+        var values: [String?] = []
+        for column in batch.columns {
+          let str = column.array as! AsString
+          if column.data.isNull(i) {
+            values.append(nil)
+          } else {
+            values.append(str.asString(i))
+          }
+        }
+        result.append(values)
+      }
+    }
+
+    return result
   }
 
   /// Execute the plan and show the result.
diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index 0ab1e0e..b9b6e34 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -275,9 +275,11 @@ public actor SparkConnectClient {
     let expressions: [Spark_Connect_Expression.SortOrder] = cols.map {
       var expression = Spark_Connect_Expression.SortOrder()
       expression.child.exprType = .unresolvedAttribute($0.toUnresolvedAttribute)
+      expression.direction = .ascending
       return expression
     }
     sort.order = expressions
+    sort.isGlobal = true
     var relation = Relation()
     relation.sort = sort
     var plan = Plan()
diff --git a/Sources/SparkConnect/SparkSession.swift b/Sources/SparkConnect/SparkSession.swift
index 9eec0de..660a70a 100644
--- a/Sources/SparkConnect/SparkSession.swift
+++ b/Sources/SparkConnect/SparkSession.swift
@@ -45,12 +45,10 @@ public actor SparkSession {
   ///   - userID: an optional user ID. If absent, `SPARK_USER` environment or ``ProcessInfo.processInfo.userName`` is used.
   init(_ connection: String, _ userID: String? = nil) {
     let processInfo = ProcessInfo.processInfo
-#if os(iOS) || os(watchOS) || os(tvOS)
-    let userName = processInfo.environment["SPARK_USER"] ?? ""
-#elseif os(macOS) || os(Linux)
+#if os(macOS) || os(Linux)
     let userName = processInfo.environment["SPARK_USER"] ?? processInfo.userName
 #else
-    assert(false, "Unsupported platform")
+    let userName = processInfo.environment["SPARK_USER"] ?? ""
 #endif
     self.client = SparkConnectClient(remote: connection, user: userID ?? userName)
     self.conf = RuntimeConf(self.client)
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift b/Tests/SparkConnectTests/DataFrameTests.swift
index 2b57d2d..c44254b 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -125,19 +125,25 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+#if !os(Linux)
   @Test
   func sort() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
-    #expect(try await spark.range(10).sort("id").count() == 10)
+    let expected = (1...10).map{ [String($0)] }
+    #expect(try await spark.range(10, 0, -1).sort("id").collect() == expected)
     await spark.stop()
   }
 
   @Test
   func orderBy() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
-    #expect(try await spark.range(10).orderBy("id").count() == 10)
+    let expected = (1...10).map{ [String($0)] }
+    print(expected)
+    print(try await spark.range(10, 0, -1).orderBy("id").collect())
+    #expect(try await spark.range(10, 0, -1).orderBy("id").collect() == expected)
     await spark.stop()
   }
+#endif
 
   @Test
   func table() async throws {
@@ -153,6 +159,17 @@ struct DataFrameTests {
   }
 
 #if !os(Linux)
+  @Test
+  func collect() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    #expect(try await spark.range(0).collect().isEmpty)
+    #expect(
+      try await spark.sql(
+        "SELECT * FROM VALUES (1, true, 'abc'), (null, null, null), (3, false, 'def')"
+      ).collect() == [["1", "true", "abc"], [nil, nil, nil], ["3", "false", "def"]])
+    await spark.stop()
+  }
+
   @Test
   func show() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
diff --git a/Tests/SparkConnectTests/SparkSessionTests.swift b/Tests/SparkConnectTests/SparkSessionTests.swift
index 921cecc..4a2a549 100644
--- a/Tests/SparkConnectTests/SparkSessionTests.swift
+++ b/Tests/SparkConnectTests/SparkSessionTests.swift
@@ -41,7 +41,11 @@ struct SparkSessionTests {
 
   @Test func userContext() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
+#if os(macOS) || os(Linux)
     let defaultUserContext = ProcessInfo.processInfo.userName.toUserContext
+#else
+    let defaultUserContext = "".toUserContext
+#endif
     #expect(await spark.client.userContext == defaultUserContext)
     await spark.stop()
   }

From b0282833cf9eb8edb44518fb2004975d58c9a53d Mon Sep 17 00:00:00 2001
From: Dongjoon Hyun <dongjoon@apache.org>
Date: Thu, 13 Mar 2025 19:45:31 -0700
Subject: [PATCH 2/2] clean up print statement

---
 Tests/SparkConnectTests/DataFrameTests.swift | 2 --
 1 file changed, 2 deletions(-)

diff --git a/Tests/SparkConnectTests/DataFrameTests.swift b/Tests/SparkConnectTests/DataFrameTests.swift
index c44254b..a1c7e7e 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -138,8 +138,6 @@ struct DataFrameTests {
   func orderBy() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
     let expected = (1...10).map{ [String($0)] }
-    print(expected)
-    print(try await spark.range(10, 0, -1).orderBy("id").collect())
     #expect(try await spark.range(10, 0, -1).orderBy("id").collect() == expected)
     await spark.stop()
   }