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

import contextlib
import hashlib
import json
import logging
import platform
from collections.abc import Iterable, Iterator
from datetime import timedelta
from functools import cached_property
from typing import Any, Literal

import attrs
import lance
import lancedb
import numpy as np
import pyarrow as pa
from lancedb import AsyncConnection, connect_async
from lancedb._lancedb import MergeResult
from lancedb.common import DATA, VECTOR_COLUMN_NAME, Credential
from lancedb.index import IndexConfig
from lancedb.merge import LanceMergeInsertBuilder
from lancedb.namespace import AsyncLanceNamespaceDBConnection
from lancedb.query import LanceQueryBuilder, LanceTakeQueryBuilder
from lancedb.query import Query as LanceQuery
from lancedb.table import IndexStatistics, TableStatistics, Tags
from lancedb.table import LanceTable as LanceLocalTable
from lancedb.table import Table as LanceTable
from lancedb.types import OnBadVectorsType

# Python 3.10 compatibility
from typing_extensions import TYPE_CHECKING, Never, Optional, override  # noqa: UP035
from yarl import URL

from geneva._context import get_current_context
from geneva.checkpoint import (
    CheckpointStore,
    InMemoryCheckpointStore,
)
from geneva.db import (
    Connection,
    _get_db_uri,
    connect,
    infer_system_db_uri,
    infer_system_namespace,
    resolve_table_physical_uri,
    uses_namespace_for_system_tables,
)
from geneva.partitioning import _VECTOR_INDEX_TYPE_URL
from geneva.query import (
    MATVIEW_META_BASE_DBURI,
    MATVIEW_META_BASE_TABLE,
    MATVIEW_META_BASE_VERSION,
    MATVIEW_META_NAMESPACE_PATH,
    MATVIEW_META_SOURCE_QUERY,
    MATVIEW_META_UDTF,
    MATVIEW_META_VERSION,
    MATVIEW_VERSION_SCALAR_UDTF,
    GenevaQueryBuilder,
)
from geneva.transformer import UDF, UDFArgType
from geneva.utils import status_updates
from geneva.utils.batch_size import resolve_batch_size

if TYPE_CHECKING:
    from lance_namespace import LanceNamespace

_LOG = logging.getLogger(__name__)

_IVF_PARTITION_COL = "_ivf_partition"

_IMPLICIT_LOCAL_RAY_WARNING = (
    "You are implicitly using a local Ray cluster for %s. "
    "If you want this to continue to work and ensure proper cleanup, "
    "wrap your call in a context: conn.local_ray_context()."
)
# Metadata key for tracking the last successfully refreshed source version
MATVIEW_LAST_REFRESHED_VERSION = "geneva::mv::last_refreshed_version"


def _get_last_refreshed_version(table: "Table") -> int | None:
    """Read last refreshed version from __source_row_id column metadata.

    Returns None if the table doesn't have __source_row_id column or
    if the metadata key doesn't exist (backwards compatibility).
    """
    schema = table.schema
    if "__source_row_id" not in schema.names:
        return None
    field = schema.field("__source_row_id")
    if field.metadata is None:
        return None
    version_bytes = field.metadata.get(MATVIEW_LAST_REFRESHED_VERSION.encode())
    if version_bytes is None:
        return None
    return int(version_bytes.decode())


def _set_last_refreshed_version(table: "Table", version: int) -> None:
    """Write last refreshed version to __source_row_id column metadata."""
    field = table.schema.field("__source_row_id")
    # Convert existing bytes metadata to string dict for update
    existing: dict[str, str] = {}
    if field.metadata:
        for k, v in field.metadata.items():
            key = k.decode() if isinstance(k, bytes) else k
            val = v.decode() if isinstance(v, bytes) else v
            existing[key] = val
    existing[MATVIEW_LAST_REFRESHED_VERSION] = str(version)
    table._native_tbl.replace_field_metadata("__source_row_id", existing)  # type: ignore[attr-defined]


def _plain_api_key(api_key: Any) -> Any:
    """Convert masked Credential wrappers back to their raw string value."""
    if isinstance(api_key, Credential):
        return api_key[:]
    return api_key


def _make_udtf_processor_actor() -> Any:
    """Lazily define the UDTFProcessorActor Ray actor class.

    We wrap the definition in a factory to defer the ``import ray`` until
    Ray is actually needed (table.py is imported even in non-Ray contexts).
    """
    import attrs
    import ray

    @ray.remote  # type: ignore[misc]
    @attrs.define
    class UDTFProcessorActor:  # pyright: ignore[reportRedeclaration]
        """Ray actor that processes a single UDTF partition.

        Expensive initialisation (deserialise the UDTF, open source connection,
        open checkpoint store) happens once in ``__ray_ready__`` so it is
        amortised across all partitions dispatched to this actor.
        """

        udtf_pickle_bytes: bytes
        source_uri: str
        source_name: str
        source_query_json_bytes: bytes | None
        ckp_store_uri: str
        source_api_key: Any = attrs.field(default=None, repr=False)
        source_host_override: str | None = attrs.field(default=None)
        error_handling_pickle_bytes: bytes | None = attrs.field(default=None)
        job_id: str = attrs.field(default="")
        dest_table_uri: str = attrs.field(default="")
        dest_table_name: str = attrs.field(default="")
        source_version: int | None = attrs.field(default=None)
        # Namespace fields for source table connection
        namespace_impl: str | None = attrs.field(default=None)
        namespace_properties_json: str | None = attrs.field(default=None)
        namespace_path: list[str] | None = attrs.field(default=None)

        # Initialised lazily on first process_partition call
        _udtf: Any = attrs.field(init=False, default=None, repr=False)
        _conn: Any = attrs.field(init=False, default=None, repr=False)
        _tbl: Any = attrs.field(init=False, default=None, repr=False)
        _store: Any = attrs.field(init=False, default=None, repr=False)
        _error_config: Any = attrs.field(init=False, default=None, repr=False)
        _initialized: bool = attrs.field(init=False, default=False, repr=False)

        def __ray_ready__(self) -> None:
            self._ensure_initialized()

        def _ensure_initialized(self) -> None:
            if self._initialized:
                return
            import json as _json

            import geneva.checkpoint as _ckp
            import geneva.cloudpickle as _cp
            from geneva.db import connect as _connect

            self._udtf = _cp.loads(self.udtf_pickle_bytes)

            # Use namespace connection if available, otherwise fallback to db_uri
            if self.namespace_impl is not None and self.namespace_properties_json:
                namespace_props = _json.loads(self.namespace_properties_json)
                self._conn = _connect(
                    namespace_impl=self.namespace_impl,
                    namespace_properties=namespace_props,
                )
                self._tbl = self._conn.open_table(
                    self.source_name, namespace=self.namespace_path
                )
            else:
                self._conn = _connect(
                    self.source_uri,
                    api_key=self.source_api_key,
                    host_override=self.source_host_override,
                )
                self._tbl = self._conn.open_table(self.source_name)

            if self.source_version is not None:
                self._tbl.checkout(self.source_version)
            self._store = _ckp.CheckpointStore.from_uri(self.ckp_store_uri)
            if self.error_handling_pickle_bytes is not None:
                self._error_config = _cp.loads(self.error_handling_pickle_bytes)
            self._initialized = True

        def __repr__(self) -> str:
            try:
                return (
                    f"UDTFProcessorActor(source={self.source_name!r}, "
                    f"store={self.ckp_store_uri!r})"
                )
            except Exception:
                return "<UDTFProcessorActor (repr failed)>"

        def _get_existing_batch_indices(self, partition_prefix: str) -> set[int]:
            """Return batch indices already checkpointed for this partition."""
            import geneva.checkpoint_utils as _ckp_utils

            return {
                idx
                for k in self._store.list_keys(partition_prefix)
                if (idx := _ckp_utils.parse_udtf_batch_index(k)) is not None
            }

        def _build_source_query(self, partition_filter: str | None) -> Any:
            """Build the source query, optionally scoped to a partition."""
            from geneva.query import GenevaQuery, GenevaQueryBuilder

            if self.source_query_json_bytes is not None:
                _query = GenevaQuery.model_validate_json(
                    self.source_query_json_bytes.decode()
                )
                qb = GenevaQueryBuilder.from_query_object(self._tbl, _query)
            else:
                qb = self._tbl.search(None)

            if partition_filter is not None:
                qb = qb.where(partition_filter)
            return qb

        def _execute_with_checkpoints(
            self,
            qb: Any,
            partition_prefix: str,
            existing_batch_indices: set[int],
        ) -> tuple[list, dict]:
            """Run the UDTF over *qb*, checkpointing each batch."""
            import time

            import pyarrow as pa

            import geneva.checkpoint_utils as _ckp_utils

            store = self._store
            # TODO(OOM): all output batches for a partition are accumulated in
            # memory.  For large partitions, consider streaming results to the
            # checkpoint store without retaining them in this list.
            batches: list[pa.RecordBatch] = []
            total_rows = 0
            t_exec_start = time.monotonic()
            checkpoint_time = 0.0
            for batch_idx, batch in enumerate(self._udtf.execute(qb)):
                total_rows += batch.num_rows
                bk = _ckp_utils.format_udtf_batch_key(partition_prefix, batch_idx)
                if batch_idx in existing_batch_indices:
                    batches.append(store[bk])
                else:
                    t_ckp = time.monotonic()
                    store[bk] = batch
                    checkpoint_time += time.monotonic() - t_ckp
                    batches.append(batch)

            done_key = f"{partition_prefix}_{_ckp_utils.UDTF_PARTITION_DONE_MARKER}"
            t_ckp = time.monotonic()
            store[done_key] = pa.RecordBatch.from_pydict({"_done": [True]})
            checkpoint_time += time.monotonic() - t_ckp

            stats = {
                "rows": total_rows,
                "batches": len(batches),
                "execute_time_s": time.monotonic() - t_exec_start,
                "checkpoint_time_s": checkpoint_time,
            }
            return (batches, stats)

        def _build_index_source(
            self,
            index_partition_info: tuple[int, str, str],
        ) -> Any:
            """Build a source from an index partition ordinal.

            Opens a ``VectorIndexReader`` on the remote side, reads the
            partition's row IDs, and uses ``take_row_ids`` for efficient
            direct row access — bypassing SQL filter construction.

            Parameters
            ----------
            index_partition_info : tuple[int, str, str]
                A ``(partition_ordinal, index_name, column)`` triple where
                *partition_ordinal* is the zero-based IVF partition ID,
                *index_name* is the Lance index identifier, and *column*
                is the indexed vector column name.
            """
            from lance.dataset import VectorIndexReader

            ordinal, index_name, _column = index_partition_info
            lance_ds = self._tbl.to_lance()
            reader = VectorIndexReader(lance_ds, index_name)
            part = reader.read_partition(ordinal)
            # TODO(OOM): all row IDs for the partition are materialized into a
            # Python list.  For very large partitions, consider chunked takes.
            row_ids = part.column("_rowid").to_pylist()
            qb = self._tbl.take_row_ids(row_ids)
            # Project to input_columns only to avoid fetching unnecessary data
            if self._udtf.input_columns:
                qb = qb.select(self._udtf.input_columns)
            return qb

        def process_partition(
            self,
            partition_filter: str | None,
            partition_prefix: str,
            index_partition_info: tuple[int, str, str] | None = None,
        ) -> tuple[list, list[dict], dict]:
            """Execute the UDTF for one partition with checkpointing.

            Parameters
            ----------
            partition_filter : str | None
                SQL filter expression to scope the source scan to a
                single partition (e.g. ``"group = 'a'"``), or *None*
                for unpartitioned execution.
            partition_prefix : str
                Checkpoint key prefix that uniquely identifies this
                partition within the checkpoint store.
            index_partition_info : tuple[int, str, str] | None
                Optional ``(partition_ordinal, index_name, column)``
                tuple.  When provided, the actor loads row IDs from
                the index locally instead of using a SQL filter.

            Returns
            -------
            tuple[list, list[dict], dict]
                A triple of ``(batches, error_dicts, stats)`` where
                *batches* is a list of ``pa.RecordBatch``,
                *error_dicts* contains serialised ``ErrorRecord``
                dicts (empty on success), and *stats* contains
                execution timing information.
            """
            self._ensure_initialized()

            existing = self._get_existing_batch_indices(partition_prefix)
            if index_partition_info is not None:
                qb = self._build_index_source(index_partition_info)
            else:
                qb = self._build_source_query(partition_filter)

            _empty_stats: dict = {
                "rows": 0,
                "batches": 0,
                "execute_time_s": 0.0,
                "checkpoint_time_s": 0.0,
            }

            def _do_execute() -> tuple[list, dict]:
                return self._execute_with_checkpoints(qb, partition_prefix, existing)

            # No error handling configured — fail fast (current behaviour)
            if self._error_config is None:
                batches, stats = _do_execute()
                return (batches, [], stats)

            try:
                # If retry is configured, wrap _do_execute with tenacity
                from tenacity import Retrying
                from tenacity.stop import stop_after_attempt as _stop1

                retry_cfg = self._error_config.retry_config
                if retry_cfg.stop != _stop1(1):  # has non-trivial retry
                    retrying = Retrying(
                        retry=retry_cfg.retry,
                        stop=retry_cfg.stop,
                        wait=retry_cfg.wait,
                        reraise=True,
                    )
                    batches, stats = retrying(_do_execute)
                    return (batches, [], stats)
                else:
                    batches, stats = _do_execute()
                    return (batches, [], stats)
            except Exception as exc:
                from geneva.debug.error_store import (
                    Outcome,
                    get_exception_outcome,
                    make_error_record_from_exception,
                )

                outcome = get_exception_outcome(exc, self._error_config)
                if outcome == Outcome.SKIP:
                    logging.warning(
                        "UDTF: skipping partition %s due to %s: %s",
                        partition_prefix,
                        type(exc).__name__,
                        exc,
                    )
                    error_record = make_error_record_from_exception(
                        exc,
                        job_id=self.job_id,
                        table_uri=self.dest_table_uri,
                        table_name=self.dest_table_name,
                        table_version=None,
                        column_name="*",
                        udf_name=self._udtf.name,
                        udf_version=str(self._udtf.version),
                        batch_index=0,
                    )
                    return ([], [attrs.asdict(error_record)], _empty_stats)
                raise  # FAIL or unmatched → propagate

    return UDTFProcessorActor


# Lazily initialised by _refresh_udtf_matview on first use
_UDTFProcessorActor: Any = None


def _get_udtf_processor_actor() -> Any:
    """Return the UDTFProcessorActor class, creating it on first call."""
    global _UDTFProcessorActor  # noqa: PLW0603
    if _UDTFProcessorActor is None:
        _UDTFProcessorActor = _make_udtf_processor_actor()
    return _UDTFProcessorActor


@attrs.frozen
class _IndexPartitionInfo:
    """Lightweight metadata for an index-based partition.

    Passed to the remote actor so it can load row IDs locally
    instead of receiving a giant ``_rowid IN (...)`` filter string.
    """

    partition_ordinal: int
    index_name: str
    column: str


def _build_column_partition_work_items(
    src_tbl: "Table",
    partition_col: str,
    top_prefix: str,
    checkpoint_store: CheckpointStore,
) -> tuple[list[tuple[str | None, str, _IndexPartitionInfo | None]], list[Any]]:
    """Build work items by scanning distinct values of a partition column.

    Returns ``(work_items, distinct_values)`` where each work item
    is a ``(sql_filter, checkpoint_prefix, None)`` triple.
    Completed partitions (those with a ``__done__`` marker in the
    checkpoint store) are skipped.
    """
    from geneva.checkpoint_utils import (
        UDTF_PARTITION_DONE_MARKER,
        format_udtf_partition_prefix,
    )

    distinct_values: list[Any] = sorted(
        set(
            src_tbl.to_lance()
            .to_table(columns=[partition_col])
            .column(partition_col)
            .to_pylist()
        )
    )

    work_items: list[tuple[str | None, str, _IndexPartitionInfo | None]] = []
    for v in distinct_values:
        part_prefix = format_udtf_partition_prefix(
            top_prefix,
            partition_col=partition_col,
            partition_value=v,
        )
        done_key = f"{part_prefix}_{UDTF_PARTITION_DONE_MARKER}"
        if done_key in checkpoint_store:
            _LOG.info(
                "UDTF partition %s=%s completed, loading from checkpoint",
                partition_col,
                v,
            )
        else:
            if isinstance(v, str):
                escaped = str(v).replace("'", "''")
                pf = f"{partition_col} = '{escaped}'"
            else:
                pf = f"{partition_col} = {v}"
            work_items.append((pf, part_prefix, None))

    return work_items, distinct_values


def _build_index_partition_work_items(
    src_tbl: "Table",
    column: str,
    top_prefix: str,
    checkpoint_store: CheckpointStore,
) -> tuple[list[tuple[str | None, str, _IndexPartitionInfo | None]], list[int]]:
    """Build work items from an existing IVF vector index.

    Returns ``(work_items, distinct_partition_ids)`` where each work item
    is a ``(None, checkpoint_prefix, _IndexPartitionInfo)`` triple.
    Completed partitions (those with a ``__done__`` marker in the
    checkpoint store) are skipped.
    """
    from lance.dataset import VectorIndexReader

    from geneva.checkpoint_utils import (
        UDTF_PARTITION_DONE_MARKER,
        format_udtf_partition_prefix,
    )

    lance_ds = src_tbl.to_lance()
    # TODO: describe_indices should not require filtering by type_url;
    # this is a workaround for a current limitation in the API.
    try:
        idx_info = next(
            idx
            for idx in lance_ds.describe_indices()
            if column in idx.field_names and idx.type_url == _VECTOR_INDEX_TYPE_URL
        )
    except StopIteration:
        raise ValueError(
            f"No IVF vector index found on column '{column}'. "
            f"Create one first (e.g. create_ivf_flat_index())."
        ) from None

    reader = VectorIndexReader(lance_ds, idx_info.name)

    partition_col = _IVF_PARTITION_COL
    work_items: list[tuple[str | None, str, _IndexPartitionInfo | None]] = []
    distinct_values: list[int] = []

    for pid in range(reader.num_partitions()):
        part = reader.read_partition(pid)
        if part.num_rows == 0:
            continue

        distinct_values.append(pid)
        part_prefix = format_udtf_partition_prefix(
            top_prefix,
            partition_col=partition_col,
            partition_value=pid,
        )
        done_key = f"{part_prefix}_{UDTF_PARTITION_DONE_MARKER}"
        if done_key in checkpoint_store:
            _LOG.info(
                "UDTF partition %s=%d completed, loading from checkpoint",
                partition_col,
                pid,
            )
        else:
            info = _IndexPartitionInfo(
                partition_ordinal=pid,
                index_name=idx_info.name,
                column=column,
            )
            work_items.append((None, part_prefix, info))

    return work_items, distinct_values


def _get_udf_name_from_field(field: pa.Field) -> str | None:
    """Extract UDF name from column field metadata."""
    if field.metadata is None:
        return None
    udf_name = field.metadata.get(b"virtual_column.udf_name")
    if udf_name is None:
        return None
    return udf_name.decode() if isinstance(udf_name, bytes) else udf_name


@attrs.define
class JobFuture:
    job_id: str

    def done(self, timeout: float | None = None) -> bool:
        raise NotImplementedError("JobFuture.done() must be implemented in subclasses")

    def result(self, timeout: float | None = None) -> Any:
        raise NotImplementedError(
            "JobFuture.result() must be implemented in subclasses"
        )

    def status(self, timeout: float | None = None) -> None:
        raise NotImplementedError(
            "JobFuture.status() must be implemented in subclasses"
        )


@attrs.define(order=True)
class TableReference:
    """
    Serializable reference to a Geneva Table.

    Used to pass through ray.remote calls
    """

    table_id: list[str]
    version: int | None

    db_uri: str | None

    api_key: Any = attrs.field(default=None, repr=False)
    host_override: str | None = attrs.field(default=None)
    namespace_impl: str | None = attrs.field(default=None)
    namespace_properties: dict[str, str] | None = attrs.field(default=None)
    system_namespace: list[str] | None = attrs.field(default=None)
    is_system_table: bool = attrs.field(default=False)
    co_located_system_tables: bool = attrs.field(default=False)

    @property
    def table_name(self) -> str:
        """Return the table name (last element of table_id)."""
        return self.table_id[-1] if self.table_id else ""

    def _uses_namespace_for_system_tables(self) -> bool:
        return uses_namespace_for_system_tables(
            self.db_uri,
            self.namespace_impl,
        )

    @property
    def inferred_system_namespace(self) -> list[str]:
        return infer_system_namespace(
            uri=self.db_uri,
            namespace_impl=self.namespace_impl,
            system_namespace=self.system_namespace,
        )

    @property
    def inferred_system_db_uri(self) -> str | None:
        return infer_system_db_uri(
            uri=self.db_uri,
            namespace_impl=self.namespace_impl,
        )

    def _connection_args(
        self,
        *,
        use_system_db_root: bool,
    ) -> tuple[str | None, list[str], bool]:
        system_namespace = self.inferred_system_namespace
        db_uri = self.db_uri
        is_system_db_root = False
        if use_system_db_root and not self._uses_namespace_for_system_tables():
            db_uri = self.inferred_system_db_uri
            is_system_db_root = True
        return db_uri, system_namespace, is_system_db_root

    def as_system_table(self, table_name: str) -> "TableReference":
        """Create a sibling reference that targets a system table."""
        if self.inferred_system_namespace:
            table_id = self.inferred_system_namespace + [table_name]
        else:
            table_id = [table_name]
        return TableReference(
            table_id=table_id,
            version=None,
            db_uri=self.db_uri,
            api_key=self.api_key,
            host_override=self.host_override,
            namespace_impl=self.namespace_impl,
            namespace_properties=self.namespace_properties,
            system_namespace=self.inferred_system_namespace,
            is_system_table=not self.co_located_system_tables,
            co_located_system_tables=self.co_located_system_tables,
        )

    @property
    def system_table_namespace(self) -> list[str]:
        return self.inferred_system_namespace

    def open_checkpoint_store(self) -> CheckpointStore:
        """Open a Lance checkpoint store for this table."""
        try:
            namespace_client = self.connect_namespace()
            table_location = resolve_table_physical_uri(
                self.table_id,
                db_uri=self.db_uri,
                api_key=self.api_key,
                host_override=self.host_override,
                namespace_impl=self.namespace_impl,
                namespace_properties=self.namespace_properties,
            )
            checkpoint_uri = str(URL(table_location) / "_ckp")

            # Create checkpoint store with namespace support
            from geneva.checkpoint import LanceCheckpointStore

            return LanceCheckpointStore(
                checkpoint_uri,
                namespace_client=namespace_client,
                table_id=self.table_id,
            )
        except Exception:
            # Fallback to in-memory checkpoint store if Lance store fails
            return InMemoryCheckpointStore()

    def open_db(self) -> Connection:
        """Open a connection to the Lance database.
        Set read consistency interval to 0 for strongly consistent reads."""
        cp = self.open_checkpoint_store()
        interval = timedelta(0)

        if self.namespace_impl is not None:
            assert self.namespace_properties is not None, (
                "namespace_properties must be set when namespace_impl is set"
            )
            _, system_namespace, _ = self._connection_args(use_system_db_root=False)
            return connect(
                namespace_impl=self.namespace_impl,
                namespace_properties=self.namespace_properties,
                api_key=self.api_key,
                host_override=self.host_override,
                system_namespace=system_namespace,
                checkpoint=cp,
                initialize_upload_dir=False,
                read_consistency_interval=interval,
            )

        db_uri, system_namespace, is_system_db_root = self._connection_args(
            use_system_db_root=self.is_system_table
        )
        assert db_uri is not None, "db_uri must be set"
        connect_kwargs: dict[str, Any] = {
            "checkpoint": cp,
            "initialize_upload_dir": False,
            "read_consistency_interval": interval,
            "system_namespace": system_namespace,
            "is_system_db_root": is_system_db_root,
        }
        if str(db_uri).startswith("db://"):
            connect_kwargs["api_key"] = self.api_key
            connect_kwargs["host_override"] = self.host_override
        return connect(db_uri, **connect_kwargs)

    async def open_db_async(
        self,
    ) -> AsyncConnection | AsyncLanceNamespaceDBConnection:
        """Open an async connection to the Lance database.
        This uses native lancedb AsyncConnection and doesn't support checkpoint store.
        Currently used by JobTracker only.
        """
        if namespace := self.connect_namespace():
            return AsyncLanceNamespaceDBConnection(
                namespace,
                read_consistency_interval=timedelta(0),
            )

        db_uri, _, _ = self._connection_args(use_system_db_root=False)
        assert db_uri is not None, "db_uri must be set"
        return await connect_async(
            db_uri,
            api_key=self.api_key,
            host_override=self.host_override,
            read_consistency_interval=timedelta(0),
        )

    async def open_system_db_async(
        self,
    ) -> AsyncConnection | AsyncLanceNamespaceDBConnection:
        """Open an async connection suitable for system-table operations."""
        if namespace := self.connect_namespace():
            return AsyncLanceNamespaceDBConnection(
                namespace,
                read_consistency_interval=timedelta(0),
            )

        system_db_uri, _, _ = self._connection_args(use_system_db_root=True)
        assert system_db_uri is not None, "db_uri must be set"
        return await connect_async(
            system_db_uri,
            api_key=self.api_key,
            host_override=self.host_override,
            read_consistency_interval=timedelta(0),
        )

    def open(self) -> "Table":
        # Extract namespace from table_id (everything except the last element)
        namespace = self.table_id[:-1] if len(self.table_id) > 1 else []
        return self.open_db().open_table(
            self.table_name, version=self.version, namespace=namespace
        )

    def connect_namespace(self) -> Optional["LanceNamespace"]:
        """Connect using the Lance namespace if configured"""
        if self.namespace_impl and self.namespace_properties:
            from lance_namespace import connect as namespace_connect

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


class Table(LanceTable):
    """Table in Geneva.

    A Table is a Lance dataset
    """

    def __init__(
        self,
        conn: Connection,
        name: str,
        *,
        namespace: list[str] | None = None,
        version: int | None = None,
        storage_options: dict[str, str] | None = None,
        index_cache_size: int | None = None,
        **kwargs,
    ) -> None:
        self._conn_uri = conn.uri
        self._name = name

        if namespace is None:
            namespace = []
        self._namespace = namespace
        self._table_id = namespace + [name]

        self._conn = conn

        self._version: int | None = version
        self._index_cache_size = index_cache_size
        self._storage_options = storage_options

        # Load table
        self._ltbl  # noqa

    def __repr__(self) -> str:
        return f"<Table {self._table_id}>"

    # TODO: This annotation sucks
    def __reduce__(self):  # noqa: ANN204
        return (self.__class__, (self._conn, self._name))

    def get_reference(self) -> TableReference:
        physical_db_uri = _get_db_uri(self.uri)
        co_located_system_tables = (
            self._conn._host_override is not None
            and self._conn.effective_system_namespace == []
        )
        return TableReference(
            table_id=self._table_id,
            version=self._version,
            db_uri=physical_db_uri,
            api_key=_plain_api_key(self._conn._api_key),
            host_override=self._conn._host_override,
            namespace_impl=self._conn.namespace_impl,
            namespace_properties=self._conn.namespace_properties,
            system_namespace=self._conn.effective_system_namespace,
            is_system_table=self._conn._is_system_table_db_root(),
            co_located_system_tables=co_located_system_tables,
        )

    def _is_namespace_table(self) -> bool:
        """
        Check if this table is backed by a namespace implementation.

        A table is considered namespace-backed if both namespace_impl and
        namespace_properties are configured on the connection.

        Returns
        -------
        bool
            True if this is a namespace table, False otherwise.
        """
        return (
            self._conn.namespace_impl is not None
            and self._conn.namespace_properties is not None
        )

    def get_fragments(self) -> list[lance.LanceFragment]:
        return self.to_lance().get_fragments()

    @cached_property
    def _ltbl(self) -> lancedb.table.Table:
        """Return the inner LanceDB table.
        This might be a LanceTable or a RemoteTable depending on the connection
        configuration
        """
        inner = self._conn._connect

        # remote db, open table directly
        if self._conn.is_remote():
            tbl = inner.open_table(self._name, namespace=self._namespace)
        else:
            _LOG.debug(
                f"opening table {self._table_id} {type(self)=} {self._namespace=} "
            )
            tbl = inner.open_table(self.name, namespace=self._namespace)

        # Check out the specified version regardless of database type
        if self._version:
            tbl.checkout(self._version)
        return tbl

    @cached_property
    def _native_tbl(self) -> lancedb.table.Table:
        """Return a native lancedb table. If the connection is configured for
        remote tables (phalanx) this will convert the remote table to a native
        table pointing directly at object storage.
        This is required for some advanced operations like add_columns,
        materialized views, etc.
        Usage of this table requires client object storage permissions.
        """
        tbl = self._ltbl
        if not self._conn.is_remote():
            return tbl

        tbl_uri: str = tbl.uri  # pyright: ignore[reportAttributeAccessIssue]
        db_uri = _get_db_uri(tbl_uri)

        # create a new direct lance connection
        conn = self._conn._get_conn(uri=db_uri)

        new_tbl = conn.open_table(tbl.name, namespace=self._namespace)
        if self._version:
            new_tbl.checkout(self._version)
        return new_tbl

    @property
    def name(self) -> str:
        """Get the name of the table."""
        return self._name

    @property
    def version(self) -> int:
        """Get the current version of the table"""
        return self._native_tbl.version

    @property
    def schema(self) -> pa.Schema:
        """The Arrow Schema of the Table."""
        return self._native_tbl.schema

    @cached_property
    def uri(self) -> str:
        conn = self._conn
        if conn.namespace_impl is not None and conn.namespace_properties is not None:
            # Get the actual table location from the namespace
            from lance_namespace import DescribeTableRequest

            ns = conn.namespace_client
            if ns is not None:
                response = ns.describe_table(DescribeTableRequest(id=self._table_id))
                if response.location is None:
                    raise ValueError(
                        f"Table location is None for table {self._table_id}"
                    )
                return response.location
            else:
                return f"{self.name}.lance"

        inner_uri = self._ltbl.uri  # pyright: ignore[reportAttributeAccessIssue]
        inner_uri = inner_uri.rstrip("?")
        return inner_uri

    @property
    def embedding_functions(self) -> Never:
        raise NotImplementedError("Embedding functions are not supported.")

    def add(
        self,
        data,
        mode: str = "append",
        on_bad_vectors: str = "error",
        fill_value: float = 0.0,
    ) -> None:
        self._ltbl.add(
            data,
            mode=mode,  # type: ignore[arg-type]
            on_bad_vectors=on_bad_vectors,  # type: ignore[arg-type]
            fill_value=fill_value,
        )

    def checkout(self, version: int) -> None:
        self._version = version
        self._ltbl.checkout(version)

    def checkout_latest(self) -> None:
        self._version = None
        self._ltbl.checkout_latest()

    def add_columns(
        self, transforms: dict[str, str | UDF | tuple[UDF, list[str]]], *args, **kwargs
    ) -> None:
        """
        Add columns or UDF-based columns to the Geneva table.

        For UDF columns, this method validates that:
        - All input columns exist in the table schema
        - Column types are compatible with UDF type annotations (if present)
        - RecordBatch UDFs do not have input_columns defined

        This early validation helps catch configuration errors before job execution.

        Parameters
        ----------
        transforms : dict[str, str | UDF | tuple[UDF, list[str]]]
            The key is the column name to add and the value is a
            specification of the column type/value.

            * If the spec is a string, it is expected to be a datafusion
              sql expression. (e.g "cast(null as string)")
            * If the spec is a UDF, a virtual column is added with input
              columns inferred from the UDF's argument names.
            * If the spec is a tuple, the first element is a UDF and the
              second element is a list of input column names.

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

        Warns
        -----
        UserWarning
            If type validation is skipped due to missing type annotations

        Examples
        --------
        >>> @udf(data_type=pa.int32())
        ... def double(a: int) -> int:
        ...     return a * 2
        >>> table.add_columns({"doubled": double})  # Validates 'a' column exists

        """
        t = self._native_tbl

        # handle basic columns
        basic_cols = {k: v for k, v in transforms.items() if isinstance(v, str)}
        if len(basic_cols) > 0:
            t.add_columns(basic_cols, *args)

        # handle UDF virtual columns
        udf_cols = {k: v for k, v in transforms.items() if not isinstance(v, str)}
        for k, v in udf_cols.items():
            if isinstance(v, UDF):
                # infer column names from udf arguments
                udf = v
                self._add_virtual_columns(
                    {k: udf}, *args, input_columns=udf.input_columns, **kwargs
                )
            else:
                # explicitly specify input columns
                (udf, cols) = v
                self._add_virtual_columns({k: udf}, *args, input_columns=cols, **kwargs)

    def _add_virtual_columns(
        self,
        mapping: dict[str, UDF],  # this breaks the non udf mapping
        *args,
        input_columns: list[str] | None = None,
        **kwargs,
    ) -> None:
        """
        This is an internal method and not intended to be called directly.

        Add udf based virtual columns to the Geneva table.
        """

        if len(mapping) != 1:
            raise ValueError("Only one UDF is supported for now.")

        _LOG.info("Adding column: udf=%s", mapping)
        col_name = next(iter(mapping))
        udf = mapping[col_name]

        if not isinstance(udf, UDF):
            # Stateful udfs are implemenated as Callable classses, and look
            # like partial functions here.  Instantiate to get the return
            # data_type annotations.
            udf = udf()

        # Validate input columns exist in table schema before adding the column
        self._validate_udf_input_columns(udf, input_columns)

        # Check for circular dependencies before adding the column
        cols_to_check = (
            input_columns if input_columns is not None else udf.input_columns
        )
        if (
            udf.arg_type != UDFArgType.RECORD_BATCH
            and cols_to_check
            and col_name in cols_to_check
        ):
            raise ValueError(
                f"UDF output column {col_name} is not allowed to be in"
                f" input {cols_to_check}"
            )

        self._native_tbl.add_columns(pa.field(col_name, udf.data_type))
        self._configure_computed_column(col_name, udf, input_columns)

    def _validate_udf_input_columns(
        self, udf: UDF, input_columns: list[str] | None
    ) -> None:
        """
        Validate that UDF input columns exist in the table schema.

        This method delegates to the UDF's validate_against_schema() method
        for consolidated validation logic.

        Parameters
        ----------
        udf: UDF
            The UDF to validate
        input_columns: list[str] | None
            The input column names to validate

        Raises
        ------
        ValueError: If input columns don't exist in table schema or have type mismatches
        """
        # Delegate to UDF's consolidated validation method
        udf.validate_against_schema(self._native_tbl.schema, input_columns)

    def refresh(
        self,
        *,
        where: str | None = None,
        src_version: int | None = None,
        max_rows_per_fragment: int | None = None,
        concurrency: int = 8,
        intra_applier_concurrency: int = 1,
        _admission_check: bool | None = None,
        _admission_strict: bool | None = None,
        **kwargs,
    ) -> None:
        """
        Refresh the specified materialized view.

        If run without a context, a warning is emitted and local Ray is used
        implicitly. For proper cleanup, wrap in conn.local_ray_context().

        Parameters
        ----------
        where: str | None
            TODO: sql expression filter used to only backfill selected rows
        src_version: int | None
            Optional source table version to refresh from. If None (default),
            uses the latest version of the source table.
        max_rows_per_fragment: int | None
            Optional maximum number of rows per destination fragment when adding
            placeholder rows for new source data. If None, uses LanceDB's default
            (1 million rows). Use smaller values to control fragment granularity.
        concurrency: int
            (default = 8) This controls the number of processes that tasks run
            concurrently. For max throughput, ideally this is larger than the number
            of nodes in the k8s cluster. This is the number of Ray actor processes
            that are started.
        intra_applier_concurrency: int
            (default = 1) This controls the number of threads used to execute tasks
            within a process. Multiplying this times `concurrency` roughly corresponds
            to the number of cpu's being used.
        _admission_check: bool | None
            Whether to run admission control to validate cluster resources before
            starting the job. If None, uses config (default: true). Set to False to
            skip the check.
            **Experimental**: Parameters starting with `_` are subject to change.
        _admission_strict: bool | None
            If True, raises ResourcesUnavailableError when resources are
            insufficient. If False, logs a warning but allows the job to proceed.
            If None, uses config (default: true).
            **Experimental**: Parameters starting with `_` are subject to change.

        Raises
        ------
        RuntimeError
            If attempting to refresh to a different version without stable row IDs
            enabled on the source table. This is because compaction may have
            invalidated the __source_row_id values, breaking incremental refresh.
        """
        if where:
            raise NotImplementedError(
                "where clauses on materialized view refresh not implemented yet."
            )

        if get_current_context() is None:
            _LOG.warning(_IMPLICIT_LOCAL_RAY_WARNING, "refresh()")

        # Check if source table has stable row IDs and validate version compatibility
        schema = self.schema
        metadata = schema.metadata or {}

        # Check if this is a UDTF-backed materialized view
        mv_version_bytes = metadata.get(MATVIEW_META_VERSION.encode(), b"1")
        mv_version_str = mv_version_bytes.decode()
        if mv_version_str == "udtf":
            # UDTF-backed view - do full refresh by re-executing UDTF
            return self._refresh_udtf_matview(
                src_version=src_version,
                concurrency=concurrency,
                _admission_check=_admission_check,
                _admission_strict=_admission_strict,
            )

        if mv_version_str == MATVIEW_VERSION_SCALAR_UDTF:
            # Scalar UDTF-backed view (1:N) — routes through run_ray_copy_table
            # which detects the scalar_udtf metadata and runs the expansion.
            from geneva.runners.ray.pipeline import run_ray_copy_table

            table_ref = self.get_reference()
            checkpoint_store = table_ref.open_checkpoint_store()
            run_ray_copy_table(
                table_ref,
                self._conn._packager,
                checkpoint_store,
                src_version=src_version,
                max_rows_per_fragment=max_rows_per_fragment,
                concurrency=concurrency,
                intra_applier_concurrency=intra_applier_concurrency,
                _admission_check=_admission_check,
                _admission_strict=_admission_strict,
            )
            self.checkout_latest()
            return

        # Get MV format version from metadata
        # Version 1: fragment+offset encoding, no stable row IDs
        # Version 2: stable row IDs enabled
        mv_version_bytes = metadata.get(MATVIEW_META_VERSION.encode(), b"1")
        mv_version = int(mv_version_bytes.decode())
        has_stable_row_ids = mv_version >= 2

        # Get the base version (version when MV was created)
        base_version_str = metadata.get(MATVIEW_META_BASE_VERSION.encode())
        # If no base version metadata, assume it's safe to proceed
        # (for backwards compatibility with MVs created before this feature)
        base_version = int(base_version_str.decode()) if base_version_str else None

        # Resolve src_version to actual version number if None (implicit latest)
        if src_version is None:
            # Get the source table to find its latest version

            source_table_name = metadata.get(MATVIEW_META_BASE_TABLE.encode())
            source_db_uri = metadata.get(MATVIEW_META_BASE_DBURI.encode())

            # Get namespace path from metadata (source table's path within namespace)
            namespace_path_b = 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 source_table_name and source_db_uri:
                # Use view's namespace connection if available
                if (
                    self._conn.namespace_impl is not None
                    and self._conn.namespace_properties is not None
                ):
                    source_table = self._conn.open_table(
                        source_table_name.decode(), namespace=namespace_path
                    )
                elif source_db_uri.decode() == self._conn._uri:
                    # Same database - reuse connection
                    source_conn = self._conn
                    source_table = source_conn.open_table(source_table_name.decode())
                else:
                    # Different database - create new connection
                    source_conn = connect(source_db_uri.decode())
                    source_table = source_conn.open_table(source_table_name.decode())

                src_version = source_table.version

        # Validate: if no stable row IDs and src_version differs from base, fail
        if (
            not has_stable_row_ids
            and base_version is not None
            and src_version is not None
            and src_version != base_version
        ):
            raise RuntimeError(
                f"Cannot refresh materialized view to version {src_version} "
                "because the source table does not have stable row IDs "
                f"enabled.\n\n"
                f"This materialized view was created from source version "
                f"{base_version}. "
                "Without stable row IDs, incremental refresh is only supported "
                "when refreshing to the SAME version it was created from.\n\n"
                "This limitation exists because compaction operations may have "
                "changed the physical row IDs between versions, which would "
                "break the materialized view's ability to track source rows.\n\n"
                "To enable refresh across all versions, recreate the source "
                "table with stable row IDs:\n"
                "  db.create_table(\n"
                "      name='table_name',\n"
                "      data=data,\n"
                "      storage_options={'new_table_enable_stable_row_ids': True}\n"
                "  )"
            )

        # Note: backwards refresh (point-in-time refresh to older versions) is now
        # supported when stable row IDs are enabled. The actual rollback logic
        # (deleting rows not in the target version) is handled in run_ray_copy_table.

        from geneva.runners.ray.pipeline import run_ray_copy_table

        # Use table-specific checkpoint store
        table_ref = self.get_reference()
        checkpoint_store = table_ref.open_checkpoint_store()
        run_ray_copy_table(
            table_ref,
            self._conn._packager,
            checkpoint_store,
            src_version=src_version,
            max_rows_per_fragment=max_rows_per_fragment,
            concurrency=concurrency,
            intra_applier_concurrency=intra_applier_concurrency,
            _admission_check=_admission_check,
            _admission_strict=_admission_strict,
        )

        # Update last refreshed version in metadata
        if src_version is not None:
            _set_last_refreshed_version(self, src_version)

        self.checkout_latest()

    def _refresh_udtf_matview(
        self,
        src_version: int | None = None,
        concurrency: int = 8,
        _admission_check: bool | None = None,
        _admission_strict: bool | None = None,
    ) -> None:
        """Refresh a UDTF-backed materialized view (full refresh).

        Re-executes the UDTF against the source table via Ray and overwrites
        the view with the new results.  Uses an ActorPool for fault-tolerant
        parallel execution and a JobTracker for progress metrics.
        Per-batch checkpointing allows failed refreshes to resume from the
        last completed batch.

        Ordering guarantees: within a single partition, batch indices are
        deterministic (0, 1, 2, ...) based on the generator yield order, and
        the source scan order is deterministic for a given pinned version.
        Across partitions, execution and assembly order is not guaranteed.
        """
        import functools
        import uuid

        import ray

        import geneva.cloudpickle as cloudpickle
        from geneva.checkpoint_utils import (
            format_udtf_checkpoint_prefix,
        )
        from geneva.packager import UDTFSpec, unmarshal_udtf
        from geneva.runners.ray.admission import (
            JOBTRACKER_MEMORY,
            JOBTRACKER_NUM_CPUS,
            validate_udtf_admission,
        )
        from geneva.runners.ray.jobtracker import JobTracker

        schema = self.to_lance().schema
        metadata = schema.metadata or {}

        # Load UDTF from metadata
        udtf_spec_json = metadata.get(MATVIEW_META_UDTF.encode())
        if not udtf_spec_json:
            raise ValueError(
                "Cannot refresh: No UDTF specification found in view metadata."
            )

        udtf_spec = UDTFSpec.from_json(udtf_spec_json.decode())
        udtf_obj = unmarshal_udtf(udtf_spec)
        if udtf_obj is None:
            raise ValueError(
                "Cannot refresh: Failed to deserialize UDTF. "
                "The UDTF module may not be available in this environment."
            )

        # Load source query
        source_query_bytes: bytes | None = metadata.get(
            MATVIEW_META_SOURCE_QUERY.encode()
        )

        # Get source table info
        source_table_name_b = metadata.get(MATVIEW_META_BASE_TABLE.encode())
        source_db_uri_b = metadata.get(MATVIEW_META_BASE_DBURI.encode())

        if not source_table_name_b or not source_db_uri_b:
            raise ValueError(
                "Cannot refresh: Source table information not found in view metadata."
            )

        source_uri = source_db_uri_b.decode()
        source_name = source_table_name_b.decode()

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

        # Open source using view's namespace connection if available
        # Also track source_api_key and source_host_override for Ray actor
        source_api_key: str | None = None
        source_host_override: str | None = None
        if (
            self._conn.namespace_impl is not None
            and self._conn.namespace_properties is not None
        ):
            # Using namespace connection - actor will use namespace_impl/properties
            # api_key/host_override are fallback values from view's connection
            src_tbl = self._conn.open_table(source_name, namespace=namespace_path)
            source_api_key = _plain_api_key(self._conn._api_key)
            source_host_override = self._conn._host_override
        elif source_uri == self._conn._uri:
            # Same URI - use view's connection credentials
            src_tbl = self._conn.open_table(source_name)
            source_api_key = _plain_api_key(self._conn._api_key)
            source_host_override = self._conn._host_override
        else:
            # Different URI - create new connection with view's credentials
            source_conn = connect(
                source_uri,
                api_key=_plain_api_key(self._conn._api_key),
                host_override=self._conn._host_override,
            )
            src_tbl = source_conn.open_table(source_name)
            source_api_key = _plain_api_key(source_conn._api_key)
            source_host_override = source_conn._host_override

        # Determine actual source version
        if src_version is not None:
            actual_source_version = src_version
        else:
            actual_source_version = src_tbl.version
        src_tbl.checkout(actual_source_version)

        # Version-aware bail-out: skip refresh if source hasn't changed
        # and the view already has data (not the initial empty state after
        # create_udtf_view).
        base_version_str = metadata.get(MATVIEW_META_BASE_VERSION.encode())
        base_version = int(base_version_str.decode()) if base_version_str else None
        if (
            base_version is not None
            and base_version == actual_source_version
            and self.count_rows() > 0
        ):
            _LOG.info(
                "UDTF source unchanged (version %d), skipping refresh",
                base_version,
            )
            return

        # Ensure Ray is initialized
        if get_current_context() is None:
            _LOG.warning(_IMPLICIT_LOCAL_RAY_WARNING, "_refresh_udtf_matview()")
            ray.init(ignore_reinit_error=True)

        # --- Admission control ---
        udtf_num_cpus = udtf_obj.num_cpus or 1.0
        udtf_num_gpus = udtf_obj.num_gpus or 0.0
        udtf_memory = udtf_obj.memory or 0
        validate_udtf_admission(
            udtf_num_cpus=udtf_num_cpus,
            udtf_num_gpus=udtf_num_gpus,
            udtf_memory=udtf_memory,
            concurrency=concurrency,
            check=_admission_check,
            strict=_admission_strict,
        )

        # Open checkpoint store
        table_ref = self.get_reference()
        checkpoint_store = table_ref.open_checkpoint_store()
        top_prefix = format_udtf_checkpoint_prefix(
            udtf_name=udtf_obj.name,
            udtf_version=udtf_obj.version,
            source_version=actual_source_version,
        )

        udtf_pickle = cloudpickle.dumps(udtf_obj)
        error_handling_pickle: bytes | None = (
            cloudpickle.dumps(udtf_obj.error_handling)
            if udtf_obj.error_handling is not None
            else None
        )
        ckp_uri = checkpoint_store.uri()

        # --- JobTracker ---
        job_id = uuid.uuid4().hex
        job_tracker = JobTracker.options(
            name=f"jobtracker-udtf-{job_id}",
            num_cpus=JOBTRACKER_NUM_CPUS,
            memory=JOBTRACKER_MEMORY,
            max_restarts=-1,
        ).remote(job_id, table_ref, enable_saves=False)  # type: ignore[call-arg]

        # Register with cluster exit bookkeeping so the cluster won't be
        # torn down while the UDTF matview refresh is still running.
        from geneva.runners.ray.raycluster import RayCluster

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

        # --- Actor resource kwargs ---
        actor_resource_kwargs: dict[str, Any] = {
            "num_cpus": udtf_num_cpus,
            "num_gpus": udtf_num_gpus,
        }
        if udtf_obj.memory is not None:
            actor_resource_kwargs["memory"] = udtf_obj.memory

        # --- Actor factory ---
        # Use self.uri instead of constructing from self._conn._uri directly
        # so that namespace-backed connections resolve the correct table
        # location via the catalog.
        dest_table_uri = self.uri
        actor_cls = _get_udtf_processor_actor().options(**actor_resource_kwargs)

        # Serialize namespace_properties as JSON string for Ray
        namespace_properties_json: str | None = None
        if self._conn.namespace_properties is not None:
            namespace_properties_json = json.dumps(self._conn.namespace_properties)

        actor_factory = functools.partial(
            actor_cls.remote,
            udtf_pickle_bytes=udtf_pickle,
            source_uri=source_uri,
            source_api_key=source_api_key,
            source_host_override=source_host_override,
            source_name=source_name,
            source_query_json_bytes=source_query_bytes,
            ckp_store_uri=ckp_uri,
            error_handling_pickle_bytes=error_handling_pickle,
            job_id=job_id,
            dest_table_uri=dest_table_uri,
            dest_table_name=self.name,
            source_version=actual_source_version,
            namespace_impl=self._conn.namespace_impl,
            namespace_properties_json=namespace_properties_json,
            namespace_path=namespace_path,
        )

        try:
            self._run_udtf_refresh(
                udtf_obj=udtf_obj,
                job_tracker=job_tracker,
                checkpoint_store=checkpoint_store,
                metadata=metadata,
                actual_source_version=actual_source_version,
                concurrency=concurrency,
                top_prefix=top_prefix,
                src_tbl=src_tbl,
                table_ref=table_ref,
                actor_factory=actor_factory,
                dest_table_uri=dest_table_uri,
            )
        finally:
            # Signal the JobTracker that this job is finished so that
            # _wait_for_tracked_jobs can detect completion.
            with contextlib.suppress(Exception):
                job_tracker.mark_job_done.remote()

        self.checkout_latest()

    def _run_udtf_refresh(  # noqa: C901
        self,
        *,
        udtf_obj: Any,
        job_tracker: Any,
        checkpoint_store: "CheckpointStore",
        metadata: dict[bytes, bytes],
        actual_source_version: int,
        concurrency: int,
        top_prefix: str,
        src_tbl: Any,
        table_ref: Any,
        actor_factory: Any,
        dest_table_uri: str,
    ) -> None:
        """Execute UDTF dispatch + assembly.

        Separated from ``_refresh_udtf_matview`` so the caller can guarantee
        ``mark_job_done`` in a ``finally`` block.
        """
        import lance
        import pyarrow as pa

        from geneva.checkpoint_utils import (
            UDTF_PARTITION_DONE_MARKER,
            format_udtf_partition_prefix,
            parse_udtf_batch_index,
        )
        from geneva.runners.ray.actor_pool import ActorPool
        from geneva.tqdm import Colors, fmt, tqdm

        # Helper: load all batches for a completed partition from checkpoint
        def _load_partition_batches(
            store: CheckpointStore, part_prefix: str
        ) -> list[pa.RecordBatch]:
            batch_keys = sorted(
                k
                for k in store.list_keys(part_prefix)
                if parse_udtf_batch_index(k) is not None
            )
            return [store[k] for k in batch_keys]

        # --- Build work items ---
        # work_items: list of (partition_filter, partition_prefix, index_info)
        work_items: list[tuple[str | None, str, _IndexPartitionInfo | None]] = []

        if udtf_obj.partition_by_indexed_column:
            work_items, distinct_values = _build_index_partition_work_items(
                src_tbl,
                udtf_obj.partition_by_indexed_column,
                top_prefix,
                checkpoint_store,
            )
        elif udtf_obj.partition_by:
            work_items, distinct_values = _build_column_partition_work_items(
                src_tbl, udtf_obj.partition_by, top_prefix, checkpoint_store
            )
        else:
            part_prefix = format_udtf_partition_prefix(
                top_prefix, partition_col=None, partition_value=None
            )
            done_key = f"{part_prefix}_{UDTF_PARTITION_DONE_MARKER}"
            if done_key in checkpoint_store:
                _LOG.info(
                    "UDTF (unpartitioned) already completed, loading from checkpoint"
                )
            else:
                work_items.append((None, part_prefix, None))

        # --- ActorPool dispatch ---
        all_error_dicts: list[dict] = []
        if work_items:
            num_actors = min(len(work_items), concurrency)
            job_tracker.set_total.remote("partitions", len(work_items))
            job_tracker.set_desc.remote("partitions", "Partitions processed")

            pool = ActorPool(
                actor_factory,
                num_actors,
                job_tracker=job_tracker,
                worker_metric="workers",
            )
            skipped_count = 0
            empty_count = 0
            total_rows_produced = 0
            bar_partitions = tqdm(
                total=len(work_items),
                desc=fmt("udtf | partitions", Colors.CYAN, bold=True),
                unit="part",
            )
            bar_rows = tqdm(
                total=0,
                bar_format="{desc}: {n_fmt} [{elapsed}]",
                desc=fmt("udtf | rows produced", Colors.CYAN, bold=True),
            )
            try:
                for result in pool.map_unordered(
                    lambda actor, item: actor.process_partition.remote(
                        item[0],
                        item[1],
                        (
                            (
                                item[2].partition_ordinal,
                                item[2].index_name,
                                item[2].column,
                            )
                            if item[2] is not None
                            else None
                        ),
                    ),
                    work_items,
                ):
                    batch_list, error_dicts, stats = result
                    if error_dicts:
                        skipped_count += 1
                    elif not batch_list:
                        empty_count += 1
                    else:
                        # Track the prefix for checkpoint-based assembly.
                        # The prefix is embedded in the work_items but we
                        # don't know which item produced this result. Instead
                        # we note that successful batches means the __done__
                        # marker was written, so we'll discover them when we
                        # re-scan the checkpoint store below.
                        pass
                    total_rows_produced += stats.get("rows", 0)
                    all_error_dicts.extend(error_dicts)
                    job_tracker.increment.remote("partitions", 1)
                    bar_partitions.update(1)
                    bar_rows.n = total_rows_produced
                    bar_rows.refresh()
            finally:
                bar_partitions.close()
                bar_rows.close()
                pool.shutdown()
            if skipped_count:
                _LOG.warning(
                    "UDTF refresh: %d/%d partitions skipped due to errors",
                    skipped_count,
                    len(work_items),
                )
            if empty_count:
                _LOG.info(
                    "UDTF refresh: %d/%d partitions produced no output rows",
                    empty_count,
                    len(work_items),
                )

        # --- ErrorStore bulk logging ---
        if all_error_dicts:
            from geneva.debug.error_store import ErrorRecord, ErrorStore

            error_store = ErrorStore(self._conn, namespace=self._conn.system_namespace)
            errors = [ErrorRecord(**d) for d in all_error_dicts]
            error_store.log_errors(errors)

        # --- Checkpoint-based final assembly ---
        # Re-scan checkpoint store for all completed partitions (covers both
        # previously-checkpointed and newly-processed partitions).
        all_completed_prefixes: list[str] = []
        if udtf_obj.partition_by_indexed_column:
            partition_col = _IVF_PARTITION_COL
            for v in distinct_values:
                pfx = format_udtf_partition_prefix(
                    top_prefix,
                    partition_col=partition_col,
                    partition_value=v,
                )
                done_key = f"{pfx}_{UDTF_PARTITION_DONE_MARKER}"
                if done_key in checkpoint_store:
                    all_completed_prefixes.append(pfx)
        elif udtf_obj.partition_by:
            partition_col = udtf_obj.partition_by
            # Reuse distinct_values computed during work-item build above
            # (source is pinned to the same version, so values are identical).
            for v in distinct_values:
                pfx = format_udtf_partition_prefix(
                    top_prefix,
                    partition_col=partition_col,
                    partition_value=v,
                )
                done_key = f"{pfx}_{UDTF_PARTITION_DONE_MARKER}"
                if done_key in checkpoint_store:
                    all_completed_prefixes.append(pfx)
        else:
            pfx = format_udtf_partition_prefix(
                top_prefix, partition_col=None, partition_value=None
            )
            done_key = f"{pfx}_{UDTF_PARTITION_DONE_MARKER}"
            if done_key in checkpoint_store:
                all_completed_prefixes.append(pfx)

        # --- Fragment-per-partition final write ---
        # Write each partition as its own Lance fragment so that only one
        # partition's batches are in driver memory at a time (O(largest
        # partition) instead of O(total output)).
        from geneva.namespace import get_storage_options_provider

        _, storage_options = get_storage_options_provider(
            table_ref.namespace_impl,
            table_ref.namespace_properties,
            table_ref.table_id,
        )

        new_metadata: dict[str | bytes, str | bytes] = dict(metadata)  # type: ignore[arg-type]
        new_metadata[MATVIEW_META_BASE_VERSION.encode()] = str(
            actual_source_version
        ).encode()
        new_schema = udtf_obj.output_schema.with_metadata(new_metadata)

        first_write = True
        bar_assembly = tqdm(
            total=len(all_completed_prefixes),
            desc=fmt("udtf | assembling fragments", Colors.CYAN, bold=True),
            unit="frag",
        )
        for pfx in all_completed_prefixes:
            batches = _load_partition_batches(checkpoint_store, pfx)
            if not batches:
                bar_assembly.update(1)
                continue
            partition_data = pa.Table.from_batches(batches, schema=new_schema)
            lance.write_dataset(
                partition_data,
                dest_table_uri,
                mode="overwrite" if first_write else "append",
                storage_options=storage_options,
            )
            first_write = False
            bar_assembly.update(1)
        bar_assembly.close()

        if first_write:
            # No partitions produced data — write empty table with metadata
            lance.write_dataset(
                new_schema.empty_table(),
                dest_table_uri,
                mode="overwrite",
                storage_options=storage_options,
            )

    def backfill_async(
        self,
        col_name: str,
        *,
        udf: UDF | None = None,
        where: str | None = None,
        concurrency: int = 8,
        intra_applier_concurrency: int = 1,
        _admission_check: bool | None = None,
        _admission_strict: bool | None = None,
        min_checkpoint_size: int | None = None,
        max_checkpoint_size: int | None = None,
        batch_checkpoint_flush_interval_seconds: float | None = None,
        _enable_job_tracker_saves: bool = True,
        **kwargs,
    ) -> JobFuture:
        """
        Backfills the specified column asynchronously.

        Returns job future. Call .result() to wait for completion.

        If run without a context, a warning is emitted and local Ray is used
        implicitly. For proper cleanup, wrap in conn.local_ray_context().

        Parameters
        ----------
        col_name: str
            Target column name to backfill
        udf: UDF | None
            Optionally override the UDF used to backfill the column.
        where: str | None
            SQL expression filter to select rows to backfill. Defaults to
            '<col_name> IS NULL' to skip already-computed rows. Use where="1=1"
            to force reprocessing all rows.
        concurrency: int
            (default = 8) This controls the number of processes that tasks run
            concurrently. For max throughput, ideally this is larger than the number
            of nodes in the k8s cluster.   This is the number of Ray actor processes
            are started.
        intra_applier_concurrency: int
            (default = 1) This controls the number of threads used to execute tasks
            within a process. Multiplying this times `concurrency` roughly corresponds
            to the number of cpu's being used.
        _admission_check: bool | None
            Whether to run admission control to validate cluster resources before
            starting the job. If None, uses GENEVA_ADMISSION__CHECK env var
            (default: true). Set to False to skip the check.
            **Experimental**: Parameters starting with `_` are subject to change.
        _admission_strict: bool | None
            If True, raises ResourcesUnavailableError when resources are
            insufficient. If False, logs a warning but allows the job to proceed.
            If None, uses GENEVA_ADMISSION__STRICT env var (default: true).
            **Experimental**: Parameters starting with `_` are subject to change.
        commit_granularity: int | None
            (default = 64) Show a partial result everytime this number of fragments
            are completed. If None, the entire result is committed at once.
        read_version: int | None
            (default = None) The version of the table to read from.  If None, the
            latest version is used.
        task_shuffle_diversity: int | None
            (default = 8) ??
        batch_size: int | None (deprecated)
            (default = 10240) Legacy alias for checkpoint_size. Prefer checkpoint_size.
        checkpoint_size: int | None
            The max number of rows per checkpoint.
            This influences how often progress and proof of life is presented.
            When adaptive sizing is enabled, an explicit checkpoint_size seeds the
            initial checkpoint size; otherwise the initial size defaults to
            min_checkpoint_size.
        min_checkpoint_size: int | None
            Minimum adaptive checkpoint size (lower bound).
        max_checkpoint_size: int | None
            Maximum adaptive checkpoint size (upper bound). This also caps the
            largest read batch and thus the maximum memory footprint per batch.
        batch_checkpoint_flush_interval_seconds: float | None
            Controls how frequently in-progress results are persisted as batch
            checkpoints. Larger values usually improve throughput, but if the
            job stops unexpectedly, more recently computed work may need to be
            redone. Smaller values checkpoint progress sooner, which improves
            durability and resume behavior, but can reduce throughput. Set to
            ``0`` to persist every batch as soon as it is produced.
        task_size: int | None
            Controls read-task sizing (rows per worker task). Defaults to
            ``table.count_rows() // num_workers // 2`` when omitted.
        num_frags: int | None
            (default = None) The number of table fragments to process.  If None,
            process all fragments.
        _enable_job_tracker_saves: bool
            (default = False) Experimentally enable persistence of job metrics to the
            database. When disabled, metrics are tracked in-memory only.
        """
        if get_current_context() is None:
            _LOG.warning(_IMPLICIT_LOCAL_RAY_WARNING, "backfill_async()")

        from geneva.apply.utils import (
            _any_checkpoint_has_srcfiles_mismatch,
            _any_checkpoint_has_udf_mismatch,
        )
        from geneva.runners.ray.pipeline import (
            dispatch_run_ray_add_column,
            fetch_udf,
            validate_backfill_args,
        )

        if min_checkpoint_size is not None:
            kwargs["min_checkpoint_size"] = min_checkpoint_size
        if max_checkpoint_size is not None:
            kwargs["max_checkpoint_size"] = max_checkpoint_size
        if batch_checkpoint_flush_interval_seconds is not None:
            kwargs["batch_checkpoint_flush_interval_seconds"] = (
                batch_checkpoint_flush_interval_seconds
            )

        self._normalize_backfill_batch_kwargs(kwargs)

        read_version = kwargs.get("read_version")
        if read_version is None:
            read_version = self.version
            kwargs["read_version"] = read_version

        validate_backfill_args(self, col_name, udf, read_version=read_version)

        # Get UDF for version check and admission control
        current_udf = udf
        if current_udf is None:
            try:
                udf_spec = fetch_udf(self, col_name)
                current_udf = self._conn._packager.unmarshal(udf_spec)
            except Exception as e:
                _LOG.debug("Could not fetch UDF: %s", e)
                current_udf = None

        # Check for UDF version or srcfiles mismatch (always, even with explicit where)
        # This allows us to warn users when they provide a filter that may skip rows
        # that should be recomputed due to UDF or input data changes.
        col_schema = self._native_tbl.schema
        col_field = col_schema.field(col_name)
        udf_mismatch = False
        srcfiles_mismatch = False

        if current_udf is not None:
            udf_name = _get_udf_name_from_field(col_field)
            udf_version = current_udf.version
            if udf_name and udf_version:
                try:
                    checkpoint_store = self.get_reference().open_checkpoint_store()
                    udf_mismatch = _any_checkpoint_has_udf_mismatch(
                        udf_name, col_name, checkpoint_store, udf_version
                    )
                except Exception as e:
                    _LOG.debug("Error checking checkpoint UDF versions: %s", e)

            # Check for srcfiles mismatch (indicates input column data changed)
            # This handles cases like: column c depends on b, b is re-backfilled
            if not udf_mismatch and current_udf.input_columns:
                try:
                    import lance

                    from geneva.runners.ray.pipeline import _get_relevant_field_ids

                    checkpoint_store = self.get_reference().open_checkpoint_store()
                    dataset = lance.dataset(self.uri, version=read_version)
                    input_field_ids = _get_relevant_field_ids(
                        dataset, current_udf.input_columns
                    )
                    if input_field_ids:
                        srcfiles_mismatch = _any_checkpoint_has_srcfiles_mismatch(
                            col_name,
                            checkpoint_store,
                            dataset,
                            input_field_ids,
                        )
                except Exception as e:
                    _LOG.debug("Error checking srcfiles hash: %s", e)

        # Handle where filter based on mismatch detection
        has_mismatch = udf_mismatch or srcfiles_mismatch
        user_provided_where = where is not None

        if user_provided_where and has_mismatch:
            # User provided explicit filter but UDF/input data changed
            mismatch_type = "UDF version" if udf_mismatch else "input column data"
            _LOG.warning(
                "Column %s has %s changes but explicit where filter provided. "
                "Some rows computed with old UDF/data may not be reprocessed. "
                "Use where='1=1' to force reprocessing all rows.",
                col_name,
                mismatch_type,
            )
        elif where is None:
            # Default to filtering for NULL values to skip already-computed rows.
            # We skip the filter if:
            # 1. UDF/srcfiles mismatch detected - need to recompute all rows
            # 2. Struct column type (IS NULL doesn't work for structs with NULL fields)
            if has_mismatch:
                mismatch_type = "UDF version" if udf_mismatch else "input column data"
                _LOG.info(
                    "%s changed for column %s, processing all rows",
                    mismatch_type.capitalize(),
                    col_name,
                )
            elif col_field is not None and pa.types.is_struct(col_field.type):
                # Struct columns: can't use IS NULL filter effectively because
                # a struct with NULL fields is not the same as a NULL struct.
                # Just process all rows (less efficient, but correct).
                pass
            else:
                where = f"{col_name} IS NULL"

        # Admission control: validate cluster has sufficient resources
        # Always call validate_admission when there's a UDF - it handles check=None
        # by reading from config (GENEVA_ADMISSION__CHECK env var, default True)
        from geneva.runners.ray.admission import validate_admission

        if current_udf is not None:
            validate_admission(
                current_udf,
                concurrency=concurrency,
                intra_applier_concurrency=intra_applier_concurrency,
                check=_admission_check,
                strict=_admission_strict,
            )

        fut = dispatch_run_ray_add_column(
            self.get_reference(),
            col_name,
            udf=udf,
            where=where,
            concurrency=concurrency,
            intra_applier_concurrency=intra_applier_concurrency,
            enable_job_tracker_saves=_enable_job_tracker_saves,
            **kwargs,
        )
        return fut

    def backfill(
        self,
        col_name,
        *,
        udf: UDF | None = None,
        where: str | None = None,
        concurrency: int = 8,
        intra_applier_concurrency: int = 1,
        _admission_check: bool | None = None,
        _admission_strict: bool | None = None,
        refresh_status_secs: float = 2.0,
        min_checkpoint_size: int | None = None,
        max_checkpoint_size: int | None = None,
        batch_checkpoint_flush_interval_seconds: float | None = None,
        _enable_job_tracker_saves: bool = True,
        **kwargs,
    ) -> str:
        """
        Backfills the specified column.

        Returns job_id string

        If run without a context, a warning is emitted and local Ray is used
        implicitly. For proper cleanup, wrap in conn.local_ray_context().

        Parameters
        ----------
        col_name: str
            Target column name to backfill
        udf: UDF | None
            Optionally override the UDF used to backfill the column.
        where: str | None
            SQL expression filter to select rows to backfill. Defaults to
            '<col_name> IS NULL' to skip already-computed rows. Use where="1=1"
            to force reprocessing all rows.
        concurrency: int
            (default = 8) This controls the number of processes that tasks run
            concurrently. For max throughput, ideally this is larger than the number
            of nodes in the k8s cluster.   This is the number of Ray actor processes
            are started.
        intra_applier_concurrency: int
            (default = 1) This controls the number of threads used to execute tasks
            within a process. Multiplying this times `concurrency` roughly corresponds
            to the number of cpu's being used.
        _admission_check: bool | None
            Whether to run admission control to validate cluster resources before
            starting the job. If None, uses GENEVA_ADMISSION__CHECK env var
            (default: true). Set to False to skip the check.
            **Experimental**: Parameters starting with `_` are subject to change.
        _admission_strict: bool | None
            If True, raises ResourcesUnavailableError when resources are
            insufficient. If False, logs a warning but allows the job to proceed.
            If None, uses GENEVA_ADMISSION__STRICT env var (default: true).
            **Experimental**: Parameters starting with `_` are subject to change.
        commit_granularity: int | None
            (default = 64) Show a partial result everytime this number of fragments
            are completed. If None, the entire result is committed at once.
        read_version: int | None
            (default = None) The version of the table to read from.  If None, the
            latest version is used.
        task_shuffle_diversity: int | None
            (default = 8) ??
        batch_size: int | None (deprecated)
            (default = 100) Legacy alias for checkpoint_size. Prefer checkpoint_size.
            If 0, the batch will be the total number of rows from a fragment.
        checkpoint_size: int | None
            The max number of rows per checkpoint.
            This influences how often progress and proof of life is presented.
            When adaptive sizing is enabled, an explicit checkpoint_size seeds the
            initial checkpoint size; otherwise the initial size defaults to
            min_checkpoint_size.
        min_checkpoint_size: int | None
            Minimum adaptive checkpoint size (lower bound).
        max_checkpoint_size: int | None
            Maximum adaptive checkpoint size (upper bound). This also caps the
            largest read batch and thus the maximum memory footprint per batch.
        batch_checkpoint_flush_interval_seconds: float | None
            Controls how frequently in-progress results are persisted as batch
            checkpoints. Larger values usually improve throughput, but if the
            job stops unexpectedly, more recently computed work may need to be
            redone. Smaller values checkpoint progress sooner, which improves
            durability and resume behavior, but can reduce throughput. Set to
            ``0`` to persist every batch as soon as it is produced.
        task_size: int | None
            Controls read-task sizing (rows per worker task). Defaults to
            ``table.count_rows() // num_workers // 2`` when omitted.
        num_frags: int | None
            (default = None) The number of table fragments to process.  If None,
            process all fragments.
        _enable_job_tracker_saves: bool
            (default = False) Experimentally enable persistence of job metrics to the
            database. When disabled, metrics are tracked in-memory only.
        """
        if get_current_context() is None:
            _LOG.warning(_IMPLICIT_LOCAL_RAY_WARNING, "backfill()")

        # Input validation
        from geneva.runners.ray.pipeline import validate_backfill_args

        if min_checkpoint_size is not None:
            kwargs["min_checkpoint_size"] = min_checkpoint_size
        if max_checkpoint_size is not None:
            kwargs["max_checkpoint_size"] = max_checkpoint_size
        if batch_checkpoint_flush_interval_seconds is not None:
            kwargs["batch_checkpoint_flush_interval_seconds"] = (
                batch_checkpoint_flush_interval_seconds
            )

        self._normalize_backfill_batch_kwargs(kwargs)

        read_version = kwargs.get("read_version")
        if read_version is None:
            read_version = self.version
            kwargs["read_version"] = read_version

        validate_backfill_args(self, col_name, udf, read_version=read_version)

        # get cluster status
        from geneva.runners.ray.raycluster import ClusterStatus

        cs = ClusterStatus()
        try:
            with status_updates(cs.get_status, refresh_status_secs):
                # Kick off the job
                fut = self.backfill_async(
                    col_name,
                    udf=udf,
                    where=where,
                    concurrency=concurrency,
                    intra_applier_concurrency=intra_applier_concurrency,
                    _admission_check=_admission_check,
                    _admission_strict=_admission_strict,
                    batch_checkpoint_flush_interval_seconds=(
                        batch_checkpoint_flush_interval_seconds
                    ),
                    _enable_job_tracker_saves=_enable_job_tracker_saves,
                    **kwargs,
                )

            while not fut.done(timeout=refresh_status_secs):
                # wait for the backfill to complete, updating statuses
                cs.get_status()
                fut.status()

            cs.get_status()
            fut.status()

            # Check for errors - this will raise if the job failed
            fut.result()

            # updates came from an external writer, so get the latest version.
            self._native_tbl.checkout_latest()
            return fut.job_id
        finally:
            with contextlib.suppress(Exception):
                cs.close()

    @staticmethod
    def _normalize_backfill_batch_kwargs(kwargs: dict[str, Any]) -> None:
        """Normalize batch-size kwargs for backfill calls."""

        checkpoint_size = kwargs.pop("checkpoint_size", None)
        batch_size = kwargs.pop("batch_size", None)
        task_size = kwargs.pop("task_size", None)

        resolved = resolve_batch_size(
            batch_size=batch_size,
            checkpoint_size=checkpoint_size,
        )

        if task_size is not None:
            kwargs["task_size"] = task_size

        kwargs["checkpoint_size"] = resolved

    def alter_columns(self, *alterations: dict[str, Any], **kwargs) -> None:
        """
        Alter columns in the table.  This can change the computed columns' udf

        Parameters
        ----------
        alterations:  Iterable[dict[str, Any]]
            This is a list of alterations to apply to the table.

        Examples
        --------

            table.alter_columns(
                { "path": "col1", "udf": col1_udf_v2, },
                { "path": "col2", "udf": col2_udf})

        """
        basic_column_alterations = []
        for alter in alterations:
            if "path" not in alter:
                raise ValueError("path is required to alter computed column's udf")

            if "virtual_column" in alter:  # deprecated
                udf = alter.get("virtual_column")
                if not isinstance(udf, UDF):
                    raise ValueError("virtual_column must be a UDF")
                _LOG.warning(
                    "alter_columns 'virtual_column' is deprecated, use 'udf' instead."
                )
            elif "udf" in alter:
                udf = alter.get("udf")
                if not isinstance(udf, UDF):
                    raise ValueError("udf must be a UDF")
            else:
                basic_column_alterations.append(alter)
                continue

            col_name = alter["path"]

            input_cols = alter.get("input_columns", None)
            if input_cols is None:
                input_cols = udf.input_columns

            self._configure_computed_column(col_name, udf, input_cols)

        if len(basic_column_alterations) > 0:
            self._native_tbl.alter_columns(*basic_column_alterations)

    def _upload_udf(
        self,
        udf_payload: bytes,
        udf_location: str,
    ) -> None:
        """
        Upload UDF payload to the dataset 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
        ----------
        udf_payload : bytes
            The serialized UDF package to upload
        udf_location : str
            Relative path within the dataset (e.g., "_udfs/checksum")
        """
        import os
        import tempfile

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

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

        try:
            session.upload_file(tmp_path, udf_location)
            _LOG.debug(f"Uploaded UDF to {udf_location} ({len(udf_payload)} bytes)")
        except Exception as e:
            raise RuntimeError(f"Failed to upload UDF at {udf_location}: {e}") from e
        finally:
            os.unlink(tmp_path)

    def _configure_computed_column(
        self,
        col_name: str,
        udf: UDF,
        input_cols: list[str] | None,
    ) -> None:
        """
        Configure a column to be a computed column for the given UDF.

        This procedure includes:
        - Packaging the UDF
        - Uploading the UDF to the dataset
        - Updating the field metadata to include the UDF information

        Note that the column should already exist on the table.
        """
        # record batch udf's don't specify inputs
        if (
            udf.arg_type != UDFArgType.RECORD_BATCH
            and udf.input_columns
            and col_name in udf.input_columns
        ):
            raise ValueError(
                f"UDF output column {col_name} is not allowed to be in"
                f" input {udf.input_columns}"
            )

        udf_spec = self._conn._packager.marshal(udf, table_ref=self.get_reference())

        # upload the UDF to the dataset URL
        if not isinstance(self._native_tbl, LanceLocalTable):
            raise TypeError(
                "adding udf column is currently only supported for local tables"
            )

        # upload the packaged UDF to some location inside the dataset:
        checksum = hashlib.sha256(udf_spec.udf_payload).hexdigest()
        udf_location = f"_udfs/{checksum}"

        self._upload_udf(udf_spec.udf_payload, udf_location)

        # TODO rename this from virtual_column to computed column
        field_metadata = udf.field_metadata | {
            "virtual_column": "true",
            "virtual_column.udf_backend": udf_spec.backend,
            "virtual_column.udf_name": udf_spec.name,
            "virtual_column.udf": "_udfs/" + checksum,
            "virtual_column.udf_inputs": json.dumps(input_cols),
            "virtual_column.platform.system": platform.system(),
            "virtual_column.platform.arch": platform.machine(),
            "virtual_column.platform.python_version": platform.python_version(),
            "virtual_column.auto_backfill": "true" if udf.auto_backfill else "false",
        }

        # Add the column metadata:
        self._native_tbl.replace_field_metadata(col_name, field_metadata)

    def create_index(
        self,
        metric: str = "L2",
        num_partitions: int | None = None,
        num_sub_vectors: int | None = None,
        vector_column_name: str = VECTOR_COLUMN_NAME,
        replace: bool = True,
        accelerator=None,
        index_cache_size=None,
        *,
        index_type: Literal[
            "IVF_FLAT",
            "IVF_PQ",
            "IVF_HNSW_SQ",
            "IVF_HNSW_PQ",
        ] = "IVF_PQ",
        num_bits: int = 8,
        max_iterations: int = 50,
        sample_rate: int = 256,
        m: int = 20,
        ef_construction: int = 300,
    ) -> None:
        """Create Vector Index"""
        self._ltbl.create_index(
            metric,
            num_partitions or 256,
            num_sub_vectors or 96,
            vector_column_name,
            replace,
            accelerator,
            index_cache_size,
            index_type=index_type,
            num_bits=num_bits,
            max_iterations=max_iterations,
            sample_rate=sample_rate,
            m=m,
            ef_construction=ef_construction,
        )

    @override
    def create_fts_index(
        self,
        field_names: str | list[str],
        *,
        ordering_field_names: str | list[str] | None = None,
        replace: bool = False,
        writer_heap_size: int | None = None,
        tokenizer_name: str | None = None,
        with_position: bool = True,
        base_tokenizer: Literal["simple", "raw", "whitespace"] = "simple",
        language: str = "English",
        max_token_length: int | None = 40,
        lower_case: bool = True,
        stem: bool = False,
        remove_stop_words: bool = False,
        ascii_folding: bool = False,
        **_kwargs,
    ) -> None:
        self._ltbl.create_fts_index(
            field_names,
            ordering_field_names=ordering_field_names,
            replace=replace,
            writer_heap_size=writer_heap_size,
            tokenizer_name=tokenizer_name,
            with_position=with_position,
            base_tokenizer=base_tokenizer,
            language=language,
            max_token_length=max_token_length,
            lower_case=lower_case,
            stem=stem,
            remove_stop_words=remove_stop_words,
            ascii_folding=ascii_folding,
            use_tantivy=False,
        )

    @override
    def create_scalar_index(
        self,
        column: str,
        *,
        replace: bool = True,
        index_type: Literal["BTREE", "BITMAP", "LABEL_LIST"] = "BTREE",
    ) -> None:
        self._ltbl.create_scalar_index(
            column,
            replace=replace,
            index_type=index_type,
        )

    @override
    def _do_merge(
        self,
        merge: LanceMergeInsertBuilder,
        new_data: DATA,
        on_bad_vectors: OnBadVectorsType,
        fill_value: float,
    ) -> MergeResult:
        return self._ltbl._do_merge(merge, new_data, on_bad_vectors, fill_value)

    @override
    def _execute_query(
        self,
        query: LanceQuery,
        batch_size: int | None = None,
    ) -> pa.RecordBatchReader:
        return self._ltbl._execute_query(query, batch_size=batch_size)

    def list_versions(self) -> list[dict[str, Any]]:
        return self._ltbl.list_versions()

    @override
    def cleanup_old_versions(
        self,
        older_than: timedelta | None = None,
        *,
        delete_unverified=False,
    ) -> Any:  # lance.CleanupStats not available in type stubs
        return self._ltbl.cleanup_old_versions(
            older_than,
            delete_unverified=delete_unverified,
        )

    def to_batches(self, batch_size: int | None = None) -> Iterator[pa.RecordBatch]:
        from .query import Query

        if isinstance(self._ltbl, Query):
            return self._ltbl.to_batches(batch_size)  # type: ignore[attr-defined]
        return self.to_lance().to_batches(batch_size=batch_size)  # type: ignore[arg-type]

    """This is the signature for the standard LanceDB table.search call"""

    def search(  # type: ignore[override]
        self,
        query: list | pa.Array | pa.ChunkedArray | np.ndarray | None = None,
        vector_column_name: str | None = None,
        query_type: Literal["vector", "fts", "hybrid", "auto"] = "auto",
        ordering_field_name: str | None = None,
        fts_columns: str | list[str] | None = None,
    ) -> GenevaQueryBuilder | LanceQueryBuilder:
        if query is None:
            return GenevaQueryBuilder(self)
        else:
            return self._ltbl.search(
                query, vector_column_name, query_type, ordering_field_name, fts_columns
            )

    @override
    def drop_columns(self, columns: Iterable[str]) -> None:
        self._ltbl.drop_columns(columns)

    @override
    def to_arrow(self) -> pa.Table:
        return self._native_tbl.to_arrow()

    @override
    def count_rows(self, filter: str | None = None) -> int:
        return self._ltbl.count_rows(filter)

    @override
    def update(
        self,
        where: str | None = None,
        values: dict | None = None,
        *,
        values_sql: dict[str, str] | None = None,
    ) -> None:
        self._ltbl.update(where, values, values_sql=values_sql)

    @override
    def delete(self, where: str) -> None:
        self._ltbl.delete(where)

    @override
    def list_indices(self) -> Iterable[IndexConfig]:
        return self._ltbl.list_indices()

    @override
    def index_stats(self, index_name: str) -> IndexStatistics | None:
        return self._ltbl.index_stats(index_name)

    @override
    def optimize(
        self,
        *,
        cleanup_older_than: timedelta | None = None,
        delete_unverified: bool = False,
    ) -> None:
        return self._ltbl.optimize(
            cleanup_older_than=cleanup_older_than,
            delete_unverified=delete_unverified,
        )

    @override
    def compact_files(self) -> None:
        self._ltbl.compact_files()

    @override
    def restore(self, *args, **kwargs) -> None:
        self._ltbl.restore(*args, **kwargs)

    # TODO: This annotation sucks
    # NOTE: When using blob columns with stable row IDs enabled (e.g., for
    # materialized views), pylance >= 1.1.0b2 is required. Earlier versions
    # have a bug where take_blobs fails on fragments created via DataReplacement.
    def take_blobs(self, indices: list[int] | pa.Array, column: str):  # noqa: ANN201
        return self.to_lance().take_blobs(blob_column=column, indices=indices)

    def to_lance(self) -> lance.LanceDataset:
        return self._native_tbl.to_lance()  # type: ignore[attr-defined]

    def uses_v2_manifest_paths(self) -> bool:
        return self._ltbl.uses_v2_manifest_paths()

    def migrate_v2_manifest_paths(self) -> None:
        return self._ltbl.migrate_v2_manifest_paths()

    def _analyze_plan(self, query: LanceQuery) -> str:
        return self._ltbl._analyze_plan(query)

    def _explain_plan(self, query: LanceQuery, verbose: bool | None = False) -> str:
        return self._ltbl._explain_plan(query, verbose=verbose)

    def stats(self) -> TableStatistics:
        return self._ltbl.stats()

    @property
    def tags(self) -> Tags:
        return self._ltbl.tags

    def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder:
        return self._ltbl.take_offsets(offsets)

    def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder:
        return self._ltbl.take_row_ids(row_ids)

    def get_errors(
        self,
        job_id: str | None = None,
        column_name: str | None = None,
        error_type: str | None = None,
    ) -> list[Any]:
        """Get error records for this table.

        Parameters
        ----------
        job_id : str, optional
            Filter errors by job ID
        column_name : str, optional
            Filter errors by column name
        error_type : str, optional
            Filter errors by exception type

        Returns
        -------
        list[ErrorRecord]
            List of error records matching the filters

        Examples
        --------
        >>> # Get all errors for this table
        >>> errors = table.get_errors()
        >>>
        >>> # Get errors for a specific job
        >>> errors = table.get_errors(job_id="abc123")
        >>>
        >>> # Get errors for a specific column
        >>> errors = table.get_errors(column_name="my_column")
        """
        from geneva.debug.error_store import ErrorStore

        error_store = ErrorStore(self._conn)
        return error_store.get_errors(
            job_id=job_id,
            table_name=self._name,
            column_name=column_name,
            error_type=error_type,
        )

    def get_failed_row_addresses(self, job_id: str, column_name: str) -> list[int]:
        """Get row addresses for all failed rows in a job.

        Parameters
        ----------
        job_id : str
            Job ID to query
        column_name : str
            Column name to filter by

        Returns
        -------
        list[int]
            List of row addresses that failed

        Examples
        --------

            # Get failed row addresses
            failed_rows = table.get_failed_row_addresses(
                job_id="abc123", column_name="my_col"
            )

            # Retry processing only failed rows
            row_ids = ','.join(map(str, failed_rows))
            table.backfill("my_col", where=f"_rowaddr IN ({row_ids})")
        """
        from geneva.debug.error_store import ErrorStore

        error_store = ErrorStore(self._conn)
        return error_store.get_failed_row_addresses(
            job_id=job_id, column_name=column_name
        )

    @override
    def _output_schema(self, query: LanceQuery) -> pa.Schema:
        return self._ltbl._output_schema(query)
