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 65e60b2 [SPARK-58608] Support `Observation` to retrieve observed
metrics
65e60b2 is described below
commit 65e60b228e019f79e315023ef49715d88c9ae8e8
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Wed Aug 5 20:14:17 2026 -0700
[SPARK-58608] Support `Observation` to retrieve observed metrics
### What changes were proposed in this pull request?
This PR aims to support `Observation` to retrieve the observed metrics of
`DataFrame.observe`.
```swift
let observation = Observation("my_metrics")
let observedDf = try await df.observe(observation, count(col("*")),
max(col("id")))
try await observedDf.count()
let metrics = try await observation.get // ["count(1)": Int64(10),
"max(id)": ...]
```
- Add a new `Observation` actor whose `get` property returns the observed
metrics keyed by the metric column names, or throws
`SparkConnectError.invalidState` if no action has been performed yet.
- Add a `DataFrame.observe(_ observation: Observation, _ expr: Column, _
exprs: Column...)` overload which registers the observation to
`SparkConnectClient`.
- Collect `ExecutePlanResponse.observed_metrics` inside
`executePlanWithReattach` and deliver the values to the registered observations
by matching `ObservedMetrics.name`, like the PySpark Connect client. Since a
reattached stream resumes after `last_response_id`, no response is delivered
twice.
- Add an `ExpressionLiteral.toSwiftValue` helper converting metric value
literals back to Swift values.
### Why are the changes needed?
SPARK-58425 added `DataFrame.observe` which builds the `CollectMetrics`
relation, but the observed metric values returned by the server via
`ExecutePlanResponse.observed_metrics` were discarded, so users had no way to
read them. This provides the `Observation` API with the same usage as Scala and
PySpark.
### Does this PR introduce _any_ user-facing change?
No. This is a new feature addition.
### How was this patch tested?
Pass the CIs with a newly added test case.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Fable 5
Closes #484 from dongjoon-hyun/SPARK-58608.
Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../SparkConnect/DataFrame+Transformations.swift | 25 ++++++++
Sources/SparkConnect/Extension.swift | 23 +++++++
Sources/SparkConnect/Observation.swift | 70 ++++++++++++++++++++++
.../SparkConnectClient+ReattachableExecute.swift | 3 +
Sources/SparkConnect/SparkConnectClient.swift | 20 +++++++
Tests/SparkConnectTests/DataFrameTests.swift | 17 ++++++
6 files changed, 158 insertions(+)
diff --git a/Sources/SparkConnect/DataFrame+Transformations.swift
b/Sources/SparkConnect/DataFrame+Transformations.swift
index 0134181..ccb83c9 100644
--- a/Sources/SparkConnect/DataFrame+Transformations.swift
+++ b/Sources/SparkConnect/DataFrame+Transformations.swift
@@ -102,6 +102,31 @@ extension DataFrame {
spark: self.spark, plan:
SparkConnectClient.getCollectMetrics(self.plan.root, name, metrics))
}
+ /// Defines (named) metrics to observe on the ``DataFrame`` through an
``Observation`` instance.
+ /// The metric values are delivered to the ``Observation`` when an action is
performed on the
+ /// returned ``DataFrame``.
+ ///
+ /// ```swift
+ /// let observation = Observation("my_metrics")
+ /// let observedDf = try await df.observe(observation, count(col("*")),
max(col("id")))
+ /// try await observedDf.count()
+ /// let metrics = try await observation.get
+ /// ```
+ /// - Parameters:
+ /// - observation: An ``Observation`` instance to receive the observed
metrics.
+ /// - expr: A metric ``Column`` expression to compute.
+ /// - exprs: Additional metric ``Column`` expressions.
+ /// - Returns: An observed ``DataFrame``.
+ public func observe(_ observation: Observation, _ expr: Column, _ exprs:
Column...) async
+ -> DataFrame
+ {
+ await spark.client.addObservation(observation)
+ let metrics = ([expr] + exprs).map { $0.expr }
+ return DataFrame(
+ spark: self.spark,
+ plan: SparkConnectClient.getCollectMetrics(self.plan.root,
observation.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/Extension.swift
b/Sources/SparkConnect/Extension.swift
index f7e038f..7983ba7 100644
--- a/Sources/SparkConnect/Extension.swift
+++ b/Sources/SparkConnect/Extension.swift
@@ -210,6 +210,29 @@ extension ExpressionLiteral {
}
}
}
+
+ /// A Swift value converted from this `ExpressionLiteral`, or nil for null
and unsupported types.
+ var toSwiftValue: Sendable? {
+ switch self.literalType {
+ case .boolean(let value): value
+ case .byte(let value): Int8(value)
+ case .short(let value): Int16(value)
+ case .integer(let value): value
+ case .long(let value): value
+ case .float(let value): value
+ case .double(let value): value
+ case .decimal(let value): SwiftDecimal(string: value.value)
+ case .string(let value): value
+ case .binary(let value): value
+ // Date in units of days since the UNIX epoch.
+ case .date(let value): Date(timeIntervalSince1970: TimeInterval(value) *
86_400)
+ // Timestamp in units of microseconds since the UNIX epoch.
+ case .timestamp(let value): Date(timeIntervalSince1970:
TimeInterval(value) / 1_000_000)
+ case .timestampNtz(let value): Date(timeIntervalSince1970:
TimeInterval(value) / 1_000_000)
+ case .time(let value): LocalTime(nanoOfDay: value.nano)
+ default: nil
+ }
+ }
}
extension [String: String] {
diff --git a/Sources/SparkConnect/Observation.swift
b/Sources/SparkConnect/Observation.swift
new file mode 100644
index 0000000..b808227
--- /dev/null
+++ b/Sources/SparkConnect/Observation.swift
@@ -0,0 +1,70 @@
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+//
+
+#if canImport(FoundationEssentials)
+import FoundationEssentials
+#else
+import Foundation
+#endif
+
+/// A helper class to observe metrics of a ``DataFrame`` while an action is
performed on it.
+///
+/// ```swift
+/// let observation = Observation("my_metrics")
+/// let observedDf = try await df.observe(observation, count(col("*")),
max(col("id")))
+/// try await observedDf.count()
+/// let metrics = try await observation.get
+/// ```
+public actor Observation {
+ /// The name of the observation.
+ public nonisolated let name: String
+
+ private var values: [String: Sendable]? = nil
+
+ /// Create an `Observation` instance with a random name.
+ public init() {
+ self.name = UUID().uuidString
+ }
+
+ /// Create an `Observation` instance with the given name.
+ /// - Parameter name: The name of the observation.
+ public init(_ name: String) {
+ self.name = name
+ }
+
+ /// The observed metrics keyed by the metric column names. Metrics whose
observed values are
+ /// null are not included.
+ /// - Throws: `SparkConnectError.invalidState` if no action has been
performed on the observed
+ /// ``DataFrame`` yet.
+ public var get: [String: Sendable] {
+ get throws {
+ guard let values else {
+ throw SparkConnectError.invalidState(
+ SparkConnectError.Details(
+ message: "No metrics are observed yet. Perform an action on the
observed DataFrame."))
+ }
+ return values
+ }
+ }
+
+ /// Update the observed metric values delivered by an `ExecutePlanResponse`.
+ func setValues(_ values: [String: Sendable]) {
+ self.values = values
+ }
+}
diff --git a/Sources/SparkConnect/SparkConnectClient+ReattachableExecute.swift
b/Sources/SparkConnect/SparkConnectClient+ReattachableExecute.swift
index 2fc05a6..c1b3817 100644
--- a/Sources/SparkConnect/SparkConnectClient+ReattachableExecute.swift
+++ b/Sources/SparkConnect/SparkConnectClient+ReattachableExecute.swift
@@ -121,6 +121,9 @@ extension SparkConnectClient {
) async throws {
for try await m in response.messages {
try state.record(m)
+ if !m.observedMetrics.isEmpty {
+ await self.updateObservations(m.observedMetrics)
+ }
try await processResponse(m)
}
}
diff --git a/Sources/SparkConnect/SparkConnectClient.swift
b/Sources/SparkConnect/SparkConnectClient.swift
index caf8ca9..bff6214 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -857,6 +857,26 @@ public actor SparkConnectClient {
self.result.append(response)
}
+ private var observations: [String: Observation] = [:]
+
+ /// Register an ``Observation`` to receive the observed metrics of its name.
+ func addObservation(_ observation: Observation) {
+ observations[observation.name] = observation
+ }
+
+ /// Update the registered ``Observation``s with the observed metrics of an
+ /// `ExecutePlanResponse`.
+ func updateObservations(_ metrics: [ExecutePlanResponse.ObservedMetrics])
async {
+ for metric in metrics {
+ guard let observation = observations[metric.name] else { continue }
+ var values: [String: Sendable] = [:]
+ for (key, value) in zip(metric.keys, metric.values) {
+ values[key] = value.toSwiftValue
+ }
+ await observation.setValues(values)
+ }
+ }
+
@discardableResult
func execute(_ sessionID: String, _ command: Command) async throws ->
[ExecutePlanResponse] {
self.result.removeAll()
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift
b/Tests/SparkConnectTests/DataFrameTests.swift
index 3ee9e67..261b58b 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -367,6 +367,23 @@ struct DataFrameTests {
await spark.stop()
}
+ @Test
+ func observation() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let observation = Observation("my_metrics")
+ let df = try await spark.range(10)
+ .observe(observation, SparkConnect.count(col("id")), sum(col("id")),
max(col("id")))
+ await #expect(throws: SparkConnectError.self) {
+ try await observation.get
+ }
+ #expect(try await df.count() == 10)
+ let metrics = try await observation.get
+ #expect(metrics["count(id)"] as? Int64 == 10)
+ #expect(metrics["sum(id)"] as? Int64 == 45)
+ #expect(metrics["max(id)"] as? Int64 == 9)
+ await spark.stop()
+ }
+
@Test
func withColumnRenamed() async throws {
let spark = try await SparkSession.builder.getOrCreate()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
