Add limit

Author: dongjoon-hyun

Sources/SparkConnect/DataFrame.swift

@@ -197,4 +197,9 @@ public actor DataFrame: Sendable {
     let plan = SparkConnectClient.getProject(self.plan.root, cols)
     return DataFrame(spark: self.spark, plan: plan)
   }
+
+  public func limit(_ n: Int32) -> DataFrame {
+    let plan = SparkConnectClient.getLimit(self.plan.root, n)
+    return DataFrame(spark: self.spark, plan: plan)
+  }
 }

Sources/SparkConnect/SparkConnectClient.swift

@@ -268,4 +268,15 @@ public actor SparkConnectClient {
     plan.opType = .root(relation)
     return plan
   }
+
+  static func getLimit(_ child: Relation, _ n: Int32) -> Plan {
+    var limit = Limit()
+    limit.input = child
+    limit.limit = n
+    var relation = Relation()
+    relation.limit = limit
+    var plan = Plan()
+    plan.opType = .root(relation)
+    return plan
+  }
 }

Sources/SparkConnect/TypeAliases.swift

@@ -24,6 +24,7 @@ typealias ExecutePlanRequest = Spark_Connect_ExecutePlanRequest
 typealias Plan = Spark_Connect_Plan
 typealias Project = Spark_Connect_Project
 typealias KeyValue = Spark_Connect_KeyValue
+typealias Limit = Spark_Connect_Limit
 typealias Range = Spark_Connect_Range
 typealias Relation = Spark_Connect_Relation
 typealias SparkConnectService = Spark_Connect_SparkConnectService

Tests/SparkConnectTests/DataFrameTests.swift

@@ -107,6 +107,15 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func limit() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    #expect(try await spark.range(10).limit(0).count() == 0)
+    #expect(try await spark.range(10).limit(1).count() == 1)
+    #expect(try await spark.range(10).limit(2).count() == 2)
+    await spark.stop()
+  }
+
   @Test
   func table() async throws {
     let spark = try await SparkSession.builder.getOrCreate()