[FLINK-40191][python] Add aggregation support to PyFlink DataFrame API

flink-python/docs/reference/pyflink.dataframe/dataframe.rst

@@ -54,6 +54,19 @@ Transformations
     DataFrame.filter
     DataFrame.__getitem__
 
+Aggregations
+------------
+
+.. currentmodule:: pyflink.dataframe
+
+.. autosummary::
+    :toctree: api/
+
+    DataFrame.group_by
+    DataFrame.agg
+    GroupedDataFrame
+    GroupedDataFrame.agg
+
 Results
 -------
 

flink-python/pyflink/dataframe/__init__.py

@@ -44,11 +44,12 @@
     get_table_environment,
     set_table_environment,
 )
-from pyflink.dataframe.dataframe import DataFrame, col, lit
+from pyflink.dataframe.dataframe import DataFrame, GroupedDataFrame, col, lit
 from pyflink.dataframe.datatype import DataType
 
 __all__ = [
     "DataFrame",
+    "GroupedDataFrame",
     "DataType",
     "col",
     "lit",

flink-python/pyflink/dataframe/dataframe.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 ################################################################################
 
-from typing import Any, Callable, List, Optional, Tuple, Union, overload
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, overload
 
 from pyflink.common import Row
 from pyflink.dataframe.datatype import DataType
@@ -31,7 +31,7 @@
 from pyflink.table.types import DataTypes as TableDataTypes
 from pyflink.util.api_stability_decorators import PublicEvolving
 
-__all__ = ["DataFrame", "col", "lit"]
+__all__ = ["DataFrame", "GroupedDataFrame", "col", "lit"]
 
 
 @PublicEvolving()
@@ -116,6 +116,8 @@ class DataFrame:
     def __init__(self, table: Table):
         self._table = table
 
+    # ======================== Core Operations ========================
+
     @PublicEvolving()
     def filter(
         self,
@@ -271,6 +273,87 @@ def select(
 
         return DataFrame(self._table.select(*expressions))
 
+    # ======================== Aggregation ========================
+
+    @PublicEvolving()
+    def group_by(self, *columns: Union[str, Expression]) -> "GroupedDataFrame":
+        """
+        Group rows by one or more columns for aggregation.
+
+        String column names are converted to column expressions. Grouping keys are retained in
+        their supplied order and are included first in the result of
+        :meth:`GroupedDataFrame.agg`.
+
+        :param columns: Column names or expressions used as grouping keys.
+        :return: A grouped DataFrame that can be aggregated.
+        :raises TypeError: If a grouping key is not a string or expression.
+        :raises ValueError: If no grouping keys are provided.
+
+        Example::
+
+            >>> import pyflink.dataframe as pf
+            >>> df = pf.from_records([
+            ...     ("engineering", 10),
+            ...     ("engineering", 20),
+            ...     ("sales", 5),
+            ... ], schema=["department", "amount"])
+            >>> totals = df.group_by("department").agg(
+            ...     pf.col("amount").sum.alias("total_amount"),
+            ...     row_count=pf.col("amount").count,
+            ... )
+            >>> # totals schema: [department: STRING, total_amount: BIGINT,
+            >>> #                 row_count: BIGINT NOT NULL]
+
+        .. versionadded:: 2.4.0
+        """
+        if not columns:
+            raise ValueError("group_by() requires at least one grouping key")
+
+        grouping_keys: List[Expression] = []
+        for column in columns:
+            if isinstance(column, str):
+                grouping_keys.append(table_col(column))
+            elif isinstance(column, Expression):
+                grouping_keys.append(column)
+            else:
+                raise TypeError(
+                    "group_by() grouping keys must be strings or Expression instances"
+                )
+        return GroupedDataFrame(self, grouping_keys)
+
+    @PublicEvolving()
+    def agg(self, *aggs: Expression, **named_aggs: Expression) -> "DataFrame":
+        """
+        Aggregate all rows in this DataFrame.
+
+        Positional aggregation expressions are followed by named aggregations in the result.
+        Each named aggregation is aliased to its keyword name.
+
+        :param aggs: Aggregation expressions.
+        :param named_aggs: Aggregation expressions keyed by their result column names.
+        :return: A DataFrame containing the global aggregation results.
+        :raises TypeError: If an aggregation is not an expression.
+        :raises ValueError: If no aggregations are provided.
+
+        Example::
+
+            >>> import pyflink.dataframe as pf
+            >>> df = pf.from_records([
+            ...     (1, 10), (2, 20)
+            ... ], schema=["order_id", "amount"])
+            >>> summary = df.agg(
+            ...     pf.col("order_id").count.alias("order_count"),
+            ...     total_amount=pf.col("amount").sum,
+            ... )
+            >>> # summary schema: [order_count: BIGINT NOT NULL, total_amount: BIGINT]
+
+        .. versionadded:: 2.4.0
+        """
+        aggregations = _normalize_aggregations(aggs, named_aggs)
+        return DataFrame(self._table.group_by().select(*aggregations))
+
+    # ======================== Special Methods ========================
+
     @overload
     def __getitem__(self, key: str) -> Expression:
         ...
@@ -329,6 +412,8 @@ def __getitem__(
             return self.filter(key)
         raise TypeError("key must be a string, list, tuple, or Expression")
 
+    # ======================== Conversion ========================
+
     @PublicEvolving()
     def collect(self) -> List[Row]:
         """
@@ -348,3 +433,75 @@ def collect(self) -> List[Row]:
         """
         with self._table.execute().collect() as rows:
             return list(rows)
+
+
+@PublicEvolving()
+class GroupedDataFrame:
+    """
+    A DataFrame grouped by one or more keys and ready for aggregation.
+
+    Instances are created by :meth:`DataFrame.group_by`.
+
+    .. versionadded:: 2.4.0
+    """
+
+    def __init__(self, dataframe: DataFrame, grouping_keys: List[Expression]):
+        self._dataframe = dataframe
+        self._grouping_keys = grouping_keys
+
+    @PublicEvolving()
+    def agg(self, *aggs: Expression, **named_aggs: Expression) -> DataFrame:
+        """
+        Aggregate the rows in each group.
+
+        Grouping keys are included first in their supplied order, followed by positional
+        aggregation expressions and then named aggregations. Each named aggregation is aliased to
+        its keyword name.
+
+        :param aggs: Aggregation expressions.
+        :param named_aggs: Aggregation expressions keyed by their result column names.
+        :return: A DataFrame containing the grouping keys and aggregation results.
+        :raises TypeError: If an aggregation is not an expression.
+        :raises ValueError: If no aggregations are provided.
+
+        Example::
+
+            >>> import pyflink.dataframe as pf
+            >>> df = pf.from_records([
+            ...     ("engineering", 10),
+            ...     ("engineering", 20),
+            ...     ("sales", 5),
+            ... ], schema=["department", "amount"])
+            >>> totals = df.group_by("department").agg(
+            ...     pf.col("amount").sum.alias("total_amount"),
+            ...     row_count=pf.col("amount").count,
+            ... )
+            >>> # totals schema: [department: STRING, total_amount: BIGINT,
+            >>> #                 row_count: BIGINT NOT NULL]
+
+        .. versionadded:: 2.4.0
+        """
+        aggregations = _normalize_aggregations(aggs, named_aggs)
+        grouped_table = self._dataframe._table.group_by(*self._grouping_keys)
+        return DataFrame(grouped_table.select(*self._grouping_keys, *aggregations))
+
+
+# ======================== Internal Helpers ========================
+
+
+def _normalize_aggregations(
+    aggs: Tuple[Expression, ...], named_aggs: Dict[str, Expression]
+) -> List[Expression]:
+    if not aggs and not named_aggs:
+        raise ValueError("agg() requires at least one aggregation")
+
+    aggregations: List[Expression] = []
+    for aggregation in aggs:
+        if not isinstance(aggregation, Expression):
+            raise TypeError("agg() aggregations must be expressions")
+        aggregations.append(aggregation)
+    for name, aggregation in named_aggs.items():
+        if not isinstance(aggregation, Expression):
+            raise TypeError("agg() aggregations must be expressions")
+        aggregations.append(aggregation.alias(name))
+    return aggregations

flink-python/pyflink/dataframe/tests/test_dataframe.py

@@ -468,6 +468,81 @@ def test_lit_rejects_non_dataframe_data_type(self):
             pf.lit(1, object())
 
 
+class DataFrameAggregationTests(PyFlinkDataFrameUTTestCase):
+    def setUp(self):
+        super().setUp()
+        self.dataframe = pf.from_records(
+            [
+                ("engineering", "east", 10),
+                ("engineering", "west", 20),
+                ("sales", "east", 5),
+            ],
+            schema=["department", "region", "amount"],
+        )
+
+    def test_global_aggregation_preserves_positional_and_named_order(self):
+        result = self.dataframe.agg(
+            pf.col("amount").sum.alias("total_amount"),
+            row_count=pf.col("amount").count,
+        )
+
+        self.assert_dataframe_schema(
+            result,
+            ["total_amount", "row_count"],
+            [TableDataTypes.BIGINT(), TableDataTypes.BIGINT().not_null()],
+        )
+
+    def test_grouped_aggregation_emits_string_and_expression_keys_first(self):
+        grouped = self.dataframe.group_by("department", pf.col("region"))
+
+        self.assertIsInstance(grouped, pf.GroupedDataFrame)
+        result = grouped.agg(
+            pf.col("amount").sum.alias("total_amount"),
+            row_count=pf.col("amount").count,
+        )
+
+        self.assert_dataframe_schema(
+            result,
+            ["department", "region", "total_amount", "row_count"],
+            [
+                TableDataTypes.STRING(),
+                TableDataTypes.STRING(),
+                TableDataTypes.BIGINT(),
+                TableDataTypes.BIGINT().not_null(),
+            ],
+        )
+
+    def test_aggregation_python_contract_validation(self):
+        with self.assertRaisesRegex(ValueError, "requires at least one grouping key"):
+            self.dataframe.group_by()
+        with self.assertRaisesRegex(TypeError, "grouping keys must be strings"):
+            self.dataframe.group_by(42)
+        with self.assertRaisesRegex(ValueError, "requires at least one aggregation"):
+            self.dataframe.agg()
+        with self.assertRaisesRegex(TypeError, "aggregations must be expressions"):
+            self.dataframe.agg(42)
+        with self.assertRaisesRegex(TypeError, "aggregations must be expressions"):
+            self.dataframe.agg(total=42)
+
+        grouped = self.dataframe.group_by("department")
+        with self.assertRaisesRegex(ValueError, "requires at least one aggregation"):
+            grouped.agg()
+        with self.assertRaisesRegex(TypeError, "aggregations must be expressions"):
+            grouped.agg(42)
+        with self.assertRaisesRegex(TypeError, "aggregations must be expressions"):
+            grouped.agg(total=42)
+
+    def test_global_aggregation_delegates_expression_legality_to_planner(self):
+        with self.assertRaisesRegex(Py4JJavaError, "ValidationException"):
+            self.dataframe.agg(pf.col("amount"))
+
+    def test_grouped_aggregation_delegates_ambiguous_output_to_planner(self):
+        with self.assertRaisesRegex(Py4JJavaError, "ValidationException"):
+            self.dataframe.group_by("department").agg(
+                department=pf.col("amount").sum
+            )
+
+
 class DataFrameITTests(PyFlinkStreamDataFrameTestCase):
     def test_from_records(self):
         dataframe = pf.from_records(
@@ -570,6 +645,26 @@ def test_from_records_with_batch_table_environment(self):
 
         self.assertEqual(result.collect(), [Row(2, "Bob")])
 
+    def test_grouped_aggregation_with_batch_table_environment(self):
+        pf.set_table_environment(self.t_env)
+
+        result = pf.from_records(
+            [
+                ("engineering", 10),
+                ("engineering", 20),
+                ("sales", 5),
+            ],
+            schema=["department", "amount"],
+        ).group_by("department").agg(
+            total_amount=pf.col("amount").sum,
+            row_count=pf.col("amount").count,
+        )
+
+        self.assertCountEqual(
+            result.collect(),
+            [Row("engineering", 30, 2), Row("sales", 5, 1)],
+        )
+
 
 if __name__ == "__main__":
     unittest.main()