[Data] Clean up ray.data.aggregate namespace (ray-project#36468)

This cleans up code completion for the ray.data.aggregate namespace (I ran into this while using ray.data):

Before:
<img width="712" alt="Screen Shot 2023-06-15 at 1 09 45 PM" src="https://github.com/ray-project/ray/assets/113316/38313256-735b-4ae1-b8f0-78c5729683d6">


After:
<img width="498" alt="Screen Shot 2023-06-15 at 1 10 39 PM" src="https://github.com/ray-project/ray/assets/113316/01e4cc8a-36ed-466e-8b4c-83b003e9994a">

python/ray/data/_internal/planner/exchange/aggregate_task_spec.py

@@ -2,7 +2,8 @@
 
 from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
 from ray.data._internal.table_block import TableBlockAccessor
-from ray.data.aggregate import AggregateFn, Count, _AggregateOnKeyBase
+from ray.data.aggregate import AggregateFn, Count
+from ray.data.aggregate._aggregate import _AggregateOnKeyBase
 from ray.data.block import Block, BlockAccessor, BlockExecStats, BlockMetadata, KeyType
 
 

python/ray/data/aggregate/__init__.py

@@ -0,0 +1,23 @@
+from ray.data.aggregate._aggregate import (
+ AbsMax,
+ AggregateFn,
+ Count,
+ Max,
+ Mean,
+ Min,
+ Quantile,
+ Std,
+ Sum,
+)
+
+__all__ = [
+ "AbsMax",
+ "AggregateFn",
+ "Count",
+ "Max",
+ "Mean",
+ "Min",
+ "Quantile",
+ "Std",
+ "Sum",
+]

python/ray/data/grouped_data.py

@@ -10,16 +10,8 @@
 from ray.data._internal.plan import AllToAllStage
 from ray.data._internal.push_based_shuffle import PushBasedShufflePlan
 from ray.data._internal.shuffle import ShuffleOp, SimpleShufflePlan
-from ray.data.aggregate import (
- AggregateFn,
- Count,
- Max,
- Mean,
- Min,
- Std,
- Sum,
- _AggregateOnKeyBase,
-)
+from ray.data.aggregate import AggregateFn, Count, Max, Mean, Min, Std, Sum
+from ray.data.aggregate._aggregate import _AggregateOnKeyBase
 from ray.data.block import (
  Block,
  BlockAccessor,

release/nightly_tests/dataset/aggregate_benchmark.py

@@ -1,7 +1,8 @@
 from typing import Tuple
 
 import ray
-from ray.data.aggregate import _AggregateOnKeyBase, Max, Mean, Min, Sum
+from ray.data.aggregate import Max, Mean, Min, Sum
+from ray.data.aggregate._aggregate import _AggregateOnKeyBase
 from ray.data.block import Block
 from ray.data.dataset import Dataset
 import pyarrow.compute as pac