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

Code before:

        assert set(all_ids) == expected_ids

    # Create dataset with 2 blocks of 2 rows each
    expected = {0, 1, (1 << 33) + 0, (1 << 33) + 1}

    ds = ray.data.from_items(
        [{"a": 1}, {"a": 2}, {"a": 1}, {"a": 2}], override_num_blocks=2
    )
    ds = ds.with_column("uid", monotonically_increasing_id())
    check_ids(ds, expected)

    # pandas
    ds_pandas = ray.data.from_pandas([pd.DataFrame({"a": [1, 2]}) for _ in range(2)])
    ds_pandas = ds_pandas.with_column("uid", monotonically_increasing_id())
    check_ids(ds_pandas, expected, batch_format="pandas")

    # pyarrow
    ds_arrow = ray.data.range(4, override_num_blocks=2)
    ds_arrow = ds_arrow.with_column("uid", monotonically_increasing_id())
    check_ids(ds_arrow, expected, batch_format="pyarrow")

    # numpy batch
    ds_numpy = ray.data.from_numpy([np.ones((2, 1)) for _ in range(2)])
    ds_numpy = ds_numpy.with_column("uid", monotonically_increasing_id())
    check_ids(ds_numpy, expected, batch_format="numpy")


def test_monotonically_increasing_id_multiple_expressions(ray_start_regular_shared):
    """Test that two monotonically_increasing_id() expressions are isolated."""
    ds = ray.data.range(10, override_num_blocks=5)

    # Two monotonically_increasing_id() expressions should have isolated row counts
    ds = ds.with_column("uid1", monotonically_increasing_id())
    ds = ds.with_column("uid2", monotonically_increasing_id())

    result = ds.to_pandas()

    assert list(result["uid1"]) == list(result["uid2"])


def test_monotonically_increasing_id_structurally_equals_always_false():
    """Test that structurally_equals() is False for monotonically_increasing_id() expressions."""
    expr1 = monotonically_increasing_id()
    expr2 = monotonically_increasing_id()

    # Should always be false (even to itself) due to non-determinism
    assert not expr1.structurally_equals(expr2)
    assert not expr1.structurally_equals(expr1)



Review comment (suggestion):
Rather than having multiple cases, can you use `pytest.mark.parmaetrize`? Motivation is reasons described here: https://learn.microsoft.com/en-us/dotnet/core/testing/unit-testing-best-practices#avoid-multiple-act-tasks

Code after:
    Test that two monotonically_increasing_id() expressions are isolated
    if executed by the same task.
    """
    ds = ray.data.range(10, override_num_blocks=5)

    # Two monotonically_increasing_id() expressions should have isolated row counts
    ds = ds.with_column("uid1", monotonically_increasing_id())
    ds = ds.with_column("uid2", monotonically_increasing_id())

    result = ds.to_pandas()

    assert list(result["uid1"]) == list(result["uid2"])


def test_monotonically_increasing_id_multi_block_per_task(ray_start_regular_shared):
    """Test that IDs are unique when a single task processes multiple blocks."""
    ctx = ray.data.DataContext.get_current()
    original_max_block_size = ctx.target_max_block_size
    try:
        # Set max block size to 32 bytes ~ 4 int64 rows per block.
        # With 5 read tasks of 20 rows each every task should see 5 blocks.
        ctx.target_max_block_size = 32

        ds = ray.data.range(100, override_num_blocks=5)
        ds = ds.with_column("uid", monotonically_increasing_id())
        result = ds.take_all()

        uids = [row["uid"] for row in result]
        assert len(uids) == 100, f"expected 100 rows, got {len(uids)}"
        assert len(uids) == len(set(uids)), "IDs are not unique across blocks"
    finally:
        ctx.target_max_block_size = original_max_block_size


def test_monotonically_increasing_id_structurally_equals_always_false():
    """Test that structurally_equals() is False for monotonically_increasing_id() expressions."""
    expr1 = monotonically_increasing_id()
    expr2 = monotonically_increasing_id()

    # Should always be false (even to itself) due to non-determinism
    assert not expr1.structurally_equals(expr2)
    assert not expr1.structurally_equals(expr1)


def test_monotonically_increasing_id_shuffle_and_sort(ray_start_regular_shared):
    """Test monotonically_increasing_id() in shuffle and sort."""
    ds = ray.data.range(20, override_num_blocks=5)
    ds = ds.with_column("uid", monotonically_increasing_id())
    ds = ds.random_shuffle()
    ds = ds.sort("uid")