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

import itertools

import pandas as pd
import pyarrow as pa
import pytest

import ray
from ray.data import Schema
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.util import column_udf, named_values
from ray.tests.conftest import *  # noqa


@pytest.mark.parametrize("num_datasets", [2, 3, 4, 5, 10])
def test_zip_multiple_datasets(ray_start_regular_shared, num_datasets):
    # Create multiple datasets with different transformations
    datasets = []
    for i in range(num_datasets):
        ds = ray.data.range(5, override_num_blocks=5)
        if i > 0:  # Apply transformation to all but the first dataset
            ds = ds.map(column_udf("id", lambda x, offset=i: x + offset))
        datasets.append(ds)

    ds = datasets[0].zip(*datasets[1:])

    # Verify schema names
    expected_names = ["id"] + [f"id_{i}" for i in range(1, num_datasets)]
    assert ds.schema().names == expected_names

    # Verify data
    expected_data = []
    for row_idx in range(5):
        row_data = tuple(row_idx + i for i in range(num_datasets))
        expected_data.append(row_data)

    assert ds.take() == named_values(expected_names, expected_data)


@pytest.mark.parametrize(
    "num_blocks1,num_blocks2",
    list(itertools.combinations_with_replacement([1, 2, 4, 16], 2)),
)
def test_zip_different_num_blocks_combinations(
    ray_start_regular_shared, num_blocks1, num_blocks2
):
    n = 12
    ds1 = ray.data.range(n, override_num_blocks=num_blocks1)
    ds2 = ray.data.range(n, override_num_blocks=num_blocks2).map(
        column_udf("id", lambda x: x + 1)
    )
    ds = ds1.zip(ds2)
    assert ds.schema().names == ["id", "id_1"]
    assert ds.take() == named_values(
        ["id", "id_1"], list(zip(range(n), range(1, n + 1)))
    )


@pytest.mark.parametrize(
    "num_cols1,num_cols2,should_invert",
    [
        (1, 1, False),
        (4, 1, False),
        (1, 4, True),
        (1, 10, True),
        (10, 10, False),
    ],
)
def test_zip_different_num_blocks_split_smallest(
    ray_start_regular_shared,
    num_cols1,
    num_cols2,
    should_invert,
):
    n = 12
    num_blocks1 = 4
    num_blocks2 = 2
    ds1 = ray.data.from_items(
        [{str(i): i for i in range(num_cols1)}] * n, override_num_blocks=num_blocks1
    )
    ds2 = ray.data.from_items(
        [{str(i): i for i in range(num_cols1, num_cols1 + num_cols2)}] * n,
        override_num_blocks=num_blocks2,
    )
    ds = ds1.zip(ds2).materialize()
    bundles = ds.iter_internal_ref_bundles()
    num_blocks = sum(len(b.block_refs) for b in bundles)
    assert ds.take() == [{str(i): i for i in range(num_cols1 + num_cols2)}] * n
    if should_invert:
        assert num_blocks == num_blocks2
    else:
        assert num_blocks == num_blocks1


def test_zip_pandas(ray_start_regular_shared):
    ds1 = ray.data.from_pandas(pd.DataFrame({"col1": [1, 2], "col2": [4, 5]}))
    ds2 = ray.data.from_pandas(pd.DataFrame({"col3": ["a", "b"], "col4": ["d", "e"]}))
    ds = ds1.zip(ds2)
    assert ds.count() == 2
    result = list(ds.take())
    assert result[0] == {"col1": 1, "col2": 4, "col3": "a", "col4": "d"}

    ds3 = ray.data.from_pandas(pd.DataFrame({"col2": ["a", "b"], "col4": ["d", "e"]}))
    ds = ds1.zip(ds3)
    assert ds.count() == 2
    result = list(ds.take())
    assert result[0] == {"col1": 1, "col2": 4, "col2_1": "a", "col4": "d"}


def test_zip_arrow(ray_start_regular_shared):
    ds1 = ray.data.range(5).map(lambda r: {"id": r["id"]})
    ds2 = ray.data.range(5).map(lambda r: {"a": r["id"] + 1, "b": r["id"] + 2})
    ds = ds1.zip(ds2)
    assert ds.count() == 5
    assert ds.schema() == Schema(
        pa.schema([("id", pa.int64()), ("a", pa.int64()), ("b", pa.int64())])
    )
    result = list(ds.take())
    assert result[0] == {"id": 0, "a": 1, "b": 2}

    # Test duplicate column names.
    ds = ds1.zip(ds1).zip(ds1)
    assert ds.count() == 5
    assert ds.schema() == Schema(
        pa.schema([("id", pa.int64()), ("id_1", pa.int64()), ("id_2", pa.int64())])
    )
    result = list(ds.take())
    assert result[0] == {"id": 0, "id_1": 0, "id_2": 0}


def test_zip_multiple_block_types(ray_start_regular_shared):
    df = pd.DataFrame({"spam": [0]})
    ds_pd = ray.data.from_pandas(df)
    ds2_arrow = ray.data.from_items([{"ham": [0]}])
    assert ds_pd.zip(ds2_arrow).take_all() == [{"spam": 0, "ham": [0]}]


def test_zip_preserve_order(ray_start_regular_shared):
    def foo(x):
        import time

        if x["item"] < 5:
            time.sleep(1)
        return x

    num_items = 10
    items = list(range(num_items))
    ds1 = ray.data.from_items(items, override_num_blocks=num_items)
    ds2 = ray.data.from_items(items, override_num_blocks=num_items)
    ds2 = ds2.map_batches(foo, batch_size=1)
    result = ds1.zip(ds2).take_all()
    assert result == named_values(
        ["item", "item_1"], list(zip(range(num_items), range(num_items)))
    ), result


def test_zip_does_not_free_shared_materialized_blocks(ray_start_regular_shared):
    """Regression test: ZipOperator should not free blocks from a materialized
    dataset that is shared with another consumer.

    Previously, ZipOperator._zip() called _split_at_indices() without specifying
    owned_by_consumer, which defaulted to True. This caused ray.internal.free()
    to be called on blocks that were shared with other operators in the DAG,
    leading to ObjectFreedError.
    """
    # Create a dataset with 3 blocks (rows [7, 7, 6]) and materialize it.
    # The materialized blocks have owns_blocks=False.
    ds = ray.data.range(20, override_num_blocks=3).materialize()
    assert not ds._execute().owns_blocks

    # Consumer 1: a map_batches that uses the same materialized dataset.
    mapped_ds = ds.map_batches(lambda batch: batch, batch_format="pandas")

    # Consumer 2: zip the same materialized dataset with another dataset.
    # This triggers _split_at_indices inside ZipOperator._zip().
    # Use 2 blocks (rows [10, 10]) so that block boundaries are NOT aligned
    # with ds's blocks (rows [7, 7, 6]). This forces actual block splitting
    # (e.g., the first 10-row block gets split at row 7), which exercises
    # the owned_by_consumer code path in _split_all_blocks.
    other_ds = ray.data.range(20, override_num_blocks=2)
    zipped = other_ds.zip(ds)

    # Consuming the zipped result should not raise ObjectFreedError.
    result = zipped.take_all()
    assert len(result) == 20

    # The mapped_ds should also work fine (blocks not freed by the zip).
    result2 = mapped_ds.take_all()
    assert len(result2) == 20


if __name__ == "__main__":
    import sys

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