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 e41476a [SPARK-57542] Support `nearestByJoin` for `DataFrame`
e41476a is described below
commit e41476ae2184e1b6764290cc756600b53f475a41
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Thu Jun 18 14:51:52 2026 -0700
[SPARK-57542] Support `nearestByJoin` for `DataFrame`
### What changes were proposed in this pull request?
This PR adds the public `DataFrame.nearestByJoin` API .
### Why are the changes needed?
`NearestByJoin` is a new relation in Apache Spark `4.2.0`. This keeps the
Swift
client at API parity with PySpark/Spark SQL.
### Does this PR introduce _any_ user-facing change?
No behavior change because this is a new method.
### 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 Opus 4.8
Closes #424 from dongjoon-hyun/SPARK-57542.
Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
.../SparkConnect/DataFrame+Transformations.swift | 34 ++++++++++++++++++++
Sources/SparkConnect/SparkConnectClient.swift | 16 ++++++++++
Sources/SparkConnect/TypeAliases.swift | 1 +
Tests/SparkConnectTests/DataFrameTests.swift | 36 ++++++++++++++++++++++
4 files changed, 87 insertions(+)
diff --git a/Sources/SparkConnect/DataFrame+Transformations.swift
b/Sources/SparkConnect/DataFrame+Transformations.swift
index 2f21d29..e4511eb 100644
--- a/Sources/SparkConnect/DataFrame+Transformations.swift
+++ b/Sources/SparkConnect/DataFrame+Transformations.swift
@@ -474,6 +474,40 @@ extension DataFrame {
return DataFrame(spark: self.spark, plan: plan)
}
+ /// Nearest-by top-K ranking join with another ``DataFrame``.
+ ///
+ /// For each row on the left (query) side, returns up to `numResults` rows
from `right`
+ /// (the base side), ranked by `rankingExpression`.
+ ///
+ /// - Parameters:
+ /// - right: The base ``DataFrame`` to search for matches.
+ /// - rankingExpression: A scalar expression string used to rank candidate
rows.
+ /// - numResults: Maximum number of matches per left row. Must be between
1 and 100000.
+ /// - mode: Search algorithm contract. One of `approx`, `exact`.
+ /// - direction: Ranking direction. One of `distance`, `similarity`.
+ /// - joinType: One of `inner` (default), `leftouter`.
+ /// - Returns: A ``DataFrame``.
+ public func nearestByJoin(
+ _ right: DataFrame,
+ rankingExpression: String,
+ numResults: Int32,
+ mode: String,
+ direction: String,
+ joinType: String = "inner"
+ ) async -> DataFrame {
+ let rightPlan = await (right.getPlan() as! Plan).root
+ let plan = SparkConnectClient.getNearestByJoin(
+ self.plan.root,
+ rightPlan,
+ rankingExpression: rankingExpression,
+ numResults: numResults,
+ mode: mode,
+ direction: direction,
+ joinType: joinType
+ )
+ return DataFrame(spark: self.spark, plan: plan)
+ }
+
// MARK: - Set Operations
/// Returns a new `DataFrame` containing rows in this `DataFrame` but not in
another `DataFrame`.
diff --git a/Sources/SparkConnect/SparkConnectClient.swift
b/Sources/SparkConnect/SparkConnectClient.swift
index d424a63..a68a465 100644
--- a/Sources/SparkConnect/SparkConnectClient.swift
+++ b/Sources/SparkConnect/SparkConnectClient.swift
@@ -1035,6 +1035,22 @@ public actor SparkConnectClient {
return createPlan { $0.lateralJoin = lateralJoin }
}
+ static func getNearestByJoin(
+ _ left: Relation, _ right: Relation,
+ rankingExpression: String, numResults: Int32,
+ mode: String, direction: String, joinType: String
+ ) -> Plan {
+ var nearestByJoin = NearestByJoin()
+ nearestByJoin.left = left
+ nearestByJoin.right = right
+ nearestByJoin.rankingExpression = rankingExpression.toExpression
+ nearestByJoin.numResults = numResults
+ nearestByJoin.mode = mode
+ nearestByJoin.direction = direction
+ nearestByJoin.joinType = joinType
+ return createPlan { $0.nearestByJoin = nearestByJoin }
+ }
+
static func getSetOperation(
_ left: Relation, _ right: Relation, _ opType: SetOpType, isAll: Bool =
false,
byName: Bool = false, allowMissingColumns: Bool = false
diff --git a/Sources/SparkConnect/TypeAliases.swift
b/Sources/SparkConnect/TypeAliases.swift
index d70d161..315d2e0 100644
--- a/Sources/SparkConnect/TypeAliases.swift
+++ b/Sources/SparkConnect/TypeAliases.swift
@@ -42,6 +42,7 @@ typealias MapType = Spark_Connect_DataType.Map
typealias MergeAction = Spark_Connect_MergeAction
typealias MergeIntoTableCommand = Spark_Connect_MergeIntoTableCommand
typealias NamedTable = Spark_Connect_Read.NamedTable
+typealias NearestByJoin = Spark_Connect_NearestByJoin
typealias OneOf_Analyze = AnalyzePlanRequest.OneOf_Analyze
typealias OneOf_CatType = Spark_Connect_Catalog.OneOf_CatType
typealias OutputType = Spark_Connect_OutputType
diff --git a/Tests/SparkConnectTests/DataFrameTests.swift
b/Tests/SparkConnectTests/DataFrameTests.swift
index 139c0ed..b54e459 100644
--- a/Tests/SparkConnectTests/DataFrameTests.swift
+++ b/Tests/SparkConnectTests/DataFrameTests.swift
@@ -630,6 +630,42 @@ struct DataFrameTests {
await spark.stop()
}
+ @Test
+ func nearestByJoin() async throws {
+ let spark = try await SparkSession.builder.getOrCreate()
+ let version = await spark.version
+ if version >= "4.2", !version.contains("preview") {
+ let users = try await spark.sql(
+ "SELECT * FROM VALUES (1, 10.0), (2, 20.0), (3, 30.0) AS T(user_id,
score)")
+ let products = try await spark.sql(
+ "SELECT * FROM VALUES ('A', 11.0), ('B', 22.0), ('C', 5.0) AS
S(product, pscore)")
+
+ let distance = await users.nearestByJoin(
+ products,
+ rankingExpression: "abs(score - pscore)",
+ numResults: 2,
+ mode: "approx",
+ direction: "distance"
+ ).select("user_id", "product").orderBy("user_id", "product")
+ #expect(
+ try await distance.collect() == [
+ Row(1, "A"), Row(1, "C"),
+ Row(2, "A"), Row(2, "B"),
+ Row(3, "A"), Row(3, "B"),
+ ])
+
+ let similarity = await users.nearestByJoin(
+ products,
+ rankingExpression: "-abs(score - pscore)",
+ numResults: 1,
+ mode: "approx",
+ direction: "similarity"
+ ).select("user_id", "product").orderBy("user_id")
+ #expect(try await similarity.collect() == [Row(1, "A"), Row(2, "B"),
Row(3, "B")])
+ }
+ await spark.stop()
+ }
+
@Test
func except() async throws {
let spark = try await SparkSession.builder.getOrCreate()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
