From ad251c1c2d02f74495a10eb61d73aa1293db5d87 Mon Sep 17 00:00:00 2001 From: auroflow Date: Thu, 6 Aug 2026 17:21:52 +0800 Subject: [PATCH 1/2] [FLINK-40191][python] Add DataFrame aggregation support Add global and grouped aggregation APIs to PyFlink DataFrame, including planner-backed validation tests, one batch execution test, and reference documentation. Generated-by: OpenAI Codex (GPT-5.6) --- .../reference/pyflink.dataframe/dataframe.rst | 13 ++ flink-python/pyflink/dataframe/__init__.py | 3 +- flink-python/pyflink/dataframe/dataframe.py | 155 +++++++++++++++++- .../pyflink/dataframe/tests/test_dataframe.py | 95 +++++++++++ 4 files changed, 263 insertions(+), 3 deletions(-) diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst index 2caa7503eba132..19ebb0ad9c0d23 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 8ad43bcbbe2560..5412c50adf9308 100644 --- a/flink-python/pyflink/dataframe/__init__.py +++ b/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", diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 66c2020ee4db74..312e6159e78f19 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.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,83 @@ 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( + ... total_amount=pf.col("amount").sum + ... ) + + .. 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, + ... ) + + .. 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 +408,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 +429,73 @@ 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( + ... total_amount=pf.col("amount").sum, + ... row_count=pf.col("amount").count, + ... ) + + .. 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 98714e9fbb6a47..6384936911c786 100644 --- a/flink-python/pyflink/dataframe/tests/test_dataframe.py +++ b/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() From 98bc2a62d54eff1d8499da48cc8801e5f4c4fa61 Mon Sep 17 00:00:00 2001 From: auroflow Date: Fri, 7 Aug 2026 10:04:05 +0800 Subject: [PATCH 2/2] [FLINK-40191][python] Improve aggregation examples Show positional and named aggregations together and document the returned DataFrame schemas. Generated-by: OpenAI Codex (GPT-5.6) --- flink-python/pyflink/dataframe/dataframe.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/flink-python/pyflink/dataframe/dataframe.py b/flink-python/pyflink/dataframe/dataframe.py index 312e6159e78f19..0cf0f1eb80d2ef 100644 --- a/flink-python/pyflink/dataframe/dataframe.py +++ b/flink-python/pyflink/dataframe/dataframe.py @@ -298,8 +298,11 @@ def group_by(self, *columns: Union[str, Expression]) -> "GroupedDataFrame": ... ("sales", 5), ... ], schema=["department", "amount"]) >>> totals = df.group_by("department").agg( - ... total_amount=pf.col("amount").sum + ... 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 """ @@ -342,6 +345,7 @@ def agg(self, *aggs: Expression, **named_aggs: Expression) -> "DataFrame": ... 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 """ @@ -469,9 +473,11 @@ def agg(self, *aggs: Expression, **named_aggs: Expression) -> DataFrame: ... ("sales", 5), ... ], schema=["department", "amount"]) >>> totals = df.group_by("department").agg( - ... total_amount=pf.col("amount").sum, + ... 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 """