chore: import upstream snapshot with attribution

python/ray/data/tests/test_consumption.py

@@ -0,0 +1,884 @@
+import logging
+import os
+import sys
+import time
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.block_builder import BlockBuilder
+from ray.data._internal.datasource.csv_datasink import CSVDatasink
+from ray.data._internal.datasource.csv_datasource import CSVDatasource
+from ray.data._internal.datasource.range_datasource import RangeDatasource
+from ray.data._internal.execution.interfaces.ref_bundle import (
+    _ref_bundles_iterator_to_block_refs_list,
+)
+from ray.data.block import BlockAccessor
+from ray.data.dataset import Dataset, MaterializedDataset
+from ray.data.datasource.util import (
+    _validate_head_node_resources_for_local_scheduling,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.conftest import (
+    CoreExecutionMetrics,
+    assert_core_execution_metrics_equals,
+    get_initial_core_execution_metrics_snapshot,
+)
+from ray.data.tests.util import column_udf, extract_values
+from ray.tests.conftest import *  # noqa
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+
+def test_schema(ray_start_regular):
+    last_snapshot = get_initial_core_execution_metrics_snapshot()
+
+    ds2 = ray.data.range(10, override_num_blocks=10)
+    ds3 = ds2.repartition(5)
+    ds3 = ds3.materialize()
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                "ReadRange": 10,
+                "reduce": 5,
+            }
+        ),
+        last_snapshot,
+    )
+
+    ds4 = ds3.map(lambda x: {"a": "hi", "b": 1.0}).limit(5).repartition(1)
+    ds4 = ds4.materialize()
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                "Map(<lambda>)": lambda count: count <= 5,
+                "slice_fn": 1,
+                "reduce": 1,
+            }
+        ),
+        last_snapshot,
+    )
+
+    ds2_schema = ds2.schema(fetch_if_missing=False)
+    assert ds2_schema is not None
+    assert ds2_schema.names == ["id"]
+    assert not isinstance(ds2, MaterializedDataset)
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(task_count={}), last_snapshot
+    )
+
+    ds3_schema = ds3.schema(fetch_if_missing=False)
+    assert ds3_schema is not None
+    assert ds3_schema.names == ["id"]
+    assert isinstance(ds3, MaterializedDataset)
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(task_count={}), last_snapshot
+    )
+    ds4_schema = ds4.schema(fetch_if_missing=False)
+    assert ds4_schema is not None
+    assert ds4_schema.names == ["a", "b"]
+    assert isinstance(ds4, MaterializedDataset)
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(task_count={}), last_snapshot
+    )
+
+
+def test_schema_no_execution(ray_start_regular):
+    last_snapshot = get_initial_core_execution_metrics_snapshot()
+    ds = ray.data.range(100, override_num_blocks=10)
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(task_count={}),
+        last_snapshot,
+    )
+    # We do not kick off the read task by default.
+    schema = ds.schema()
+    assert schema.names == ["id"]
+
+    # Fetching the schema does not trigger execution, since
+    # the schema is known beforehand for RangeDatasource.
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(task_count={}), last_snapshot
+    )
+    # Fetching the schema should not trigger execution of extra read tasks.
+
+
+def test_schema_cached(ray_start_regular):
+    def check_schema_cached(ds, expected_task_count, last_snapshot):
+        schema = ds.schema()
+        last_snapshot = assert_core_execution_metrics_equals(
+            CoreExecutionMetrics(expected_task_count), last_snapshot
+        )
+        assert schema.names == ["a"]
+        cached_schema = ds.schema(fetch_if_missing=False)
+        assert cached_schema is not None
+        assert schema == cached_schema
+        last_snapshot = assert_core_execution_metrics_equals(
+            CoreExecutionMetrics({}), last_snapshot
+        )
+        return last_snapshot
+
+    last_snapshot = get_initial_core_execution_metrics_snapshot()
+    ds = ray.data.from_items([{"a": i} for i in range(100)], override_num_blocks=10)
+    last_snapshot = check_schema_cached(ds, {}, last_snapshot)
+
+    # Add a map_batches operator so that we are forced to compute the schema.
+    ds = ds.map_batches(lambda x: x)
+    last_snapshot = check_schema_cached(
+        ds,
+        {
+            "MapBatches(<lambda>)": lambda count: count <= 5,
+            "slice_fn": 1,
+        },
+        last_snapshot,
+    )
+
+
+def test_avoid_placement_group_capture(shutdown_only):
+    ray.init(num_cpus=2)
+
+    @ray.remote
+    def run():
+        ds = ray.data.range(5)
+        assert sorted(
+            extract_values("id", ds.map(column_udf("id", lambda x: x + 1)).take())
+        ) == [1, 2, 3, 4, 5]
+        assert ds.count() == 5
+        assert sorted(extract_values("id", ds.iter_rows())) == [0, 1, 2, 3, 4]
+
+    pg = ray.util.placement_group([{"CPU": 1}])
+    ray.get(
+        run.options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=pg, placement_group_capture_child_tasks=True
+            )
+        ).remote()
+    )
+
+
+@pytest.fixture
+def remove_named_placement_groups():
+    yield
+    for info in ray.util.placement_group_table().values():
+        if info["name"]:
+            pg = ray.util.get_placement_group(info["name"])
+            ray.util.remove_placement_group(pg)
+
+
+def test_ray_remote_args_fn(shutdown_only, remove_named_placement_groups):
+    ray.init()
+
+    pg = ray.util.placement_group([{"CPU": 1}], name="test_pg")
+
+    def ray_remote_args_fn():
+        scheduling_strategy = PlacementGroupSchedulingStrategy(placement_group=pg)
+        return {"scheduling_strategy": scheduling_strategy}
+
+    class ActorClass:
+        def __call__(self, batch):
+            assert ray.util.get_current_placement_group() == pg
+            return batch
+
+    ray.data.range(1).map_batches(
+        ActorClass, concurrency=1, ray_remote_args_fn=ray_remote_args_fn
+    ).take_all()
+
+
+def test_dataset_lineage_serialization(shutdown_only):
+    ray.init()
+    ds = ray.data.range(10)
+    ds = ds.map(column_udf("id", lambda x: x + 1))
+    ds = ds.map(column_udf("id", lambda x: x + 1))
+    ds = ds.random_shuffle()
+    uuid = ds._get_uuid()
+    plan_uuid = ds._uuid
+
+    serialized_ds = ds.serialize_lineage()
+
+    ray.shutdown()
+    ray.init()
+
+    ds = Dataset.deserialize_lineage(serialized_ds)
+    # Check Dataset state.
+    assert ds._get_uuid() == uuid
+    assert ds._uuid == plan_uuid
+    # Check Dataset content.
+    assert ds.count() == 10
+    assert sorted(extract_values("id", ds.take())) == list(range(2, 12))
+
+
+def test_dataset_lineage_serialization_unsupported(shutdown_only):
+    ray.init()
+    # In-memory data sources not supported.
+    ds = ray.data.from_items(list(range(10)))
+    ds = ds.map(column_udf("item", lambda x: x + 1))
+    ds = ds.map(column_udf("item", lambda x: x + 1))
+
+    with pytest.raises(ValueError):
+        ds.serialize_lineage()
+
+    # In-memory data source unions not supported.
+    ds = ray.data.from_items(list(range(10)))
+    ds1 = ray.data.from_items(list(range(10, 20)))
+    ds2 = ds.union(ds1)
+
+    with pytest.raises(ValueError):
+        ds2.serialize_lineage()
+
+    # Lazy read unions supported.
+    ds = ray.data.range(10)
+    ds1 = ray.data.range(20)
+    ds2 = ds.union(ds1)
+
+    serialized_ds = ds2.serialize_lineage()
+    ds3 = Dataset.deserialize_lineage(serialized_ds)
+    assert set(extract_values("id", ds3.take(30))) == set(
+        list(range(10)) + list(range(20))
+    )
+
+    # Zips not supported.
+    ds = ray.data.from_items(list(range(10)))
+    ds1 = ray.data.from_items(list(range(10, 20)))
+    ds2 = ds.zip(ds1)
+
+    with pytest.raises(ValueError):
+        ds2.serialize_lineage()
+
+
+def test_basic(ray_start_regular_shared):
+    ds = ray.data.range(5)
+    assert sorted(
+        extract_values("id", ds.map(column_udf("id", lambda x: x + 1)).take())
+    ) == [1, 2, 3, 4, 5]
+    assert ds.count() == 5
+    assert sorted(extract_values("id", ds.iter_rows())) == [0, 1, 2, 3, 4]
+
+
+def test_range(ray_start_regular_shared):
+    ds = ray.data.range(10, override_num_blocks=10)
+    assert ds._logical_plan.initial_num_blocks() == 10
+    assert ds.count() == 10
+    assert ds.take() == [{"id": i} for i in range(10)]
+
+    ds = ray.data.range(10, override_num_blocks=2)
+    assert ds._logical_plan.initial_num_blocks() == 2
+    assert ds.count() == 10
+    assert ds.take() == [{"id": i} for i in range(10)]
+
+
+def test_empty_dataset(ray_start_regular_shared):
+    ds = ray.data.range(0)
+    assert ds.count() == 0
+    assert ds.size_bytes() == 0
+    assert ds.schema() is None
+
+    ds = ray.data.range(1)
+    ds = ds.filter(lambda x: x["id"] > 1)
+    ds = ds.materialize()
+    assert (
+        str(ds)
+        == "MaterializedDataset(num_blocks=1, num_rows=0, schema=Unknown schema)"
+    )
+
+    # Test map on empty dataset.
+    ds = ray.data.from_items([])
+    ds = ds.map(lambda x: x)
+    ds = ds.materialize()
+    assert ds.count() == 0
+
+    # Test filter on empty dataset.
+    ds = ray.data.from_items([])
+    ds = ds.filter(lambda: True)
+    ds = ds.materialize()
+    assert ds.count() == 0
+
+
+@ray.remote
+class Counter:
+    def __init__(self):
+        self.value = 0
+
+    def increment(self):
+        self.value += 1
+        return self.value
+
+
+def test_cache_dataset(ray_start_regular_shared):
+
+    c = Counter.remote()
+
+    def inc(x):
+        ray.get(c.increment.remote())
+        return x
+
+    ds = ray.data.range(1)
+    ds = ds.map(inc)
+    assert not isinstance(ds, MaterializedDataset)
+    ds2 = ds.materialize()
+    assert isinstance(ds2, MaterializedDataset)
+    assert not isinstance(ds, MaterializedDataset)
+
+    # Tests standard iteration uses the materialized blocks.
+    for _ in range(10):
+        ds2.take_all()
+
+    assert ray.get(c.increment.remote()) == 2
+
+    # Tests streaming iteration uses the materialized blocks.
+    for _ in range(10):
+        list(ds2.streaming_split(1)[0].iter_batches())
+
+    assert ray.get(c.increment.remote()) == 3
+
+
+def test_columns(ray_start_regular_shared):
+    ds = ray.data.range(1)
+    assert ds.columns() == ds.schema().names
+    assert ds.columns() == ["id"]
+
+    ds = ds.map(lambda x: x)
+    assert ds.columns(fetch_if_missing=False) is None
+
+
+def test_schema_repr(ray_start_regular_shared):
+    ds = ray.data.from_items([{"text": "spam", "number": 0}])
+    # fmt: off
+    expected_repr = (
+        "Column  Type\n"
+        "------  ----\n"
+        "text    string\n"
+        "number  int64"
+    )
+    # fmt:on
+    assert repr(ds.schema()) == expected_repr
+
+    ds = ray.data.from_items([{"long_column_name": "spam"}])
+    # fmt: off
+    expected_repr = (
+        "Column            Type\n"
+        "------            ----\n"
+        "long_column_name  string"
+    )
+    # fmt: on
+    assert repr(ds.schema()) == expected_repr
+
+
+def _check_none_computed(ds):
+    # In streaming executor, ds.take() will not invoke partial execution
+    # in LazyBlocklist.
+    assert not ds._has_computed_output()
+
+
+def test_lazy_loading_exponential_rampup(ray_start_regular_shared):
+    ds = ray.data.range(100, override_num_blocks=20)
+    _check_none_computed(ds)
+    assert extract_values("id", ds.take(10)) == list(range(10))
+    _check_none_computed(ds)
+    assert extract_values("id", ds.take(20)) == list(range(20))
+    _check_none_computed(ds)
+    assert extract_values("id", ds.take(30)) == list(range(30))
+    _check_none_computed(ds)
+    assert extract_values("id", ds.take(50)) == list(range(50))
+    _check_none_computed(ds)
+    assert extract_values("id", ds.take(100)) == list(range(100))
+    _check_none_computed(ds)
+
+
+def test_dataset_repr_not_materialized(ray_start_regular_shared, restore_data_context):
+    ds = ray.data.range(5)
+    assert repr(ds) == (
+        "shape: (5, 1)\n"
+        "╭───────╮\n"
+        "│ id    │\n"
+        "│ ---   │\n"
+        "│ int64 │\n"
+        "╰───────╯\n"
+        "(Dataset isn't materialized)"
+    )
+
+
+def test_dataset_repr_materialized(ray_start_regular_shared, restore_data_context):
+    materialized = ray.data.range(5).materialize()
+    assert repr(materialized) == (
+        "shape: (5, 1)\n"
+        "╭───────╮\n"
+        "│ id    │\n"
+        "│ ---   │\n"
+        "│ int64 │\n"
+        "╞═══════╡\n"
+        "│ 0     │\n"
+        "│ 1     │\n"
+        "│ 2     │\n"
+        "│ 3     │\n"
+        "│ 4     │\n"
+        "╰───────╯\n"
+        "(Showing 5 of 5 rows)"
+    )
+
+
+def test_dataset_repr_gap(ray_start_regular_shared, restore_data_context):
+    ds_with_gap = ray.data.range(20).materialize()
+    assert repr(ds_with_gap) == (
+        "shape: (20, 1)\n"
+        "╭───────╮\n"
+        "│ id    │\n"
+        "│ ---   │\n"
+        "│ int64 │\n"
+        "╞═══════╡\n"
+        "│ 0     │\n"
+        "│ 1     │\n"
+        "│ 2     │\n"
+        "│ 3     │\n"
+        "│ 4     │\n"
+        "│ …     │\n"
+        "│ 15    │\n"
+        "│ 16    │\n"
+        "│ 17    │\n"
+        "│ 18    │\n"
+        "│ 19    │\n"
+        "╰───────╯\n"
+        "(Showing 10 of 20 rows)"
+    )
+
+
+def test_dataset_explain(ray_start_regular_shared, capsys):
+    ds = ray.data.range(10, override_num_blocks=10)
+    ds = ds.map(lambda x: x)
+
+    ds.explain()
+    captured = capsys.readouterr()
+    assert captured.out.strip() == (
+        "-------- Logical Plan --------\n"
+        "MapRows[Map(<lambda>)]\n"
+        "+- Read[ReadRange]\n"
+        "\n-------- Logical Plan (Optimized) --------\n"
+        "MapRows[Map(<lambda>)]\n"
+        "+- Read[ReadRange]\n"
+        "\n-------- Physical Plan --------\n"
+        "TaskPoolMapOperator[Map(<lambda>)]\n"
+        "+- TaskPoolMapOperator[ReadRange]\n"
+        "   +- InputDataBuffer[Input]\n"
+        "\n-------- Physical Plan (Optimized) --------\n"
+        "TaskPoolMapOperator[ReadRange->Map(<lambda>)]\n"
+        "+- InputDataBuffer[Input]"
+    )
+
+    ds = ds.filter(lambda x: x["id"] > 0)
+    ds.explain()
+    captured = capsys.readouterr()
+    assert captured.out.strip() == (
+        "-------- Logical Plan --------\n"
+        "Filter[Filter(<lambda>)]\n"
+        "+- MapRows[Map(<lambda>)]\n"
+        "   +- Read[ReadRange]\n"
+        "\n-------- Logical Plan (Optimized) --------\n"
+        "Filter[Filter(<lambda>)]\n"
+        "+- MapRows[Map(<lambda>)]\n"
+        "   +- Read[ReadRange]\n"
+        "\n-------- Physical Plan --------\n"
+        "TaskPoolMapOperator[Filter(<lambda>)]\n"
+        "+- TaskPoolMapOperator[Map(<lambda>)]\n"
+        "   +- TaskPoolMapOperator[ReadRange]\n"
+        "      +- InputDataBuffer[Input]\n"
+        "\n-------- Physical Plan (Optimized) --------\n"
+        "TaskPoolMapOperator[ReadRange->Map(<lambda>)->Filter(<lambda>)]\n"
+        "+- InputDataBuffer[Input]"
+    )
+    ds = ds.random_shuffle().map(lambda x: x)
+    ds.explain()
+    captured = capsys.readouterr()
+    assert captured.out.strip() == (
+        "-------- Logical Plan --------\n"
+        "MapRows[Map(<lambda>)]\n"
+        "+- RandomShuffle[RandomShuffle]\n"
+        "   +- Filter[Filter(<lambda>)]\n"
+        "      +- MapRows[Map(<lambda>)]\n"
+        "         +- Read[ReadRange]\n"
+        "\n-------- Logical Plan (Optimized) --------\n"
+        "MapRows[Map(<lambda>)]\n"
+        "+- RandomShuffle[RandomShuffle]\n"
+        "   +- Filter[Filter(<lambda>)]\n"
+        "      +- MapRows[Map(<lambda>)]\n"
+        "         +- Read[ReadRange]\n"
+        "\n-------- Physical Plan --------\n"
+        "TaskPoolMapOperator[Map(<lambda>)]\n"
+        "+- AllToAllOperator[RandomShuffle]\n"
+        "   +- TaskPoolMapOperator[Filter(<lambda>)]\n"
+        "      +- TaskPoolMapOperator[Map(<lambda>)]\n"
+        "         +- TaskPoolMapOperator[ReadRange]\n"
+        "            +- InputDataBuffer[Input]\n"
+        "\n-------- Physical Plan (Optimized) --------\n"
+        "TaskPoolMapOperator[Map(<lambda>)]\n"
+        "+- AllToAllOperator[ReadRange->Map(<lambda>)->Filter(<lambda>)->RandomShuffle]\n"
+        "   +- InputDataBuffer[Input]"
+    )
+
+
+def test_convert_types(ray_start_regular_shared):
+    plain_ds = ray.data.range(1)
+    arrow_ds = plain_ds.map(lambda x: {"a": x["id"]})
+    assert arrow_ds.take() == [{"a": 0}]
+    assert "dict" in str(arrow_ds.map(lambda x: {"out": str(type(x))}).take()[0])
+
+    arrow_ds = ray.data.range(1)
+    assert arrow_ds.map(lambda x: {"out": "plain_{}".format(x["id"])}).take() == [
+        {"out": "plain_0"}
+    ]
+    assert arrow_ds.map(lambda x: {"a": (x["id"],)}).take() == [{"a": [0]}]
+
+
+def test_take_batch(ray_start_regular_shared):
+    ds = ray.data.range(10, override_num_blocks=2)
+    assert ds.take_batch(3)["id"].tolist() == [0, 1, 2]
+    assert ds.take_batch(6)["id"].tolist() == [0, 1, 2, 3, 4, 5]
+    assert ds.take_batch(100)["id"].tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+    assert isinstance(ds.take_batch(3, batch_format="pandas"), pd.DataFrame)
+    assert isinstance(ds.take_batch(3, batch_format="numpy"), dict)
+
+    ds = ray.data.range_tensor(10, override_num_blocks=2)
+    assert np.all(ds.take_batch(3)["data"] == np.array([[0], [1], [2]]))
+    assert isinstance(ds.take_batch(3, batch_format="pandas"), pd.DataFrame)
+    assert isinstance(ds.take_batch(3, batch_format="numpy"), dict)
+
+    with pytest.raises(ValueError):
+        ray.data.range(0).take_batch()
+
+
+def test_take_all(ray_start_regular_shared):
+    assert extract_values("id", ray.data.range(5).take_all()) == [0, 1, 2, 3, 4]
+
+    with pytest.raises(ValueError):
+        assert ray.data.range(5).take_all(4)
+
+
+def test_union(ray_start_regular_shared, restore_data_context):
+    # Set aggregator CPU to 0 to avoid deadlock in resource-constrained test env.
+    # Without this, the shuffle task (1 CPU) + aggregator actor (~0.25 CPU) would
+    # exceed the 1 CPU available in the test cluster, causing scheduler deadlock.
+    restore_data_context.hash_aggregate_operator_actor_num_cpus_override = 0
+
+    ds = ray.data.range(20, override_num_blocks=10).materialize()
+
+    # Test lazy union.
+    ds = ds.union(ds, ds, ds, ds)
+    assert ds._logical_plan.initial_num_blocks() == 50
+    assert ds.count() == 100
+    assert ds.sum() == 950
+
+    ds = ds.union(ds)
+    assert ds.count() == 200
+    assert ds.sum() == (950 * 2)
+
+    # Test materialized union.
+    ds2 = ray.data.from_items([1, 2, 3, 4, 5])
+    assert ds2.count() == 5
+    ds2 = ds2.union(ds2)
+    assert ds2.count() == 10
+    ds2 = ds2.union(ds)
+    assert ds2.count() == 210
+
+
+def test_block_builder_for_block(ray_start_regular_shared):
+    # pandas dataframe
+    builder = BlockBuilder.for_block(pd.DataFrame())
+    b1 = pd.DataFrame({"A": [1], "B": ["a"]})
+    builder.add_block(b1)
+    assert builder.build().equals(b1)
+    b2 = pd.DataFrame({"A": [2, 3], "B": ["c", "d"]})
+    builder.add_block(b2)
+    expected = pd.DataFrame({"A": [1, 2, 3], "B": ["a", "c", "d"]})
+    assert builder.build().equals(expected)
+
+    # pyarrow table
+    builder = BlockBuilder.for_block(pa.Table.from_arrays(list()))
+    b1 = pa.Table.from_pydict({"A": [1], "B": ["a"]})
+    builder.add_block(b1)
+    builder.build().equals(b1)
+    b2 = pa.Table.from_pydict({"A": [2, 3], "B": ["c", "d"]})
+    builder.add_block(b2)
+    expected = pa.Table.from_pydict({"A": [1, 2, 3], "B": ["a", "c", "d"]})
+    builder.build().equals(expected)
+
+    # wrong type
+    with pytest.raises(TypeError):
+        BlockBuilder.for_block(str())
+
+
+def test_len(ray_start_regular_shared):
+    ds = ray.data.range(1)
+    with pytest.raises(AttributeError):
+        len(ds)
+
+
+def test_pandas_block_select():
+    df = pd.DataFrame({"one": [10, 11, 12], "two": [11, 12, 13], "three": [14, 15, 16]})
+    block_accessor = BlockAccessor.for_block(df)
+
+    block = block_accessor.select(["two"])
+    assert block.equals(df[["two"]])
+
+    block = block_accessor.select(["two", "one"])
+    assert block.equals(df[["two", "one"]])
+
+    with pytest.raises(ValueError):
+        block = block_accessor.select([lambda x: x % 3, "two"])
+
+
+# NOTE: All tests above share a Ray cluster, while the tests below do not. These
+# tests should only be carefully reordered to retain this invariant!
+
+
+def test_read_write_local_node_ray_client(ray_start_cluster_enabled):
+    cluster = ray_start_cluster_enabled
+    cluster.add_node(num_cpus=4)
+    cluster.head_node._ray_params.ray_client_server_port = "10004"
+    cluster.head_node.start_ray_client_server()
+    address = "ray://localhost:10004"
+
+    import tempfile
+
+    data_path = tempfile.mkdtemp()
+    df = pd.DataFrame({"one": list(range(0, 10)), "two": list(range(10, 20))})
+    path = os.path.join(data_path, "test.parquet")
+    df.to_parquet(path)
+
+    # Read/write from Ray Client will result in error.
+    ray.init(address)
+    with pytest.raises(ValueError):
+        ds = ray.data.read_parquet("local://" + path).materialize()
+    ds = ray.data.from_pandas(df)
+    with pytest.raises(ValueError):
+        ds.write_parquet("local://" + data_path).materialize()
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="No tensorflow for Python 3.12+"
+)
+def test_read_warning_large_parallelism(
+    ray_start_regular_shared, propagate_logs, caplog
+):
+    with caplog.at_level(logging.WARNING, logger="ray.data.read_api"):
+        ray.data.range(5000, override_num_blocks=5000).materialize()
+    assert (
+        "The requested number of read blocks of 5000 is "
+        "more than 4x the number of available CPU slots in the cluster" in caplog.text
+    ), caplog.text
+
+
+def test_read_write_local_node(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(
+        resources={"bar:1": 100},
+        num_cpus=10,
+        _system_config={"max_direct_call_object_size": 0},
+    )
+    cluster.add_node(resources={"bar:2": 100}, num_cpus=10)
+    cluster.add_node(resources={"bar:3": 100}, num_cpus=10)
+
+    ray.shutdown()
+    ray.init(cluster.address)
+
+    import os
+    import tempfile
+
+    data_path = tempfile.mkdtemp()
+    num_files = 5
+    for idx in range(num_files):
+        df = pd.DataFrame(
+            {"one": list(range(idx, idx + 10)), "two": list(range(idx + 10, idx + 20))}
+        )
+        path = os.path.join(data_path, f"test{idx}.parquet")
+        df.to_parquet(path)
+
+    ctx = ray.data.context.DataContext.get_current()
+    ctx.read_write_local_node = True
+
+    def check_dataset_is_local(ds):
+        bundles = ds.iter_internal_ref_bundles()
+        block_refs = _ref_bundles_iterator_to_block_refs_list(bundles)
+        ray.wait(block_refs, num_returns=len(block_refs), fetch_local=False)
+        location_data = ray.experimental.get_object_locations(block_refs)
+        locations = []
+        for block in block_refs:
+            locations.extend(location_data[block]["node_ids"])
+        assert set(locations) == {ray.get_runtime_context().get_node_id()}
+
+    local_path = "local://" + data_path
+    # Plain read.
+    ds = ray.data.read_parquet(local_path).materialize()
+    check_dataset_is_local(ds)
+
+    # SPREAD scheduling got overridden when read local scheme.
+    ds = ray.data.read_parquet(
+        local_path, ray_remote_args={"scheduling_strategy": "SPREAD"}
+    ).materialize()
+    check_dataset_is_local(ds)
+
+    # With fusion.
+    ds = (
+        ray.data.read_parquet(local_path, override_num_blocks=1)
+        .map(lambda x: x)
+        .materialize()
+    )
+    check_dataset_is_local(ds)
+
+    # Write back to local scheme.
+    output = os.path.join(local_path, "test_read_write_local_node")
+    ds.write_parquet(output)
+    assert "1 nodes used" in ds.stats(), ds.stats()
+    ray.data.read_parquet(output).take_all() == ds.take_all()
+
+    # Mixing paths of local and non-local scheme is invalid.
+    with pytest.raises(ValueError):
+        ds = ray.data.read_parquet(
+            [local_path + "/test1.parquet", data_path + "/test2.parquet"]
+        ).materialize()
+    with pytest.raises(ValueError):
+        ds = ray.data.read_parquet(
+            [local_path + "/test1.parquet", "example://iris.parquet"]
+        ).materialize()
+    with pytest.raises(ValueError):
+        ds = ray.data.read_parquet(
+            ["example://iris.parquet", local_path + "/test1.parquet"]
+        ).materialize()
+
+
+def test_validate_head_node_resources_zero_head_cpu(ray_start_cluster):
+
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)
+    cluster.wait_for_nodes()
+
+    ray.shutdown()
+    ray.init(address=cluster.address)
+
+    with pytest.raises(ValueError, match=r"head node doesn't have enough resources"):
+        _validate_head_node_resources_for_local_scheduling(
+            {}, op_description="read local"
+        )
+
+
+class FlakyCSVDatasource(CSVDatasource):
+    def __init__(self, paths, **csv_datasource_kwargs):
+        super().__init__(paths, **csv_datasource_kwargs)
+
+        self.counter = Counter.remote()
+
+    def _read_stream(self, f: "pa.NativeFile", path: str):
+        count = self.counter.increment.remote()
+        if ray.get(count) == 1:
+            raise ValueError("oops")
+        else:
+            for block in CSVDatasource._read_stream(self, f, path):
+                yield block
+
+
+class FlakyCSVDatasink(CSVDatasink):
+    def __init__(self, path, **csv_datasink_kwargs):
+        super().__init__(path, **csv_datasink_kwargs)
+
+        self.counter = Counter.remote()
+
+    def write_block_to_file(self, block: BlockAccessor, file):
+        count = self.counter.increment.remote()
+        if ray.get(count) == 1:
+            raise ValueError("oops")
+        else:
+            super().write_block_to_file(block, file)
+
+
+def test_datasource(ray_start_regular):
+    source = ray.data.datasource.RandomIntRowDatasource(n=10, num_columns=2)
+    assert len(ray.data.read_datasource(source).take()) == 10
+    source = RangeDatasource(n=10)
+    assert extract_values(
+        "value",
+        ray.data.read_datasource(source).take(),
+    ) == list(range(10))
+
+
+@pytest.mark.skip(reason="")
+def test_polars_lazy_import(shutdown_only):
+    import sys
+
+    ctx = ray.data.context.DataContext.get_current()
+
+    try:
+        original_use_polars = ctx.use_polars
+        ctx.use_polars = True
+
+        num_items = 100
+        parallelism = 4
+        ray.init(num_cpus=4)
+
+        @ray.remote
+        def f(should_import_polars):
+            # Sleep to spread the tasks.
+            time.sleep(1)
+            polars_imported = "polars" in sys.modules.keys()
+            return polars_imported == should_import_polars
+
+        # We should not use polars for non-Arrow sort.
+        _ = ray.data.range(num_items, override_num_blocks=parallelism).sort()
+        assert all(ray.get([f.remote(False) for _ in range(parallelism)]))
+
+        a = range(100)
+        dfs = []
+        partition_size = num_items // parallelism
+        for i in range(parallelism):
+            dfs.append(
+                pd.DataFrame({"a": a[i * partition_size : (i + 1) * partition_size]})
+            )
+        # At least one worker should have imported polars.
+        _ = (
+            ray.data.from_pandas(dfs)
+            .map_batches(lambda t: t, batch_format="pyarrow", batch_size=None)
+            .sort(key="a")
+            .materialize()
+        )
+        assert any(ray.get([f.remote(True) for _ in range(parallelism)]))
+
+    finally:
+        ctx.use_polars = original_use_polars
+
+
+def test_batch_formats(shutdown_only):
+    ds = ray.data.range(100)
+    assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
+
+    ds = ray.data.range_tensor(100)
+    assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
+
+    df = pd.DataFrame({"foo": ["a", "b"], "bar": [0, 1]})
+    ds = ray.data.from_pandas(df)
+    assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pd.DataFrame)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
+    assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
+
+
+def test_dataset_schema_after_read_stats(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+    ray.init(cluster.address)
+    cluster.add_node(num_cpus=1, resources={"foo": 1})
+    ds = ray.data.read_csv(
+        "example://iris.csv", ray_remote_args={"resources": {"foo": 1}}
+    )
+    schema = ds.schema()
+    ds.stats()
+    assert schema == ds.schema()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))