chore: import upstream snapshot with attribution

python/ray/data/_internal/planner/aggregate.py

@@ -0,0 +1,111 @@
+from typing import List, Optional, Union
+
+from ray.data._internal.execution.interfaces import (
+    AllToAllTransformFn,
+    RefBundle,
+    TaskContext,
+)
+from ray.data._internal.execution.interfaces.transform_fn import (
+    AllToAllTransformFnResult,
+)
+from ray.data._internal.execution.util import merge_label_selector
+from ray.data._internal.planner.exchange.aggregate_task_spec import (
+    SortAggregateTaskSpec,
+)
+from ray.data._internal.planner.exchange.pull_based_shuffle_task_scheduler import (
+    PullBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+from ray.data._internal.planner.exchange.sort_task_spec import SortKey, SortTaskSpec
+from ray.data._internal.util import unify_ref_bundles_schema
+from ray.data.aggregate import AggregateFn
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_aggregate_fn(
+    key: Optional[Union[str, List[str]]],
+    aggs: List[AggregateFn],
+    data_context: DataContext,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to aggregate blocks by the specified key column or key
+    function.
+    """
+    assert data_context.shuffle_strategy in [
+        ShuffleStrategy.SORT_SHUFFLE_PULL_BASED,
+        ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED,
+    ]
+
+    if len(aggs) == 0:
+        raise ValueError("Aggregate requires at least one aggregation")
+
+    def fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        blocks = []
+        metadata = []
+        for ref_bundle in refs:
+            blocks.extend(ref_bundle.block_refs)
+            metadata.extend(ref_bundle.metadata)
+        if len(blocks) == 0:
+            return (blocks, {})
+
+        unified_schema = unify_ref_bundles_schema(refs)
+        for agg_fn in aggs:
+            agg_fn._validate(unified_schema)
+
+        num_mappers = len(blocks)
+
+        sort_key = SortKey(key)
+
+        label_selector = data_context.execution_options.label_selector
+        if key is None:
+            num_outputs = 1
+            boundaries = []
+        else:
+            # Use same number of output partitions.
+            num_outputs = num_mappers
+            sample_bar = ctx.sub_progress_bar_dict[
+                SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME
+            ]
+            # Sample boundaries for aggregate key.
+            boundaries = SortTaskSpec.sample_boundaries(
+                blocks,
+                sort_key,
+                num_outputs,
+                sample_bar,
+                label_selector=label_selector,
+            )
+
+        agg_spec = SortAggregateTaskSpec(
+            boundaries=boundaries,
+            key=sort_key,
+            aggs=aggs,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            scheduler = PushBasedShuffleTaskScheduler(agg_spec)
+        elif data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PULL_BASED:
+            scheduler = PullBasedShuffleTaskScheduler(agg_spec)
+        else:
+            raise ValueError(
+                f"Invalid shuffle strategy '{data_context.shuffle_strategy}'"
+            )
+
+        map_ray_remote_args = merge_label_selector({}, label_selector)
+        reduce_ray_remote_args = merge_label_selector({}, label_selector)
+        return scheduler.execute(
+            refs,
+            num_outputs,
+            ctx,
+            map_ray_remote_args=map_ray_remote_args,
+            reduce_ray_remote_args=reduce_ray_remote_args,
+            _debug_limit_execution_to_num_blocks=(
+                _debug_limit_shuffle_execution_to_num_blocks
+            ),
+        )
+
+    return fn