################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
# limitations under the License.
################################################################################

"""Tests for _read_table() — reads from a local Paimon filesystem catalog.

All tests run without Docker or external services; pypaimon is used directly
to create and populate test tables, then _read_table() is validated.
"""

from __future__ import annotations

import asyncio
import decimal

import pyarrow as pa
import pytest

pypaimon = pytest.importorskip("pypaimon")
daft = pytest.importorskip("daft")

from daft import col, lit

from pypaimon.daft.daft_paimon import _read_table
from pypaimon.daft.daft_predicate_visitor import convert_filters_to_paimon


def _contains_expr(py_expr):
    from daft.expressions import Expression

    expr_text = str(Expression._from_pyexpr(py_expr))
    return "contains" in expr_text


def _predicate_leaves(predicate):
    if predicate is None:
        return []
    if predicate.method in ("and", "or"):
        result = []
        for child in predicate.literals:
            result.extend(_predicate_leaves(child))
        return result
    return [(predicate.method, predicate.field, tuple(predicate.literals or []))]


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------


class _UnserializableFileIoMarker:
    def __reduce__(self):
        raise TypeError("file io marker should not be serialized")


class _UnserializableStorageConfig:
    multithreaded_io = False

    def __reduce__(self):
        raise TypeError("storage config marker should not be serialized")


def _write_to_paimon(table, arrow_table, mode="append", overwrite_partition=None):
    write_builder = table.new_batch_write_builder()
    if mode == "overwrite":
        write_builder.overwrite(overwrite_partition or {})
    table_write = write_builder.new_write()
    table_commit = write_builder.new_commit()
    try:
        table_write.write_arrow(arrow_table)
        commit_messages = table_write.prepare_commit()
        table_commit.commit(commit_messages)
    finally:
        table_write.close()
        table_commit.close()


async def _collect_paimon_source_batches(source, pushdowns):
    batches = []
    fallback_task_count = 0
    async for task in source.get_tasks(pushdowns):
        if type(task).__name__ == "_PaimonPKSplitTask":
            fallback_task_count += 1
        async for batch in task.read():
            batches.append(batch.to_pydict())
    assert fallback_task_count > 0
    return batches


async def _read_paimon_source_batches(
    table,
    filter_expr=None,
    columns=None,
    limit=None,
):
    from daft import context, runners
    from daft.daft import StorageConfig
    from daft.io.pushdowns import Pushdowns

    from pypaimon.daft.daft_datasource import PaimonDataSource

    io_config = context.get_context().daft_planning_config.default_io_config
    storage_config = StorageConfig(runners.get_or_create_runner().name != "ray", io_config)
    source = PaimonDataSource(table, storage_config=storage_config, catalog_options={})

    pushdowns = Pushdowns(filters=filter_expr, columns=columns, limit=limit)
    return await _collect_paimon_source_batches(source, pushdowns)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture(scope="function")
def local_paimon_catalog(tmp_path):
    catalog = pypaimon.CatalogFactory.create({"warehouse": str(tmp_path)})
    catalog.create_database("test_db", ignore_if_exists=True)
    return catalog, tmp_path


@pytest.fixture
def append_only_table(local_paimon_catalog):
    catalog, tmp_path = local_paimon_catalog
    schema = pypaimon.Schema.from_pyarrow_schema(
        pa.schema([
            pa.field("id", pa.int64()),
            pa.field("name", pa.string()),
            pa.field("value", pa.float64()),
            pa.field("dt", pa.string()),
        ]),
        partition_keys=["dt"],
        options={"bucket": "1", "file.format": "parquet"},
    )
    catalog.create_table("test_db.append_table", schema, ignore_if_exists=False)
    return catalog.get_table("test_db.append_table"), tmp_path


@pytest.fixture
def pk_table(local_paimon_catalog):
    catalog, tmp_path = local_paimon_catalog
    schema = pypaimon.Schema.from_pyarrow_schema(
        pa.schema([
            pa.field("id", pa.int64()),
            pa.field("name", pa.string()),
            pa.field("dt", pa.string()),
        ]),
        partition_keys=["dt"],
        primary_keys=["id", "dt"],
        options={"bucket": "1", "file.format": "parquet"},
    )
    catalog.create_table("test_db.pk_table", schema, ignore_if_exists=False)
    return catalog.get_table("test_db.pk_table"), tmp_path


# ---------------------------------------------------------------------------
# Basic read roundtrip
# ---------------------------------------------------------------------------


def test_read_paimon_basic(append_only_table):
    """Write data via pypaimon, read back via _read_table(table), verify correctness."""
    table, tmp_path = append_only_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3], pa.int64()),
            "name": pa.array(["alice", "bob", "charlie"], pa.string()),
            "value": pa.array([1.1, 2.2, 3.3], pa.float64()),
            "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table)
    result = df.sort("id").to_arrow()

    assert result.num_rows == 3
    assert result.schema.field("id").type == pa.int64()
    assert result.schema.field("name").type in (pa.string(), pa.large_string())
    assert result.column("id").to_pylist() == [1, 2, 3]
    assert result.column("name").to_pylist() == ["alice", "bob", "charlie"]


def test_read_paimon_empty_table(append_only_table):
    """Reading an empty table should return a DataFrame with the correct schema but zero rows."""
    table, _ = append_only_table
    df = _read_table(table)
    result = df.to_arrow()

    assert result.num_rows == 0
    assert "id" in result.schema.names
    assert "name" in result.schema.names
    assert "dt" in result.schema.names


def test_read_paimon_schema_matches(append_only_table):
    """The schema reported by _read_table(table) should match the Paimon table schema."""
    table, _ = append_only_table
    df = _read_table(table)
    schema = df.schema()

    assert "id" in schema.column_names()
    assert "name" in schema.column_names()
    assert "value" in schema.column_names()
    assert "dt" in schema.column_names()


def test_read_paimon_source_is_serializable(append_only_table):
    """The Daft source must not serialize live table/file_io/storage objects."""
    from daft.pickle import dumps, loads

    from pypaimon.daft.daft_datasource import PaimonDataSource

    table, _ = append_only_table
    table.file_io._unserializable_marker = _UnserializableFileIoMarker()

    source = PaimonDataSource(
        table,
        storage_config=_UnserializableStorageConfig(),
        catalog_options={},
    )

    restored = loads(dumps(source))

    assert restored is not source
    assert restored.schema.column_names() == source.schema.column_names()
    assert restored._table is not table
    assert restored._table.identifier.get_full_name() == table.identifier.get_full_name()
    assert restored._storage_config.multithreaded_io is False


def test_source_serialization_preserves_explicit_io_config_for_native_task(
    append_only_table, monkeypatch
):
    from daft.io import IOConfig, S3Config
    from daft.io.pushdowns import Pushdowns
    from daft.io.source import DataSourceTask
    from daft.pickle import dumps, loads

    from pypaimon.daft.daft_paimon import _source_for_table

    table, warehouse = append_only_table
    _write_to_paimon(
        table,
        pa.table(
            {
                "id": [1],
                "name": ["a"],
                "value": [1.0],
                "dt": ["2024-01-01"],
            }
        ),
    )

    explicit_values = (
        "https://explicit.example",
        "explicit-key",
        "explicit-secret",
        "explicit-token",
    )
    explicit_io_config = IOConfig(
        s3=S3Config(
            endpoint_url=explicit_values[0],
            key_id=explicit_values[1],
            access_key=explicit_values[2],
            session_token=explicit_values[3],
        )
    )
    catalog_options = {
        "warehouse": str(warehouse),
        "fs.s3.endpoint": "https://catalog.example",
        "fs.s3.accessKeyId": "catalog-key",
        "fs.s3.accessKeySecret": "catalog-secret",
        "fs.s3.securityToken": "catalog-token",
    }

    source = _source_for_table(
        table,
        catalog_options=catalog_options,
        io_config=explicit_io_config,
    )
    restored = loads(dumps(source))

    captured = {}
    sentinel = object()

    def capture_parquet(**kwargs):
        captured.update(kwargs)
        return sentinel

    monkeypatch.setattr(DataSourceTask, "parquet", capture_parquet)

    async def first_task():
        async for task in restored.get_tasks(Pushdowns()):
            return task
        raise AssertionError("Expected at least one native task")

    assert asyncio.run(first_task()) is sentinel
    assert captured["storage_config"] is restored._storage_config

    for storage_config in (source._storage_config, restored._storage_config):
        s3 = storage_config.io_config.s3
        assert (
            s3.endpoint_url,
            s3.key_id,
            s3.access_key,
            s3.session_token,
        ) == explicit_values


def test_read_paimon_serialized_source_uses_current_filter_for_fallback(local_paimon_catalog):
    """A serialized source must plan from the filters in the current request."""
    from daft import context, runners
    from daft.daft import StorageConfig
    from daft.io.pushdowns import Pushdowns
    from daft.pickle import dumps, loads

    from pypaimon.daft.daft_datasource import PaimonDataSource

    catalog, _ = local_paimon_catalog
    schema = pypaimon.Schema.from_pyarrow_schema(
        pa.schema([
            pa.field("id", pa.int64()),
            pa.field("name", pa.string()),
        ]),
        options={
            "file.format": "avro",
            "source.split.target-size": "800b",
            "source.split.open-file-cost": "600b",
        },
    )
    catalog.create_table("test_db.avro_serialized_pushdown_filter", schema, ignore_if_exists=False)
    table = catalog.get_table("test_db.avro_serialized_pushdown_filter")
    _write_to_paimon(table, pa.table({"id": [1], "name": ["first"]}))
    _write_to_paimon(table, pa.table({"id": [999], "name": ["match"]}))

    io_config = context.get_context().daft_planning_config.default_io_config
    storage_config = StorageConfig(runners.get_or_create_runner().name != "ray", io_config)
    source = PaimonDataSource(table, storage_config=storage_config, catalog_options={})

    restored = loads(dumps(source))
    batches = asyncio.run(
        _collect_paimon_source_batches(
            restored,
            Pushdowns(filters=col("id") == 999, limit=1),
        )
    )

    assert batches == [{"id": [999], "name": ["match"]}]


def test_read_paimon_remote_ray_task_is_serializable(pk_table, monkeypatch):
    """A fallback PK split task must reopen the table from metadata on Ray workers.

    Splits that need an LSM merge (here, overlapping primary-key writes) are read
    by the pypaimon reader task. Under the Ray runner that task is pickled to
    remote workers, so it must serialize only rebuildable metadata -- never the
    live table / file_io / storage objects.
    """
    from daft import runners
    from daft.io.pushdowns import Pushdowns
    from daft.pickle import dumps, loads

    from pypaimon.daft.daft_datasource import PaimonDataSource

    class _RayRunner:
        name = "ray"

    table, _ = pk_table
    # Two overlapping writes on id=1 create non-raw-convertible splits that
    # require the pypaimon merge reader (the fallback _PaimonPKSplitTask).
    _write_to_paimon(
        table,
        pa.table(
            {
                "id": pa.array([1, 2], pa.int64()),
                "name": pa.array(["old_a", "old_b"], pa.string()),
                "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()),
            }
        ),
    )
    _write_to_paimon(
        table,
        pa.table(
            {
                "id": pa.array([1], pa.int64()),
                "name": pa.array(["new_a"], pa.string()),
                "dt": pa.array(["2024-01-01"], pa.string()),
            }
        ),
    )
    table.file_io._unserializable_marker = _UnserializableFileIoMarker()

    source = PaimonDataSource(
        table,
        storage_config=_UnserializableStorageConfig(),
        catalog_options={},
    )
    monkeypatch.setattr(runners, "get_or_create_runner", lambda: _RayRunner())

    async def first_task():
        async for task in source.get_tasks(Pushdowns()):
            return task
        raise AssertionError("Expected at least one task")

    async def read_task(task):
        rows = []
        async for batch in task.read():
            rows.append(batch.to_pydict())
        return rows

    task = asyncio.run(first_task())
    assert type(task).__name__ == "_PaimonPKSplitTask"

    restored_task = loads(dumps(task))
    batches = asyncio.run(read_task(restored_task))

    merged = {
        _id: name
        for batch in batches
        for _id, name in zip(batch["id"], batch["name"])
    }
    assert merged == {1: "new_a", 2: "old_b"}


# ---------------------------------------------------------------------------
# Multi-partition reads
# ---------------------------------------------------------------------------


def test_read_paimon_multiple_partitions(append_only_table):
    """Data spread across multiple partitions should all be read back."""
    table, _ = append_only_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3, 4], pa.int64()),
            "name": pa.array(["a", "b", "c", "d"], pa.string()),
            "value": pa.array([1.0, 2.0, 3.0, 4.0], pa.float64()),
            "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table)
    result = df.sort("id").to_arrow()

    assert result.num_rows == 4
    assert sorted(result.column("id").to_pylist()) == [1, 2, 3, 4]
    assert set(result.column("dt").to_pylist()) == {"2024-01-01", "2024-01-02"}


def test_read_paimon_partition_column_present(append_only_table):
    """Partition columns must appear in the result schema and have correct values."""
    table, _ = append_only_table
    data = pa.table(
        {
            "id": pa.array([10], pa.int64()),
            "name": pa.array(["x"], pa.string()),
            "value": pa.array([9.9], pa.float64()),
            "dt": pa.array(["2024-03-15"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table)
    result = df.to_arrow()

    assert "dt" in result.schema.names
    assert result.column("dt").to_pylist() == ["2024-03-15"]


# ---------------------------------------------------------------------------
# Column projection
# ---------------------------------------------------------------------------


def test_read_paimon_column_projection(append_only_table):
    """Requesting a subset of columns should work without errors."""
    table, _ = append_only_table
    data = pa.table(
        {
            "id": pa.array([1, 2], pa.int64()),
            "name": pa.array(["a", "b"], pa.string()),
            "value": pa.array([1.0, 2.0], pa.float64()),
            "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table).select("id", "name")
    result = df.sort("id").to_arrow()

    assert result.schema.names == ["id", "name"]
    assert result.num_rows == 2
    assert result.column("id").to_pylist() == [1, 2]


def test_read_paimon_multiple_batches(local_paimon_catalog):
    """Writing data in two separate batches (two snapshots) should read all rows."""
    catalog, tmp_path = local_paimon_catalog
    schema = pypaimon.Schema.from_pyarrow_schema(
        pa.schema(
            [
                pa.field("id", pa.int64()),
                pa.field("v", pa.string()),
                pa.field("dt", pa.string()),
            ]
        ),
        partition_keys=["dt"],
        options={"bucket": "1", "file.format": "parquet"},
    )
    catalog.create_table("test_db.multi_batch", schema, ignore_if_exists=False)
    table = catalog.get_table("test_db.multi_batch")

    batch1 = pa.table({"id": [1, 2], "v": ["a", "b"], "dt": ["2024-01-01", "2024-01-01"]})
    batch2 = pa.table({"id": [3, 4], "v": ["c", "d"], "dt": ["2024-01-02", "2024-01-02"]})

    _write_to_paimon(table, batch1)
    _write_to_paimon(table, batch2)

    df = _read_table(table)
    result = df.sort("id").to_arrow()

    assert result.num_rows == 4
    assert result.column("id").to_pylist() == [1, 2, 3, 4]


# ---------------------------------------------------------------------------
# Primary-key table (LSM merge fall-back)
# ---------------------------------------------------------------------------


def test_read_paimon_pk_table_basic(pk_table):
    """Primary-key table: falls back to pypaimon reader, data should be correct."""
    table, _ = pk_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3], pa.int64()),
            "name": pa.array(["x", "y", "z"], pa.string()),
            "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table)
    result = df.sort("id").to_arrow()

    assert result.num_rows == 3
    assert result.column("id").to_pylist() == [1, 2, 3]
    assert result.column("name").to_pylist() == ["x", "y", "z"]


def test_read_paimon_pk_table_deduplication(pk_table):
    """Writing the same PK twice should result in the latest value being visible."""
    table, _ = pk_table
    batch1 = pa.table(
        {
            "id": pa.array([1, 2], pa.int64()),
            "name": pa.array(["old_a", "old_b"], pa.string()),
            "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch1)

    batch2 = pa.table(
        {
            "id": pa.array([1], pa.int64()),
            "name": pa.array(["new_a"], pa.string()),
            "dt": pa.array(["2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch2)

    df = _read_table(table)
    result = df.sort("id").to_arrow()

    assert result.num_rows == 2
    id1_row = [row for row in zip(result.column("id").to_pylist(), result.column("name").to_pylist()) if row[0] == 1]
    assert len(id1_row) == 1
    assert id1_row[0][1] == "new_a"


def test_read_paimon_pk_fallback_applies_pushed_filter(pk_table):
    """Fallback PK reads must apply filters pushed into the Paimon source."""
    table, _ = pk_table
    batch1 = pa.table(
        {
            "id": pa.array([1, 2], pa.int64()),
            "name": pa.array(["old_a", "old_b"], pa.string()),
            "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch1)

    batch2 = pa.table(
        {
            "id": pa.array([1], pa.int64()),
            "name": pa.array(["new_a"], pa.string()),
            "dt": pa.array(["2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch2)

    batches = asyncio.run(_read_paimon_source_batches(table, filter_expr=col("id") == 1))
    ids = [value for batch in batches for value in batch["id"]]
    names = [value for batch in batches for value in batch["name"]]

    assert ids == [1]
    assert names == ["new_a"]


def test_read_paimon_pk_fallback_filters_before_projection(pk_table):
    """Fallback reads keep filter columns available for Daft's upper filter."""
    table, _ = pk_table
    batch1 = pa.table(
        {
            "id": pa.array([1, 2], pa.int64()),
            "name": pa.array(["old_a", "old_b"], pa.string()),
            "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch1)

    batch2 = pa.table(
        {
            "id": pa.array([1], pa.int64()),
            "name": pa.array(["new_a"], pa.string()),
            "dt": pa.array(["2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch2)

    batches = asyncio.run(
        _read_paimon_source_batches(
            table,
            filter_expr=col("id") == 1,
            columns=["name"],
        )
    )

    assert batches == [{"name": ["new_a"], "id": [1]}]


def test_read_paimon_fallback_plans_current_pushdown_filter(local_paimon_catalog):
    """Fallback planning must use the filter from the current request."""
    catalog, _ = local_paimon_catalog
    schema = pypaimon.Schema.from_pyarrow_schema(
        pa.schema([
            pa.field("id", pa.int64()),
            pa.field("name", pa.string()),
        ]),
        options={
            "file.format": "avro",
            "source.split.target-size": "800b",
            "source.split.open-file-cost": "600b",
        },
    )
    catalog.create_table("test_db.avro_pushdown_filter", schema, ignore_if_exists=False)
    table = catalog.get_table("test_db.avro_pushdown_filter")
    _write_to_paimon(table, pa.table({"id": [1], "name": ["first"]}))
    _write_to_paimon(table, pa.table({"id": [999], "name": ["match"]}))

    batches = asyncio.run(
        _read_paimon_source_batches(
            table,
            filter_expr=col("id") == 999,
            limit=1,
        )
    )

    assert batches == [{"id": [999], "name": ["match"]}]


def test_read_paimon_fallback_not_in_filter_excludes_nulls_before_limit(local_paimon_catalog):
    """Fallback datasource tasks must satisfy pushed NOT IN filters before limit."""
    catalog, _ = local_paimon_catalog
    pa_schema = pa.schema([
        pa.field("id", pa.int64()),
        pa.field("name", pa.string()),
    ])
    schema = pypaimon.Schema.from_pyarrow_schema(
        pa_schema,
        options={"file.format": "row"},
    )
    catalog.create_table("test_db.row_not_in_filter_limit", schema, ignore_if_exists=False)
    table = catalog.get_table("test_db.row_not_in_filter_limit")
    _write_to_paimon(table, pa.table({"id": [None], "name": ["null-row"]}, schema=pa_schema))
    _write_to_paimon(table, pa.table({"id": [3], "name": ["match"]}, schema=pa_schema))

    batches = asyncio.run(
        _read_paimon_source_batches(
            table,
            filter_expr=~col("id").is_in([1, 2]),
            limit=1,
        )
    )

    assert batches == [{"id": [3], "name": ["match"]}]


def test_read_paimon_fallback_keeps_limit_above_remaining_filter(local_paimon_catalog):
    """Fallback reads must not apply limit before Daft evaluates remaining filters."""
    catalog, _ = local_paimon_catalog
    schema = pypaimon.Schema.from_pyarrow_schema(
        pa.schema([
            pa.field("id", pa.int64()),
            pa.field("name", pa.string()),
        ]),
        options={"file.format": "avro"},
    )
    catalog.create_table("test_db.avro_remaining_filter_limit", schema, ignore_if_exists=False)
    table = catalog.get_table("test_db.avro_remaining_filter_limit")
    _write_to_paimon(table, pa.table({"id": [1, 999], "name": ["first", "match"]}))

    result = _read_table(table).where(~(col("id") == 1)).limit(1).to_pydict()

    assert result == {"id": [999], "name": ["match"]}


def test_read_paimon_keeps_limit_above_partition_filter(append_only_table):
    """Scan planning must not apply limit before datasource partition pruning."""
    table, _ = append_only_table
    _write_to_paimon(
        table,
        pa.table(
            {
                "id": pa.array([1], pa.int64()),
                "name": pa.array(["first"], pa.string()),
                "value": pa.array([1.0], pa.float64()),
                "dt": pa.array(["2024-01-01"], pa.string()),
            }
        ),
    )
    _write_to_paimon(
        table,
        pa.table(
            {
                "id": pa.array([2], pa.int64()),
                "name": pa.array(["match"], pa.string()),
                "value": pa.array([2.0], pa.float64()),
                "dt": pa.array(["2024-01-02"], pa.string()),
            }
        ),
    )

    result = _read_table(table).where(col("dt") == "2024-01-02").limit(1).to_pydict()

    assert result["id"] == [2]


def test_read_paimon_pk_fallback_filter_then_project_dataframe(pk_table):
    """Daft may keep the filter above the source while pushing projection."""
    table, _ = pk_table
    batch1 = pa.table(
        {
            "id": pa.array([1, 2], pa.int64()),
            "name": pa.array(["old_a", "old_b"], pa.string()),
            "dt": pa.array(["2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch1)

    batch2 = pa.table(
        {
            "id": pa.array([1], pa.int64()),
            "name": pa.array(["new_a"], pa.string()),
            "dt": pa.array(["2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, batch2)

    result = _read_table(table).where(col("id") == 1).select("name").to_pydict()

    assert result == {"name": ["new_a"]}


def test_read_paimon_pk_fallback_applies_limit(pk_table):
    """Fallback PK reads must use the same limit as the planning read builder."""
    table, _ = pk_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3], pa.int64()),
            "name": pa.array(["a", "b", "c"], pa.string()),
            "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, data)
    _write_to_paimon(table, data)

    batches = asyncio.run(_read_paimon_source_batches(table, limit=1))
    ids = [value for batch in batches for value in batch["id"]]

    assert len(ids) == 1


# ---------------------------------------------------------------------------
# Filter pushdown
# ---------------------------------------------------------------------------


def test_read_paimon_partition_filter(append_only_table):
    """Partition filter should prune partitions at scan time."""
    table, _ = append_only_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3, 4], pa.int64()),
            "name": pa.array(["a", "b", "c", "d"], pa.string()),
            "value": pa.array([1.0, 2.0, 3.0, 4.0], pa.float64()),
            "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table).where(col("dt") == "2024-01-01")
    result = df.sort("id").to_arrow()

    assert result.num_rows == 2
    assert result.column("id").to_pylist() == [1, 2]
    assert all(dt == "2024-01-01" for dt in result.column("dt").to_pylist())


def test_partition_filter_prunes_at_plan_level(append_only_table):
    """Regression: Daft routes partition predicates to pushdowns.partition_filters
    (a separate channel from pushdowns.filters). They must become a plan-time
    predicate so plan() prunes partitions, instead of planning every split and
    skipping in Python (a full-table plan)."""
    from daft import context, runners
    from daft.daft import StorageConfig
    from daft.io.pushdowns import Pushdowns

    from pypaimon.daft.daft_datasource import PaimonDataSource

    table, _ = append_only_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3], pa.int64()),
            "name": pa.array(["a", "b", "c"], pa.string()),
            "value": pa.array([1.0, 2.0, 3.0], pa.float64()),
            "dt": pa.array(["2024-01-01", "2024-01-02", "2024-01-03"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    io_config = context.get_context().daft_planning_config.default_io_config
    storage_config = StorageConfig(runners.get_or_create_runner().name != "ray", io_config)
    source = PaimonDataSource(table, storage_config=storage_config, catalog_options={})

    pushdowns = Pushdowns(partition_filters=(col("dt") == "2024-01-02"))

    state = source._read_pushdown_state(table, pushdowns)
    # partition filter must become a plan-time predicate
    assert state.planning_predicate is not None

    pruned = source._scan_read_builder(table, state).new_scan().plan().splits()
    all_splits = table.new_read_builder().new_scan().plan().splits()

    # pruning reduced the planned splits, and only the matching partition remains
    assert 0 < len(pruned) < len(all_splits)
    assert all(s.partition.to_dict().get("dt") == "2024-01-02" for s in pruned)


def test_row_and_partition_filters_both_reach_planning_predicate(append_only_table):
    """filters and partition_filters arrive on separate Daft channels; both must
    fold into the planning predicate (via _and_predicates)."""
    from daft import context, runners
    from daft.daft import StorageConfig
    from daft.io.pushdowns import Pushdowns

    from pypaimon.daft.daft_datasource import PaimonDataSource

    table, _ = append_only_table
    _write_to_paimon(table, pa.table(
        {
            "id": pa.array([1, 2, 3], pa.int64()),
            "name": pa.array(["a", "b", "c"], pa.string()),
            "value": pa.array([1.0, 2.0, 3.0], pa.float64()),
            "dt": pa.array(["2024-01-01", "2024-01-02", "2024-01-02"], pa.string()),
        }
    ))

    io_config = context.get_context().daft_planning_config.default_io_config
    storage_config = StorageConfig(runners.get_or_create_runner().name != "ray", io_config)
    source = PaimonDataSource(table, storage_config=storage_config, catalog_options={})
    pushdowns = Pushdowns(filters=(col("id") > 1),
                          partition_filters=(col("dt") == "2024-01-02"))

    state = source._read_pushdown_state(table, pushdowns)
    # both channels combined into one AND planning predicate
    assert state.planning_predicate is not None
    assert state.planning_predicate.method == "and"

    pruned = source._scan_read_builder(table, state).new_scan().plan().splits()
    assert pruned and all(s.partition.to_dict().get("dt") == "2024-01-02" for s in pruned)


def test_isnull_partition_filter_keeps_null_partition(append_only_table):
    """isNull must NOT be pushed to plan pruning: plan pruning would drop the
    null partition and the Python post-filter can't restore it. It must be left
    to the post-filter so col(dt).is_null() still returns the null-partition row."""
    table, _ = append_only_table
    _write_to_paimon(table, pa.table(
        {
            "id": pa.array([1, 2], pa.int64()),
            "name": pa.array(["a", "b"], pa.string()),
            "value": pa.array([1.0, 2.0], pa.float64()),
            "dt": pa.array([None, "2024-01-02"], pa.string()),
        }
    ))

    result = _read_table(table).where(col("dt").is_null()).to_arrow()

    assert result.num_rows == 1
    assert result.column("id").to_pylist() == [1]
    assert result.column("dt").to_pylist() == [None]


def test_isnull_partition_filter_not_pushed_even_for_pk_table(pk_table):
    """isNull must be excluded from plan pushdown for ALL tables. The row gate
    (_can_plan_predicate) would allow isNull on a PK table without deletion
    vectors; the partition path must be stricter (_predicate_contains_is_null),
    else a PK table would lose its null partition."""
    from daft import context, runners
    from daft.daft import StorageConfig
    from daft.io.pushdowns import Pushdowns

    from pypaimon.daft.daft_datasource import PaimonDataSource

    table, _ = pk_table
    io_config = context.get_context().daft_planning_config.default_io_config
    storage_config = StorageConfig(runners.get_or_create_runner().name != "ray", io_config)
    source = PaimonDataSource(table, storage_config=storage_config, catalog_options={})

    st_null = source._read_pushdown_state(
        table, Pushdowns(partition_filters=col("dt").is_null()))
    assert st_null.planning_predicate is None  # would be non-None via _can_plan_predicate
    st_eq = source._read_pushdown_state(
        table, Pushdowns(partition_filters=col("dt") == "2024-01-02"))
    assert st_eq.planning_predicate is not None


def test_read_paimon_row_filter(append_only_table):
    """Row-level filter should be applied after reading data."""
    table, _ = append_only_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3, 4], pa.int64()),
            "name": pa.array(["alice", "bob", "charlie", "dave"], pa.string()),
            "value": pa.array([10.0, 20.0, 30.0, 40.0], pa.float64()),
            "dt": pa.array(["2024-01-01", "2024-01-01", "2024-01-01", "2024-01-01"], pa.string()),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table).where(col("id") > 2)
    result = df.sort("id").to_arrow()

    assert result.num_rows == 2
    assert result.column("id").to_pylist() == [3, 4]


def test_read_paimon_combined_filter(append_only_table):
    """Combined partition + row filter should work together."""
    table, _ = append_only_table
    data = pa.table(
        {
            "id": pa.array([1, 2, 3, 4, 5, 6], pa.int64()),
            "name": pa.array(["a", "b", "c", "d", "e", "f"], pa.string()),
            "value": pa.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], pa.float64()),
            "dt": pa.array(
                ["2024-01-01", "2024-01-01", "2024-01-02", "2024-01-02", "2024-01-03", "2024-01-03"],
                pa.string(),
            ),
        }
    )
    _write_to_paimon(table, data)

    df = _read_table(table).where((col("dt") == "2024-01-01") & (col("id") == 2))
    result = df.to_arrow()

    assert result.num_rows == 1
    assert result.column("id").to_pylist() == [2]
    assert result.column("dt").to_pylist() == ["2024-01-01"]


# ---------------------------------------------------------------------------
# Filter pushdown tests
# ---------------------------------------------------------------------------


class TestFilterPushdown:
    """Tests for filter pushdown to Paimon."""

    @pytest.fixture
    def filter_table(self, local_paimon_catalog):
        """Create a table for filter pushdown tests."""
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([("id", pa.int64()), ("value", pa.string())])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.filter_test", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.filter_test")

        data = pa.table({"id": [1, 2, 3, 4, 5], "value": ["a", "b", "c", "d", "e"]})
        _write_to_paimon(table, data)
        return table

    def test_filter_pushdown_equal(self, filter_table):
        df = _read_table(filter_table).where(col("id") == 3)
        result = df.to_pydict()
        assert result["id"] == [3]
        assert result["value"] == ["c"]

    def test_filter_pushdown_comparison(self, filter_table):
        df = _read_table(filter_table).where(col("id") > 3)
        result = df.to_pydict()
        assert result["id"] == [4, 5]

        df = _read_table(filter_table).where(col("id") <= 2)
        result = df.to_pydict()
        assert result["id"] == [1, 2]

    def test_filter_pushdown_is_in(self, filter_table):
        df = _read_table(filter_table).where(col("id").is_in([2, 4]))
        result = df.to_pydict()
        assert result["id"] == [2, 4]

    def test_filter_pushdown_is_null(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([("id", pa.int64()), ("value", pa.string())])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.filter_null", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.filter_null")

        data = pa.table({"id": [1, 2, 3], "value": ["a", None, "c"]})
        _write_to_paimon(table, data)

        df = _read_table(table).where(col("value").is_null())
        result = df.to_pydict()
        assert result["id"] == [2]

    def test_filter_pushdown_string(self, filter_table):
        df = _read_table(filter_table).where(col("value").startswith("a"))
        result = df.to_pydict()
        assert result["value"] == ["a"]

        df = _read_table(filter_table).where(col("value").contains("b"))
        result = df.to_pydict()
        assert result["value"] == ["b"]

    def test_filter_pushdown_combined(self, filter_table):
        df = _read_table(filter_table).where((col("id") >= 2) & (col("id") <= 4))
        result = df.to_pydict()
        assert result["id"] == [2, 3, 4]

    def test_filter_pushdown_splits_supported_conjuncts(self, filter_table):
        unsupported = col("value").contains(col("id"))
        cases = [
            ((col("id") == 1) & unsupported, [("equal", "id", (1,))]),
            (unsupported & (col("id") == 1), [("equal", "id", (1,))]),
            (
                (col("id") == 1) & (unsupported & (col("value") == "a")),
                [("equal", "id", (1,)), ("equal", "value", ("a",))],
            ),
        ]

        for expr, expected_leaves in cases:
            pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr)

            assert len(pushed_filters) == len(expected_leaves)
            assert len(remaining_filters) == 1
            assert _contains_expr(remaining_filters[0])
            assert _predicate_leaves(predicate) == expected_leaves

    def test_filter_pushdown_does_not_split_or_with_unsupported_branch(self, filter_table):
        expr = (col("id") == 1) | col("value").contains(col("id"))

        pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr)

        assert pushed_filters == []
        assert remaining_filters == [expr._expr]
        assert predicate is None

    def test_filter_pushdown_pushes_supported_or_conjunct(self, filter_table):
        supported_or = (col("id") == 1) | (col("id") == 2)
        expr = supported_or & col("value").contains(col("id"))

        pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr)

        assert len(pushed_filters) == 1
        assert len(remaining_filters) == 1
        assert _contains_expr(remaining_filters[0])
        assert predicate is not None
        assert predicate.method == "or"
        assert _predicate_leaves(predicate) == [("equal", "id", (1,)), ("equal", "id", (2,))]

    def test_filter_pushdown_rewrites_supported_not_predicates(self, filter_table):
        cases = [
            (~(col("id") == 1), [("notEqual", "id", (1,))]),
            (~col("id").is_in([1, 2]), [("notIn", "id", (1, 2))]),
            (~col("id").between(1, 3), [("notBetween", "id", (1, 3))]),
            (~col("value").is_null(), [("isNotNull", "value", ())]),
            (~col("value").not_null(), [("isNull", "value", ())]),
        ]

        for expr, expected_leaves in cases:
            pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr)

            assert len(pushed_filters) == 1
            assert remaining_filters == []
            assert _predicate_leaves(predicate) == expected_leaves

    def test_filter_pushdown_supported_not_predicates_read_path(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([
            ("id", pa.int64()),
            ("value", pa.string()),
        ])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.filter_not_read", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.filter_not_read")

        data = pa.table(
            {
                "id": [1, 2, 3, 4, 5],
                "value": ["a", "b", None, "d", None],
            },
            schema=pa_schema,
        )
        _write_to_paimon(table, data)

        cases = [
            (~(col("id") == 1), [2, 3, 4, 5]),
            (~col("id").is_in([1, 3]), [2, 4, 5]),
            (~col("id").between(2, 4), [1, 5]),
            (~col("value").is_null(), [1, 2, 4]),
            (~col("value").not_null(), [3, 5]),
        ]

        for expr, expected_ids in cases:
            result = _read_table(table).where(expr).select("id").sort("id").to_pydict()

            assert result["id"] == expected_ids

    def test_filter_pushdown_does_not_demorgan_not_compound_predicates(self, filter_table):
        expressions = [
            ~((col("id") == 1) & (col("value") == "a")),
            ~((col("id") == 1) | (col("value") == "a")),
        ]

        for expr in expressions:
            pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr)

            assert pushed_filters == []
            assert remaining_filters == [expr._expr]
            assert predicate is None

    def test_unsupported_expression_remains_in_daft(self, filter_table):
        expressions = [
            col("id") == lit(1).cast("int64"),
            col("id") == lit(None),
            col("value").contains(col("id")),
            col("value").startswith(col("id")),
            col("value").endswith(col("id")),
            col("value").contains(lit("a").cast("string")),
            col("id").between(col("id"), 3),
            col("id").is_in([1, None]),
        ]

        for expr in expressions:
            pushed_filters, remaining_filters, predicate = convert_filters_to_paimon(filter_table, expr._expr)

            assert pushed_filters == []
            assert remaining_filters == [expr._expr]
            assert predicate is None

    def test_mixed_string_expression_is_filtered_by_daft(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([
            ("id", pa.int64()),
            ("value", pa.string()),
            ("pattern", pa.string()),
        ])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.filter_mixed_string", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.filter_mixed_string")

        data = pa.table(
            {
                "id": [1, 2, 3],
                "value": ["alpha", "bravo", "charlie"],
                "pattern": ["lp", "zz", "lie"],
            }
        )
        _write_to_paimon(table, data)

        df = _read_table(table).where(col("value").contains(col("pattern")))
        result = df.sort("id").to_pydict()

        assert result["id"] == [1, 3]

    def test_mixed_conjunctive_expression_is_filtered_by_daft(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([
            ("id", pa.int64()),
            ("value", pa.string()),
            ("pattern", pa.string()),
        ])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.filter_mixed_conjunctive", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.filter_mixed_conjunctive")

        data = pa.table(
            {
                "id": [1, 1, 2, 3],
                "value": ["alpha", "bravo", "alps", "charlie"],
                "pattern": ["lp", "zz", "lp", "lie"],
            }
        )
        _write_to_paimon(table, data)

        df = _read_table(table).where((col("id") == 1) & col("value").contains(col("pattern")))
        result = df.sort("value").to_pydict()

        assert result["id"] == [1]
        assert result["value"] == ["alpha"]


# ---------------------------------------------------------------------------
# Advanced data types
# ---------------------------------------------------------------------------


class TestNestedTypes:
    """Tests for nested data types (list, map, struct)."""

    def test_read_nested_list(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([("id", pa.int64()), ("list_col", pa.list_(pa.int64()))])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.nested_list", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.nested_list")

        data = pa.table({"id": [1, 2], "list_col": [[1, 2, 3], [4, 5]]}, schema=pa_schema)
        _write_to_paimon(table, data)

        df = _read_table(table)
        result = df.to_pydict()
        assert len(result["id"]) == 2
        assert result["list_col"][0] == [1, 2, 3]

    def test_read_nested_map(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([("id", pa.int64()), ("map_col", pa.map_(pa.string(), pa.int64()))])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.nested_map", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.nested_map")

        data = pa.table(
            {"id": [1, 2], "map_col": [[("k1", 1), ("k2", 2)], [("k3", 3)]]},
            schema=pa_schema,
        )
        _write_to_paimon(table, data)

        df = _read_table(table)
        result = df.to_pydict()
        assert len(result["id"]) == 2


class TestDecimalType:
    """Tests for decimal type support."""

    def test_read_decimal(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([("id", pa.int64()), ("amount", pa.decimal128(10, 2))])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.decimal_table", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.decimal_table")

        data = pa.table(
            {
                "id": [1, 2, 3],
                "amount": [decimal.Decimal("100.50"), decimal.Decimal("200.75"), decimal.Decimal("300.00")],
            },
            schema=pa_schema,
        )
        _write_to_paimon(table, data)

        df = _read_table(table)
        result = df.to_pydict()
        assert len(result["id"]) == 3


class TestNullValues:
    """Tests for null value handling."""

    def test_read_with_nulls(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([("id", pa.int64()), ("name", pa.string()), ("value", pa.float64())])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(pa_schema)
        catalog.create_table("test_db.null_table", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.null_table")

        data = pa.table({"id": [1, 2, 3], "name": ["alice", None, "charlie"], "value": [1.1, 2.2, None]})
        _write_to_paimon(table, data)

        df = _read_table(table)
        result = df.to_pydict()
        assert len(result["id"]) == 3
        assert result["name"][1] is None
        assert result["value"][2] is None


class TestCompositePrimaryKey:
    """Tests for composite primary key tables."""

    def test_read_composite_pk_table(self, local_paimon_catalog):
        catalog, tmp_path = local_paimon_catalog
        pa_schema = pa.schema([("pk1", pa.int64()), ("pk2", pa.string()), ("value", pa.float64())])
        paimon_schema = pypaimon.Schema.from_pyarrow_schema(
            pa_schema, primary_keys=["pk1", "pk2"], options={"bucket": "2"}
        )
        catalog.create_table("test_db.composite_pk", paimon_schema, ignore_if_exists=True)
        table = catalog.get_table("test_db.composite_pk")

        data = pa.table({"pk1": [1, 1, 2], "pk2": ["a", "b", "a"], "value": [1.1, 2.2, 3.3]})
        _write_to_paimon(table, data)

        df = _read_table(table)
        result = df.to_pydict()
        assert len(result["pk1"]) == 3
