[SPARK-52168] Support to for DataFrame

Author: dongjoon-hyun

Sources/SparkConnect/DataFrame.swift

@@ -94,6 +94,7 @@ import Synchronization
 /// - ``show(_:_:_:)``
 ///
 /// ### Transformation Operations
+/// - ``to(_:)``
 /// - ``toDF(_:)``
 /// - ``toJSON()``
 /// - ``select(_:)``
@@ -478,6 +479,19 @@ public actor DataFrame: Sendable {
     return df
   }
 
+  /// Returns a new DataFrame where each row is reconciled to match the specified schema.
+  /// - Parameter schema: The given schema.
+  /// - Returns: A ``DataFrame`` with the given schema.
+  public func to(_ schema: String) async throws -> DataFrame {
+    // Validate by parsing.
+    do {
+      let dataType = try await sparkSession.client.ddlParse(schema)
+      return DataFrame(spark: self.spark, plan: SparkConnectClient.getToSchema(self.plan.root, dataType))
+    } catch {
+      throw SparkConnectError.InvalidTypeException
+    }
+  }
+
   /// Returns the content of the Dataset as a Dataset of JSON strings.
   /// - Returns: A ``DataFrame`` with a single string column whose content is JSON.
   public func toJSON() -> DataFrame {

Sources/SparkConnect/SparkConnectClient.swift

@@ -505,6 +505,17 @@ public actor SparkConnectClient {
     return plan
   }
 
+  static func getToSchema(_ child: Relation, _ schema: Spark_Connect_DataType) -> Plan {
+    var toSchema = Spark_Connect_ToSchema()
+    toSchema.input = child
+    toSchema.schema = schema
+    var relation = Relation()
+    relation.toSchema = toSchema
+    var plan = Plan()
+    plan.opType = .root(relation)
+    return plan
+  }
+
   static func getProjectExprs(_ child: Relation, _ exprs: [String]) -> Plan {
     var project = Project()
     project.input = child

Tests/SparkConnectTests/DataFrameTests.swift

@@ -201,6 +201,26 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func to() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+
+    let schema1 = try await spark.range(1).to("shortID SHORT").schema
+    #expect(
+      schema1
+      == #"{"struct":{"fields":[{"name":"shortID","dataType":{"short":{}},"nullable":true}]}}"#
+    )
+
+    let schema2 = try await spark.sql("SELECT '1'").to("id INT").schema
+    print(schema2)
+    #expect(
+      schema2
+      == #"{"struct":{"fields":[{"name":"id","dataType":{"integer":{}},"nullable":true}]}}"#
+    )
+
+    await spark.stop()
+  }
+
   @Test
   func selectMultipleColumns() async throws {
     let spark = try await SparkSession.builder.getOrCreate()