chore: import upstream snapshot with attribution

python/ray/data/tests/test_splitblocks.py

@@ -0,0 +1,192 @@
+import numpy as np
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.execution.operators.map_operator import (
+    _split_blocks,
+    _splitrange,
+)
+from ray.data.block import BlockAccessor
+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.tests.conftest import *  # noqa
+
+
+def test_splitrange():
+    def f(n, k):
+        assert _splitrange(n, k) == [len(a) for a in np.array_split(range(n), k)]
+
+    f(0, 1)
+    f(5, 1)
+    f(5, 3)
+    f(5, 5)
+    f(5, 10)
+    f(50, 1)
+    f(50, 2)
+    f(50, 3)
+    f(50, 4)
+    f(50, 5)
+
+
+def test_split_blocks():
+    def f(n, k):
+        table = pa.Table.from_arrays([np.arange(n)], names=["value"])
+        in_blocks = [table]
+        out_blocks = list(_split_blocks(in_blocks, k))
+        sizes = [BlockAccessor.for_block(b).num_rows() for b in out_blocks]
+        expected = [len(a) for a in np.array_split(range(n), min(k, n))]
+        assert sizes == expected
+
+    f(5, 1)
+    f(5, 3)
+    f(5, 5)
+    f(5, 10)
+    f(50, 1)
+    f(50, 2)
+    f(50, 3)
+    f(50, 4)
+    f(50, 5)
+
+
+def test_small_file_split(ray_start_10_cpus_shared, restore_data_context):
+    last_snapshot = get_initial_core_execution_metrics_snapshot()
+
+    ds = ray.data.read_csv("example://iris.csv", override_num_blocks=1)
+    materialized_ds = ds.materialize()
+    assert materialized_ds._logical_plan.initial_num_blocks() == 1
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                "ReadCSV": 1,
+            },
+        ),
+        last_snapshot,
+    )
+
+    materialized_ds = ds.map_batches(lambda x: x).materialize()
+    assert materialized_ds._logical_plan.initial_num_blocks() == 1
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                "ReadCSV->MapBatches(<lambda>)": 1,
+            },
+        ),
+        last_snapshot,
+    )
+
+    stats = materialized_ds.stats()
+    assert "Operator 1 ReadCSV->MapBatches" in stats, stats
+
+    ds = ray.data.read_csv("example://iris.csv", override_num_blocks=10)
+    assert ds._logical_plan.initial_num_blocks() == 1
+    assert (
+        ds.map_batches(lambda x: x).materialize()._logical_plan.initial_num_blocks()
+        == 10
+    )
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                "MapBatches(<lambda>)": 10,
+                "ReadCSV->SplitBlocks(10)": 1,
+            },
+        ),
+        last_snapshot,
+    )
+
+    assert ds.materialize()._logical_plan.initial_num_blocks() == 10
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                "ReadCSV->SplitBlocks(10)": 1,
+            },
+        ),
+        last_snapshot,
+    )
+
+    ds = ray.data.read_csv("example://iris.csv", override_num_blocks=100)
+    assert ds._logical_plan.initial_num_blocks() == 1
+    assert (
+        ds.map_batches(lambda x: x).materialize()._logical_plan.initial_num_blocks()
+        == 100
+    )
+    assert ds.materialize()._logical_plan.initial_num_blocks() == 100
+
+    ds = ds.map_batches(lambda x: x).materialize()
+    stats = ds.stats()
+    assert "Operator 1 ReadCSV->SplitBlocks(100)" in stats, stats
+    assert "Operator 2 MapBatches" in stats, stats
+
+    # Smaller than a single row.
+    ds.context.target_max_block_size = 1
+    ds = ds.map_batches(lambda x: x).materialize()
+    # 150 rows.
+    assert ds._logical_plan.initial_num_blocks() == 150
+    print(ds.stats())
+
+
+def test_large_file_additional_split(ray_start_10_cpus_shared, tmp_path):
+    ctx = ray.data.context.DataContext.get_current()
+    if ctx.use_datasource_v2:
+        pytest.skip(
+            "V2 defers file listing to execution time, so "
+            "``LogicalPlan.initial_num_blocks()`` can't report a "
+            "file-count-based estimate pre-materialization. The "
+            "post-materialize block-split assertions this test also "
+            "makes are still covered by V1."
+        )
+    ctx.target_max_block_size = 10 * 1024 * 1024
+
+    # ~100MiB of tensor data
+    ds = ray.data.range_tensor(1000, shape=(10000,))
+    ds.repartition(1).write_parquet(tmp_path)
+
+    ds = ray.data.read_parquet(tmp_path, override_num_blocks=1)
+    assert ds._logical_plan.initial_num_blocks() == 1
+    print(ds.materialize().stats())
+    assert (
+        5 < ds.materialize()._logical_plan.initial_num_blocks() < 20
+    )  # Size-based block split
+
+    ds = ray.data.read_parquet(tmp_path, override_num_blocks=10)
+    assert ds._logical_plan.initial_num_blocks() == 1
+    assert 5 < ds.materialize()._logical_plan.initial_num_blocks() < 20
+
+    ds = ray.data.read_parquet(tmp_path, override_num_blocks=100)
+    assert ds._logical_plan.initial_num_blocks() == 1
+    assert 50 < ds.materialize()._logical_plan.initial_num_blocks() < 200
+
+    ds = ray.data.read_parquet(tmp_path, override_num_blocks=1000)
+    assert ds._logical_plan.initial_num_blocks() == 1
+    assert 500 < ds.materialize()._logical_plan.initial_num_blocks() < 2000
+
+
+def test_map_batches_split(ray_start_10_cpus_shared, restore_data_context):
+    ds = ray.data.range(1000, override_num_blocks=1).map_batches(
+        lambda x: x, batch_size=1000
+    )
+    assert ds.materialize()._logical_plan.initial_num_blocks() == 1
+
+    ctx = ray.data.context.DataContext.get_current()
+    # 100 integer rows per block.
+    ctx.target_max_block_size = 800
+
+    ds = ray.data.range(1000, override_num_blocks=1).map_batches(
+        lambda x: x, batch_size=1000
+    )
+    assert ds.materialize()._logical_plan.initial_num_blocks() == 10
+
+    # A single row is already larger than the target block
+    # size.
+    ds.context.target_max_block_size = 4
+    assert ds.materialize()._logical_plan.initial_num_blocks() == 1000
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))