[SPARK-51911] Support lateralJoin in DataFrame

Sources/SparkConnect/DataFrame.swift

@@ -595,6 +595,82 @@ public actor DataFrame: Sendable {
     return DataFrame(spark: self.spark, plan: plan)
   }
 
+  /// Lateral join with another ``DataFrame``.
+  ///
+  /// Behaves as an JOIN LATERAL.
+  ///
+  /// - Parameters:
+  ///   - right: Right side of the join operation.
+  /// - Returns: A ``DataFrame``.
+  public func lateralJoin(_ right: DataFrame) async -> DataFrame {
+    let rightPlan = await (right.getPlan() as! Plan).root
+    let plan = SparkConnectClient.getLateralJoin(
+      self.plan.root,
+      rightPlan,
+      JoinType.inner
+    )
+    return DataFrame(spark: self.spark, plan: plan)
+  }
+
+  /// Lateral join with another ``DataFrame``.
+  ///
+  /// Behaves as an JOIN LATERAL.
+  ///
+  /// - Parameters:
+  ///   - right: Right side of the join operation.
+  ///   - joinType: One of `inner` (default), `cross`, `left`, `leftouter`, `left_outer`.
+  /// - Returns: A ``DataFrame``.
+  public func lateralJoin(_ right: DataFrame, joinType: String) async -> DataFrame {
+    let rightPlan = await (right.getPlan() as! Plan).root
+    let plan = SparkConnectClient.getLateralJoin(
+      self.plan.root,
+      rightPlan,
+      joinType.toJoinType
+    )
+    return DataFrame(spark: self.spark, plan: plan)
+  }
+
+  /// Lateral join with another ``DataFrame``.
+  ///
+  /// Behaves as an JOIN LATERAL.
+  ///
+  /// - Parameters:
+  ///   - right: Right side of the join operation.
+  ///   - joinExprs: A join expression string.
+  /// - Returns: A ``DataFrame``.
+  public func lateralJoin(_ right: DataFrame, joinExprs: String) async -> DataFrame {
+    let rightPlan = await (right.getPlan() as! Plan).root
+    let plan = SparkConnectClient.getLateralJoin(
+      self.plan.root,
+      rightPlan,
+      JoinType.inner,
+      joinCondition: joinExprs
+    )
+    return DataFrame(spark: self.spark, plan: plan)
+  }
+
+  /// Lateral join with another ``DataFrame``.
+  ///
+  /// Behaves as an JOIN LATERAL.
+  ///
+  /// - Parameters:
+  ///   - right: Right side of the join operation.
+  ///   - joinType: One of `inner` (default), `cross`, `left`, `leftouter`, `left_outer`.
+  ///   - joinExprs: A join expression string.
+  /// - Returns: A ``DataFrame``.
+  public func lateralJoin(
+    _ right: DataFrame, joinExprs: String, joinType: String = "inner"
+  ) async -> DataFrame {
+    let rightPlan = await (right.getPlan() as! Plan).root
+    let plan = SparkConnectClient.getLateralJoin(
+      self.plan.root,
+      rightPlan,
+      joinType.toJoinType,
+      joinCondition: joinExprs
+    )
+    return DataFrame(spark: self.spark, plan: plan)
+  }
+
   /// Returns a new `DataFrame` containing rows in this `DataFrame` but not in another `DataFrame`.
   /// This is equivalent to `EXCEPT DISTINCT` in SQL.
   /// - Parameter other: A `DataFrame` to exclude.

Sources/SparkConnect/SparkConnectClient.swift

@@ -590,6 +590,24 @@ public actor SparkConnectClient {
     return plan
   }
 
+  static func getLateralJoin(
+    _ left: Relation, _ right: Relation, _ joinType: JoinType,
+    joinCondition: String? = nil
+  ) -> Plan {
+    var lateralJoin = LateralJoin()
+    lateralJoin.left = left
+    lateralJoin.right = right
+    lateralJoin.joinType = joinType
+    if let joinCondition {
+      lateralJoin.joinCondition.expressionString = joinCondition.toExpressionString
+    }
+    var relation = Relation()
+    relation.lateralJoin = lateralJoin
+    var plan = Plan()
+    plan.opType = .root(relation)
+    return plan
+  }
+
   static func getSetOperation(
     _ left: Relation, _ right: Relation, _ opType: SetOpType, isAll: Bool = false,
     byName: Bool = false, allowMissingColumns: Bool = false

Sources/SparkConnect/TypeAliases.swift

@@ -34,6 +34,7 @@ typealias GroupType = Spark_Connect_Aggregate.GroupType
 typealias Join = Spark_Connect_Join
 typealias JoinType = Spark_Connect_Join.JoinType
 typealias KeyValue = Spark_Connect_KeyValue
+typealias LateralJoin = Spark_Connect_LateralJoin
 typealias Limit = Spark_Connect_Limit
 typealias MapType = Spark_Connect_DataType.Map
 typealias NamedTable = Spark_Connect_Read.NamedTable

Tests/SparkConnectTests/DataFrameTests.swift

@@ -466,6 +466,26 @@ struct DataFrameTests {
     await spark.stop()
   }
 
+  @Test
+  func lateralJoin() async throws {
+    let spark = try await SparkSession.builder.getOrCreate()
+    let df1 = try await spark.sql("SELECT * FROM VALUES ('a', '1'), ('b', '2') AS T(a, b)")
+    let df2 = try await spark.sql("SELECT * FROM VALUES ('c', '2'), ('d', '3') AS S(c, b)")
+    let expectedCross = [
+      Row("a", "1", "c", "2"),
+      Row("a", "1", "d", "3"),
+      Row("b", "2", "c", "2"),
+      Row("b", "2", "d", "3"),
+    ]
+    #expect(try await df1.lateralJoin(df2).collect() == expectedCross)
+    #expect(try await df1.lateralJoin(df2, joinType: "inner").collect() == expectedCross)
+
+    let expected = [Row("b", "2", "c", "2")]
+    #expect(try await df1.lateralJoin(df2, joinExprs: "T.b = S.b").collect() == expected)
+    #expect(try await df1.lateralJoin(df2, joinExprs: "T.b = S.b", joinType: "inner").collect() == expected)
+    await spark.stop()
+  }
+
   @Test
   func except() async throws {
     let spark = try await SparkSession.builder.getOrCreate()