ray-project--ray/python/ray/data/tests/test_map_operator.py

import time
from typing import Iterable
from unittest.mock import MagicMock

import numpy as np
import pandas as pd
import pytest

import ray
from ray._common.test_utils import wait_for_condition
from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
from ray.data._internal.execution.interfaces import (
    ExecutionOptions,
)
from ray.data._internal.execution.interfaces.task_context import TaskContext
from ray.data._internal.execution.operators.actor_pool_map_operator import (
    ActorPoolMapOperator,
)
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.task_pool_map_operator import (
    TaskPoolMapOperator,
)
from ray.data._internal.execution.util import make_ref_bundles
from ray.data._internal.output_buffer import OutputBlockSizeOption
from ray.data._internal.stats import Timer
from ray.data.block import Block
from ray.data.context import (
    DataContext,
)
from ray.data.tests.conftest import noop_counter
from ray.data.tests.util import (
    _get_blocks,
    _mul2_transform,
    _take_outputs,
    create_map_transformer_from_block_fn,
    run_one_op_task,
    run_op_tasks_sync,
)
from ray.tests.conftest import *  # noqa

_mul2_map_data_prcessor = create_map_transformer_from_block_fn(_mul2_transform)


def _run_map_operator_test(
    ray_start_regular_shared,
    use_actors,
    preserve_order,
    transform_fn,
    output_block_size_option,
    expected_blocks,
    test_name="TestMapper",
):
    """Shared test function for MapOperator output unbundling tests."""
    # Create with inputs.
    input_op = InputDataBuffer(
        DataContext.get_current(), make_ref_bundles([[i] for i in range(10)])
    )
    compute_strategy = ActorPoolStrategy() if use_actors else TaskPoolStrategy()

    transformer = create_map_transformer_from_block_fn(
        transform_fn,
        output_block_size_option=output_block_size_option,
    )

    op = MapOperator.create(
        transformer,
        input_op=input_op,
        data_context=DataContext.get_current(),
        name=test_name,
        compute_strategy=compute_strategy,
        # Send everything in a single bundle of 10 blocks.
        min_rows_per_bundle=10,
    )

    # Feed data and block on exec.
    op.start(ExecutionOptions(preserve_order=preserve_order), noop_counter())
    if use_actors:
        # Wait for actors to be ready before adding inputs.
        run_op_tasks_sync(op, only_existing=True)

    while input_op.has_next():
        assert op.can_add_input()
        op.add_input(input_op.get_next(), 0)

    op.all_inputs_done()

    run_op_tasks_sync(op)

    # Check that bundles are unbundled in the output queue.
    outputs = []
    while op.has_next():
        outputs.append(op.get_next())
    assert len(outputs) == expected_blocks
    assert op.has_completed()


@pytest.mark.parametrize("use_actors", [False, True])
def test_map_operator_streamed(ray_start_regular_shared, use_actors):
    # Create with inputs.
    input_op = InputDataBuffer(
        DataContext.get_current(),
        make_ref_bundles([[np.ones(1024) * i] for i in range(100)]),
    )
    compute_strategy = ActorPoolStrategy() if use_actors else TaskPoolStrategy()
    op = MapOperator.create(
        _mul2_map_data_prcessor,
        input_op,
        DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
    )

    # Feed data and implement streaming exec.
    output = []
    # Use preserve_order so output order matches input order (required for
    # actor pool, which otherwise returns results in completion order).
    op.start(
        ExecutionOptions(actor_locality_enabled=True, preserve_order=True),
        noop_counter(),
    )

    if use_actors:
        # Wait for actors to be ready before adding inputs.
        run_op_tasks_sync(op, only_existing=True)

    while input_op.has_next():
        # If actor pool at capacity run 1 task and allow it to copmlete
        while not op.can_add_input():
            run_one_op_task(op)

        op.add_input(input_op.get_next(), 0)

    # Complete ingesting inputs
    op.all_inputs_done()
    run_op_tasks_sync(op)

    assert op.has_execution_finished()
    # NOTE: Op is not considered completed until its outputs are drained
    assert not op.has_completed()

    # Fetch all outputs
    while op.has_next():
        ref = op.get_next()
        assert ref.owns_blocks, ref
        _get_blocks(ref, output)

    assert op.has_completed()

    expected = [[np.ones(1024) * i * 2] for i in range(100)]
    output_sorted = sorted(output, key=lambda x: np.asarray(x[0]).flat[0])
    expected_sorted = sorted(expected, key=lambda x: np.asarray(x[0]).flat[0])
    assert np.array_equal(output_sorted, expected_sorted)
    metrics = op.metrics.as_dict()
    assert metrics["obj_store_mem_freed"] == pytest.approx(832200, 0.5), metrics
    if use_actors:
        assert "locality_hits" in metrics, metrics
        assert "locality_misses" in metrics, metrics
    else:
        assert "locality_hits" not in metrics, metrics
        assert "locality_misses" not in metrics, metrics


def test_map_operator_actor_locality_stats(ray_start_regular_shared):
    # Create with inputs.
    input_op = InputDataBuffer(
        DataContext.get_current(),
        make_ref_bundles([[np.ones(100) * i] for i in range(100)]),
    )
    compute_strategy = ActorPoolStrategy()
    op = MapOperator.create(
        _mul2_map_data_prcessor,
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
        min_rows_per_bundle=None,
    )

    # Feed data and implement streaming exec.
    output = []
    options = ExecutionOptions()
    options.preserve_order = True
    options.actor_locality_enabled = True
    op.start(options, noop_counter())
    # Wait for actors to be ready before adding inputs.
    run_op_tasks_sync(op, only_existing=True)

    while input_op.has_next():
        # If actor pool at capacity run 1 task and allow it to copmlete
        while not op.can_add_input():
            run_one_op_task(op)

        op.add_input(input_op.get_next(), 0)

    # Complete ingesting inputs
    op.all_inputs_done()
    run_op_tasks_sync(op)

    assert op.has_execution_finished()
    # NOTE: Op is not considered completed until its outputs are drained
    assert not op.has_completed()

    # Fetch all outputs
    while op.has_next():
        ref = op.get_next()
        assert ref.owns_blocks, ref
        _get_blocks(ref, output)

    assert op.has_completed()

    # Check equivalent to bulk execution in order.
    assert np.array_equal(output, [[np.ones(100) * i * 2] for i in range(100)])
    metrics = op.metrics.as_dict()
    assert metrics["obj_store_mem_freed"] == pytest.approx(92900, 0.5), metrics
    # Check e2e locality manager working.
    assert metrics["locality_hits"] == 100, metrics
    assert metrics["locality_misses"] == 0, metrics


@pytest.mark.parametrize("use_actors", [False, True])
def test_map_operator_min_rows_per_bundle(ray_start_regular_shared, use_actors):
    # Simple sanity check of batching behavior.
    def _check_batch(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
        block_iter = list(block_iter)
        assert len(block_iter) == 5, block_iter
        data = [block["id"][0] for block in block_iter]
        assert data == list(range(5)) or data == list(range(5, 10)), data
        for block in block_iter:
            yield block

    # Create with inputs.
    input_op = InputDataBuffer(
        DataContext.get_current(), make_ref_bundles([[i] for i in range(10)])
    )
    compute_strategy = ActorPoolStrategy() if use_actors else TaskPoolStrategy()
    op = MapOperator.create(
        create_map_transformer_from_block_fn(_check_batch),
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
        min_rows_per_bundle=5,
    )

    # Feed data and block on exec.
    op.start(ExecutionOptions(), noop_counter())
    if use_actors:
        # Wait for actors to be ready before adding inputs.
        run_op_tasks_sync(op, only_existing=True)

    while input_op.has_next():
        # Should be able to launch 2 tasks:
        #   - Input: 10 blocks of 1 row each
        #   - Bundled into 2 bundles (5 rows each)
        assert op.can_add_input()
        op.add_input(input_op.get_next(), 0)

    op.all_inputs_done()
    run_op_tasks_sync(op)

    _take_outputs(op)
    assert op.has_completed()


@pytest.mark.parametrize("use_actors", [False, True])
@pytest.mark.parametrize("preserve_order", [False, True])
@pytest.mark.parametrize(
    "target_max_block_size,num_expected_blocks", [(1, 10), (2**20, 1), (None, 1)]
)
def test_map_operator_output_unbundling(
    ray_start_regular_shared,
    use_actors,
    preserve_order,
    target_max_block_size,
    num_expected_blocks,
):
    """Test that MapOperator's output queue unbundles bundles from tasks."""

    def noop(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
        for block in block_iter:
            yield block

    _run_map_operator_test(
        ray_start_regular_shared,
        use_actors,
        preserve_order,
        noop,
        OutputBlockSizeOption.of(target_max_block_size=target_max_block_size),
        num_expected_blocks,
    )


@pytest.mark.parametrize("preserve_order", [False, True])
@pytest.mark.parametrize(
    "output_block_size_option,expected_blocks",
    [
        # Test target_max_block_size
        (OutputBlockSizeOption.of(target_max_block_size=1), 10),
        (OutputBlockSizeOption.of(target_max_block_size=2**20), 1),
        (OutputBlockSizeOption.of(target_max_block_size=None), 1),
        # Test target_num_rows_per_block
        (OutputBlockSizeOption.of(target_num_rows_per_block=1), 10),
        (OutputBlockSizeOption.of(target_num_rows_per_block=5), 2),
        (OutputBlockSizeOption.of(target_num_rows_per_block=10), 1),
        (OutputBlockSizeOption.of(target_num_rows_per_block=None), 1),
        # Test disable_block_shaping
        (OutputBlockSizeOption.of(disable_block_shaping=True), 10),
        (OutputBlockSizeOption.of(disable_block_shaping=False), 1),
        # Test combinations
        (
            OutputBlockSizeOption.of(
                target_max_block_size=1, target_num_rows_per_block=5
            ),
            10,
        ),
        (
            OutputBlockSizeOption.of(
                target_max_block_size=2**20, disable_block_shaping=True
            ),
            10,
        ),
        (
            OutputBlockSizeOption.of(
                target_num_rows_per_block=5, disable_block_shaping=True
            ),
            10,
        ),
    ],
)
def test_map_operator_output_block_size_options(
    ray_start_regular_shared,
    preserve_order,
    output_block_size_option,
    expected_blocks,
):
    """Test MapOperator with various OutputBlockSizeOption configurations."""

    def noop(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
        for block in block_iter:
            yield block

    _run_map_operator_test(
        ray_start_regular_shared,
        use_actors=False,
        preserve_order=preserve_order,
        transform_fn=noop,
        output_block_size_option=output_block_size_option,
        expected_blocks=expected_blocks,
    )


@pytest.mark.parametrize("preserve_order", [False, True])
def test_map_operator_disable_block_shaping_with_batches(
    ray_start_regular_shared,
    preserve_order,
):
    """Test MapOperator with disable_block_shaping=True using batch operations."""

    def batch_transform(batch_iter, ctx):
        for batch in batch_iter:
            # Simple transformation: add 1 to each value
            if hasattr(batch, "to_pandas"):
                df = batch.to_pandas()
                df = df + 1
                yield df
            else:
                yield batch

    _run_map_operator_test(
        ray_start_regular_shared,
        use_actors=False,
        preserve_order=preserve_order,
        transform_fn=batch_transform,
        output_block_size_option=OutputBlockSizeOption.of(disable_block_shaping=True),
        expected_blocks=10,  # With disable_block_shaping=True, we expect 10 blocks
        test_name="TestBatchMapper",
    )


@pytest.mark.parametrize("use_actors", [False, True])
def test_map_operator_ray_args(shutdown_only, use_actors):
    ray.shutdown()
    ray.init(num_cpus=0, num_gpus=1)
    # Create with inputs.
    input_op = InputDataBuffer(
        DataContext.get_current(), make_ref_bundles([[i] for i in range(10)])
    )
    compute_strategy = ActorPoolStrategy(size=1) if use_actors else TaskPoolStrategy()
    op = MapOperator.create(
        _mul2_map_data_prcessor,
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
        ray_remote_args={"num_cpus": 0, "num_gpus": 1},
    )

    # Feed data and block on exec.
    op.start(ExecutionOptions(), noop_counter())
    if use_actors:
        # Wait for the actor to start.
        run_op_tasks_sync(op)

    while input_op.has_next():
        if use_actors:
            # For actors, we need to check capacity before adding input
            # and process tasks when the actor pool is at capacity.
            while not op.can_add_input():
                run_one_op_task(op)

        assert op.can_add_input()
        op.add_input(input_op.get_next(), 0)

    op.all_inputs_done()
    run_op_tasks_sync(op)

    # Check we don't hang and complete with num_gpus=1.
    outputs = _take_outputs(op)
    expected = [[i * 2] for i in range(10)]
    assert sorted(outputs) == expected, f"Expected {expected}, got {outputs}"
    assert op.has_completed()


@pytest.mark.parametrize("use_actors", [False, True])
def test_map_operator_shutdown(shutdown_only, use_actors):
    ray.shutdown()
    ray.init(num_cpus=0, num_gpus=1)

    def _sleep(block_iter: Iterable[Block]) -> Iterable[Block]:
        time.sleep(999)

    # Create with inputs.
    input_op = InputDataBuffer(
        DataContext.get_current(), make_ref_bundles([[i] for i in range(10)])
    )
    compute_strategy = ActorPoolStrategy(size=1) if use_actors else TaskPoolStrategy()
    op = MapOperator.create(
        create_map_transformer_from_block_fn(_sleep),
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
        ray_remote_args={"num_cpus": 0, "num_gpus": 1},
    )

    # Start one task and then cancel.
    op.start(ExecutionOptions(), noop_counter())
    if use_actors:
        # Wait for the actor to start.
        run_op_tasks_sync(op)
    op.add_input(input_op.get_next(), 0)
    assert op.num_active_tasks() == 1
    # Regular Ray tasks can be interrupted/cancelled, so graceful shutdown works.
    # Actors running time.sleep() cannot be interrupted gracefully and need ray.kill() to release resources.
    # After proper shutdown, both should return the GPU to ray.available_resources().
    force_shutdown = use_actors
    op.shutdown(timer=Timer(), force=force_shutdown)

    # Tasks/actors should be cancelled/killed.
    wait_for_condition(lambda: (ray.available_resources().get("GPU", 0) == 1.0))


@pytest.mark.parametrize(
    "compute,expected",
    [
        (TaskPoolStrategy(), TaskPoolMapOperator),
        (ActorPoolStrategy(), ActorPoolMapOperator),
    ],
)
def test_map_operator_pool_delegation(compute, expected):
    # Test that the MapOperator factory delegates to the appropriate pool
    # implementation.
    input_op = InputDataBuffer(
        DataContext.get_current(), make_ref_bundles([[i] for i in range(100)])
    )
    op = MapOperator.create(
        _mul2_map_data_prcessor,
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute,
    )
    assert isinstance(op, expected)


@pytest.mark.parametrize("use_actors", [False, True])
def test_map_kwargs(ray_start_regular_shared, use_actors):
    """Test propagating additional kwargs to map tasks."""
    foo = 1
    bar = np.random.random(1024 * 1024)
    kwargs = {
        "foo": foo,  # Pass by value
        "bar": ray.put(bar),  # Pass by ObjectRef
    }

    def map_fn(block_iter: Iterable[Block], ctx: TaskContext) -> Iterable[Block]:
        nonlocal foo, bar
        assert ctx.kwargs["foo"] == foo
        # bar should be automatically deref'ed.
        assert np.array_equal(ctx.kwargs["bar"], bar)

        yield from block_iter

    input_op = InputDataBuffer(
        DataContext.get_current(),
        make_ref_bundles([[i] for i in range(10)]),
    )
    compute_strategy = ActorPoolStrategy() if use_actors else TaskPoolStrategy()
    op = MapOperator.create(
        create_map_transformer_from_block_fn(map_fn),
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestMapper",
        compute_strategy=compute_strategy,
    )
    op.add_map_task_kwargs_fn(lambda: kwargs)
    op.start(ExecutionOptions(), noop_counter())
    if use_actors:
        # Wait for the actor to start.
        run_op_tasks_sync(op)

    while input_op.has_next():
        if use_actors:
            # For actors, we need to check capacity before adding input
            # and process tasks when the actor pool is at capacity.
            while not op.can_add_input():
                run_one_op_task(op)

        assert op.can_add_input()
        op.add_input(input_op.get_next(), 0)
    op.all_inputs_done()
    run_op_tasks_sync(op)

    _take_outputs(op)
    assert op.has_completed()


@pytest.mark.parametrize(
    "target_max_block_size, expected_num_outputs_per_task",
    [
        # 5 blocks (8b each) // 1 = 5 outputs / task
        [1, 5],
        # 5 blocks (8b each) // 1024 = 1 output / task
        [1024, 1],
        # All outputs combined in a single output
        [None, 1],
    ],
)
def test_map_estimated_num_output_bundles(
    target_max_block_size,
    expected_num_outputs_per_task,
):
    # Test map operator estimation
    input_op = InputDataBuffer(
        DataContext.get_current(), make_ref_bundles([[i] for i in range(100)])
    )

    def yield_five(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
        for i in range(5):
            yield pd.DataFrame({"id": [i]})

    min_rows_per_bundle = 10
    # 100 inputs -> 100 / 10 = 10 tasks
    num_tasks = 10

    op = MapOperator.create(
        create_map_transformer_from_block_fn(
            yield_five,
            # Limit single block to hold no more than 1 byte
            output_block_size_option=OutputBlockSizeOption.of(
                target_max_block_size=target_max_block_size,
            ),
        ),
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestEstimatedNumBlocks",
        min_rows_per_bundle=min_rows_per_bundle,
    )

    op.start(ExecutionOptions(), noop_counter())
    while input_op.has_next():
        op.add_input(input_op.get_next(), 0)
        if op.metrics.num_inputs_received % min_rows_per_bundle == 0:
            # enough inputs for a task bundle
            run_op_tasks_sync(op)
            assert (
                op._estimated_num_output_bundles
                == expected_num_outputs_per_task * num_tasks
            )

    op.all_inputs_done()

    assert op._estimated_num_output_bundles == expected_num_outputs_per_task * num_tasks


def test_map_estimated_blocks_split():
    # Test read output splitting

    min_rows_per_bundle = 10
    input_op = InputDataBuffer(
        DataContext.get_current(),
        make_ref_bundles(
            [[i, i + 1] for i in range(100)]
        ),  # create 2-row blocks so split_blocks can split into 2 blocks
    )

    def yield_five(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
        for i in range(5):
            yield pd.DataFrame({"id": [i]})

    op = MapOperator.create(
        create_map_transformer_from_block_fn(
            yield_five,
            # NOTE: Disable output block-shaping to keep blocks from being
            #       combined
            disable_block_shaping=True,
        ),
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestEstimatedNumBlocksSplit",
        min_rows_per_bundle=min_rows_per_bundle,
    )
    op.set_additional_split_factor(2)

    op.start(ExecutionOptions(), noop_counter())
    while input_op.has_next():
        op.add_input(input_op.get_next(), 0)
        if op.metrics.num_inputs_received % min_rows_per_bundle == 0:
            # enough inputs for a task bundle
            run_op_tasks_sync(op)
            assert op._estimated_num_output_bundles == 100

    op.all_inputs_done()
    # Each output block is split in 2, so the number of blocks double.
    assert op._estimated_num_output_bundles == 100


def test_operator_metrics():
    NUM_INPUTS = 100
    NUM_BLOCKS_PER_TASK = 5
    MIN_ROWS_PER_BUNDLE = 10

    inputs = make_ref_bundles([[i] for i in range(NUM_INPUTS)])
    input_op = InputDataBuffer(DataContext.get_current(), inputs)

    def map_fn(block_iter: Iterable[Block], ctx) -> Iterable[Block]:
        for i in range(NUM_BLOCKS_PER_TASK):
            yield pd.DataFrame({"id": [i]})

    op = MapOperator.create(
        create_map_transformer_from_block_fn(
            map_fn,
            output_block_size_option=OutputBlockSizeOption.of(
                target_max_block_size=1,
            ),
        ),
        input_op=input_op,
        data_context=DataContext.get_current(),
        name="TestEstimatedNumBlocks",
        min_rows_per_bundle=MIN_ROWS_PER_BUNDLE,
    )

    op.start(ExecutionOptions(), noop_counter())
    num_outputs_taken = 0
    bytes_outputs_taken = 0
    for i in range(len(inputs)):
        # Add an input, run all tasks, and take all outputs.
        op.add_input(input_op.get_next(), 0)
        run_op_tasks_sync(op)
        while op.has_next():
            output = op.get_next()
            num_outputs_taken += 1
            bytes_outputs_taken += output.size_bytes()

        num_tasks_submitted = (i + 1) // MIN_ROWS_PER_BUNDLE

        metrics = op.metrics
        # Check input metrics
        assert metrics.num_inputs_received == i + 1, i
        assert metrics.bytes_inputs_received == sum(
            inputs[k].size_bytes() for k in range(i + 1)
        ), i
        assert (
            metrics.num_task_inputs_processed
            == num_tasks_submitted * MIN_ROWS_PER_BUNDLE
        ), i
        assert metrics.bytes_task_inputs_processed == sum(
            inputs[k].size_bytes()
            for k in range(num_tasks_submitted * MIN_ROWS_PER_BUNDLE)
        ), i

        # Check outputs metrics
        assert num_outputs_taken == num_tasks_submitted * NUM_BLOCKS_PER_TASK, i
        assert metrics.num_task_outputs_generated == num_outputs_taken, i
        assert metrics.bytes_task_outputs_generated == bytes_outputs_taken, i
        assert metrics.num_outputs_taken == num_outputs_taken, i
        assert metrics.bytes_outputs_taken == bytes_outputs_taken, i
        assert metrics.num_outputs_of_finished_tasks == num_outputs_taken, i
        assert metrics.bytes_outputs_of_finished_tasks == bytes_outputs_taken, i

        # Check task metrics
        assert metrics.num_tasks_submitted == num_tasks_submitted, i
        assert metrics.num_tasks_running == 0, i
        assert metrics.num_tasks_have_outputs == num_tasks_submitted, i
        assert metrics.num_tasks_finished == num_tasks_submitted, i

        # Check object store metrics
        assert metrics.obj_store_mem_freed == metrics.bytes_task_inputs_processed, i


@pytest.mark.parametrize(
    "ray_remote_args", [{}, {"num_cpus": 0}, {"num_cpus": 0.5}, {"num_cpus": 1}]
)
@pytest.mark.parametrize(
    "compute_strategy",
    [ray.data.TaskPoolStrategy(), ray.data.ActorPoolStrategy(size=1)],
)
def test_map_operator_specifies_default_memory(
    ray_start_regular_shared, ray_remote_args, compute_strategy
):
    data_context = ray.data.DataContext.get_current()
    data_context.default_map_logical_memory_enabled = True
    op = MapOperator.create(
        map_transformer=MagicMock(),
        input_op=InputDataBuffer(data_context, input_data=MagicMock()),
        data_context=data_context,
        compute_strategy=compute_strategy,
        ray_remote_args=ray_remote_args,
    )

    # If Ray Data doesn't specify a default memory, then the system can oversubscribe
    # tasks and actors even if the user has correctly specified memory for some UDFs.
    #
    # This assertion just checks that map operators default to *something*, without
    # making assumptions about the actual heuristic.
    assert op.min_scheduling_resources().memory > 0


@pytest.mark.parametrize(
    "compute_strategy",
    [ray.data.TaskPoolStrategy(), ray.data.ActorPoolStrategy(size=1)],
)
def test_map_operator_no_default_memory_when_disabled(
    ray_start_regular_shared, compute_strategy
):
    data_context = ray.data.DataContext.get_current()
    op = MapOperator.create(
        map_transformer=MagicMock(),
        input_op=InputDataBuffer(data_context, input_data=MagicMock()),
        data_context=data_context,
        compute_strategy=compute_strategy,
        ray_remote_args={},
    )

    # When the flag is disabled (the default), map operators shouldn't assign a default
    # logical memory unless the user explicitly requested it.
    assert not op.min_scheduling_resources().memory


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))