From 1544ad65a80ab952f888e74ff7a121bf4e284bdd Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 3 Nov 2026 18:17:19 +0800 Subject: [PATCH] feat: expose structured job progress to SDK callers (#4808) Job callers currently lose the server's structured progress when describing, listing, or waiting for a job. Expose Rust and Python progress snapshots and optional wait observers using the existing backend poller. Python supports callbacks, a caller-owned tqdm-like bar, or `progress=True`; `status()` still returns a string and the default wait behavior is unchanged. Preserve unknown progress fields and phases from newer servers, and report omitted counters as unknown. Each wait calibrates a caller-owned bar from its first complete snapshot, then uses a high-water mark within the same execution epoch; a new epoch or smaller plan resets the bar. ETA is best-effort and omitted for missing, stale, paused, or incompatible samples. Newly submitted index and refresh handles forward observers through their freshness wrapper while retaining successful-wait read consistency. The Jobs guide includes draft Python/Rust snapshot and observer examples for release; TypeScript progress support remains tracked under ENT-3500. Validation on the development EC2 host: the initial 11 Python progress tests, REST fixture caller, and two Rust progress regressions passed. After fixing freshness-wrapper forwarding, rebuilt the Python extension and ran two submitted-handle Python tests plus the new Rust refresh and index regressions (one test each); all passed. Strict Clippy, rustfmt, and targeted Ruff passed. The final commit only wraps the new guide section in an MDX draft and clarifies language support. Repository-wide Ruff reports 48 existing findings in six unchanged docs/web-tests files. This publishes the SDK portion of ENT-3500. Server-side live-split and terminal-progress changes and the combined server/SDK end-to-end scenario remain in progress. Refs ENT-3500. BREAKING CHANGE: Rust struct literals for `JobInfo` must add `progress: None`, `phase: None`, and `paused: false`. Literals for `JobDescription` must add `progress: None`, `phase: None`, `queue_reason: None`, `waiting_for_job_id: None`, and `paused_at_millis: None` when that metadata is unavailable. Existing Python calls remain compatible. --- docs/web/basics/jobs.mdx | 64 ++++ python/python/lancedb/_lancedb.pyi | 24 +- python/python/lancedb/job.py | 249 +++++++++++++-- python/python/tests/test_job.py | 432 +++++++++++++++++++++++++- python/python/tests/test_remote_db.py | 134 +++++++- python/src/job.rs | 54 +++- rust/lancedb/src/database.rs | 75 +++++ rust/lancedb/src/job.rs | 42 ++- rust/lancedb/src/remote/db.rs | 14 +- rust/lancedb/src/remote/job.rs | 120 ++++++- rust/lancedb/src/remote/table.rs | 175 ++++++++++- 11 files changed, 1335 insertions(+), 48 deletions(-) diff --git a/docs/web/basics/jobs.mdx b/docs/web/basics/jobs.mdx index e6330e124..5a8a178fd 100644 --- a/docs/web/basics/jobs.mdx +++ b/docs/web/basics/jobs.mdx @@ -117,6 +117,70 @@ Event responses have [row limits](/basics/limits#jobs). An explicit limit can request more than the default, but it does not guarantee a complete event history. +{/* +Draft for the Functions release. Add this section before "Embedded and +Enterprise execution": + +### Progress snapshots and observers + +Enterprise registry handles expose a structured progress snapshot. This +includes index and refresh handles returned directly by their async submission +methods, as well as handles reopened with `open_job()`. Call `progress()` for a +point-in-time snapshot, or pass a callback to `wait()` to observe the same +server polls that drive completion. The callback receives running and terminal +observations; it does not start another polling loop. Progress snapshots and +wait observers are currently available in Python and Rust; TypeScript does not +yet expose this API. + +```python +def show_progress(progress): + print( + progress.get("state"), + progress.get("phase"), + progress.get("rows_computed"), + progress.get("rows_planned"), + ) + +job = tbl.refresh_column_async("embedding") +snapshot = job.progress() +if snapshot is not None: + show_progress(snapshot) + +result = job.wait(progress=show_progress) +``` + +```rust +let job = table.refresh_column_async("embedding").await?; +if let Some(progress) = job.progress().await? { + println!( + "state={:?} rows={:?}/{:?}", + progress.get("state"), + progress.get("rows_computed"), + progress.get("rows_planned"), + ); +} + +let result = job + .wait_with_progress(|description| { + let progress = description.progress_snapshot(); + println!( + "state={:?} rows={:?}/{:?}", + progress.get("state"), + progress.get("rows_computed"), + progress.get("rows_planned"), + ); + }) + .await?; +``` + +Job-type-specific counters are optional. A missing `rows_computed`, +`rows_planned`, or `rows_remaining` field means the server did not report that +counter; it does not mean zero, complete, or failed. Embedded handles and +SQL-backed query handles do not have registry progress records. Use their +`status()` and `wait()` APIs, and use registry progress only on handles backed +by the jobs API. +*/} + ## Embedded and Enterprise execution {#in-process} The following table describes LanceDB Embedded operations and registry jobs on diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index b8a90c08e..fcad61e71 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -1,6 +1,17 @@ from datetime import date, datetime, timedelta from decimal import Decimal -from typing import Dict, List, Optional, Sequence, Tuple, Any, TypedDict, Union, Literal +from typing import ( + Callable, + Dict, + List, + Optional, + Sequence, + Tuple, + Any, + TypedDict, + Union, + Literal, +) from uuid import UUID import pyarrow as pa @@ -388,7 +399,10 @@ class Job: @property def _description(self) -> Optional[JobDescription]: ... async def status(self) -> str: ... - async def wait(self) -> Optional[str]: ... + async def wait( + self, + progress: Optional[Callable[[dict[str, Any]], None]] = None, + ) -> Optional[str]: ... async def cancel(self) -> None: ... async def refresh(self) -> None: ... async def events( @@ -409,6 +423,10 @@ class JobInfo: def state(self) -> str: ... @property def created_at_millis(self) -> int: ... + @property + def progress(self) -> Optional[dict[str, Any]]: ... + @property + def paused(self) -> bool: ... class SessionStatus: @property @@ -467,6 +485,8 @@ class JobDescription: @property def result(self) -> Optional[Any]: ... @property + def progress(self) -> Optional[dict[str, Any]]: ... + @property def failure(self) -> Optional[JobFailureInfo]: ... class SqlQuery: diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index 582724892..92b545ce7 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -5,6 +5,7 @@ import asyncio import json +import time from datetime import timedelta from typing import Any, Callable, Generic, Optional, TypeVar, cast @@ -49,6 +50,140 @@ def _timeout_message( ) +def _normalize_job_progress(progress): + """Normalize ``wait(progress=...)`` using the same convention as writes.""" + if progress is True: + from tqdm.auto import tqdm + + return tqdm(unit=" rows"), True + if progress is False or progress is None: + return None, False + return progress, False + + +_ETA_MAX_SAMPLE_INTERVAL_MILLIS = 10 * 60 * 1_000 + + +def _progress_eta(previous, current): + """Best-effort remaining time from two same-execution observations.""" + if not previous or current.get("state") not in ("running", "finished"): + return None + remaining = current.get("rows_remaining") + if current.get("state") == "finished": + return 0 if remaining == 0 else None + if current.get("paused") or current.get("phase") in ( + "queued", + "starting", + "paused", + ): + return None + epoch = current.get("execution_epoch") + if epoch is None or previous.get("execution_epoch") != epoch: + return None + sample_at = current.get("sample_at_millis") + previous_at = previous.get("sample_at_millis") + if not isinstance(sample_at, int) or not isinstance(previous_at, int): + return None + elapsed = sample_at - previous_at + if not 0 < elapsed <= _ETA_MAX_SAMPLE_INTERVAL_MILLIS: + return None + if time.time() * 1_000 - sample_at > _ETA_MAX_SAMPLE_INTERVAL_MILLIS: + return None + planned = current.get("rows_planned") + if not isinstance(planned, int) or not isinstance(remaining, int): + return None + if remaining == 0: + return 0 + if planned != previous.get("rows_planned"): + return None + if current.get("fragments_planned") != previous.get("fragments_planned"): + return None + accounted = current.get("rows_computed", 0) + current.get("rows_skipped", 0) + previous_accounted = previous.get("rows_computed", 0) + previous.get( + "rows_skipped", 0 + ) + advanced = accounted - previous_accounted + previous_remaining = previous.get("rows_remaining") + if advanced <= 0 or not isinstance(previous_remaining, int): + return None + if remaining > previous_remaining: + return None + return remaining * elapsed // advanced + + +def _with_progress_eta(previous, current): + snapshot = dict(current) + estimate = _progress_eta(previous, snapshot) + if estimate is not None: + snapshot["estimated_remaining_time_millis"] = estimate + return snapshot + + +def _job_progress_callback(progress, enrich): + """Adapt a callback or tqdm-like object to structured job snapshots.""" + if progress is None: + return None + if callable(progress): + return lambda snapshot: progress(enrich(snapshot)) + + # A caller-owned bar can outlive one wait call. Resume from its displayed + # position instead of adding the same completed rows on the next call. + last_rows = getattr(progress, "n", 0) or 0 + last_planned = getattr(progress, "total", None) + last_epoch = None + first_snapshot = True + + def update(snapshot): + nonlocal first_snapshot, last_epoch, last_planned, last_rows + snapshot = enrich(snapshot) + planned = snapshot.get("rows_planned") + remaining = snapshot.get("rows_remaining") + computed = snapshot.get("rows_computed") + if isinstance(planned, int): + epoch = snapshot.get("execution_epoch") + reset = (isinstance(last_planned, int) and planned < last_planned) or ( + isinstance(last_epoch, int) + and isinstance(epoch, int) + and epoch != last_epoch + ) + if reset: + last_rows = 0 + if hasattr(progress, "reset"): + progress.reset(total=planned) + elif hasattr(progress, "n"): + progress.n = 0 + progress.total = planned + last_planned = planned + if isinstance(epoch, int): + last_epoch = epoch + completed = None + if isinstance(planned, int) and isinstance(remaining, int): + completed = max(planned - remaining, 0) + elif computed is not None: + completed = computed + if completed is not None: + if isinstance(planned, int): + completed = min(completed, planned) + if first_snapshot and completed < last_rows: + # The server may have started a new execution while a + # caller-owned bar retained its old position. + last_rows = completed + if hasattr(progress, "n"): + progress.n = completed + if hasattr(progress, "refresh"): + progress.refresh() + if completed > last_rows: + delta = completed - last_rows + last_rows = completed + progress.update(delta) + first_snapshot = False + phase = snapshot.get("phase") or snapshot.get("state") + if phase is not None and hasattr(progress, "set_postfix_str"): + progress.set_postfix_str(phase) + + return update + + class AsyncJob(Generic[T]): """A handle to an operation that may still be running. @@ -64,6 +199,12 @@ class AsyncJob(Generic[T]): ): self._inner = inner self._result_decoder = result_decoder + self._last_progress_snapshot = None + + def _enrich_progress(self, snapshot): + enriched = _with_progress_eta(self._last_progress_snapshot, snapshot) + self._last_progress_snapshot = enriched + return enriched @property def id(self) -> Optional[str]: @@ -88,12 +229,44 @@ class AsyncJob(Generic[T]): return "finished" return await self._inner.status() - async def wait(self, timeout: Optional[timedelta] = None) -> T: + async def progress(self) -> Optional[dict[str, Any]]: + """Fetch a structured progress snapshot. + + Registry jobs always report lifecycle fields. Queue metadata and + counters are included when supplied by the server. Older servers may + omit counters, and in-process operations return ``None``. A missing + counter is unknown rather than zero. This is additive to :meth:`status`, + whose string contract is unchanged. + """ + if self._inner is None: + return None + await self._inner.status() + description = getattr(self._inner, "_description", None) + if description is None or description.progress is None: + return None + return self._enrich_progress(description.progress) + + async def wait( + self, + timeout: Optional[timedelta] = None, + progress: Any = False, + ) -> T: """Wait until the operation reaches a terminal state. Returns the endpoint's typed result, or ``None`` for a unit-result job. + ``progress=True`` renders a progress bar, ``progress=callable`` + receives each structured snapshot, and a tqdm-like object is updated in + place. Observations use the backend's existing wait poller and never + start a second job. The bar advances by remaining rows when available, + because completed computation can include repeated work after a plan + change. On each wait call, the first complete observation synchronizes + a caller-owned bar with the server's current completed rows and total; + later observations in a stable execution epoch use a high-water mark. A + new execution epoch or smaller plan resets the bar to that plan. + ``False`` preserves the original wait behavior. + Raises `JobFailedError` if the operation failed, `JobCancelledError` if it was cancelled, and `TimeoutError` if `timeout` elapses first. Transient connection and read errors are retried with capped exponential @@ -109,6 +282,8 @@ class AsyncJob(Generic[T]): """ if self._inner is None: return cast(T, None) + progress_obj, owns = _normalize_job_progress(progress) + progress_callback = _job_progress_callback(progress_obj, self._enrich_progress) last_error = None async def wait_with_retries(): @@ -122,7 +297,9 @@ class AsyncJob(Generic[T]): await self._inner.status() last_error = None delay = 0.25 - return await self._inner.wait() + if progress_callback is None: + return await self._inner.wait() + return await self._inner.wait(progress=progress_callback) except (HttpError, RetryError) as error: cause = error.__cause__ if isinstance(error, RetryError) else error if not isinstance(cause, HttpError) or not cause.is_transport_error: @@ -132,32 +309,38 @@ class AsyncJob(Generic[T]): delay = min(delay * 2, 30.0) try: - if timeout is None: - result = await wait_with_retries() - else: - result = await asyncio.wait_for( - wait_with_retries(), timeout.total_seconds() - ) - except asyncio.TimeoutError as error: - if last_error is None: + try: if timeout is None: - raise - raise TimeoutError( - _timeout_message(self.id, timeout, self.job_type) - ) from error - identity = f"job {self.id!r}" if self.id is not None else "the operation" - if self.id is not None and self.job_type != "query": - reconnect = ( - f"Use db.open_job({self.id!r}) to reconnect, then call wait() " - "on the returned job." + result = await wait_with_retries() + else: + result = await asyncio.wait_for( + wait_with_retries(), timeout.total_seconds() + ) + except asyncio.TimeoutError as error: + if last_error is None: + if timeout is None: + raise + raise TimeoutError( + _timeout_message(self.id, timeout, self.job_type) + ) from error + identity = ( + f"job {self.id!r}" if self.id is not None else "the operation" ) - else: - reconnect = "Call wait() on this job handle to reconnect." - raise ConnectionError( - f"Lost connection while waiting and the timeout expired; " - f"{identity} may still be running. " - f"{reconnect}" - ) from last_error + if self.id is not None and self.job_type != "query": + reconnect = ( + f"Use db.open_job({self.id!r}) to reconnect, then call wait() " + "on the returned job." + ) + else: + reconnect = "Call wait() on this job handle to reconnect." + raise ConnectionError( + f"Lost connection while waiting and the timeout expired; " + f"{identity} may still be running. " + f"{reconnect}" + ) from last_error + finally: + if owns: + progress_obj.close() if self._result_decoder is not None: return self._result_decoder(result) return cast(T, result) @@ -343,7 +526,11 @@ class Job(Generic[T]): return "finished" return LOOP.run(self._inner.status()) - def wait(self, timeout: Optional[timedelta] = None) -> T: + def wait( + self, + timeout: Optional[timedelta] = None, + progress: Any = False, + ) -> T: """Block until the operation reaches a terminal state. Returns the endpoint's typed result, or ``None`` for a unit-result @@ -364,7 +551,13 @@ class Job(Generic[T]): """ if self._inner is None: return cast(T, None) - return LOOP.run(self._inner.wait(timeout)) + return LOOP.run(self._inner.wait(timeout, progress=progress)) + + def progress(self) -> Optional[dict[str, Any]]: + """Fetch a structured progress snapshot. See :meth:`AsyncJob.progress`.""" + if self._inner is None: + return None + return LOOP.run(self._inner.progress()) def cancel(self): """Request cancellation. Cancelling a finished operation is a no-op.""" diff --git a/python/python/tests/test_job.py b/python/python/tests/test_job.py index cbff31e43..c3189cbad 100644 --- a/python/python/tests/test_job.py +++ b/python/python/tests/test_job.py @@ -2,13 +2,14 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import asyncio +import time from datetime import timedelta from unittest.mock import AsyncMock, MagicMock import pytest from lancedb.exceptions import JobCancelledError, JobFailedError -from lancedb.job import AsyncJob, Job +from lancedb.job import AsyncJob, Job, _progress_eta from lancedb.remote.errors import HttpError, RetryError @@ -238,6 +239,20 @@ class _PendingQuery: await asyncio.sleep(60) +class _ProgressJob: + id = "job-progress" + job_type = "refresh_column" + + async def status(self): + return "finished" + + async def wait(self, progress=None): + if progress is not None: + progress({"state": "running", "rows_planned": 4, "rows_computed": 2}) + progress({"state": "finished", "rows_planned": 4, "rows_computed": 4}) + return None + + def _assert_timeout(error: TimeoutError) -> None: message = str(error) assert "job 'job-timeout' still running after 0.01 seconds" in message @@ -264,6 +279,421 @@ def test_job_wait_timeout_explains_reattach(): _assert_timeout(exc_info.value) +@pytest.mark.asyncio +async def test_progress_is_additive_to_string_status_and_returns_none_locally(): + inner = MagicMock(id="job-progress", job_type="refresh_column") + inner.status = AsyncMock(return_value="running") + inner._description = type( + "Description", (), {"progress": {"state": "running", "rows_computed": 2}} + )() + job = AsyncJob(inner) + + assert await job.status() == "running" + assert await job.progress() == {"state": "running", "rows_computed": 2} + assert await AsyncJob(None).progress() is None + assert await AsyncJob(_ProgressJob()).progress() is None + assert Job(AsyncJob(None)).progress() is None + + +def test_progress_eta_uses_two_stable_advancing_same_epoch_samples(): + now = int(time.time() * 1_000) + previous = { + "state": "running", + "execution_epoch": 7, + "sample_at_millis": now - 1_000, + "rows_planned": 100, + "rows_computed": 20, + "rows_skipped": 0, + "rows_remaining": 80, + "fragments_planned": 4, + } + current = { + **previous, + "sample_at_millis": now, + "rows_computed": 40, + "rows_remaining": 60, + } + assert _progress_eta(previous, current) == 3_000 + assert _progress_eta(previous, {**current, "rows_planned": 120}) is None + assert _progress_eta(previous, {**current, "execution_epoch": 8}) is None + assert _progress_eta(previous, {**current, "rows_computed": 20}) is None + stale = { + **current, + "sample_at_millis": now - 11 * 60 * 1_000, + } + assert ( + _progress_eta({**previous, "sample_at_millis": now - 12 * 60 * 1_000}, stale) + is None + ) + assert _progress_eta(previous, {**current, "state": "finished"}) is None + assert ( + _progress_eta(previous, {**current, "state": "finished", "rows_remaining": 0}) + == 0 + ) + + +@pytest.mark.asyncio +async def test_progress_retains_eta_from_same_handle_observations(): + now = int(time.time() * 1_000) + first = { + "state": "running", + "execution_epoch": 7, + "sample_at_millis": now - 1_000, + "rows_planned": 100, + "rows_computed": 20, + "rows_skipped": 0, + "rows_remaining": 80, + "fragments_planned": 4, + } + second = { + **first, + "sample_at_millis": now, + "rows_computed": 40, + "rows_remaining": 60, + } + inner = MagicMock(id="job-progress", job_type="refresh_column") + inner.status = AsyncMock(return_value="running") + description = type("Description", (), {"progress": first})() + inner._description = description + job = AsyncJob(inner) + + assert await job.progress() == first + description.progress = second + assert (await job.progress())["estimated_remaining_time_millis"] == 3_000 + + +@pytest.mark.asyncio +async def test_wait_progress_callback_observes_backend_poll_updates(): + updates = [] + + result = await AsyncJob(_ProgressJob()).wait(progress=updates.append) + + assert result is None + assert updates == [ + {"state": "running", "rows_planned": 4, "rows_computed": 2}, + {"state": "finished", "rows_planned": 4, "rows_computed": 4}, + ] + + +@pytest.mark.asyncio +async def test_wait_progress_callback_uses_the_public_eta_enrichment(monkeypatch): + now = 10_000_000 + monkeypatch.setattr("lancedb.job.time.time", lambda: now / 1_000) + + class EtaJob: + id = "job-eta" + job_type = "refresh_column" + + async def wait(self, progress): + first = { + "state": "running", + "execution_epoch": 7, + "sample_at_millis": now - 1_000, + "rows_planned": 100, + "rows_computed": 20, + "rows_skipped": 0, + "rows_remaining": 80, + "fragments_planned": 4, + } + progress(first) + progress( + { + **first, + "sample_at_millis": now, + "rows_computed": 40, + "rows_remaining": 60, + } + ) + + updates = [] + assert await AsyncJob(EtaJob()).wait(progress=updates.append) is None + + assert "estimated_remaining_time_millis" not in updates[0] + assert updates[1]["estimated_remaining_time_millis"] == 3_000 + + +@pytest.mark.asyncio +async def test_wait_progress_true_updates_and_closes_only_the_owned_bar(monkeypatch): + tqdm_module = pytest.importorskip("tqdm.auto") + + class Bar: + def __init__(self): + self.total = None + self.updates = [] + self.phases = [] + self.closed = False + + def update(self, delta): + self.updates.append(delta) + + def set_postfix_str(self, phase): + self.phases.append(phase) + + def close(self): + self.closed = True + + bar = Bar() + monkeypatch.setattr(tqdm_module, "tqdm", lambda **kwargs: bar) + + assert await AsyncJob(_ProgressJob()).wait(progress=True) is None + assert bar.total == 4 + assert bar.updates == [2, 2] + assert bar.phases == ["running", "finished"] + assert bar.closed + + +@pytest.mark.asyncio +async def test_wait_progress_bar_resumes_caller_owned_bar_after_timeout(): + class ResumableJob: + id = "job-resume" + job_type = "refresh_column" + + def __init__(self, bar): + self.calls = 0 + self.bar = bar + + async def wait(self, progress): + self.calls += 1 + if self.calls == 1: + progress( + { + "state": "running", + "execution_epoch": 1, + "rows_planned": 4, + "rows_remaining": 2, + } + ) + await asyncio.sleep(0.1) + else: + progress( + { + "state": "running", + "execution_epoch": 2, + "rows_planned": 4, + "rows_remaining": 3, + } + ) + assert self.bar.n == 1 + progress( + { + "state": "finished", + "execution_epoch": 2, + "rows_planned": 4, + "rows_remaining": 0, + } + ) + + class Bar: + def __init__(self): + self.total = None + self.n = 0 + self.updates = [] + + def update(self, delta): + self.updates.append(delta) + self.n += delta + + def set_postfix_str(self, _phase): + pass + + bar = Bar() + job = AsyncJob(ResumableJob(bar)) + + with pytest.raises(TimeoutError): + await job.wait(timeout=timedelta(seconds=0.01), progress=bar) + + assert await job.wait(progress=bar) is None + assert bar.total == 4 + assert bar.n == 4 + assert bar.updates == [2, 3] + + +@pytest.mark.asyncio +async def test_wait_progress_bar_resets_display_for_new_plan_or_epoch(): + class ReplannedJob: + id = "job-replanned" + job_type = "refresh_column" + + async def wait(self, progress): + progress( + { + "state": "running", + "execution_epoch": 1, + "rows_planned": 6, + "rows_remaining": 2, + } + ) + progress( + { + "state": "running", + "execution_epoch": 1, + "rows_planned": 4, + "rows_remaining": 3, + } + ) + progress( + { + "state": "running", + "execution_epoch": 2, + "rows_planned": 4, + "rows_remaining": 3, + } + ) + progress( + { + "state": "finished", + "execution_epoch": 2, + "rows_planned": 4, + "rows_remaining": 0, + } + ) + + class Bar: + def __init__(self): + self.total = None + self.n = 0 + self.updates = [] + self.resets = [] + + def update(self, delta): + self.updates.append(delta) + self.n += delta + + def reset(self, total=None): + self.resets.append(total) + self.total = total + self.n = 0 + + def set_postfix_str(self, _phase): + pass + + bar = Bar() + + assert await AsyncJob(ReplannedJob()).wait(progress=bar) is None + assert bar.total == 4 + assert bar.n == 4 + assert bar.n <= bar.total + assert bar.resets == [4, 4] + assert bar.updates == [4, 1, 1, 3] + + +@pytest.mark.asyncio +async def test_wait_progress_bar_uses_remaining_rows_when_computation_is_repeated( + monkeypatch, +): + tqdm_module = pytest.importorskip("tqdm.auto") + + class Bar: + def __init__(self): + self.total = None + self.updates = [] + self.closed = False + + def update(self, delta): + self.updates.append(delta) + + def set_postfix_str(self, _phase): + pass + + def close(self): + self.closed = True + + class RepeatedComputationJob: + id = "job-recomputed" + job_type = "refresh_column" + + async def wait(self, progress): + progress( + { + "state": "running", + "rows_planned": 4, + "rows_computed": 6, + "rows_remaining": 2, + } + ) + progress( + { + "state": "finished", + "rows_planned": 4, + "rows_computed": 8, + "rows_remaining": 0, + } + ) + + bar = Bar() + monkeypatch.setattr(tqdm_module, "tqdm", lambda **kwargs: bar) + + assert await AsyncJob(RepeatedComputationJob()).wait(progress=True) is None + assert bar.total == 4 + assert bar.updates == [2, 2] + assert bar.closed + + +@pytest.mark.asyncio +async def test_wait_progress_bar_tracks_late_planning_without_regressing( + monkeypatch, +): + tqdm_module = pytest.importorskip("tqdm.auto") + + class Bar: + def __init__(self): + self.total = None + self.updates = [] + self.closed = False + + def update(self, delta): + self.updates.append(delta) + + def set_postfix_str(self, _phase): + pass + + def close(self): + self.closed = True + + class RegressingPlanJob: + id = "job-regressing-plan" + job_type = "refresh_column" + + async def wait(self, progress): + samples = [ + (0, 0), + (4, 2), + (4, 3), + (4, 2), + (6, 1), + (6, 0), + ] + for planned, remaining in samples: + progress( + { + "state": "running", + "rows_planned": planned, + "rows_remaining": remaining, + } + ) + + bar = Bar() + monkeypatch.setattr(tqdm_module, "tqdm", lambda **kwargs: bar) + + assert await AsyncJob(RegressingPlanJob()).wait(progress=True) is None + assert bar.total == 6 + # The remaining-row dip at (4, 3) must not rewind the high-water mark. + assert bar.updates == [2, 3, 1] + assert bar.closed + + +@pytest.mark.asyncio +async def test_wait_progress_false_keeps_legacy_inner_wait_call(): + class NoProgressArgument: + id = "job-legacy" + job_type = None + + async def wait(self): + return None + + assert await AsyncJob(NoProgressArgument()).wait(progress=False) is None + + @pytest.mark.asyncio async def test_query_wait_timeout_requires_keeping_original_handle(): job = AsyncJob(_PendingQuery()) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index a57e85100..5deaf4d0a 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1738,12 +1738,16 @@ def test_remote_create_index_async_returns_job(): assert json.loads(body)["job_id"] == "job-1" describe_calls.append(1) state = "IN_PROGRESS" if len(describe_calls) == 1 else "DONE" + response = dict(job_id="job-1", job_state=state) + if len(describe_calls) == 1: + response.update( + progress=dict(rows_planned=10, rows_computed=3), + phase="building", + ) request.send_response(200) request.send_header("Content-Type", "application/json") request.end_headers() - request.wfile.write( - json.dumps(dict(job_id="job-1", job_state=state)).encode() - ) + request.wfile.write(json.dumps(response).encode()) elif request.path == "/v1/jobs/cancel": assert json.loads(body)["job_id"] == "job-1" request.send_response(200) @@ -1779,8 +1783,13 @@ def test_remote_create_index_async_returns_job(): table = db.create_table("test", [{"id": 1}]) job = table.create_index_async("id", config=BTree()) assert job.id == "job-1" - assert job.wait(timeout=timedelta(seconds=30)) is None + observed = [] + assert job.wait(timeout=timedelta(seconds=30), progress=observed.append) is None assert len(describe_calls) == 2 + assert observed[0]["state"] == "running" + assert observed[0]["rows_planned"] == 10 + assert observed[0]["rows_computed"] == 3 + assert observed[-1]["state"] == "finished" job.cancel() @@ -1886,6 +1895,91 @@ async def test_remote_refresh_async_returns_typed_terminal_result( assert result.version == 8 +def test_remote_refresh_async_forwards_progress_from_submitted_handle(): + terminal_result = { + "rows_assigned": 12, + "rows_failed": 0, + "rows_remaining": 0, + "source_version": 7, + "published_version": 8, + } + describe_calls = [] + + def handler(request): + content_len = int(request.headers.get("Content-Length", 0)) + body = request.rfile.read(content_len) if content_len > 0 else b"" + if request.path == "/v1/table/test/backfill_column": + assert json.loads(body) == {"column": "derived"} + request.send_response(202) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"job_id": "refresh-progress"}') + elif request.path == "/v1/jobs/describe": + assert json.loads(body)["job_id"] == "refresh-progress" + describe_calls.append(1) + if len(describe_calls) == 1: + response = { + "job_id": "refresh-progress", + "job_type": "function_refresh", + "job_state": "IN_PROGRESS", + "progress": {"rows_planned": 12, "rows_computed": 4}, + "phase": "running", + } + else: + response = { + "job_id": "refresh-progress", + "job_type": "function_refresh", + "job_state": "DONE", + "result": terminal_result, + } + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(response).encode()) + elif request.path == "/v1/table/test/create/?mode=create": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b"{}") + elif request.path == "/v1/table/test/describe/": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write( + json.dumps( + { + "version": 1, + "schema": { + "fields": [ + { + "name": "id", + "type": {"type": "int64"}, + "nullable": False, + } + ] + }, + } + ).encode() + ) + else: + request.send_response(404) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + table = db.create_table("test", [{"id": 1}]) + job = table.refresh_column_async("derived") + observed = [] + result = job.wait(timeout=timedelta(seconds=30), progress=observed.append) + + assert result.rows_filled == 12 + assert result.version == 8 + assert len(describe_calls) == 2 + assert observed[0]["state"] == "running" + assert observed[0]["rows_planned"] == 12 + assert observed[0]["rows_computed"] == 4 + assert observed[-1]["state"] == "finished" + + @pytest.mark.parametrize( "refresh_kwargs,message", [ @@ -4961,6 +5055,9 @@ def test_remote_connection_jobs_surface(): job_type="create_index", state="in_progress", created_at_millis=1000, + progress=dict(rows_planned=10, rows_computed=4), + phase="running", + paused=True, ) ], page_token="next", @@ -5014,6 +5111,11 @@ def test_remote_connection_jobs_surface(): job_state="FAILED", creation_ms=1000, spec=dict(column="vec"), + progress=dict(rows_planned=10, rows_computed=4), + phase="initializing", + queue_reason="waiting for capacity", + waiting_for_job_id="job-maintenance", + paused_at_millis=1200, failure=dict( phase="execute", message="worker died", retryable=True ), @@ -5064,6 +5166,16 @@ def test_remote_connection_jobs_surface(): assert [job.job_id for job in jobs] == ["job-1", "job-2"] assert jobs[0].state == "running" assert jobs[0].table == "t1" + assert jobs[0].progress == { + "job_type": "create_index", + "state": "running", + "creation_ms": 1000, + "paused": True, + "phase": "running", + "rows_planned": 10, + "rows_computed": 4, + } + assert jobs[0].paused is True assert jobs[1].state == "finished" assert db.cancel_job("job-1") is True @@ -5084,13 +5196,25 @@ def test_remote_connection_jobs_surface(): # Opening already populated the handle. assert job.state == "failed" assert job.spec == {"column": "vec"} + progress = job.progress() + assert progress["state"] == "failed" + assert progress["rows_computed"] == 4 + assert progress["paused"] is True + assert progress["phase"] == "initializing" + assert progress["queue_reason"] == "waiting for capacity" + assert progress["waiting_for_job_id"] == "job-maintenance" + assert progress["paused_at_millis"] == 1200 assert job.failure.message == "worker died" assert repr(job.failure) == ( 'JobFailureInfo(phase="execute", message="worker died", retryable=True)' ) assert job.status() == "failed" + observed = [] with pytest.raises(JobFailedError, match="worker died"): - job.wait(timeout=timedelta(seconds=5)) + job.wait(timeout=timedelta(seconds=5), progress=observed.append) + assert observed[0]["state"] == "failed" + assert observed[0]["rows_computed"] == 4 + assert finished.progress()["paused"] is False def test_remote_job_handle_reports_its_own_detail(): diff --git a/python/src/job.rs b/python/src/job.rs index 0a01d9faf..fd5ef64be 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -10,7 +10,7 @@ use arrow::{ }; use lancedb::job::JobEventsRequest; use pyo3::{ - Bound, PyAny, PyRef, PyResult, Python, + Bound, Py, PyAny, PyRef, PyResult, Python, exceptions::PyValueError, pyclass, pymethods, types::{PyAnyMethods, PyDict, PyDictMethods}, @@ -94,10 +94,30 @@ impl Job { ) } - pub fn wait(self_: PyRef<'_, Self>) -> PyResult> { + #[pyo3(signature = (progress=None))] + pub fn wait(self_: PyRef<'_, Self>, progress: Option>) -> PyResult> { let inner = self_.inner.clone(); future_into_py(self_.py(), async move { - let result = inner.wait().await.infer_error()?; + let result = if let Some(progress) = progress { + let observer = move |description: &lancedb::database::JobDescription| { + Python::attach(|py| { + let snapshot = description.progress_snapshot(); + let callback_result: PyResult<()> = (|| { + let raw = snapshot.to_string(); + let parsed = parse_json_payload(py, Some(&raw))? + .ok_or_else(|| PyValueError::new_err("missing progress"))?; + progress.bind(py).call1((parsed,))?; + Ok(()) + })(); + if let Err(error) = callback_result { + log::warn!("job progress callback error: {error}"); + } + }); + }; + inner.wait_with_progress(observer).await.infer_error()? + } else { + inner.wait().await.infer_error()? + }; result .map_err(|message| lancedb::Error::Runtime { message }) .infer_error() @@ -164,6 +184,8 @@ pub struct JobInfo { job_type: String, state: String, created_at_millis: i64, + _progress_json: Option, + paused: bool, } #[pymethods] @@ -174,16 +196,32 @@ impl JobInfo { self.job_id, self.table, self.job_type, self.state, self.created_at_millis ) } + + /// Structured progress for job types that report it. Missing fields are + /// unknown, not zero. + #[getter] + fn progress<'py>(&self, py: Python<'py>) -> PyResult>> { + parse_json_payload(py, self._progress_json.as_deref()) + } + + /// Whether an operator has parked this job. + #[getter] + fn paused(&self) -> bool { + self.paused + } } impl From for JobInfo { fn from(info: lancedb::database::JobInfo) -> Self { + let progress = info.progress_snapshot().to_string(); Self { job_id: info.job_id, table: info.table, job_type: info.job_type, state: info.state, created_at_millis: info.created_at_millis, + _progress_json: Some(progress), + paused: info.paused, } } } @@ -229,6 +267,7 @@ pub struct JobDescription { _spec_json: Option, /// Internal: the wire form behind the `result` property. _result_json: Option, + _progress_json: Option, failure: Option, } @@ -246,6 +285,13 @@ impl JobDescription { parse_json_payload(py, self._result_json.as_deref()) } + /// Structured progress snapshot. Lifecycle and queue fields are present + /// even when a job type or older server supplies no counters. + #[getter] + fn progress<'py>(&self, py: Python<'py>) -> PyResult>> { + parse_json_payload(py, self._progress_json.as_deref()) + } + fn __repr__(&self, py: Python<'_>) -> PyResult { let mut fields = vec![ format!("job_id={:?}", self.job_id), @@ -273,6 +319,7 @@ impl JobDescription { impl From for JobDescription { fn from(description: lancedb::database::JobDescription) -> Self { + let progress_json = description.progress_snapshot().to_string(); Self { job_id: description.job_id, job_type: description.job_type, @@ -283,6 +330,7 @@ impl From for JobDescription { .result .filter(|result| !result.is_null()) .map(|result| result.to_string()), + _progress_json: Some(progress_json), failure: description.failure.map(|failure| JobFailureInfo { phase: failure.phase, message: failure.message, diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 02c412c22..daa7f3f29 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -249,6 +249,34 @@ pub struct JobInfo { pub state: String, /// When the job was created, in milliseconds since the epoch. pub created_at_millis: i64, + /// Latest progress sample for a job type that reports structured progress. + pub progress: Option, + /// Coarse server-derived phase, including phases introduced by newer + /// servers. + pub phase: Option, + /// Whether an operator has parked this job. + pub paused: bool, +} + +impl JobInfo { + /// A flat progress mapping for list callers. Lifecycle and pause metadata + /// remain available when a job type or older server has no counters. + pub fn progress_snapshot(&self) -> serde_json::Value { + let mut snapshot = serde_json::Map::new(); + snapshot.insert("job_type".to_string(), self.job_type.clone().into()); + snapshot.insert("state".to_string(), self.state.clone().into()); + snapshot.insert("creation_ms".to_string(), self.created_at_millis.into()); + snapshot.insert("paused".to_string(), self.paused.into()); + if let Some(phase) = &self.phase { + snapshot.insert("phase".to_string(), phase.clone().into()); + } + if let Some(serde_json::Value::Object(progress)) = &self.progress { + for (key, value) in progress { + snapshot.insert(key.clone(), value.clone()); + } + } + serde_json::Value::Object(snapshot) + } } /// The server-side record behind a [`crate::job::Job`] handle: lifecycle @@ -261,6 +289,17 @@ pub struct JobDescription { pub state: String, /// When the job was created, in milliseconds since the epoch. pub creation_ms: i64, + /// Latest progress sample for a job type that reports structured progress. + pub progress: Option, + /// Coarse server-derived phase, including phases introduced by newer + /// servers. + pub phase: Option, + /// Why this job is queued before execution. Absent once work can start. + pub queue_reason: Option, + /// Maintenance job this queued job is waiting for, if any. + pub waiting_for_job_id: Option, + /// When the job was paused, if it is. + pub paused_at_millis: Option, /// The job-type-specific specification. Null when the server omits it. pub spec: serde_json::Value, /// The job-type-specific terminal result, for job types that define one. @@ -272,6 +311,42 @@ pub struct JobDescription { pub failure: Option, } +impl JobDescription { + /// A flat, job-type-neutral progress snapshot for SDK callers. + /// + /// Lifecycle metadata is always available. Queue metadata and structured + /// progress fields are included when supplied by the server; older servers + /// simply omit those fields. + pub fn progress_snapshot(&self) -> serde_json::Value { + let mut snapshot = serde_json::Map::new(); + snapshot.insert("job_type".to_string(), self.job_type.clone().into()); + snapshot.insert("state".to_string(), self.state.clone().into()); + snapshot.insert("creation_ms".to_string(), self.creation_ms.into()); + if let Some(phase) = &self.phase { + snapshot.insert("phase".to_string(), phase.clone().into()); + } + if let Some(queue_reason) = &self.queue_reason { + snapshot.insert("queue_reason".to_string(), queue_reason.clone().into()); + } + if let Some(waiting_for_job_id) = &self.waiting_for_job_id { + snapshot.insert( + "waiting_for_job_id".to_string(), + waiting_for_job_id.clone().into(), + ); + } + snapshot.insert("paused".to_string(), self.paused_at_millis.is_some().into()); + if let Some(paused_at_millis) = self.paused_at_millis { + snapshot.insert("paused_at_millis".to_string(), paused_at_millis.into()); + } + if let Some(serde_json::Value::Object(progress)) = &self.progress { + for (key, value) in progress { + snapshot.insert(key.clone(), value.clone()); + } + } + serde_json::Value::Object(snapshot) + } +} + /// The server's answer to a pause request. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PauseJobStatus { diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 4bb4d96dc..7b801bc0a 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -5,7 +5,7 @@ use std::future::Future; use std::pin::Pin; -use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; +use std::sync::{Arc, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard}; use arrow_array::RecordBatch; use async_trait::async_trait; @@ -19,6 +19,11 @@ use crate::database::JobDescription; use crate::docs::DocLink; use crate::error::{Error, JobFailure, Result}; +/// Observer called with each description fetched by a job's existing wait +/// poller. Keeping this callback on the poller avoids a second lifecycle or +/// retry loop just to render progress. +pub type JobProgressCallback = Arc>; + /// Which of a job's events [`Job::events`] returns. /// /// The handle already knows which job to ask about, so this narrows the @@ -66,6 +71,14 @@ pub(crate) trait JobHandle: Send + Sync { } async fn status(&self) -> Result; async fn wait(&self) -> Result; + /// Wait while observing descriptions from the same backend poller. + /// Backends without a server-side record keep the default behavior. + async fn wait_with_progress( + &self, + _progress: Option, + ) -> Result { + self.wait().await + } async fn cancel(&self) -> Result<()>; /// The job's full server-side record. Backends that run the operation in /// this process have none and keep the default. @@ -367,6 +380,15 @@ where self.refresh_state().await } + /// The current structured progress snapshot, or `None` for an in-process + /// operation. This refreshes the server-side record once, matching + /// [`Job::status`]; older servers return lifecycle and queue fields without + /// job-type-specific counters. + pub async fn progress(&self) -> Result> { + self.refresh_state().await?; + Ok(self.with_description(JobDescription::progress_snapshot)) + } + /// The last lifecycle state this handle observed, without contacting the /// backend. `None` until the handle has. pub fn state(&self) -> Option { @@ -448,9 +470,25 @@ where /// Returns [`crate::Error::JobFailed`] if the operation failed and /// [`crate::Error::JobCancelled`] if it was cancelled. pub async fn wait(&self) -> Result { + self.wait_with_progress_opt(None).await + } + + /// Wait for the operation while observing progress from the backend's + /// existing wait poller. The callback is called once for each description + /// the poller fetches, including the terminal observation when the server + /// supplies one. It does not create another lifecycle or retry loop. + pub async fn wait_with_progress(&self, progress: F) -> Result + where + F: FnMut(&JobDescription) + Send + 'static, + { + self.wait_with_progress_opt(Some(Arc::new(Mutex::new(progress)))) + .await + } + + async fn wait_with_progress_opt(&self, progress: Option) -> Result { match &self.inner { JobInner::Handle { handle, decode } => { - let settled = handle.wait().await; + let settled = handle.wait_with_progress(progress).await; // Waiting already established a terminal state; record it so // the detail accessors do not need another round trip for it. if let Some(state) = terminal_state(&settled) { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index d53838401..229cdd28c 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -699,6 +699,12 @@ struct RemoteListJobRow { state: String, #[serde(default)] created_at_millis: i64, + #[serde(default)] + progress: Option, + #[serde(default)] + phase: Option, + #[serde(default)] + paused: bool, } #[derive(serde::Deserialize)] @@ -1211,6 +1217,9 @@ impl Database for RemoteDatabase { job_type: row.job_type, state: job_state_to_client(&row.state), created_at_millis: row.created_at_millis, + progress: row.progress, + phase: row.phase, + paused: row.paused, }) .collect(); Ok(ListingPage { @@ -3965,7 +3974,7 @@ mod tests { http::Response::builder() .status(200) .body( - r#"{"jobs": [{"job_id": "job-1", "table": "t1", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1000}], "page_token": "next"}"#, + r#"{"jobs": [{"job_id": "job-1", "table": "t1", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1000, "progress": {"rows_planned": 10, "rows_computed": 4}, "phase": "running", "paused": true}], "page_token": "next"}"#, ) .unwrap() } @@ -3989,6 +3998,9 @@ mod tests { assert_eq!(jobs[0].job_id, "job-1"); assert_eq!(jobs[0].table, "t1"); assert_eq!(jobs[0].state, "running"); + assert_eq!(jobs[0].progress.as_ref().unwrap()["rows_computed"], 4); + assert_eq!(jobs[0].phase.as_deref(), Some("running")); + assert!(jobs[0].paused); assert_eq!(jobs[1].job_id, "job-2"); assert_eq!(jobs[1].state, "finished"); assert_eq!(jobs[1].created_at_millis, 2000); diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 7d2e70ac9..22ae938ed 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -13,7 +13,7 @@ use serde::Deserialize; use crate::database::JobDescription; use crate::error::{Error, JobFailure, Result}; -use crate::job::{JobEventsRequest, JobHandle, TerminalResult}; +use crate::job::{JobEventsRequest, JobHandle, JobProgressCallback, TerminalResult}; use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. @@ -67,7 +67,7 @@ pub(super) fn job_state_to_client(state: &str) -> String { /// The server's account of why a job failed. Absent from older servers, which /// report only the terminal state. -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] pub(super) struct ReportedFailure { #[serde(default)] phase: Option, @@ -100,7 +100,7 @@ pub(super) struct ResumeJobResponse { } /// Forward-compatible `/v1/jobs/describe` wire envelope. -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] pub(super) struct DescribeJobResponse { #[serde(default)] pub(super) job_id: String, @@ -110,6 +110,16 @@ pub(super) struct DescribeJobResponse { #[serde(default)] pub(super) creation_ms: i64, #[serde(default)] + pub(super) progress: Option, + #[serde(default)] + pub(super) phase: Option, + #[serde(default)] + pub(super) queue_reason: Option, + #[serde(default)] + pub(super) waiting_for_job_id: Option, + #[serde(default)] + pub(super) paused_at_millis: Option, + #[serde(default)] pub(super) spec: serde_json::Value, #[serde(default)] pub(super) result: Option, @@ -144,6 +154,11 @@ impl DescribeJobResponse { job_type: self.job_type, state: JobState::from(self.job_state.as_str()).client_label(), creation_ms: self.creation_ms, + progress: self.progress, + phase: self.phase, + queue_reason: self.queue_reason, + waiting_for_job_id: self.waiting_for_job_id, + paused_at_millis: self.paused_at_millis, spec: self.spec, result: self.result, failure: self.failure.map(ReportedFailure::into_job_failure), @@ -226,9 +241,23 @@ impl JobHandle for RemoteJob { } async fn wait(&self) -> Result { + self.wait_with_progress(None).await + } + + async fn wait_with_progress( + &self, + progress: Option, + ) -> Result { let mut interval = INITIAL_POLL_INTERVAL; loop { let (request_id, description) = self.fetch_description().await?; + if let Some(progress) = &progress { + if let Ok(mut callback) = progress.lock() { + callback(&description.clone().into_description()); + } else { + log::warn!("job {} progress callback lock was poisoned", self.job_id); + } + } match description.state() { JobState::Done => return Ok(description.into_terminal_result(request_id)), JobState::Failed => { @@ -270,15 +299,15 @@ impl JobHandle for RemoteJob { #[cfg(test)] mod tests { - use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; use std::time::Duration; use async_trait::async_trait; use crate::Result; use crate::function::{FunctionVersion, RefreshColumnResult}; - use crate::job::{Job, JobHandle, TerminalResult}; + use crate::job::{Job, JobHandle, JobProgressCallback, TerminalResult}; use crate::remote::client::test_utils::client_with_handler; use super::{DescribeJobResponse, RemoteJob}; @@ -371,4 +400,85 @@ mod tests { serde_json::from_str(FUNCTION_JOB).expect("function job fixture"); assert_eq!(response.job_state, "DONE"); } + + #[test] + fn describe_decodes_progress_queue_and_pause_fields() { + let response: DescribeJobResponse = serde_json::from_str( + r#"{ + "job_id":"job-1", + "job_type":"refresh_column", + "job_state":"IN_PROGRESS", + "creation_ms":10, + "progress":{"rows_planned":10,"rows_computed":4}, + "phase":"initializing", + "queue_reason":"waiting for capacity", + "waiting_for_job_id":"job-0", + "paused_at_millis":20 + }"#, + ) + .expect("progress fields"); + let description = response.into_description(); + let snapshot = description.progress_snapshot(); + assert_eq!(snapshot["state"], "running"); + assert_eq!(snapshot["phase"], "initializing"); + assert_eq!(snapshot["rows_planned"], 10); + assert_eq!(snapshot["rows_computed"], 4); + assert_eq!(snapshot["queue_reason"], "waiting for capacity"); + assert_eq!(snapshot["waiting_for_job_id"], "job-0"); + assert_eq!(snapshot["paused"], true); + assert_eq!(snapshot["paused_at_millis"], 20); + } + + #[test] + fn describe_progress_reports_stable_unpaused_state() { + let response: DescribeJobResponse = serde_json::from_str( + r#"{ + "job_id":"job-1", + "job_type":"refresh_column", + "job_state":"IN_PROGRESS", + "creation_ms":10 + }"#, + ) + .expect("job description"); + let snapshot = response.into_description().progress_snapshot(); + assert_eq!(snapshot["paused"], false); + assert!(snapshot.get("paused_at_millis").is_none()); + } + + #[tokio::test(start_paused = true)] + async fn wait_progress_observes_the_existing_poll_loop() { + let describes = Arc::new(AtomicUsize::new(0)); + let seen = describes.clone(); + let client = client_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/jobs/describe"); + let n = seen.fetch_add(1, Ordering::SeqCst); + let body = if n == 0 { + r#"{"job_state":"IN_PROGRESS","job_type":"refresh_column","progress":{"rows_planned":10,"rows_computed":4}}"# + } else { + r#"{"job_state":"DONE","job_type":"refresh_column","progress":{"rows_planned":10,"rows_computed":10},"result":null}"# + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let snapshots = Arc::new(Mutex::new(Vec::new())); + let observed = snapshots.clone(); + let progress: JobProgressCallback = Arc::new(Mutex::new( + move |description: &crate::database::JobDescription| { + observed + .lock() + .expect("snapshot lock") + .push(description.progress_snapshot()); + }, + )); + + RemoteJob::new(client, "job-1".to_string()) + .wait_with_progress(Some(progress)) + .await + .expect("unit job completes"); + + let snapshots = snapshots.lock().expect("snapshot lock"); + assert_eq!(snapshots.len(), 2); + assert_eq!(snapshots[0]["rows_computed"], 4); + assert_eq!(snapshots[1]["state"], "finished"); + assert_eq!(snapshots[1]["rows_computed"], 10); + } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 461cf713a..a46fb8a11 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -289,7 +289,14 @@ impl crate::job::JobHandle for FreshnessJob { } async fn wait(&self) -> Result { - let result = crate::job::JobHandle::wait(&self.inner).await?; + crate::job::JobHandle::wait_with_progress(self, None).await + } + + async fn wait_with_progress( + &self, + progress: Option, + ) -> Result { + let result = crate::job::JobHandle::wait_with_progress(&self.inner, progress).await?; let version = self.version.read().await; if version.is_none() { let result_version = match self.tracked_result { @@ -10915,6 +10922,172 @@ mod tests { assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); } + /// A newly submitted refresh returns a freshness-wrapped handle. Waiting + /// through it must forward progress observations and still establish the + /// published-version read fence after success. + #[tokio::test] + async fn test_new_refresh_handle_forwards_progress_and_keeps_read_freshness() { + let polls = Arc::new(AtomicUsize::new(0)); + let polls_in_handler = polls.clone(); + let saw_version = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let saw_version_in_handler = saw_version.clone(); + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/backfill_column" => http::Response::builder() + .status(202) + .body(r#"{"job_id": "j-progress"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => { + let response = if polls_in_handler.fetch_add(1, Ordering::SeqCst) == 0 { + json!({ + "job_id": "j-progress", + "job_type": "refresh_column", + "job_state": "IN_PROGRESS", + "progress": {"rows_planned": 12, "rows_computed": 4}, + "phase": "running", + }) + } else { + serde_json::from_str(&refresh_done("j-progress")).unwrap() + }; + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + } + "/v1/table/my_table/count_rows/" => { + saw_version_in_handler.store( + request + .headers() + .get("x-lancedb-min-read-version") + .and_then(|value| value.to_str().ok()) + == Some("8"), + Ordering::SeqCst, + ); + http::Response::builder() + .status(200) + .body("1".to_string()) + .unwrap() + } + path => panic!("unexpected path {path}"), + }); + + let job = table.refresh_column_async("doubled").await.unwrap(); + let observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let observed_in_callback = observed.clone(); + let result = job + .wait_with_progress(move |description| { + observed_in_callback + .lock() + .unwrap() + .push(description.progress_snapshot()); + }) + .await + .unwrap(); + assert_eq!(result.rows_assigned, 12); + + { + let observed = observed.lock().unwrap(); + assert_eq!(observed.len(), 2); + assert_eq!(observed[0]["state"], "running"); + assert_eq!(observed[0]["rows_planned"], 12); + assert_eq!(observed[0]["rows_computed"], 4); + assert_eq!(observed[1]["state"], "finished"); + } + + table.count_rows(None).await.unwrap(); + assert!( + saw_version.load(Ordering::SeqCst), + "new refresh handle lost its published-version freshness fence" + ); + } + + /// Index submissions use the same freshness wrapper. Its wait must forward + /// observations and retain the successful-wait timestamp bookkeeping. + #[tokio::test] + async fn test_new_index_handle_forwards_progress_and_keeps_read_freshness() { + let polls = Arc::new(AtomicUsize::new(0)); + let polls_in_handler = polls.clone(); + let saw_timestamp = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let saw_timestamp_in_handler = saw_timestamp.clone(); + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => http::Response::builder() + .status(200) + .body(r#"{"job_id": "j-index-progress"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => { + let response = if polls_in_handler.fetch_add(1, Ordering::SeqCst) == 0 { + json!({ + "job_id": "j-index-progress", + "job_type": "create_index", + "job_state": "IN_PROGRESS", + "progress": {"rows_planned": 10, "rows_computed": 3}, + "phase": "building", + }) + } else { + json!({ + "job_id": "j-index-progress", + "job_type": "create_index", + "job_state": "DONE", + }) + }; + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + } + "/v1/table/my_table/count_rows/" => { + saw_timestamp_in_handler.store( + request.headers().contains_key("x-lancedb-min-timestamp"), + Ordering::SeqCst, + ); + http::Response::builder() + .status(200) + .body("1".to_string()) + .unwrap() + } + path => panic!("unexpected path {path}"), + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + let observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let observed_in_callback = observed.clone(); + job.wait_with_progress(move |description| { + observed_in_callback + .lock() + .unwrap() + .push(description.progress_snapshot()); + }) + .await + .unwrap(); + + { + let observed = observed.lock().unwrap(); + assert_eq!(observed.len(), 2); + assert_eq!(observed[0]["state"], "running"); + assert_eq!(observed[0]["rows_planned"], 10); + assert_eq!(observed[0]["rows_computed"], 3); + assert_eq!(observed[1]["state"], "finished"); + } + + table.count_rows(None).await.unwrap(); + assert!( + saw_timestamp.load(Ordering::SeqCst), + "new index handle lost its successful-wait timestamp bookkeeping" + ); + } + #[tokio::test] async fn test_refresh_submission_uses_add_columns_version_fence() { let table = Table::new_with_handler("my_table", |request| match request.url().path() {