Skip to content

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

Open
auroflow wants to merge 3 commits into
apache:masterfrom
auroflow:codex/FLINK-40191-dataframe-aggregation
Open

[FLINK-40191][python] Add aggregation support to PyFlink DataFrame API#28937
auroflow wants to merge 3 commits into
apache:masterfrom
auroflow:codex/FLINK-40191-dataframe-aggregation

Conversation

@auroflow

@auroflow auroflow commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

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 with DataFrame.group_by() and aggregate them with GroupedDataFrame.agg().

Brief change log

  • Add DataFrame.agg() for global aggregations.
  • Add DataFrame.group_by() and GroupedDataFrame.agg() for grouped aggregations.
  • Export GroupedDataFrame from pyflink.dataframe.
  • Add API documentation and usage examples.

Verifying this change

This change added tests and can be verified as follows:

  • Added unit tests for global and grouped aggregation schemas and output ordering.
  • Added unit tests for Python argument validation and planner-delegated validation.
  • Added a batch integration test covering grouped sum and count aggregations.

Does this pull request potentially affect one of the following parts:

  • Dependencies (does it add or upgrade a dependency): no
  • The public API, i.e., is any changed class annotated with @Public(Evolving): yes
  • The serializers: no
  • The runtime per-record code paths (performance sensitive): no
  • Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Kubernetes/Yarn, ZooKeeper: no
  • The S3 file system connector: no

Documentation

  • Does this pull request introduce a new feature? yes
  • If yes, how is the feature documented? PyDocs and Python API docstrings

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

Generated-by: OpenAI Codex (GPT-5.6)

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)
@flinkbot

flinkbot commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

CI report:

Bot commands The @flinkbot bot supports the following commands:
  • @flinkbot run azure re-run the last Azure build

Show positional and named aggregations together and document the returned DataFrame schemas.

Generated-by: OpenAI Codex (GPT-5.6)
@auroflow
auroflow marked this pull request as ready for review August 7, 2026 03:09
aggregations: List[Expression] = []
for aggregation in aggs:
if not isinstance(aggregation, Expression):
raise TypeError("agg() aggregations must be expressions")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I standardized it to "expressions."

],
)

def test_aggregation_python_contract_validation(self):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, I split it into separate tests.

@github-actions github-actions Bot added the community-reviewed PR has been reviewed by the community. label Aug 11, 2026

def __init__(self, dataframe: DataFrame, grouping_keys: List[Expression]):
self._dataframe = dataframe
self._grouping_keys = grouping_keys

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

group_by(pf.col("amount") + 1).agg(...) raises ValidationException because the grouped projection re-evaluates the expression after amount is no longer available.

@auroflow auroflow Aug 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-reviewed PR has been reviewed by the community.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants