Add sort, orderBy

Author: dongjoon-hyun

Sources/SparkConnect/DataFrame.swift

@@ -194,12 +194,18 @@ public actor DataFrame: Sendable {
   }
 
   public func select(_ cols: String...) -> DataFrame {
-    let plan = SparkConnectClient.getProject(self.plan.root, cols)
-    return DataFrame(spark: self.spark, plan: plan)
+    return DataFrame(spark: self.spark, plan: SparkConnectClient.getProject(self.plan.root, cols))
+  }
+
+  public func sort(_ cols: String...) -> DataFrame {
+    return DataFrame(spark: self.spark, plan: SparkConnectClient.getSort(self.plan.root, cols))
+  }
+
+  public func orderBy(_ cols: String...) -> DataFrame {
+    return DataFrame(spark: self.spark, plan: SparkConnectClient.getSort(self.plan.root, cols))
   }
 
   public func limit(_ n: Int32) -> DataFrame {
-    let plan = SparkConnectClient.getLimit(self.plan.root, n)
-    return DataFrame(spark: self.spark, plan: plan)
+    return DataFrame(spark: self.spark, plan: SparkConnectClient.getLimit(self.plan.root, n))
   }
 }

Sources/SparkConnect/SparkConnectClient.swift

@@ -269,6 +269,22 @@ public actor SparkConnectClient {
     return plan
   }
 
+  static func getSort(_ child: Relation, _ cols: [String]) -> Plan {
+    var sort = Sort()
+    sort.input = child
+    let expressions: [Spark_Connect_Expression.SortOrder] = cols.map {
+      var expression = Spark_Connect_Expression.SortOrder()
+      expression.child.exprType = .unresolvedAttribute($0.toUnresolvedAttribute)
+      return expression
+    }
+    sort.order = expressions
+    var relation = Relation()
+    relation.sort = sort
+    var plan = Plan()
+    plan.opType = .root(relation)
+    return plan
+  }
+
   static func getLimit(_ child: Relation, _ n: Int32) -> Plan {
     var limit = Limit()
     limit.input = child

Sources/SparkConnect/TypeAliases.swift

@@ -28,5 +28,6 @@ typealias Limit = Spark_Connect_Limit
 typealias Range = Spark_Connect_Range
 typealias Relation = Spark_Connect_Relation
 typealias SparkConnectService = Spark_Connect_SparkConnectService
+typealias Sort = Spark_Connect_Sort
 typealias UserContext = Spark_Connect_UserContext
 typealias UnresolvedAttribute = Spark_Connect_Expression.UnresolvedAttribute

Tests/SparkConnectTests/DataFrameTests.swift

@@ -116,6 +116,20 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func sort() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    #expect(try await spark.range(10).sort("id").count() == 10)
+    await spark.stop()
+  }
+
+  @Test
+  func orderBy() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    #expect(try await spark.range(10).orderBy("id").count() == 10)
+    await spark.stop()
+  }
+
   @Test
   func table() async throws {
     let spark = try await SparkSession.builder.getOrCreate()