snowpark-connect 0.24.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

snowflake/snowpark_connect/expression/map_unresolved_star.py

@@ -34,6 +34,7 @@ def map_unresolved_star(
     column_mapping: ColumnNameMap,
     typer: ExpressionTyper,
 ) -> tuple[list[str], TypedColumn]:
+
     if exp.unresolved_star.HasField("unparsed_target"):
         unparsed_target = exp.unresolved_star.unparsed_target
         name_parts = split_fully_qualified_spark_name(unparsed_target)
@@ -102,7 +103,7 @@ def map_unresolved_star(
                 prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
             prefix_candidate = (
                 column_mapping.get_snowpark_column_name_from_spark_column_name(
-                    prefix_candidate_str, allow_non_exists=True
+                    prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
                 )
             )
             if prefix_candidate is None:
@@ -180,7 +181,7 @@ def map_unresolved_star_struct(
             prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
         prefix_candidate = (
             column_mapping.get_snowpark_column_name_from_spark_column_name(
-                prefix_candidate_str, allow_non_exists=True
+                prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
             )
         )
         if prefix_candidate is None:

snowflake/snowpark_connect/hidden_column.py

@@ -0,0 +1,39 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+
+class HiddenColumn:
+    """
+    Represents a hidden column in a Snowflake table.
+
+    Hidden columns are not visible in standard queries but can be accessed
+    directly if needed. This class provides a way to reference such columns
+    in Snowpark operations
+    """
+
+    def __init__(
+        self,
+        hidden_snowpark_name: str,
+        spark_name: str,
+        visible_snowpark_name: str,
+        qualifiers: list[str] | None = None,
+        original_position: int | None = None,
+    ) -> None:
+        """
+        Initializes a HiddenColumn instance.
+
+        Args:
+            name (str): The name of the hidden column.
+        """
+
+        # The Snowpark internal name for the hidden column
+        self.hidden_snowpark_name = hidden_snowpark_name
+        # The Spark name for the hidden column
+        self.spark_name = spark_name
+        # The left side visible Snowpark name for the dropped right side column
+        self.visible_snowpark_name = visible_snowpark_name
+        # Qualifiers for the hidden column (e.g., table or schema names)
+        self.qualifiers = qualifiers if qualifiers is not None else []
+        # The position of the hidden column in the original schema
+        self.original_position = original_position

snowflake/snowpark_connect/relation/map_column_ops.py

@@ -735,10 +735,17 @@ def map_with_columns(
     """
     input_container = map_relation(rel.with_columns.input)
     input_df = input_container.dataframe
-    with_columns = [
-        map_alias(alias, input_container.column_map, ExpressionTyper(input_df))
-        for alias in rel.with_columns.aliases
-    ]
+    with_columns = []
+    for alias in rel.with_columns.aliases:
+        spark_names, typed_alias = map_alias(
+            alias, input_container.column_map, ExpressionTyper(input_df)
+        )
+        register_lca_alias(spark_names[0], typed_alias)
+        with_columns.append((spark_names, typed_alias))
+
+    # we don't need lateral aliases anymore
+    clear_lca_alias_map()
+
     # TODO: This list needs to contain all unique column names, but the code below doesn't
     # guarantee that.
     with_columns_names = []
@@ -806,6 +813,12 @@ def map_with_columns(
         with_columns_names_deduped, with_columns_exprs_deduped
     ).select(*new_snowpark_columns)
 
+    # SNOW-2306644: the next projection after a withColumn call can completely remove the added column
+    # df.withColumn("new").select("foo").filter("new") will fail with a missing column error
+    # the column will be preserved if flattening is disabled
+    if hasattr(result, "_select_statement"):
+        result._select_statement.flatten_disabled = True
+
     snowpark_name_to_type = dict(
         [(f.name, f.datatype) for f in input_df.schema.fields]
         + list(zip(with_columns_names, with_columns_types))

snowflake/snowpark_connect/relation/map_extension.py

@@ -350,15 +350,51 @@ def map_aggregate(
     if not is_group_by_all:
         raw_groupings = [_map_column(exp) for exp in aggregate.grouping_expressions]
 
-    # Set the current grouping columns in context for grouping_id() function
-    grouping_spark_columns = [spark_name for spark_name, _ in raw_groupings]
-    set_current_grouping_columns(grouping_spark_columns)
+    # Determine grouping columns for context
+    # For GROUPING SETS, we need to extract the columns from the sets
+    grouping_columns_for_context = []
+    if aggregate.group_type == snowflake_proto.Aggregate.GROUP_TYPE_GROUPING_SETS:
+        # Use a list to preserve order, avoiding duplicates
+        for grouping_set in aggregate.grouping_sets:
+            for exp in grouping_set.grouping_set:
+                spark_names, _ = map_expression(exp, input_container.column_map, typer)
+                # map_expression always returns a list, get the first element
+                col_name = spark_names[0]
+                if col_name not in grouping_columns_for_context:
+                    grouping_columns_for_context.append(col_name)
+    else:
+        grouping_columns_for_context = [spark_name for spark_name, _ in raw_groupings]
+
+    # Set grouping columns context for processing aggregate expressions
+    # This context is needed for resolving grouping__id references
+    # TODO: This should properly handle nested queries with GROUP BY using push/pop
+    # Currently, nested queries may interfere with parent queries
+    set_current_grouping_columns(grouping_columns_for_context)
+
+    # LCA Support for aggregate expressions: Use the LCA alias map
+    # Note: We don't clear the map here to preserve any parent context aliases
+    from snowflake.snowpark_connect.utils.context import register_lca_alias
 
     agg_count = get_sql_aggregate_function_count()
     for exp in aggregate.aggregate_expressions:
         col = _map_column(exp)
         raw_aggregations.append(col)
 
+        # If this is an alias, register it in the LCA map for subsequent expressions
+        if (
+            exp.WhichOneof("expr_type") == "alias"
+            and exp.alias.name
+            and len(exp.alias.name) > 0
+        ):
+            alias_name = exp.alias.name[0]
+            spark_name, snowpark_column = col
+
+            # Register the alias pointing to the result of its expression
+            # This handles both simple aliases (k as lca) and complex ones (lca + 1 as col)
+            # The snowpark_column already contains the computed expression with its alias wrapper,
+            # which is fine - when referenced later, the column's value is what gets used
+            register_lca_alias(alias_name, snowpark_column)
+
         if is_group_by_all:
             new_agg_count = get_sql_aggregate_function_count()
             if new_agg_count == agg_count:
@@ -404,15 +440,20 @@ def map_aggregate(
         case snowflake_proto.Aggregate.GROUP_TYPE_CUBE:
             result = input_df.cube(groupings)
         case snowflake_proto.Aggregate.GROUP_TYPE_GROUPING_SETS:
-            # TODO: What do we do about groupings?
-            sets = (
-                [
-                    map_expression(exp, input_container.column_map, typer)[1].col
-                    for exp in grouping_sets.grouping_set
-                ]
-                for grouping_sets in aggregate.grouping_sets
+            # Map each grouping set to columns
+            sets_mapped = []
+            for grouping_set in aggregate.grouping_sets:
+                set_cols = []
+                for exp in grouping_set.grouping_set:
+                    _, typed_col = map_expression(
+                        exp, input_container.column_map, typer
+                    )
+                    set_cols.append(typed_col.col)
+                sets_mapped.append(set_cols)
+
+            result = input_df.group_by_grouping_sets(
+                snowpark.GroupingSets(*sets_mapped)
             )
-            result = input_df.group_by_grouping_sets(snowpark.GroupingSets(*sets))
         case other:
             raise SnowparkConnectNotImplementedError(
                 f"Unsupported GROUP BY type: {other}"

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_map_partitions.py

@@ -46,9 +46,10 @@ def map_map_partitions(
     udf_check(udf_proto)
 
     # Check if this is mapInArrow (eval_type == 207)
-    eval_type = udf_proto.python_udf.eval_type
-
-    if eval_type == MAP_IN_ARROW_EVAL_TYPE:
+    if (
+        udf_proto.WhichOneof("function") == "python_udf"
+        and udf_proto.python_udf.eval_type == MAP_IN_ARROW_EVAL_TYPE
+    ):
         return _map_in_arrow_with_pandas_udtf(input_container, udf_proto)
     else:
         return _map_partitions_with_udf(input_df, udf_proto)
@@ -126,7 +127,11 @@ def _map_partitions_with_udf(
         "udf_name": "spark_map_partitions_udf",
         "input_column_names": input_column_names,
         "replace": True,
-        "return_type": proto_to_snowpark_type(udf_proto.python_udf.output_type),
+        "return_type": proto_to_snowpark_type(
+            udf_proto.python_udf.output_type
+            if udf_proto.WhichOneof("function") == "python_udf"
+            else udf_proto.scalar_scala_udf.outputType
+        ),
         "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
         "udf_imports": get_python_udxf_import_files(input_df.session),
     }