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

Code before:
        ds = ds.repartition(target_num_rows_per_block=target_num_rows)
        expected_fused_name = f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_num_rows}]"
    else:  # sr_then_map
        ds = ds.repartition(target_num_rows_per_block=target_num_rows)
        ds = ds.map_batches(lambda x: x, batch_size=batch_size)
        expected_fused_name = f"StreamingRepartition[num_rows_per_block={target_num_rows}]->MapBatches(<lambda>)"

    assert len(ds.take_all()) == n

    stats = ds.stats()
    if should_fuse:
        assert (
            expected_fused_name in stats
        ), f"Expected '{expected_fused_name}' in stats: {stats}"
    else:
        assert (
            expected_fused_name not in stats
        ), f"Did not expect '{expected_fused_name}' in stats: {stats}"


def test_streaming_repartition_no_further_fuse(
    ray_start_regular_shared_2_cpus,
):
    """Test that fused streaming_repartition operators don't fuse further.

    Case 1: map_batches -> map_batches -> streaming_repartition -> map_batches -> map_batches
            Result: map -> (map -> s_r)-> (map -> map)
            The fused (map -> s_r) doesn't fuse further with surrounding maps.
    """
    n = 100
    target_rows = 20

    # Case 1: map_batches -> map_batches -> streaming_repartition -> map_batches -> map_batches
    # Result: map -> (map -> s_r)-> (map -> map)
    ds1 = ray.data.range(n, override_num_blocks=2)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)
    ds1 = ds1.repartition(target_num_rows_per_block=target_rows)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)

    assert len(ds1.take_all()) == n
    stats1 = ds1.stats()

    assert (
        f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_rows}]"
        in stats1
    ), stats1
    assert "MapBatches(<lambda>)->MapBatches(<lambda>)" in stats1



Review comment (suggestion):
This shouldn't fuse irrespective of whether it's strict or not (otherwise we might decrease parallelism regardless of whether it's strict or not) 

Code after:
        ds = ds.repartition(target_num_rows_per_block=target_num_rows, strict=True)
        expected_fused_name = f"MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_num_rows},strict=True]"
    else:  # sr_then_map
        ds = ds.repartition(target_num_rows_per_block=target_num_rows, strict=True)
        ds = ds.map_batches(lambda x: x, batch_size=batch_size)
        expected_fused_name = f"StreamingRepartition[num_rows_per_block={target_num_rows},strict=True]->MapBatches(<lambda>)"

    assert len(ds.take_all()) == n

    stats = ds.stats()
    if should_fuse:
        assert (
            expected_fused_name in stats
        ), f"Expected '{expected_fused_name}' in stats: {stats}"
    else:
        assert (
            expected_fused_name not in stats
        ), f"Did not expect '{expected_fused_name}' in stats: {stats}"


def test_streaming_repartition_no_further_fuse(
    ray_start_regular_shared_2_cpus,
):
    """Test that streaming_repartition (strict mode) blocks fusion with downstream operators.

    Case 1: map_batches -> map_batches -> streaming_repartition(strict=True) -> map_batches -> map_batches
            Result: (map -> map -> s_r) -> (map -> map)
            SR can fuse with upstream maps but not with downstream maps to preserve parallelism.
    """
    n = 100
    target_rows = 20

    # Case 1: map_batches -> map_batches -> streaming_repartition(strict=True) -> map_batches -> map_batches
    # Result: (map -> map -> s_r) -> (map -> map)
    ds1 = ray.data.range(n, override_num_blocks=2)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)
    ds1 = ds1.repartition(target_num_rows_per_block=target_rows, strict=True)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)
    ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows)

    assert len(ds1.take_all()) == n
    stats1 = ds1.stats()

    assert (
        f"MapBatches(<lambda>)->MapBatches(<lambda>)->StreamingRepartition[num_rows_per_block={target_rows},strict=True]"
        in stats1
    ), stats1
    assert "MapBatches(<lambda>)->MapBatches(<lambda>)" in stats1