[SPARK-51971] Improve DataFrame.collect to return the original values

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

This PR aims to improve `DataFrame.collect` to return the original values.

Note that this PR provides simple value types first. More types like `Decimal` will be added later.

### Why are the changes needed?

The initial implementation has a limitation to return rows of `String` values.

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

No, because there is no released versions yet.

### How was this patch tested?

Pass the CIs.

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

No.

Closes #99 from dongjoon-hyun/SPARK-51971.

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

Author: dongjoon-hyun

Sources/SparkConnect/Catalog.swift

@@ -252,7 +252,7 @@ public actor Catalog: Sendable {
       catalog.tableExists = tableExists
       return catalog
     })
-    return "true" == (try await df.collect().first!.get(0) as! String)
+    return try await df.collect()[0].getAsBool(0)
   }
 
   /// Check if the table or view with the specified name exists. This can either be a temporary
@@ -270,7 +270,7 @@ public actor Catalog: Sendable {
       catalog.tableExists = tableExists
       return catalog
     })
-    return "true" == (try await df.collect().first!.get(0) as! String)
+    return try await df.collect()[0].getAsBool(0)
   }
 
   /// Check if the function with the specified name exists. This can either be a temporary function
@@ -287,7 +287,7 @@ public actor Catalog: Sendable {
       catalog.functionExists = functionExists
       return catalog
     })
-    return "true" == (try await df.collect().first!.get(0) as! String)
+    return try await df.collect()[0].getAsBool(0)
   }
 
   /// Check if the function with the specified name exists in the specified database under the Hive
@@ -305,7 +305,7 @@ public actor Catalog: Sendable {
       catalog.functionExists = functionExists
       return catalog
     })
-    return "true" == (try await df.collect().first!.get(0) as! String)
+    return try await df.collect()[0].getAsBool(0)
   }
 
   /// Caches the specified table in-memory.
@@ -338,7 +338,7 @@ public actor Catalog: Sendable {
       catalog.isCached = isCached
       return catalog
     })
-    return "true" == (try await df.collect().first!.get(0) as! String)
+    return try await df.collect()[0].getAsBool(0)
   }
 
   /// Invalidates and refreshes all the cached data and metadata of the given table.
@@ -407,7 +407,7 @@ public actor Catalog: Sendable {
       catalog.dropTempView = dropTempView
       return catalog
     })
-    return "true" == (try await df.collect().first!.get(0) as! String)
+    return try await df.collect().first!.getAsBool(0)
   }
 
   /// Drops the global temporary view with the given view name in the catalog. If the view has been
@@ -423,6 +423,6 @@ public actor Catalog: Sendable {
       catalog.dropGlobalTempView = dropGlobalTempView
       return catalog
     })
-    return "true" == (try await df.collect().first!.get(0) as! String)
+    return try await df.collect()[0].getAsBool(0)
   }
 }

Sources/SparkConnect/DataFrame.swift

@@ -208,14 +208,34 @@ public actor DataFrame: Sendable {
       for i in 0..<batch.length {
         var values: [Sendable?] = []
         for column in batch.columns {
-          let str = column.array as! AsString
           if column.data.isNull(i) {
             values.append(nil)
-          } else if column.data.type.info == ArrowType.ArrowBinary {
-            let binary = str.asString(i).utf8.map { String(format: "%02x", $0) }.joined(separator: " ")
-            values.append("[\(binary)]")
           } else {
-            values.append(str.asString(i))
+            let array = column.array
+            switch column.data.type.info {
+            case .primitiveInfo(.boolean):
+              values.append(array.asAny(i) as? Bool)
+            case .primitiveInfo(.int8):
+              values.append(array.asAny(i) as? Int8)
+            case .primitiveInfo(.int16):
+              values.append(array.asAny(i) as? Int16)
+            case .primitiveInfo(.int32):
+              values.append(array.asAny(i) as? Int32)
+            case .primitiveInfo(.int64):
+              values.append(array.asAny(i) as! Int64)
+            case .primitiveInfo(.float):
+              values.append(array.asAny(i) as? Float)
+            case .primitiveInfo(.double):
+              values.append(array.asAny(i) as? Double)
+            case .primitiveInfo(.date32):
+              values.append(array.asAny(i) as! Date)
+            case ArrowType.ArrowBinary:
+              values.append((array as! AsString).asString(i).utf8)
+            case .complexInfo(.strct):
+              values.append((array as! AsString).asString(i))
+            default:
+              values.append(array.asAny(i) as? String)
+            }
           }
         }
         result.append(Row(valueArray: values))

Sources/SparkConnect/Row.swift

@@ -50,6 +50,10 @@ public struct Row: Sendable, Equatable {
     return values[i]
   }
 
+  public func getAsBool(_ i: Int) throws -> Bool {
+    return try get(i) as! Bool
+  }
+
   public static func == (lhs: Row, rhs: Row) -> Bool {
     if lhs.values.count != rhs.values.count {
       return false
@@ -59,16 +63,8 @@ public struct Row: Sendable, Equatable {
         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? any FixedWidthInteger, let b = y as? any FixedWidthInteger {
+        return Int64(a) == Int64(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 {