Repository: ray-project/ray
File: python/ray/data/tests/test_read_datasource.py (Python)

Code before:
import time
from typing import TYPE_CHECKING, List, Optional

import pytest

import ray
from ray.data import ActorPoolStrategy, TaskPoolStrategy
from ray.data._internal.datasource.range_datasource import RangeDatasource
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
from ray.data._internal.execution.operators.map_operator import MapOperator
from ray.data._internal.logical.interfaces import LogicalPlan
from ray.data._internal.logical.operators.map_operator import MapBatches
from ray.data._internal.logical.operators.read_operator import Read
from ray.data._internal.logical.optimizers import PhysicalOptimizer
from ray.data._internal.planner import create_planner
from ray.data.block import Block, BlockMetadata
from ray.data.context import DataContext
from ray.data.datasource import Datasource, ReadTask
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.test_util import ConcurrencyCounter
from ray.tests.conftest import *  # noqa

if TYPE_CHECKING:
    from ray.actor import ActorHandle
    from ray.data._internal.logical.interfaces import LogicalOperator


Review comment (performance):
Can simplify the tests:

```
import time
from typing import TYPE_CHECKING, List, Optional

import pytest

import ray
from ray.data import ActorPoolStrategy, TaskPoolStrategy
from ray.data._internal.datasource.range_datasource import RangeDatasource
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
from ray.data._internal.execution.operators.map_operator import MapOperator
from ray.data._internal.logical.interfaces import LogicalPlan
from ray.data._internal.logical.operators.map_operator import MapBatches
from ray.data._internal.logical.operators.read_operator import Read
from ray.data._internal.logical.optimizers import PhysicalOptimizer
from ray.data._internal.planner import create_planner
from ray.data.block import Block, BlockMetadata
from ray.data.context import DataContext
from ray.data.datasource import Datasource, ReadTask
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.test_util import ConcurrencyCounter
from ray.tests.conftest import *  # noqa

if TYPE_CHECKING:
    from ray.actor import ActorHandle


def get_read_op(ds) -> Read:
    """Get the Read operator from a dataset's logical plan."""
    op = ds._plan._logical_plan.dag
    while not isinstance(op, Read) and hasattr(op, "input_dependencies"):
        op = op.input_dependencies[0]
    assert isinstance(op, Read), "Could not find Read operator"
    return op


class TestDatasource(Datasource):
    """Datasource that captures actor_id and optionally tracks concurrency."""

    def __init__(self, n: int, concurrency_counter: Optional["ActorHandle"] = None):
        super().__init__()
        self._n, self._counter = int(n), concurrency_counter

    def get_name(self) -> str:
        return "TestDatasource"

    def estimate_inmemory_data_size(self) -> Optional[int]:
        return 16 * self._n  # 2 columns, 8 bytes each

    def get_read_tasks(self, parallelism: int, **kwargs) -> List[ReadTask]:
        if self._n == 0:
            return []

        block_size = max(1, self._n // parallelism)
        tasks = []

        for start in range(0, self._n, block_size):
            count = min(block_size, self._n - start)
            counter = self._counter

            def read_fn(s=start, c=count, ctr=counter):
                import pyarrow as pa

                if ctr:
                    ray.get(ctr.inc.remote())
                try:
                    if ctr:
                        time.sleep(0.1)
                    actor_id = ray.get_runtime_context().get_actor_id()
                    yield pa.Table.from_pydict({
                        "value": list(range(s, s + c)),
                        "actor_id": [actor_id] * c,
                    })
                finally:
                    if ctr:
                        ray.get(ctr.decr.remote())

            tasks.append(ReadTask(
                read_fn,
                BlockMetadata(num_rows=count, size_bytes=16 * count, input_files=None, exec_stats=None),
            ))
        return tasks


@pytest.mark.parametrize("concurrency", [1, 2, 4])
def test_read_datasource_concurrency_limit(
    ray_start_regular_shared_2_cpus, concurrency
):
    """Test that concurrency parameter limits concurrent tasks."""
    counter = ConcurrencyCounter.remote()
    ds = ray.data.read_datasource(
        TestDatasource(n=100, concurrency_counter=counter),
        concurrency=concurrency,
        override_num_blocks=8,
    )
    ds.materialize()

    max_seen = ray.get(counter.get_max_concurrency.remote())
    assert max_seen <= concurrency, f"Max concurrency {max_seen} exceeded limit {concurrency}"


@pytest.mark.parametrize("concurrency,compute,expected_type,expected_size", [
    (None, None, TaskPoolStrategy, None),
    (2, None, TaskPoolStrategy, 2),
    (None, TaskPoolStrategy(size=4), TaskPoolStrategy, 4),
    (None, ActorPoolStrategy(size=2), ActorPoolStrategy, None),
    (1, ActorPoolStrategy(size=4), TaskPoolStrategy, 1),  # concurrency takes precedence
])
def test_read_datasource_compute_strategy(
    ray_start_regular_shared_2_cpus, concurrency, compute, expected_type, expected_size
):
    """Test compute strategy is correctly set based on concurrency and compute parameters."""
    ds = ray.data.read_datasource(
        RangeDatasource(n=100), concurrency=concurrency, compute=compute, override_num_blocks=4
    )
    read_op = get_read_op(ds)

    assert isinstance(read_op._compute, expected_type)
    if expected_size is not None:
        assert read_op._compute.size == expected_size


@pytest.mark.parametrize("compute,expect_actors", [
    (TaskPoolStrategy(), False),
    (ActorPoolStrategy(size=2), True),
])
def test_read_datasource_actor_execution(
    ray_start_regular_shared_2_cpus, compute, expect_actors, target_max_block_size_infinite_or_default
):
    """Test ReadTasks execute in actors when using ActorPoolStrategy."""
    ds = ray.data.read_datasource(TestDatasource(n=100), compute=compute, override_num_blocks=4)
    actor_ids = {row["actor_id"] for row in ds.take_all()}

    if expect_actors:
        assert None not in actor_ids, "Expected all actor_ids to be non-None"
        if isinstance(compute, ActorPoolStrategy) and compute.min_size == compute.max_size:
            assert len(actor_ids) <= compute.min_size
    else:
        assert actor_ids == {None}, f"Expected all actor_ids to be None, got {actor_ids}"


@pytest.mark.parametrize("read_compute,map_compute,should_fuse", [
    (TaskPoolStrategy(), ActorPoolStrategy(size=2), True),
    (ActorPoolStrategy(size=2), ActorPoolStrategy(size=2), False),
    (TaskPoolStrategy(size=4), TaskPoolStrategy(size=4), True),
    (TaskPoolStrategy(size=2), TaskPoolStrategy(size=4), False),
])
def test_read_datasource_map_batches_fusion(
    ray_start_regular_shared_2_cpus, read_compute, map_compute, should_fuse
):
    """Test operator fusion behavior with different compute strategies."""
    ctx = DataContext.get_current()
    read_op = Read(
        datasource=RangeDatasource(n=100),
        datasource_or_legacy_reader=RangeDatasource(n=100),
        parallelism=4, num_outputs=4, compute=read_compute,
    )
    map_op = MapBatches(read_op, lambda x: x, compute=map_compute)

    planner = create_planner()
    physical_plan = PhysicalOptimizer().optimize(planner.plan(LogicalPlan(map_op, ctx)))
    physical_op = physical_plan.dag

    assert isinstance(physical_op, MapOperator)
    if should_fuse:
        assert "->" in physical_op.name, f"Expected fused operator, got {physical_op.name}"
        assert isinstance(physical_op.input_dependencies[0], InputDataBuffer)
        assert any(isinstance(op, Read) for op in physical_op._logical_operators)
        assert any(isinstance(op, MapBatches) for op in physical_op._logical_operators)
    else:
        assert physical_op.name == "MapBatches(<lambda>)"
        assert isinstance(physical_op.input_dependencies[0], MapOperator)


@pytest.mark.parametrize("compute", [
    TaskPoolStrategy(),
    TaskPoolStrategy(size=2),
    ActorPoolStrategy(size=2),
])
def test_read_datasource_basic_functionality(
    ray_start_regular_shared_2_cpus, compute
):
    """Test read_datasource works correctly with different compute strategies."""
    result = ray.data.read_datasource(
        RangeDatasource(n=100), compute=compute, override_num_blocks=4
    ).take_all()

    assert len(result) == 100
    assert sorted(row["value"] for row in result) == list(range(100))


if __name__ == "__main__":
    import sys
    sys.exit(pytest.main(["-v", __file__]))
```

Code after:
import time
from typing import TYPE_CHECKING, List, Optional

import pandas as pd
import pytest

import ray
from ray.data import ActorPoolStrategy, TaskPoolStrategy
from ray.data._internal.datasource.range_datasource import RangeDatasource
from ray.data._internal.logical.operators.read_operator import Read
from ray.data._internal.util import rows_same
from ray.data.block import Block, BlockMetadata
from ray.data.context import DataContext
from ray.data.datasource import Datasource, ReadTask
from ray.data.tests.conftest import *  # noqa
from ray.tests.conftest import *  # noqa

if TYPE_CHECKING:
    from ray.actor import ActorHandle
    from ray.data._internal.logical.interfaces import LogicalOperator


def find_read_op(op: "LogicalOperator") -> Optional[Read]:
    """Find the Read operator in the logical plan."""
    if isinstance(op, Read):
        return op