chore: import upstream snapshot with attribution

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

@@ -0,0 +1,101 @@
+from typing import Any, Dict, List, Optional
+
+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.operators.map_transformer import MapTransformer
+from ray.data._internal.execution.util import merge_label_selector
+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.shuffle_task_spec import ShuffleTaskSpec
+from ray.data._internal.random_config import (
+    RandomSeedConfig,
+    get_single_integer_random_seed,
+)
+from ray.data.context import DataContext, ShuffleStrategy
+
+
+def generate_random_shuffle_fn(
+    data_context: DataContext,
+    seed_config: RandomSeedConfig,
+    num_outputs: Optional[int] = None,
+    ray_remote_args: Optional[Dict[str, Any]] = None,
+    _debug_limit_shuffle_execution_to_num_blocks: Optional[int] = None,
+) -> AllToAllTransformFn:
+    """Generate function to randomly shuffle each records of blocks."""
+
+    # If no seed has been specified, pin timestamp based one
+    # so that task could be safely retried (w/o changing their output)
+    seed = get_single_integer_random_seed(seed_config, data_context)
+
+    def fn(
+        refs: List[RefBundle],
+        ctx: TaskContext,
+    ) -> AllToAllTransformFnResult:
+        num_input_blocks = sum(len(r.blocks) for r in refs)
+
+        # If map_transformer is specified (e.g. from fusing
+        # MapOperator->AllToAllOperator), we pass a map function which
+        # is applied to each block before shuffling.
+        map_transformer: Optional[MapTransformer] = ctx.upstream_map_transformer
+        upstream_map_fn = None
+        nonlocal ray_remote_args
+        if map_transformer:
+            # NOTE: We override target max-block sizing of the previous
+            #       transformation to avoid unnecessary block shaping (if any)
+            map_transformer.override_target_max_block_size(None)
+
+            def upstream_map_fn(blocks):
+                DataContext._set_current(data_context)
+                return map_transformer.apply_transform(blocks, ctx)
+
+            # If there is a fused upstream operator,
+            # also use the ray_remote_args from the fused upstream operator.
+            ray_remote_args = ctx.upstream_map_ray_remote_args
+
+        shuffle_spec = ShuffleTaskSpec(
+            target_shuffle_max_block_size=(
+                ctx.target_max_block_size_override or data_context.target_max_block_size
+            ),
+            random_shuffle=True,
+            random_seed=seed,
+            upstream_map_fn=upstream_map_fn,
+        )
+
+        if data_context.shuffle_strategy == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
+            if num_outputs is not None:
+                raise NotImplementedError(
+                    "Push-based shuffle doesn't support setting num_blocks yet."
+                )
+            scheduler = PushBasedShuffleTaskScheduler(shuffle_spec)
+        else:
+            scheduler = PullBasedShuffleTaskScheduler(shuffle_spec)
+
+        label_selector = data_context.execution_options.label_selector
+        map_ray_remote_args = merge_label_selector(
+            ray_remote_args or {}, label_selector
+        )
+        reduce_ray_remote_args = merge_label_selector(
+            ray_remote_args or {}, label_selector
+        )
+        return scheduler.execute(
+            refs,
+            num_outputs or num_input_blocks,
+            task_ctx=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