Function _run_map_operator_test

python/ray/data/tests/test_map_operator.py:46–96

Shared test function for MapOperator output unbundling tests.

(
    ray_start_regular_shared,
    use_actors,
    preserve_order,
    transform_fn,
    output_block_size_option,
    expected_blocks,
    test_name="TestMapper",
)

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))
    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])