chore: import upstream snapshot with attribution

python/ray/data/tests/test_block_sizing.py

@@ -0,0 +1,259 @@
+import pytest
+
+import ray
+from ray.data.context import DataContext
+from ray.data.dataset import Dataset
+from ray.data.tests.conftest import *  # noqa
+from ray.data.tests.conftest import (
+    assert_blocks_expected_in_plasma,
+    get_initial_core_execution_metrics_snapshot,
+)
+from ray.tests.conftest import *  # noqa
+
+
+def _assert_num_blocks(ds, expected, tolerance=0.5):
+    actual = ds.num_blocks()
+    assert (
+        expected * (1 - tolerance) <= actual <= expected * (1 + tolerance)
+    ), f"Expected ~{expected} blocks (±{tolerance*100}%), got {actual}"
+
+
+def test_map(shutdown_only, restore_data_context):
+    ray.init(
+        _system_config={
+            "max_direct_call_object_size": 10_000,
+        },
+        num_cpus=2,
+        object_store_memory=int(100e6),
+    )
+
+    ctx = DataContext.get_current()
+    ctx.target_min_block_size = 10_000 * 8
+    ctx.target_max_block_size = 10_000 * 8
+    num_blocks_expected = 10
+
+    # Test read.
+    ds = ray.data.range(100_000, override_num_blocks=1).materialize()
+    _assert_num_blocks(ds, num_blocks_expected)
+
+    # Test read -> map.
+    # NOTE(swang): For some reason BlockBuilder's estimated memory usage when a
+    # map fn is used is 2x the actual memory usage.
+    ds = (
+        ray.data.range(100_000, override_num_blocks=1)
+        .map(lambda row: row)
+        .materialize()
+    )
+    _assert_num_blocks(ds, num_blocks_expected * 2)
+
+    # Test adjusted block size.
+    ctx.target_max_block_size *= 2
+    num_blocks_expected //= 2
+
+    # Test read.
+    ds = ray.data.range(100_000, override_num_blocks=1).materialize()
+    _assert_num_blocks(ds, num_blocks_expected)
+
+    # Test read -> map.
+    ds = (
+        ray.data.range(100_000, override_num_blocks=1)
+        .map(lambda row: row)
+        .materialize()
+    )
+    _assert_num_blocks(ds, num_blocks_expected * 2)
+
+    # Setting the shuffle block size prints a warning and actually resets
+    # target_max_block_size
+    ctx.target_shuffle_max_block_size = ctx.target_max_block_size / 2
+    num_blocks_expected *= 2
+
+    # Test read.
+    ds = ray.data.range(100_000, override_num_blocks=1).materialize()
+    _assert_num_blocks(ds, num_blocks_expected)
+
+    # Test read -> map.
+    ds = (
+        ray.data.range(100_000, override_num_blocks=1)
+        .map(lambda row: row)
+        .materialize()
+    )
+    _assert_num_blocks(ds, num_blocks_expected * 2)
+
+
+# TODO: Test that map stage output blocks are the correct size for groupby and
+# repartition. Currently we only have access to the reduce stage output block
+# size.
+SHUFFLE_ALL_TO_ALL_OPS = [
+    (Dataset.random_shuffle, {}, True),
+    (Dataset.sort, {"key": "id"}, False),
+]
+
+
+@pytest.mark.parametrize(
+    "shuffle_op",
+    SHUFFLE_ALL_TO_ALL_OPS,
+)
+def test_shuffle(shutdown_only, restore_data_context, shuffle_op):
+    ray.init(
+        _system_config={
+            "max_direct_call_object_size": 250,
+        },
+        num_cpus=2,
+        object_store_memory=int(100e6),
+    )
+
+    # Test AllToAll and Map -> AllToAll Datasets. Check that Map inherits
+    # AllToAll's target block size.
+    ctx = DataContext.get_current()
+    ctx.read_op_min_num_blocks = 1
+    ctx.target_min_block_size = 1
+
+    N = 100_000
+    mem_size = 800_000
+
+    shuffle_fn, kwargs, fusion_supported = shuffle_op
+
+    ctx.target_max_block_size = 10_000 * 8
+    num_blocks_expected = mem_size // ctx.target_max_block_size
+    last_snapshot = get_initial_core_execution_metrics_snapshot()
+
+    ds = shuffle_fn(ray.data.range(N), **kwargs).materialize()
+    assert (
+        num_blocks_expected
+        <= ds._logical_plan.initial_num_blocks()
+        <= num_blocks_expected * 1.5
+    )
+
+    def _estimate_intermediate_blocks(fusion_supported: bool, num_blocks_expected: int):
+        return num_blocks_expected**2 + num_blocks_expected * (
+            2 if fusion_supported else 4
+        )
+
+    # map * reduce intermediate blocks + 1 metadata ref per map/reduce task.
+    # If fusion is not supported, the un-fused map stage produces 1 data and 1
+    # metadata per task.
+    num_intermediate_blocks = _estimate_intermediate_blocks(
+        fusion_supported, num_blocks_expected
+    )
+
+    print(f">>> Asserting {num_intermediate_blocks} blocks are in plasma")
+
+    last_snapshot = assert_blocks_expected_in_plasma(
+        last_snapshot,
+        # Dataset.sort produces some empty intermediate blocks because the
+        # input range is already partially sorted.
+        num_intermediate_blocks,
+    )
+
+    ds = shuffle_fn(ray.data.range(N).map(lambda x: x), **kwargs).materialize()
+    if not fusion_supported:
+        # TODO(swang): For some reason BlockBuilder's estimated
+        # memory usage for range(1000)->map is 2x the actual memory usage.
+        # Remove once https://github.com/ray-project/ray/issues/40246 is fixed.
+        num_blocks_expected = int(num_blocks_expected * 2.2)
+    assert (
+        num_blocks_expected
+        <= ds._logical_plan.initial_num_blocks()
+        <= num_blocks_expected * 1.5
+    )
+    num_intermediate_blocks = _estimate_intermediate_blocks(
+        fusion_supported, num_blocks_expected
+    )
+    last_snapshot = assert_blocks_expected_in_plasma(
+        last_snapshot,
+        # Dataset.sort produces some empty intermediate blocks because the
+        # input range is already partially sorted.
+        num_intermediate_blocks,
+    )
+
+    ctx.target_max_block_size //= 2
+    num_blocks_expected = mem_size // ctx.target_max_block_size
+    block_size_expected = ctx.target_max_block_size
+
+    ds = shuffle_fn(ray.data.range(N), **kwargs).materialize()
+    assert (
+        num_blocks_expected
+        <= ds._logical_plan.initial_num_blocks()
+        <= num_blocks_expected * 1.5
+    )
+    num_intermediate_blocks = _estimate_intermediate_blocks(
+        fusion_supported, num_blocks_expected
+    )
+    last_snapshot = assert_blocks_expected_in_plasma(
+        last_snapshot,
+        num_intermediate_blocks,
+    )
+
+    ds = shuffle_fn(ray.data.range(N).map(lambda x: x), **kwargs).materialize()
+    if not fusion_supported:
+        num_blocks_expected = int(num_blocks_expected * 2.2)
+        block_size_expected //= 2.2
+    assert (
+        num_blocks_expected
+        <= ds._logical_plan.initial_num_blocks()
+        <= num_blocks_expected * 1.5
+    )
+    num_intermediate_blocks = _estimate_intermediate_blocks(
+        fusion_supported, num_blocks_expected
+    )
+    last_snapshot = assert_blocks_expected_in_plasma(
+        last_snapshot,
+        num_intermediate_blocks,
+    )
+
+    # Setting target max block size does not affect map ops when there is a
+    # shuffle downstream.
+    ctx.target_max_block_size = ctx.target_max_block_size * 2
+    num_blocks_expected //= 2
+
+    ds = shuffle_fn(ray.data.range(N).map(lambda x: x), **kwargs).materialize()
+    assert (
+        num_blocks_expected
+        <= ds._logical_plan.initial_num_blocks()
+        <= num_blocks_expected * 1.5
+    )
+
+    num_intermediate_blocks = _estimate_intermediate_blocks(
+        fusion_supported, num_blocks_expected
+    )
+
+    assert_blocks_expected_in_plasma(
+        last_snapshot,
+        num_intermediate_blocks,
+    )
+
+
+def test_target_max_block_size_infinite_or_default_disables_splitting_globally(
+    shutdown_only, restore_data_context
+):
+    """Test that setting target_max_block_size to None disables block splitting globally."""
+    ray.init(num_cpus=2)
+
+    # Create a large dataset that would normally trigger block splitting
+    N = 1_000_000  # ~8MB worth of data
+
+    # First, test with normal target_max_block_size (should split into multiple blocks)
+    ctx = DataContext.get_current()
+    ctx.target_max_block_size = 1_000_000  # ~1MB
+
+    ds_with_limit = ray.data.range(N, override_num_blocks=1).materialize()
+    blocks_with_limit = ds_with_limit._logical_plan.initial_num_blocks()
+
+    # Now test with target_max_block_size = None (should not split)
+    ctx.target_max_block_size = None  # Disable block size limit
+
+    ds_unlimited = (
+        ray.data.range(N, override_num_blocks=1).map(lambda x: x).materialize()
+    )
+    blocks_unlimited = ds_unlimited._logical_plan.initial_num_blocks()
+
+    # Verify that unlimited creates fewer blocks (no splitting)
+    assert blocks_unlimited <= blocks_with_limit
+    # With target_max_block_size=None, it should maintain the original block structure
+    assert blocks_unlimited == 1
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))