snowpark-connect 0.23.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

snowflake/snowpark_connect/relation/map_join.py

@@ -1,13 +1,18 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
+from collections import Counter
 from functools import reduce
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+from pyspark.errors.exceptions.base import AnalysisException
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+    unquote_if_quoted,
+)
 from snowflake.snowpark_connect.column_name_handler import JoinColumnNameMap
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.constants import COLUMN_METADATA_COLLISION_KEY
@@ -17,6 +22,7 @@ from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
 )
 from snowflake.snowpark_connect.expression.typer import JoinExpressionTyper
+from snowflake.snowpark_connect.hidden_column import HiddenColumn
 from snowflake.snowpark_connect.relation.map_relation import (
     NATURAL_JOIN_TYPE_BASE,
     map_relation,
@@ -24,7 +30,6 @@ from snowflake.snowpark_connect.relation.map_relation import (
 from snowflake.snowpark_connect.utils.context import (
     push_evaluating_join_condition,
     push_sql_scope,
-    set_plan_id_map,
     set_sql_plan_name,
 )
 from snowflake.snowpark_connect.utils.telemetry import (
@@ -33,6 +38,9 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 USING_COLUMN_NOT_FOUND_ERROR = "[UNRESOLVED_USING_COLUMN_FOR_JOIN] USING column `{0}` not found on the {1} side of the join. The {1}-side columns: {2}"
 
+DUPLICATED_JOIN_COL_LSUFFIX = "_left"
+DUPLICATED_JOIN_COL_RSUFFIX = "_right"
+
 
 def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
     left_container: DataFrameContainer = map_relation(rel.join.left)
@@ -74,6 +82,13 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
 
     # This handles case sensitivity for using_columns
     case_corrected_right_columns: list[str] = []
+    hidden_columns = set()
+    # Propagate the hidden columns from left/right inputs to the result in case of chained joins
+    if left_container.column_map.hidden_columns:
+        hidden_columns.update(left_container.column_map.hidden_columns)
+
+    if right_container.column_map.hidden_columns:
+        hidden_columns.update(right_container.column_map.hidden_columns)
 
     if rel.join.HasField("join_condition"):
         assert not using_columns
@@ -105,8 +120,8 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
             right=right_input,
             on=join_expression.col,
             how=join_type,
-            lsuffix="_left",
-            rsuffix="_right",
+            lsuffix=DUPLICATED_JOIN_COL_LSUFFIX,
+            rsuffix=DUPLICATED_JOIN_COL_RSUFFIX,
         )
     elif using_columns:
         if any(
@@ -156,12 +171,24 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
                 )
             )
 
+        using_columns_snowpark_names = (
+            left_container.column_map.get_snowpark_column_names_from_spark_column_names(
+                list(using_columns), return_first=True
+            )
+        )
+
+        using_columns_snowpark_types = [
+            left_container.dataframe.schema.fields[idx].datatype
+            for idx, col in enumerate(left_container.column_map.get_snowpark_columns())
+            if col in using_columns_snowpark_names
+        ]
+
         # Round trip the using columns through the column map to get the correct names
         # in order to support case sensitivity.
         # TODO: case_corrected_left_columns / case_corrected_right_columns may no longer be required as Snowpark dataframe preserves the column casing now.
-        case_corrected_left_columns = left_container.column_map.get_spark_column_names_from_snowpark_column_names(
-            left_container.column_map.get_snowpark_column_names_from_spark_column_names(
-                list(using_columns), return_first=True
+        case_corrected_left_columns = (
+            left_container.column_map.get_spark_column_names_from_snowpark_column_names(
+                using_columns_snowpark_names
             )
         )
         case_corrected_right_columns = right_container.column_map.get_spark_column_names_from_snowpark_column_names(
@@ -195,28 +222,141 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
                 (left == right for left, right in snowpark_using_columns),
             ),
             how=join_type,
+            rsuffix=DUPLICATED_JOIN_COL_RSUFFIX,
         )
+        # If we disambiguated the snowpark_using_columns during the join, we need to update 'snowpark_using_columns' to
+        # use the disambiguated names.
+        disambiguated_snowpark_using_columns = []
+
+        # Ignore disambiguation for LEFT SEMI JOIN and LEFT ANTI JOIN because they drop the right columns, so it'll never disambiguate.
+        if join_type in ["leftsemi", "leftanti"]:
+            disambiguated_snowpark_using_columns = snowpark_using_columns
+        else:
+            normalized_joined_columns = [
+                unquote_if_quoted(col) for col in joined_df.columns
+            ]
+            # snowpark_using_columns is a list of tuples of snowpark columns, joined_df.columns is a list of strings of column names
+            for (left, right) in snowpark_using_columns:
+                normalized_left_name = unquote_if_quoted(left.getName())
+                normalized_right_name = unquote_if_quoted(right.getName())
+
+                # are both left and right in joined_df? if not, it's been disambiguated
+                if (
+                    normalized_left_name in normalized_joined_columns
+                    and normalized_right_name in normalized_joined_columns
+                ):
+                    # we want to just add this
+                    disambiguated_snowpark_using_columns.append((left, right))
+                else:
+                    # we need to figure out the disambiguated names and add those - it only disambiguates if left == right
+                    disambiguated_left: snowpark.Column | None = None
+                    disambiguated_right: snowpark.Column | None = None
+
+                    for col in normalized_joined_columns:
+                        quoted_col = f'"{col}"'
+                        # get the column name and cross check it to see if it ends with the og name
+                        if col.endswith(normalized_left_name) and col.startswith("l_"):
+                            disambiguated_left = joined_df[quoted_col]
+                        elif col.endswith(normalized_right_name) and col.startswith(
+                            "r_"
+                        ):
+                            disambiguated_right = joined_df[quoted_col]
+
+                        # If we have both disambiguated columns, we can break out of the loop to save processing time
+                        if (
+                            disambiguated_left is not None
+                            and disambiguated_right is not None
+                        ):
+                            break
+                    if disambiguated_left is None or disambiguated_right is None:
+                        raise AnalysisException(
+                            f"Disambiguated columns not found for {normalized_left_name} and {normalized_right_name}."
+                        )
+                    disambiguated_snowpark_using_columns.append(
+                        (disambiguated_left, disambiguated_right)
+                    )
+
         # For outer joins, we need to preserve join keys from both sides using COALESCE
+        """
+        CHANGES:
+            - IF CASE
+                - Need to drop the using columns
+                - Need to create the hidden_columns DF with the using columns from right and left
+            - ELSE CASE
+                - Need to drop the right side using columns
+                - Need to create the hidden_columns DF with the using columns from right
+        """
         if join_type == "full_outer":
             coalesced_columns = []
-            columns_to_drop = []
-            for i, (left_col, right_col) in enumerate(snowpark_using_columns):
+            for i, (left_col, _right_col) in enumerate(snowpark_using_columns):
                 # Use the original user-specified column name to preserve case sensitivity
-                original_column_name = rel.join.using_columns[i]
-                coalesced_col = snowpark_fn.coalesce(left_col, right_col).alias(
-                    original_column_name
-                )
+                # Use the disambiguated columns for coalescing
+                disambiguated_left_col = disambiguated_snowpark_using_columns[i][0]
+                disambiguated_right_col = disambiguated_snowpark_using_columns[i][1]
+
+                coalesced_col = snowpark_fn.coalesce(
+                    disambiguated_left_col, disambiguated_right_col
+                ).alias(left_col.get_name())
                 coalesced_columns.append(coalesced_col)
-                columns_to_drop.extend([left_col, right_col])
 
+                # Create HiddenColumn objects for each hidden column
+                hidden_left = HiddenColumn(
+                    hidden_snowpark_name=disambiguated_left_col.getName(),
+                    spark_name=case_corrected_left_columns[i],
+                    visible_snowpark_name=left_col.get_name(),
+                    qualifiers=left_container.column_map.get_qualifier_for_spark_column(
+                        case_corrected_left_columns[i]
+                    ),
+                    original_position=left_container.column_map.get_spark_columns().index(
+                        case_corrected_left_columns[i]
+                    ),
+                )
+
+                hidden_right = HiddenColumn(
+                    hidden_snowpark_name=disambiguated_right_col.getName(),
+                    spark_name=case_corrected_right_columns[i],
+                    visible_snowpark_name=left_col.get_name(),
+                    qualifiers=right_container.column_map.get_qualifier_for_spark_column(
+                        case_corrected_right_columns[i]
+                    ),
+                    original_position=right_container.column_map.get_spark_columns().index(
+                        case_corrected_right_columns[i]
+                    ),
+                )
+                hidden_columns.update(
+                    [
+                        hidden_left,
+                        hidden_right,
+                    ]
+                )
+
+            # All non-hidden columns (not including the coalesced columns)
             other_columns = [
                 snowpark_fn.col(col_name)
                 for col_name in joined_df.columns
-                if col_name not in [col.getName() for col in columns_to_drop]
+                if col_name not in [col.hidden_snowpark_name for col in hidden_columns]
             ]
             result = joined_df.select(coalesced_columns + other_columns)
+
         else:
             result = joined_df.drop(*(right for _, right in snowpark_using_columns))
+            # We never run into the disambiguation case unless it's a full outer join.
+            for i, (left_col, right_col) in enumerate(
+                disambiguated_snowpark_using_columns
+            ):
+                # Only right side columns are hidden
+                hidden_col = HiddenColumn(
+                    hidden_snowpark_name=right_col.getName(),
+                    spark_name=case_corrected_right_columns[i],
+                    visible_snowpark_name=left_col.getName(),
+                    qualifiers=right_container.column_map.get_qualifier_for_spark_column(
+                        case_corrected_right_columns[i]
+                    ),
+                    original_position=right_container.column_map.get_spark_columns().index(
+                        case_corrected_right_columns[i]
+                    ),
+                )
+                hidden_columns.add(hidden_col)
     else:
         if join_type != "cross" and not global_config.spark_sql_crossJoin_enabled:
             raise SparkException.implicit_cartesian_product("inner")
@@ -230,35 +370,110 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
         # - LEFT SEMI JOIN: Returns left rows that have matches in right table (no right columns)
         # - LEFT ANTI JOIN: Returns left rows that have NO matches in right table (no right columns)
         # Both preserve only the columns from the left DataFrame without adding any columns from the right.
-        spark_cols_after_join: list[str] = left_container.column_map.get_spark_columns()
+        spark_cols_after_join = left_container.column_map.get_spark_columns()
+        snowpark_cols_after_join = left_container.column_map.get_snowpark_columns()
+        snowpark_col_types = [
+            f.datatype for f in left_container.dataframe.schema.fields
+        ]
         qualifiers = left_container.column_map.get_qualifiers()
+    elif join_type == "full_outer" and using_columns:
+        # We want the coalesced columns to be first, followed by all the left and right columns (excluding using columns)
+        spark_cols_after_join: list[str] = []
+        snowpark_cols_after_join: list[str] = []
+        snowpark_col_types: list[str] = []
+
+        left_container_snowpark_columns = (
+            left_container.column_map.get_snowpark_columns()
+        )
+        right_container_snowpark_columns = (
+            right_container.column_map.get_snowpark_columns()
+        )
+
+        qualifiers = []
+        for i in range(len(case_corrected_left_columns)):
+            spark_cols_after_join.append(case_corrected_left_columns[i])
+            snowpark_cols_after_join.append(using_columns_snowpark_names[i])
+            snowpark_col_types.append(using_columns_snowpark_types[i])
+            qualifiers.append([])
+
+        # Handle adding left and right columns, excluding the using columns
+        for i, spark_col in enumerate(left_container.column_map.get_spark_columns()):
+            if (
+                spark_col not in case_corrected_left_columns
+                or spark_col in left_container.column_map.get_spark_columns()[:i]
+            ):
+                spark_cols_after_join.append(spark_col)
+                snowpark_cols_after_join.append(left_container_snowpark_columns[i])
+                qualifiers.append(
+                    left_container.column_map.get_qualifier_for_spark_column(spark_col)
+                )
+
+                snowpark_col_types.append(
+                    left_container.dataframe.schema.fields[i].datatype
+                )
+
+        for i, spark_col in enumerate(right_container.column_map.get_spark_columns()):
+            if (
+                spark_col not in case_corrected_right_columns
+                or spark_col in right_container.column_map.get_spark_columns()[:i]
+            ):
+                spark_cols_after_join.append(spark_col)
+                snowpark_cols_after_join.append(right_container_snowpark_columns[i])
+                qualifiers.append(
+                    right_container.column_map.get_qualifier_for_spark_column(spark_col)
+                )
+
+                snowpark_col_types.append(
+                    right_container.dataframe.schema.fields[i].datatype
+                )
+
     else:
-        # Add Spark columns and plan_ids from left DF
-        spark_cols_after_join: list[str] = list(
-            left_container.column_map.get_spark_columns()
-        ) + [
-            spark_col
-            for i, spark_col in enumerate(
-                right_container.column_map.get_spark_columns()
-            )
-            if spark_col not in case_corrected_right_columns
-            or spark_col
-            in right_container.column_map.get_spark_columns()[
-                :i
-            ]  # this is to make sure we only remove the column once
+        spark_cols_after_join = left_container.column_map.get_spark_columns()
+        snowpark_cols_after_join = left_container.column_map.get_snowpark_columns()
+        snowpark_col_types = [
+            f.datatype for f in left_container.dataframe.schema.fields
         ]
 
-        qualifiers = list(left_container.column_map.get_qualifiers()) + [
-            right_container.column_map.get_qualifier_for_spark_column(spark_col)
-            for i, spark_col in enumerate(
-                right_container.column_map.get_spark_columns()
+        qualifiers = left_container.column_map.get_qualifiers()
+
+        right_df_snowpark_columns = right_container.column_map.get_snowpark_columns()
+
+        for i, spark_col in enumerate(right_container.column_map.get_spark_columns()):
+            if (
+                spark_col not in case_corrected_right_columns
+                or spark_col in right_container.column_map.get_spark_columns()[:i]
+            ):
+                spark_cols_after_join.append(spark_col)
+                snowpark_cols_after_join.append(right_df_snowpark_columns[i])
+                snowpark_col_types.append(
+                    right_container.dataframe.schema.fields[i].datatype
+                )
+
+                qualifiers.append(
+                    right_container.column_map.get_qualifier_for_spark_column(spark_col)
+                )
+
+    snowpark_cols_after_join_deduplicated = []
+    snowpark_cols_after_join_counter = Counter(snowpark_cols_after_join)
+    seen_duplicated_columns = set()
+
+    for col in snowpark_cols_after_join:
+        if snowpark_cols_after_join_counter[col] == 2:
+            # This means that the same column exists twice in the joined df, likely due to a self-join and
+            # we need to lsuffix and rsuffix to the names of both columns, similar to what Snowpark did under the hood.
+
+            suffix = (
+                DUPLICATED_JOIN_COL_RSUFFIX
+                if col in seen_duplicated_columns
+                else DUPLICATED_JOIN_COL_LSUFFIX
             )
-            if spark_col not in case_corrected_right_columns
-            or spark_col
-            in right_container.column_map.get_spark_columns()[
-                :i
-            ]  # this is to make sure we only remove the column once]
-        ]
+            unquoted_col = unquote_if_quoted(col)
+            quoted = quote_name_without_upper_casing(unquoted_col + suffix)
+            snowpark_cols_after_join_deduplicated.append(quoted)
+
+            seen_duplicated_columns.add(col)
+        else:
+            snowpark_cols_after_join_deduplicated.append(col)
 
     column_metadata = {}
     if left_container.column_map.column_metadata:
@@ -287,33 +502,13 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
     result_container = DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=spark_cols_after_join,
-        snowpark_column_names=result.columns,
+        snowpark_column_names=snowpark_cols_after_join_deduplicated,
         column_metadata=column_metadata,
         column_qualifiers=qualifiers,
+        hidden_columns=hidden_columns,
+        snowpark_column_types=snowpark_col_types,
     )
 
-    # Fix for USING join column references with different plan IDs
-    # After a USING join, references to the right dataframe's columns should resolve
-    # to the result dataframe that contains the merged columns
-    if (
-        using_columns
-        and rel.join.right.HasField("common")
-        and rel.join.right.common.HasField("plan_id")
-    ):
-        right_plan_id = rel.join.right.common.plan_id
-        set_plan_id_map(right_plan_id, result_container)
-
-    # For FULL OUTER joins, we also need to map the left dataframe's plan_id
-    # since both columns are replaced with a coalesced column
-    if (
-        using_columns
-        and join_type == "full_outer"
-        and rel.join.left.HasField("common")
-        and rel.join.left.common.HasField("plan_id")
-    ):
-        left_plan_id = rel.join.left.common.plan_id
-        set_plan_id_map(left_plan_id, result_container)
-
     if rel.join.using_columns:
         # When join 'using_columns', the 'join columns' should go first in result DF.
         idxs_to_shift = [
@@ -345,6 +540,7 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
             cached_schema_getter=lambda: snowpark.types.StructType(
                 reorder(original_df.schema.fields)
             ),
+            hidden_columns=hidden_columns,
         )
 
     return result_container

snowflake/snowpark_connect/relation/map_row_ops.py

@@ -1,7 +1,7 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-from copy import copy
+
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
@@ -9,7 +9,6 @@ from pyspark.errors.exceptions.base import AnalysisException, IllegalArgumentExc
 
 import snowflake.snowpark_connect.relation.utils as utils
 from snowflake import snowpark
-from snowflake.snowpark._internal.analyzer.binary_expression import And
 from snowflake.snowpark.functions import col, expr as snowpark_expr
 from snowflake.snowpark.types import (
     BooleanType,
@@ -31,7 +30,6 @@ from snowflake.snowpark_connect.expression.map_expression import (
 )
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
-from snowflake.snowpark_connect.relation.utils import can_filter_be_flattened
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -555,32 +553,7 @@ def map_filter(
         rel.filter.condition, input_container.column_map, typer
     )
 
-    select_statement = getattr(input_df, "_select_statement", None)
-    condition_exp = condition.col._expression
-    if (
-        can_filter_be_flattened(select_statement, condition_exp)
-        and input_df._ops_after_agg is None
-    ):
-        new = copy(select_statement)
-        new.from_ = select_statement.from_.to_subqueryable()
-        new.pre_actions = new.from_.pre_actions
-        new.post_actions = new.from_.post_actions
-        new.column_states = select_statement.column_states
-        new.where = (
-            And(select_statement.where, condition_exp)
-            if select_statement.where is not None
-            else condition_exp
-        )
-        new._merge_projection_complexity_with_subquery = False
-        new.df_ast_ids = (
-            select_statement.df_ast_ids.copy()
-            if select_statement.df_ast_ids is not None
-            else None
-        )
-        new.attributes = select_statement.attributes
-        result = input_df._with_plan(new)
-    else:
-        result = input_df.filter(condition.col)
+    result = input_df.filter(condition.col)
 
     return DataFrameContainer(
         result,

snowflake/snowpark_connect/relation/map_sql.py

@@ -77,6 +77,7 @@ from ..expression.map_sql_expression import (
 from ..utils.identifiers import spark_to_sf_single_id
 
 _ctes = ContextVar[dict[str, relation_proto.Relation]]("_ctes", default={})
+_cte_definitions = ContextVar[dict[str, any]]("_cte_definitions", default={})
 _having_condition = ContextVar[expressions_proto.Expression | None](
     "_having_condition", default=None
 )
@@ -117,12 +118,15 @@ def _push_cte_scope():
     """
     Creates a new CTE scope when evaluating nested WITH clauses.
     """
-    cur = _ctes.get()
-    token = _ctes.set(cur.copy())
+    cur_ctes = _ctes.get()
+    cur_definitions = _cte_definitions.get()
+    cte_token = _ctes.set(cur_ctes.copy())
+    def_token = _cte_definitions.set(cur_definitions.copy())
     try:
         yield
     finally:
-        _ctes.reset(token)
+        _ctes.reset(cte_token)
+        _cte_definitions.reset(def_token)
 
 
 @contextmanager
@@ -405,15 +409,21 @@ def map_sql_to_pandas_df(
                 else:
                     # TODO: Unset the schema
                     pass
-            case "CreateTable":
+            case "CreateTable" | "ReplaceTable":
+                if class_name == "ReplaceTable":
+                    replace_table = " OR REPLACE "
+                    if_not_exists = ""
+                else:
+                    replace_table = ""
+                    if_not_exists = (
+                        "IF NOT EXISTS " if logical_plan.ignoreIfExists() else ""
+                    )
+
                 name = get_relation_identifier_name(logical_plan.name())
                 columns = ", ".join(
                     _spark_field_to_sql(f, True)
                     for f in logical_plan.tableSchema().fields()
                 )
-                if_not_exists = (
-                    "IF NOT EXISTS " if logical_plan.ignoreIfExists() else ""
-                )
                 comment_opt = logical_plan.tableSpec().comment()
                 comment = (
                     f"COMMENT = '{_escape_sql_comment(str(comment_opt.get()))}'"
@@ -422,7 +432,7 @@ def map_sql_to_pandas_df(
                 )
                 # NOTE: We are intentionally ignoring any FORMAT=... parameters here.
                 session.sql(
-                    f"CREATE TABLE {if_not_exists}{name} ({columns}) {comment}"
+                    f"CREATE {replace_table} TABLE {if_not_exists}{name} ({columns}) {comment}"
                 ).collect()
             case "CreateTableAsSelect":
                 mode = "ignore" if logical_plan.ignoreIfExists() else "errorifexists"
@@ -1351,8 +1361,28 @@ def map_logical_plan_relation(
             value_column_names = [e for e in as_java_list(rel.valueColumnNames())]
             variable_column_name = rel.variableColumnName()
 
+            # Check for multi-column UNPIVOT which Snowflake doesn't support
+            if len(value_column_names) > 1:
+                raise UnsupportedOperationException(
+                    f"Multi-column UNPIVOT is not supported. Snowflake SQL does not support unpivoting "
+                    f"multiple value columns ({', '.join(value_column_names)}) in a single operation. "
+                    f"Workaround: Use separate UNPIVOT operations for each value column and join the results, "
+                    f"or restructure your query to unpivot columns individually."
+                )
+
             values = []
-            for e1 in as_java_list(rel.values().get()):
+            values_groups = as_java_list(rel.values().get())
+
+            # Check if we have multi-column groups in the IN clause
+            if values_groups and len(as_java_list(values_groups[0])) > 1:
+                group_sizes = [len(as_java_list(group)) for group in values_groups]
+                raise UnsupportedOperationException(
+                    f"Multi-column UNPIVOT is not supported. Snowflake SQL does not support unpivoting "
+                    f"multiple columns together in groups. Found groups with {max(group_sizes)} columns. "
+                    f"Workaround: Unpivot each column separately and then join/union the results as needed."
+                )
+
+            for e1 in values_groups:
                 for e in as_java_list(e1):
                     values.append(map_logical_plan_expression(e))
 
@@ -1468,7 +1498,50 @@ def map_logical_plan_relation(
                 # The name corresponds to a `WITH` alias rather than a table.
                 # TODO: We currently evaluate the query each time its alias is used;
                 # we should eventually start using `WITH` in Snowflake SQL.
-                proto = cte_proto
+                # Each CTE reference should get completely fresh evaluation to prevent ambiguity
+                # when the same CTE is joined multiple times. Instead of reusing the same cte_proto,
+                # re-evaluate the CTE definition to get fresh column identifiers.
+
+                # Re-evaluate the CTE definition to get fresh column identifiers
+                cte_definition = _cte_definitions.get().get(name)
+                if cte_definition is not None:
+                    # Get the original column names for consistency across CTE references
+                    original_container = map_relation(cte_proto)
+                    original_spark_columns = (
+                        original_container.column_map.get_spark_columns()
+                    )
+
+                    # Re-evaluate the CTE definition with a fresh plan_id
+                    fresh_plan_id = gen_sql_plan_id()
+                    fresh_cte_proto = map_logical_plan_relation(
+                        cte_definition, fresh_plan_id
+                    )
+
+                    # Use SubqueryColumnAliases to ensure consistent column names across CTE references
+                    # This is crucial for CTEs that reference other CTEs
+                    any_proto = Any()
+                    any_proto.Pack(
+                        snowflake_proto.Extension(
+                            subquery_column_aliases=snowflake_proto.SubqueryColumnAliases(
+                                input=fresh_cte_proto,
+                                aliases=original_spark_columns,
+                            )
+                        )
+                    )
+                    column_aliased_proto = relation_proto.Relation(extension=any_proto)
+                    column_aliased_proto.common.plan_id = gen_sql_plan_id()
+
+                    # Wrap in SubqueryAlias with the CTE name
+                    proto = relation_proto.Relation(
+                        subquery_alias=relation_proto.SubqueryAlias(
+                            input=column_aliased_proto,
+                            alias=name,
+                        )
+                    )
+                    proto.common.plan_id = gen_sql_plan_id()
+                else:
+                    # Fallback to stored CTE if definition not found
+                    proto = cte_proto
             else:
                 tmp_views = _get_current_temp_objects()
                 current_schema = session.connection.schema
@@ -1600,7 +1673,11 @@ def map_logical_plan_relation(
             with _push_cte_scope():
                 for cte in as_java_list(rel.cteRelations()):
                     name = str(cte._1())
-                    cte_proto = map_logical_plan_relation(cte._2())
+                    # Store the original CTE definition for re-evaluation
+                    _cte_definitions.get()[name] = cte._2()
+                    # Process CTE definition with a unique plan_id to ensure proper column naming
+                    cte_plan_id = gen_sql_plan_id()
+                    cte_proto = map_logical_plan_relation(cte._2(), cte_plan_id)
                     _ctes.get()[name] = cte_proto
 
                 proto = map_logical_plan_relation(rel.child())

snowflake/snowpark_connect/relation/map_udtf.py

@@ -27,7 +27,7 @@ from snowflake.snowpark_connect.expression.map_expression import (
 )
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.type_mapping import (
-    parse_ddl_string,
+    map_type_string_to_proto,
     proto_to_snowpark_type,
 )
 from snowflake.snowpark_connect.utils.context import push_udtf_context
@@ -125,7 +125,9 @@ def process_return_type(
 ) -> tuple[list[tuple[str, Any]], DataType, StructType, list[str]]:
     try:
         if return_type.HasField("unparsed"):
-            parsed_return = parse_ddl_string(return_type.unparsed.data_type_string)
+            parsed_return = map_type_string_to_proto(
+                return_type.unparsed.data_type_string
+            )
         else:
             parsed_return = return_type
     except ValueError as e: