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

Code before:
        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.
        dataset_uuid: UUID of the dataset for this execution.
        preserve_order: Whether to preserve order in execution.

    Returns:
        The output as a RefBundle.
    """
    dag, stats = _get_execution_dag(
        executor,
        plan,
        preserve_order,
    )
    bundles = executor.execute(dag, initial_stats=stats)

Review comment (refactor):
In general, I think we should avoid breaking abstraction barriers and accessing internal attributes `plan._logical_plan`. Given this is PR is sort of a temporary state, I think it might be okay 

Code after:
                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.

    Returns:
        The output as a RefBundle.
    """
    dag, stats = _get_execution_dag(
        executor,
        plan,
        preserve_order,
    )
    bundles = executor.execute(dag, initial_stats=stats)
    ref_bundle = RefBundle.merge_ref_bundles(bundles)
    # Set the stats UUID after execution finishes.
    _set_stats_uuid_recursive(executor.get_stats(), dataset_uuid)