import os
import re
from typing import Any, List

import lance
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import pytest
from packaging.version import Version, parse as version_parse
from pytest_lazy_fixtures import lf as lazy_fixture

import ray
from ray.data import Dataset
from ray.data._internal.logical.operators import (
    Filter,
    Limit,
    Project,
    Repartition,
    Sort,
)
from ray.data._internal.logical.optimizers import LogicalOptimizer
from ray.data._internal.util import rows_same
from ray.data.datasource.partitioning import Partitioning
from ray.data.datasource.path_util import _unwrap_protocol
from ray.data.datatype import DataType
from ray.data.expressions import col, lit, udf
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.test_execution_optimizer_limit_pushdown import (
    _check_valid_plan_and_result,
)
from ray.data.tests.test_util import (
    get_operator_types,
    get_operators_of_type,
    plan_has_operator,
    plan_operator_comes_before,
)
from ray.tests.conftest import *  # noqa

# Pattern to match read operators in logical plans.
# Matches V1 ``Read[Read<Format>]`` or the V2 ``ListFiles → ReadFiles``
# chain where the consumer is named ``ReadFiles<Format>`` (e.g.
# ``ReadFilesParquetV2``).
READ_OPERATOR_PATTERN = (
    r"^(Read\[Read\w+\]" r"|ListFiles\[ListFiles\] -> ReadFiles\[ReadFiles\w*\])"
)


def _check_plan_with_flexible_read(
    ds: Dataset, expected_plan_suffix: str, expected_result: List[Any]
):
    """Check the logical plan with flexible read operator matching.

    This function allows flexibility in the read operator part of the plan
    by using a configurable pattern (READ_OPERATOR_PATTERN).

    Args:
        ds: The dataset to check.
        expected_plan_suffix: The expected plan after the read operator(s).
            If empty string, only the read operator is expected.
        expected_result: The expected result data.
    """
    # Optimize the logical plan before checking
    logical_plan = ds._logical_plan
    optimized_plan = LogicalOptimizer().optimize(logical_plan)
    actual_plan = optimized_plan.dag.dag_str

    match = re.match(READ_OPERATOR_PATTERN, actual_plan)
    assert match, f"Expected plan to start with read operator, got: {actual_plan}"

    # Check if there's a suffix expected
    if expected_plan_suffix:
        # The suffix should appear after the read operator
        expected_full_pattern = (
            f"{READ_OPERATOR_PATTERN} -> {re.escape(expected_plan_suffix)}"
        )
        assert re.match(expected_full_pattern, actual_plan), (
            f"Expected plan to match pattern with suffix '{expected_plan_suffix}', "
            f"got: {actual_plan}"
        )
    # If no suffix, the plan should be just the read operator
    else:
        assert actual_plan == match.group(
            1
        ), f"Expected plan to be just the read operator, got: {actual_plan}"

    # Check the result
    assert ds.take_all() == expected_result


@pytest.fixture
def parquet_ds(ray_start_regular_shared):
    """Fixture to load the Parquet dataset for testing."""
    ds = ray.data.read_parquet("example://iris.parquet")
    assert ds.count() == 150
    return ds


@pytest.fixture
def csv_ds(ray_start_regular_shared):
    """Fixture to load the CSV dataset for testing."""
    ds = ray.data.read_csv("example://iris.csv")
    assert ds.count() == 150
    return ds


def test_filter_with_udfs(parquet_ds):
    """Test filtering with UDFs where predicate pushdown does not occur."""
    filtered_udf_ds = parquet_ds.filter(lambda r: r["sepal.length"] > 5.0)
    filtered_udf_data = filtered_udf_ds.take_all()
    assert filtered_udf_ds.count() == 118
    assert all(record["sepal.length"] > 5.0 for record in filtered_udf_data)
    _check_plan_with_flexible_read(
        filtered_udf_ds,
        "Filter[Filter(<lambda>)]",  # UDF filter doesn't push down
        filtered_udf_data,
    )


def test_filter_with_expressions(parquet_ds):
    """Test filtering with expressions where predicate pushdown occurs."""
    filtered_udf_data = parquet_ds.filter(lambda r: r["sepal.length"] > 5.0).take_all()
    filtered_expr_ds = parquet_ds.filter(expr="sepal.length > 5.0")
    _check_plan_with_flexible_read(
        filtered_expr_ds,
        "",  # Pushed down to read, no additional operators
        filtered_udf_data,
    )


def test_filter_with_udf_expression_not_pushed_down(parquet_ds):
    """A UDF based filter expression must not be pushed into the datasource."""

    @udf(return_dtype=DataType.bool())
    def gt5(value: pa.Array) -> pa.Array:
        return pa.compute.greater(value, 5.0)

    expected = parquet_ds.filter(lambda r: r["sepal.length"] > 5.0).take_all()
    filtered_ds = parquet_ds.filter(expr=gt5(col("sepal.length")))

    # The whole UDF predicate stays as a Filter above the Read
    _check_plan_with_flexible_read(
        filtered_ds,
        "Filter[Filter(gt5(col('sepal.length')))]",
        expected,
    )


@pytest.mark.parametrize("udf_first", [True, False])
def test_filter_mixed_udf_and_expression(parquet_ds, udf_first):
    """Convertible conjuncts push down while the UDF stays as a Filter."""

    @udf(return_dtype=DataType.bool())
    def gt5(value: pa.Array) -> pa.Array:
        return pa.compute.greater(value, 5.0)

    udf_filter = gt5(col("sepal.length"))
    expr_filter = col("sepal.width") > lit(3.0)

    if udf_first:
        filtered_ds = parquet_ds.filter(expr=udf_filter).filter(expr=expr_filter)
    else:
        filtered_ds = parquet_ds.filter(expr=expr_filter).filter(expr=udf_filter)

    expected = parquet_ds.filter(
        lambda r: r["sepal.length"] > 5.0 and r["sepal.width"] > 3.0
    ).take_all()

    # The convertible ``sepal.width > 3.0`` conjunct is pushed into the Read
    # and only the UDF conjunct survives as a residual Filter
    _check_plan_with_flexible_read(
        filtered_ds,
        "Filter[Filter(gt5(col('sepal.length')))]",
        expected,
    )


def test_or_of_pyarrow_and_udf_not_pushed_down(parquet_ds):
    """``pyarrow_expr | udf_expr`` must stay as a Filter above the Read."""

    @udf(return_dtype=DataType.bool())
    def gt5(value: pa.Array) -> pa.Array:
        return pa.compute.greater(value, 5.0)

    or_predicate = (col("sepal.width") > 3.0) | gt5(col("sepal.length"))
    filtered_ds = parquet_ds.filter(expr=or_predicate)
    expected = parquet_ds.filter(
        lambda r: r["sepal.width"] > 3.0 or r["sepal.length"] > 5.0
    ).take_all()

    # The entire OR predicate stays as a single Filter above the Read
    _check_plan_with_flexible_read(
        filtered_ds,
        "Filter[Filter((col('sepal.width') > 3.0) | gt5(col('sepal.length')))]",
        expected,
    )


def test_and_wrapping_or_with_udf_splits(parquet_ds):
    """An AND of a convertible conjunct and a UDF containing OR splits correctly."""

    @udf(return_dtype=DataType.bool())
    def gt5(value: pa.Array) -> pa.Array:
        return pa.compute.greater(value, 5.0)

    or_conjunct = (col("sepal.width") > 3.0) | gt5(col("sepal.length"))
    predicate = or_conjunct & (col("sepal.length") > 4.0)
    filtered_ds = parquet_ds.filter(expr=predicate)
    expected = parquet_ds.filter(
        lambda r: (r["sepal.width"] > 3.0 or r["sepal.length"] > 5.0)
        and r["sepal.length"] > 4.0
    ).take_all()

    # ``sepal.length > 4.0`` is pushed into the Read and only the OR conjunct
    # survives as a residual Filter
    _check_plan_with_flexible_read(
        filtered_ds,
        "Filter[Filter((col('sepal.width') > 3.0) | gt5(col('sepal.length')))]",
        expected,
    )


def test_not_udf_not_pushed_down(parquet_ds):
    """A negation of a UDF must stay as a Filter."""

    @udf(return_dtype=DataType.bool())
    def gt5(value: pa.Array) -> pa.Array:
        return pa.compute.greater(value, 5.0)

    filtered_ds = parquet_ds.filter(expr=~gt5(col("sepal.length")))
    expected = parquet_ds.filter(lambda r: r["sepal.length"] <= 5.0).take_all()

    # The negated UDF stays as a Filter above the Read
    _check_plan_with_flexible_read(
        filtered_ds,
        "Filter[Filter(~gt5(col('sepal.length')))]",
        expected,
    )


def test_filter_pushdown_source_and_op(ray_start_regular_shared):
    """Test filtering when expressions are provided both in source and operator."""
    filter_expr = "sepal.width > 3.0"

    ds = (
        ray.data.read_parquet("example://iris.parquet")
        .filter(expr=col("sepal.length") > lit(5.0))
        .filter(expr=filter_expr)
    )
    result = ds.take_all()
    assert all(r["sepal.length"] > 5.0 and r["sepal.width"] > 3.0 for r in result)
    _check_plan_with_flexible_read(
        ds,
        "",  # Both filters pushed down to read
        result,
    )


def test_chained_filter_with_expressions(parquet_ds):
    """Test chained filtering with expressions where combined pushdown occurs."""
    filtered_expr_chained_ds = (
        parquet_ds.filter(expr=col("sepal.length") > 1.0)
        .filter(expr=col("sepal.length") > 2.0)
        .filter(expr=col("sepal.length") > 3.0)
        .filter(expr=col("sepal.length") > 3.0)
        .filter(expr=col("sepal.length") > 5.0)
    )
    filtered_udf_data = parquet_ds.filter(lambda r: r["sepal.length"] > 5.0).take_all()
    _check_plan_with_flexible_read(
        filtered_expr_chained_ds,
        "",  # All filters combined and pushed down to read
        filtered_udf_data,
    )


@pytest.mark.parametrize(
    "fs,data_path",
    [
        (None, lazy_fixture("local_path")),
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        # NOTE: an ``s3_fs`` parametrization was previously listed here, but
        # it didn't actually exercise S3 — the test uses ``_unwrap_protocol``
        # to derive a local FS path, which on the moto-mocked ``s3_path``
        # fixture returns a *relative* path (the fixture strips the leading
        # ``/`` so the first segment becomes a moto bucket name). The
        # resulting ``lance.write_dataset(<relative-path>)`` wrote into the
        # current working directory, polluting the repo on every run.
        # If S3 lance pushdown ever needs end-to-end coverage, add it
        # separately and pass an actual ``s3://`` URI plus storage options
        # through to ``lance.write_dataset``/``ray.data.read_lance``.
    ],
)
# Same pylance version gate as tests/datasource/test_lance.py
@pytest.mark.skipif(
    Version(lance.__version__) <= Version("0.3.19"),
    reason=f"pylance {lance.__version__} <= 0.3.19; API incompatible",
)
def test_pushdown_filter_lance(ray_start_regular_shared, fs, data_path):
    """Test that Lance predicate pushdown absorbs expression filters into Read."""

    df1 = pa.table({"a": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
    setup_data_path = _unwrap_protocol(data_path)
    path = os.path.join(setup_data_path, "test.lance")
    lance.write_dataset(df1, path)
    # Both filters specified on read_lance() and .filter() should be applied
    lance_ds = ray.data.read_lance(path, filter="a <= 5")
    filtered_expr_ds = lance_ds.filter(expr=col("a") >= 1.0)

    filtered_expr_data = lance_ds.filter(
        lambda r: r["a"] <= 5.0 and r["a"] >= 1.0
    ).take_all()
    _check_plan_with_flexible_read(
        filtered_expr_ds,
        "",  # Pushed down to read, no additional Filter operator
        filtered_expr_data,
    )


@pytest.mark.parametrize(
    "filter_fn,expected_suffix",
    [
        (
            lambda ds: ds.filter(lambda r: r["sepal.length"] > 5.0),
            "Filter[Filter(<lambda>)]",  # UDF filter doesn't push down
        ),
        (
            lambda ds: ds.filter(expr=col("sepal.length") > 5.0),
            "Filter[Filter(col('sepal.length') > 5.0)]",  # CSV doesn't support predicate pushdown
        ),
    ],
)
def test_filter_pushdown_csv(csv_ds, filter_fn, expected_suffix):
    """Test filtering on CSV files (CSV doesn't support predicate pushdown)."""
    filtered_ds = filter_fn(csv_ds)
    filtered_data = filtered_ds.take_all()
    assert filtered_ds.count() == 118
    assert all(record["sepal.length"] > 5.0 for record in filtered_data)
    _check_plan_with_flexible_read(
        filtered_ds,
        expected_suffix,
        filtered_data,
    )


def test_filter_mixed(csv_ds):
    """Test that mixed function and expressions work (CSV doesn't support predicate pushdown)."""
    csv_ds = csv_ds.filter(lambda r: r["sepal.length"] < 5.0)
    csv_ds = csv_ds.filter(expr="sepal.length > 3.0")
    csv_ds = csv_ds.filter(expr="sepal.length > 4.0")
    csv_ds = csv_ds.map(lambda x: x)
    csv_ds = csv_ds.filter(expr="sepal.length > 2.0")
    csv_ds = csv_ds.filter(expr="sepal.length > 1.0")
    filtered_expr_data = csv_ds.take_all()
    assert csv_ds.count() == 22
    assert all(record["sepal.length"] < 5.0 for record in filtered_expr_data)
    assert all(record["sepal.length"] > 4.0 for record in filtered_expr_data)
    # After optimization: expression filters before map get fused, expression filters after map get fused
    # CSV doesn't support predicate pushdown, so filters stay after Read
    _check_plan_with_flexible_read(
        csv_ds,
        "Filter[Filter(<lambda>)] -> Filter[Filter((col('sepal.length') > 4.0) & (col('sepal.length') > 3.0))] -> "
        "MapRows[Map(<lambda>)] -> Filter[Filter((col('sepal.length') > 1.0) & (col('sepal.length') > 2.0))]",
        filtered_expr_data,
    )


def test_filter_mixed_expression_first_parquet(ray_start_regular_shared):
    """Test that mixed functional and expressions work with Parquet (supports predicate pushdown)."""
    ds = ray.data.read_parquet("example://iris.parquet")
    ds = ds.filter(expr="sepal.length > 3.0")
    ds = ds.filter(expr="sepal.length > 4.0")
    ds = ds.filter(lambda r: r["sepal.length"] < 5.0)
    filtered_expr_data = ds.take_all()
    assert ds.count() == 22
    assert all(record["sepal.length"] < 5.0 for record in filtered_expr_data)
    assert all(record["sepal.length"] > 4.0 for record in filtered_expr_data)
    _check_plan_with_flexible_read(
        ds,
        "Filter[Filter(<lambda>)]",  # Expressions pushed down, UDF remains
        filtered_expr_data,
    )


def test_filter_mixed_expression_first_csv(ray_start_regular_shared):
    """Test that mixed functional and expressions work with CSV (doesn't support predicate pushdown)."""
    ds = ray.data.read_csv("example://iris.csv")
    ds = ds.filter(expr="sepal.length > 3.0")
    ds = ds.filter(expr="sepal.length > 4.0")
    ds = ds.filter(lambda r: r["sepal.length"] < 5.0)
    filtered_expr_data = ds.take_all()
    assert ds.count() == 22
    assert all(record["sepal.length"] < 5.0 for record in filtered_expr_data)
    assert all(record["sepal.length"] > 4.0 for record in filtered_expr_data)
    # CSV doesn't support predicate pushdown, so expression filters get fused but not pushed down
    _check_plan_with_flexible_read(
        ds,
        "Filter[Filter((col('sepal.length') > 4.0) & (col('sepal.length') > 3.0))] -> Filter[Filter(<lambda>)]",
        filtered_expr_data,
    )


def test_filter_mixed_expression_not_readfiles(ray_start_regular_shared):
    """Test that mixed functional and expressions work."""
    ds = ray.data.range(100).filter(expr="id > 1.0")
    ds = ds.filter(expr="id > 2.0")
    ds = ds.filter(lambda r: r["id"] < 5.0)
    filtered_expr_data = ds.take_all()
    assert ds.count() == 2
    assert all(record["id"] < 5.0 for record in filtered_expr_data)
    assert all(record["id"] > 2.0 for record in filtered_expr_data)
    _check_valid_plan_and_result(
        ds,
        "Read[ReadRange] -> Filter[Filter((col('id') > 2.0) & (col('id') > 1.0))] -> "
        "Filter[Filter(<lambda>)]",
        filtered_expr_data,
    )


@pytest.mark.parametrize(
    "operations,output_rename_map,expected_filter_expr,test_id",
    [
        (
            # rename("sepal.length" -> a).filter(a)
            lambda ds: ds.rename_columns({"sepal.length": "a"}).filter(
                expr=col("a") > 2.0
            ),
            {"a": "sepal.length"},
            col("sepal.length") > 2.0,
            "rename_filter",
        ),
        (
            # rename("sepal.length" -> a).filter(a).rename(a -> b)
            lambda ds: (
                ds.rename_columns({"sepal.length": "a"})
                .filter(expr=col("a") > 2.0)
                .rename_columns({"a": "b"})
            ),
            {"b": "sepal.length"},
            col("sepal.length") > 2.0,
            "rename_filter_rename",
        ),
        (
            # rename("sepal.length" -> a).filter(a).rename(a -> b).filter(b)
            lambda ds: (
                ds.rename_columns({"sepal.length": "a"})
                .filter(expr=col("a") > 2.0)
                .rename_columns({"a": "b"})
                .filter(expr=col("b") < 5.0)
            ),
            {"b": "sepal.length"},
            (col("sepal.length") > 2.0) & (col("sepal.length") < 5.0),
            "rename_filter_rename_filter",
        ),
        (
            # rename("sepal.length" -> a).filter(a).rename(a -> b).filter(b).rename("sepal.width" -> a)
            # Here column a is referred multiple times in rename
            lambda ds: (
                ds.rename_columns({"sepal.length": "a"})
                .filter(expr=col("a") > 2.0)
                .rename_columns({"a": "b"})
                .filter(expr=col("b") < 5.0)
                .rename_columns({"sepal.width": "a"})
            ),
            {"b": "sepal.length", "a": "sepal.width"},
            (col("sepal.length") > 2.0) & (col("sepal.length") < 5.0),
            "rename_filter_rename_filter_rename",
        ),
    ],
    ids=lambda x: x if isinstance(x, str) else "",
)
def test_pushdown_with_rename_and_filter(
    ray_start_regular_shared,
    operations,
    output_rename_map,
    expected_filter_expr,
    test_id,
):
    """Test predicate pushdown with various combinations of rename and filter operations."""
    path = "example://iris.parquet"
    ds = operations(ray.data.read_parquet(path))
    result = ds.take_all()

    # Filters are pushed into the scan; renames stay as a ``Project`` of
    # ``AliasExpr``s above the pruned scan.
    _check_plan_with_flexible_read(ds, "Project[Project]", result)

    ds1 = ray.data.read_parquet(path).filter(expr=expected_filter_expr)
    # Convert to pandas to ensure both datasets are fully executed
    df = ds.to_pandas().rename(columns=output_rename_map)
    df1 = ds1.to_pandas()
    assert len(df) == len(df1), f"Expected {len(df)} rows, got {len(df1)} rows"
    pd.testing.assert_frame_equal(df, df1)


def _get_optimized_plan(ds: Dataset) -> str:
    """Get the optimized logical plan as a string."""
    logical_plan = ds._logical_plan
    optimized_plan = LogicalOptimizer().optimize(logical_plan)
    return optimized_plan.dag.dag_str


def _check_plan_matches_pattern(ds: Dataset, expected_pattern: str):
    """Check that the optimized plan matches the expected regex pattern."""
    actual_plan = _get_optimized_plan(ds)
    assert re.match(expected_pattern, actual_plan), (
        f"Plan mismatch:\n"
        f"Expected pattern: {expected_pattern}\n"
        f"Actual plan: {actual_plan}"
    )


class TestPredicatePushdownIntoRead:
    """Tests for pushing predicates into Read operators.

    When a data source supports predicate pushdown (like Parquet),
    the filter should be absorbed into the Read operator itself.
    """

    @pytest.fixture
    def parquet_ds(self, ray_start_regular_shared):
        return ray.data.read_parquet("example://iris.parquet")

    def test_complex_pipeline_all_filters_push_to_read(self, parquet_ds):
        """Complex pipeline: filters should push through all operators into Read.

        Pipeline: Read -> Filter -> Rename -> Filter -> Sort -> Repartition
                  -> Filter -> Limit -> Filter

        All filters should fuse, push through all operators, rebind through rename,
        and be absorbed into the Read operator.
        """
        ds = (
            parquet_ds.filter(expr=col("sepal.length") > 4.0)
            .rename_columns({"sepal.length": "len", "sepal.width": "width"})
            .filter(expr=col("len") < 7.0)
            .sort("len")
            .repartition(3)
            .filter(expr=col("width") > 2.5)
            .limit(100)
            .filter(expr=col("len") > 4.5)
        )

        # Verify correctness: should apply all filters correctly
        expected = (
            parquet_ds.filter(
                expr=(col("sepal.length") > 4.0)
                & (col("sepal.length") < 7.0)
                & (col("sepal.width") > 2.5)
                & (col("sepal.length") > 4.5)
            )
            .rename_columns({"sepal.length": "len", "sepal.width": "width"})
            .sort("len")
            .repartition(3)
            .limit(100)
        )

        assert rows_same(ds.to_pandas(), expected.to_pandas())

        # Verify plan: all filters pushed into Read, passthrough ops remain
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "No Filter operators should remain after pushdown into Read"


class TestPassthroughBehavior:
    """Tests for PASSTHROUGH behavior operators.

    Operators: Sort, Repartition, RandomShuffle, Limit
    Predicates pass through unchanged - operators don't affect filtering.
    """

    @pytest.fixture
    def base_ds(self, ray_start_regular_shared):
        return ray.data.range(100)

    @pytest.mark.parametrize(
        "transform,expected_op_type",
        [
            (lambda ds: ds.sort("id"), "Sort"),
            (lambda ds: ds.repartition(10), "Repartition"),
            (
                lambda ds: ds.repartition(target_num_rows_per_block=10),
                "StreamingRepartition",
            ),
            (lambda ds: ds.random_shuffle(), "RandomShuffle"),
            (lambda ds: ds.limit(50), "Limit"),
        ],
        ids=["sort", "repartition", "streaming_repartition", "random_shuffle", "limit"],
    )
    def test_filter_pushes_through_operator(self, base_ds, transform, expected_op_type):
        """Filter should push through passthrough operators."""
        ds = transform(base_ds).filter(expr=col("id") < 10)

        # Verify correctness against expected result
        expected = base_ds.filter(expr=col("id") < 10)
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        # Filter pushed down, operator remains
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert plan_has_operator(
            optimized_plan, Filter
        ), "Filter should exist after pushdown"

        # Verify the passthrough operator is still present
        op_types = get_operator_types(optimized_plan)
        assert expected_op_type in op_types, f"{expected_op_type} should remain in plan"

    def test_filter_pushes_through_multiple_ops(self, base_ds):
        """Filter should push through multiple passthrough operators."""
        ds = base_ds.sort("id").repartition(5).limit(50).filter(expr=col("id") < 10)

        # Verify correctness against expected result
        expected = base_ds.filter(expr=col("id") < 10)
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        # Verify plan: filter pushed down, all operators remain
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert plan_has_operator(optimized_plan, Filter), "Filter should exist"
        assert plan_has_operator(optimized_plan, Sort), "Sort should remain"
        assert plan_has_operator(
            optimized_plan, Repartition
        ), "Repartition should remain"
        assert plan_has_operator(optimized_plan, Limit), "Limit should remain"

    def test_multiple_filters_fuse_and_push_through(self, base_ds):
        """Multiple filters should fuse and push through passthrough operators."""
        ds = base_ds.filter(expr=col("id") > 5).sort("id").filter(expr=col("id") < 20)

        # Verify correctness against expected result
        expected = base_ds.filter(expr=(col("id") > 5) & (col("id") < 20))
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        # Verify plan: filters fused and pushed, Sort remains
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        filters = get_operators_of_type(optimized_plan, Filter)
        assert len(filters) == 1, "Multiple filters should be fused into one"
        assert plan_has_operator(optimized_plan, Sort), "Sort should remain"
        assert plan_operator_comes_before(
            optimized_plan, Filter, Sort
        ), "Fused filter should come before Sort"


class TestPassthroughWithSubstitutionBehavior:
    """Tests for PASSTHROUGH_WITH_SUBSTITUTION behavior operators.

    Operator: Project (used by rename_columns, select, with_column)
    Predicates push through but column names must be rebound.
    """

    @pytest.fixture
    def parquet_ds(self, ray_start_regular_shared):
        return ray.data.read_parquet("example://iris.parquet")

    def test_simple_rename_with_filter(self, parquet_ds):
        """Filter after rename should rebind columns and push down."""
        ds = parquet_ds.rename_columns({"sepal.length": "len"}).filter(
            expr=col("len") > 5.0
        )

        # Verify correctness against expected result
        expected = parquet_ds.filter(expr=col("sepal.length") > 5.0).rename_columns(
            {"sepal.length": "len"}
        )
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        # Filter rebound and pushed to Read (no Filter operators should remain)
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "Filter should be pushed into Read, no Filter operators should remain"

    def test_chained_renames_with_filter(self, parquet_ds):
        """Multiple renames should track through filter pushdown."""
        ds = (
            parquet_ds.rename_columns({"sepal.length": "a"})
            .rename_columns({"a": "b"})
            .filter(expr=col("b") > 5.0)
        )

        # Verify correctness against expected result
        expected = (
            parquet_ds.filter(expr=col("sepal.length") > 5.0)
            .rename_columns({"sepal.length": "a"})
            .rename_columns({"a": "b"})
        )
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        # Filter should be pushed into Read after column rebinding
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "Filter should be pushed into Read after rebinding through renames"

    def test_multiple_filters_with_renames(self, parquet_ds):
        """Multiple filters with renames should all rebind and push."""
        ds = (
            parquet_ds.rename_columns({"sepal.length": "a"})
            .filter(expr=col("a") > 2.0)
            .rename_columns({"a": "b"})
            .filter(expr=col("b") < 5.0)
        )

        # Verify correctness against expected result
        expected = (
            parquet_ds.filter(
                expr=(col("sepal.length") > 2.0) & (col("sepal.length") < 5.0)
            )
            .rename_columns({"sepal.length": "a"})
            .rename_columns({"a": "b"})
        )
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        # Multiple filters should be fused, rebound, and pushed into Read
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "All filters should be fused, rebound, and pushed into Read"

    def test_rename_with_partition_residual_filter(
        self, ray_start_regular_shared, tmp_path
    ):
        """Residual Filter ends up below a rename Project, in original names.

        When a predicate mixes partition and data columns under OR, the
        unsplittable part is wrapped in a Filter above ReadFiles by
        ``ReadFiles.apply_predicate``. The read stage never renames
        columns (renaming is always carried by a ``Project`` above the
        read), so the residual Filter — which sits between the rename
        ``Project`` and ``ReadFiles`` after predicate pushdown — must
        reference the original on-disk column names that the scanner
        produces, not the renamed ones the user wrote.
        """
        table = pa.table(
            {
                "partition_col": [1, 1, 2, 2, 3, 3],
                "data1": [10, 20, 30, 40, 50, 60],
                "data2": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            }
        )
        pq.write_to_dataset(
            table, root_path=str(tmp_path), partition_cols=["partition_col"]
        )

        ds = (
            ray.data.read_parquet(
                str(tmp_path),
                partitioning=Partitioning("hive", field_types={"partition_col": int}),
            )
            .rename_columns({"data1": "D1", "data2": "D2"})
            .filter(
                expr=((col("D1") > 25) | (col("partition_col") == 1))
                & (col("D2") < 5.5)
            )
        )

        # Equivalent plan without the rename, to validate row-level correctness.
        expected = ray.data.read_parquet(
            str(tmp_path),
            partitioning=Partitioning("hive", field_types={"partition_col": int}),
        ).filter(
            expr=((col("data1") > 25) | (col("partition_col") == 1))
            & (col("data2") < 5.5)
        )

        # The dataset must execute without binding errors and produce the
        # expected rows (with renamed column names).
        result = ds.to_pandas().rename(columns={"D1": "data1", "D2": "data2"})
        assert rows_same(result, expected.to_pandas())

        # A residual Filter should remain below the rename ``Project``,
        # and its predicate must reference the original on-disk column
        # names (``data1``), not the renamed ones the user wrote.
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        residual_filters = get_operators_of_type(optimized_plan, Filter)
        assert (
            len(residual_filters) == 1
        ), f"Expected one residual Filter, got plan: {optimized_plan.dag.dag_str}"
        residual_expr_str = str(residual_filters[0].predicate_expr)
        assert "data1" in residual_expr_str and "D1" not in residual_expr_str, (
            f"Residual Filter predicate should reference original column 'data1', "
            f"got: {residual_expr_str}"
        )


class TestProjectionWithFilterEdgeCases:
    """Tests for edge cases with select_columns and with_column followed by filters.

    These tests verify that filters correctly handle:
    - Columns that are kept by select (should push through)
    - Columns that are removed by select (should NOT push through)
    - Computed columns from with_column (should NOT push through)
    """

    @pytest.fixture
    def base_ds(self, ray_start_regular_shared):
        return ray.data.from_items(
            [
                {"a": 1, "b": 2, "c": 3},
                {"a": 2, "b": 5, "c": 8},
                {"a": 3, "b": 6, "c": 9},
            ]
        )

    def test_select_then_filter_on_selected_column(self, base_ds):
        """Filter on selected column should push through select."""
        ds = base_ds.select_columns(["a", "b"]).filter(expr=col("a") > 1)

        # Verify correctness
        result_df = ds.to_pandas()
        expected_df = pd.DataFrame(
            [
                {"a": 2, "b": 5},
                {"a": 3, "b": 6},
            ]
        )
        # Sort columns before comparison
        result_df = result_df[sorted(result_df.columns)]
        expected_df = expected_df[sorted(expected_df.columns)]
        assert rows_same(result_df, expected_df)

        # Verify plan: filter pushed through select
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert plan_operator_comes_before(
            optimized_plan, Filter, Project
        ), "Filter should be pushed before Project"

    def test_select_then_filter_on_removed_column(self, base_ds):
        """Filter on removed column should fail, not push through."""
        ds = base_ds.select_columns(["a"])

        with pytest.raises((KeyError, ray.exceptions.RayTaskError)):
            ds.filter(expr=col("b") == 2).take_all()

    def test_with_column_then_filter_on_computed_column(self, base_ds):
        """Filter on computed column should not push through."""

        from ray.data.expressions import lit

        ds = base_ds.with_column("d", lit(4)).filter(expr=col("d") == 4)

        # Verify correctness - all rows should pass (d is always 4)
        result_df = ds.to_pandas()
        expected_df = pd.DataFrame(
            [
                {"a": 1, "b": 2, "c": 3, "d": 4},
                {"a": 2, "b": 5, "c": 8, "d": 4},
                {"a": 3, "b": 6, "c": 9, "d": 4},
            ]
        )
        # Sort columns before comparison
        result_df = result_df[sorted(result_df.columns)]
        expected_df = expected_df[sorted(expected_df.columns)]
        assert rows_same(result_df, expected_df)

        # Verify plan: filter should NOT push through (stays after with_column)
        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert plan_has_operator(
            optimized_plan, Filter
        ), "Filter should remain (not pushed through)"

    def test_rename_then_filter_on_old_column_name(self, base_ds):
        """Filter using old column name after rename should fail."""
        ds = base_ds.rename_columns({"b": "B"})

        with pytest.raises((KeyError, ray.exceptions.RayTaskError)):
            ds.filter(expr=col("b") == 2).take_all()

    @pytest.mark.parametrize(
        "ds_factory,rename_map,filter_col,filter_value,expected_rows",
        [
            # In-memory dataset: rename a->b, b->b_old
            (
                lambda: ray.data.from_items(
                    [
                        {"a": 1, "b": 2, "c": 3},
                        {"a": 2, "b": 5, "c": 8},
                        {"a": 3, "b": 6, "c": 9},
                    ]
                ),
                {"a": "b", "b": "b_old"},
                "b",
                1,
                [{"b": 2, "b_old": 5, "c": 8}, {"b": 3, "b_old": 6, "c": 9}],
            ),
            # Parquet dataset: rename sepal.length->sepal.width, sepal.width->old_width
            (
                lambda: ray.data.read_parquet("example://iris.parquet"),
                {"sepal.length": "sepal.width", "sepal.width": "old_width"},
                "sepal.width",
                5.0,
                None,  # Will verify via alternative computation
            ),
        ],
        ids=["in_memory", "parquet"],
    )
    def test_rename_chain_with_name_reuse(
        self,
        ray_start_regular_shared,
        ds_factory,
        rename_map,
        filter_col,
        filter_value,
        expected_rows,
    ):
        """Test rename chains where an output name matches another rename's input name.

        This tests the fix for a bug where rename(a->b, b->c) followed by filter(b>5)
        would incorrectly block pushdown, even though 'b' is a valid output column
        (created by a->b).

        Example: rename({'a': 'b', 'b': 'temp'}) creates 'b' from 'a' and 'temp' from 'b'.
        A filter on 'b' should be able to push through.
        """
        ds = ds_factory()

        # Apply rename and filter
        ds_renamed_filtered = ds.rename_columns(rename_map).filter(
            expr=col(filter_col) > filter_value
        )

        # Verify correctness
        if expected_rows is not None:
            # For in-memory, compare against expected rows
            result_df = ds_renamed_filtered.to_pandas()
            expected_df = pd.DataFrame(expected_rows)
            result_df = result_df[sorted(result_df.columns)]
            expected_df = expected_df[sorted(expected_df.columns)]
            assert rows_same(result_df, expected_df)
        else:
            # For parquet, compare against alternative computation
            # Filter on original column, then rename
            original_col = next(k for k, v in rename_map.items() if v == filter_col)
            expected = ds.filter(expr=col(original_col) > filter_value).rename_columns(
                rename_map
            )
            assert rows_same(ds_renamed_filtered.to_pandas(), expected.to_pandas())

        # Verify plan optimization
        optimized_plan = LogicalOptimizer().optimize(ds_renamed_filtered._logical_plan)

        # Determine if the data source supports predicate pushdown by checking
        # if the filter was completely eliminated (pushed into the read operator)
        has_filter = plan_has_operator(optimized_plan, Filter)
        has_project = plan_has_operator(optimized_plan, Project)

        # Three valid post-optimization shapes:
        #   1. ``has_filter=False, has_project=False`` — both pushed into a
        #      legacy Read (rare; happens when neither rename nor filter
        #      survives optimization).
        #   2. ``has_filter=False, has_project=True`` - file-based reads
        #      can push the filter into the scan and leave the rename
        #      ``Project`` above it.
        #   3. ``has_filter=True, has_project=True`` — source doesn't
        #      support predicate pushdown (e.g. in-memory); filter at
        #      least pushed below the rename ``Project``.
        if not has_filter and not has_project:
            # Filter was pushed into Read - this is the optimal case
            pass
        elif not has_filter and has_project:
            pass
        elif has_filter and has_project:
            # For in-memory datasets, filter should at least push through projection
            assert plan_operator_comes_before(
                optimized_plan, Filter, Project
            ), "Filter should be pushed before Project after rebinding through rename chain"
        else:
            # Unexpected state - either filter or project but not both
            raise AssertionError(
                f"Unexpected optimization state: has_filter={has_filter}, has_project={has_project}"
            )


class TestPyArrowComputeUDFPushdown:
    """Tests for predicate pushdown of PyArrow-compute-backed UDF expressions.

    String namespace methods (str.match_regex, str.starts_with, etc.) and
    numeric helpers (ceil, abs, etc.) produce PyArrowComputeUDFExpr nodes
    that should be pushed into Read operators just like plain comparison
    expressions.
    """

    @pytest.fixture
    def parquet_ds(self, ray_start_regular_shared):
        return ray.data.read_parquet("example://iris.parquet")

    @pytest.mark.parametrize(
        "build_filter,equivalent_fn",
        [
            pytest.param(
                lambda: ~col("variety").str.match_regex("Set.*"),
                lambda r: not bool(re.search("Set.*", r["variety"])),
                id="negated_match_regex",
            ),
            pytest.param(
                lambda: col("variety").str.starts_with("Vir"),
                lambda r: r["variety"].startswith("Vir"),
                id="starts_with",
            ),
            pytest.param(
                lambda: col("variety").str.contains("set"),
                lambda r: "set" in r["variety"],
                id="contains",
            ),
        ],
    )
    @pytest.mark.skipif(
        version_parse(pa.__version__) < version_parse("15.0.0"),
        reason="Requires PyArrow >= 15 for string compute UDF pushdown",
    )
    def test_string_udf_pushdown_into_parquet(
        self, parquet_ds, build_filter, equivalent_fn
    ):
        """String UDF predicates should be pushed into the Parquet reader."""
        ds = parquet_ds.filter(expr=build_filter())
        expected = parquet_ds.filter(fn=equivalent_fn)
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "PyArrow-compute UDF filter should be pushed into Read"

    @pytest.mark.skipif(
        version_parse(pa.__version__) < version_parse("15.0.0"),
        reason="Requires PyArrow >= 15 for string compute UDF pushdown",
    )
    def test_udf_combined_with_comparison_pushdown(self, parquet_ds):
        """UDF predicate combined with comparison should both push down."""
        ds = parquet_ds.filter(
            expr=col("variety").str.starts_with("Vir") & (col("sepal.length") > 6.0)
        )
        expected = parquet_ds.filter(
            fn=lambda r: r["variety"].startswith("Vir") and r["sepal.length"] > 6.0
        )
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "Combined UDF + comparison filter should be pushed into Read"

    @pytest.mark.skipif(
        version_parse(pa.__version__) < version_parse("15.0.0"),
        reason="Requires PyArrow >= 15 for string compute UDF pushdown",
    )
    def test_chained_udf_filters_fuse_and_push(self, parquet_ds):
        """Multiple UDF filters should fuse and push into Read."""
        ds = parquet_ds.filter(expr=col("variety").str.contains("set")).filter(
            expr=col("sepal.length") > 5.0
        )
        expected = parquet_ds.filter(
            fn=lambda r: "set" in r["variety"] and r["sepal.length"] > 5.0
        )
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "Chained UDF filters should fuse and push into Read"

    @pytest.mark.skipif(
        version_parse(pa.__version__) < version_parse("15.0.0"),
        reason="Requires PyArrow >= 15 for complex type UDF pushdown",
    )
    @pytest.mark.parametrize(
        "table,build_filter,equivalent_fn",
        [
            pytest.param(
                pa.table(
                    {
                        "tags": [[1, 2, 3], [4, 5], [6], [7, 8, 9, 10]],
                        "value": [10, 20, 30, 40],
                    }
                ),
                lambda: col("tags").list.len() > 2,
                lambda r: len(r["tags"]) > 2,
                id="list_len",
            ),
            pytest.param(
                pa.table(
                    {
                        "info": pa.array(
                            [
                                {"name": "Alice", "age": 30},
                                {"name": "Bob", "age": 25},
                                {"name": "Carol", "age": 35},
                                {"name": "Dave", "age": 20},
                            ],
                            type=pa.struct(
                                [
                                    pa.field("name", pa.string()),
                                    pa.field("age", pa.int64()),
                                ]
                            ),
                        ),
                        "id": [1, 2, 3, 4],
                    }
                ),
                lambda: col("info").struct.field("age") > 25,
                lambda r: r["info"]["age"] > 25,
                id="struct_field",
            ),
        ],
    )
    def test_complex_type_udf_pushdown(
        self, ray_start_regular_shared, tmp_path, table, build_filter, equivalent_fn
    ):
        """List/struct UDF predicates should be pushed into the Parquet reader."""
        path = str(tmp_path / "data.parquet")
        pq.write_table(table, path)

        ds = ray.data.read_parquet(path).filter(expr=build_filter())
        expected = ray.data.read_parquet(path).filter(fn=equivalent_fn)
        assert rows_same(ds.to_pandas(), expected.to_pandas())

        optimized_plan = LogicalOptimizer().optimize(ds._logical_plan)
        assert not plan_has_operator(
            optimized_plan, Filter
        ), "Complex-type UDF filter should be pushed into Read"


class TestPushIntoBranchesBehavior:
    """Tests for PUSH_INTO_BRANCHES behavior operators.

    Operator: Union
    Predicates are duplicated and pushed into each branch.
    """

    def test_simple_union_with_filter(self, ray_start_regular_shared):
        """Filter after union should push into both branches."""
        ds1 = ray.data.range(100, parallelism=2)
        ds2 = ray.data.range(100, parallelism=2)
        ds = ds1.union(ds2).filter(expr=col("id") >= 50)

        # Verify correctness: should have duplicates from both branches
        base = ray.data.range(100)
        expected = base.filter(expr=col("id") >= 50).union(
            base.filter(expr=col("id") >= 50)
        )
        assert rows_same(ds.to_pandas(), expected.to_pandas())

    def test_multiple_unions_with_filter(self, ray_start_regular_shared):
        """Filter should push into all branches of multiple unions."""
        ds1 = ray.data.read_parquet("example://iris.parquet")
        ds2 = ray.data.read_parquet("example://iris.parquet")
        ds3 = ray.data.read_parquet("example://iris.parquet")
        ds = ds1.union(ds2).union(ds3).filter(expr=col("sepal.length") > 5.0)

        # Verify correctness: should have 3x the filtered results
        single_filtered = ray.data.read_parquet("example://iris.parquet").filter(
            expr=col("sepal.length") > 5.0
        )
        expected = single_filtered.union(single_filtered).union(single_filtered)
        assert rows_same(ds.to_pandas(), expected.to_pandas())

    def test_branch_filters_plus_union_filter(self, ray_start_regular_shared):
        """Individual branch filters plus union filter should all push."""
        parquet_ds = ray.data.read_parquet("example://iris.parquet")
        ds1 = parquet_ds.filter(expr=col("sepal.width") > 2.0)
        ds2 = parquet_ds.filter(expr=col("sepal.width") > 2.0)
        ds = ds1.union(ds2).filter(expr=col("sepal.length") < 5.0)

        # Verify correctness: both filters applied
        expected_single = parquet_ds.filter(
            expr=(col("sepal.width") > 2.0) & (col("sepal.length") < 5.0)
        )
        expected = expected_single.union(expected_single)
        assert rows_same(ds.to_pandas(), expected.to_pandas())


if __name__ == "__main__":
    import sys

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