# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

import hashlib
import json
import logging
from collections.abc import Iterator
from pathlib import Path
from typing import Any, NoReturn

import lance
import pyarrow as pa
import pyarrow.compute as pc
import pytest
from lance.blob import BlobArray, BlobType
from yarl import URL

import geneva.apply as apply_module
import geneva.runners.ray.pipeline as pipeline_module
from geneva import CheckpointStore, connect, udf
from geneva.apply import (
    CheckpointingApplier,
    MapBatchCheckpoint,
    _check_fragment_data_file_exists,
    _find_output_data_file_in_fragment,
    _legacy_fragment_dedupe_key,
    plan_read,
)
from geneva.apply.multiprocess import (
    MultiProcessBatchApplier,
    _batch_to_buf,
    _buf_to_batch,
    _restore_extension_types,
    _strip_extension_types,
)
from geneva.apply.task import (
    DEFAULT_CHECKPOINT_ROWS,
    BackfillUDFTask,
    ReadTask,
    ScanTask,
)
from geneva.apply.utils import _index_checkpoint_ranges, _parse_checkpoint_range_key
from geneva.checkpoint import InMemoryCheckpointStore
from geneva.checkpoint_utils import hash_source_files
from geneva.debug.error_store import ErrorStore
from geneva.debug.logger import TableErrorLogger
from geneva.jobs.config import JobConfig
from geneva.runners.ray.pipeline import (
    ColumnAddPipelineJob,
    FragmentWriterManager,
    _get_fragment_dedupe_key,
    _get_relevant_field_ids,
    get_source_data_files,
)
from geneva.runners.ray.writer import FragmentWriter
from geneva.table import TableReference
from geneva.utils.object_store_retry import is_retryable_object_store_error
from geneva.utils.parse_rust_debug import extract_field_ids

_LOG = logging.getLogger(__name__)


@pytest.fixture
def tbl_ref(tmp_path) -> TableReference:
    return TableReference(table_id=["tbl"], version=None, db_uri=str(tmp_path))


class _NoopRemote:
    def remote(self, *args: object, **kwargs: object) -> None:
        return None


class _NoopJobTracker:
    def __init__(self) -> None:
        self.increment = _NoopRemote()
        self.set_total = _NoopRemote()
        self.set_desc = _NoopRemote()
        self.mark_done = _NoopRemote()


class _DummyReadTask(ReadTask):
    def __init__(
        self,
        *,
        frag_id: int,
        rows: int,
        offset: int = 0,
        batches: list[pa.RecordBatch] | None = None,
    ) -> None:
        self._frag_id = frag_id
        self._rows = rows
        self._offset = offset
        self._batches = batches or []

    def to_batches(
        self, *, batch_size=DEFAULT_CHECKPOINT_ROWS
    ) -> Iterator[pa.RecordBatch]:
        yield from self._batches

    def checkpoint_key(self) -> str:
        return f"dummy-{self._frag_id}-{self._offset}-{self._rows}"

    def dest_frag_id(self) -> int:
        return self._frag_id

    def dest_offset(self) -> int:
        return self._offset

    def num_rows(self) -> int:
        return self._rows

    def table_uri(self) -> str:
        return "memory://dummy"


def _src_files_hash_for_cols(tbl, cols: list[str]) -> str:
    dataset = tbl.to_lance()
    relevant_field_ids = _get_relevant_field_ids(dataset, cols)
    frag = dataset.get_fragment(0)
    return hash_source_files(get_source_data_files(frag, relevant_field_ids))


def test_create_plan(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    plans = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=16)[0])
    assert len(plans) == 1
    plan = plans[0]
    assert plan.uri == tbl.uri
    assert plan.offset == 0
    assert plan.limit == 3


def test_create_plan_with_diverse_shuffle(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    ds = lance.write_dataset(
        pa.table({"a": range(1024)}), tmp_path / "tbl", max_rows_per_file=16
    )

    plans = list(
        plan_read(ds.uri, tbl_ref, ["a"], batch_size=1, task_shuffle_diversity=4)[0]
    )
    assert len(plans) == 1024
    plan = plans[0]
    assert plan.uri == ds.uri
    assert plan.offset == 0
    assert plan.limit == 1


@udf(input_columns=["a"])
def one(*args, **kwargs) -> int:
    return 1


def test_applier(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    plans = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=16)[0])
    assert len(plans) == 1
    plan = plans[0]
    assert plan.uri == tbl.uri
    assert plan.offset == 0
    assert plan.limit == 3

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
    )
    results, cnt_udf_computed = applier.run(plan)
    assert len(results) == 1
    batch = store[results[0].checkpoint_key]
    assert len(batch) == 3
    assert batch.to_pydict() == {"one": [1, 1, 1], "_rowaddr": [0, 1, 2]}
    assert cnt_udf_computed == 3


class _StatusCodeError(RuntimeError):
    def __init__(self, message: str, status_code: int) -> None:
        super().__init__(message)
        self.status_code = status_code


@pytest.mark.parametrize(
    ("exc", "expected"),
    [
        pytest.param(
            RuntimeError("Failed to get AWS credentials: HttpTimeoutError"),
            True,
            id="aws-timeout-message",
        ),
        pytest.param(
            OSError("gcs object store temporarily unavailable"),
            True,
            id="gcs-temporary-unavailable",
        ),
        pytest.param(
            RuntimeError("azure object store connection reset by peer"),
            True,
            id="azure-connection-reset",
        ),
        pytest.param(
            _StatusCodeError("object store too many requests", 429),
            True,
            id="http-429",
        ),
        pytest.param(
            _StatusCodeError("object store service unavailable", 503),
            True,
            id="http-503",
        ),
        pytest.param(
            _StatusCodeError("rate limited", 429),
            False,
            id="http-429-without-object-store-context",
        ),
        pytest.param(
            RuntimeError("table was not found"),
            False,
            id="table-not-found",
        ),
        pytest.param(
            ValueError("table does not exist"),
            False,
            id="table-does-not-exist",
        ),
        pytest.param(
            _StatusCodeError("unauthorized", 401),
            False,
            id="http-401",
        ),
        pytest.param(
            _StatusCodeError("forbidden", 403),
            False,
            id="http-403",
        ),
        pytest.param(
            _StatusCodeError("not found", 404),
            False,
            id="http-404",
        ),
        pytest.param(
            RuntimeError("Error running task: invalid schema"),
            False,
            id="invalid-schema",
        ),
        pytest.param(
            TimeoutError("timed out while calling external udf"),
            False,
            id="udf-timeout-without-object-store-context",
        ),
        pytest.param(
            ConnectionError("connection reset by peer from openai"),
            False,
            id="udf-connection-error-without-object-store-context",
        ),
    ],
)
def test_retryable_object_store_error_classification(
    exc: BaseException, expected: bool
) -> None:
    assert is_retryable_object_store_error(exc) is expected


def test_retryable_object_store_error_detects_nested_timeout() -> None:
    inner = ValueError(
        "LanceError(IO): provider connector timeout; connection reset by peer"
    )
    outer = RuntimeError("Error running task ScanTask(...)")
    outer.__cause__ = inner

    assert is_retryable_object_store_error(outer) is True


def test_retryable_object_store_error_prefers_nested_non_retryable_error() -> None:
    inner = _StatusCodeError("forbidden", 403)
    outer = RuntimeError("temporary timeout opening object store")
    outer.__cause__ = inner

    assert is_retryable_object_store_error(outer) is False


def test_retryable_object_store_error_does_not_false_match_eof_substring() -> None:
    exc = RuntimeError("geoff reported an invalid schema")

    assert is_retryable_object_store_error(exc) is False


def test_setup_inputplans_counts_readtasks_for_expected_tasks(
    tbl_ref: TableReference,
) -> None:
    task = _DummyReadTask(frag_id=0, rows=10)
    job = ColumnAddPipelineJob(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            override_batch_size=4,
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_store=CheckpointStore.from_uri("memory"),
        error_store=None,
        config=JobConfig(),
        dst=tbl_ref,
        input_plan=iter([task]),
        job_id="job-expected-tasks",
        job_tracker=_NoopJobTracker(),
    )

    plans, tasks_by_frag, total_tasks = job.setup_inputplans()

    assert tasks_by_frag == {0: 1}
    assert total_tasks == 1
    assert len(list(plans)) == 1


def test_setup_inputplans_counts_multiple_readtasks_per_fragment(
    tbl_ref: TableReference,
) -> None:
    tasks = [
        _DummyReadTask(frag_id=0, rows=5, offset=0),
        _DummyReadTask(frag_id=0, rows=7, offset=5),
    ]
    job = ColumnAddPipelineJob(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            override_batch_size=2,
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_store=CheckpointStore.from_uri("memory"),
        error_store=None,
        config=JobConfig(),
        dst=tbl_ref,
        input_plan=iter(tasks),
        job_id="job-multi-tasks",
        job_tracker=_NoopJobTracker(),
    )

    _, tasks_by_frag, total_tasks = job.setup_inputplans()

    assert tasks_by_frag == {0: 2}
    assert total_tasks == 2


def test_applier_synthesizes_checkpoints_when_no_batches() -> None:
    task = _DummyReadTask(frag_id=0, rows=5, offset=0)
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            override_batch_size=2,
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri="memory",
    )

    checkpoints, cnt_udf_computed = applier.run(task)

    assert cnt_udf_computed == 0
    assert [checkpoint.span for checkpoint in checkpoints] == [2, 2, 1]
    assert [
        applier.checkpoint_store[checkpoint.checkpoint_key].num_rows
        for checkpoint in checkpoints
    ] == [2, 2, 1]


def test_checkpoint_subranges_cover_gaps(tmp_path: Path) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 3]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"one": one},
        override_batch_size=4,
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    applier = CheckpointingApplier(map_task=map_task, checkpoint_uri=store.root)

    class _Task(ScanTask):
        pass

    task = ScanTask(
        uri=tbl.uri,
        table_ref=tbl.get_reference(),
        columns=["a"],
        frag_id=0,
        offset=0,
        limit=4,
        version=tbl.version,
        where=None,
        with_row_address=True,
    )

    batch = pa.record_batch(
        [pa.array([1, 3]), pa.array([0, 3], type=pa.uint64())],
        names=["a", "_rowaddr"],
    )

    result = applier._checkpoint_single_batch(
        task,
        batch,
        dataset_uri=tbl.uri,
        dataset_version=tbl.version,
        where=None,
        udf_rows=None,
        start=0,
        checkpoint_size=4,
    )

    assert result.offset == 0
    # span should expand to checkpoint_size (=4) even with gaps, capped by task_end
    assert result.span == 4
    assert store[result.checkpoint_key].num_rows == 2


def test_applier_chain_spans_respects_task_end(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4, 5]}))

    plan = next(iter(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=5)[0]))
    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            override_batch_size=2,
            min_checkpoint_size=2,
            max_checkpoint_size=2,
        ),
        checkpoint_uri=store.root,
        batch_checkpoint_flush_interval_seconds=0,
    )

    results, _ = applier.run(plan)

    assert [r.offset for r in results] == [0, 2, 4]
    assert [r.span for r in results] == [2, 2, 1]


def test_plan_read_skips_checkpointed_ranges(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4, 5, 6]}))
    dataset = tbl.to_lance()

    map_task = BackfillUDFTask(
        udfs={"one": one},
        override_batch_size=2,
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))

    # Pre-populate checkpoints for [0,2), [2,4), [4,6)
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    for start in (0, 2, 4):
        end = start + 2
        key = map_task.checkpoint_key(
            dataset_uri=tbl.uri,
            dataset_version=dataset.version,
            frag_id=0,
            start=start,
            end=end,
            where=None,
            src_files_hash=src_files_hash,
        )
        store[key] = pa.record_batch([], names=[])

    tasks, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=2,
        map_task=map_task,
        checkpoint_store=store,
    )

    assert list(tasks) == []
    assert pipeline_args["skipped_stats"]["rows"] == 6


def test_plan_read_builds_gaps_and_chunks(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": list(range(10))}))

    map_task = BackfillUDFTask(
        udfs={"one": one},
        override_batch_size=2,
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))

    # Checkpoints cover [0,5) and [7,10), leaving gap [5,7)
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    for start, end in [(0, 5), (7, 10)]:
        key = map_task.checkpoint_key(
            dataset_uri=tbl.uri,
            dataset_version=tbl.version,
            frag_id=0,
            start=start,
            end=end,
            where=None,
            src_files_hash=src_files_hash,
        )
        store[key] = pa.record_batch([], names=[])

    tasks, _ = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=3,  # task_size
        map_task=map_task,
        checkpoint_store=store,
    )

    task_list = list(tasks)
    assert len(task_list) == 1
    task = task_list[0]
    assert task.offset == 5
    assert task.limit == 2


def test_plan_read_chunks_large_gap(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": list(range(20))}))

    map_task = BackfillUDFTask(
        udfs={"one": one},
        override_batch_size=2,
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))

    # Covered [0,2) and [10,20); gap [2,10) length 8.
    # task_size=3 -> chunks [2,5), (5,8), (8,10).
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    for start, end in [(0, 2), (10, 20)]:
        key = map_task.checkpoint_key(
            dataset_uri=tbl.uri,
            dataset_version=tbl.version,
            frag_id=0,
            start=start,
            end=end,
            where=None,
            src_files_hash=src_files_hash,
        )
        store[key] = pa.record_batch([], names=[])

    tasks, _ = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=3,
        map_task=map_task,
        checkpoint_store=store,
    )

    task_list = list(tasks)
    offsets_limits = [(t.offset, t.limit) for t in task_list]
    assert offsets_limits == [(2, 3), (5, 3), (8, 2)]


def test_applier_checkpoints_each_map_batch(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4, 5]}))

    plans = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=5)[0])
    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))

    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            override_batch_size=2,
            min_checkpoint_size=2,
            max_checkpoint_size=2,
        ),
        checkpoint_uri=store.root,
        batch_checkpoint_flush_interval_seconds=0,
    )

    results, cnt_udf = applier.run(plans[0])

    assert [r.offset for r in results] == [0, 2, 4]
    assert [r.span for r in results] == [2, 2, 1]
    assert cnt_udf == 5

    # Only per-map-batch checkpoints are stored (no task-level aggregate)


def test_applier_concats_checkpoints_with_time_flush(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4, 5]}))

    plans = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=5)[0])
    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))

    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            override_batch_size=2,
            min_checkpoint_size=2,
            max_checkpoint_size=2,
        ),
        checkpoint_uri=store.root,
    )

    results, cnt_udf = applier.run(plans[0])

    assert len(results) == 1
    assert results[0].offset == 0
    assert results[0].span == 5
    assert store[results[0].checkpoint_key].num_rows == 5
    assert cnt_udf == 5


def test_checkpoint_key_format(tmp_path: Path, tbl_ref: TableReference) -> None:
    @udf(version="v0", input_columns=["a"])
    def times_two(a: int) -> int:
        return a * 2

    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4]}))

    where = "a > 1"
    read_version = tbl.version

    map_task = BackfillUDFTask(
        udfs={"b": times_two},
        where=where,
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    plans = list(
        plan_read(
            tbl.uri,
            tbl_ref,
            ["a"],
            batch_size=2,
            where=where,
            read_version=read_version,
            map_task=map_task,
        )[0]
    )

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(map_task=map_task, checkpoint_uri=store.root)

    results, cnt_udf_computed = applier.run(plans[0])
    assert len(results) == 1
    key = results[0].checkpoint_key

    where_hash = hashlib.md5(where.encode()).hexdigest()
    uri_hash = hashlib.md5(tbl.uri.encode()).hexdigest()
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    prefix = (
        f"udf-{times_two.name}_ver-{times_two.version}"
        f"_col-b_where-{where_hash}_uri-{uri_hash}_srcfiles-{src_files_hash}"
    )
    expected_key = (
        f"{prefix}_frag-{plans[0].dest_frag_id()}_"
        f"range-{plans[0].dest_offset()}-{plans[0].dest_offset() + plans[0].num_rows()}"
    )

    assert key == expected_key
    assert key in store
    # First batch has a values [1,2]; only 2 satisfies where
    assert cnt_udf_computed == 1


def test_plan_read_with_legacy_checkpoint_and_partial(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    """
    Legacy fragment checkpoints are still honored even when srcfiles hashes are
    available. Partial old batch checkpoints should not skip.
    """
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))
    tbl.add(pa.table({"a": [4, 5, 6]}))  # second fragment

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"one": one},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )

    # Insert legacy fragment-level checkpoint for fragment 0
    legacy_frag_key = _legacy_fragment_dedupe_key(tbl.uri, 0, map_task)
    store[legacy_frag_key] = pa.RecordBatch.from_pydict({"file": ["fragment0.lance"]})
    assert legacy_frag_key in store
    staging_dir = tmp_path / "data"
    staging_dir.mkdir(exist_ok=True)
    (staging_dir / "fragment0.lance").touch()

    # Insert an old-format partial batch checkpoint for fragment 1 (should be ignored)
    store["fragment_1_batch_0_50:old"] = pa.RecordBatch.from_pydict(
        {"file": ["partial_batch"]}
    )

    # Legacy checkpoint should be detectable
    assert _check_fragment_data_file_exists(
        tbl.uri, 0, map_task, store, dataset_version=tbl.version
    )

    plans, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=16,
        map_task=map_task,
        checkpoint_store=store,
    )

    task_list = list(plans)

    # fragment 0 should be skipped via legacy checkpoint
    # fragment 1 should still have tasks
    assert "skipped_fragments" in pipeline_args
    assert 0 in pipeline_args["skipped_fragments"]
    frag_ids = {t.dest_frag_id() for t in task_list}
    assert 1 in frag_ids
    assert len(task_list) > 0


def test_applier_with_where(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4, 5, 6, 7, 8]}))

    plans = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=3, where="a%2=0")[0])

    assert len(plans) == 3  # 1-3, 4-6, and 7-8
    plan = plans[0]
    assert plan.uri == tbl.uri
    assert plan.offset == 0
    assert plan.limit == 3

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
    )

    # Lance forces us to eithe write the entire column or write an entire row.  This
    # applier writes the whole col.  So we actually do all the scans and filter at udf
    # execution time.  When the udf is not executed we return None.

    expected = [
        {"one": [None, 1, None], "_rowaddr": [0, 1, 2]},
        {"one": [1, None, 1], "_rowaddr": [3, 4, 5]},
        {"one": [None, 1], "_rowaddr": [6, 7]},
    ]

    expected_counts = [1, 2, 1]
    for i, plan in enumerate(plans):
        results, cnt_udf_computed = applier.run(plan)
        assert len(results) == 1
        batch = store[results[0].checkpoint_key]
        assert batch.to_pydict() == expected[i]
        assert cnt_udf_computed == expected_counts[i]


def test_applier_with_where2(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4, 5, 6, 7, 8]}))

    plans = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=1, where="a%2=0")[0])

    assert len(plans) == 8  # 1-3, 4-6, and 7-8
    plan = plans[0]
    assert plan.uri == tbl.uri
    assert plan.offset == 0
    assert plan.limit == 1

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
    )

    expected = [
        {"one": [None], "_rowaddr": [0]},
        {"one": [1], "_rowaddr": [1]},
        {"one": [None], "_rowaddr": [2]},
        {"one": [1], "_rowaddr": [3]},
        {"one": [None], "_rowaddr": [4]},
        {"one": [1], "_rowaddr": [5]},
        {"one": [None], "_rowaddr": [6]},
        {"one": [1], "_rowaddr": [7]},
    ]

    expected_counts = [0, 1, 0, 1, 0, 1, 0, 1]
    for i, plan in enumerate(plans):
        results, cnt_udf_computed = applier.run(plan)
        assert len(results) == 1
        batch = store[results[0].checkpoint_key]
        assert batch.to_pydict() == expected[i]
        assert cnt_udf_computed == expected_counts[i]


def test_applier_with_incremental(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    tbl = db.create_table(
        "tbl",
        pa.table(
            {
                "a": [1, 2, 3, 4, 5, 6, 7, 8],
                "one": [
                    None,
                    1,
                    None,
                    1,
                    None,
                    1,
                    None,
                    1,
                ],
            }
        ),
    )

    # apply a update plan that covers the rest
    plans = list(
        plan_read(
            tbl.uri,
            tbl_ref,
            ["a", "one"],  # input col and carry forward the output cols
            batch_size=1,
            carry_forward_cols=["one"],
            where="one is Null",
        )[0]
    )
    _LOG.debug(plans)

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": one},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
    )

    expected = [
        {"one": [1], "_rowaddr": [0]},
        {"one": [1], "_rowaddr": [1]},
        {"one": [1], "_rowaddr": [2]},
        {"one": [1], "_rowaddr": [3]},
        {"one": [1], "_rowaddr": [4]},
        {"one": [1], "_rowaddr": [5]},
        {"one": [1], "_rowaddr": [6]},
        {"one": [1], "_rowaddr": [7]},
    ]

    expected_counts = [1, 0, 1, 0, 1, 0, 1, 0]
    for i, plan in enumerate(plans):
        results, cnt_udf_computed = applier.run(plan)
        assert len(results) == 1
        batch = store[results[0].checkpoint_key]
        assert batch.to_pydict() == expected[i]
        assert cnt_udf_computed == expected_counts[i]


@udf()
def errors_on_three(a: int) -> int:
    if a == 3:
        raise ValueError("This is an error")
    return 1


def test_applier_error_logging(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    plans = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=16)[0])
    assert len(plans) == 1
    plan = plans[0]
    assert plan.uri == tbl.uri
    assert plan.offset == 0
    assert plan.limit == 3

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    error_store = ErrorStore(db, "test_errors")
    error_logger = TableErrorLogger(error_store=error_store, table_ref=tbl_ref)
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"one": errors_on_three},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
        error_logger=error_logger,
    )
    with pytest.raises(RuntimeError):
        applier.run(plan)

    # Verify error was logged to error store
    errors = error_store.get_errors()
    assert len(errors) == 1
    error = errors[0]
    assert error.error_message == "This is an error"
    assert error.batch_index == 0


def test_error_store(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)
    error_store = ErrorStore(db, "test_errors")
    logger = TableErrorLogger(error_store=error_store, table_ref=tbl_ref)
    logger._store.get_errors()


def test_plan_with_where(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)

    tbl = db.create_table("t", pa.table({"a": range(100)}))
    tbl.add(pa.table({"a": range(100, 200)}))
    tbl.add(pa.table({"a": range(200, 300)}))
    tbl.add(pa.table({"a": range(300, 400)}))

    fragments = tbl.get_fragments()
    assert len(fragments) == 4

    # even though we have a filter, we still have to read all the fragments
    # batch size 0 means one task per  fragment
    tasks = list(
        plan_read(
            tbl.uri, tbl_ref, ["a"], where="a > 100 AND a % 2 == 0", batch_size=0
        )[0]
    )
    # there are only 3 tasks because we skip the first fragment due to the where clause.
    assert len(tasks) == 3


def test_plan_with_row_address(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)

    tbl = db.create_table("tbl", pa.table({"a": range(100)}))

    fragments = tbl.get_fragments()
    assert len(fragments) == 1

    tasks = list(plan_read(tbl.uri, tbl_ref, ["a"], batch_size=1000)[0])
    assert len(tasks) == 1

    for batch in tasks[0].to_batches():
        assert "_rowaddr" in batch.column_names


def test_plan_with_num_frags(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)

    tbl = db.create_table("t", pa.table({"a": range(100)}))
    tbl.add(pa.table({"a": range(100, 200)}))
    tbl.add(pa.table({"a": range(200, 300)}))
    tbl.add(pa.table({"a": range(300, 400)}))

    fragments = tbl.get_fragments()
    assert len(fragments) == 4

    # even though we have a filter, we still have to read all the fragments
    tasks = list(plan_read(tbl.uri, tbl_ref, ["a"], num_frags=2)[0])
    # there are only 2 tasks because we set num_frags=2
    assert len(tasks) == 2


def test_plan_with_skip_frags(tmp_path: Path, tbl_ref: TableReference) -> None:
    db = connect(tmp_path)

    tbl = db.create_table("t", pa.table({"a": range(100)}))
    tbl.add(pa.table({"a": range(100, 200)}))
    tbl.add(pa.table({"a": range(200, 300)}))
    tbl.add(pa.table({"a": range(300, 400)}))

    fragments = tbl.get_fragments()
    assert len(fragments) == 4

    # skip_frags=2 skips the first 2 fragments, returns remaining 2
    tasks = list(plan_read(tbl.uri, tbl_ref, ["a"], skip_frags=2)[0])
    assert len(tasks) == 2

    # Verify the tasks correspond to the correct fragments (fragments 2 and 3)
    frag_ids = {t.dest_frag_id() for t in tasks}
    expected_frag_ids = {fragments[2].fragment_id, fragments[3].fragment_id}
    assert frag_ids == expected_frag_ids


def test_plan_with_skip_frags_and_num_frags(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)

    tbl = db.create_table("t", pa.table({"a": range(100)}))
    tbl.add(pa.table({"a": range(100, 200)}))
    tbl.add(pa.table({"a": range(200, 300)}))
    tbl.add(pa.table({"a": range(300, 400)}))

    fragments = tbl.get_fragments()
    assert len(fragments) == 4

    # skip_frags=1, num_frags=2 → fragments 1 and 2 only
    tasks = list(plan_read(tbl.uri, tbl_ref, ["a"], skip_frags=1, num_frags=2)[0])
    assert len(tasks) == 2

    frag_ids = {t.dest_frag_id() for t in tasks}
    expected_frag_ids = {fragments[1].fragment_id, fragments[2].fragment_id}
    assert frag_ids == expected_frag_ids


def test_plan_with_skip_frags_beyond_end(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)

    tbl = db.create_table("t", pa.table({"a": range(100)}))
    tbl.add(pa.table({"a": range(100, 200)}))

    fragments = tbl.get_fragments()
    assert len(fragments) == 2

    # skip_frags beyond total fragments → no tasks
    tasks = list(plan_read(tbl.uri, tbl_ref, ["a"], skip_frags=5)[0])
    assert len(tasks) == 0


def test_udf_with_arrow_params(tmp_path: Path, tbl_ref: TableReference) -> None:
    @udf(data_type=pa.int32())
    def batch_udf(a: pa.Array, b: pa.Array) -> pa.Array:
        assert a == pa.array([1, 2, 3])
        assert b == pa.array([4, 5, 6])
        return pc.cast(pc.add(a, b), pa.int32())

    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"c": batch_udf},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
    )
    results, cnt_udf_computed = applier.run(
        next(plan_read(tbl.uri, tbl_ref, ["a", "b"], batch_size=16)[0])
    )
    assert len(results) == 1
    batch = store[results[0].checkpoint_key]
    assert batch == pa.RecordBatch.from_pydict(
        {
            "c": pa.array([5, 7, 9], type=pa.int32()),
            "_rowaddr": pa.array([0, 1, 2], pa.uint64()),
        },
    )
    assert cnt_udf_computed == 3


def test_udf_with_arrow_struct(tmp_path: Path, tbl_ref: TableReference) -> None:
    struct_type = pa.struct([("rpad", pa.string()), ("lpad", pa.string())])

    @udf(data_type=struct_type)
    def struct_udf(a: pa.Array, b: pa.Array) -> pa.Array:
        assert a == pa.array([1, 2, 3])
        assert b == pa.array([4, 5, 6])
        rpad = pc.ascii_rpad(pc.cast(a, target_type="string"), 4, padding="0")
        lpad = pc.ascii_lpad(pc.cast(a, target_type="string"), 4, padding="0")
        return pc.make_struct(rpad, lpad, field_names=["rpad", "lpad"])

    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"c": struct_udf},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
    )
    results, cnt_udf_computed = applier.run(
        next(plan_read(tbl.uri, tbl_ref, ["a", "b"], batch_size=16)[0])
    )
    assert len(results) == 1
    batch = store[results[0].checkpoint_key]
    # Build the expected RecordBatch
    # The function calls produce ["1000", "2000", "3000"] for rpad
    # and ["0001", "0002", "0003"] for lpad
    expected_batch = pa.RecordBatch.from_arrays(
        [
            pa.StructArray.from_arrays(
                [
                    pa.array(["1000", "2000", "3000"]),
                    pa.array(["0001", "0002", "0003"]),
                ],
                names=["rpad", "lpad"],
            ),
            pa.array([0, 1, 2], pa.uint64()),
        ],
        ["c", "_rowaddr"],
    )

    assert batch == expected_batch
    assert cnt_udf_computed == 3


def test_plan_read_supports_struct_field_projection(tmp_path: Path) -> None:
    # Create a dataset with a struct column
    struct_type = pa.struct([("left", pa.string()), ("right", pa.string())])
    tbl = pa.table(
        {
            "info": pa.array(
                [
                    {"left": "alpha", "right": "one"},
                    {"left": "beta", "right": "two"},
                ],
                type=struct_type,
            )
        }
    )

    ds_path = tmp_path / "ds.lance"
    lance.write_dataset(tbl, ds_path, max_rows_per_file=16)

    tbl_ref = TableReference(table_id=["ds"], version=None, db_uri=str(tmp_path))

    # columns includes a dotted struct field
    plan, _ = plan_read(str(ds_path), tbl_ref, columns=["info.left"], batch_size=1024)

    task = next(iter(plan))
    batches = list(task.to_batches())
    assert len(batches) == 1
    batch = batches[0]
    # Should project only the requested sub-field plus row address
    assert "info.left" in batch.schema.names
    assert "info" not in batch.schema.names


def test_plan_read_with_udf_projects_struct_field(tmp_path: Path) -> None:
    struct_type = pa.struct([("left", pa.string()), ("right", pa.string())])
    tbl = pa.table(
        {
            "info": pa.array(
                [
                    {"left": "alpha", "right": "one"},
                    {"left": "beta", "right": "two"},
                ],
                type=struct_type,
            )
        }
    )

    ds_path = tmp_path / "ds_udf.lance"
    lance.write_dataset(tbl, ds_path, max_rows_per_file=16)

    tbl_ref = TableReference(table_id=["ds_udf"], version=None, db_uri=str(tmp_path))

    @udf(data_type=pa.string(), input_columns=["info.left"])
    def left_upper(left: pa.Array) -> pa.Array:  # pyright: ignore[reportReturnType]
        return pc.utf8_upper(left)

    map_task = BackfillUDFTask(
        udfs={"c": left_upper},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )

    plan, _ = plan_read(
        str(ds_path),
        tbl_ref,
        columns=["info.left"],
        batch_size=1024,
        map_task=map_task,
    )

    task = next(iter(plan))
    batches = list(task.to_batches())
    assert len(batches) == 1
    names = set(batches[0].schema.names)
    assert "info.left" in names
    assert "info" not in names


def test_udf_with_arrow_array(tmp_path: Path, tbl_ref: TableReference) -> None:
    array_type = pa.list_(pa.int64())

    @udf(data_type=array_type)
    def array_udf(a: pa.Array, b: pa.Array) -> pa.Array:
        assert a == pa.array([1, 2, 3])
        assert b == pa.array([4, 5, 6])
        arr = [
            [val] * cnt for val, cnt in zip(a.to_pylist(), b.to_pylist(), strict=True)
        ]
        c = pa.array(arr, type=pa.list_(pa.int64()))
        return c

    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    applier = CheckpointingApplier(
        map_task=BackfillUDFTask(
            udfs={"c": array_udf},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_uri=store.root,
    )
    results, cnt_udf_computed = applier.run(
        next(plan_read(tbl.uri, tbl_ref, ["a", "b"], batch_size=16)[0])
    )
    assert len(results) == 1
    batch = store[results[0].checkpoint_key]

    # Build the expected RecordBatch
    expected_c = pa.array(
        [[1, 1, 1, 1], [2, 2, 2, 2, 2], [3, 3, 3, 3, 3, 3]], type=pa.list_(pa.int64())
    )

    expected_batch = pa.RecordBatch.from_arrays(
        [expected_c, pa.array([0, 1, 2], pa.uint64())], ["c", "_rowaddr"]
    )
    assert batch == expected_batch
    assert cnt_udf_computed == 3


# Tests for fragment-level checkpoint functionality


def test_parse_checkpoint_range_key_valid() -> None:
    prefix = "udf-test_frag-x_ver-1_col-a_where-abc_uri-def"
    key = f"{prefix}_frag-12_range-3-10"
    parsed = _parse_checkpoint_range_key(key)
    assert parsed == (prefix, 12, 3, 10)


@pytest.mark.parametrize(
    "key",
    [
        "no-frag-or-range",
        "prefix_frag-1_range-10-10",  # empty range
        "prefix_frag-1_range-10-5",  # inverted range
        "prefix_frag-abc_range-0-1",  # invalid frag id
        "prefix_frag-1_range-foo-2",  # invalid start
        "prefix_range-0-1",  # missing frag marker
        "prefix_frag-1_range-0-1-extra",  # trailing junk
    ],
)
def test_parse_checkpoint_range_key_invalid(key: str) -> None:
    assert _parse_checkpoint_range_key(key) is None


def test_index_checkpoint_ranges_in_memory_store() -> None:
    store = CheckpointStore.from_uri("memory")
    prefix_a = "udf-a_ver-1_col-x_where-h_uri-u"
    prefix_b = "udf-b_ver-2_col-y_where-h_uri-u"

    valid_keys = [
        f"{prefix_a}_frag-0_range-0-10",
        f"{prefix_a}_frag-0_range-10-20",
        f"{prefix_a}_frag-2_range-5-9",
        f"{prefix_b}_frag-1_range-0-4",
    ]
    invalid_keys = [
        "garbage",
        f"{prefix_a}_frag-0_range-5-5",
        f"{prefix_a}_frag-foo_range-0-1",
        f"{prefix_a}_range-0-1",
    ]

    batch = pa.RecordBatch.from_pydict({"a": [1]})
    for key in valid_keys + invalid_keys:
        store[key] = batch

    all_keys, ranges_by_prefix = _index_checkpoint_ranges(checkpoint_store=store)

    assert all_keys == set(valid_keys + invalid_keys)
    assert sorted(ranges_by_prefix[prefix_a][0]) == [(0, 10), (10, 20)]
    assert ranges_by_prefix[prefix_a][2] == [(5, 9)]
    assert ranges_by_prefix[prefix_b][1] == [(0, 4)]


def test_check_fragment_data_file_exists_no_checkpoint(tmp_path: Path) -> None:
    """Test _check_fragment_data_file_exists when fragment is not checkpointed."""
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"one": one},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])

    # Fragment 0 should not exist in checkpoint store yet
    exists = _check_fragment_data_file_exists(
        tbl.uri, 0, map_task, store, src_files_hash=src_files_hash
    )
    assert not exists


def test_check_fragment_data_file_exists_with_staging_file(tmp_path: Path) -> None:
    """Test _check_fragment_data_file_exists when fragment file exists in staging."""
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"one": one},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])

    # Create a checkpoint entry for fragment 0
    dedupe_key = _get_fragment_dedupe_key(
        tbl.uri, 0, map_task, src_files_hash=src_files_hash
    )
    fake_file_path = "test_fragment_0.lance"
    store[dedupe_key] = pa.RecordBatch.from_pydict({"file": [fake_file_path]})

    # Create the staging file
    staging_dir = tmp_path / "data"
    staging_dir.mkdir(exist_ok=True)
    staging_file = staging_dir / fake_file_path
    staging_file.touch()

    # Should return True since file exists in staging
    exists = _check_fragment_data_file_exists(
        tbl.uri, 0, map_task, store, src_files_hash=src_files_hash
    )
    assert exists


def test_check_fragment_data_file_exists_with_cloud_url() -> None:
    """Test _check_fragment_data_file_exists with cloud URLs."""
    # Create a mock checkpoint store
    store = {}
    map_task = BackfillUDFTask(
        udfs={"one": one},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )

    # Test with S3 URL - should not crash but return False since no real file
    s3_uri = "s3://test-bucket/dataset"
    exists = _check_fragment_data_file_exists(s3_uri, 0, map_task, store)
    assert not exists

    # Test with GCS URL - should not crash but return False since no real file
    gcs_uri = "gs://test-bucket/dataset"
    exists = _check_fragment_data_file_exists(gcs_uri, 0, map_task, store)
    assert not exists


def test_find_output_data_file_requires_full_output_field_coverage() -> None:
    class _FakeFragment:
        def __init__(self, data_files: list[object]) -> None:
            self._data_files = data_files

        def data_files(self) -> list[object]:
            return self._data_files

    partial_match = type("FakeDataFile", (), {"fields": [3, 5]})()
    full_match = type("FakeDataFile", (), {"fields": [5, 6, 7]})()
    frag = _FakeFragment([partial_match, full_match])

    assert _find_output_data_file_in_fragment(frag, frozenset({5, 6})) is full_match
    assert _find_output_data_file_in_fragment(frag, frozenset({5, 8})) is None


def test_plan_read_validates_checkpoint_payload_before_skipping_committed_fragment(
    tmp_path: Path, tbl_ref: TableReference, monkeypatch: pytest.MonkeyPatch
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "a2": [2, 4, 6]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"a2": double_a},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    dataset = tbl.to_lance()
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    dedupe_key = _get_fragment_dedupe_key(
        tbl.uri,
        0,
        map_task,
        dataset_version=dataset.version,
        src_files_hash=src_files_hash,
    )
    expected_path = next(
        df.path
        for df in dataset.get_fragment(0).data_files()
        if set(df.fields) >= set(extract_field_ids(dataset.lance_schema, "a2"))
    )
    store[dedupe_key] = pa.RecordBatch.from_pydict(
        {
            "file": [expected_path],
            "output_field_ids": [
                json.dumps(sorted(extract_field_ids(dataset.lance_schema, "a2")))
            ],
        }
    )

    checkpoint_probe_calls: list[tuple[object, ...]] = []
    original_probe = apply_module._check_fragment_data_file_exists

    def _record_checkpoint_probe(*args: object, **kwargs: object) -> str | None:
        checkpoint_probe_calls.append(args)
        return original_probe(*args, **kwargs)

    monkeypatch.setattr(
        apply_module,
        "_check_fragment_data_file_exists",
        _record_checkpoint_probe,
    )

    plans, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=16,
        map_task=map_task,
        checkpoint_store=store,
    )

    assert list(plans) == []
    skipped_fragments = pipeline_args["skipped_fragments"]
    assert 0 in skipped_fragments
    data_file, row_count = skipped_fragments[0]
    assert row_count == 3
    assert checkpoint_probe_calls
    assert set(data_file.fields) >= set(extract_field_ids(dataset.lance_schema, "a2"))


def test_plan_read_does_not_skip_committed_fragment_with_stale_checkpoint_payload(
    tmp_path: Path, tbl_ref: TableReference, monkeypatch: pytest.MonkeyPatch
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "a2": [2, 4, 6]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"a2": double_a},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    dataset = tbl.to_lance()
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    dedupe_key = _get_fragment_dedupe_key(
        tbl.uri,
        0,
        map_task,
        dataset_version=dataset.version,
        src_files_hash=src_files_hash,
    )
    store[dedupe_key] = pa.RecordBatch.from_pydict({"file": ["stale-file.lance"]})

    checkpoint_probe_calls = 0

    def _stale_checkpoint_payload(*args: object, **kwargs: object) -> None:
        nonlocal checkpoint_probe_calls
        checkpoint_probe_calls += 1
        return None

    monkeypatch.setattr(
        apply_module,
        "_check_fragment_data_file_exists",
        _stale_checkpoint_payload,
    )

    plans, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=16,
        map_task=map_task,
        checkpoint_store=store,
    )

    task_list = list(plans)
    assert checkpoint_probe_calls == 1
    assert len(task_list) == 1
    assert pipeline_args["skipped_fragments"] == {}


def test_plan_read_counts_full_fragment_rows_for_committed_incremental_fragments(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "a2": [2, 4, 6]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"a2": double_a},
        where="a2 IS NULL",
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    dataset = tbl.to_lance()
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    dedupe_key = _get_fragment_dedupe_key(
        tbl.uri,
        0,
        map_task,
        dataset_version=dataset.version,
        src_files_hash=src_files_hash,
    )
    expected_path = next(
        df.path
        for df in dataset.get_fragment(0).data_files()
        if set(df.fields) >= set(extract_field_ids(dataset.lance_schema, "a2"))
    )
    store[dedupe_key] = pa.RecordBatch.from_pydict(
        {
            "file": [expected_path],
            "output_field_ids": [
                json.dumps(sorted(extract_field_ids(dataset.lance_schema, "a2")))
            ],
        }
    )

    plans, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        where="a2 IS NULL",
        batch_size=16,
        map_task=map_task,
        checkpoint_store=store,
    )

    assert list(plans) == []
    assert pipeline_args["skipped_fragments"][0][1] == 3
    assert pipeline_args["skipped_stats"]["rows"] == 3


@pytest.mark.parametrize("use_legacy_key", [False, True], ids=["primary", "legacy"])
def test_plan_read_falls_back_when_checkpoint_indexing_fails(
    tmp_path: Path,
    tbl_ref: TableReference,
    monkeypatch: pytest.MonkeyPatch,
    use_legacy_key: bool,
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "a2": [2, 4, 6]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"a2": double_a},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    dataset = tbl.to_lance()
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    if use_legacy_key:
        checkpoint_key = _legacy_fragment_dedupe_key(tbl.uri, 0, map_task)
    else:
        checkpoint_key = _get_fragment_dedupe_key(
            tbl.uri,
            0,
            map_task,
            dataset_version=dataset.version,
            src_files_hash=src_files_hash,
        )
    expected_path = next(
        df.path
        for df in dataset.get_fragment(0).data_files()
        if set(df.fields) >= set(extract_field_ids(dataset.lance_schema, "a2"))
    )
    store[checkpoint_key] = pa.RecordBatch.from_pydict(
        {
            "file": [expected_path],
            "output_field_ids": [
                json.dumps(sorted(extract_field_ids(dataset.lance_schema, "a2")))
            ],
        }
    )

    def _raise_index_error(*args: object, **kwargs: object) -> NoReturn:
        raise RuntimeError("checkpoint listing unavailable")

    checkpoint_probe_calls = 0

    original_probe = apply_module._check_fragment_data_file_exists

    def _record_checkpoint_probe(*args: object, **kwargs: object) -> str | None:
        nonlocal checkpoint_probe_calls
        checkpoint_probe_calls += 1
        return original_probe(*args, **kwargs)

    monkeypatch.setattr(apply_module, "_index_checkpoint_ranges", _raise_index_error)
    monkeypatch.setattr(
        apply_module,
        "_check_fragment_data_file_exists",
        _record_checkpoint_probe,
    )

    plans, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=16,
        map_task=map_task,
        checkpoint_store=store,
    )

    assert list(plans) == []
    assert checkpoint_probe_calls == 1
    assert pipeline_args["skipped_fragments"][0][1] == 3


def test_plan_read_with_skipped_fragments(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    """Test that plan_read correctly identifies and skips fragments."""
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))
    tbl.add(pa.table({"a": [4, 5, 6]}))  # Add second fragment

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"one": one},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])

    # Create checkpoint for fragment 0 only
    dedupe_key = _get_fragment_dedupe_key(
        tbl.uri, 0, map_task, src_files_hash=src_files_hash
    )
    fake_file_path = "fragment_0.lance"
    store[dedupe_key] = pa.RecordBatch.from_pydict({"file": [fake_file_path]})

    # Create the staging file for fragment 0
    staging_dir = tmp_path / "data"
    staging_dir.mkdir(exist_ok=True)
    staging_file = staging_dir / fake_file_path
    staging_file.touch()

    # Plan read with checkpoint information
    plans, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=16,
        map_task=map_task,
        checkpoint_store=store,
    )

    # Should still have tasks for fragment 1 (not checkpointed)
    task_list = list(plans)
    assert len(task_list) > 0  # Fragment 1 should have tasks

    # Check that skipped_fragments contains fragment 0
    assert "skipped_fragments" in pipeline_args
    skipped_fragments = pipeline_args["skipped_fragments"]
    assert 0 in skipped_fragments
    assert 1 not in skipped_fragments  # Fragment 1 should not be skipped

    # Verify the DataFile uses the dataset's data_storage_version
    data_file, row_count = skipped_fragments[0]
    assert data_file.file_major_version == 2
    assert data_file.file_minor_version == 0
    assert row_count > 0


@pytest.mark.parametrize("data_storage_version", ["2.0", "2.1"])
def test_plan_read_skipped_fragment_uses_dataset_version(
    tmp_path: Path, tbl_ref: TableReference, data_storage_version: str
) -> None:
    """Skipped-fragment DataFile must match the dataset's data_storage_version."""
    # Create dataset with explicit data_storage_version
    data = pa.table({"a": [1, 2, 3]})
    ds_path = str(tmp_path / "tbl.lance")
    lance.write_dataset(data, ds_path, data_storage_version=data_storage_version)
    lance.write_dataset(
        pa.table({"a": [4, 5, 6]}),
        ds_path,
        mode="append",
    )

    db = connect(tmp_path)
    tbl = db.open_table("tbl")

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    tbl.add_columns({"a2": double_a})

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"a2": double_a},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])

    # Checkpoint fragment 0
    dedupe_key = _get_fragment_dedupe_key(
        tbl.uri, 0, map_task, src_files_hash=src_files_hash
    )
    fake_file = "fragment_0.lance"
    store[dedupe_key] = pa.RecordBatch.from_pydict({"file": [fake_file]})
    (tmp_path / "data").mkdir(exist_ok=True)
    (tmp_path / "data" / fake_file).touch()

    _, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        batch_size=16,
        map_task=map_task,
        checkpoint_store=store,
    )

    skipped = pipeline_args["skipped_fragments"]
    assert 0 in skipped
    df, row_count = skipped[0]
    major, minor = data_storage_version.split(".")
    assert df.file_major_version == int(major)
    assert df.file_minor_version == int(minor)
    assert row_count > 0


def test_plan_read_no_checkpointing_params(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    """Test that plan_read works normally when no checkpointing params are provided."""
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    # Plan read without checkpoint information
    plans, pipeline_args = plan_read(tbl.uri, tbl_ref, ["a"], batch_size=16)

    # Should have tasks for all fragments
    task_list = list(plans)
    assert len(task_list) == 1

    # Should have empty skipped_fragments
    assert "skipped_fragments" in pipeline_args
    skipped_fragments = pipeline_args["skipped_fragments"]
    assert len(skipped_fragments) == 0


def test_plan_read_caches_dataset_version(
    tbl_ref: TableReference, monkeypatch: pytest.MonkeyPatch
) -> None:
    class _FakeFragment:
        def __init__(self, fragment_id: int, rows: int) -> None:
            self.fragment_id = fragment_id
            self.physical_rows = rows
            self.num_deletions = 0

    class _FakeDataset:
        def __init__(self) -> None:
            self._version_reads = 0
            self._fragments = [_FakeFragment(0, 3), _FakeFragment(1, 2)]

        @property
        def version(self) -> int:
            self._version_reads += 1
            return 17

        def get_fragments(self) -> list[_FakeFragment]:
            return self._fragments

    fake_dataset = _FakeDataset()

    monkeypatch.setattr(apply_module.lance, "dataset", lambda uri: fake_dataset)

    plans, _ = plan_read("memory://fake", tbl_ref, ["a"], batch_size=2)

    task_list = list(plans)
    assert [task.version for task in task_list] == [17, 17, 17]
    assert fake_dataset._version_reads == 1


def test_plan_read_uses_logical_rows_for_fragments_with_deletes(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [0, 1, 2, 3]}))
    tbl.delete("a % 2 == 1")

    plans, pipeline_args = plan_read(tbl.uri, tbl_ref, ["a"], batch_size=1)

    task_list = list(plans)
    assert len(task_list) == 2
    assert [task.offset for task in task_list] == [0, 1]
    assert len(pipeline_args["skipped_fragments"]) == 0


def test_plan_read_len_matches_actual_filtered_tasks(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "a2": [2, 4, 6]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    plans, _ = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        where="a2 IS NULL",
        batch_size=16,
        map_task=BackfillUDFTask(
            udfs={"a2": double_a},
            where="a2 IS NULL",
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_store=InMemoryCheckpointStore(),
        default_where_generated=True,
    )

    assert len(plans) == 0
    task_list = list(plans)
    assert len(task_list) == 0


def test_plan_read_skips_generated_default_where_filter_count(
    tmp_path: Path, tbl_ref: TableReference, monkeypatch: pytest.MonkeyPatch
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    tbl.add_columns({"a2": double_a})
    dataset = tbl.to_lance()
    frag_type = type(dataset.get_fragment(0))
    original_count_rows = frag_type.count_rows
    filter_calls: list[str] = []

    def _count_rows(self, *args: object, **kwargs: object) -> int:
        filter_expr = kwargs.get("filter")
        if args:
            filter_expr = args[0]
        if filter_expr is not None:
            filter_calls.append(str(filter_expr))
        return original_count_rows(self, *args, **kwargs)

    monkeypatch.setattr(frag_type, "count_rows", _count_rows)

    plans, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        where="a2 IS NULL",
        batch_size=16,
        map_task=BackfillUDFTask(
            udfs={"a2": double_a},
            where="a2 IS NULL",
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_store=InMemoryCheckpointStore(),
        default_where_generated=True,
    )

    assert len(list(plans)) == 1
    assert filter_calls == []


def test_plan_read_keeps_explicit_where_filter_count(
    tmp_path: Path, tbl_ref: TableReference, monkeypatch: pytest.MonkeyPatch
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    tbl.add_columns({"a2": double_a})
    dataset = tbl.to_lance()
    frag_type = type(dataset.get_fragment(0))
    original_count_rows = frag_type.count_rows
    filter_calls: list[str] = []

    def _count_rows(self, *args: object, **kwargs: object) -> int:
        filter_expr = kwargs.get("filter")
        if args:
            filter_expr = args[0]
        if filter_expr is not None:
            filter_calls.append(str(filter_expr))
        return original_count_rows(self, *args, **kwargs)

    monkeypatch.setattr(frag_type, "count_rows", _count_rows)

    plans, _ = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        where="a2 IS NULL",
        batch_size=16,
        map_task=BackfillUDFTask(
            udfs={"a2": double_a},
            where="a2 IS NULL",
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_store=InMemoryCheckpointStore(),
        default_where_generated=False,
    )

    assert len(list(plans)) == 1
    assert filter_calls == ["a2 IS NULL"]


def test_plan_read_generated_default_where_keeps_output_file_safety(
    tmp_path: Path, tbl_ref: TableReference, monkeypatch: pytest.MonkeyPatch
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3], "a2": [2, 4, 6]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    dataset = tbl.to_lance()
    frag_type = type(dataset.get_fragment(0))
    original_count_rows = frag_type.count_rows
    filter_calls: list[str] = []

    def _count_rows(self, *args: object, **kwargs: object) -> int:
        filter_expr = kwargs.get("filter")
        if args:
            filter_expr = args[0]
        if filter_expr is not None:
            filter_calls.append(str(filter_expr))
        return original_count_rows(self, *args, **kwargs)

    monkeypatch.setattr(frag_type, "count_rows", _count_rows)

    plans, _ = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        where="a2 IS NULL",
        batch_size=16,
        map_task=BackfillUDFTask(
            udfs={"a2": double_a},
            where="a2 IS NULL",
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_store=InMemoryCheckpointStore(),
        default_where_generated=True,
    )

    assert list(plans) == []
    assert filter_calls == ["a2 IS NULL"]


def test_plan_read_skips_hash_and_dedupe_for_empty_checkpoint_index(
    tmp_path: Path, tbl_ref: TableReference, monkeypatch: pytest.MonkeyPatch
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    tbl.add_columns({"a2": double_a})

    def _unexpected_hash(*args: object, **kwargs: object) -> NoReturn:
        raise AssertionError("empty checkpoint path should not hash source files")

    def _unexpected_dedupe(*args: object, **kwargs: object) -> NoReturn:
        raise AssertionError("empty checkpoint path should not build dedupe keys")

    monkeypatch.setattr(apply_module, "hash_source_files", _unexpected_hash)
    monkeypatch.setattr(pipeline_module, "_get_fragment_dedupe_key", _unexpected_dedupe)

    plans, _ = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        where="a2 IS NULL",
        batch_size=16,
        map_task=BackfillUDFTask(
            udfs={"a2": double_a},
            where="a2 IS NULL",
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        ),
        checkpoint_store=InMemoryCheckpointStore(),
        default_where_generated=True,
    )

    assert len(list(plans)) == 1


def test_plan_read_generated_default_where_preserves_partial_checkpoint_gaps(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3, 4, 5, 6]}))

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    tbl.add_columns({"a2": double_a})
    dataset = tbl.to_lance()
    map_task = BackfillUDFTask(
        udfs={"a2": double_a},
        where="a2 IS NULL",
        override_batch_size=2,
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )
    store = InMemoryCheckpointStore()
    src_files_hash = _src_files_hash_for_cols(tbl, ["a"])
    for start, end in [(0, 2), (2, 4)]:
        key = map_task.checkpoint_key(
            dataset_uri=tbl.uri,
            dataset_version=dataset.version,
            frag_id=0,
            start=start,
            end=end,
            where="a2 IS NULL",
            src_files_hash=src_files_hash,
        )
        store[key] = pa.record_batch([], names=[])

    tasks, pipeline_args = plan_read(
        tbl.uri,
        tbl_ref,
        ["a"],
        where="a2 IS NULL",
        batch_size=2,
        map_task=map_task,
        checkpoint_store=store,
        default_where_generated=True,
    )

    task_list = list(tasks)
    assert [(task.offset, task.limit) for task in task_list] == [(4, 2)]
    assert pipeline_args["skipped_stats"]["rows"] == 0


@pytest.mark.ray
def test_fragment_writer_manager_with_skipped_fragments(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    """Test that FragmentWriterManager correctly handles skipped fragments."""
    import lance.fragment
    import ray

    # Start Ray if not already started
    if not ray.is_initialized():
        ray.init(local_mode=True, ignore_reinit_error=True)

    try:
        db = connect(tmp_path)
        tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

        store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
        map_task = BackfillUDFTask(
            udfs={"one": one},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        )

        # Create a mock skipped fragment data file
        skipped_data_file = lance.fragment.DataFile(
            "skipped_fragment.lance",
            [],  # field_ids
            [],  # field_id_to_column_indices
            2,  # major_version
            0,  # minor_version
        )

        skipped_fragments = {0: (skipped_data_file, 3)}

        # Create FragmentWriterManager with skipped fragments
        fwm = FragmentWriterManager(
            dst_read_version=tbl.version,
            ds_uri=tbl.uri,
            job_tracker=None,
            map_task=map_task,
            checkpoint_store=store,
            where=None,
            commit_granularity=1,
            expected_tasks={1: 1},  # Only fragment 1 has expected tasks
            skipped_fragments=skipped_fragments,
        )

        # Check that skipped fragment is immediately in to_commit
        assert len(fwm.to_commit) == 1
        frag_id, data_file, row_count = fwm.to_commit[0]
        assert frag_id == 0
        assert data_file == skipped_data_file
        assert row_count >= 0  # Row count should be determined

    finally:
        if ray.is_initialized():
            ray.shutdown()


@pytest.mark.ray
def test_fragment_writer_manager_no_skipped_fragments(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    """Test that FragmentWriterManager works normally with no skipped fragments."""
    import ray

    # Start Ray if not already started
    if not ray.is_initialized():
        ray.init(local_mode=True, ignore_reinit_error=True)

    try:
        db = connect(tmp_path)
        tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

        store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
        map_task = BackfillUDFTask(
            udfs={"one": one},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        )

        # Create FragmentWriterManager with no skipped fragments
        fwm = FragmentWriterManager(
            dst_read_version=tbl.version,
            ds_uri=tbl.uri,
            map_task=map_task,
            checkpoint_store=store,
            where=None,
            job_tracker=None,
            commit_granularity=1,
            expected_tasks={0: 1},  # Fragment 0 has expected tasks
            skipped_fragments={},  # No skipped fragments
        )

        # Should have no items in to_commit initially
        assert len(fwm.to_commit) == 0

    finally:
        if ray.is_initialized():
            ray.shutdown()


@pytest.mark.ray
def test_fragment_writer_manager_mixed_fragments(
    tmp_path: Path, tbl_ref: TableReference
) -> None:
    """Test FragmentWriterManager with both skipped and normal fragments."""
    import lance.fragment
    import ray

    # Start Ray if not already started
    if not ray.is_initialized():
        ray.init(local_mode=True, ignore_reinit_error=True)

    try:
        db = connect(tmp_path)
        tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))
        tbl.add(pa.table({"a": [4, 5, 6]}))  # Add second fragment

        store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
        map_task = BackfillUDFTask(
            udfs={"one": one},
            min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
            max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        )

        # Create a mock skipped fragment data file for fragment 0
        skipped_data_file = lance.fragment.DataFile(
            "skipped_fragment_0.lance",
            [],  # field_ids
            [],  # field_id_to_column_indices
            2,  # major_version
            0,  # minor_version
        )

        skipped_fragments = {0: (skipped_data_file, 3)}

        # Create FragmentWriterManager with mixed fragments
        fwm = FragmentWriterManager(
            dst_read_version=tbl.version,
            ds_uri=tbl.uri,
            map_task=map_task,
            checkpoint_store=store,
            where=None,
            job_tracker=None,
            commit_granularity=1,
            expected_tasks={1: 1},  # Only fragment 1 has expected tasks to process
            skipped_fragments=skipped_fragments,
        )

        # Should have 1 item in to_commit (the skipped fragment)
        assert len(fwm.to_commit) == 1
        frag_id, data_file, row_count = fwm.to_commit[0]
        assert frag_id == 0
        assert data_file == skipped_data_file

        # Fragment 1 should still be tracked in remaining_tasks
        assert 1 in fwm.remaining_tasks
        assert fwm.remaining_tasks[1] == 1

    finally:
        if ray.is_initialized():
            ray.shutdown()


def test_fragment_writer_manager_dedupes_fragment_records(tmp_path: Path) -> None:
    """Ensure duplicate fragment records don't double count or re-enqueue."""
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"one": one},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )

    fwm = FragmentWriterManager(
        dst_read_version=tbl.version,
        ds_uri=tbl.uri,
        job_tracker=None,
        map_task=map_task,
        checkpoint_store=store,
        where=None,
        commit_granularity=999,
        expected_tasks={0: 1},
        skipped_fragments={},
    )

    data_file = lance.fragment.DataFile(
        "fragment_0.lance",
        [],
        [],
        2,
        0,
    )
    fwm.rows_input_by_frag[0] = 3

    fwm._record_fragment(0, data_file, commit_granularity=999, rows_written=3)
    fwm._record_fragment(0, data_file, commit_granularity=999, rows_written=3)

    assert len(fwm.to_commit) == 1
    assert fwm._reconciled_rows_ready_total == 3


def test_fragment_writer_manager_dedupes_duplicate_task_ingest(tmp_path: Path) -> None:
    """Ensure replayed read-task results are not ingested twice."""
    db = connect(tmp_path)
    tbl = db.create_table("tbl", pa.table({"a": [1, 2, 3]}))

    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"one": one},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )

    fwm = FragmentWriterManager(
        dst_read_version=tbl.version,
        ds_uri=tbl.uri,
        job_tracker=None,
        map_task=map_task,
        checkpoint_store=store,
        where=None,
        commit_granularity=999,
        expected_tasks={0: 1},
        skipped_fragments={},
    )

    class _DummySession:
        failed = False
        failure_reason = None

        def __init__(self) -> None:
            self.sealed = False
            self.ingested: list[tuple[int, str]] = []

        def ingest_task(self, offset: int, checkpoint_key: str) -> None:
            self.ingested.append((offset, checkpoint_key))

        def seal(self) -> None:
            self.sealed = True

    sess = _DummySession()
    fwm.sessions[0] = sess

    task = ScanTask(
        uri=tbl.uri,
        table_ref=tbl.get_reference(),
        columns=["a"],
        frag_id=0,
        offset=0,
        limit=3,
        version=tbl.version,
        with_row_address=True,
    )
    checkpoints = [
        MapBatchCheckpoint(
            checkpoint_key="frag-0_range-0-3",
            offset=0,
            num_rows=3,
            span=3,
            udf_rows=3,
        )
    ]

    fwm.ingest_task(task, checkpoints)
    fwm.ingest_task(task, checkpoints)

    assert sess.ingested == [(0, "frag-0_range-0-3")]
    assert fwm._reconciled_rows_checkpointed_total == 3
    assert fwm.rows_input_by_frag[0] == 3
    assert fwm.remaining_tasks[0] == 0
    assert sess.sealed is True


def test_fragment_writer_uses_precomputed_metadata(tmp_path: Path, monkeypatch) -> None:
    """Ensure writer avoids opening dataset when metadata is provided."""
    from collections import deque

    tbl_path = tmp_path / "tbl.lance"
    ds = lance.write_dataset(pa.table({"a": [1, 2, 3]}), tbl_path)
    frag = ds.get_fragment(0)
    assert frag is not None

    num_physical_rows = frag.physical_rows
    num_logical_rows = frag.count_rows()

    store = CheckpointStore.from_uri(str(tmp_path / "ckp"))
    checkpoint_key = "fragment_0_batch_0"
    rowaddr = pa.array([0, 1, 2], type=pa.uint64())
    batch = pa.record_batch([pa.array([10, 11, 12]), rowaddr], names=["a", "_rowaddr"])
    store[checkpoint_key] = batch

    class _Queue:
        def __init__(self, items) -> None:
            self._items = deque(items)

        def get(self) -> tuple[int, str]:
            return self._items.popleft()

    queue = _Queue([(0, checkpoint_key), (-1, "")])

    base_schema = ds.schema
    fields = [base_schema.field("a"), pa.field("_rowaddr", pa.uint64())]
    filler_schema = pa.schema(fields)
    field_ids = extract_field_ids(ds.lance_schema, "a")

    def _boom(*args, **kwargs) -> NoReturn:
        raise AssertionError("lance.dataset should not be called")

    monkeypatch.setattr(lance, "dataset", _boom)

    writer_cls = FragmentWriter.__ray_actor_class__
    writer = writer_cls(
        ds.uri,
        ["a"],
        store.uri(),
        0,
        queue,
        read_version=ds.version,
        filler_schema=filler_schema,
        field_ids=field_ids,
        num_physical_rows=num_physical_rows,
        num_logical_rows=num_logical_rows,
        data_storage_version=ds.data_storage_version,
    )

    result = writer.write()
    assert result.frag_id == 0
    assert result.rows_written == num_physical_rows


def test_fragment_writer_falls_back_to_open_dataset(
    tmp_path: Path, monkeypatch
) -> None:
    """Ensure writer opens dataset when metadata is missing."""
    from collections import deque

    tbl_path = tmp_path / "tbl.lance"
    ds = lance.write_dataset(pa.table({"a": [1, 2, 3]}), tbl_path)

    store = CheckpointStore.from_uri(str(tmp_path / "ckp"))
    checkpoint_key = "fragment_0_batch_0"
    rowaddr = pa.array([0, 1, 2], type=pa.uint64())
    batch = pa.record_batch([pa.array([10, 11, 12]), rowaddr], names=["a", "_rowaddr"])
    store[checkpoint_key] = batch

    class _Queue:
        def __init__(self, items) -> None:
            self._items = deque(items)

        def get(self) -> tuple[int, str]:
            return self._items.popleft()

    queue = _Queue([(0, checkpoint_key), (-1, "")])

    call_count = {"n": 0}
    original_dataset = lance.dataset

    def _wrapped(*args, **kwargs) -> Any:
        call_count["n"] += 1
        return original_dataset(*args, **kwargs)

    monkeypatch.setattr(lance, "dataset", _wrapped)

    writer_cls = FragmentWriter.__ray_actor_class__
    writer = writer_cls(
        ds.uri,
        ["a"],
        store.uri(),
        0,
        queue,
        read_version=ds.version,
        filler_schema=None,
        field_ids=None,
        num_physical_rows=None,
        num_logical_rows=None,
    )

    result = writer.write()
    assert result.frag_id == 0
    assert result.rows_written == 3
    assert call_count["n"] == 1


@pytest.mark.parametrize(
    ("version_str", "expected"),
    [("2.0", (2, 0)), ("2.1", (2, 1)), ("3.0", (3, 0))],
)
def test_parse_data_storage_version(
    version_str: str, expected: tuple[int, int]
) -> None:
    from geneva.utils import parse_data_storage_version

    assert parse_data_storage_version(version_str) == expected


@pytest.mark.parametrize("data_storage_version", ["2.0", "2.1"])
def test_collect_skipped_fragments_uses_dataset_version(
    tmp_path: Path, data_storage_version: str
) -> None:
    """_collect_skipped_fragments must create DataFile with the dataset's version."""
    from geneva.runners.ray.pipeline import _collect_skipped_fragments

    # Create a 2-fragment dataset with explicit data_storage_version
    ds_path = str(tmp_path / "ds.lance")
    lance.write_dataset(
        pa.table({"a": [1, 2, 3]}),
        ds_path,
        data_storage_version=data_storage_version,
    )
    lance.write_dataset(
        pa.table({"a": [4, 5, 6]}),
        ds_path,
        mode="append",
    )

    db = connect(tmp_path)
    tbl = db.open_table("ds")

    @udf(data_type=pa.int64())
    def double_a(a: int) -> int:
        return a * 2

    tbl.add_columns({"a2": double_a})

    ds = lance.dataset(ds_path)
    store = CheckpointStore.from_uri(str(URL(str(tmp_path)) / "ckp"))
    map_task = BackfillUDFTask(
        udfs={"a2": double_a},
        min_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
        max_checkpoint_size=DEFAULT_CHECKPOINT_ROWS,
    )

    # Checkpoint fragment 0 — use the same dedupe key that _collect_skipped_fragments
    # computes internally (no src_files_hash, no dataset_version)
    dedupe_key = _get_fragment_dedupe_key(ds.uri, 0, map_task)
    fake_file = "fragment_0.lance"
    store[dedupe_key] = pa.RecordBatch.from_pydict({"file": [fake_file]})

    skipped, stats = _collect_skipped_fragments(
        dst_dataset=ds,
        dst_frags_with_checkpoint={0},
        map_task=map_task,
        checkpoint_store=store,
    )

    assert 0 in skipped
    df, row_count = skipped[0]
    major, minor = data_storage_version.split(".")
    assert df.file_major_version == int(major)
    assert df.file_minor_version == int(minor)
    assert row_count > 0


# ---------------------------------------------------------------------------
# Helpers for extension-type / multiprocess tests
# ---------------------------------------------------------------------------


def _make_blob_batch(n_rows: int = 3) -> pa.RecordBatch:
    """Create a RecordBatch with a ``_rowaddr`` column and a BlobType column."""
    blob_arr = BlobArray.from_pylist([f"data_{i}".encode() for i in range(n_rows)])
    return pa.record_batch(
        {
            "_rowaddr": pa.array(list(range(n_rows)), type=pa.uint64()),
            "blob_col": blob_arr,
        }
    )


def _make_standard_batch(n_rows: int = 3) -> pa.RecordBatch:
    """Create a simple RecordBatch with no extension types."""
    return pa.record_batch(
        {
            "_rowaddr": pa.array(list(range(n_rows)), type=pa.uint64()),
            "value": pa.array(list(range(n_rows)), type=pa.int64()),
        }
    )


# ---------------------------------------------------------------------------
# Test Group 1: _strip_extension_types / _restore_extension_types
# ---------------------------------------------------------------------------


def test_strip_extension_types_no_extension_columns() -> None:
    """Standard batch passes through unchanged (same object)."""
    batch = _make_standard_batch()
    result = _strip_extension_types(batch)
    assert result is batch  # identity — no copy


def test_strip_extension_types_with_blob_column() -> None:
    """BlobType column is replaced with its struct storage type."""
    batch = _make_blob_batch()
    stripped = _strip_extension_types(batch)

    # The blob column should now be the storage struct type
    blob_field = stripped.schema.field("blob_col")
    assert not isinstance(blob_field.type, pa.ExtensionType)
    assert pa.types.is_struct(blob_field.type)

    # Metadata should contain the extension type name.  Lance only defines
    # BlobType with extension_name "lance.blob.v2" — there is no v1 variant.
    # _get_extension_type uses startswith("lance.blob") so future versions
    # (e.g. v3) would be handled without code changes.
    assert blob_field.metadata is not None
    assert blob_field.metadata[b"__ext_type_name"] == b"lance.blob.v2"
    assert b"__ext_type_serialized" in blob_field.metadata

    # Non-extension column is unchanged
    assert stripped.schema.field("_rowaddr").type == pa.uint64()


def test_restore_extension_types_roundtrip() -> None:
    """Strip then restore preserves type and data values."""
    batch = _make_blob_batch()
    stripped = _strip_extension_types(batch)
    restored = _restore_extension_types(stripped)

    # Type must be restored
    assert isinstance(restored.schema.field("blob_col").type, BlobType)

    # Data must be preserved
    assert (
        restored.column("blob_col").to_pylist() == batch.column("blob_col").to_pylist()
    )
    assert restored.column("_rowaddr") == batch.column("_rowaddr")

    # Metadata markers must be cleaned up
    restored_meta = restored.schema.field("blob_col").metadata
    assert restored_meta is None or b"__ext_type_name" not in restored_meta


def test_strip_restore_with_nulls() -> None:
    """Null values in a BlobType column survive the roundtrip."""
    blob_arr = BlobArray.from_pylist([b"hello", None, b"world"])
    batch = pa.record_batch(
        {
            "_rowaddr": pa.array([0, 1, 2], type=pa.uint64()),
            "blob_col": blob_arr,
        }
    )
    restored = _restore_extension_types(_strip_extension_types(batch))

    assert isinstance(restored.schema.field("blob_col").type, BlobType)
    py_vals = restored.column("blob_col").to_pylist()
    assert py_vals[1] is None
    assert py_vals[0] is not None
    assert py_vals[2] is not None


def test_strip_restore_with_multiple_extension_columns() -> None:
    """Two BlobType columns are both correctly stripped and restored."""
    blob_a = BlobArray.from_pylist([b"a1", b"a2"])
    blob_b = BlobArray.from_pylist([b"b1", b"b2"])
    batch = pa.record_batch(
        {
            "_rowaddr": pa.array([0, 1], type=pa.uint64()),
            "blob_a": blob_a,
            "blob_b": blob_b,
        }
    )

    stripped = _strip_extension_types(batch)
    assert not isinstance(stripped.schema.field("blob_a").type, pa.ExtensionType)
    assert not isinstance(stripped.schema.field("blob_b").type, pa.ExtensionType)

    restored = _restore_extension_types(stripped)
    assert isinstance(restored.schema.field("blob_a").type, BlobType)
    assert isinstance(restored.schema.field("blob_b").type, BlobType)
    assert restored.column("blob_a").to_pylist() == batch.column("blob_a").to_pylist()
    assert restored.column("blob_b").to_pylist() == batch.column("blob_b").to_pylist()


def test_strip_restore_with_mixed_columns() -> None:
    """Only the extension column is affected; int and string columns are untouched."""
    blob_arr = BlobArray.from_pylist([b"x"])
    batch = pa.record_batch(
        {
            "_rowaddr": pa.array([0], type=pa.uint64()),
            "name": pa.array(["alice"], type=pa.utf8()),
            "blob_col": blob_arr,
            "score": pa.array([42], type=pa.int32()),
        }
    )

    restored = _restore_extension_types(_strip_extension_types(batch))
    assert restored.schema.field("name").type == pa.utf8()
    assert restored.schema.field("score").type == pa.int32()
    assert isinstance(restored.schema.field("blob_col").type, BlobType)
    assert restored.column("name").to_pylist() == ["alice"]
    assert restored.column("score").to_pylist() == [42]


def test_strip_restore_preserves_existing_metadata() -> None:
    """Pre-existing field metadata survives the roundtrip."""
    blob_arr = BlobArray.from_pylist([b"hi"])
    field = pa.field(
        "blob_col",
        BlobType(),
        metadata={b"lance-encoding:blob": b"true", b"custom_key": b"custom_val"},
    )
    batch = pa.RecordBatch.from_arrays(
        [pa.array([0], type=pa.uint64()), blob_arr],
        schema=pa.schema([pa.field("_rowaddr", pa.uint64()), field]),
    )

    restored = _restore_extension_types(_strip_extension_types(batch))
    meta = restored.schema.field("blob_col").metadata
    assert meta[b"lance-encoding:blob"] == b"true"
    assert meta[b"custom_key"] == b"custom_val"
    # Internal markers should be cleaned
    assert b"__ext_type_name" not in meta


# ---------------------------------------------------------------------------
# Test Group 2: _batch_to_buf / _buf_to_batch IPC round-trip
# ---------------------------------------------------------------------------


def test_ipc_roundtrip_standard_batch() -> None:
    """Regression: standard batch survives IPC round-trip unchanged."""
    batch = _make_standard_batch(5)
    buf = _batch_to_buf(batch)
    batches = _buf_to_batch(buf, coalesce=False)
    assert len(batches) == 1
    result = batches[0]
    assert result.schema == batch.schema
    assert result.to_pydict() == batch.to_pydict()


def test_ipc_roundtrip_blob_batch() -> None:
    """Blob batch survives IPC round-trip with extension type preserved."""
    batch = _make_blob_batch(4)
    buf = _batch_to_buf(batch)
    batches = _buf_to_batch(buf, coalesce=False)
    assert len(batches) == 1
    result = batches[0]

    assert isinstance(result.schema.field("blob_col").type, BlobType)
    assert result.column("blob_col").to_pylist() == batch.column("blob_col").to_pylist()
    assert result.column("_rowaddr").to_pylist() == batch.column("_rowaddr").to_pylist()


def test_ipc_roundtrip_blob_batch_coalesce() -> None:
    """Blob batch survives IPC round-trip with coalesce=True."""
    batch = _make_blob_batch(3)
    buf = _batch_to_buf(batch)
    result = _buf_to_batch(buf, coalesce=True)
    assert isinstance(result, pa.RecordBatch)
    assert isinstance(result.schema.field("blob_col").type, BlobType)
    assert result.column("blob_col").to_pylist() == batch.column("blob_col").to_pylist()


def test_ipc_roundtrip_empty_blob_batch() -> None:
    """Zero-row batch with blob schema survives IPC round-trip."""
    blob_arr = BlobArray.from_pylist([])
    batch = pa.record_batch(
        {
            "_rowaddr": pa.array([], type=pa.uint64()),
            "blob_col": blob_arr,
        }
    )
    buf = _batch_to_buf(batch)

    batches = _buf_to_batch(buf, coalesce=False)
    if batches:
        assert batches[0].num_rows == 0

    result = _buf_to_batch(buf, coalesce=True)
    assert result.num_rows == 0


def test_ipc_roundtrip_mixed_columns() -> None:
    """Batch with both blob and standard columns survives IPC."""
    blob_arr = BlobArray.from_pylist([b"one", b"two"])
    batch = pa.record_batch(
        {
            "_rowaddr": pa.array([0, 1], type=pa.uint64()),
            "label": pa.array(["a", "b"], type=pa.utf8()),
            "blob_col": blob_arr,
            "count": pa.array([10, 20], type=pa.int64()),
        }
    )
    buf = _batch_to_buf(batch)
    result = _buf_to_batch(buf, coalesce=True)

    assert isinstance(result.schema.field("blob_col").type, BlobType)
    assert result.schema.field("label").type == pa.utf8()
    assert result.schema.field("count").type == pa.int64()
    assert result.to_pydict()["label"] == ["a", "b"]
    assert result.to_pydict()["count"] == [10, 20]


# ---------------------------------------------------------------------------
# Test Group 3: MultiProcessBatchApplier integration
# ---------------------------------------------------------------------------


def test_multiprocess_applier_standard_batch() -> None:
    """Regression: standard batches through MultiProcessBatchApplier."""
    from geneva.debug.logger import NoOpErrorLogger

    @udf(data_type=pa.int64(), input_columns=["value"])
    def double_value(value: int) -> int:
        return value * 2

    n_rows = 6
    batch = pa.record_batch(
        {
            "_rowaddr": pa.array(list(range(n_rows)), type=pa.uint64()),
            "value": pa.array(list(range(n_rows)), type=pa.int64()),
        }
    )

    read_task = _DummyReadTask(frag_id=0, rows=n_rows, batches=[batch])
    map_task = BackfillUDFTask(udfs={"result": double_value})

    applier = MultiProcessBatchApplier(num_processes=2, job_id="test")
    results = list(applier.run(read_task, map_task, NoOpErrorLogger()))

    assert len(results) >= 1
    combined = pa.concat_tables([pa.table(b) for b in results]).to_pydict()
    assert combined["result"] == [0, 2, 4, 6, 8, 10]


def test_multiprocess_applier_blob_batch() -> None:
    """Blob-containing batch through MultiProcessBatchApplier.

    The UDF receives a RecordBatch where the blob column has its BlobType
    extension type intact.  It extracts the inline data length.
    """
    from geneva.debug.logger import NoOpErrorLogger

    @udf(data_type=pa.int64())
    def blob_data_len(batch: pa.RecordBatch) -> pa.Array:
        # The blob column should be a BlobType extension array.
        # Access the underlying struct's "data" field to get the binary data.
        blob_col = batch.column("blob_col")
        storage = blob_col.storage if hasattr(blob_col, "storage") else blob_col
        data_arr = storage.field("data")
        return pa.array(
            [len(v.as_py()) if v.is_valid else 0 for v in data_arr],
            type=pa.int64(),
        )

    batch = _make_blob_batch(4)
    read_task = _DummyReadTask(frag_id=0, rows=4, batches=[batch])
    map_task = BackfillUDFTask(udfs={"data_len": blob_data_len})

    applier = MultiProcessBatchApplier(num_processes=2, job_id="test")
    results = list(applier.run(read_task, map_task, NoOpErrorLogger()))

    assert len(results) >= 1
    combined = pa.concat_tables([pa.table(b) for b in results]).to_pydict()
    expected_lens = [len(f"data_{i}".encode()) for i in range(4)]
    assert combined["data_len"] == expected_lens


def test_multiprocess_applier_multiple_blob_batches() -> None:
    """Multiple blob batches exercise the backpressure path."""
    from geneva.debug.logger import NoOpErrorLogger

    @udf(data_type=pa.int64())
    def blob_size(batch: pa.RecordBatch) -> pa.Array:
        blob_col = batch.column("blob_col")
        storage = blob_col.storage if hasattr(blob_col, "storage") else blob_col
        data_arr = storage.field("data")
        return pa.array(
            [len(v.as_py()) if v.is_valid else 0 for v in data_arr],
            type=pa.int64(),
        )

    batches = [_make_blob_batch(3) for _ in range(5)]
    total_rows = sum(b.num_rows for b in batches)
    read_task = _DummyReadTask(frag_id=0, rows=total_rows, batches=batches)
    map_task = BackfillUDFTask(udfs={"size": blob_size})

    applier = MultiProcessBatchApplier(num_processes=2, job_id="test")
    results = list(applier.run(read_task, map_task, NoOpErrorLogger()))

    combined = pa.concat_tables([pa.table(b) for b in results]).to_pydict()
    assert len(combined["size"]) == total_rows
    # Each batch has data_0, data_1, data_2 → lengths 6, 6, 6
    assert all(v == 6 for v in combined["size"])


def test_multiprocess_applier_blob_batch_with_error_handling() -> None:
    """Blob batch through the _worker_apply error-handling code path.

    When a UDF has an ``error_handling`` config, MultiProcessBatchApplier
    routes through ``_worker_apply`` instead of ``_apply_with_stream_buf``.
    Both paths deserialise via ``_buf_to_batch`` and must restore extension
    types.
    """
    from geneva.debug.error_store import ErrorHandlingConfig, UDFRetryConfig
    from geneva.debug.logger import NoOpErrorLogger

    @udf(
        data_type=pa.int64(),
        error_handling=ErrorHandlingConfig(
            retry_config=UDFRetryConfig.no_retry(),
            log_errors=False,
        ),
    )
    def blob_len_eh(batch: pa.RecordBatch) -> pa.Array:
        blob_col = batch.column("blob_col")
        storage = blob_col.storage if hasattr(blob_col, "storage") else blob_col
        data_arr = storage.field("data")
        return pa.array(
            [len(v.as_py()) if v.is_valid else 0 for v in data_arr],
            type=pa.int64(),
        )

    batch = _make_blob_batch(4)
    read_task = _DummyReadTask(frag_id=0, rows=4, batches=[batch])
    map_task = BackfillUDFTask(udfs={"data_len": blob_len_eh})

    applier = MultiProcessBatchApplier(num_processes=2, job_id="test")
    results = list(applier.run(read_task, map_task, NoOpErrorLogger()))

    assert len(results) >= 1
    combined = pa.concat_tables([pa.table(b) for b in results]).to_pydict()
    expected_lens = [len(f"data_{i}".encode()) for i in range(4)]
    assert combined["data_len"] == expected_lens


# ---------------------------------------------------------------------------
# Test Group 4: Edge cases and error paths
# ---------------------------------------------------------------------------


def test_restore_extension_types_no_markers() -> None:
    """Batch without extension markers passes through unchanged (same object)."""
    batch = _make_standard_batch()
    result = _restore_extension_types(batch)
    assert result is batch  # identity — no copy


def test_get_extension_type_unknown_raises() -> None:
    """_get_extension_type raises ValueError for an unrecognised type name."""
    from geneva.apply.multiprocess import _get_extension_type

    with pytest.raises(ValueError, match="Unknown Arrow extension type"):
        _get_extension_type("acme.custom_type", pa.int32(), b"")


def test_ipc_roundtrip_blob_batch_multiple_cycles() -> None:
    """Blob batch survives two consecutive serialize-deserialize cycles."""
    batch = _make_blob_batch(3)
    buf1 = _batch_to_buf(batch)
    mid = _buf_to_batch(buf1, coalesce=True)
    buf2 = _batch_to_buf(mid)
    result = _buf_to_batch(buf2, coalesce=True)

    assert isinstance(result.schema.field("blob_col").type, BlobType)
    assert result.column("blob_col").to_pylist() == batch.column("blob_col").to_pylist()
