fix

Author: sfc-gh-aling

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

@@ -3,7 +3,7 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 import uuid
-from collections import Counter, defaultdict
+from collections import defaultdict
 from typing import TYPE_CHECKING, DefaultDict, Dict, List, Optional, Union
 
 import snowflake.snowpark
@@ -151,7 +151,10 @@
 )
 from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 from snowflake.snowpark._internal.telemetry import TelemetryField
-from snowflake.snowpark._internal.utils import quote_name
+from snowflake.snowpark._internal.utils import (
+    quote_name,
+    merge_multiple_snowflake_plan_expr_to_alias,
+)
 from snowflake.snowpark.types import BooleanType, _NumericType
 
 ARRAY_BIND_THRESHOLD = 512
@@ -804,22 +807,9 @@ def do_resolve(self, logical_plan: LogicalPlan) -> SnowflakePlan:
             # Selectable doesn't have children. It already has the expr_to_alias dict.
             self.alias_maps_to_use = logical_plan.expr_to_alias.copy()
         else:
-            use_maps = {}
-            # get counts of expr_to_alias keys
-            counts = Counter()
-            for v in resolved_children.values():
-                if v.expr_to_alias:
-                    counts.update(list(v.expr_to_alias.keys()))
-
-            # Keep only non-shared expr_to_alias keys
-            # let (df1.join(df2)).join(df2.join(df3)).select(df2) report error
-            for v in resolved_children.values():
-                if v.expr_to_alias:
-                    use_maps.update(
-                        {p: q for p, q in v.expr_to_alias.items() if counts[p] < 2}
-                    )
-
-            self.alias_maps_to_use = use_maps
+            self.alias_maps_to_use = merge_multiple_snowflake_plan_expr_to_alias(
+                list(resolved_children.values())
+            )
 
         res = self.do_resolve_with_resolved_children(
             logical_plan, resolved_children, df_aliased_col_name_to_real_col_name

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

@@ -108,6 +108,7 @@
     generate_random_alphanumeric,
     get_copy_into_table_options,
     is_sql_select_statement,
+    merge_multiple_snowflake_plan_expr_to_alias,
 )
 from snowflake.snowpark.row import Row
 from snowflake.snowpark.types import StructType
@@ -580,17 +581,10 @@ def build_binary(
             right_schema_query = schema_value_statement(select_right.attributes)
             schema_query = sql_generator(left_schema_query, right_schema_query)
 
-        common_columns = set(select_left.expr_to_alias.keys()).intersection(
-            select_right.expr_to_alias.keys()
+        new_expr_to_alias = merge_multiple_snowflake_plan_expr_to_alias(
+            [select_left, select_right]
         )
-        new_expr_to_alias = {
-            k: v
-            for k, v in {
-                **select_left.expr_to_alias,
-                **select_right.expr_to_alias,
-            }.items()
-            if k not in common_columns
-        }
+
         api_calls = [*select_left.api_calls, *select_right.api_calls]
 
         # Need to do a deduplication to avoid repeated query.

src/snowflake/snowpark/_internal/utils.py

@@ -21,6 +21,7 @@
 import sys
 import threading
 import traceback
+import uuid
 import zipfile
 from enum import Enum
 from functools import lru_cache
@@ -50,12 +51,15 @@
 from snowflake.connector.description import OPERATING_SYSTEM, PLATFORM
 from snowflake.connector.options import MissingOptionalDependency, ModuleLikeObject
 from snowflake.connector.version import VERSION as connector_version
+
 from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 from snowflake.snowpark.context import _should_use_structured_type_semantics
 from snowflake.snowpark.row import Row
 from snowflake.snowpark.version import VERSION as snowpark_version
 
 if TYPE_CHECKING:
+    from snowflake.snowpark._internal.analyzer.snowflake_plan import SnowflakePlan
+
     try:
         from snowflake.connector.cursor import ResultMetadataV2
     except ImportError:
@@ -1411,3 +1415,59 @@ def next(self) -> int:
 
 
 global_counter: GlobalCounter = GlobalCounter()
+
+
+def merge_multiple_snowflake_plan_expr_to_alias(
+    snowflake_plans: List["SnowflakePlan"],
+) -> Dict[uuid.UUID, str]:
+    """
+    Merges expression-to-alias mappings from multiple Snowflake plans, resolving conflicts where possible.
+
+    Args:
+        snowflake_plans (List[SnowflakePlan]): List of SnowflakePlan objects.
+
+    Returns:
+        Dict[Any, str]: Merged expression-to-alias mapping.
+    """
+    # Collect all unique output column names from all plans
+    all_output_columns = {attr.name for plan in snowflake_plans for attr in plan.output}
+
+    # Gather all expression-to-alias mappings
+    all_expr_to_alias_dicts = [plan.expr_to_alias for plan in snowflake_plans]
+
+    # Initialize the merged dictionary
+    merged_dict = {}
+
+    # Collect all unique keys from all dictionaries
+    all_keys = set().union(*all_expr_to_alias_dicts)
+
+    for key in all_keys:
+        # Gather all aliases for the current key
+        values = [d[key] for d in all_expr_to_alias_dicts if key in d]
+
+        # Check if all aliases are identical
+        if len(set(values)) == 1:
+            merged_dict[key] = values[0]
+        else:
+            # Resolve conflicts by checking against output columns
+            candidate = None
+            for alias in values:
+                if alias in all_output_columns:
+                    if candidate is None:
+                        candidate = alias
+                    else:
+                        # Ambiguous case: multiple valid candidates
+                        candidate = None
+                        break
+
+            # Add the candidate to the merged dictionary if resolved
+            if candidate is not None:
+                merged_dict[key] = candidate
+            else:
+                # No valid candidate found
+                _logger.debug(
+                    f"Expression '{key}' is associated with multiple aliases across different plans. "
+                    f"Unable to determine which alias to use. Conflicting values: {values}"
+                )
+
+    return merged_dict

tests/integ/scala/test_dataframe_join_suite.py

@@ -1462,22 +1462,17 @@ def test_select_columns_on_join_result_with_conflict_name(
     "config.getoption('local_testing_mode', default=False)",
     reason="SNOW-1235716: match error behavior",
 )
-def test_nested_join_diamond_shape_error(
+def test_nested_join_diamond_shape(
     session,
 ):  # TODO: local testing match error behavior
-    """This is supposed to work but currently we don't handle it correctly. We should fix this with a good design."""
     df1 = session.create_dataframe([[1]], schema=["a"])
     df2 = session.create_dataframe([[1]], schema=["a"])
     df3 = df1.join(df2, df1["a"] == df2["a"])
     df4 = df3.select(df1["a"].as_("a"))
     # df1["a"] and df4["a"] has the same expr_id in map expr_to_alias. When they join, only one will be in df5's alias
     # map. It leaves the other one resolved to "a" instead of the alias.
     df5 = df1.join(df4, df1["a"] == df4["a"])  # (df1) JOIN ((df1 JOIN df2)->df4)
-    with pytest.raises(
-        SnowparkSQLAmbiguousJoinException,
-        match="The reference to the column 'A' is ambiguous.",
-    ):
-        df5.collect()
+    Utils.check_answer(df5, [Row(1, 1)])
 
 
 def test_nested_join_diamond_shape_workaround(session):