(spark-connect-swift) branch main updated: [SPARK-59416] Support `col` and subscript column accessors in `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 830b64e  [SPARK-59416] Support `col` and subscript column accessors in `DataFrame`
830b64e is described below

commit 830b64e6f708e0906e8a6ab6b4207bb1e4db5ccb
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Thu Sep 10 17:04:56 2026 -0700

    [SPARK-59416] Support `col` and subscript column accessors in `DataFrame`
    
    ### What changes were proposed in this pull request?
    
    This PR aims to support `col` and subscript column accessors in `DataFrame`.
    
    ```swift
    public nonisolated func col(_ colName: String) -> Column
    public nonisolated subscript(_ colName: String) -> Column
    ```
    
    Unlike the global `col(_:)` function, the returned `Column` carries the plan ID of the `DataFrame` (assigned since SPARK-59414), so the server can tell apart the columns with the same name in multiple `DataFrame`s.
    
    ```swift
    let df = try await spark.range(3)
    let df1 = await df.filter("id > 0")
    let df2 = await df.select("id")
    let joined = await df1.join(df2, joinExprs: df1["id"] == df2["id"])
    try await joined.select(df1["id"], df2.col("id")).show()
    ```
    
    - `"*"` becomes `UnresolvedStar` with `plan_id`.
    - A name ending with `".*"`, e.g. `"s.*"`, becomes `UnresolvedStar` with `unparsed_target` only, like the global `col(_:)`. The server rejects a star with both a target and a plan ID (`CONNECT_INVALID_PLAN.UNRESOLVED_STAR_WITH_BOTH_TARGET_AND_PLAN_ID`).
    - Other names become `UnresolvedAttribute` with `plan_id`.
    
    The plan ID of the root relation is stored as an immutable property when a `DataFrame` is created, so both accessors are `nonisolated` and don't require `await`.
    
    ### Why are the changes needed?
    
    To resolve ambiguous column references like `Dataset.col` and `Dataset.apply` in Scala and `DataFrame.__getitem__` in PySpark. Spark Connect has supported this since Apache Spark [v3.4.0](https://github.com/apache/spark/releases/tag/v3.4.0) (2023-04-07) via SPARK-41812 and SPARK-41823. No version gating is needed for the supported Apache Spark 4.x servers.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No, this adds only new public APIs, `DataFrame.col` and `DataFrame.subscript`.
    
    ### How was this patch tested?
    
    Pass the CIs with the newly added test cases.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Opus 5
    
    Closes #548 from dongjoon-hyun/SPARK-59416.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 .../SparkConnect/DataFrame+Transformations.swift   | 48 ++++++++++++++++++++++
 Sources/SparkConnect/DataFrame.swift               |  4 ++
 .../SparkConnectTests/DataFrameInternalTests.swift | 18 ++++++++
 Tests/SparkConnectTests/DataFrameTests.swift       | 32 +++++++++++++++
 4 files changed, 102 insertions(+)

diff --git a/Sources/SparkConnect/DataFrame+Transformations.swift b/Sources/SparkConnect/DataFrame+Transformations.swift
index d240883..11114f2 100644
--- a/Sources/SparkConnect/DataFrame+Transformations.swift
+++ b/Sources/SparkConnect/DataFrame+Transformations.swift
@@ -47,6 +47,54 @@ extension DataFrame {
     return DataFrame(spark: self.spark, plan: SparkConnectClient.getProject(self.plan.root, exprs))
   }
 
+  /// Selects a column based on the column name and returns it as a ``Column`` bound to this
+  /// ``DataFrame``.
+  ///
+  /// Unlike the global ``SparkConnect/col(_:)`` function, the returned ``Column`` remembers which
+  /// ``DataFrame`` it comes from. This lets the server tell apart the columns with the same name
+  /// in multiple ``DataFrame``s, e.g. both sides of a self-join.
+  ///
+  /// ```swift
+  /// let df = try await spark.range(3)
+  /// let df1 = await df.filter("id > 0")
+  /// let df2 = await df.select("id")
+  /// // `col("id") == col("id")` fails because `id` is ambiguous.
+  /// let joined = await df1.join(df2, joinExprs: df1.col("id") == df2.col("id"))
+  /// try await joined.select(df1.col("id")).show()
+  /// ```
+  ///
+  /// `"*"` selects all columns of this ``DataFrame``. A name ending with `".*"`, e.g. `"s.*"`,
+  /// expands the fields of a struct column like ``SparkConnect/col(_:)``, but it is not bound to
+  /// this ``DataFrame`` because the server doesn't allow a star with both a target and a plan ID.
+  /// - Parameter colName: A column name.
+  /// - Returns: A ``Column`` expression.
+  public nonisolated func col(_ colName: String) -> Column {
+    var expr = Spark_Connect_Expression()
+    if colName == "*" {
+      var star = Spark_Connect_Expression.UnresolvedStar()
+      star.planID = planID
+      expr.unresolvedStar = star
+    } else if colName.hasSuffix(".*") {
+      return Column(colName)
+    } else {
+      var attribute = colName.toUnresolvedAttribute
+      attribute.planID = planID
+      expr.unresolvedAttribute = attribute
+    }
+    return Column(expr)
+  }
+
+  /// Selects a column based on the column name and returns it as a ``Column`` bound to this
+  /// ``DataFrame``. This is an alias of ``col(_:)``.
+  ///
+  /// ```swift
+  /// let joined = await df1.join(df2, joinExprs: df1["id"] == df2["id"])
+  /// ```
+  /// - Parameter colName: A column name.
+  public nonisolated subscript(_ colName: String) -> Column {
+    col(colName)
+  }
+
   /// Selects column based on the column name specified as a regex and returns it as ``Column``.
   ///
   /// Only a column name enclosed in backticks is interpreted as a regex. The regex follows the
diff --git a/Sources/SparkConnect/DataFrame.swift b/Sources/SparkConnect/DataFrame.swift
index e39e686..9e7db1d 100644
--- a/Sources/SparkConnect/DataFrame.swift
+++ b/Sources/SparkConnect/DataFrame.swift
@@ -205,6 +205,9 @@ private let planIDGenerator = Atomic<Int64>(0)
 public actor DataFrame: Sendable {
   var spark: SparkSession
   var plan: Plan
+  /// The plan ID of the root relation, which binds the ``Column``s from ``col(_:)`` to this
+  /// `DataFrame`.
+  let planID: Int64
   var _schema: StructType? = nil
   var batches: [RecordBatch] = [RecordBatch]()
 
@@ -215,6 +218,7 @@ public actor DataFrame: Sendable {
   init(spark: SparkSession, plan: Plan) {
     self.spark = spark
     self.plan = DataFrame.withPlanID(plan)
+    self.planID = self.plan.root.common.planID
   }
 
   /// Create a new `DataFrame` instance with the given SparkSession and a SQL statement.
diff --git a/Tests/SparkConnectTests/DataFrameInternalTests.swift b/Tests/SparkConnectTests/DataFrameInternalTests.swift
index df5d25a..b31400c 100644
--- a/Tests/SparkConnectTests/DataFrameInternalTests.swift
+++ b/Tests/SparkConnectTests/DataFrameInternalTests.swift
@@ -123,6 +123,24 @@ struct DataFrameInternalTests {
     await spark.stop()
   }
 
+  @Test
+  func colExpression() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let df = try await spark.range(1)
+    let planID = await df.plan.root.common.planID
+    let attribute = df["id"].expr.unresolvedAttribute
+    #expect(attribute.unparsedIdentifier == "id")
+    #expect(attribute.planID == planID)
+    #expect(df.col("id").expr == df["id"].expr)
+    let star = df["*"].expr.unresolvedStar
+    #expect(!star.hasUnparsedTarget)
+    #expect(star.planID == planID)
+    let structStar = df["s.*"].expr.unresolvedStar
+    #expect(structStar.unparsedTarget == "s.*")
+    #expect(!structStar.hasPlanID)
+    await spark.stop()
+  }
+
   @Test
   func removeCachedRemoteRelation() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift b/Tests/SparkConnectTests/DataFrameTests.swift
index 1a81bee..2daae54 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -421,6 +421,20 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func dataFrameCol() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let df = try await spark.sql("SELECT 1 a, struct(2 b, 3 c) s")
+    #expect(try await df.select(df["a"], df.col("s")).columns == ["a", "s"])
+    #expect(try await df.select(df["*"]).columns == ["a", "s"])
+    #expect(try await df.select(df["s.*"]).columns == ["b", "c"])
+    // A column of an unrelated `DataFrame` cannot be resolved.
+    try await #require(throws: Error.self) {
+      try await spark.range(1).select(df["a"]).collect()
+    }
+    await spark.stop()
+  }
+
   @Test
   func colRegexWithPrefix() async throws {
     let spark = try await SparkSession.builder.getOrCreate()
@@ -939,6 +953,24 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func joinWithDataFrameColumn() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let df = try await spark.range(3)
+    let df1 = await df.filter("id > 0")
+    let df2 = await df.select("id")
+    // The global `col` function cannot tell which side `id` comes from.
+    try await #require(throws: Error.self) {
+      try await df1.join(df2, joinExprs: col("id") == col("id")).collect()
+    }
+    let joined = await df1.join(df2, joinExprs: df1["id"] == df2["id"])
+    #expect(
+      try await joined.select(df1["id"], df2.col("id")).orderBy(df1["id"]).collect()
+        == [Row(1, 1), Row(2, 2)])
+    #expect(try await joined.select(df2["*"]).columns == ["id"])
+    await spark.stop()
+  }
+
   @Test
   func alias() 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