[SPARK-51815] Add Row struct

### What changes were proposed in this pull request?

This PR aims to add `Row` struct and use it.

### Why are the changes needed?

To make `DataFrame` APIs return `Row` instead of `[String?]` in `Swift`.

```swift
- public func collect() async throws -> [[String?]] {
+ public func collect() async throws -> [Row] {

- public func head(_ n: Int32 = 1) async throws -> [[String?]] {
+ public func head(_ n: Int32 = 1) async throws -> [Row] {
```

Note that `Row` is added to support general type fields, but this PR replaces the existing API's `[String?]` signature into `Row`-based signature. The detailed one-to-one mapping among other types will be handled later.

### Does this PR introduce _any_ user-facing change?

Yes, but this is a change to the unreleased version.

### How was this patch tested?

Pass the CIs.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes #63 from dongjoon-hyun/SPARK-51815.

Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>

Author: dongjoon-hyun

Sources/SparkConnect/Catalog.swift

@@ -100,7 +100,7 @@ public actor Catalog: Sendable {
       catalog.catType = .currentCatalog(Spark_Connect_CurrentCatalog())
       return catalog
     })
-    return try await df.collect()[0][0]!
+    return try await df.collect()[0][0] as! String
   }
 
   /// Sets the current default catalog in this session.
@@ -130,7 +130,7 @@ public actor Catalog: Sendable {
       return catalog
     })
     return try await df.collect().map {
-      CatalogMetadata(name: $0[0]!, description: $0[1])
+      try CatalogMetadata(name: $0[0] as! String, description: $0[1] as? String)
     }
   }
 
@@ -142,7 +142,7 @@ public actor Catalog: Sendable {
       catalog.catType = .currentDatabase(Spark_Connect_CurrentDatabase())
       return catalog
     })
-    return try await df.collect()[0][0]!
+    return try await df.collect()[0][0] as! String
   }
 
   /// Sets the current default database in this session.
@@ -173,7 +173,7 @@ public actor Catalog: Sendable {
       return catalog
     })
     return try await df.collect().map {
-      Database(name: $0[0]!, catalog: $0[1], description: $0[2], locationUri: $0[3]!)
+      try Database(name: $0[0] as! String, catalog: $0[1] as? String, description: $0[2] as? String, locationUri: $0[3] as! String)
     }
   }
 
@@ -189,7 +189,7 @@ public actor Catalog: Sendable {
       return catalog
     })
     return try await df.collect().map {
-      Database(name: $0[0]!, catalog: $0[1], description: $0[2], locationUri: $0[3]!)
+      try Database(name: $0[0] as! String, catalog: $0[1] as? String, description: $0[2] as? String, locationUri: $0[3] as! String)
     }.first!
   }
 

Sources/SparkConnect/DataFrame.swift

@@ -197,15 +197,15 @@ public actor DataFrame: Sendable {
     }
   }
 
-  /// Execute the plan and return the result as ``[[String?]]``.
-  /// - Returns: ``[[String?]]``
-  public func collect() async throws -> [[String?]] {
+  /// Execute the plan and return the result as ``[Row]``.
+  /// - Returns: ``[Row]``
+  public func collect() async throws -> [Row] {
     try await execute()
 
-    var result: [[String?]] = []
+    var result: [Row] = []
     for batch in self.batches {
       for i in 0..<batch.length {
-        var values: [String?] = []
+        var values: [Sendable?] = []
         for column in batch.columns {
           let str = column.array as! AsString
           if column.data.isNull(i) {
@@ -217,7 +217,7 @@ public actor DataFrame: Sendable {
             values.append(str.asString(i))
           }
         }
-        result.append(values)
+        result.append(Row(valueArray: values))
       }
     }
 
@@ -377,15 +377,15 @@ public actor DataFrame: Sendable {
 
   /// Returns the first `n` rows.
   /// - Parameter n: The number of rows. (default: 1)
-  /// - Returns: ``[[String?]]``
-  public func head(_ n: Int32 = 1) async throws -> [[String?]] {
+  /// - Returns: ``[Row]``
+  public func head(_ n: Int32 = 1) async throws -> [Row] {
     return try await limit(n).collect()
   }
 
   /// Returns the last `n` rows.
   /// - Parameter n: The number of rows.
-  /// - Returns: ``[[String?]]``
-  public func tail(_ n: Int32) async throws -> [[String?]] {
+  /// - Returns: ``[Row]``
+  public func tail(_ n: Int32) async throws -> [Row] {
     let lastN = DataFrame(spark:spark, plan: SparkConnectClient.getTail(self.plan.root, n))
     return try await lastN.collect()
   }

Sources/SparkConnect/Row.swift

@@ -0,0 +1,90 @@
+//
+// 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.
+//
+import Foundation
+
+public struct Row: Sendable, Equatable {
+  let values: [Sendable?]
+
+  public init(_ values: Sendable?...) {
+    self.values = values
+  }
+
+  public init(valueArray: [Sendable?]) {
+    self.values = valueArray
+  }
+
+  public static var empty: Row {
+    return Row()
+  }
+
+  public var size: Int { return length }
+
+  public var length: Int { return values.count }
+
+  subscript(index: Int) -> Sendable {
+    get throws {
+      return try get(index)
+    }
+  }
+
+  public func get(_ i: Int) throws -> Sendable {
+    if i < 0 || i >= self.length {
+      throw SparkConnectError.InvalidArgumentException
+    }
+    return values[i]
+  }
+
+  public static func == (lhs: Row, rhs: Row) -> Bool {
+    if lhs.values.count != rhs.values.count {
+      return false
+    }
+    return lhs.values.elementsEqual(rhs.values) { (x, y) in
+      if x == nil && y == nil {
+        return true
+      } else if let a = x as? Bool, let b = y as? Bool {
+        return a == b
+      } else if let a = x as? Int, let b = y as? Int {
+        return a == b
+      } else if let a = x as? Int8, let b = y as? Int8 {
+        return a == b
+      } else if let a = x as? Int16, let b = y as? Int16 {
+        return a == b
+      } else if let a = x as? Int32, let b = y as? Int32 {
+        return a == b
+      } else if let a = x as? Int64, let b = y as? Int64 {
+        return a == b
+      } else if let a = x as? Float, let b = y as? Float {
+        return a == b
+      } else if let a = x as? Double, let b = y as? Double {
+        return a == b
+      } else if let a = x as? String, let b = y as? String {
+        return a == b
+      } else {
+        return false
+      }
+    }
+  }
+
+  public func toString() -> String {
+    return "[\(self.values.map { "\($0 ?? "null")" }.joined(separator: ","))]"
+  }
+}
+
+extension Row {
+}

Tests/SparkConnectTests/DataFrameTests.swift

@@ -249,15 +249,15 @@ struct DataFrameTests {
   @Test
   func sort() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
-    let expected = (1...10).map{ [String($0)] }
+    let expected = Array((1...10).map{ Row(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()
-    let expected = (1...10).map{ [String($0)] }
+    let expected = Array((1...10).map{ Row(String($0)) })
     #expect(try await spark.range(10, 0, -1).orderBy("id").collect() == expected)
     await spark.stop()
   }
@@ -284,28 +284,28 @@ struct DataFrameTests {
     #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"]])
+      ).collect() == [Row("1", "true", "abc"), Row(nil, nil, nil), Row("3", "false", "def")])
     await spark.stop()
   }
 
   @Test
   func head() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
     #expect(try await spark.range(0).head().isEmpty)
-    #expect(try await spark.range(2).sort("id").head() == [["0"]])
-    #expect(try await spark.range(2).sort("id").head(1) == [["0"]])
-    #expect(try await spark.range(2).sort("id").head(2) == [["0"], ["1"]])
-    #expect(try await spark.range(2).sort("id").head(3) == [["0"], ["1"]])
+    #expect(try await spark.range(2).sort("id").head() == [Row("0")])
+    #expect(try await spark.range(2).sort("id").head(1) == [Row("0")])
+    #expect(try await spark.range(2).sort("id").head(2) == [Row("0"), Row("1")])
+    #expect(try await spark.range(2).sort("id").head(3) == [Row("0"), Row("1")])
     await spark.stop()
   }
 
   @Test
   func tail() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
     #expect(try await spark.range(0).tail(1).isEmpty)
-    #expect(try await spark.range(2).sort("id").tail(1) == [["1"]])
-    #expect(try await spark.range(2).sort("id").tail(2) == [["0"], ["1"]])
-    #expect(try await spark.range(2).sort("id").tail(3) == [["0"], ["1"]])
+    #expect(try await spark.range(2).sort("id").tail(1) == [Row("1")])
+    #expect(try await spark.range(2).sort("id").tail(2) == [Row("0"), Row("1")])
+    #expect(try await spark.range(2).sort("id").tail(3) == [Row("0"), Row("1")])
     await spark.stop()
   }
 

Tests/SparkConnectTests/Resources/queries/binary.sql.answer

@@ -0,0 +1 @@
+[[61 62 63]]

Tests/SparkConnectTests/Resources/queries/binary.sql.json

@@ -0,0 +1 @@
+

Tests/SparkConnectTests/Resources/queries/cache.sql.answer

@@ -0,0 +1 @@
+