from unittest.mock import MagicMock import numpy as np import pandas as pd import pytest import ray from ray.data._internal.execution.bundle_queue import EstimateSize 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.execution.operators.map_transformer import ( BatchMapTransformFn, BlockMapTransformFn, ) from ray.data._internal.logical.interfaces import LogicalPlan from ray.data._internal.logical.operators import ( Filter, FlatMap, InputData, MapBatches, MapRows, Project, Read, Write, ) from ray.data._internal.logical.optimizers import PhysicalOptimizer, get_execution_plan from ray.data._internal.planner import create_planner from ray.data._internal.stats import DatasetStats from ray.data._internal.util import rows_same from ray.data.context import DataContext, ShuffleStrategy from ray.data.dataset import Dataset from ray.data.expressions import star from ray.data.tests.conftest import * # noqa from ray.data.tests.test_util import _check_usage_record, get_parquet_read_logical_op from ray.data.tests.util import column_udf, extract_values from ray.tests.conftest import * # noqa def test_read_map_batches_operator_fusion(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() # Test that Read is fused with MapBatches. planner = create_planner() read_op = get_parquet_read_logical_op(parallelism=1) op = MapBatches( lambda x: x, input_dependencies=[read_op], ) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert op.name == "MapBatches()" assert physical_op.name == "ReadParquet->MapBatches()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 input = physical_op.input_dependencies[0] assert isinstance(input, InputDataBuffer) assert physical_op in input.output_dependencies, input.output_dependencies assert physical_op._logical_operators == [read_op, op] def test_read_map_chain_operator_fusion(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() # Test that a chain of different map operators are fused. planner = create_planner() read_op = get_parquet_read_logical_op(parallelism=1) map1 = MapRows(lambda x: x, input_dependencies=[read_op]) map2 = MapBatches(lambda x: x, input_dependencies=[map1]) map3 = FlatMap(lambda x: x, input_dependencies=[map2]) map4 = Filter(fn=lambda x: x, input_dependencies=[map3]) logical_plan = LogicalPlan(map4, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert map4.name == "Filter()" assert ( physical_op.name == "ReadParquet->Map()->MapBatches()" "->FlatMap()->Filter()" ) assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], InputDataBuffer) assert physical_op._logical_operators == [read_op, map1, map2, map3, map4] def test_read_map_batches_operator_fusion_compatible_remote_args( ray_start_regular_shared_2_cpus, ): ctx = DataContext.get_current() # Test that map operators are stilled fused when remote args are compatible. compatiple_remote_args_pairs = [ # Empty remote args are compatible. ({}, {}), # Test `num_cpus` and `num_gpus`. ({"num_cpus": 2}, {"num_cpus": 2}), ({"num_gpus": 2}, {"num_gpus": 2}), # `num_cpus` defaults to 1, `num_gpus` defaults to 0. # The following 2 should be compatible. ({"num_cpus": 1}, {}), ({}, {"num_gpus": 0}), # Test specifying custom resources. ({"resources": {"custom": 1}}, {"resources": {"custom": 1}}), ({"resources": {"custom": 0}}, {"resources": {}}), # If the downstream op doesn't have `scheduling_strategy`, it will # inherit from the upstream op. ({"scheduling_strategy": "SPREAD"}, {}), ] for up_remote_args, down_remote_args in compatiple_remote_args_pairs: planner = create_planner() read_op = get_parquet_read_logical_op( ray_remote_args={"resources": {"non-existent": 1}}, parallelism=1, ) op = MapBatches( lambda x: x, input_dependencies=[read_op], ray_remote_args=up_remote_args ) op = MapBatches( lambda x: x, input_dependencies=[op], ray_remote_args=down_remote_args ) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) optimized_physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = optimized_physical_plan.dag assert op.name == "MapBatches()", (up_remote_args, down_remote_args) assert physical_op.name == "MapBatches()->MapBatches()", ( up_remote_args, down_remote_args, ) assert isinstance(physical_op, MapOperator), (up_remote_args, down_remote_args) assert len(physical_op.input_dependencies) == 1, ( up_remote_args, down_remote_args, ) assert physical_op.input_dependencies[0].name == "ReadParquet", ( up_remote_args, down_remote_args, ) def test_read_map_batches_operator_fusion_incompatible_remote_args( ray_start_regular_shared_2_cpus, ): ctx = DataContext.get_current() # Test that map operators won't get fused if the remote args are incompatible. incompatible_remote_args_pairs = [ # Use different resources. ({"num_cpus": 2}, {"num_gpus": 2}), # Same resource, but different values. ({"num_cpus": 3}, {"num_cpus": 2}), # Incompatible custom resources. ({"resources": {"custom": 2}}, {"resources": {"custom": 1}}), ({"resources": {"custom1": 1}}, {"resources": {"custom2": 1}}), # Different scheduling strategies. ({"scheduling_strategy": "SPREAD"}, {"scheduling_strategy": "PACK"}), # Label selectors targeting different ray.io/node-id. ( {"label_selector": {ray._raylet.RAY_NODE_ID_KEY: "node_A"}}, {"label_selector": {ray._raylet.RAY_NODE_ID_KEY: "node_B"}}, ), ] for up_remote_args, down_remote_args in incompatible_remote_args_pairs: planner = create_planner() read_op = get_parquet_read_logical_op( ray_remote_args={"resources": {"non-existent": 1}} ) op = MapBatches( lambda x: x, input_dependencies=[read_op], ray_remote_args=up_remote_args ) op = MapBatches( lambda x: x, input_dependencies=[op], ray_remote_args=down_remote_args ) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert op.name == "MapBatches()", (up_remote_args, down_remote_args) assert physical_op.name == "MapBatches()", ( up_remote_args, down_remote_args, ) assert isinstance(physical_op, MapOperator), (up_remote_args, down_remote_args) assert len(physical_op.input_dependencies) == 1, ( up_remote_args, down_remote_args, ) assert physical_op.input_dependencies[0].name == "MapBatches()", ( up_remote_args, down_remote_args, ) def test_read_map_batches_operator_fusion_compute_tasks_to_actors( ray_start_regular_shared_2_cpus, ): ctx = DataContext.get_current() # Test that a task-based map operator is fused into an actor-based map operator when # the former comes before the latter. planner = create_planner() read_op = get_parquet_read_logical_op(parallelism=1) op = MapBatches(lambda x: x, input_dependencies=[read_op]) op = MapBatches( lambda x: x, input_dependencies=[op], compute=ray.data.ActorPoolStrategy() ) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert op.name == "MapBatches()" assert physical_op.name == "ReadParquet->MapBatches()->MapBatches()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], InputDataBuffer) def test_read_map_batches_operator_fusion_compute_read_to_actors( ray_start_regular_shared_2_cpus, ): ctx = DataContext.get_current() # Test that reads fuse into an actor-based map operator. planner = create_planner() read_op = get_parquet_read_logical_op(parallelism=1) op = MapBatches( lambda x: x, input_dependencies=[read_op], compute=ray.data.ActorPoolStrategy() ) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert op.name == "MapBatches()" assert physical_op.name == "ReadParquet->MapBatches()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], InputDataBuffer) def test_read_map_batches_operator_fusion_incompatible_compute( ray_start_regular_shared_2_cpus, ): ctx = DataContext.get_current() # Test that map operators are not fused when compute strategies are incompatible. planner = create_planner() read_op = get_parquet_read_logical_op(parallelism=1) op = MapBatches( lambda x: x, input_dependencies=[read_op], compute=ray.data.ActorPoolStrategy() ) op = MapBatches(lambda x: x, input_dependencies=[op]) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert op.name == "MapBatches()" assert physical_op.name == "MapBatches()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 upstream_physical_op = physical_op.input_dependencies[0] assert isinstance(upstream_physical_op, MapOperator) # Reads should fuse into actor compute. assert upstream_physical_op.name == "ReadParquet->MapBatches()" def test_read_with_map_batches_fused_successfully( ray_start_regular_shared_2_cpus, temp_dir ): """Since MapBatches does NOT specify `batch_size`, successfully fused with ReadParquet""" # Test that fusion of map operators merges their block sizes in the expected way # (taking the max). n = 10 ds = ray.data.range(n) mapped_ds = ds.map_batches(lambda x: x).map_batches(lambda x: x) physical_plan, _ = get_execution_plan(mapped_ds._logical_plan) physical_op = physical_plan.dag assert isinstance(physical_op, MapOperator) actual_plan_str = physical_op.dag_str # All Map ops are fused with Read assert ( "InputDataBuffer[Input] -> " "TaskPoolMapOperator[ReadRange->MapBatches()->MapBatches()]" == actual_plan_str ) # # Target min-rows requirement is not set strategy = physical_op._block_ref_bundler._strategy assert isinstance(strategy, EstimateSize) assert strategy._min_rows_per_bundle is None @pytest.mark.parametrize( "input_op,fused", [ ( # No fusion (could drastically expand dataset) Read( datasource=MagicMock(name="Parquet"), datasource_or_legacy_reader=MagicMock( get_read_tasks=lambda _: [MagicMock()] ), parallelism=1, ), False, ), ( # No fusion (could drastically reduce dataset) Filter(fn=lambda x: False, input_dependencies=[InputData([])]), False, ), ( # No fusion (could drastically expand/reduce dataset) FlatMap(lambda x: x, input_dependencies=[InputData([])]), False, ), ( # Fusion MapBatches(lambda x: x, input_dependencies=[InputData([])]), True, ), ( # Fusion MapRows(lambda x: x, input_dependencies=[InputData([])]), True, ), ( # Fusion Project(exprs=[star()], input_dependencies=[InputData([])]), True, ), ], ) def test_map_batches_batch_size_fusion( ray_start_regular_shared_2_cpus, input_op, fused ): """Since MapBatches specifies `batch_size` there's no fusion with ReadParquet""" context = DataContext.get_current() # Test that fusion of map operators merges their block sizes in the expected way # (taking the max). ds = Dataset( LogicalPlan(input_op, context), context, DatasetStats(metadata={}, parent=None), ) mapped_ds = ds.map_batches(lambda x: x, batch_size=2).map_batches( lambda x: x, batch_size=5 ) physical_plan, _ = get_execution_plan(mapped_ds._logical_plan) physical_op = physical_plan.dag assert isinstance(physical_op, MapOperator) actual_plan_str = physical_op.dag_str if fused: assert ( f"InputDataBuffer[Input] -> TaskPoolMapOperator[{input_op.name}->" f"MapBatches()->MapBatches()]" == actual_plan_str ) else: assert ( f"InputDataBuffer[Input] -> TaskPoolMapOperator[{input_op.name}] -> " "TaskPoolMapOperator[MapBatches()->MapBatches()]" == actual_plan_str ) # Target min-rows requirement is set to max of upstream and downstream strategy = physical_op._block_ref_bundler._strategy assert isinstance(strategy, EstimateSize) assert strategy._min_rows_per_bundle == 5 assert len(physical_op.input_dependencies) == 1 @pytest.mark.parametrize("upstream_batch_size", [None, 1, 2]) @pytest.mark.parametrize("downstream_batch_size", [None, 1, 2]) def test_map_batches_with_batch_size_specified_fusion( ray_start_regular_shared_2_cpus, temp_dir, upstream_batch_size, downstream_batch_size, ): # Test that fusion of map operators merges their block sizes in the expected way # (taking the max). n = 10 ds = ray.data.range(n) mapped_ds = ds.map_batches( lambda x: x, batch_size=upstream_batch_size, ).map_batches( lambda x: x, batch_size=downstream_batch_size, ) physical_plan, _ = get_execution_plan(mapped_ds._logical_plan) root_op = physical_plan.dag assert isinstance(root_op, MapOperator) actual_plan_str = root_op.dag_str if upstream_batch_size is None and downstream_batch_size is None: expected_min_rows_per_bundle = None expected_plan_str = ( "InputDataBuffer[Input] -> " "TaskPoolMapOperator[ReadRange->MapBatches()->MapBatches()]" ) else: expected_min_rows_per_bundle = max( upstream_batch_size or 0, downstream_batch_size or 0 ) expected_plan_str = ( "InputDataBuffer[Input] -> TaskPoolMapOperator[ReadRange] -> " "TaskPoolMapOperator[MapBatches()->MapBatches()]" ) assert expected_plan_str == actual_plan_str # Target min-rows requirement is set to max of upstream and downstream strategy = root_op._block_ref_bundler._strategy assert isinstance(strategy, EstimateSize) assert expected_min_rows_per_bundle == strategy._min_rows_per_bundle def test_read_map_batches_operator_fusion_with_randomize_blocks_operator( ray_start_regular_shared_2_cpus, ): # Note: We currently do not fuse MapBatches->RandomizeBlocks. # This test is to ensure that we don't accidentally fuse them. def fn(batch): return {"id": [x + 1 for x in batch["id"]]} n = 10 ds = ray.data.range(n) ds = ds.randomize_block_order() ds = ds.map_batches(fn, batch_size=None) assert set(extract_values("id", ds.take_all())) == set(range(1, n + 1)) stats = ds.stats() # Ensure RandomizeBlockOrder and MapBatches are not fused. assert "RandomizeBlockOrder->MapBatches(fn)" not in stats assert "ReadRange" in stats assert "RandomizeBlockOrder" in stats assert "MapBatches(fn)" in stats # Regression tests ensuring RandomizeBlockOrder is never bypassed in the future assert "ReadRange->MapBatches(fn)->RandomizeBlockOrder" not in stats assert "ReadRange->MapBatches(fn)" not in stats # Ensure all three operators are also present in usage record _check_usage_record(["ReadRange", "MapBatches", "RandomizeBlocks"]) def test_read_map_batches_operator_fusion_with_random_shuffle_operator( ray_start_regular_shared_2_cpus, configure_shuffle_method ): # Note: we currently only support fusing MapOperator->AllToAllOperator. def fn(batch): return {"id": [x + 1 for x in batch["id"]]} n = 10 ds = ray.data.range(n) ds = ds.map_batches(fn, batch_size=None) ds = ds.random_shuffle() assert set(extract_values("id", ds.take_all())) == set(range(1, n + 1)) assert "ReadRange->MapBatches(fn)->RandomShuffle" in ds.stats() _check_usage_record(["ReadRange", "MapBatches", "RandomShuffle"]) ds = ray.data.range(n) ds = ds.random_shuffle() ds = ds.map_batches(fn, batch_size=None) assert set(extract_values("id", ds.take_all())) == set(range(1, n + 1)) # TODO(Scott): Update below assertion after supporting fusion in # the other direction (AllToAllOperator->MapOperator) assert "ReadRange->RandomShuffle->MapBatches(fn)" not in ds.stats() assert all(op in ds.stats() for op in ("ReadRange", "RandomShuffle", "MapBatches")) _check_usage_record(["ReadRange", "RandomShuffle", "MapBatches"]) # Test fusing multiple `map_batches` with multiple `random_shuffle` operations. ds = ray.data.range(n) for _ in range(5): ds = ds.map_batches(fn, batch_size=None) ds = ds.random_shuffle() assert set(extract_values("id", ds.take_all())) == set(range(5, n + 5)) assert f"ReadRange->{'MapBatches(fn)->' * 5}RandomShuffle" in ds.stats() # For interweaved map_batches and random_shuffle operations, we expect to fuse the # two pairs of MapBatches->RandomShuffle, but not the resulting # RandomShuffle operators. ds = ray.data.range(n) ds = ds.map_batches(fn, batch_size=None) ds = ds.random_shuffle() ds = ds.map_batches(fn, batch_size=None) ds = ds.random_shuffle() assert set(extract_values("id", ds.take_all())) == set(range(2, n + 2)) assert "Operator 1 ReadRange->MapBatches(fn)->RandomShuffle" in ds.stats() assert "Operator 2 MapBatches(fn)->RandomShuffle" in ds.stats() _check_usage_record(["ReadRange", "RandomShuffle", "MapBatches"]) # Check the case where the upstream map function returns multiple blocks. ctx = ray.data.DataContext.get_current() old_target_max_block_size = ctx.target_max_block_size ctx.target_max_block_size = 100 def fn(_): return {"data": np.zeros((100, 100))} ds = ray.data.range(10) ds = ds.repartition(2).map(fn).random_shuffle().materialize() assert "Operator 1 ReadRange" in ds.stats() assert "Operator 2 Repartition" in ds.stats() assert "Operator 3 Map(fn)->RandomShuffle" in ds.stats() _check_usage_record(["ReadRange", "RandomShuffle", "MapRows"]) ctx.target_max_block_size = old_target_max_block_size @pytest.mark.parametrize("shuffle", (True, False)) def test_read_map_batches_operator_fusion_with_repartition_operator( ray_start_regular_shared_2_cpus, shuffle, configure_shuffle_method ): def fn(batch): return {"id": [x + 1 for x in batch["id"]]} n = 10 ds = ray.data.range(n) ds = ds.map_batches(fn, batch_size=None) ds = ds.repartition(2, shuffle=shuffle) assert set(extract_values("id", ds.take_all())) == set(range(1, n + 1)) # Operator fusion is only supported for shuffle repartition. if shuffle: assert "ReadRange->MapBatches(fn)->Repartition" in ds.stats() else: assert "ReadRange->MapBatches(fn)->Repartition" not in ds.stats() assert "ReadRange->MapBatches(fn)" in ds.stats() assert "Repartition" in ds.stats() _check_usage_record(["ReadRange", "MapBatches", "Repartition"]) def test_fuse_map_into_shuffle_reduce( ray_start_regular_shared_2_cpus, restore_data_context ): DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE ds = ray.data.range(100).repartition(4, keys=["id"]).map_batches(lambda b: b) dag = get_execution_plan(ds._logical_plan)[0].dag assert dag.name == ( "HashShuffleReduce(keys=('id',), partitions=4)->MapBatches()" ) assert dag._fused_output_map_transformer is not None assert sorted(extract_values("id", ds.take_all())) == list(range(100)) def test_map_not_fused_into_shuffle_reduce_with_downstream_limit( ray_start_regular_shared_2_cpus, restore_data_context ): DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE ds = ( ray.data.range(100) .repartition(4, keys=["id"]) .map_batches(lambda b: b) .limit(10) ) dag = get_execution_plan(ds._logical_plan)[0].dag assert dag.name == "limit=10" map_op = dag.input_dependencies[0] assert map_op.name == "MapBatches()" reduce_op = map_op.input_dependencies[0] assert reduce_op.name == "HashShuffleReduce(keys=('id',), partitions=4)" assert reduce_op._fused_output_map_transformer is None assert len(ds.take_all()) == 10 def test_concurrency_capped_map_not_fused_into_shuffle_reduce( ray_start_regular_shared_2_cpus, restore_data_context ): """A map with a ``concurrency=`` cap is NOT fused into the reduce. The reduce runs one task per partition with no concurrency cap, so fusing would silently ignore the user's limit; keeping the map separate honors it.""" DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE ds = ( ray.data.range(100) .repartition(4, keys=["id"]) .map_batches(lambda b: b, concurrency=2) ) dag = get_execution_plan(ds._logical_plan)[0].dag assert dag.name == "MapBatches()" reduce_op = dag.input_dependencies[0] assert reduce_op.name == "HashShuffleReduce(keys=('id',), partitions=4)" assert reduce_op._fused_output_map_transformer is None def test_non_file_datasink_write_not_fused_into_shuffle_reduce( ray_start_regular_shared_2_cpus, restore_data_context ): from ray.data.datasource.datasink import Datasink DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE class _NoopDatasink(Datasink): def write(self, blocks, ctx): for _ in blocks: pass return None repartitioned = ray.data.range(100).repartition(4, keys=["id"]) write_op = Write( _NoopDatasink(), input_dependencies=[repartitioned._logical_plan.dag], ) dag = get_execution_plan(LogicalPlan(write_op, DataContext.get_current()))[0].dag # The write stays a separate root op feeding off an un-fused reduce. assert dag.name == "Write" reduce_op = dag.input_dependencies[0] assert reduce_op.name == "HashShuffleReduce(keys=('id',), partitions=4)" assert reduce_op._fused_output_map_transformer is None def test_read_map_batches_operator_fusion_with_sort_operator( ray_start_regular_shared_2_cpus, ): # Note: We currently do not fuse MapBatches->Sort. # This test is to ensure that we don't accidentally fuse them, until # we implement it later. def fn(batch): return {"id": [x + 1 for x in batch["id"]]} n = 10 ds = ray.data.range(n) ds = ds.map_batches(fn, batch_size=None) ds = ds.sort("id") assert extract_values("id", ds.take_all()) == list(range(1, n + 1)) # TODO(Scott): update the below assertions after we support fusion. assert "ReadRange->MapBatches->Sort" not in ds.stats() assert "ReadRange->MapBatches" in ds.stats() assert "Sort" in ds.stats() _check_usage_record(["ReadRange", "MapBatches", "Sort"]) def test_read_map_batches_operator_fusion_with_aggregate_operator( ray_start_regular_shared_2_cpus, configure_shuffle_method ): from ray.data.aggregate import AggregateFn # Note: We currently do not fuse MapBatches->Aggregate. # This test is to ensure that we don't accidentally fuse them, until # we implement it later. def fn(batch): return {"id": [x % 2 for x in batch["id"]]} n = 100 grouped_ds = ray.data.range(n).map_batches(fn, batch_size=None).groupby("id") agg_ds = grouped_ds.aggregate( AggregateFn( init=lambda k: [0, 0], accumulate_row=lambda a, r: [a[0] + r["id"], a[1] + 1], merge=lambda a1, a2: [a1[0] + a2[0], a1[1] + a2[1]], finalize=lambda a: a[0] / a[1], name="foo", ), ) agg_ds.take_all() == [{"id": 0, "foo": 0.0}, {"id": 1, "foo": 1.0}] # TODO(Scott): update the below assertions after we support fusion. assert "ReadRange->MapBatches->Aggregate" not in agg_ds.stats() assert "ReadRange->MapBatches" in agg_ds.stats() assert "Aggregate" in agg_ds.stats() _check_usage_record(["ReadRange", "MapBatches", "Aggregate"]) def test_read_map_chain_operator_fusion_e2e( ray_start_regular_shared_2_cpus, ): ds = ray.data.range(10, override_num_blocks=2) ds = ds.filter(fn=lambda x: x["id"] % 2 == 0) ds = ds.map(column_udf("id", lambda x: x + 1)) ds = ds.map_batches( lambda batch: {"id": [2 * x for x in batch["id"]]}, batch_size=None ) ds = ds.flat_map(lambda x: [{"id": -x["id"]}, {"id": x["id"]}]) assert extract_values("id", ds.take_all()) == [ -2, 2, -6, 6, -10, 10, -14, 14, -18, 18, ] name = ( "ReadRange->Filter()->Map()" "->MapBatches()->FlatMap():" ) assert name in ds.stats() _check_usage_record(["ReadRange", "Filter", "MapRows", "MapBatches", "FlatMap"]) def test_write_fusion(ray_start_regular_shared_2_cpus, tmp_path): ds = ray.data.range(10, override_num_blocks=2) ds.write_csv(tmp_path) assert "ReadRange->Write" in ds._write_ds.stats() _check_usage_record(["ReadRange", "WriteCSV"]) @pytest.mark.parametrize( "up_use_actor, up_concurrency, down_use_actor, down_concurrency, should_fuse", [ # === Task->Task cases === # Same concurrency set. Should fuse. (False, 1, False, 1, True), # Different concurrency set. Should not fuse. (False, 1, False, 2, False), # If one op has concurrency set, and the other doesn't, should not fuse. (False, None, False, 1, False), (False, 1, False, None, False), # === Task->Actor cases === # When Task's concurrency is not set, should fuse. (False, None, True, 2, True), (False, None, True, (1, 2), True), # When max size matches, should fuse. (False, 2, True, 2, True), (False, 2, True, (1, 2), True), # When max size doesn't match, should not fuse. (False, 1, True, 2, False), (False, 1, True, (1, 2), False), # === Actor->Task cases === # Should not fuse whatever concurrency is set. (True, 2, False, 2, False), # === Actor->Actor cases === # Should not fuse whatever concurrency is set. (True, 2, True, 2, False), ], ) def test_map_fusion_with_concurrency_arg( ray_start_regular_shared_2_cpus, up_use_actor, up_concurrency, down_use_actor, down_concurrency, should_fuse, ): """Test map operator fusion with different concurrency settings.""" class Map: def __call__(self, row): return row def map(row): return row ds = ray.data.range(10, override_num_blocks=2) if not up_use_actor: ds = ds.map(map, num_cpus=0, concurrency=up_concurrency) up_name = "Map(map)" else: ds = ds.map(Map, num_cpus=0, concurrency=up_concurrency) up_name = "Map(Map)" if not down_use_actor: ds = ds.map(map, num_cpus=0, concurrency=down_concurrency) down_name = "Map(map)" else: ds = ds.map(Map, num_cpus=0, concurrency=down_concurrency) down_name = "Map(Map)" actual_data = ds.to_pandas() expected_data = pd.DataFrame({"id": list(range(10))}) assert rows_same(actual_data, expected_data) name = f"{up_name}->{down_name}" stats = ds.stats() if should_fuse: assert name in stats, stats else: assert name not in stats, stats def check_transform_fns(op, expected_types): assert isinstance(op, MapOperator) transform_fns = op.get_map_transformer().get_transform_fns() assert len(transform_fns) == len(expected_types), transform_fns for i, transform_fn in enumerate(transform_fns): assert isinstance(transform_fn, expected_types[i]), transform_fn @pytest.mark.skip("Needs zero-copy optimization for read->map_batches.") def test_zero_copy_fusion_eliminate_build_output_blocks( ray_start_regular_shared_2_cpus, ): ctx = DataContext.get_current() # Test the EliminateBuildOutputBlocks optimization rule. planner = create_planner() read_op = get_parquet_read_logical_op() op = MapBatches(lambda x: x, input_dependencies=[read_op]) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) # Before optimization, there should be a map op and and read op. # And they should have the following transform_fns. map_op = physical_plan.dag check_transform_fns( map_op, [ BatchMapTransformFn, ], ) read_op = map_op.input_dependencies[0] check_transform_fns( read_op, [ BlockMapTransformFn, ], ) physical_plan = PhysicalOptimizer().optimize(physical_plan) fused_op = physical_plan.dag # After optimization, read and map ops should be fused as one op. # And the BuidlOutputBlocksMapTransformFn in the middle should be dropped. check_transform_fns( fused_op, [ BlockMapTransformFn, BatchMapTransformFn, ], ) @pytest.mark.parametrize( "order,target_num_rows,batch_size,should_fuse", [ # map_batches -> streaming_repartition: fuse when batch_size is a multiple of target_num_rows ("map_then_sr", 20, 20, True), ("map_then_sr", 20, 10, False), ("map_then_sr", 20, 40, True), ("map_then_sr", 20, None, False), # streaming_repartition -> map_batches: not fused ("sr_then_map", 20, 20, False), ], ) def test_streaming_repartition_map_batches_fusion_order_and_params( ray_start_regular_shared_2_cpus, order, target_num_rows, batch_size, should_fuse, ): """Test fusion of streaming_repartition and map_batches with different orders and different target_num_rows/batch_size values.""" n = 100 ds = ray.data.range(n, override_num_blocks=2) if order == "map_then_sr": ds = ds.map_batches(lambda x: x, batch_size=batch_size) ds = ds.repartition(target_num_rows_per_block=target_num_rows, strict=True) expected_fused_name = f"MapBatches()->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()" 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()->MapBatches()->StreamingRepartition[num_rows_per_block={target_rows},strict=True]" in stats1 ), stats1 assert "MapBatches()->MapBatches()" in stats1 def test_filter_operator_no_upstream_fusion(ray_start_regular_shared_2_cpus, capsys): """Test that fused filter operators doesn't fuse further with upstream maps that specify batch_size (since it filter can change the # of rows.) Case 1: map_batches -> filter -> map_batchess Result: (map -> filter) -> map The fused (map -> filter) doesn't fuse with upstream maps. """ n = 100 target_rows = 20 ds1 = ray.data.range(n, override_num_blocks=2) ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows) ds1 = ds1.filter(lambda x: True) ds1 = ds1.map_batches(lambda x: x, batch_size=target_rows) ds1.explain() captured = capsys.readouterr().out.strip() assert "TaskPoolMapOperator[MapBatches()]" in captured assert "TaskPoolMapOperator[MapBatches()->Filter()]" in captured def test_streaming_repartition_multiple_fusion_non_strict( ray_start_regular_shared_2_cpus, ): """Test that non-strict mode allows multiple operators to fuse with StreamingRepartition. Case 1: Map > Map > SR (non-strict) Case 2: Map > SR (non-strict) > SR (non-strict) """ n = 100 target_rows = 20 # Case 1: Map > Map > SR (non-strict) ds1 = ray.data.range(n, override_num_blocks=2) ds1 = ds1.map_batches(lambda x: x, batch_size=None) ds1 = ds1.map_batches(lambda x: x, batch_size=None) ds1 = ds1.repartition(target_num_rows_per_block=target_rows, strict=False) assert len(ds1.take_all()) == n stats1 = ds1.stats() # Verify all three operators are fused together assert ( f"MapBatches()->MapBatches()->StreamingRepartition[num_rows_per_block={target_rows},strict=False]" in stats1 ), f"Expected full fusion in stats: {stats1}" # Case 2: Map > SR (non-strict) > SR (non-strict) # Note: Two consecutive StreamingRepartition operators are merged into one by # CombineShuffles._combine() during logical optimization (before physical fusion). # This test verifies that Map > SR fusion still works after the SR merging. ds2 = ray.data.range(n, override_num_blocks=2) ds2 = ds2.map_batches(lambda x: x, batch_size=None) ds2 = ds2.repartition(target_num_rows_per_block=target_rows, strict=False) ds2 = ds2.repartition(target_num_rows_per_block=target_rows, strict=False) assert len(ds2.take_all()) == n stats2 = ds2.stats() # Verify Map > SR fusion (the two SRs were already merged into one) assert ( f"MapBatches()->StreamingRepartition[num_rows_per_block={target_rows},strict=False]" in stats2 ), f"Expected Map->SR fusion in stats: {stats2}" def test_combine_repartition_aggregate( ray_start_regular_shared_2_cpus, configure_shuffle_method, capsys ): ds = ray.data.range(100) # Apply repartition with shuffle ds = ds.repartition(5, shuffle=True) # Apply groupby aggregate (creates Aggregate operator) ds = ds.groupby("id").count() ds.explain() captured = capsys.readouterr().out # Verify the first shuffle (Repartition) was dropped and Aggregate connects directly to Read expected_optimized_plan = ( "-------- Logical Plan (Optimized) --------\n" "Aggregate[Aggregate]\n" "+- Read[ReadRange]" ) assert expected_optimized_plan in captured def test_combine_streaming_repartition_to_shuffle_repartition( ray_start_regular_shared_2_cpus, configure_shuffle_method, capsys ): ds = ray.data.range(100, override_num_blocks=10) # Apply StreamingRepartition (local repartition) ds = ds.repartition(target_num_rows_per_block=20) # Apply shuffle Repartition (global repartition) ds = ds.repartition(num_blocks=3, shuffle=True) ds.explain() captured = capsys.readouterr().out # Verify the first shuffle (StreamingRepartition) was dropped and Repartition connects directly to Read expected_optimized_plan = ( "-------- Logical Plan (Optimized) --------\n" "Repartition[Repartition]\n" "+- Read[ReadRange]" ) assert expected_optimized_plan in captured def test_combine_sort_sort(ray_start_regular_shared_2_cpus, capsys): data = [{"a": i, "b": 100 - i} for i in range(50)] ds = ray.data.from_items(data) # Apply first sort on column 'a' ds = ds.sort("a") # Apply second sort on column 'b' ds = ds.sort("b") ds.explain() captured = capsys.readouterr().out # Verify the first shuffle (first Sort) was dropped and only the second Sort remains expected_optimized_plan = ( "-------- Logical Plan (Optimized) --------\n" "Sort[Sort]\n" "+- FromItems[FromItems]" ) assert expected_optimized_plan in captured if __name__ == "__main__": import sys sys.exit(pytest.main(["-v", __file__]))