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

Code before:
    stats = _get_initial_stats_from_plan(plan)

    # Enforce to preserve ordering if the plan has operators
    # required to do so, such as Zip and Sort.
    if preserve_order or plan.require_preserve_order():
        executor._options.preserve_order = True

    return dag, stats


def _get_initial_stats_from_plan(plan: ExecutionPlan) -> DatasetStats:
    if plan._snapshot_bundle is not None:
        return plan._snapshot_stats
    # For Datasets created from "read_xxx", `plan._in_stats` contains useless data.
    # For Datasets created from "from_xxx", we need to use `plan._in_stats` as
    # the initial stats. Because the `FromXxx` logical operators will be translated to
    # "InputDataBuffer" physical operators, which will be ignored when generating
    # stats, see `StreamingExecutor._generate_stats`.
    # TODO(hchen): Unify the logic by saving the initial stats in `InputDataBuffer
    if plan.has_lazy_input():
        return DatasetStats(metadata={}, parent=None)
    else:
        return plan._in_stats


def _bundles_to_ref_bundle(bundles: Iterator[RefBundle]) -> RefBundle:
    blocks = []
    owns_blocks = True
    bundle_list = list(bundles)
    schema = _take_first_non_empty_schema(
        ref_bundle.schema for ref_bundle in bundle_list
    )

    for ref_bundle in bundle_list:
        if not ref_bundle.owns_blocks:
            owns_blocks = False
        blocks.extend(ref_bundle.blocks)

    return RefBundle(blocks=tuple(blocks), owns_blocks=owns_blocks, schema=schema)


def _set_stats_uuid_recursive(stats: DatasetStats, dataset_uuid: str) -> None:
    if not stats.dataset_uuid:
        stats.dataset_uuid = dataset_uuid
    for parent in stats.parents or []:
        _set_stats_uuid_recursive(parent, dataset_uuid)


Review comment (suggestion):
I think we can reuse `merge_ref_bundles` with some changes to it? Something like

```py
    def merge_ref_bundles(cls, bundles: Iterable["RefBundle"]) -> "RefBundle":
        bundles = list(bundles)
        if not bundles:
            return cls(blocks=(), owns_blocks=True, schema=None)
        merged_blocks = list(itertools.chain.from_iterable(bundle.blocks for bundle in bundles))
        merged_slices = list(itertools.chain.from_iterable(bundle.slices for bundle in bundles))
        owns_blocks = all(bundle.owns_blocks for bundle in bundles)
        schema = _take_first_non_empty_schema(bundle.schema for bundle in bundles)
        return cls(
            blocks=tuple(merged_blocks),
            schema=schema,
            owns_blocks=owns_blocks,
            slices=merged_slices,
        )
```

Code after:

    # Enforce to preserve ordering if the plan has operators
    # required to do so, such as Zip and Sort.
    if preserve_order or plan.require_preserve_order():
        executor._options.preserve_order = True

    return dag, stats


def _get_initial_stats_from_plan(plan: ExecutionPlan) -> DatasetStats:
    if plan._snapshot_bundle is not None:
        return plan._snapshot_stats
    # For Datasets created from "read_xxx", `plan._in_stats` contains useless data.
    # For Datasets created from "from_xxx", we need to use `plan._in_stats` as
    # the initial stats. Because the `FromXxx` logical operators will be translated to
    # "InputDataBuffer" physical operators, which will be ignored when generating
    # stats, see `StreamingExecutor._generate_stats`.
    # TODO(hchen): Unify the logic by saving the initial stats in `InputDataBuffer
    if plan.has_lazy_input():
        return DatasetStats(metadata={}, parent=None)
    else:
        return plan._in_stats


def _set_stats_uuid_recursive(stats: DatasetStats, dataset_uuid: str) -> None:
    if not stats.dataset_uuid:
        stats.dataset_uuid = dataset_uuid
    for parent in stats.parents or []:
        _set_stats_uuid_recursive(parent, dataset_uuid)