Fix allowing multiple aggregates in Pivot fix (#3171)

Co-authored-by: Jamison Rose <[email protected]>
Co-authored-by: Afroz Alam <[email protected]>

Author: 3 people

src/snowflake/snowpark/_internal/analyzer/analyzer.py

@@ -50,9 +50,10 @@
     BinaryExpression,
 )
 from snowflake.snowpark._internal.analyzer.binary_plan_node import (
+    FullOuter,
     Join,
     SetOperation,
-    Union as UnionPlan,
+    UsingJoin,
 )
 from snowflake.snowpark._internal.analyzer.datatype_mapper import (
     numeric_to_sql_without_cast,
@@ -1167,7 +1168,10 @@ def do_resolve_with_resolved_children(
                 pivot_values = None
 
             plan = None
+
             for agg_expr in logical_plan.aggregates:
+                # We only allow pivot on more than one aggregates when it on a groupby clause
+                join_columns: List[str] | None = None
                 if (
                     len(logical_plan.grouping_columns) != 0
                     and agg_expr.children is not None
@@ -1186,6 +1190,10 @@ def do_resolve_with_resolved_children(
                         ],  # aggregate column is first child in logical_plan.aggregates
                         logical_plan.pivot_column,
                     ]
+                    join_columns = [
+                        self.analyze(expression, df_aliased_col_name_to_real_col_name)
+                        for expression in logical_plan.grouping_columns
+                    ]
                     child = self.plan_builder.project(
                         [
                             self.analyze(col, df_aliased_col_name_to_real_col_name)
@@ -1202,9 +1210,7 @@ def do_resolve_with_resolved_children(
                         logical_plan.pivot_column, df_aliased_col_name_to_real_col_name
                     ),
                     pivot_values,
-                    self.analyze(
-                        logical_plan.aggregates[0], df_aliased_col_name_to_real_col_name
-                    ),
+                    self.analyze(agg_expr, df_aliased_col_name_to_real_col_name),
                     self.analyze(
                         logical_plan.default_on_null,
                         df_aliased_col_name_to_real_col_name,
@@ -1213,6 +1219,8 @@ def do_resolve_with_resolved_children(
                     else None,
                     child,
                     logical_plan,
+                    len(logical_plan.aggregates)
+                    > 1,  # we need to alias the names with agg function when we have more than one agg functions on the pivot
                 )
 
                 # If this is a dynamic pivot, then we can't use child.schema_query which is used in the schema_query
@@ -1225,15 +1233,35 @@ def do_resolve_with_resolved_children(
                     # table as it may not exist at later point in time when dataframe.schema is called.
                     pivot_plan.schema_query = pivot_plan.queries[-1].sql
 
-                # union multiple aggregations
-                # https://docs.snowflake.com/en/sql-reference/constructs/pivot#dynamic-pivot-with-multiple-aggregations-using-union
+                # using join here to have the output similar to what spark have
+                # both the aggregations are happening over the same set of columns and pivot values
+                # we will receive left and right both pivot table with same set of groupby columns and columns corresponding to pivot values
+                # to differentiate between columns corresponding to pivot values for aggregation function they will have name suffixed by agg fun
+                # join would keep the group by column same and append the columns corresponding to pivot values for multiple agg functions
+                # output would look similar to below for a statement like
+                # df.groupBy("name").pivot("department", ["Sales", "Marketing"]).sum("year", "salary").show()
+                # +-------+---------------+---------------+-------------------+-------------------+
+                # |   name|Sales_sum(year)|Sales_sum(year)|Marketing_sum(year)|Marketing_sum(year)|
+                # +-------+---------------+---------------+-------------------+-------------------+
+                # |  Scott|           NULL|           NULL|               NULL|               NULL|
+                # |  James|           4039|           4039|               NULL|               NULL|
+                # |    Jen|           NULL|           NULL|               NULL|               NULL|
+                # |Michael|           2020|           2020|               NULL|               NULL|
+
                 if plan is None:
                     plan = pivot_plan
-                else:
-                    union_plan = UnionPlan(plan, pivot_plan, is_all=False)
-                    plan = self.plan_builder.set_operator(
-                        plan, pivot_plan, union_plan.sql, union_plan
+                elif join_columns is not None:
+                    plan = self.plan_builder.join(
+                        plan,
+                        pivot_plan,
+                        UsingJoin(FullOuter(), join_columns),
+                        "",
+                        "",
+                        logical_plan,
+                        self.session.conf.get("use_constant_subquery_alias", False),
                     )
+                # we have a check in relational_grouped_dataframe.py which will prevent a case where there are more than one aggregate
+                # without having a grouping condition which is essential to create join_columns
 
             assert plan is not None
             return plan

src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py

@@ -46,6 +46,7 @@
     is_sql_select_statement,
     quote_name,
     random_name_for_temp_object,
+    unwrap_single_quote,
 )
 from snowflake.snowpark.row import Row
 from snowflake.snowpark.types import DataType
@@ -63,6 +64,7 @@
 LEFT_BRACKET = "["
 RIGHT_BRACKET = "]"
 AS = " AS "
+EXCLUDE = " EXCLUDE "
 AND = " AND "
 OR = " OR "
 NOT = " NOT "
@@ -1248,7 +1250,9 @@ def pivot_statement(
     aggregate: str,
     default_on_null: Optional[str],
     child: str,
+    should_alias_column_with_agg: bool,
 ) -> str:
+    select_str = STAR
     if isinstance(pivot_values, str):
         # The subexpression in this case already includes parenthesis.
         values_str = pivot_values
@@ -1258,10 +1262,24 @@ def pivot_statement(
             + (ANY if pivot_values is None else COMMA.join(pivot_values))
             + RIGHT_PARENTHESIS
         )
+        if pivot_values is not None and should_alias_column_with_agg:
+            quoted_names = [quote_name(value) for value in pivot_values]
+            # unwrap_single_quote on the value to match the output closer to what spark generates
+            aliased_names = [
+                quote_name(f"{unwrap_single_quote(value)}_{aggregate}")
+                for value in pivot_values
+            ]
+            aliased_string = [
+                f"{quoted_name}{AS}{aliased_name}"
+                for aliased_name, quoted_name in zip(aliased_names, quoted_names)
+            ]
+            exclude_str = COMMA.join(quoted_names)
+            aliased_str = COMMA.join(aliased_string)
+            select_str = f"{STAR}{EXCLUDE}{LEFT_PARENTHESIS}{exclude_str}{RIGHT_PARENTHESIS}, {aliased_str}"
 
     return (
         SELECT
-        + STAR
+        + select_str
         + FROM
         + LEFT_PARENTHESIS
         + child

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

@@ -1105,10 +1105,16 @@ def pivot(
         default_on_null: Optional[str],
         child: SnowflakePlan,
         source_plan: Optional[LogicalPlan],
+        should_alias_column_with_agg: bool,
     ) -> SnowflakePlan:
         return self.build(
             lambda x: pivot_statement(
-                pivot_column, pivot_values, aggregate, default_on_null, x
+                pivot_column,
+                pivot_values,
+                aggregate,
+                default_on_null,
+                x,
+                should_alias_column_with_agg,
             ),
             child,
             source_plan,

src/snowflake/snowpark/_internal/error_message.py

@@ -137,6 +137,14 @@ def DF_CROSS_TAB_COUNT_TOO_LARGE(
             error_code="1107",
         )
 
+    @staticmethod
+    def DF_PIVOT_ONLY_SUPPORT_ONE_AGG_EXPR() -> SnowparkDataframeException:
+        return SnowparkDataframeException(
+            "You can apply only one aggregate expression to a RelationalGroupedDataFrame "
+            "returned by the pivot() method unless the pivot is applied with a groupby clause.",
+            error_code="1109",
+        )
+
     @staticmethod
     def DF_DATAFRAME_IS_NOT_QUALIFIED_FOR_SCALAR_QUERY(
         count: int, columns: str

src/snowflake/snowpark/relational_grouped_dataframe.py

@@ -5,6 +5,7 @@
 from typing import Callable, Dict, Iterable, List, Optional, Tuple, Union
 import inspect
 
+from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages
 import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto
 import snowflake.snowpark.context as context
 from snowflake.connector.options import pandas
@@ -225,6 +226,8 @@ def _to_df(
                 self._dataframe._select_statement or self._dataframe._plan,
             )
         elif isinstance(self._group_type, _PivotType):
+            if len(agg_exprs) != 1 and len(unaliased_grouping) == 0:
+                raise SnowparkClientExceptionMessages.DF_PIVOT_ONLY_SUPPORT_ONE_AGG_EXPR()
             group_plan = Pivot(
                 unaliased_grouping,
                 self._group_type.pivot_col,

tests/integ/scala/test_dataframe_aggregate_suite.py

@@ -5,6 +5,7 @@
 
 from decimal import Decimal
 from math import sqrt
+import re
 from typing import NamedTuple
 
 import pytest
@@ -15,6 +16,7 @@
 )
 from snowflake.snowpark.column import Column
 from snowflake.snowpark.exceptions import (
+    SnowparkDataframeException,
     SnowparkSQLException,
 )
 from snowflake.snowpark.functions import (
@@ -419,39 +421,80 @@ class MonthlySales(NamedTuple):
     reason="Multiple aggregations are not supported in local testing mode",
 )
 def test_pivot_multiple_aggs(session):
-    # 1) SUM and AVG
-    Utils.check_answer(
+    with pytest.raises(
+        SnowparkDataframeException,
+        match=re.escape(
+            "You can apply only one aggregate expression to a RelationalGroupedDataFrame returned by the pivot() method unless the pivot is applied with a groupby clause."
+        ),
+    ):
+        TestData.monthly_sales(session).pivot(
+            "month", ["JAN", "FEB", "MAR", "APR"]
+        ).agg([sum(col("amount")), avg(col("amount"))]).sort(col("empid"))
+
+    df = (
         TestData.monthly_sales(session)
+        .groupBy(col("empid"))
         .pivot("month", ["JAN", "FEB", "MAR", "APR"])
         .agg([sum(col("amount")), avg(col("amount"))])
-        .sort(col("empid")),
+        .sort(col("empid"))
+    )
+
+    assert [f.name for f in df.schema.fields] == [
+        "EMPID",
+        '"JAN_sum(""AMOUNT"")"',
+        '"FEB_sum(""AMOUNT"")"',
+        '"MAR_sum(""AMOUNT"")"',
+        '"APR_sum(""AMOUNT"")"',
+        '"JAN_avg(""AMOUNT"")"',
+        '"FEB_avg(""AMOUNT"")"',
+        '"MAR_avg(""AMOUNT"")"',
+        '"APR_avg(""AMOUNT"")"',
+    ]
+
+    Utils.check_answer(
+        df,
         [
-            Row(1, 10400, 8000, 11000, 18000),
-            Row(2, 39500, 90700, 12000, 5300),
+            Row(1, 10400, 8000, 11000, 18000, 5200.0, 4000.0, 5500.0, 9000.0),
+            Row(
+                2,
+                39500,
+                90700,
+                12000,
+                5300,
+                19750.0,
+                45350.0,
+                6000.0,
+                2650.0,
+            ),
         ],
     )
 
-    # 2) MIN and MAX
-    Utils.check_answer(
+    df = (
         TestData.monthly_sales(session)
+        .groupBy(col("empid"))
         .pivot("month", ["JAN", "FEB", "MAR", "APR"])
         .agg([min(col("amount")), max(col("amount"))])
-        .sort(col("empid")),
-        [
-            Row(1, 400, 3000, 5000, 8000),
-            Row(2, 4500, 200, 2500, 800),
-        ],
+        .sort(col("empid"))
     )
 
-    # 3) AVG and COUNT_DISTINCT
+    assert [f.name for f in df.schema.fields] == [
+        "EMPID",
+        '"JAN_min(""AMOUNT"")"',
+        '"FEB_min(""AMOUNT"")"',
+        '"MAR_min(""AMOUNT"")"',
+        '"APR_min(""AMOUNT"")"',
+        '"JAN_max(""AMOUNT"")"',
+        '"FEB_max(""AMOUNT"")"',
+        '"MAR_max(""AMOUNT"")"',
+        '"APR_max(""AMOUNT"")"',
+    ]
+
+    # 2) MIN and MAX
     Utils.check_answer(
-        TestData.monthly_sales(session)
-        .pivot("month", ["JAN", "FEB", "MAR", "APR"])
-        .agg([avg(col("amount")), count_distinct(col("amount"))])
-        .sort(col("empid")),
+        df,
         [
-            Row(1, 5200, 4000, 5500, 9000),
-            Row(2, 19750, 45350, 6000, 2650),
+            Row(1, 400, 3000, 5000, 8000, 10000, 5000, 6000, 10000),
+            Row(2, 4500, 200, 2500, 800, 35000, 90500, 9500, 4500),
         ],
     )