mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-27 16:38:31 +00:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 011def461c | |||
| ed6be12ad6 | |||
| ac2b689cdb | |||
| 4fc8114871 |
Generated
+118
-411
File diff suppressed because it is too large
Load Diff
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
|||||||
rust-version = "1.91.0"
|
rust-version = "1.91.0"
|
||||||
|
|
||||||
[workspace.dependencies]
|
[workspace.dependencies]
|
||||||
lance = { "version" = "=10.1.0-beta.1", default-features = false, features = ["geo"], "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance = { "version" = "=10.0.0-beta.5", default-features = false, "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-core = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-core = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-datagen = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-datagen = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-file = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-file = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-io = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-io = { "version" = "=10.0.0-beta.5", default-features = false, "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-index = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-index = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-linalg = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-linalg = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-namespace = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-namespace = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-namespace-impls = { "version" = "=10.1.0-beta.1", default-features = false, "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-namespace-impls = { "version" = "=10.0.0-beta.5", default-features = false, "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-table = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-table = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-testing = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-testing = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-datafusion = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-datafusion = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-encoding = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-encoding = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
lance-arrow = { "version" = "=10.1.0-beta.1", "tag" = "v10.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
lance-arrow = { "version" = "=10.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||||
ahash = "0.8"
|
ahash = "0.8"
|
||||||
# Note that this one does not include pyarrow
|
# Note that this one does not include pyarrow
|
||||||
arrow = { version = "58.0.0", optional = false }
|
arrow = { version = "58.0.0", optional = false }
|
||||||
|
|||||||
@@ -25,27 +25,6 @@ the underlying connection has been closed.
|
|||||||
|
|
||||||
## Methods
|
## Methods
|
||||||
|
|
||||||
### cancelJob()
|
|
||||||
|
|
||||||
```ts
|
|
||||||
abstract cancelJob(jobId): Promise<boolean>
|
|
||||||
```
|
|
||||||
|
|
||||||
Request cancellation of a server-side job by id.
|
|
||||||
|
|
||||||
Resolves to true if the server accepted the cancellation, false if no
|
|
||||||
such job exists. Cancelling an already-terminal job is a no-op success.
|
|
||||||
|
|
||||||
#### Parameters
|
|
||||||
|
|
||||||
* **jobId**: `string`
|
|
||||||
|
|
||||||
#### Returns
|
|
||||||
|
|
||||||
`Promise`<`boolean`>
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### cloneTable()
|
### cloneTable()
|
||||||
|
|
||||||
```ts
|
```ts
|
||||||
@@ -386,26 +365,6 @@ Drop an existing table.
|
|||||||
|
|
||||||
***
|
***
|
||||||
|
|
||||||
### getJob()
|
|
||||||
|
|
||||||
```ts
|
|
||||||
abstract getJob(jobId): Promise<null | JobDescription>
|
|
||||||
```
|
|
||||||
|
|
||||||
Describe a single server-side job by id.
|
|
||||||
|
|
||||||
Resolves to `null` when the server has no such job.
|
|
||||||
|
|
||||||
#### Parameters
|
|
||||||
|
|
||||||
* **jobId**: `string`
|
|
||||||
|
|
||||||
#### Returns
|
|
||||||
|
|
||||||
`Promise`<`null` \| [`JobDescription`](../interfaces/JobDescription.md)>
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### isOpen()
|
### isOpen()
|
||||||
|
|
||||||
```ts
|
```ts
|
||||||
@@ -420,62 +379,6 @@ Return true if the connection has not been closed
|
|||||||
|
|
||||||
***
|
***
|
||||||
|
|
||||||
### job()
|
|
||||||
|
|
||||||
```ts
|
|
||||||
abstract job(jobId): Job
|
|
||||||
```
|
|
||||||
|
|
||||||
A [Job](Job.md) handle for a server-side job by id.
|
|
||||||
|
|
||||||
The handle is constructed without a server round trip; an unknown id
|
|
||||||
surfaces when the handle is used. Dropping the handle has no effect on
|
|
||||||
the job itself.
|
|
||||||
|
|
||||||
#### Parameters
|
|
||||||
|
|
||||||
* **jobId**: `string`
|
|
||||||
|
|
||||||
#### Returns
|
|
||||||
|
|
||||||
[`Job`](Job.md)
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### jobHistory()
|
|
||||||
|
|
||||||
```ts
|
|
||||||
abstract jobHistory(jobId?): Promise<Table<any>>
|
|
||||||
```
|
|
||||||
|
|
||||||
The lifecycle event history of a server-side job, as an Arrow table.
|
|
||||||
|
|
||||||
Lists history across all jobs when `jobId` is omitted.
|
|
||||||
|
|
||||||
#### Parameters
|
|
||||||
|
|
||||||
* **jobId?**: `string`
|
|
||||||
|
|
||||||
#### Returns
|
|
||||||
|
|
||||||
`Promise`<`Table`<`any`>>
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### listJobs()
|
|
||||||
|
|
||||||
```ts
|
|
||||||
abstract listJobs(): Promise<JobInfo[]>
|
|
||||||
```
|
|
||||||
|
|
||||||
List server-side jobs across the database's tables.
|
|
||||||
|
|
||||||
#### Returns
|
|
||||||
|
|
||||||
`Promise`<[`JobInfo`](../interfaces/JobInfo.md)[]>
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### listNamespaces()
|
### listNamespaces()
|
||||||
|
|
||||||
```ts
|
```ts
|
||||||
|
|||||||
@@ -1,83 +0,0 @@
|
|||||||
[**@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`>
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### status()
|
|
||||||
|
|
||||||
```ts
|
|
||||||
status(): Promise<string>
|
|
||||||
```
|
|
||||||
|
|
||||||
The operation's current lifecycle state: "running", "finished",
|
|
||||||
"failed", or "cancelled".
|
|
||||||
|
|
||||||
A point snapshot; unlike [Job.wait](Job.md#wait) it does not block or reject
|
|
||||||
on a terminal failure state. States a newer server reports that this
|
|
||||||
client version does not know pass through as-is.
|
|
||||||
|
|
||||||
#### Returns
|
|
||||||
|
|
||||||
`Promise`<`string`>
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### wait()
|
|
||||||
|
|
||||||
```ts
|
|
||||||
wait(): Promise<void>
|
|
||||||
```
|
|
||||||
|
|
||||||
Wait until the operation reaches a terminal state.
|
|
||||||
|
|
||||||
#### Returns
|
|
||||||
|
|
||||||
`Promise`<`void`>
|
|
||||||
@@ -295,29 +295,6 @@ 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()
|
### currentBranch()
|
||||||
|
|
||||||
```ts
|
```ts
|
||||||
|
|||||||
@@ -25,7 +25,6 @@
|
|||||||
- [Connection](classes/Connection.md)
|
- [Connection](classes/Connection.md)
|
||||||
- [HeaderProvider](classes/HeaderProvider.md)
|
- [HeaderProvider](classes/HeaderProvider.md)
|
||||||
- [Index](classes/Index.md)
|
- [Index](classes/Index.md)
|
||||||
- [Job](classes/Job.md)
|
|
||||||
- [MakeArrowTableOptions](classes/MakeArrowTableOptions.md)
|
- [MakeArrowTableOptions](classes/MakeArrowTableOptions.md)
|
||||||
- [MatchQuery](classes/MatchQuery.md)
|
- [MatchQuery](classes/MatchQuery.md)
|
||||||
- [MergeInsertBuilder](classes/MergeInsertBuilder.md)
|
- [MergeInsertBuilder](classes/MergeInsertBuilder.md)
|
||||||
@@ -89,9 +88,6 @@
|
|||||||
- [IvfFlatOptions](interfaces/IvfFlatOptions.md)
|
- [IvfFlatOptions](interfaces/IvfFlatOptions.md)
|
||||||
- [IvfPqOptions](interfaces/IvfPqOptions.md)
|
- [IvfPqOptions](interfaces/IvfPqOptions.md)
|
||||||
- [IvfRqOptions](interfaces/IvfRqOptions.md)
|
- [IvfRqOptions](interfaces/IvfRqOptions.md)
|
||||||
- [JobDescription](interfaces/JobDescription.md)
|
|
||||||
- [JobFailureInfo](interfaces/JobFailureInfo.md)
|
|
||||||
- [JobInfo](interfaces/JobInfo.md)
|
|
||||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
||||||
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
|
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
|
||||||
|
|||||||
@@ -1,66 +0,0 @@
|
|||||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
[@lancedb/lancedb](../globals.md) / JobDescription
|
|
||||||
|
|
||||||
# Interface: JobDescription
|
|
||||||
|
|
||||||
A described job from `Connection.getJob`.
|
|
||||||
|
|
||||||
## Properties
|
|
||||||
|
|
||||||
### creationMs
|
|
||||||
|
|
||||||
```ts
|
|
||||||
creationMs: number;
|
|
||||||
```
|
|
||||||
|
|
||||||
When the job was created, in milliseconds since the epoch.
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### failure?
|
|
||||||
|
|
||||||
```ts
|
|
||||||
optional failure: JobFailureInfo;
|
|
||||||
```
|
|
||||||
|
|
||||||
Why the job failed, when the job is failed and the server reports a
|
|
||||||
reason.
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### jobId
|
|
||||||
|
|
||||||
```ts
|
|
||||||
jobId: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### jobType
|
|
||||||
|
|
||||||
```ts
|
|
||||||
jobType: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### specJson?
|
|
||||||
|
|
||||||
```ts
|
|
||||||
optional specJson: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
The job-type-specific specification as a JSON string, when present.
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### state
|
|
||||||
|
|
||||||
```ts
|
|
||||||
state: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
Lifecycle state: "running", "finished", "failed", or "cancelled".
|
|
||||||
@@ -1,33 +0,0 @@
|
|||||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
[@lancedb/lancedb](../globals.md) / JobFailureInfo
|
|
||||||
|
|
||||||
# Interface: JobFailureInfo
|
|
||||||
|
|
||||||
The server's account of why a job failed.
|
|
||||||
|
|
||||||
## Properties
|
|
||||||
|
|
||||||
### message?
|
|
||||||
|
|
||||||
```ts
|
|
||||||
optional message: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### phase?
|
|
||||||
|
|
||||||
```ts
|
|
||||||
optional phase: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### retryable?
|
|
||||||
|
|
||||||
```ts
|
|
||||||
optional retryable: boolean;
|
|
||||||
```
|
|
||||||
@@ -1,58 +0,0 @@
|
|||||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
[@lancedb/lancedb](../globals.md) / JobInfo
|
|
||||||
|
|
||||||
# Interface: JobInfo
|
|
||||||
|
|
||||||
A row from `Connection.listJobs`: one server-side job.
|
|
||||||
|
|
||||||
## Properties
|
|
||||||
|
|
||||||
### createdAtMillis
|
|
||||||
|
|
||||||
```ts
|
|
||||||
createdAtMillis: number;
|
|
||||||
```
|
|
||||||
|
|
||||||
When the job was created, in milliseconds since the epoch.
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### jobId
|
|
||||||
|
|
||||||
```ts
|
|
||||||
jobId: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
The job id -- what `Connection.getJob` and `Connection.cancelJob`
|
|
||||||
accept.
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### jobType
|
|
||||||
|
|
||||||
```ts
|
|
||||||
jobType: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### state
|
|
||||||
|
|
||||||
```ts
|
|
||||||
state: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
Lifecycle state: "running", "finished", "failed", or "cancelled".
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### table
|
|
||||||
|
|
||||||
```ts
|
|
||||||
table: string;
|
|
||||||
```
|
|
||||||
|
|
||||||
The table the job runs against, without URI or namespace.
|
|
||||||
+1
-1
@@ -28,7 +28,7 @@
|
|||||||
<properties>
|
<properties>
|
||||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||||
<arrow.version>15.0.0</arrow.version>
|
<arrow.version>15.0.0</arrow.version>
|
||||||
<lance-core.version>10.1.0-beta.1</lance-core.version>
|
<lance-core.version>10.0.0-beta.5</lance-core.version>
|
||||||
<spotless.skip>false</spotless.skip>
|
<spotless.skip>false</spotless.skip>
|
||||||
<spotless.version>2.30.0</spotless.version>
|
<spotless.version>2.30.0</spotless.version>
|
||||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||||
|
|||||||
@@ -877,96 +877,3 @@ describe("remote connection", () => {
|
|||||||
});
|
});
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
describe("remote connection jobs surface", () => {
|
|
||||||
it("lists, describes, cancels, and reads history", async () => {
|
|
||||||
const { tableFromArrays, tableToIPC } = await import("apache-arrow");
|
|
||||||
const eventsTable = tableFromArrays({ state: ["created", "succeeded"] });
|
|
||||||
const eventsBody = Buffer.from(tableToIPC(eventsTable, "stream"));
|
|
||||||
|
|
||||||
await withMockDatabase(
|
|
||||||
(req, res) => {
|
|
||||||
let body = "";
|
|
||||||
req.on("data", (chunk) => {
|
|
||||||
body += chunk;
|
|
||||||
});
|
|
||||||
req.on("end", () => {
|
|
||||||
const payload = body.length > 0 ? JSON.parse(body) : {};
|
|
||||||
if (req.url === "/v1/jobs/list") {
|
|
||||||
if (payload["page_token"] === undefined) {
|
|
||||||
res
|
|
||||||
.writeHead(200, { "Content-Type": "application/json" })
|
|
||||||
.end(
|
|
||||||
'{"jobs": [{"job_id": "job-1", "table": "t1", ' +
|
|
||||||
'"job_type": "create_index", "state": "in_progress", ' +
|
|
||||||
'"created_at_millis": 1000}], "page_token": "next"}',
|
|
||||||
);
|
|
||||||
} else {
|
|
||||||
res
|
|
||||||
.writeHead(200, { "Content-Type": "application/json" })
|
|
||||||
.end(
|
|
||||||
'{"jobs": [{"job_id": "job-2", "table": "t2", ' +
|
|
||||||
'"job_type": "create_index", "state": "succeeded", ' +
|
|
||||||
'"created_at_millis": 2000}]}',
|
|
||||||
);
|
|
||||||
}
|
|
||||||
} else if (req.url === "/v1/jobs/describe") {
|
|
||||||
if (payload["job_id"] !== "job-1") {
|
|
||||||
res.writeHead(404).end("no such job");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
res
|
|
||||||
.writeHead(200, { "Content-Type": "application/json" })
|
|
||||||
.end(
|
|
||||||
'{"job_id": "job-1", "job_type": "create_index", ' +
|
|
||||||
'"job_state": "FAILED", "creation_ms": 1000, ' +
|
|
||||||
'"spec": {"column": "vec"}, "failure": {"phase": "execute", ' +
|
|
||||||
'"message": "worker died", "retryable": true}}',
|
|
||||||
);
|
|
||||||
} else if (req.url === "/v1/jobs/cancel") {
|
|
||||||
if (payload["job_id"] !== "job-1") {
|
|
||||||
res.writeHead(404).end("no such job");
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
res
|
|
||||||
.writeHead(200, { "Content-Type": "application/json" })
|
|
||||||
.end('{"job_id": "job-1"}');
|
|
||||||
} else if (req.url === "/v1/jobs/query_events") {
|
|
||||||
res
|
|
||||||
.writeHead(200, {
|
|
||||||
"Content-Type": "application/vnd.apache.arrow.stream",
|
|
||||||
})
|
|
||||||
.end(eventsBody);
|
|
||||||
} else {
|
|
||||||
res.writeHead(404).end();
|
|
||||||
}
|
|
||||||
});
|
|
||||||
},
|
|
||||||
async (db) => {
|
|
||||||
const jobs = await db.listJobs();
|
|
||||||
expect(jobs.map((job) => job.jobId)).toEqual(["job-1", "job-2"]);
|
|
||||||
expect(jobs[0].state).toEqual("running");
|
|
||||||
expect(jobs[1].state).toEqual("finished");
|
|
||||||
|
|
||||||
const description = await db.getJob("job-1");
|
|
||||||
expect(description?.state).toEqual("failed");
|
|
||||||
expect(JSON.parse(description?.specJson ?? "")).toEqual({
|
|
||||||
column: "vec",
|
|
||||||
});
|
|
||||||
expect(description?.failure?.message).toEqual("worker died");
|
|
||||||
expect(await db.getJob("missing")).toBeNull();
|
|
||||||
|
|
||||||
expect(await db.cancelJob("job-1")).toBe(true);
|
|
||||||
expect(await db.cancelJob("missing")).toBe(false);
|
|
||||||
|
|
||||||
const history = await db.jobHistory("job-1");
|
|
||||||
expect(history.numRows).toEqual(2);
|
|
||||||
|
|
||||||
const job = db.job("job-1");
|
|
||||||
expect(job.id).toEqual("job-1");
|
|
||||||
expect(await job.status()).toEqual("failed");
|
|
||||||
await expect(job.wait()).rejects.toThrow("worker died");
|
|
||||||
},
|
|
||||||
);
|
|
||||||
});
|
|
||||||
});
|
|
||||||
|
|||||||
@@ -851,11 +851,7 @@ describe("When creating an index", () => {
|
|||||||
afterEach(() => tmpDir.removeCallback());
|
afterEach(() => tmpDir.removeCallback());
|
||||||
|
|
||||||
it("should create a vector index on vector columns", async () => {
|
it("should create a vector index on vector columns", async () => {
|
||||||
const job = await tbl.createIndexAsync("vec");
|
await tbl.createIndex("vec");
|
||||||
expect(job.id).toBeNull();
|
|
||||||
await job.wait();
|
|
||||||
// Cancelling a job that already finished succeeds and does nothing.
|
|
||||||
await job.cancel();
|
|
||||||
|
|
||||||
// check index directory
|
// check index directory
|
||||||
const indexDir = path.join(tmpDir.name, "test.lance", "_indices");
|
const indexDir = path.join(tmpDir.name, "test.lance", "_indices");
|
||||||
|
|||||||
@@ -1,7 +1,6 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
import { tableFromIPC } from "apache-arrow";
|
|
||||||
import {
|
import {
|
||||||
Data,
|
Data,
|
||||||
SchemaLike,
|
SchemaLike,
|
||||||
@@ -21,9 +20,6 @@ import type {
|
|||||||
CreateNamespaceResponse,
|
CreateNamespaceResponse,
|
||||||
DescribeNamespaceResponse,
|
DescribeNamespaceResponse,
|
||||||
DropNamespaceResponse,
|
DropNamespaceResponse,
|
||||||
Job,
|
|
||||||
JobDescription,
|
|
||||||
JobInfo,
|
|
||||||
ListNamespacesResponse,
|
ListNamespacesResponse,
|
||||||
} from "./native";
|
} from "./native";
|
||||||
export type {
|
export type {
|
||||||
@@ -440,40 +436,6 @@ export abstract class Connection {
|
|||||||
newName: string,
|
newName: string,
|
||||||
options?: RenameTableOptions,
|
options?: RenameTableOptions,
|
||||||
): Promise<void>;
|
): Promise<void>;
|
||||||
|
|
||||||
/**
|
|
||||||
* A {@link Job} handle for a server-side job by id.
|
|
||||||
*
|
|
||||||
* The handle is constructed without a server round trip; an unknown id
|
|
||||||
* surfaces when the handle is used. Dropping the handle has no effect on
|
|
||||||
* the job itself.
|
|
||||||
*/
|
|
||||||
abstract job(jobId: string): Job;
|
|
||||||
|
|
||||||
/** List server-side jobs across the database's tables. */
|
|
||||||
abstract listJobs(): Promise<JobInfo[]>;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Describe a single server-side job by id.
|
|
||||||
*
|
|
||||||
* Resolves to `null` when the server has no such job.
|
|
||||||
*/
|
|
||||||
abstract getJob(jobId: string): Promise<JobDescription | null>;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* Request cancellation of a server-side job by id.
|
|
||||||
*
|
|
||||||
* Resolves to true if the server accepted the cancellation, false if no
|
|
||||||
* such job exists. Cancelling an already-terminal job is a no-op success.
|
|
||||||
*/
|
|
||||||
abstract cancelJob(jobId: string): Promise<boolean>;
|
|
||||||
|
|
||||||
/**
|
|
||||||
* The lifecycle event history of a server-side job, as an Arrow table.
|
|
||||||
*
|
|
||||||
* Lists history across all jobs when `jobId` is omitted.
|
|
||||||
*/
|
|
||||||
abstract jobHistory(jobId?: string): Promise<ArrowTable>;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/** @hideconstructor */
|
/** @hideconstructor */
|
||||||
@@ -760,30 +722,6 @@ export class LocalConnection extends Connection {
|
|||||||
options?.newNamespacePath,
|
options?.newNamespacePath,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
job(jobId: string): Job {
|
|
||||||
return this.inner.job(jobId);
|
|
||||||
}
|
|
||||||
|
|
||||||
async listJobs(): Promise<JobInfo[]> {
|
|
||||||
return this.inner.listJobs();
|
|
||||||
}
|
|
||||||
|
|
||||||
async getJob(jobId: string): Promise<JobDescription | null> {
|
|
||||||
return this.inner.getJob(jobId);
|
|
||||||
}
|
|
||||||
|
|
||||||
async cancelJob(jobId: string): Promise<boolean> {
|
|
||||||
return this.inner.cancelJob(jobId);
|
|
||||||
}
|
|
||||||
|
|
||||||
async jobHistory(jobId?: string): Promise<ArrowTable> {
|
|
||||||
const buf = await this.inner.jobHistory(jobId);
|
|
||||||
if (buf.length === 0) {
|
|
||||||
return new ArrowTable();
|
|
||||||
}
|
|
||||||
return tableFromIPC(buf);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|||||||
@@ -85,13 +85,7 @@ export {
|
|||||||
RenameTableOptions,
|
RenameTableOptions,
|
||||||
} from "./connection";
|
} from "./connection";
|
||||||
|
|
||||||
export {
|
export { Session } from "./native.js";
|
||||||
Job,
|
|
||||||
JobDescription,
|
|
||||||
JobFailureInfo,
|
|
||||||
JobInfo,
|
|
||||||
Session,
|
|
||||||
} from "./native.js";
|
|
||||||
|
|
||||||
export {
|
export {
|
||||||
ExecutableQuery,
|
ExecutableQuery,
|
||||||
|
|||||||
@@ -30,7 +30,6 @@ import {
|
|||||||
DropColumnsResult,
|
DropColumnsResult,
|
||||||
IndexConfig,
|
IndexConfig,
|
||||||
IndexStatistics,
|
IndexStatistics,
|
||||||
Job,
|
|
||||||
Branches as NativeBranches,
|
Branches as NativeBranches,
|
||||||
OptimizeStats,
|
OptimizeStats,
|
||||||
TableStatistics,
|
TableStatistics,
|
||||||
@@ -359,17 +358,6 @@ export abstract class Table {
|
|||||||
options?: Partial<IndexOptions>,
|
options?: Partial<IndexOptions>,
|
||||||
): Promise<void>;
|
): 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.
|
* Drop an index from the table.
|
||||||
*
|
*
|
||||||
@@ -952,22 +940,6 @@ 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> {
|
async dropIndex(name: string): Promise<void> {
|
||||||
await this.inner.dropIndex(name);
|
await this.inner.dropIndex(name);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -340,69 +340,6 @@ impl Connection {
|
|||||||
self.get_inner()?.drop_all_tables(&ns).await.default_error()
|
self.get_inner()?.drop_all_tables(&ns).await.default_error()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A `Job` handle for a server-side job by id.
|
|
||||||
///
|
|
||||||
/// The handle is constructed without a server round trip; an unknown id
|
|
||||||
/// surfaces when the handle is used.
|
|
||||||
#[napi]
|
|
||||||
pub fn job(&self, job_id: String) -> napi::Result<crate::job::Job> {
|
|
||||||
let job = self.get_inner()?.job(job_id).default_error()?;
|
|
||||||
Ok(crate::job::Job::new(job))
|
|
||||||
}
|
|
||||||
|
|
||||||
/// List server-side jobs across the database's tables.
|
|
||||||
#[napi(catch_unwind)]
|
|
||||||
pub async fn list_jobs(&self) -> napi::Result<Vec<crate::job::JobInfo>> {
|
|
||||||
let jobs = self.get_inner()?.list_jobs().await.default_error()?;
|
|
||||||
Ok(jobs.into_iter().map(Into::into).collect())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Describe a single server-side job by id. `null` when the server has
|
|
||||||
/// no such job.
|
|
||||||
#[napi(catch_unwind)]
|
|
||||||
pub async fn get_job(
|
|
||||||
&self,
|
|
||||||
job_id: String,
|
|
||||||
) -> napi::Result<Option<crate::job::JobDescription>> {
|
|
||||||
let description = self.get_inner()?.get_job(&job_id).await.default_error()?;
|
|
||||||
Ok(description.map(Into::into))
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Request cancellation of a server-side job by id. Returns true if the
|
|
||||||
/// server accepted the cancellation, false if no such job exists.
|
|
||||||
#[napi(catch_unwind)]
|
|
||||||
pub async fn cancel_job(&self, job_id: String) -> napi::Result<bool> {
|
|
||||||
self.get_inner()?.cancel_job(&job_id).await.default_error()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The lifecycle event history of a server-side job (all jobs when
|
|
||||||
/// `job_id` is null), as an Arrow IPC stream buffer. Empty when there is
|
|
||||||
/// no history.
|
|
||||||
#[napi(catch_unwind)]
|
|
||||||
pub async fn job_history(&self, job_id: Option<String>) -> napi::Result<Buffer> {
|
|
||||||
let batches = self
|
|
||||||
.get_inner()?
|
|
||||||
.job_history(job_id.as_deref())
|
|
||||||
.await
|
|
||||||
.default_error()?;
|
|
||||||
let Some(first) = batches.first() else {
|
|
||||||
return Ok(Buffer::from(Vec::<u8>::new()));
|
|
||||||
};
|
|
||||||
let mut out = Vec::new();
|
|
||||||
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut out, &first.schema())
|
|
||||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
|
||||||
for batch in &batches {
|
|
||||||
writer
|
|
||||||
.write(batch)
|
|
||||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
|
||||||
}
|
|
||||||
writer
|
|
||||||
.finish()
|
|
||||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
|
||||||
drop(writer);
|
|
||||||
Ok(Buffer::from(out))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[napi(catch_unwind)]
|
#[napi(catch_unwind)]
|
||||||
/// Describe a namespace and return its properties.
|
/// Describe a namespace and return its properties.
|
||||||
pub async fn describe_namespace(
|
pub async fn describe_namespace(
|
||||||
|
|||||||
@@ -1,123 +0,0 @@
|
|||||||
// 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)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The operation's current lifecycle state: "running", "finished",
|
|
||||||
/// "failed", or "cancelled".
|
|
||||||
///
|
|
||||||
/// A point snapshot; unlike {@link Job.wait} it does not block or reject
|
|
||||||
/// on a terminal failure state. States a newer server reports that this
|
|
||||||
/// client version does not know pass through as-is.
|
|
||||||
#[napi(catch_unwind)]
|
|
||||||
pub async fn status(&self) -> napi::Result<String> {
|
|
||||||
self.inner.status().await.default_error()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 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()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A row from `Connection.listJobs`: one server-side job.
|
|
||||||
#[napi(object)]
|
|
||||||
pub struct JobInfo {
|
|
||||||
/// The job id -- what `Connection.getJob` and `Connection.cancelJob`
|
|
||||||
/// accept.
|
|
||||||
pub job_id: String,
|
|
||||||
/// The table the job runs against, without URI or namespace.
|
|
||||||
pub table: String,
|
|
||||||
pub job_type: String,
|
|
||||||
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
|
|
||||||
pub state: String,
|
|
||||||
/// When the job was created, in milliseconds since the epoch.
|
|
||||||
pub created_at_millis: i64,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl From<lancedb::database::JobInfo> for JobInfo {
|
|
||||||
fn from(info: lancedb::database::JobInfo) -> Self {
|
|
||||||
Self {
|
|
||||||
job_id: info.job_id,
|
|
||||||
table: info.table,
|
|
||||||
job_type: info.job_type,
|
|
||||||
state: info.state,
|
|
||||||
created_at_millis: info.created_at_millis,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The server's account of why a job failed.
|
|
||||||
#[napi(object)]
|
|
||||||
pub struct JobFailureInfo {
|
|
||||||
pub phase: Option<String>,
|
|
||||||
pub message: Option<String>,
|
|
||||||
pub retryable: Option<bool>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A described job from `Connection.getJob`.
|
|
||||||
#[napi(object)]
|
|
||||||
pub struct JobDescription {
|
|
||||||
pub job_id: String,
|
|
||||||
pub job_type: String,
|
|
||||||
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
|
|
||||||
pub state: String,
|
|
||||||
/// When the job was created, in milliseconds since the epoch.
|
|
||||||
pub creation_ms: i64,
|
|
||||||
/// The job-type-specific specification as a JSON string, when present.
|
|
||||||
pub spec_json: Option<String>,
|
|
||||||
/// Why the job failed, when the job is failed and the server reports a
|
|
||||||
/// reason.
|
|
||||||
pub failure: Option<JobFailureInfo>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl From<lancedb::database::JobDescription> for JobDescription {
|
|
||||||
fn from(description: lancedb::database::JobDescription) -> Self {
|
|
||||||
Self {
|
|
||||||
job_id: description.job_id,
|
|
||||||
job_type: description.job_type,
|
|
||||||
state: description.state,
|
|
||||||
creation_ms: description.creation_ms,
|
|
||||||
spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
|
|
||||||
failure: description.failure.map(|failure| JobFailureInfo {
|
|
||||||
phase: failure.phase,
|
|
||||||
message: failure.message,
|
|
||||||
retryable: failure.retryable,
|
|
||||||
}),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -11,7 +11,6 @@ mod error;
|
|||||||
mod header;
|
mod header;
|
||||||
mod index;
|
mod index;
|
||||||
mod iterator;
|
mod iterator;
|
||||||
mod job;
|
|
||||||
pub mod merge;
|
pub mod merge;
|
||||||
pub mod otel;
|
pub mod otel;
|
||||||
pub mod permutation;
|
pub mod permutation;
|
||||||
|
|||||||
+2
-39
@@ -168,39 +168,6 @@ impl Table {
|
|||||||
builder.execute().await.default_error()
|
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)]
|
#[napi(catch_unwind)]
|
||||||
pub async fn drop_index(&self, index_name: String) -> napi::Result<()> {
|
pub async fn drop_index(&self, index_name: String) -> napi::Result<()> {
|
||||||
self.inner_ref()?
|
self.inner_ref()?
|
||||||
@@ -339,9 +306,7 @@ impl Table {
|
|||||||
let transforms = NewColumnTransform::SqlExpressions(transforms);
|
let transforms = NewColumnTransform::SqlExpressions(transforms);
|
||||||
let res = self
|
let res = self
|
||||||
.inner_ref()?
|
.inner_ref()?
|
||||||
.add_columns()
|
.add_columns(transforms, None)
|
||||||
.transform(transforms)
|
|
||||||
.execute()
|
|
||||||
.await
|
.await
|
||||||
.default_error()?;
|
.default_error()?;
|
||||||
Ok(res.into())
|
Ok(res.into())
|
||||||
@@ -358,9 +323,7 @@ impl Table {
|
|||||||
let transforms = NewColumnTransform::AllNulls(schema);
|
let transforms = NewColumnTransform::AllNulls(schema);
|
||||||
let res = self
|
let res = self
|
||||||
.inner_ref()?
|
.inner_ref()?
|
||||||
.add_columns()
|
.add_columns(transforms, None)
|
||||||
.transform(transforms)
|
|
||||||
.execute()
|
|
||||||
.await
|
.await
|
||||||
.default_error()?;
|
.default_error()?;
|
||||||
Ok(res.into())
|
Ok(res.into())
|
||||||
|
|||||||
@@ -20,7 +20,6 @@ from .remote import ClientConfig
|
|||||||
from .remote.db import RemoteDBConnection
|
from .remote.db import RemoteDBConnection
|
||||||
from .expr import Expr, col, lit, func
|
from .expr import Expr, col, lit, func
|
||||||
from .schema import blob, vector, BlobType
|
from .schema import blob, vector, BlobType
|
||||||
from .job import AsyncJob, Job
|
|
||||||
from .table import AsyncTable, Table
|
from .table import AsyncTable, Table
|
||||||
from .types import BaseTokenizerType
|
from .types import BaseTokenizerType
|
||||||
from ._lancedb import Session
|
from ._lancedb import Session
|
||||||
@@ -501,7 +500,6 @@ __all__ = [
|
|||||||
"connect_namespace",
|
"connect_namespace",
|
||||||
"connect_namespace_async",
|
"connect_namespace_async",
|
||||||
"AsyncConnection",
|
"AsyncConnection",
|
||||||
"AsyncJob",
|
|
||||||
"AsyncLanceNamespaceDBConnection",
|
"AsyncLanceNamespaceDBConnection",
|
||||||
"AsyncTable",
|
"AsyncTable",
|
||||||
"FtsToken",
|
"FtsToken",
|
||||||
@@ -515,7 +513,6 @@ __all__ = [
|
|||||||
"BlobType",
|
"BlobType",
|
||||||
"vector",
|
"vector",
|
||||||
"DBConnection",
|
"DBConnection",
|
||||||
"Job",
|
|
||||||
"LanceDBConnection",
|
"LanceDBConnection",
|
||||||
"LanceNamespaceDBConnection",
|
"LanceNamespaceDBConnection",
|
||||||
"RemoteDBConnection",
|
"RemoteDBConnection",
|
||||||
|
|||||||
@@ -14,10 +14,14 @@ import pyarrow as pa
|
|||||||
from .expr import Expr
|
from .expr import Expr
|
||||||
from .schema import blob_v2_column_paths
|
from .schema import blob_v2_column_paths
|
||||||
from .types import BlobMode, QueryProjection, QueryProjectionSpec
|
from .types import BlobMode, QueryProjection, QueryProjectionSpec
|
||||||
|
from .util import get_uri_scheme
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from _typeshed import WriteableBuffer
|
from _typeshed import WriteableBuffer
|
||||||
|
|
||||||
|
from .remote.table import RemoteTable
|
||||||
|
from .table import AsyncTable, Table
|
||||||
|
|
||||||
BLOB_MODE_TO_HANDLING = {
|
BLOB_MODE_TO_HANDLING = {
|
||||||
"lazy": "blobs_descriptions",
|
"lazy": "blobs_descriptions",
|
||||||
"bytes": "all_binary",
|
"bytes": "all_binary",
|
||||||
@@ -100,6 +104,22 @@ def validate_blob_mode(blob_mode: BlobMode) -> None:
|
|||||||
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
||||||
|
|
||||||
|
|
||||||
|
def supports_blob_auto_row_id(table: Table | AsyncTable | RemoteTable) -> bool:
|
||||||
|
"""Blob auto row-id applies to native tables, not LanceDB Cloud."""
|
||||||
|
from .remote.table import RemoteTable
|
||||||
|
|
||||||
|
if isinstance(table, RemoteTable):
|
||||||
|
return False
|
||||||
|
|
||||||
|
inner = getattr(table, "_inner", None)
|
||||||
|
if inner is not None:
|
||||||
|
uri = inner.database().uri
|
||||||
|
if isinstance(uri, str) and get_uri_scheme(uri) == "db":
|
||||||
|
return False
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
def projection_includes_blob_column(
|
def projection_includes_blob_column(
|
||||||
projection: QueryProjection,
|
projection: QueryProjection,
|
||||||
blob_columns: Iterable[str],
|
blob_columns: Iterable[str],
|
||||||
@@ -144,14 +164,16 @@ def v2_projection_needs_row_id(
|
|||||||
|
|
||||||
|
|
||||||
def blob_auto_row_id_for_scan(
|
def blob_auto_row_id_for_scan(
|
||||||
|
table: Table | AsyncTable | RemoteTable,
|
||||||
schema: pa.Schema,
|
schema: pa.Schema,
|
||||||
projection: QueryProjection,
|
projection: QueryProjection,
|
||||||
*,
|
*,
|
||||||
with_row_id: bool | None,
|
with_row_id: bool | None,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
"""Auto row-id only applies when the caller said nothing about row ids."""
|
|
||||||
if with_row_id is not None:
|
if with_row_id is not None:
|
||||||
return False
|
return False
|
||||||
|
if not supports_blob_auto_row_id(table):
|
||||||
|
return False
|
||||||
return v2_projection_needs_row_id(schema, projection, with_row_id=False)
|
return v2_projection_needs_row_id(schema, projection, with_row_id=False)
|
||||||
|
|
||||||
|
|
||||||
@@ -164,11 +186,6 @@ def finalize_blob_query_table(
|
|||||||
) -> pa.Table:
|
) -> pa.Table:
|
||||||
if user_requested_row_id or not blob_auto_row_id:
|
if user_requested_row_id or not blob_auto_row_id:
|
||||||
return tbl
|
return tbl
|
||||||
if "_rowid" not in tbl.column_names:
|
|
||||||
# A backend that ignores the row-id request leaves nothing to stash. Hand
|
|
||||||
# back the projection as-is so fetch_blobs raises the error that names the
|
|
||||||
# ways to supply row ids, rather than failing here about a hidden column.
|
|
||||||
return tbl
|
|
||||||
return stash_auto_row_ids(tbl, blob_paths)
|
return stash_auto_row_ids(tbl, blob_paths)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -146,13 +146,6 @@ class Connection(object):
|
|||||||
start_after: Optional[str],
|
start_after: Optional[str],
|
||||||
limit: Optional[int],
|
limit: Optional[int],
|
||||||
) -> list[str]: ... # Deprecated: Use list_tables instead
|
) -> list[str]: ... # Deprecated: Use list_tables instead
|
||||||
def job(self, job_id: str) -> Job: ...
|
|
||||||
async def list_jobs(self) -> List[JobInfo]: ...
|
|
||||||
async def get_job(self, job_id: str) -> Optional[JobDescription]: ...
|
|
||||||
async def cancel_job(self, job_id: str) -> bool: ...
|
|
||||||
async def job_history(
|
|
||||||
self, job_id: Optional[str] = None
|
|
||||||
) -> List[pa.RecordBatch]: ...
|
|
||||||
async def create_table(
|
async def create_table(
|
||||||
self,
|
self,
|
||||||
name: str,
|
name: str,
|
||||||
@@ -216,47 +209,6 @@ class BlobFile:
|
|||||||
def read_range(self, offset: int, length: int) -> bytes: ...
|
def read_range(self, offset: int, length: int) -> bytes: ...
|
||||||
def read_up_to(self, length: int) -> bytes: ...
|
def read_up_to(self, length: int) -> bytes: ...
|
||||||
|
|
||||||
class Job:
|
|
||||||
@property
|
|
||||||
def id(self) -> Optional[str]: ...
|
|
||||||
async def status(self) -> str: ...
|
|
||||||
async def wait(self) -> None: ...
|
|
||||||
async def cancel(self) -> None: ...
|
|
||||||
|
|
||||||
class JobInfo:
|
|
||||||
@property
|
|
||||||
def job_id(self) -> str: ...
|
|
||||||
@property
|
|
||||||
def table(self) -> str: ...
|
|
||||||
@property
|
|
||||||
def job_type(self) -> str: ...
|
|
||||||
@property
|
|
||||||
def state(self) -> str: ...
|
|
||||||
@property
|
|
||||||
def created_at_millis(self) -> int: ...
|
|
||||||
|
|
||||||
class JobFailureInfo:
|
|
||||||
@property
|
|
||||||
def phase(self) -> Optional[str]: ...
|
|
||||||
@property
|
|
||||||
def message(self) -> Optional[str]: ...
|
|
||||||
@property
|
|
||||||
def retryable(self) -> Optional[bool]: ...
|
|
||||||
|
|
||||||
class JobDescription:
|
|
||||||
@property
|
|
||||||
def job_id(self) -> str: ...
|
|
||||||
@property
|
|
||||||
def job_type(self) -> str: ...
|
|
||||||
@property
|
|
||||||
def state(self) -> str: ...
|
|
||||||
@property
|
|
||||||
def creation_ms(self) -> int: ...
|
|
||||||
@property
|
|
||||||
def spec_json(self) -> Optional[str]: ...
|
|
||||||
@property
|
|
||||||
def failure(self) -> Optional[JobFailureInfo]: ...
|
|
||||||
|
|
||||||
class Table:
|
class Table:
|
||||||
def name(self) -> str: ...
|
def name(self) -> str: ...
|
||||||
def __repr__(self) -> str: ...
|
def __repr__(self) -> str: ...
|
||||||
@@ -296,28 +248,6 @@ class Table:
|
|||||||
name: Optional[str],
|
name: Optional[str],
|
||||||
train: Optional[bool],
|
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 list_versions(self) -> List[Dict[str, Any]]: ...
|
||||||
async def version(self) -> int: ...
|
async def version(self) -> int: ...
|
||||||
async def checkout(self, version: Union[int, str]): ...
|
async def checkout(self, version: Union[int, str]): ...
|
||||||
|
|||||||
+5
-166
@@ -45,7 +45,6 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
|||||||
|
|
||||||
from . import __version__
|
from . import __version__
|
||||||
from ._lancedb import connect as lancedb_connect # type: ignore
|
from ._lancedb import connect as lancedb_connect # type: ignore
|
||||||
from .job import AsyncJob, Job
|
|
||||||
from .table import (
|
from .table import (
|
||||||
AsyncTable,
|
AsyncTable,
|
||||||
LanceTable,
|
LanceTable,
|
||||||
@@ -64,7 +63,6 @@ if TYPE_CHECKING:
|
|||||||
from .pydantic import LanceModel
|
from .pydantic import LanceModel
|
||||||
|
|
||||||
from ._lancedb import Connection as LanceDbConnection
|
from ._lancedb import Connection as LanceDbConnection
|
||||||
from ._lancedb import JobDescription, JobInfo
|
|
||||||
from .common import DATA, URI
|
from .common import DATA, URI
|
||||||
from .embeddings import EmbeddingFunctionConfig
|
from .embeddings import EmbeddingFunctionConfig
|
||||||
from ._lancedb import Session
|
from ._lancedb import Session
|
||||||
@@ -180,51 +178,6 @@ class DBConnection(EnforceOverrides):
|
|||||||
"Namespace operations are not supported for this connection type"
|
"Namespace operations are not supported for this connection type"
|
||||||
)
|
)
|
||||||
|
|
||||||
def namespace_exists(self, namespace_id: List[str]) -> bool:
|
|
||||||
"""Check if a namespace exists.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
namespace_id: List[str]
|
|
||||||
The namespace identifier to check.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
bool
|
|
||||||
True if the namespace exists, False otherwise.
|
|
||||||
|
|
||||||
Raises
|
|
||||||
------
|
|
||||||
NotImplementedError
|
|
||||||
If the connection type does not support namespace operations.
|
|
||||||
"""
|
|
||||||
raise NotImplementedError(
|
|
||||||
"Namespace operations are not supported for this connection type"
|
|
||||||
)
|
|
||||||
|
|
||||||
def table_exists(self, table_id: List[str]) -> bool:
|
|
||||||
"""Check if a table exists.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
table_id: List[str]
|
|
||||||
The table identifier to check (full path including namespace
|
|
||||||
segments and table name).
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
bool
|
|
||||||
True if the table exists, False otherwise.
|
|
||||||
|
|
||||||
Raises
|
|
||||||
------
|
|
||||||
NotImplementedError
|
|
||||||
If the connection type does not support namespace operations.
|
|
||||||
"""
|
|
||||||
raise NotImplementedError(
|
|
||||||
"Namespace operations are not supported for this connection type"
|
|
||||||
)
|
|
||||||
|
|
||||||
def list_tables(
|
def list_tables(
|
||||||
self,
|
self,
|
||||||
namespace_path: Optional[List[str]] = None,
|
namespace_path: Optional[List[str]] = None,
|
||||||
@@ -610,46 +563,6 @@ class DBConnection(EnforceOverrides):
|
|||||||
"""
|
"""
|
||||||
raise NotImplementedError("serialize is not supported for this connection type")
|
raise NotImplementedError("serialize is not supported for this connection type")
|
||||||
|
|
||||||
def job(self, job_id: str) -> Job:
|
|
||||||
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
|
|
||||||
|
|
||||||
The handle is constructed without a server round trip; an unknown id
|
|
||||||
surfaces when the handle is used. Dropping the handle has no effect
|
|
||||||
on the job itself.
|
|
||||||
"""
|
|
||||||
raise NotImplementedError("job is not supported for this connection type")
|
|
||||||
|
|
||||||
def list_jobs(self) -> List[JobInfo]:
|
|
||||||
"""List server-side jobs across the database's tables."""
|
|
||||||
raise NotImplementedError("list_jobs is not supported for this connection type")
|
|
||||||
|
|
||||||
def get_job(self, job_id: str) -> Optional[JobDescription]:
|
|
||||||
"""Describe a single server-side job by id.
|
|
||||||
|
|
||||||
Returns None when the server has no such job.
|
|
||||||
"""
|
|
||||||
raise NotImplementedError("get_job is not supported for this connection type")
|
|
||||||
|
|
||||||
def cancel_job(self, job_id: str) -> bool:
|
|
||||||
"""Request cancellation of a server-side job by id.
|
|
||||||
|
|
||||||
Returns True if the server accepted the cancellation, False if no
|
|
||||||
such job exists. Cancelling an already-terminal job is a no-op
|
|
||||||
success.
|
|
||||||
"""
|
|
||||||
raise NotImplementedError(
|
|
||||||
"cancel_job is not supported for this connection type"
|
|
||||||
)
|
|
||||||
|
|
||||||
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
|
||||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
|
||||||
|
|
||||||
Lists history across all jobs when `job_id` is None.
|
|
||||||
"""
|
|
||||||
raise NotImplementedError(
|
|
||||||
"job_history is not supported for this connection type"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class LanceDBConnection(DBConnection):
|
class LanceDBConnection(DBConnection):
|
||||||
"""
|
"""
|
||||||
@@ -775,7 +688,11 @@ class LanceDBConnection(DBConnection):
|
|||||||
return cls(None, _inner=inner)
|
return cls(None, _inner=inner)
|
||||||
|
|
||||||
def __repr__(self) -> str:
|
def __repr__(self) -> str:
|
||||||
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
|
val = f"{self.__class__.__name__}(uri={self._conn.uri!r}"
|
||||||
|
if self.read_consistency_interval is not None:
|
||||||
|
val += f", read_consistency_interval={repr(self.read_consistency_interval)}"
|
||||||
|
val += ")"
|
||||||
|
return val
|
||||||
|
|
||||||
@override
|
@override
|
||||||
def serialize(self) -> str:
|
def serialize(self) -> str:
|
||||||
@@ -1212,47 +1129,6 @@ class LanceDBConnection(DBConnection):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@override
|
|
||||||
def job(self, job_id: str) -> Job:
|
|
||||||
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
|
|
||||||
|
|
||||||
The handle is constructed without a server round trip; an unknown id
|
|
||||||
surfaces when the handle is used. Dropping the handle has no effect
|
|
||||||
on the job itself.
|
|
||||||
"""
|
|
||||||
return Job(self._conn.job(job_id))
|
|
||||||
|
|
||||||
@override
|
|
||||||
def list_jobs(self) -> List[JobInfo]:
|
|
||||||
"""List server-side jobs across the database's tables."""
|
|
||||||
return LOOP.run(self._conn.list_jobs())
|
|
||||||
|
|
||||||
@override
|
|
||||||
def get_job(self, job_id: str) -> Optional[JobDescription]:
|
|
||||||
"""Describe a single server-side job by id.
|
|
||||||
|
|
||||||
Returns None when the server has no such job.
|
|
||||||
"""
|
|
||||||
return LOOP.run(self._conn.get_job(job_id))
|
|
||||||
|
|
||||||
@override
|
|
||||||
def cancel_job(self, job_id: str) -> bool:
|
|
||||||
"""Request cancellation of a server-side job by id.
|
|
||||||
|
|
||||||
Returns True if the server accepted the cancellation, False if no
|
|
||||||
such job exists. Cancelling an already-terminal job is a no-op
|
|
||||||
success.
|
|
||||||
"""
|
|
||||||
return LOOP.run(self._conn.cancel_job(job_id))
|
|
||||||
|
|
||||||
@override
|
|
||||||
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
|
||||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
|
||||||
|
|
||||||
Lists history across all jobs when `job_id` is None.
|
|
||||||
"""
|
|
||||||
return LOOP.run(self._conn.job_history(job_id))
|
|
||||||
|
|
||||||
@override
|
@override
|
||||||
def namespace_client(self) -> LanceNamespace:
|
def namespace_client(self) -> LanceNamespace:
|
||||||
"""Get the equivalent namespace client for this connection.
|
"""Get the equivalent namespace client for this connection.
|
||||||
@@ -1962,43 +1838,6 @@ class AsyncConnection(object):
|
|||||||
namespace_path = []
|
namespace_path = []
|
||||||
await self._inner.drop_all_tables(namespace_path=namespace_path)
|
await self._inner.drop_all_tables(namespace_path=namespace_path)
|
||||||
|
|
||||||
def job(self, job_id: str) -> AsyncJob:
|
|
||||||
"""An [AsyncJob][lancedb.job.AsyncJob] handle for a server-side job
|
|
||||||
by id.
|
|
||||||
|
|
||||||
The handle is constructed without a server round trip; an unknown id
|
|
||||||
surfaces when the handle is used. Dropping the handle has no effect
|
|
||||||
on the job itself.
|
|
||||||
"""
|
|
||||||
return AsyncJob(self._inner.job(job_id))
|
|
||||||
|
|
||||||
async def list_jobs(self) -> List[JobInfo]:
|
|
||||||
"""List server-side jobs across the database's tables."""
|
|
||||||
return await self._inner.list_jobs()
|
|
||||||
|
|
||||||
async def get_job(self, job_id: str) -> Optional[JobDescription]:
|
|
||||||
"""Describe a single server-side job by id.
|
|
||||||
|
|
||||||
Returns None when the server has no such job.
|
|
||||||
"""
|
|
||||||
return await self._inner.get_job(job_id)
|
|
||||||
|
|
||||||
async def cancel_job(self, job_id: str) -> bool:
|
|
||||||
"""Request cancellation of a server-side job by id.
|
|
||||||
|
|
||||||
Returns True if the server accepted the cancellation, False if no
|
|
||||||
such job exists. Cancelling an already-terminal job is a no-op
|
|
||||||
success.
|
|
||||||
"""
|
|
||||||
return await self._inner.cancel_job(job_id)
|
|
||||||
|
|
||||||
async def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
|
||||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
|
||||||
|
|
||||||
Lists history across all jobs when `job_id` is None.
|
|
||||||
"""
|
|
||||||
return await self._inner.job_history(job_id)
|
|
||||||
|
|
||||||
async def namespace_client(self) -> LanceNamespace:
|
async def namespace_client(self) -> LanceNamespace:
|
||||||
"""Get the equivalent namespace client for this connection.
|
"""Get the equivalent namespace client for this connection.
|
||||||
|
|
||||||
|
|||||||
@@ -23,15 +23,3 @@ class MissingColumnError(KeyError):
|
|||||||
return (
|
return (
|
||||||
f"Error: Column '{self.column_name}' does not exist in the DataFrame object"
|
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
|
|
||||||
|
|||||||
@@ -1,105 +0,0 @@
|
|||||||
# 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 status(self) -> str:
|
|
||||||
"""The operation's current lifecycle state: "running", "finished",
|
|
||||||
"failed", or "cancelled".
|
|
||||||
|
|
||||||
A point snapshot; unlike `wait` it does not block or raise on a
|
|
||||||
terminal failure state. States a newer server reports that this
|
|
||||||
client version does not know pass through as-is.
|
|
||||||
"""
|
|
||||||
if self._inner is None:
|
|
||||||
return "finished"
|
|
||||||
return await self._inner.status()
|
|
||||||
|
|
||||||
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 status(self) -> str:
|
|
||||||
"""The operation's current lifecycle state: "running", "finished",
|
|
||||||
"failed", or "cancelled".
|
|
||||||
|
|
||||||
See :meth:`AsyncJob.status`.
|
|
||||||
"""
|
|
||||||
if self._inner is None:
|
|
||||||
return "finished"
|
|
||||||
return LOOP.run(self._inner.status())
|
|
||||||
|
|
||||||
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())
|
|
||||||
@@ -92,10 +92,8 @@ class LanceMergeInsertBuilder(object):
|
|||||||
self._when_not_matched_by_source_delete = True
|
self._when_not_matched_by_source_delete = True
|
||||||
if isinstance(condition, Expr):
|
if isinstance(condition, Expr):
|
||||||
self._when_not_matched_by_source_condition_expr = condition._inner
|
self._when_not_matched_by_source_condition_expr = condition._inner
|
||||||
self._when_not_matched_by_source_condition = None
|
elif condition is not None:
|
||||||
else:
|
|
||||||
self._when_not_matched_by_source_condition = condition
|
self._when_not_matched_by_source_condition = condition
|
||||||
self._when_not_matched_by_source_condition_expr = None
|
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def use_index(self, use_index: bool) -> LanceMergeInsertBuilder:
|
def use_index(self, use_index: bool) -> LanceMergeInsertBuilder:
|
||||||
|
|||||||
@@ -38,11 +38,7 @@ from lance_namespace_urllib3_client.models.query_table_request_vector import (
|
|||||||
QueryTableRequestVector,
|
QueryTableRequestVector,
|
||||||
)
|
)
|
||||||
from lance_namespace_urllib3_client.models.string_fts_query import StringFtsQuery
|
from lance_namespace_urllib3_client.models.string_fts_query import StringFtsQuery
|
||||||
from lance_namespace.errors import (
|
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||||
NamespaceNotEmptyError,
|
|
||||||
NamespaceNotFoundError,
|
|
||||||
TableNotFoundError,
|
|
||||||
)
|
|
||||||
from lancedb._lancedb import (
|
from lancedb._lancedb import (
|
||||||
connect_namespace as _connect_namespace,
|
connect_namespace as _connect_namespace,
|
||||||
connect_namespace_client as _connect_namespace_client,
|
connect_namespace_client as _connect_namespace_client,
|
||||||
@@ -57,8 +53,6 @@ from lance_namespace import (
|
|||||||
DropNamespaceResponse,
|
DropNamespaceResponse,
|
||||||
ListNamespacesResponse,
|
ListNamespacesResponse,
|
||||||
ListTablesResponse,
|
ListTablesResponse,
|
||||||
NamespaceExistsRequest,
|
|
||||||
TableExistsRequest,
|
|
||||||
)
|
)
|
||||||
from lancedb.table import AsyncTable, LanceTable, Table
|
from lancedb.table import AsyncTable, LanceTable, Table
|
||||||
from lancedb.util import validate_table_name
|
from lancedb.util import validate_table_name
|
||||||
@@ -786,51 +780,6 @@ class LanceNamespaceDBConnection(DBConnection):
|
|||||||
"""
|
"""
|
||||||
return LOOP.run(self._inner.describe_namespace(namespace_path))
|
return LOOP.run(self._inner.describe_namespace(namespace_path))
|
||||||
|
|
||||||
@override
|
|
||||||
def namespace_exists(self, namespace_id: List[str]) -> bool:
|
|
||||||
"""
|
|
||||||
Check if a namespace exists.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
namespace_id : List[str]
|
|
||||||
The namespace identifier to check.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
bool
|
|
||||||
True if the namespace exists, False otherwise.
|
|
||||||
"""
|
|
||||||
request = NamespaceExistsRequest(id=namespace_id)
|
|
||||||
try:
|
|
||||||
self._namespace_client.namespace_exists(request)
|
|
||||||
return True
|
|
||||||
except NamespaceNotFoundError:
|
|
||||||
return False
|
|
||||||
|
|
||||||
@override
|
|
||||||
def table_exists(self, table_id: List[str]) -> bool:
|
|
||||||
"""
|
|
||||||
Check if a table exists.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
table_id : List[str]
|
|
||||||
The table identifier to check (full path including namespace
|
|
||||||
segments and table name).
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
bool
|
|
||||||
True if the table exists, False otherwise.
|
|
||||||
"""
|
|
||||||
request = TableExistsRequest(id=table_id)
|
|
||||||
try:
|
|
||||||
self._namespace_client.table_exists(request)
|
|
||||||
return True
|
|
||||||
except TableNotFoundError:
|
|
||||||
return False
|
|
||||||
|
|
||||||
@override
|
@override
|
||||||
def list_tables(
|
def list_tables(
|
||||||
self,
|
self,
|
||||||
@@ -1284,49 +1233,6 @@ class AsyncLanceNamespaceDBConnection:
|
|||||||
"""
|
"""
|
||||||
return await self._inner.describe_namespace(namespace_path)
|
return await self._inner.describe_namespace(namespace_path)
|
||||||
|
|
||||||
async def namespace_exists(self, namespace_id: List[str]) -> bool:
|
|
||||||
"""
|
|
||||||
Check if a namespace exists.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
namespace_id : List[str]
|
|
||||||
The namespace identifier to check.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
bool
|
|
||||||
True if the namespace exists, False otherwise.
|
|
||||||
"""
|
|
||||||
request = NamespaceExistsRequest(id=namespace_id)
|
|
||||||
try:
|
|
||||||
self._namespace_client.namespace_exists(request)
|
|
||||||
return True
|
|
||||||
except NamespaceNotFoundError:
|
|
||||||
return False
|
|
||||||
|
|
||||||
async def table_exists(self, table_id: List[str]) -> bool:
|
|
||||||
"""
|
|
||||||
Check if a table exists.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
table_id : List[str]
|
|
||||||
The table identifier to check (full path including namespace
|
|
||||||
segments and table name).
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
bool
|
|
||||||
True if the table exists, False otherwise.
|
|
||||||
"""
|
|
||||||
request = TableExistsRequest(id=table_id)
|
|
||||||
try:
|
|
||||||
self._namespace_client.table_exists(request)
|
|
||||||
return True
|
|
||||||
except TableNotFoundError:
|
|
||||||
return False
|
|
||||||
|
|
||||||
async def list_tables(
|
async def list_tables(
|
||||||
self,
|
self,
|
||||||
namespace_path: Optional[List[str]] = None,
|
namespace_path: Optional[List[str]] = None,
|
||||||
|
|||||||
@@ -52,6 +52,7 @@ from ._blob import (
|
|||||||
finalize_blob_query_table,
|
finalize_blob_query_table,
|
||||||
replace_v2_blob_columns_with_bytes,
|
replace_v2_blob_columns_with_bytes,
|
||||||
replace_v2_blob_columns_with_bytes_sync,
|
replace_v2_blob_columns_with_bytes_sync,
|
||||||
|
supports_blob_auto_row_id,
|
||||||
validate_blob_mode,
|
validate_blob_mode,
|
||||||
)
|
)
|
||||||
from .types import BlobMode, QueryProjection
|
from .types import BlobMode, QueryProjection
|
||||||
@@ -1279,7 +1280,10 @@ class LanceQueryBuilder(ABC):
|
|||||||
return self._with_row_id is True
|
return self._with_row_id is True
|
||||||
|
|
||||||
def _blob_auto_row_id_enabled(self) -> bool:
|
def _blob_auto_row_id_enabled(self) -> bool:
|
||||||
|
if not supports_blob_auto_row_id(self._table):
|
||||||
|
return False
|
||||||
return blob_auto_row_id_for_scan(
|
return blob_auto_row_id_for_scan(
|
||||||
|
self._table,
|
||||||
self._table.schema,
|
self._table.schema,
|
||||||
self._columns,
|
self._columns,
|
||||||
with_row_id=self._with_row_id,
|
with_row_id=self._with_row_id,
|
||||||
@@ -2697,7 +2701,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
|||||||
self._fts_query.phrase_query(True)
|
self._fts_query.phrase_query(True)
|
||||||
if self._distance_type:
|
if self._distance_type:
|
||||||
self._vector_query.metric(self._distance_type)
|
self._vector_query.metric(self._distance_type)
|
||||||
if self._minimum_nprobes is not None:
|
if self._minimum_nprobes:
|
||||||
self._vector_query.minimum_nprobes(self._minimum_nprobes)
|
self._vector_query.minimum_nprobes(self._minimum_nprobes)
|
||||||
if self._maximum_nprobes is not None:
|
if self._maximum_nprobes is not None:
|
||||||
self._vector_query.maximum_nprobes(self._maximum_nprobes)
|
self._vector_query.maximum_nprobes(self._maximum_nprobes)
|
||||||
@@ -2770,7 +2774,7 @@ class AsyncQueryBase(object):
|
|||||||
)
|
)
|
||||||
|
|
||||||
async def _maybe_add_blob_row_id(self) -> None:
|
async def _maybe_add_blob_row_id(self) -> None:
|
||||||
if self._table is None:
|
if self._table is None or not supports_blob_auto_row_id(self._table):
|
||||||
self._blob_auto_row_id = False
|
self._blob_auto_row_id = False
|
||||||
self._blob_paths = ()
|
self._blob_paths = ()
|
||||||
return
|
return
|
||||||
@@ -2778,6 +2782,7 @@ class AsyncQueryBase(object):
|
|||||||
req = self._inner.to_query_request()
|
req = self._inner.to_query_request()
|
||||||
schema = await self._table.schema()
|
schema = await self._table.schema()
|
||||||
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||||
|
self._table,
|
||||||
schema,
|
schema,
|
||||||
req.select,
|
req.select,
|
||||||
with_row_id=self._with_row_id,
|
with_row_id=self._with_row_id,
|
||||||
@@ -3029,6 +3034,7 @@ class AsyncQueryBase(object):
|
|||||||
|
|
||||||
schema = await self._table.schema()
|
schema = await self._table.schema()
|
||||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||||
|
self._table,
|
||||||
schema,
|
schema,
|
||||||
query.columns,
|
query.columns,
|
||||||
with_row_id=self._with_row_id,
|
with_row_id=self._with_row_id,
|
||||||
@@ -3874,9 +3880,10 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
|||||||
req = fts_query._inner.to_query_request()
|
req = fts_query._inner.to_query_request()
|
||||||
blob_auto_row_id = False
|
blob_auto_row_id = False
|
||||||
blob_paths: tuple[str, ...] = ()
|
blob_paths: tuple[str, ...] = ()
|
||||||
if self._table is not None:
|
if self._table is not None and supports_blob_auto_row_id(self._table):
|
||||||
schema = await self._table.schema()
|
schema = await self._table.schema()
|
||||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||||
|
self._table,
|
||||||
schema,
|
schema,
|
||||||
req.select,
|
req.select,
|
||||||
with_row_id=self._with_row_id,
|
with_row_id=self._with_row_id,
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ import json
|
|||||||
import logging
|
import logging
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
import sys
|
import sys
|
||||||
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
|
from typing import Any, Dict, Iterable, List, Optional, Union
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
import warnings
|
import warnings
|
||||||
|
|
||||||
@@ -23,10 +23,6 @@ import pyarrow as pa
|
|||||||
|
|
||||||
from ..common import DATA
|
from ..common import DATA
|
||||||
from ..db import DBConnection, LOOP
|
from ..db import DBConnection, LOOP
|
||||||
from ..job import Job
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from .._lancedb import JobDescription, JobInfo
|
|
||||||
from ..embeddings import EmbeddingFunctionConfig
|
from ..embeddings import EmbeddingFunctionConfig
|
||||||
from lance_namespace import (
|
from lance_namespace import (
|
||||||
LanceNamespace,
|
LanceNamespace,
|
||||||
@@ -419,11 +415,6 @@ class RemoteDBConnection(DBConnection):
|
|||||||
|
|
||||||
if namespace_path is None:
|
if namespace_path is None:
|
||||||
namespace_path = []
|
namespace_path = []
|
||||||
if storage_options is not None:
|
|
||||||
logging.info(
|
|
||||||
"storage_options is ignored in LanceDb Cloud"
|
|
||||||
" (storage is managed; set storage_options on connect() instead)"
|
|
||||||
)
|
|
||||||
if index_cache_size is not None:
|
if index_cache_size is not None:
|
||||||
logging.info(
|
logging.info(
|
||||||
"index_cache_size is ignored in LanceDb Cloud"
|
"index_cache_size is ignored in LanceDb Cloud"
|
||||||
@@ -693,47 +684,6 @@ class RemoteDBConnection(DBConnection):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@override
|
|
||||||
def job(self, job_id: str) -> Job:
|
|
||||||
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
|
|
||||||
|
|
||||||
The handle is constructed without a server round trip; an unknown id
|
|
||||||
surfaces when the handle is used. Dropping the handle has no effect
|
|
||||||
on the job itself.
|
|
||||||
"""
|
|
||||||
return Job(self._conn.job(job_id))
|
|
||||||
|
|
||||||
@override
|
|
||||||
def list_jobs(self) -> List["JobInfo"]:
|
|
||||||
"""List server-side jobs across the database's tables."""
|
|
||||||
return LOOP.run(self._conn.list_jobs())
|
|
||||||
|
|
||||||
@override
|
|
||||||
def get_job(self, job_id: str) -> Optional["JobDescription"]:
|
|
||||||
"""Describe a single server-side job by id.
|
|
||||||
|
|
||||||
Returns None when the server has no such job.
|
|
||||||
"""
|
|
||||||
return LOOP.run(self._conn.get_job(job_id))
|
|
||||||
|
|
||||||
@override
|
|
||||||
def cancel_job(self, job_id: str) -> bool:
|
|
||||||
"""Request cancellation of a server-side job by id.
|
|
||||||
|
|
||||||
Returns True if the server accepted the cancellation, False if no
|
|
||||||
such job exists. Cancelling an already-terminal job is a no-op
|
|
||||||
success.
|
|
||||||
"""
|
|
||||||
return LOOP.run(self._conn.cancel_job(job_id))
|
|
||||||
|
|
||||||
@override
|
|
||||||
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
|
||||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
|
||||||
|
|
||||||
Lists history across all jobs when `job_id` is None.
|
|
||||||
"""
|
|
||||||
return LOOP.run(self._conn.job_history(job_id))
|
|
||||||
|
|
||||||
@override
|
@override
|
||||||
def namespace_client(self) -> LanceNamespace:
|
def namespace_client(self) -> LanceNamespace:
|
||||||
"""Get the equivalent namespace client for this connection.
|
"""Get the equivalent namespace client for this connection.
|
||||||
|
|||||||
@@ -20,7 +20,6 @@ from typing import (
|
|||||||
import warnings
|
import warnings
|
||||||
|
|
||||||
from lancedb import __version__
|
from lancedb import __version__
|
||||||
from lancedb._blob import BlobFile
|
|
||||||
|
|
||||||
from lancedb._lancedb import (
|
from lancedb._lancedb import (
|
||||||
AddColumnsResult,
|
AddColumnsResult,
|
||||||
@@ -48,7 +47,6 @@ from lancedb.index import (
|
|||||||
IvfSq,
|
IvfSq,
|
||||||
LabelList,
|
LabelList,
|
||||||
)
|
)
|
||||||
from lancedb.job import Job
|
|
||||||
from lancedb.remote.db import LOOP
|
from lancedb.remote.db import LOOP
|
||||||
from lancedb.table import IndexConfigType, KNOWN_METRICS
|
from lancedb.table import IndexConfigType, KNOWN_METRICS
|
||||||
import pyarrow as pa
|
import pyarrow as pa
|
||||||
@@ -542,34 +540,6 @@ 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(
|
def _is_legacy_create_index_call(
|
||||||
self,
|
self,
|
||||||
first_arg: str,
|
first_arg: str,
|
||||||
@@ -1069,22 +1039,22 @@ class RemoteTable(Table):
|
|||||||
)
|
)
|
||||||
|
|
||||||
def blob_columns(self) -> list[str]:
|
def blob_columns(self) -> list[str]:
|
||||||
return LOOP.run(self._table.blob_columns())
|
raise NotImplementedError(
|
||||||
|
"blob_columns() is not yet supported on the LanceDB Cloud"
|
||||||
|
)
|
||||||
|
|
||||||
def fetch_blobs(
|
def fetch_blobs(self, column: str, row_ids) -> pa.LargeBinaryArray:
|
||||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
raise NotImplementedError("fetch_blobs() is not supported on LanceDB Cloud")
|
||||||
) -> pa.LargeBinaryArray:
|
|
||||||
return LOOP.run(self._table.fetch_blobs(column, row_ids))
|
|
||||||
|
|
||||||
def fetch_blob_ranges(self, column: str, requests) -> pa.LargeBinaryArray:
|
def fetch_blob_ranges(self, column: str, requests) -> pa.LargeBinaryArray:
|
||||||
raise NotImplementedError(
|
raise NotImplementedError(
|
||||||
"fetch_blob_ranges() is not supported on LanceDB Cloud"
|
"fetch_blob_ranges() is not supported on LanceDB Cloud"
|
||||||
)
|
)
|
||||||
|
|
||||||
def fetch_blob_files(
|
def fetch_blob_files(self, column: str, row_ids):
|
||||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
raise NotImplementedError(
|
||||||
) -> "list[Optional[BlobFile]]":
|
"fetch_blob_files() is not supported on LanceDB Cloud"
|
||||||
return LOOP.run(self._table.fetch_blob_files(column, row_ids))
|
)
|
||||||
|
|
||||||
def head(self, n=5) -> pa.Table:
|
def head(self, n=5) -> pa.Table:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -40,7 +40,6 @@ from ._blob import (
|
|||||||
from .types import BlobMode
|
from .types import BlobMode
|
||||||
from lancedb.arrow import peek_reader
|
from lancedb.arrow import peek_reader
|
||||||
from lancedb.background_loop import LOOP, embedding_executor
|
from lancedb.background_loop import LOOP, embedding_executor
|
||||||
from lancedb.job import AsyncJob, Job
|
|
||||||
from .dependencies import (
|
from .dependencies import (
|
||||||
_check_for_hugging_face,
|
_check_for_hugging_face,
|
||||||
_check_for_lance,
|
_check_for_lance,
|
||||||
@@ -978,24 +977,6 @@ class Table(ABC):
|
|||||||
"""
|
"""
|
||||||
raise NotImplementedError
|
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:
|
def drop_index(self, name: str) -> None:
|
||||||
"""
|
"""
|
||||||
Drop an index from the table.
|
Drop an index from the table.
|
||||||
@@ -1593,10 +1574,8 @@ class Table(ABC):
|
|||||||
"""Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles.
|
"""Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles.
|
||||||
|
|
||||||
Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is
|
Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is
|
||||||
a ``list[int]`` or a query ``pyarrow.Table`` carrying row identity via
|
a ``list[int]`` or query ``pyarrow.Table`` with ``_rowid`` (or stashed
|
||||||
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
|
row-id metadata). Null rows are ``None``. Local tables only.
|
||||||
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
|
|
||||||
newer.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
@@ -2490,7 +2469,13 @@ class LanceTable(Table):
|
|||||||
return LOOP.run(self._table.count_rows(filter))
|
return LOOP.run(self._table.count_rows(filter))
|
||||||
|
|
||||||
def __repr__(self) -> str:
|
def __repr__(self) -> str:
|
||||||
return f"{self.__class__.__name__}(name={self.name!r}, _conn={self._conn!r})"
|
val = f"{self.__class__.__name__}(name={self.name!r}"
|
||||||
|
if self._conn.read_consistency_interval is not None:
|
||||||
|
val += ", read_consistency_interval={!r}".format(
|
||||||
|
self._conn.read_consistency_interval
|
||||||
|
)
|
||||||
|
val += f", _conn={self._conn!r})"
|
||||||
|
return val
|
||||||
|
|
||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
return self.__repr__()
|
return self.__repr__()
|
||||||
@@ -2799,34 +2784,6 @@ 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(
|
def _is_legacy_create_index_call(
|
||||||
self,
|
self,
|
||||||
first_arg: str,
|
first_arg: str,
|
||||||
@@ -4914,46 +4871,6 @@ class AsyncTable:
|
|||||||
)
|
)
|
||||||
raise e
|
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:
|
async def drop_index(self, name: str) -> None:
|
||||||
"""
|
"""
|
||||||
Drop an index from the table.
|
Drop an index from the table.
|
||||||
|
|||||||
@@ -226,13 +226,13 @@ def test_fetch_blob_ranges_validates_requests():
|
|||||||
table = _blob_table("range_validation", [{"id": 1, "image": b"abc"}])
|
table = _blob_table("range_validation", [{"id": 1, "image": b"abc"}])
|
||||||
row_id = _row_ids_by_id(table)[1]
|
row_id = _row_ids_by_id(table)[1]
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="exceeds blob size"):
|
with pytest.raises(RuntimeError, match="exceeds blob size"):
|
||||||
table.fetch_blob_ranges("image", [(row_id, 2, 2)])
|
table.fetch_blob_ranges("image", [(row_id, 2, 2)])
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="offset \\+ length overflowed"):
|
with pytest.raises(RuntimeError, match="offset \\+ length overflowed"):
|
||||||
table.fetch_blob_ranges("image", [(row_id, 2**64 - 1, 1)])
|
table.fetch_blob_ranges("image", [(row_id, 2**64 - 1, 1)])
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="row IDs"):
|
with pytest.raises(ValueError, match="row ids"):
|
||||||
table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)])
|
table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)])
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -62,21 +62,6 @@ def test_basic(tmp_path):
|
|||||||
assert db.open_table("test").name == db["test"].name
|
assert db.open_table("test").name == db["test"].name
|
||||||
|
|
||||||
|
|
||||||
def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
|
|
||||||
from lancedb.background_loop import LOOP
|
|
||||||
|
|
||||||
db = lancedb.connect(tmp_path)
|
|
||||||
table = db.create_table("test", data=[{"id": 1}])
|
|
||||||
|
|
||||||
def fail_run(*args, **kwargs):
|
|
||||||
raise AssertionError("repr should not use the Python background loop")
|
|
||||||
|
|
||||||
monkeypatch.setattr(LOOP, "run", fail_run)
|
|
||||||
|
|
||||||
assert repr(db) == f"LanceDBConnection(uri={str(tmp_path)!r})"
|
|
||||||
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
|
|
||||||
|
|
||||||
|
|
||||||
def test_ingest_pd(tmp_path):
|
def test_ingest_pd(tmp_path):
|
||||||
db = lancedb.connect(tmp_path)
|
db = lancedb.connect(tmp_path)
|
||||||
|
|
||||||
|
|||||||
@@ -123,19 +123,6 @@ async def test_async_hybrid_query_default_limit(table: AsyncTable):
|
|||||||
assert texts.count("a") == 1
|
assert texts.count("a") == 1
|
||||||
|
|
||||||
|
|
||||||
def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table):
|
|
||||||
# minimum_nprobes(0) must raise the same validation error a plain vector
|
|
||||||
# query raises, not silently no-op because 0 is falsy.
|
|
||||||
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
|
|
||||||
(
|
|
||||||
sync_table.search(query_type="hybrid")
|
|
||||||
.vector([0.0, 0.4])
|
|
||||||
.text("dog")
|
|
||||||
.minimum_nprobes(0)
|
|
||||||
.to_arrow()
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def test_hybrid_query_distance_range(sync_table: Table):
|
def test_hybrid_query_distance_range(sync_table: Table):
|
||||||
reranker = RRFReranker(return_score="all")
|
reranker = RRFReranker(return_score="all")
|
||||||
result = (
|
result = (
|
||||||
|
|||||||
@@ -84,15 +84,6 @@ 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
|
@pytest.mark.asyncio
|
||||||
async def test_create_scalar_index(some_table: AsyncTable):
|
async def test_create_scalar_index(some_table: AsyncTable):
|
||||||
# Can create
|
# Can create
|
||||||
|
|||||||
@@ -18,7 +18,6 @@ Tests verify:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import copy
|
import copy
|
||||||
import os
|
|
||||||
import shutil
|
import shutil
|
||||||
import sys
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
@@ -240,7 +239,7 @@ def create_tracking_namespace(
|
|||||||
|
|
||||||
dir_props = {f"storage.{k}": v for k, v in storage_options_with_refresh.items()}
|
dir_props = {f"storage.{k}": v for k, v in storage_options_with_refresh.items()}
|
||||||
|
|
||||||
if os.path.isabs(bucket_name) or bucket_name.startswith("file://"):
|
if bucket_name.startswith("/") or bucket_name.startswith("file://"):
|
||||||
dir_props["root"] = f"{bucket_name}/namespace_root"
|
dir_props["root"] = f"{bucket_name}/namespace_root"
|
||||||
else:
|
else:
|
||||||
dir_props["root"] = f"s3://{bucket_name}/namespace_root"
|
dir_props["root"] = f"s3://{bucket_name}/namespace_root"
|
||||||
@@ -768,70 +767,3 @@ def test_namespace_with_schema_only(s3_bucket: str, use_custom: bool):
|
|||||||
|
|
||||||
# Verify data was added
|
# Verify data was added
|
||||||
assert table.count_rows() == 2
|
assert table.count_rows() == 2
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
|
||||||
def test_namespace_exists(use_custom: bool):
|
|
||||||
"""
|
|
||||||
Test namespace_exists returns True for existing and False for non-existent.
|
|
||||||
"""
|
|
||||||
temp_dir = tempfile.mkdtemp()
|
|
||||||
try:
|
|
||||||
ns_client, _ = create_tracking_namespace(
|
|
||||||
bucket_name=temp_dir,
|
|
||||||
storage_options={},
|
|
||||||
credential_expires_in_seconds=3600,
|
|
||||||
use_custom=use_custom,
|
|
||||||
)
|
|
||||||
db = LanceNamespaceDBConnection(ns_client)
|
|
||||||
|
|
||||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
|
||||||
db.create_namespace([namespace_name])
|
|
||||||
|
|
||||||
# Existing namespace should return True
|
|
||||||
assert db.namespace_exists(namespace_id=[namespace_name]) is True
|
|
||||||
|
|
||||||
# Non-existent namespace should return False
|
|
||||||
assert db.namespace_exists(namespace_id=["nonexistent_ns"]) is False
|
|
||||||
finally:
|
|
||||||
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("use_custom", [False, True], ids=["DirectoryNS", "CustomNS"])
|
|
||||||
def test_table_exists(use_custom: bool):
|
|
||||||
"""
|
|
||||||
Test table_exists returns True for existing table and False for non-existent.
|
|
||||||
"""
|
|
||||||
temp_dir = tempfile.mkdtemp()
|
|
||||||
try:
|
|
||||||
ns_client, _ = create_tracking_namespace(
|
|
||||||
bucket_name=temp_dir,
|
|
||||||
storage_options={},
|
|
||||||
credential_expires_in_seconds=3600,
|
|
||||||
use_custom=use_custom,
|
|
||||||
)
|
|
||||||
db = LanceNamespaceDBConnection(ns_client)
|
|
||||||
|
|
||||||
namespace_name = f"test_ns_{uuid.uuid4().hex[:8]}"
|
|
||||||
db.create_namespace([namespace_name])
|
|
||||||
|
|
||||||
table_name = f"test_table_{uuid.uuid4().hex}"
|
|
||||||
namespace_path = [namespace_name]
|
|
||||||
schema = pa.schema(
|
|
||||||
[
|
|
||||||
pa.field("id", pa.int64()),
|
|
||||||
pa.field("vector", pa.list_(pa.float32(), 2)),
|
|
||||||
pa.field("text", pa.string()),
|
|
||||||
]
|
|
||||||
)
|
|
||||||
|
|
||||||
db.create_table(table_name, schema=schema, namespace_path=namespace_path)
|
|
||||||
|
|
||||||
# Existing table should return True
|
|
||||||
table_id = namespace_path + [table_name]
|
|
||||||
assert db.table_exists(table_id=table_id) is True
|
|
||||||
|
|
||||||
# Non-existent table should return False
|
|
||||||
assert db.table_exists(table_id=namespace_path + ["nonexistent_table"]) is False
|
|
||||||
finally:
|
|
||||||
shutil.rmtree(temp_dir, ignore_errors=True)
|
|
||||||
|
|||||||
@@ -812,121 +812,6 @@ def test_table_create_indices():
|
|||||||
table.drop_index("custom_fts_idx")
|
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():
|
def test_remote_create_index_new_api():
|
||||||
received_requests = []
|
received_requests = []
|
||||||
|
|
||||||
@@ -1135,7 +1020,7 @@ def query_test_table(query_handler, *, server_version=Version("0.1.0")):
|
|||||||
request.send_header("Content-Type", "application/json")
|
request.send_header("Content-Type", "application/json")
|
||||||
request.send_header("phalanx-version", str(server_version))
|
request.send_header("phalanx-version", str(server_version))
|
||||||
request.end_headers()
|
request.end_headers()
|
||||||
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
|
request.wfile.write(b"{}")
|
||||||
elif request.path == "/v1/table/test/query/":
|
elif request.path == "/v1/table/test/query/":
|
||||||
content_len = int(request.headers.get("Content-Length"))
|
content_len = int(request.headers.get("Content-Length"))
|
||||||
body = request.rfile.read(content_len)
|
body = request.rfile.read(content_len)
|
||||||
@@ -1973,330 +1858,3 @@ def test_inherited_remote_table_reopens_after_fork():
|
|||||||
finally:
|
finally:
|
||||||
server.shutdown()
|
server.shutdown()
|
||||||
server_thread.join()
|
server_thread.join()
|
||||||
|
|
||||||
|
|
||||||
BLOB_DESCRIBE_RESPONSE = {
|
|
||||||
"table": "test",
|
|
||||||
"version": 1,
|
|
||||||
"schema": {
|
|
||||||
"fields": [
|
|
||||||
{"name": "id", "type": {"type": "int64"}, "nullable": False},
|
|
||||||
{
|
|
||||||
"name": "image",
|
|
||||||
"type": {
|
|
||||||
"type": "struct",
|
|
||||||
"fields": [
|
|
||||||
{
|
|
||||||
"name": "data",
|
|
||||||
"type": {"type": "large_binary"},
|
|
||||||
"nullable": True,
|
|
||||||
},
|
|
||||||
{"name": "uri", "type": {"type": "string"}, "nullable": True},
|
|
||||||
],
|
|
||||||
},
|
|
||||||
"nullable": True,
|
|
||||||
"metadata": {
|
|
||||||
"ARROW:extension:name": "lance.blob.v2",
|
|
||||||
"ARROW:extension:metadata": "",
|
|
||||||
},
|
|
||||||
},
|
|
||||||
]
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def blob_query_response_table():
|
|
||||||
image_field = pa.field(
|
|
||||||
"image",
|
|
||||||
pa.struct(
|
|
||||||
[
|
|
||||||
pa.field("kind", pa.uint8(), nullable=False),
|
|
||||||
pa.field("position", pa.uint64(), nullable=False),
|
|
||||||
pa.field("size", pa.uint64(), nullable=False),
|
|
||||||
pa.field("blob_id", pa.uint32(), nullable=False),
|
|
||||||
pa.field("blob_uri", pa.string(), nullable=False),
|
|
||||||
]
|
|
||||||
),
|
|
||||||
metadata={"lance-encoding:blob": "true"},
|
|
||||||
)
|
|
||||||
images = pa.StructArray.from_arrays(
|
|
||||||
[
|
|
||||||
pa.array([1, 0, 0], type=pa.uint8()),
|
|
||||||
pa.array([0, 0, 0], type=pa.uint64()),
|
|
||||||
pa.array([5, 0, 5], type=pa.uint64()),
|
|
||||||
pa.array([1, 0, 2], type=pa.uint32()),
|
|
||||||
pa.array(["", "", ""], type=pa.string()),
|
|
||||||
],
|
|
||||||
fields=image_field.type,
|
|
||||||
mask=pa.array([False, True, False]),
|
|
||||||
)
|
|
||||||
return pa.Table.from_arrays(
|
|
||||||
[
|
|
||||||
pa.array([1, 2, 3], type=pa.int64()),
|
|
||||||
images,
|
|
||||||
pa.array([10, 20, 30], type=pa.uint64()),
|
|
||||||
],
|
|
||||||
schema=pa.schema(
|
|
||||||
[
|
|
||||||
pa.field("id", pa.int64(), nullable=False),
|
|
||||||
image_field,
|
|
||||||
pa.field("_rowid", pa.uint64()),
|
|
||||||
]
|
|
||||||
),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@contextlib.contextmanager
|
|
||||||
def blob_remote_table(*, server_version=Version("0.5.0")):
|
|
||||||
def handler(request):
|
|
||||||
if request.path == "/v1/table/test/describe/":
|
|
||||||
request.send_response(200)
|
|
||||||
request.send_header("Content-Type", "application/json")
|
|
||||||
request.send_header("phalanx-version", str(server_version))
|
|
||||||
request.end_headers()
|
|
||||||
request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
|
|
||||||
elif request.path.startswith("/v1/table/test/blob/image/"):
|
|
||||||
path = request.path.partition("?")[0]
|
|
||||||
row_id = int(path.split("/")[-2])
|
|
||||||
payload = {10: b"alpha", 20: None, 30: b"gamma"}[row_id]
|
|
||||||
if payload is None:
|
|
||||||
request.send_response(204)
|
|
||||||
request.end_headers()
|
|
||||||
return
|
|
||||||
byte_range = request.headers["Range"].removeprefix("bytes=")
|
|
||||||
start_text, end_text = byte_range.split("-", maxsplit=1)
|
|
||||||
start = int(start_text)
|
|
||||||
end = int(end_text) if end_text else len(payload) - 1
|
|
||||||
chunk = payload[start : end + 1]
|
|
||||||
request.send_response(206)
|
|
||||||
request.send_header("Content-Range", f"bytes {start}-{end}/{len(payload)}")
|
|
||||||
request.send_header("Content-Length", str(len(chunk)))
|
|
||||||
request.end_headers()
|
|
||||||
request.wfile.write(chunk)
|
|
||||||
elif request.path == "/v1/table/test/query/":
|
|
||||||
content_len = int(request.headers.get("Content-Length", 0))
|
|
||||||
body = json.loads(request.rfile.read(content_len))
|
|
||||||
assert body["columns"] == ["id", "image"]
|
|
||||||
assert body["with_row_id"] is True
|
|
||||||
response_table = blob_query_response_table()
|
|
||||||
request.send_response(200)
|
|
||||||
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
|
|
||||||
request.end_headers()
|
|
||||||
with pa.ipc.new_file(request.wfile, response_table.schema) as writer:
|
|
||||||
writer.write_table(response_table)
|
|
||||||
elif request.path == "/v1/table/test/fetch_blobs/":
|
|
||||||
content_len = int(request.headers.get("Content-Length", 0))
|
|
||||||
body = json.loads(request.rfile.read(content_len))
|
|
||||||
assert body["column"] == "image"
|
|
||||||
assert body["row_ids"] == [10, 20, 30]
|
|
||||||
response_table = pa.table(
|
|
||||||
{"image": pa.array([b"alpha", None, b"gamma"], type=pa.large_binary())}
|
|
||||||
)
|
|
||||||
request.send_response(200)
|
|
||||||
request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
|
|
||||||
request.end_headers()
|
|
||||||
with pa.ipc.new_stream(request.wfile, response_table.schema) as writer:
|
|
||||||
writer.write_table(response_table)
|
|
||||||
else:
|
|
||||||
request.send_response(404)
|
|
||||||
request.end_headers()
|
|
||||||
|
|
||||||
with mock_lancedb_connection(handler) as db:
|
|
||||||
yield db.open_table("test")
|
|
||||||
|
|
||||||
|
|
||||||
def test_remote_blob_columns_and_fetch():
|
|
||||||
with blob_remote_table() as table:
|
|
||||||
assert table.blob_columns() == ["image"]
|
|
||||||
blobs = table.fetch_blobs("image", [10, 20, 30])
|
|
||||||
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_remote_blob_files_are_lazy_seekable_handles():
|
|
||||||
with blob_remote_table() as table:
|
|
||||||
files = table.fetch_blob_files("image", [10, 20, 30])
|
|
||||||
|
|
||||||
assert len(files) == 3
|
|
||||||
alpha, null_row, gamma = files
|
|
||||||
assert null_row is None
|
|
||||||
assert alpha is not None
|
|
||||||
assert gamma is not None
|
|
||||||
assert alpha.size() == 5
|
|
||||||
assert alpha.read_range(1, 3) == b"lph"
|
|
||||||
gamma.seek(2)
|
|
||||||
assert gamma.read() == b"mma"
|
|
||||||
|
|
||||||
|
|
||||||
def test_remote_blob_fetch_accepts_query_table():
|
|
||||||
hits = pa.table({"_rowid": pa.array([10, 20, 30], type=pa.uint64())})
|
|
||||||
|
|
||||||
with blob_remote_table() as table:
|
|
||||||
blobs = table.fetch_blobs("image", hits)
|
|
||||||
|
|
||||||
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_remote_blob_query_stashes_row_ids_for_fetch():
|
|
||||||
with blob_remote_table() as table:
|
|
||||||
hits = table.search().select(["id", "image"]).limit(3).to_arrow()
|
|
||||||
assert "_rowid" not in hits.column_names
|
|
||||||
assert "_lance_row_id" in hits.schema.field("image").type.names
|
|
||||||
blobs = table.fetch_blobs("image", hits)
|
|
||||||
|
|
||||||
assert blobs.to_pylist() == [b"alpha", None, b"gamma"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_remote_blob_query_survives_a_server_that_ignores_the_row_id_request():
|
|
||||||
def handler(request):
|
|
||||||
if request.path == "/v1/table/test/describe/":
|
|
||||||
request.send_response(200)
|
|
||||||
request.send_header("Content-Type", "application/json")
|
|
||||||
request.send_header("phalanx-version", "0.5.0")
|
|
||||||
request.end_headers()
|
|
||||||
request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode())
|
|
||||||
elif request.path == "/v1/table/test/query/":
|
|
||||||
content_len = int(request.headers.get("Content-Length", 0))
|
|
||||||
assert json.loads(request.rfile.read(content_len))["with_row_id"] is True
|
|
||||||
response_table = blob_query_response_table().drop_columns(["_rowid"])
|
|
||||||
request.send_response(200)
|
|
||||||
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
|
|
||||||
request.end_headers()
|
|
||||||
with pa.ipc.new_file(request.wfile, response_table.schema) as writer:
|
|
||||||
writer.write_table(response_table)
|
|
||||||
else:
|
|
||||||
request.send_response(404)
|
|
||||||
request.end_headers()
|
|
||||||
|
|
||||||
with mock_lancedb_connection(handler) as db:
|
|
||||||
table = db.open_table("test")
|
|
||||||
hits = table.search().select(["id", "image"]).limit(3).to_arrow()
|
|
||||||
|
|
||||||
assert hits.column_names == ["id", "image"]
|
|
||||||
assert "_lance_row_id" not in hits.schema.field("image").type.names
|
|
||||||
with pytest.raises(ValueError, match="pass a list of row ids"):
|
|
||||||
table.fetch_blobs("image", hits)
|
|
||||||
|
|
||||||
|
|
||||||
def test_remote_blob_byte_apis_not_supported_on_old_server():
|
|
||||||
with blob_remote_table(server_version=Version("0.1.0")) as table:
|
|
||||||
assert table.blob_columns() == ["image"]
|
|
||||||
with pytest.raises(NotImplementedError, match="not supported"):
|
|
||||||
table.fetch_blobs("image", [1])
|
|
||||||
with pytest.raises(NotImplementedError, match="not supported"):
|
|
||||||
table.fetch_blob_files("image", [1])
|
|
||||||
|
|
||||||
|
|
||||||
def test_remote_connection_jobs_surface():
|
|
||||||
from lancedb.exceptions import JobFailedError
|
|
||||||
|
|
||||||
schema = pa.schema([("state", pa.string())])
|
|
||||||
batch = pa.record_batch([pa.array(["created", "done"])], schema=schema)
|
|
||||||
sink = pa.BufferOutputStream()
|
|
||||||
with pa.ipc.new_stream(sink, schema) as writer:
|
|
||||||
writer.write_batch(batch)
|
|
||||||
events_body = sink.getvalue().to_pybytes()
|
|
||||||
|
|
||||||
def handler(request):
|
|
||||||
content_len = int(request.headers.get("Content-Length", 0))
|
|
||||||
body = request.rfile.read(content_len) if content_len > 0 else b""
|
|
||||||
payload = json.loads(body) if body else {}
|
|
||||||
if request.path == "/v1/jobs/list":
|
|
||||||
if payload.get("page_token") is None:
|
|
||||||
rsp = dict(
|
|
||||||
jobs=[
|
|
||||||
dict(
|
|
||||||
job_id="job-1",
|
|
||||||
table="t1",
|
|
||||||
job_type="create_index",
|
|
||||||
state="in_progress",
|
|
||||||
created_at_millis=1000,
|
|
||||||
)
|
|
||||||
],
|
|
||||||
page_token="next",
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
assert payload["page_token"] == "next"
|
|
||||||
rsp = dict(
|
|
||||||
jobs=[
|
|
||||||
dict(
|
|
||||||
job_id="job-2",
|
|
||||||
table="t2",
|
|
||||||
job_type="create_index",
|
|
||||||
state="succeeded",
|
|
||||||
created_at_millis=2000,
|
|
||||||
)
|
|
||||||
]
|
|
||||||
)
|
|
||||||
request.send_response(200)
|
|
||||||
request.send_header("Content-Type", "application/json")
|
|
||||||
request.end_headers()
|
|
||||||
request.wfile.write(json.dumps(rsp).encode())
|
|
||||||
elif request.path == "/v1/jobs/describe":
|
|
||||||
if payload["job_id"] != "job-1":
|
|
||||||
request.send_response(404)
|
|
||||||
request.end_headers()
|
|
||||||
return
|
|
||||||
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_type="create_index",
|
|
||||||
job_state="FAILED",
|
|
||||||
creation_ms=1000,
|
|
||||||
spec=dict(column="vec"),
|
|
||||||
failure=dict(
|
|
||||||
phase="execute", message="worker died", retryable=True
|
|
||||||
),
|
|
||||||
)
|
|
||||||
).encode()
|
|
||||||
)
|
|
||||||
elif request.path == "/v1/jobs/cancel":
|
|
||||||
if payload["job_id"] != "job-1":
|
|
||||||
request.send_response(404)
|
|
||||||
request.end_headers()
|
|
||||||
return
|
|
||||||
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/query_events":
|
|
||||||
assert payload["job_id"] == "job-1"
|
|
||||||
request.send_response(200)
|
|
||||||
request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
|
|
||||||
request.end_headers()
|
|
||||||
request.wfile.write(events_body)
|
|
||||||
else:
|
|
||||||
request.send_response(404)
|
|
||||||
request.end_headers()
|
|
||||||
|
|
||||||
with mock_lancedb_connection(handler) as db:
|
|
||||||
jobs = db.list_jobs()
|
|
||||||
assert [job.job_id for job in jobs] == ["job-1", "job-2"]
|
|
||||||
assert jobs[0].state == "running"
|
|
||||||
assert jobs[0].table == "t1"
|
|
||||||
assert jobs[1].state == "finished"
|
|
||||||
|
|
||||||
description = db.get_job("job-1")
|
|
||||||
assert description.job_type == "create_index"
|
|
||||||
assert description.state == "failed"
|
|
||||||
assert json.loads(description.spec_json) == {"column": "vec"}
|
|
||||||
assert description.failure.message == "worker died"
|
|
||||||
assert description.failure.retryable is True
|
|
||||||
assert db.get_job("missing") is None
|
|
||||||
|
|
||||||
assert db.cancel_job("job-1") is True
|
|
||||||
assert db.cancel_job("missing") is False
|
|
||||||
|
|
||||||
batches = db.job_history("job-1")
|
|
||||||
assert len(batches) == 1
|
|
||||||
assert batches[0].num_rows == 2
|
|
||||||
assert batches[0].column("state").to_pylist() == ["created", "done"]
|
|
||||||
|
|
||||||
job = db.job("job-1")
|
|
||||||
assert job.id == "job-1"
|
|
||||||
assert job.status() == "failed"
|
|
||||||
with pytest.raises(JobFailedError, match="worker died"):
|
|
||||||
job.wait(timeout=timedelta(seconds=5))
|
|
||||||
|
|||||||
@@ -1402,15 +1402,6 @@ async def test_async_open_table_with_branch_version(tmp_path):
|
|||||||
assert await pinned.count_rows() == 4 # writable again
|
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")
|
@patch("lancedb.table.AsyncTable.create_index")
|
||||||
def test_create_index_method(mock_create_index, mem_db: DBConnection):
|
def test_create_index_method(mock_create_index, mem_db: DBConnection):
|
||||||
table = mem_db.create_table(
|
table = mem_db.create_table(
|
||||||
@@ -2364,29 +2355,6 @@ def test_merge_insert_by_source_delete_expr(mem_db: DBConnection):
|
|||||||
assert table.to_arrow().sort_by("a") == expected
|
assert table.to_arrow().sort_by("a") == expected
|
||||||
|
|
||||||
|
|
||||||
def test_merge_insert_by_source_delete_reconfigure(mem_db: DBConnection):
|
|
||||||
# Calling when_not_matched_by_source_delete() again with no condition must
|
|
||||||
# widen the delete to unconditional, not keep the earlier condition around.
|
|
||||||
table = mem_db.create_table(
|
|
||||||
"my_table",
|
|
||||||
data=pa.table({"a": [1, 2, 3], "b": ["a", "b", "c"]}),
|
|
||||||
)
|
|
||||||
new_data = pa.table({"a": [2, 4], "b": ["x", "z"]})
|
|
||||||
|
|
||||||
merge_insert_res = (
|
|
||||||
table.merge_insert("a")
|
|
||||||
.when_matched_update_all()
|
|
||||||
.when_not_matched_insert_all()
|
|
||||||
.when_not_matched_by_source_delete("a > 2")
|
|
||||||
.when_not_matched_by_source_delete()
|
|
||||||
.execute(new_data)
|
|
||||||
)
|
|
||||||
assert merge_insert_res.num_deleted_rows == 2
|
|
||||||
|
|
||||||
expected = pa.table({"a": [2, 4], "b": ["x", "z"]})
|
|
||||||
assert table.to_arrow().sort_by("a") == expected
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_merge_insert_by_source_delete_expr_async(
|
async def test_merge_insert_by_source_delete_expr_async(
|
||||||
mem_db_async: AsyncConnection,
|
mem_db_async: AsyncConnection,
|
||||||
@@ -3119,6 +3087,9 @@ def test_consistency(tmp_path, consistency_interval):
|
|||||||
|
|
||||||
db2 = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
|
db2 = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
|
||||||
table2 = db2.open_table("my_table")
|
table2 = db2.open_table("my_table")
|
||||||
|
if consistency_interval is not None:
|
||||||
|
assert "read_consistency_interval=datetime.timedelta(" in repr(db2)
|
||||||
|
assert "read_consistency_interval=datetime.timedelta(" in repr(table2)
|
||||||
assert table2.version == table.version
|
assert table2.version == table.version
|
||||||
|
|
||||||
table.add([{"id": 1}])
|
table.add([{"id": 1}])
|
||||||
|
|||||||
@@ -13,11 +13,7 @@ use crate::{
|
|||||||
runtime::future_into_py,
|
runtime::future_into_py,
|
||||||
table::Table,
|
table::Table,
|
||||||
};
|
};
|
||||||
use arrow::{
|
use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow};
|
||||||
datatypes::Schema,
|
|
||||||
ffi_stream::ArrowArrayStreamReader,
|
|
||||||
pyarrow::{FromPyArrow, ToPyArrow},
|
|
||||||
};
|
|
||||||
use lancedb::{
|
use lancedb::{
|
||||||
connection::Connection as LanceConnection,
|
connection::Connection as LanceConnection,
|
||||||
connection::NamespaceClientPushdownOperation,
|
connection::NamespaceClientPushdownOperation,
|
||||||
@@ -28,7 +24,7 @@ use pyo3::{
|
|||||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||||
exceptions::{PyRuntimeError, PyValueError},
|
exceptions::{PyRuntimeError, PyValueError},
|
||||||
pyclass, pyfunction, pymethods,
|
pyclass, pyfunction, pymethods,
|
||||||
types::{PyDict, PyDictMethods, PyList, PyListMethods},
|
types::{PyDict, PyDictMethods},
|
||||||
};
|
};
|
||||||
|
|
||||||
#[pyclass]
|
#[pyclass]
|
||||||
@@ -540,55 +536,6 @@ impl Connection {
|
|||||||
})
|
})
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn job(&self, job_id: String) -> PyResult<crate::job::Job> {
|
|
||||||
let inner = self.get_inner()?.clone();
|
|
||||||
Ok(crate::job::Job::new(inner.job(job_id).infer_error()?))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
|
||||||
let inner = self_.get_inner()?.clone();
|
|
||||||
future_into_py(self_.py(), async move {
|
|
||||||
let jobs = inner.list_jobs().await.infer_error()?;
|
|
||||||
Ok(jobs
|
|
||||||
.into_iter()
|
|
||||||
.map(crate::job::JobInfo::from)
|
|
||||||
.collect::<Vec<_>>())
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn get_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
|
||||||
let inner = self_.get_inner()?.clone();
|
|
||||||
future_into_py(self_.py(), async move {
|
|
||||||
let description = inner.get_job(&job_id).await.infer_error()?;
|
|
||||||
Ok(description.map(crate::job::JobDescription::from))
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
|
||||||
let inner = self_.get_inner()?.clone();
|
|
||||||
future_into_py(self_.py(), async move {
|
|
||||||
inner.cancel_job(&job_id).await.infer_error()
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
#[pyo3(signature = (job_id=None))]
|
|
||||||
pub fn job_history(
|
|
||||||
self_: PyRef<'_, Self>,
|
|
||||||
job_id: Option<String>,
|
|
||||||
) -> PyResult<Bound<'_, PyAny>> {
|
|
||||||
let inner = self_.get_inner()?.clone();
|
|
||||||
future_into_py(self_.py(), async move {
|
|
||||||
let batches = inner.job_history(job_id.as_deref()).await.infer_error()?;
|
|
||||||
Python::attach(|py| {
|
|
||||||
let list = PyList::empty(py);
|
|
||||||
for batch in batches {
|
|
||||||
list.append(batch.to_pyarrow(py)?)?;
|
|
||||||
}
|
|
||||||
Ok(list.unbind())
|
|
||||||
})
|
|
||||||
})
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[pyfunction]
|
#[pyfunction]
|
||||||
|
|||||||
@@ -102,18 +102,6 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
|
|||||||
err.setattr(intern!(py, "__cause__"), cause_err)?;
|
err.setattr(intern!(py, "__cause__"), cause_err)?;
|
||||||
Err(PyErr::from_value(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(),
|
_ => self.runtime_error(),
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,145 +0,0 @@
|
|||||||
// 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 status(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
|
||||||
let inner = self_.inner.clone();
|
|
||||||
future_into_py(
|
|
||||||
self_.py(),
|
|
||||||
async move { inner.status().await.infer_error() },
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
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(())
|
|
||||||
})
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A row from `Connection.list_jobs`: one server-side job.
|
|
||||||
#[pyclass(get_all, skip_from_py_object)]
|
|
||||||
#[derive(Clone)]
|
|
||||||
pub struct JobInfo {
|
|
||||||
job_id: String,
|
|
||||||
table: String,
|
|
||||||
job_type: String,
|
|
||||||
state: String,
|
|
||||||
created_at_millis: i64,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[pymethods]
|
|
||||||
impl JobInfo {
|
|
||||||
fn __repr__(&self) -> String {
|
|
||||||
format!(
|
|
||||||
"JobInfo(job_id={:?}, table={:?}, job_type={:?}, state={:?}, created_at_millis={})",
|
|
||||||
self.job_id, self.table, self.job_type, self.state, self.created_at_millis
|
|
||||||
)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl From<lancedb::database::JobInfo> for JobInfo {
|
|
||||||
fn from(info: lancedb::database::JobInfo) -> Self {
|
|
||||||
Self {
|
|
||||||
job_id: info.job_id,
|
|
||||||
table: info.table,
|
|
||||||
job_type: info.job_type,
|
|
||||||
state: info.state,
|
|
||||||
created_at_millis: info.created_at_millis,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The server's account of why a job failed.
|
|
||||||
#[pyclass(get_all, skip_from_py_object)]
|
|
||||||
#[derive(Clone)]
|
|
||||||
pub struct JobFailureInfo {
|
|
||||||
phase: Option<String>,
|
|
||||||
message: Option<String>,
|
|
||||||
retryable: Option<bool>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[pymethods]
|
|
||||||
impl JobFailureInfo {
|
|
||||||
fn __repr__(&self) -> String {
|
|
||||||
format!(
|
|
||||||
"JobFailureInfo(phase={:?}, message={:?}, retryable={:?})",
|
|
||||||
self.phase, self.message, self.retryable
|
|
||||||
)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A described job from `Connection.get_job`.
|
|
||||||
#[pyclass(get_all, skip_from_py_object)]
|
|
||||||
#[derive(Clone)]
|
|
||||||
pub struct JobDescription {
|
|
||||||
job_id: String,
|
|
||||||
job_type: String,
|
|
||||||
state: String,
|
|
||||||
creation_ms: i64,
|
|
||||||
spec_json: Option<String>,
|
|
||||||
failure: Option<JobFailureInfo>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[pymethods]
|
|
||||||
impl JobDescription {
|
|
||||||
fn __repr__(&self) -> String {
|
|
||||||
format!(
|
|
||||||
"JobDescription(job_id={:?}, job_type={:?}, state={:?}, creation_ms={})",
|
|
||||||
self.job_id, self.job_type, self.state, self.creation_ms
|
|
||||||
)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl From<lancedb::database::JobDescription> for JobDescription {
|
|
||||||
fn from(description: lancedb::database::JobDescription) -> Self {
|
|
||||||
Self {
|
|
||||||
job_id: description.job_id,
|
|
||||||
job_type: description.job_type,
|
|
||||||
state: description.state,
|
|
||||||
creation_ms: description.creation_ms,
|
|
||||||
spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
|
|
||||||
failure: description.failure.map(|failure| JobFailureInfo {
|
|
||||||
phase: failure.phase,
|
|
||||||
message: failure.message,
|
|
||||||
retryable: failure.retryable,
|
|
||||||
}),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -25,7 +25,6 @@ pub mod error;
|
|||||||
pub mod expr;
|
pub mod expr;
|
||||||
pub mod header;
|
pub mod header;
|
||||||
pub mod index;
|
pub mod index;
|
||||||
pub mod job;
|
|
||||||
pub mod namespace;
|
pub mod namespace;
|
||||||
pub mod oauth;
|
pub mod oauth;
|
||||||
pub mod otel;
|
pub mod otel;
|
||||||
@@ -45,10 +44,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|||||||
m.add_class::<Connection>()?;
|
m.add_class::<Connection>()?;
|
||||||
m.add_class::<Session>()?;
|
m.add_class::<Session>()?;
|
||||||
m.add_class::<Table>()?;
|
m.add_class::<Table>()?;
|
||||||
m.add_class::<crate::job::Job>()?;
|
|
||||||
m.add_class::<crate::job::JobInfo>()?;
|
|
||||||
m.add_class::<crate::job::JobDescription>()?;
|
|
||||||
m.add_class::<crate::job::JobFailureInfo>()?;
|
|
||||||
m.add_class::<PyBlobFile>()?;
|
m.add_class::<PyBlobFile>()?;
|
||||||
m.add_class::<IndexConfig>()?;
|
m.add_class::<IndexConfig>()?;
|
||||||
m.add_class::<Query>()?;
|
m.add_class::<Query>()?;
|
||||||
|
|||||||
+13
-68
@@ -426,11 +426,9 @@ pub struct PyBlobFile {
|
|||||||
impl PyBlobFile {
|
impl PyBlobFile {
|
||||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||||
let inner = self_.inner.clone();
|
let inner = self_.inner.clone();
|
||||||
let py = self_.py();
|
let bytes = block_on(async move { inner.read().await })
|
||||||
let bytes = py
|
|
||||||
.detach(move || block_on(async move { inner.read().await }))
|
|
||||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||||
@@ -446,32 +444,24 @@ impl PyBlobFile {
|
|||||||
|
|
||||||
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
|
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
|
||||||
let inner = self_.inner.clone();
|
let inner = self_.inner.clone();
|
||||||
self_
|
block_on(async move { inner.close().await })
|
||||||
.py()
|
|
||||||
.detach(move || block_on(async move { inner.close().await }))
|
|
||||||
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
||||||
let inner = self_.inner.clone();
|
let inner = self_.inner.clone();
|
||||||
self_
|
block_on(async move { inner.is_closed().await })
|
||||||
.py()
|
|
||||||
.detach(move || block_on(async move { inner.is_closed().await }))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
|
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
|
||||||
let inner = self_.inner.clone();
|
let inner = self_.inner.clone();
|
||||||
self_
|
block_on(async move { inner.seek(position).await })
|
||||||
.py()
|
|
||||||
.detach(move || block_on(async move { inner.seek(position).await }))
|
|
||||||
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
||||||
let inner = self_.inner.clone();
|
let inner = self_.inner.clone();
|
||||||
self_
|
block_on(async move { inner.tell().await })
|
||||||
.py()
|
|
||||||
.detach(move || block_on(async move { inner.tell().await }))
|
|
||||||
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -485,20 +475,16 @@ impl PyBlobFile {
|
|||||||
.checked_add(length as u64)
|
.checked_add(length as u64)
|
||||||
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
|
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
|
||||||
let inner = self_.inner.clone();
|
let inner = self_.inner.clone();
|
||||||
let py = self_.py();
|
let bytes = block_on(async move { inner.read_range(offset..end).await })
|
||||||
let bytes = py
|
|
||||||
.detach(move || block_on(async move { inner.read_range(offset..end).await }))
|
|
||||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
|
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
|
||||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
|
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
|
||||||
let inner = self_.inner.clone();
|
let inner = self_.inner.clone();
|
||||||
let py = self_.py();
|
let bytes = block_on(async move { inner.read_up_to(length).await })
|
||||||
let bytes = py
|
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||||
.detach(move || block_on(async move { inner.read_up_to(length).await }))
|
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?;
|
|
||||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -819,37 +805,6 @@ 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>> {
|
pub fn drop_index(self_: PyRef<'_, Self>, index_name: String) -> PyResult<Bound<'_, PyAny>> {
|
||||||
let inner = self_.inner_ref()?.clone();
|
let inner = self_.inner_ref()?.clone();
|
||||||
future_into_py(self_.py(), async move {
|
future_into_py(self_.py(), async move {
|
||||||
@@ -1375,12 +1330,7 @@ impl Table {
|
|||||||
|
|
||||||
let inner = self_.inner_ref()?.clone();
|
let inner = self_.inner_ref()?.clone();
|
||||||
future_into_py(self_.py(), async move {
|
future_into_py(self_.py(), async move {
|
||||||
let result = inner
|
let result = inner.add_columns(definitions, None).await.infer_error()?;
|
||||||
.add_columns()
|
|
||||||
.transform(definitions)
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.infer_error()?;
|
|
||||||
Ok(AddColumnsResult::from(result))
|
Ok(AddColumnsResult::from(result))
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -1394,12 +1344,7 @@ impl Table {
|
|||||||
|
|
||||||
let inner = self_.inner_ref()?.clone();
|
let inner = self_.inner_ref()?.clone();
|
||||||
future_into_py(self_.py(), async move {
|
future_into_py(self_.py(), async move {
|
||||||
let result = inner
|
let result = inner.add_columns(transform, None).await.infer_error()?;
|
||||||
.add_columns()
|
|
||||||
.transform(transform)
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.infer_error()?;
|
|
||||||
Ok(AddColumnsResult::from(result))
|
Ok(AddColumnsResult::from(result))
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|||||||
Generated
+1066
-1168
File diff suppressed because it is too large
Load Diff
+1
-1
@@ -1,2 +1,2 @@
|
|||||||
[toolchain]
|
[toolchain]
|
||||||
channel = "1.97.0"
|
channel = "1.95.0"
|
||||||
|
|||||||
+2
-198
@@ -9,7 +9,6 @@
|
|||||||
//!
|
//!
|
||||||
//! Blob tables require Lance file format >= 2.2 and stable row ids at create.
|
//! Blob tables require Lance file format >= 2.2 and stable row ids at create.
|
||||||
|
|
||||||
use std::ops::Range;
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
|
||||||
use arrow_array::LargeBinaryArray;
|
use arrow_array::LargeBinaryArray;
|
||||||
@@ -18,202 +17,10 @@ use arrow_schema::{DataType, Field, Schema};
|
|||||||
use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams};
|
use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams};
|
||||||
use lance_arrow::FieldExt;
|
use lance_arrow::FieldExt;
|
||||||
use lance_encoding::version::LanceFileVersion;
|
use lance_encoding::version::LanceFileVersion;
|
||||||
use lance_io::object_store::ObjectStore;
|
|
||||||
use object_store::path::Path;
|
|
||||||
|
|
||||||
use crate::error::{Error, Result};
|
use crate::error::{Error, Result};
|
||||||
|
|
||||||
/// Seekable handle for one blob value, backed by local storage or a remote
|
pub use lance::dataset::BlobFile;
|
||||||
/// HTTP byte-range endpoint.
|
|
||||||
#[derive(Debug)]
|
|
||||||
pub struct BlobFile {
|
|
||||||
inner: BlobFileInner,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Debug)]
|
|
||||||
enum BlobFileInner {
|
|
||||||
Native(lance::dataset::BlobFile),
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
Remote(Box<crate::remote::table::blobs::RemoteBlobFile>),
|
|
||||||
}
|
|
||||||
|
|
||||||
impl From<lance::dataset::BlobFile> for BlobFile {
|
|
||||||
fn from(value: lance::dataset::BlobFile) -> Self {
|
|
||||||
Self {
|
|
||||||
inner: BlobFileInner::Native(value),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
impl From<crate::remote::table::blobs::RemoteBlobFile> for BlobFile {
|
|
||||||
fn from(value: crate::remote::table::blobs::RemoteBlobFile) -> Self {
|
|
||||||
Self {
|
|
||||||
inner: BlobFileInner::Remote(Box::new(value)),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl BlobFile {
|
|
||||||
/// Inline reader over a data-file slice.
|
|
||||||
pub fn new_inline(
|
|
||||||
object_store: Arc<ObjectStore>,
|
|
||||||
path: Path,
|
|
||||||
position: u64,
|
|
||||||
size: u64,
|
|
||||||
) -> Self {
|
|
||||||
lance::dataset::BlobFile::new_inline(object_store, path, position, size).into()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Dedicated sidecar-file reader.
|
|
||||||
pub fn new_dedicated(object_store: Arc<ObjectStore>, path: Path, size: u64) -> Self {
|
|
||||||
lance::dataset::BlobFile::new_dedicated(object_store, path, size).into()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Packed reader for a slice in a shared sidecar.
|
|
||||||
pub fn new_packed(
|
|
||||||
object_store: Arc<ObjectStore>,
|
|
||||||
path: Path,
|
|
||||||
position: u64,
|
|
||||||
size: u64,
|
|
||||||
) -> Self {
|
|
||||||
lance::dataset::BlobFile::new_packed(object_store, path, position, size).into()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// External reader at a resolved object location.
|
|
||||||
pub fn new_external(
|
|
||||||
object_store: Arc<ObjectStore>,
|
|
||||||
path: Path,
|
|
||||||
uri: String,
|
|
||||||
position: u64,
|
|
||||||
size: u64,
|
|
||||||
) -> Self {
|
|
||||||
lance::dataset::BlobFile::new_external(object_store, path, uri, position, size).into()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Close the handle.
|
|
||||||
pub async fn close(&self) -> lance_core::Result<()> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.close().await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.close().await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Whether the handle is closed.
|
|
||||||
pub async fn is_closed(&self) -> bool {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.is_closed().await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.is_closed(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Read a range without moving the cursor.
|
|
||||||
pub async fn read_range(&self, range: Range<u64>) -> lance_core::Result<bytes::Bytes> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.read_range(range).await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.read_range(range).await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Read ranges without moving the cursor.
|
|
||||||
pub async fn read_ranges(
|
|
||||||
&self,
|
|
||||||
ranges: &[Range<u64>],
|
|
||||||
) -> lance_core::Result<Vec<bytes::Bytes>> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.read_ranges(ranges).await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.read_ranges(ranges).await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Read from the cursor to the end.
|
|
||||||
pub async fn read(&self) -> lance_core::Result<bytes::Bytes> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.read().await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.read().await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Read up to `len` bytes and advance the cursor.
|
|
||||||
pub async fn read_up_to(&self, len: usize) -> lance_core::Result<bytes::Bytes> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.read_up_to(len).await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.read_up_to(len).await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Move the cursor to `new_cursor`.
|
|
||||||
pub async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.seek(new_cursor).await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.seek(new_cursor).await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Current cursor position.
|
|
||||||
pub async fn tell(&self) -> lance_core::Result<u64> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.tell().await,
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.tell().await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Blob length in bytes.
|
|
||||||
pub fn size(&self) -> u64 {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.size(),
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(file) => file.size(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Physical byte offset in the data file. `None` on remote handles. The
|
|
||||||
/// Cloud byte-range route does not expose storage layout.
|
|
||||||
pub fn position(&self) -> Option<u64> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => Some(file.position()),
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(_) => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Path of the data file holding the blob. `None` on remote handles. The
|
|
||||||
/// Cloud byte-range route does not expose storage layout.
|
|
||||||
pub fn data_path(&self) -> Option<&Path> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => Some(file.data_path()),
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(_) => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Native storage layout. `None` on remote handles. The Cloud byte-range
|
|
||||||
/// route does not expose layout.
|
|
||||||
pub fn kind(&self) -> Option<lance_core::datatypes::BlobKind> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => Some(file.kind()),
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(_) => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// External URI for native handles. Remote handles do not expose storage URIs.
|
|
||||||
pub fn uri(&self) -> Option<&str> {
|
|
||||||
match &self.inner {
|
|
||||||
BlobFileInner::Native(file) => file.uri(),
|
|
||||||
#[cfg(feature = "remote")]
|
|
||||||
BlobFileInner::Remote(_) => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// One row-specific blob range read request.
|
/// One row-specific blob range read request.
|
||||||
///
|
///
|
||||||
@@ -457,10 +264,7 @@ pub(crate) async fn take_blob_files_aligned(
|
|||||||
|
|
||||||
let handles = dataset.take_blobs(row_ids, column).await?;
|
let handles = dataset.take_blobs(row_ids, column).await?;
|
||||||
ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?;
|
ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?;
|
||||||
Ok(handles
|
Ok(handles)
|
||||||
.into_iter()
|
|
||||||
.map(|handle| handle.map(Into::into))
|
|
||||||
.collect())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
|
|||||||
@@ -23,8 +23,8 @@ use crate::connection::create_table::CreateTableBuilder;
|
|||||||
use crate::data::scannable::Scannable;
|
use crate::data::scannable::Scannable;
|
||||||
use crate::database::listing::ListingDatabase;
|
use crate::database::listing::ListingDatabase;
|
||||||
use crate::database::{
|
use crate::database::{
|
||||||
CloneTableRequest, Database, DatabaseOptions, JobDescription, JobInfo, OpenTableRequest,
|
CloneTableRequest, Database, DatabaseOptions, OpenTableRequest, ReadConsistency,
|
||||||
ReadConsistency, TableNamesRequest,
|
TableNamesRequest,
|
||||||
};
|
};
|
||||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||||
use crate::error::{Error, Result};
|
use crate::error::{Error, Result};
|
||||||
@@ -456,10 +456,6 @@ impl Connection {
|
|||||||
///
|
///
|
||||||
/// # Returns
|
/// # Returns
|
||||||
/// Created [`TableRef`], or [`Error::TableNotFound`] if the table does not exist.
|
/// Created [`TableRef`], or [`Error::TableNotFound`] if the table does not exist.
|
||||||
/// If the table's storage is present but holds no readable dataset (for example a
|
|
||||||
/// `<name>.lance` directory left behind by an interrupted drop and re-create, which
|
|
||||||
/// [`Self::table_names`] still lists) this returns [`Error::TableCorrupted`]
|
|
||||||
/// instead.
|
|
||||||
pub fn open_table(&self, name: impl Into<String>) -> OpenTableBuilder {
|
pub fn open_table(&self, name: impl Into<String>) -> OpenTableBuilder {
|
||||||
OpenTableBuilder::new(
|
OpenTableBuilder::new(
|
||||||
self.internal.clone(),
|
self.internal.clone(),
|
||||||
@@ -517,39 +513,6 @@ impl Connection {
|
|||||||
self.internal.read_consistency().await
|
self.internal.read_consistency().await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A [`crate::job::Job`] handle for a server-side job by id, suitable for
|
|
||||||
/// waiting on or cancelling the job.
|
|
||||||
///
|
|
||||||
/// The handle is constructed without a server round trip; an unknown id
|
|
||||||
/// surfaces when the handle is used. Only server-backed databases support
|
|
||||||
/// job handles by id.
|
|
||||||
pub fn job(&self, job_id: impl AsRef<str>) -> Result<crate::job::Job> {
|
|
||||||
self.internal.job(job_id.as_ref())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// List server-side jobs across the database's tables.
|
|
||||||
pub async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
|
||||||
self.internal.list_jobs().await
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Describe a single server-side job by id. `None` when the server has no
|
|
||||||
/// such job.
|
|
||||||
pub async fn get_job(&self, job_id: impl AsRef<str>) -> Result<Option<JobDescription>> {
|
|
||||||
self.internal.get_job(job_id.as_ref()).await
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Request cancellation of a server-side job by id. Returns true if the
|
|
||||||
/// server accepted the cancellation, false if no such job exists.
|
|
||||||
pub async fn cancel_job(&self, job_id: impl AsRef<str>) -> Result<bool> {
|
|
||||||
self.internal.cancel_job(job_id.as_ref()).await
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The lifecycle event history of a server-side job (all jobs when
|
|
||||||
/// `job_id` is `None`), as recorded Arrow batches.
|
|
||||||
pub async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<RecordBatch>> {
|
|
||||||
self.internal.job_history(job_id).await
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Drop a table in the database.
|
/// Drop a table in the database.
|
||||||
///
|
///
|
||||||
/// # Arguments
|
/// # Arguments
|
||||||
|
|||||||
@@ -18,8 +18,6 @@ use std::collections::HashMap;
|
|||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use arrow_array::RecordBatch;
|
|
||||||
|
|
||||||
use lance::dataset::ReadParams;
|
use lance::dataset::ReadParams;
|
||||||
use lance_namespace::LanceNamespace;
|
use lance_namespace::LanceNamespace;
|
||||||
use lance_namespace::models::{
|
use lance_namespace::models::{
|
||||||
@@ -202,45 +200,6 @@ pub enum ReadConsistency {
|
|||||||
Strong,
|
Strong,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A row from [`Database::list_jobs`]: one server-side job (index build,
|
|
||||||
/// compaction, column refresh, ...).
|
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
pub struct JobInfo {
|
|
||||||
/// The job id -- what [`Database::get_job`] and [`Database::cancel_job`]
|
|
||||||
/// accept.
|
|
||||||
pub job_id: String,
|
|
||||||
/// The table the job runs against, without URI or namespace.
|
|
||||||
pub table: String,
|
|
||||||
pub job_type: String,
|
|
||||||
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
|
|
||||||
pub state: String,
|
|
||||||
/// When the job was created, in milliseconds since the epoch.
|
|
||||||
pub created_at_millis: i64,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A described job from [`Database::get_job`]: lifecycle state plus the
|
|
||||||
/// job-type-specific specification.
|
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
pub struct JobDescription {
|
|
||||||
pub job_id: String,
|
|
||||||
pub job_type: String,
|
|
||||||
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
|
|
||||||
pub state: String,
|
|
||||||
/// When the job was created, in milliseconds since the epoch.
|
|
||||||
pub creation_ms: i64,
|
|
||||||
/// The job-type-specific specification. Null when the server omits it.
|
|
||||||
pub spec: serde_json::Value,
|
|
||||||
/// Why the job failed, when the job is failed and the server reports a
|
|
||||||
/// reason.
|
|
||||||
pub failure: Option<crate::error::JobFailure>,
|
|
||||||
}
|
|
||||||
|
|
||||||
fn job_op_not_supported<T>(what: &str) -> Result<T> {
|
|
||||||
Err(crate::error::Error::NotSupported {
|
|
||||||
message: format!("{} is not supported by this database", what),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The `Database` trait defines the interface for database implementations.
|
/// The `Database` trait defines the interface for database implementations.
|
||||||
///
|
///
|
||||||
/// A database is responsible for managing tables and their metadata.
|
/// A database is responsible for managing tables and their metadata.
|
||||||
@@ -286,31 +245,6 @@ pub trait Database:
|
|||||||
///
|
///
|
||||||
/// See [`CloneTableRequest`] for detailed documentation and examples.
|
/// See [`CloneTableRequest`] for detailed documentation and examples.
|
||||||
async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>>;
|
async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||||
/// A [`crate::job::Job`] handle for a server-side job by id, suitable for
|
|
||||||
/// waiting on or cancelling the job. The handle is constructed without a
|
|
||||||
/// server round trip; an unknown id surfaces when the handle is used.
|
|
||||||
fn job(&self, _job_id: &str) -> Result<crate::job::Job> {
|
|
||||||
job_op_not_supported("job")
|
|
||||||
}
|
|
||||||
/// List server-side jobs across the database's tables.
|
|
||||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
|
||||||
job_op_not_supported("list_jobs")
|
|
||||||
}
|
|
||||||
/// Describe a single job by id. `None` when the server has no such job.
|
|
||||||
async fn get_job(&self, _job_id: &str) -> Result<Option<JobDescription>> {
|
|
||||||
job_op_not_supported("get_job")
|
|
||||||
}
|
|
||||||
/// Request cancellation of a job by id. Returns true if the server
|
|
||||||
/// accepted the cancellation, false if no such job exists. Cancelling an
|
|
||||||
/// already-terminal job is a no-op success.
|
|
||||||
async fn cancel_job(&self, _job_id: &str) -> Result<bool> {
|
|
||||||
job_op_not_supported("cancel_job")
|
|
||||||
}
|
|
||||||
/// The lifecycle event history of a job (all jobs when `job_id` is
|
|
||||||
/// `None`), as recorded Arrow batches.
|
|
||||||
async fn job_history(&self, _job_id: Option<&str>) -> Result<Vec<RecordBatch>> {
|
|
||||||
job_op_not_supported("job_history")
|
|
||||||
}
|
|
||||||
/// Open a table in the database
|
/// Open a table in the database
|
||||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
|
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||||
/// Rename a table in the database
|
/// Rename a table in the database
|
||||||
|
|||||||
@@ -11,9 +11,7 @@ use lance_core::{cache::LanceCache, utils::futures::FinallyStreamExt};
|
|||||||
use lance_encoding::decoder::{DecoderPlugins, FilterExpression};
|
use lance_encoding::decoder::{DecoderPlugins, FilterExpression};
|
||||||
use lance_file::{
|
use lance_file::{
|
||||||
reader::{FileReader, FileReaderOptions},
|
reader::{FileReader, FileReaderOptions},
|
||||||
version::ConcreteFileVersion,
|
writer::{FileWriter, FileWriterOptions},
|
||||||
versions,
|
|
||||||
writer::FileWriterOptions,
|
|
||||||
};
|
};
|
||||||
use lance_io::{
|
use lance_io::{
|
||||||
ReadBatchParams,
|
ReadBatchParams,
|
||||||
@@ -154,12 +152,8 @@ impl Shuffler {
|
|||||||
source: None,
|
source: None,
|
||||||
})?;
|
})?;
|
||||||
let object_writer = object_store.create(&path).await?;
|
let object_writer = object_store.create(&path).await?;
|
||||||
let writer = versions::create_writer(
|
let writer =
|
||||||
ConcreteFileVersion::V2_1,
|
FileWriter::try_new(object_writer, schema.clone(), FileWriterOptions::default())?;
|
||||||
object_writer,
|
|
||||||
schema.clone(),
|
|
||||||
FileWriterOptions::default(),
|
|
||||||
)?;
|
|
||||||
file_writers.push(writer);
|
file_writers.push(writer);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -264,7 +264,7 @@ pub fn compute_output_schema(
|
|||||||
let field_name = ed
|
let field_name = ed
|
||||||
.dest_column
|
.dest_column
|
||||||
.clone()
|
.clone()
|
||||||
.unwrap_or_else(|| format!("{}_embedding", ed.source_column));
|
.unwrap_or_else(|| format!("{}_embedding", &ed.source_column));
|
||||||
|
|
||||||
sb.push(Field::new(
|
sb.push(Field::new(
|
||||||
field_name,
|
field_name,
|
||||||
@@ -291,7 +291,7 @@ pub fn compute_embeddings_for_batch(
|
|||||||
let dst_field_name = fld
|
let dst_field_name = fld
|
||||||
.dest_column
|
.dest_column
|
||||||
.clone()
|
.clone()
|
||||||
.unwrap_or_else(|| format!("{}_embedding", fld.source_column));
|
.unwrap_or_else(|| format!("{}_embedding", &fld.source_column));
|
||||||
|
|
||||||
let dst_field = Field::new(
|
let dst_field = Field::new(
|
||||||
dst_field_name,
|
dst_field_name,
|
||||||
@@ -315,7 +315,7 @@ impl<R: RecordBatchReader> WithEmbeddings<R> {
|
|||||||
let field_name = ed
|
let field_name = ed
|
||||||
.dest_column
|
.dest_column
|
||||||
.clone()
|
.clone()
|
||||||
.unwrap_or_else(|| format!("{}_embedding", ed.source_column));
|
.unwrap_or_else(|| format!("{}_embedding", &ed.source_column));
|
||||||
Ok(Field::new(
|
Ok(Field::new(
|
||||||
field_name,
|
field_name,
|
||||||
func.dest_type()?.into_owned(),
|
func.dest_type()?.into_owned(),
|
||||||
|
|||||||
@@ -1,8 +1,7 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
use std::fmt::{self, Display, Formatter};
|
use std::sync::PoisonError;
|
||||||
use std::sync::{Arc, PoisonError};
|
|
||||||
|
|
||||||
use arrow_schema::ArrowError;
|
use arrow_schema::ArrowError;
|
||||||
use datafusion_common::DataFusionError;
|
use datafusion_common::DataFusionError;
|
||||||
@@ -10,46 +9,6 @@ use snafu::Snafu;
|
|||||||
|
|
||||||
pub(crate) type BoxError = Box<dyn std::error::Error + Send + Sync>;
|
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)]
|
#[derive(Debug, Snafu)]
|
||||||
#[snafu(visibility(pub(crate)))]
|
#[snafu(visibility(pub(crate)))]
|
||||||
pub enum Error {
|
pub enum Error {
|
||||||
@@ -59,10 +18,6 @@ pub enum Error {
|
|||||||
InvalidInput { message: String },
|
InvalidInput { message: String },
|
||||||
#[snafu(display("Table '{name}' was not found"))]
|
#[snafu(display("Table '{name}' was not found"))]
|
||||||
TableNotFound { name: String, source: BoxError },
|
TableNotFound { name: String, source: BoxError },
|
||||||
#[snafu(display(
|
|
||||||
"Table '{name}' exists but could not be loaded (it may be corrupt or incomplete): {source}"
|
|
||||||
))]
|
|
||||||
TableCorrupted { name: String, source: BoxError },
|
|
||||||
#[snafu(display("Database '{name}' was not found"))]
|
#[snafu(display("Database '{name}' was not found"))]
|
||||||
DatabaseNotFound { name: String },
|
DatabaseNotFound { name: String },
|
||||||
#[snafu(display("Database '{name}' already exists."))]
|
#[snafu(display("Database '{name}' already exists."))]
|
||||||
@@ -85,13 +40,6 @@ pub enum Error {
|
|||||||
Runtime { message: String },
|
Runtime { message: String },
|
||||||
#[snafu(display("Timeout error: {message}"))]
|
#[snafu(display("Timeout error: {message}"))]
|
||||||
Timeout { message: String },
|
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
|
// 3rd party / external errors
|
||||||
#[snafu(display("object_store error: {source}"))]
|
#[snafu(display("object_store error: {source}"))]
|
||||||
@@ -173,9 +121,6 @@ impl From<lance::Error> for Error {
|
|||||||
match source {
|
match source {
|
||||||
lance::Error::Wrapped { error, .. } => Self::from_box_error(error),
|
lance::Error::Wrapped { error, .. } => Self::from_box_error(error),
|
||||||
lance::Error::External { source } => Self::from_box_error(source),
|
lance::Error::External { source } => Self::from_box_error(source),
|
||||||
lance::Error::InvalidInput { source, .. } => Self::InvalidInput {
|
|
||||||
message: source.to_string(),
|
|
||||||
},
|
|
||||||
_ => Self::Lance { source },
|
_ => Self::Lance { source },
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ use std::time::Duration;
|
|||||||
use vector::IvfFlatIndexBuilder;
|
use vector::IvfFlatIndexBuilder;
|
||||||
|
|
||||||
use crate::index::vector::IvfRqIndexBuilder;
|
use crate::index::vector::IvfRqIndexBuilder;
|
||||||
use crate::{DistanceType, Error, Result, job::Job, table::BaseTable};
|
use crate::{DistanceType, Error, Result, table::BaseTable};
|
||||||
|
|
||||||
use self::{
|
use self::{
|
||||||
scalar::{BTreeIndexBuilder, BitmapIndexBuilder, FmIndexBuilder, LabelListIndexBuilder},
|
scalar::{BTreeIndexBuilder, BitmapIndexBuilder, FmIndexBuilder, LabelListIndexBuilder},
|
||||||
@@ -305,14 +305,6 @@ impl IndexBuilder {
|
|||||||
pub async fn execute(self) -> Result<()> {
|
pub async fn execute(self) -> Result<()> {
|
||||||
self.parent.clone().create_index(self).await
|
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)]
|
#[derive(Debug, Clone, PartialEq, Deserialize)]
|
||||||
|
|||||||
@@ -1,182 +0,0 @@
|
|||||||
// 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 status(&self) -> Result<String>;
|
|
||||||
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())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The operation's current lifecycle state: "running", "finished",
|
|
||||||
/// "failed", or "cancelled".
|
|
||||||
///
|
|
||||||
/// A point snapshot; unlike [`Job::wait`] it does not block, raise on a
|
|
||||||
/// terminal failure state, or retry. States a newer server reports that
|
|
||||||
/// this client version does not know pass through as-is.
|
|
||||||
pub async fn status(&self) -> Result<String> {
|
|
||||||
match &self.handle {
|
|
||||||
None => Ok("finished".to_string()),
|
|
||||||
Some(handle) => handle.status().await,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 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 status(&self) -> Result<String> {
|
|
||||||
let label = match &*self.outcome.borrow() {
|
|
||||||
None => "running",
|
|
||||||
Some(Outcome::Succeeded) => "finished",
|
|
||||||
Some(Outcome::Failed(_)) => "failed",
|
|
||||||
Some(Outcome::Cancelled) => "cancelled",
|
|
||||||
};
|
|
||||||
Ok(label.to_string())
|
|
||||||
}
|
|
||||||
|
|
||||||
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,7 +184,6 @@ pub mod expr;
|
|||||||
pub mod index;
|
pub mod index;
|
||||||
pub mod io;
|
pub mod io;
|
||||||
pub mod ipc;
|
pub mod ipc;
|
||||||
pub mod job;
|
|
||||||
#[cfg(feature = "metrics-otel")]
|
#[cfg(feature = "metrics-otel")]
|
||||||
pub mod metrics_otel;
|
pub mod metrics_otel;
|
||||||
#[cfg(feature = "polars")]
|
#[cfg(feature = "polars")]
|
||||||
@@ -204,8 +203,7 @@ use serde::{Deserialize, Serialize};
|
|||||||
|
|
||||||
pub use blob::{BlobRangeRequest, blob, is_blob};
|
pub use blob::{BlobRangeRequest, blob, is_blob};
|
||||||
pub use connection::{ConnectNamespaceBuilder, Connection};
|
pub use connection::{ConnectNamespaceBuilder, Connection};
|
||||||
pub use error::{Error, JobFailure, Result};
|
pub use error::{Error, Result};
|
||||||
pub use job::Job;
|
|
||||||
use lance_index::vector::ApproxMode as LanceApproxMode;
|
use lance_index::vector::ApproxMode as LanceApproxMode;
|
||||||
use lance_linalg::distance::DistanceType as LanceDistanceType;
|
use lance_linalg::distance::DistanceType as LanceDistanceType;
|
||||||
/// Re-export of the [`metrics`](https://docs.rs/metrics) crate facade. Enable
|
/// Re-export of the [`metrics`](https://docs.rs/metrics) crate facade. Enable
|
||||||
|
|||||||
@@ -32,7 +32,6 @@ use crate::{
|
|||||||
};
|
};
|
||||||
|
|
||||||
mod hybrid;
|
mod hybrid;
|
||||||
mod sql;
|
|
||||||
|
|
||||||
pub(crate) const DEFAULT_TOP_K: usize = 10;
|
pub(crate) const DEFAULT_TOP_K: usize = 10;
|
||||||
|
|
||||||
@@ -562,8 +561,8 @@ impl<T: HasQuery> QueryBase for T {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn only_if(mut self, filter: impl AsRef<str>) -> Self {
|
fn only_if(mut self, filter: impl AsRef<str>) -> Self {
|
||||||
let filter = sql::rewrite_st_dwithin(filter.as_ref());
|
self.mut_query()
|
||||||
self.mut_query().add_filter(QueryFilter::Sql(filter));
|
.add_filter(QueryFilter::Sql(filter.as_ref().to_string()));
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1649,8 +1648,8 @@ mod tests {
|
|||||||
use super::*;
|
use super::*;
|
||||||
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
|
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
|
||||||
use arrow_array::{
|
use arrow_array::{
|
||||||
FixedSizeListArray, Float32Array, Float64Array, Int32Array, RecordBatch, StringArray,
|
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray,
|
||||||
cast::AsArray, types::Float32Type,
|
types::Float32Type,
|
||||||
};
|
};
|
||||||
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
|
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
|
||||||
use futures::{StreamExt, TryStreamExt};
|
use futures::{StreamExt, TryStreamExt};
|
||||||
@@ -2025,46 +2024,6 @@ mod tests {
|
|||||||
assert!(result.is_err());
|
assert!(result.is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_st_dwithin_filter() {
|
|
||||||
let tmp_dir = tempdir().unwrap();
|
|
||||||
let uri = tmp_dir.path().to_str().unwrap();
|
|
||||||
let batch = RecordBatch::try_from_iter([
|
|
||||||
("id", Arc::new(Int32Array::from_iter_values([1, 2, 3])) as _),
|
|
||||||
(
|
|
||||||
"x",
|
|
||||||
Arc::new(Float64Array::from_iter_values([0.0, 1.0, 3.0])) as _,
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"y",
|
|
||||||
Arc::new(Float64Array::from_iter_values([0.0, 1.0, 4.0])) as _,
|
|
||||||
),
|
|
||||||
])
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let table = connect(uri)
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.unwrap()
|
|
||||||
.create_table("points", batch)
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
let batches = table
|
|
||||||
.query()
|
|
||||||
.only_if("ST_DWithin(ST_Point(x, y), ST_Point(0.0, 0.0), 2.0)")
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.unwrap()
|
|
||||||
.try_collect::<Vec<_>>()
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
|
|
||||||
let ids = batch["id"].as_primitive::<Int32Type>();
|
|
||||||
|
|
||||||
assert_eq!(ids.values(), &[1, 2]);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn make_non_empty_batches() -> Box<dyn arrow_array::RecordBatchReader + Send> {
|
fn make_non_empty_batches() -> Box<dyn arrow_array::RecordBatchReader + Send> {
|
||||||
let vec = Box::new(RandomVector::new().named("vector".to_string()));
|
let vec = Box::new(RandomVector::new().named("vector".to_string()));
|
||||||
let id = Box::new(IncrementingInt32::new().named("id".to_string()));
|
let id = Box::new(IncrementingInt32::new().named("id".to_string()));
|
||||||
|
|||||||
@@ -1,150 +0,0 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
|
||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
||||||
|
|
||||||
//! SQL compatibility helpers for query filters.
|
|
||||||
|
|
||||||
use std::{any::TypeId, ops::ControlFlow};
|
|
||||||
|
|
||||||
use datafusion_sql::sqlparser::{
|
|
||||||
ast::{
|
|
||||||
BinaryOperator, Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, ObjectName,
|
|
||||||
SelectItem, SetExpr, Statement, visit_expressions_mut,
|
|
||||||
},
|
|
||||||
dialect::{Dialect, GenericDialect},
|
|
||||||
parser::Parser,
|
|
||||||
tokenizer::{Token, Tokenizer},
|
|
||||||
};
|
|
||||||
|
|
||||||
#[derive(Debug, Default)]
|
|
||||||
struct LanceDialect(GenericDialect);
|
|
||||||
|
|
||||||
impl Dialect for LanceDialect {
|
|
||||||
fn dialect(&self) -> TypeId {
|
|
||||||
self.0.dialect()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_identifier_start(&self, ch: char) -> bool {
|
|
||||||
self.0.is_identifier_start(ch)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_identifier_part(&self, ch: char) -> bool {
|
|
||||||
self.0.is_identifier_part(ch)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_delimited_identifier_start(&self, ch: char) -> bool {
|
|
||||||
ch == '`'
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Rewrite `ST_DWithin(a, b, distance)` as `ST_Distance(a, b) <= distance`.
|
|
||||||
///
|
|
||||||
/// GeoDataFusion does not currently expose `ST_DWithin`, but its definition is
|
|
||||||
/// exactly this distance comparison. Invalid SQL and unsupported function
|
|
||||||
/// shapes are left untouched so the query planner can report the usual error.
|
|
||||||
pub(super) fn rewrite_st_dwithin(filter: &str) -> String {
|
|
||||||
if !filter.to_ascii_lowercase().contains("st_dwithin") {
|
|
||||||
return filter.to_string();
|
|
||||||
}
|
|
||||||
|
|
||||||
let Some(mut expr) = parse_filter(filter) else {
|
|
||||||
return filter.to_string();
|
|
||||||
};
|
|
||||||
let mut rewritten = false;
|
|
||||||
let _ = visit_expressions_mut(&mut expr, |expr| {
|
|
||||||
let Expr::Function(function) = expr else {
|
|
||||||
return ControlFlow::<()>::Continue(());
|
|
||||||
};
|
|
||||||
if !function.name.to_string().eq_ignore_ascii_case("st_dwithin") {
|
|
||||||
return ControlFlow::Continue(());
|
|
||||||
}
|
|
||||||
let FunctionArguments::List(arguments) = &function.args else {
|
|
||||||
return ControlFlow::Continue(());
|
|
||||||
};
|
|
||||||
let [
|
|
||||||
FunctionArg::Unnamed(FunctionArgExpr::Expr(first)),
|
|
||||||
FunctionArg::Unnamed(FunctionArgExpr::Expr(second)),
|
|
||||||
FunctionArg::Unnamed(FunctionArgExpr::Expr(distance)),
|
|
||||||
] = arguments.args.as_slice()
|
|
||||||
else {
|
|
||||||
return ControlFlow::Continue(());
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut st_distance = function.clone();
|
|
||||||
st_distance.name = ObjectName::from(vec![Ident::new("ST_Distance")]);
|
|
||||||
if let FunctionArguments::List(arguments) = &mut st_distance.args {
|
|
||||||
arguments.args = vec![
|
|
||||||
FunctionArg::Unnamed(FunctionArgExpr::Expr(first.clone())),
|
|
||||||
FunctionArg::Unnamed(FunctionArgExpr::Expr(second.clone())),
|
|
||||||
];
|
|
||||||
}
|
|
||||||
*expr = Expr::BinaryOp {
|
|
||||||
left: Box::new(Expr::Function(st_distance)),
|
|
||||||
op: BinaryOperator::LtEq,
|
|
||||||
right: Box::new(distance.clone()),
|
|
||||||
};
|
|
||||||
rewritten = true;
|
|
||||||
ControlFlow::Continue(())
|
|
||||||
});
|
|
||||||
|
|
||||||
if rewritten {
|
|
||||||
expr.to_string()
|
|
||||||
} else {
|
|
||||||
filter.to_string()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn parse_filter(filter: &str) -> Option<Expr> {
|
|
||||||
let statement = format!("SELECT 1 FROM t WHERE {filter}");
|
|
||||||
let dialect = LanceDialect::default();
|
|
||||||
let mut token_iter = Tokenizer::new(&dialect, &statement)
|
|
||||||
.tokenize()
|
|
||||||
.ok()?
|
|
||||||
.into_iter();
|
|
||||||
let mut previous = token_iter.next()?;
|
|
||||||
let mut tokens = Vec::new();
|
|
||||||
|
|
||||||
// Match Lance's support for `==` as an equality operator.
|
|
||||||
for next in token_iter {
|
|
||||||
if let (Token::Eq, Token::Eq) = (&previous, &next) {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
tokens.push(std::mem::replace(&mut previous, next));
|
|
||||||
}
|
|
||||||
tokens.push(previous);
|
|
||||||
|
|
||||||
let statement = Parser::new(&dialect)
|
|
||||||
.with_tokens(tokens)
|
|
||||||
.parse_statement()
|
|
||||||
.ok()?;
|
|
||||||
if let Statement::Query(query) = statement
|
|
||||||
&& let SetExpr::Select(select) = *query.body
|
|
||||||
&& let Some(expr) = select.selection
|
|
||||||
&& matches!(select.projection.as_slice(), [SelectItem::UnnamedExpr(_)])
|
|
||||||
{
|
|
||||||
Some(expr)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rewrites_nested_and_case_insensitive_dwithin() {
|
|
||||||
assert_eq!(
|
|
||||||
rewrite_st_dwithin("id == 1 AND st_dwithin(ST_Point(x, y), ST_Point(0, 0), radius)"),
|
|
||||||
"id = 1 AND ST_Distance(ST_Point(x, y), ST_Point(0, 0)) <= radius"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn leaves_other_filters_unchanged() {
|
|
||||||
assert_eq!(rewrite_st_dwithin("id == 1"), "id == 1");
|
|
||||||
assert_eq!(
|
|
||||||
rewrite_st_dwithin("ST_DWithin(point, origin)"),
|
|
||||||
"ST_DWithin(point, origin)"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -8,13 +8,13 @@
|
|||||||
|
|
||||||
pub(crate) mod client;
|
pub(crate) mod client;
|
||||||
pub(crate) mod db;
|
pub(crate) mod db;
|
||||||
pub(crate) mod job;
|
|
||||||
pub mod oauth;
|
pub mod oauth;
|
||||||
mod retry;
|
mod retry;
|
||||||
pub(crate) mod table;
|
pub(crate) mod table;
|
||||||
pub(crate) mod util;
|
pub(crate) mod util;
|
||||||
|
|
||||||
const ARROW_STREAM_CONTENT_TYPE: &str = "application/vnd.apache.arrow.stream";
|
const ARROW_STREAM_CONTENT_TYPE: &str = "application/vnd.apache.arrow.stream";
|
||||||
|
#[cfg(test)]
|
||||||
const ARROW_FILE_CONTENT_TYPE: &str = "application/vnd.apache.arrow.file";
|
const ARROW_FILE_CONTENT_TYPE: &str = "application/vnd.apache.arrow.file";
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
const JSON_CONTENT_TYPE: &str = "application/json";
|
const JSON_CONTENT_TYPE: &str = "application/json";
|
||||||
|
|||||||
@@ -706,7 +706,7 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
|
|||||||
.err_to_http(request_id.clone())?;
|
.err_to_http(request_id.clone())?;
|
||||||
debug!(
|
debug!(
|
||||||
"Received response for request_id={}: {:?}",
|
"Received response for request_id={}: {:?}",
|
||||||
request_id, response
|
request_id, &response
|
||||||
);
|
);
|
||||||
Ok((request_id, response))
|
Ok((request_id, response))
|
||||||
}
|
}
|
||||||
@@ -768,7 +768,7 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
|
|||||||
Ok((status, response)) if status.is_success() => {
|
Ok((status, response)) if status.is_success() => {
|
||||||
debug!(
|
debug!(
|
||||||
"Received response for request_id={}: {:?}",
|
"Received response for request_id={}: {:?}",
|
||||||
retry_counter.request_id, response
|
retry_counter.request_id, &response
|
||||||
);
|
);
|
||||||
return Ok((retry_counter.request_id, response));
|
return Ok((retry_counter.request_id, response));
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ use lance_namespace::models::{
|
|||||||
use crate::Error;
|
use crate::Error;
|
||||||
use crate::database::{
|
use crate::database::{
|
||||||
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
|
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
|
||||||
JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest,
|
OpenTableRequest, ReadConsistency, TableNamesRequest,
|
||||||
};
|
};
|
||||||
use crate::error::Result;
|
use crate::error::Result;
|
||||||
use crate::remote::util::stream_as_body;
|
use crate::remote::util::stream_as_body;
|
||||||
@@ -79,10 +79,6 @@ impl ServerVersion {
|
|||||||
pub fn support_multipart_write(&self) -> bool {
|
pub fn support_multipart_write(&self) -> bool {
|
||||||
self.0 >= semver::Version::new(0, 4, 0)
|
self.0 >= semver::Version::new(0, 4, 0)
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn support_blobs(&self) -> bool {
|
|
||||||
self.0 >= semver::Version::new(0, 5, 0)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub const OPT_REMOTE_PREFIX: &str = "remote_database_";
|
pub const OPT_REMOTE_PREFIX: &str = "remote_database_";
|
||||||
@@ -432,73 +428,6 @@ fn build_cache_key(name: &str, namespace: &[String]) -> String {
|
|||||||
key.iter().map(|b| format!("{:02x}", b)).collect()
|
key.iter().map(|b| format!("{:02x}", b)).collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(serde::Deserialize)]
|
|
||||||
struct RemoteListJobRow {
|
|
||||||
job_id: String,
|
|
||||||
#[serde(default)]
|
|
||||||
table: String,
|
|
||||||
#[serde(default)]
|
|
||||||
job_type: String,
|
|
||||||
#[serde(default)]
|
|
||||||
state: String,
|
|
||||||
#[serde(default)]
|
|
||||||
created_at_millis: i64,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(serde::Deserialize)]
|
|
||||||
struct RemoteListJobsResponse {
|
|
||||||
#[serde(default)]
|
|
||||||
jobs: Vec<RemoteListJobRow>,
|
|
||||||
#[serde(default)]
|
|
||||||
page_token: Option<String>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The server's account of why a job failed. Absent from older servers,
|
|
||||||
/// which report only the terminal state.
|
|
||||||
#[derive(serde::Deserialize)]
|
|
||||||
struct RemoteReportedFailure {
|
|
||||||
#[serde(default)]
|
|
||||||
phase: Option<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
message: Option<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
retryable: Option<bool>,
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(serde::Deserialize)]
|
|
||||||
struct RemoteDescribeJobResponse {
|
|
||||||
job_id: String,
|
|
||||||
#[serde(default)]
|
|
||||||
job_type: String,
|
|
||||||
job_state: String,
|
|
||||||
#[serde(default)]
|
|
||||||
creation_ms: i64,
|
|
||||||
#[serde(default)]
|
|
||||||
spec: serde_json::Value,
|
|
||||||
#[serde(default)]
|
|
||||||
failure: Option<RemoteReportedFailure>,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Server job states -> the client vocabulary ("running" / "finished" /
|
|
||||||
/// "failed" / "cancelled"). Covers both the describe enum (IN_PROGRESS /
|
|
||||||
/// DONE / FAILED / CANCELLED) and the registry's lowercase list-row states
|
|
||||||
/// (in_progress / succeeded / failed / canceled / timed_out). States this
|
|
||||||
/// client version does not know (e.g. created, queued) pass through as-is.
|
|
||||||
fn job_state_to_client(state: &str) -> String {
|
|
||||||
match state {
|
|
||||||
"IN_PROGRESS" | "in_progress" => "running",
|
|
||||||
"DONE" | "done" | "succeeded" => "finished",
|
|
||||||
"FAILED" | "failed" | "TIMED_OUT" | "timed_out" => "failed",
|
|
||||||
"CANCELLED" | "cancelled" | "canceled" => "cancelled",
|
|
||||||
other => other,
|
|
||||||
}
|
|
||||||
.to_string()
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Bound on `list_jobs` page walking; a warning is logged when the listing
|
|
||||||
/// is truncated at this many pages.
|
|
||||||
const MAX_LIST_JOBS_PAGES: usize = 100;
|
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl<S: HttpSend> Database for RemoteDatabase<S> {
|
impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||||
fn uri(&self) -> &str {
|
fn uri(&self) -> &str {
|
||||||
@@ -512,108 +441,6 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn job(&self, job_id: &str) -> Result<crate::job::Job> {
|
|
||||||
Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new(
|
|
||||||
self.client.clone(),
|
|
||||||
job_id.to_string(),
|
|
||||||
))))
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
|
||||||
let mut out = Vec::new();
|
|
||||||
let mut page_token: Option<String> = None;
|
|
||||||
for page in 0..MAX_LIST_JOBS_PAGES {
|
|
||||||
let mut body = serde_json::json!({});
|
|
||||||
if let Some(token) = &page_token {
|
|
||||||
body["page_token"] = serde_json::Value::String(token.clone());
|
|
||||||
}
|
|
||||||
let req = self.client.post("/v1/jobs/list").json(&body);
|
|
||||||
let (request_id, rsp) = self.client.send(req).await?;
|
|
||||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
|
||||||
let body: RemoteListJobsResponse = rsp.json().await.err_to_http(request_id)?;
|
|
||||||
out.extend(body.jobs.into_iter().map(|row| JobInfo {
|
|
||||||
job_id: row.job_id,
|
|
||||||
table: row.table,
|
|
||||||
job_type: row.job_type,
|
|
||||||
state: job_state_to_client(&row.state),
|
|
||||||
created_at_millis: row.created_at_millis,
|
|
||||||
}));
|
|
||||||
page_token = body.page_token;
|
|
||||||
if page_token.is_none() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
if page + 1 == MAX_LIST_JOBS_PAGES {
|
|
||||||
log::warn!(
|
|
||||||
"list_jobs truncated after {} pages ({} jobs)",
|
|
||||||
MAX_LIST_JOBS_PAGES,
|
|
||||||
out.len()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Ok(out)
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn get_job(&self, job_id: &str) -> Result<Option<JobDescription>> {
|
|
||||||
let req = self
|
|
||||||
.client
|
|
||||||
.post("/v1/jobs/describe")
|
|
||||||
.json(&serde_json::json!({ "job_id": job_id }));
|
|
||||||
let (request_id, rsp) = self.client.send(req).await?;
|
|
||||||
let rsp = match self.client.check_response(&request_id, rsp).await {
|
|
||||||
Ok(rsp) => rsp,
|
|
||||||
Err(Error::Http {
|
|
||||||
status_code: Some(StatusCode::NOT_FOUND),
|
|
||||||
..
|
|
||||||
}) => return Ok(None),
|
|
||||||
Err(err) => return Err(err),
|
|
||||||
};
|
|
||||||
let body: RemoteDescribeJobResponse = rsp.json().await.err_to_http(request_id)?;
|
|
||||||
Ok(Some(JobDescription {
|
|
||||||
job_id: body.job_id,
|
|
||||||
job_type: body.job_type,
|
|
||||||
state: job_state_to_client(&body.job_state),
|
|
||||||
creation_ms: body.creation_ms,
|
|
||||||
spec: body.spec,
|
|
||||||
failure: body.failure.map(|reported| crate::error::JobFailure {
|
|
||||||
phase: reported.phase,
|
|
||||||
message: reported.message,
|
|
||||||
retryable: reported.retryable,
|
|
||||||
source: None,
|
|
||||||
}),
|
|
||||||
}))
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
|
||||||
let req = self
|
|
||||||
.client
|
|
||||||
.post("/v1/jobs/cancel")
|
|
||||||
.json(&serde_json::json!({ "job_id": job_id }));
|
|
||||||
let (request_id, rsp) = self.client.send(req).await?;
|
|
||||||
match self.client.check_response(&request_id, rsp).await {
|
|
||||||
Ok(_) => Ok(true),
|
|
||||||
Err(Error::Http {
|
|
||||||
status_code: Some(StatusCode::NOT_FOUND),
|
|
||||||
..
|
|
||||||
}) => Ok(false),
|
|
||||||
Err(err) => Err(err),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<arrow_array::RecordBatch>> {
|
|
||||||
let mut body = serde_json::json!({});
|
|
||||||
if let Some(job_id) = job_id {
|
|
||||||
body["job_id"] = serde_json::Value::String(job_id.to_string());
|
|
||||||
}
|
|
||||||
let req = self.client.post("/v1/jobs/query_events").json(&body);
|
|
||||||
let (request_id, rsp) = self.client.send(req).await?;
|
|
||||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
|
||||||
let bytes = rsp.bytes().await.err_to_http(request_id)?;
|
|
||||||
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None)?;
|
|
||||||
reader
|
|
||||||
.collect::<std::result::Result<Vec<_>, _>>()
|
|
||||||
.map_err(Into::into)
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn table_names(&self, request: TableNamesRequest) -> Result<Vec<String>> {
|
async fn table_names(&self, request: TableNamesRequest) -> Result<Vec<String>> {
|
||||||
let mut req = if !request.namespace_path.is_empty() {
|
let mut req = if !request.namespace_path.is_empty() {
|
||||||
let namespace_id =
|
let namespace_id =
|
||||||
@@ -834,7 +661,6 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
|||||||
RemoteTable::<S>::handle_table_not_found(&request.name, rsp, &request_id).await?;
|
RemoteTable::<S>::handle_table_not_found(&request.name, rsp, &request_id).await?;
|
||||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||||
let version = parse_server_version(&request_id, &rsp)?;
|
let version = parse_server_version(&request_id, &rsp)?;
|
||||||
let describe_body = rsp.text().await.ok();
|
|
||||||
let table_identifier = build_table_identifier(
|
let table_identifier = build_table_identifier(
|
||||||
&request.name,
|
&request.name,
|
||||||
&request.namespace_path,
|
&request.namespace_path,
|
||||||
@@ -847,12 +673,6 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
|||||||
table_identifier,
|
table_identifier,
|
||||||
version,
|
version,
|
||||||
));
|
));
|
||||||
// This describe already carries the schema, so hand it to the table
|
|
||||||
// instead of making the first schema read fetch it again. A version or
|
|
||||||
// branch pin applied after this invalidates the cache.
|
|
||||||
if let Some(body) = &describe_body {
|
|
||||||
table.seed_schema(body);
|
|
||||||
}
|
|
||||||
let cache_key = build_cache_key(&request.name, &request.namespace_path);
|
let cache_key = build_cache_key(&request.name, &request.namespace_path);
|
||||||
self.table_cache.insert(cache_key, table.clone()).await;
|
self.table_cache.insert(cache_key, table.clone()).await;
|
||||||
Ok(table)
|
Ok(table)
|
||||||
@@ -1103,7 +923,6 @@ impl From<StorageOptions> for RemoteOptions {
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::{NamespaceHeaderProviderContext, build_cache_key};
|
use super::{NamespaceHeaderProviderContext, build_cache_key};
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
||||||
use std::sync::{Arc, OnceLock};
|
use std::sync::{Arc, OnceLock};
|
||||||
|
|
||||||
use arrow_array::{Int32Array, RecordBatch};
|
use arrow_array::{Int32Array, RecordBatch};
|
||||||
@@ -1254,46 +1073,6 @@ mod tests {
|
|||||||
assert_eq!(table.name(), "table1");
|
assert_eq!(table.name(), "table1");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_open_table_seeds_the_schema_from_its_describe() {
|
|
||||||
let describe_calls = Arc::new(AtomicUsize::new(0));
|
|
||||||
let counted = describe_calls.clone();
|
|
||||||
let conn = Connection::new_with_handler(move |request| {
|
|
||||||
assert_eq!(request.url().path(), "/v1/table/table1/describe/");
|
|
||||||
counted.fetch_add(1, Ordering::SeqCst);
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(
|
|
||||||
r#"{"version": 1, "schema": {"fields": [
|
|
||||||
{"name": "id", "type": {"type": "int64"}, "nullable": false}
|
|
||||||
]}}"#
|
|
||||||
.to_string(),
|
|
||||||
)
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
|
|
||||||
let table = conn.open_table("table1").execute().await.unwrap();
|
|
||||||
let schema = table.schema().await.unwrap();
|
|
||||||
|
|
||||||
assert_eq!(schema.field(0).name(), "id");
|
|
||||||
assert_eq!(describe_calls.load(Ordering::SeqCst), 1);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_open_table_survives_a_describe_body_it_cannot_parse() {
|
|
||||||
let conn = Connection::new_with_handler(|request| {
|
|
||||||
assert_eq!(request.url().path(), "/v1/table/table1/describe/");
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(r#"{"table": "table1"}"#.to_string())
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
|
|
||||||
let table = conn.open_table("table1").execute().await.unwrap();
|
|
||||||
|
|
||||||
assert_eq!(table.name(), "table1");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_open_table_branch_and_version() {
|
async fn test_open_table_branch_and_version() {
|
||||||
let conn = Connection::new_with_handler(|request| {
|
let conn = Connection::new_with_handler(|request| {
|
||||||
@@ -2263,165 +2042,4 @@ mod tests {
|
|||||||
assert!(list_response.tables.contains(&"table3".to_string()));
|
assert!(list_response.tables.contains(&"table3".to_string()));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_list_jobs_paginates() {
|
|
||||||
let page = Arc::new(AtomicUsize::new(0));
|
|
||||||
let conn = Connection::new_with_handler(move |request| {
|
|
||||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
|
||||||
assert_eq!(request.url().path(), "/v1/jobs/list");
|
|
||||||
let body: serde_json::Value =
|
|
||||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
|
||||||
match page.fetch_add(1, Ordering::SeqCst) {
|
|
||||||
0 => {
|
|
||||||
assert!(body.get("page_token").is_none());
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(
|
|
||||||
r#"{"jobs": [{"job_id": "job-1", "table": "t1", "job_type": "create_index", "state": "in_progress", "created_at_millis": 1000}], "page_token": "next"}"#,
|
|
||||||
)
|
|
||||||
.unwrap()
|
|
||||||
}
|
|
||||||
_ => {
|
|
||||||
assert_eq!(body["page_token"], "next");
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(
|
|
||||||
r#"{"jobs": [{"job_id": "job-2", "table": "t2", "job_type": "create_index", "state": "succeeded", "created_at_millis": 2000}, {"job_id": "job-3", "table": "t3", "job_type": "create_index", "state": "timed_out", "created_at_millis": 3000}]}"#,
|
|
||||||
)
|
|
||||||
.unwrap()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
let jobs = conn.list_jobs().await.unwrap();
|
|
||||||
assert_eq!(jobs.len(), 3);
|
|
||||||
assert_eq!(jobs[0].job_id, "job-1");
|
|
||||||
assert_eq!(jobs[0].table, "t1");
|
|
||||||
assert_eq!(jobs[0].state, "running");
|
|
||||||
assert_eq!(jobs[1].job_id, "job-2");
|
|
||||||
assert_eq!(jobs[1].state, "finished");
|
|
||||||
assert_eq!(jobs[1].created_at_millis, 2000);
|
|
||||||
assert_eq!(jobs[2].job_id, "job-3");
|
|
||||||
assert_eq!(jobs[2].state, "failed");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_get_job() {
|
|
||||||
let conn = Connection::new_with_handler(|request| {
|
|
||||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
|
||||||
assert_eq!(request.url().path(), "/v1/jobs/describe");
|
|
||||||
let body: serde_json::Value =
|
|
||||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
|
||||||
assert_eq!(body["job_id"], "job-1");
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(
|
|
||||||
r#"{"job_id": "job-1", "job_type": "create_index", "job_state": "FAILED", "creation_ms": 1000, "spec": {"column": "vec"}, "failure": {"phase": "execute", "message": "worker died", "retryable": true}}"#,
|
|
||||||
)
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
let job = conn.get_job("job-1").await.unwrap().unwrap();
|
|
||||||
assert_eq!(job.job_id, "job-1");
|
|
||||||
assert_eq!(job.job_type, "create_index");
|
|
||||||
assert_eq!(job.state, "failed");
|
|
||||||
assert_eq!(job.creation_ms, 1000);
|
|
||||||
assert_eq!(job.spec["column"], "vec");
|
|
||||||
let failure = job.failure.unwrap();
|
|
||||||
assert_eq!(failure.phase.as_deref(), Some("execute"));
|
|
||||||
assert_eq!(failure.message.as_deref(), Some("worker died"));
|
|
||||||
assert_eq!(failure.retryable, Some(true));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_get_job_missing_is_none() {
|
|
||||||
let conn = Connection::new_with_handler(|_| {
|
|
||||||
http::Response::builder()
|
|
||||||
.status(404)
|
|
||||||
.body("no such job")
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
assert!(conn.get_job("nope").await.unwrap().is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_cancel_job() {
|
|
||||||
let conn = Connection::new_with_handler(|request| {
|
|
||||||
assert_eq!(request.url().path(), "/v1/jobs/cancel");
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(r#"{"job_id": "job-1"}"#)
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
assert!(conn.cancel_job("job-1").await.unwrap());
|
|
||||||
|
|
||||||
let conn = Connection::new_with_handler(|_| {
|
|
||||||
http::Response::builder()
|
|
||||||
.status(404)
|
|
||||||
.body("no such job")
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
assert!(!conn.cancel_job("nope").await.unwrap());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_job_history_parses_arrow_stream() {
|
|
||||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
|
||||||
"state",
|
|
||||||
DataType::Utf8,
|
|
||||||
false,
|
|
||||||
)]));
|
|
||||||
let batch = RecordBatch::try_new(
|
|
||||||
schema.clone(),
|
|
||||||
vec![Arc::new(arrow_array::StringArray::from(vec![
|
|
||||||
"created", "done",
|
|
||||||
]))],
|
|
||||||
)
|
|
||||||
.unwrap();
|
|
||||||
let mut body = Vec::new();
|
|
||||||
{
|
|
||||||
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &schema).unwrap();
|
|
||||||
writer.write(&batch).unwrap();
|
|
||||||
writer.finish().unwrap();
|
|
||||||
}
|
|
||||||
let conn = Connection::new_with_handler(move |request| {
|
|
||||||
assert_eq!(request.url().path(), "/v1/jobs/query_events");
|
|
||||||
let req_body: serde_json::Value =
|
|
||||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
|
||||||
assert_eq!(req_body["job_id"], "job-1");
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(body.clone())
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
let batches = conn.job_history(Some("job-1")).await.unwrap();
|
|
||||||
assert_eq!(batches.len(), 1);
|
|
||||||
assert_eq!(batches[0].num_rows(), 2);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_conn_job_waits_to_done() {
|
|
||||||
let polls = Arc::new(AtomicUsize::new(0));
|
|
||||||
let polls_ref = polls.clone();
|
|
||||||
let conn = Connection::new_with_handler(move |request| {
|
|
||||||
assert_eq!(request.url().path(), "/v1/jobs/describe");
|
|
||||||
let state = if polls_ref.fetch_add(1, Ordering::SeqCst) == 0 {
|
|
||||||
"IN_PROGRESS"
|
|
||||||
} else {
|
|
||||||
"DONE"
|
|
||||||
};
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(format!(
|
|
||||||
r#"{{"job_id": "job-1", "job_type": "create_index", "job_state": "{}", "creation_ms": 1}}"#,
|
|
||||||
state
|
|
||||||
))
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
let job = conn.job("job-1").unwrap();
|
|
||||||
assert_eq!(job.id(), Some("job-1"));
|
|
||||||
assert_eq!(job.status().await.unwrap(), "running");
|
|
||||||
job.wait().await.unwrap();
|
|
||||||
assert_eq!(job.status().await.unwrap(), "finished");
|
|
||||||
assert!(polls.load(Ordering::SeqCst) >= 3);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,170 +0,0 @@
|
|||||||
// 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 JobState {
|
|
||||||
/// The client vocabulary label for this state.
|
|
||||||
fn client_label(&self) -> String {
|
|
||||||
match self {
|
|
||||||
Self::InProgress => "running".to_string(),
|
|
||||||
Self::Done => "finished".to_string(),
|
|
||||||
Self::Failed => "failed".to_string(),
|
|
||||||
Self::Cancelled => "cancelled".to_string(),
|
|
||||||
Self::Other(state) => state.clone(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl From<&str> for JobState {
|
|
||||||
fn from(state: &str) -> Self {
|
|
||||||
match state {
|
|
||||||
"IN_PROGRESS" => Self::InProgress,
|
|
||||||
"CANCELLED" => Self::Cancelled,
|
|
||||||
// The server reports a timed-out job as FAILED on describe;
|
|
||||||
// accept the raw registry state too in case a future server
|
|
||||||
// stops folding it.
|
|
||||||
"FAILED" | "TIMED_OUT" => 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 {
|
|
||||||
#[serde(default)]
|
|
||||||
phase: Option<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
message: Option<String>,
|
|
||||||
#[serde(default)]
|
|
||||||
retryable: Option<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 status(&self) -> Result<String> {
|
|
||||||
Ok(self.describe().await?.job_state.client_label())
|
|
||||||
}
|
|
||||||
|
|
||||||
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: reported.phase,
|
|
||||||
message: reported.message,
|
|
||||||
retryable: 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(|_| ())
|
|
||||||
}
|
|
||||||
}
|
|
||||||
+254
-1027
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,8 +1,7 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
// SPDX-License-Identifier: Apache-2.0
|
||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
//! DataFusion ExecutionPlan for streaming writes (add / merge_insert) to
|
//! DataFusion ExecutionPlan for inserting data into remote LanceDB tables.
|
||||||
//! remote LanceDB tables.
|
|
||||||
|
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
use std::time::{Duration, Instant};
|
use std::time::{Duration, Instant};
|
||||||
@@ -24,57 +23,28 @@ use lance::io::exec::utils::InstrumentedRecordBatchStreamAdapter;
|
|||||||
use crate::Error;
|
use crate::Error;
|
||||||
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
||||||
use crate::remote::client::{HttpSend, RestfulLanceDbClient, Sender};
|
use crate::remote::client::{HttpSend, RestfulLanceDbClient, Sender};
|
||||||
use crate::remote::table::{MergeInsertRequest, REQUEST_TIMEOUT_HEADER, RemoteTable};
|
use crate::remote::table::RemoteTable;
|
||||||
|
use crate::table::AddResult;
|
||||||
use crate::table::datafusion::insert::COUNT_SCHEMA;
|
use crate::table::datafusion::insert::COUNT_SCHEMA;
|
||||||
use crate::table::write_progress::WriteProgressTracker;
|
use crate::table::write_progress::WriteProgressTracker;
|
||||||
use crate::table::{AddResult, MergeResult};
|
|
||||||
|
|
||||||
/// The write operation a [`RemoteWriteExec`] performs. Both variants share the
|
/// ExecutionPlan for inserting data into a remote LanceDB table.
|
||||||
/// same Arrow-IPC streaming body and error side-channel; only the target
|
|
||||||
/// endpoint, query parameters, and parsed result type differ.
|
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
pub(crate) enum WriteOp {
|
|
||||||
/// `add`: stream to `/v1/table/{id}/insert/`, optionally overwriting.
|
|
||||||
Insert { overwrite: bool },
|
|
||||||
/// `merge_insert`: stream to `/v1/table/{id}/merge_insert/` with the merge
|
|
||||||
/// parameters carried as query params. Multipart is not supported for this
|
|
||||||
/// operation (the server has no multipart merge_insert endpoint), so an
|
|
||||||
/// `upload_id` combined with this op is a programming error.
|
|
||||||
MergeInsert {
|
|
||||||
query: MergeInsertRequest,
|
|
||||||
timeout: Option<Duration>,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
/// The parsed server response for a completed write, discriminated by the
|
|
||||||
/// operation that produced it.
|
|
||||||
#[derive(Debug, Clone)]
|
|
||||||
pub(crate) enum WriteResult {
|
|
||||||
Add(AddResult),
|
|
||||||
Merge(MergeResult),
|
|
||||||
}
|
|
||||||
|
|
||||||
/// ExecutionPlan for streaming a write (add or merge_insert) to a remote
|
|
||||||
/// LanceDB table.
|
|
||||||
///
|
///
|
||||||
/// Streams data as Arrow IPC to the endpoint selected by [`WriteOp`]. Both
|
/// Streams data as Arrow IPC to `/v1/table/{id}/insert/` endpoint.
|
||||||
/// operations reuse the same error side-channel so an input stream error (e.g.
|
|
||||||
/// NaN rejection) surfaces with its original message rather than the masked
|
|
||||||
/// HTTP error Hyper produces when a request body stream fails under HTTP2.
|
|
||||||
///
|
///
|
||||||
/// When `upload_id` is set, inserts are staged as part of a multipart write
|
/// When `upload_id` is set, inserts are staged as part of a multipart write
|
||||||
/// session and the plan supports multiple partitions for parallel uploads.
|
/// session and the plan supports multiple partitions for parallel uploads.
|
||||||
/// Without `upload_id`, the plan requires a single partition and commits
|
/// Without `upload_id`, the plan requires a single partition and commits
|
||||||
/// immediately. Multipart applies to `add` only.
|
/// immediately.
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub struct RemoteWriteExec<S: HttpSend = Sender> {
|
pub struct RemoteInsertExec<S: HttpSend = Sender> {
|
||||||
table_name: String,
|
table_name: String,
|
||||||
identifier: String,
|
identifier: String,
|
||||||
client: RestfulLanceDbClient<S>,
|
client: RestfulLanceDbClient<S>,
|
||||||
input: Arc<dyn ExecutionPlan>,
|
input: Arc<dyn ExecutionPlan>,
|
||||||
op: WriteOp,
|
overwrite: bool,
|
||||||
properties: Arc<PlanProperties>,
|
properties: Arc<PlanProperties>,
|
||||||
result: Arc<Mutex<Option<WriteResult>>>,
|
add_result: Arc<Mutex<Option<AddResult>>>,
|
||||||
metrics: ExecutionPlanMetricsSet,
|
metrics: ExecutionPlanMetricsSet,
|
||||||
upload_id: Option<String>,
|
upload_id: Option<String>,
|
||||||
tracker: Option<Arc<WriteProgressTracker>>,
|
tracker: Option<Arc<WriteProgressTracker>>,
|
||||||
@@ -91,28 +61,27 @@ pub struct RemoteWriteExec<S: HttpSend = Sender> {
|
|||||||
max_request_duration: Option<Duration>,
|
max_request_duration: Option<Duration>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||||
/// Create a new single-partition RemoteWriteExec.
|
/// Create a new single-partition RemoteInsertExec.
|
||||||
pub fn new(
|
pub fn new(
|
||||||
table_name: String,
|
table_name: String,
|
||||||
identifier: String,
|
identifier: String,
|
||||||
client: RestfulLanceDbClient<S>,
|
client: RestfulLanceDbClient<S>,
|
||||||
input: Arc<dyn ExecutionPlan>,
|
input: Arc<dyn ExecutionPlan>,
|
||||||
op: WriteOp,
|
overwrite: bool,
|
||||||
tracker: Option<Arc<WriteProgressTracker>>,
|
tracker: Option<Arc<WriteProgressTracker>>,
|
||||||
branch: Option<String>,
|
branch: Option<String>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self::new_inner(
|
Self::new_inner(
|
||||||
table_name, identifier, client, input, op, None, tracker, branch, None, None,
|
table_name, identifier, client, input, overwrite, None, tracker, branch, None, None,
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Create a multi-partition RemoteWriteExec for use with multipart writes.
|
/// Create a multi-partition RemoteInsertExec for use with multipart writes.
|
||||||
///
|
///
|
||||||
/// Each partition's insert is staged under the given `upload_id` without
|
/// Each partition's insert is staged under the given `upload_id` without
|
||||||
/// committing. The caller is responsible for calling the complete (or abort)
|
/// committing. The caller is responsible for calling the complete (or abort)
|
||||||
/// endpoint after all partitions finish. Multipart is insert-only, so the
|
/// endpoint after all partitions finish.
|
||||||
/// op is fixed to [`WriteOp::Insert`].
|
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub fn new_multipart(
|
pub fn new_multipart(
|
||||||
table_name: String,
|
table_name: String,
|
||||||
@@ -131,7 +100,7 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
|||||||
identifier,
|
identifier,
|
||||||
client,
|
client,
|
||||||
input,
|
input,
|
||||||
WriteOp::Insert { overwrite },
|
overwrite,
|
||||||
Some(upload_id),
|
Some(upload_id),
|
||||||
tracker,
|
tracker,
|
||||||
branch,
|
branch,
|
||||||
@@ -146,7 +115,7 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
|||||||
identifier: String,
|
identifier: String,
|
||||||
client: RestfulLanceDbClient<S>,
|
client: RestfulLanceDbClient<S>,
|
||||||
input: Arc<dyn ExecutionPlan>,
|
input: Arc<dyn ExecutionPlan>,
|
||||||
op: WriteOp,
|
overwrite: bool,
|
||||||
upload_id: Option<String>,
|
upload_id: Option<String>,
|
||||||
tracker: Option<Arc<WriteProgressTracker>>,
|
tracker: Option<Arc<WriteProgressTracker>>,
|
||||||
branch: Option<String>,
|
branch: Option<String>,
|
||||||
@@ -171,9 +140,9 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
|||||||
identifier,
|
identifier,
|
||||||
client,
|
client,
|
||||||
input,
|
input,
|
||||||
op,
|
overwrite,
|
||||||
properties: Arc::new(properties),
|
properties: Arc::new(properties),
|
||||||
result: Arc::new(Mutex::new(None)),
|
add_result: Arc::new(Mutex::new(None)),
|
||||||
metrics: ExecutionPlanMetricsSet::new(),
|
metrics: ExecutionPlanMetricsSet::new(),
|
||||||
upload_id,
|
upload_id,
|
||||||
tracker,
|
tracker,
|
||||||
@@ -183,30 +152,14 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Get the add result after execution, if this exec ran an insert.
|
/// Get the add result after execution.
|
||||||
|
// TODO: this will be used when we wire this up to Table::add().
|
||||||
|
#[allow(dead_code)]
|
||||||
pub fn add_result(&self) -> Option<AddResult> {
|
pub fn add_result(&self) -> Option<AddResult> {
|
||||||
match self
|
self.add_result
|
||||||
.result
|
|
||||||
.lock()
|
.lock()
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
.unwrap_or_else(|e| e.into_inner())
|
||||||
.clone()
|
.clone()
|
||||||
{
|
|
||||||
Some(WriteResult::Add(r)) => Some(r),
|
|
||||||
_ => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Get the merge result after execution, if this exec ran a merge_insert.
|
|
||||||
pub fn merge_result(&self) -> Option<MergeResult> {
|
|
||||||
match self
|
|
||||||
.result
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|e| e.into_inner())
|
|
||||||
.clone()
|
|
||||||
{
|
|
||||||
Some(WriteResult::Merge(r)) => Some(r),
|
|
||||||
_ => None,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Stream the input into an HTTP body as an Arrow IPC stream, capturing any
|
/// Stream the input into an HTTP body as an Arrow IPC stream, capturing any
|
||||||
@@ -511,24 +464,24 @@ impl<S: HttpSend + 'static> PartRequestCtx<'_, S> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<S: HttpSend + 'static> DisplayAs for RemoteWriteExec<S> {
|
impl<S: HttpSend + 'static> DisplayAs for RemoteInsertExec<S> {
|
||||||
fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
match t {
|
match t {
|
||||||
DisplayFormatType::Default | DisplayFormatType::Verbose => {
|
DisplayFormatType::Default | DisplayFormatType::Verbose => {
|
||||||
write!(f, "RemoteWriteExec: table={}, op=", self.table_name)?;
|
write!(
|
||||||
match &self.op {
|
f,
|
||||||
WriteOp::Insert { overwrite } => write!(f, "insert, overwrite={}", overwrite),
|
"RemoteInsertExec: table={}, overwrite={}",
|
||||||
WriteOp::MergeInsert { .. } => write!(f, "merge_insert"),
|
self.table_name, self.overwrite
|
||||||
}
|
)
|
||||||
}
|
}
|
||||||
DisplayFormatType::TreeRender => {
|
DisplayFormatType::TreeRender => {
|
||||||
write!(f, "RemoteWriteExec")
|
write!(f, "RemoteInsertExec")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
impl<S: HttpSend + 'static> ExecutionPlan for RemoteInsertExec<S> {
|
||||||
fn name(&self) -> &str {
|
fn name(&self) -> &str {
|
||||||
Self::static_name()
|
Self::static_name()
|
||||||
}
|
}
|
||||||
@@ -563,18 +516,15 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
|
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
|
||||||
if children.len() != 1 {
|
if children.len() != 1 {
|
||||||
return Err(DataFusionError::Internal(
|
return Err(DataFusionError::Internal(
|
||||||
"RemoteWriteExec requires exactly one child".to_string(),
|
"RemoteInsertExec requires exactly one child".to_string(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
// Building a fresh exec (with a new, empty `result`) is what makes the
|
|
||||||
// outer rescannable retry loop work: `reset_state()` clears the captured
|
|
||||||
// result so a re-execution starts clean.
|
|
||||||
Ok(Arc::new(Self::new_inner(
|
Ok(Arc::new(Self::new_inner(
|
||||||
self.table_name.clone(),
|
self.table_name.clone(),
|
||||||
self.identifier.clone(),
|
self.identifier.clone(),
|
||||||
self.client.clone(),
|
self.client.clone(),
|
||||||
children[0].clone(),
|
children[0].clone(),
|
||||||
self.op.clone(),
|
self.overwrite,
|
||||||
self.upload_id.clone(),
|
self.upload_id.clone(),
|
||||||
self.tracker.clone(),
|
self.tracker.clone(),
|
||||||
self.branch.clone(),
|
self.branch.clone(),
|
||||||
@@ -590,19 +540,11 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
) -> DataFusionResult<SendableRecordBatchStream> {
|
) -> DataFusionResult<SendableRecordBatchStream> {
|
||||||
if self.upload_id.is_none() && partition != 0 {
|
if self.upload_id.is_none() && partition != 0 {
|
||||||
return Err(DataFusionError::Internal(
|
return Err(DataFusionError::Internal(
|
||||||
"RemoteWriteExec only supports single partition execution without upload_id"
|
"RemoteInsertExec only supports single partition execution without upload_id"
|
||||||
.to_string(),
|
.to_string(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Multipart is insert-only: the server has no multipart merge_insert
|
|
||||||
// endpoint, so a merge_insert with an upload_id is a programming error.
|
|
||||||
if self.upload_id.is_some() && matches!(self.op, WriteOp::MergeInsert { .. }) {
|
|
||||||
return Err(DataFusionError::Internal(
|
|
||||||
"merge_insert does not support multipart (upload_id) writes".to_string(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
let input_stream = self.input.execute(partition, context)?;
|
let input_stream = self.input.execute(partition, context)?;
|
||||||
let input_schema = input_stream.schema();
|
let input_schema = input_stream.schema();
|
||||||
let input_stream: SendableRecordBatchStream =
|
let input_stream: SendableRecordBatchStream =
|
||||||
@@ -614,8 +556,8 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
));
|
));
|
||||||
let client = self.client.clone();
|
let client = self.client.clone();
|
||||||
let identifier = self.identifier.clone();
|
let identifier = self.identifier.clone();
|
||||||
let op = self.op.clone();
|
let overwrite = self.overwrite;
|
||||||
let result_slot = self.result.clone();
|
let add_result = self.add_result.clone();
|
||||||
let table_name = self.table_name.clone();
|
let table_name = self.table_name.clone();
|
||||||
let upload_id = self.upload_id.clone();
|
let upload_id = self.upload_id.clone();
|
||||||
let tracker = self.tracker.clone();
|
let tracker = self.tracker.clone();
|
||||||
@@ -626,12 +568,10 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
let stream = futures::stream::once(async move {
|
let stream = futures::stream::once(async move {
|
||||||
// Multipart writes with a byte budget split the partition into
|
// Multipart writes with a byte budget split the partition into
|
||||||
// several bounded, still-streamed requests so no single request
|
// several bounded, still-streamed requests so no single request
|
||||||
// stays open long enough to hit the client read timeout. This path
|
// stays open long enough to hit the client read timeout.
|
||||||
// is insert-only (guarded above).
|
|
||||||
if let (Some(upload_id), Some(max_bytes)) =
|
if let (Some(upload_id), Some(max_bytes)) =
|
||||||
(upload_id.as_deref(), max_bytes_per_request)
|
(upload_id.as_deref(), max_bytes_per_request)
|
||||||
{
|
{
|
||||||
let overwrite = matches!(op, WriteOp::Insert { overwrite: true });
|
|
||||||
let ctx = PartRequestCtx {
|
let ctx = PartRequestCtx {
|
||||||
client: &client,
|
client: &client,
|
||||||
identifier: &identifier,
|
identifier: &identifier,
|
||||||
@@ -652,36 +592,16 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
)?);
|
)?);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Build the request for the selected operation. Both endpoints take
|
let mut request = client
|
||||||
// an Arrow-IPC streaming body and reuse the same error side-channel.
|
.post(&format!("/v1/table/{}/insert/", identifier))
|
||||||
let mut request = match &op {
|
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE);
|
||||||
WriteOp::Insert { overwrite } => {
|
|
||||||
let mut request = client
|
|
||||||
.post(&format!("/v1/table/{}/insert/", identifier))
|
|
||||||
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE);
|
|
||||||
if *overwrite {
|
|
||||||
request = request.query(&[("mode", "overwrite")]);
|
|
||||||
}
|
|
||||||
if let Some(ref uid) = upload_id {
|
|
||||||
request = request.query(&[("upload_id", uid.as_str())]);
|
|
||||||
}
|
|
||||||
request
|
|
||||||
}
|
|
||||||
WriteOp::MergeInsert { query, timeout } => {
|
|
||||||
let mut request = client
|
|
||||||
.post(&format!("/v1/table/{}/merge_insert/", identifier))
|
|
||||||
.query(query)
|
|
||||||
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE);
|
|
||||||
if let Some(timeout) = timeout {
|
|
||||||
// (If it doesn't fit into u64, it's not worth sending anyways.)
|
|
||||||
if let Ok(timeout_ms) = u64::try_from(timeout.as_millis()) {
|
|
||||||
request = request.header(REQUEST_TIMEOUT_HEADER, timeout_ms);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
request
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
|
if overwrite {
|
||||||
|
request = request.query(&[("mode", "overwrite")]);
|
||||||
|
}
|
||||||
|
if let Some(ref uid) = upload_id {
|
||||||
|
request = request.query(&[("upload_id", uid.as_str())]);
|
||||||
|
}
|
||||||
if let Some(ref b) = branch {
|
if let Some(ref b) = branch {
|
||||||
request = request.query(&[("branch", b.as_str())]);
|
request = request.query(&[("branch", b.as_str())]);
|
||||||
}
|
}
|
||||||
@@ -715,8 +635,6 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
|
|
||||||
// If the request failed due to an input stream error, surface the
|
// If the request failed due to an input stream error, surface the
|
||||||
// original error (e.g. NaN rejection) instead of the HTTP error.
|
// original error (e.g. NaN rejection) instead of the HTTP error.
|
||||||
// This is the crux of the #2339 fix: Hyper silently swallows body
|
|
||||||
// stream errors under HTTP2, so we recover the original here.
|
|
||||||
if let Ok(stream_err) = error_rx.try_recv() {
|
if let Ok(stream_err) = error_rx.try_recv() {
|
||||||
return Err(stream_err);
|
return Err(stream_err);
|
||||||
}
|
}
|
||||||
@@ -724,7 +642,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
let (request_id, response) = result?;
|
let (request_id, response) = result?;
|
||||||
|
|
||||||
// For multipart writes, the staging response is not the final
|
// For multipart writes, the staging response is not the final
|
||||||
// version. Only parse the result for non-multipart writes.
|
// version. Only parse AddResult for non-multipart inserts.
|
||||||
if upload_id.is_none() {
|
if upload_id.is_none() {
|
||||||
let body_text = response.text().await.map_err(|e| {
|
let body_text = response.text().await.map_err(|e| {
|
||||||
DataFusionError::External(Box::new(Error::Http {
|
DataFusionError::External(Box::new(Error::Http {
|
||||||
@@ -734,44 +652,21 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
}))
|
}))
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let parsed_result = match &op {
|
let parsed_result = if body_text.trim().is_empty() {
|
||||||
WriteOp::Insert { .. } => {
|
// Backward compatible with old servers
|
||||||
let add = if body_text.trim().is_empty() {
|
AddResult { version: 0 }
|
||||||
// Backward compatible with old servers
|
} else {
|
||||||
AddResult { version: 0 }
|
serde_json::from_str(&body_text).map_err(|e| {
|
||||||
} else {
|
DataFusionError::External(Box::new(Error::Http {
|
||||||
serde_json::from_str(&body_text).map_err(|e| {
|
source: format!("Failed to parse add response: {}", e).into(),
|
||||||
DataFusionError::External(Box::new(Error::Http {
|
request_id: request_id.clone(),
|
||||||
source: format!("Failed to parse add response: {}", e).into(),
|
status_code: None,
|
||||||
request_id: request_id.clone(),
|
}))
|
||||||
status_code: None,
|
})?
|
||||||
}))
|
|
||||||
})?
|
|
||||||
};
|
|
||||||
WriteResult::Add(add)
|
|
||||||
}
|
|
||||||
WriteOp::MergeInsert { .. } => {
|
|
||||||
let merge = if body_text.trim().is_empty() {
|
|
||||||
// Backward compatible with old servers
|
|
||||||
MergeResult::default()
|
|
||||||
} else {
|
|
||||||
serde_json::from_str(&body_text).map_err(|e| {
|
|
||||||
DataFusionError::External(Box::new(Error::Http {
|
|
||||||
source: format!("Failed to parse merge_insert response: {}", e)
|
|
||||||
.into(),
|
|
||||||
request_id: request_id.clone(),
|
|
||||||
status_code: None,
|
|
||||||
}))
|
|
||||||
})?
|
|
||||||
};
|
|
||||||
WriteResult::Merge(merge)
|
|
||||||
}
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut res_lock = result_slot.lock().map_err(|_| {
|
let mut res_lock = add_result.lock().map_err(|_| {
|
||||||
DataFusionError::Execution(
|
DataFusionError::Execution("Failed to acquire lock for add_result".to_string())
|
||||||
"Failed to acquire lock for write result".to_string(),
|
|
||||||
)
|
|
||||||
})?;
|
})?;
|
||||||
*res_lock = Some(parsed_result);
|
*res_lock = Some(parsed_result);
|
||||||
} else {
|
} else {
|
||||||
@@ -785,7 +680,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
|||||||
})?;
|
})?;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Return a single batch with count 0 (actual count is tracked in result)
|
// Return a single batch with count 0 (actual count is tracked in add_result)
|
||||||
let count_array: ArrayRef = Arc::new(UInt64Array::from(vec![0u64]));
|
let count_array: ArrayRef = Arc::new(UInt64Array::from(vec![0u64]));
|
||||||
let batch = RecordBatch::try_new(COUNT_SCHEMA.clone(), vec![count_array])?;
|
let batch = RecordBatch::try_new(COUNT_SCHEMA.clone(), vec![count_array])?;
|
||||||
Ok::<_, DataFusionError>(batch)
|
Ok::<_, DataFusionError>(batch)
|
||||||
@@ -816,11 +711,9 @@ mod tests {
|
|||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
use super::RemoteWriteExec;
|
use super::RemoteInsertExec;
|
||||||
use super::WriteOp;
|
|
||||||
use crate::Table;
|
use crate::Table;
|
||||||
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
||||||
use crate::remote::table::MergeInsertRequest;
|
|
||||||
use crate::table::datafusion::BaseTableAdapter;
|
use crate::table::datafusion::BaseTableAdapter;
|
||||||
|
|
||||||
fn schema_json() -> &'static str {
|
fn schema_json() -> &'static str {
|
||||||
@@ -1135,7 +1028,7 @@ mod tests {
|
|||||||
let input = input_plan_from_batches(schema, batches).await;
|
let input = input_plan_from_batches(schema, batches).await;
|
||||||
|
|
||||||
// A 1-byte budget forces every batch into its own part.
|
// A 1-byte budget forces every batch into its own part.
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
@@ -1175,7 +1068,7 @@ mod tests {
|
|||||||
|
|
||||||
// A large byte budget and no time limit keep the whole partition in a
|
// A large byte budget and no time limit keep the whole partition in a
|
||||||
// single part.
|
// single part.
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
@@ -1216,7 +1109,7 @@ mod tests {
|
|||||||
// A large byte budget but a tiny duration budget: writing and sending
|
// A large byte budget but a tiny duration budget: writing and sending
|
||||||
// one batch already takes longer than the limit, so each batch is cut
|
// one batch already takes longer than the limit, so each batch is cut
|
||||||
// into its own part on the time check rather than the byte check.
|
// into its own part on the time check rather than the byte check.
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
@@ -1251,7 +1144,7 @@ mod tests {
|
|||||||
// write relies on another partition having data to commit.
|
// write relies on another partition having data to commit.
|
||||||
let input = input_plan_from_batches(schema, vec![]).await;
|
let input = input_plan_from_batches(schema, vec![]).await;
|
||||||
|
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
@@ -1291,7 +1184,7 @@ mod tests {
|
|||||||
let input = input_plan_from_batches(schema, batches).await;
|
let input = input_plan_from_batches(schema, batches).await;
|
||||||
|
|
||||||
// A 1-byte budget forces every batch into its own part.
|
// A 1-byte budget forces every batch into its own part.
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
@@ -1340,7 +1233,7 @@ mod tests {
|
|||||||
];
|
];
|
||||||
let input = input_plan_from_partitions(schema, partitions).await;
|
let input = input_plan_from_partitions(schema, partitions).await;
|
||||||
|
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
@@ -1374,7 +1267,7 @@ mod tests {
|
|||||||
// A large byte budget keeps the good batch and the following error in
|
// A large byte budget keeps the good batch and the following error in
|
||||||
// the same part, exercising the mid-part abort path.
|
// the same part, exercising the mid-part abort path.
|
||||||
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
|
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
@@ -1404,66 +1297,6 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_merge_insert_input_error_surfaces_original() {
|
|
||||||
// Regression test for #2339 on the single-request merge_insert path.
|
|
||||||
// When the input stream errors mid-body, Hyper masks it under HTTP2 as a
|
|
||||||
// generic "stream error sent by user" message. The error side-channel
|
|
||||||
// must recover and surface the original DataFusion error instead.
|
|
||||||
use futures::StreamExt;
|
|
||||||
|
|
||||||
let client = crate::remote::client::test_utils::client_with_handler(|request| {
|
|
||||||
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
|
|
||||||
http::Response::builder()
|
|
||||||
.status(200)
|
|
||||||
.body(
|
|
||||||
r#"{"version": 2, "num_updated_rows": 0, "num_inserted_rows": 0, "num_deleted_rows": 0}"#
|
|
||||||
.to_string(),
|
|
||||||
)
|
|
||||||
.unwrap()
|
|
||||||
});
|
|
||||||
|
|
||||||
let query = MergeInsertRequest {
|
|
||||||
on: "id".to_string(),
|
|
||||||
when_matched_update_all: false,
|
|
||||||
when_matched_update_all_filt: None,
|
|
||||||
when_not_matched_insert_all: false,
|
|
||||||
when_not_matched_by_source_delete: false,
|
|
||||||
when_not_matched_by_source_delete_filt: None,
|
|
||||||
use_index: true,
|
|
||||||
use_lsm: None,
|
|
||||||
};
|
|
||||||
|
|
||||||
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
|
|
||||||
let exec = RemoteWriteExec::new(
|
|
||||||
"my_table".to_string(),
|
|
||||||
"my_table".to_string(),
|
|
||||||
client,
|
|
||||||
input,
|
|
||||||
WriteOp::MergeInsert {
|
|
||||||
query,
|
|
||||||
timeout: None,
|
|
||||||
},
|
|
||||||
None,
|
|
||||||
None,
|
|
||||||
);
|
|
||||||
|
|
||||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
|
||||||
let mut err = None;
|
|
||||||
while let Some(item) = stream.next().await {
|
|
||||||
if let Err(e) = item {
|
|
||||||
err = Some(e);
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let err = err.expect("expected the input stream error to surface");
|
|
||||||
assert!(
|
|
||||||
err.to_string().contains("boom"),
|
|
||||||
"expected original input error, got: {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_multipart_records_progress_within_a_part() {
|
async fn test_multipart_records_progress_within_a_part() {
|
||||||
use crate::table::write_progress::{ProgressCallback, WriteProgress, WriteProgressTracker};
|
use crate::table::write_progress::{ProgressCallback, WriteProgress, WriteProgressTracker};
|
||||||
@@ -1494,7 +1327,7 @@ mod tests {
|
|||||||
// A large byte budget keeps all three batches in one part; smooth
|
// A large byte budget keeps all three batches in one part; smooth
|
||||||
// progress therefore requires bytes to be reported per chunk rather than
|
// progress therefore requires bytes to be reported per chunk rather than
|
||||||
// once when the part completes.
|
// once when the part completes.
|
||||||
let exec = RemoteWriteExec::new_multipart(
|
let exec = RemoteInsertExec::new_multipart(
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
"my_table".to_string(),
|
"my_table".to_string(),
|
||||||
client,
|
client,
|
||||||
|
|||||||
+37
-177
@@ -3,7 +3,6 @@
|
|||||||
|
|
||||||
//! LanceDB Table APIs
|
//! LanceDB Table APIs
|
||||||
|
|
||||||
use crate::blob::BlobFile;
|
|
||||||
use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
|
use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
|
||||||
use arrow_schema::{Schema, SchemaRef};
|
use arrow_schema::{Schema, SchemaRef};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
@@ -13,6 +12,7 @@ use datafusion_physical_plan::ExecutionPlan;
|
|||||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||||
use futures::StreamExt;
|
use futures::StreamExt;
|
||||||
use futures::stream::FuturesUnordered;
|
use futures::stream::FuturesUnordered;
|
||||||
|
use lance::dataset::BlobFile;
|
||||||
pub use lance::dataset::ColumnAlteration;
|
pub use lance::dataset::ColumnAlteration;
|
||||||
pub use lance::dataset::NewColumnTransform;
|
pub use lance::dataset::NewColumnTransform;
|
||||||
pub use lance::dataset::ReadParams;
|
pub use lance::dataset::ReadParams;
|
||||||
@@ -50,14 +50,12 @@ use crate::DistanceType;
|
|||||||
use crate::blob::BlobRangeRequest;
|
use crate::blob::BlobRangeRequest;
|
||||||
use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
|
use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
|
||||||
use crate::database::Database;
|
use crate::database::Database;
|
||||||
use crate::database::listing::LANCE_FILE_EXTENSION;
|
|
||||||
use crate::database::read_freshness::TableFreshness;
|
use crate::database::read_freshness::TableFreshness;
|
||||||
use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry};
|
use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry};
|
||||||
use crate::error::{Error, Result};
|
use crate::error::{Error, Result};
|
||||||
use crate::index::IndexStatistics;
|
use crate::index::IndexStatistics;
|
||||||
use crate::index::{Index, IndexBuilder};
|
use crate::index::{Index, IndexBuilder};
|
||||||
use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType};
|
use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType};
|
||||||
use crate::job::Job;
|
|
||||||
use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery};
|
use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery};
|
||||||
use crate::table::datafusion::insert::InsertExec;
|
use crate::table::datafusion::insert::InsertExec;
|
||||||
use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
||||||
@@ -65,7 +63,6 @@ use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
|||||||
use self::dataset::DatasetConsistencyWrapper;
|
use self::dataset::DatasetConsistencyWrapper;
|
||||||
use self::merge::MergeInsertBuilder;
|
use self::merge::MergeInsertBuilder;
|
||||||
|
|
||||||
pub mod add_columns;
|
|
||||||
mod add_data;
|
mod add_data;
|
||||||
pub mod branch_merge;
|
pub mod branch_merge;
|
||||||
mod create_index;
|
mod create_index;
|
||||||
@@ -80,7 +77,6 @@ pub mod schema_evolution;
|
|||||||
pub mod update;
|
pub mod update;
|
||||||
pub mod write_progress;
|
pub mod write_progress;
|
||||||
use crate::index::waiter::wait_for_index;
|
use crate::index::waiter::wait_for_index;
|
||||||
pub use add_columns::AddColumnsBuilder;
|
|
||||||
#[cfg(feature = "remote")]
|
#[cfg(feature = "remote")]
|
||||||
pub(crate) use add_data::PreprocessingOutput;
|
pub(crate) use add_data::PreprocessingOutput;
|
||||||
pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
|
pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior};
|
||||||
@@ -150,55 +146,6 @@ pub(crate) fn map_namespace_lance_error(err: lance::Error, table_name: &str) ->
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Map a `lance::Error::DatasetNotFound` for the table at `uri` into a `lancedb::Error`.
|
|
||||||
///
|
|
||||||
/// Lance reports "there is nothing at this location" and "there is a table directory
|
|
||||||
/// here but nothing loadable inside it" with the same error. Only the first is a
|
|
||||||
/// `TableNotFound`: a `<name>.lance` directory left behind by an interrupted drop and
|
|
||||||
/// re-create is still reported by `Connection::table_names`, so callers need to be able
|
|
||||||
/// to tell "never existed" from "exists but is broken".
|
|
||||||
///
|
|
||||||
/// See <https://github.com/lancedb/lancedb/issues/3127>.
|
|
||||||
async fn map_dataset_not_found(
|
|
||||||
uri: &str,
|
|
||||||
name: &str,
|
|
||||||
params: ReadParams,
|
|
||||||
err: lance::Error,
|
|
||||||
) -> Error {
|
|
||||||
let name = name.to_string();
|
|
||||||
let source = Box::new(err);
|
|
||||||
if table_dir_exists(uri, params).await.unwrap_or(false) {
|
|
||||||
Error::TableCorrupted { name, source }
|
|
||||||
} else {
|
|
||||||
Error::TableNotFound { name, source }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Whether a table directory is present at `uri`, even though no dataset could be
|
|
||||||
/// loaded from it.
|
|
||||||
///
|
|
||||||
/// This looks for a `<name>.lance` entry in the parent directory, which is exactly what
|
|
||||||
/// `ListingDatabase::table_names` lists, so the two APIs agree on whether a table is
|
|
||||||
/// present. Probing `uri` itself would not work: object stores have no empty
|
|
||||||
/// directories to probe, and on a local filesystem the interesting case is precisely an
|
|
||||||
/// empty directory.
|
|
||||||
async fn table_dir_exists(uri: &str, params: ReadParams) -> Result<bool> {
|
|
||||||
let (object_store, path, _) = DatasetBuilder::from_uri(uri)
|
|
||||||
.with_read_params(params)
|
|
||||||
.build_object_store()
|
|
||||||
.await?;
|
|
||||||
// Only `*.lance` entries are ever reported as tables, so nothing else can produce
|
|
||||||
// the list-then-open mismatch this guards against.
|
|
||||||
if path.extension() != Some(LANCE_FILE_EXTENSION) {
|
|
||||||
return Ok(false);
|
|
||||||
}
|
|
||||||
let (Some(parent), Some(dir_name)) = (path.parent(), path.filename()) else {
|
|
||||||
return Ok(false);
|
|
||||||
};
|
|
||||||
let entries = object_store.read_dir(parent).await?;
|
|
||||||
Ok(entries.iter().any(|entry| entry.as_str() == dir_name))
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Defines the type of column
|
/// Defines the type of column
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
pub enum ColumnKind {
|
pub enum ColumnKind {
|
||||||
@@ -616,9 +563,6 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
|||||||
async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult>;
|
async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult>;
|
||||||
/// Create an index on the provided column(s).
|
/// Create an index on the provided column(s).
|
||||||
async fn create_index(&self, index: IndexBuilder) -> Result<()>;
|
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.
|
/// List the indices on the table.
|
||||||
async fn list_indices(&self) -> Result<Vec<IndexConfig>>;
|
async fn list_indices(&self) -> Result<Vec<IndexConfig>>;
|
||||||
/// Drop an index from the table.
|
/// Drop an index from the table.
|
||||||
@@ -1622,8 +1566,12 @@ impl Table {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Add new columns to the table, providing values to fill in.
|
/// Add new columns to the table, providing values to fill in.
|
||||||
pub fn add_columns(&self) -> AddColumnsBuilder {
|
pub async fn add_columns(
|
||||||
AddColumnsBuilder::new(self.inner.clone())
|
&self,
|
||||||
|
transforms: NewColumnTransform,
|
||||||
|
read_columns: Option<Vec<String>>,
|
||||||
|
) -> Result<AddColumnsResult> {
|
||||||
|
self.inner.add_columns(transforms, read_columns).await
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Change a column's name or nullability.
|
/// Change a column's name or nullability.
|
||||||
@@ -2292,8 +2240,6 @@ impl NativeTable {
|
|||||||
None => false,
|
None => false,
|
||||||
};
|
};
|
||||||
|
|
||||||
// Kept so that a `DatasetNotFound` can be re-checked against storage below.
|
|
||||||
let recovery_params = params.clone();
|
|
||||||
let mut builder = DatasetBuilder::from_uri(uri).with_read_params(params);
|
let mut builder = DatasetBuilder::from_uri(uri).with_read_params(params);
|
||||||
|
|
||||||
// Set up commit handler when managed_versioning is enabled
|
// Set up commit handler when managed_versioning is enabled
|
||||||
@@ -2309,13 +2255,13 @@ impl NativeTable {
|
|||||||
builder = builder.with_commit_handler(commit_handler);
|
builder = builder.with_commit_handler(commit_handler);
|
||||||
}
|
}
|
||||||
|
|
||||||
let dataset = match builder.load().await {
|
let dataset = builder.load().await.map_err(|e| match e {
|
||||||
Ok(dataset) => dataset,
|
lance::Error::DatasetNotFound { .. } => Error::TableNotFound {
|
||||||
Err(e @ lance::Error::DatasetNotFound { .. }) => {
|
name: name.to_string(),
|
||||||
return Err(map_dataset_not_found(uri, name, recovery_params, e).await);
|
source: Box::new(e),
|
||||||
}
|
},
|
||||||
Err(e) => return Err(e.into()),
|
e => e.into(),
|
||||||
};
|
})?;
|
||||||
|
|
||||||
let dataset = DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval);
|
let dataset = DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval);
|
||||||
let id = Self::build_id(&namespace, name);
|
let id = Self::build_id(&namespace, name);
|
||||||
@@ -3063,18 +3009,29 @@ impl BaseTable for NativeTable {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn create_index(&self, opts: IndexBuilder) -> Result<()> {
|
async fn create_index(&self, opts: IndexBuilder) -> Result<()> {
|
||||||
let prepared = self.prepare_index(&opts).await?;
|
if opts.columns.len() != 1 {
|
||||||
self.build_index(opts, prepared).await
|
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])?;
|
||||||
|
|
||||||
async fn create_index_async(&self, opts: IndexBuilder) -> Result<Job> {
|
let lance_idx_params = self.make_index_params(&field, opts.index.clone()).await?;
|
||||||
// Prepare before spawning so bad input is reported by this call rather
|
let index_type = self.get_index_type_for_field(&field, &opts.index);
|
||||||
// than only by the job.
|
let columns = [column.as_str()];
|
||||||
let prepared = self.prepare_index(&opts).await?;
|
let mut builder = dataset
|
||||||
let table = self.clone();
|
.create_index_builder(&columns, index_type, lance_idx_params.as_ref())
|
||||||
Ok(Job::spawned(tokio::spawn(async move {
|
.train(opts.train)
|
||||||
table.build_index(opts, prepared).await
|
.replace(opts.replace);
|
||||||
})))
|
|
||||||
|
if let Some(name) = opts.name {
|
||||||
|
builder = builder.name(name);
|
||||||
|
}
|
||||||
|
builder.await?;
|
||||||
|
self.dataset.update(dataset);
|
||||||
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn drop_index(&self, index_name: &str) -> Result<()> {
|
async fn drop_index(&self, index_name: &str) -> Result<()> {
|
||||||
@@ -3627,103 +3584,6 @@ mod tests {
|
|||||||
assert!(matches!(table.unwrap_err(), Error::TableNotFound { .. }));
|
assert!(matches!(table.unwrap_err(), Error::TableNotFound { .. }));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_open_not_found_missing_lance_dir() {
|
|
||||||
let tmp_dir = tempdir().unwrap();
|
|
||||||
let dataset_path = tmp_dir.path().join("test.lance");
|
|
||||||
|
|
||||||
let err = NativeTable::open(dataset_path.to_str().unwrap())
|
|
||||||
.await
|
|
||||||
.unwrap_err();
|
|
||||||
assert!(
|
|
||||||
matches!(&err, Error::TableNotFound { name, .. } if name == "test"),
|
|
||||||
"got {err:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Write a table and then break it, leaving the `<name>.lance` directory in place.
|
|
||||||
///
|
|
||||||
/// `remove_all` reproduces an interrupted drop + re-create (the directory is left
|
|
||||||
/// empty); otherwise only the manifests are removed, leaving the data files behind.
|
|
||||||
async fn write_then_corrupt_table(dir: &std::path::Path, remove_all: bool) -> String {
|
|
||||||
let dataset_path = dir.join("test.lance");
|
|
||||||
let uri = dataset_path.to_str().unwrap().to_string();
|
|
||||||
|
|
||||||
let batch = make_test_batches();
|
|
||||||
let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
|
|
||||||
Dataset::write(reader, &uri, None).await.unwrap();
|
|
||||||
|
|
||||||
if remove_all {
|
|
||||||
for entry in std::fs::read_dir(&dataset_path).unwrap() {
|
|
||||||
let entry = entry.unwrap();
|
|
||||||
if entry.file_type().unwrap().is_dir() {
|
|
||||||
std::fs::remove_dir_all(entry.path()).unwrap();
|
|
||||||
} else {
|
|
||||||
std::fs::remove_file(entry.path()).unwrap();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
assert_eq!(std::fs::read_dir(&dataset_path).unwrap().count(), 0);
|
|
||||||
} else {
|
|
||||||
let versions = dataset_path.join("_versions");
|
|
||||||
assert!(versions.is_dir(), "expected manifests under {versions:?}");
|
|
||||||
std::fs::remove_dir_all(&versions).unwrap();
|
|
||||||
assert!(std::fs::read_dir(&dataset_path).unwrap().count() > 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
uri
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_open_corrupt_empty_dir() {
|
|
||||||
let tmp_dir = tempdir().unwrap();
|
|
||||||
let uri = write_then_corrupt_table(tmp_dir.path(), true).await;
|
|
||||||
|
|
||||||
let err = NativeTable::open(&uri).await.unwrap_err();
|
|
||||||
assert!(
|
|
||||||
matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
|
|
||||||
"got {err:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_open_corrupt_missing_manifest() {
|
|
||||||
let tmp_dir = tempdir().unwrap();
|
|
||||||
let uri = write_then_corrupt_table(tmp_dir.path(), false).await;
|
|
||||||
|
|
||||||
let err = NativeTable::open(&uri).await.unwrap_err();
|
|
||||||
assert!(
|
|
||||||
matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
|
|
||||||
"got {err:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A table listed by `table_names()` must not be reported as missing by
|
|
||||||
/// `open_table()`. See <https://github.com/lancedb/lancedb/issues/3127>.
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_open_table_corrupt_is_still_listed() {
|
|
||||||
let tmp_dir = tempdir().unwrap();
|
|
||||||
let db = connect(tmp_dir.path().to_str().unwrap())
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
write_then_corrupt_table(tmp_dir.path(), true).await;
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
db.table_names().execute().await.unwrap(),
|
|
||||||
vec!["test".to_string()]
|
|
||||||
);
|
|
||||||
let err = db.open_table("test").execute().await.unwrap_err();
|
|
||||||
assert!(
|
|
||||||
matches!(&err, Error::TableCorrupted { name, .. } if name == "test"),
|
|
||||||
"got {err:?}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
err.to_string().contains("exists but could not be loaded"),
|
|
||||||
"got {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
#[cfg(not(windows))]
|
#[cfg(not(windows))]
|
||||||
fn test_object_store_path() {
|
fn test_object_store_path() {
|
||||||
@@ -4731,7 +4591,7 @@ mod tests {
|
|||||||
.set_lsm_write_spec(LsmWriteSpec::bucket("id", bad))
|
.set_lsm_write_spec(LsmWriteSpec::bucket("id", bad))
|
||||||
.await
|
.await
|
||||||
.expect_err("should reject");
|
.expect_err("should reject");
|
||||||
assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
|
assert!(matches!(err, Error::Lance { .. }), "got {:?}", err);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Happy path: install spec; verify MemWAL details record it.
|
// Happy path: install spec; verify MemWAL details record it.
|
||||||
|
|||||||
@@ -1,161 +0,0 @@
|
|||||||
// SPDX-License-Identifier: Apache-2.0
|
|
||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
||||||
|
|
||||||
//! Builder for adding columns to a table.
|
|
||||||
|
|
||||||
use std::sync::Arc;
|
|
||||||
|
|
||||||
use lance::dataset::NewColumnTransform;
|
|
||||||
|
|
||||||
use super::BaseTable;
|
|
||||||
use super::schema_evolution::AddColumnsResult;
|
|
||||||
use crate::{Error, Result};
|
|
||||||
|
|
||||||
/// Adds columns to a table. See [`Table::add_columns`](super::Table::add_columns).
|
|
||||||
pub struct AddColumnsBuilder {
|
|
||||||
parent: Arc<dyn BaseTable>,
|
|
||||||
transform: Option<NewColumnTransform>,
|
|
||||||
read_columns: Option<Vec<String>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::fmt::Debug for AddColumnsBuilder {
|
|
||||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
||||||
f.debug_struct("AddColumnsBuilder")
|
|
||||||
.field("parent", &self.parent)
|
|
||||||
.field("has_transform", &self.transform.is_some())
|
|
||||||
.field("read_columns", &self.read_columns)
|
|
||||||
.finish()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl AddColumnsBuilder {
|
|
||||||
pub(crate) fn new(parent: Arc<dyn BaseTable>) -> Self {
|
|
||||||
Self {
|
|
||||||
parent,
|
|
||||||
transform: None,
|
|
||||||
read_columns: None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Set how the new columns' values are produced. Required.
|
|
||||||
pub fn transform(mut self, transform: NewColumnTransform) -> Self {
|
|
||||||
self.transform = Some(transform);
|
|
||||||
self
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper
|
|
||||||
/// receives. Every other transform determines what it reads, so setting
|
|
||||||
/// this alongside one is an error rather than a silent no-op.
|
|
||||||
pub fn read_columns(mut self, columns: impl IntoIterator<Item = impl Into<String>>) -> Self {
|
|
||||||
self.read_columns = Some(columns.into_iter().map(Into::into).collect());
|
|
||||||
self
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Add the columns.
|
|
||||||
pub async fn execute(self) -> Result<AddColumnsResult> {
|
|
||||||
let Self {
|
|
||||||
parent,
|
|
||||||
transform,
|
|
||||||
read_columns,
|
|
||||||
} = self;
|
|
||||||
|
|
||||||
let Some(transform) = transform else {
|
|
||||||
return Err(Error::InvalidInput {
|
|
||||||
message: "add_columns requires a transform".into(),
|
|
||||||
});
|
|
||||||
};
|
|
||||||
|
|
||||||
if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) {
|
|
||||||
return Err(Error::InvalidInput {
|
|
||||||
message: "read_columns applies only to a BatchUDF transform; \
|
|
||||||
every other transform determines what it reads"
|
|
||||||
.into(),
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
parent.add_columns(transform, read_columns).await
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use std::sync::Arc;
|
|
||||||
|
|
||||||
use arrow_array::{Int32Array, RecordBatch, record_batch};
|
|
||||||
use arrow_schema::{DataType, Field, Schema};
|
|
||||||
use lance::dataset::{BatchUDF, NewColumnTransform};
|
|
||||||
|
|
||||||
use crate::Table;
|
|
||||||
use crate::connect;
|
|
||||||
|
|
||||||
async fn table_with_two_columns(name: &str) -> Table {
|
|
||||||
let conn = connect("memory://").execute().await.unwrap();
|
|
||||||
let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [10, 20, 30])).unwrap();
|
|
||||||
conn.create_table(name, batch).execute().await.unwrap()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_requires_a_transform() {
|
|
||||||
let table = table_with_two_columns("no_transform").await;
|
|
||||||
let err = table.add_columns().execute().await.unwrap_err();
|
|
||||||
assert!(
|
|
||||||
err.to_string().contains("requires a transform"),
|
|
||||||
"got: {err}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_read_columns_with_sql_expressions_is_rejected() {
|
|
||||||
let table = table_with_two_columns("read_cols_sql").await;
|
|
||||||
let err = table
|
|
||||||
.add_columns()
|
|
||||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
|
||||||
"doubled".into(),
|
|
||||||
"x * 2".into(),
|
|
||||||
)]))
|
|
||||||
.read_columns(["x"])
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.unwrap_err();
|
|
||||||
assert!(err.to_string().contains("BatchUDF"), "got: {err}");
|
|
||||||
|
|
||||||
let schema = table.schema().await.unwrap();
|
|
||||||
assert!(
|
|
||||||
schema.field_with_name("doubled").is_err(),
|
|
||||||
"a rejected call must not commit"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_read_columns_limits_what_a_batch_udf_sees() {
|
|
||||||
let table = table_with_two_columns("read_cols_udf").await;
|
|
||||||
|
|
||||||
let output_schema = Arc::new(Schema::new(vec![Field::new("sum", DataType::Int32, true)]));
|
|
||||||
let mapper_schema = output_schema.clone();
|
|
||||||
let udf = BatchUDF {
|
|
||||||
mapper: Box::new(move |batch: &RecordBatch| {
|
|
||||||
assert!(batch.column_by_name("x").is_some());
|
|
||||||
assert!(batch.column_by_name("y").is_none(), "y was not requested");
|
|
||||||
let x = batch["x"].as_any().downcast_ref::<Int32Array>().unwrap();
|
|
||||||
let doubled: Int32Array = x.iter().map(|v| v.map(|v| v * 2)).collect();
|
|
||||||
Ok(RecordBatch::try_new(
|
|
||||||
mapper_schema.clone(),
|
|
||||||
vec![Arc::new(doubled)],
|
|
||||||
)?)
|
|
||||||
}),
|
|
||||||
output_schema,
|
|
||||||
result_checkpoint: None,
|
|
||||||
};
|
|
||||||
|
|
||||||
table
|
|
||||||
.add_columns()
|
|
||||||
.transform(NewColumnTransform::BatchUDF(udf))
|
|
||||||
.read_columns(["x"])
|
|
||||||
.execute()
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let schema = table.schema().await.unwrap();
|
|
||||||
assert!(schema.field_with_name("sum").is_ok());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -576,12 +576,10 @@ mod tests {
|
|||||||
|
|
||||||
// Add a new physical column AFTER the embedding column.
|
// Add a new physical column AFTER the embedding column.
|
||||||
table
|
table
|
||||||
.add_columns()
|
.add_columns(
|
||||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
NewColumnTransform::SqlExpressions(vec![("score".into(), "42.0".into())]),
|
||||||
"score".into(),
|
None,
|
||||||
"42.0".into(),
|
)
|
||||||
)]))
|
|
||||||
.execute()
|
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
@@ -685,9 +683,7 @@ mod tests {
|
|||||||
true,
|
true,
|
||||||
)]));
|
)]));
|
||||||
table
|
table
|
||||||
.add_columns()
|
.add_columns(NewColumnTransform::AllNulls(nested_schema), None)
|
||||||
.transform(NewColumnTransform::AllNulls(nested_schema))
|
|
||||||
.execute()
|
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
|
|||||||
@@ -21,9 +21,6 @@ use lance_index::vector::pq::PQBuildParams;
|
|||||||
use lance_index::vector::sq::builder::SQBuildParams;
|
use lance_index::vector::sq::builder::SQBuildParams;
|
||||||
|
|
||||||
use crate::error::{Error, Result};
|
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::Index;
|
||||||
use crate::index::vector::{VectorIndex, suggested_num_sub_vectors};
|
use crate::index::vector::{VectorIndex, suggested_num_sub_vectors};
|
||||||
use crate::utils::{
|
use crate::utils::{
|
||||||
@@ -108,47 +105,6 @@ 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(
|
pub(super) fn resolve_index_field(
|
||||||
schema: &lance_core::datatypes::Schema,
|
schema: &lance_core::datatypes::Schema,
|
||||||
column: &str,
|
column: &str,
|
||||||
@@ -519,208 +475,6 @@ mod tests {
|
|||||||
assert_eq!(table.list_indices().await.unwrap().len(), 0);
|
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]
|
#[tokio::test]
|
||||||
async fn test_ivf_pq_uses_default_partition_size_for_num_partitions() {
|
async fn test_ivf_pq_uses_default_partition_size_for_num_partitions() {
|
||||||
use crate::index::vector::IvfPqIndexBuilder;
|
use crate::index::vector::IvfPqIndexBuilder;
|
||||||
|
|||||||
@@ -193,12 +193,10 @@ mod tests {
|
|||||||
|
|
||||||
// Add a computed column
|
// Add a computed column
|
||||||
let result = table
|
let result = table
|
||||||
.add_columns()
|
.add_columns(
|
||||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
NewColumnTransform::SqlExpressions(vec![("doubled".into(), "id * 2".into())]),
|
||||||
"doubled".into(),
|
None,
|
||||||
"id * 2".into(),
|
)
|
||||||
)]))
|
|
||||||
.execute()
|
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
@@ -253,12 +251,13 @@ mod tests {
|
|||||||
|
|
||||||
// Add multiple columns at once
|
// Add multiple columns at once
|
||||||
table
|
table
|
||||||
.add_columns()
|
.add_columns(
|
||||||
.transform(NewColumnTransform::SqlExpressions(vec![
|
NewColumnTransform::SqlExpressions(vec![
|
||||||
("y".into(), "x + 1".into()),
|
("y".into(), "x + 1".into()),
|
||||||
("z".into(), "x * x".into()),
|
("z".into(), "x * x".into()),
|
||||||
]))
|
]),
|
||||||
.execute()
|
None,
|
||||||
|
)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
@@ -284,12 +283,10 @@ mod tests {
|
|||||||
|
|
||||||
// Add a column with a constant value
|
// Add a column with a constant value
|
||||||
table
|
table
|
||||||
.add_columns()
|
.add_columns(
|
||||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
NewColumnTransform::SqlExpressions(vec![("constant".into(), "42".into())]),
|
||||||
"constant".into(),
|
None,
|
||||||
"42".into(),
|
)
|
||||||
)]))
|
|
||||||
.execute()
|
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|
||||||
@@ -662,12 +659,10 @@ mod tests {
|
|||||||
|
|
||||||
// Add column increments version
|
// Add column increments version
|
||||||
let add_result = table
|
let add_result = table
|
||||||
.add_columns()
|
.add_columns(
|
||||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + b".into())]),
|
||||||
"c".into(),
|
None,
|
||||||
"a + b".into(),
|
)
|
||||||
)]))
|
|
||||||
.execute()
|
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
assert!(add_result.version > v1);
|
assert!(add_result.version > v1);
|
||||||
|
|||||||
@@ -9,17 +9,14 @@ use arrow_array::{
|
|||||||
};
|
};
|
||||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||||
use futures::TryStreamExt;
|
use futures::TryStreamExt;
|
||||||
use lance::Dataset;
|
|
||||||
use lance_encoding::version::LanceFileVersion;
|
use lance_encoding::version::LanceFileVersion;
|
||||||
use lancedb::{
|
use lancedb::{
|
||||||
Connection, Error, Result, Table,
|
Connection, Error, Result, Table,
|
||||||
blob::{BlobRangeRequest, blob},
|
blob::{BlobRangeRequest, blob},
|
||||||
connect, connect_namespace,
|
connect, connect_namespace,
|
||||||
database::listing::{
|
database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||||
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
|
||||||
},
|
|
||||||
query::{ExecutableQuery, QueryBase},
|
query::{ExecutableQuery, QueryBase},
|
||||||
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats},
|
table::{AddDataMode, CompactionOptions, OptimizeAction},
|
||||||
};
|
};
|
||||||
use tempfile::tempdir;
|
use tempfile::tempdir;
|
||||||
|
|
||||||
@@ -650,7 +647,7 @@ async fn fetch_blob_ranges_validates_requests() -> Result<()> {
|
|||||||
.await
|
.await
|
||||||
.unwrap_err();
|
.unwrap_err();
|
||||||
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
||||||
assert!(err.to_string().contains("row IDs"));
|
assert!(err.to_string().contains("row ids"));
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -683,7 +680,7 @@ async fn fetch_blobs_out_of_range_id_errors_without_panic() -> Result<()> {
|
|||||||
let table = create_inline_blob_table(&db, "t", &[1], &[Some(b"x".as_slice())]).await?;
|
let table = create_inline_blob_table(&db, "t", &[1], &[Some(b"x".as_slice())]).await?;
|
||||||
|
|
||||||
let err = table.fetch_blobs("image", &[u64::MAX]).await.unwrap_err();
|
let err = table.fetch_blobs("image", &[u64::MAX]).await.unwrap_err();
|
||||||
assert!(err.to_string().contains("row IDs"));
|
assert!(err.to_string().contains("row ids"));
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -697,11 +694,11 @@ async fn fetch_blob_apis_reject_mixed_valid_and_missing_row_ids() -> Result<()>
|
|||||||
|
|
||||||
let err = table.fetch_blobs("image", &row_ids).await.unwrap_err();
|
let err = table.fetch_blobs("image", &row_ids).await.unwrap_err();
|
||||||
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
||||||
assert!(err.to_string().contains("row IDs"));
|
assert!(err.to_string().contains("row ids"));
|
||||||
|
|
||||||
let err = table.fetch_blob_files("image", &row_ids).await.unwrap_err();
|
let err = table.fetch_blob_files("image", &row_ids).await.unwrap_err();
|
||||||
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
||||||
assert!(err.to_string().contains("row IDs"));
|
assert!(err.to_string().contains("row ids"));
|
||||||
|
|
||||||
let requests = row_ids.map(|row_id| BlobRangeRequest::new(row_id, 0, 1));
|
let requests = row_ids.map(|row_id| BlobRangeRequest::new(row_id, 0, 1));
|
||||||
let err = table
|
let err = table
|
||||||
@@ -709,7 +706,7 @@ async fn fetch_blob_apis_reject_mixed_valid_and_missing_row_ids() -> Result<()>
|
|||||||
.await
|
.await
|
||||||
.unwrap_err();
|
.unwrap_err();
|
||||||
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
||||||
assert!(err.to_string().contains("row IDs"));
|
assert!(err.to_string().contains("row ids"));
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1078,252 +1075,3 @@ async fn fetch_blob_files_aligns_across_fragments_with_nulls_and_dups() -> Resul
|
|||||||
}
|
}
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Rows exercising the null/empty interleavings from
|
|
||||||
/// <https://github.com/lancedb/lancedb/issues/3744>: a payload, a null, a valid
|
|
||||||
/// empty value, then payloads whose descriptors a fragment rewrite used to zero.
|
|
||||||
fn null_empty_input_batch() -> RecordBatch {
|
|
||||||
let owned = [
|
|
||||||
Some(dedicated_blob_bytes(1)),
|
|
||||||
None,
|
|
||||||
Some(Vec::new()),
|
|
||||||
Some(dedicated_blob_bytes(4)),
|
|
||||||
Some(dedicated_blob_bytes(5)),
|
|
||||||
Some(dedicated_blob_bytes(6)),
|
|
||||||
];
|
|
||||||
let payloads: Vec<Option<&[u8]>> = owned.iter().map(|payload| payload.as_deref()).collect();
|
|
||||||
binary_input_batch(&[1, 2, 3, 4, 5, 6], &payloads)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// One `(id, Some((payload length, first byte)))` per live row, or `(id, None)`
|
|
||||||
/// for a null blob. Comparing lengths and first bytes keeps failure output
|
|
||||||
/// readable where comparing whole payloads would not.
|
|
||||||
type BlobSummary = Vec<(i64, Option<(usize, Option<u8>)>)>;
|
|
||||||
|
|
||||||
/// The rows [`null_empty_input_batch`] leaves behind after `id IN (1, 4)` is
|
|
||||||
/// deleted: a null, a valid empty value, and the two payloads that follow them.
|
|
||||||
fn expected_null_empty_survivors() -> BlobSummary {
|
|
||||||
vec![
|
|
||||||
(2, None),
|
|
||||||
(3, Some((0, None))),
|
|
||||||
(5, Some((DEDICATED_BLOB_LEN, Some(5)))),
|
|
||||||
(6, Some((DEDICATED_BLOB_LEN, Some(6)))),
|
|
||||||
]
|
|
||||||
}
|
|
||||||
|
|
||||||
/// `optimize()` only rewrites a fragment when lance's compaction planner selects
|
|
||||||
/// it — here because the delete pushes the fragment past
|
|
||||||
/// `materialize_deletions_threshold` (0.1 by default; these tests delete 2 of 6
|
|
||||||
/// rows). Without this check, a planner or threshold change upstream would leave
|
|
||||||
/// both regression tests green while no rewrite happened at all.
|
|
||||||
fn assert_compacted(stats: &OptimizeStats) {
|
|
||||||
let metrics = stats
|
|
||||||
.compaction
|
|
||||||
.as_ref()
|
|
||||||
.expect("OptimizeAction::All runs compaction");
|
|
||||||
assert!(
|
|
||||||
metrics.fragments_removed >= 1,
|
|
||||||
"optimize() rewrote no fragment, so this test proves nothing: {metrics:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn summarize(rows: &[(i64, Option<Vec<u8>>)]) -> BlobSummary {
|
|
||||||
rows.iter()
|
|
||||||
.map(|(id, payload)| {
|
|
||||||
(
|
|
||||||
*id,
|
|
||||||
payload
|
|
||||||
.as_ref()
|
|
||||||
.map(|bytes| (bytes.len(), bytes.first().copied())),
|
|
||||||
)
|
|
||||||
})
|
|
||||||
.collect()
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn sorted_id_rowid(table: &Table) -> Result<Vec<(i64, u64)>> {
|
|
||||||
let mut pairs = collect_id_rowid(table).await?;
|
|
||||||
pairs.sort_by_key(|(id, _)| *id);
|
|
||||||
Ok(pairs)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// `{position, size}` descriptors of a legacy v1 blob column, keyed by `id`.
|
|
||||||
async fn v1_blob_descriptors(table: &Table) -> Result<Vec<(i64, Option<(u64, u64)>)>> {
|
|
||||||
let batches = table
|
|
||||||
.query()
|
|
||||||
.execute()
|
|
||||||
.await?
|
|
||||||
.try_collect::<Vec<_>>()
|
|
||||||
.await?;
|
|
||||||
let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap();
|
|
||||||
let ids = batch
|
|
||||||
.column_by_name("id")
|
|
||||||
.unwrap()
|
|
||||||
.as_any()
|
|
||||||
.downcast_ref::<Int64Array>()
|
|
||||||
.unwrap();
|
|
||||||
let descriptors = batch
|
|
||||||
.column_by_name("image")
|
|
||||||
.unwrap()
|
|
||||||
.as_any()
|
|
||||||
.downcast_ref::<StructArray>()
|
|
||||||
.expect("v1 blob column reads back as a descriptor struct");
|
|
||||||
let position = descriptors
|
|
||||||
.column_by_name("position")
|
|
||||||
.unwrap()
|
|
||||||
.as_any()
|
|
||||||
.downcast_ref::<UInt64Array>()
|
|
||||||
.unwrap();
|
|
||||||
let size = descriptors
|
|
||||||
.column_by_name("size")
|
|
||||||
.unwrap()
|
|
||||||
.as_any()
|
|
||||||
.downcast_ref::<UInt64Array>()
|
|
||||||
.unwrap();
|
|
||||||
let mut rows: Vec<(i64, Option<(u64, u64)>)> = (0..batch.num_rows())
|
|
||||||
.map(|row| {
|
|
||||||
let descriptor =
|
|
||||||
(!descriptors.is_null(row)).then(|| (position.value(row), size.value(row)));
|
|
||||||
(ids.value(row), descriptor)
|
|
||||||
})
|
|
||||||
.collect();
|
|
||||||
rows.sort_by_key(|(id, _)| *id);
|
|
||||||
Ok(rows)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Payload bytes of every live row of a legacy v1 blob column, keyed by `id`.
|
|
||||||
/// [`Table::fetch_blobs`] rejects v1 columns, so read them through lance.
|
|
||||||
async fn v1_blob_payloads(dataset_uri: &str, table: &Table) -> Result<Vec<(i64, Option<Vec<u8>>)>> {
|
|
||||||
let pairs = sorted_id_rowid(table).await?;
|
|
||||||
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
|
|
||||||
let dataset = Arc::new(Dataset::open(dataset_uri).await?);
|
|
||||||
let files = dataset.take_blobs(&row_ids, "image").await?;
|
|
||||||
assert_eq!(
|
|
||||||
files.len(),
|
|
||||||
pairs.len(),
|
|
||||||
"take_blobs returned {} handles for {} live rows",
|
|
||||||
files.len(),
|
|
||||||
pairs.len()
|
|
||||||
);
|
|
||||||
let mut rows = Vec::with_capacity(pairs.len());
|
|
||||||
for ((id, _), file) in pairs.iter().zip(files) {
|
|
||||||
let payload = match file {
|
|
||||||
Some(file) => Some(file.read().await?.to_vec()),
|
|
||||||
None => None,
|
|
||||||
};
|
|
||||||
rows.push((*id, payload));
|
|
||||||
}
|
|
||||||
Ok(rows)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Length and first byte of every live blob v2 value, keyed by `id`.
|
|
||||||
async fn blob_v2_values(table: &Table) -> Result<BlobSummary> {
|
|
||||||
let pairs = sorted_id_rowid(table).await?;
|
|
||||||
let row_ids: Vec<u64> = pairs.iter().map(|(_, row_id)| *row_id).collect();
|
|
||||||
let bytes = table.fetch_blobs("image", &row_ids).await?;
|
|
||||||
Ok(pairs
|
|
||||||
.iter()
|
|
||||||
.enumerate()
|
|
||||||
.map(|(slot, (id, _))| {
|
|
||||||
let value = (!bytes.is_null(slot))
|
|
||||||
.then(|| (bytes.value(slot).len(), bytes.value(slot).first().copied()));
|
|
||||||
(*id, value)
|
|
||||||
})
|
|
||||||
.collect())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Regression test for [#3744]: on storage 2.0 (legacy v1 descriptors),
|
|
||||||
/// compaction rewrote every payload following a null or empty value in the same
|
|
||||||
/// fragment as `{position: 0, size: 0}`, so the payload bytes read back as `b""`
|
|
||||||
/// and the new fragment no longer referenced them at all.
|
|
||||||
///
|
|
||||||
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
|
|
||||||
#[tokio::test]
|
|
||||||
async fn optimize_preserves_v1_blob_payloads_with_null_and_empty() -> Result<()> {
|
|
||||||
let tmp = tempdir().unwrap();
|
|
||||||
let db_uri = tmp.path().to_str().unwrap().to_string();
|
|
||||||
let db = connect(&db_uri)
|
|
||||||
.database_options(&ListingDatabaseOptions {
|
|
||||||
new_table_config: NewTableConfig {
|
|
||||||
data_storage_version: Some(LanceFileVersion::V2_0),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
..Default::default()
|
|
||||||
})
|
|
||||||
.execute()
|
|
||||||
.await?;
|
|
||||||
let legacy = Field::new("image", DataType::LargeBinary, true).with_metadata(
|
|
||||||
std::collections::HashMap::from([("lance-encoding:blob".to_string(), "true".to_string())]),
|
|
||||||
);
|
|
||||||
let schema = Arc::new(Schema::new(vec![
|
|
||||||
Field::new("id", DataType::Int64, false),
|
|
||||||
legacy,
|
|
||||||
]));
|
|
||||||
let table = db.create_empty_table("t", schema).execute().await?;
|
|
||||||
table.add(null_empty_input_batch()).execute().await?;
|
|
||||||
assert_eq!(
|
|
||||||
storage_format_version(&table).await,
|
|
||||||
LanceFileVersion::V2_0.resolve(),
|
|
||||||
"v1 blob descriptors only exist below storage 2.2"
|
|
||||||
);
|
|
||||||
let dataset_uri = table.uri().await?;
|
|
||||||
|
|
||||||
// Any rewrite triggers it; deleting rows is the shape from the issue.
|
|
||||||
table.delete("id IN (1, 4)").await?;
|
|
||||||
let descriptors_before = v1_blob_descriptors(&table).await?;
|
|
||||||
let before = v1_blob_payloads(&dataset_uri, &table).await?;
|
|
||||||
assert_eq!(
|
|
||||||
summarize(&before),
|
|
||||||
expected_null_empty_survivors(),
|
|
||||||
"test setup no longer produces the null/empty/payload mix"
|
|
||||||
);
|
|
||||||
|
|
||||||
let stats = table.optimize(OptimizeAction::All).await?;
|
|
||||||
assert_compacted(&stats);
|
|
||||||
|
|
||||||
let descriptors_after = v1_blob_descriptors(&table).await?;
|
|
||||||
let after = v1_blob_payloads(&dataset_uri, &table).await?;
|
|
||||||
assert_eq!(
|
|
||||||
summarize(&after),
|
|
||||||
summarize(&before),
|
|
||||||
"optimize() lost blob payloads; descriptors before={descriptors_before:?} after={descriptors_after:?}"
|
|
||||||
);
|
|
||||||
assert!(after == before, "optimize() changed blob payload bytes");
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Regression test for the blob v2 half of [#3744]: compaction rewrote a valid
|
|
||||||
/// empty value as null, destroying the null-vs-empty distinction.
|
|
||||||
///
|
|
||||||
/// [#3744]: https://github.com/lancedb/lancedb/issues/3744
|
|
||||||
#[tokio::test]
|
|
||||||
async fn optimize_preserves_blob_v2_null_and_empty_distinction() -> Result<()> {
|
|
||||||
let tmp = tempdir().unwrap();
|
|
||||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
|
||||||
let table = db
|
|
||||||
.create_empty_table("t", blob_table_schema())
|
|
||||||
.execute()
|
|
||||||
.await?;
|
|
||||||
table.add(null_empty_input_batch()).execute().await?;
|
|
||||||
assert!(
|
|
||||||
storage_format_version(&table).await >= LanceFileVersion::V2_2,
|
|
||||||
"blob v2 columns require storage >= 2.2"
|
|
||||||
);
|
|
||||||
|
|
||||||
table.delete("id IN (1, 4)").await?;
|
|
||||||
let before = blob_v2_values(&table).await?;
|
|
||||||
assert_eq!(
|
|
||||||
before,
|
|
||||||
expected_null_empty_survivors(),
|
|
||||||
"test setup no longer produces the null/empty/payload mix"
|
|
||||||
);
|
|
||||||
|
|
||||||
let stats = table.optimize(OptimizeAction::All).await?;
|
|
||||||
assert_compacted(&stats);
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
blob_v2_values(&table).await?,
|
|
||||||
before,
|
|
||||||
"optimize() changed blob v2 values"
|
|
||||||
);
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|||||||
Reference in New Issue
Block a user