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

Code before:
        original RefBundle. Only after the entire iterator is exhausted,
        we cache the resulting metadata to the execution plan."""

        def __init__(self, base_iterator: OutputIterator, dag: "LogicalOperator"):
            # 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
            self._dag = dag

        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
                # collecting metadata. We can add this cached metadata to the plan.

                # Traverse the topology backwards and find the first available schema
                schema = next(reversed(topology.values()))._schema

                plan._cache.set_from_iteration(
                    self._dag, schema, self._num_rows, self._size_bytes
                )
                raise

        def _collect_metadata(self, bundle: RefBundle) -> RefBundle:
            """Collect the metadata from each output bundle and accumulate
            results, so we can access important information, such as
            row count, schema, etc., after iteration completes."""
            self._num_rows += bundle.num_rows()
            self._size_bytes += bundle.size_bytes()
            return bundle

    return CacheMetadataIterator(bundle_iter, plan._logical_plan.dag)


def execute_to_ref_bundle(
    executor: Executor,
    plan: ExecutionPlan,
    dataset_uuid: str,
    preserve_order: bool,
) -> RefBundle:
    """Execute a plan with the new executor and return the output as a RefBundle.

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

Review comment (suggestion):
I don't think you need to pass `dag` as a constructor argument for this class, because the class is defined within the `execute_to_legacy_bundle_iterator` function. I think you can just access the `plan` directly.
```suggestion
                    plan._logical_plan.dag, schema, self._num_rows, self._size_bytes
```

Code after:
        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
                # collecting metadata. We can add this cached metadata to the plan.

                # Traverse the topology backwards and find the first available schema
                schema = next(reversed(topology.values()))._schema

                dag = plan._logical_plan.dag
                plan._cache.set_num_rows(dag, self._num_rows)
                plan._cache.set_size_bytes(dag, self._size_bytes)
                plan._cache.set_schema(dag, schema)
                raise

        def _collect_metadata(self, bundle: RefBundle) -> RefBundle:
            """Collect the metadata from each output bundle and accumulate
            results, so we can access important information, such as
            row count, schema, etc., after iteration completes."""
            self._num_rows += bundle.num_rows()
            self._size_bytes += bundle.size_bytes()
            return bundle

    return CacheMetadataIterator(bundle_iter)


def execute_to_ref_bundle(
    executor: Executor,
    plan: ExecutionPlan,
    dataset_uuid: str,
    preserve_order: bool,
) -> RefBundle:
    """Execute a plan with the new executor and return the output as a RefBundle.

    Args:
        executor: The executor to use.
        plan: The legacy plan to execute.
        dataset_uuid: UUID of the dataset for this execution.
        preserve_order: Whether to preserve order in execution.