[SPARK-51676] Support printSchema for DataFrame

Author: dongjoon-hyun

Sources/SparkConnect/DataFrame.swift

@@ -355,4 +355,24 @@ public actor DataFrame: Sendable {
       print(response.explain.explainString)
     }
   }
+
+  /// Prints the schema to the console in a nice tree format.
+  public func printSchema() async throws {
+    try await printSchema(Int32.max)
+  }
+
+  /// Prints the schema up to the given level to the console in a nice tree format.
+  /// - Parameter level: A level to be printed.
+  public func printSchema(_ level: Int32) async throws {
+    try await withGRPCClient(
+      transport: .http2NIOPosix(
+        target: .dns(host: spark.client.host, port: spark.client.port),
+        transportSecurity: .plaintext
+      )
+    ) { client in
+      let service = Spark_Connect_SparkConnectService.Client(wrapping: client)
+      let response = try await service.analyzePlan(spark.client.getTreeString(spark.sessionID, plan, level))
+      print(response.treeString.treeString)
+    }
+  }
 }

Sources/SparkConnect/SparkConnectClient.swift

@@ -294,6 +294,18 @@ public actor SparkConnectClient {
       })
   }
 
+  func getTreeString(_ sessionID: String, _ plan: Plan, _ level: Int32) async -> AnalyzePlanRequest
+  {
+    return analyze(
+      sessionID,
+      {
+        var treeString = AnalyzePlanRequest.TreeString()
+        treeString.plan = plan
+        treeString.level = level
+        return OneOf_Analyze.treeString(treeString)
+      })
+  }
+
   static func getProject(_ child: Relation, _ cols: [String]) -> Plan {
     var project = Project()
     project.input = child

Tests/SparkConnectTests/DataFrameTests.swift

@@ -70,6 +70,14 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func printSchema() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    try await spark.sql("SELECT struct(1, 2)").printSchema()
+    try await spark.sql("SELECT struct(1, 2)").printSchema(1)
+    await spark.stop()
+  }
+
   @Test
   func explain() async throws {
     let spark = try await SparkSession.builder.getOrCreate()