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

Code before:
    def check_ids(ds, expected_ids, batch_format="default"):
        all_ids = []
        for block in ds.iter_batches(batch_size=None, batch_format=batch_format):
            if batch_format == "pandas":
                assert isinstance(block, pd.DataFrame)
                block_ids = block["uid"].tolist()
            elif batch_format == "pyarrow":
                assert isinstance(block, pa.Table)
                block_ids = block["uid"].to_pylist()
            elif batch_format == "numpy":
                assert isinstance(block, dict)
                assert isinstance(block["uid"], np.ndarray)
                block_ids = block["uid"].tolist()
            else:
                block_ids = list(block["uid"])

            print(f"block_ids: {block_ids}")
            all_ids.extend(block_ids)
            assert block_ids == sorted(block_ids), "block IDs are not monotonic"
            if len(block_ids) > 1:
                diffs = [
                    block_ids[i + 1] - block_ids[i] for i in range(len(block_ids) - 1)
                ]
                assert all(d == 1 for d in diffs), "block IDs are not consecutive"

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


Review comment (performance):
I think we can just do something like this:
```python

all_ids = []
for block in ds.iter_batches(batch_size=None, batch_format="pyarrow"):
    block_ids = block["uid"].to_pylist() 
    assert block_ids == sorted(block_ids)

    all_ids.extend(block_ids)

assert set(all_ids) == expected_ids
```

I don't think testing the different iter batches formats is relevant for testing this expression

Code after:
def test_monotonically_increasing_id(ray_start_regular_shared, block_type):
    """Test monotonically_increasing_id() expression produces monotonically increasing IDs."""

    if block_type == "arrow":
        blocks = [pa.table({"a": [1, 2]}), pa.table({"a": [3, 4]})]
    else:
        blocks = [pd.DataFrame({"a": [1, 2]}), pd.DataFrame({"a": [3, 4]})]

    # Create dataset with 2 blocks of 2 rows each
    ds = ray.data.from_blocks(blocks)
    ds = ds.with_column("uid", monotonically_increasing_id())
    expected = {0, 1, (1 << 33) + 0, (1 << 33) + 1}

    all_ids = []
    for batch in ds.iter_batches(batch_size=None, batch_format="pyarrow"):
        block_ids = batch["uid"].to_pylist()
        all_ids.extend(block_ids)
        assert block_ids == sorted(block_ids), "block IDs are not monotonic"

    assert set(all_ids) == expected


def test_monotonically_increasing_id_multiple_expressions(ray_start_regular_shared):
    """
    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()