snowpark-connect 0.27.0__py3-none-any.whl → 0.28.1__py3-none-any.whl

snowflake/snowpark_connect/relation/map_join.py

@@ -1,18 +1,13 @@
 #
 # 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
@@ -22,7 +17,6 @@ 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,
@@ -30,6 +24,7 @@ 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 (
@@ -38,9 +33,6 @@ 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)
@@ -82,13 +74,6 @@ 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
@@ -120,8 +105,8 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
             right=right_input,
             on=join_expression.col,
             how=join_type,
-            lsuffix=DUPLICATED_JOIN_COL_LSUFFIX,
-            rsuffix=DUPLICATED_JOIN_COL_RSUFFIX,
+            lsuffix="_left",
+            rsuffix="_right",
         )
     elif using_columns:
         if any(
@@ -171,24 +156,12 @@ 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(
-                using_columns_snowpark_names
+        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_right_columns = right_container.column_map.get_spark_column_names_from_snowpark_column_names(
@@ -222,141 +195,28 @@ 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 = []
-            for i, (left_col, _right_col) in enumerate(snowpark_using_columns):
+            columns_to_drop = []
+            for i, (left_col, right_col) in enumerate(snowpark_using_columns):
                 # Use the original user-specified column name to preserve case sensitivity
-                # 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)
-
-                # 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,
-                    ]
+                original_column_name = rel.join.using_columns[i]
+                coalesced_col = snowpark_fn.coalesce(left_col, right_col).alias(
+                    original_column_name
                 )
+                coalesced_columns.append(coalesced_col)
+                columns_to_drop.extend([left_col, right_col])
 
-            # 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.hidden_snowpark_name for col in hidden_columns]
+                if col_name not in [col.getName() for col in columns_to_drop]
             ]
             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")
@@ -370,110 +230,35 @@ 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 = 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
-        ]
+        spark_cols_after_join: list[str] = left_container.column_map.get_spark_columns()
         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:
-        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
+        # 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
         ]
 
-        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
+        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()
             )
-            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)
+            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]
+        ]
 
     column_metadata = {}
     if left_container.column_map.column_metadata:
@@ -502,13 +287,33 @@ 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=snowpark_cols_after_join_deduplicated,
+        snowpark_column_names=result.columns,
         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 = [
@@ -540,7 +345,6 @@ 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_relation.py

@@ -15,6 +15,7 @@ from snowflake.snowpark_connect.utils.cache import (
 from snowflake.snowpark_connect.utils.context import (
     get_plan_id_map,
     get_session_id,
+    not_resolving_fun_args,
     push_map_partitions,
     push_operation_scope,
     set_is_aggregate_function,
@@ -150,7 +151,10 @@ def map_relation(
             case "drop_na":
                 result = map_row_ops.map_dropna(rel)
             case "extension":
-                result = map_extension.map_extension(rel)
+                # Extensions can be passed as function args, and we need to reset the context here.
+                # Matters only for resolving alias expressions in the extensions rel.
+                with not_resolving_fun_args():
+                    result = map_extension.map_extension(rel)
             case "fill_na":
                 result = map_row_ops.map_fillna(rel)
             case "filter":