ray-project--ray/python/ray/data/tests/unit/expressions/test_schema_resolution.py

"""Tests for schema-aware expression resolution.

Covers ``Expr.get_type``, ``Expr.nullable``, ``Expr.to_field``, and the
``exprlist_to_fields`` helper. These are the building blocks Phase 1
operators (``Project``, ``Aggregate``, ``Join``, etc.) use to compute
``infer_schema()`` without executing the plan.
"""

import pyarrow as pa
import pytest

from ray.data.datatype import DataType
from ray.data.expressions import (
    DownloadExpr,
    MonotonicallyIncreasingIdExpr,
    RandomExpr,
    UDFExpr,
    UUIDExpr,
    col,
    expand_star_exprs,
    exprlist_to_fields,
    lit,
    star,
    udf,
)
from ray.data.tests.util import assert_exprs_equal


@pytest.fixture
def schema():
    return pa.schema(
        [
            pa.field("a", pa.int32(), nullable=True),
            pa.field("b", pa.float32(), nullable=False),
            pa.field("name", pa.string(), nullable=True),
            pa.field("flag", pa.bool_(), nullable=True),
        ]
    )


class TestColumnExpr:
    def test_to_field_returns_input_field_verbatim(self):
        expected = pa.field("x", pa.int64(), nullable=False, metadata={b"k": b"v"})
        in_schema = pa.schema([expected])
        assert col("x").to_field(in_schema) == expected

    def test_to_field_int_column(self, schema):
        assert col("a").to_field(schema) == pa.field("a", pa.int32(), nullable=True)

    def test_to_field_non_nullable_column(self, schema):
        assert col("b").to_field(schema) == pa.field("b", pa.float32(), nullable=False)

    def test_to_field_missing(self, schema):
        assert col("missing").to_field(schema) is None


class TestLiteralExpr:
    @pytest.mark.parametrize(
        "value,expected",
        [
            (5, pa.field("v", pa.int64(), nullable=False)),
            (5.0, pa.field("v", pa.float64(), nullable=False)),
            ("hello", pa.field("v", pa.string(), nullable=False)),
            (True, pa.field("v", pa.bool_(), nullable=False)),
            (None, pa.field("v", pa.null(), nullable=True)),
        ],
    )
    def test_to_field(self, schema, value, expected):
        # ``LiteralExpr`` has no name, so ``to_field`` returns ``None``;
        # we wrap in an alias to get a named field.
        assert lit(value).alias("v").to_field(schema) == expected


class TestBinaryExpr:
    @pytest.mark.parametrize(
        "expr_factory,expected",
        [
            # int32 + int64 -> int64 (PyArrow promotion); a is nullable.
            (lambda: col("a") + lit(5), pa.field("out", pa.int64(), nullable=True)),
            # int32 + float32 -> float; a is nullable, so output is nullable.
            (lambda: col("a") + col("b"), pa.field("out", pa.float32(), nullable=True)),
            # float32 * float64 -> double; both operands are non-nullable
            # (b is non-null, literal is not None) -> output is non-nullable.
            (
                lambda: col("b") * lit(2.0),
                pa.field("out", pa.float64(), nullable=False),
            ),
            # comparisons -> bool
            (lambda: col("a") > lit(0), pa.field("out", pa.bool_(), nullable=True)),
            (lambda: col("a") == lit(1), pa.field("out", pa.bool_(), nullable=True)),
            # logical -> bool
            (
                lambda: col("flag") & col("flag"),
                pa.field("out", pa.bool_(), nullable=True),
            ),
            (
                lambda: col("flag") | col("flag"),
                pa.field("out", pa.bool_(), nullable=True),
            ),
            # in/not_in -> bool
            (
                lambda: col("a").is_in([1, 2, 3]),
                pa.field("out", pa.bool_(), nullable=True),
            ),
            (
                lambda: col("a").not_in([1, 2, 3]),
                pa.field("out", pa.bool_(), nullable=True),
            ),
            # string concat
            (
                lambda: col("name") + lit("!"),
                pa.field("out", pa.string(), nullable=True),
            ),
        ],
    )
    def test_to_field(self, schema, expr_factory, expected):
        assert expr_factory().alias("out").to_field(schema) == expected

    def test_unresolvable_returns_none(self, schema):
        assert (col("missing") + lit(1)).alias("x").to_field(schema) is None


class TestUnaryExpr:
    def test_is_null(self, schema):
        expected = pa.field("isnull", pa.bool_(), nullable=False)
        assert col("a").is_null().alias("isnull").to_field(schema) == expected

    def test_is_not_null(self, schema):
        expected = pa.field("isnotnull", pa.bool_(), nullable=False)
        assert col("a").is_not_null().alias("isnotnull").to_field(schema) == expected

    def test_not_bool(self, schema):
        expected = pa.field("neg", pa.bool_(), nullable=True)
        assert (~col("flag")).alias("neg").to_field(schema) == expected


class TestAliasExpr:
    def test_to_field_renames(self, schema):
        # Wraps a column field; alias swaps the name, preserves type/nullability.
        assert col("a").alias("renamed").to_field(schema) == pa.field(
            "renamed", pa.int32(), nullable=True
        )

    def test_to_field_around_binary(self, schema):
        assert (col("a") + col("b")).alias("sum").to_field(schema) == pa.field(
            "sum", pa.float32(), nullable=True
        )


class TestSelfContainedExprs:
    def test_udf_uses_return_dtype(self, schema):
        @udf(return_dtype=DataType.float64())  # pyrefly: ignore[missing-attribute]
        def double(x):
            return x

        assert double(col("a")).alias("d").to_field(  # pyrefly: ignore[not-callable]
            schema
        ) == pa.field("d", pa.float64(), nullable=True)

    def test_download_is_binary(self, schema):
        assert DownloadExpr("uri").alias("bytes").to_field(schema) == pa.field(
            "bytes", pa.binary(), nullable=True
        )

    def test_monotonically_increasing_id(self, schema):
        assert MonotonicallyIncreasingIdExpr().alias("id").to_field(schema) == pa.field(
            "id", pa.int64(), nullable=False
        )

    def test_random(self, schema):
        assert RandomExpr().alias("r").to_field(schema) == pa.field(
            "r", pa.float64(), nullable=False
        )

    def test_uuid(self, schema):
        assert UUIDExpr().alias("u").to_field(schema) == pa.field(
            "u", pa.string(), nullable=False
        )


class TestStarExpr:
    def test_to_field_returns_none(self, schema):
        # ``StarExpr`` represents many columns; ``exprlist_to_fields``
        # expands it inline rather than calling ``to_field`` on it.
        assert star().to_field(schema) is None


class TestExprlistToFields:
    def test_simple_columns(self, schema):
        expected = pa.schema(
            [
                pa.field("a", pa.int32(), nullable=True),
                pa.field("b", pa.float32(), nullable=False),
            ]
        )
        result = pa.schema(exprlist_to_fields([col("a"), col("b")], schema))
        assert result == expected

    def test_star_expansion(self, schema):
        # Star expands to all input fields verbatim.
        result = pa.schema(exprlist_to_fields([star()], schema))
        assert result == schema

    def test_star_with_rename(self, schema):
        result = pa.schema(
            exprlist_to_fields([star(), col("a")._rename("renamed_a")], schema)
        )
        # Renaming "a" -> "renamed_a" substitutes the renamed field at
        # "a"'s position (matching runtime ``eval_projection``).
        expected = pa.schema(
            [
                pa.field("renamed_a", pa.int32(), nullable=True),
                pa.field("b", pa.float32(), nullable=False),
                pa.field("name", pa.string(), nullable=True),
                pa.field("flag", pa.bool_(), nullable=True),
            ]
        )
        assert result == expected

    def test_star_with_rename_missing_source_returns_none(self, schema):
        # Renaming an absent column must fail resolution (matching the
        # runtime, which raises "column not found"), not silently append.
        assert exprlist_to_fields([star(), col("missing")._rename("x")], schema) is None

    def test_star_with_with_column(self, schema):
        # with_column-style: [star(), expr.alias(name)] preserves all input
        # columns and appends the new computed column.
        result = pa.schema(
            exprlist_to_fields([star(), (col("a") + col("b")).alias("sum")], schema)
        )
        expected = pa.schema(
            list(schema) + [pa.field("sum", pa.float32(), nullable=True)]
        )
        assert result == expected

    def test_returns_none_on_unresolvable(self, schema):
        assert exprlist_to_fields([col("missing")], schema) is None

    def test_with_column_overrides_existing_column(self, schema):
        # ``with_column("a", new_expr)`` builds ``[star(), new_expr.alias("a")]``.
        # The override should replace the existing "a" in place (last-wins,
        # matching runtime ``eval_projection``'s upsert behavior), not
        # produce a duplicate.
        result = pa.schema(
            exprlist_to_fields([star(), (col("a") + lit(10)).alias("a")], schema)
        )
        expected = pa.schema(
            [
                pa.field("a", pa.int64(), nullable=True),  # new type from a + 10
                pa.field("b", pa.float32(), nullable=False),
                pa.field("name", pa.string(), nullable=True),
                pa.field("flag", pa.bool_(), nullable=True),
            ]
        )
        assert result == expected

    def test_returns_none_on_udf_without_return_dtype(self, schema):
        # Construct a synthetic UDFExpr with object data_type to simulate
        # the "untyped UDF" case.
        e = UDFExpr(
            fn=lambda x: x,
            args=[col("a")],
            kwargs={},
            data_type=DataType(object),
        )
        assert exprlist_to_fields([e.alias("out")], schema) is None


class TestExpandStarExprs:
    """Tests for ``expand_star_exprs`` (eager expansion in ``Project``)."""

    def test_passthrough_without_star(self, schema):
        # No StarExpr -> input list returned unchanged.
        exprs = [col("a"), (col("a") + col("b")).alias("sum")]
        assert expand_star_exprs(exprs, schema) is exprs

    def test_passthrough_when_schema_is_none(self):
        exprs = [star(), col("a")]
        assert expand_star_exprs(exprs, None) is exprs

    def test_simple_star(self, schema):
        # ``[star()]`` -> one ``col()`` per input column.
        result = expand_star_exprs([star()], schema)
        assert_exprs_equal(result, [col("a"), col("b"), col("name"), col("flag")])

    def test_star_with_with_column(self, schema):
        # ``with_column``-style: ``[star(), expr.alias("new")]`` expands
        # to ``[col(a), col(b), col(name), col(flag), expr.alias("new")]``.
        new_expr = (col("a") + col("b")).alias("new")
        result = expand_star_exprs([star(), new_expr], schema)
        assert_exprs_equal(
            result, [col("a"), col("b"), col("name"), col("flag"), new_expr]
        )

    def test_star_with_rename(self, schema):
        # ``rename_columns({"a": "renamed_a"})``: the rename substitutes for
        # its source column *in place* (matching runtime ``eval_projection``
        # / ``exprlist_to_fields``), so ``a`` becomes ``renamed_a`` at
        # position 0 rather than moving to the end.
        rename = col("a")._rename("renamed_a")
        result = expand_star_exprs([star(), rename], schema)
        assert_exprs_equal(result, [rename, col("b"), col("name"), col("flag")])

    def test_star_with_rename_source_missing(self, schema):
        # A rename whose source column isn't in the input schema stays in
        # its trailing position so it still errors ("column not found") at
        # runtime instead of being silently dropped.
        rename = col("missing")._rename("renamed")
        result = expand_star_exprs([star(), rename], schema)
        assert_exprs_equal(
            result, [col("a"), col("b"), col("name"), col("flag"), rename]
        )


if __name__ == "__main__":
    import sys

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