From ced4799ad51f7ee4dd4bdff40c061ef59d8bb014 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Thu, 12 Nov 2026 01:21:58 -0800 Subject: [PATCH] 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 | 10 +- python/python/lancedb/job.py | 7 ++ python/python/tests/test_remote_db.py | 138 ++++++++++++++++++++++++++ python/src/connection.rs | 11 +- rust/lancedb/src/job.rs | 51 ++++++++++ rust/lancedb/src/remote/job.rs | 23 +++++ 6 files changed, 236 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 84ebb41a1..e2138520a 100644 --- a/python/python/lancedb/db.py +++ b/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 diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index 094a7b8d5..bf9605567 100644 --- a/python/python/lancedb/job.py +++ b/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) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index b0098aca3..8fde1209c 100644 --- a/python/python/tests/test_remote_db.py +++ b/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 diff --git a/python/src/connection.rs b/python/src/connection.rs index 0be933cf2..69cf54c99 100644 --- a/python/src/connection.rs +++ b/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::().infer_error()?, + )) + } + _ => Ok(crate::job::Job::new(job)), + } }) } diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index d178542b7..092ece8b5 100644 --- a/rust/lancedb/src/job.rs +++ b/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 { + /// let job = db.open_job(id).await? + /// .into_typed::()?; + /// job.wait().await + /// # } + /// ``` + pub fn into_typed(self) -> Result> + where + T: Clone + DeserializeOwned + Send + Sync + 'static, + { + let inner = match self.inner { + JobInner::Handle { handle, .. } => JobInner::Handle { + handle, + decode: Arc::new(TerminalResult::decode::), + }, + JobInner::Completed(()) => { + JobInner::Completed(TerminalResult::local(Value::Null).decode::()?) + } + }; + Ok(Job { + inner, + cache: self.cache, + }) + } } impl Job @@ -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::(), + 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); diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index d27c9c666..5a0afa7e8 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/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::().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::() + .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)]