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

Code before:
    """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 (refactor):
Though I don't think this is possible given the current Ray Data implementation, this row ordering isn't guaranteed by the interface. This has historically been the cause of a lot of Ray Data's flaky tests.

Could you refactor this test so that it doesn't depend on a particular test ordering?

Code after:
    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}

    # Rows with a=1.5 should have float a, with b=None
    assert rows[2]["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__]))