This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 7b84e3fa54f [FLINK-40191][python] Add aggregation support to PyFlink
DataFrame API (#28937)
7b84e3fa54f is described below
commit 7b84e3fa54f08a9ad401231ab4bceb6e308ac44e
Author: Liu Liu <[email protected]>
AuthorDate: Fri Aug 14 14:28:10 2026 +0800
[FLINK-40191][python] Add aggregation support to PyFlink DataFrame API
(#28937)
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 13 ++
flink-python/pyflink/dataframe/__init__.py | 3 +-
flink-python/pyflink/dataframe/dataframe.py | 161 ++++++++++++++++++++-
.../pyflink/dataframe/tests/test_dataframe.py | 112 ++++++++++++++
4 files changed, 286 insertions(+), 3 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index 2caa7503eba..19ebb0ad9c0 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/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
-------
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 8ad43bcbbe2..5412c50adf9 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -44,11 +44,12 @@ from pyflink.dataframe.context import (
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",
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index 66c2020ee4d..4d386a1580c 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/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.table import Table
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 @@ class DataFrame:
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 expressions"
+ )
+ 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 @@ class DataFrame:
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 @@ class DataFrame:
"""
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
diff --git a/flink-python/pyflink/dataframe/tests/test_dataframe.py
b/flink-python/pyflink/dataframe/tests/test_dataframe.py
index 98714e9fbb6..8f5b7620cab 100644
--- a/flink-python/pyflink/dataframe/tests/test_dataframe.py
+++ b/flink-python/pyflink/dataframe/tests/test_dataframe.py
@@ -468,6 +468,98 @@ class DataFrameLiteralTests(PyFlinkDataFrameUTTestCase):
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_group_by_requires_grouping_key(self):
+ with self.assertRaisesRegex(ValueError, "requires at least one
grouping key"):
+ self.dataframe.group_by()
+
+ def test_group_by_rejects_unsupported_key_type(self):
+ with self.assertRaisesRegex(
+ TypeError, "grouping keys must be strings or expressions"
+ ):
+ self.dataframe.group_by(42)
+
+ def test_global_aggregation_requires_aggregation(self):
+ with self.assertRaisesRegex(ValueError, "requires at least one
aggregation"):
+ self.dataframe.agg()
+
+ def test_global_aggregation_rejects_unsupported_positional_type(self):
+ with self.assertRaisesRegex(TypeError, "aggregations must be
expressions"):
+ self.dataframe.agg(42)
+
+ def test_global_aggregation_rejects_unsupported_named_type(self):
+ with self.assertRaisesRegex(TypeError, "aggregations must be
expressions"):
+ self.dataframe.agg(total=42)
+
+ def test_grouped_aggregation_requires_aggregation(self):
+ grouped = self.dataframe.group_by("department")
+ with self.assertRaisesRegex(ValueError, "requires at least one
aggregation"):
+ grouped.agg()
+
+ def test_grouped_aggregation_rejects_unsupported_positional_type(self):
+ grouped = self.dataframe.group_by("department")
+ with self.assertRaisesRegex(TypeError, "aggregations must be
expressions"):
+ grouped.agg(42)
+
+ def test_grouped_aggregation_rejects_unsupported_named_type(self):
+ grouped = self.dataframe.group_by("department")
+ 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 +662,26 @@ class DataFrameBatchITTests(PyFlinkITTestCase):
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()
