[SPARK-51504] Support DataFrame.select

Author: dongjoon-hyun

Sources/SparkConnect/DataFrame.swift

@@ -36,7 +36,7 @@ public actor DataFrame: Sendable {
   /// - Parameters:
   ///   - spark: A ``SparkSession`` instance to use.
   ///   - plan: A plan to execute.
-  init(spark: SparkSession, plan: Plan) async throws {
+  init(spark: SparkSession, plan: Plan) {
     self.spark = spark
     self.plan = plan
   }
@@ -192,4 +192,9 @@ public actor DataFrame: Sendable {
       print(table.render())
     }
   }
+
+  public func select(_ cols: String...) -> DataFrame {
+    let plan = SparkConnectClient.getProject(self.plan.root, cols)
+    return DataFrame(spark: self.spark, plan: plan)
+  }
 }

Sources/SparkConnect/Extension.swift

@@ -45,6 +45,12 @@ extension String {
     keyValue.key = self
     return keyValue
   }
+
+  var toUnresolvedAttribute: UnresolvedAttribute {
+    var attribute = UnresolvedAttribute()
+    attribute.unparsedIdentifier = self
+    return attribute
+  }
 }
 
 extension [String: String] {

Sources/SparkConnect/SparkConnectClient.swift

@@ -252,4 +252,20 @@ public actor SparkConnectClient {
     request.analyze = .schema(schema)
     return request
   }
+
+  static func getProject(_ child: Relation, _ cols: [String]) -> Plan {
+    var project = Project()
+    project.input = child
+    let expressions: [Spark_Connect_Expression] = cols.map {
+      var expression = Spark_Connect_Expression()
+      expression.exprType = .unresolvedAttribute($0.toUnresolvedAttribute)
+      return expression
+    }
+    project.expressions = expressions
+    var relation = Relation()
+    relation.project = project
+    var plan = Plan()
+    plan.opType = .root(relation)
+    return plan
+  }
 }

Sources/SparkConnect/TypeAliases.swift

@@ -22,8 +22,10 @@ typealias ConfigRequest = Spark_Connect_ConfigRequest
 typealias DataType = Spark_Connect_DataType
 typealias ExecutePlanRequest = Spark_Connect_ExecutePlanRequest
 typealias Plan = Spark_Connect_Plan
+typealias Project = Spark_Connect_Project
 typealias KeyValue = Spark_Connect_KeyValue
 typealias Range = Spark_Connect_Range
 typealias Relation = Spark_Connect_Relation
 typealias SparkConnectService = Spark_Connect_SparkConnectService
 typealias UserContext = Spark_Connect_UserContext
+typealias UnresolvedAttribute = Spark_Connect_Expression.UnresolvedAttribute

Tests/SparkConnectTests/DataFrameTests.swift

@@ -68,6 +68,45 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func selectNone() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let emptySchema = try await spark.range(1).select().schema()
+    #expect(emptySchema == #"{"struct":{}}"#)
+    await spark.stop()
+  }
+
+  @Test
+  func select() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let schema = try await spark.range(1).select("id").schema()
+    #expect(
+      schema
+        == #"{"struct":{"fields":[{"name":"id","dataType":{"long":{}}}]}}"#
+    )
+    await spark.stop()
+  }
+
+  @Test
+  func selectMultipleColumns() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let schema = try await spark.sql("SELECT * FROM VALUES (1, 2)").select("col2", "col1").schema()
+    #expect(
+      schema
+        == #"{"struct":{"fields":[{"name":"col2","dataType":{"integer":{}}},{"name":"col1","dataType":{"integer":{}}}]}}"#
+    )
+    await spark.stop()
+  }
+
+  @Test
+  func selectInvalidColumn() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    try await #require(throws: Error.self) {
+      let _ = try await spark.range(1).select("invalid").schema()
+    }
+    await spark.stop()
+  }
+
   @Test
   func table() async throws {
     let spark = try await SparkSession.builder.getOrCreate()