propagate through plan

Author: sfc-gh-joshi

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

@@ -996,24 +996,30 @@ def snowflake_supported_join_statement(
     left_uuid: Optional[str] = None,
     right_uuid: Optional[str] = None,
     directed: bool = False,
+    left_is_join: bool = False,
 ) -> str:
     LEFT_UUID = format_uuid(left_uuid)
     RIGHT_UUID = format_uuid(right_uuid)
 
-    # If left is a simple SELECT * FROM (\n<join_source>\n) wrapper from a
-    # previous join, flatten into a multi-way join by appending the new right
-    # operand directly to the existing join source. This avoids nested
+    # If left is the output of a previous join, flatten into a multi-way join
+    # by unwrapping the SELECT * FROM (...) envelope and appending the new
+    # right operand directly to the existing join source. This avoids nested
     # SELECT * layers that inflate query text without changing semantics.
     #
     # Though it is technically less efficient than constructing the join sub-queries
     # without the SELECT in the first place, the structure of our SQL processing code
-    # top-level projections to be wrapped by a select.
-    unwrapped_left = _unwrap_select_star_from(left)
+    # needs top-level projections to be wrapped by a select to be well-formed, so we
+    # must strip it here instead.
+    #
+    # We only unwrap the left side because it is simpler to deal with than unwrapping
+    # both left and right, and left-deep chains are more common, as they're produced
+    # by calls like df1.join(df2).join(df3) etc.
+    unwrapped_left = _unwrap_select_star_from(left) if left_is_join else None
     right_alias = (
         "SNOWPARK_RIGHT"
+        if use_constant_subquery_alias and unwrapped_left is None
         # Multi-way join: right alias must be unique to avoid collisions
         # with aliases already present in the flattened join source.
-        if use_constant_subquery_alias and unwrapped_left is None
         else random_name_for_temp_object(TempObjectType.TABLE)
     )
 
@@ -1105,6 +1111,7 @@ def join_statement(
     left_uuid: Optional[str] = None,
     right_uuid: Optional[str] = None,
     directed: bool = False,
+    left_is_join: bool = False,
 ) -> str:
     if isinstance(join_type, (LeftSemi, LeftAnti)):
         return left_semi_or_anti_join_statement(
@@ -1132,6 +1139,7 @@ def join_statement(
         left_uuid=left_uuid,
         right_uuid=right_uuid,
         directed=directed,
+        left_is_join=left_is_join,
     )
 
 

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

@@ -451,6 +451,13 @@ def get_snowflake_plan(self, skip_schema_query) -> SnowflakePlan:
                 # Add the last df ast id to the snowflake plan as the most recent
                 # dataframe operation to create this plan.
                 self._snowflake_plan.df_ast_ids = self.df_ast_ids
+            # Propagate join output flag through passthrough SelectStatements
+            # so chained joins can flatten into multi-way joins.
+            if not self.has_clause and not self.has_projection:
+                if isinstance(self.from_, SelectSnowflakePlan):
+                    self._snowflake_plan._is_join_output = (
+                        self.from_._snowflake_plan._is_join_output
+                    )
         return self._snowflake_plan
 
     @property

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

@@ -464,6 +464,7 @@ def __init__(
         self.session = session
         self.source_plan = source_plan
         self.is_ddl_on_temp_object = is_ddl_on_temp_object
+        self._is_join_output = False
         # We need to copy this list since we don't want to change it for the
         # previous SnowflakePlan objects
         self.api_calls = api_calls.copy() if api_calls else []
@@ -1237,7 +1238,8 @@ def join(
         use_constant_subquery_alias: bool,
         directed: bool = False,
     ):
-        return self.build_binary(
+        left_is_join = left._is_join_output
+        result = self.build_binary(
             lambda x, y: join_statement(
                 x,
                 y,
@@ -1254,11 +1256,14 @@ def join(
                     else None
                 ),
                 directed=directed,
+                left_is_join=left_is_join,
             ),
             left,
             right,
             source_plan,
         )
+        result._is_join_output = True
+        return result
 
     def save_as_table(
         self,

tests/unit/test_analyzer_util_suite.py

@@ -465,14 +465,16 @@ def test_join_statement_flattens_chained_joins():
         True,
     )
 
-    # Second join uses first join's output as left operand
+    # Second join uses first join's output as left operand.
+    # left_is_join=True signals that the left operand is a join result.
     second_join = join_statement(
         first_join,
         "SELECT * FROM table_c",
         join_type,
         "",
         "",
         True,
+        left_is_join=True,
     )
 
     # Should NOT have nested SELECT * FROM (SELECT * FROM (...))
@@ -485,6 +487,45 @@ def test_join_statement_flattens_chained_joins():
     assert "table_c" in second_join
 
 
+def test_join_statement_does_not_flatten_user_generated_select_star():
+    """A user-generated SELECT * that coincidentally matches the internal
+    pattern must NOT be flattened when left_is_join is False (the default)."""
+    join_type = UsingJoin(Inner(), ["key"])
+
+    # Craft a SQL string that matches the internal _SELECT_STAR_FROM_PREFIX/SUFFIX
+    # pattern exactly — this simulates what session.sql("SELECT * FROM (...)") or
+    # a similar user-provided query might produce.
+    user_sql = (
+        " SELECT  * \n FROM (\n"
+        "(SELECT id, key FROM user_table) AS t1"
+        " INNER  JOIN \n(SELECT id, key FROM other_table) AS t2\n"
+        " USING (key)"
+        "\n)"
+    )
+
+    # join_statement with left_is_join=False (default) — should NOT unwrap because
+    # left_is_join is only True when a plan object is constructed from a dataframe
+    # join operation
+    result = join_statement(
+        user_sql,
+        "SELECT * FROM table_c",
+        join_type,
+        "",
+        "",
+        True,
+        left_is_join=False,
+    )
+
+    # The user SQL should be preserved as a nested subquery, producing
+    # two SELECT levels (the outer wrapper + the user's original SELECT *)
+    assert result.count(" SELECT ") >= 2
+
+    # The user's original SQL should appear within the output intact
+    assert "user_table" in result
+    assert "other_table" in result
+    assert "table_c" in result
+
+
 def test_create_iceberg_table_statement():
     assert create_table_statement(
         table_name="test_table",