SNOW-2268207: Support lateral join (#3971)

Author: sfc-gh-mayliu

CHANGELOG.md

@@ -6,6 +6,7 @@
 
 #### New Features
 
+- Added support for `DataFrame.lateral_join`
 - Added support for `Session.client_telemetry`.
 - Added support for `Session.udf_profiler`.
 - Added support for `functions.ai_translate`.

docs/source/snowpark/dataframe.rst

@@ -61,6 +61,7 @@ DataFrame
     DataFrame.intersect
     DataFrame.join
     DataFrame.join_table_function
+    DataFrame.lateral_join
     DataFrame.limit
     DataFrame.minus
     DataFrame.natural_join

src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py

@@ -26,6 +26,7 @@
     LeftAnti,
     LeftSemi,
     NaturalJoin,
+    LateralJoin,
     UsingJoin,
 )
 from snowflake.snowpark._internal.analyzer.datatype_mapper import (
@@ -169,6 +170,7 @@
 OUTER = " OUTER "
 RECURSIVE = " RECURSIVE "
 MODE = " MODE "
+INNER = " INNER "
 LATERAL = " LATERAL "
 PUT = " PUT "
 GET = " GET "
@@ -879,6 +881,65 @@ def asof_join_statement(
     )
 
 
+def lateral_join_statement(
+    left: str,
+    right: str,
+    join_condition: str,
+    use_constant_subquery_alias: bool,
+) -> str:
+    left_alias = (
+        "SNOWPARK_LEFT"
+        if use_constant_subquery_alias
+        else random_name_for_temp_object(TempObjectType.TABLE)
+    )
+    right_alias = (
+        "SNOWPARK_RIGHT"
+        if use_constant_subquery_alias
+        else random_name_for_temp_object(TempObjectType.TABLE)
+    )
+
+    # wrap the right side in subquery with WHERE clause if condition exists
+    if join_condition:
+        right_with_condition = (
+            LEFT_PARENTHESIS
+            + NEW_LINE
+            + SELECT
+            + STAR
+            + FROM
+            + LEFT_PARENTHESIS
+            + right
+            + RIGHT_PARENTHESIS
+            + WHERE
+            + join_condition
+            + NEW_LINE
+            + RIGHT_PARENTHESIS
+        )
+    else:
+        right_with_condition = LEFT_PARENTHESIS + right + RIGHT_PARENTHESIS
+
+    return (
+        SELECT
+        + STAR
+        + NEW_LINE
+        + FROM
+        + LEFT_PARENTHESIS
+        + NEW_LINE
+        + left
+        + NEW_LINE
+        + RIGHT_PARENTHESIS
+        + AS
+        + left_alias
+        + NEW_LINE
+        + INNER
+        + JOIN
+        + LATERAL
+        + NEW_LINE
+        + right_with_condition
+        + AS
+        + right_alias
+    )
+
+
 def snowflake_supported_join_statement(
     left: str,
     right: str,
@@ -983,6 +1044,10 @@ def join_statement(
         return asof_join_statement(
             left, right, join_condition, match_condition, use_constant_subquery_alias
         )
+    if isinstance(join_type, LateralJoin):
+        return lateral_join_statement(
+            left, right, join_condition, use_constant_subquery_alias
+        )
     if isinstance(join_type, UsingJoin) and isinstance(
         join_type.tpe, (LeftSemi, LeftAnti)
     ):

src/snowflake/snowpark/_internal/analyzer/binary_plan_node.py

@@ -137,6 +137,10 @@ class AsOf(JoinType):
     sql = "ASOF"
 
 
+class LateralJoin(JoinType):
+    sql = "INNER JOIN LATERAL"
+
+
 class NaturalJoin(JoinType):
     def __init__(self, tpe: JoinType) -> None:
         if not isinstance(