diff --git a/docs/src/js/classes/Job.md b/docs/src/js/classes/Job.md new file mode 100644 index 000000000..c46b9651f --- /dev/null +++ b/docs/src/js/classes/Job.md @@ -0,0 +1,64 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / Job + +# Class: Job + +A handle to an operation that may still be running. + +## Constructors + +### new Job() + +```ts +new Job(): Job +``` + +#### Returns + +[`Job`](Job.md) + +## Accessors + +### id + +```ts +get id(): null | string +``` + +Identifies the operation on the server that is running it. Operations +that run in this process have no server id. The value is opaque. + +#### Returns + +`null` \| `string` + +## Methods + +### cancel() + +```ts +cancel(): Promise +``` + +Request cancellation. Cancelling a finished operation is a no-op. + +#### Returns + +`Promise`<`void`> + +*** + +### wait() + +```ts +wait(): Promise +``` + +Wait until the operation reaches a terminal state. + +#### Returns + +`Promise`<`void`> diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 05636657e..11fca32d0 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -295,6 +295,29 @@ await table.createIndex("my_float_col"); *** +### createIndexAsync() + +```ts +abstract createIndexAsync(column, options?): Promise +``` + +Create an index, returning a handle to the indexing job. + +The job may already be complete when returned; callers must not assume +the index exists until [Job.wait](Job.md#wait) resolves. + +#### Parameters + +* **column**: `string` + +* **options?**: `Partial`<[`IndexOptions`](../interfaces/IndexOptions.md)> + +#### Returns + +`Promise`<[`Job`](Job.md)> + +*** + ### currentBranch() ```ts diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index c214a5f13..323a24a49 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -25,6 +25,7 @@ - [Connection](classes/Connection.md) - [HeaderProvider](classes/HeaderProvider.md) - [Index](classes/Index.md) +- [Job](classes/Job.md) - [MakeArrowTableOptions](classes/MakeArrowTableOptions.md) - [MatchQuery](classes/MatchQuery.md) - [MergeInsertBuilder](classes/MergeInsertBuilder.md) diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 41025047e..4cad365af 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -851,7 +851,11 @@ describe("When creating an index", () => { afterEach(() => tmpDir.removeCallback()); it("should create a vector index on vector columns", async () => { - await tbl.createIndex("vec"); + const job = await tbl.createIndexAsync("vec"); + expect(job.id).toBeNull(); + await job.wait(); + // Cancelling a job that already finished succeeds and does nothing. + await job.cancel(); // check index directory const indexDir = path.join(tmpDir.name, "test.lance", "_indices"); diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 63a4b08d0..2ce031458 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -85,7 +85,7 @@ export { RenameTableOptions, } from "./connection"; -export { Session } from "./native.js"; +export { Job, Session } from "./native.js"; export { ExecutableQuery, diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index d814e3e55..3359a2643 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -30,6 +30,7 @@ import { DropColumnsResult, IndexConfig, IndexStatistics, + Job, Branches as NativeBranches, OptimizeStats, TableStatistics, @@ -358,6 +359,17 @@ export abstract class Table { options?: Partial, ): Promise; + /** + * Create an index, returning a handle to the indexing job. + * + * The job may already be complete when returned; callers must not assume + * the index exists until {@link Job.wait} resolves. + */ + abstract createIndexAsync( + column: string, + options?: Partial, + ): Promise; + /** * Drop an index from the table. * @@ -940,6 +952,22 @@ export class LocalTable extends Table { ); } + async createIndexAsync( + column: string, + options?: Partial, + ): Promise { + // biome-ignore lint/suspicious/noExplicitAny: skip + const nativeIndex = (options?.config as any)?.inner; + return await this.inner.createIndexAsync( + nativeIndex, + column, + options?.replace, + options?.waitTimeoutSeconds, + options?.name, + options?.train, + ); + } + async dropIndex(name: string): Promise { await this.inner.dropIndex(name); } diff --git a/nodejs/src/job.rs b/nodejs/src/job.rs new file mode 100644 index 000000000..52d082beb --- /dev/null +++ b/nodejs/src/job.rs @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::sync::Arc; + +use napi_derive::napi; + +use crate::error::NapiErrorExt; + +/// A handle to an operation that may still be running. +#[napi] +pub struct Job { + inner: Arc, +} + +impl Job { + pub(crate) fn new(inner: lancedb::Job) -> Self { + Self { + inner: Arc::new(inner), + } + } +} + +#[napi] +impl Job { + /// Identifies the operation on the server that is running it. Operations + /// that run in this process have no server id. The value is opaque. + #[napi(getter)] + pub fn id(&self) -> Option { + self.inner.id().map(str::to_string) + } + + /// Wait until the operation reaches a terminal state. + #[napi(catch_unwind)] + pub async fn wait(&self) -> napi::Result<()> { + self.inner.wait().await.default_error() + } + + /// Request cancellation. Cancelling a finished operation is a no-op. + #[napi(catch_unwind)] + pub async fn cancel(&self) -> napi::Result<()> { + self.inner.cancel().await.default_error() + } +} diff --git a/nodejs/src/lib.rs b/nodejs/src/lib.rs index 93bb2044a..312b675bd 100644 --- a/nodejs/src/lib.rs +++ b/nodejs/src/lib.rs @@ -11,6 +11,7 @@ mod error; mod header; mod index; mod iterator; +mod job; pub mod merge; pub mod otel; pub mod permutation; diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index d1489034f..7dd902e19 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -168,6 +168,39 @@ impl Table { builder.execute().await.default_error() } + #[napi(catch_unwind)] + pub async fn create_index_async( + &self, + index: Option<&Index>, + column: String, + replace: Option, + wait_timeout_s: Option, + name: Option, + train: Option, + ) -> napi::Result { + let lancedb_index = if let Some(index) = index { + index.consume()? + } else { + lancedb::index::Index::Auto + }; + let mut builder = self.inner_ref()?.create_index(&[column], lancedb_index); + if let Some(replace) = replace { + builder = builder.replace(replace); + } + if let Some(timeout) = wait_timeout_s { + builder = + builder.wait_timeout(std::time::Duration::from_secs(timeout.try_into().unwrap())); + } + if let Some(name) = name { + builder = builder.name(name); + } + if let Some(train) = train { + builder = builder.train(train); + } + let job = builder.execute_async().await.default_error()?; + Ok(crate::job::Job::new(job)) + } + #[napi(catch_unwind)] pub async fn drop_index(&self, index_name: String) -> napi::Result<()> { self.inner_ref()? diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 71349b8a5..235049f97 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -20,6 +20,7 @@ from .remote import ClientConfig from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector, BlobType +from .job import AsyncJob, Job from .table import AsyncTable, Table from .types import BaseTokenizerType from ._lancedb import Session @@ -500,6 +501,7 @@ __all__ = [ "connect_namespace", "connect_namespace_async", "AsyncConnection", + "AsyncJob", "AsyncLanceNamespaceDBConnection", "AsyncTable", "FtsToken", @@ -513,6 +515,7 @@ __all__ = [ "BlobType", "vector", "DBConnection", + "Job", "LanceDBConnection", "LanceNamespaceDBConnection", "RemoteDBConnection", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 67325e10b..f5c59c152 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -209,6 +209,12 @@ class BlobFile: def read_range(self, offset: int, length: int) -> bytes: ... def read_up_to(self, length: int) -> bytes: ... +class Job: + @property + def id(self) -> Optional[str]: ... + async def wait(self) -> None: ... + async def cancel(self) -> None: ... + class Table: def name(self) -> str: ... def __repr__(self) -> str: ... @@ -248,6 +254,28 @@ class Table: name: Optional[str], train: Optional[bool], ): ... + async def create_index_async( + self, + column: str, + index: Union[ + IvfFlat, + IvfSq, + IvfPq, + HnswPq, + HnswSq, + HnswFlat, + BTree, + Bitmap, + LabelList, + Fm, + FTS, + ], + replace: Optional[bool], + wait_timeout: Optional[object], + *, + name: Optional[str], + train: Optional[bool], + ) -> Job: ... async def list_versions(self) -> List[Dict[str, Any]]: ... async def version(self) -> int: ... async def checkout(self, version: Union[int, str]): ... diff --git a/python/python/lancedb/exceptions.py b/python/python/lancedb/exceptions.py index 0ac0971e3..daa98ee6e 100644 --- a/python/python/lancedb/exceptions.py +++ b/python/python/lancedb/exceptions.py @@ -23,3 +23,15 @@ class MissingColumnError(KeyError): return ( f"Error: Column '{self.column_name}' does not exist in the DataFrame object" ) + + +class JobFailedError(RuntimeError): + """Exception raised when an asynchronous job reaches the failed state.""" + + pass + + +class JobCancelledError(RuntimeError): + """Exception raised when an asynchronous job was cancelled.""" + + pass diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py new file mode 100644 index 000000000..82911119b --- /dev/null +++ b/python/python/lancedb/job.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Handles to operations a server may run asynchronously.""" + +import asyncio +from datetime import timedelta +from typing import Optional + +from lancedb.background_loop import LOOP + +from . import _lancedb + + +class AsyncJob: + """A handle to an operation that may still be running. + + The operation may already be complete when the handle is created. + """ + + def __init__(self, inner: Optional["_lancedb.Job"]): + self._inner = inner + + @property + def id(self) -> Optional[str]: + """Identifies the operation on the server that is running it. + + Returned for correlating with server logs or the jobs API. Operations + that run in this process have no server id and return `None`. The value + is opaque: parsing it or storing it to resume the job later is not + supported. + """ + return self._inner.id if self._inner is not None else None + + async def wait(self, timeout: Optional[timedelta] = None): + """Wait until the operation reaches a terminal state. + + Raises `JobFailedError` if the operation failed, `JobCancelledError` + if it was cancelled, and `TimeoutError` if `timeout` elapses first. + """ + if self._inner is None: + return + if timeout is None: + await self._inner.wait() + else: + await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()) + + async def cancel(self): + """Request cancellation. Cancelling a finished operation is a no-op.""" + if self._inner is None: + return + await self._inner.cancel() + + +class Job: + """Synchronous counterpart of `AsyncJob`.""" + + def __init__(self, inner: Optional[AsyncJob]): + self._inner = inner + + @property + def id(self) -> Optional[str]: + """Identifies the operation on the server that is running it. + + See :attr:`AsyncJob.id`. + """ + return self._inner.id if self._inner is not None else None + + def wait(self, timeout: Optional[timedelta] = None): + """Block until the operation reaches a terminal state. + + Raises `JobFailedError` if the operation failed, `JobCancelledError` + if it was cancelled, and `TimeoutError` if `timeout` elapses first. + """ + if self._inner is None: + return + LOOP.run(self._inner.wait(timeout)) + + def cancel(self): + """Request cancellation. Cancelling a finished operation is a no-op.""" + if self._inner is None: + return + LOOP.run(self._inner.cancel()) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 268e0d5a4..acc2f4c9d 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -48,6 +48,7 @@ from lancedb.index import ( IvfSq, LabelList, ) +from lancedb.job import Job from lancedb.remote.db import LOOP from lancedb.table import IndexConfigType, KNOWN_METRICS import pyarrow as pa @@ -541,6 +542,34 @@ class RemoteTable(Table): ) ) + def create_index_async( + self, + column: str, + *, + config: IndexConfigType, + replace: Optional[bool] = None, + wait_timeout: Optional[timedelta] = None, + name: Optional[str] = None, + train: bool = True, + ) -> Job: + """Create an index, returning a handle to the indexing job. + + The job may already be complete when returned; callers must not assume + the index exists until :meth:`Job.wait` returns. + """ + return Job( + LOOP.run( + self._table.create_index_async( + column, + replace=replace, + config=config, + wait_timeout=wait_timeout, + name=name, + train=train, + ) + ) + ) + def _is_legacy_create_index_call( self, first_arg: str, diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 399b075f0..ea0f97621 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -40,6 +40,7 @@ from ._blob import ( from .types import BlobMode from lancedb.arrow import peek_reader from lancedb.background_loop import LOOP, embedding_executor +from lancedb.job import AsyncJob, Job from .dependencies import ( _check_for_hugging_face, _check_for_lance, @@ -977,6 +978,24 @@ class Table(ABC): """ raise NotImplementedError + def create_index_async( + self, + column: str, + *, + config: IndexConfigType, + replace: Optional[bool] = None, + wait_timeout: Optional[timedelta] = None, + name: Optional[str] = None, + train: bool = True, + ) -> Job: + """Create an index, returning a handle to the indexing job. + + Takes the same arguments as :meth:`create_index`. The job may already + be complete when returned; callers must not assume the index exists + until :meth:`Job.wait` returns. + """ + raise NotImplementedError + def drop_index(self, name: str) -> None: """ Drop an index from the table. @@ -2780,6 +2799,34 @@ class LanceTable(Table): ) ) + def create_index_async( + self, + column: str, + *, + config: IndexConfigType, + replace: Optional[bool] = None, + wait_timeout: Optional[timedelta] = None, + name: Optional[str] = None, + train: bool = True, + ) -> Job: + """Create an index, returning a handle to the indexing job. + + The job may already be complete when returned; callers must not assume + the index exists until :meth:`Job.wait` returns. + """ + return Job( + LOOP.run( + self._table.create_index_async( + column, + replace=replace, + config=config, + wait_timeout=wait_timeout, + name=name, + train=train, + ) + ) + ) + def _is_legacy_create_index_call( self, first_arg: str, @@ -4867,6 +4914,46 @@ class AsyncTable: ) raise e + async def create_index_async( + self, + column: str, + *, + replace: Optional[bool] = None, + config: Optional[ + Union[ + IvfFlat, + IvfPq, + IvfRq, + HnswPq, + HnswSq, + HnswFlat, + BTree, + Bitmap, + LabelList, + Fm, + FTS, + ] + ] = None, + wait_timeout: Optional[timedelta] = None, + name: Optional[str] = None, + train: bool = True, + ) -> AsyncJob: + """Create an index, returning a handle to the indexing job. + + Takes the same arguments as :meth:`create_index`. The job may already + be complete when returned; callers must not assume the index exists + until :meth:`AsyncJob.wait` resolves. + """ + job = await self._inner.create_index_async( + column, + index=config, + replace=replace, + wait_timeout=wait_timeout, + name=name, + train=train, + ) + return AsyncJob(job) + async def drop_index(self, name: str) -> None: """ Drop an index from the table. diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 65a244c5b..1cf2c733c 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -84,6 +84,15 @@ async def binary_table(db_async): ) +@pytest.mark.asyncio +async def test_create_index_async_returns_done_job(some_table: AsyncTable): + job = await some_table.create_index_async("id", config=BTree()) + assert job.id is None + await job.wait() + assert len(await some_table.list_indices()) == 1 + await job.cancel() + + @pytest.mark.asyncio async def test_create_scalar_index(some_table: AsyncTable): # Can create diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index dd46628f1..9bb162b15 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -812,6 +812,121 @@ def test_table_create_indices(): table.drop_index("custom_fts_idx") +def test_remote_create_index_async_returns_job(): + from lancedb.index import BTree + + 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/create_index/": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"job_id": "job-1"}') + elif request.path == "/v1/jobs/describe": + assert json.loads(body)["job_id"] == "job-1" + describe_calls.append(1) + state = "IN_PROGRESS" if len(describe_calls) == 1 else "DONE" + 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() + ) + elif request.path == "/v1/jobs/cancel": + assert json.loads(body)["job_id"] == "job-1" + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b"{}") + 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( + dict( + version=1, + schema=dict( + fields=[ + dict(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.create_index_async("id", config=BTree()) + assert job.id == "job-1" + job.wait(timeout=timedelta(seconds=30)) + assert len(describe_calls) == 2 + job.cancel() + + +def test_remote_job_wait_raises_on_failure(): + from lancedb.exceptions import JobFailedError + from lancedb.index import BTree + + 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/create_index/": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"job_id": "job-2"}') + elif request.path == "/v1/jobs/describe": + assert json.loads(body)["job_id"] == "job-2" + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write( + json.dumps(dict(job_id="job-2", job_state="FAILED")).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( + dict( + version=1, + schema=dict( + fields=[ + dict(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.create_index_async("id", config=BTree()) + with pytest.raises(JobFailedError, match="job-2"): + job.wait() + + def test_remote_create_index_new_api(): received_requests = [] diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 2a34084cf..a355cf101 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1402,6 +1402,15 @@ async def test_async_open_table_with_branch_version(tmp_path): assert await pinned.count_rows() == 4 # writable again +def test_create_index_async_returns_done_job(mem_db: DBConnection): + table = mem_db.create_table("job_test", [{"id": i} for i in range(10)]) + job = table.create_index_async("id", config=BTree()) + assert job.id is None + job.wait() + assert len(table.list_indices()) == 1 + job.cancel() + + @patch("lancedb.table.AsyncTable.create_index") def test_create_index_method(mock_create_index, mem_db: DBConnection): table = mem_db.create_table( diff --git a/python/src/error.rs b/python/src/error.rs index 062a1dc96..b66afe47b 100644 --- a/python/src/error.rs +++ b/python/src/error.rs @@ -102,6 +102,18 @@ impl PythonErrorExt for std::result::Result { err.setattr(intern!(py, "__cause__"), cause_err)?; Err(PyErr::from_value(err)) }), + LanceError::JobFailed { .. } => Python::attach(|py| { + let cls = py + .import(intern!(py, "lancedb.exceptions"))? + .getattr(intern!(py, "JobFailedError"))?; + Err(PyErr::from_value(cls.call1((err.to_string(),))?)) + }), + LanceError::JobCancelled { .. } => Python::attach(|py| { + let cls = py + .import(intern!(py, "lancedb.exceptions"))? + .getattr(intern!(py, "JobCancelledError"))?; + Err(PyErr::from_value(cls.call1((err.to_string(),))?)) + }), _ => self.runtime_error(), }, } diff --git a/python/src/job.rs b/python/src/job.rs new file mode 100644 index 000000000..83303ddea --- /dev/null +++ b/python/src/job.rs @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::sync::Arc; + +use crate::runtime::future_into_py; +use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods}; + +use crate::error::PythonErrorExt; + +#[pyclass] +pub struct Job { + inner: Arc, +} + +impl Job { + pub(crate) fn new(inner: lancedb::Job) -> Self { + Self { + inner: Arc::new(inner), + } + } +} + +#[pymethods] +impl Job { + #[getter] + pub fn id(&self) -> Option { + self.inner.id().map(str::to_string) + } + + pub fn wait(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.wait().await.infer_error()?; + Ok(()) + }) + } + + pub fn cancel(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.cancel().await.infer_error()?; + Ok(()) + }) + } +} diff --git a/python/src/lib.rs b/python/src/lib.rs index d480f9613..6c9f62df9 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -25,6 +25,7 @@ pub mod error; pub mod expr; pub mod header; pub mod index; +pub mod job; pub mod namespace; pub mod oauth; pub mod otel; @@ -44,6 +45,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/table.rs b/python/src/table.rs index 84ba91f39..a4e50c183 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -805,6 +805,37 @@ impl Table { }) } + #[pyo3(signature = (column, index=None, replace=None, wait_timeout=None, *, name=None, train=None))] + pub fn create_index_async<'a>( + self_: PyRef<'a, Self>, + column: String, + index: Option>, + replace: Option, + wait_timeout: Option>, + name: Option, + train: Option, + ) -> PyResult> { + let index = extract_index_params(&index)?; + let timeout = wait_timeout.map(|t| t.extract::().unwrap()); + let mut op = self_ + .inner_ref()? + .create_index_with_timeout(&[column], index, timeout); + if let Some(replace) = replace { + op = op.replace(replace); + } + if let Some(name) = name { + op = op.name(name); + } + if let Some(train) = train { + op = op.train(train); + } + + future_into_py(self_.py(), async move { + let job = op.execute_async().await.infer_error()?; + Ok(crate::job::Job::new(job)) + }) + } + pub fn drop_index(self_: PyRef<'_, Self>, index_name: String) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index 9e4dd1f8c..f6f596f3d 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::sync::PoisonError; +use std::fmt::{self, Display, Formatter}; +use std::sync::{Arc, PoisonError}; use arrow_schema::ArrowError; use datafusion_common::DataFusionError; @@ -9,6 +10,46 @@ use snafu::Snafu; pub(crate) type BoxError = Box; +/// Why a job failed, to whatever precision the backend provides. +/// +/// A job run in this process carries the error it failed with in [`Self::source`]. +/// A job run remotely carries whatever the server reported, which older servers +/// do not report at all. Every field is absent rather than invented when the +/// backend does not supply it. +#[derive(Debug, Clone, Default)] +pub struct JobFailure { + /// The stage the job was in, when known. + pub phase: Option, + /// A human-readable reason, when known. + pub message: Option, + /// Whether a retry could clear the failure, when known. + pub retryable: Option, + /// The error the job failed with, when it ran in this process. + pub source: Option>, +} + +impl JobFailure { + /// A failure whose only known detail is the error that caused it. + pub(crate) fn from_source(source: Arc) -> Self { + Self { + message: Some(source.to_string()), + source: Some(source), + ..Default::default() + } + } +} + +impl Display for JobFailure { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match (&self.message, &self.phase) { + (Some(message), Some(phase)) => write!(f, ": {message} (in {phase})"), + (Some(message), None) => write!(f, ": {message}"), + (None, Some(phase)) => write!(f, " in {phase}"), + (None, None) => Ok(()), + } + } +} + #[derive(Debug, Snafu)] #[snafu(visibility(pub(crate)))] pub enum Error { @@ -44,6 +85,13 @@ pub enum Error { Runtime { message: String }, #[snafu(display("Timeout error: {message}"))] Timeout { message: String }, + #[snafu(display("Job{} failed{failure}", job_id.as_ref().map(|id| format!(" {id}")).unwrap_or_default()))] + JobFailed { + job_id: Option, + failure: JobFailure, + }, + #[snafu(display("Job{} was cancelled", job_id.as_ref().map(|id| format!(" {id}")).unwrap_or_default()))] + JobCancelled { job_id: Option }, // 3rd party / external errors #[snafu(display("object_store error: {source}"))] diff --git a/rust/lancedb/src/index.rs b/rust/lancedb/src/index.rs index 308a64ca1..c693dc056 100644 --- a/rust/lancedb/src/index.rs +++ b/rust/lancedb/src/index.rs @@ -10,7 +10,7 @@ use std::time::Duration; use vector::IvfFlatIndexBuilder; use crate::index::vector::IvfRqIndexBuilder; -use crate::{DistanceType, Error, Result, table::BaseTable}; +use crate::{DistanceType, Error, Result, job::Job, table::BaseTable}; use self::{ scalar::{BTreeIndexBuilder, BitmapIndexBuilder, FmIndexBuilder, LabelListIndexBuilder}, @@ -305,6 +305,14 @@ impl IndexBuilder { pub async fn execute(self) -> Result<()> { self.parent.clone().create_index(self).await } + + /// Creates the index, returning a [`Job`] tracking the operation. + /// + /// The job may already be complete when returned, and callers must not + /// assume the index exists until [`Job::wait`] resolves. + pub async fn execute_async(self) -> Result { + self.parent.clone().create_index_async(self).await + } } #[derive(Debug, Clone, PartialEq, Deserialize)] diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs new file mode 100644 index 000000000..4c2c76d95 --- /dev/null +++ b/rust/lancedb/src/job.rs @@ -0,0 +1,158 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Handles to operations a server may run asynchronously. + +use std::sync::Arc; + +use async_trait::async_trait; +use tokio::sync::watch; +use tokio::task::{AbortHandle, JoinHandle}; + +use crate::error::{Error, JobFailure, Result}; + +/// Backend-specific tracking for an asynchronous operation. +#[async_trait] +pub(crate) trait JobHandle: Send + Sync { + /// Server-assigned id, when the backend has one. + fn id(&self) -> Option<&str> { + None + } + async fn wait(&self) -> Result<()>; + async fn cancel(&self) -> Result<()>; +} + +/// A handle to an operation that may still be running. +/// +/// The operation may already be complete when the handle is created. +pub struct Job { + handle: Option>, +} + +impl std::fmt::Debug for Job { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Job") + .field("id", &self.id()) + .field("done", &self.handle.is_none()) + .finish() + } +} + +impl Job { + /// A job whose operation finished before the handle was created. + pub(crate) fn new_done() -> Self { + Self { handle: None } + } + + pub(crate) fn new(handle: Box) -> Self { + Self { + handle: Some(handle), + } + } + + /// A job running as a task in this process. + pub(crate) fn spawned(task: JoinHandle>) -> Self { + Self::new(Box::new(SpawnedJob::new(task))) + } + + /// Identifies the operation on the server that is running it. + /// + /// Returned for correlating with server logs or the jobs API. Operations + /// that run in this process have no server id and return `None`. The + /// value is opaque: parsing it or storing it to resume the job later is + /// not supported. + pub fn id(&self) -> Option<&str> { + self.handle.as_ref().and_then(|handle| handle.id()) + } + + /// Waits until the operation reaches a terminal state. + /// + /// Returns [`crate::Error::JobFailed`] if the operation failed and + /// [`crate::Error::JobCancelled`] if it was cancelled. + pub async fn wait(&self) -> Result<()> { + match &self.handle { + None => Ok(()), + Some(handle) => handle.wait().await, + } + } + + /// Requests cancellation of the operation. + /// + /// Cancelling an operation that already finished is a no-op. + pub async fn cancel(&self) -> Result<()> { + match &self.handle { + None => Ok(()), + Some(handle) => handle.cancel().await, + } + } +} + +/// How an in-process operation ended. Cloneable so every waiter can be given +/// the outcome; [`Error`] is not, so failures share one behind an [`Arc`]. +#[derive(Clone)] +enum Outcome { + Succeeded, + Failed(Arc), + Cancelled, +} + +impl Outcome { + fn into_result(self) -> Result<()> { + match self { + Self::Succeeded => Ok(()), + Self::Failed(source) => Err(Error::JobFailed { + job_id: None, + failure: JobFailure::from_source(source), + }), + Self::Cancelled => Err(Error::JobCancelled { job_id: None }), + } + } +} + +/// Tracks an operation running as a task in this process. A second task +/// watches the first so that aborting it still produces an outcome, and so +/// that every caller of `wait` observes the same one. +struct SpawnedJob { + outcome: watch::Receiver>, + abort: AbortHandle, +} + +impl SpawnedJob { + fn new(task: JoinHandle>) -> Self { + let abort = task.abort_handle(); + let (tx, outcome) = watch::channel(None); + tokio::spawn(async move { + let outcome = match task.await { + Ok(Ok(())) => Outcome::Succeeded, + Ok(Err(err)) => Outcome::Failed(Arc::new(err)), + Err(err) if err.is_cancelled() => Outcome::Cancelled, + Err(err) => Outcome::Failed(Arc::new(Error::Runtime { + message: format!("index job task failed: {err}"), + })), + }; + let _ = tx.send(Some(outcome)); + }); + Self { outcome, abort } + } +} + +#[async_trait] +impl JobHandle for SpawnedJob { + async fn wait(&self) -> Result<()> { + let mut outcome = self.outcome.clone(); + let settled = outcome + .wait_for(|outcome| outcome.is_some()) + .await + .map_err(|_| Error::Runtime { + message: "index job outcome was dropped before it completed".to_string(), + })? + .clone() + .expect("wait_for returns once an outcome is set"); + settled.into_result() + } + + async fn cancel(&self) -> Result<()> { + self.abort.abort(); + Ok(()) + } +} diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 3cdf33615..70d023ccc 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -184,6 +184,7 @@ pub mod expr; pub mod index; pub mod io; pub mod ipc; +pub mod job; #[cfg(feature = "metrics-otel")] pub mod metrics_otel; #[cfg(feature = "polars")] @@ -203,7 +204,8 @@ use serde::{Deserialize, Serialize}; pub use blob::{BlobRangeRequest, blob, is_blob}; pub use connection::{ConnectNamespaceBuilder, Connection}; -pub use error::{Error, Result}; +pub use error::{Error, JobFailure, Result}; +pub use job::Job; use lance_index::vector::ApproxMode as LanceApproxMode; use lance_linalg::distance::DistanceType as LanceDistanceType; /// Re-export of the [`metrics`](https://docs.rs/metrics) crate facade. Enable diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 657ae1a47..4b5f8832f 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -8,6 +8,7 @@ pub(crate) mod client; pub(crate) mod db; +pub(crate) mod job; pub mod oauth; mod retry; pub(crate) mod table; diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs new file mode 100644 index 000000000..26bbfbbf9 --- /dev/null +++ b/rust/lancedb/src/remote/job.rs @@ -0,0 +1,147 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Tracking for server-side jobs through the `/v1/jobs` API. + +use std::time::Duration; + +use async_trait::async_trait; +use tokio::time::sleep; + +use serde::{Deserialize, Deserializer}; + +use crate::error::{Error, JobFailure, Result}; +use crate::job::JobHandle; +use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; + +/// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. +const INITIAL_POLL_INTERVAL: Duration = Duration::from_millis(200); +const MAX_POLL_INTERVAL: Duration = Duration::from_secs(5); + +#[derive(Debug, Clone, PartialEq, Eq)] +enum JobState { + InProgress, + Cancelled, + Failed, + Done, + /// A state this client version does not know; treated as still running + /// and reported as-is if the job never settles. + Other(String), +} + +impl<'de> Deserialize<'de> for JobState { + fn deserialize>(deserializer: D) -> std::result::Result { + Ok(Self::from(String::deserialize(deserializer)?.as_str())) + } +} + +impl From<&str> for JobState { + fn from(state: &str) -> Self { + match state { + "IN_PROGRESS" => Self::InProgress, + "CANCELLED" => Self::Cancelled, + "FAILED" => Self::Failed, + "DONE" => Self::Done, + other => Self::Other(other.to_string()), + } + } +} + +/// The server's account of why a job failed. Absent from older servers, which +/// report only the terminal state. +#[derive(Deserialize)] +struct ReportedFailure { + phase: String, + message: String, + retryable: bool, +} + +#[derive(Deserialize)] +struct DescribeJobResponse { + job_state: JobState, + #[serde(default)] + failure: Option, +} + +pub struct RemoteJob { + client: RestfulLanceDbClient, + job_id: String, +} + +impl RemoteJob { + pub fn new(client: RestfulLanceDbClient, job_id: String) -> Self { + Self { client, job_id } + } + + /// One `/v1/jobs/describe` round trip. + async fn describe(&self) -> Result { + let request = self + .client + .post("/v1/jobs/describe") + .json(&serde_json::json!({ "job_id": self.job_id })); + let (request_id, response) = self.client.send(request).await?; + let response = self.client.check_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + let description: DescribeJobResponse = + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("failed to parse job description: {}", e).into(), + request_id, + status_code: None, + })?; + Ok(description) + } +} + +#[async_trait] +impl JobHandle for RemoteJob { + fn id(&self) -> Option<&str> { + Some(&self.job_id) + } + + async fn wait(&self) -> Result<()> { + let mut interval = INITIAL_POLL_INTERVAL; + loop { + let description = self.describe().await?; + match description.job_state { + JobState::Done => return Ok(()), + JobState::Failed => { + return Err(Error::JobFailed { + job_id: Some(self.job_id.clone()), + failure: description + .failure + .map(|reported| JobFailure { + phase: Some(reported.phase), + message: Some(reported.message), + retryable: Some(reported.retryable), + source: None, + }) + .unwrap_or_default(), + }); + } + JobState::Cancelled => { + return Err(Error::JobCancelled { + job_id: Some(self.job_id.clone()), + }); + } + JobState::InProgress => {} + JobState::Other(ref state) => { + log::debug!("job {} is in unrecognized state {state}", self.job_id) + } + } + sleep(interval).await; + interval = (interval * 2).min(MAX_POLL_INTERVAL); + } + } + + async fn cancel(&self) -> Result<()> { + let request = self + .client + .post("/v1/jobs/cancel") + .json(&serde_json::json!({ "job_id": self.job_id })); + let (request_id, response) = self.client.send(request).await?; + self.client + .check_response(&request_id, response) + .await + .map(|_| ()) + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index b668f6d67..e9c200521 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -15,7 +15,9 @@ use crate::expr::expr_to_sql_string; use crate::index::Index; use crate::index::IndexStatistics; use crate::index::waiter::wait_for_index; +use crate::job::Job; use crate::query::{QueryFilter, QueryRequest, Select, VectorQueryRequest}; +use crate::remote::job::RemoteJob; use crate::table::AddColumnsResult; use crate::table::AddResult; use crate::table::BranchDiff; @@ -290,6 +292,120 @@ impl std::fmt::Debug for RemoteTable { } impl RemoteTable { + async fn submit_create_index(&self, mut index: IndexBuilder) -> Result> { + self.check_mutable().await?; + let request = self + .client + .post(&format!("/v1/table/{}/create_index/", self.identifier)); + + let column = match index.columns.len() { + 0 => { + return Err(Error::InvalidInput { + message: "No columns specified".into(), + }); + } + 1 => index.columns.pop().unwrap(), + _ => { + return Err(Error::NotSupported { + message: "Indices over multiple columns not yet supported".into(), + }); + } + }; + let schema = self.schema().await?; + let (canonical_column, field) = resolve_arrow_field_path(&schema, &column)?; + let mut body = serde_json::json!({ + "column": canonical_column + }); + + // Add name parameter if provided (for backwards compatibility, only include if Some) + if let Some(ref name) = index.name { + body["name"] = serde_json::Value::String(name.clone()); + } + + // Warn if train=false is specified since it's not meaningful + if !index.train { + log::warn!( + "train=false has no effect remote tables. The index will be created empty and automatically populated in the background." + ); + } + + fn to_json(params: &impl serde::Serialize) -> crate::Result { + serde_json::to_value(params).map_err(|e| Error::InvalidInput { + message: format!("failed to serialize index params {:?}", e), + }) + } + + // Map each Index variant to its wire type name and serializable params. + // Auto is special-cased since it needs schema inspection. + let (index_type_str, params) = match &index.index { + Index::IvfFlat(p) => ("IVF_FLAT", Some(to_json(p)?)), + Index::IvfPq(p) => ("IVF_PQ", Some(to_json(p)?)), + Index::IvfSq(p) => ("IVF_SQ", Some(to_json(p)?)), + Index::IvfHnswSq(p) => ("IVF_HNSW_SQ", Some(to_json(p)?)), + Index::IvfHnswFlat(p) => ("IVF_HNSW_FLAT", Some(to_json(p)?)), + Index::IvfRq(p) => ("IVF_RQ", Some(to_json(p)?)), + Index::BTree(p) => ("BTREE", Some(to_json(p)?)), + Index::Bitmap(p) => ("BITMAP", Some(to_json(p)?)), + Index::LabelList(p) => ("LABEL_LIST", Some(to_json(p)?)), + Index::Fm(p) => ("FM", Some(to_json(p)?)), + Index::FTS(p) => ("FTS", Some(to_json(p)?)), + Index::Auto => { + if supported_vector_data_type(field.data_type()) { + body[METRIC_TYPE_KEY] = + serde_json::Value::String(DistanceType::L2.to_string().to_lowercase()); + ("IVF_PQ", None) + } else if supported_btree_data_type(field.data_type()) { + ("BTREE", None) + } else { + return Err(Error::NotSupported { + message: format!( + "there are no indices supported for the field `{}` with the data type {}", + field.name(), + field.data_type() + ), + }); + } + } + _ => { + return Err(Error::NotSupported { + message: "Index type not supported".into(), + }); + } + }; + + body[INDEX_TYPE_KEY] = index_type_str.into(); + if let Some(params) = params { + for (key, value) in params.as_object().expect("params should be a JSON object") { + body[key] = value.clone(); + } + } + self.apply_branch_body(&mut body); + + let request = request.json(&body); + + let (request_id, response) = self.send(request, true).await?; + + let response = self.check_table_response(&request_id, response).await?; + let job_id = response + .text() + .await + .ok() + .and_then(|body| serde_json::from_str::(&body).ok()) + .and_then(|value| { + value + .get("job_id") + .and_then(|id| id.as_str()) + .map(str::to_string) + }); + + if let Some(wait_timeout) = index.wait_timeout { + let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); + self.wait_for_index(&[&index_name], wait_timeout).await?; + } + + Ok(job_id) + } + pub fn new( client: RestfulLanceDbClient, name: String, @@ -2259,107 +2375,15 @@ impl BaseTable for RemoteTable { Ok(delete_response) } - async fn create_index(&self, mut index: IndexBuilder) -> Result<()> { - self.check_mutable().await?; - let request = self - .client - .post(&format!("/v1/table/{}/create_index/", self.identifier)); + async fn create_index(&self, index: IndexBuilder) -> Result<()> { + self.submit_create_index(index).await.map(|_| ()) + } - let column = match index.columns.len() { - 0 => { - return Err(Error::InvalidInput { - message: "No columns specified".into(), - }); - } - 1 => index.columns.pop().unwrap(), - _ => { - return Err(Error::NotSupported { - message: "Indices over multiple columns not yet supported".into(), - }); - } - }; - let schema = self.schema().await?; - let (canonical_column, field) = resolve_arrow_field_path(&schema, &column)?; - let mut body = serde_json::json!({ - "column": canonical_column - }); - - // Add name parameter if provided (for backwards compatibility, only include if Some) - if let Some(ref name) = index.name { - body["name"] = serde_json::Value::String(name.clone()); - } - - // Warn if train=false is specified since it's not meaningful - if !index.train { - log::warn!( - "train=false has no effect remote tables. The index will be created empty and automatically populated in the background." - ); - } - - fn to_json(params: &impl serde::Serialize) -> crate::Result { - serde_json::to_value(params).map_err(|e| Error::InvalidInput { - message: format!("failed to serialize index params {:?}", e), - }) - } - - // Map each Index variant to its wire type name and serializable params. - // Auto is special-cased since it needs schema inspection. - let (index_type_str, params) = match &index.index { - Index::IvfFlat(p) => ("IVF_FLAT", Some(to_json(p)?)), - Index::IvfPq(p) => ("IVF_PQ", Some(to_json(p)?)), - Index::IvfSq(p) => ("IVF_SQ", Some(to_json(p)?)), - Index::IvfHnswSq(p) => ("IVF_HNSW_SQ", Some(to_json(p)?)), - Index::IvfHnswFlat(p) => ("IVF_HNSW_FLAT", Some(to_json(p)?)), - Index::IvfRq(p) => ("IVF_RQ", Some(to_json(p)?)), - Index::BTree(p) => ("BTREE", Some(to_json(p)?)), - Index::Bitmap(p) => ("BITMAP", Some(to_json(p)?)), - Index::LabelList(p) => ("LABEL_LIST", Some(to_json(p)?)), - Index::Fm(p) => ("FM", Some(to_json(p)?)), - Index::FTS(p) => ("FTS", Some(to_json(p)?)), - Index::Auto => { - if supported_vector_data_type(field.data_type()) { - body[METRIC_TYPE_KEY] = - serde_json::Value::String(DistanceType::L2.to_string().to_lowercase()); - ("IVF_PQ", None) - } else if supported_btree_data_type(field.data_type()) { - ("BTREE", None) - } else { - return Err(Error::NotSupported { - message: format!( - "there are no indices supported for the field `{}` with the data type {}", - field.name(), - field.data_type() - ), - }); - } - } - _ => { - return Err(Error::NotSupported { - message: "Index type not supported".into(), - }); - } - }; - - body[INDEX_TYPE_KEY] = index_type_str.into(); - if let Some(params) = params { - for (key, value) in params.as_object().expect("params should be a JSON object") { - body[key] = value.clone(); - } - } - self.apply_branch_body(&mut body); - - let request = request.json(&body); - - let (request_id, response) = self.send(request, true).await?; - - self.check_table_response(&request_id, response).await?; - - if let Some(wait_timeout) = index.wait_timeout { - let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); - self.wait_for_index(&[&index_name], wait_timeout).await?; - } - - Ok(()) + async fn create_index_async(&self, index: IndexBuilder) -> Result { + Ok(match self.submit_create_index(index).await? { + Some(job_id) => Job::new(Box::new(RemoteJob::new(self.client.clone(), job_id))), + None => Job::new_done(), + }) } /// Poll until the columns are fully indexed. Will return Error::Timeout if the columns @@ -5154,6 +5178,224 @@ mod tests { } } + #[tokio::test] + async fn test_create_index_returns_job() { + let describe_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let describe_calls_in_handler = describe_calls.clone(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + 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": "job-123"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["job_id"], "job-123"); + let state = if describe_calls_in_handler + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0 + { + "IN_PROGRESS" + } else { + "DONE" + }; + http::Response::builder() + .status(200) + .body(format!( + r#"{{"job_id": "job-123", "job_state": "{state}"}}"# + )) + .unwrap() + } + "/v1/jobs/cancel" => { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["job_id"], "job-123"); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + assert_eq!(job.id(), Some("job-123")); + job.wait().await.unwrap(); + assert_eq!(describe_calls.load(std::sync::atomic::Ordering::SeqCst), 2); + job.cancel().await.unwrap(); + } + + /// An unrecognized state is treated as still running, so the client keeps + /// polling rather than reporting a wrong terminal outcome. + #[tokio::test] + async fn test_job_wait_treats_unknown_state_as_running() { + let calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let calls_in_handler = calls.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": "job-unknown"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => { + let state = if calls_in_handler + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0 + { + "SOMETHING_NEW" + } else { + "DONE" + }; + http::Response::builder() + .status(200) + .body(format!( + r#"{{"job_id": "job-unknown", "job_state": "{state}"}}"# + )) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + job.wait().await.unwrap(); + assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn test_job_wait_surfaces_failure() { + 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": "job-err"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body(r#"{"job_id": "job-err", "job_state": "FAILED"}"#.to_string()) + .unwrap(), + path => panic!("Unexpected path: {}", path), + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + let err = job.wait().await.unwrap_err(); + let crate::Error::JobFailed { failure, .. } = &err else { + panic!("expected JobFailed, got {err:?}"); + }; + // The server said only that it failed, so nothing may be invented. + assert!(failure.message.is_none(), "{failure:?}"); + assert!(failure.phase.is_none(), "{failure:?}"); + assert!(failure.retryable.is_none(), "{failure:?}"); + assert_eq!(err.to_string(), "Job job-err failed"); + } + + /// A server that reports why the job failed has that reason surfaced + /// verbatim rather than replaced with a generic message. + #[tokio::test] + async fn test_job_wait_reports_the_server_failure_reason() { + 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": "job-err"}"#.to_string()) + .unwrap(), + "/v1/jobs/describe" => http::Response::builder() + .status(200) + .body( + r#"{"job_id": "job-err", "job_state": "FAILED", "failure": {"phase": "commit", "message": "preempted", "retryable": true}}"# + .to_string(), + ) + .unwrap(), + path => panic!("Unexpected path: {}", path), + } + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + let err = job.wait().await.unwrap_err(); + let crate::Error::JobFailed { failure, .. } = &err else { + panic!("expected JobFailed, got {err:?}"); + }; + assert_eq!(failure.message.as_deref(), Some("preempted")); + assert_eq!(failure.phase.as_deref(), Some("commit")); + assert_eq!(failure.retryable, Some(true)); + assert_eq!(err.to_string(), "Job job-err failed: preempted (in commit)"); + } + + /// Servers that return no job id (e.g. an empty create-index response) + /// yield an already-done job. + #[tokio::test] + async fn test_create_index_without_job_id_is_done() { + 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("{}".to_string()) + .unwrap(), + path => panic!("Unexpected path: {}", path), + }); + + let job = table + .create_index(&["a"], Index::BTree(Default::default())) + .execute_async() + .await + .unwrap(); + job.wait().await.unwrap(); + job.cancel().await.unwrap(); + } + #[tokio::test] async fn test_create_index_nested_field_paths() { let schema = nested_index_schema(); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index e0f12e2ff..1c9d68f4c 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -57,6 +57,7 @@ use crate::error::{Error, Result}; use crate::index::IndexStatistics; use crate::index::{Index, IndexBuilder}; use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType}; +use crate::job::Job; use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery}; use crate::table::datafusion::insert::InsertExec; use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path}; @@ -613,6 +614,9 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { async fn update(&self, update: UpdateBuilder) -> Result; /// Create an index on the provided column(s). async fn create_index(&self, index: IndexBuilder) -> Result<()>; + + /// Starts index creation, returning a handle to the resulting job. + async fn create_index_async(&self, index: IndexBuilder) -> Result; /// List the indices on the table. async fn list_indices(&self) -> Result>; /// Drop an index from the table. @@ -3061,29 +3065,18 @@ impl BaseTable for NativeTable { } async fn create_index(&self, opts: IndexBuilder) -> Result<()> { - if opts.columns.len() != 1 { - return Err(Error::Schema { - message: "Multi-column (composite) indices are not yet supported".to_string(), - }); - } - self.dataset.ensure_mutable()?; - let mut dataset = (*self.dataset.get().await?).clone(); - let (column, field) = Self::resolve_index_field(dataset.schema(), &opts.columns[0])?; + let prepared = self.prepare_index(&opts).await?; + self.build_index(opts, prepared).await + } - let lance_idx_params = self.make_index_params(&field, opts.index.clone()).await?; - let index_type = self.get_index_type_for_field(&field, &opts.index); - let columns = [column.as_str()]; - let mut builder = dataset - .create_index_builder(&columns, index_type, lance_idx_params.as_ref()) - .train(opts.train) - .replace(opts.replace); - - if let Some(name) = opts.name { - builder = builder.name(name); - } - builder.await?; - self.dataset.update(dataset); - Ok(()) + async fn create_index_async(&self, opts: IndexBuilder) -> Result { + // Prepare before spawning so bad input is reported by this call rather + // than only by the job. + let prepared = self.prepare_index(&opts).await?; + let table = self.clone(); + Ok(Job::spawned(tokio::spawn(async move { + table.build_index(opts, prepared).await + }))) } async fn drop_index(&self, index_name: &str) -> Result<()> { diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index ab412a502..144c6dbfb 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -21,6 +21,9 @@ use lance_index::vector::pq::PQBuildParams; use lance_index::vector::sq::builder::SQBuildParams; use crate::error::{Error, Result}; + +/// Resolved column, index parameters and index type for one build. +pub(super) type PreparedIndex = (String, Box, IndexType); use crate::index::Index; use crate::index::vector::{VectorIndex, suggested_num_sub_vectors}; use crate::utils::{ @@ -105,6 +108,47 @@ impl NativeTable { } } + /// Resolves the target column and index parameters, erroring on input the + /// build would reject. + pub(super) async fn prepare_index( + &self, + opts: &crate::index::IndexBuilder, + ) -> Result { + if opts.columns.len() != 1 { + return Err(Error::Schema { + message: "Multi-column (composite) indices are not yet supported".to_string(), + }); + } + self.dataset.ensure_mutable()?; + let dataset = self.dataset.get().await?; + let (column, field) = Self::resolve_index_field(dataset.schema(), &opts.columns[0])?; + let params = self.make_index_params(&field, opts.index.clone()).await?; + let index_type = self.get_index_type_for_field(&field, &opts.index); + Ok((column, params, index_type)) + } + + /// Builds a prepared index and publishes the new dataset version. + pub(super) async fn build_index( + &self, + opts: crate::index::IndexBuilder, + prepared: PreparedIndex, + ) -> Result<()> { + let (column, lance_idx_params, index_type) = prepared; + let mut dataset = (*self.dataset.get().await?).clone(); + let columns = [column.as_str()]; + let mut builder = dataset + .create_index_builder(&columns, index_type, lance_idx_params.as_ref()) + .train(opts.train) + .replace(opts.replace); + + if let Some(name) = opts.name { + builder = builder.name(name); + } + builder.await?; + self.dataset.update(dataset); + Ok(()) + } + pub(super) fn resolve_index_field( schema: &lance_core::datatypes::Schema, column: &str, @@ -475,6 +519,208 @@ mod tests { assert_eq!(table.list_indices().await.unwrap().len(), 0); } + #[tokio::test] + async fn test_execute_async_job_waits_for_local_build() { + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let batch = record_batch!(("id", Int32, (0..512).collect::>())).unwrap(); + let table = conn.create_table("t", batch).execute().await.unwrap(); + + let job = table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .execute_async() + .await + .unwrap(); + // Local jobs run in this process and have no server id. + assert_eq!(job.id(), None); + // The build runs as a task, so the index need not exist yet; it must + // once the job resolves. + job.wait().await.unwrap(); + assert_eq!(table.list_indices().await.unwrap().len(), 1); + // Cancelling a finished job is a no-op. + job.cancel().await.unwrap(); + } + + /// Concurrent waiters, and a wait issued after the job settled, all + /// succeed once the build does. + #[tokio::test] + async fn test_execute_async_job_reports_success_to_every_waiter() { + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let batch = record_batch!(("id", Int32, (0..512).collect::>())).unwrap(); + let table = conn.create_table("t", batch).execute().await.unwrap(); + + let job = Arc::new( + table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .execute_async() + .await + .unwrap(), + ); + let waiters = (0..4) + .map(|_| { + let job = job.clone(); + tokio::spawn(async move { job.wait().await }) + }) + .collect::>(); + for waiter in waiters { + waiter.await.unwrap().unwrap(); + } + // A wait after the job settled still reports the same outcome. + job.wait().await.unwrap(); + assert_eq!(table.list_indices().await.unwrap().len(), 1); + } + + /// Every waiter sees a failure, not just the first: a waiter that missed + /// the outcome would be told the job succeeded. + #[tokio::test] + async fn test_execute_async_job_reports_failure_to_every_waiter() { + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let batch = record_batch!(("id", Int32, (0..512).collect::>())).unwrap(); + let table = conn.create_table("t", batch).execute().await.unwrap(); + table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .execute() + .await + .unwrap(); + + // Rebuilding the same index without replace fails once the build + // starts, so the failure reaches the job rather than execute_async. + let job = Arc::new( + table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .replace(false) + .execute_async() + .await + .unwrap(), + ); + let waiters = (0..3) + .map(|_| { + let job = job.clone(); + tokio::spawn(async move { job.wait().await }) + }) + .collect::>(); + for waiter in waiters { + waiter + .await + .unwrap() + .expect_err("every waiter must see the failure"); + } + job.wait().await.expect_err("a later wait still fails"); + } + + /// A local failure keeps the error it failed with, so a caller can match on + /// the original variant rather than parse a message. + #[tokio::test] + async fn test_execute_async_failure_keeps_the_source_error() { + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let batch = record_batch!(("id", Int32, (0..512).collect::>())).unwrap(); + let table = conn.create_table("t", batch).execute().await.unwrap(); + table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .execute() + .await + .unwrap(); + + let job = table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .replace(false) + .execute_async() + .await + .unwrap(); + + let crate::Error::JobFailed { failure, .. } = job.wait().await.unwrap_err() else { + panic!("a failed job reports JobFailed"); + }; + let source = failure.source.expect("a local failure carries its error"); + assert_eq!( + failure.message.as_deref(), + Some(source.to_string()).as_deref() + ); + // Nothing local can report these, so they must be absent rather than invented. + assert!(failure.phase.is_none()); + assert!(failure.retryable.is_none()); + } + + /// Every waiter sees the cancellation, including ones that were already + /// waiting when the cancel landed. + #[tokio::test] + async fn test_execute_async_job_reports_cancellation_to_every_waiter() { + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let batch = record_batch!(("id", Int32, (0..512).collect::>())).unwrap(); + let table = conn.create_table("t", batch).execute().await.unwrap(); + + let job = Arc::new( + table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .execute_async() + .await + .unwrap(), + ); + // Cancel before yielding, so the build cannot have started and the + // outcome is always the cancellation. + job.cancel().await.unwrap(); + + let waiters = (0..2) + .map(|_| { + let job = job.clone(); + tokio::spawn(async move { job.wait().await }) + }) + .collect::>(); + for waiter in waiters { + match waiter.await.unwrap() { + Err(crate::Error::JobCancelled { .. }) => {} + other => panic!("expected the cancellation, got {other:?}"), + } + } + match job.wait().await { + Err(crate::Error::JobCancelled { .. }) => {} + other => panic!("expected the cancellation, got {other:?}"), + } + } + + #[tokio::test] + async fn test_execute_async_job_cancel_stops_local_build() { + let tmp_dir = tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let batch = record_batch!(("id", Int32, (0..512).collect::>())).unwrap(); + let table = conn.create_table("t", batch).execute().await.unwrap(); + + let job = table + .create_index(&["id"], Index::BTree(BTreeIndexBuilder::default())) + .execute_async() + .await + .unwrap(); + job.cancel().await.unwrap(); + match job.wait().await { + Err(crate::Error::JobCancelled { .. }) => {} + // The build may finish before the abort lands. + Ok(()) => {} + other => panic!("unexpected job outcome: {other:?}"), + } + } + #[tokio::test] async fn test_ivf_pq_uses_default_partition_size_for_num_partitions() { use crate::index::vector::IvfPqIndexBuilder;