(spark-connect-swift) branch main updated: [SPARK-52182] Support `withWatermark` for `DataFrame`

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/spark-connect-swift.git


The following commit(s) were added to refs/heads/main by this push:
     new e346c9e  [SPARK-52182] Support `withWatermark` for `DataFrame`
e346c9e is described below

commit e346c9ee99132592872821bdb88c7f86fcb2fb8b
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Fri May 16 05:51:08 2025 -0700

    [SPARK-52182] Support `withWatermark` for `DataFrame`
    
    ### What changes were proposed in this pull request?
    
    This PR aims to ad `withWatermark` for `DataFrame`.
    
    ### Why are the changes needed?
    
    For feature parity.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Pass the CIs.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    No.
    
    Closes #158 from dongjoon-hyun/SPARK-52182.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 Sources/SparkConnect/DataFrame.swift          | 14 ++++++++++++++
 Sources/SparkConnect/SparkConnectClient.swift | 17 +++++++++++++++++
 Tests/SparkConnectTests/DataFrameTests.swift  | 16 ++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/Sources/SparkConnect/DataFrame.swift b/Sources/SparkConnect/DataFrame.swift
index 75f555f..945651e 100644
--- a/Sources/SparkConnect/DataFrame.swift
+++ b/Sources/SparkConnect/DataFrame.swift
@@ -117,6 +117,7 @@ import Synchronization
 /// - ``transpose()``
 /// - ``transpose(_:)``
 /// - ``hint(_:_:)``
+/// - ``withWatermark(_:_:)``
 ///
 /// ### Join Operations
 /// - ``join(_:)``
@@ -1444,6 +1445,19 @@ public actor DataFrame: Sendable {
     try await checkpoint(eager, false, storageLevel)
   }
 
+  /// Defines an event time watermark for this ``DataFrame``.  A watermark tracks a point in time
+  /// before which we assume no more late data is going to arrive.
+  /// - Parameters:
+  ///   - eventTime: the name of the column that contains the event time of the row.
+  ///   - delayThreshold: the minimum delay to wait to data to arrive late, relative to
+  ///   the latest record that has been processed in the form of an interval (e.g. "1 minute" or "5 hours").
+  ///   NOTE: This should not be negative.
+  /// - Returns: A ``DataFrame`` instance.
+  public func withWatermark(_ eventTime: String, _ delayThreshold: String) -> DataFrame {
+    let plan = SparkConnectClient.getWithWatermark(self.plan.root, eventTime, delayThreshold)
+    return DataFrame(spark: self.spark, plan: plan)
+  }
+
   /// Returns a ``DataFrameWriter`` that can be used to write non-streaming data.
   public var write: DataFrameWriter {
     get {
diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index c69888f..a5306dc 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -1022,6 +1022,23 @@ public actor SparkConnectClient {
     return plan
   }
 
+  static func getWithWatermark(
+    _ child: Relation,
+    _ eventTime: String,
+    _ delayThreshold: String
+  ) -> Plan {
+    var withWatermark = Spark_Connect_WithWatermark()
+    withWatermark.input = child
+    withWatermark.eventTime = eventTime
+    withWatermark.delayThreshold = delayThreshold
+
+    var relation = Relation()
+    relation.withWatermark = withWatermark
+    var plan = Plan()
+    plan.opType = .root(relation)
+    return plan
+  }
+
   func createTempView(
     _ child: Relation, _ viewName: String, replace: Bool, isGlobal: Bool
   ) async throws {
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift b/Tests/SparkConnectTests/DataFrameTests.swift
index fd15496..09aab96 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -744,6 +744,22 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func withWatermark() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let df = try await spark
+      .sql("""
+        SELECT * FROM VALUES
+          (1, now()),
+          (1, now() - INTERVAL 1 HOUR),
+          (1, now() - INTERVAL 2 HOUR)
+          T(data, eventTime)
+        """)
+      .withWatermark("eventTime", "1 minute") // This tests only API for now
+    #expect(try await df.dropDuplicatesWithinWatermark("data").count() == 1)
+    await spark.stop()
+  }
+
   @Test
   func describe() async throws {
     let spark = try await SparkSession.builder.getOrCreate()


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org