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

snowflake/snowpark_connect/relation/map_aggregate.py

@@ -2,9 +2,8 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
-import re
+import copy
 from dataclasses import dataclass
-from typing import Optional
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
@@ -16,18 +15,26 @@ from snowflake.snowpark.types import DataType
 from snowflake.snowpark_connect.column_name_handler import (
     make_column_names_snowpark_compatible,
 )
-from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
+from snowflake.snowpark_connect.dataframe_container import (
+    AggregateMetadata,
+    DataFrameContainer,
+)
 from snowflake.snowpark_connect.expression.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_expression import (
     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,
+    map_pivot_value_to_spark_column_name,
+)
 from snowflake.snowpark_connect.typed_column import TypedColumn
+from snowflake.snowpark_connect.utils import expression_transformer
 from snowflake.snowpark_connect.utils.context import (
-    get_is_evaluating_sql,
+    grouping_by_scala_udf_key,
     set_current_grouping_columns,
-    temporary_pivot_expression,
 )
 
 
@@ -49,6 +56,20 @@ def map_group_by_aggregate(
         result = input_df_actual.group_by(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
+
+    # Store aggregate metadata for ORDER BY resolution
+    aggregate_metadata = AggregateMetadata(
+        input_column_map=input_df_container.column_map,
+        input_dataframe=input_df_actual,
+        grouping_expressions=list(rel.aggregate.grouping_expressions),
+        aggregate_expressions=list(rel.aggregate.aggregate_expressions),
+        spark_columns=columns.spark_names(),
+        raw_aggregations=[
+            (col.spark_name, TypedColumn(col.expression, col.data_type))
+            for col in columns.aggregation_columns
+        ],
+    )
+
     return DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=columns.spark_names(),
@@ -56,6 +77,8 @@ def map_group_by_aggregate(
         snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
         parent_column_name_map=input_df_container.column_map,
+        equivalent_snowpark_names=columns.get_equivalent_snowpark_names(),
+        aggregate_metadata=aggregate_metadata,
     )
 
 
@@ -77,6 +100,11 @@ def map_rollup_aggregate(
         result = input_df_actual.rollup(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
+
+    # NOTE: Do NOT attach aggregate_metadata for ROLLUP
+    # Spark does not allow ORDER BY to reference pre-aggregation columns for ROLLUP
+    # Only regular GROUP BY supports this
+
     return DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=columns.spark_names(),
@@ -84,6 +112,7 @@ def map_rollup_aggregate(
         snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
         parent_column_name_map=input_container.column_map,
+        equivalent_snowpark_names=columns.get_equivalent_snowpark_names(),
     )
 
 
@@ -105,6 +134,11 @@ def map_cube_aggregate(
         result = input_df_actual.cube(*columns.grouping_expressions()).agg(
             *columns.aggregation_expressions()
         )
+
+    # NOTE: Do NOT attach aggregate_metadata for CUBE
+    # Spark does not allow ORDER BY to reference pre-aggregation columns for CUBE
+    # Only regular GROUP BY supports this
+
     return DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=columns.spark_names(),
@@ -112,6 +146,7 @@ def map_cube_aggregate(
         snowpark_column_types=columns.data_types(),
         column_qualifiers=columns.get_qualifiers(),
         parent_column_name_map=input_container.column_map,
+        equivalent_snowpark_names=columns.get_equivalent_snowpark_names(),
     )
 
 
@@ -136,220 +171,111 @@ def map_pivot_aggregate(
         get_literal_field_and_name(lit)[0] for lit in rel.aggregate.pivot.values
     ]
 
-    used_columns = {pivot_column[1].col._expression.name}
-    if get_is_evaluating_sql():
-        # When evaluating SQL spark doesn't trim columns from the result
-        used_columns = {"*"}
-    else:
-        for expression in rel.aggregate.aggregate_expressions:
-            matched_identifiers = re.findall(
-                r'unparsed_identifier: "(.*)"', expression.__str__()
-            )
-            for identifier in matched_identifiers:
-                mapped_col = input_container.column_map.spark_to_col.get(
-                    identifier, None
-                )
-                if mapped_col:
-                    used_columns.add(mapped_col[0].snowpark_name)
+    if not pivot_values:
+        distinct_col_values = (
+            input_df_actual.select(pivot_column[1].col)
+            .distinct()
+            .sort(snowpark_fn.asc_nulls_first(pivot_column[1].col))
+            .collect()
+        )
+        pivot_values = [
+            row[0].as_dict() if isinstance(row[0], snowpark.Row) else row[0]
+            for row in distinct_col_values
+        ]
 
-    if len(columns.grouping_expressions()) == 0:
-        # Snowpark doesn't support multiple aggregations in pivot without groupBy
-        # So we need to perform each aggregation separately and then combine results
-        if len(columns.aggregation_expressions(unalias=True)) > 1:
-            agg_expressions = columns.aggregation_expressions(unalias=True)
-            agg_metadata = columns.aggregation_columns
-            num_agg_functions = len(agg_expressions)
-
-            spark_names = []
-            pivot_results = []
-            for i, agg_expr in enumerate(agg_expressions):
-                pivot_result = (
-                    input_df_actual.select(*used_columns)
-                    .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
-                    .agg(agg_expr)
+    agg_expressions = columns.aggregation_expressions(unalias=True)
+
+    spark_col_names = []
+    aggregations = []
+    final_pivot_names = []
+    grouping_columns_qualifiers = []
+    grouping_eq_snowpark_names = []
+
+    grouping_columns = columns.grouping_expressions()
+    if grouping_columns:
+        for col in grouping_columns:
+            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)
+            grouping_eq_snowpark_names.append(
+                input_container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                    snowpark_name
                 )
-                for col_name in pivot_result.columns:
-                    spark_names.append(
-                        f"{pivot_column_name(col_name)}_{agg_metadata[i].spark_name}"
-                    )
-                pivot_results.append(pivot_result)
-
-            result = pivot_results[0]
-            for pivot_result in pivot_results[1:]:
-                result = result.cross_join(pivot_result)
-
-            pivot_columns_per_agg = len(pivot_results[0].columns)
-            reordered_spark_names = []
-            reordered_snowpark_names = []
-            reordered_types = []
-            column_selectors = []
-
-            for pivot_idx in range(pivot_columns_per_agg):
-                for agg_idx in range(num_agg_functions):
-                    current_pos = agg_idx * pivot_columns_per_agg + pivot_idx
-                    if current_pos < len(spark_names):
-                        idx = current_pos + 1  # 1-based indexing for Snowpark
-                        reordered_spark_names.append(spark_names[current_pos])
-                        reordered_snowpark_names.append(f"${idx}")
-                        reordered_types.append(
-                            result.schema.fields[current_pos].datatype
-                        )
-                        column_selectors.append(snowpark_fn.col(f"${idx}"))
-
-            return DataFrameContainer.create_with_column_mapping(
-                dataframe=result.select(*column_selectors),
-                spark_column_names=reordered_spark_names,
-                snowpark_column_names=reordered_snowpark_names,
-                column_qualifiers=[[]] * len(reordered_spark_names),
-                parent_column_name_map=input_container.column_map,
-                snowpark_column_types=reordered_types,
             )
-        else:
-            result = (
-                input_df_actual.select(*used_columns)
-                .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
-                .agg(*columns.aggregation_expressions(unalias=True))
+
+    for pv_value in pivot_values:
+        pv_value_spark, pv_is_null = map_pivot_value_to_spark_column_name(pv_value)
+
+        for i, agg_expression in enumerate(agg_expressions):
+            agg_fun_expr = copy.deepcopy(agg_expression._expr1)
+
+            condition = create_pivot_column_condition(
+                pivot_column[1].col,
+                pv_value,
+                pv_is_null,
+                pivot_column[1].typ if isinstance(pv_value, (list, dict)) else None,
             )
-    else:
-        result = (
-            input_df_actual.group_by(*columns.grouping_expressions())
-            .pivot(pivot_column[1].col, pivot_values if pivot_values else None)
-            .agg(*columns.aggregation_expressions(unalias=True))
-        )
 
-    agg_name_list = [c.spark_name for c in columns.grouping_columns]
+            expression_transformer.inject_condition_to_all_agg_functions(
+                agg_fun_expr, condition
+            )
 
-    # Calculate number of pivot values for proper Spark-compatible indexing
-    total_pivot_columns = len(result.columns) - len(agg_name_list)
-    num_pivot_values = (
-        total_pivot_columns // len(columns.aggregation_columns)
-        if len(columns.aggregation_columns) > 0
-        else 1
-    )
+            curr_expression = Column(agg_fun_expr)
 
-    def _get_agg_exp_alias_for_col(col_index: int) -> Optional[str]:
-        if col_index < len(agg_name_list) or len(columns.aggregation_columns) <= 1:
-            return None
-        else:
-            index = (col_index - len(agg_name_list)) // num_pivot_values
-            return columns.aggregation_columns[index].spark_name
-
-    spark_columns = []
-    for col in [
-        pivot_column_name(c, _get_agg_exp_alias_for_col(i))
-        for i, c in enumerate(result.columns)
-    ]:
-        spark_col = (
-            input_container.column_map.get_spark_column_name_from_snowpark_column_name(
-                col, allow_non_exists=True
+            spark_col_name = (
+                f"{pv_value_spark}_{columns.aggregation_columns[i].spark_name}"
+                if len(agg_expressions) > 1
+                else f"{pv_value_spark}"
             )
-        )
 
-        if spark_col is not None:
-            spark_columns.append(spark_col)
-        else:
-            # Handle NULL column names to match Spark behavior (lowercase 'null')
-            if col == "NULL":
-                spark_columns.append(col.lower())
-            else:
-                spark_columns.append(col)
-
-    grouping_cols_count = len(agg_name_list)
-    pivot_cols = result.columns[grouping_cols_count:]
-    spark_pivot_cols = spark_columns[grouping_cols_count:]
-
-    num_agg_functions = len(columns.aggregation_columns)
-    num_pivot_values = len(pivot_cols) // num_agg_functions
-
-    reordered_snowpark_cols = []
-    reordered_spark_cols = []
-    column_indices = []  # 1-based indexing
-
-    for i in range(grouping_cols_count):
-        reordered_snowpark_cols.append(result.columns[i])
-        reordered_spark_cols.append(spark_columns[i])
-        column_indices.append(i + 1)
-
-    for pivot_idx in range(num_pivot_values):
-        for agg_idx in range(num_agg_functions):
-            current_pos = agg_idx * num_pivot_values + pivot_idx
-            if current_pos < len(pivot_cols):
-                reordered_snowpark_cols.append(pivot_cols[current_pos])
-                reordered_spark_cols.append(spark_pivot_cols[current_pos])
-                original_index = grouping_cols_count + current_pos
-                column_indices.append(original_index + 1)
-
-    reordered_result = result.select(
-        *[snowpark_fn.col(f"${idx}") for idx in column_indices]
+            snowpark_col_name = make_column_names_snowpark_compatible(
+                [spark_col_name],
+                rel.common.plan_id,
+                len(grouping_columns) + len(agg_expressions),
+            )[0]
+
+            curr_expression = curr_expression.alias(snowpark_col_name)
+
+            aggregations.append(curr_expression)
+            spark_col_names.append(spark_col_name)
+            final_pivot_names.append(snowpark_col_name)
+
+    result_df = (
+        input_df_actual.group_by(*grouping_columns)
+        .agg(*aggregations)
+        .select(*grouping_columns, *final_pivot_names)
     )
 
     return DataFrameContainer.create_with_column_mapping(
-        dataframe=reordered_result,
-        spark_column_names=reordered_spark_cols,
-        snowpark_column_names=[f"${idx}" for idx in column_indices],
-        column_qualifiers=(
-            columns.get_qualifiers()[: len(agg_name_list)]
-            + [[]] * (len(reordered_spark_cols) - len(agg_name_list))
-        ),
-        parent_column_name_map=input_container.column_map,
+        dataframe=result_df,
+        spark_column_names=spark_col_names,
+        snowpark_column_names=result_df.columns,
         snowpark_column_types=[
-            result.schema.fields[idx - 1].datatype for idx in column_indices
+            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,
+        equivalent_snowpark_names=grouping_eq_snowpark_names
+        + [set() for _ in final_pivot_names],
     )
 
 
-def pivot_column_name(snowpark_cname, opt_alias: Optional[str] = None) -> Optional[str]:
-    # For values that are used as pivoted columns, the input and output are in the following format (outermost double quotes are part of the input):
-
-    # 1. "'Java'" -> Java
-    # 2. "'""C++""'" -> "C++"
-    # 3. "'""""''Scala''""""'" -> ""'Scala'""
-
-    # As we can see:
-    # 1. the whole content is always nested in a double quote followed by a single quote ("'<content>'").
-    # 2. the string content is nested in single quotes ('<string_content>')
-    # 3. double quote is escased by another double quote, this is snowflake behavior
-    # 4. if there is a single quote followed by a single quote, the first single quote needs to be preserved in the output
-
-    try:
-        # handling values that are used as pivoted columns
-        match = re.match(r'^"\'(.*)\'"$', snowpark_cname)
-        # extract the content between the outermost double quote followed by a single quote "'
-        content = match.group(1)
-        # convert the escaped double quote to the actual double quote
-        content = content.replace('""', '"')
-        escape_single_quote_placeholder = "__SAS_PLACEHOLDER_ESCAPE_SINGLE_QUOTE__"
-        # replace two consecutive single quote in the content with a placeholder, the first single quote needs to be preserved
-        content = re.sub(r"''", escape_single_quote_placeholder, content)
-        # remove the solo single quote, they are not part of the string content
-        content = re.sub(r"'", "", content)
-        # replace the placeholder with the single quote which we want to preserve
-        result = content.replace(escape_single_quote_placeholder, "'")
-        return f"{result}_{opt_alias}" if opt_alias else result
-    except Exception:
-        # fallback to the original logic, handling aliased column names
-        double_quote_list = re.findall(r'"(.*?)"', snowpark_cname)
-        spark_string = ""
-        for entry in list(filter(None, double_quote_list)):
-            if "'" in entry:
-                entry = entry.replace("'", "")
-                if len(entry) > 0:
-                    spark_string += entry
-            elif entry.isdigit() or re.compile(r"^\d+?\.\d+?$").match(entry):
-                # skip quoting digits or decimal numbers as column names.
-                spark_string += entry
-            else:
-                spark_string += '"' + entry + '"'
-        return snowpark_cname if spark_string == "" else spark_string
-
-
 @dataclass(frozen=True)
 class _ColumnMetadata:
     expression: snowpark.Column
     spark_name: str
     snowpark_name: str
     data_type: DataType
-    qualifiers: list[str]
+    qualifiers: set[ColumnQualifier]
+    equivalent_snowpark_names: set[str]
 
 
 @dataclass(frozen=True)
@@ -385,7 +311,7 @@ class _Columns:
             col.spark_name for col in self.grouping_columns + self.aggregation_columns
         ]
 
-    def get_qualifiers(self) -> list[list[str]]:
+    def get_qualifiers(self) -> list[set[ColumnQualifier]]:
         return [
             col.qualifiers for col in self.grouping_columns + self.aggregation_columns
         ]
@@ -399,6 +325,12 @@ class _Columns:
             if col.data_type is not None
         ]
 
+    def get_equivalent_snowpark_names(self) -> list[set[str]]:
+        return [
+            col.equivalent_snowpark_names
+            for col in self.grouping_columns + self.aggregation_columns
+        ]
+
 
 def map_aggregate_helper(
     rel: relation_proto.Relation, pivot: bool = False, skip_alias: bool = False
@@ -413,71 +345,80 @@ def map_aggregate_helper(
     typer = ExpressionTyper(input_df)
     schema_inferrable = True
 
-    with temporary_pivot_expression(pivot):
-        for exp in grouping_expressions:
+    for exp in grouping_expressions:
+        with grouping_by_scala_udf_key(
+            exp.WhichOneof("expr_type") == "common_inline_user_defined_function"
+            and exp.common_inline_user_defined_function.scalar_scala_udf is not None
+        ):
             new_name, snowpark_column = map_single_column_expression(
                 exp, input_container.column_map, typer
             )
-            alias = make_column_names_snowpark_compatible(
-                [new_name], rel.common.plan_id, len(groupings)
-            )[0]
-            groupings.append(
-                _ColumnMetadata(
-                    snowpark_column.col
-                    if skip_alias
-                    else snowpark_column.col.alias(alias),
-                    new_name,
-                    None if skip_alias else alias,
-                    None if pivot else snowpark_column.typ,
-                    snowpark_column.get_qualifiers(),
-                )
-            )
 
-        grouping_cols = [g.spark_name for g in groupings]
-        set_current_grouping_columns(grouping_cols)
+        alias = make_column_names_snowpark_compatible(
+            [new_name], rel.common.plan_id, len(groupings)
+        )[0]
 
-        for exp in expressions:
-            new_name, snowpark_column = map_single_column_expression(
-                exp, input_container.column_map, typer
+        equivalent_snowpark_names = (
+            input_container.column_map.get_equivalent_snowpark_names_for_snowpark_name(
+                snowpark_column.col.get_name()
             )
-            alias = make_column_names_snowpark_compatible(
-                [new_name], rel.common.plan_id, len(groupings) + len(aggregations)
-            )[0]
+        )
 
-            def type_agg_expr(
-                agg_exp: TypedColumn, schema_inferrable: bool
-            ) -> DataType | None:
-                if pivot or not schema_inferrable:
-                    return None
-                try:
-                    return agg_exp.typ
-                except Exception:
-                    # This type used for schema inference optimization purposes.
-                    # typer may not be able to infer the type of some expressions
-                    # in that case we return None, and the optimization will not be applied.
-                    return None
-
-            agg_col_typ = type_agg_expr(snowpark_column, schema_inferrable)
-            if agg_col_typ is None:
-                schema_inferrable = False
-
-            aggregations.append(
-                _ColumnMetadata(
-                    snowpark_column.col
-                    if skip_alias
-                    else snowpark_column.col.alias(alias),
-                    new_name,
-                    None if skip_alias else alias,
-                    agg_col_typ,
-                    [],
-                )
+        groupings.append(
+            _ColumnMetadata(
+                snowpark_column.col if skip_alias else snowpark_column.col.alias(alias),
+                new_name,
+                None if skip_alias else alias,
+                None if pivot else snowpark_column.typ,
+                qualifiers=snowpark_column.get_qualifiers(),
+                equivalent_snowpark_names=equivalent_snowpark_names,
             )
+        )
+
+    grouping_cols = [g.spark_name for g in groupings]
+    set_current_grouping_columns(grouping_cols)
 
-        return (
-            input_container,
-            _Columns(
-                grouping_columns=groupings,
-                aggregation_columns=aggregations,
-                can_infer_schema=schema_inferrable,
-            ),
+    for exp in expressions:
+        new_name, snowpark_column = map_single_column_expression(
+            exp, input_container.column_map, typer
+        )
+        alias = make_column_names_snowpark_compatible(
+            [new_name], rel.common.plan_id, len(groupings) + len(aggregations)
+        )[0]
+
+        def type_agg_expr(
+            agg_exp: TypedColumn, schema_inferrable: bool
+        ) -> DataType | None:
+            if pivot or not schema_inferrable:
+                return None
+            try:
+                return agg_exp.typ
+            except Exception:
+                # This type used for schema inference optimization purposes.
+                # typer may not be able to infer the type of some expressions
+                # in that case we return None, and the optimization will not be applied.
+                return None
+
+        agg_col_typ = type_agg_expr(snowpark_column, schema_inferrable)
+        if agg_col_typ is None:
+            schema_inferrable = False
+
+        aggregations.append(
+            _ColumnMetadata(
+                snowpark_column.col if skip_alias else snowpark_column.col.alias(alias),
+                new_name,
+                None if skip_alias else alias,
+                agg_col_typ,
+                qualifiers=set(),
+                equivalent_snowpark_names=set(),
+            )
         )
+
+    return (
+        input_container,
+        _Columns(
+            grouping_columns=groupings,
+            aggregation_columns=aggregations,
+            can_infer_schema=schema_inferrable,
+        ),
+    )

snowflake/snowpark_connect/relation/map_catalog.py

@@ -8,6 +8,8 @@ import pandas
 import pyspark.sql.connect.proto.catalog_pb2 as catalog_proto
 
 from snowflake.snowpark_connect.dataframe_container import 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.relation.catalogs import CATALOGS
 from snowflake.snowpark_connect.relation.catalogs.utils import (
     CURRENT_CATALOG_NAME,
@@ -148,4 +150,6 @@ def map_catalog(
             return get_current_catalog().uncacheTable(rel.uncache_table.table_name)
         case other:
             # TODO: list_function implementation is blocked on SNOW-1787268
-            raise SnowparkConnectNotImplementedError(f"Other Relation {other}")
+            exception = SnowparkConnectNotImplementedError(f"Other Relation {other}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception