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

Code before:
)
@pytest.mark.parametrize(
    "expr, target_type, expected_value, expected_type",
    [
        # Basic type conversions using Ray Data's DataType
        (col("id"), DataType.int64(), 0, int),  # int64 -> int64 (no change)
        (col("id"), DataType.float64(), 0.0, float),  # int64 -> float64
        (col("id"), DataType.string(), "0", str),  # int64 -> string
        (col("id") / 2, DataType.int64(), 0, int),  # float -> int64
        (col("id") / 2, DataType.float64(), 0.0, float),  # float -> float64 (no change)
    ],
)
def test_cast_expression_basic(
    ray_start_regular_shared,
    expr,
    target_type,
    expected_value,
    expected_type,
    target_max_block_size_infinite_or_default,
):
    """Test basic type casting with cast() method."""
    ds = ray.data.range(5).with_column("result", expr.cast(target_type))
    result = ds.take(1)[0]
    assert result["id"] == 0
    assert result["result"] == expected_value
    assert isinstance(result["result"], expected_type)


@pytest.mark.skipif(
    get_pyarrow_version() < parse_version("20.0.0"),
    reason="with_column requires PyArrow >= 20.0.0",
)
def test_cast_expression_usecase(
    ray_start_regular_shared, target_max_block_size_infinite_or_default
):
    """Test the user use case: converting float result from modulo to int64."""
    ds = ray.data.range(10)
    # The modulo operation returns float, cast it to int64
    ds = ds.with_column("part", (col("id") % 2).cast(DataType.int64()))
    results = ds.take_all()

    # Verify all part values are integers
    for row in results:
        assert isinstance(row["part"], int)
        assert row["part"] in [0, 1]

    # Verify the schema shows int64 type
    schema = ds.schema()
    assert "part" in schema.names
    # The type should be int64
    part_type = schema.types[schema.names.index("part")]

Review comment (suggestion):
Please use `rows_same` with pandas. Ensure to do the same across other tests

Code after:
    reason="with_column requires PyArrow >= 20.0.0",
)
@pytest.mark.parametrize(
    "expr, target_type, expected_rows",
    [
        # Basic type conversions using Ray Data's DataType
        (col("id"), DataType.int64(), [{"id": i, "result": i} for i in range(5)]),
        (
            col("id"),
            DataType.float64(),
            [{"id": i, "result": float(i)} for i in range(5)],
        ),
        (
            col("id"),
            DataType.string(),
            [{"id": i, "result": str(i)} for i in range(5)],
        ),
        (
            col("id") / 2,
            DataType.int64(),
            [{"id": i, "result": i // 2} for i in range(5)],
        ),
        # col("id")/2 uses integer division in expression layer, then cast to float64
        (
            col("id") / 2,
            DataType.float64(),
            [{"id": i, "result": float(i // 2)} for i in range(5)],
        ),
    ],
)
def test_cast_expression_basic(
    ray_start_regular_shared,
    expr,
    target_type,
    expected_rows,
    target_max_block_size_infinite_or_default,
):
    """Test basic type casting with cast() method."""
    ds = ray.data.range(5).with_column("result", expr.cast(target_type))
    actual = ds.take_all()
    assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))


@pytest.mark.skipif(
    get_pyarrow_version() < parse_version("20.0.0"),
    reason="with_column requires PyArrow >= 20.0.0",
)
def test_cast_expression_usecase(
    ray_start_regular_shared, target_max_block_size_infinite_or_default
):
    """Test the user use case: converting float result from modulo to int64."""