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

Code before:
    block_ref = ray.put(block)
    metadata = BlockAccessor.for_block(block).get_metadata()
    schema = BlockAccessor.for_block(block).schema()
    op = InputDataBuffer(
        DataContext.get_current(),
        input_data=[
            RefBundle([(block_ref, metadata)], owns_blocks=False, schema=schema)
        ],
    )

    op.get_next()
    gc.collect()

    # `InputDataBuffer` should still hold a reference to the input block even after
    # `get_next` is called.
    assert len(gc.get_referrers(block_ref)) > 0


def test_union_operator_throttling_disabled(ray_start_regular_shared):
    """Test that UnionOperator has throttling disabled.

    UnionOperator only manipulates bundle metadata and doesn't launch any tasks,
    so it should not be allocated resources.
    """
    ctx = DataContext.get_current()

    # Create input buffers with some test data
    input_bundles_1 = make_ref_bundles([[1, 2, 3]])
    input_bundles_2 = make_ref_bundles([[4, 5, 6]])

    input_op1 = InputDataBuffer(ctx, input_bundles_1)
    input_op2 = InputDataBuffer(ctx, input_bundles_2)

    # Create union operator
    union_op = UnionOperator(ctx, input_op1, input_op2)

    # Verify that throttling_disabled() returns True
    assert union_op.throttling_disabled() is True


def test_union_operator_throttling_matches_similar_operators(ray_start_regular_shared):
    """Test that UnionOperator's throttling behavior matches other metadata-only operators."""
    from ray.data._internal.execution.operators.limit_operator import LimitOperator
    from ray.data._internal.execution.operators.output_splitter import OutputSplitter

    ctx = DataContext.get_current()

    # Create input buffer
    input_bundles = make_ref_bundles([[1, 2, 3]])
    input_op = InputDataBuffer(ctx, input_bundles)


Review comment (refactor):
In addition to testing the state of union operator (i.e., that throttling is disabled), I think it'd be also useful to test that it doesn't get allocated resources. 

You can use `create_resource_allocator` to make the test agnostic to the specific allocator implementation, and the pseudocode could look like this:
```python
allocator = create_resource_allocator(...)
allocator.update_budgets(...)

allocation = allocator.get_allocation(union_op)

assert allocation == ExecutionResources.zero()
```

This sort of test would be more robust to refactors (e.g., if we remove `throttling_disabled)

Code after:
    assert op2._estimated_num_output_bundles is None
    assert op2.num_outputs_total() == estimated_output_blocks


def test_input_data_buffer_does_not_free_inputs():
    # Tests https://github.com/ray-project/ray/issues/46282
    block = pd.DataFrame({"id": [0]})
    block_ref = ray.put(block)
    metadata = BlockAccessor.for_block(block).get_metadata()
    schema = BlockAccessor.for_block(block).schema()
    op = InputDataBuffer(
        DataContext.get_current(),
        input_data=[
            RefBundle([(block_ref, metadata)], owns_blocks=False, schema=schema)
        ],
    )

    op.get_next()
    gc.collect()

    # `InputDataBuffer` should still hold a reference to the input block even after
    # `get_next` is called.
    assert len(gc.get_referrers(block_ref)) > 0


def test_union_operator_throttling_disabled(ray_start_regular_shared):
    """Test that UnionOperator has throttling disabled.

    UnionOperator only manipulates bundle metadata and doesn't launch any tasks,
    so it should not be allocated resources.
    """
    ctx = DataContext.get_current()
    union_op = UnionOperator(ctx)

    # Verify that throttling_disabled() returns True
    assert union_op.throttling_disabled() is True


def test_union_operator_not_allocated_resources(
    ray_start_regular_shared, restore_data_context
):
    ctx = DataContext.get_current()
    ctx.op_resource_reservation_enabled = True
    input_op1 = InputDataBuffer(ctx, make_ref_bundles([[1, 2, 3]]))
    input_op2 = InputDataBuffer(ctx, make_ref_bundles([[4, 5, 6]]))
    union_op = UnionOperator(ctx, input_op1, input_op2)

    topo = build_streaming_topology(union_op, ExecutionOptions())
    resource_manager = ResourceManager(topo, ExecutionOptions(), MagicMock(), ctx)
    allocator = create_resource_allocator(resource_manager, ctx)
    assert allocator is not None