from __future__ import annotations import abc from collections import deque from typing import TYPE_CHECKING, Any, Deque, List, Optional, Tuple from typing_extensions import override from ray.data._internal.execution.bundle_queue import ( BaseBundleQueue, ) if TYPE_CHECKING: from ray.data._internal.execution.interfaces import RefBundle class RebundlingStrategy(abc.ABC): """Base class for strategies describing how to rebundle queues.""" @abc.abstractmethod def can_build_ready_bundle(self, num_pending_rows: int) -> bool: """Signifies whether we can build a ready bundle. A ready bundle is a bundle that will be returned from `get_next()` calls. Pending bundles merge into Ready bundles.""" ... @abc.abstractmethod def rows_needed_from_last_pending_bundle( self, total_pending_rows: int, last_pending_bundle: RefBundle ) -> int: """Used to determine how to rebundle and slice an existing bundle. Args: total_pending_rows: The number of rows in a batch of pending bundles that will be merged to form a ready bundle, including the last_pending_bundle. last_pending_bundle: The last pending bundles in that batch ^. The term *last* means the bundle that caused `can_build_ready_bundle(num_pending_rows)` to be `True` for the first time. Returns: The # of rows needed from the last pending bundle. This should be > 0, unless bundle.num_rows() is None. """ ... class EstimateSize(RebundlingStrategy): """Rebundles RefBundles to get them close to a particular number of rows.""" def __init__(self, min_rows_per_bundle: Optional[int]): """Creates a strategy for combining bundles close to a particular row count. Args: min_rows_per_bundle: The target number of rows per bundle. Note that we bundle up to this target, but only exceed it if not doing so would result in an empty bundle. If None, this behaves like a normal queue. """ assert ( min_rows_per_bundle is None or min_rows_per_bundle >= 0 ), "Min rows per bundle has to be non-negative" self._min_rows_per_bundle: Optional[int] = min_rows_per_bundle @override def can_build_ready_bundle(self, num_pending_rows: int) -> bool: return num_pending_rows > 0 and ( self._min_rows_per_bundle is None or num_pending_rows >= self._min_rows_per_bundle ) @override def rows_needed_from_last_pending_bundle( self, total_pending_rows: int, last_pending_bundle: RefBundle ) -> int: """Returns all the rows in the pending bundle, since we only care about an estimate""" return last_pending_bundle.num_rows() or 0 class ExactMultipleSize(RebundlingStrategy): def __init__(self, target_num_rows_per_block: int): assert ( target_num_rows_per_block > 0 ), "target_num_rows_per_block must be positive for streaming repartition." self._target_num_rows = target_num_rows_per_block @override def can_build_ready_bundle(self, num_pending_rows: int) -> bool: return num_pending_rows >= self._target_num_rows @override def rows_needed_from_last_pending_bundle( self, total_pending_rows: int, last_pending_bundle: RefBundle ) -> int: """Returns an exact MULTIPLE of target_num_rows from the last pending bundle.""" pending_rows = last_pending_bundle.num_rows() or 0 assert total_pending_rows - pending_rows < self._target_num_rows, ( f"Total pending rows={total_pending_rows} should be less than target_num_rows={self._target_num_rows}, " "because last_pending_bundle should trigger building ready bundles" ) extra_rows = total_pending_rows % self._target_num_rows assert extra_rows < pending_rows return pending_rows - extra_rows """**For `ExactMultipleSize` strategy ONLY** Streaming repartition builds fixed-size outputs from a stream of inputs. We construct batches here to produce exactly sized outputs from arbitrary [start, end) slices across input blocks. The task builder submits a map task only after the total number of rows accumulated across pending blocks reaches target num rows (except during the final flush, which may emit a smaller tail block). This allows us to create target-sized batches without materializing entire large blocks on the driver. Detailed Implementation: 1. When a new bundle arrives, buffer it in the pending list. 2. Whenever the total number of rows in the pending bundles reaches the target row count, try to build a ready bundle. 3. Determine the slice needed from the final bundle so the ready bundle holds an exact multiple of the target rows, and add the remaining bundle to the pending bundles for the next iteration. 4. Submit that ready bundle to a remote map task; the task slices each block according to the slice metadata stored in the RefBundle (the bundle now contains n * target rows for n ≥ 1). 5. We configured the `OutputBlockSizeOption.target_num_rows_per_block` to the target number of rows per block in plan_streaming_repartition_op so the output buffer further splits the n * target rows into n blocks of exactly the target size. 6. Once upstream input is exhausted, flush any leftover pending bundles and repeat steps 1-5 for the tail. 7. The resulting blocks have lengths `[target, …, target, (total_rows % target)]`; ordering isn't guaranteed, but the remainder block should appear near the end. """ class RebundleQueue(BaseBundleQueue): """Incrementally builds task inputs to produce multiples of target-sized outputs.""" def __init__(self, strategy: RebundlingStrategy): super().__init__() self._strategy = strategy self._pending_bundles: Deque[RefBundle] = deque() self._ready_bundles: Deque[RefBundle] = deque() self._curr_consumed_bundles: List[RefBundle] = [] # The original bundles that formed a ready bundle self._consumed_bundles_list: Deque[List[RefBundle]] = deque() self._total_pending_rows: int = 0 def _merge_bundles(self): """Combine *ALL* pending_bundles into a single, ready bundle.""" from ray.data._internal.execution.interfaces import RefBundle merged_bundle = RefBundle.merge_ref_bundles(self._pending_bundles) # Update the metrics self._ready_bundles.append(merged_bundle) self._on_enqueue_bundle(merged_bundle) # Clear the pending queue since all bundles have been processed for bundle in self._pending_bundles: self._on_dequeue_bundle(bundle) self._pending_bundles.clear() self._total_pending_rows = 0 def _try_build_ready_bundle(self, flush_remaining: bool) -> int: """Attempts to build a ready bundle from a list of pending bundles by: - Checking the threshold to build a ready bundle defined by `RebundlingStrategy` - Appropiately keeping track of queue metrics Returns `True` if ready bundle built, otherwise `False` """ ready_bundles_built: int = 0 if self._pending_bundles and self._strategy.can_build_ready_bundle( self._total_pending_rows ): last_pending_bundle = self._pending_bundles.pop() # We now know `pending_bundle` is the bundle that enabled us to # build a ready bundle. Therefore, we may need to slice the bundle. rows_needed = self._strategy.rows_needed_from_last_pending_bundle( total_pending_rows=self._total_pending_rows, last_pending_bundle=last_pending_bundle, ) assert rows_needed > 0, ( "A refbundle has zero row-count but triggered building a ready bundle" "This is a bug in the Ray Data code." ) remaining_bundle: Optional[RefBundle] = None last_num_rows = last_pending_bundle.num_rows() or 0 if rows_needed < last_num_rows: sliced_bundle, remaining_bundle = last_pending_bundle.slice(rows_needed) # The original bundle was enqueued in add(). We need to dequeue it # and enqueue the sliced portion, since _merge_bundles will dequeue # sliced_bundle (which has different metrics than the original). self._on_dequeue_bundle(last_pending_bundle) self._on_enqueue_bundle(sliced_bundle) self._pending_bundles.append(sliced_bundle) else: assert rows_needed == last_num_rows self._pending_bundles.append(last_pending_bundle) self._merge_bundles() ready_bundles_built += 1 if remaining_bundle is not None: # Add back remaining sliced bundle that was not included to build # a ready bundle. self._pending_bundles.appendleft(remaining_bundle) self._total_pending_rows += remaining_bundle.num_rows() or 0 self._on_enqueue_bundle(remaining_bundle) # If we're flushing and have leftover bundles, convert them to a ready bundle. # Note: add() eagerly calls _try_build_ready_bundle after every insertion, so # pending rows are always below the threshold when finalize() is called. This # means at most one ready bundle is built per call (only the flush path fires). if flush_remaining and self._pending_bundles: self._merge_bundles() ready_bundles_built += 1 return ready_bundles_built @override def add(self, bundle: RefBundle, **kwargs: Any): from ray.data._internal.execution.interfaces import RefBundle num_rows = bundle.num_rows() or 0 if num_rows == 0: if self._pending_bundles: last = self._pending_bundles.pop() self._on_dequeue_bundle(last) merged = RefBundle.merge_ref_bundles([last, bundle]) self._pending_bundles.append(merged) self._on_enqueue_bundle(merged) else: self._pending_bundles.append(bundle) self._on_enqueue_bundle(bundle) return self._total_pending_rows += num_rows self._pending_bundles.append(bundle) self._on_enqueue_bundle(bundle) self._curr_consumed_bundles.append(bundle) ready_bundles_built = self._try_build_ready_bundle(flush_remaining=False) if ready_bundles_built > 0: assert ready_bundles_built == 1 self._consumed_bundles_list.append(self._curr_consumed_bundles) self._curr_consumed_bundles = [] @override def has_next(self) -> bool: return len(self._ready_bundles) > 0 @override def _get_next_inner(self) -> RefBundle: if not self.has_next(): raise ValueError("You can't pop from empty queue") ready_bundle = self._ready_bundles.popleft() # discard the original bundle self._consumed_bundles_list.popleft() return ready_bundle def get_next_with_original(self) -> Tuple[RefBundle, List[RefBundle]]: if not self.has_next(): raise ValueError("You can't pop from empty queue") ready_bundle = self._ready_bundles.popleft() self._on_dequeue_bundle(ready_bundle) consumed_bundle = self._consumed_bundles_list.popleft() return ready_bundle, consumed_bundle @override def peek_next(self) -> Optional[RefBundle]: if not self.has_next(): return None return self._ready_bundles[0] @override def finalize(self, **kwargs: Any): if len(self._pending_bundles) > 0: ready_bundles_built = self._try_build_ready_bundle(flush_remaining=True) assert ready_bundles_built == 1 self._consumed_bundles_list.append(self._curr_consumed_bundles) self._curr_consumed_bundles = [] @override def clear(self): self._reset_metrics() self._pending_bundles.clear() self._ready_bundles.clear() self._curr_consumed_bundles.clear() self._consumed_bundles_list.clear() self._total_pending_rows = 0