Repository: ray-project/ray
File: python/ray/data/_internal/execution/legacy_compat.py (Python)

Code before:
# Warn about tasks larger than this.
TASK_SIZE_WARN_THRESHOLD_BYTES = 100000

logger = logging.getLogger(__name__)


def execute_to_legacy_bundle_iterator(
    executor: StreamingExecutor,
    plan: ExecutionPlan,
) -> Iterator[RefBundle]:
    """Execute a plan with the new executor and return a bundle iterator.

    Args:
        executor: The executor to use.
        plan: The legacy plan to execute.

    Returns:
        The output as a bundle iterator.
    """
    dag, stats = _get_execution_dag(
        executor,
        plan,
        preserve_order=False,
    )

    callbacks = get_execution_callbacks(executor._data_context)
    bundle_iter = executor.execute(dag, initial_stats=stats, callbacks=callbacks)

    topology: "Topology" = executor._topology

    class CacheMetadataIterator(OutputIterator):
        """Wrapper for `bundle_iterator` above.

        For a given iterator which yields output RefBundles,
        collect the metadata from each output bundle, and yield the
        original RefBundle. Only after the entire iterator is exhausted,
        we cache the resulting metadata to the execution plan."""

        def __init__(self, base_iterator: OutputIterator):
            # Note: the base_iterator should be of type StreamIterator,
            # defined within `StreamingExecutor.execute()`. It must
            # support the `get_next()` method.
            self._base_iterator = base_iterator
            self._collected_metadata = BlockMetadata(
                num_rows=0,
                size_bytes=0,
                input_files=None,
                exec_stats=None,
            )

        def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:

Review comment (refactor):
Can we pass in the `DataContext` directly as a parameter to this function and `execute_to_iterator`? In general, I think we should avoid breaking abstraction barriers 

Code after:


def execute_to_legacy_bundle_iterator(
    executor: Executor,
    plan: ExecutionPlan,
    data_context: DataContext,
) -> Iterator[RefBundle]:
    """Execute a plan with the new executor and return a bundle iterator.

    Args:
        executor: The executor to use.
        plan: The legacy plan to execute.
        data_context: The DataContext to use for retrieving execution callbacks.

    Returns:
        The output as a bundle iterator.
    """
    dag, stats = _get_execution_dag(
        executor,
        plan,
        preserve_order=False,
    )

    callbacks = get_execution_callbacks(data_context)
    bundle_iter = executor.execute(dag, initial_stats=stats, callbacks=callbacks)

    topology: "Topology" = executor._topology

    class CacheMetadataIterator(OutputIterator):
        """Wrapper for `bundle_iterator` above.

        For a given iterator which yields output RefBundles,
        collect the metadata from each output bundle, and yield the
        original RefBundle. Only after the entire iterator is exhausted,
        we cache the resulting metadata to the execution plan."""

        def __init__(self, base_iterator: OutputIterator):
            # Note: the base_iterator should be of type StreamIterator,
            # defined within `StreamingExecutor.execute()`. It must
            # support the `get_next()` method.
            self._base_iterator = base_iterator
            self._num_rows = 0
            self._size_bytes = 0

        def get_next(self, output_split_idx: Optional[int] = None) -> RefBundle:
            try:
                bundle = self._base_iterator.get_next(output_split_idx)
                self._collect_metadata(bundle)
                return bundle
            except StopIteration:
                # Once the iterator is completely exhausted, we are done