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

import contextlib
import functools
import json
import logging
import os
import random
import time
import uuid
from collections import Counter
from collections.abc import Generator, Iterable, Iterator
from typing import TYPE_CHECKING, Any, NamedTuple, Optional, cast

import attrs
import cloudpickle
import lance
import pyarrow as pa
import ray.actor
import ray.exceptions
import ray.util.queue
from ray.actor import ActorHandle
from tenacity import (
    Retrying,
    retry_if_exception,
    stop_after_attempt,
    wait_exponential_jitter,
)
from tqdm.std import tqdm as TqdmType  # noqa: N812

from geneva.apply import (
    DEFAULT_BATCH_CHECKPOINT_FLUSH_INTERVAL_SECONDS,
    CheckpointingApplier,
    MapBatchCheckpoint,
    plan_copy,
    plan_read,
)
from geneva.apply.applier import BatchApplier
from geneva.apply.multiprocess import MultiProcessBatchApplier
from geneva.apply.simple import SimpleApplier
from geneva.apply.table_cache import (
    TableCache,
    bind_tables_for_task,
    clear_bound_tables,
)
from geneva.apply.task import BackfillUDFTask, CopyTableTask, MapTask, ReadTask
from geneva.checkpoint import CheckpointStore
from geneva.checkpoint_utils import hash_source_files
from geneva.db import Connection
from geneva.debug.error_store import GENEVA_ERRORS_TABLE_NAME, ErrorStore
from geneva.debug.logger import TableErrorLogger
from geneva.jobs.config import JobConfig
from geneva.namespace import get_storage_options_provider
from geneva.packager import UDFPackager, UDFSpec
from geneva.query import (
    MATVIEW_META_BASE_DBURI,
    MATVIEW_META_BASE_TABLE,
    MATVIEW_META_NAMESPACE_PATH,
    MATVIEW_META_QUERY,
    MATVIEW_META_SCALAR_UDTF,
    MATVIEW_META_SOURCE_QUERY,
    MATVIEW_META_VERSION,
    GenevaQuery,
    GenevaQueryBuilder,
)

if TYPE_CHECKING:
    from lance_namespace import LanceNamespace

# Namespace imports for storage_options_provider

from geneva.runners.ray.actor_pool import ActorPool
from geneva.runners.ray.admission import (
    DRIVER_NUM_CPUS,
    FRAGMENT_WRITER_MEMORY,
    FRAGMENT_WRITER_NUM_CPUS,
    JOBTRACKER_MEMORY,
    JOBTRACKER_NUM_CPUS,
)
from geneva.runners.ray.jobtracker import (
    METRIC_AVG_BATCH_NUM_ROWS,
    METRIC_AVG_BATCH_SIZE,
    METRIC_BATCH_CHECKPOINTING_TIME,
    METRIC_CHECKPOINT_EXISTS_TIME,
    METRIC_CHECKPOINT_LIST_TIME,
    METRIC_CHECKPOINT_LOAD_TIME,
    METRIC_COMMIT_BACKOFF_TIME,
    METRIC_COMMIT_CONCURRENT_WRITER_RETRIES,
    METRIC_COMMIT_CONFLICT_RETRIES,
    METRIC_COMMIT_TIME,
    METRIC_FRAGMENT_CHECKPOINTING_TIME,
    METRIC_PLAN_READ_TIME,
    METRIC_READ_IO_TIME,
    METRIC_READ_TASK_TOTAL_TIME,
    METRIC_UDF_PROCESSING_TIME,
    METRIC_WRITER_ALIGN_TIME,
    METRIC_WRITER_BUFFER_SORT_TIME,
    METRIC_WRITER_CHECKPOINT_READ_TIME,
    METRIC_WRITER_QUEUE_WAIT_TIME,
    METRIC_WRITER_WRITE_TIME,
    job_tracker_options,
)
from geneva.runners.ray.kuberay import PodStatus, _ray_status, k8s_status
from geneva.runners.ray.raycluster import (
    ray_tqdm,
)
from geneva.runners.ray.writer import FragmentWriter, FragmentWriteResult
from geneva.table import JobFuture, Table, TableReference
from geneva.tqdm import (
    Colors,
    fmt,
    fmt_numeric,
    fmt_pending,
    tqdm,
)
from geneva.transformer import UDF
from geneva.utils import parse_data_storage_version
from geneva.utils.batch_size import resolve_batch_size, resolve_task_size
from geneva.utils.parse_rust_debug import extract_field_ids
from geneva.utils.ray import CPU_ONLY_NODE

_LOG = logging.getLogger(__name__)


# Max retries for Lance "Too many concurrent writers" errors during commit.
# With exponential backoff (1s, 2s, 4s, 8s, 16s, then 16s capped), 12 retries
# gives ~2.5 minutes total wait time before giving up.
GENEVA_COMMIT_MAX_RETRIES = int(os.environ.get("GENEVA_COMMIT_MAX_RETRIES", "12"))


REFRESH_EVERY_SECONDS = 5.0

# How often to poll for results and log k8s pod status when the pipeline stalls.
POLL_INTERVAL_S = 30.0

# How long to wait with zero task completions before raising a TimeoutError.
PIPELINE_STALL_TIMEOUT_S = float(
    os.environ.get("GENEVA_PIPELINE_STALL_TIMEOUT_S", "600")
)

# Maximum number of times to restart a writer actor before giving up
MAX_WRITER_RESTARTS = 5

# Number of idle rounds (5s each) before considering a writer stalled during drain.
# With many concurrent backfills, writers may be slow due to resource contention,
# not actually stalled. 6 rounds = 30 seconds gives more time before restarting.
GENEVA_WRITER_STALL_IDLE_ROUNDS = int(
    os.environ.get("GENEVA_WRITER_STALL_IDLE_ROUNDS", "6")
)

# Maximum retries for version conflicts during commit. Version conflicts occur when
# concurrent backfills commit to the same fragments. We attempt column merging on
# conflict, but limit retries to prevent infinite loops.
GENEVA_VERSION_CONFLICT_MAX_RETRIES = int(
    os.environ.get("GENEVA_VERSION_CONFLICT_MAX_RETRIES", "10")
)

# Lance uses -2 as a "tombstone" field_id marker. When a field is marked as -2 in a
# DataFile's field_ids list, Lance will not read that field from that file. This is
# used in Merge operations to mask columns that are being replaced by a new file.
# See: lance/rust/lance/src/dataset/fragment.rs:1685
LANCE_FIELD_ID_TOMBSTONE = -2

CNT_WORKERS_PENDING = "cnt_geneva_workers_pending"
CNT_WORKERS_ACTIVE = "cnt_geneva_workers_active"
CNT_RAY_NODES = "cnt_ray_nodes"
CNT_K8S_NODES = "k8s_nodes_provisioned"
CNT_K8S_PHASE = "k8s_cluster_phase"
METRIC_UDF_VALUES = "udf_values_computed"

# Maximum retries for transient AWS credential provider timeouts in appliers.
# A value of 3 means up to 3 retries after the initial failure.
GENEVA_APPLIER_TRANSIENT_RETRIES = int(
    os.environ.get("GENEVA_APPLIER_TRANSIENT_RETRIES", "3")
)
# Exponential backoff base for transient applier retries (seconds).
GENEVA_APPLIER_RETRY_BASE_BACKOFF_SECONDS = float(
    os.environ.get("GENEVA_APPLIER_RETRY_BASE_BACKOFF_SECONDS", "1.0")
)
# Cap backoff growth to avoid very long sleeps for one task.
GENEVA_APPLIER_RETRY_MAX_BACKOFF_SECONDS = float(
    os.environ.get("GENEVA_APPLIER_RETRY_MAX_BACKOFF_SECONDS", "8.0")
)

_RETRYABLE_AWS_CREDENTIAL_ERROR_MARKERS = (
    "failed to get aws credentials",
    "httptimeouterror",
    "connectorerror",
    "providererror",
)


def _iter_exception_chain(exc: BaseException) -> Iterator[BaseException]:
    """Yield exception and its nested causes/contexts without looping."""
    seen: set[int] = set()
    current: BaseException | None = exc
    while current is not None and id(current) not in seen:
        seen.add(id(current))
        yield current
        current = cast("BaseException | None", current.__cause__ or current.__context__)


def _is_retryable_aws_credential_error(exc: BaseException) -> bool:
    """Return True for transient AWS credential provider timeout failures."""
    for candidate in _iter_exception_chain(exc):
        msg = str(candidate).lower()
        if not msg:
            continue
        if "failed to get aws credentials" in msg:
            return True
        if "aws credentials" in msg and (
            "httptimeouterror" in msg
            or "connectorerror" in msg
            or "providererror" in msg
            or "timeout" in msg
        ):
            return True
        if all(marker in msg for marker in _RETRYABLE_AWS_CREDENTIAL_ERROR_MARKERS):
            return True
    return False


@ray.remote  # type: ignore[misc]
@attrs.define
class ApplierActor:  # pyright: ignore[reportRedeclaration]
    applier: CheckpointingApplier
    table_cache: TableCache = attrs.field(factory=TableCache, init=False, repr=False)

    def __ray_ready__(self) -> None:
        pass

    def __repr__(self) -> str:
        """Custom repr that safely handles missing attributes during unpickling.

        This is necessary because attrs-generated __repr__ can fail when called
        during exception handling in Ray if the object hasn't been fully unpickled yet.
        """
        try:
            # Try to get all attrs fields safely
            field_strs = []
            for field in attrs.fields(self.__class__):
                # Check if attribute exists first before accessing it
                if hasattr(self, field.name):
                    value = getattr(self, field.name)
                    field_strs.append(f"{field.name}={value!r}")
                else:
                    field_strs.append(f"{field.name}=<not set>")

            return f"{self.__class__.__qualname__}({', '.join(field_strs)})"
        except Exception:
            # Fallback if even that fails
            return f"<{self.__class__.__name__} (repr failed)>"

    def run(
        self, task
    ) -> tuple[
        ReadTask,
        list[MapBatchCheckpoint],
        int,
        int,
        int,
        int,
        int,
        int,
        int,
        int,
    ]:
        max_retries = max(0, int(GENEVA_APPLIER_TRANSIENT_RETRIES))
        max_attempts = max_retries + 1
        start = time.perf_counter()

        def _before_sleep(retry_state) -> None:  # noqa: ANN001
            exc = (
                retry_state.outcome.exception()
                if retry_state.outcome is not None
                else None
            )
            backoff = (
                retry_state.next_action.sleep
                if retry_state.next_action is not None
                else 0.0
            )
            _LOG.warning(
                (
                    "Transient AWS credential error for task %s "
                    "(retry %d/%d). Sleeping %.2fs before retry. Error: %s"
                ),
                task,
                retry_state.attempt_number,
                max_retries,
                backoff,
                exc,
            )

        retrier = Retrying(
            retry=retry_if_exception(_is_retryable_aws_credential_error),
            stop=stop_after_attempt(max_attempts),
            wait=wait_exponential_jitter(
                initial=GENEVA_APPLIER_RETRY_BASE_BACKOFF_SECONDS,
                max=GENEVA_APPLIER_RETRY_MAX_BACKOFF_SECONDS,
            ),
            reraise=True,
            before_sleep=_before_sleep,
        )

        bind_tables_for_task(task, self.table_cache)
        try:
            for attempt in retrier:
                with attempt:
                    checkpoints, cnt_udf_computed = self.applier.run(task)
                    # Wall time for the full read task execution (overlaps with
                    # read IO + UDF + checkpointing); keep separate from
                    # sub-metrics.
                    read_task_total_time_ms = int((time.perf_counter() - start) * 1000)
                    udf_processing_time_ms = int(
                        getattr(self.applier.batch_applier, "udf_processing_time_ms", 0)
                        or 0
                    )
                    read_io_time_ms = int(
                        getattr(self.applier.batch_applier, "read_io_time_ms", 0) or 0
                    )
                    batch_checkpointing_time_ms = int(
                        getattr(self.applier, "batch_checkpointing_time_ms", 0) or 0
                    )
                    checkpoint_load_time_ms = int(
                        getattr(self.applier, "checkpoint_load_time_ms", 0) or 0
                    )
                    checkpoint_exists_time_ms = int(
                        getattr(self.applier, "checkpoint_exists_time_ms", 0) or 0
                    )
                    checkpoint_list_time_ms = int(
                        getattr(self.applier, "checkpoint_list_time_ms", 0) or 0
                    )
                    return (
                        task,
                        checkpoints,
                        cnt_udf_computed,
                        udf_processing_time_ms,
                        batch_checkpointing_time_ms,
                        read_io_time_ms,
                        checkpoint_load_time_ms,
                        checkpoint_exists_time_ms,
                        checkpoint_list_time_ms,
                        read_task_total_time_ms,
                    )
        except Exception as exc:
            raise _picklable_remote_error(exc) from None
        finally:
            clear_bound_tables(task)

        # Unreachable with reraise=True, but keeps static analyzers satisfied.
        raise RuntimeError("Applier task failed after retry exhaustion")


ApplierActor: ray.actor.ActorClass = cast("ray.actor.ActorClass", ApplierActor)  # type: ignore[no-redef]


def _get_fragment_dedupe_key(
    uri: str,
    frag_id: int,
    map_task: MapTask,
    dataset_version: int | str | None = None,
    src_files_hash: str | None = None,
) -> str:
    prefix = map_task.checkpoint_prefix(
        dataset_uri=uri,
        where=getattr(map_task, "where", None),
        column=None,
        src_files_hash=src_files_hash,
    )
    return f"{prefix}_frag-{frag_id}"


def get_source_data_files(
    src_frag,
    relevant_field_ids: frozenset[int] | None = None,
) -> frozenset[str]:
    """Get set of data file paths in this fragment for relevant fields.

    This is used to detect when source data has been modified via backfill.
    When a backfill runs (even on the same field), new data files are created
    with new UUIDs. Tracking file paths detects ALL data changes, including:
    - New fields being backfilled
    - Existing fields being re-backfilled with different UDF

    Args:
        src_frag: A Lance fragment object
        relevant_field_ids: If provided, only include files containing these
            field IDs. If None, include all files (backward compatible).

    Returns:
        frozenset of data file paths in this fragment
    """
    if relevant_field_ids is None:
        return frozenset(df.path for df in src_frag.data_files())

    # Filter to files containing at least one relevant field
    return frozenset(
        df.path for df in src_frag.data_files() if set(df.fields) & relevant_field_ids
    )


def get_combined_source_data_files(
    src_dataset: lance.LanceDataset,
    src_frag_ids: set[int],
    relevant_field_ids: frozenset[int] | None = None,
) -> frozenset[str]:
    """Get union of data file paths across multiple source fragments.

    Why union? We want to detect if ANY data file changes in ANY source fragment
    that contributes to a destination fragment. When any backfill runs on any
    source fragment, new data files are created, and the union changes.

    This is more robust than field coverage tracking because it detects:
    - New fields being backfilled (new files for new field IDs)
    - Existing fields being re-backfilled (new files replace old files)

    Note: Compaction is handled separately via stable_row_id mapping.

    Args:
        src_dataset: The source Lance dataset
        src_frag_ids: Set of source fragment IDs that contribute to a destination
        relevant_field_ids: If provided, only include files containing these
            field IDs. If None, include all files.

    Returns:
        frozenset of all data file paths from all source fragments (union)
    """
    all_files: set[str] = set()
    for frag_id in src_frag_ids:
        frag = src_dataset.get_fragment(frag_id)
        if frag is not None:
            all_files.update(get_source_data_files(frag, relevant_field_ids))
    return frozenset(all_files)


def _get_relevant_field_ids(
    src_dataset: lance.LanceDataset,
    input_cols: Iterable[str],
) -> frozenset[int]:
    """Get field IDs for columns the MV reads from the source.

    This is used to filter data file tracking to only files containing
    columns that the MV actually uses. If a new column is added to the
    source that isn't in the MV, changes to that column won't trigger
    MV refresh.

    Args:
        src_dataset: The source Lance dataset
        input_cols: Column names the MV reads (list or set)

    Returns:
        frozenset of field IDs for the relevant columns
    """
    from geneva.utils.parse_rust_debug import extract_field_ids

    field_ids: set[int] = set()
    for col_name in input_cols:
        try:
            ids = extract_field_ids(src_dataset.lance_schema, col_name)
            field_ids.update(ids)
        except Exception:  # noqa: PERF203
            # Column may not exist in source (e.g., computed column)
            _LOG.debug(f"Column {col_name} not found in source schema, skipping")
    return frozenset(field_ids)


def _get_fragments_missing_column_data(
    dataset: lance.LanceDataset, col_name: str
) -> set[int]:
    """Get fragment IDs that don't have data files for a column.

    When a column has no data files in a fragment, Lance imputes NULL values.
    This is useful for struct columns where IS NULL filter doesn't work correctly.

    Args:
        dataset: The Lance dataset to check
        col_name: Column name to check for data files

    Returns:
        Set of fragment IDs that are missing data files for the column
        (these fragments need processing)
    """
    field_ids = _get_relevant_field_ids(dataset, [col_name])
    if not field_ids:
        # Column has no field IDs - return all fragments
        return {frag.fragment_id for frag in dataset.get_fragments()}

    missing = set()
    for frag in dataset.get_fragments():
        data_files = get_source_data_files(frag, field_ids)
        if not data_files:
            missing.add(frag.fragment_id)
    return missing


def _run_column_adding_pipeline(
    map_task: MapTask,
    checkpoint_store: CheckpointStore,
    error_store: ErrorStore,
    config: JobConfig,
    dst: TableReference,
    input_plan: Iterator[ReadTask],
    job_id: str | None,
    applier_concurrency: int = 8,
    *,
    intra_applier_concurrency: int = 1,
    batch_checkpoint_flush_interval_seconds: float = (
        DEFAULT_BATCH_CHECKPOINT_FLUSH_INTERVAL_SECONDS
    ),
    use_cpu_only_pool: bool = False,
    job_tracker=None,
    where=None,
    skipped_fragments: dict | None = None,
    skipped_stats: dict | None = None,
    enable_job_tracker_saves: bool = True,
    src_data_files_by_dst: dict[int, frozenset[str]] | None = None,
) -> None:
    """
    Run the column adding pipeline.

    Args:
    * use_cpu_only_pool: If True will force schedule cpu-only actors on cpu-only nodes.

    """
    job_id = job_id or uuid.uuid4().hex

    job_tracker = job_tracker or job_tracker_options(
        name=f"jobtracker-{job_id}",
        num_cpus=JOBTRACKER_NUM_CPUS,
        memory=JOBTRACKER_MEMORY,
        max_restarts=-1,
    ).remote(job_id, dst, enable_saves=enable_job_tracker_saves)  # type: ignore[call-arg]

    job = ColumnAddPipelineJob(
        map_task=map_task,
        checkpoint_store=checkpoint_store,
        error_store=error_store,
        config=config,
        dst=dst,
        input_plan=input_plan,
        job_id=job_id,
        applier_concurrency=applier_concurrency,
        intra_applier_concurrency=intra_applier_concurrency,
        batch_checkpoint_flush_interval_seconds=(
            batch_checkpoint_flush_interval_seconds
        ),
        use_cpu_only_pool=use_cpu_only_pool,
        job_tracker=job_tracker,  # type: ignore[arg-type]
        where=where,
        skipped_fragments=skipped_fragments or {},
        skipped_stats=skipped_stats or {},
        src_data_files_by_dst=src_data_files_by_dst or {},
    )
    job.run()


@attrs.define
class ColumnAddPipelineJob:
    """ColumnAddPipeline drives batches of rows to commits in the dataset.

    ReadTasks are defined wrapped for tracking, and then dispatched for udf exeuction
    in the ActorPool.  The results are sent to the FragmentWriterManager which
    manages fragment checkpoints and incremental commits.
    """

    map_task: MapTask
    checkpoint_store: CheckpointStore
    error_store: ErrorStore
    config: JobConfig
    dst: TableReference
    input_plan: Iterator[ReadTask]
    job_id: str
    applier_concurrency: int = 8
    intra_applier_concurrency: int = 1
    batch_checkpoint_flush_interval_seconds: float = (
        DEFAULT_BATCH_CHECKPOINT_FLUSH_INTERVAL_SECONDS
    )
    use_cpu_only_pool: bool = False
    job_tracker: ActorHandle | None = None
    where: str | None = None
    skipped_fragments: dict = attrs.field(factory=dict)
    skipped_stats: dict = attrs.field(factory=dict)
    src_data_files_by_dst: dict[int, frozenset[str]] = attrs.field(factory=dict)
    _total_rows: int = attrs.field(default=0, init=False)
    _last_status_refresh: float = attrs.field(factory=lambda: 0.0, init=False)

    def setup_inputplans(self) -> tuple[Iterator[ReadTask], Counter[int], int]:
        all_tasks = list(self.input_plan)
        self.job_tracker = (
            self.job_tracker
            or job_tracker_options(  # type: ignore[assignment]
                name=f"jobtracker-{self.job_id}",
                num_cpus=JOBTRACKER_NUM_CPUS,
                memory=JOBTRACKER_MEMORY,
                max_restarts=-1,
            ).remote(self.job_id, self.dst)  # type: ignore[call-arg]
        )

        self._total_rows = sum(t.num_rows() for t in all_tasks)

        # Count ReadTasks per fragment. FragmentWriterManager
        # decrements remaining_tasks once per ReadTask in ingest_task(), so this
        # must match. A single ReadTask may produce multiple checkpoints when
        # checkpoint_size < task_size, but seal() should be called once all
        # ReadTasks for a fragment have completed.
        readtasks_by_frag: Counter[int] = Counter()
        for t in all_tasks:
            frag_id = t.dest_frag_id()
            readtasks_by_frag[frag_id] += 1

        total_readtasks = sum(readtasks_by_frag.values())

        # fragments
        if self.job_tracker is not None:
            self.job_tracker.set_total.remote("fragments", total_readtasks)
            self.job_tracker.set_desc.remote(
                "fragments",
                f"[{self.dst.table_name} - {self.map_task.name()}] Tasks scheduled",
            )

        # this reports # of read tasks started, not completed.
        return (
            ray_tqdm(all_tasks, self.job_tracker, metric="fragments"),
            readtasks_by_frag,
            total_readtasks,
        )

    def setup_actor(self) -> ray.actor.ActorHandle:
        actor = ApplierActor

        # actor.options can only be called once, we must pass all override args
        # in one shot
        num_cpus = self.map_task.num_cpus()
        args = {
            "num_cpus": (num_cpus or 1) * self.intra_applier_concurrency,
        }
        num_gpus = self.map_task.num_gpus()
        if num_gpus and num_gpus > 0:
            args["num_gpus"] = num_gpus
        elif self.use_cpu_only_pool:
            _LOG.info("Using CPU only pool for applier, setting %s to 1", CPU_ONLY_NODE)
            args["resources"] = {CPU_ONLY_NODE: 1}  # type: ignore[assignment]
        memory = self.map_task.memory()
        if memory:
            args["memory"] = memory * self.intra_applier_concurrency
        actor = actor.options(**args)  # type: ignore[attr-defined]
        return actor  # type: ignore[return-value]

    def setup_batchapplier(self) -> BatchApplier:
        if self.intra_applier_concurrency > 1:
            return MultiProcessBatchApplier(
                num_processes=self.intra_applier_concurrency, job_id=self.job_id
            )
        else:
            return SimpleApplier(job_id=self.job_id)

    def setup_actorpool(self) -> ActorPool:
        batch_applier = self.setup_batchapplier()

        errors_tbl = self.dst.as_system_table(GENEVA_ERRORS_TABLE_NAME)

        applier = CheckpointingApplier(
            map_task=self.map_task,
            batch_applier=batch_applier,
            checkpoint_uri=self.checkpoint_store.uri(),
            error_logger=TableErrorLogger(table_ref=errors_tbl),
            batch_checkpoint_flush_interval_seconds=(
                self.batch_checkpoint_flush_interval_seconds
            ),
        )

        actor = self.setup_actor()
        if self.job_tracker is not None:
            self.job_tracker.set_total.remote("workers", self.applier_concurrency)
            self.job_tracker.set_desc.remote("workers", "Workers started")

        pool = ActorPool(
            functools.partial(actor.remote, applier=applier),
            self.applier_concurrency,
            job_tracker=self.job_tracker,
            worker_metric="workers",
        )
        return pool

    def setup_writertracker(self) -> tuple[lance.LanceDataset, int]:
        ds = self.dst.open().to_lance()
        fragments = ds.get_fragments()
        len_frags = len(fragments)

        if self.job_tracker is not None:
            self.job_tracker.set_total.remote("writer_fragments", len_frags)
            self.job_tracker.set_desc.remote("writer_fragments", "Fragments written")
        ray_tqdm(fragments, self.job_tracker, metric="writer_fragments")

        return ds, len_frags

    def _refresh_cluster_status(self) -> None:
        # cluster metrics
        try:
            ray_status = _ray_status()

            # TODO batch this.
            if self.job_tracker is not None:
                m_rn = CNT_RAY_NODES
                cnt_workers = ray_status.get(m_rn, 0)
                self.job_tracker.set_desc.remote(m_rn, "ray nodes provisioned")
                self.job_tracker.set.remote(m_rn, cnt_workers)

                # TODO separate metrics for gpu and cpu workers?
                m_caa = CNT_WORKERS_ACTIVE
                cnt_active = ray_status.get(m_caa, 0)
                self.job_tracker.set_desc.remote(m_caa, "active workers")
                self.job_tracker.set_total.remote(m_caa, self.applier_concurrency)
                self.job_tracker.set.remote(m_caa, cnt_active)

                m_cpa = CNT_WORKERS_PENDING
                cnt_pending = ray_status.get(m_cpa, 0)
                self.job_tracker.set_desc.remote(m_cpa, "pending workers")
                self.job_tracker.set_total.remote(m_cpa, self.applier_concurrency)
                self.job_tracker.set.remote(m_cpa, cnt_pending)

        except Exception:
            _LOG.debug("refresh: failed to get ray status", exc_info=True)
            # do nothing

    def _try_refresh_cluster_status(self) -> None:
        now = time.monotonic()
        if now - self._last_status_refresh >= REFRESH_EVERY_SECONDS:
            self._refresh_cluster_status()
            self._last_status_refresh = now

    def _get_k8s_pod_statuses(self) -> list[PodStatus] | None:
        """Return worker pod statuses from the k8s API, or None on failure."""
        try:
            from geneva._context import get_current_context
            from geneva.runners.ray.raycluster import RayCluster

            ctx = get_current_context()
            if not isinstance(ctx, RayCluster):
                return None
            return k8s_status(ctx.clients, ctx.namespace)
        except Exception:
            _LOG.debug("Failed to get k8s pod status", exc_info=True)
            return None

    def _log_k8s_pod_status(self) -> None:
        """Log the current k8s worker pod status."""
        pods = self._get_k8s_pod_statuses()
        if pods is None:
            return
        worker_pods = [p for p in pods if p.get("node_type") == "worker"]
        if not worker_pods:
            _LOG.warning("Pipeline stall: no worker pods found in namespace")
            return
        for pod in worker_pods:
            reasons = dict(pod["waiting_reasons"]) if pod["waiting_reasons"] else {}
            _LOG.info(
                "Worker pod %s: phase=%s ready=%s waiting=%s",
                pod["name"],
                pod["phase"],
                pod["ready"],
                reasons or "none",
            )

    def _k8s_pod_diagnostics(self) -> str:
        """Return a formatted string of pod status for error messages."""
        pods = self._get_k8s_pod_statuses()
        if pods is None:
            return "(k8s pod status unavailable)"
        worker_pods = [p for p in pods if p.get("node_type") == "worker"]
        if not worker_pods:
            return "No worker pods found in namespace"
        lines: list[str] = []
        for pod in worker_pods:
            reasons = dict(pod["waiting_reasons"]) if pod["waiting_reasons"] else {}
            lines.append(
                f"  {pod['name']}: phase={pod['phase']} "
                f"ready={pod['ready']} waiting={reasons or 'none'}"
            )
        return "\n".join(lines)

    def run(self) -> None:
        plans, tasks_by_frag, cnt_tasks = self.setup_inputplans()
        pool = self.setup_actorpool()
        ds, cnt_fragments = self.setup_writertracker()

        prefix = (
            f"[{self.dst.table_name} - {self.map_task.name()} "
            f"({cnt_fragments} fragments)]"
        )

        try:
            self._refresh_cluster_status()
        except Exception:
            _LOG.debug("initial cluster status refresh failed", exc_info=True)
            # do nothing

        # formatting to show fragments
        try:
            cg = (
                int(self.config.commit_granularity)
                if self.config.commit_granularity is not None
                else 0
            )
        except Exception:
            cg = 0
        cg = max(cg, 0)
        cgstr = (
            "(commit at completion)"
            if cg == 0
            else f"(every {cg} fragment{'s' if cg != 1 else ''})"
        )
        # rows metrics (all cumulative)
        if self.job_tracker is not None:
            # Get the full dataset total (including skipped fragments)
            dataset_total_rows = self._total_rows
            # skipped_stats is {'fragments': count, 'rows': count} for
            # progress tracking.
            if self.skipped_stats:
                dataset_total_rows += self.skipped_stats.get("rows", 0)

            skipped_rows = self.skipped_stats.get("rows", 0)

            previously_completed = (
                f" ({skipped_rows} previously completed)" if skipped_rows > 0 else ""
            )
            for m, desc in [
                (
                    "rows_checkpointed",
                    f"{prefix} Rows checkpointed{previously_completed}",
                ),
                (
                    "rows_ready_for_commit",
                    f"{prefix} Rows ready for commit",
                ),
                (
                    "rows_committed",
                    f"{prefix} Rows committed {cgstr}",
                ),
            ]:
                self.job_tracker.set_total.remote(m, dataset_total_rows)
                self.job_tracker.set_desc.remote(m, desc)
                # Initialize with skipped rows as already completed
                if skipped_rows > 0:
                    self.job_tracker.set.remote(m, skipped_rows)
            # Track the number of rows processed by UDF excluding skipped rows.
            self.job_tracker.set_desc.remote(
                METRIC_UDF_VALUES,
                f"{prefix} Rows processed by UDF",
            )
            # Performance metrics are cumulative across all tasks/fragments.
            # Time metrics are in milliseconds and may overlap; they are not
            # additive. For example, read_task_total_time_ms includes read IO,
            # UDF compute, and checkpointing time.
            for m, desc in [
                # UDF/read path (ApplierActor + CheckpointingApplier).
                (
                    METRIC_UDF_PROCESSING_TIME,
                    f"{prefix} Total UDF processing time (ms)",
                ),
                (
                    METRIC_BATCH_CHECKPOINTING_TIME,
                    f"{prefix} Total batch checkpointing time (ms)",
                ),
                (
                    METRIC_READ_IO_TIME,
                    f"{prefix} Total read IO time (ms)",
                ),
                (
                    METRIC_CHECKPOINT_LOAD_TIME,
                    f"{prefix} Total checkpoint load time (ms)",
                ),
                (
                    METRIC_CHECKPOINT_EXISTS_TIME,
                    f"{prefix} Total checkpoint exists check time (ms)",
                ),
                (
                    METRIC_CHECKPOINT_LIST_TIME,
                    f"{prefix} Total checkpoint list time (ms)",
                ),
                (
                    METRIC_READ_TASK_TOTAL_TIME,
                    f"{prefix} Total read task time (ms)",
                ),
                # Planning (per job).
                (
                    METRIC_PLAN_READ_TIME,
                    f"{prefix} Plan read total time (ms)",
                ),
                # Writer path (FragmentWriter + FragmentWriterManager).
                (
                    METRIC_FRAGMENT_CHECKPOINTING_TIME,
                    f"{prefix} Total fragment checkpointing time (ms)",
                ),
                (
                    METRIC_WRITER_ALIGN_TIME,
                    f"{prefix} Total writer align time (ms)",
                ),
                (
                    METRIC_WRITER_WRITE_TIME,
                    f"{prefix} Total writer write time (ms)",
                ),
                (
                    METRIC_WRITER_QUEUE_WAIT_TIME,
                    f"{prefix} Total writer queue wait time (ms)",
                ),
                (
                    METRIC_WRITER_CHECKPOINT_READ_TIME,
                    f"{prefix} Total writer checkpoint read time (ms)",
                ),
                (
                    METRIC_WRITER_BUFFER_SORT_TIME,
                    f"{prefix} Deprecated writer buffer sort time (ms, always 0)",
                ),
                # Commit path (dataset commit attempts and retries).
                (
                    METRIC_COMMIT_TIME,
                    f"{prefix} Total commit time (ms)",
                ),
                (
                    METRIC_COMMIT_BACKOFF_TIME,
                    f"{prefix} Total commit backoff time (ms)",
                ),
                (
                    METRIC_COMMIT_CONFLICT_RETRIES,
                    f"{prefix} Commit version conflict retries",
                ),
                (
                    METRIC_COMMIT_CONCURRENT_WRITER_RETRIES,
                    f"{prefix} Commit concurrent writer retries",
                ),
                # Batch shape (writer-reconciled averages).
                (
                    METRIC_AVG_BATCH_NUM_ROWS,
                    f"{prefix} Avg checkpoint batch rows (unweighted by fragment)",
                ),
                (
                    METRIC_AVG_BATCH_SIZE,
                    f"{prefix} Avg checkpoint batch size (bytes, "
                    f"unweighted by fragment)",
                ),
            ]:
                self.job_tracker.set_desc.remote(m, desc)

        _LOG.info(
            f"Pipeline executing on {cnt_tasks} tasks over "
            f"{cnt_fragments} table fragments"
        )

        # Submit tasks to the actor pool with timeout-based polling so we can
        # log k8s pod status and detect stalls instead of blocking forever.
        submit_fn = lambda actor, value: actor.run.remote(value)  # noqa: E731
        plan_iter = iter(plans)

        def _maybe_submit() -> bool:
            try:
                v = next(plan_iter)
            except StopIteration:
                return False
            pool.submit(submit_fn, v)
            return True

        # Prime the pool: num_actors + 1 so there's always a queued task ready.
        priming = pool._num_actors + 1
        while priming and _maybe_submit():
            priming -= 1

        fwm = FragmentWriterManager(
            ds.version,
            ds_uri=ds.uri,
            map_task=self.map_task,
            checkpoint_store=self.checkpoint_store,
            where=self.where,
            job_tracker=self.job_tracker,
            commit_granularity=self.config.commit_granularity,
            expected_tasks=dict(tasks_by_frag),
            skipped_fragments=self.skipped_fragments or {},
            namespace_impl=self.dst.namespace_impl,
            namespace_properties=self.dst.namespace_properties,
            table_id=self.dst.table_id,
            storage_options=None,
            src_data_files_by_dst=self.src_data_files_by_dst,
        )

        stall_deadline = time.monotonic() + PIPELINE_STALL_TIMEOUT_S
        while pool.has_next():
            try:
                (
                    task,
                    checkpoints,
                    _cnt_udf_computed,
                    udf_processing_time_ms,
                    batch_checkpointing_time_ms,
                    read_io_time_ms,
                    checkpoint_load_time_ms,
                    checkpoint_exists_time_ms,
                    checkpoint_list_time_ms,
                    read_task_total_time_ms,
                ) = pool.get_next_unordered(timeout=POLL_INTERVAL_S)
            except TimeoutError:
                self._log_k8s_pod_status()
                if time.monotonic() > stall_deadline:
                    diag = self._k8s_pod_diagnostics()
                    raise TimeoutError(
                        f"Pipeline stalled: no task completed in "
                        f"{PIPELINE_STALL_TIMEOUT_S}s. "
                        f"Worker pod status:\n{diag}"
                    ) from None
                continue

            # Got a result — reset stall deadline and submit next task.
            stall_deadline = time.monotonic() + PIPELINE_STALL_TIMEOUT_S
            _maybe_submit()

            fwm.ingest_task(task, checkpoints)
            if self.job_tracker:
                metric_deltas = {
                    METRIC_UDF_PROCESSING_TIME: int(udf_processing_time_ms),
                    METRIC_BATCH_CHECKPOINTING_TIME: int(batch_checkpointing_time_ms),
                    METRIC_READ_IO_TIME: int(read_io_time_ms),
                    METRIC_CHECKPOINT_LOAD_TIME: int(checkpoint_load_time_ms),
                    METRIC_CHECKPOINT_EXISTS_TIME: int(checkpoint_exists_time_ms),
                    METRIC_CHECKPOINT_LIST_TIME: int(checkpoint_list_time_ms),
                    METRIC_READ_TASK_TOTAL_TIME: int(read_task_total_time_ms),
                }
                self.job_tracker.batch_increment.remote(metric_deltas)
            # ensure we discover any frgments that finished writing even if the
            # current task belongs to another fragment.
            fwm.poll_all()
            self._try_refresh_cluster_status()

        pool.shutdown()
        fwm.cleanup()
        # Ensure final metrics reflect all processed rows even if some tracker
        # updates were dropped or delayed.
        fwm.finalize_metrics()
        with contextlib.suppress(Exception):
            self._refresh_cluster_status()


@attrs.define
class FragmentWriterSession:
    """This tracks all the batch tasks for a single fragment.

    It is responsible for managing the fragment writer's life cycle and does the
    bookkeeping of inflight tasks, completed tasks, and the queue of tasks to write.
    These are locally tracked and accounted for before the fragment is considered
    complete and ready to be commited to the dataset.

    It expects to be initialized and then fed with `ingest_task` calls. After all tasks
    have been added, it is `seal`ed meaning no more input tasks are expected.  Then it
    can be `drain`ed to yield all completed tasks.
    """

    frag_id: int
    ds_uri: str
    output_columns: list[str]
    checkpoint_store: CheckpointStore
    where: str | None
    data_storage_version: str | None = None
    read_version: int | None = None
    namespace_impl: Optional[str] = None
    namespace_properties: Optional[dict[str, str]] = None
    table_id: Optional[list[str]] = None
    filler_schema: pa.Schema | None = None
    field_ids: list[int] | None = None
    num_physical_rows: int | None = None
    num_logical_rows: int | None = None

    # runtime state.  This is single-threaded and is not thread-safe.
    queue: ray.util.queue.Queue = attrs.field(init=False)
    actor: ActorHandle = attrs.field(init=False)
    cached_tasks: list[tuple[int, Any]] = attrs.field(factory=list, init=False)
    inflight: dict[ray.ObjectRef, int] = attrs.field(factory=dict, init=False)
    _shutdown: bool = attrs.field(default=False, init=False)

    sealed: bool = attrs.field(default=False, init=False)  # no more tasks will be added
    enqueued: int = attrs.field(default=0, init=False)  # total expected tasks
    completed: int = attrs.field(default=0, init=False)  # total compelted tasks
    _restart_count: int = attrs.field(default=0, init=False)  # restart attempts
    # a seal signal is sent if no more batches will be enqueued for this fragment.
    # this is needed because the map task would produce less checkpoints if
    # there are deletions or where filters.
    _seal_signal_sent: bool = attrs.field(default=False, init=False)

    # Graceful degradation: track if this fragment failed (e.g., fragment not found)
    failed: bool = attrs.field(default=False, init=False)
    failure_reason: str | None = attrs.field(default=None, init=False)

    def __attrs_post_init__(self) -> None:
        # Create queue with num_cpus=0 so it doesn't consume scheduling resources
        self.queue = ray.util.queue.Queue(actor_options={"num_cpus": 0})
        self._start_writer()

    def _start_writer(self) -> None:
        self.actor = FragmentWriter.options(  # type: ignore[assignment]
            num_cpus=FRAGMENT_WRITER_NUM_CPUS,  # make it cheap to schedule
            memory=FRAGMENT_WRITER_MEMORY,
        ).remote(
            self.ds_uri,
            self.output_columns,
            self.checkpoint_store.uri(),
            self.frag_id,
            self.queue,
            where=self.where,
            read_version=self.read_version,
            namespace_impl=self.namespace_impl,
            namespace_properties=self.namespace_properties,
            table_id=self.table_id,
            filler_schema=self.filler_schema,
            field_ids=self.field_ids,
            num_physical_rows=self.num_physical_rows,
            num_logical_rows=self.num_logical_rows,
            data_storage_version=self.data_storage_version,
        )
        # prime one future so we can detect when it finishes
        fut = self.actor.write.remote()  # type: ignore[call-arg]
        self.inflight[fut] = self.frag_id  # type: ignore[assignment]

    def shutdown(self) -> None:
        len_inflight = len(self.inflight)
        if len_inflight > 0:
            try:
                is_empty = self.queue.empty()
            except (ray.exceptions.RayError, Exception):
                # queue actor died or unavailble.  assume empty
                is_empty = True
                # queue should be empty and inflight should be 0.
                _LOG.warning(
                    "Shutting down frag %s - queue empty %s, inflight: %d",
                    self.frag_id,
                    is_empty,
                    len_inflight,
                )

        if self._shutdown:
            return  # idempotent
        self.queue.shutdown()
        ray.kill(self.actor)
        self._shutdown = True

    def _restart(self) -> None:
        self._restart_count += 1
        if self._restart_count > MAX_WRITER_RESTARTS:
            raise RuntimeError(
                f"Writer actor for frag {self.frag_id} died "
                f"{self._restart_count} times, exceeding max restarts "
                f"({MAX_WRITER_RESTARTS}). Giving up."
            )

        self.shutdown()

        # Queue uses no resources (0 CPU, no memory requirement)
        self.queue = ray.util.queue.Queue(actor_options={"num_cpus": 0})
        self._seal_signal_sent = False
        self.inflight.clear()
        self.cached_tasks, old_tasks = [], self.cached_tasks
        self._shutdown = False  # Reset shutdown flag to allow restart
        self.__attrs_post_init__()  # recreates writer & first future

        # replay tasks
        for off, res in old_tasks:
            self.queue.put((off, res))
        if self.sealed and not self._seal_signal_sent:
            self.queue.put((-1, ""))
            self._seal_signal_sent = True

    def ingest_task(self, offset: int, result: Any) -> None:
        """Called by manager when a new (offset, result) arrives."""
        self.cached_tasks.append((offset, result))
        self.enqueued += 1
        try:
            self.queue.put((offset, result))
        except (ray.exceptions.ActorDiedError, ray.exceptions.ActorUnavailableError):
            _LOG.warning(
                "Writer actor for frag %s died – restarting (attempt %d/%d)",
                self.frag_id,
                self._restart_count + 1,
                MAX_WRITER_RESTARTS,
            )
            self._restart()

    def poll_ready(self) -> list[FragmentWriteResult]:
        """Non‑blocking check for any finished futures.
        Returns list of FragmentWriteResult that completed."""
        # If already failed, don't poll anymore
        if self.failed:
            return []

        ready, _ = ray.wait(list(self.inflight.keys()), timeout=0.0)
        completed: list[FragmentWriteResult] = []

        for fut in ready:
            try:
                res: FragmentWriteResult = ray.get(fut)
                completed.append(res)
            except (
                ray.exceptions.ActorDiedError,
                ray.exceptions.ActorUnavailableError,
            ):
                _LOG.warning(
                    "Writer actor for frag %s unavailable – restarting (attempt %d/%d)",
                    self.frag_id,
                    self._restart_count + 1,
                    MAX_WRITER_RESTARTS,
                )
                self._restart()
                return []  # will show up next poll
            except Exception as e:
                # Graceful degradation: mark fragment as failed, don't crash pipeline
                error_msg = f"{type(e).__name__}: {e}"
                _LOG.warning(
                    "Fragment %s write failed: %s. Marking as failed and continuing.",
                    self.frag_id,
                    error_msg,
                )
                self.failed = True
                self.failure_reason = error_msg
                self.inflight.pop(fut)
                self.shutdown()
                return []
            assert res.frag_id == self.frag_id
            self.completed += 1
            self.inflight.pop(fut)

        return completed

    def seal(self) -> None:
        if self.sealed:
            return
        self.sealed = True
        if self._seal_signal_sent:
            return
        try:
            # In-band signal to the writer that no more checkpoints will be enqueued.
            self.queue.put((-1, ""))
            self._seal_signal_sent = True
        except (ray.exceptions.ActorDiedError, ray.exceptions.ActorUnavailableError):
            _LOG.warning(
                "Writer queue for frag %s died while sealing – restarting",
                self.frag_id,
            )
            self._restart()

    def drain(self) -> Generator[FragmentWriteResult, None, None]:
        """Yield FragmentWriteResult as futures complete."""
        # If already failed, nothing to drain
        if self.failed:
            return

        idle_rounds = 0
        while self.inflight:
            ready, _ = ray.wait(list(self.inflight.keys()), timeout=5.0)
            if not ready:
                idle_rounds += 1
                # If we've been sealed (no more tasks will arrive) and the writer
                # future isn't making progress, assume the actor is stalled or died
                # without surfacing an exception. Restarting is safe here because
                # any partial output from the dead actor is lost, and cached tasks
                # can be replayed into a fresh writer.
                if self.sealed and idle_rounds >= GENEVA_WRITER_STALL_IDLE_ROUNDS:
                    _LOG.warning(
                        "Writer actor for frag %s appears stalled during drain; "
                        "restarting",
                        self.frag_id,
                    )
                    self._restart()
                    idle_rounds = 0
                continue
            idle_rounds = 0

            for fut in ready:
                try:
                    res: FragmentWriteResult = ray.get(fut)
                    yield res
                    self.completed += 1
                except (
                    ray.exceptions.ActorDiedError,
                    ray.exceptions.ActorUnavailableError,
                ):
                    _LOG.warning(
                        "Writer actor for frag %s died during drain—restarting "
                        "(attempt %d/%d)",
                        self.frag_id,
                        self._restart_count + 1,
                        MAX_WRITER_RESTARTS,
                    )
                    # clear out any old futures, spin up a fresh actor & queue
                    self._restart()
                    # break out to re-enter the while loop with a clean slate
                    break
                except Exception as e:
                    # Graceful degradation: mark fragment as failed, don't crash
                    error_msg = f"{type(e).__name__}: {e}"
                    _LOG.warning(
                        "Fragment %s write failed during drain: %s. "
                        "Marking as failed and continuing.",
                        self.frag_id,
                        error_msg,
                    )
                    self.failed = True
                    self.failure_reason = error_msg
                    self.inflight.pop(fut)
                    self.shutdown()
                    return  # Exit drain since we're failed
                # sucessful write
                self.inflight.pop(fut)


@attrs.define
class FragmentWriterManager:
    """FragmentWriterManager is responsible for writing out fragments
    from the ReadTasks to the destination dataset.

    There is one instance so that we can track pending completed fragments and do
    partial commits.
    """

    dst_read_version: int
    ds_uri: str
    map_task: MapTask
    checkpoint_store: CheckpointStore
    where: str | None
    job_tracker: ActorHandle | None
    commit_granularity: int
    expected_tasks: dict[int, int]  # frag_id, # read tasks
    # int key is frag_id
    skipped_fragments: dict[int, lance.fragment.DataFile] = attrs.field(factory=dict)
    namespace_impl: Optional[str] = None
    namespace_properties: Optional[dict[str, str]] = None
    table_id: Optional[list[str]] = None
    storage_options: Optional[dict[str, str]] = None
    # Source data files per destination fragment (for MV checkpoint validation)
    src_data_files_by_dst: dict[int, frozenset[str]] = attrs.field(factory=dict)

    # internal state
    sessions: dict[int, FragmentWriterSession] = attrs.field(factory=dict, init=False)
    remaining_tasks: dict[int, int] = attrs.field(init=False)
    output_columns: list[str] = attrs.field(init=False)
    output_field_ids: frozenset[int] | None = attrs.field(default=None, init=False)
    writer_filler_schema: pa.Schema | None = attrs.field(default=None, init=False)
    writer_field_ids: list[int] | None = attrs.field(default=None, init=False)
    writer_data_storage_version: str | None = attrs.field(default=None, init=False)
    _dataset: lance.LanceDataset | None = attrs.field(default=None, init=False)
    _fragment_row_counts: dict[int, tuple[int, int]] = attrs.field(
        factory=dict, init=False
    )
    # Track ingested read-task identities to avoid counting replayed task results.
    _ingested_task_keys: set[str] = attrs.field(factory=set, init=False)
    # Track fragment IDs that are skipped to avoid double-counting in progress
    _skipped_fragment_ids: set[int] = attrs.field(factory=set, init=False)
    # Track fragment IDs already recorded to avoid duplicate commit/metrics
    _recorded_fragment_ids: set[int] = attrs.field(factory=set, init=False)
    # Graceful degradation: track failed fragments (frag_id -> error message)
    failed_fragments: dict[int, str] = attrs.field(factory=dict, init=False)
    # (frag_id, lance.fragment.DataFile, # rows)
    rows_input_by_frag: dict[int, int] = attrs.field(factory=dict, init=False)
    # Local reconciled totals (authoritative at teardown even if tracker updates drop)
    _reconciled_rows_checkpointed_total: int = attrs.field(default=0, init=False)
    _reconciled_rows_ready_total: int = attrs.field(default=0, init=False)
    _reconciled_rows_committed_total: int = attrs.field(default=0, init=False)
    _reconciled_written_fragments_total: int = attrs.field(default=0, init=False)
    _reconciled_avg_batch_num_rows_sum: int = attrs.field(default=0, init=False)
    _reconciled_avg_batch_size_sum: int = attrs.field(default=0, init=False)
    to_commit: list[tuple[int, lance.fragment.DataFile, int]] = attrs.field(
        factory=list, init=False
    )

    @property
    def namespace_client(self) -> "LanceNamespace | None":
        """Create namespace client from impl and properties if available."""
        if self.namespace_impl is not None and self.namespace_properties is not None:
            from lance_namespace import connect as namespace_connect

            return namespace_connect(self.namespace_impl, self.namespace_properties)
        return None

    def __attrs_post_init__(self) -> None:
        # all output cols except for _rowaddr because it is implicit since the
        # lancedatafile is writing out in sequential order
        self.output_columns = [
            f.name for f in self.map_task.output_schema() if f.name != "_rowaddr"
        ]
        self.remaining_tasks = dict(self.expected_tasks)

        dataset = self._open_dataset_for_metadata()
        self._dataset = dataset

        if self._dataset is not None:
            self.writer_data_storage_version = self._dataset.data_storage_version

            from geneva.utils.parse_rust_debug import extract_field_ids

            try:
                base_schema = self._dataset.schema
                fields: list[pa.Field] = [
                    base_schema.field(name)
                    for name in self.output_columns
                    if name in base_schema.names
                ]
                fields.append(pa.field("_rowaddr", pa.uint64()))
                self.writer_filler_schema = pa.schema(fields)
            except Exception:  # noqa: PERF203
                self.writer_filler_schema = None

            try:
                writer_field_ids: list[int] = []
                for field in self.output_columns:
                    writer_field_ids.extend(
                        extract_field_ids(self._dataset.lance_schema, field)
                    )
                self.writer_field_ids = writer_field_ids
            except Exception:  # noqa: PERF203
                self.writer_field_ids = None

            try:
                output_field_id_set: set[int] = set()
                for field in self.map_task.output_schema():
                    if field.name == "_rowaddr":
                        continue
                    try:
                        output_field_id_set.update(
                            extract_field_ids(self._dataset.lance_schema, field.name)
                        )
                    except Exception:  # noqa: PERF203
                        _LOG.debug(
                            "Output column %s not found in schema, skipping",
                            field.name,
                        )
                if output_field_id_set:
                    self.output_field_ids = frozenset(output_field_id_set)
            except Exception:  # noqa: PERF203
                self.output_field_ids = None

        # Immediately add skipped fragments to commit list and update progress tracking
        total_skipped_rows = 0
        for frag_id, data_file in self.skipped_fragments.items():
            # Estimate row count from original fragment if available
            try:
                original_fragment = (
                    self._dataset.get_fragment(frag_id) if self._dataset else None
                )
                if original_fragment is None:
                    row_count = 0
                else:
                    row_count = original_fragment.count_rows()
                    # Note: we ignore where filters since we care about rows in the
                    # fragment
            except Exception:
                row_count = 0  # Default if we can't get the count

            self.to_commit.append((frag_id, data_file, row_count))
            self._skipped_fragment_ids.add(
                frag_id
            )  # Track that this fragment is skipped
            self._recorded_fragment_ids.add(frag_id)
            total_skipped_rows += row_count
            _LOG.info(
                f"Added skipped fragment {frag_id} to commit list with {row_count} rows"
            )
        # Skipped rows are effectively "already ready/committed"
        self._reconciled_rows_checkpointed_total += total_skipped_rows
        self._reconciled_rows_ready_total += total_skipped_rows
        self._reconciled_rows_committed_total += total_skipped_rows

    def _open_dataset_for_metadata(self) -> lance.LanceDataset | None:
        try:
            if self.namespace_client and self.table_id:
                return lance.dataset(
                    namespace=self.namespace_client,
                    table_id=self.table_id,
                    storage_options=self.storage_options,
                    version=self.dst_read_version,
                )
            else:
                return lance.dataset(
                    self.ds_uri,
                    storage_options=self.storage_options,
                    version=self.dst_read_version,
                )
        except Exception:
            return None

    def _get_fragment_row_counts(self, frag_id: int) -> tuple[int, int] | None:
        if frag_id in self._fragment_row_counts:
            return self._fragment_row_counts[frag_id]
        if self._dataset is None:
            return None
        try:
            frag = self._dataset.get_fragment(frag_id)
            if frag is None:
                _LOG.warning(
                    "Fragment %d not found in dataset %s", frag_id, self._dataset.uri
                )
                raise ValueError(f"Fragment {frag_id} not found")
            stats = (frag.physical_rows, frag.count_rows())
            self._fragment_row_counts[frag_id] = stats
            return stats
        except Exception as e:
            _LOG.warning(
                "Error getting fragment row counts for fragment %d: %s", frag_id, e
            )
            # return None then the fragment writer will open the dataset,
            # and get the stats from the dataset itself again.
            return None

    def poll_all(self) -> None:
        for frag_id, sess in list(self.sessions.items()):
            # Check for newly failed sessions (graceful degradation)
            if sess.failed and frag_id not in self.failed_fragments:
                if sess.failure_reason:
                    self.failed_fragments[frag_id] = sess.failure_reason
                    _LOG.warning(
                        "Fragment %d marked as failed during poll: %s",
                        frag_id,
                        sess.failure_reason,
                    )
                continue  # Skip polling failed sessions

            for result in sess.poll_ready():
                self._record_fragment(
                    result.frag_id,
                    result.new_file,
                    self.commit_granularity,
                    result.rows_written,
                    buffer_sort_ms=result.buffer_sort_ms,
                    align_ms=result.align_ms,
                    write_ms=result.write_ms,
                    queue_wait_ms=result.queue_wait_ms,
                    checkpoint_read_ms=result.checkpoint_read_ms,
                    avg_batch_num_rows=result.avg_batch_num_rows,
                    avg_batch_size=result.avg_batch_size,
                )

            # Check if session failed during poll_ready
            if (
                sess.failed
                and frag_id not in self.failed_fragments
                and sess.failure_reason
            ):
                self.failed_fragments[frag_id] = sess.failure_reason

    def _session_for_frag(self, frag_id: int) -> FragmentWriterSession:
        sess = self.sessions.get(frag_id)
        if sess is not None:
            return sess
        _LOG.debug("Creating writer for fragment %d", frag_id)
        row_counts = self._get_fragment_row_counts(frag_id)
        num_physical_rows = row_counts[0] if row_counts else None
        num_logical_rows = row_counts[1] if row_counts else None
        sess = FragmentWriterSession(
            frag_id=frag_id,
            ds_uri=self.ds_uri,
            output_columns=self.output_columns,
            checkpoint_store=self.checkpoint_store,
            where=self.where,
            read_version=self.dst_read_version,
            namespace_impl=self.namespace_impl,
            namespace_properties=self.namespace_properties,
            table_id=self.table_id,
            filler_schema=self.writer_filler_schema,
            field_ids=self.writer_field_ids,
            num_physical_rows=num_physical_rows,
            num_logical_rows=num_logical_rows,
            data_storage_version=self.writer_data_storage_version,
        )
        self.sessions[frag_id] = sess
        return sess

    def ingest_task(
        self, task: ReadTask, checkpoints: list[MapBatchCheckpoint]
    ) -> None:
        """Ingest all checkpoints produced for a single ReadTask.

        `expected_tasks` / `remaining_tasks` are tracked in units of ReadTasks. A
        single ReadTask can yield multiple checkpoints when `checkpoint_size` is
        smaller than `task_size`, so we decrement `remaining_tasks` once per ReadTask
        (not once per checkpoint) and seal the fragment only when all ReadTasks
        for that fragment have been processed.
        """
        frag_id = task.dest_frag_id()
        task_key = task.checkpoint_key()

        if task_key in self._ingested_task_keys:
            _LOG.info(
                "Read task %s for fragment %d already ingested; skipping duplicate",
                task_key,
                frag_id,
            )
            return

        # Skip ingesting to already-failed fragments (graceful degradation)
        if frag_id in self.failed_fragments:
            _LOG.debug(
                "Skipping ingest for fragment %d - already marked as failed", frag_id
            )
            return

        sess = self._session_for_frag(frag_id)

        # Check if session is already failed (might have failed during a previous poll)
        if sess.failed:
            if sess.failure_reason and frag_id not in self.failed_fragments:
                self.failed_fragments[frag_id] = sess.failure_reason
            _LOG.debug("Skipping ingest for failed session fragment %d", frag_id)
            return

        for result in checkpoints:
            sess.ingest_task(result.offset, result.checkpoint_key)
            # Track progress in the planner's offset domain (span).
            if result.span > 0:
                self.rows_input_by_frag[frag_id] = self.rows_input_by_frag.get(
                    frag_id, 0
                ) + int(result.span)
                self._reconciled_rows_checkpointed_total += int(result.span)
                if self.job_tracker:
                    try:
                        self.job_tracker.increment.remote(
                            "rows_checkpointed", result.span
                        )
                    except Exception:
                        _LOG.exception(
                            "Failed to update rows_checkpointed for task %s "
                            "(checkpoint %s)",
                            task,
                            result.checkpoint_key,
                        )
            if self.job_tracker and result.udf_rows > 0:
                try:
                    self.job_tracker.increment.remote(
                        METRIC_UDF_VALUES, result.udf_rows
                    )
                except Exception:
                    _LOG.exception(
                        "Failed to update UDF metrics for task %s (checkpoint %s)",
                        task,
                        result.checkpoint_key,
                    )

        # One read task completed for this fragment.
        self.remaining_tasks[frag_id] = self.remaining_tasks.get(frag_id, 0) - 1
        if self.remaining_tasks[frag_id] <= 0:
            sess.seal()
        self._ingested_task_keys.add(task_key)

        # TODO check if previously checkpointed fragment exists

    def _record_fragment(
        self,
        frag_id: int,
        new_file,
        commit_granularity: int,
        rows_written: int,
        *,
        buffer_sort_ms: int = 0,
        align_ms: int = 0,
        write_ms: int = 0,
        queue_wait_ms: int = 0,
        checkpoint_read_ms: int = 0,
        avg_batch_num_rows: int = 0,
        avg_batch_size: int = 0,
    ) -> None:
        if frag_id in self._recorded_fragment_ids:
            _LOG.info(
                "Fragment %d already recorded, skipping duplicate write record",
                frag_id,
            )
            return
        src_files = self.src_data_files_by_dst.get(frag_id)
        src_files_hash = hash_source_files(src_files) if src_files is not None else None
        dedupe_key = _get_fragment_dedupe_key(
            self.ds_uri,
            frag_id,
            self.map_task,
            dataset_version=self.dst_read_version,
            src_files_hash=src_files_hash,
        )
        # Store file name and source data file paths in checkpoint.
        # The source data files are the UNION of file paths across all source
        # fragments that contributed rows to this destination fragment.
        # When any backfill runs (even re-running UDF on same field), new data
        # files with new UUIDs are created, changing the union and triggering
        # checkpoint invalidation.
        checkpoint_data: dict[str, list] = {"file": [new_file.path]}

        # Include source data file paths if available (for MV refresh)
        if src_files is not None:
            checkpoint_data["src_data_files"] = [json.dumps(sorted(src_files))]
            _LOG.debug(
                f"Storing data files in checkpoint for frag {frag_id}: "
                f"{len(src_files)} files"
            )
        if self.output_field_ids is not None:
            checkpoint_data["output_field_ids"] = [
                json.dumps(sorted(self.output_field_ids))
            ]

        # Store UDF version for version change detection (survives compaction)
        udf_version = self.map_task.udf_version()
        if udf_version is not None:
            checkpoint_data["udf_version"] = [udf_version]

        ckpt_start = time.perf_counter()
        self.checkpoint_store[dedupe_key] = pa.RecordBatch.from_pydict(checkpoint_data)
        ckpt_elapsed_ms = int((time.perf_counter() - ckpt_start) * 1000)
        if self.job_tracker:
            writer_metric_deltas = {
                METRIC_FRAGMENT_CHECKPOINTING_TIME: ckpt_elapsed_ms,
                METRIC_WRITER_ALIGN_TIME: int(align_ms),
                METRIC_WRITER_WRITE_TIME: int(write_ms),
                METRIC_WRITER_QUEUE_WAIT_TIME: int(queue_wait_ms),
                METRIC_WRITER_CHECKPOINT_READ_TIME: int(checkpoint_read_ms),
                "writer_fragments": 1,
            }
            self.job_tracker.batch_increment.remote(writer_metric_deltas)

        # Maintain unweighted averages across writer fragments.
        self._reconciled_written_fragments_total += 1
        self._reconciled_avg_batch_num_rows_sum += int(avg_batch_num_rows)
        self._reconciled_avg_batch_size_sum += int(avg_batch_size)
        if self.job_tracker and self._reconciled_written_fragments_total > 0:
            cur_avg_rows = int(
                round(
                    self._reconciled_avg_batch_num_rows_sum
                    / self._reconciled_written_fragments_total
                )
            )
            cur_avg_size = int(
                round(
                    self._reconciled_avg_batch_size_sum
                    / self._reconciled_written_fragments_total
                )
            )
            self.job_tracker.set.remote(METRIC_AVG_BATCH_NUM_ROWS, cur_avg_rows)
            self.job_tracker.set.remote(METRIC_AVG_BATCH_SIZE, cur_avg_size)

        input_rows = int(self.rows_input_by_frag.get(frag_id, 0))

        # Check if this fragment was already committed as a skipped fragment
        # Use _skipped_fragment_ids instead of searching to_commit, since to_commit
        # gets cleared after each commit but _skipped_fragment_ids persists
        if frag_id in self._skipped_fragment_ids:
            _LOG.info(
                f"Fragment {frag_id} was already committed as skipped fragment, "
                f"not re-adding to commit list"
            )
            return  # Don't add it again to avoid double-commit

        # New fragment, add it normally
        _LOG.info(f"Adding new fragment {frag_id} to commit list")
        self.to_commit.append((frag_id, new_file, input_rows))
        self._recorded_fragment_ids.add(frag_id)
        if input_rows > 0 and self.job_tracker:
            self.job_tracker.batch_increment.remote(
                {"rows_ready_for_commit": int(input_rows)}
            )
        # Track totals locally so we can reconcile metrics even if tracker updates drop
        self._reconciled_rows_ready_total += input_rows

        # Track processed writes and hybrid-shutdown
        sess = self.sessions.get(frag_id)
        if sess and sess.sealed and not sess.inflight:
            # flush any pending commit for this fragment
            sess.shutdown()
            self.sessions.pop(frag_id, None)

        self._commit_if_n_fragments(commit_granularity)

    def _record_commit_elapsed(self, start_time: float) -> None:
        if self.job_tracker is None:
            return
        commit_elapsed_ms = int((time.perf_counter() - start_time) * 1000)
        if commit_elapsed_ms > 0:
            self.job_tracker.batch_increment.remote(
                {METRIC_COMMIT_TIME: commit_elapsed_ms}
            )

    # aka _try_commit
    def _commit_if_n_fragments(
        self, commit_granularity: int, robust: bool = False
    ) -> None:
        """Commit fragments if we have enough to meet granularity threshold.

        Args:
            commit_granularity: Minimum number of fragments to commit
            robust: If True, retry RuntimeError with 'Too many concurrent writers'
        """
        n = max(1, int(commit_granularity))
        if len(self.to_commit) < n:
            return

        to_commit = self.to_commit
        self.to_commit = []
        version = self.dst_read_version
        operation = lance.LanceOperation.DataReplacement(
            replacements=[
                lance.LanceOperation.DataReplacementGroup(
                    fragment_id=frag_id,
                    new_file=new_file,
                )
                for frag_id, new_file, _rows in to_commit
            ]
        )

        retry_attempt = 0
        version_conflict_attempts = 0
        max_retries = GENEVA_COMMIT_MAX_RETRIES if robust else 0
        commit_type = "Robust" if robust else "Standard"

        while True:
            attempt_start = time.perf_counter()
            try:
                _LOG.info(
                    "%s commit: %d fragments to %s at version %d%s",
                    commit_type,
                    len(to_commit),
                    self.ds_uri,
                    version,
                    f" (attempt {retry_attempt + 1})"
                    if robust and retry_attempt > 0
                    else "",
                )

                # Reconstruct storage_options_provider if we have namespace config
                # and the namespace provides storage_options (e.g., cloud credentials)
                storage_options_provider, storage_options = (
                    get_storage_options_provider(
                        self.namespace_impl, self.namespace_properties, self.table_id
                    )
                )
                # Pass namespace and table_id for managed versioning
                # This triggers CreateTableVersion API calls through phalanx
                lance.LanceDataset.commit(
                    self.ds_uri,
                    operation,
                    read_version=version,
                    storage_options=storage_options,
                    storage_options_provider=storage_options_provider,
                    namespace=self.namespace_client,
                    table_id=self.table_id,
                )
                self._record_commit_elapsed(attempt_start)
                # rows committed == sum(input rows for fragments just committed)
                # Exclude skipped fragments since they were already counted in
                # "ready for commit"
                committed_rows = sum(
                    _rows
                    for _fid, _new_file, _rows in to_commit
                    if _fid not in self._skipped_fragment_ids
                )
                if committed_rows and self.job_tracker:
                    self.job_tracker.increment.remote("rows_committed", committed_rows)
                self._reconciled_rows_committed_total += committed_rows

                if robust and retry_attempt > 0:
                    _LOG.info(
                        "%s commit succeeded after %d attempts",
                        commit_type,
                        retry_attempt + 1,
                    )
                break
            except OSError as e:
                self._record_commit_elapsed(attempt_start)
                # Handle post-compaction case: DataReplacement fails because
                # compaction merged column files into combined files.
                #
                # Why this happens:
                # 1. Partial backfill creates separate column file (field_ids=[10])
                # 2. Compaction merges all column files into one (field_ids=[0,10])
                # 3. alter_columns changes UDF version
                # 4. Resume backfill recomputes data, creates new file (field_ids=[10])
                # 5. DataReplacement looks for file with field_ids=[10] to replace
                # 6. No such file exists - only merged file with field_ids=[0,10]
                # 7. Error: "no changes were made"
                #
                # The Merge fallback handles this by directly setting fragment
                # metadata with masked original files + new column files.
                #
                # TODO: Investigate behavior when compaction materializes deletions.
                # If deletions are materialized, the row count changes and the new
                # column file may have different row alignment. Lance doesn't
                # validate file row counts (see transaction.rs TODO at line 2090).
                # Current testing shows this works, but more investigation needed.
                if "no changes were made" in str(e):
                    _LOG.info(
                        "DataReplacement failed (post-compaction merged files), "
                        "falling back to Merge operation: %s",
                        e,
                    )
                    fallback_start = time.perf_counter()
                    self._commit_with_merge_fallback(
                        to_commit, version, storage_options, storage_options_provider
                    )
                    self._record_commit_elapsed(fallback_start)
                    break

                # Conflict error has this message:
                # OSError: Commit conflict for version 6: This DataReplacement \
                # transaction is incompatible with concurrent transaction \
                # DataReplacement at version 6.,
                if "Commit conflict for version" not in str(e):
                    # only handle version conflict
                    raise e

                # Version conflict: another backfill may have committed columns to
                # the same fragments. Try to merge our columns with current data.
                version_conflict_attempts += 1
                if self.job_tracker:
                    self.job_tracker.increment.remote(METRIC_COMMIT_CONFLICT_RETRIES, 1)
                if version_conflict_attempts >= GENEVA_VERSION_CONFLICT_MAX_RETRIES:
                    _LOG.error(
                        "%s commit failed after %d version conflict retries: %s",
                        commit_type,
                        version_conflict_attempts,
                        e,
                    )
                    raise e

                _LOG.info(
                    "%s commit failed with version conflict at version %d "
                    "(attempt %d/%d): %s",
                    commit_type,
                    version,
                    version_conflict_attempts,
                    GENEVA_VERSION_CONFLICT_MAX_RETRIES,
                    e,
                )

                # Get latest version from dataset - concurrent commits may have
                # advanced it by more than 1
                storage_options_provider, storage_options = (
                    get_storage_options_provider(
                        self.namespace_impl, self.namespace_properties, self.table_id
                    )
                )
                latest_ds = lance.dataset(self.ds_uri, storage_options=storage_options)
                latest_version = latest_ds.version

                _LOG.info(
                    "Version conflict at %d, latest version is %d, retrying",
                    version,
                    latest_version,
                )

                # Retry with updated read_version - DataReplacement adds our
                # column's data file to the fragment without replacing other
                # columns' files (each file has different field IDs)
                version = latest_version
            except RuntimeError as e:
                self._record_commit_elapsed(attempt_start)
                # Handle "Too many concurrent writers" errors from Lance
                # Only retry in robust mode
                if not robust or "Too many concurrent writers" not in str(e):
                    # only handle concurrent writers error in robust mode
                    raise e

                retry_attempt += 1
                if self.job_tracker:
                    self.job_tracker.increment.remote(
                        METRIC_COMMIT_CONCURRENT_WRITER_RETRIES, 1
                    )
                if retry_attempt >= max_retries:
                    _LOG.error(
                        (
                            "%s commit failed after %d attempts with concurrent "
                            "writers error: %s"
                        ),
                        commit_type,
                        retry_attempt,
                        e,
                    )
                    raise e

                _LOG.info(
                    (
                        "%s commit failed with concurrent writers (attempt %d/%d): %s. "
                        "Retrying with backoff."
                    ),
                    commit_type,
                    retry_attempt,
                    max_retries,
                    e,
                )
                # Use exponential backoff with jitter for concurrent writers
                backoff = min(30.0, 0.5 * (2 ** min(5, retry_attempt)))
                backoff += random.uniform(0, backoff * 0.1)  # Add 10% jitter
                time.sleep(backoff)
                if self.job_tracker:
                    self.job_tracker.increment.remote(
                        METRIC_COMMIT_BACKOFF_TIME, int(backoff * 1000)
                    )

    def _commit_with_merge_fallback(
        self,
        to_commit: list[tuple[int, lance.fragment.DataFile, int]],
        version: int,
        storage_options: dict[str, str] | None,
        storage_options_provider,
    ) -> None:
        """Commit using Merge operation when DataReplacement fails post-compaction.

        Why DataReplacement fails after compaction:
        - DataReplacement requires exact field_ids match to replace/add files
        - After compaction, column files are merged (e.g., [0] + [10] → [0,10])
        - New column file has field_ids=[10], but no existing file has just [10]
        - Lance error: "no changes were made" (can't find file to replace)

        How Merge fallback works:
        - Uses LanceOperation.Merge to directly set fragment metadata
        - Masks our field_ids in original files (set to LANCE_FIELD_ID_TOMBSTONE = -2)
        - Adds new column file with correct field_ids
        - Effectively overlays new column data on existing fragment

        Caveats:
        - Merge requires ALL fragments to be provided (not just modified ones)
        - This can conflict with concurrent writes to other fragments
        - Use only as fallback when DataReplacement fails

        TODO: Investigate behavior when compaction materializes deletions.
        If deletions are materialized (rows physically removed), the fragment's
        physical_rows changes. Our new column file was computed from the
        pre-compaction row count. Lance doesn't validate that file row counts
        match (see transaction.rs:2090 TODO). Current tests pass, but potential
        for data misalignment needs further investigation.
        """
        from lance.fragment import FragmentMetadata

        # Get current dataset state
        ds = lance.dataset(self.ds_uri, storage_options=storage_options)

        # Build map of fragments we're updating
        frag_updates: dict[int, lance.fragment.DataFile] = {
            frag_id: new_file for frag_id, new_file, _rows in to_commit
        }

        # Get field IDs we're writing (to mask in original files)
        field_ids_to_mask = (
            set(self.output_field_ids) if self.output_field_ids else set()
        )

        # Build fragment metadata for ALL fragments
        all_frags = []
        for frag in ds.get_fragments():
            if frag.fragment_id in frag_updates:
                # Modified fragment: mask our columns in existing files, add new file
                new_files = []
                for df in frag.data_files():
                    # Create masked version of this file - tombstone our field_ids
                    masked_field_ids = [
                        LANCE_FIELD_ID_TOMBSTONE if fid in field_ids_to_mask else fid
                        for fid in df.field_ids()
                    ]
                    # Only include if there are non-tombstoned fields
                    if any(fid != LANCE_FIELD_ID_TOMBSTONE for fid in masked_field_ids):
                        masked_df = lance.fragment.DataFile(
                            df.path,
                            masked_field_ids,
                            df.column_indices,
                            df.file_major_version,
                            df.file_minor_version,
                        )
                        new_files.append(masked_df)

                # Add our new column file
                new_files.append(frag_updates[frag.fragment_id])

                new_frag = FragmentMetadata(
                    id=frag.fragment_id,
                    files=new_files,
                    physical_rows=frag.physical_rows,
                    deletion_file=frag.metadata.deletion_file,
                )
                all_frags.append(new_frag)
            else:
                # Unmodified fragment: keep existing metadata
                all_frags.append(frag.metadata)

        # Commit with Merge operation
        op = lance.LanceOperation.Merge(
            fragments=all_frags,
            schema=ds.lance_schema,
        )
        # Pass namespace and table_id for managed versioning
        lance.LanceDataset.commit(
            self.ds_uri,
            op,
            read_version=version,
            storage_options=storage_options,
            storage_options_provider=storage_options_provider,
            namespace=self.namespace_client,
            table_id=self.table_id,
        )

        _LOG.info(
            "Merge fallback committed %d fragments to %s",
            len(frag_updates),
            self.ds_uri,
        )

    def cleanup(self) -> None:
        _LOG.debug("draining & shutting down any leftover sessions")

        # 1) Commit any top‑of‑buffer fragments with robust retry logic
        self._commit_if_n_fragments(1, robust=True)

        # 2) Drain & shutdown whatever sessions remain.
        #
        # During planning we count the expected number of read tasks per fragment.
        # With filters, deletes, or adaptive checkpoint spans, a single read task may
        # emit fewer checkpoints than expected, but we decrement remaining_tasks once
        # per read task (not per checkpoint). At job teardown we know no more applier
        # results will arrive, so it is safe to seal any remaining sessions to
        # guarantee termination and idempotent cleanup.
        for _frag_id, sess in list(self.sessions.items()):
            # Collect failed fragments before draining (graceful degradation)
            if sess.failed and sess.failure_reason:
                self.failed_fragments[_frag_id] = sess.failure_reason
                _LOG.warning(
                    "Fragment %d failed: %s. Skipping drain.",
                    _frag_id,
                    sess.failure_reason,
                )
                sess.shutdown()
                continue

            if not sess.sealed:
                sess.seal()
            for result in sess.drain():
                # this may in turn pop more sessions via _record_fragment
                self._record_fragment(
                    result.frag_id,
                    result.new_file,
                    self.commit_granularity,
                    result.rows_written,
                    buffer_sort_ms=result.buffer_sort_ms,
                    align_ms=result.align_ms,
                    write_ms=result.write_ms,
                    queue_wait_ms=result.queue_wait_ms,
                    checkpoint_read_ms=result.checkpoint_read_ms,
                    avg_batch_num_rows=result.avg_batch_num_rows,
                    avg_batch_size=result.avg_batch_size,
                )

            # Check if session failed during drain
            if sess.failed and sess.failure_reason:
                self.failed_fragments[_frag_id] = sess.failure_reason
                _LOG.warning(
                    "Fragment %d failed during drain: %s",
                    _frag_id,
                    sess.failure_reason,
                )

            sess.shutdown()

        # 3) Clear out any sessions that finished in the loop above
        self.sessions.clear()

        # 4) Final safety commit of anything left with robust retry logic
        self._commit_if_n_fragments(1, robust=True)

        # 5) Log summary of failed fragments for graceful degradation
        if self.failed_fragments:
            _LOG.warning(
                "Graceful degradation: %d fragment(s) failed. "
                "Re-run backfill to process failed fragments.",
                len(self.failed_fragments),
            )
            for frag_id, reason in self.failed_fragments.items():
                _LOG.warning("  Fragment %d: %s", frag_id, reason)

    def finalize_metrics(self) -> None:
        """Ensure row progress metrics are fully populated at job end.

        Ray actor messages can be dropped or delayed under heavy load; by the time the
        job finishes the tracker may not have seen every incremental update.  We keep
        local running totals and overwrite the tracker metrics with the reconciled
        values so finished jobs always report the correct counts.
        """
        if self.job_tracker is None:
            return

        try:
            total_checkpointed = self._reconciled_rows_checkpointed_total
            total_ready = self._reconciled_rows_ready_total
            total_committed = self._reconciled_rows_committed_total
            total_written_frags = self._reconciled_written_fragments_total
            avg_rows = (
                int(
                    round(self._reconciled_avg_batch_num_rows_sum / total_written_frags)
                )
                if total_written_frags
                else 0
            )
            avg_size = (
                int(round(self._reconciled_avg_batch_size_sum / total_written_frags))
                if total_written_frags
                else 0
            )

            refs = []
            refs.append(
                self.job_tracker.finalize_rows.remote(
                    total_checkpointed,
                    total_ready,
                    total_committed,
                )
            )

            # Best-effort: these are derived metrics (no totals).
            for name, value in (
                (METRIC_AVG_BATCH_NUM_ROWS, avg_rows),
                (METRIC_AVG_BATCH_SIZE, avg_size),
            ):
                refs.append(self.job_tracker.set.remote(name, value))
                refs.append(self.job_tracker.mark_done.remote(name))

            if refs:
                # Best‑effort wait so we don't hang shutdown indefinitely
                try:
                    ray.get(refs, timeout=5.0)
                except ray.exceptions.GetTimeoutError:
                    _LOG.warning(
                        "Final metric reconciliation timed out; metrics will still be "
                        "correct on the next tracker flush if the actor stays alive"
                    )
        except Exception:
            _LOG.exception("Failed to finalize row metrics")


def _fetch_udf(table: Table, udf_path: str) -> bytes:
    """
    Fetch UDF payload from table storage using Lance file sessions.

    Uses Lance's native object_store layer which handles credential
    management for all storage backends (S3, GCS, Azure, local).

    Parameters
    ----------
    table : Table
        The Geneva table containing the UDF
    udf_path : str
        Relative path within the dataset (e.g., "_udfs/checksum")

    Returns
    -------
    bytes
        The serialized UDF package
    """
    import os
    import tempfile

    ds = table.to_lance()
    session = ds.new_file_session()

    with tempfile.NamedTemporaryFile(
        mode="rb", delete=False, suffix=".udf"
    ) as tmp_file:
        tmp_path = tmp_file.name

    try:
        session.download_file(udf_path, tmp_path)
        with open(tmp_path, "rb") as f:
            udf_payload = f.read()
        _LOG.debug(f"Downloaded UDF from {udf_path} ({len(udf_payload)} bytes)")
        return udf_payload
    except Exception as e:
        raise RuntimeError(f"Failed to download UDF at {udf_path}: {e}") from e
    finally:
        os.unlink(tmp_path)


def fetch_udf(table: Table, column_name: str) -> UDFSpec:
    schema = table._native_tbl.schema
    field = schema.field(column_name)
    if field is None:
        raise ValueError(f"Column {column_name} not found in table {table}")

    udf_path = metadata_value("virtual_column.udf", field.metadata)

    udf_payload = _fetch_udf(table, udf_path)

    udf_name = metadata_value("virtual_column.udf_name", field.metadata)
    udf_backend = metadata_value("virtual_column.udf_backend", field.metadata)

    return UDFSpec(
        name=udf_name,
        backend=udf_backend,
        udf_payload=udf_payload,
    )


def metadata_value(key: str, metadata: dict[bytes, bytes] | None) -> str:
    if metadata is None:
        raise ValueError(f"Metadata is None, cannot find key {key}")
    value = metadata.get(key.encode("utf-8"))
    if value is None:
        raise ValueError(f"Metadata key {key} not found in metadata {metadata}")
    return value.decode("utf-8")


def _get_matview_version(schema: pa.Schema) -> int:
    """Get materialized view version from schema metadata.

    Args:
        schema: PyArrow schema from materialized view table

    Returns:
        int: Version number (1 for legacy fragment+offset, 2 for direct row IDs,
             3 for scalar UDTF views).
             Defaults to 1 for backwards compatibility with v0.7.x MVs.
    """
    metadata = schema.metadata or {}
    version_bytes = metadata.get(MATVIEW_META_VERSION.encode(), b"1")
    version_str = version_bytes.decode()
    if version_str == "scalar_udtf":
        # Scalar UDTF views use stable row IDs from the source, like version 2.
        return 3
    return int(version_str)


# Maximum number of row IDs to include in a single DELETE statement
# to avoid SQL statement size limits
MAX_DELETE_BATCH_SIZE = 10000


def _get_valid_source_row_ids_at_version(
    src_db_uri: str,
    src_table_name: str,
    target_version: int,
    query: GenevaQuery,
    namespace_impl: str | None = None,
    namespace_properties: dict[str, str] | None = None,
) -> set[int]:
    """Get all row IDs that exist in source table at target version.

    Applies the MV's WHERE filter to only return matching rows.

    Args:
        src_db_uri: URI of the source database
        src_table_name: Name of the source table
        target_version: Version of the source table to query
        query: The GenevaQuery containing the WHERE filter
        namespace_impl: Optional namespace implementation
        namespace_properties: Optional namespace properties

    Returns:
        set: Row IDs that exist in the source table at target_version and match filter
    """
    from geneva.db import connect

    # Connect to source database (use context manager to ensure connection is closed)
    with connect(
        src_db_uri,
        namespace_impl=namespace_impl,
        namespace_properties=namespace_properties,
    ) as src_db:
        src_table = src_db.open_table(src_table_name)
        src_dataset = src_table.to_lance().checkout_version(target_version)

        # Build scanner with the MV's WHERE filter
        filter_expr = query.base.filter if query.base and query.base.filter else None

        scanner = src_dataset.scanner(
            columns=[],  # Only need row IDs, not data columns
            filter=filter_expr,
            with_row_id=True,
        )

        # Collect all valid row IDs
        valid_row_ids: set[int] = set()
        for batch in scanner.to_batches():
            if "_rowid" in batch.schema.names:
                valid_row_ids.update(
                    rid for rid in batch["_rowid"].to_pylist() if rid is not None
                )

        _LOG.info(
            f"Found {len(valid_row_ids)} valid row IDs in source table "
            f"at version {target_version}"
        )
        return valid_row_ids


def _delete_rows_not_in_source_version(
    dst_table: Table,
    valid_source_row_ids: set[int],
    batch_size: int = MAX_DELETE_BATCH_SIZE,
) -> int:
    """Delete destination rows whose __source_row_id is not in the valid set.

    Args:
        dst_table: The destination (MV) table
        valid_source_row_ids: Set of row IDs that are valid in the target source version
        batch_size: Number of rows to delete per batch (default: MAX_DELETE_BATCH_SIZE)

    Returns:
        int: Count of deleted rows
    """
    # Get current destination row IDs
    dst_data = dst_table.to_lance().scanner(columns=["__source_row_id"]).to_table()

    current_row_ids = {
        rid for rid in dst_data["__source_row_id"].to_pylist() if rid is not None
    }

    # Find rows to delete (in destination but NOT in source)
    rows_to_delete = current_row_ids - valid_source_row_ids

    if not rows_to_delete:
        _LOG.info("No rows to delete - all destination rows exist in target version")
        return 0

    _LOG.info(f"Deleting {len(rows_to_delete)} rows not present in target version")

    # Batch deletions for large sets to avoid SQL statement size limits
    rows_list = list(rows_to_delete)
    total_deleted = 0

    for i in range(0, len(rows_list), batch_size):
        batch = rows_list[i : i + batch_size]
        row_ids_str = ",".join(str(rid) for rid in batch)
        dst_table.delete(where=f"__source_row_id IN ({row_ids_str})")
        total_deleted += len(batch)
        _LOG.info(f"Deleted batch of {len(batch)} rows (total: {total_deleted})")

    return total_deleted


def _delete_stale_mv_rows(
    dst: TableReference,
    dst_table: Table,
    src_dburi: str,
    src_name: str,
    src_version: int,
    query: GenevaQuery,
    delete_batch_size: int,
    reason: str,
    existing_source_row_ids: set[int] | None = None,
) -> tuple[Table, lance.LanceDataset, int, set[int]]:
    """Delete MV rows whose source rows don't exist at src_version.

    This is used by both backward refresh (rollback) and forward refresh
    (delete sync) to remove stale rows from the materialized view.

    Args:
        dst: The destination table reference
        dst_table: The destination (MV) table
        src_dburi: URI of the source database
        src_name: Name of the source table
        src_version: Version of the source table to query
        query: The GenevaQuery containing the WHERE filter
        delete_batch_size: Number of rows to delete per batch
        reason: Description for logging (e.g., "Backward refresh", "Forward refresh")
        existing_source_row_ids: If provided, first checks if any deletions needed
                                 before querying source. Returns early if none.

    Returns:
        Tuple of (dst_table, dst_dataset, deleted_count, valid_row_ids)
    """
    # Get valid row IDs at target version (with MV filter applied)
    valid_row_ids = _get_valid_source_row_ids_at_version(
        src_dburi, src_name, src_version, query
    )

    # Early exit if caller knows there are no deletions
    if existing_source_row_ids is not None:
        rows_to_delete = existing_source_row_ids - valid_row_ids
        if not rows_to_delete:
            _LOG.info(f"{reason}: no rows to delete")
            return dst_table, dst_table.to_lance(), 0, valid_row_ids

    # Delete destination rows not in target version
    deleted_count = _delete_rows_not_in_source_version(
        dst_table, valid_row_ids, batch_size=delete_batch_size
    )
    _LOG.info(f"{reason}: deleted {deleted_count} rows from MV")

    # Re-open destination table after deletion to get fresh state
    dst_table = dst.open()
    dst_dataset = dst_table.to_lance()

    return dst_table, dst_dataset, deleted_count, valid_row_ids


def _extract_fragment_ids_from_row_ids(row_ids: list[int], mv_version: int) -> set[int]:
    """Extract source fragment IDs from __source_row_id values.

    Args:
        row_ids: List of __source_row_id values from materialized view
        mv_version: Materialized view version (1 or 2)

    Returns:
        set: Source fragment IDs referenced by these row IDs

    Note:
        Version 1: Fragment+offset encoding (fragment << 32 | offset)
          - Used for v0.7.x and earlier (always)
          - Used for v0.8.x+ without stable row IDs (Lance's _rowid encoding)
        Version 2: Stable row IDs (v0.8.x+ with stable row IDs enabled)
          - Row IDs are stable logical IDs that persist across compaction
          - Fragment IDs cannot be reliably extracted from row IDs (may change)
          - For now, conservatively return empty set to force full processing
          - TODO: Query Lance to map stable row IDs to current fragment IDs
    """
    if mv_version == 1:
        # Fragment+offset encoding: extract fragment ID from upper 32 bits
        return {row_id >> 32 for row_id in row_ids if row_id is not None}
    else:  # version 2 (stable row IDs) or 3 (scalar UDTF)
        # Stable row IDs: Cannot extract fragment IDs from row IDs.
        # Return empty set to indicate we can't determine source fragments
        # from row IDs alone. This will cause the refresh logic to process
        # all source fragments (conservative but correct).
        return set()


def _validate_checkpoint_data_files(
    checkpoint_store: CheckpointStore,
    dedupe_key: str,
    current_data_files: frozenset[str],
) -> bool:
    """Validate that stored data files match current source data files.

    This detects when source data has been modified via backfill. When any
    backfill runs (including re-running a UDF on the same column), new data
    files with new UUIDs are created. Comparing file paths detects ALL changes.

    Args:
        checkpoint_store: The checkpoint store
        dedupe_key: The checkpoint key for this fragment
        current_data_files: Current data file paths from source fragments

    Returns:
        True if checkpoint is valid (data files match or not stored),
        False if checkpoint is invalid (data files changed)
    """
    _LOG.debug(
        f"_validate_checkpoint_data_files: key={dedupe_key}, "
        f"current={sorted(current_data_files)}"
    )
    if dedupe_key not in checkpoint_store:
        _LOG.info(
            f"_validate_checkpoint_data_files: checkpoint key not found {dedupe_key}, "
            "invalidating to force reprocess"
        )
        return False  # Key not found - force reprocess to ensure correct checkpoint

    try:
        checkpointed_data = checkpoint_store[dedupe_key]
        _LOG.debug(
            f"_validate_checkpoint_data_files: checkpoint schema="
            f"{checkpointed_data.schema.names}"
        )

        # If no src_data_files stored, invalidate checkpoint to ensure refresh.
        # This handles legacy checkpoints created before data file tracking.
        if "src_data_files" not in checkpointed_data.schema.names:
            _LOG.info(
                "_validate_checkpoint_data_files: no src_data_files in "
                "checkpoint, invalidating to ensure refresh"
            )
            return False

        stored_files_json = checkpointed_data["src_data_files"][0].as_py()
        _LOG.debug(
            f"_validate_checkpoint_data_files: stored_files_json={stored_files_json}"
        )
        if stored_files_json is None:
            _LOG.info(
                "_validate_checkpoint_data_files: src_data_files is null, "
                "invalidating to ensure refresh"
            )
            return False

        stored_files = frozenset(json.loads(stored_files_json))
        if stored_files != current_data_files:
            _LOG.info(
                f"Checkpoint data files mismatch for {dedupe_key}: "
                f"stored={sorted(stored_files)}, "
                f"current={sorted(current_data_files)}"
            )
            return False

        _LOG.debug(
            "_validate_checkpoint_data_files: data files match, checkpoint valid"
        )
        return True
    except Exception as e:
        _LOG.debug(f"Failed to validate data files for {dedupe_key}: {e}")
        return True  # On error, assume valid (conservative)


class DstToSrcMappingResult(NamedTuple):
    """Result of mapping destination fragments to source fragments.

    Attributes:
        dst_to_src_map: Mapping from dst fragment ID to set of src fragment IDs
        dst_frags_with_checkpoint: Set of dst fragment IDs with valid checkpoints
        existing_source_row_ids: Set of all source row IDs already in destination
        src_data_files_by_dst: Mapping from dst fragment ID to union of data file
            paths from all contributing source fragments
    """

    dst_to_src_map: dict[int, set[int]]
    dst_frags_with_checkpoint: set[int]
    existing_source_row_ids: set[int]
    src_data_files_by_dst: dict[int, frozenset[str]]


def _build_dst_to_src_mapping(
    dst_dataset: lance.LanceDataset,
    dst_uri: str,
    map_task,
    checkpoint_store: CheckpointStore,
    src_dataset: lance.LanceDataset | None = None,
    relevant_field_ids: frozenset[int] | None = None,
) -> DstToSrcMappingResult:
    """Build mapping from destination to source fragments, track checkpoints.

    This function iterates over destination fragments and:
    1. Extracts which source fragments contributed to each destination fragment
    2. Computes the UNION of source data file paths for checkpoint validation
    3. Validates existing checkpoints against current source data files

    Why track data file paths? When any backfill runs (including re-running a UDF
    on the same column), new data files with new UUIDs are created. By tracking
    the union of all data file paths from source fragments, we detect ANY change.

    Args:
        dst_dataset: The destination Lance dataset
        dst_uri: URI of the destination dataset
        map_task: The map task being applied
        checkpoint_store: The checkpoint store
        src_dataset: The source Lance dataset (optional, for data file validation)
        relevant_field_ids: If provided, only track data files containing these
            field IDs. This filters out unrelated columns from tracking.

    Returns:
        DstToSrcMappingResult with fragment mappings and checkpoint info
    """
    from geneva.apply import _check_fragment_data_file_exists

    # Get MV version for backwards compatibility
    mv_version = _get_matview_version(dst_dataset.schema)
    _LOG.info(f"Materialized view version: {mv_version}")

    dst_to_src_map: dict[int, set[int]] = {}
    dst_frags_with_checkpoint: set[int] = set()
    existing_source_row_ids: set[int] = set()
    src_data_files_by_dst: dict[int, frozenset[str]] = {}
    output_field_ids: frozenset[int] | None = None
    try:
        from geneva.utils.parse_rust_debug import extract_field_ids

        field_id_set: set[int] = set()
        for field in map_task.output_schema():
            if field.name == "_rowaddr":
                continue
            try:
                field_id_set.update(
                    extract_field_ids(dst_dataset.lance_schema, field.name)
                )
            except Exception:  # noqa: PERF203
                _LOG.debug("Output column %s not found in schema, skipping", field.name)
        if field_id_set:
            output_field_ids = frozenset(field_id_set)
    except Exception:  # noqa: PERF203
        output_field_ids = None

    _LOG.debug(
        f"_build_dst_to_src_mapping: src_dataset={src_dataset is not None}, "
        f"dst frags={[f.fragment_id for f in dst_dataset.get_fragments()]}"
    )
    for dst_frag in dst_dataset.get_fragments():
        dst_frag_id = dst_frag.fragment_id

        # Read __source_row_id to map destination to source fragments FIRST
        # so we can validate data files
        source_frag_ids: set[int] = set()
        try:
            scanner = dst_dataset.scanner(
                columns=["__source_row_id"],
                fragments=[dst_frag],
            )
            dst_frag_data = scanner.to_table()
            if len(dst_frag_data) > 0:
                source_row_ids = cast(
                    "list[int]", dst_frag_data["__source_row_id"].to_pylist()
                )
                # Build fragment ID mapping using version-aware extraction
                source_frag_ids = _extract_fragment_ids_from_row_ids(
                    [rid for rid in source_row_ids if rid is not None],
                    mv_version,
                )
                dst_to_src_map[dst_frag_id] = source_frag_ids
                _LOG.debug(
                    f"_build_dst_to_src_mapping: dst_frag={dst_frag_id} -> "
                    f"src_frags={source_frag_ids}"
                )
                # Collect all source row IDs for deduplication
                existing_source_row_ids.update(
                    row_id for row_id in source_row_ids if row_id is not None
                )
        except Exception as e:
            _LOG.warning(
                f"Failed to read __source_row_id from destination fragment "
                f"{dst_frag_id}: {e}"
            )

        # Compute source data files for this destination fragment.
        # We compute the UNION of data file paths from all source fragments
        # that contributed rows to this destination. Union detects when ANY
        # data file changes in ANY source fragment (new backfill creates new files).
        current_data_files: frozenset[str] = frozenset()
        if src_dataset is not None:
            if source_frag_ids:
                # Have specific source fragment IDs - compute union of data files
                current_data_files = get_combined_source_data_files(
                    src_dataset, source_frag_ids, relevant_field_ids
                )
                _LOG.debug(
                    f"Computed source data files for dst frag {dst_frag_id}: "
                    f"src_frags={source_frag_ids}, "
                    f"file_count={len(current_data_files)}"
                )
            else:
                # Can't determine source fragments (e.g., MV v2 with stable row IDs)
                # Use ALL source fragments' data files (conservative but correct)
                all_src_frag_ids = {f.fragment_id for f in src_dataset.get_fragments()}
                current_data_files = get_combined_source_data_files(
                    src_dataset, all_src_frag_ids, relevant_field_ids
                )
                _LOG.debug(
                    f"Computed source data files for dst frag {dst_frag_id} "
                    f"(all source frags): file_count={len(current_data_files)}"
                )
            src_data_files_by_dst[dst_frag_id] = current_data_files
        src_files_hash = (
            hash_source_files(current_data_files) if src_dataset is not None else None
        )

        # Check if this destination fragment has a valid checkpoint
        checkpoint_exists = _check_fragment_data_file_exists(
            dst_uri,
            dst_frag_id,
            map_task,
            checkpoint_store,
            dataset_version=dst_dataset.version,
            src_files_hash=src_files_hash,
            current_output_field_ids=output_field_ids,
        )

        # Validate data files if checkpoint exists and source dataset is provided
        # Note: We validate even if source_frag_ids is empty (MV v2 with stable row IDs)
        # because current_data_files covers all source fragments in that case
        if checkpoint_exists and src_dataset is not None:
            dedupe_key = _get_fragment_dedupe_key(
                dst_uri,
                dst_frag_id,
                map_task,
                dataset_version=dst_dataset.version,
                src_files_hash=src_files_hash,
            )
            data_files_valid = _validate_checkpoint_data_files(
                checkpoint_store, dedupe_key, current_data_files
            )
            if not data_files_valid:
                _LOG.info(
                    f"Destination fragment {dst_frag_id} checkpoint invalidated: "
                    f"source data files changed"
                )
                checkpoint_exists = False

        if checkpoint_exists:
            dst_frags_with_checkpoint.add(dst_frag_id)
            _LOG.info(f"Destination fragment {dst_frag_id} has valid checkpoint")
        else:
            _LOG.info(f"Destination fragment {dst_frag_id} has no valid checkpoint")

    return DstToSrcMappingResult(
        dst_to_src_map=dst_to_src_map,
        dst_frags_with_checkpoint=dst_frags_with_checkpoint,
        existing_source_row_ids=existing_source_row_ids,
        src_data_files_by_dst=src_data_files_by_dst,
    )


def _identify_new_source_fragments(
    src_dataset: lance.LanceDataset,
    dst_to_src_map: dict[int, set[int]],
) -> list[int]:
    """Identify source fragments that have no representation in destination yet.

    These are truly NEW source fragments that need placeholder rows added.

    Returns:
        list: Source fragment IDs not present in any destination fragment
    """
    # Collect all source fragment IDs that exist in destination
    source_frags_in_destination = set()
    for src_frags in dst_to_src_map.values():
        source_frags_in_destination.update(src_frags)

    # Find source fragments not yet in destination
    new_source_fragments = []
    for frag in src_dataset.get_fragments():
        if frag.fragment_id not in source_frags_in_destination:
            _LOG.info(
                f"Found new source fragment {frag.fragment_id} - needs placeholder rows"
            )
            new_source_fragments.append(frag.fragment_id)

    return new_source_fragments


def _extract_new_row_ids_from_source_fragment(
    src_table,
    query: GenevaQuery,
    frag_id: int,
    existing_source_row_ids: set[int],
) -> list[int]:
    """Extract new row IDs from a source fragment that match the query filter.

    Returns:
        list: New row IDs from this fragment (not already in destination)
    """
    # Build query for this fragment with filters applied
    fragment_query_obj = GenevaQuery(
        fragment_ids=[frag_id],
        base=query.base.model_copy(deep=True),
    )
    fragment_query_obj.base.with_row_id = True
    fragment_query_obj.base.columns = []  # Only get row IDs
    fragment_query_obj.column_udfs = None
    fragment_query_obj.with_row_address = None

    fragment_query_builder = GenevaQueryBuilder.from_query_object(
        src_table, fragment_query_obj
    )

    # Get row IDs from this fragment (with filters applied)
    fragment_data = fragment_query_builder.to_arrow()
    _LOG.info(f"  Fragment {frag_id} query returned {len(fragment_data)} rows")

    if len(fragment_data) == 0:
        return []

    row_ids = cast("list[int]", fragment_data["_rowid"].to_pylist())
    _LOG.info(f"  Row IDs from fragment {frag_id}: {row_ids}")
    # Only add row IDs that don't already exist in destination
    new_row_ids = [
        row_id for row_id in row_ids if row_id not in existing_source_row_ids
    ]

    return new_row_ids


def _append_placeholder_fragments(
    dst_table,
    new_fragment_row_ids: list[int],
    max_rows_per_fragment: int | None = None,
) -> set[int] | None:
    """Append placeholder rows for new source data as new destination fragment(s).

    Args:
        dst_table: The destination table to append to
        new_fragment_row_ids: List of source row IDs to create placeholders for
        max_rows_per_fragment: Optional max rows per fragment. If the number of rows
            exceeds this, multiple fragments will be created by splitting the data
            and calling add() for each chunk. Must be at least 1.

    Returns:
        set[int] | None: The new destination fragment IDs, or None if no rows to add
    """
    if not new_fragment_row_ids:
        _LOG.info(
            "No new placeholder rows needed - all source row IDs already in destination"
        )
        return None

    # Validate max_rows_per_fragment if provided
    if max_rows_per_fragment is not None and max_rows_per_fragment < 1:
        raise ValueError(
            f"max_rows_per_fragment must be at least 1 (got {max_rows_per_fragment})"
        )

    _LOG.info(
        f"Adding {len(new_fragment_row_ids)} placeholder rows for new source data"
        + (
            f" (max_rows_per_fragment={max_rows_per_fragment})"
            if max_rows_per_fragment
            else ""
        )
    )

    # Get the current destination fragment count before adding
    dst_dataset_before = dst_table.to_lance()
    existing_fragment_ids = {
        frag.fragment_id for frag in dst_dataset_before.get_fragments()
    }
    _LOG.info(
        f"Existing destination fragment IDs before append: {existing_fragment_ids}"
    )

    # Split the row IDs into chunks based on max_rows_per_fragment
    # Each chunk will be added separately to create separate fragments
    if max_rows_per_fragment is not None and max_rows_per_fragment > 0:
        chunks = [
            new_fragment_row_ids[i : i + max_rows_per_fragment]
            for i in range(0, len(new_fragment_row_ids), max_rows_per_fragment)
        ]
    else:
        chunks = [new_fragment_row_ids]

    _LOG.info(f"Splitting into {len(chunks)} chunks for separate fragments")

    # Add each chunk as a separate fragment
    for i, chunk in enumerate(chunks):
        placeholder_data = pa.table(
            [
                pa.array(chunk, type=pa.int64()),
                pa.array([False] * len(chunk), type=pa.bool_()),
            ],
            names=["__source_row_id", "__is_set"],
        )
        dst_table.add(placeholder_data)
        _LOG.debug(f"Added chunk {i + 1}/{len(chunks)} with {len(chunk)} rows")

    # Identify the newly created fragment IDs
    dst_dataset_after = dst_table.to_lance()
    new_fragment_ids = {frag.fragment_id for frag in dst_dataset_after.get_fragments()}
    added_fragments = new_fragment_ids - existing_fragment_ids

    if len(added_fragments) >= 1:
        _LOG.info(
            f"Successfully added {len(added_fragments)} placeholder fragment(s) "
            f"{added_fragments} with {len(new_fragment_row_ids)} total rows"
        )
        return added_fragments
    else:
        _LOG.warning(f"Expected at least 1 new fragment but got {len(added_fragments)}")
        return None


def _make_scalar_udtf_expand_actor() -> Any:
    """Lazily define the ScalarUDTFExpandActor Ray actor class.

    Defers ``import ray`` until Ray is actually needed.
    """
    import ray

    @ray.remote
    class ScalarUDTFExpandActor:
        """Ray actor that expands source batches via a ScalarUDTF.

        Each actor deserialises the pickled ScalarUDTF once in ``__init__``
        and reuses it across ``expand_batch`` calls.  Batches are transferred
        as IPC bytes to avoid Arrow serialisation edge-cases in Ray.
        """

        def __init__(self, scalar_udtf_obj: Any) -> None:
            self._scalar_udtf = scalar_udtf_obj

        def expand_batch(
            self,
            batch_ipc: bytes,
            inherited_col_names: list[str],
            view_schema_ipc: bytes,
        ) -> tuple[bytes | None, int, int]:
            """Expand one source batch. Returns (result_ipc, src_rows, out_rows)."""
            from typing import cast as _cast

            import pyarrow as pa

            reader = pa.ipc.open_stream(batch_ipc)
            batch = reader.read_all().to_batches()[0]
            src_rows = batch.num_rows

            expanded = self._scalar_udtf.execute_on_record_batch(batch)
            out_rows = expanded.num_rows
            if out_rows == 0:
                return (None, src_rows, 0)

            # Join inherited columns from source batch
            src_row_ids_in = batch["__source_row_id"]
            row_id_to_idx: dict[int, int] = {}
            for i in range(batch.num_rows):
                row_id_to_idx[src_row_ids_in[i].as_py()] = i

            expanded_src_ids = _cast(
                "list[int]", expanded["__source_row_id"].to_pylist()
            )
            source_indices = pa.array(
                [row_id_to_idx[sid] for sid in expanded_src_ids],
                type=pa.int32(),
            )

            view_schema = pa.ipc.open_stream(view_schema_ipc).schema

            result_arrays: dict[str, pa.Array] = {
                "__source_row_id": expanded["__source_row_id"],
                "__child_index": expanded["__child_index"],
            }
            for col_name in inherited_col_names:
                result_arrays[col_name] = batch[col_name].take(source_indices)
            for field in self._scalar_udtf.output_schema:
                result_arrays[field.name] = expanded[field.name]

            result_batch = pa.RecordBatch.from_pydict(result_arrays, schema=view_schema)

            sink = pa.BufferOutputStream()
            writer = pa.ipc.new_stream(sink, result_batch.schema)
            writer.write_batch(result_batch)
            writer.close()
            return (sink.getvalue().to_pybytes(), src_rows, out_rows)

    return ScalarUDTFExpandActor


_ScalarUDTFExpandActor: Any = None


def _get_scalar_udtf_expand_actor() -> Any:
    """Return the ScalarUDTFExpandActor class, creating it on first call."""
    global _ScalarUDTFExpandActor  # noqa: PLW0603
    if _ScalarUDTFExpandActor is None:
        _ScalarUDTFExpandActor = _make_scalar_udtf_expand_actor()
    return _ScalarUDTFExpandActor


def _serialize_batch_ipc(batch: pa.RecordBatch) -> bytes:
    """Serialize a RecordBatch to IPC bytes for Ray transfer."""
    sink = pa.BufferOutputStream()
    writer = pa.ipc.new_stream(sink, batch.schema)
    writer.write_batch(batch)
    writer.close()
    return sink.getvalue().to_pybytes()


def _serialize_schema_ipc(schema: pa.Schema) -> bytes:
    """Serialize a Schema to IPC bytes for Ray transfer."""
    sink = pa.BufferOutputStream()
    writer = pa.ipc.new_stream(sink, schema)
    writer.close()
    return sink.getvalue().to_pybytes()


def _append_expanded_fragments(
    src_table,
    dst_table,
    scalar_udtf_obj: Any,
    source_query: GenevaQuery,
    new_source_fragments: list[int],
    existing_source_row_ids: set[int],
    max_rows_per_fragment: int | None = None,
    concurrency: int = 4,
    dst_ref: TableReference | None = None,
    job_id: str | None = None,
) -> set[int] | None:
    """Expand new source rows via a scalar UDTF and append to the destination.

    For each new source fragment, reads source rows (filtered to those not
    already in destination), distributes batches to Ray actors running the
    scalar UDTF, collects expanded results, and appends to the destination.

    Returns:
        set[int] | None: The new destination fragment IDs, or None if nothing added.
    """
    if not new_source_fragments:
        _LOG.info("No new source fragments to expand for scalar UDTF")
        return None

    # Get existing destination fragment IDs for comparison
    dst_dataset_before = dst_table.to_lance()
    existing_fragment_ids = {
        frag.fragment_id for frag in dst_dataset_before.get_fragments()
    }

    # Determine inherited column names (everything in view except meta + UDTF output)
    view_schema = dst_table.schema
    udtf_output_names = {f.name for f in scalar_udtf_obj.output_schema}
    inherited_col_names = [
        f.name
        for f in view_schema
        if f.name not in udtf_output_names
        and f.name not in ("__source_row_id", "__child_index")
    ]

    # Read all source batches up front, filtering to new rows only
    work_items: list[bytes] = []
    for frag_id in new_source_fragments:
        fragment_query = GenevaQuery(
            fragment_ids=[frag_id],
            base=source_query.base.model_copy(deep=True),
        )
        fragment_query.base.with_row_id = True
        fragment_query.column_udfs = None
        fragment_query.with_row_address = None

        qb = GenevaQueryBuilder.from_query_object(src_table, fragment_query)
        source_data = qb.to_arrow()

        if source_data.num_rows == 0:
            continue

        # Filter to only new source row IDs
        row_ids = source_data["_rowid"]
        mask = pa.array(
            [
                row_ids[i].as_py() not in existing_source_row_ids
                for i in range(source_data.num_rows)
            ],
            type=pa.bool_(),
        )
        source_data = source_data.filter(mask)
        if source_data.num_rows == 0:
            continue

        # Rename _rowid to __source_row_id for the UDTF
        col_names = source_data.column_names
        if "_rowid" in col_names:
            idx = col_names.index("_rowid")
            source_data = source_data.rename_columns(
                col_names[:idx] + ["__source_row_id"] + col_names[idx + 1 :]
            )

        work_items.extend(
            _serialize_batch_ipc(b) for b in source_data.to_batches() if b.num_rows > 0
        )

    if not work_items:
        _LOG.info("Scalar UDTF expansion: no new source rows to expand")
        return None

    _LOG.info(
        "Scalar UDTF expansion: dispatching %d batch(es) to %d actors",
        len(work_items),
        min(len(work_items), concurrency),
    )

    # Create actor pool — pass the ScalarUDTF object directly and let Ray
    # handle serialisation so client/worker cloudpickle versions stay consistent.
    actor_resource_kwargs: dict[str, Any] = {
        "num_cpus": scalar_udtf_obj.num_cpus or 1.0,
    }
    if scalar_udtf_obj.num_gpus:
        actor_resource_kwargs["num_gpus"] = scalar_udtf_obj.num_gpus
    if scalar_udtf_obj.memory:
        actor_resource_kwargs["memory"] = scalar_udtf_obj.memory

    actor_cls = _get_scalar_udtf_expand_actor().options(**actor_resource_kwargs)
    actor_factory = functools.partial(actor_cls.remote, scalar_udtf_obj)

    num_actors = min(len(work_items), concurrency)
    jt_name = (
        f"scalar-udtf-{job_id}" if job_id else f"scalar-udtf-{uuid.uuid4().hex[:8]}"
    )
    job_tracker = job_tracker_options(
        name=f"jobtracker-{jt_name}",
        num_cpus=JOBTRACKER_NUM_CPUS,
        memory=JOBTRACKER_MEMORY,
    ).remote(jt_name, dst_ref, enable_saves=dst_ref is not None)  # type: ignore[call-arg]

    # Register so _wait_for_tracked_jobs keeps the cluster alive during expansion.
    # Use a sentinel ObjectRef that resolves immediately — the real gate is the
    # synchronous pool.map_unordered loop below, but registration makes the UDTF
    # visible to cluster-exit bookkeeping (consistent with the standard UDF path).
    from geneva._context import get_current_context
    from geneva.runners.ray.raycluster import RayCluster

    ctx = get_current_context()
    if isinstance(ctx, RayCluster):
        sentinel_ref = ray.put(True)
        ctx.register_tracked_job(jt_name, sentinel_ref, job_tracker)

    pool = ActorPool(actor_factory, num_actors, job_tracker=job_tracker)

    view_schema_ipc = _serialize_schema_ipc(view_schema)
    total_expanded = 0

    job_tracker.set_total.remote("batches", len(work_items))  # type: ignore[attr-defined]
    job_tracker.set_desc.remote("batches", "UDTF batches expanded")  # type: ignore[attr-defined]
    job_tracker.set.remote("batches", 0)  # type: ignore[attr-defined]
    ray_tqdm([], job_tracker, "batches")

    try:
        for result in pool.map_unordered(
            lambda actor, item: actor.expand_batch.remote(
                item, inherited_col_names, view_schema_ipc
            ),
            work_items,
        ):
            result_ipc, src_rows, out_rows = result
            total_expanded += out_rows
            job_tracker.increment.remote("batches", 1)  # type: ignore[attr-defined]

            if result_ipc is None:
                continue

            reader = pa.ipc.open_stream(result_ipc)
            result_table = reader.read_all()

            if (
                max_rows_per_fragment is not None
                and max_rows_per_fragment > 0
                and result_table.num_rows > max_rows_per_fragment
            ):
                for start in range(0, result_table.num_rows, max_rows_per_fragment):
                    chunk = result_table.slice(start, max_rows_per_fragment)
                    dst_table.add(chunk)
            else:
                dst_table.add(result_table)
    finally:
        job_tracker.mark_done.remote("batches")  # type: ignore[attr-defined]
        pool.shutdown()
        with contextlib.suppress(Exception):
            job_tracker.mark_job_done.remote()
        ray.kill(job_tracker)

    if total_expanded == 0:
        _LOG.info("Scalar UDTF expansion produced no output rows")
        return None

    # Identify new fragment IDs
    dst_dataset_after = dst_table.to_lance()
    new_fragment_ids = {frag.fragment_id for frag in dst_dataset_after.get_fragments()}
    added_fragments = new_fragment_ids - existing_fragment_ids

    _LOG.info(
        "Scalar UDTF expansion: %d batches → %d expanded rows across %d fragment(s)",
        len(work_items),
        total_expanded,
        len(added_fragments),
    )
    return added_fragments if added_fragments else None


def _determine_fragments_to_process(
    dst_to_src_map: dict[int, set[int]],
    dst_frags_with_checkpoint: set[int],
    new_dst_frag_ids: set[int] | None,
) -> set[int] | None:
    """Determine which destination fragments need processing (destination-driven).

    We process destination fragments that **do not** already have a fragment-level
    checkpointed DataFile.

    IMPORTANT: This helper previously returned ``None`` when there was nothing to
    process. Downstream, ``None`` is interpreted as "process all fragments", which
    caused us to re-run work even when every fragment was checkpointed and, worse,
    could deadlock the writer because fragments marked "skipped" don't necessarily
    produce any checkpoint batches.
    """
    # Find all destination fragments without checkpoints
    fragments_to_process = set(dst_to_src_map.keys()) - dst_frags_with_checkpoint

    # Add newly created placeholder fragments (not in dst_to_src_map yet)
    if new_dst_frag_ids:
        fragments_to_process.update(new_dst_frag_ids)
        _LOG.info(
            f"Adding {len(new_dst_frag_ids)} newly created placeholder "
            f"fragment(s) {new_dst_frag_ids} to processing list"
        )

    if not fragments_to_process:
        _LOG.info("All destination fragments have checkpoints - nothing to process")
        return set()

    _LOG.info(f"Will process destination fragments: {fragments_to_process}")
    return fragments_to_process


def _safe_extract_field_ids(
    lance_schema, field_name: str, failed_fields: list[str]
) -> list[int]:
    """Extract field IDs, recording failures instead of raising exceptions."""
    try:
        return extract_field_ids(lance_schema, field_name)
    except Exception:
        failed_fields.append(field_name)
        return []


def _collect_skipped_fragments(
    dst_dataset: lance.LanceDataset,
    dst_frags_with_checkpoint: set[int],
    map_task,
    checkpoint_store: CheckpointStore,
    src_data_files_by_dst: dict[int, frozenset[str]] | None = None,
) -> tuple[dict[int, lance.fragment.DataFile], dict[str, int]]:
    """Collect checkpointed fragments for commit inclusion.

    Returns:
        tuple: (skipped_fragments, skipped_stats)
            - skipped_fragments: dict mapping frag ID -> DataFile object
            - skipped_stats: dict with 'fragments' and 'rows' counts
    """

    skipped_fragments = {}
    skipped_stats = {"fragments": 0, "rows": 0}

    dsv_major, dsv_minor = parse_data_storage_version(dst_dataset.data_storage_version)

    for frag_id in dst_frags_with_checkpoint:
        try:
            fragment = dst_dataset.get_fragment(frag_id)
            if fragment is None:
                _LOG.warning(f"Fragment {frag_id} not found")
                continue
            frag_rows = fragment.count_rows()
        except Exception as e:
            _LOG.warning(f"Could not get fragment {frag_id}: {e}")
            continue

        skipped_stats["fragments"] += 1
        skipped_stats["rows"] += frag_rows

        # Get checkpoint data
        src_files = (
            src_data_files_by_dst.get(frag_id)
            if src_data_files_by_dst is not None
            else None
        )
        src_files_hash = hash_source_files(src_files) if src_files is not None else None
        dedupe_key = _get_fragment_dedupe_key(
            dst_dataset.uri,
            frag_id,
            map_task,
            src_files_hash=src_files_hash,
        )
        checkpointed_data = checkpoint_store[dedupe_key]
        file_list = checkpointed_data["file"].to_pylist()
        file_path = "".join(str(f) for f in file_list if f is not None)

        # Extract field_ids for all columns in the output schema
        # For CopyTableTask (materialized views), the checkpoint contains all columns
        # from the view schema, not just the UDF outputs
        field_ids = []
        output_schema = map_task.output_schema()
        failed_fields = []

        for field_name in output_schema.names:
            extracted = _safe_extract_field_ids(
                dst_dataset.lance_schema, field_name, failed_fields
            )
            field_ids.extend(extracted)

        if failed_fields:
            _LOG.warning(
                f"Could not extract field_ids for fields {failed_fields} "
                f"in checkpointed fragment {frag_id}"
            )

        # Create DataFile object for the checkpointed fragment
        existing_data_file = lance.fragment.DataFile(
            file_path,
            field_ids,
            list(range(len(field_ids))),
            dsv_major,
            dsv_minor,
        )
        skipped_fragments[frag_id] = existing_data_file

    _LOG.info(
        f"Collected {skipped_stats['fragments']} checkpointed fragments "
        f"with {skipped_stats['rows']} rows"
    )

    return skipped_fragments, skipped_stats


def _validate_mv_source_schema(
    query: GenevaQuery,
    src_table: "Table",
    mv_name: str,
    packager: UDFPackager,
) -> None:
    """
    Validate that source table has all columns required by materialized view query.

    This prevents cryptic errors during refresh when source table schema has evolved
    and dropped columns that the MV depends on.

    Parameters
    ----------
    query : GenevaQuery
        The materialized view query from metadata
    src_table : Table
        The source table to validate against
    mv_name : str
        Name of the materialized view (for error messages)
    packager : UDFPackager
        Packager to unmarshal UDFs to extract input columns

    Raises
    ------
    ValueError
        If source table is missing columns required by the MV query
    """
    required_columns = set()

    # Extract columns from base query SELECT clause
    if query.base.columns:
        if isinstance(query.base.columns, list):
            # Simple column list: ['col1', 'col2']
            required_columns.update(query.base.columns)
        elif isinstance(query.base.columns, dict):
            # Column mapping: {'output': 'source_col'} or {'output': <expression>}
            # Extract string values (column names), skip UDFs (handled separately)
            source_col_names = set(src_table.schema.names)
            for value in query.base.columns.values():
                if isinstance(value, str):
                    if value in source_col_names:
                        # Direct column reference
                        required_columns.add(value)
                    else:
                        # SQL expression - extract referenced columns
                        for col_name in source_col_names:
                            if col_name in value:
                                required_columns.add(col_name)

    # Extract input columns from UDFs
    if query.column_udfs:
        for column_udf in query.column_udfs:
            try:
                # Unmarshal UDF to access input_columns
                udf = packager.unmarshal(column_udf.udf.to_attrs())
                if udf.input_columns is not None:  # type: ignore[reportOptionalMemberAccess]
                    required_columns.update(udf.input_columns)  # type: ignore[reportOptionalMemberAccess]
            except Exception as e:  # noqa: PERF203
                # Note: Performance overhead acceptable for rare schema validation
                _LOG.warning(
                    f"Failed to unmarshal UDF '{column_udf.output_name}' "
                    f"for schema validation: {e}"
                )
                # Continue validation with other columns

    # Compare against source table schema
    source_columns = set(src_table.schema.names)
    missing_columns = required_columns - source_columns

    if missing_columns:
        missing_list = sorted(missing_columns)
        raise ValueError(
            f"Cannot refresh materialized view '{mv_name}': "
            f"Source table is missing required columns: {missing_list}\n\n"
            f"The materialized view query references columns that no longer exist "
            f"in the source table.\n\n"
            f"Options:\n"
            f"  1. Restore the missing columns to the source table\n"
            f"  2. Drop and recreate the materialized view with an updated query\n"
            f"  3. If the column was renamed, update the source table accordingly"
        )


def run_ray_copy_table(
    dst: TableReference,
    packager: UDFPackager,
    checkpoint_store: CheckpointStore | None = None,
    *,
    job_id: str | None = None,
    concurrency: int = 8,
    intra_applier_concurrency: int = 1,
    batch_size: int | None = None,
    checkpoint_size: int | None = None,
    task_size: int | None = None,
    task_shuffle_diversity: int | None = None,
    commit_granularity: int | None = None,
    delete_batch_size: int | None = None,
    src_version: int | None = None,
    max_rows_per_fragment: int | None = None,
    _admission_check: bool | None = None,
    _admission_strict: bool | None = None,
    **kwargs,
) -> None:
    # prepare job parameters
    base_config = JobConfig.get()
    map_checkpoint_size = resolve_batch_size(
        batch_size=batch_size,
        checkpoint_size=checkpoint_size,
    )
    if map_checkpoint_size is None:
        map_checkpoint_size = base_config.batch_size

    dst_table = dst.open()

    try:
        dst_row_count = dst_table.count_rows()
    except Exception:
        _LOG.warning("Failed to count rows for %s; using fallback task_size", dst)
        dst_row_count = None

    task_size = resolve_task_size(
        task_size=task_size,
        row_count=dst_row_count,
        num_workers=concurrency,
    )

    # A map batch (checkpoint) cannot span more rows than a single read task
    # window. When both sizes are positive, cap the checkpoint size to the read
    # task size so the effective batching is explicit and predictable. We do
    # not cap when `task_size <= 0`, since that value is used to request
    # "one task per fragment" semantics.
    if map_checkpoint_size > 0 and task_size > 0 and map_checkpoint_size > task_size:
        _LOG.debug(
            "Capping checkpoint_size from %s to task_size %s",
            map_checkpoint_size,
            task_size,
        )
        map_checkpoint_size = task_size

    config = base_config.with_overrides(
        batch_size=map_checkpoint_size,
        task_size=task_size,
        task_shuffle_diversity=task_shuffle_diversity,
        commit_granularity=commit_granularity,
        delete_batch_size=delete_batch_size,
    )

    checkpoint_store = checkpoint_store or config.make_checkpoint_store()

    # Create error store from dst connection
    dst_db = dst.open_db()
    error_store = _make_error_store(dst_db)

    # Initialize the JobStateManager to ensure Jobs table is created
    dst_db._history  # noqa: B018

    # Open destination table once and reuse
    dst_schema = dst_table.schema
    if dst_schema.metadata is None:
        raise Exception("Destination dataset must have view metadata.")
    src_dburi = dst_schema.metadata[MATVIEW_META_BASE_DBURI.encode("utf-8")].decode(
        "utf-8"
    )
    src_name = dst_schema.metadata[MATVIEW_META_BASE_TABLE.encode("utf-8")].decode(
        "utf-8"
    )

    # Get namespace path from metadata (source table's namespace path)
    namespace_path_b = dst_schema.metadata.get(MATVIEW_META_NAMESPACE_PATH.encode())
    namespace_path: list[str] | None = None
    if namespace_path_b:
        namespace_path = namespace_path_b.decode().split("$")

    # If src_version not specified, fetch the latest version of the source table
    if src_version is None:
        from geneva.db import connect

        # Use view's namespace connection if available, otherwise fallback to db_uri
        if dst.namespace_impl is not None and dst.namespace_properties is not None:
            src_db = connect(
                namespace_impl=dst.namespace_impl,
                namespace_properties=dst.namespace_properties,
            )
            src_table_tmp = src_db.open_table(src_name, namespace=namespace_path)
        elif src_dburi == dst.db_uri:
            src_db = dst.open_db()
            src_table_tmp = src_db.open_table(src_name)
        else:
            # Source is in a different database, connect to it
            src_db = connect(src_dburi, checkpoint=dst.open_checkpoint_store())
            src_table_tmp = src_db.open_table(src_name)

        src_version = src_table_tmp.version
        _LOG.info(
            f"Materialized view refresh: fetched latest source version={src_version} "
            f"for table={src_name} from db={src_dburi}"
        )

    # Create table_id for source table - include namespace path if present
    src_table_id = (namespace_path or []) + [src_name]

    src = attrs.evolve(
        dst,
        table_id=src_table_id,
        version=src_version,
        db_uri=src_dburi,
        is_system_table=False,
    )

    # Detect scalar UDTF view
    scalar_udtf_spec_json = dst_schema.metadata.get(
        MATVIEW_META_SCALAR_UDTF.encode("utf-8")
    )
    is_scalar_udtf = scalar_udtf_spec_json is not None
    scalar_udtf_obj = None

    if is_scalar_udtf:
        from geneva.packager import ScalarUDTFSpec, unmarshal_scalar_udtf

        spec = ScalarUDTFSpec.from_json(scalar_udtf_spec_json.decode())
        scalar_udtf_obj = unmarshal_scalar_udtf(spec)
        if scalar_udtf_obj is None:
            raise ValueError("Cannot refresh: Failed to deserialize ScalarUDTF.")

    # Load source query — scalar UDTF views use MATVIEW_META_SOURCE_QUERY,
    # regular MVs use MATVIEW_META_QUERY.
    if is_scalar_udtf:
        query_json = dst_schema.metadata[MATVIEW_META_SOURCE_QUERY.encode("utf-8")]
    else:
        query_json = dst_schema.metadata[MATVIEW_META_QUERY.encode("utf-8")]
    query = GenevaQuery.model_validate_json(query_json)

    # Check for point-in-time refresh (rollback to older version)
    from geneva.table import _get_last_refreshed_version

    last_refreshed = _get_last_refreshed_version(dst_table)
    if (
        src_version is not None
        and last_refreshed is not None
        and src_version < last_refreshed
    ):
        # Point-in-time refresh - delete rows not in target version
        _LOG.info(
            f"[EXPERIMENTAL] Point-in-time refresh: rolling back from version "
            f"{last_refreshed} to version {src_version}"
        )
        dst_table, _, _, _ = _delete_stale_mv_rows(
            dst,
            dst_table,
            src_dburi,
            src_name,
            src_version,
            query,
            config.delete_batch_size,
            "Backward refresh",
        )

    src_table = src.open()

    # Validate that source table schema is compatible with MV query
    if not is_scalar_udtf:
        _validate_mv_source_schema(query, src_table, dst.table_name, packager)

    if is_scalar_udtf:
        # For scalar UDTF views, use the destination (view) schema.
        schema = dst_table.schema
    else:
        schema = GenevaQueryBuilder.from_query_object(src_table, query).schema

    job_id = job_id or uuid.uuid4().hex

    # Scalar UDTF views have no column_udfs — the expansion produces all output
    column_udfs = [] if is_scalar_udtf else query.extract_column_udfs(packager)

    # Admission control: validate cluster has sufficient resources
    # For matview refresh with UDFs, check each UDF's resource requirements.
    # Note: column_udfs from extract_column_udfs already has unmarshaled UDFs
    if column_udfs:
        from geneva.runners.ray.admission import validate_admission

        for column_udf in column_udfs:
            validate_admission(
                column_udf.udf,
                concurrency=concurrency,
                intra_applier_concurrency=intra_applier_concurrency,
                check=_admission_check,
                strict=_admission_strict,
            )

    # Build input_cols based on query's select clause
    # For dict selects, pass directly to Lance which handles both:
    # - Column renames: {"my_id": "id"} - Lance evaluates "id" as expression
    # - SQL expressions: {"doubled": "value * 2"} - Lance evaluates the expression
    # Note: No quoting needed - Lance's expression parser treats bare identifiers
    # as column references (e.g., "abs" is column, "abs()" is function call)
    input_cols: list[str] | dict[str, str]
    # Track which source columns are required (for field ID tracking)
    # Use set to avoid duplicates when multiple outputs reference the same source column
    src_cols_required: set[str] = set()
    src_col_names = set(src_table.schema.names)

    if query.base.columns and isinstance(query.base.columns, dict):
        # Dict select - pass through to Lance for expression evaluation
        input_cols_dict: dict[str, str] = {}
        for dst_col_name, dst_col_expr in query.base.columns.items():
            if isinstance(dst_col_expr, str):
                input_cols_dict[dst_col_name] = dst_col_expr
                # Track source columns for field ID lookup
                if dst_col_expr in src_col_names:
                    src_cols_required.add(dst_col_expr)

        # Add columns needed for UDFs that aren't already in input_cols
        for udf in column_udfs:
            if udf.udf.input_columns:
                for src_col_name in udf.udf.input_columns:
                    if src_col_name not in input_cols_dict:
                        input_cols_dict[src_col_name] = src_col_name
                        if src_col_name in src_col_names:
                            src_cols_required.add(src_col_name)

        input_cols = input_cols_dict
    elif query.base.columns and isinstance(query.base.columns, list):
        # Simple column list
        input_cols = list(query.base.columns)
        src_cols_required = {c for c in input_cols if c in src_col_names}
        # Also include UDF input columns that may not be in the select list
        for udf_transform in column_udfs:
            if udf_transform.udf.input_columns:
                for src_col_name in udf_transform.udf.input_columns:
                    if src_col_name not in input_cols:
                        input_cols.append(src_col_name)
                        if src_col_name in src_col_names:
                            src_cols_required.add(src_col_name)
    else:
        # Fallback: no explicit select
        input_cols = [
            n
            for n in src_table.schema.names
            if n not in ["__is_set", "__source_row_id"]
        ]
        src_cols_required = set(input_cols)
        # Use set for O(1) membership checks
        input_cols_set = set(input_cols)
        # Also include UDF input columns that may not be in the select list
        for udf_transform in column_udfs:
            if udf_transform.udf.input_columns:
                for src_col_name in udf_transform.udf.input_columns:
                    if src_col_name not in input_cols_set:
                        input_cols.append(src_col_name)
                        input_cols_set.add(src_col_name)
                        if src_col_name in src_col_names:
                            src_cols_required.add(src_col_name)

    # Create map_task before plan_copy so we can use it for checkpoint checking
    map_task = CopyTableTask(
        column_udfs=column_udfs,
        view_name=dst.table_name,
        schema=schema,
        override_batch_size=config.batch_size,
    )

    # === PASS 1: Identify new source fragments and add placeholder rows ===
    # This ensures CopyTask can read __source_row_id from destination during refresh
    _LOG.info("=== PASS 1: Identifying new source fragments ===")

    src_dataset = src_table.to_lance()
    _LOG.info(
        f"Source table has {len(list(src_dataset.get_fragments()))} fragments "
        f"at version={src.version} (Lance version={src_dataset.version})"
    )

    # Compute field IDs for columns the MV reads - used to filter data file tracking
    # so that changes to unrelated columns don't trigger MV refresh.
    # Use src_cols_required (actual col names), not input_cols (may be dict)
    relevant_field_ids = _get_relevant_field_ids(src_dataset, src_cols_required)
    _LOG.debug(
        f"Relevant field IDs for MV columns {src_cols_required}: {relevant_field_ids}"
    )

    dst_dataset = dst_table.to_lance()

    # Build destination-to-source mapping, identify checkpointed fragments,
    # and collect existing source row IDs (all in one pass for efficiency)
    # Also validates source data files to detect backfill changes
    (
        dst_to_src_map,
        dst_frags_with_checkpoint,
        existing_source_row_ids,
        src_data_files_by_dst,
    ) = _build_dst_to_src_mapping(
        dst_dataset,
        dst_dataset.uri,
        map_task,
        checkpoint_store,
        src_dataset,
        relevant_field_ids,
    )

    # === HANDLE SOURCE DELETIONS (FORWARD REFRESH) ===
    # When doing a forward refresh, check if any source rows have been deleted
    # and remove them from the MV
    if (
        last_refreshed is not None
        and src_version is not None
        and src_version >= last_refreshed  # Forward or same version refresh
        and existing_source_row_ids  # MV has rows
    ):
        dst_table, dst_dataset, deleted_count, valid_row_ids = _delete_stale_mv_rows(
            dst,
            dst_table,
            src_dburi,
            src_name,
            src_version,
            query,
            config.delete_batch_size,
            "Forward refresh",
            existing_source_row_ids,
        )
        if deleted_count > 0:
            existing_source_row_ids = existing_source_row_ids & valid_row_ids

    # Identify truly NEW source fragments (not yet in destination at all)
    new_source_fragments = _identify_new_source_fragments(src_dataset, dst_to_src_map)

    if is_scalar_udtf:
        # Scalar UDTF: expand new source rows via the UDTF and write to dest.
        # The expansion produces fully populated rows (no Pass 2 needed).
        assert scalar_udtf_obj is not None
        new_dst_frag_ids = _append_expanded_fragments(
            src_table=src_table,
            dst_table=dst_table,
            scalar_udtf_obj=scalar_udtf_obj,
            source_query=query,
            new_source_fragments=new_source_fragments,
            existing_source_row_ids=existing_source_row_ids,
            max_rows_per_fragment=max_rows_per_fragment,
            concurrency=concurrency,
            dst_ref=dst,
            job_id=job_id,
        )

        if new_dst_frag_ids:
            dst_table.checkout_latest()

        from geneva.table import _set_last_refreshed_version

        if src_version is not None:
            _set_last_refreshed_version(dst_table, src_version)

        _LOG.info(
            f"Scalar UDTF refresh complete for {dst.table_name}: "
            f"source_version={src_version}"
        )
        return

    # --- Standard 1:1 MV path below ---

    # Extract row IDs from new source fragments and add placeholder rows
    new_fragment_row_ids = []
    for frag_id in new_source_fragments:
        new_row_ids = _extract_new_row_ids_from_source_fragment(
            src_table, query, frag_id, existing_source_row_ids
        )
        new_fragment_row_ids.extend(new_row_ids)
        existing_source_row_ids.update(new_row_ids)

    # Append placeholder rows for new data
    new_dst_frag_ids = _append_placeholder_fragments(
        dst_table, new_fragment_row_ids, max_rows_per_fragment
    )

    # After appending placeholder fragments, update dst reference to use latest version.
    # This is needed because _append_placeholder_fragments creates a new table version.
    if new_dst_frag_ids:
        dst_table.checkout_latest()
        dst = attrs.evolve(dst, version=dst_table.version)
        _LOG.info(f"Updated destination reference to version {dst.version}")

        # Compute source data files for newly created destination fragments.
        # This is needed so that checkpoints store data file paths for new fragments,
        # enabling detection of source backfill changes on subsequent refreshes.
        dst_dataset = dst_table.to_lance()
        mv_version = _get_matview_version(dst_dataset.schema)
        for new_dst_frag_id in new_dst_frag_ids:
            new_frag = dst_dataset.get_fragment(new_dst_frag_id)
            if new_frag is None:
                continue
            try:
                scanner = dst_dataset.scanner(
                    columns=["__source_row_id"],
                    fragments=[new_frag],
                )
                frag_data = scanner.to_table()
                if len(frag_data) > 0:
                    source_row_ids = cast(
                        "list[int]", frag_data["__source_row_id"].to_pylist()
                    )
                    src_frag_ids = _extract_fragment_ids_from_row_ids(
                        [rid for rid in source_row_ids if rid is not None],
                        mv_version,
                    )
                    data_files = get_combined_source_data_files(
                        src_dataset, src_frag_ids, relevant_field_ids
                    )
                    src_data_files_by_dst[new_dst_frag_id] = data_files
                    _LOG.info(
                        f"Computed data files for new dst fragment "
                        f"{new_dst_frag_id}: source_frags={src_frag_ids}, "
                        f"file_count={len(data_files)}"
                    )
            except Exception as e:
                _LOG.warning(
                    f"Failed to compute data files for new dst fragment "
                    f"{new_dst_frag_id}: {e}"
                )

    # Collect checkpointed fragments for commit inclusion
    skipped_fragments, skipped_stats = _collect_skipped_fragments(
        dst_dataset,
        dst_frags_with_checkpoint,
        map_task,
        checkpoint_store,
        src_data_files_by_dst=src_data_files_by_dst,
    )

    # Determine which destination fragments need processing (destination-driven)
    # Simply process all destination fragments without checkpoints
    fragments_to_process = _determine_fragments_to_process(
        dst_to_src_map, dst_frags_with_checkpoint, new_dst_frag_ids
    )

    # Create plan - process only fragments that need processing
    # (or all if none specified)
    plan = plan_copy(
        src,
        dst,
        input_cols,
        task_size=task_size,
        task_shuffle_diversity=config.task_shuffle_diversity,
        only_fragment_ids=fragments_to_process,  # None means process all fragments
        src_data_files_by_dst=src_data_files_by_dst,
    )

    if fragments_to_process is not None:
        _LOG.info(f"Planning to process destination fragments: {fragments_to_process}")
    else:
        _LOG.info("Planning to process ALL destination fragments")

    _run_column_adding_pipeline(
        map_task,
        checkpoint_store,
        error_store,
        config,
        dst,
        plan,
        job_id,
        concurrency,
        skipped_fragments=skipped_fragments,
        skipped_stats=skipped_stats,
        src_data_files_by_dst=src_data_files_by_dst,
        **kwargs,
    )

    # Update materialized view metadata to track the source version
    # that was just refreshed
    _LOG.info(
        f"Updating materialized view {dst.table_name} metadata: "
        f"base_table_version={src_version}"
    )
    # Note: Schema metadata updates would require re-creating the table or using
    # Lance operations. For now, we log the version that was refreshed.
    # In future iterations, we could store this in a separate metadata table or
    # use schema evolution to track refresh history.


def dispatch_run_ray_add_column(
    table_ref: TableReference,
    col_name: str,
    *,
    read_version: int | None = None,
    concurrency: int = 8,
    batch_size: int | None = None,
    checkpoint_size: int | None = None,
    min_checkpoint_size: int | None = None,
    max_checkpoint_size: int | None = None,
    task_size: int | None = None,
    task_shuffle_diversity: int | None = None,
    commit_granularity: int | None = None,
    batch_checkpoint_flush_interval_seconds: float = (
        DEFAULT_BATCH_CHECKPOINT_FLUSH_INTERVAL_SECONDS
    ),
    where: str | None = None,
    enable_job_tracker_saves: bool = True,
    **kwargs,
) -> JobFuture:
    """
    Dispatch the Ray add column operation to a remote function.
    This is a convenience function to allow calling the remote function directly.
    """
    if "task_size" in kwargs and task_size is None:
        task_size = kwargs.pop("task_size")
    else:
        kwargs.pop("task_size", None)

    checkpoint_size = resolve_batch_size(
        batch_size=batch_size,
        checkpoint_size=checkpoint_size,
    )
    batch_size = None

    from datetime import datetime

    from geneva._context import get_current_context
    from geneva.manifest.mgr import ManifestConfigManager
    from geneva.utils import current_user

    db = table_ref.open_db()
    hist = db._history

    # Extract manifest info from current context
    manifest_id = None
    manifest_checksum = None
    ctx = get_current_context()
    if ctx is not None and ctx.manifest is not None:
        manifest = ctx.manifest
        manifest_checksum = manifest.checksum

        # Auto-register manifest if not already persisted
        if not manifest.name:
            # Generate auto name
            user = current_user()
            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
            checksum_short = manifest.checksum[:8] if manifest.checksum else "unknown"
            manifest.name = f"auto-{user}-{timestamp}-{checksum_short}"
            _LOG.info(f"Auto-registering manifest as '{manifest.name}'")

        manifest_id = manifest.name

        # Ensure manifest is registered in the database
        manifest_mgr = ManifestConfigManager(db)
        existing = manifest_mgr.load(manifest.name)
        if existing is None:
            _LOG.info(f"Upserting manifest '{manifest.name}' to database")
            manifest_mgr.upsert(manifest)

    job = hist.launch(
        table_ref.table_name,
        col_name,
        where=where,
        manifest_id=manifest_id,
        manifest_checksum=manifest_checksum,
        **kwargs,
    )

    job_tracker = job_tracker_options(
        name=f"jobtracker-{job.job_id}",
        num_cpus=JOBTRACKER_NUM_CPUS,
        memory=JOBTRACKER_MEMORY,
        max_restarts=-1,
    ).remote(job.job_id, table_ref, enable_saves=enable_job_tracker_saves)  # type: ignore[call-arg]

    # Run on small cpu so that it's always easy to schedule.
    remote_runner = cast(
        "Any",
        run_ray_add_column_remote.options(num_cpus=DRIVER_NUM_CPUS).remote,
    )
    obj_ref = remote_runner(
        table_ref,
        col_name,
        read_version=read_version,  # type: ignore[call-arg]
        job_id=job.job_id,  # type: ignore[call-arg]
        job_tracker=job_tracker,  # type: ignore[call-arg]
        concurrency=concurrency,  # type: ignore[call-arg]
        checkpoint_size=checkpoint_size,  # type: ignore[call-arg]
        min_checkpoint_size=min_checkpoint_size,  # type: ignore[call-arg]
        max_checkpoint_size=max_checkpoint_size,  # type: ignore[call-arg]
        task_size=task_size,  # type: ignore[call-arg]
        task_shuffle_diversity=task_shuffle_diversity,  # type: ignore[call-arg]
        commit_granularity=commit_granularity,  # type: ignore[call-arg]
        batch_checkpoint_flush_interval_seconds=(
            batch_checkpoint_flush_interval_seconds
        ),  # type: ignore[call-arg]
        where=where,  # type: ignore[call-arg]
        **kwargs,
    )
    # object ref is only available here
    hist.set_object_ref(job.job_id, cloudpickle.dumps(obj_ref))

    # Register with the current context so _wait_for_tracked_jobs can await it
    from geneva.runners.ray.raycluster import RayCluster

    ctx = get_current_context()
    if isinstance(ctx, RayCluster):
        ctx.register_tracked_job(job.job_id, obj_ref, job_tracker)

    return RayJobFuture(
        job_id=job.job_id,
        ray_obj_ref=obj_ref,
        job_tracker=job_tracker,  # type: ignore[arg-type]
    )


def _schema_for_validation(tbl: Table, read_version: int | None) -> pa.Schema:
    """
    Get the table schema to validate UDF inputs against.

    If `read_version` is provided, we validate against that historical version of the
    table. Otherwise, we use the current schema.
    """
    if read_version is None:
        return tbl._native_tbl.schema

    try:
        dataset = tbl.to_lance().checkout_version(read_version)
        return dataset.schema
    except Exception as exc:  # noqa: BLE001
        raise ValueError(
            f"Failed to load schema for read_version={read_version}. "
            "Ensure the version exists before running backfill."
        ) from exc


def _resolve_read_version(tbl: Table, read_version: int | None) -> int:
    """
    Determine the concrete version a job should read.

    If the caller did not provide read_version, we snapshot the table's current
    version at job launch time. This pins the job to a consistent schema and guards
    against concurrent column drops during execution.
    """
    return read_version if read_version is not None else tbl.version


def validate_backfill_args(
    tbl: Table,
    col_name: str,
    udf: UDF | None = None,
    input_columns: list[str] | None = None,
    *,
    read_version: int | None = None,
) -> None:
    """
    Validate the arguments for the backfill operation.

    This function performs validation before starting a backfill job to catch
    configuration errors early. It validates:

    1. Target column exists in the table
    2. UDF input columns exist in the table schema (for the requested read_version)
    3. Column types are compatible with UDF type annotations (if present)

    All validation is delegated to UDF.validate_against_schema() for consistency.

    Parameters
    ----------
    tbl : Table
        The table to backfill
    col_name : str
        The column name to backfill
    udf : UDF | None
        The UDF to use (if None, will be loaded from column metadata)
    input_columns : list[str] | None
        The input columns for the UDF (if None, will be loaded from column metadata)
    read_version : int | None
        If provided, validate UDF inputs against this historical table version instead
        of the current schema. This catches jobs that try to read from a version that
        lacks required input columns.

    Raises
    ------
    ValueError
        If validation fails (missing columns, type mismatches, etc.)

    Warns
    -----
    UserWarning
        If type validation is skipped due to missing type annotations
    """
    current_schema = tbl._native_tbl.schema
    if col_name not in current_schema.names:
        raise ValueError(
            f"Column {col_name} is not defined this table.  "
            "Use add_columns to register it first"
        )

    if udf is None:
        udf_spec = fetch_udf(tbl, col_name)
        udf = tbl._conn._packager.unmarshal(udf_spec)

    # Get input_columns from column metadata if not provided
    if input_columns is None:
        field = tbl._native_tbl.schema.field(col_name)
        metadata = field.metadata or {}
        input_columns = json.loads(metadata.get(b"virtual_column.udf_inputs", "null"))

    # Delegate to UDF's consolidated validation method if UDF was successfully
    # unmarshaled. If unmarshal returned None (modules not available), skip validation
    if udf is not None:
        schema = _schema_for_validation(tbl, read_version)
        try:
            udf.validate_against_schema(schema, input_columns)
        except ValueError as exc:
            if read_version is not None:
                raise ValueError(
                    f"{exc} (while validating against read_version {read_version})"
                ) from exc
            raise


@ray.remote
def run_ray_add_column_remote(
    table_ref: TableReference,
    col_name: str,
    *,
    job_id: str | None = None,
    udf: UDF | None = None,
    read_version: int | None = None,
    concurrency: int = 8,
    checkpoint_size: int | None = None,
    min_checkpoint_size: int | None = None,
    max_checkpoint_size: int | None = None,
    task_size: int | None = None,
    task_shuffle_diversity: int | None = None,
    commit_granularity: int | None = None,
    batch_checkpoint_flush_interval_seconds: float = (
        DEFAULT_BATCH_CHECKPOINT_FLUSH_INTERVAL_SECONDS
    ),
    where: str | None = None,
    job_tracker: ActorHandle | None = None,
    **kwargs,
) -> None:
    """
    Remote function to run the Ray add column operation.
    This is a wrapper around `run_ray_add_column` to allow it to be called as a Ray
    task.
    """
    legacy_batch_size = kwargs.pop("batch_size", None)
    checkpoint_size = resolve_batch_size(
        batch_size=legacy_batch_size,
        checkpoint_size=checkpoint_size,
    )

    import geneva  # noqa: F401  Force so that we have the same env in next level down

    hist = None
    try:
        tbl = table_ref.open()
        read_version = _resolve_read_version(tbl, read_version)
        hist = tbl._conn._history
        if job_id:
            hist.set_running(job_id)

        # Get input_columns from column metadata first
        field = tbl._native_tbl.schema.field(col_name)
        metadata = field.metadata or {}
        input_columns = json.loads(metadata.get(b"virtual_column.udf_inputs", "null"))

        validate_backfill_args(
            tbl, col_name, udf, input_columns, read_version=read_version
        )
        if udf is None:
            udf_spec = fetch_udf(tbl, col_name)
            udf = tbl._conn._packager.unmarshal(udf_spec)

            # If unmarshal still returns None, we cannot proceed
            if udf is None:
                raise RuntimeError(
                    f"Failed to unmarshal UDF for column '{col_name}'. "
                    "The UDF modules may not be available in the Ray worker "
                    "environment. Ensure the manifest's py_modules includes all "
                    "required dependencies."
                )

        # Apply input_columns override to the UDF if needed
        # This handles the case where add_columns was called with explicit column
        # mapping e.g., table.add_columns({"col": (udf, ["seq"])})
        if input_columns is not None and udf.input_columns != input_columns:
            udf.input_columns = input_columns

        from geneva.runners.ray.pipeline import run_ray_add_column

        # Use table-specific checkpoint store
        checkpoint_store = table_ref.open_checkpoint_store()
        run_ray_add_column(
            table_ref,
            input_columns,
            {col_name: udf},
            checkpoint_store=checkpoint_store,
            read_version=read_version,
            job_id=job_id,
            concurrency=concurrency,
            checkpoint_size=checkpoint_size,
            min_checkpoint_size=min_checkpoint_size,
            max_checkpoint_size=max_checkpoint_size,
            task_size=task_size,
            task_shuffle_diversity=task_shuffle_diversity,
            commit_granularity=commit_granularity,
            batch_checkpoint_flush_interval_seconds=(
                batch_checkpoint_flush_interval_seconds
            ),
            where=where,
            job_tracker=job_tracker,
            **kwargs,
        )
        if job_id:
            hist.set_completed(job_id)
    except Exception as e:
        _LOG.exception("Error running Ray add column operation")
        if job_id and hist is not None:
            with contextlib.suppress(Exception):
                hist.set_failed(job_id, f"{type(e).__name__}: {e}")
        raise _picklable_remote_error(e) from None
    finally:
        # Signal the JobTracker that this job is finished so that
        # _wait_for_tracked_jobs can detect completion.
        if job_tracker is not None:
            with contextlib.suppress(Exception):
                job_tracker.mark_job_done.remote()


def _materialize_read_columns(columns: list[str]) -> list[str]:
    """Return the minimal set of columns to read, preserving dotted paths.

    LanceDB can project struct subfields directly (e.g., "info.left"), so we keep
    dotted paths instead of expanding to the parent struct column. We only de-dupe
    exact column strings while preserving order.
    """

    seen = set()
    physical: list[str] = []
    for col in columns:
        if col in seen:
            continue
        seen.add(col)
        physical.append(col)
    return physical


def _picklable_remote_error(exc: Exception) -> RuntimeError:
    """Convert worker exceptions into a Ray-picklable error payload."""
    return RuntimeError(f"{type(exc).__module__}.{type(exc).__name__}: {exc}")


def run_ray_add_column(
    table_ref: TableReference,
    columns: list[str] | None,
    transforms: dict[str, UDF],
    checkpoint_store: CheckpointStore | None = None,
    *,
    read_version: int | None = None,
    job_id: str | None = None,
    concurrency: int = 8,
    batch_size: int | None = None,
    checkpoint_size: int | None = None,
    min_checkpoint_size: int | None = None,
    max_checkpoint_size: int | None = None,
    task_size: int | None = None,
    task_shuffle_diversity: int | None = None,
    commit_granularity: int | None = None,
    batch_checkpoint_flush_interval_seconds: float = (
        DEFAULT_BATCH_CHECKPOINT_FLUSH_INTERVAL_SECONDS
    ),
    where: str | None = None,
    job_tracker=None,
    **kwargs,
) -> None:
    # prepare job parameters
    base_config = JobConfig.get()

    map_checkpoint_size = resolve_batch_size(
        batch_size=batch_size,
        checkpoint_size=checkpoint_size,
    )
    explicit_checkpoint_size = checkpoint_size is not None or batch_size is not None
    if map_checkpoint_size is None:
        map_checkpoint_size = base_config.batch_size

    # Open table early for row-count based defaults
    table = table_ref.open()
    read_version = _resolve_read_version(table, read_version)

    # UDF-level task_size can act as a hint when no job override is provided
    if task_size is None:
        udf_task_sizes = [
            udf.task_size for udf in transforms.values() if udf.task_size is not None
        ]
        if udf_task_sizes:
            task_size = min(int(ts) for ts in udf_task_sizes)

    try:
        row_count = table.count_rows()
    except Exception:
        _LOG.warning("Failed to count rows for %s; using fallback task_size", table)
        row_count = None

    task_size = resolve_task_size(
        task_size=task_size,
        row_count=row_count,
        num_workers=concurrency,
    )

    # See run_ray_copy_table for rationale on capping checkpoint size.
    if map_checkpoint_size > 0 and task_size > 0 and map_checkpoint_size > task_size:
        _LOG.debug(
            "Capping checkpoint_size from %s to task_size %s",
            map_checkpoint_size,
            task_size,
        )
        map_checkpoint_size = task_size

    config = base_config.with_overrides(
        batch_size=map_checkpoint_size,
        task_size=task_size,
        task_shuffle_diversity=task_shuffle_diversity,
        commit_granularity=commit_granularity,
    )

    checkpoint_store = checkpoint_store or config.make_checkpoint_store()

    # Create error store from table connection
    db = table_ref.open_db()
    error_store = _make_error_store(db)
    uri = table.uri

    # Re‑validate UDF inputs against the schema we will actually read.
    # This narrows the window for concurrent schema changes that could otherwise
    # surface as late, opaque runtime errors.
    schema_for_validation = _schema_for_validation(table, read_version)
    for udf in transforms.values():
        cols_for_udf = columns if columns is not None else udf.input_columns
        udf.validate_against_schema(schema_for_validation, cols_for_udf)

    # add pre-existing col if carrying previous values forward
    carry_forward_cols = list(
        set(transforms.keys()) & set(table._native_tbl.schema.names)
    )
    _LOG.debug(f"carry_forward_cols {carry_forward_cols}")
    # this copy is necessary because the array extending updates inplace and this
    # columns array is directly referenced by the udf instance earlier
    cols = table._native_tbl.schema.names.copy() if columns is None else columns.copy()
    cols = _materialize_read_columns(cols)

    for cfcol in carry_forward_cols:
        # only append if cf col is not in col list already
        if cfcol not in cols:
            cols.append(cfcol)

    # Respect backfill's batch_size override by passing it into the task.
    map_task = BackfillUDFTask(
        udfs=transforms,
        where=where,
        override_batch_size=config.batch_size,
        explicit_checkpoint_size=explicit_checkpoint_size,
        min_checkpoint_size=min_checkpoint_size,
        max_checkpoint_size=max_checkpoint_size,
    )

    plan_start = time.perf_counter()
    plan, pipeline_args = plan_read(
        uri,
        table_ref,
        cols,
        task_size=task_size,
        read_version=read_version,
        task_shuffle_diversity=config.task_shuffle_diversity,
        where=where,
        map_task=map_task,
        checkpoint_store=checkpoint_store,
        **kwargs,
    )
    if job_tracker is not None:
        plan_elapsed_ms = int((time.perf_counter() - plan_start) * 1000)
        if plan_elapsed_ms > 0:
            job_tracker.increment.remote(METRIC_PLAN_READ_TIME, plan_elapsed_ms)

    _LOG.info(
        f"starting backfill pipeline for {transforms} where='{where}'"
        f" with carry_forward_cols={carry_forward_cols}"
    )
    _run_column_adding_pipeline(
        map_task,
        checkpoint_store,
        error_store,
        config,
        table_ref,
        plan,
        job_id,
        concurrency,
        where=where,
        job_tracker=job_tracker,
        batch_checkpoint_flush_interval_seconds=(
            batch_checkpoint_flush_interval_seconds
        ),
        **pipeline_args,
    )


def _make_error_store(db: Connection) -> ErrorStore:
    return ErrorStore(db)


@attrs.define
class RayJobFuture(JobFuture):
    ray_obj_ref: ActorHandle = attrs.field()
    job_tracker: ActorHandle | None = attrs.field(default=None)
    _pbars: dict[str, TqdmType] = attrs.field(factory=dict)
    _RAY_LINE_KEY: str = "_ray_summary_line"

    def _sync_bars(self, snapshot: dict[str, dict]) -> None:
        # single line ray summary
        wa = snapshot.get(CNT_WORKERS_ACTIVE)
        wp = snapshot.get(CNT_WORKERS_PENDING)

        if wa or wp:
            bar = self._pbars.get(self._RAY_LINE_KEY)
            if bar is None:
                # text-only line, like k8s/kr bars
                bar = tqdm(total=0, bar_format="{desc} {bar:0}[{elapsed}]")
                self._pbars[self._RAY_LINE_KEY] = bar

            active_count = wa.get("n", 0) if wa else 0
            pending_count = wp.get("n", 0) if wp else 0
            label = fmt(
                "geneva | workers (active/pending): ", Colors.BRIGHT_MAGENTA, bold=True
            )
            bar.desc = (
                f"{label}({fmt_numeric(active_count)}/{fmt_pending(pending_count)})"
            )
            bar.refresh()

            # close when all are done (harmless if left open)
            if all(m and m.get("done") for m in (wa, wp)):
                bar.close()

        for name, m in snapshot.items():
            if name in {
                CNT_RAY_NODES,
                CNT_WORKERS_PENDING,
                CNT_WORKERS_ACTIVE,
            }:
                continue

            n, total, done, desc = m["n"], m["total"], m["done"], m.get("desc", name)
            bar = self._pbars.get(name)
            if bar is None:
                # Only make bars for the known core metrics (skip "fragments",
                # "writer_fragments", and other randoms)
                if name not in {
                    "rows_checkpointed",
                    "rows_ready_for_commit",
                    "rows_committed",
                }:
                    continue
                bar = tqdm(total=total, desc=fmt(desc, Colors.CYAN, bold=True))
                self._pbars[name] = bar
            bar.total = total
            bar.n = n
            bar.refresh()
            if done:
                bar.close()

    def status(self, timeout: float | None = 0.05) -> None:
        if self.job_tracker is None:
            return
        try:
            snapshot = ray.get(self.job_tracker.get_all.remote(), timeout=timeout)  # type: ignore[call-arg,arg-type]
            self._sync_bars(snapshot)  # type: ignore[arg-type]

        except ray.exceptions.GetTimeoutError:
            _LOG.debug("JobTracker not ready? skip this tick")
            return

    def done(self, timeout: float | None = None) -> bool:
        self.status()
        _LOG.debug("Waiting for Ray job %s to complete", self.ray_obj_ref)
        ready, _ = ray.wait([self.ray_obj_ref], timeout=timeout)
        done = bool(ready)

        _LOG.debug(f"Ray jobs ready to complete: {ready}")

        if done:
            # force final update of progress bars
            self.status(timeout=None)
            _LOG.debug(f"RayJobFuture complete. {done=} {ready=} {self._pbars=}")

        return done

    def result(self, timeout: float | None = None) -> Any:
        # TODO this can throw a ray.exceptions.GetTimeoutError if the task
        # does not complete in time, we should create a new exception type to
        # encapsulate Ray specifics
        self.status()
        return ray.get(self.ray_obj_ref, timeout=timeout)  # type: ignore[call-overload,arg-type]


def get_imported_packages() -> dict:
    import importlib
    import sys

    packages = {}
    for name, module in list(sys.modules.items()):
        if module is None:
            continue
        try:
            mod = importlib.import_module(name.split(".")[0])
            version = getattr(mod, "__version__", None)
            if version:
                packages[mod.__name__] = version
        except Exception as e:
            _LOG.error(e)
            continue
    return packages


@ray.remote
def get_imported() -> dict:
    """A simple utility to return the names and versions of python packages
    installed in the current Ray worker environment."""
    return get_imported_packages()
