[Data] Streaming Partition enforce row_num per block

python/ray/data/_internal/execution/interfaces/__init__.py

@@ -2,7 +2,7 @@
 from .execution_options import ExecutionOptions, ExecutionResources
 from .executor import Executor, OutputIterator
 from .physical_operator import PhysicalOperator, ReportsExtraResourceUsage
-from .ref_bundle import RefBundle
+from .ref_bundle import BlockSlice, RefBundle
 from .task_context import TaskContext
 from .transform_fn import AllToAllTransformFn
 
@@ -15,6 +15,7 @@
     "OutputIterator",
     "PhysicalOperator",
     "RefBundle",
+    "BlockSlice",
     "ReportsExtraResourceUsage",
     "TaskContext",
 ]

python/ray/data/_internal/execution/interfaces/ref_bundle.py

@@ -1,15 +1,31 @@
+import math
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Dict, Iterator, List, Optional, Tuple
+from typing import Dict, Iterable, Iterator, List, Optional, Tuple
 
 import ray
 from .common import NodeIdStr
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
 from ray.data._internal.memory_tracing import trace_deallocation
-from ray.data.block import Block, BlockMetadata, Schema
+from ray.data.block import Block, BlockAccessor, BlockMetadata, Schema
 from ray.data.context import DataContext
 from ray.types import ObjectRef
 
 
+@dataclass
+class BlockSlice:
+    """A slice of a block."""
+
+    # Starting row offset (inclusive) within the block.
+    start_offset: int
+    # Ending row offset (exclusive) within the block.
+    end_offset: int
+
+    @property
+    def num_rows(self) -> int:
+        return self.end_offset - self.start_offset
+
+
 @dataclass
 class RefBundle:
     """A group of data block references and their metadata.
@@ -38,6 +54,10 @@ class RefBundle:
     # Whether we own the blocks (can safely destroy them).
     owns_blocks: bool
 
+    # The slices of the blocks in this bundle. This is optional, and may be None
+    # if the blocks are not sliced.
+    slices: Optional[List[BlockSlice]] = None
+
     # This attribute is used by the split() operator to assign bundles to logical
     # output splits. It is otherwise None.
     output_split_idx: Optional[int] = None
@@ -53,6 +73,10 @@ class RefBundle:
     def __post_init__(self):
         if not isinstance(self.blocks, tuple):
             object.__setattr__(self, "blocks", tuple(self.blocks))
+        if self.slices is not None:
+            assert len(self.blocks) == len(
+                self.slices
+            ), "Number of blocks and slices must match"
         for b in self.blocks:
             assert isinstance(b, tuple), b
             assert len(b) == 2, b
@@ -80,16 +104,29 @@ def metadata(self) -> List[BlockMetadata]:
 
     def num_rows(self) -> Optional[int]:
         """Number of rows present in this bundle, if known."""
+        slice_total: Optional[int] = None
+        if self.slices is not None:
+            slice_total = sum(block_slice.num_rows for block_slice in self.slices)
+            return slice_total
+
         total = 0
         for m in self.metadata:
             if m.num_rows is None:
                 return None
-            else:
-                total += m.num_rows
+            total += m.num_rows
         return total
 
     def size_bytes(self) -> int:
         """Size of the blocks of this bundle in bytes."""
+        if self.slices is not None:
+            total = 0
+            for (_, metadata), block_slice in zip(self.blocks, self.slices):
+                if metadata.num_rows and metadata.num_rows != block_slice.num_rows:
+                    per_row = metadata.size_bytes / metadata.num_rows
+                    total += max(1, int(math.ceil(per_row * block_slice.num_rows)))
+                else:
+                    total += metadata.size_bytes
+            return total
         return sum(m.size_bytes for m in self.metadata)
 
     def destroy_if_owned(self) -> int:
@@ -170,3 +207,41 @@ def _ref_bundles_iterator_to_block_refs_list(
     return [
         block_ref for ref_bundle in ref_bundles for block_ref in ref_bundle.block_refs
     ]
+
+
+def _slice_block_metadata(
+    metadata: BlockMetadata, num_rows_in_slice: int
+) -> BlockMetadata:
+    assert (
+        num_rows_in_slice > 0
+    ), "num_rows_in_slice must be positive for slicing block metadata."
+    size_bytes = metadata.size_bytes
+    if metadata.size_bytes is not None and metadata.num_rows:
+        per_row = metadata.size_bytes / metadata.num_rows
+        size_bytes = max(1, int(math.ceil(per_row * num_rows_in_slice)))
+    return BlockMetadata(
+        num_rows=num_rows_in_slice if metadata.num_rows is not None else None,
+        size_bytes=size_bytes,
+        exec_stats=None,
+        input_files=list(metadata.input_files),
+    )
+
+
+def _iter_sliced_blocks(
+    blocks: Iterable[Block],
+    slices: List[BlockSlice],
+) -> Iterator[Block]:
+    blocks_list = list(blocks)
+    builder = DelegatingBlockBuilder()
+    for block, block_slice in zip(blocks_list, slices):
+        accessor = BlockAccessor.for_block(block)
+        start = block_slice.start_offset
+        end = block_slice.end_offset
+
+        if start == 0 and end >= accessor.num_rows():
+            builder.add_block(block)
+        else:
+            builder.add_block(accessor.slice(start, end, copy=False))
+
+    if builder is not None:
+        yield builder.build()

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

@@ -20,6 +20,7 @@
 from ray.data._internal.execution.bundle_queue import create_bundle_queue
 from ray.data._internal.execution.bundle_queue.bundle_queue import BundleQueue
 from ray.data._internal.execution.interfaces import (
+    BlockSlice,
     ExecutionOptions,
     ExecutionResources,
     NodeIdStr,
@@ -32,7 +33,11 @@
     ActorLocationTracker,
     get_or_create_actor_location_tracker,
 )
-from ray.data._internal.execution.operators.map_operator import MapOperator, _map_task
+from ray.data._internal.execution.operators.map_operator import (
+    BaseRefBundler,
+    MapOperator,
+    _map_task,
+)
 from ray.data._internal.execution.operators.map_transformer import MapTransformer
 from ray.data._internal.execution.util import locality_string
 from ray.data._internal.remote_fn import _add_system_error_to_retry_exceptions
@@ -71,6 +76,7 @@ def __init__(
         compute_strategy: ActorPoolStrategy,
         name: str = "ActorPoolMap",
         min_rows_per_bundle: Optional[int] = None,
+        ref_bundler: Optional[BaseRefBundler] = None,
         supports_fusion: bool = True,
         map_task_kwargs: Optional[Dict[str, Any]] = None,
         ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
@@ -91,6 +97,7 @@ def __init__(
                 transform_fn, or None to use the block size. Setting the batch size is
                 important for the performance of GPU-accelerated transform functions.
                 The actual rows passed may be less if the dataset is small.
+            ref_bundler: The ref bundler to use for this operator.
             supports_fusion: Whether this operator supports fusion with other operators.
             map_task_kwargs: A dictionary of kwargs to pass to the map task. You can
                 access these kwargs through the `TaskContext.kwargs` dictionary.
@@ -113,6 +120,7 @@ def __init__(
             name,
             target_max_block_size_override,
             min_rows_per_bundle,
+            ref_bundler,
             supports_fusion,
             map_task_kwargs,
             ray_remote_args_fn,
@@ -296,13 +304,15 @@ def _task_done_callback(res_ref):
         )
         return actor, res_ref
 
-    def _add_bundled_input(self, bundle: RefBundle):
+    def _add_bundled_input(
+        self, bundle: RefBundle, slices: Optional[List[BlockSlice]] = None
+    ):
         self._bundle_queue.add(bundle)
         self._metrics.on_input_queued(bundle)
         # Try to dispatch all bundles in the queue, including this new bundle.
-        self._dispatch_tasks()
+        self._dispatch_tasks(slices)
 
-    def _dispatch_tasks(self):
+    def _dispatch_tasks(self, slices: Optional[List[BlockSlice]] = None):
         """Try to dispatch tasks from the bundle buffer to the actor pool.
 
         This is called when:
@@ -330,6 +340,7 @@ def _dispatch_tasks(self):
                 self.data_context,
                 ctx,
                 *input_blocks,
+                slices=bundle.slices,
                 **self.get_map_task_kwargs(),
             )
 
@@ -571,13 +582,15 @@ def submit(
         data_context: DataContext,
         ctx: TaskContext,
         *blocks: Block,
+        slices: Optional[List[BlockSlice]] = None,
         **kwargs: Dict[str, Any],
     ) -> Iterator[Union[Block, List[BlockMetadata]]]:
         yield from _map_task(
             self._map_transformer,
             data_context,
             ctx,
             *blocks,
+            slices=slices,
             **kwargs,
         )