fix(python): improve function failure diagnostics (#4790)

## Summary

- Make `AsyncJob.wait(timeout)` and `Job.wait(timeout)` raise a
descriptive `TimeoutError` containing the job ID and elapsed timeout.
Registry jobs explain that work continues on the server and include the
`db.open_job(id).wait()` reattach call. SQL query handles instead say to
keep the original handle because query IDs cannot be reopened through
the registry jobs API.
- Render `JobFailureInfo` and `FunctionErrorRecord` with Python values
instead of Rust `Option` wrappers.
- Decode and expose the server's `row_address` on `FunctionErrorRecord`,
and document how to retry recorded rows with `REFRESH COLUMN ... WHERE
_rowaddr IN (...)` before compaction or rewrite.
- Warn when an import inside a packaged Function body resolves to
undeclared local source. The `udf` documentation explains why arbitrary
base-image coverage and distribution-to-import-name mapping cannot be
validated reliably by the SDK.

`row_address` is optional, so records from older servers still decode
with `None`. Rust code that constructs `FunctionErrorRecord` literals
must now supply this new field.

Part of ENT-3448 and ENT-3442. This consumes the server-side diagnostics
from lancedb/sophon#9509 and logical row addresses from
lancedb/sophon#9512.

## Validation

- Rust `function_errors` remote decoding and
legacy-missing-`row_address` tests.
- Rust/PyO3 clippy for `lancedb` and `lancedb-python`.
- Python extension build, pyright, and ruff.
- Targeted Python job, Function packaging, Function error, and job
detail tests: 62 passed, plus timeout-handle tests: 3 passed.

python/python/lancedb/_lancedb.pyi

@@ -947,6 +947,7 @@ class FunctionErrorRecord:
     job_id: str
     fragment_id: int
     row_offset: Optional[int]
+    row_address: Optional[int]
     column: str
     function: str
     function_version: str

python/python/lancedb/functions.py

@@ -16,7 +16,7 @@ import builtins
 import base64
 import functools
 import hashlib
-import importlib
+import importlib.util
 import inspect
 import symtable
 import json
@@ -1464,6 +1464,28 @@ def _module_locations(module: types.ModuleType) -> tuple[str, ...]:
     return tuple(os.path.realpath(path) for path in getattr(module, "__path__", ()))
 
 
+def _module_locations_for_name(module_name: str) -> tuple[str, ...]:
+    """Where a top-level module would be imported from, without importing it
+    when it is not already loaded."""
+    import os
+
+    top = module_name.partition(".")[0]
+    module = sys.modules.get(top)
+    if module is not None:
+        return _module_locations(module)
+    try:
+        spec = importlib.util.find_spec(top)
+    except (AttributeError, ImportError, ValueError):
+        return ()
+    if spec is None:
+        return ()
+    if spec.submodule_search_locations:
+        return tuple(os.path.realpath(path) for path in spec.submodule_search_locations)
+    if spec.origin and spec.origin not in {"built-in", "frozen"}:
+        return (os.path.realpath(spec.origin),)
+    return ()
+
+
 class _SourcePackager:
     """Builds the entry module of a Function artifact.
 
@@ -1480,6 +1502,7 @@ class _SourcePackager:
         self._definitions: list[str] = []
         self._packaged: set[int] = set()
         self._names: dict[str, Any] = {}
+        self._warned_imports: set[str] = set()
 
     def package(self, target: Any, signature: FunctionSignature) -> str:
         entry = self._definition(target, entry=True, signature=signature)
@@ -1512,6 +1535,7 @@ class _SourcePackager:
             raise ValueError(
                 f"@udf cannot package dynamic namespace access: {dynamic!r}"
             )
+        self._warn_unshipped_direct_imports(definition)
         self._packaged.add(id(target))
         self._claim(target.__name__, ("definition", id(target)))
         if reader_name is not None:
@@ -1652,13 +1676,14 @@ class _SourcePackager:
             or top in sys.stdlib_module_names
             or top in sys.builtin_module_names
             or _normalized_distribution(top) in self._requirements
+            or top in self._warned_imports
         ):
             return
-        module = sys.modules.get(top)
-        locations = () if module is None else _module_locations(module)
+        locations = _module_locations_for_name(top)
         roots = _installed_roots()
         if not locations or any(location.startswith(roots) for location in locations):
             return
+        self._warned_imports.add(top)
         warnings.warn(
             f"@udf source imports {module_name!r} from {locations[0]}, a local "
             "module the Function's environment may not be able to import; ship "
@@ -1667,6 +1692,22 @@ class _SourcePackager:
             UserWarning,
         )
 
+    def _warn_unshipped_direct_imports(
+        self, definition: Union[ast.FunctionDef, ast.ClassDef]
+    ) -> None:
+        """Warn for body imports that resolve to local, undeclared source.
+
+        Imports from site-packages are deliberately not guessed about: the
+        declared distribution name need not match the imported module, and the
+        client cannot inspect the remote base image.
+        """
+        for node in ast.walk(definition):
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    self._require_importable(alias.name)
+            elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module:
+                self._require_importable(node.module)
+
 
 def _code_files(modules: Sequence[Any]) -> tuple[dict[str, str], frozenset[str]]:
     """The Python source files of the modules and packages in ``code=``,
@@ -2147,9 +2188,15 @@ def udf(
       imported normally, so helpers shared by several Functions live in one
       place.
 
-    A reference to a module in a local source tree that is neither shipped
-    with ``code`` nor named by a declared package draws a warning at
-    registration, because the worker will most likely fail to import it.
+    A module imported by the packaged source -- including a direct ``import``
+    inside a function or method body -- draws a warning at registration when
+    it resolves to local source outside ``code`` and no declared requirement
+    matches the imported module name. The client cannot reliably validate
+    imports provided only by the base image: ``pip``/``conda`` distribution
+    names need not match the modules they install, and the client does not
+    know the contents of a built or custom image. Declare every third-party
+    import in ``pip`` or ``conda`` even when it happens to be installed on the
+    registration host.
     Closures, lambdas, nested definitions, mutable module-level
     objects, code that reaches the module namespace another way
     (``globals()``/``eval``, ``sys.modules``, ``builtins``), and a

python/python/lancedb/job.py

@@ -26,6 +26,28 @@ __all__ = [
 ]
 
 
+def _timeout_message(
+    job_id: Optional[str], timeout: timedelta, job_type: Optional[str]
+) -> str:
+    seconds = f"{timeout.total_seconds():g}"
+    unit = "second" if seconds == "1" else "seconds"
+    if job_id is None:
+        return (
+            f"operation still running after {seconds} {unit}; "
+            "it continues on the server"
+        )
+    if job_type == "query":
+        return (
+            f"query {job_id!r} still running after {seconds} {unit}; "
+            "it continues on the server; keep this handle to continue waiting"
+        )
+    return (
+        f"job {job_id!r} still running after {seconds} {unit}; "
+        "it continues on the server; "
+        f"reattach with db.open_job({job_id!r}).wait()"
+    )
+
+
 class AsyncJob(Generic[T]):
     """A handle to an operation that may still be running.
 
@@ -73,13 +95,24 @@ class AsyncJob(Generic[T]):
 
         Raises `JobFailedError` if the operation failed, `JobCancelledError`
         if it was cancelled, and `TimeoutError` if `timeout` elapses first.
+        The timeout error identifies the job and explains that it continues on
+        the server. Registry jobs include the ``db.open_job(id).wait()``
+        reattach call; SQL query handles must be kept because their IDs cannot
+        be reopened through the registry jobs API.
         """
         if self._inner is None:
             return cast(T, None)
         if timeout is None:
             result = await self._inner.wait()
         else:
-            result = await asyncio.wait_for(self._inner.wait(), timeout.total_seconds())
+            try:
+                result = await asyncio.wait_for(
+                    self._inner.wait(), timeout.total_seconds()
+                )
+            except asyncio.TimeoutError as error:
+                raise TimeoutError(
+                    _timeout_message(self.id, timeout, self.job_type)
+                ) from error
         if self._result_decoder is not None:
             return self._result_decoder(result)
         return cast(T, result)
@@ -270,6 +303,10 @@ class Job(Generic[T]):
 
         Raises `JobFailedError` if the operation failed, `JobCancelledError`
         if it was cancelled, and `TimeoutError` if `timeout` elapses first.
+        The timeout error identifies the job and explains that it continues on
+        the server. Registry jobs include the ``db.open_job(id).wait()``
+        reattach call; SQL query handles must be kept because their IDs cannot
+        be reopened through the registry jobs API.
         """
         if self._inner is None:
             return cast(T, None)

python/python/lancedb/table.py

@@ -3191,6 +3191,30 @@ class Table(ABC):
             ``records``, ``fragments`` and ``truncated``, the last saying
             whether the listing stopped at its limit.
 
+        Each record identifies its row with ``fragment_id``, ``row_offset``
+        and the physical ``row_address`` (``(fragment_id << 32) |
+        row_offset``). To retry the recorded rows, collect the non-null
+        addresses before the table is compacted or rewritten and run the
+        SQL refresh:
+
+        .. code-block:: python
+
+            errors = table.function_errors(column="embedding")
+            addresses = sorted(
+                record.row_address
+                for record in errors.records
+                if record.row_address is not None
+            )
+            if addresses:
+                address_list = ", ".join(str(address) for address in addresses)
+                database.execute_query(
+                    f"REFRESH COLUMN embedding ON {table.name()} "
+                    f"WHERE _rowaddr IN ({address_list})"
+                )
+
+        Records whose per-row detail was capped have no address and appear
+        only in ``fragments``; they cannot be selected individually here.
+
         Examples
         --------
         >>> errors = table.function_errors(column="embedding")  # doctest: +SKIP
@@ -7515,6 +7539,30 @@ class AsyncTable:
             ``records``, ``fragments`` and ``truncated``, the last saying
             whether the listing stopped at its limit.
 
+        Each record identifies its row with ``fragment_id``, ``row_offset``
+        and the physical ``row_address`` (``(fragment_id << 32) |
+        row_offset``). To retry the recorded rows, collect the non-null
+        addresses before the table is compacted or rewritten and run the
+        SQL refresh:
+
+        .. code-block:: python
+
+            errors = await table.function_errors(column="embedding")
+            addresses = sorted(
+                record.row_address
+                for record in errors.records
+                if record.row_address is not None
+            )
+            if addresses:
+                address_list = ", ".join(str(address) for address in addresses)
+                await database.execute_query(
+                    f"REFRESH COLUMN embedding ON {table.name()} "
+                    f"WHERE _rowaddr IN ({address_list})"
+                )
+
+        Records whose per-row detail was capped have no address and appear
+        only in ``fragments``; they cannot be selected individually here.
+
         Examples
         --------
         >>> errors = await table.function_errors(column="embedding")  # doctest: +SKIP

python/python/tests/test_function_packaging.py

@@ -388,6 +388,7 @@ def test_code_packages_ship_as_a_bundle_and_version_with_their_source(
 
     def define():
         sys.modules.pop("fx_helpers_functions", None)
+        importlib.invalidate_caches()
         return importlib.import_module("fx_helpers_functions").limited
 
     first = define()
@@ -454,6 +455,7 @@ def test_local_module_outside_code_warns_at_registration(tmp_path, helper_packag
                 """
             )
         )
+        importlib.invalidate_caches()
         return getattr(importlib.import_module(name), name)
 
     with pytest.warns(UserWarning, match=r"local module.*code=\[fx_local\]"):
@@ -467,6 +469,44 @@ def test_local_module_outside_code_warns_at_registration(tmp_path, helper_packag
     assert "from fx_local import normalize" in artifact(installed).decode()
 
 
+def test_local_module_imported_in_function_body_warns_at_registration(
+    tmp_path, helper_package
+):
+    """A direct import inside the packaged definition is checked too."""
+    helper_package(
+        "fx_body_local",
+        {"__init__.py": "def normalize(value):\n    return value\n"},
+    )
+    sys.modules.pop("fx_body_local", None)
+
+    def define(name, options):
+        (tmp_path / f"{name}.py").write_text(
+            textwrap.dedent(
+                f"""
+                from lancedb import udf
+
+
+                @udf({options})
+                def {name}(value: int) -> int:
+                    import fx_body_local
+
+                    return fx_body_local.normalize(value)
+                """
+            )
+        )
+        importlib.invalidate_caches()
+        return getattr(importlib.import_module(name), name)
+
+    with pytest.warns(UserWarning, match=r"local module.*code=\[fx_body_local\]"):
+        local = define("fx_body_uses_local", "")
+    assert "import fx_body_local" in artifact(local).decode()
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        installed = define("fx_body_uses_installed", "pip=['fx-body-local==1.0']")
+    assert "import fx_body_local" in artifact(installed).decode()
+
+
 MMLB_HELPERS = {
     "__init__.py": "",
     "rate.py": """
@@ -527,6 +567,7 @@ def test_shared_helpers_construct_state_on_every_path(tmp_path, helper_package):
     initialized through the same helper, so neither path can skip it."""
     helper_package("fx_mmlb", MMLB_HELPERS)
     (tmp_path / "fx_mmlb_functions.py").write_text(MMLB_FUNCTIONS)
+    importlib.invalidate_caches()
     functions = importlib.import_module("fx_mmlb_functions")
 
     signatures = [

python/python/tests/test_job.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+import asyncio
+from datetime import timedelta
+
+import pytest
+
+from lancedb.job import AsyncJob, Job
+
+
+class _PendingJob:
+    id = "job-timeout"
+    job_type = None
+
+    async def wait(self):
+        await asyncio.sleep(60)
+
+
+class _PendingQuery:
+    id = "query-timeout"
+    job_type = "query"
+
+    async def wait(self):
+        await asyncio.sleep(60)
+
+
+def _assert_timeout(error: TimeoutError) -> None:
+    message = str(error)
+    assert "job 'job-timeout' still running after 0.01 seconds" in message
+    assert "continues on the server" in message
+    assert "db.open_job('job-timeout').wait()" in message
+
+
+@pytest.mark.asyncio
+async def test_async_job_wait_timeout_explains_reattach():
+    job = AsyncJob(_PendingJob())
+
+    with pytest.raises(TimeoutError) as exc_info:
+        await job.wait(timeout=timedelta(seconds=0.01))
+
+    _assert_timeout(exc_info.value)
+
+
+def test_job_wait_timeout_explains_reattach():
+    job = Job(AsyncJob(_PendingJob()))
+
+    with pytest.raises(TimeoutError) as exc_info:
+        job.wait(timeout=timedelta(seconds=0.01))
+
+    _assert_timeout(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_query_wait_timeout_requires_keeping_original_handle():
+    job = AsyncJob(_PendingQuery())
+
+    with pytest.raises(TimeoutError) as exc_info:
+        await job.wait(timeout=timedelta(seconds=0.01))
+
+    message = str(exc_info.value)
+    assert "query 'query-timeout' still running after 0.01 seconds" in message
+    assert "continues on the server" in message
+    assert "keep this handle to continue waiting" in message
+    assert "db.open_job" not in message

python/python/tests/test_remote_db.py

@@ -1795,6 +1795,7 @@ def test_remote_function_errors_lists_the_rows_a_refresh_skipped():
                 "job_id": "j-7",
                 "fragment_id": 3,
                 "row_offset": 9,
+                "row_address": 12884901897,
                 "column": "embedding",
                 "function": "embed",
                 "function_version": "2",
@@ -1845,7 +1846,13 @@ def test_remote_function_errors_lists_the_rows_a_refresh_skipped():
     assert errors.truncated is True
     assert [r.error_message for r in errors.records] == ["bad input 'x'"]
     assert errors.records[0].row_offset == 9
+    assert errors.records[0].row_address == 12884901897
     assert errors.records[0].function_version == "2"
+    assert repr(errors.records[0]) == (
+        'FunctionErrorRecord(job_id="j-7", fragment_id=3, row_offset=9, '
+        'row_address=12884901897, column="embedding", '
+        'error_type="ValueError", error_message="bad input \'x\'")'
+    )
     assert (errors.fragments[0].rows_skipped, errors.fragments[0].rows_recorded) == (
         500,
         100,
@@ -4812,6 +4819,9 @@ def test_remote_connection_jobs_surface():
         assert job.state == "failed"
         assert job.spec == {"column": "vec"}
         assert job.failure.message == "worker died"
+        assert repr(job.failure) == (
+            'JobFailureInfo(phase="execute", message="worker died", retryable=True)'
+        )
         assert job.status() == "failed"
         with pytest.raises(JobFailedError, match="worker died"):
             job.wait(timeout=timedelta(seconds=5))

python/src/job.rs

@@ -200,10 +200,20 @@ pub struct JobFailureInfo {
 #[pymethods]
 impl JobFailureInfo {
     fn __repr__(&self) -> String {
-        format!(
-            "JobFailureInfo(phase={:?}, message={:?}, retryable={:?})",
-            self.phase, self.message, self.retryable
-        )
+        let phase = self
+            .phase
+            .as_ref()
+            .map_or_else(|| "None".to_string(), |value| format!("{value:?}"));
+        let message = self
+            .message
+            .as_ref()
+            .map_or_else(|| "None".to_string(), |value| format!("{value:?}"));
+        let retryable = match self.retryable {
+            Some(true) => "True",
+            Some(false) => "False",
+            None => "None",
+        };
+        format!("JobFailureInfo(phase={phase}, message={message}, retryable={retryable})")
     }
 }
 

python/src/table.rs

@@ -596,6 +596,7 @@ pub struct FunctionErrorRecord {
     pub job_id: String,
     pub fragment_id: u64,
     pub row_offset: Option<u32>,
+    pub row_address: Option<u64>,
     pub column: String,
     pub function: String,
     pub function_version: String,
@@ -608,12 +609,19 @@ pub struct FunctionErrorRecord {
 #[pymethods]
 impl FunctionErrorRecord {
     pub fn __repr__(&self) -> String {
+        let row_offset = self
+            .row_offset
+            .map_or_else(|| "None".to_string(), |value| value.to_string());
+        let row_address = self
+            .row_address
+            .map_or_else(|| "None".to_string(), |value| value.to_string());
         format!(
-            "FunctionErrorRecord(job_id={:?}, fragment_id={}, row_offset={:?}, column={:?}, \
-             error_type={:?}, error_message={:?})",
+            "FunctionErrorRecord(job_id={:?}, fragment_id={}, row_offset={}, row_address={}, \
+             column={:?}, error_type={:?}, error_message={:?})",
             self.job_id,
             self.fragment_id,
-            self.row_offset,
+            row_offset,
+            row_address,
             self.column,
             self.error_type,
             self.error_message
@@ -627,6 +635,7 @@ impl From<lancedb::function::FunctionErrorRecord> for FunctionErrorRecord {
             job_id: record.job_id,
             fragment_id: record.fragment_id,
             row_offset: record.row_offset,
+            row_address: record.row_address,
             column: record.column,
             function: record.function,
             function_version: record.function_version,