mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
feat: create_index returns a Job handle (#3742)
IndexBuilder::execute now returns a Job with wait and cancel methods. Local tables build the index synchronously and return an already-done job. Remote tables read the job id the server returns from create_index and track it through the /v1/jobs API: wait polls describe until the job reaches a terminal state and cancel posts a cancellation. Servers that return no job id yield a done job, so behavior against older servers is unchanged. The job id is not exposed on the handle. The Python and TypeScript bindings keep their current signatures and discard the handle; exposing Job there is left to follow-ups. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -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<void>
|
||||
```
|
||||
|
||||
Request cancellation. Cancelling a finished operation is a no-op.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`void`>
|
||||
|
||||
***
|
||||
|
||||
### wait()
|
||||
|
||||
```ts
|
||||
wait(): Promise<void>
|
||||
```
|
||||
|
||||
Wait until the operation reaches a terminal state.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`void`>
|
||||
@@ -295,6 +295,29 @@ await table.createIndex("my_float_col");
|
||||
|
||||
***
|
||||
|
||||
### createIndexAsync()
|
||||
|
||||
```ts
|
||||
abstract createIndexAsync(column, options?): Promise<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 [Job.wait](Job.md#wait) resolves.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **column**: `string`
|
||||
|
||||
* **options?**: `Partial`<[`IndexOptions`](../interfaces/IndexOptions.md)>
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`Job`](Job.md)>
|
||||
|
||||
***
|
||||
|
||||
### currentBranch()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -85,7 +85,7 @@ export {
|
||||
RenameTableOptions,
|
||||
} from "./connection";
|
||||
|
||||
export { Session } from "./native.js";
|
||||
export { Job, Session } from "./native.js";
|
||||
|
||||
export {
|
||||
ExecutableQuery,
|
||||
|
||||
@@ -30,6 +30,7 @@ import {
|
||||
DropColumnsResult,
|
||||
IndexConfig,
|
||||
IndexStatistics,
|
||||
Job,
|
||||
Branches as NativeBranches,
|
||||
OptimizeStats,
|
||||
TableStatistics,
|
||||
@@ -358,6 +359,17 @@ export abstract class Table {
|
||||
options?: Partial<IndexOptions>,
|
||||
): Promise<void>;
|
||||
|
||||
/**
|
||||
* 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<IndexOptions>,
|
||||
): Promise<Job>;
|
||||
|
||||
/**
|
||||
* Drop an index from the table.
|
||||
*
|
||||
@@ -940,6 +952,22 @@ export class LocalTable extends Table {
|
||||
);
|
||||
}
|
||||
|
||||
async createIndexAsync(
|
||||
column: string,
|
||||
options?: Partial<IndexOptions>,
|
||||
): Promise<Job> {
|
||||
// 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<void> {
|
||||
await this.inner.dropIndex(name);
|
||||
}
|
||||
|
||||
@@ -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<lancedb::Job>,
|
||||
}
|
||||
|
||||
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<String> {
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ mod error;
|
||||
mod header;
|
||||
mod index;
|
||||
mod iterator;
|
||||
mod job;
|
||||
pub mod merge;
|
||||
pub mod otel;
|
||||
pub mod permutation;
|
||||
|
||||
@@ -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<bool>,
|
||||
wait_timeout_s: Option<i64>,
|
||||
name: Option<String>,
|
||||
train: Option<bool>,
|
||||
) -> napi::Result<crate::job::Job> {
|
||||
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()?
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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]): ...
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 = []
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -102,6 +102,18 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
|
||||
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(),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -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<lancedb::Job>,
|
||||
}
|
||||
|
||||
impl Job {
|
||||
pub(crate) fn new(inner: lancedb::Job) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(inner),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Job {
|
||||
#[getter]
|
||||
pub fn id(&self) -> Option<String> {
|
||||
self.inner.id().map(str::to_string)
|
||||
}
|
||||
|
||||
pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.cancel().await.infer_error()?;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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::<Connection>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<crate::job::Job>()?;
|
||||
m.add_class::<PyBlobFile>()?;
|
||||
m.add_class::<IndexConfig>()?;
|
||||
m.add_class::<Query>()?;
|
||||
|
||||
@@ -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<Bound<'_, PyAny>>,
|
||||
replace: Option<bool>,
|
||||
wait_timeout: Option<Bound<'_, PyAny>>,
|
||||
name: Option<String>,
|
||||
train: Option<bool>,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let index = extract_index_params(&index)?;
|
||||
let timeout = wait_timeout.map(|t| t.extract::<std::time::Duration>().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<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
|
||||
@@ -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<dyn std::error::Error + Send + Sync>;
|
||||
|
||||
/// 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<String>,
|
||||
/// A human-readable reason, when known.
|
||||
pub message: Option<String>,
|
||||
/// Whether a retry could clear the failure, when known.
|
||||
pub retryable: Option<bool>,
|
||||
/// The error the job failed with, when it ran in this process.
|
||||
pub source: Option<Arc<Error>>,
|
||||
}
|
||||
|
||||
impl JobFailure {
|
||||
/// A failure whose only known detail is the error that caused it.
|
||||
pub(crate) fn from_source(source: Arc<Error>) -> 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<String>,
|
||||
failure: JobFailure,
|
||||
},
|
||||
#[snafu(display("Job{} was cancelled", job_id.as_ref().map(|id| format!(" {id}")).unwrap_or_default()))]
|
||||
JobCancelled { job_id: Option<String> },
|
||||
|
||||
// 3rd party / external errors
|
||||
#[snafu(display("object_store error: {source}"))]
|
||||
|
||||
@@ -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<Job> {
|
||||
self.parent.clone().create_index_async(self).await
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Deserialize)]
|
||||
|
||||
@@ -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<Box<dyn JobHandle>>,
|
||||
}
|
||||
|
||||
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<dyn JobHandle>) -> Self {
|
||||
Self {
|
||||
handle: Some(handle),
|
||||
}
|
||||
}
|
||||
|
||||
/// A job running as a task in this process.
|
||||
pub(crate) fn spawned(task: JoinHandle<Result<()>>) -> 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<Error>),
|
||||
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<Option<Outcome>>,
|
||||
abort: AbortHandle,
|
||||
}
|
||||
|
||||
impl SpawnedJob {
|
||||
fn new(task: JoinHandle<Result<()>>) -> 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(())
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
|
||||
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<ReportedFailure>,
|
||||
}
|
||||
|
||||
pub struct RemoteJob<S: HttpSend> {
|
||||
client: RestfulLanceDbClient<S>,
|
||||
job_id: String,
|
||||
}
|
||||
|
||||
impl<S: HttpSend> RemoteJob<S> {
|
||||
pub fn new(client: RestfulLanceDbClient<S>, job_id: String) -> Self {
|
||||
Self { client, job_id }
|
||||
}
|
||||
|
||||
/// One `/v1/jobs/describe` round trip.
|
||||
async fn describe(&self) -> Result<DescribeJobResponse> {
|
||||
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<S: HttpSend> JobHandle for RemoteJob<S> {
|
||||
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(|_| ())
|
||||
}
|
||||
}
|
||||
+342
-100
@@ -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<S: HttpSend> std::fmt::Debug for RemoteTable<S> {
|
||||
}
|
||||
|
||||
impl<S: HttpSend> RemoteTable<S> {
|
||||
async fn submit_create_index(&self, mut index: IndexBuilder) -> Result<Option<String>> {
|
||||
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::Value> {
|
||||
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::<serde_json::Value>(&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<S>,
|
||||
name: String,
|
||||
@@ -2259,107 +2375,15 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
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::Value> {
|
||||
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<Job> {
|
||||
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();
|
||||
|
||||
+15
-22
@@ -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<UpdateResult>;
|
||||
/// 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<Job>;
|
||||
/// List the indices on the table.
|
||||
async fn list_indices(&self) -> Result<Vec<IndexConfig>>;
|
||||
/// 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<Job> {
|
||||
// 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<()> {
|
||||
|
||||
@@ -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<dyn lance::index::IndexParams>, 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<PreparedIndex> {
|
||||
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::<Vec<_>>())).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::<Vec<_>>())).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::<Vec<_>>();
|
||||
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::<Vec<_>>())).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::<Vec<_>>();
|
||||
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::<Vec<_>>())).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::<Vec<_>>())).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::<Vec<_>>();
|
||||
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::<Vec<_>>())).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;
|
||||
|
||||
Reference in New Issue
Block a user