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

snowflake/snowpark_connect/relation/map_extension.py

@@ -2,6 +2,9 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import copy
+from typing import Any
+
 import cloudpickle as pkl
 import pyspark.sql.connect.proto.expressions_pb2 as expression_proto
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
@@ -11,22 +14,40 @@ import snowflake.snowpark.functions as snowpark_fn
 import snowflake.snowpark.types as snowpark_types
 import snowflake.snowpark_connect.proto.snowflake_relation_ext_pb2 as snowflake_proto
 from snowflake import snowpark
+from snowflake.snowpark import Column
 from snowflake.snowpark_connect.column_name_handler import (
     ColumnNameMap,
     make_column_names_snowpark_compatible,
 )
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 from snowflake.snowpark_connect.config import get_boolean_session_config_param
-from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
-from snowflake.snowpark_connect.expression.map_expression import map_expression
+from snowflake.snowpark_connect.dataframe_container import (
+    AggregateMetadata,
+    DataFrameContainer,
+)
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
+from snowflake.snowpark_connect.expression.map_expression import (
+    map_expression,
+    map_single_column_expression,
+)
 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 (
+    create_pivot_column_condition,
+    get_all_dependent_column_names,
+    map_pivot_value_to_spark_column_name,
+)
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
     get_sql_aggregate_function_count,
-    not_resolving_fun_args,
     push_outer_dataframe,
     set_current_grouping_columns,
 )
+from snowflake.snowpark_connect.utils.expression_transformer import (
+    inject_condition_to_all_agg_functions,
+)
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
@@ -85,16 +106,19 @@ def map_extension(
             input_df = result.dataframe
             snowpark_col_names = result.column_map.get_snowpark_columns()
             if len(subquery_aliases.aliases) != len(snowpark_col_names):
-                raise AnalysisException(
+                exception = AnalysisException(
                     "Number of column aliases does not match number of columns. "
                     f"Number of column aliases: {len(subquery_aliases.aliases)}; "
                     f"number of columns: {len(snowpark_col_names)}."
                 )
+                attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                raise exception
             return DataFrameContainer.create_with_column_mapping(
                 dataframe=input_df,
                 spark_column_names=subquery_aliases.aliases,
                 snowpark_column_names=snowpark_col_names,
                 column_qualifiers=result.column_map.get_qualifiers(),
+                equivalent_snowpark_names=result.column_map.get_equivalent_snowpark_names(),
             )
         case "lateral_join":
             lateral_join = extension.lateral_join
@@ -109,18 +133,22 @@ def map_extension(
 
             left_queries = left_df.queries["queries"]
             if len(left_queries) != 1:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Unexpected number of queries: {len(left_queries)}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
             left_query = left_queries[0]
             with push_outer_dataframe(left_result):
                 right_result = map_relation(lateral_join.right)
                 right_df = right_result.dataframe
             right_queries = right_df.queries["queries"]
             if len(right_queries) != 1:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Unexpected number of queries: {len(right_queries)}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
             right_query = right_queries[0]
             input_df_sql = f"WITH __left AS ({left_query}) SELECT * FROM __left INNER JOIN LATERAL ({right_query})"
             session = snowpark.Session.get_active_session()
@@ -133,6 +161,8 @@ def map_extension(
                 + right_result.column_map.get_snowpark_columns(),
                 column_qualifiers=left_result.column_map.get_qualifiers()
                 + right_result.column_map.get_qualifiers(),
+                equivalent_snowpark_names=left_result.column_map.get_equivalent_snowpark_names()
+                + right_result.column_map.get_equivalent_snowpark_names(),
             )
 
         case "udtf_with_table_arguments":
@@ -140,7 +170,11 @@ def map_extension(
         case "aggregate":
             return map_aggregate(extension.aggregate, rel.common.plan_id)
         case other:
-            raise SnowparkConnectNotImplementedError(f"Unexpected extension {other}")
+            exception = SnowparkConnectNotImplementedError(
+                f"Unexpected extension {other}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def get_udtf_project(relation: relation_proto.Relation) -> bool:
@@ -167,7 +201,7 @@ def get_udtf_project(relation: relation_proto.Relation) -> bool:
 
 def handle_udtf_with_table_arguments(
     udtf_info: snowflake_proto.UDTFWithTableArguments,
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Handle UDTF with one or more table arguments using Snowpark's join_table_function.
     For multiple table arguments, this creates a Cartesian product of all input tables.
@@ -175,7 +209,9 @@ def handle_udtf_with_table_arguments(
     session = snowpark.Session.get_active_session()
     udtf_name_lower = udtf_info.function_name.lower()
     if udtf_name_lower not in session._udtfs:
-        raise ValueError(f"UDTF '{udtf_info.function_name}' not found.")
+        exception = ValueError(f"UDTF '{udtf_info.function_name}' not found.")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
     _udtf_obj, udtf_spark_output_names = session._udtfs[udtf_name_lower]
 
     table_containers = []
@@ -189,10 +225,12 @@ def handle_udtf_with_table_arguments(
         if not get_boolean_session_config_param(
             "spark.sql.tvf.allowMultipleTableArguments.enabled"
         ):
-            raise AnalysisException(
+            exception = AnalysisException(
                 "[TABLE_VALUED_FUNCTION_TOO_MANY_TABLE_ARGUMENTS] Multiple table arguments are not enabled. "
                 "Please set `spark.sql.tvf.allowMultipleTableArguments.enabled` to `true`"
             )
+            attach_custom_error_code(exception, ErrorCodes.CONFIG_NOT_ENABLED)
+            raise exception
 
         base_df = table_containers[0][0].dataframe
         first_table_col_count = len(base_df.columns)
@@ -271,7 +309,7 @@ def handle_lateral_join_with_udtf(
     left_result: DataFrameContainer,
     udtf_relation: relation_proto.Relation,
     udtf_info: tuple[snowpark.udtf.UserDefinedTableFunction, list],
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     """
     Handle lateral join with UDTF on the right side using join_table_function.
     """
@@ -298,13 +336,15 @@ def handle_lateral_join_with_udtf(
         + udtf_spark_output_names,
         snowpark_column_names=result_df.columns,
         column_qualifiers=left_result.column_map.get_qualifiers()
-        + [[]] * len(udtf_spark_output_names),
+        + [set() for _ in udtf_spark_output_names],
+        equivalent_snowpark_names=left_result.column_map.get_equivalent_snowpark_names()
+        + [set() for _ in udtf_spark_output_names],
     )
 
 
 def map_aggregate(
     aggregate: snowflake_proto.Aggregate, plan_id: int
-) -> snowpark.DataFrame:
+) -> DataFrameContainer:
     input_container = map_relation(aggregate.input)
     input_df: snowpark.DataFrame = input_container.dataframe
 
@@ -336,18 +376,19 @@ def map_aggregate(
     typer = ExpressionTyper(input_df)
 
     def _map_column(exp: expression_proto.Expression) -> tuple[str, TypedColumn]:
-        with not_resolving_fun_args():
-            new_names, snowpark_column = map_expression(
-                exp, input_container.column_map, typer
+        new_names, snowpark_column = map_expression(
+            exp, input_container.column_map, typer
+        )
+        if len(new_names) != 1:
+            exception = SnowparkConnectNotImplementedError(
+                "Multi-column aggregate expressions are not supported"
             )
-            if len(new_names) != 1:
-                raise SnowparkConnectNotImplementedError(
-                    "Multi-column aggregate expressions are not supported"
-                )
-            return new_names[0], snowpark_column
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
+        return new_names[0], snowpark_column
 
     raw_groupings: list[tuple[str, TypedColumn]] = []
-    raw_aggregations: list[tuple[str, TypedColumn]] = []
+    raw_aggregations: list[tuple[str, TypedColumn, set[ColumnQualifier]]] = []
 
     if not is_group_by_all:
         raw_groupings = [_map_column(exp) for exp in aggregate.grouping_expressions]
@@ -377,10 +418,22 @@ def map_aggregate(
     # Note: We don't clear the map here to preserve any parent context aliases
     from snowflake.snowpark_connect.utils.context import register_lca_alias
 
+    # If it's an unresolved attribute when its in aggregate.aggregate_expressions, we know it came from the parent map straight away
+    # in this case, we should see if the parent map has a qualifier for it and propagate that here, in case the order by references it in
+    # a qualified way later.
     agg_count = get_sql_aggregate_function_count()
     for exp in aggregate.aggregate_expressions:
         col = _map_column(exp)
-        raw_aggregations.append(col)
+        if exp.WhichOneof("expr_type") == "unresolved_attribute":
+            qualifiers: set[
+                ColumnQualifier
+            ] = input_container.column_map.get_qualifiers_for_snowpark_column(
+                col[1].col.get_name()
+            )
+        else:
+            qualifiers = set()
+
+        raw_aggregations.append((col[0], col[1], qualifiers))
 
         # If this is an alias, register it in the LCA map for subsequent expressions
         if (
@@ -411,18 +464,20 @@ def map_aggregate(
     spark_columns: list[str] = []
     snowpark_columns: list[str] = []
     snowpark_column_types: list[snowpark_types.DataType] = []
+    all_qualifiers: list[set[ColumnQualifier]] = []
 
     # Use grouping columns directly without aliases
-    groupings = [col.col for _, col in raw_groupings]
+    groupings: list[Column] = [tc.col for _, tc in raw_groupings]
 
     # Create aliases only for aggregation columns
     aggregations = []
-    for i, (spark_name, snowpark_column) in enumerate(raw_aggregations):
+    for i, (spark_name, snowpark_column, qualifiers) in enumerate(raw_aggregations):
         alias = make_column_names_snowpark_compatible([spark_name], plan_id, i)[0]
 
         spark_columns.append(spark_name)
         snowpark_columns.append(alias)
         snowpark_column_types.append(snowpark_column.typ)
+        all_qualifiers.append(qualifiers)
 
         aggregations.append(snowpark_column.col.alias(alias))
 
@@ -431,12 +486,18 @@ def map_aggregate(
             if groupings:
                 # Normal GROUP BY with explicit grouping columns
                 result = input_df.group_by(groupings)
-            else:
+            elif not is_group_by_all:
                 # No explicit GROUP BY - this is an aggregate over the entire table
                 # Use a dummy constant that will be excluded from the final result
                 result = input_df.with_column(
                     "__dummy_group__", snowpark_fn.lit(1)
                 ).group_by("__dummy_group__")
+            else:
+                # GROUP BY ALL with only one aggregate column
+                # Snowpark doesn't support GROUP BY ALL
+                # TODO: Change in future with Snowpark Supported arguments or API for GROUP BY ALL
+                result = input_df.group_by()
+
         case snowflake_proto.Aggregate.GROUP_TYPE_ROLLUP:
             result = input_df.rollup(groupings)
         case snowflake_proto.Aggregate.GROUP_TYPE_CUBE:
@@ -456,10 +517,148 @@ def map_aggregate(
             result = input_df.group_by_grouping_sets(
                 snowpark.GroupingSets(*sets_mapped)
             )
+        case snowflake_proto.Aggregate.GROUP_TYPE_PIVOT:
+            pivot_typed_columns: list[TypedColumn] = [
+                map_single_column_expression(
+                    pivot_col,
+                    input_container.column_map,
+                    ExpressionTyper(input_df),
+                )[1]
+                for pivot_col in aggregate.pivot.pivot_columns
+            ]
+
+            pivot_columns = [col.col for col in pivot_typed_columns]
+            pivot_column_types = [col.typ for col in pivot_typed_columns]
+
+            pivot_values: list[list[Any]] = []
+            pivot_aliases: list[str] = []
+
+            for pivot_value in aggregate.pivot.pivot_values:
+                current_values = [
+                    get_literal_field_and_name(val)[0] for val in pivot_value.values
+                ]
+                pivot_values.append(current_values)
+
+                if pivot_value.alias:
+                    pivot_aliases.append(pivot_value.alias)
+
+            spark_col_names = []
+            final_pivot_names = []
+            grouping_columns_qualifiers = []
+            aggregations_pivot = []
+
+            pivot_col_names: set[str] = {col.get_name() for col in pivot_columns}
+
+            agg_columns = get_all_dependent_column_names(aggregations)
+
+            if groupings:
+                for col in groupings:
+                    snowpark_name = col.get_name()
+                    spark_col_name = input_container.column_map.get_spark_column_name_from_snowpark_column_name(
+                        snowpark_name
+                    )
+                    qualifiers = (
+                        input_container.column_map.get_qualifiers_for_snowpark_column(
+                            snowpark_name
+                        )
+                    )
+                    grouping_columns_qualifiers.append(qualifiers)
+                    spark_col_names.append(spark_col_name)
+            else:
+                for col in input_container.column_map.columns:
+                    if (
+                        col.snowpark_name not in pivot_col_names
+                        and col.snowpark_name not in agg_columns
+                    ):
+                        groupings.append(snowpark_fn.col(col.snowpark_name))
+                        grouping_columns_qualifiers.append(col.qualifiers)
+                        spark_col_names.append(col.spark_name)
+
+            for pivot_value_idx, pivot_value_group in enumerate(pivot_values):
+                pivot_values_spark_names = []
+                pivot_value_is_null = []
+
+                for val in pivot_value_group:
+                    spark_name, is_null = map_pivot_value_to_spark_column_name(val)
+
+                    pivot_values_spark_names.append(spark_name)
+                    pivot_value_is_null.append(is_null)
+
+                for agg_idx, agg_expression in enumerate(aggregations):
+                    agg_fun_expr = copy.deepcopy(agg_expression._expr1)
+
+                    condition = None
+                    for pivot_col_idx, (pivot_col, pivot_val) in enumerate(
+                        zip(pivot_columns, pivot_value_group)
+                    ):
+                        current_condition = create_pivot_column_condition(
+                            pivot_col,
+                            pivot_val,
+                            pivot_value_is_null[pivot_col_idx],
+                            pivot_column_types[pivot_col_idx]
+                            if isinstance(pivot_val, (list, dict))
+                            else None,
+                        )
+
+                        condition = (
+                            current_condition
+                            if condition is None
+                            else condition & current_condition
+                        )
+
+                    inject_condition_to_all_agg_functions(agg_fun_expr, condition)
+                    curr_expression = Column(agg_fun_expr)
+
+                    if pivot_aliases and not any(pivot_value_is_null):
+                        aliased_pivoted_column_spark_name = pivot_aliases[
+                            pivot_value_idx
+                        ]
+                    elif len(pivot_values_spark_names) > 1:
+                        aliased_pivoted_column_spark_name = (
+                            "{" + ", ".join(pivot_values_spark_names) + "}"
+                        )
+                    else:
+                        aliased_pivoted_column_spark_name = pivot_values_spark_names[0]
+
+                    spark_col_name = (
+                        f"{aliased_pivoted_column_spark_name}_{raw_aggregations[agg_idx][0]}"
+                        if len(aggregations) > 1
+                        else f"{aliased_pivoted_column_spark_name}"
+                    )
+
+                    snowpark_col_name = make_column_names_snowpark_compatible(
+                        [spark_col_name],
+                        plan_id,
+                        len(aggregations) + len(groupings),
+                    )[0]
+
+                    curr_expression = curr_expression.alias(snowpark_col_name)
+
+                    aggregations_pivot.append(curr_expression)
+                    spark_col_names.append(spark_col_name)
+                    final_pivot_names.append(snowpark_col_name)
+
+            result_df = input_df.group_by(*groupings).agg(*aggregations_pivot)
+
+            return DataFrameContainer.create_with_column_mapping(
+                dataframe=result_df,
+                spark_column_names=spark_col_names,
+                snowpark_column_names=result_df.columns,
+                snowpark_column_types=[
+                    result_df.schema.fields[idx].datatype
+                    for idx, _ in enumerate(result_df.columns)
+                ],
+                column_qualifiers=grouping_columns_qualifiers
+                + [set() for _ in final_pivot_names],
+                parent_column_name_map=input_container.column_map,
+            )
+
         case other:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported GROUP BY type: {other}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
     result = result.agg(*aggregations, exclude_grouping_columns=True)
 
@@ -479,6 +678,13 @@ def map_aggregate(
             spark_column_names=spark_columns,
             snowpark_column_names=snowpark_columns,
             snowpark_column_types=snowpark_column_types,
+            column_qualifiers=all_qualifiers,
+            equivalent_snowpark_names=[
+                input_container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                    new_name
+                )
+                for new_name in snowpark_columns
+            ],
         ).column_map
 
         # Create hybrid column map that can resolve both input and aggregate contexts
@@ -490,7 +696,9 @@ def map_aggregate(
             aggregate_expressions=list(aggregate.aggregate_expressions),
             grouping_expressions=list(aggregate.grouping_expressions),
             spark_columns=spark_columns,
-            raw_aggregations=raw_aggregations,
+            raw_aggregations=[
+                (spark_name, col) for spark_name, col, _ in raw_aggregations
+            ],
         )
 
         # Map the HAVING condition using hybrid resolution
@@ -504,11 +712,37 @@ def map_aggregate(
         # grouping sets don't allow ORDER BY with columns that aren't in the aggregate list.
         result = result.select(result.columns[-len(aggregations) :])
 
+    # Store aggregate metadata for ORDER BY resolution
+    # Only for regular GROUP BY - ROLLUP, CUBE, and GROUPING_SETS should NOT allow
+    # ORDER BY to reference pre-aggregation columns (Spark compatibility)
+    # This enables ORDER BY to resolve expressions that reference pre-aggregation columns
+    # (e.g., ORDER BY year(date) when only 'year' alias exists in aggregated result)
+    aggregate_metadata = None
+    if aggregate.group_type == snowflake_proto.Aggregate.GROUP_TYPE_GROUPBY:
+        aggregate_metadata = AggregateMetadata(
+            input_column_map=input_container.column_map,
+            input_dataframe=input_df,
+            grouping_expressions=list(aggregate.grouping_expressions),
+            aggregate_expressions=list(aggregate.aggregate_expressions),
+            spark_columns=spark_columns,
+            raw_aggregations=[
+                (spark_name, col) for spark_name, col, _ in raw_aggregations
+            ],
+        )
+
     # Return only aggregation columns in the column map
     return DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=spark_columns,
         snowpark_column_names=snowpark_columns,
         snowpark_column_types=snowpark_column_types,
-        parent_column_name_map=input_df._column_map,
+        parent_column_name_map=input_container.column_map,
+        column_qualifiers=all_qualifiers,
+        equivalent_snowpark_names=[
+            input_container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                new_name
+            )
+            for new_name in snowpark_columns
+        ],
+        aggregate_metadata=aggregate_metadata,
     )