fix(data): add missing Optional and return type annotations to aggreg…

[AyushKashyapII]

Summary

Fixes incorrect and missing type annotations in the Ray Data aggregation API (aggregate.py and grouped_data.py), addressing issue #64369.

These are not just style additions — several parameters had provably wrong types that mypy would flag: they defaulted to None but None was not included in the type union.

Changes in aggregate.py

• AggregateFn.__init__: wrapped accumulate_row and accumulate_block
  with Optional[...] — both defaulted to None but were typed as bare
  Callable[...], which excludes None
• Added -> None to __init__ of all 16 public aggregate classes:
  AggregateFn, AggregateFnV2, Count, AsList, Sum, Min, Max,
  Mean, Std, AbsMax, Quantile, Unique, CountDistinct,
  ValueCounter, MissingValuePercentage, ZeroPercentage,
  ApproximateQuantile, ApproximateTopK

Changes in grouped_data.py

• GroupedData.__init__: added -> None
• GroupedData._aggregate_on: corrected on param to
  Optional[Union[str, List[str]]] and added missing -> Dataset return type
• GroupedData.map_groups: corrected compute param to
  Optional[Union[str, ComputeStrategy]]
• GroupedData.sum, min, max, mean, std: corrected on param to
  Optional[Union[str, List[str]]] in all five methods

Why these are real errors, not just style

# Before — mypy error: Incompatible default for argument "on"
def sum(self, on: Union[str, List[str]] = None, ...) -> Dataset: ...
#                                         ^^^^
#            None is not assignable to Union[str, List[str]]

# After — correct
def sum(self, on: Optional[Union[str, List[str]]] = None, ...) -> Dataset: ...

Zero runtime impact
Python ignores annotations at runtime. All changed classes were
imported and instantiated locally to confirm no regressions.

[gemini-code-assist]

Code Review

This pull request improves type annotations across python/ray/data/aggregate.py and python/ray/data/grouped_data.py. Specifically, it adds explicit -> None return type annotations to various __init__ constructors and wraps parameters that default to None with Optional types (such as accumulate_row, accumulate_block, on, and compute). There are no review comments, so I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.