[Data] - Make AggregateFnV2 generic over accumulator/result (Generic[AggType, U]) #57281

ArthurBook · 2025-10-07T21:30:18Z

Why are these changes needed?

The current Generic types in AggregateFnV2 are not tied to the class, so they are not picked up properly by static type checkers such as mypy.

By adding the Generic[] in the class definition, we get full type checking support.

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run pre-commit jobs to lint the changes in this PR. (pre-commit setup)
[N/A] I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
[N/A] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

gemini-code-assist

Code Review

This pull request correctly makes AggregateFnV2 a generic class, which is a great improvement for type safety. By inheriting from Generic[AggType, U], subclasses can now specify concrete types for the accumulator and result, allowing static type checkers like mypy to perform more robust validation.

To fully leverage this change, I recommend specifying the concrete generic types for all subclasses of AggregateFnV2 in ray/data/aggregate.py. For example, Count could be updated to class Count(AggregateFnV2[int, int]):. Applying this to all subclasses (Sum, Min, Max, etc.) would complete this typing enhancement. As this is outside the lines of the current diff, I'm mentioning it here as a suggestion for a follow-up.

The changes in this PR are correct and a good step forward. No further comments on the current changes.

…gType, U]) Signed-off-by: Arthur <atte.book@gmail.com>

python/ray/data/aggregate.py

goutamvenkat-anyscale · 2025-10-10T20:55:51Z

python/ray/data/aggregate.py


 @PublicAPI
-class Min(AggregateFnV2):
+class Min(AggregateFnV2[Union[int, float], Union[int, float]]):


So for min, max aggregates it's not guaranteed that we will have int and float types for the result/accumulator cause it ultimately depends on the underlying pyarrow type of the column.

I recommend leaving this as a generic type and on instantiation specify the type.

It looks like the value must at least be comparable be with a float because the (hard coded) default factory is a float.

Sorry, which default factory are you referring to here?

The

zero_factory=lambda: float("+inf"),

on line 428

^ @ArthurBook Can you please update this thread? Are you referring to the zero_factory? We should ideally change that to reflect the actual type of the column

Hi @goutamvenkat-anyscale I can adjust it, but I'm currently not fully in the clear how we want to address this without changing the actual runtime code.

I believe that current implementation would fail for other than floats and ints because they will initially be compared with a float.

My intention with this PR was only to fix type hints and leave runtime untouched.

Let me know what you think and I'll happily add it in though!

If we want to support any type, we can move the zero_factory to an __init__ kwarg and then tie the TypeVar to a SupportsRichComparison protocol that we define ourselves.

python/ray/data/aggregate.py

Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook · 2025-10-19T13:55:59Z

Hey @goutamvenkat-anyscale, I was out on vacation but now back!
LMK if this one is good to go?

Signed-off-by: Arthur <atte.book@gmail.com>

goutamvenkat-anyscale

lgtm

goutamvenkat-anyscale · 2025-10-24T19:17:32Z

Can you please address the merge conflict @ArthurBook ?

Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook · 2025-10-25T20:46:29Z

@goutamvenkat-anyscale done!

Signed-off-by: Arthur <atte.book@gmail.com>

alexeykudinkin

Thank you for your contribution @ArthurBook!

angelinalg

stamp

…AggType, U]) (ray-project#57281) ## Why are these changes needed? The current Generic types in `AggregateFnV2` are not tied to the class, so they are not picked up properly by static type checkers such as mypy.  By adding the Generic[] in the class definition, we get full type checking support. ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run pre-commit jobs to lint the changes in this PR. ([pre-commit setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting)) - [N/A] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [N/A] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Arthur <atte.book@gmail.com> Co-authored-by: Goutam <goutam@anyscale.com>

…AggType, U]) (ray-project#57281) ## Why are these changes needed? The current Generic types in `AggregateFnV2` are not tied to the class, so they are not picked up properly by static type checkers such as mypy.  By adding the Generic[] in the class definition, we get full type checking support. ## Checks - [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [x] I've run pre-commit jobs to lint the changes in this PR. ([pre-commit setup](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#lint-and-formatting)) - [N/A] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [N/A] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Arthur <atte.book@gmail.com> Co-authored-by: Goutam <goutam@anyscale.com> Signed-off-by: Aydin Abiar <aydin@anyscale.com>

ArthurBook requested a review from a team as a code owner October 7, 2025 21:30

gemini-code-assist bot reviewed Oct 7, 2025

View reviewed changes

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from db34450 to b2e09a6 Compare October 7, 2025 21:37

This comment was marked as outdated.

Sign in to view

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from b2e09a6 to 16b0bc3 Compare October 7, 2025 21:48

ray-gardener bot added data Ray Data-related issues community-contribution Contributed by the community labels Oct 8, 2025

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch 4 times, most recently from 9a5517b to 8251e55 Compare October 10, 2025 18:40

ArthurBook requested a review from a team as a code owner October 10, 2025 18:40

style: make AggregateFnV2 generic over accumulator/result (Generic[Ag…

3a05b2a

…gType, U]) Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from 8251e55 to 3a05b2a Compare October 10, 2025 19:08

goutamvenkat-anyscale reviewed Oct 10, 2025

View reviewed changes

python/ray/data/aggregate.py Outdated Show resolved Hide resolved

goutamvenkat-anyscale reviewed Oct 10, 2025

View reviewed changes

python/ray/data/aggregate.py Outdated Show resolved Hide resolved

goutamvenkat-anyscale reviewed Oct 10, 2025

View reviewed changes

python/ray/data/aggregate.py Outdated Show resolved Hide resolved

This comment was marked as outdated.

Sign in to view

goutamvenkat-anyscale reviewed Oct 10, 2025

View reviewed changes

python/ray/data/aggregate.py Outdated Show resolved Hide resolved

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from 8c08234 to 28fd4bf Compare October 10, 2025 21:19

style: address PR comments

bf2686b

Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from 28fd4bf to bf2686b Compare October 10, 2025 21:20

goutamvenkat-anyscale changed the title ~~style: make AggregateFnV2 generic over accumulator/result (Generic[AggType, U])~~ [Data] - make AggregateFnV2 generic over accumulator/result (Generic[AggType, U]) Oct 10, 2025

goutamvenkat-anyscale changed the title ~~[Data] - make AggregateFnV2 generic over accumulator/result (Generic[AggType, U])~~ [Data] - Make AggregateFnV2 generic over accumulator/result (Generic[AggType, U]) Oct 10, 2025

This comment was marked as outdated.

Sign in to view

ArthurBook added a commit to ArthurBook/ray that referenced this pull request Oct 10, 2025

address ray-project#57281 (comment)

9949a02

Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook added a commit to ArthurBook/ray that referenced this pull request Oct 10, 2025

address ray-project#57281 (comment)

e87c298

Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from 9949a02 to e87c298 Compare October 10, 2025 21:28

address ray-project#57281 (comment)

fc00f22

Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from e87c298 to fc00f22 Compare October 10, 2025 21:29

This comment was marked as outdated.

Sign in to view

refactor: use SupportsRichComparison in ordering functions

b443078

Signed-off-by: Arthur <atte.book@gmail.com>

ArthurBook force-pushed the arthurbook/aggregatefnv2-typehints branch from 9600612 to b443078 Compare October 23, 2025 22:27

goutamvenkat-anyscale approved these changes Oct 24, 2025

View reviewed changes

Merge branch 'master' into arthurbook/aggregatefnv2-typehints

beedb4b

Signed-off-by: Arthur <atte.book@gmail.com>

goutamvenkat-anyscale added the go add ONLY when ready to merge, run all tests label Oct 27, 2025

ArthurBook and others added 2 commits October 27, 2025 10:57

docs: add zero_factory to AggregationFnV2 docs

4303b14

Signed-off-by: Arthur <atte.book@gmail.com>

Merge branch 'master' into arthurbook/aggregatefnv2-typehints

f9b2171

alexeykudinkin approved these changes Oct 28, 2025

View reviewed changes

alexeykudinkin enabled auto-merge (squash) October 28, 2025 18:48

angelinalg approved these changes Oct 28, 2025

View reviewed changes

alexeykudinkin merged commit 9e8b291 into ray-project:master Oct 28, 2025
7 checks passed

[Data] - Make AggregateFnV2 generic over accumulator/result (Generic[AggType, U]) #57281

[Data] - Make AggregateFnV2 generic over accumulator/result (Generic[AggType, U]) #57281

Uh oh!

Conversation

ArthurBook commented Oct 7, 2025

Why are these changes needed?

Checks

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

goutamvenkat-anyscale Oct 10, 2025

Choose a reason for hiding this comment

Uh oh!

ArthurBook Oct 10, 2025

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 10, 2025

Choose a reason for hiding this comment

Uh oh!

ArthurBook Oct 19, 2025

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 20, 2025

Choose a reason for hiding this comment

Uh oh!

ArthurBook Oct 20, 2025

Choose a reason for hiding this comment

Uh oh!

ArthurBook Oct 20, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

ArthurBook commented Oct 19, 2025

Uh oh!

This comment was marked as outdated.

Uh oh!

goutamvenkat-anyscale left a comment

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale commented Oct 24, 2025

Uh oh!

ArthurBook commented Oct 25, 2025

Uh oh!

alexeykudinkin left a comment

Choose a reason for hiding this comment

Uh oh!

angelinalg left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants