MattBelle commented on code in PR #28937:
URL: https://github.com/apache/flink/pull/28937#discussion_r3752740250
##########
flink-python/pyflink/dataframe/dataframe.py:
##########
@@ -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")
Review Comment:
The wording of this is inconsistent with the [group_by
error](https://github.com/apache/flink/pull/28937/changes#diff-6150ce6d3ad3555400ab39e466dec4adcf4af65f6fb7c4be0068fdb1fae0e9afR320)
("expressions" vs "Expression instances"). Could we please standardize how we
refer to expressions in error messages?
##########
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):
Review Comment:
Can we either split these assertations up or make them a subtest? The way it
is currently structured, a failure on `group_by` would mask a failure on `agg`.
I'd personally prefer the verbosity of making them each their own separate
test, but if we went the subtest route it'd look something like:
```python
def test_aggregation_python_contract_validation(self):
test_cases = [
("group_by_no_keys",
lambda df: df.group_by(),
ValueError("group_by() requires at least one grouping key")),
("group_by_invalid_type",
lambda df: df.group_by(42),
TypeError("group_by() grouping keys must be strings or Expression
instances")),
# ... 6 more cases
]
for name, func, expected_error in test_cases:
with self.subTest(validation=name):
with self.assertRaisesRegex(type(expected_error),
str(expected_error)):
func(self.dataframe)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
