[Data] Streaming Partition enforce row_num per block

python/ray/data/_internal/execution/operators/map_operator.py

@@ -4,6 +4,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections import defaultdict, deque
+from dataclasses import dataclass
 from typing import (
     Any,
     Callable,
@@ -65,6 +66,12 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class _TaskInput:
+    bundle: RefBundle
+    task_kwargs: Optional[Dict[str, Any]] = None
+
+
 class MapOperator(OneToOneOperator, InternalQueueOperatorMixin, ABC):
     """A streaming operator that maps input bundles 1:1 to output bundles.
 
@@ -117,6 +124,9 @@ def __init__(
         # All active `MetadataOpTask`s.
         self._metadata_tasks: Dict[int, MetadataOpTask] = {}
         self._next_metadata_task_idx = 0
+        # Optional helper that can turn incoming ref bundles into ready-to-run
+        # task inputs (bundle + per-task kwargs).
+        self._task_input_builder = None
         # Keep track of all finished streaming generators.
         super().__init__(name, input_op, data_context, target_max_block_size_override)
 
@@ -153,6 +163,9 @@ def get_additional_split_factor(self) -> int:
     def set_additional_split_factor(self, k: int):
         self._additional_split_factor = k
 
+    def set_task_input_builder(self, builder: Optional[Any]) -> None:
+        self._task_input_builder = builder
+
     def internal_queue_size(self) -> int:
         return self._block_ref_bundler.num_bundles()
 
@@ -328,6 +341,14 @@ def _warn_large_udf(self):
     def _add_input_inner(self, refs: RefBundle, input_index: int):
         assert input_index == 0, input_index
 
+        if self._task_input_builder is not None:
+            self._metrics.on_input_queued(refs)
+            task_inputs: List[_TaskInput] = self._task_input_builder.add_input(refs)
+            self._metrics.on_input_dequeued(refs)
+            for task_input in task_inputs:
+                self._submit_task_input(task_input)
+            return
+
         # Add RefBundle to the bundler.
         self._block_ref_bundler.add_bundle(refs)
         self._metrics.on_input_queued(refs)
@@ -380,8 +401,14 @@ def _get_dynamic_ray_remote_args(
             return self._ray_remote_args_factory_actor_locality(ray_remote_args)
         return ray_remote_args
 
+    def _submit_task_input(self, task_input: _TaskInput) -> None:
+        """Submit a ready-to-run task input produced by a task input builder."""
+        self._add_bundled_input(task_input.bundle, task_input.task_kwargs)
+
     @abstractmethod
-    def _add_bundled_input(self, refs: RefBundle):
+    def _add_bundled_input(
+        self, refs: RefBundle, task_kwargs: Optional[Dict[str, Any]] = None
+    ):
         """Add a pre-bundled upstream output to this operator.
 
         Unlike the add_input() arg, this RefBundle has already been further bundled by
@@ -392,6 +419,8 @@ def _add_bundled_input(self, refs: RefBundle):
 
         Args:
             refs: The fully-bundled ref bundle that should be added as input.
+            task_kwargs: A dictionary of kwargs to pass to the map task. You can
+                access these kwargs through the `TaskContext.kwargs` dictionary.
         """
         raise NotImplementedError
 
@@ -469,6 +498,12 @@ def get_active_tasks(self) -> List[OpTask]:
         return list(self._metadata_tasks.values()) + list(self._data_tasks.values())
 
     def all_inputs_done(self):
+        if self._task_input_builder is not None:
+            for task_input in self._task_input_builder.finish():
+                self._submit_task_input(task_input)
+            super().all_inputs_done()
+            return
+
         self._block_ref_bundler.done_adding_bundles()
         if self._block_ref_bundler.has_bundle():
             # Handle any leftover bundles in the bundler.

python/ray/data/_internal/execution/operators/task_pool_map_operator.py

@@ -79,7 +79,9 @@ def __init__(
 
         self._map_task = cached_remote_fn(_map_task, **ray_remote_static_args)
 
-    def _add_bundled_input(self, bundle: RefBundle):
+    def _add_bundled_input(
+        self, bundle: RefBundle, task_kwargs: Optional[Dict[str, Any]] = None
+    ):
         # Submit the task as a normal Ray task.
         ctx = TaskContext(
             task_idx=self._next_data_task_idx,
@@ -102,13 +104,16 @@ def _add_bundled_input(self, bundle: RefBundle):
             )
 
         data_context = self.data_context
+        per_task_kwargs = self.get_map_task_kwargs().copy()
+        if task_kwargs:
+            per_task_kwargs.update(task_kwargs)
 
         gen = self._map_task.options(**dynamic_ray_remote_args).remote(
             self._map_transformer_ref,
             data_context,
             ctx,
             *bundle.block_refs,
-            **self.get_map_task_kwargs(),
+            **per_task_kwargs,
         )
         self._submit_data_task(gen, bundle)
 

python/ray/data/_internal/logical/operators/map_operator.py

@@ -344,19 +344,26 @@ class StreamingRepartition(AbstractMap):
     Args:
         target_num_rows_per_block: The target number of rows per block granularity for
            streaming repartition.
+        enforce_target_num_rows_per_block: Whether to enforce the target number of rows per block. Default to False.
     """
 
     def __init__(
         self,
         input_op: LogicalOperator,
         target_num_rows_per_block: int,
+        enforce_target_num_rows_per_block: bool = False,
     ):
         super().__init__("StreamingRepartition", input_op)
         self._target_num_rows_per_block = target_num_rows_per_block
+        self._enforce_target_num_rows_per_block = enforce_target_num_rows_per_block
 
     @property
     def target_num_rows_per_block(self) -> int:
         return self._target_num_rows_per_block
 
+    @property
+    def enforce_target_num_rows_per_block(self) -> bool:
+        return self._enforce_target_num_rows_per_block
+
     def can_modify_num_rows(self) -> bool:
         return False

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

@@ -48,6 +48,10 @@
 )
 from ray.data._internal.numpy_support import _is_valid_column_values
 from ray.data._internal.output_buffer import OutputBlockSizeOption
+from ray.data._internal.streaming_repartition import (
+    StreamingRepartitionTaskBuilder,
+    streaming_repartition_block_fn,
+)
 from ray.data._internal.util import _truncated_repr
 from ray.data.block import (
     Block,
@@ -151,19 +155,25 @@ def plan_streaming_repartition_op(
     input_physical_dag = physical_children[0]
     compute = get_compute(op._compute)
 
-    # Create a no-op transform that is just coalescing/slicing the incoming
-    # blocks
-    transform_fn = BlockMapTransformFn(
-        lambda blocks, ctx: blocks,
-        output_block_size_option=OutputBlockSizeOption.of(
-            target_num_rows_per_block=op.target_num_rows_per_block
-        ),
-    )
-
-    map_transformer = MapTransformer([transform_fn])
+    if op.enforce_target_num_rows_per_block:
+        transform_fn = BlockMapTransformFn(
+            streaming_repartition_block_fn,
+            disable_block_shaping=True,
+        )
+        map_transformer = MapTransformer([transform_fn])
+    else:
+        # Create a no-op transform that is just coalescing/slicing the incoming
+        # blocks
+        transform_fn = BlockMapTransformFn(
+            lambda blocks, ctx: blocks,
+            output_block_size_option=OutputBlockSizeOption.of(
+                target_num_rows_per_block=op.target_num_rows_per_block
+            ),
+        )
+        map_transformer = MapTransformer([transform_fn])
 
     # Disable fusion for streaming repartition with the downstream op.
-    return MapOperator.create(
+    operator = MapOperator.create(
         map_transformer,
         input_physical_dag,
         data_context,
@@ -174,6 +184,13 @@ def plan_streaming_repartition_op(
         supports_fusion=False,
     )
 
+    if op.enforce_target_num_rows_per_block:
+        operator.set_task_input_builder(
+            StreamingRepartitionTaskBuilder(op.target_num_rows_per_block)
+        )
+
+    return operator
+
 
 def plan_filter_op(
     op: Filter,