[FLINK-40191][python] Add aggregation support to PyFlink DataFrame API - #28937
[FLINK-40191][python] Add aggregation support to PyFlink DataFrame API#28937auroflow wants to merge 3 commits into
Conversation
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)
Show positional and named aggregations together and document the returned DataFrame schemas. Generated-by: OpenAI Codex (GPT-5.6)
| aggregations: List[Expression] = [] | ||
| for aggregation in aggs: | ||
| if not isinstance(aggregation, Expression): | ||
| raise TypeError("agg() aggregations must be expressions") |
There was a problem hiding this comment.
The wording of this is inconsistent with the group_by error ("expressions" vs "Expression instances"). Could we please standardize how we refer to expressions in error messages?
There was a problem hiding this comment.
Yes, I standardized it to "expressions."
| ], | ||
| ) | ||
|
|
||
| def test_aggregation_python_contract_validation(self): |
There was a problem hiding this 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:
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)There was a problem hiding this comment.
Agreed, I split it into separate tests.
|
|
||
| def __init__(self, dataframe: DataFrame, grouping_keys: List[Expression]): | ||
| self._dataframe = dataframe | ||
| self._grouping_keys = grouping_keys |
There was a problem hiding this comment.
group_by(pf.col("amount") + 1).agg(...) raises ValidationException because the grouped projection re-evaluates the expression after amount is no longer available.
There was a problem hiding this comment.
This seems to be a bug in the Table API planner, which affects Table API for Python and Java as well. I have opened a separate PR for this: FLINK-40344
Standardize validation terminology and isolate Python contract checks so failures are reported independently. Generated-by: OpenAI Codex (GPT-5.6)
What is the purpose of the change
This pull request adds global and grouped aggregation support to the PyFlink DataFrame API, as described in FLINK-40191.
It allows users to aggregate an entire DataFrame with
DataFrame.agg(), or group rows withDataFrame.group_by()and aggregate them withGroupedDataFrame.agg().Brief change log
DataFrame.agg()for global aggregations.DataFrame.group_by()andGroupedDataFrame.agg()for grouped aggregations.GroupedDataFramefrompyflink.dataframe.Verifying this change
This change added tests and can be verified as follows:
sumandcountaggregations.Does this pull request potentially affect one of the following parts:
@Public(Evolving): yesDocumentation
Was generative AI tooling used to co-author this PR?
Generated-by: OpenAI Codex (GPT-5.6)