"""Integration tests for ``LogicalOperator.infer_schema()`` (Phase 1).

Asserts that ``Dataset.schema()`` resolves the output schema **without**
falling back to a ``limit(1)`` execution for every non-UDF chain. UDF
chains (``map``, ``map_batches``, ``flat_map``) correctly return ``None``.

The headline guarantee is verified by calling ``ds.schema(fetch_if_missing=False)``
through the public API: this disables the ``limit(1)`` fallback in
``_base_schema``, so a non-``None`` return proves the static
``LogicalOperator.infer_schema()`` path resolved the schema without
materializing any blocks.
"""

import pyarrow as pa
import pyarrow.parquet as pq
import pytest

import ray
from ray.data._internal.logical.operators import Project
from ray.data.aggregate import Count, Max, Mean, Min, Sum
from ray.data.expressions import col, lit, star
from ray.data.tests.conftest import *  # noqa: F401,F403
from ray.data.tests.util import assert_exprs_equal


@pytest.fixture(scope="module")
def parquet_path(tmp_path_factory):
    tmp = tmp_path_factory.mktemp("infer_schema")
    table = pa.table(
        {
            "a": pa.array([1, 2, 3, 4, 5], type=pa.int32()),
            "b": pa.array([1.0, 2.0, 3.0, 4.0, 5.0], type=pa.float32()),
            "k": pa.array(["x", "x", "y", "y", "z"], type=pa.string()),
        }
    )
    pq.write_table(table, tmp / "data.parquet")
    return tmp


def _static_schema(ds: ray.data.Dataset) -> pa.Schema:
    """Return the dataset's schema via the static path only.

    ``fetch_if_missing=False`` disables ``_base_schema``'s ``limit(1)``
    fallback, so a non-``None`` return proves the chain's
    ``infer_schema()`` resolved without materializing any blocks. Returns
    ``None`` if the static path failed.
    """
    schema = ds.schema(fetch_if_missing=False)
    return None if schema is None else schema.base_schema


class TestSourceAndPassthroughs:
    def test_read_parquet(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = ray.data.read_parquet(str(parquet_path))
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
            ]
        )

    def test_filter_with_expr_passthrough(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        ds = ray.data.read_parquet(str(parquet_path)).filter(expr=col("a") > 0)
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
            ]
        )

    def test_filter_with_fn_passthrough(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        # Filter with a callable still preserves the input schema.
        ds = ray.data.read_parquet(str(parquet_path)).filter(lambda row: row["a"] > 0)
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
            ]
        )

    def test_limit_passthrough(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = ray.data.read_parquet(str(parquet_path)).limit(2)
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
            ]
        )

    def test_sort_passthrough(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = ray.data.read_parquet(str(parquet_path)).sort("a")
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
            ]
        )


class TestProject:
    def test_select_columns(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = ray.data.read_parquet(str(parquet_path)).select_columns(["a", "b"])
        assert _static_schema(ds) == pa.schema(
            [pa.field("a", pa.int32()), pa.field("b", pa.float32())]
        )

    def test_with_column(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = ray.data.read_parquet(str(parquet_path)).with_column(
            "s", col("a") + col("b")
        )
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
                pa.field("s", pa.float32()),
            ]
        )

    def test_rename_columns(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = ray.data.read_parquet(str(parquet_path)).rename_columns({"a": "x"})
        # Renamed columns occupy the source column's position so the static
        # schema matches the materialized column order.
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("x", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
            ]
        )

    def test_with_column_chain(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .with_column("s", col("a") + col("b"))
            .with_column("d", col("s") * lit(2.0))
        )
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
                pa.field("s", pa.float32()),
                pa.field("d", pa.float64()),
            ]
        )


class TestDropColumns:
    """Phase 3: ``Dataset.drop_columns`` reshapes into a ``Project`` over
    the surviving columns when the input schema is known, so the typed
    chain stays intact through it."""

    def test_drop_columns_static_schema(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        ds = ray.data.read_parquet(str(parquet_path)).drop_columns(["b"])
        # Reshaped to a ``Project`` (not a UDF ``MapBatches``), so the
        # typed schema chain survives.
        assert isinstance(ds._logical_plan.dag, Project)
        assert _static_schema(ds) == pa.schema(
            [pa.field("a", pa.int32()), pa.field("k", pa.string())]
        )

    def test_drop_columns_missing_raises_eagerly(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        ds = ray.data.read_parquet(str(parquet_path))
        with pytest.raises(KeyError, match="not found in dataset schema"):
            ds.drop_columns(["does_not_exist"])

    def test_drop_columns_after_map_batches_falls_back(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        # Input schema is unknown downstream of ``map_batches``, so the
        # implementation falls back to a ``MapBatches`` closure and the
        # typed-chain guarantee no longer holds.
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .map_batches(lambda b: b)
            .drop_columns(["b"])
        )
        assert ds.schema(fetch_if_missing=False) is None


class TestAggregate:
    def test_groupby_multi_aggs(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .groupby("k")
            .aggregate(Sum("a"), Mean("b"), Count("a"), Max("a"), Min("a"))
        )
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("k", pa.string()),
                pa.field("sum(a)", pa.int64()),
                pa.field("mean(b)", pa.float64()),
                pa.field("count(a)", pa.int64(), nullable=False),
                pa.field("max(a)", pa.int32()),
                pa.field("min(a)", pa.int32()),
            ]
        )

    def test_groupby_multi_key(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .groupby(["k", "a"])
            .aggregate(Count("b"))
        )
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("k", pa.string()),
                pa.field("a", pa.int32()),
                pa.field("count(b)", pa.int64(), nullable=False),
            ]
        )

    def test_groupby_then_sort(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .groupby("k")
            .aggregate(Sum("a"))
            .sort("k")
        )
        assert _static_schema(ds) == pa.schema(
            [pa.field("k", pa.string()), pa.field("sum(a)", pa.int64())]
        )


class TestNAry:
    def test_union(self, ray_start_regular_shared_2_cpus, parquet_path):
        # Disjoint-but-overlapping column sets exercise the schema
        # unification: ``a`` only in the first input, ``k`` only in the
        # second, ``b`` shared. The output is the merged superset.
        ds_a = ray.data.read_parquet(str(parquet_path)).select_columns(["a", "b"])
        ds_b = ray.data.read_parquet(str(parquet_path)).select_columns(["b", "k"])
        ds = ds_a.union(ds_b)
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("a", pa.int32()),
                pa.field("b", pa.float32()),
                pa.field("k", pa.string()),
            ]
        )

    def test_zip_disjoint_columns(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds_a = ray.data.read_parquet(str(parquet_path)).select_columns(["a"])
        ds_b = ray.data.read_parquet(str(parquet_path)).select_columns(["b"])
        ds = ds_a.zip(ds_b)
        assert _static_schema(ds) == pa.schema(
            [pa.field("a", pa.int32()), pa.field("b", pa.float32())]
        )

    def test_zip_overlapping_columns_get_suffixed(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        # When both inputs have a column named "a", the second input's
        # column is renamed to "a_1" to match runtime ``_zip`` behavior.
        ds_a = ray.data.read_parquet(str(parquet_path)).select_columns(["a"])
        ds_b = ray.data.read_parquet(str(parquet_path)).select_columns(["a"])
        ds = ds_a.zip(ds_b)
        assert _static_schema(ds) == pa.schema(
            [pa.field("a", pa.int32()), pa.field("a_1", pa.int32())]
        )

    def test_union_incompatible_column_types_returns_none(
        self, ray_start_regular_shared_2_cpus, tmp_path
    ):
        # Same column name, irreconcilable types: ``infer_schema`` must
        # return None (and ``Dataset.schema(fetch_if_missing=False)`` along
        # with it), not surface an ArrowTypeError from unify_schemas.
        a_path = tmp_path / "a.parquet"
        b_path = tmp_path / "b.parquet"
        pq.write_table(pa.table({"x": pa.array([1, 2], type=pa.int32())}), a_path)
        pq.write_table(
            pa.table({"x": pa.array([[1.0], [2.0]], type=pa.list_(pa.float64()))}),
            b_path,
        )
        ds = ray.data.read_parquet(str(a_path)).union(
            ray.data.read_parquet(str(b_path))
        )
        assert ds.schema(fetch_if_missing=False) is None


class TestJoin:
    def test_inner_join(self, ray_start_regular_shared_2_cpus, tmp_path):
        left = pa.table(
            {
                "k": pa.array(["a", "b", "c"]),
                "lval": pa.array([1, 2, 3], type=pa.int32()),
            }
        )
        right = pa.table(
            {
                "k": pa.array(["a", "b", "c"]),
                "rval": pa.array([10.0, 20.0, 30.0], type=pa.float32()),
            }
        )
        l_path = tmp_path / "left.parquet"
        r_path = tmp_path / "right.parquet"
        pq.write_table(left, l_path)
        pq.write_table(right, r_path)
        ds = ray.data.read_parquet(str(l_path)).join(
            ray.data.read_parquet(str(r_path)),
            on=("k",),
            join_type="inner",
            num_partitions=2,
        )
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("k", pa.string()),
                pa.field("lval", pa.int32()),
                pa.field("rval", pa.float32()),
            ]
        )


class TestUDFFallback:
    def test_map_batches_returns_none(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        # The headline guarantee inverse: UDF maps return None from
        # infer_schema, so ds.schema(fetch_if_missing=False) returns None
        # (it would only return a value via the limit(1) fallback, which
        # we've disabled).
        ds = ray.data.read_parquet(str(parquet_path)).map_batches(lambda b: b)
        assert ds.schema(fetch_if_missing=False) is None

    def test_select_after_map_batches_also_returns_none(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        # The transitive break: anything downstream of a UDF can't infer
        # schema either.
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .map_batches(lambda b: b)
            .select_columns(["a"])
        )
        assert ds.schema(fetch_if_missing=False) is None


class TestEndToEndStaticResolution:
    """The headline Phase 1 guarantee: a complex non-UDF chain resolves
    via ``Dataset.schema()`` without triggering the ``limit(1)`` fallback.

    The check is:
    - ``ds.schema(fetch_if_missing=False)`` returns the expected schema.
      Because ``fetch_if_missing=False`` disables the ``limit(1)``
      fallback in ``_base_schema``, a non-``None`` return proves that
      ``LogicalOperator.infer_schema()`` resolved the schema statically.
    """

    def test_complex_typed_chain(self, ray_start_regular_shared_2_cpus, parquet_path):
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .filter(expr=col("a") > 0)
            .select_columns(["a", "b", "k"])
            .with_column("s", col("a") + col("b"))
            .groupby("k")
            .aggregate(Sum("a"), Mean("b"))
            .sort("k")
        )
        assert _static_schema(ds) == pa.schema(
            [
                pa.field("k", pa.string()),
                pa.field("sum(a)", pa.int64()),
                pa.field("mean(b)", pa.float64()),
            ]
        )


class TestEagerStarExpansion:
    """``Project.__post_init__`` should expand ``StarExpr`` to
    explicit ``col()`` references when the input schema is known, so
    downstream optimizer rules never see ``StarExpr`` on typed chains."""

    def test_with_column_expands_star(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        ds = ray.data.read_parquet(str(parquet_path)).with_column(
            "new", col("a") + lit(10)
        )
        project = ds._logical_plan.dag
        assert isinstance(project, Project)
        # Star expanded to explicit input cols + the new aliased expr; no
        # ``StarExpr`` remains on this typed chain.
        assert_exprs_equal(
            project.exprs,
            [col("a"), col("b"), col("k"), (col("a") + lit(10)).alias("new")],
        )

    def test_rename_columns_expands_star(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        ds = ray.data.read_parquet(str(parquet_path)).rename_columns({"a": "A"})
        project = ds._logical_plan.dag
        assert isinstance(project, Project)
        # The rename AliasExpr substitutes for its source column "a" *in
        # place* (position preserved), matching runtime ``eval_projection``
        # / ``exprlist_to_fields`` ordering.
        assert_exprs_equal(project.exprs, [col("a")._rename("A"), col("b"), col("k")])

    def test_udf_chain_preserves_star(
        self, ray_start_regular_shared_2_cpus, parquet_path
    ):
        # Input schema is unknown after map_batches, so StarExpr must
        # stay for runtime ``eval_projection`` to expand per-block.
        ds = (
            ray.data.read_parquet(str(parquet_path))
            .map_batches(lambda b: b)
            .with_column("new", col("a") + lit(10))
        )
        project = ds._logical_plan.dag
        assert isinstance(project, Project)
        assert_exprs_equal(project.exprs, [star(), (col("a") + lit(10)).alias("new")])


if __name__ == "__main__":
    import sys

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