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 9dbb31e [SPARK-58336] Support `Column` expression type and
`functions` API
9dbb31e is described below
commit 9dbb31e793e47a4010dfc372cc7f4e3b3fa55ac0
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Fri Jul 24 17:14:21 2026 -0700
[SPARK-58336] Support `Column` expression type and `functions` API
### What changes were proposed in this pull request?
This PR adds a `Column` expression type and a minimal set of `functions` to
`SparkConnect`
module, like Apache Spark's `org.apache.spark.sql.Column` and
`org.apache.spark.sql.functions`.
- `Column`: a `Sendable` struct wrapping `Spark_Connect_Expression`,
providing
`alias`, `asc`, `desc`, and `cast`.
- Functions: `col`, `column`, `lit`, `asc`, `desc`, `count`, `sum`, `avg`,
`mean`, `min`, `max`.
- New overloads accepting `Column` expressions: `DataFrame.select` and
`GroupedData.agg`.
```swift
let df2 = df.select(col("name"), col("age").cast("long").alias("age_long"))
let df3 = await df.groupBy("department")
.agg(count(col("*")).alias("employee_count"), avg(col("salary")))
```
In addition, the unused `Column` struct in `Catalog.swift` is renamed to
`CatalogColumn`
to free the name, and the `groupBy` documentation example is fixed.
### Why are the changes needed?
To provide a type-safe column expression DSL consistent with the other
Spark Connect
clients, instead of relying only on SQL string expressions.
### Does this PR introduce _any_ user-facing change?
Yes, the unused `Column` struct is renamed to `CatalogColumn`. The existing
APIs
behave identically, and the rest are API additions.
### How was this patch tested?
Pass the CIs including the new `FunctionsTests` suite.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Fable 5
Closes #441 from dongjoon-hyun/SPARK-58336.
Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
Sources/SparkConnect/Catalog.swift | 2 +-
Sources/SparkConnect/Column.swift | 101 +++++++++++
.../SparkConnect/DataFrame+Transformations.swift | 24 ++-
Sources/SparkConnect/Functions.swift | 184 +++++++++++++++++++++
Sources/SparkConnect/GroupedData.swift | 20 ++-
Sources/SparkConnect/SparkConnectClient.swift | 7 +
Tests/SparkConnectTests/FunctionsTests.swift | 128 ++++++++++++++
7 files changed, 459 insertions(+), 7 deletions(-)
diff --git a/Sources/SparkConnect/Catalog.swift
b/Sources/SparkConnect/Catalog.swift
index d846e4c..743515d 100644
--- a/Sources/SparkConnect/Catalog.swift
+++ b/Sources/SparkConnect/Catalog.swift
@@ -55,7 +55,7 @@ public struct SparkTable: Sendable, Equatable {
}
}
-public struct Column: Sendable, Equatable {
+public struct CatalogColumn: Sendable, Equatable {
public var name: String
public var description: String?
public var dataType: String
diff --git a/Sources/SparkConnect/Column.swift
b/Sources/SparkConnect/Column.swift
new file mode 100644
index 0000000..7bb76cc
--- /dev/null
+++ b/Sources/SparkConnect/Column.swift
@@ -0,0 +1,101 @@
+//
+// 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.
+//
+
+/// A column expression that can be used in ``DataFrame`` operations.
+///
+/// A ``Column`` wraps a Spark Connect `Expression` and is created via
+/// ``col(_:)``, `lit`, or other functions.
+///
+/// ```swift
+/// let df2 = df.select(col("name"), col("age").cast("long").alias("age_long"))
+/// ```
+public struct Column: Sendable {
+ let expr: Spark_Connect_Expression
+
+ init(_ expr: Spark_Connect_Expression) {
+ self.expr = expr
+ }
+
+ /// Creates a ``Column`` referencing the given column name.
+ ///
+ /// `"*"` and names ending with `".*"` are expanded as stars.
+ /// - Parameter name: A column name.
+ public init(_ name: String) {
+ var expr = Spark_Connect_Expression()
+ if name == "*" {
+ expr.unresolvedStar = Spark_Connect_Expression.UnresolvedStar()
+ } else if name.hasSuffix(".*") {
+ var star = Spark_Connect_Expression.UnresolvedStar()
+ star.unparsedTarget = name
+ expr.unresolvedStar = star
+ } else {
+ expr.unresolvedAttribute = name.toUnresolvedAttribute
+ }
+ self.expr = expr
+ }
+
+ /// Returns this column aliased with the given name.
+ /// - Parameter name: An alias name.
+ /// - Returns: A ``Column`` with the alias applied.
+ public func alias(_ name: String) -> Column {
+ var alias = Spark_Connect_Expression.Alias()
+ alias.expr = self.expr
+ alias.name = [name]
+ var expr = Spark_Connect_Expression()
+ expr.alias = alias
+ return Column(expr)
+ }
+
+ /// Returns a sort expression based on the ascending order of this column.
+ /// - Returns: A ``Column`` with ascending sort order.
+ public func asc() -> Column {
+ return sortOrder(.ascending, .sortNullsFirst)
+ }
+
+ /// Returns a sort expression based on the descending order of this column.
+ /// - Returns: A ``Column`` with descending sort order.
+ public func desc() -> Column {
+ return sortOrder(.descending, .sortNullsLast)
+ }
+
+ /// Casts this column to a different data type.
+ /// - Parameter to: A data type name like `"int"`, `"long"`, or `"string"`.
+ /// - Returns: A ``Column`` casted to the given type.
+ public func cast(_ to: String) -> Column {
+ var cast = Spark_Connect_Expression.Cast()
+ cast.expr = self.expr
+ cast.typeStr = to
+ var expr = Spark_Connect_Expression()
+ expr.cast = cast
+ return Column(expr)
+ }
+
+ private func sortOrder(
+ _ direction: Spark_Connect_Expression.SortOrder.SortDirection,
+ _ nullOrdering: Spark_Connect_Expression.SortOrder.NullOrdering
+ ) -> Column {
+ var sortOrder = Spark_Connect_Expression.SortOrder()
+ sortOrder.child = self.expr
+ sortOrder.direction = direction
+ sortOrder.nullOrdering = nullOrdering
+ var expr = Spark_Connect_Expression()
+ expr.sortOrder = sortOrder
+ return Column(expr)
+ }
+}
diff --git a/Sources/SparkConnect/DataFrame+Transformations.swift
b/Sources/SparkConnect/DataFrame+Transformations.swift
index e4511eb..fdc5f7f 100644
--- a/Sources/SparkConnect/DataFrame+Transformations.swift
+++ b/Sources/SparkConnect/DataFrame+Transformations.swift
@@ -33,6 +33,20 @@ extension DataFrame {
return DataFrame(spark: self.spark, plan:
SparkConnectClient.getProject(self.plan.root, cols))
}
+ /// Projects a set of ``Column`` expressions and returns a new ``DataFrame``.
+ ///
+ /// ```swift
+ /// let df2 = df.select(col("name"),
col("age").cast("long").alias("age_long"))
+ /// ```
+ /// - Parameters:
+ /// - col: A ``Column`` expression.
+ /// - cols: Additional ``Column`` expressions.
+ /// - Returns: A ``DataFrame`` with subset of columns.
+ public func select(_ col: Column, _ cols: Column...) -> DataFrame {
+ let exprs = ([col] + cols).map { $0.expr }
+ return DataFrame(spark: self.spark, plan:
SparkConnectClient.getProject(self.plan.root, exprs))
+ }
+
/// Selects a subset of existing columns using column names.
/// - Parameter cols: Column names
/// - Returns: A ``DataFrame`` with subset of columns.
@@ -781,14 +795,14 @@ extension DataFrame {
///
/// ```swift
/// // Group by single column
- /// let byDept = df.groupBy("department")
- /// .agg(count("*").alias("employee_count"))
+ /// let byDept = await df.groupBy("department")
+ /// .agg("count(*) AS employee_count")
///
/// // Group by multiple columns
- /// let byDeptAndLocation = df.groupBy("department", "location")
+ /// let byDeptAndLocation = await df.groupBy("department", "location")
/// .agg(
- /// avg("salary").alias("avg_salary"),
- /// max("salary").alias("max_salary")
+ /// "avg(salary) AS avg_salary",
+ /// "max(salary) AS max_salary"
/// )
/// ```
///
diff --git a/Sources/SparkConnect/Functions.swift
b/Sources/SparkConnect/Functions.swift
new file mode 100644
index 0000000..b05fd5a
--- /dev/null
+++ b/Sources/SparkConnect/Functions.swift
@@ -0,0 +1,184 @@
+//
+// 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.
+//
+
+/// Returns a ``Column`` based on the given column name.
+/// - Parameter name: A column name.
+/// - Returns: A ``Column``.
+public func col(_ name: String) -> Column {
+ return Column(name)
+}
+
+/// Returns a ``Column`` based on the given column name. This is an alias of
``col(_:)``.
+/// - Parameter name: A column name.
+/// - Returns: A ``Column``.
+public func column(_ name: String) -> Column {
+ return Column(name)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Bool) -> Column {
+ var literal = ExpressionLiteral()
+ literal.boolean = value
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Int8) -> Column {
+ var literal = ExpressionLiteral()
+ literal.byte = Int32(value)
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Int16) -> Column {
+ var literal = ExpressionLiteral()
+ literal.short = Int32(value)
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Int32) -> Column {
+ var literal = ExpressionLiteral()
+ literal.integer = value
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Int64) -> Column {
+ var literal = ExpressionLiteral()
+ literal.long = value
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Int) -> Column {
+ var literal = ExpressionLiteral()
+ literal.long = Int64(value)
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Float) -> Column {
+ var literal = ExpressionLiteral()
+ literal.float = value
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: Double) -> Column {
+ var literal = ExpressionLiteral()
+ literal.double = value
+ return litColumn(literal)
+}
+
+/// Creates a ``Column`` of literal value.
+/// - Parameter value: A literal value.
+/// - Returns: A ``Column``.
+public func lit(_ value: String) -> Column {
+ var literal = ExpressionLiteral()
+ literal.string = value
+ return litColumn(literal)
+}
+
+/// Returns a sort expression based on the ascending order of the given column
name.
+/// - Parameter name: A column name.
+/// - Returns: A ``Column`` with ascending sort order.
+public func asc(_ name: String) -> Column {
+ return Column(name).asc()
+}
+
+/// Returns a sort expression based on the descending order of the given
column name.
+/// - Parameter name: A column name.
+/// - Returns: A ``Column`` with descending sort order.
+public func desc(_ name: String) -> Column {
+ return Column(name).desc()
+}
+
+/// Returns the number of items in a group.
+/// - Parameter col: A ``Column`` to count.
+/// - Returns: A ``Column``.
+public func count(_ col: Column) -> Column {
+ return fn("count", col)
+}
+
+/// Returns the sum of all values in the expression.
+/// - Parameter col: A ``Column`` to aggregate.
+/// - Returns: A ``Column``.
+public func sum(_ col: Column) -> Column {
+ return fn("sum", col)
+}
+
+/// Returns the average of the values in a group.
+/// - Parameter col: A ``Column`` to aggregate.
+/// - Returns: A ``Column``.
+public func avg(_ col: Column) -> Column {
+ return fn("avg", col)
+}
+
+/// Returns the average of the values in a group. This is an alias of
``avg(_:)``.
+/// - Parameter col: A ``Column`` to aggregate.
+/// - Returns: A ``Column``.
+public func mean(_ col: Column) -> Column {
+ return fn("avg", col)
+}
+
+/// Returns the minimum value of the expression in a group.
+/// - Parameter col: A ``Column`` to aggregate.
+/// - Returns: A ``Column``.
+public func min(_ col: Column) -> Column {
+ return fn("min", col)
+}
+
+/// Returns the maximum value of the expression in a group.
+/// - Parameter col: A ``Column`` to aggregate.
+/// - Returns: A ``Column``.
+public func max(_ col: Column) -> Column {
+ return fn("max", col)
+}
+
+private func fn(_ name: String, _ args: Column...) -> Column {
+ var function = Spark_Connect_Expression.UnresolvedFunction()
+ function.functionName = name
+ function.arguments = args.map { $0.expr }
+ var expr = Spark_Connect_Expression()
+ expr.unresolvedFunction = function
+ return Column(expr)
+}
+
+private func litColumn(_ literal: ExpressionLiteral) -> Column {
+ var expr = Spark_Connect_Expression()
+ expr.literal = literal
+ return Column(expr)
+}
diff --git a/Sources/SparkConnect/GroupedData.swift
b/Sources/SparkConnect/GroupedData.swift
index 6a09fd5..b4a2467 100644
--- a/Sources/SparkConnect/GroupedData.swift
+++ b/Sources/SparkConnect/GroupedData.swift
@@ -29,11 +29,29 @@ public actor GroupedData {
}
public func agg(_ exprs: String...) async -> DataFrame {
+ return await buildAggregate(exprs.map { $0.toExpression })
+ }
+
+ /// Computes aggregates using the given ``Column`` expressions.
+ ///
+ /// ```swift
+ /// let df2 = await df.groupBy("department")
+ /// .agg(count(col("*")).alias("employee_count"), avg(col("salary")))
+ /// ```
+ /// - Parameters:
+ /// - expr: A ``Column`` expression.
+ /// - exprs: Additional ``Column`` expressions.
+ /// - Returns: A ``DataFrame``.
+ public func agg(_ expr: Column, _ exprs: Column...) async -> DataFrame {
+ return await buildAggregate(([expr] + exprs).map { $0.expr })
+ }
+
+ private func buildAggregate(_ exprs: [Spark_Connect_Expression]) async ->
DataFrame {
var aggregate = Aggregate()
aggregate.input = await (self.df.getPlan() as! Plan).root
aggregate.groupType = self.groupType
aggregate.groupingExpressions = self.groupingCols.map { $0.toExpression }
- aggregate.aggregateExpressions = exprs.map { $0.toExpression }
+ aggregate.aggregateExpressions = exprs
var relation = Relation()
relation.aggregate = aggregate
var plan = Plan()
diff --git a/Sources/SparkConnect/SparkConnectClient.swift
b/Sources/SparkConnect/SparkConnectClient.swift
index cb67026..aecbc1b 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -544,6 +544,13 @@ public actor SparkConnectClient {
return createPlan { $0.project = project }
}
+ static func getProject(_ child: Relation, _ exprs:
[Spark_Connect_Expression]) -> Plan {
+ var project = Project()
+ project.input = child
+ project.expressions = exprs
+ return createPlan { $0.project = project }
+ }
+
static func getToSchema(_ child: Relation, _ schema: Spark_Connect_DataType)
-> Plan {
var toSchema = Spark_Connect_ToSchema()
toSchema.input = child
diff --git a/Tests/SparkConnectTests/FunctionsTests.swift
b/Tests/SparkConnectTests/FunctionsTests.swift
new file mode 100644
index 0000000..e1440a0
--- /dev/null
+++ b/Tests/SparkConnectTests/FunctionsTests.swift
@@ -0,0 +1,128 @@
+//
+// 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.
+//
+
+import Testing
+
+@testable import SparkConnect
+
+/// A test suite for `Column` and `Functions`
+@Suite(.serialized)
+struct FunctionsTests {
+
+ @Test
+ func colFunction() throws {
+ let expr = col("id").expr
+ #expect(expr.unresolvedAttribute.unparsedIdentifier == "id")
+ }
+
+ @Test
+ func columnFunction() throws {
+ let expr = column("id").expr
+ #expect(expr.unresolvedAttribute.unparsedIdentifier == "id")
+ }
+
+ @Test
+ func star() throws {
+ let expr = col("*").expr
+ #expect(expr.unresolvedStar.hasUnparsedTarget == false)
+
+ let target = col("t.*").expr
+ #expect(target.unresolvedStar.unparsedTarget == "t.*")
+ }
+
+ @Test
+ func litFunction() throws {
+ #expect(lit(true).expr.literal.boolean == true)
+ #expect(lit(Int8(1)).expr.literal.byte == 1)
+ #expect(lit(Int16(1)).expr.literal.short == 1)
+ #expect(lit(Int32(1)).expr.literal.integer == 1)
+ #expect(lit(Int64(1)).expr.literal.long == 1)
+ #expect(lit(1).expr.literal.long == 1)
+ #expect(lit(Float(1.0)).expr.literal.float == 1.0)
+ #expect(lit(1.0).expr.literal.double == 1.0)
+ #expect(lit("a").expr.literal.string == "a")
+ }
+
+ @Test
+ func alias() throws {
+ let expr = col("id").alias("x").expr
+ #expect(expr.alias.name == ["x"])
+ #expect(expr.alias.expr.unresolvedAttribute.unparsedIdentifier == "id")
+ }
+
+ @Test
+ func ascFunction() throws {
+ let expr = asc("id").expr
+ #expect(expr.sortOrder.child.unresolvedAttribute.unparsedIdentifier ==
"id")
+ #expect(expr.sortOrder.direction == .ascending)
+ #expect(expr.sortOrder.nullOrdering == .sortNullsFirst)
+ }
+
+ @Test
+ func descFunction() throws {
+ let expr = desc("id").expr
+ #expect(expr.sortOrder.child.unresolvedAttribute.unparsedIdentifier ==
"id")
+ #expect(expr.sortOrder.direction == .descending)
+ #expect(expr.sortOrder.nullOrdering == .sortNullsLast)
+ }
+
+ @Test
+ func cast() throws {
+ let expr = col("id").cast("string").expr
+ #expect(expr.cast.expr.unresolvedAttribute.unparsedIdentifier == "id")
+ #expect(expr.cast.typeStr == "string")
+ }
+
+ @Test
+ func aggregateFunctions() throws {
+ for (aggColumn, name) in [
+ (count(col("id")), "count"),
+ (sum(col("id")), "sum"),
+ (avg(col("id")), "avg"),
+ (mean(col("id")), "avg"),
+ (min(col("id")), "min"),
+ (max(col("id")), "max"),
+ ] {
+ let expr = aggColumn.expr
+ #expect(expr.unresolvedFunction.functionName == name)
+ #expect(expr.unresolvedFunction.arguments.count == 1)
+ #expect(
+
expr.unresolvedFunction.arguments[0].unresolvedAttribute.unparsedIdentifier ==
"id")
+ }
+ }
+
+ @Test
+ func selectColumns() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let df = try await spark.range(3).select(col("id"),
col("id").cast("string").alias("id_string"))
+ #expect(try await df.columns == ["id", "id_string"])
+ #expect(try await df.count() == 3)
+ await spark.stop()
+ }
+
+ @Test
+ func aggColumns() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let rows = try await spark.range(3).groupBy("id")
+ .agg(count(col("*")).alias("cnt"), sum(col("id")), avg(col("id")))
+ .orderBy("id").collect()
+ #expect(rows == [Row(0, 1, 0, 0.0), Row(1, 1, 1, 1.0), Row(2, 1, 2, 2.0)])
+ await spark.stop()
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
