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

Code before:
    n = 8
    ds = ray.data.range(n, override_num_blocks=n)
    ds = ds.map_batches(
        AsyncActor,
        batch_size=n,
        compute=ray.data.ActorPoolStrategy(size=1, max_tasks_in_flight_per_actor=n),
        concurrency=1,
        max_concurrency=n,
    )

    output = ds.take_all()
    expected_output = [{"id": i} for i in range(n)]
    # Because all tasks are submitted almost simultaneously,
    # the output order may be different compared to the original input.
    assert len(output) == len(expected_output), (len(output), len(expected_output))


def test_map_batches_struct_field_type_divergence(shutdown_only):
    """Test map_batches with struct fields that have diverging primitive types."""

    def generator_fn(batch):
        for i, row_id in enumerate(batch["id"]):
            if i % 2 == 0:
                # Yield struct with fields (a: int64, b: string)
                yield {"data": [{"a": 1, "b": "hello"}]}
            else:
                # Yield struct with fields (a: float64, c: int32)
                # Field 'a' has different type, field 'b' missing, field 'c' new
                yield {"data": [{"a": 1.5, "c": 100}]}

    ds = ray.data.range(4, override_num_blocks=1)
    ds = ds.map_batches(generator_fn, batch_size=4)
    result = ds.materialize()

    rows = result.take_all()
    assert len(rows) == 4

    # Rows 0 and 2 should have int cast to float, with c=None
    assert rows[0]["data"] == {"a": 1.0, "b": "hello", "c": None}
    assert rows[2]["data"] == {"a": 1.0, "b": "hello", "c": None}

    # Rows 1 and 3 should have float a, with b=None
    assert rows[1]["data"] == {"a": 1.5, "b": None, "c": 100}
    assert rows[3]["data"] == {"a": 1.5, "b": None, "c": 100}


if __name__ == "__main__":
    import sys

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


Review comment (bug):
In addition to this E2E test, I think we should add a unit test for this bug as well (maybe at the `concat` function layer of abstraction). Unit tests are not only much faster to run, but also serve as documentation

Code after:
            tasks = [asyncio.create_task(task_yield(row)) for row in rows]
            for task in tasks:
                yield await task

    n = 8
    ds = ray.data.range(n, override_num_blocks=n)
    ds = ds.map_batches(
        AsyncActor,
        batch_size=n,
        compute=ray.data.ActorPoolStrategy(size=1, max_tasks_in_flight_per_actor=n),
        concurrency=1,
        max_concurrency=n,
    )

    output = ds.take_all()
    expected_output = [{"id": i} for i in range(n)]
    # Because all tasks are submitted almost simultaneously,
    # the output order may be different compared to the original input.
    assert len(output) == len(expected_output), (len(output), len(expected_output))


@pytest.mark.skipif(
    get_pyarrow_version() < MIN_PYARROW_VERSION_TYPE_PROMOTION,
    reason="Requires PyArrow >= 14.0.0 for type promotion in nested struct fields",
)
def test_map_batches_struct_field_type_divergence(shutdown_only):
    """Test map_batches with struct fields that have diverging primitive types."""

    def generator_fn(batch):
        for i, row_id in enumerate(batch["id"]):
            if i % 2 == 0:
                # Yield struct with fields (a: int64, b: string)
                yield {"data": [{"a": 1, "b": "hello"}]}
            else:
                # Yield struct with fields (a: float64, c: int32)
                # Field 'a' has different type, field 'b' missing, field 'c' new
                yield {"data": [{"a": 1.5, "c": 100}]}

    ds = ray.data.range(4, override_num_blocks=1)
    ds = ds.map_batches(generator_fn, batch_size=4)
    result = ds.materialize()

    rows = result.take_all()
    assert len(rows) == 4

    # Sort to make the order deterministic.
    rows.sort(key=lambda r: (r["data"]["a"], str(r["data"]["b"])))

    # Rows with a=1.0 (originally int) should have int cast to float, with c=None
    assert rows[0]["data"] == {"a": 1.0, "b": "hello", "c": None}
    assert rows[1]["data"] == {"a": 1.0, "b": "hello", "c": None}