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

[do...@apache.org]

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

dongjoon-hyun 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 5cb7b48  [SPARK-58425] Support `observe` for `DataFrame`
5cb7b48 is described below

commit 5cb7b4836603d221d22584d8ab9ea079afd75b05
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Wed Jul 29 13:05:05 2026 -0700

    [SPARK-58425] Support `observe` for `DataFrame`
    
    ### What changes were proposed in this pull request?
    
    This PR aims to support `DataFrame.observe` API which defines named metrics to observe on a `DataFrame` by building a `CollectMetrics` relation.
    
    ```swift
    let observedDf = df.observe("my_metrics", sum(col("id")), max(col("id")))
    ```
    
    - `DataFrame.observe(_ name: String, _ expr: Column, _ exprs: Column...)` returns an *observed* `DataFrame` that returns the same result as the input. At least one metric column is required at compile time, like Scala's `Dataset.observe(name:expr:exprs:)`.
    - `SparkConnectClient.getCollectMetrics` assembles the `CollectMetrics` relation following the existing static `get*` helper pattern.
    
    Retrieving the observed metric values from `ExecutePlanResponse.observed_metrics` will be handled in a follow-up.
    
    ### Why are the changes needed?
    
    To provide a way to attach observable metrics (literals or aggregate expressions) to a `DataFrame` like Apache Spark's Scala/PySpark clients. This leverages the recently added `Column` type and aggregate functions ([SPARK-58336]).
    
    ### Does this PR introduce _any_ user-facing change?
    
    No. This is a new API addition.
    
    ### How was this patch tested?
    
    Pass the CIs.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Fable 5
    
    Closes #449 from dongjoon-hyun/SPARK-58425.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 Sources/SparkConnect/DataFrame+Transformations.swift | 19 +++++++++++++++++++
 Sources/SparkConnect/SparkConnectClient.swift        | 10 ++++++++++
 Sources/SparkConnect/TypeAliases.swift               |  1 +
 Tests/SparkConnectTests/DataFrameTests.swift         | 10 ++++++++++
 4 files changed, 40 insertions(+)

diff --git a/Sources/SparkConnect/DataFrame+Transformations.swift b/Sources/SparkConnect/DataFrame+Transformations.swift
index fdc5f7f..b024719 100644
--- a/Sources/SparkConnect/DataFrame+Transformations.swift
+++ b/Sources/SparkConnect/DataFrame+Transformations.swift
@@ -83,6 +83,25 @@ extension DataFrame {
       spark: self.spark, plan: SparkConnectClient.getProjectExprs(self.plan.root, exprs))
   }
 
+  /// Defines (named) metrics to observe on the ``DataFrame``. This method returns an 'observed'
+  /// ``DataFrame`` that returns the same result as the input, and computes the defined aggregates
+  /// (metrics) on all the data that is flowing through the ``DataFrame`` at that point.
+  /// Each metric ``Column`` must either contain a literal, or one or more aggregate functions.
+  ///
+  /// ```swift
+  /// let observedDf = df.observe("my_metrics", count(col("*")), max(col("id")))
+  /// ```
+  /// - Parameters:
+  ///   - name: The name of the observed metrics.
+  ///   - expr: A metric ``Column`` expression to compute.
+  ///   - exprs: Additional metric ``Column`` expressions.
+  /// - Returns: An observed ``DataFrame``.
+  public func observe(_ name: String, _ expr: Column, _ exprs: Column...) -> DataFrame {
+    let metrics = ([expr] + exprs).map { $0.expr }
+    return DataFrame(
+      spark: self.spark, plan: SparkConnectClient.getCollectMetrics(self.plan.root, name, metrics))
+  }
+
   /// Returns a new Dataset with a column dropped. This is a no-op if schema doesn't contain column name.
   /// - Parameter cols: Column names
   /// - Returns: A ``DataFrame`` with subset of columns.
diff --git a/Sources/SparkConnect/SparkConnectClient.swift b/Sources/SparkConnect/SparkConnectClient.swift
index 93ef80f..7789cae 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -566,6 +566,16 @@ public actor SparkConnectClient {
     return createPlan { $0.project = project }
   }
 
+  static func getCollectMetrics(
+    _ child: Relation, _ name: String, _ metrics: [Spark_Connect_Expression]
+  ) -> Plan {
+    var collectMetrics = CollectMetrics()
+    collectMetrics.input = child
+    collectMetrics.name = name
+    collectMetrics.metrics = metrics
+    return createPlan { $0.collectMetrics = collectMetrics }
+  }
+
   static func getWithColumnRenamed(_ child: Relation, _ colsMap: [String: String]) -> Plan {
     var withColumnsRenamed = WithColumnsRenamed()
     withColumnsRenamed.input = child
diff --git a/Sources/SparkConnect/TypeAliases.swift b/Sources/SparkConnect/TypeAliases.swift
index 315d2e0..f3027df 100644
--- a/Sources/SparkConnect/TypeAliases.swift
+++ b/Sources/SparkConnect/TypeAliases.swift
@@ -20,6 +20,7 @@ typealias ActionType = Spark_Connect_MergeAction.ActionType
 typealias Aggregate = Spark_Connect_Aggregate
 typealias AnalyzePlanRequest = Spark_Connect_AnalyzePlanRequest
 typealias AnalyzePlanResponse = Spark_Connect_AnalyzePlanResponse
+typealias CollectMetrics = Spark_Connect_CollectMetrics
 typealias Command = Spark_Connect_Command
 typealias ConfigRequest = Spark_Connect_ConfigRequest
 typealias DataSource = Spark_Connect_Read.DataSource
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift b/Tests/SparkConnectTests/DataFrameTests.swift
index f417d13..7a6ea90 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -282,6 +282,16 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func observe() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let df = try await spark.range(10).observe("my_metrics", sum(col("id")), max(col("id")))
+    #expect(try await df.columns == ["id"])
+    #expect(try await df.count() == 10)
+    #expect(try await df.collect().count == 10)
+    await spark.stop()
+  }
+
   @Test
   func withColumnRenamed() 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