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

snowflake/snowpark_connect/relation/map_column_ops.py

@@ -28,12 +28,19 @@ from snowflake.snowpark._internal.utils import generate_random_alphanumeric
 from snowflake.snowpark.column import Column
 from snowflake.snowpark.table_function import _ExplodeFunctionCall
 from snowflake.snowpark.types import DataType, StructField, StructType, _NumericType
+from snowflake.snowpark_connect import tcm
 from snowflake.snowpark_connect.column_name_handler import (
+    ColumnQualifier,
     make_column_names_snowpark_compatible,
+    make_unique_snowpark_name,
 )
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
-from snowflake.snowpark_connect.error.error_utils import SparkException
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import (
+    SparkException,
+    attach_custom_error_code,
+)
 from snowflake.snowpark_connect.expression.map_expression import (
     map_alias,
     map_expression,
@@ -42,6 +49,9 @@ from snowflake.snowpark_connect.expression.map_expression import (
 from snowflake.snowpark_connect.expression.map_unresolved_function import unwrap_literal
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.relation.read.metadata_utils import (
+    without_internal_columns,
+)
 from snowflake.snowpark_connect.relation.utils import (
     TYPE_MAP_FOR_TO_SCHEMA,
     snowpark_functions_col,
@@ -56,6 +66,9 @@ from snowflake.snowpark_connect.utils.context import (
     clear_lca_alias_map,
     register_lca_alias,
 )
+from snowflake.snowpark_connect.utils.expression_transformer import (
+    is_child_agg_function_expression,
+)
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
@@ -73,7 +86,7 @@ def map_drop(
 
     The drop is a list of expressions that is applied to the DataFrame.
     """
-    input_container = map_relation(rel.drop.input)
+    input_container = map_relation(rel.drop.input).without_hidden_columns()
     input_df = input_container.dataframe
     typer = ExpressionTyper(input_df)
     columns_to_drop_with_names = []
@@ -129,20 +142,19 @@ def map_drop(
     # object to handle these cases.
     try:
         column_map = input_container.column_map
-        new_columns_names = column_map.get_snowpark_columns_after_drop(
+        new_columns = column_map.get_snowpark_columns_after_drop(
             _get_column_names_to_drop()
         )
         result: snowpark.DataFrame = input_df.drop(*columns_to_drop)
         return DataFrameContainer.create_with_column_mapping(
             dataframe=result,
-            spark_column_names=column_map.get_spark_column_names_from_snowpark_column_names(
-                new_columns_names
-            ),
-            snowpark_column_names=new_columns_names,
-            column_qualifiers=column_map.get_qualifiers_for_columns_after_drop(
-                _get_column_names_to_drop()
-            ),
+            spark_column_names=[c.spark_name for c in new_columns],
+            snowpark_column_names=[c.snowpark_name for c in new_columns],
+            column_qualifiers=[c.qualifiers for c in new_columns],
             parent_column_name_map=column_map,
+            equivalent_snowpark_names=[
+                c.equivalent_snowpark_names for c in new_columns
+            ],
         )
     except snowpark.exceptions.SnowparkColumnException:
         from snowflake.snowpark_connect.empty_dataframe import EmptyDataFrame
@@ -160,7 +172,6 @@ def map_project(
     """
     if rel.project.HasField("input"):
         input_container = map_relation(rel.project.input)
-        input_df = input_container.dataframe
     else:
         # Create a dataframe to represent a OneRowRelation AST node.
         # XXX: Snowflake does not support 0-column tables, so create a dummy column;
@@ -200,20 +211,33 @@ def map_project(
 
     # Collect aliases to batch process them
     pending_aliases = []  # List of (spark_name, snowpark_col, aliased_col, alias_types)
+    # Track columns that might need aliasing if multi-column generators are present
+    # Format: (index_in_select_list, snowpark_column_name, mapper.col)
+    conditional_aliases = []
+
+    def _is_attribute(exp: expressions_proto.Expression) -> bool:
+        return exp.WhichOneof("expr_type") == "unresolved_attribute"
 
     # Detect if this is a simple projection (only unqualified column references, no aliases or functions)
     # Qualified column references (with plan_id) should NOT be considered simple projections
     # because they've already been resolved to specific DataFrames
     def _is_simple_projection(exp: expressions_proto.Expression) -> bool:
-        return exp.WhichOneof(
-            "expr_type"
-        ) == "unresolved_attribute" and not exp.unresolved_attribute.HasField(
+        return _is_attribute(exp) and not exp.unresolved_attribute.HasField(
             "plan_id"
         )  # No DataFrame qualification
 
+    has_agg_in_input_df = (
+        input_df._select_statement
+        and input_df._select_statement.projection
+        and any(
+            is_child_agg_function_expression(exp)
+            for exp in input_df._select_statement.projection
+        )
+    )
     column_types = []
     has_multi_column_alias = False
-    qualifiers = []
+    qualifiers: list[set[ColumnQualifier]] = []
+    equivalent_snowpark_names: list[set[str]] = []
 
     typer = ExpressionTyper(input_df)
 
@@ -233,10 +257,11 @@ def map_project(
             # Check if this was a qualified column reference (like df_alias.column)
             # by checking if the original expression was an alias lookup
             is_qualified_reference = (
-                exp.WhichOneof("expr_type") == "unresolved_attribute"
+                _is_attribute(exp)
                 and "." in exp.unresolved_attribute.unparsed_identifier
             )
 
+            existing_snowpark_name = None
             if (
                 _is_simple_projection(exp)
                 and not is_qualified_reference
@@ -264,16 +289,47 @@ def map_project(
                     [spark_name], rel.common.plan_id, len(new_snowpark_columns)
                 )[0]
 
-            aliased_col = mapper.col.alias(snowpark_column)
-            select_list.append(aliased_col)
+            # since unresolved attributes get aliased, we need to retain their original name
+            # so that we're able to resolve it later
+            if (
+                _is_attribute(exp)
+                and mapper.col.get_name()
+                and mapper.col.get_name() not in new_snowpark_columns
+            ):
+                old_name = mapper.col.get_name()
+                eq_names = set()
+                eq_names.update(
+                    input_container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                        old_name
+                    )
+                )
+                if old_name != snowpark_column:
+                    eq_names.update({old_name})
+                equivalent_snowpark_names.append(eq_names)
+            else:
+                equivalent_snowpark_names.append(set())
+
+            # TODO: go back to using alias always once SNOW-2203826 is done
+            if existing_snowpark_name == snowpark_column and not has_agg_in_input_df:
+                aliased_col = mapper.col
+                # Store info to potentially re-alias later if multi-column generators appear
+                conditional_aliases.append(
+                    (len(select_list), snowpark_column, mapper.col)
+                )
+            else:
+                aliased_col = mapper.col.alias(snowpark_column)
 
+            select_list.append(aliased_col)
             new_snowpark_columns.append(snowpark_column)
             new_spark_columns.append(spark_name)
             column_types.extend(mapper.types)
             qualifiers.append(mapper.get_qualifiers())
 
             # Only update the DataFrame and register LCA for explicit aliases
-            if exp.WhichOneof("expr_type") == "alias":
+            if (
+                exp.WhichOneof("expr_type") == "alias"
+                and not context.is_resolving_subquery_exp()
+            ):
                 # Collect alias for batch processing
                 pending_aliases.append(
                     (spark_name, snowpark_column, aliased_col, mapper.types)
@@ -283,6 +339,20 @@ def map_project(
                 alias_types = mapper.types
                 typed_alias = TypedColumn(aliased_col, lambda types=alias_types: types)
                 register_lca_alias(spark_name, typed_alias)
+
+                # Also register with the original qualified name if this is an alias of a column reference
+                # This handles ORDER BY referencing the original name: SELECT o.date AS order_date ... ORDER BY o.date
+                if (
+                    exp.alias.HasField("expr")
+                    and exp.alias.expr.WhichOneof("expr_type") == "unresolved_attribute"
+                ):
+                    original_name = (
+                        exp.alias.expr.unresolved_attribute.unparsed_identifier
+                    )
+                    if (
+                        original_name != spark_name
+                    ):  # Don't register twice with the same name
+                        register_lca_alias(original_name, typed_alias)
         else:
             # Multi-column case ('select *', posexplode, explode, inline, etc.)
             has_multi_column_alias = True
@@ -292,6 +362,21 @@ def map_project(
             new_spark_columns.extend(new_spark_names)
             column_types.extend(mapper.types)
             qualifiers.extend(mapper.get_multi_col_qualifiers(len(new_spark_names)))
+            equivalent_snowpark_names.extend(
+                [
+                    input_container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                        snowpark_name
+                    )
+                    for snowpark_name in result_columns
+                ]
+            )
+
+    # TODO: go back to using alias always once SNOW-2203826 is done
+    # If we have multi-column aliases, we need to ensure ALL columns are aliased
+    # to avoid Snowpark SQL generation issues
+    if has_multi_column_alias and conditional_aliases:
+        for idx, snowpark_col_name, col_obj in conditional_aliases:
+            select_list[idx] = col_obj.alias(snowpark_col_name)
 
     if pending_aliases:
         # LCA case: create intermediate DataFrame with aliases, then do final projection
@@ -311,6 +396,11 @@ def map_project(
         final_snowpark_columns = make_column_names_snowpark_compatible(
             new_spark_columns, rel.common.plan_id
         )
+        # if there are duplicate snowpark column names, we need to disambiguate them by their index
+        if len(new_spark_columns) != len(set(new_spark_columns)):
+            result = result.select(
+                [f"${i}" for i in range(1, len(new_spark_columns) + 1)]
+            )
         result = result.toDF(*final_snowpark_columns)
         new_snowpark_columns = final_snowpark_columns
 
@@ -324,6 +414,7 @@ def map_project(
         parent_column_name_map=input_container.column_map,
         table_name=input_container.table_name,
         alias=input_container.alias,
+        equivalent_snowpark_names=equivalent_snowpark_names,
     )
 
 
@@ -369,56 +460,87 @@ def map_sort(
                 for col in input_container.column_map.get_spark_columns()
             ]
 
-    for so in sort_order:
-        if so.child.HasField("literal"):
-            column_index = unwrap_literal(so.child)
-            try:
-                if column_index <= 0:
-                    raise IndexError
-                col = input_df[column_index - 1]
-            except IndexError:
-                raise AnalysisException(
-                    f"""[ORDER_BY_POS_OUT_OF_RANGE] ORDER BY position {column_index} is not in select list (valid range is [1, {len(input_df.columns)})])."""
-                )
-        else:
-            _, typed_column = map_single_column_expression(
-                so.child, input_container.column_map, typer
-            )
-            col = typed_column.col
+    # Process ORDER BY expressions with a context flag to enable column reuse optimization
+    from snowflake.snowpark_connect.utils.context import push_processing_order_by_scope
+
+    with push_processing_order_by_scope():
+        for so in sort_order:
+            if so.child.HasField("literal"):
+                column_index = unwrap_literal(so.child)
+                try:
+                    if column_index <= 0:
+                        exception = IndexError()
+                        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                        raise exception
+                    col = input_df[column_index - 1]
+                except IndexError:
+                    exception = AnalysisException(
+                        f"""[ORDER_BY_POS_OUT_OF_RANGE] ORDER BY position {column_index} is not in select list (valid range is [1, {len(input_df.columns)})])."""
+                    )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
+            else:
+                # Check if input came from an aggregate - if so, use hybrid resolution
+                # This handles ORDER BY expressions that reference pre-aggregation columns
+                # e.g., ORDER BY year(date) when the aggregated result only has 'year' alias
+                if (
+                    hasattr(input_container, "_aggregate_metadata")
+                    and input_container._aggregate_metadata is not None
+                ):
+                    from snowflake.snowpark_connect.expression.hybrid_column_map import (
+                        create_hybrid_column_map_for_order_by,
+                    )
 
-        match (so.direction, so.null_ordering):
-            case (
-                expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING,
-                expressions_proto.Expression.SortOrder.SORT_NULLS_FIRST,
-            ):
-                col = col.asc_nulls_first()
-            case (
-                expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING,
-                expressions_proto.Expression.SortOrder.SORT_NULLS_LAST,
-            ):
-                col = col.asc_nulls_last()
-            case (
-                expressions_proto.Expression.SortOrder.SORT_DIRECTION_DESCENDING,
-                expressions_proto.Expression.SortOrder.SORT_NULLS_FIRST,
-            ):
-                col = col.desc_nulls_first()
-            case (
-                expressions_proto.Expression.SortOrder.SORT_DIRECTION_DESCENDING,
-                expressions_proto.Expression.SortOrder.SORT_NULLS_LAST,
-            ):
-                col = col.desc_nulls_last()
+                    # Create hybrid map for resolving ORDER BY expressions
+                    hybrid_map = create_hybrid_column_map_for_order_by(
+                        aggregate_metadata=input_container._aggregate_metadata,
+                        aggregated_df=input_df,
+                        aggregated_column_map=input_container.column_map,
+                    )
 
-        cols.append(col)
+                    # Resolve using hybrid context (can access both input and aggregated columns)
+                    _, typed_column = hybrid_map.resolve_expression(so.child)
+                    col = typed_column.col
+                else:
+                    # Normal resolution for non-aggregate inputs
+                    _, typed_column = map_single_column_expression(
+                        so.child, input_container.column_map, typer
+                    )
+                    col = typed_column.col
 
-        ascending.append(
-            so.direction
-            == expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING
-        )
-        if (
-            so.direction
-            != expressions_proto.Expression.SortOrder.SORT_DIRECTION_UNSPECIFIED
-        ):
-            order_specified = True
+            match (so.direction, so.null_ordering):
+                case (
+                    expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING,
+                    expressions_proto.Expression.SortOrder.SORT_NULLS_FIRST,
+                ):
+                    col = col.asc_nulls_first()
+                case (
+                    expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING,
+                    expressions_proto.Expression.SortOrder.SORT_NULLS_LAST,
+                ):
+                    col = col.asc_nulls_last()
+                case (
+                    expressions_proto.Expression.SortOrder.SORT_DIRECTION_DESCENDING,
+                    expressions_proto.Expression.SortOrder.SORT_NULLS_FIRST,
+                ):
+                    col = col.desc_nulls_first()
+                case (
+                    expressions_proto.Expression.SortOrder.SORT_DIRECTION_DESCENDING,
+                    expressions_proto.Expression.SortOrder.SORT_NULLS_LAST,
+                ):
+                    col = col.desc_nulls_last()
+
+            cols.append(col)
+
+            ascending.append(
+                so.direction
+                == expressions_proto.Expression.SortOrder.SORT_DIRECTION_ASCENDING
+            )
+            if (
+                so.direction
+                != expressions_proto.Expression.SortOrder.SORT_DIRECTION_UNSPECIFIED
+            ):
+                order_specified = True
 
     # TODO: sort.isglobal.
     if not order_specified:
@@ -440,15 +562,17 @@ def map_to_df(
     """
     Transform the column names of the input DataFrame and return a container.
     """
-    input_container = map_relation(rel.to_df.input)
+    input_container = without_internal_columns(map_relation(rel.to_df.input))
     input_df = input_container.dataframe
 
     new_column_names = list(rel.to_df.column_names)
     if len(new_column_names) != len(input_container.column_map.columns):
         # TODO: Check error type here
-        raise ValueError(
+        exception = ValueError(
             "Number of column names must match number of columns in DataFrame"
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
     snowpark_new_column_names = make_column_names_snowpark_compatible(
         new_column_names, rel.common.plan_id
     )
@@ -475,6 +599,7 @@ def map_to_df(
         table_name=input_container.table_name,
         alias=input_container.alias,
         cached_schema_getter=_get_schema,
+        equivalent_snowpark_names=[set()] * len(new_column_names),
     )
     context.set_df_before_projection(result_container)
     return result_container
@@ -507,9 +632,11 @@ def map_to_schema(
     for field in rel.to_schema.schema.struct.fields:
         if field.name in already_existing_columns:
             if count_case_insensitive_column_names[field.name.lower()] > 1:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"[AMBIGUOUS_COLUMN_OR_FIELD] Column or field `{field.name}` is ambiguous and has {len(input_container.column_map.spark_to_col[field.name])} matches."
                 )
+                attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+                raise exception
             snowpark_name = None
             for name in input_container.column_map.spark_to_col:
                 if name.lower() == field.name.lower():
@@ -526,17 +653,23 @@ def map_to_schema(
                         and snowpark_field.nullable
                         and not isinstance(snowpark_field.datatype, StructType)
                     ):
-                        raise AnalysisException(
+                        exception = AnalysisException(
                             f"[NULLABLE_COLUMN_OR_FIELD] Column or field `{field.name}` is nullable while it's required to be non-nullable."
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_OPERATION
+                        )
+                        raise exception
 
                     # Check type casting validation
                     if not _can_cast_column_in_schema(
                         snowpark_field.datatype, proto_to_snowpark_type(field.data_type)
                     ):
-                        raise AnalysisException(
+                        exception = AnalysisException(
                             f"""[INVALID_COLUMN_OR_FIELD_DATA_TYPE] Column or field `{field.name}` is of type "{map_snowpark_to_pyspark_types(proto_to_snowpark_type(field.data_type))}" while it's required to be "{map_snowpark_to_pyspark_types(snowpark_field.datatype)}"."""
                         )
+                        attach_custom_error_code(exception, ErrorCodes.TYPE_MISMATCH)
+                        raise exception
     if len(already_existing_columns) == len(new_column_names):
         # All columns already exist, we're doing a simple update.
         snowpark_new_column_names = []
@@ -615,6 +748,7 @@ def map_to_schema(
         snowpark_column_types=[field.datatype for field in snowpark_schema.fields],
         column_metadata=column_metadata,
         parent_column_name_map=input_container.column_map,
+        equivalent_snowpark_names=[set()] * len(new_column_names),
     )
 
 
@@ -624,7 +758,9 @@ def map_with_columns_renamed(
     """
     Rename columns in a DataFrame and return a container.
     """
-    input_container = map_relation(rel.with_columns_renamed.input)
+    input_container = without_internal_columns(
+        map_relation(rel.with_columns_renamed.input)
+    )
     input_df = input_container.dataframe
     rename_columns_map = dict(rel.with_columns_renamed.rename_columns_map)
 
@@ -660,8 +796,6 @@ def map_with_columns_renamed(
                 # This just copies the renames from previous computed dataframe
                 rename_columns_map[key] = value
 
-    existing_columns = input_container.column_map.get_spark_columns()
-
     def _column_exists_error(name: str) -> AnalysisException:
         return AnalysisException(
             f"[COLUMN_ALREADY_EXISTS] The column `{name}` already exists. Consider to choose another name or rename the existing column."
@@ -700,27 +834,42 @@ def map_with_columns_renamed(
             raise _column_exists_error(new_name)
         seen.add(new_name)
 
-    new_columns = []
-    for c in existing_columns:
+    new_spark_names = []
+    new_snowpark_names = []
+    qualifiers = []
+    equivalent_snowpark_names = []
+    for c in column_map.columns:
+        spark_name = c.spark_name
+        new_spark_name = None
         if global_config.spark_sql_caseSensitive:
-            new_columns.append(rename_columns_map.get(c, c))
-        elif rename_columns_map.get(c.lower(), None) is not None:
-            new_columns.append(
-                rename_columns_map_original.get(rename_columns_map.get(c.lower()))
+            new_spark_name = rename_columns_map.get(spark_name, None)
+        elif rename_columns_map.get(spark_name.lower(), None) is not None:
+            new_spark_name = rename_columns_map_original.get(
+                rename_columns_map.get(spark_name.lower())
             )
+
+        if new_spark_name:
+            new_spark_names.append(new_spark_name)
+            new_snowpark_names.append(make_unique_snowpark_name(new_spark_name))
+            qualifiers.append(set())
+            equivalent_snowpark_names.append(set())
         else:
-            new_columns.append(c)
+            new_spark_names.append(c.spark_name)
+            new_snowpark_names.append(c.snowpark_name)
+            qualifiers.append(c.qualifiers)
+            equivalent_snowpark_names.append(c.equivalent_snowpark_names)
 
     # Creating a new df to avoid updating the state of cached dataframe.
     new_df = input_df.select("*")
     result_container = DataFrameContainer.create_with_column_mapping(
         dataframe=new_df,
-        spark_column_names=new_columns,
+        spark_column_names=new_spark_names,
         snowpark_column_names=input_container.column_map.get_snowpark_columns(),
         column_qualifiers=input_container.column_map.get_qualifiers(),
         parent_column_name_map=input_container.column_map.get_parent_column_name_map(),
         table_name=input_container.table_name,
         alias=input_container.alias,
+        equivalent_snowpark_names=equivalent_snowpark_names,
     )
     result_container.column_map.rename_chains = rename_columns_map
 
@@ -733,7 +882,7 @@ def map_with_columns(
     """
     Add columns to a DataFrame and return a container.
     """
-    input_container = map_relation(rel.with_columns.input)
+    input_container = without_internal_columns(map_relation(rel.with_columns.input))
     input_df = input_container.dataframe
     with_columns = []
     for alias in rel.with_columns.aliases:
@@ -761,9 +910,11 @@ def map_with_columns(
         name = names_list[0]
         name_normalized = input_container.column_map._normalized_spark_name(name)
         if name_normalized in seen_columns:
-            raise ValueError(
+            exception = ValueError(
                 f"[COLUMN_ALREADY_EXISTS] The column `{name}` already exists."
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
         seen_columns.add(name_normalized)
         # If the column name is already in the DataFrame, we replace it, so we use the
         # mapping to get the correct column name.
@@ -772,7 +923,9 @@ def map_with_columns(
                 [name]
             )
             if len(all_instances_of_spark_column_name) == 0:
-                raise KeyError(f"Spark column name {name} does not exist")
+                exception = KeyError(f"Spark column name {name} does not exist")
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
             with_columns_names.extend(all_instances_of_spark_column_name)
             with_columns_exprs.extend(
                 [expr.col] * len(all_instances_of_spark_column_name)
@@ -796,6 +949,7 @@ def map_with_columns(
         new_spark_columns,
         new_snowpark_columns,
         qualifiers,
+        equivalent_snowpark_names,
     ) = input_container.column_map.with_columns(new_spark_names, with_columns_names)
 
     # dedup the change in columns at snowpark name level, this is required by the with columns functions
@@ -843,6 +997,7 @@ def map_with_columns(
         parent_column_name_map=input_container.column_map,
         table_name=input_container.table_name,
         alias=input_container.alias,
+        equivalent_snowpark_names=equivalent_snowpark_names,
     )
 
 
@@ -852,7 +1007,9 @@ def map_unpivot(
     # Spark API:    df.unpivot([id_columns], [unpivot_columns], var_column, val_column)
     # Snowpark API: df.unpivot(val_column, var_column, [unpivot_columns])
     if rel.unpivot.HasField("values") and len(rel.unpivot.values.values) == 0:
-        raise SparkException.unpivot_requires_value_columns()
+        exception = SparkException.unpivot_requires_value_columns()
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
 
     input_container = map_relation(rel.unpivot.input)
     input_df = input_container.dataframe
@@ -893,7 +1050,7 @@ def map_unpivot(
         )
         if not get_lease_common_ancestor_classes(type_list):
             # TODO: match exactly how spark shows mismatched columns
-            raise SparkException.unpivot_value_data_type_mismatch(
+            exception = SparkException.unpivot_value_data_type_mismatch(
                 ", ".join(
                     [
                         f"{dtype} {column_name}"
@@ -901,6 +1058,8 @@ def map_unpivot(
                     ]
                 )
             )
+            attach_custom_error_code(exception, ErrorCodes.TYPE_MISMATCH)
+            raise exception
         return not is_same_type and contains_numeric_type
 
     def get_column_names(
@@ -984,7 +1143,7 @@ def map_unpivot(
     column_project = []
     column_reverse_project = []
     snowpark_columns = []
-    qualifiers = []
+    qualifiers: list[set[ColumnQualifier]] = []
     for c in input_container.column_map.get_snowpark_columns():
         c_name = snowpark_functions_col(c, input_container.column_map).get_name()
         if c_name in unpivot_col_names:
@@ -1012,7 +1171,7 @@ def map_unpivot(
             )
             snowpark_columns.append(c)
             qualifiers.append(
-                input_container.column_map.get_qualifier_for_spark_column(c)
+                input_container.column_map.get_qualifiers_for_snowpark_column(c)
             )
 
     # Without the case when postprocessing, the result Spark dataframe is:
@@ -1057,7 +1216,7 @@ def map_unpivot(
         snowpark_functions_col(snowpark_value_column_name, input_container.column_map)
     )
     snowpark_columns.append(snowpark_value_column_name)
-    qualifiers.extend([[]] * 2)
+    qualifiers.extend([set() for _ in range(2)])
 
     result = (
         input_df.select(*column_project)
@@ -1075,6 +1234,7 @@ def map_unpivot(
         snowpark_column_names=snowpark_columns,
         column_qualifiers=qualifiers,
         parent_column_name_map=input_container.column_map,
+        equivalent_snowpark_names=[set()] * len(snowpark_columns),
     )
 
 
@@ -1097,7 +1257,9 @@ def map_group_map(
         snowpark_grouping_expressions.append(snowpark_column.col)
         group_name_list.append(new_name)
     if rel.group_map.func.python_udf is None:
-        raise ValueError("group_map relation without python udf is not supported")
+        exception = ValueError("group_map relation without python udf is not supported")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     python_major, python_minor = rel.group_map.func.python_udf.python_ver.split(".")
     is_compatible_python = sys.version_info.major == int(
@@ -1106,7 +1268,7 @@ def map_group_map(
 
     output_type = proto_to_snowpark_type(rel.group_map.func.python_udf.output_type)
 
-    if not is_compatible_python or TEST_FLAG_FORCE_CREATE_SPROC:
+    if not is_compatible_python or TEST_FLAG_FORCE_CREATE_SPROC or tcm.TCM_MODE:
         original_columns = None
         if input_container.column_map is not None:
             original_columns = [
@@ -1124,10 +1286,15 @@ def map_group_map(
         group_by_df = input_df.group_by(*snowpark_grouping_expressions)
         inner_df = group_by_df._dataframe
 
-        result = inner_df.select(
-            snowpark_fn.call_table_function(
-                apply_udtf_temp_name, *inner_df.columns
-            ).over(partition_by=snowpark_grouping_expressions)
+        renamed_columns = [f"snowflake_jtf_{column}" for column in input_df.columns]
+        tfc = snowpark_fn.call_table_function(
+            apply_udtf_temp_name, *renamed_columns
+        ).over(partition_by=snowpark_grouping_expressions)
+
+        result = (
+            inner_df.to_df(renamed_columns)
+            .join_table_function(tfc)
+            .drop(*renamed_columns)
         )
     else:
         (
@@ -1146,6 +1313,7 @@ def map_group_map(
         snowpark_column_names=result.columns,
         column_qualifiers=None,
         parent_column_name_map=input_container.column_map,
+        equivalent_snowpark_names=None,
     )