SNOW-2205628: Exclude grouping columns in aggregation (#3585)

Author: sfc-gh-jdu

src/snowflake/snowpark/relational_grouped_dataframe.py

@@ -180,21 +180,25 @@ def _to_df(
         agg_exprs: List[Expression],
         _ast_stmt: Optional[proto.Bind] = None,
         _emit_ast: bool = False,
+        **kwargs,
     ) -> DataFrame:
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
         aliased_agg = []
-        for grouping_expr in self._grouping_exprs:
-            if isinstance(grouping_expr, GroupingSetsExpression):
-                # avoid doing list(set(grouping_expr.args)) because it will change the order
-                gr_used = set()
-                gr_uniq = [
-                    a
-                    for arg in grouping_expr.args
-                    for a in arg
-                    if a not in gr_used and (gr_used.add(a) or True)
-                ]
-                aliased_agg.extend(gr_uniq)
-            else:
-                aliased_agg.append(grouping_expr)
+
+        if not exclude_grouping_columns:
+            for grouping_expr in self._grouping_exprs:
+                if isinstance(grouping_expr, GroupingSetsExpression):
+                    # avoid doing list(set(grouping_expr.args)) because it will change the order
+                    gr_used = set()
+                    gr_uniq = [
+                        a
+                        for arg in grouping_expr.args
+                        for a in arg
+                        if a not in gr_used and (gr_used.add(a) or True)
+                    ]
+                    aliased_agg.extend(gr_uniq)
+                else:
+                    aliased_agg.append(grouping_expr)
 
         aliased_agg.extend(agg_exprs)
 
@@ -263,6 +267,7 @@ def agg(
         *exprs: Union[Column, Tuple[ColumnOrName, str], Dict[str, str]],
         _ast_stmt: Optional[proto.Bind] = None,
         _emit_ast: bool = True,
+        **kwargs,
     ) -> DataFrame:
         """Returns a :class:`DataFrame` with computed aggregates. See examples in :meth:`DataFrame.group_by`.
 
@@ -283,6 +288,7 @@ def agg(
             - :meth:`DataFrame.agg`
             - :meth:`DataFrame.group_by`
         """
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
 
         exprs, is_variadic = parse_positional_args_to_list_variadic(*exprs)
 
@@ -323,7 +329,11 @@ def agg(
                     )
                     agg_exprs.append(_str_to_expr(e[1], _emit_ast)(col_expr))
 
-        df = self._to_df(agg_exprs, _emit_ast=False)
+        df = self._to_df(
+            agg_exprs,
+            exclude_grouping_columns=exclude_grouping_columns,
+            _emit_ast=False,
+        )
         df._ops_after_agg = set()
 
         if _emit_ast:
@@ -649,40 +659,93 @@ def pivot(
 
     @relational_group_df_api_usage
     @publicapi
-    def avg(self, *cols: ColumnOrName, _emit_ast: bool = True) -> DataFrame:
-        """Return the average for the specified numeric columns."""
-        return self._non_empty_argument_function("avg", *cols, _emit_ast=_emit_ast)
+    def avg(self, *cols: ColumnOrName, _emit_ast: bool = True, **kwargs) -> DataFrame:
+        """Return the average for the specified numeric columns.
+
+        Args:
+            cols: The columns to calculate average for.
+        """
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
+        return self._non_empty_argument_function(
+            "avg",
+            *cols,
+            exclude_grouping_columns=exclude_grouping_columns,
+            _emit_ast=_emit_ast,
+        )
 
     mean = avg
 
     @relational_group_df_api_usage
     @publicapi
-    def sum(self, *cols: ColumnOrName, _emit_ast: bool = True) -> DataFrame:
-        """Return the sum for the specified numeric columns."""
-        return self._non_empty_argument_function("sum", *cols, _emit_ast=_emit_ast)
+    def sum(self, *cols: ColumnOrName, _emit_ast: bool = True, **kwargs) -> DataFrame:
+        """Return the sum for the specified numeric columns.
+
+        Args:
+            cols: The columns to calculate sum for.
+        """
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
+        return self._non_empty_argument_function(
+            "sum",
+            *cols,
+            exclude_grouping_columns=exclude_grouping_columns,
+            _emit_ast=_emit_ast,
+        )
 
     @relational_group_df_api_usage
     @publicapi
-    def median(self, *cols: ColumnOrName, _emit_ast: bool = True) -> DataFrame:
-        """Return the median for the specified numeric columns."""
-        return self._non_empty_argument_function("median", *cols, _emit_ast=_emit_ast)
+    def median(
+        self, *cols: ColumnOrName, _emit_ast: bool = True, **kwargs
+    ) -> DataFrame:
+        """Return the median for the specified numeric columns.
+
+        Args:
+            cols: The columns to calculate median for.
+        """
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
+        return self._non_empty_argument_function(
+            "median",
+            *cols,
+            exclude_grouping_columns=exclude_grouping_columns,
+            _emit_ast=_emit_ast,
+        )
 
     @relational_group_df_api_usage
     @publicapi
-    def min(self, *cols: ColumnOrName, _emit_ast: bool = True) -> DataFrame:
-        """Return the min for the specified numeric columns."""
-        return self._non_empty_argument_function("min", *cols, _emit_ast=_emit_ast)
+    def min(self, *cols: ColumnOrName, _emit_ast: bool = True, **kwargs) -> DataFrame:
+        """Return the min for the specified numeric columns.
+
+        Args:
+            cols: The columns to calculate min for.
+        """
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
+        return self._non_empty_argument_function(
+            "min",
+            *cols,
+            exclude_grouping_columns=exclude_grouping_columns,
+            _emit_ast=_emit_ast,
+        )
 
     @relational_group_df_api_usage
     @publicapi
-    def max(self, *cols: ColumnOrName, _emit_ast: bool = True) -> DataFrame:
-        """Return the max for the specified numeric columns."""
-        return self._non_empty_argument_function("max", *cols, _emit_ast=_emit_ast)
+    def max(self, *cols: ColumnOrName, _emit_ast: bool = True, **kwargs) -> DataFrame:
+        """Return the max for the specified numeric columns.
+
+        Args:
+            cols: The columns to calculate max for.
+        """
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
+        return self._non_empty_argument_function(
+            "max",
+            *cols,
+            exclude_grouping_columns=exclude_grouping_columns,
+            _emit_ast=_emit_ast,
+        )
 
     @relational_group_df_api_usage
     @publicapi
-    def count(self, _emit_ast: bool = True) -> DataFrame:
+    def count(self, _emit_ast: bool = True, **kwargs) -> DataFrame:
         """Return the number of rows for each group."""
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
         df = self._to_df(
             [
                 Alias(
@@ -692,6 +755,7 @@ def count(self, _emit_ast: bool = True) -> DataFrame:
                     "count",
                 )
             ],
+            exclude_grouping_columns=exclude_grouping_columns,
             _emit_ast=False,
         )
         df._ops_after_agg = set()
@@ -709,27 +773,38 @@ def count(self, _emit_ast: bool = True) -> DataFrame:
         return df
 
     @publicapi
-    def function(self, agg_name: str, _emit_ast: bool = True) -> Callable:
+    def function(self, agg_name: str, _emit_ast: bool = True, **kwargs) -> Callable:
         """Computes the builtin aggregate ``agg_name`` over the specified columns. Use
         this function to invoke any aggregates not explicitly listed in this class.
         See examples in :meth:`DataFrame.group_by`.
+
+        Args:
+            agg_name: The name of the aggregate function.
         """
-        return lambda *cols: self._function(agg_name, *cols, _emit_ast=_emit_ast)
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
+        return lambda *cols: self._function(
+            agg_name,
+            *cols,
+            exclude_grouping_columns=exclude_grouping_columns,
+            _emit_ast=_emit_ast,
+        )
 
     builtin = function
 
     @publicapi
     def _function(
-        self, agg_name: str, *cols: ColumnOrName, _emit_ast: bool = True
+        self, agg_name: str, *cols: ColumnOrName, _emit_ast: bool = True, **kwargs
     ) -> DataFrame:
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
         agg_exprs = []
         for c in cols:
             c_expr = Column(c)._expression if isinstance(c, str) else c._expression
             expr = functions._call_function(
                 agg_name, c_expr, _emit_ast=False
             )._expression
             agg_exprs.append(expr)
-        df = self._to_df(agg_exprs)
+
+        df = self._to_df(agg_exprs, exclude_grouping_columns=exclude_grouping_columns)
         df._ops_after_agg = set()
 
         if _emit_ast:
@@ -750,14 +825,19 @@ def _function(
 
     @publicapi
     def _non_empty_argument_function(
-        self, func_name: str, *cols: ColumnOrName, _emit_ast: bool = True
+        self, func_name: str, *cols: ColumnOrName, _emit_ast: bool = True, **kwargs
     ) -> DataFrame:
+        exclude_grouping_columns = kwargs.get("exclude_grouping_columns", False)
         if not cols:
             raise ValueError(
                 f"You must pass a list of one or more Columns to function: {func_name}"
             )
         else:
-            return self.builtin(func_name, _emit_ast=_emit_ast)(*cols)
+            return self.builtin(
+                func_name,
+                exclude_grouping_columns=exclude_grouping_columns,
+                _emit_ast=_emit_ast,
+            )(*cols)
 
     def _set_ast_ref(self, expr_builder: proto.Expr) -> None:
         """

tests/integ/test_df_aggregate.py

@@ -928,3 +928,143 @@ def test_filter_sort_limit_snowpark_connect_compatible(session):
 
     finally:
         context._is_snowpark_connect_compatible_mode = original_value
+
+
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="exclude_grouping_columns is not supported",
+)
+def test_group_by_exclude_grouping_columns(session):
+    """Test the exclude_grouping_columns parameter for all aggregate functions."""
+
+    # Create test data
+    df = session.create_dataframe(
+        [
+            ("A", "X", 1, 100),
+            ("A", "X", 2, 200),
+            ("A", "Y", 3, 300),
+            ("B", "X", 4, 400),
+            ("B", "Y", 5, 500),
+            ("B", "Y", 6, 600),
+        ],
+        schema=["group1", "group2", "value1", "value2"],
+    )
+
+    # Test agg() with exclude_grouping_columns
+    # Default behavior (include grouping columns)
+    result_default = df.group_by("group1").agg(sum_("value1").alias("sum_v1")).collect()
+    assert len(result_default[0]) == 2  # group1 + sum_v1
+    Utils.check_answer(result_default, [Row("A", 6), Row("B", 15)])
+
+    # Exclude grouping columns
+    result_exclude = (
+        df.group_by("group1")
+        .agg(sum_("value1").alias("sum_v1"), exclude_grouping_columns=True)
+        .collect()
+    )
+    assert len(result_exclude[0]) == 1  # only sum_v1
+    print(result_exclude)
+    Utils.check_answer(result_exclude, [Row(6), Row(15)])
+
+    # Test with multiple grouping columns
+    result_multi_default = (
+        df.group_by("group1", "group2").agg(sum_("value1").alias("sum_v1")).collect()
+    )
+    assert len(result_multi_default[0]) == 3  # group1 + group2 + sum_v1
+
+    result_multi_exclude = (
+        df.group_by("group1", "group2")
+        .agg(sum_("value1").alias("sum_v1"), exclude_grouping_columns=True)
+        .collect()
+    )
+    assert len(result_multi_exclude[0]) == 1  # only sum_v1
+    # Group by produces [('A', 'X', 3), ('A', 'Y', 3), ('B', 'X', 4), ('B', 'Y', 11)]
+    Utils.check_answer(result_multi_exclude, [Row(3), Row(3), Row(4), Row(11)])
+
+    # Test with multiple aggregations
+    result_multi_agg = (
+        df.group_by("group1")
+        .agg(
+            sum_("value1").alias("sum_v1"),
+            avg("value2").alias("avg_v2"),
+            exclude_grouping_columns=True,
+        )
+        .collect()
+    )
+    assert len(result_multi_agg[0]) == 2  # sum_v1 + avg_v2
+    Utils.check_answer(result_multi_agg, [Row(6, 200.0), Row(15, 500.0)])
+
+    # Test count() with exclude_grouping_columns
+    result_count_default = df.group_by("group1").count().collect()
+    assert len(result_count_default[0]) == 2  # group1 + count
+    Utils.check_answer(result_count_default, [Row("A", 3), Row("B", 3)])
+
+    result_count_exclude = (
+        df.group_by("group1").count(exclude_grouping_columns=True).collect()
+    )
+    assert len(result_count_exclude[0]) == 1  # only count
+    Utils.check_answer(result_count_exclude, [Row(3), Row(3)])
+
+    # Test avg() with exclude_grouping_columns
+    result_avg_default = df.group_by("group1").avg("value1").collect()
+    assert len(result_avg_default[0]) == 2  # group1 + avg
+
+    result_avg_exclude = (
+        df.group_by("group1").avg("value1", exclude_grouping_columns=True).collect()
+    )
+    assert len(result_avg_exclude[0]) == 1  # only avg
+    Utils.check_answer(result_avg_exclude, [Row(2.0), Row(5.0)])
+
+    # Test sum() with exclude_grouping_columns
+    result_sum_default = df.group_by("group1").sum("value1", "value2").collect()
+    assert len(result_sum_default[0]) == 3  # group1 + sum(value1) + sum(value2)
+
+    result_sum_exclude = (
+        df.group_by("group1")
+        .sum("value1", "value2", exclude_grouping_columns=True)
+        .collect()
+    )
+    assert len(result_sum_exclude[0]) == 2  # only sums
+    Utils.check_answer(result_sum_exclude, [Row(6, 600), Row(15, 1500)])
+
+    # Test min() with exclude_grouping_columns
+    result_min_default = df.group_by("group1").min("value1").collect()
+    assert len(result_min_default[0]) == 2  # group1 + min
+
+    result_min_exclude = (
+        df.group_by("group1").min("value1", exclude_grouping_columns=True).collect()
+    )
+    assert len(result_min_exclude[0]) == 1  # only min
+    Utils.check_answer(result_min_exclude, [Row(1), Row(4)])
+
+    # Test max() with exclude_grouping_columns
+    result_max_default = df.group_by("group1").max("value1").collect()
+    assert len(result_max_default[0]) == 2  # group1 + max
+
+    result_max_exclude = (
+        df.group_by("group1").max("value1", exclude_grouping_columns=True).collect()
+    )
+    assert len(result_max_exclude[0]) == 1  # only max
+    Utils.check_answer(result_max_exclude, [Row(3), Row(6)])
+
+    # Test median() with exclude_grouping_columns
+    result_median_default = df.group_by("group1").median("value1").collect()
+    assert len(result_median_default[0]) == 2  # group1 + median
+
+    result_median_exclude = (
+        df.group_by("group1").median("value1", exclude_grouping_columns=True).collect()
+    )
+    assert len(result_median_exclude[0]) == 1  # only median
+    Utils.check_answer(result_median_exclude, [Row(2.0), Row(5.0)])
+
+    # Test function() / builtin() with exclude_grouping_columns
+    result_builtin_default = df.group_by("group1").builtin("sum")("value1").collect()
+    assert len(result_builtin_default[0]) == 2  # group1 + sum
+
+    result_builtin_exclude = (
+        df.group_by("group1")
+        .builtin("sum", exclude_grouping_columns=True)("value1")
+        .collect()
+    )
+    assert len(result_builtin_exclude[0]) == 1  # only sum
+    Utils.check_answer(result_builtin_exclude, [Row(6), Row(15)])