[Data] Implement Operators for union() (ray-project#36242)

Implement the `LogicalOperator` and `PhysicalOperator` for `Dataset.union()`, and make `union()` lazy. This PR also introduces `Nary` and `NaryOperator` Logical/Physical Operators to support abstraction for `Union` and `Zip` operators.

---------

Signed-off-by: Scott Lee <[email protected]>

python/ray/data/_internal/execution/bulk_executor.py

@@ -62,7 +62,7 @@ def execute_recursive(op: PhysicalOperator) -> List[RefBundle]:
  for i, ref_bundles in enumerate(inputs):
  for r in ref_bundles:
  op.add_input(r, input_index=i)
- op.inputs_done()
+ op.all_inputs_done()
  output = _naive_run_until_complete(op)
  finally:
  op.shutdown()

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

@@ -396,7 +396,15 @@ def add_input(self, refs: RefBundle, input_index: int) -> None:
  """
  raise NotImplementedError
 
- def inputs_done(self) -> None:
+ def input_done(self, input_index: int) -> None:
+ """Called when the upstream operator at index `input_index` has completed().
+
+ After this is called, the executor guarantees that no more inputs will be added
+ via `add_input` for the given input index.
+ """
+ pass
+
+ def all_inputs_done(self) -> None:
  """Called when all upstream operators have completed().
 
  After this is called, the executor guarantees that no more inputs will be added

python/ray/data/_internal/execution/legacy_compat.py

@@ -14,10 +14,12 @@
  RefBundle,
  TaskContext,
 )
-from ray.data._internal.execution.operators.all_to_all_operator import AllToAllOperator
+from ray.data._internal.execution.operators.base_physical_operator import (
+ AllToAllOperator,
+)
 from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
+from ray.data._internal.execution.operators.limit_operator import LimitOperator
 from ray.data._internal.execution.operators.map_operator import MapOperator
-from ray.data._internal.execution.operators.one_to_one_operator import LimitOperator
 from ray.data._internal.execution.util import make_callable_class_concurrent
 from ray.data._internal.lazy_block_list import LazyBlockList
 from ray.data._internal.logical.optimizers import get_execution_plan

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

@@ -230,10 +230,10 @@ def notify_work_completed(
  # For either a completed task or ready worker, we try to dispatch queued tasks.
  self._dispatch_tasks()
 
- def inputs_done(self):
+ def all_inputs_done(self):
  # Call base implementation to handle any leftover bundles. This may or may not
  # trigger task dispatch.
- super().inputs_done()
+ super().all_inputs_done()
 
  # Mark inputs as done so future task dispatch will kill all inactive workers
  # once the bundle queue is exhausted.

python/ray/data/_internal/execution/operators/all_to_all_operator.py → python/ray/data/_internal/execution/operators/base_physical_operator.py

@@ -6,10 +6,34 @@
  RefBundle,
  TaskContext,
 )
+from ray.data._internal.logical.interfaces import LogicalOperator
 from ray.data._internal.progress_bar import ProgressBar
 from ray.data._internal.stats import StatsDict
 
 
+class OneToOneOperator(PhysicalOperator):
+ """An operator that has one input and one output dependency.
+
+ This operator serves as the base for map, filter, limit, etc.
+ """
+
+ def __init__(
+ self,
+ name: str,
+ input_op: PhysicalOperator,
+ ):
+ """Create a OneToOneOperator.
+ Args:
+ input_op: Operator generating input data for this op.
+ name: The name of this operator.
+ """
+ super().__init__(name, [input_op])
+
+ @property
+ def input_dependency(self) -> PhysicalOperator:
+ return self.input_dependencies[0]
+
+
 class AllToAllOperator(PhysicalOperator):
  """A blocking operator that executes once its inputs are complete.
 
@@ -25,7 +49,6 @@ def __init__(
  name: str = "AllToAll",
  ):
  """Create an AllToAllOperator.
-
  Args:
  bulk_fn: The blocking transformation function to run. The inputs are the
  list of input ref bundles, and the outputs are the output ref bundles
@@ -57,15 +80,15 @@ def add_input(self, refs: RefBundle, input_index: int) -> None:
  assert input_index == 0, input_index
  self._input_buffer.append(refs)
 
- def inputs_done(self) -> None:
+ def all_inputs_done(self) -> None:
  ctx = TaskContext(
  task_idx=self._next_task_index,
  sub_progress_bar_dict=self._sub_progress_bar_dict,
  )
  self._output_buffer, self._stats = self._bulk_fn(self._input_buffer, ctx)
  self._next_task_index += 1
  self._input_buffer.clear()
- super().inputs_done()
+ super().all_inputs_done()
 
  def has_next(self) -> bool:
  return len(self._output_buffer) > 0
@@ -102,3 +125,23 @@ def close_sub_progress_bars(self):
  if self._sub_progress_bar_dict is not None:
  for sub_bar in self._sub_progress_bar_dict.values():
  sub_bar.close()
+
+
+class NAryOperator(PhysicalOperator):
+ """An operator that has multiple input dependencies and one output.
+
+ This operator serves as the base for union, zip, etc.
+ """
+
+ def __init__(
+ self,
+ *input_ops: LogicalOperator,
+ ):
+ """Create a OneToOneOperator.
+ Args:
+ input_op: Operator generating input data for this op.
+ name: The name of this operator.
+ """
+ input_names = ", ".join([op._name for op in input_ops])
+ op_name = f"{self.__class__.__name__}({input_names})"
+ super().__init__(op_name, list(input_ops))

python/ray/data/_internal/execution/operators/one_to_one_operator.py → python/ray/data/_internal/execution/operators/limit_operator.py

@@ -4,35 +4,15 @@
 
 import ray
 from ray.data._internal.execution.interfaces import PhysicalOperator, RefBundle
+from ray.data._internal.execution.operators.base_physical_operator import (
+ OneToOneOperator,
+)
 from ray.data._internal.remote_fn import cached_remote_fn
 from ray.data._internal.stats import StatsDict
 from ray.data.block import Block, BlockAccessor, BlockMetadata
 from ray.types import ObjectRef
 
 
-class OneToOneOperator(PhysicalOperator):
- """An operator that has one input and one output dependency.
- This operator serves as the base for map, filter, limit, etc.
- """
-
- def __init__(
- self,
- name: str,
- input_op: PhysicalOperator,
- ):
- """Create a OneToOneOperator.
-
- Args:
- input_op: Operator generating input data for this op.
- name: The name of this operator.
- """
- super().__init__(name, [input_op])
-
- @property
- def input_dependency(self) -> PhysicalOperator:
- return self.input_dependencies[0]
-
-
 class LimitOperator(OneToOneOperator):
  """Physical operator for limit."""
 
@@ -49,7 +29,7 @@ def __init__(
  self._cur_output_bundles = 0
  super().__init__(self._name, input_op)
  if self._limit <= 0:
- self.inputs_done()
+ self.all_inputs_done()
 
  def _limit_reached(self) -> bool:
  return self._consumed_rows >= self._limit
@@ -99,7 +79,7 @@ def slice_fn(block, metadata, num_rows) -> Tuple[Block, BlockMetadata]:
  )
  self._buffer.append(out_refs)
  if self._limit_reached():
- self.inputs_done()
+ self.all_inputs_done()
 
  def has_next(self) -> bool:
  return len(self._buffer) > 0

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

@@ -19,7 +19,9 @@
  RefBundle,
  TaskContext,
 )
-from ray.data._internal.execution.operators.one_to_one_operator import OneToOneOperator
+from ray.data._internal.execution.operators.base_physical_operator import (
+ OneToOneOperator,
+)
 from ray.data._internal.memory_tracing import trace_allocation
 from ray.data._internal.stats import StatsDict
 from ray.data.block import Block, BlockAccessor, BlockExecStats, BlockMetadata
@@ -280,13 +282,13 @@ def _handle_task_done(self, task: "_TaskState"):
  if self._metrics.cur > self._metrics.peak:
  self._metrics.peak = self._metrics.cur
 
- def inputs_done(self):
+ def all_inputs_done(self):
  self._block_ref_bundler.done_adding_bundles()
  if self._block_ref_bundler.has_bundle():
  # Handle any leftover bundles in the bundler.
  bundle = self._block_ref_bundler.get_next_bundle()
  self._add_bundled_input(bundle)
- super().inputs_done()
+ super().all_inputs_done()
 
  def has_next(self) -> bool:
  assert self._started

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

@@ -101,8 +101,8 @@ def add_input(self, bundle, input_index) -> None:
  self._buffer.append(bundle)
  self._dispatch_bundles()
 
- def inputs_done(self) -> None:
- super().inputs_done()
+ def all_inputs_done(self) -> None:
+ super().all_inputs_done()
  if not self._equal:
  self._dispatch_bundles(dispatch_all=True)
  assert not self._buffer, "Should have dispatched all bundles."

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

@@ -0,0 +1,107 @@
+from typing import List, Optional
+
+from ray.data._internal.execution.interfaces import (
+ ExecutionOptions,
+ PhysicalOperator,
+ RefBundle,
+)
+from ray.data._internal.execution.operators.base_physical_operator import NAryOperator
+from ray.data._internal.stats import StatsDict
+
+
+class UnionOperator(NAryOperator):
+ """An operator that combines output blocks from
+ two or more input operators into a single output."""
+
+ def __init__(
+ self,
+ *input_ops: PhysicalOperator,
+ ):
+ """Create a UnionOperator.
+
+ Args:
+ input_ops: Operators generating input data for this operator to union.
+ """
+
+ # By default, union does not preserve the order of output blocks.
+ # To preserve the order, configure ExecutionOptions accordingly.
+ self._preserve_order = False
+
+ # Intermediary buffers used to store blocks from each input dependency.
+ # Only used when `self._prserve_order` is True.
+ self._input_buffers: List[List[RefBundle]] = [[] for _ in range(len(input_ops))]
+
+ # The index of the input dependency that is currently the source of
+ # the output buffer. New inputs from this input dependency will be added
+ # directly to the output buffer. Only used when `self._preserve_order` is True.
+ self._input_idx_to_output = 0
+
+ self._output_buffer: List[RefBundle] = []
+ self._stats: StatsDict = {}
+ super().__init__(*input_ops)
+
+ def start(self, options: ExecutionOptions):
+ # Whether to preserve the order of the input data (both the
+ # order of the input operators and the order of the blocks within).
+ self._preserve_order = options.preserve_order
+ super().start(options)
+
+ def num_outputs_total(self) -> Optional[int]:
+ num_outputs = 0
+ for input_op in self.input_dependencies:
+ op_num_outputs = input_op.num_outputs_total()
+ # If at least one of the input ops has an unknown number of outputs,
+ # the number of outputs of the union operator is unknown.
+ if op_num_outputs is None:
+ return None
+ num_outputs += op_num_outputs
+ return num_outputs
+
+ def add_input(self, refs: RefBundle, input_index: int) -> None:
+ assert not self.completed()
+ assert 0 <= input_index <= len(self._input_dependencies), input_index
+
+ if not self._preserve_order:
+ self._output_buffer.append(refs)
+ else:
+ if input_index == self._input_idx_to_output:
+ self._output_buffer.append(refs)
+ else:
+ self._input_buffers[input_index].append(refs)
+
+ def input_done(self, input_index: int) -> None:
+ """When `self._preserve_order` is True, change the
+ output buffer source to the next input dependency
+ once the current input dependency calls `input_done()`."""
+ if not self._preserve_order:
+ return
+ if not input_index == self._input_idx_to_output:
+ return
+ next_input_idx = self._input_idx_to_output + 1
+ if next_input_idx < len(self._input_buffers):
+ self._output_buffer.extend(self._input_buffers[next_input_idx])
+ self._input_buffers[next_input_idx].clear()
+ self._input_idx_to_output = next_input_idx
+ super().input_done(input_index)
+
+ def all_inputs_done(self) -> None:
+ # Note that in the case where order is not preserved, all inputs
+ # are directly added to the output buffer as soon as they are received,
+ # so there is no need to check any intermediary buffers.
+ if self._preserve_order:
+ for idx, input_buffer in enumerate(self._input_buffers):
+ assert len(input_buffer) == 0, (
+ f"Input at index {idx} still has "
+ f"{len(input_buffer)} blocks remaining."
+ )
+ super().all_inputs_done()
+
+ def has_next(self) -> bool:
+ # Check if the output buffer still contains at least one block.
+ return len(self._output_buffer) > 0
+
+ def get_next(self) -> RefBundle:
+ return self._output_buffer.pop(0)
+
+ def get_stats(self) -> StatsDict:
+ return self._stats

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

@@ -59,13 +59,13 @@ def add_input(self, refs: RefBundle, input_index: int) -> None:
  else:
  self._right_buffer.append(refs)
 
- def inputs_done(self) -> None:
+ def all_inputs_done(self) -> None:
  self._output_buffer, self._stats = self._zip(
  self._left_buffer, self._right_buffer
  )
  self._left_buffer.clear()
  self._right_buffer.clear()
- super().inputs_done()
+ super().all_inputs_done()
 
  def has_next(self) -> bool:
  return len(self._output_buffer) > 0

python/ray/data/_internal/execution/streaming_executor.py

@@ -27,6 +27,7 @@
  build_streaming_topology,
  process_completed_tasks,
  select_operator_to_run,
+ update_operator_states,
 )
 from ray.data._internal.progress_bar import ProgressBar
 from ray.data._internal.stats import DatasetStats
@@ -257,6 +258,8 @@ def _scheduling_loop_step(self, topology: Topology) -> bool:
  autoscaling_state=self._autoscaling_state,
  )
 
+ update_operator_states(topology)
+
  # Update the progress bar to reflect scheduling decisions.
  for op_state in topology.values():
  op_state.refresh_progress_bar()