fix(python): preserve typed results for reopened refresh jobs (#4981)

Restore typed results for reopened Function and SQL-expression refresh jobs while preserving their handles and cached descriptions.

Fixes https://linear.app/lancedb/issue/ENT-3441.

python/python/lancedb/db.py

@@ -67,7 +67,7 @@ from .graph import (
     PropertyGraph,
     _definition_json,
 )
-from .job import AsyncJob, Job, _typed_job
+from .job import AsyncJob, Job, _opened_job, _typed_job
 from .storage_integration import SecretReference, StorageIntegrationInfo
 from .sql import AsyncQuery as AsyncSqlQuery
 from .sql import Query as SqlQuery
@@ -1450,6 +1450,12 @@ class DBConnection(EnforceOverrides):
         specification, result, failure and event history, so there is no
         separate connection-level call for any of them.
 
+        For a column refresh, ``wait()`` returns a
+        [RefreshColumnResult][lancedb.functions.RefreshColumnResult], including
+        when the refresh was started through SQL. Other job types return
+        ``None`` from ``wait()``; their persisted result is available through
+        ``job.result`` after ``job.refresh()``.
+
         Raises `JobNotFoundError` when the server has no such job, the way
         `open_table` does for a missing table.
         """
@@ -3626,7 +3632,7 @@ class AsyncConnection(object):
         """Open a server-side job by id. See
         [DBConnection.open_job][lancedb.db.DBConnection.open_job].
         """
-        return AsyncJob(await self._inner.open_job(job_id))
+        return _opened_job(await self._inner.open_job(job_id))
 
     async def list_gpu_classes(self) -> List[GpuClass]:
         """List configured Function GPU classes. See

python/python/lancedb/job.py

@@ -17,6 +17,7 @@ from lancedb.background_loop import LOOP
 
 from . import _lancedb
 from ._lancedb import JobDescription, JobFailureInfo, JobInfo, _append_doc_url
+from .functions import RefreshColumnResult
 from .remote.errors import HttpError, RetryError
 
 T = TypeVar("T")
@@ -692,3 +693,9 @@ def _typed_job(
 ) -> AsyncJob[T]:
     """Bind an internal JSON-producing job to its public result model."""
     return AsyncJob(inner, result_decoder)
+
+
+def _opened_job(inner: "_lancedb.Job") -> AsyncJob:
+    if inner.job_type in ("function_refresh", "computed_refresh", "refresh_column"):
+        return _typed_job(inner, RefreshColumnResult.from_json)
+    return AsyncJob(inner)

python/python/tests/test_remote_db.py

@@ -2557,6 +2557,144 @@ async def test_remote_job_wait_disconnection_timeout_can_be_resumed(sync):
     assert exc.value.__cause__.is_transport_error is True
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [False, True])
+@pytest.mark.parametrize(
+    "job_type", ["function_refresh", "computed_refresh", "refresh_column"]
+)
+@pytest.mark.parametrize("running", [False, True])
+@pytest.mark.parametrize("rows_failed", [0, 8])
+async def test_reopened_refresh_job_wait_returns_typed_result(
+    sync, job_type, running, rows_failed
+):
+    terminal_result = {
+        "rows_assigned": 400 - rows_failed,
+        "rows_failed": rows_failed,
+        "rows_remaining": 0,
+        "source_version": 7,
+        "published_version": 8,
+        "summary": {"future_field": True},
+    }
+    polls = 0
+
+    def handler(request):
+        nonlocal polls
+        assert request.path == "/v1/jobs/describe"
+        payload = json.loads(request.rfile.read(int(request.headers["Content-Length"])))
+        assert payload == {"job_id": "refresh-1"}
+        polls += 1
+        in_progress = running and polls <= 2
+        response = {
+            "job_id": "refresh-1",
+            "job_type": job_type,
+            "job_state": "IN_PROGRESS" if in_progress else "DONE",
+            "creation_ms": 1000,
+            "spec": {"column": "derived"},
+        }
+        if not in_progress:
+            response["result"] = terminal_result
+        request.send_response(200)
+        request.send_header("Content-Type", "application/json")
+        request.end_headers()
+        request.wfile.write(json.dumps(response).encode())
+
+    observed = []
+    if sync:
+        with mock_lancedb_connection(handler) as db:
+            job = db.open_job("refresh-1")
+            assert job.id == "refresh-1"
+            assert job.spec == {"column": "derived"}
+            result = job.wait(timeout=timedelta(seconds=5), progress=observed.append)
+            assert job.wait() == result
+            job.refresh()
+            assert job.result == terminal_result
+    else:
+        async with mock_lancedb_connection_async(handler) as db:
+            job = await db.open_job("refresh-1")
+            assert job.id == "refresh-1"
+            assert job.spec == {"column": "derived"}
+            result = await job.wait(
+                timeout=timedelta(seconds=5), progress=observed.append
+            )
+            assert await job.wait() == result
+            await job.refresh()
+            assert job.result == terminal_result
+
+    assert isinstance(result, lancedb.RefreshColumnResult)
+    assert result == lancedb.RefreshColumnResult.model_validate(terminal_result)
+    assert result.rows_failed == rows_failed
+    assert observed[-1]["state"] == "finished"
+    if running:
+        assert observed[0]["state"] == "running"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [False, True])
+@pytest.mark.parametrize("terminal_result", [None, {}])
+async def test_reopened_refresh_job_rejects_missing_or_invalid_result(
+    sync, terminal_result
+):
+    def handler(request):
+        assert request.path == "/v1/jobs/describe"
+        request.rfile.read(int(request.headers["Content-Length"]))
+        request.send_response(200)
+        request.end_headers()
+        request.wfile.write(
+            json.dumps(
+                {
+                    "job_id": "refresh-1",
+                    "job_type": "function_refresh",
+                    "job_state": "DONE",
+                    "result": terminal_result,
+                }
+            ).encode()
+        )
+
+    if sync:
+        with mock_lancedb_connection(handler) as db:
+            job = db.open_job("refresh-1")
+            with pytest.raises(HttpError, match="typed job"):
+                job.wait()
+    else:
+        async with mock_lancedb_connection_async(handler) as db:
+            job = await db.open_job("refresh-1")
+            with pytest.raises(HttpError, match="typed job"):
+                await job.wait()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [False, True])
+async def test_reopened_unit_job_ignores_additive_result_payload(sync):
+    terminal_result = {"future_field": True}
+
+    def handler(request):
+        assert request.path == "/v1/jobs/describe"
+        request.rfile.read(int(request.headers["Content-Length"]))
+        request.send_response(200)
+        request.end_headers()
+        request.wfile.write(
+            json.dumps(
+                {
+                    "job_id": "index-1",
+                    "job_type": "create_index",
+                    "job_state": "DONE",
+                    "result": terminal_result,
+                }
+            ).encode()
+        )
+
+    if sync:
+        with mock_lancedb_connection(handler) as db:
+            job = db.open_job("index-1")
+            assert job.wait() is None
+            assert job.result == terminal_result
+    else:
+        async with mock_lancedb_connection_async(handler) as db:
+            job = await db.open_job("index-1")
+            assert await job.wait() is None
+            assert job.result == terminal_result
+
+
 @pytest.mark.parametrize("payload", [b"{", b"{}"])
 def test_remote_job_wait_does_not_retry_malformed_description(payload):
     polls = 0

python/src/connection.rs

@@ -24,7 +24,7 @@ use lancedb::{
     connection::NamespaceClientPushdownOperation,
     database::namespace::LanceNamespaceDatabase,
     database::{CreateTableMode, Database, ReadConsistency},
-    function::GpuClass as LanceGpuClass,
+    function::{GpuClass as LanceGpuClass, RefreshColumnResult},
     graph::PropertyGraphDefinition,
 };
 use pyo3::{
@@ -854,7 +854,14 @@ impl Connection {
         let inner = self_.get_inner()?.clone();
         future_into_py(self_.py(), async move {
             let job = inner.open_job(&job_id).await.infer_error()?;
-            Ok(crate::job::Job::new(job))
+            match job.job_type().as_deref() {
+                Some("function_refresh" | "computed_refresh" | "refresh_column") => {
+                    Ok(crate::job::Job::new_typed(
+                        job.into_typed::<RefreshColumnResult>().infer_error()?,
+                    ))
+                }
+                _ => Ok(crate::job::Job::new(job)),
+            }
         })
     }
 

rust/lancedb/src/job.rs

@@ -277,6 +277,42 @@ impl Job<()> {
         }
         job
     }
+
+    /// Decode this job's terminal payload as `T` when waiting.
+    ///
+    /// Use this after [`crate::Connection::open_job`] when the job's result
+    /// type is known. The job identity and cached description are preserved.
+    /// Missing or incompatible payloads are reported by [`Job::wait`].
+    /// A job that completed before its handle was created has only a unit
+    /// result; converting it fails immediately if `T` cannot decode JSON null.
+    ///
+    /// ```no_run
+    /// # async fn reopen(
+    /// #     db: &lancedb::Connection, id: &str,
+    /// # ) -> lancedb::Result<lancedb::function::RefreshColumnResult> {
+    /// let job = db.open_job(id).await?
+    ///     .into_typed::<lancedb::function::RefreshColumnResult>()?;
+    /// job.wait().await
+    /// # }
+    /// ```
+    pub fn into_typed<T>(self) -> Result<Job<T>>
+    where
+        T: Clone + DeserializeOwned + Send + Sync + 'static,
+    {
+        let inner = match self.inner {
+            JobInner::Handle { handle, .. } => JobInner::Handle {
+                handle,
+                decode: Arc::new(TerminalResult::decode::<T>),
+            },
+            JobInner::Completed(()) => {
+                JobInner::Completed(TerminalResult::local(Value::Null).decode::<T>()?)
+            }
+        };
+        Ok(Job {
+            inner,
+            cache: self.cache,
+        })
+    }
 }
 
 impl<T> Job<T>
@@ -750,6 +786,21 @@ mod tests {
         ));
     }
 
+    #[tokio::test]
+    async fn completed_unit_job_can_be_retyped_as_unit() {
+        let job = Job::new_done().into_typed::<()>().unwrap();
+        job.wait().await.unwrap();
+        assert_eq!(job.state().as_deref(), Some("finished"));
+    }
+
+    #[test]
+    fn completed_unit_job_rejects_an_incompatible_result_type() {
+        assert!(matches!(
+            Job::new_done().into_typed::<u64>(),
+            Err(Error::Runtime { .. })
+        ));
+    }
+
     #[tokio::test]
     async fn mapped_spawned_job_reuses_outcome() {
         let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1);

rust/lancedb/src/remote/job.rs

@@ -403,6 +403,29 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn reopened_job_preserves_its_description_and_decodes_the_terminal_result() {
+        let description: DescribeJobResponse = serde_json::from_str(REFRESH_JOB).unwrap();
+        let job = Job::opened(
+            Box::new(FixtureRemoteJob(REFRESH_JOB)),
+            description.into_description(),
+        );
+        let raw_result = job.result();
+        let job = job.into_typed::<RefreshColumnResult>().unwrap();
+
+        assert_eq!(job.job_type().as_deref(), Some("refresh_function_columns"));
+        assert_eq!(job.state().as_deref(), Some("finished"));
+        assert_eq!(job.creation_ms(), Some(1787270400001));
+        assert_eq!(job.result(), raw_result);
+        assert_eq!(job.wait().await.unwrap().rows_assigned, 999_998_800);
+        assert_eq!(job.wait().await.unwrap().published_version, Some(919));
+
+        let missing = Job::new(Box::new(FixtureRemoteJob(MISSING_RESULT_JOB)))
+            .into_typed::<RefreshColumnResult>()
+            .unwrap();
+        assert!(missing.wait().await.is_err());
+    }
+
     /// Five in-progress polls then done: 0.2 + 0.4 + 0.8 + 1 + 1 s of waiting.
     /// Under the old 5 s cap the same job would have waited 10 s.
     #[tokio::test(start_paused = true)]