mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-26 16:08:43 +00:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 011def461c | |||
| ed6be12ad6 | |||
| ac2b689cdb | |||
| 4fc8114871 |
@@ -276,14 +276,38 @@ jobs:
|
||||
# unreadable outside their own branch anyway, since GitHub scopes
|
||||
# caches to the creating ref.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
- name: Downgrade dependencies that exceed our MSRV
|
||||
# Re-resolve the lockfile against `rust-version` instead of hand-pinning
|
||||
# every crate that raises its MSRV. Hand-pinning drifts: the pins keep
|
||||
# ratcheting further back than needed and eventually contradict a real
|
||||
# requirement elsewhere in the graph.
|
||||
env:
|
||||
CARGO_RESOLVER_INCOMPATIBLE_RUST_VERSIONS: fallback
|
||||
run: cargo update
|
||||
- name: Downgrade dependencies
|
||||
# These packages have newer requirements for MSRV
|
||||
run: |
|
||||
cargo update -p aws-sdk-bedrockruntime --precise 1.77.0
|
||||
cargo update -p aws-sdk-dynamodb --precise 1.68.0
|
||||
cargo update -p aws-config --precise 1.6.0
|
||||
cargo update -p aws-sdk-kms --precise 1.63.0
|
||||
cargo update -p aws-sdk-s3 --precise 1.79.0
|
||||
cargo update -p aws-sdk-sso --precise 1.62.0
|
||||
cargo update -p aws-sdk-ssooidc --precise 1.63.0
|
||||
cargo update -p aws-sdk-sts --precise 1.63.0
|
||||
# aws-runtime/sigv4/credential-types/types and the aws-smithy-*
|
||||
# crates bumped their MSRV to 1.91.1 in late 2026; pin to the last
|
||||
# 1.91.0-compatible versions. The order matters — each downgrade
|
||||
# only succeeds once everything that still pins it at a higher
|
||||
# version has itself been downgraded.
|
||||
cargo update -p aws-runtime --precise 1.5.12
|
||||
cargo update -p aws-types --precise 1.3.9
|
||||
cargo update -p aws-sigv4 --precise 1.3.5
|
||||
cargo update -p aws-credential-types --precise 1.2.8
|
||||
cargo update -p aws-smithy-checksums --precise 0.63.9
|
||||
cargo update -p aws-smithy-runtime --precise 1.9.3
|
||||
cargo update -p aws-smithy-http --precise 0.62.4
|
||||
cargo update -p aws-smithy-eventstream --precise 0.60.12
|
||||
cargo update -p aws-smithy-http-client --precise 1.1.3
|
||||
cargo update -p aws-smithy-observability --precise 0.1.4
|
||||
cargo update -p aws-smithy-query --precise 0.60.8
|
||||
cargo update -p aws-smithy-runtime-api --precise 1.9.1
|
||||
cargo update -p aws-smithy-async --precise 1.2.6
|
||||
cargo update -p aws-smithy-types --precise 1.3.5
|
||||
cargo update -p aws-smithy-xml --precise 0.60.11
|
||||
cargo update -p home --precise 0.5.9
|
||||
- name: cargo +${{ matrix.msrv }} check
|
||||
env:
|
||||
RUSTUP_TOOLCHAIN: ${{ matrix.msrv }}
|
||||
|
||||
Generated
+301
-321
File diff suppressed because it is too large
Load Diff
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=11.0.0-beta.1", default-features = false, "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=11.0.0-beta.1", default-features = false, "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=11.0.0-beta.1", default-features = false, "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=11.0.0-beta.1", "tag" = "v11.0.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.0.0-beta.5", "tag" = "v10.0.0-beta.5", "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.0.0-beta.5", "tag" = "v10.0.0-beta.5", "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.0.0-beta.5", "tag" = "v10.0.0-beta.5", "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.0.0-beta.5", "tag" = "v10.0.0-beta.5", "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.0.0-beta.5", "tag" = "v10.0.0-beta.5", "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.0.0-beta.5", "tag" = "v10.0.0-beta.5", "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.0.0-beta.5", "tag" = "v10.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
|
||||
@@ -25,27 +25,6 @@ the underlying connection has been closed.
|
||||
|
||||
## 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()
|
||||
|
||||
```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()
|
||||
|
||||
```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()
|
||||
|
||||
```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()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -25,7 +25,6 @@
|
||||
- [Connection](classes/Connection.md)
|
||||
- [HeaderProvider](classes/HeaderProvider.md)
|
||||
- [Index](classes/Index.md)
|
||||
- [Job](classes/Job.md)
|
||||
- [MakeArrowTableOptions](classes/MakeArrowTableOptions.md)
|
||||
- [MatchQuery](classes/MatchQuery.md)
|
||||
- [MergeInsertBuilder](classes/MergeInsertBuilder.md)
|
||||
@@ -89,9 +88,6 @@
|
||||
- [IvfFlatOptions](interfaces/IvfFlatOptions.md)
|
||||
- [IvfPqOptions](interfaces/IvfPqOptions.md)
|
||||
- [IvfRqOptions](interfaces/IvfRqOptions.md)
|
||||
- [JobDescription](interfaces/JobDescription.md)
|
||||
- [JobFailureInfo](interfaces/JobFailureInfo.md)
|
||||
- [JobInfo](interfaces/JobInfo.md)
|
||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.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>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>11.0.0-beta.1</lance-core.version>
|
||||
<lance-core.version>10.0.0-beta.5</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.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());
|
||||
|
||||
it("should create a vector index on vector columns", async () => {
|
||||
const job = await tbl.createIndexAsync("vec");
|
||||
expect(job.id).toBeNull();
|
||||
await job.wait();
|
||||
// Cancelling a job that already finished succeeds and does nothing.
|
||||
await job.cancel();
|
||||
await tbl.createIndex("vec");
|
||||
|
||||
// check index directory
|
||||
const indexDir = path.join(tmpDir.name, "test.lance", "_indices");
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { tableFromIPC } from "apache-arrow";
|
||||
import {
|
||||
Data,
|
||||
SchemaLike,
|
||||
@@ -21,9 +20,6 @@ import type {
|
||||
CreateNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
Job,
|
||||
JobDescription,
|
||||
JobInfo,
|
||||
ListNamespacesResponse,
|
||||
} from "./native";
|
||||
export type {
|
||||
@@ -440,40 +436,6 @@ export abstract class Connection {
|
||||
newName: string,
|
||||
options?: RenameTableOptions,
|
||||
): 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 */
|
||||
@@ -760,30 +722,6 @@ export class LocalConnection extends Connection {
|
||||
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,
|
||||
} from "./connection";
|
||||
|
||||
export {
|
||||
Job,
|
||||
JobDescription,
|
||||
JobFailureInfo,
|
||||
JobInfo,
|
||||
Session,
|
||||
} from "./native.js";
|
||||
export { Session } from "./native.js";
|
||||
|
||||
export {
|
||||
ExecutableQuery,
|
||||
|
||||
@@ -30,7 +30,6 @@ import {
|
||||
DropColumnsResult,
|
||||
IndexConfig,
|
||||
IndexStatistics,
|
||||
Job,
|
||||
Branches as NativeBranches,
|
||||
OptimizeStats,
|
||||
TableStatistics,
|
||||
@@ -359,17 +358,6 @@ export abstract class Table {
|
||||
options?: Partial<IndexOptions>,
|
||||
): Promise<void>;
|
||||
|
||||
/**
|
||||
* Create an index, returning a handle to the indexing job.
|
||||
*
|
||||
* The job may already be complete when returned; callers must not assume
|
||||
* the index exists until {@link Job.wait} resolves.
|
||||
*/
|
||||
abstract createIndexAsync(
|
||||
column: string,
|
||||
options?: Partial<IndexOptions>,
|
||||
): Promise<Job>;
|
||||
|
||||
/**
|
||||
* Drop an index from the table.
|
||||
*
|
||||
@@ -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> {
|
||||
await this.inner.dropIndex(name);
|
||||
}
|
||||
|
||||
@@ -340,69 +340,6 @@ impl Connection {
|
||||
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)]
|
||||
/// Describe a namespace and return its properties.
|
||||
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 index;
|
||||
mod iterator;
|
||||
mod job;
|
||||
pub mod merge;
|
||||
pub mod otel;
|
||||
pub mod permutation;
|
||||
|
||||
+2
-39
@@ -168,39 +168,6 @@ impl Table {
|
||||
builder.execute().await.default_error()
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn create_index_async(
|
||||
&self,
|
||||
index: Option<&Index>,
|
||||
column: String,
|
||||
replace: Option<bool>,
|
||||
wait_timeout_s: Option<i64>,
|
||||
name: Option<String>,
|
||||
train: Option<bool>,
|
||||
) -> napi::Result<crate::job::Job> {
|
||||
let lancedb_index = if let Some(index) = index {
|
||||
index.consume()?
|
||||
} else {
|
||||
lancedb::index::Index::Auto
|
||||
};
|
||||
let mut builder = self.inner_ref()?.create_index(&[column], lancedb_index);
|
||||
if let Some(replace) = replace {
|
||||
builder = builder.replace(replace);
|
||||
}
|
||||
if let Some(timeout) = wait_timeout_s {
|
||||
builder =
|
||||
builder.wait_timeout(std::time::Duration::from_secs(timeout.try_into().unwrap()));
|
||||
}
|
||||
if let Some(name) = name {
|
||||
builder = builder.name(name);
|
||||
}
|
||||
if let Some(train) = train {
|
||||
builder = builder.train(train);
|
||||
}
|
||||
let job = builder.execute_async().await.default_error()?;
|
||||
Ok(crate::job::Job::new(job))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn drop_index(&self, index_name: String) -> napi::Result<()> {
|
||||
self.inner_ref()?
|
||||
@@ -339,9 +306,7 @@ impl Table {
|
||||
let transforms = NewColumnTransform::SqlExpressions(transforms);
|
||||
let res = self
|
||||
.inner_ref()?
|
||||
.add_columns()
|
||||
.transform(transforms)
|
||||
.execute()
|
||||
.add_columns(transforms, None)
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(res.into())
|
||||
@@ -358,9 +323,7 @@ impl Table {
|
||||
let transforms = NewColumnTransform::AllNulls(schema);
|
||||
let res = self
|
||||
.inner_ref()?
|
||||
.add_columns()
|
||||
.transform(transforms)
|
||||
.execute()
|
||||
.add_columns(transforms, None)
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(res.into())
|
||||
|
||||
@@ -20,7 +20,6 @@ from .remote import ClientConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
from .expr import Expr, col, lit, func
|
||||
from .schema import blob, vector, BlobType
|
||||
from .job import AsyncJob, Job
|
||||
from .table import AsyncTable, Table
|
||||
from .types import BaseTokenizerType
|
||||
from ._lancedb import Session
|
||||
@@ -501,7 +500,6 @@ __all__ = [
|
||||
"connect_namespace",
|
||||
"connect_namespace_async",
|
||||
"AsyncConnection",
|
||||
"AsyncJob",
|
||||
"AsyncLanceNamespaceDBConnection",
|
||||
"AsyncTable",
|
||||
"FtsToken",
|
||||
@@ -515,7 +513,6 @@ __all__ = [
|
||||
"BlobType",
|
||||
"vector",
|
||||
"DBConnection",
|
||||
"Job",
|
||||
"LanceDBConnection",
|
||||
"LanceNamespaceDBConnection",
|
||||
"RemoteDBConnection",
|
||||
|
||||
@@ -14,10 +14,14 @@ import pyarrow as pa
|
||||
from .expr import Expr
|
||||
from .schema import blob_v2_column_paths
|
||||
from .types import BlobMode, QueryProjection, QueryProjectionSpec
|
||||
from .util import get_uri_scheme
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import WriteableBuffer
|
||||
|
||||
from .remote.table import RemoteTable
|
||||
from .table import AsyncTable, Table
|
||||
|
||||
BLOB_MODE_TO_HANDLING = {
|
||||
"lazy": "blobs_descriptions",
|
||||
"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}")
|
||||
|
||||
|
||||
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(
|
||||
projection: QueryProjection,
|
||||
blob_columns: Iterable[str],
|
||||
@@ -144,14 +164,16 @@ def v2_projection_needs_row_id(
|
||||
|
||||
|
||||
def blob_auto_row_id_for_scan(
|
||||
table: Table | AsyncTable | RemoteTable,
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
*,
|
||||
with_row_id: bool | None,
|
||||
) -> bool:
|
||||
"""Auto row-id only applies when the caller said nothing about row ids."""
|
||||
if with_row_id is not None:
|
||||
return False
|
||||
if not supports_blob_auto_row_id(table):
|
||||
return False
|
||||
return v2_projection_needs_row_id(schema, projection, with_row_id=False)
|
||||
|
||||
|
||||
@@ -164,11 +186,6 @@ def finalize_blob_query_table(
|
||||
) -> pa.Table:
|
||||
if user_requested_row_id or not blob_auto_row_id:
|
||||
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)
|
||||
|
||||
|
||||
|
||||
@@ -146,13 +146,6 @@ class Connection(object):
|
||||
start_after: Optional[str],
|
||||
limit: Optional[int],
|
||||
) -> 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(
|
||||
self,
|
||||
name: str,
|
||||
@@ -216,47 +209,6 @@ class BlobFile:
|
||||
def read_range(self, offset: int, length: int) -> bytes: ...
|
||||
def read_up_to(self, length: int) -> bytes: ...
|
||||
|
||||
class Job:
|
||||
@property
|
||||
def id(self) -> Optional[str]: ...
|
||||
async def 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:
|
||||
def name(self) -> str: ...
|
||||
def __repr__(self) -> str: ...
|
||||
@@ -296,28 +248,6 @@ class Table:
|
||||
name: Optional[str],
|
||||
train: Optional[bool],
|
||||
): ...
|
||||
async def create_index_async(
|
||||
self,
|
||||
column: str,
|
||||
index: Union[
|
||||
IvfFlat,
|
||||
IvfSq,
|
||||
IvfPq,
|
||||
HnswPq,
|
||||
HnswSq,
|
||||
HnswFlat,
|
||||
BTree,
|
||||
Bitmap,
|
||||
LabelList,
|
||||
Fm,
|
||||
FTS,
|
||||
],
|
||||
replace: Optional[bool],
|
||||
wait_timeout: Optional[object],
|
||||
*,
|
||||
name: Optional[str],
|
||||
train: Optional[bool],
|
||||
) -> Job: ...
|
||||
async def list_versions(self) -> List[Dict[str, Any]]: ...
|
||||
async def version(self) -> int: ...
|
||||
async def checkout(self, version: Union[int, str]): ...
|
||||
|
||||
+5
-166
@@ -45,7 +45,6 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
|
||||
from . import __version__
|
||||
from ._lancedb import connect as lancedb_connect # type: ignore
|
||||
from .job import AsyncJob, Job
|
||||
from .table import (
|
||||
AsyncTable,
|
||||
LanceTable,
|
||||
@@ -64,7 +63,6 @@ if TYPE_CHECKING:
|
||||
from .pydantic import LanceModel
|
||||
|
||||
from ._lancedb import Connection as LanceDbConnection
|
||||
from ._lancedb import JobDescription, JobInfo
|
||||
from .common import DATA, URI
|
||||
from .embeddings import EmbeddingFunctionConfig
|
||||
from ._lancedb import Session
|
||||
@@ -180,51 +178,6 @@ class DBConnection(EnforceOverrides):
|
||||
"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(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
@@ -610,46 +563,6 @@ class DBConnection(EnforceOverrides):
|
||||
"""
|
||||
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):
|
||||
"""
|
||||
@@ -775,7 +688,11 @@ class LanceDBConnection(DBConnection):
|
||||
return cls(None, _inner=inner)
|
||||
|
||||
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
|
||||
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
|
||||
def namespace_client(self) -> LanceNamespace:
|
||||
"""Get the equivalent namespace client for this connection.
|
||||
@@ -1962,43 +1838,6 @@ class AsyncConnection(object):
|
||||
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:
|
||||
"""Get the equivalent namespace client for this connection.
|
||||
|
||||
|
||||
@@ -23,15 +23,3 @@ class MissingColumnError(KeyError):
|
||||
return (
|
||||
f"Error: Column '{self.column_name}' does not exist in the DataFrame object"
|
||||
)
|
||||
|
||||
|
||||
class JobFailedError(RuntimeError):
|
||||
"""Exception raised when an asynchronous job reaches the failed state."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class JobCancelledError(RuntimeError):
|
||||
"""Exception raised when an asynchronous job was cancelled."""
|
||||
|
||||
pass
|
||||
|
||||
@@ -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
|
||||
if isinstance(condition, Expr):
|
||||
self._when_not_matched_by_source_condition_expr = condition._inner
|
||||
self._when_not_matched_by_source_condition = None
|
||||
else:
|
||||
elif condition is not None:
|
||||
self._when_not_matched_by_source_condition = condition
|
||||
self._when_not_matched_by_source_condition_expr = None
|
||||
return self
|
||||
|
||||
def use_index(self, use_index: bool) -> LanceMergeInsertBuilder:
|
||||
|
||||
@@ -38,11 +38,7 @@ from lance_namespace_urllib3_client.models.query_table_request_vector import (
|
||||
QueryTableRequestVector,
|
||||
)
|
||||
from lance_namespace_urllib3_client.models.string_fts_query import StringFtsQuery
|
||||
from lance_namespace.errors import (
|
||||
NamespaceNotEmptyError,
|
||||
NamespaceNotFoundError,
|
||||
TableNotFoundError,
|
||||
)
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
from lancedb._lancedb import (
|
||||
connect_namespace as _connect_namespace,
|
||||
connect_namespace_client as _connect_namespace_client,
|
||||
@@ -57,8 +53,6 @@ from lance_namespace import (
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
NamespaceExistsRequest,
|
||||
TableExistsRequest,
|
||||
)
|
||||
from lancedb.table import AsyncTable, LanceTable, Table
|
||||
from lancedb.util import validate_table_name
|
||||
@@ -786,51 +780,6 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
"""
|
||||
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
|
||||
def list_tables(
|
||||
self,
|
||||
@@ -1284,49 +1233,6 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""
|
||||
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(
|
||||
self,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
|
||||
@@ -52,6 +52,7 @@ from ._blob import (
|
||||
finalize_blob_query_table,
|
||||
replace_v2_blob_columns_with_bytes,
|
||||
replace_v2_blob_columns_with_bytes_sync,
|
||||
supports_blob_auto_row_id,
|
||||
validate_blob_mode,
|
||||
)
|
||||
from .types import BlobMode, QueryProjection
|
||||
@@ -1279,7 +1280,10 @@ class LanceQueryBuilder(ABC):
|
||||
return self._with_row_id is True
|
||||
|
||||
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(
|
||||
self._table,
|
||||
self._table.schema,
|
||||
self._columns,
|
||||
with_row_id=self._with_row_id,
|
||||
@@ -2697,7 +2701,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._fts_query.phrase_query(True)
|
||||
if 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)
|
||||
if self._maximum_nprobes is not None:
|
||||
self._vector_query.maximum_nprobes(self._maximum_nprobes)
|
||||
@@ -2770,7 +2774,7 @@ class AsyncQueryBase(object):
|
||||
)
|
||||
|
||||
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_paths = ()
|
||||
return
|
||||
@@ -2778,6 +2782,7 @@ class AsyncQueryBase(object):
|
||||
req = self._inner.to_query_request()
|
||||
schema = await self._table.schema()
|
||||
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
req.select,
|
||||
with_row_id=self._with_row_id,
|
||||
@@ -3029,6 +3034,7 @@ class AsyncQueryBase(object):
|
||||
|
||||
schema = await self._table.schema()
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
query.columns,
|
||||
with_row_id=self._with_row_id,
|
||||
@@ -3874,9 +3880,10 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
req = fts_query._inner.to_query_request()
|
||||
blob_auto_row_id = False
|
||||
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()
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
req.select,
|
||||
with_row_id=self._with_row_id,
|
||||
|
||||
@@ -7,7 +7,7 @@ import json
|
||||
import logging
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
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
|
||||
import warnings
|
||||
|
||||
@@ -23,10 +23,6 @@ import pyarrow as pa
|
||||
|
||||
from ..common import DATA
|
||||
from ..db import DBConnection, LOOP
|
||||
from ..job import Job
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .._lancedb import JobDescription, JobInfo
|
||||
from ..embeddings import EmbeddingFunctionConfig
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
@@ -419,11 +415,6 @@ class RemoteDBConnection(DBConnection):
|
||||
|
||||
if namespace_path is None:
|
||||
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:
|
||||
logging.info(
|
||||
"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
|
||||
def namespace_client(self) -> LanceNamespace:
|
||||
"""Get the equivalent namespace client for this connection.
|
||||
|
||||
@@ -20,7 +20,6 @@ from typing import (
|
||||
import warnings
|
||||
|
||||
from lancedb import __version__
|
||||
from lancedb._blob import BlobFile
|
||||
|
||||
from lancedb._lancedb import (
|
||||
AddColumnsResult,
|
||||
@@ -48,7 +47,6 @@ from lancedb.index import (
|
||||
IvfSq,
|
||||
LabelList,
|
||||
)
|
||||
from lancedb.job import Job
|
||||
from lancedb.remote.db import LOOP
|
||||
from lancedb.table import IndexConfigType, KNOWN_METRICS
|
||||
import pyarrow as pa
|
||||
@@ -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(
|
||||
self,
|
||||
first_arg: str,
|
||||
@@ -1069,22 +1039,22 @@ class RemoteTable(Table):
|
||||
)
|
||||
|
||||
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(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> pa.LargeBinaryArray:
|
||||
return LOOP.run(self._table.fetch_blobs(column, row_ids))
|
||||
def fetch_blobs(self, column: str, row_ids) -> pa.LargeBinaryArray:
|
||||
raise NotImplementedError("fetch_blobs() is not supported on LanceDB Cloud")
|
||||
|
||||
def fetch_blob_ranges(self, column: str, requests) -> pa.LargeBinaryArray:
|
||||
raise NotImplementedError(
|
||||
"fetch_blob_ranges() is not supported on LanceDB Cloud"
|
||||
)
|
||||
|
||||
def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
return LOOP.run(self._table.fetch_blob_files(column, row_ids))
|
||||
def fetch_blob_files(self, column: str, row_ids):
|
||||
raise NotImplementedError(
|
||||
"fetch_blob_files() is not supported on LanceDB Cloud"
|
||||
)
|
||||
|
||||
def head(self, n=5) -> pa.Table:
|
||||
"""
|
||||
|
||||
@@ -40,7 +40,6 @@ from ._blob import (
|
||||
from .types import BlobMode
|
||||
from lancedb.arrow import peek_reader
|
||||
from lancedb.background_loop import LOOP, embedding_executor
|
||||
from lancedb.job import AsyncJob, Job
|
||||
from .dependencies import (
|
||||
_check_for_hugging_face,
|
||||
_check_for_lance,
|
||||
@@ -978,24 +977,6 @@ class Table(ABC):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def create_index_async(
|
||||
self,
|
||||
column: str,
|
||||
*,
|
||||
config: IndexConfigType,
|
||||
replace: Optional[bool] = None,
|
||||
wait_timeout: Optional[timedelta] = None,
|
||||
name: Optional[str] = None,
|
||||
train: bool = True,
|
||||
) -> Job:
|
||||
"""Create an index, returning a handle to the indexing job.
|
||||
|
||||
Takes the same arguments as :meth:`create_index`. The job may already
|
||||
be complete when returned; callers must not assume the index exists
|
||||
until :meth:`Job.wait` returns.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def drop_index(self, name: str) -> None:
|
||||
"""
|
||||
Drop an index from the table.
|
||||
@@ -1593,10 +1574,8 @@ class Table(ABC):
|
||||
"""Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles.
|
||||
|
||||
Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is
|
||||
a ``list[int]`` or a query ``pyarrow.Table`` carrying row identity via
|
||||
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
|
||||
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
|
||||
newer.
|
||||
a ``list[int]`` or query ``pyarrow.Table`` with ``_rowid`` (or stashed
|
||||
row-id metadata). Null rows are ``None``. Local tables only.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -2490,7 +2469,13 @@ class LanceTable(Table):
|
||||
return LOOP.run(self._table.count_rows(filter))
|
||||
|
||||
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:
|
||||
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(
|
||||
self,
|
||||
first_arg: str,
|
||||
@@ -4914,46 +4871,6 @@ class AsyncTable:
|
||||
)
|
||||
raise e
|
||||
|
||||
async def create_index_async(
|
||||
self,
|
||||
column: str,
|
||||
*,
|
||||
replace: Optional[bool] = None,
|
||||
config: Optional[
|
||||
Union[
|
||||
IvfFlat,
|
||||
IvfPq,
|
||||
IvfRq,
|
||||
HnswPq,
|
||||
HnswSq,
|
||||
HnswFlat,
|
||||
BTree,
|
||||
Bitmap,
|
||||
LabelList,
|
||||
Fm,
|
||||
FTS,
|
||||
]
|
||||
] = None,
|
||||
wait_timeout: Optional[timedelta] = None,
|
||||
name: Optional[str] = None,
|
||||
train: bool = True,
|
||||
) -> AsyncJob:
|
||||
"""Create an index, returning a handle to the indexing job.
|
||||
|
||||
Takes the same arguments as :meth:`create_index`. The job may already
|
||||
be complete when returned; callers must not assume the index exists
|
||||
until :meth:`AsyncJob.wait` resolves.
|
||||
"""
|
||||
job = await self._inner.create_index_async(
|
||||
column,
|
||||
index=config,
|
||||
replace=replace,
|
||||
wait_timeout=wait_timeout,
|
||||
name=name,
|
||||
train=train,
|
||||
)
|
||||
return AsyncJob(job)
|
||||
|
||||
async def drop_index(self, name: str) -> None:
|
||||
"""
|
||||
Drop an index from the table.
|
||||
|
||||
@@ -226,13 +226,13 @@ def test_fetch_blob_ranges_validates_requests():
|
||||
table = _blob_table("range_validation", [{"id": 1, "image": b"abc"}])
|
||||
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)])
|
||||
|
||||
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)])
|
||||
|
||||
with pytest.raises(ValueError, match="row IDs"):
|
||||
with pytest.raises(ValueError, match="row ids"):
|
||||
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
|
||||
|
||||
|
||||
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):
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
|
||||
@@ -123,19 +123,6 @@ async def test_async_hybrid_query_default_limit(table: AsyncTable):
|
||||
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):
|
||||
reranker = RRFReranker(return_score="all")
|
||||
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
|
||||
async def test_create_scalar_index(some_table: AsyncTable):
|
||||
# Can create
|
||||
|
||||
@@ -18,7 +18,6 @@ Tests verify:
|
||||
"""
|
||||
|
||||
import copy
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
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()}
|
||||
|
||||
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"
|
||||
else:
|
||||
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
|
||||
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")
|
||||
|
||||
|
||||
def test_remote_create_index_async_returns_job():
|
||||
from lancedb.index import BTree
|
||||
|
||||
describe_calls = []
|
||||
|
||||
def handler(request):
|
||||
content_len = int(request.headers.get("Content-Length", 0))
|
||||
body = request.rfile.read(content_len) if content_len > 0 else b""
|
||||
if request.path == "/v1/table/test/create_index/":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"job_id": "job-1"}')
|
||||
elif request.path == "/v1/jobs/describe":
|
||||
assert json.loads(body)["job_id"] == "job-1"
|
||||
describe_calls.append(1)
|
||||
state = "IN_PROGRESS" if len(describe_calls) == 1 else "DONE"
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
json.dumps(dict(job_id="job-1", job_state=state)).encode()
|
||||
)
|
||||
elif request.path == "/v1/jobs/cancel":
|
||||
assert json.loads(body)["job_id"] == "job-1"
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b"{}")
|
||||
elif request.path == "/v1/table/test/create/?mode=create":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b"{}")
|
||||
elif request.path == "/v1/table/test/describe/":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
json.dumps(
|
||||
dict(
|
||||
version=1,
|
||||
schema=dict(
|
||||
fields=[
|
||||
dict(name="id", type={"type": "int64"}, nullable=False),
|
||||
]
|
||||
),
|
||||
)
|
||||
).encode()
|
||||
)
|
||||
else:
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
table = db.create_table("test", [{"id": 1}])
|
||||
job = table.create_index_async("id", config=BTree())
|
||||
assert job.id == "job-1"
|
||||
job.wait(timeout=timedelta(seconds=30))
|
||||
assert len(describe_calls) == 2
|
||||
job.cancel()
|
||||
|
||||
|
||||
def test_remote_job_wait_raises_on_failure():
|
||||
from lancedb.exceptions import JobFailedError
|
||||
from lancedb.index import BTree
|
||||
|
||||
def handler(request):
|
||||
content_len = int(request.headers.get("Content-Length", 0))
|
||||
body = request.rfile.read(content_len) if content_len > 0 else b""
|
||||
if request.path == "/v1/table/test/create_index/":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"job_id": "job-2"}')
|
||||
elif request.path == "/v1/jobs/describe":
|
||||
assert json.loads(body)["job_id"] == "job-2"
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
json.dumps(dict(job_id="job-2", job_state="FAILED")).encode()
|
||||
)
|
||||
elif request.path == "/v1/table/test/create/?mode=create":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b"{}")
|
||||
elif request.path == "/v1/table/test/describe/":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
json.dumps(
|
||||
dict(
|
||||
version=1,
|
||||
schema=dict(
|
||||
fields=[
|
||||
dict(name="id", type={"type": "int64"}, nullable=False),
|
||||
]
|
||||
),
|
||||
)
|
||||
).encode()
|
||||
)
|
||||
else:
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
table = db.create_table("test", [{"id": 1}])
|
||||
job = table.create_index_async("id", config=BTree())
|
||||
with pytest.raises(JobFailedError, match="job-2"):
|
||||
job.wait()
|
||||
|
||||
|
||||
def test_remote_create_index_new_api():
|
||||
received_requests = []
|
||||
|
||||
@@ -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("phalanx-version", str(server_version))
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
|
||||
request.wfile.write(b"{}")
|
||||
elif request.path == "/v1/table/test/query/":
|
||||
content_len = int(request.headers.get("Content-Length"))
|
||||
body = request.rfile.read(content_len)
|
||||
@@ -1973,330 +1858,3 @@ def test_inherited_remote_table_reopens_after_fork():
|
||||
finally:
|
||||
server.shutdown()
|
||||
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
|
||||
|
||||
|
||||
def test_create_index_async_returns_done_job(mem_db: DBConnection):
|
||||
table = mem_db.create_table("job_test", [{"id": i} for i in range(10)])
|
||||
job = table.create_index_async("id", config=BTree())
|
||||
assert job.id is None
|
||||
job.wait()
|
||||
assert len(table.list_indices()) == 1
|
||||
job.cancel()
|
||||
|
||||
|
||||
@patch("lancedb.table.AsyncTable.create_index")
|
||||
def test_create_index_method(mock_create_index, mem_db: DBConnection):
|
||||
table = mem_db.create_table(
|
||||
@@ -2364,29 +2355,6 @@ def test_merge_insert_by_source_delete_expr(mem_db: DBConnection):
|
||||
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
|
||||
async def test_merge_insert_by_source_delete_expr_async(
|
||||
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)
|
||||
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
|
||||
|
||||
table.add([{"id": 1}])
|
||||
|
||||
@@ -13,11 +13,7 @@ use crate::{
|
||||
runtime::future_into_py,
|
||||
table::Table,
|
||||
};
|
||||
use arrow::{
|
||||
datatypes::Schema,
|
||||
ffi_stream::ArrowArrayStreamReader,
|
||||
pyarrow::{FromPyArrow, ToPyArrow},
|
||||
};
|
||||
use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow};
|
||||
use lancedb::{
|
||||
connection::Connection as LanceConnection,
|
||||
connection::NamespaceClientPushdownOperation,
|
||||
@@ -28,7 +24,7 @@ use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
pyclass, pyfunction, pymethods,
|
||||
types::{PyDict, PyDictMethods, PyList, PyListMethods},
|
||||
types::{PyDict, PyDictMethods},
|
||||
};
|
||||
|
||||
#[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]
|
||||
|
||||
@@ -102,18 +102,6 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
|
||||
err.setattr(intern!(py, "__cause__"), cause_err)?;
|
||||
Err(PyErr::from_value(err))
|
||||
}),
|
||||
LanceError::JobFailed { .. } => Python::attach(|py| {
|
||||
let cls = py
|
||||
.import(intern!(py, "lancedb.exceptions"))?
|
||||
.getattr(intern!(py, "JobFailedError"))?;
|
||||
Err(PyErr::from_value(cls.call1((err.to_string(),))?))
|
||||
}),
|
||||
LanceError::JobCancelled { .. } => Python::attach(|py| {
|
||||
let cls = py
|
||||
.import(intern!(py, "lancedb.exceptions"))?
|
||||
.getattr(intern!(py, "JobCancelledError"))?;
|
||||
Err(PyErr::from_value(cls.call1((err.to_string(),))?))
|
||||
}),
|
||||
_ => self.runtime_error(),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -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 header;
|
||||
pub mod index;
|
||||
pub mod job;
|
||||
pub mod namespace;
|
||||
pub mod oauth;
|
||||
pub mod otel;
|
||||
@@ -45,10 +44,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<crate::job::Job>()?;
|
||||
m.add_class::<crate::job::JobInfo>()?;
|
||||
m.add_class::<crate::job::JobDescription>()?;
|
||||
m.add_class::<crate::job::JobFailureInfo>()?;
|
||||
m.add_class::<PyBlobFile>()?;
|
||||
m.add_class::<IndexConfig>()?;
|
||||
m.add_class::<Query>()?;
|
||||
|
||||
+13
-68
@@ -426,11 +426,9 @@ pub struct PyBlobFile {
|
||||
impl PyBlobFile {
|
||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read().await }))
|
||||
let bytes = block_on(async move { inner.read().await })
|
||||
.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>> {
|
||||
@@ -446,32 +444,24 @@ impl PyBlobFile {
|
||||
|
||||
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.close().await }))
|
||||
block_on(async move { inner.close().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
||||
}
|
||||
|
||||
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.is_closed().await }))
|
||||
block_on(async move { inner.is_closed().await })
|
||||
}
|
||||
|
||||
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.seek(position).await }))
|
||||
block_on(async move { inner.seek(position).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
||||
}
|
||||
|
||||
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
||||
let inner = self_.inner.clone();
|
||||
self_
|
||||
.py()
|
||||
.detach(move || block_on(async move { inner.tell().await }))
|
||||
block_on(async move { inner.tell().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
||||
}
|
||||
|
||||
@@ -485,20 +475,16 @@ impl PyBlobFile {
|
||||
.checked_add(length as u64)
|
||||
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_range(offset..end).await }))
|
||||
let bytes = block_on(async move { inner.read_range(offset..end).await })
|
||||
.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>> {
|
||||
let inner = self_.inner.clone();
|
||||
let py = self_.py();
|
||||
let bytes = py
|
||||
.detach(move || block_on(async move { inner.read_up_to(length).await }))
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?;
|
||||
Ok(PyBytes::new(py, bytes.as_ref()).unbind())
|
||||
let bytes = block_on(async move { inner.read_up_to(length).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.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>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -1375,12 +1330,7 @@ impl Table {
|
||||
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner
|
||||
.add_columns()
|
||||
.transform(definitions)
|
||||
.execute()
|
||||
.await
|
||||
.infer_error()?;
|
||||
let result = inner.add_columns(definitions, None).await.infer_error()?;
|
||||
Ok(AddColumnsResult::from(result))
|
||||
})
|
||||
}
|
||||
@@ -1394,12 +1344,7 @@ impl Table {
|
||||
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner
|
||||
.add_columns()
|
||||
.transform(transform)
|
||||
.execute()
|
||||
.await
|
||||
.infer_error()?;
|
||||
let result = inner.add_columns(transform, None).await.infer_error()?;
|
||||
Ok(AddColumnsResult::from(result))
|
||||
})
|
||||
}
|
||||
|
||||
Generated
+1066
-1168
File diff suppressed because it is too large
Load Diff
+1
-1
@@ -1,2 +1,2 @@
|
||||
[toolchain]
|
||||
channel = "1.97.0"
|
||||
channel = "1.95.0"
|
||||
|
||||
@@ -49,8 +49,8 @@ lance-namespace = { workspace = true }
|
||||
lance-namespace-impls = { workspace = true }
|
||||
metrics = { workspace = true, optional = true }
|
||||
metrics-util = { workspace = true, optional = true }
|
||||
# Pin the GooseFS SDK to the version required by Lance's OpenDAL dependency.
|
||||
goosefs-sdk = { version = "=0.1.9", optional = true }
|
||||
# Pin the transitive GooseFS SDK until the 0.1.6 compile break is fixed upstream.
|
||||
goosefs-sdk = { version = "=0.1.5", optional = true }
|
||||
moka = { workspace = true }
|
||||
pin-project = { workspace = true }
|
||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||
|
||||
+3
-199
@@ -9,7 +9,6 @@
|
||||
//!
|
||||
//! Blob tables require Lance file format >= 2.2 and stable row ids at create.
|
||||
|
||||
use std::ops::Range;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::LargeBinaryArray;
|
||||
@@ -17,203 +16,11 @@ use arrow_array::builder::LargeBinaryBuilder;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams};
|
||||
use lance_arrow::FieldExt;
|
||||
use lance_file::version::LanceFileVersion;
|
||||
use lance_io::object_store::ObjectStore;
|
||||
use object_store::path::Path;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Seekable handle for one blob value, backed by local storage or a remote
|
||||
/// 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
pub use lance::dataset::BlobFile;
|
||||
|
||||
/// 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?;
|
||||
ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?;
|
||||
Ok(handles
|
||||
.into_iter()
|
||||
.map(|handle| handle.map(Into::into))
|
||||
.collect())
|
||||
Ok(handles)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -23,8 +23,8 @@ use crate::connection::create_table::CreateTableBuilder;
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::database::listing::ListingDatabase;
|
||||
use crate::database::{
|
||||
CloneTableRequest, Database, DatabaseOptions, JobDescription, JobInfo, OpenTableRequest,
|
||||
ReadConsistency, TableNamesRequest,
|
||||
CloneTableRequest, Database, DatabaseOptions, OpenTableRequest, ReadConsistency,
|
||||
TableNamesRequest,
|
||||
};
|
||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
@@ -34,7 +34,7 @@ use crate::remote::{
|
||||
db::{OPT_REMOTE_API_KEY, OPT_REMOTE_HOST_OVERRIDE, OPT_REMOTE_REGION},
|
||||
};
|
||||
use lance::io::ObjectStoreParams;
|
||||
pub use lance_file::version::LanceFileVersion;
|
||||
pub use lance_encoding::version::LanceFileVersion;
|
||||
#[cfg(feature = "remote")]
|
||||
use lance_io::object_store::StorageOptions;
|
||||
use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider};
|
||||
@@ -456,10 +456,6 @@ impl Connection {
|
||||
///
|
||||
/// # Returns
|
||||
/// 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 {
|
||||
OpenTableBuilder::new(
|
||||
self.internal.clone(),
|
||||
@@ -517,39 +513,6 @@ impl Connection {
|
||||
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.
|
||||
///
|
||||
/// # Arguments
|
||||
|
||||
@@ -18,8 +18,6 @@ use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use arrow_array::RecordBatch;
|
||||
|
||||
use lance::dataset::ReadParams;
|
||||
use lance_namespace::LanceNamespace;
|
||||
use lance_namespace::models::{
|
||||
@@ -202,45 +200,6 @@ pub enum ReadConsistency {
|
||||
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.
|
||||
///
|
||||
/// A database is responsible for managing tables and their metadata.
|
||||
@@ -286,31 +245,6 @@ pub trait Database:
|
||||
///
|
||||
/// See [`CloneTableRequest`] for detailed documentation and examples.
|
||||
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
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
/// Rename a table in the database
|
||||
|
||||
@@ -12,7 +12,7 @@ use lance::dataset::refs::Ref;
|
||||
use lance::dataset::{ReadParams, WriteMode, builder::DatasetBuilder};
|
||||
use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore};
|
||||
use lance_datafusion::utils::StreamingWriteSource;
|
||||
use lance_file::version::LanceFileVersion;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider};
|
||||
use lance_table::io::commit::commit_handler_from_url;
|
||||
use object_store::local::LocalFileSystem;
|
||||
|
||||
@@ -201,7 +201,7 @@ impl LanceNamespaceDatabase {
|
||||
&self,
|
||||
request: &DbCreateTableRequest,
|
||||
) -> Result<(
|
||||
Option<lance_file::version::LanceFileVersion>,
|
||||
Option<lance_encoding::version::LanceFileVersion>,
|
||||
Option<bool>,
|
||||
Option<bool>,
|
||||
)> {
|
||||
@@ -214,7 +214,7 @@ impl LanceNamespaceDatabase {
|
||||
|
||||
let storage_version_override = storage_options
|
||||
.and_then(|opts| opts.get(OPT_NEW_TABLE_STORAGE_VERSION))
|
||||
.map(|s| s.parse::<lance_file::version::LanceFileVersion>())
|
||||
.map(|s| s.parse::<lance_encoding::version::LanceFileVersion>())
|
||||
.transpose()?;
|
||||
|
||||
let v2_manifest_override = storage_options
|
||||
|
||||
@@ -11,9 +11,7 @@ use lance_core::{cache::LanceCache, utils::futures::FinallyStreamExt};
|
||||
use lance_encoding::decoder::{DecoderPlugins, FilterExpression};
|
||||
use lance_file::{
|
||||
reader::{FileReader, FileReaderOptions},
|
||||
version::ConcreteFileVersion,
|
||||
versions,
|
||||
writer::FileWriterOptions,
|
||||
writer::{FileWriter, FileWriterOptions},
|
||||
};
|
||||
use lance_io::{
|
||||
ReadBatchParams,
|
||||
@@ -154,12 +152,8 @@ impl Shuffler {
|
||||
source: None,
|
||||
})?;
|
||||
let object_writer = object_store.create(&path).await?;
|
||||
let writer = versions::create_writer(
|
||||
ConcreteFileVersion::V2_1,
|
||||
object_writer,
|
||||
schema.clone(),
|
||||
FileWriterOptions::default(),
|
||||
)?;
|
||||
let writer =
|
||||
FileWriter::try_new(object_writer, schema.clone(), FileWriterOptions::default())?;
|
||||
file_writers.push(writer);
|
||||
}
|
||||
|
||||
|
||||
@@ -264,7 +264,7 @@ pub fn compute_output_schema(
|
||||
let field_name = ed
|
||||
.dest_column
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("{}_embedding", ed.source_column));
|
||||
.unwrap_or_else(|| format!("{}_embedding", &ed.source_column));
|
||||
|
||||
sb.push(Field::new(
|
||||
field_name,
|
||||
@@ -291,7 +291,7 @@ pub fn compute_embeddings_for_batch(
|
||||
let dst_field_name = fld
|
||||
.dest_column
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("{}_embedding", fld.source_column));
|
||||
.unwrap_or_else(|| format!("{}_embedding", &fld.source_column));
|
||||
|
||||
let dst_field = Field::new(
|
||||
dst_field_name,
|
||||
@@ -315,7 +315,7 @@ impl<R: RecordBatchReader> WithEmbeddings<R> {
|
||||
let field_name = ed
|
||||
.dest_column
|
||||
.clone()
|
||||
.unwrap_or_else(|| format!("{}_embedding", ed.source_column));
|
||||
.unwrap_or_else(|| format!("{}_embedding", &ed.source_column));
|
||||
Ok(Field::new(
|
||||
field_name,
|
||||
func.dest_type()?.into_owned(),
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::fmt::{self, Display, Formatter};
|
||||
use std::sync::{Arc, PoisonError};
|
||||
use std::sync::PoisonError;
|
||||
|
||||
use arrow_schema::ArrowError;
|
||||
use datafusion_common::DataFusionError;
|
||||
@@ -10,46 +9,6 @@ use snafu::Snafu;
|
||||
|
||||
pub(crate) type BoxError = Box<dyn std::error::Error + Send + Sync>;
|
||||
|
||||
/// Why a job failed, to whatever precision the backend provides.
|
||||
///
|
||||
/// A job run in this process carries the error it failed with in [`Self::source`].
|
||||
/// A job run remotely carries whatever the server reported, which older servers
|
||||
/// do not report at all. Every field is absent rather than invented when the
|
||||
/// backend does not supply it.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct JobFailure {
|
||||
/// The stage the job was in, when known.
|
||||
pub phase: Option<String>,
|
||||
/// A human-readable reason, when known.
|
||||
pub message: Option<String>,
|
||||
/// Whether a retry could clear the failure, when known.
|
||||
pub retryable: Option<bool>,
|
||||
/// The error the job failed with, when it ran in this process.
|
||||
pub source: Option<Arc<Error>>,
|
||||
}
|
||||
|
||||
impl JobFailure {
|
||||
/// A failure whose only known detail is the error that caused it.
|
||||
pub(crate) fn from_source(source: Arc<Error>) -> Self {
|
||||
Self {
|
||||
message: Some(source.to_string()),
|
||||
source: Some(source),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for JobFailure {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
|
||||
match (&self.message, &self.phase) {
|
||||
(Some(message), Some(phase)) => write!(f, ": {message} (in {phase})"),
|
||||
(Some(message), None) => write!(f, ": {message}"),
|
||||
(None, Some(phase)) => write!(f, " in {phase}"),
|
||||
(None, None) => Ok(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Snafu)]
|
||||
#[snafu(visibility(pub(crate)))]
|
||||
pub enum Error {
|
||||
@@ -59,10 +18,6 @@ pub enum Error {
|
||||
InvalidInput { message: String },
|
||||
#[snafu(display("Table '{name}' was not found"))]
|
||||
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"))]
|
||||
DatabaseNotFound { name: String },
|
||||
#[snafu(display("Database '{name}' already exists."))]
|
||||
@@ -85,13 +40,6 @@ pub enum Error {
|
||||
Runtime { message: String },
|
||||
#[snafu(display("Timeout error: {message}"))]
|
||||
Timeout { message: String },
|
||||
#[snafu(display("Job{} failed{failure}", job_id.as_ref().map(|id| format!(" {id}")).unwrap_or_default()))]
|
||||
JobFailed {
|
||||
job_id: Option<String>,
|
||||
failure: JobFailure,
|
||||
},
|
||||
#[snafu(display("Job{} was cancelled", job_id.as_ref().map(|id| format!(" {id}")).unwrap_or_default()))]
|
||||
JobCancelled { job_id: Option<String> },
|
||||
|
||||
// 3rd party / external errors
|
||||
#[snafu(display("object_store error: {source}"))]
|
||||
@@ -173,9 +121,6 @@ impl From<lance::Error> for Error {
|
||||
match source {
|
||||
lance::Error::Wrapped { error, .. } => Self::from_box_error(error),
|
||||
lance::Error::External { source } => Self::from_box_error(source),
|
||||
lance::Error::InvalidInput { source, .. } => Self::InvalidInput {
|
||||
message: source.to_string(),
|
||||
},
|
||||
_ => Self::Lance { source },
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@ use std::time::Duration;
|
||||
use vector::IvfFlatIndexBuilder;
|
||||
|
||||
use crate::index::vector::IvfRqIndexBuilder;
|
||||
use crate::{DistanceType, Error, Result, job::Job, table::BaseTable};
|
||||
use crate::{DistanceType, Error, Result, table::BaseTable};
|
||||
|
||||
use self::{
|
||||
scalar::{BTreeIndexBuilder, BitmapIndexBuilder, FmIndexBuilder, LabelListIndexBuilder},
|
||||
@@ -305,14 +305,6 @@ impl IndexBuilder {
|
||||
pub async fn execute(self) -> Result<()> {
|
||||
self.parent.clone().create_index(self).await
|
||||
}
|
||||
|
||||
/// Creates the index, returning a [`Job`] tracking the operation.
|
||||
///
|
||||
/// The job may already be complete when returned, and callers must not
|
||||
/// assume the index exists until [`Job::wait`] resolves.
|
||||
pub async fn execute_async(self) -> Result<Job> {
|
||||
self.parent.clone().create_index_async(self).await
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Deserialize)]
|
||||
|
||||
@@ -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 io;
|
||||
pub mod ipc;
|
||||
pub mod job;
|
||||
#[cfg(feature = "metrics-otel")]
|
||||
pub mod metrics_otel;
|
||||
#[cfg(feature = "polars")]
|
||||
@@ -204,8 +203,7 @@ use serde::{Deserialize, Serialize};
|
||||
|
||||
pub use blob::{BlobRangeRequest, blob, is_blob};
|
||||
pub use connection::{ConnectNamespaceBuilder, Connection};
|
||||
pub use error::{Error, JobFailure, Result};
|
||||
pub use job::Job;
|
||||
pub use error::{Error, Result};
|
||||
use lance_index::vector::ApproxMode as LanceApproxMode;
|
||||
use lance_linalg::distance::DistanceType as LanceDistanceType;
|
||||
/// Re-export of the [`metrics`](https://docs.rs/metrics) crate facade. Enable
|
||||
|
||||
@@ -8,13 +8,13 @@
|
||||
|
||||
pub(crate) mod client;
|
||||
pub(crate) mod db;
|
||||
pub(crate) mod job;
|
||||
pub mod oauth;
|
||||
mod retry;
|
||||
pub(crate) mod table;
|
||||
pub(crate) mod util;
|
||||
|
||||
const ARROW_STREAM_CONTENT_TYPE: &str = "application/vnd.apache.arrow.stream";
|
||||
#[cfg(test)]
|
||||
const ARROW_FILE_CONTENT_TYPE: &str = "application/vnd.apache.arrow.file";
|
||||
#[cfg(test)]
|
||||
const JSON_CONTENT_TYPE: &str = "application/json";
|
||||
|
||||
@@ -706,7 +706,7 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
|
||||
.err_to_http(request_id.clone())?;
|
||||
debug!(
|
||||
"Received response for request_id={}: {:?}",
|
||||
request_id, response
|
||||
request_id, &response
|
||||
);
|
||||
Ok((request_id, response))
|
||||
}
|
||||
@@ -768,7 +768,7 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
|
||||
Ok((status, response)) if status.is_success() => {
|
||||
debug!(
|
||||
"Received response for request_id={}: {:?}",
|
||||
retry_counter.request_id, response
|
||||
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::database::{
|
||||
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
|
||||
JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest,
|
||||
OpenTableRequest, ReadConsistency, TableNamesRequest,
|
||||
};
|
||||
use crate::error::Result;
|
||||
use crate::remote::util::stream_as_body;
|
||||
@@ -79,10 +79,6 @@ impl ServerVersion {
|
||||
pub fn support_multipart_write(&self) -> bool {
|
||||
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_";
|
||||
@@ -432,73 +428,6 @@ fn build_cache_key(name: &str, namespace: &[String]) -> String {
|
||||
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]
|
||||
impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
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>> {
|
||||
let mut req = if !request.namespace_path.is_empty() {
|
||||
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?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let version = parse_server_version(&request_id, &rsp)?;
|
||||
let describe_body = rsp.text().await.ok();
|
||||
let table_identifier = build_table_identifier(
|
||||
&request.name,
|
||||
&request.namespace_path,
|
||||
@@ -847,12 +673,6 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
table_identifier,
|
||||
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);
|
||||
self.table_cache.insert(cache_key, table.clone()).await;
|
||||
Ok(table)
|
||||
@@ -1103,7 +923,6 @@ impl From<StorageOptions> for RemoteOptions {
|
||||
mod tests {
|
||||
use super::{NamespaceHeaderProviderContext, build_cache_key};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use arrow_array::{Int32Array, RecordBatch};
|
||||
@@ -1254,46 +1073,6 @@ mod tests {
|
||||
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]
|
||||
async fn test_open_table_branch_and_version() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
@@ -2263,165 +2042,4 @@ mod tests {
|
||||
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-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! DataFusion ExecutionPlan for streaming writes (add / merge_insert) to
|
||||
//! remote LanceDB tables.
|
||||
//! DataFusion ExecutionPlan for inserting data into remote LanceDB tables.
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
@@ -24,57 +23,28 @@ use lance::io::exec::utils::InstrumentedRecordBatchStreamAdapter;
|
||||
use crate::Error;
|
||||
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
||||
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::write_progress::WriteProgressTracker;
|
||||
use crate::table::{AddResult, MergeResult};
|
||||
|
||||
/// The write operation a [`RemoteWriteExec`] performs. Both variants share the
|
||||
/// same Arrow-IPC streaming body and error side-channel; only the target
|
||||
/// endpoint, query parameters, and parsed result type differ.
|
||||
#[derive(Debug, Clone)]
|
||||
pub 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 enum WriteResult {
|
||||
Add(AddResult),
|
||||
Merge(MergeResult),
|
||||
}
|
||||
|
||||
/// ExecutionPlan for streaming a write (add or merge_insert) to a remote
|
||||
/// LanceDB table.
|
||||
/// ExecutionPlan for inserting data into a remote LanceDB table.
|
||||
///
|
||||
/// Streams data as Arrow IPC to the endpoint selected by [`WriteOp`]. Both
|
||||
/// 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.
|
||||
/// Streams data as Arrow IPC to `/v1/table/{id}/insert/` endpoint.
|
||||
///
|
||||
/// When `upload_id` is set, inserts are staged as part of a multipart write
|
||||
/// session and the plan supports multiple partitions for parallel uploads.
|
||||
/// Without `upload_id`, the plan requires a single partition and commits
|
||||
/// immediately. Multipart applies to `add` only.
|
||||
/// immediately.
|
||||
#[derive(Debug)]
|
||||
pub struct RemoteWriteExec<S: HttpSend = Sender> {
|
||||
pub struct RemoteInsertExec<S: HttpSend = Sender> {
|
||||
table_name: String,
|
||||
identifier: String,
|
||||
client: RestfulLanceDbClient<S>,
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
op: WriteOp,
|
||||
overwrite: bool,
|
||||
properties: Arc<PlanProperties>,
|
||||
result: Arc<Mutex<Option<WriteResult>>>,
|
||||
add_result: Arc<Mutex<Option<AddResult>>>,
|
||||
metrics: ExecutionPlanMetricsSet,
|
||||
upload_id: Option<String>,
|
||||
tracker: Option<Arc<WriteProgressTracker>>,
|
||||
@@ -91,28 +61,27 @@ pub struct RemoteWriteExec<S: HttpSend = Sender> {
|
||||
max_request_duration: Option<Duration>,
|
||||
}
|
||||
|
||||
impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
||||
/// Create a new single-partition RemoteWriteExec.
|
||||
impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
/// Create a new single-partition RemoteInsertExec.
|
||||
pub fn new(
|
||||
table_name: String,
|
||||
identifier: String,
|
||||
client: RestfulLanceDbClient<S>,
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
op: WriteOp,
|
||||
overwrite: bool,
|
||||
tracker: Option<Arc<WriteProgressTracker>>,
|
||||
branch: Option<String>,
|
||||
) -> Self {
|
||||
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
|
||||
/// committing. The caller is responsible for calling the complete (or abort)
|
||||
/// endpoint after all partitions finish. Multipart is insert-only, so the
|
||||
/// op is fixed to [`WriteOp::Insert`].
|
||||
/// endpoint after all partitions finish.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new_multipart(
|
||||
table_name: String,
|
||||
@@ -131,7 +100,7 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
||||
identifier,
|
||||
client,
|
||||
input,
|
||||
WriteOp::Insert { overwrite },
|
||||
overwrite,
|
||||
Some(upload_id),
|
||||
tracker,
|
||||
branch,
|
||||
@@ -146,7 +115,7 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
||||
identifier: String,
|
||||
client: RestfulLanceDbClient<S>,
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
op: WriteOp,
|
||||
overwrite: bool,
|
||||
upload_id: Option<String>,
|
||||
tracker: Option<Arc<WriteProgressTracker>>,
|
||||
branch: Option<String>,
|
||||
@@ -171,9 +140,9 @@ impl<S: HttpSend + 'static> RemoteWriteExec<S> {
|
||||
identifier,
|
||||
client,
|
||||
input,
|
||||
op,
|
||||
overwrite,
|
||||
properties: Arc::new(properties),
|
||||
result: Arc::new(Mutex::new(None)),
|
||||
add_result: Arc::new(Mutex::new(None)),
|
||||
metrics: ExecutionPlanMetricsSet::new(),
|
||||
upload_id,
|
||||
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> {
|
||||
match self
|
||||
.result
|
||||
self.add_result
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.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
|
||||
@@ -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 {
|
||||
match t {
|
||||
DisplayFormatType::Default | DisplayFormatType::Verbose => {
|
||||
write!(f, "RemoteWriteExec: table={}, op=", self.table_name)?;
|
||||
match &self.op {
|
||||
WriteOp::Insert { overwrite } => write!(f, "insert, overwrite={}", overwrite),
|
||||
WriteOp::MergeInsert { .. } => write!(f, "merge_insert"),
|
||||
}
|
||||
write!(
|
||||
f,
|
||||
"RemoteInsertExec: table={}, overwrite={}",
|
||||
self.table_name, self.overwrite
|
||||
)
|
||||
}
|
||||
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 {
|
||||
Self::static_name()
|
||||
}
|
||||
@@ -563,18 +516,15 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
|
||||
if children.len() != 1 {
|
||||
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(
|
||||
self.table_name.clone(),
|
||||
self.identifier.clone(),
|
||||
self.client.clone(),
|
||||
children[0].clone(),
|
||||
self.op.clone(),
|
||||
self.overwrite,
|
||||
self.upload_id.clone(),
|
||||
self.tracker.clone(),
|
||||
self.branch.clone(),
|
||||
@@ -590,19 +540,11 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
) -> DataFusionResult<SendableRecordBatchStream> {
|
||||
if self.upload_id.is_none() && partition != 0 {
|
||||
return Err(DataFusionError::Internal(
|
||||
"RemoteWriteExec only supports single partition execution without upload_id"
|
||||
"RemoteInsertExec only supports single partition execution without upload_id"
|
||||
.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_schema = input_stream.schema();
|
||||
let input_stream: SendableRecordBatchStream =
|
||||
@@ -614,8 +556,8 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
));
|
||||
let client = self.client.clone();
|
||||
let identifier = self.identifier.clone();
|
||||
let op = self.op.clone();
|
||||
let result_slot = self.result.clone();
|
||||
let overwrite = self.overwrite;
|
||||
let add_result = self.add_result.clone();
|
||||
let table_name = self.table_name.clone();
|
||||
let upload_id = self.upload_id.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 {
|
||||
// Multipart writes with a byte budget split the partition into
|
||||
// several bounded, still-streamed requests so no single request
|
||||
// stays open long enough to hit the client read timeout. This path
|
||||
// is insert-only (guarded above).
|
||||
// stays open long enough to hit the client read timeout.
|
||||
if let (Some(upload_id), Some(max_bytes)) =
|
||||
(upload_id.as_deref(), max_bytes_per_request)
|
||||
{
|
||||
let overwrite = matches!(op, WriteOp::Insert { overwrite: true });
|
||||
let ctx = PartRequestCtx {
|
||||
client: &client,
|
||||
identifier: &identifier,
|
||||
@@ -652,36 +592,16 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
)?);
|
||||
}
|
||||
|
||||
// Build the request for the selected operation. Both endpoints take
|
||||
// an Arrow-IPC streaming body and reuse the same error side-channel.
|
||||
let mut request = match &op {
|
||||
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
|
||||
}
|
||||
};
|
||||
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())]);
|
||||
}
|
||||
if let Some(ref b) = branch {
|
||||
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
|
||||
// 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() {
|
||||
return Err(stream_err);
|
||||
}
|
||||
@@ -724,7 +642,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
let (request_id, response) = result?;
|
||||
|
||||
// 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() {
|
||||
let body_text = response.text().await.map_err(|e| {
|
||||
DataFusionError::External(Box::new(Error::Http {
|
||||
@@ -734,44 +652,21 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
}))
|
||||
})?;
|
||||
|
||||
let parsed_result = match &op {
|
||||
WriteOp::Insert { .. } => {
|
||||
let add = if body_text.trim().is_empty() {
|
||||
// Backward compatible with old servers
|
||||
AddResult { version: 0 }
|
||||
} else {
|
||||
serde_json::from_str(&body_text).map_err(|e| {
|
||||
DataFusionError::External(Box::new(Error::Http {
|
||||
source: format!("Failed to parse add response: {}", e).into(),
|
||||
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 parsed_result = if body_text.trim().is_empty() {
|
||||
// Backward compatible with old servers
|
||||
AddResult { version: 0 }
|
||||
} else {
|
||||
serde_json::from_str(&body_text).map_err(|e| {
|
||||
DataFusionError::External(Box::new(Error::Http {
|
||||
source: format!("Failed to parse add response: {}", e).into(),
|
||||
request_id: request_id.clone(),
|
||||
status_code: None,
|
||||
}))
|
||||
})?
|
||||
};
|
||||
|
||||
let mut res_lock = result_slot.lock().map_err(|_| {
|
||||
DataFusionError::Execution(
|
||||
"Failed to acquire lock for write result".to_string(),
|
||||
)
|
||||
let mut res_lock = add_result.lock().map_err(|_| {
|
||||
DataFusionError::Execution("Failed to acquire lock for add_result".to_string())
|
||||
})?;
|
||||
*res_lock = Some(parsed_result);
|
||||
} 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 batch = RecordBatch::try_new(COUNT_SCHEMA.clone(), vec![count_array])?;
|
||||
Ok::<_, DataFusionError>(batch)
|
||||
@@ -816,11 +711,9 @@ mod tests {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use super::RemoteWriteExec;
|
||||
use super::WriteOp;
|
||||
use super::RemoteInsertExec;
|
||||
use crate::Table;
|
||||
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
||||
use crate::remote::table::MergeInsertRequest;
|
||||
use crate::table::datafusion::BaseTableAdapter;
|
||||
|
||||
fn schema_json() -> &'static str {
|
||||
@@ -1135,7 +1028,7 @@ mod tests {
|
||||
let input = input_plan_from_batches(schema, batches).await;
|
||||
|
||||
// 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(),
|
||||
client,
|
||||
@@ -1175,7 +1068,7 @@ mod tests {
|
||||
|
||||
// A large byte budget and no time limit keep the whole partition in a
|
||||
// single part.
|
||||
let exec = RemoteWriteExec::new_multipart(
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
@@ -1216,7 +1109,7 @@ mod tests {
|
||||
// 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
|
||||
// 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(),
|
||||
client,
|
||||
@@ -1251,7 +1144,7 @@ mod tests {
|
||||
// write relies on another partition having data to commit.
|
||||
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(),
|
||||
client,
|
||||
@@ -1291,7 +1184,7 @@ mod tests {
|
||||
let input = input_plan_from_batches(schema, batches).await;
|
||||
|
||||
// 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(),
|
||||
client,
|
||||
@@ -1340,7 +1233,7 @@ mod tests {
|
||||
];
|
||||
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(),
|
||||
client,
|
||||
@@ -1374,7 +1267,7 @@ mod tests {
|
||||
// A large byte budget keeps the good batch and the following error in
|
||||
// the same part, exercising the mid-part abort path.
|
||||
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(),
|
||||
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]
|
||||
async fn test_multipart_records_progress_within_a_part() {
|
||||
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
|
||||
// progress therefore requires bytes to be reported per chunk rather than
|
||||
// once when the part completes.
|
||||
let exec = RemoteWriteExec::new_multipart(
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
|
||||
+37
-177
@@ -3,7 +3,6 @@
|
||||
|
||||
//! LanceDB Table APIs
|
||||
|
||||
use crate::blob::BlobFile;
|
||||
use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader};
|
||||
use arrow_schema::{Schema, SchemaRef};
|
||||
use async_trait::async_trait;
|
||||
@@ -13,6 +12,7 @@ use datafusion_physical_plan::ExecutionPlan;
|
||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||
use futures::StreamExt;
|
||||
use futures::stream::FuturesUnordered;
|
||||
use lance::dataset::BlobFile;
|
||||
pub use lance::dataset::ColumnAlteration;
|
||||
pub use lance::dataset::NewColumnTransform;
|
||||
pub use lance::dataset::ReadParams;
|
||||
@@ -50,14 +50,12 @@ use crate::DistanceType;
|
||||
use crate::blob::BlobRangeRequest;
|
||||
use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
|
||||
use crate::database::Database;
|
||||
use crate::database::listing::LANCE_FILE_EXTENSION;
|
||||
use crate::database::read_freshness::TableFreshness;
|
||||
use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
use crate::index::IndexStatistics;
|
||||
use crate::index::{Index, IndexBuilder};
|
||||
use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType};
|
||||
use crate::job::Job;
|
||||
use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery};
|
||||
use crate::table::datafusion::insert::InsertExec;
|
||||
use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
||||
@@ -65,7 +63,6 @@ use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
||||
use self::dataset::DatasetConsistencyWrapper;
|
||||
use self::merge::MergeInsertBuilder;
|
||||
|
||||
pub mod add_columns;
|
||||
mod add_data;
|
||||
pub mod branch_merge;
|
||||
mod create_index;
|
||||
@@ -80,7 +77,6 @@ pub mod schema_evolution;
|
||||
pub mod update;
|
||||
pub mod write_progress;
|
||||
use crate::index::waiter::wait_for_index;
|
||||
pub use add_columns::AddColumnsBuilder;
|
||||
#[cfg(feature = "remote")]
|
||||
pub(crate) use add_data::PreprocessingOutput;
|
||||
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
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
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>;
|
||||
/// Create an index on the provided column(s).
|
||||
async fn create_index(&self, index: IndexBuilder) -> Result<()>;
|
||||
|
||||
/// Starts index creation, returning a handle to the resulting job.
|
||||
async fn create_index_async(&self, index: IndexBuilder) -> Result<Job>;
|
||||
/// List the indices on the table.
|
||||
async fn list_indices(&self) -> Result<Vec<IndexConfig>>;
|
||||
/// Drop an index from the table.
|
||||
@@ -1622,8 +1566,12 @@ impl Table {
|
||||
}
|
||||
|
||||
/// Add new columns to the table, providing values to fill in.
|
||||
pub fn add_columns(&self) -> AddColumnsBuilder {
|
||||
AddColumnsBuilder::new(self.inner.clone())
|
||||
pub async fn add_columns(
|
||||
&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.
|
||||
@@ -2292,8 +2240,6 @@ impl NativeTable {
|
||||
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);
|
||||
|
||||
// Set up commit handler when managed_versioning is enabled
|
||||
@@ -2309,13 +2255,13 @@ impl NativeTable {
|
||||
builder = builder.with_commit_handler(commit_handler);
|
||||
}
|
||||
|
||||
let dataset = match builder.load().await {
|
||||
Ok(dataset) => dataset,
|
||||
Err(e @ lance::Error::DatasetNotFound { .. }) => {
|
||||
return Err(map_dataset_not_found(uri, name, recovery_params, e).await);
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
let dataset = builder.load().await.map_err(|e| match e {
|
||||
lance::Error::DatasetNotFound { .. } => Error::TableNotFound {
|
||||
name: name.to_string(),
|
||||
source: Box::new(e),
|
||||
},
|
||||
e => e.into(),
|
||||
})?;
|
||||
|
||||
let dataset = DatasetConsistencyWrapper::new_latest(dataset, read_consistency_interval);
|
||||
let id = Self::build_id(&namespace, name);
|
||||
@@ -3063,18 +3009,29 @@ impl BaseTable for NativeTable {
|
||||
}
|
||||
|
||||
async fn create_index(&self, opts: IndexBuilder) -> Result<()> {
|
||||
let prepared = self.prepare_index(&opts).await?;
|
||||
self.build_index(opts, prepared).await
|
||||
}
|
||||
if opts.columns.len() != 1 {
|
||||
return Err(Error::Schema {
|
||||
message: "Multi-column (composite) indices are not yet supported".to_string(),
|
||||
});
|
||||
}
|
||||
self.dataset.ensure_mutable()?;
|
||||
let mut dataset = (*self.dataset.get().await?).clone();
|
||||
let (column, field) = Self::resolve_index_field(dataset.schema(), &opts.columns[0])?;
|
||||
|
||||
async fn create_index_async(&self, opts: IndexBuilder) -> Result<Job> {
|
||||
// Prepare before spawning so bad input is reported by this call rather
|
||||
// than only by the job.
|
||||
let prepared = self.prepare_index(&opts).await?;
|
||||
let table = self.clone();
|
||||
Ok(Job::spawned(tokio::spawn(async move {
|
||||
table.build_index(opts, prepared).await
|
||||
})))
|
||||
let lance_idx_params = self.make_index_params(&field, opts.index.clone()).await?;
|
||||
let index_type = self.get_index_type_for_field(&field, &opts.index);
|
||||
let columns = [column.as_str()];
|
||||
let mut builder = dataset
|
||||
.create_index_builder(&columns, index_type, lance_idx_params.as_ref())
|
||||
.train(opts.train)
|
||||
.replace(opts.replace);
|
||||
|
||||
if let Some(name) = opts.name {
|
||||
builder = builder.name(name);
|
||||
}
|
||||
builder.await?;
|
||||
self.dataset.update(dataset);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_index(&self, index_name: &str) -> Result<()> {
|
||||
@@ -3627,103 +3584,6 @@ mod tests {
|
||||
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]
|
||||
#[cfg(not(windows))]
|
||||
fn test_object_store_path() {
|
||||
@@ -4731,7 +4591,7 @@ mod tests {
|
||||
.set_lsm_write_spec(LsmWriteSpec::bucket("id", bad))
|
||||
.await
|
||||
.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.
|
||||
|
||||
@@ -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.
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"score".into(),
|
||||
"42.0".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("score".into(), "42.0".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -685,9 +683,7 @@ mod tests {
|
||||
true,
|
||||
)]));
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::AllNulls(nested_schema))
|
||||
.execute()
|
||||
.add_columns(NewColumnTransform::AllNulls(nested_schema), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
|
||||
@@ -21,9 +21,6 @@ use lance_index::vector::pq::PQBuildParams;
|
||||
use lance_index::vector::sq::builder::SQBuildParams;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
/// Resolved column, index parameters and index type for one build.
|
||||
pub(super) type PreparedIndex = (String, Box<dyn lance::index::IndexParams>, IndexType);
|
||||
use crate::index::Index;
|
||||
use crate::index::vector::{VectorIndex, suggested_num_sub_vectors};
|
||||
use crate::utils::{
|
||||
@@ -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(
|
||||
schema: &lance_core::datatypes::Schema,
|
||||
column: &str,
|
||||
@@ -519,208 +475,6 @@ mod tests {
|
||||
assert_eq!(table.list_indices().await.unwrap().len(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_execute_async_job_waits_for_local_build() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let conn = connect(tmp_dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = record_batch!(("id", Int32, (0..512).collect::<Vec<_>>())).unwrap();
|
||||
let table = conn.create_table("t", batch).execute().await.unwrap();
|
||||
|
||||
let job = table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap();
|
||||
// Local jobs run in this process and have no server id.
|
||||
assert_eq!(job.id(), None);
|
||||
// The build runs as a task, so the index need not exist yet; it must
|
||||
// once the job resolves.
|
||||
job.wait().await.unwrap();
|
||||
assert_eq!(table.list_indices().await.unwrap().len(), 1);
|
||||
// Cancelling a finished job is a no-op.
|
||||
job.cancel().await.unwrap();
|
||||
}
|
||||
|
||||
/// Concurrent waiters, and a wait issued after the job settled, all
|
||||
/// succeed once the build does.
|
||||
#[tokio::test]
|
||||
async fn test_execute_async_job_reports_success_to_every_waiter() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let conn = connect(tmp_dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = record_batch!(("id", Int32, (0..512).collect::<Vec<_>>())).unwrap();
|
||||
let table = conn.create_table("t", batch).execute().await.unwrap();
|
||||
|
||||
let job = Arc::new(
|
||||
table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let waiters = (0..4)
|
||||
.map(|_| {
|
||||
let job = job.clone();
|
||||
tokio::spawn(async move { job.wait().await })
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for waiter in waiters {
|
||||
waiter.await.unwrap().unwrap();
|
||||
}
|
||||
// A wait after the job settled still reports the same outcome.
|
||||
job.wait().await.unwrap();
|
||||
assert_eq!(table.list_indices().await.unwrap().len(), 1);
|
||||
}
|
||||
|
||||
/// Every waiter sees a failure, not just the first: a waiter that missed
|
||||
/// the outcome would be told the job succeeded.
|
||||
#[tokio::test]
|
||||
async fn test_execute_async_job_reports_failure_to_every_waiter() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let conn = connect(tmp_dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = record_batch!(("id", Int32, (0..512).collect::<Vec<_>>())).unwrap();
|
||||
let table = conn.create_table("t", batch).execute().await.unwrap();
|
||||
table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Rebuilding the same index without replace fails once the build
|
||||
// starts, so the failure reaches the job rather than execute_async.
|
||||
let job = Arc::new(
|
||||
table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.replace(false)
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let waiters = (0..3)
|
||||
.map(|_| {
|
||||
let job = job.clone();
|
||||
tokio::spawn(async move { job.wait().await })
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for waiter in waiters {
|
||||
waiter
|
||||
.await
|
||||
.unwrap()
|
||||
.expect_err("every waiter must see the failure");
|
||||
}
|
||||
job.wait().await.expect_err("a later wait still fails");
|
||||
}
|
||||
|
||||
/// A local failure keeps the error it failed with, so a caller can match on
|
||||
/// the original variant rather than parse a message.
|
||||
#[tokio::test]
|
||||
async fn test_execute_async_failure_keeps_the_source_error() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let conn = connect(tmp_dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = record_batch!(("id", Int32, (0..512).collect::<Vec<_>>())).unwrap();
|
||||
let table = conn.create_table("t", batch).execute().await.unwrap();
|
||||
table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let job = table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.replace(false)
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let crate::Error::JobFailed { failure, .. } = job.wait().await.unwrap_err() else {
|
||||
panic!("a failed job reports JobFailed");
|
||||
};
|
||||
let source = failure.source.expect("a local failure carries its error");
|
||||
assert_eq!(
|
||||
failure.message.as_deref(),
|
||||
Some(source.to_string()).as_deref()
|
||||
);
|
||||
// Nothing local can report these, so they must be absent rather than invented.
|
||||
assert!(failure.phase.is_none());
|
||||
assert!(failure.retryable.is_none());
|
||||
}
|
||||
|
||||
/// Every waiter sees the cancellation, including ones that were already
|
||||
/// waiting when the cancel landed.
|
||||
#[tokio::test]
|
||||
async fn test_execute_async_job_reports_cancellation_to_every_waiter() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let conn = connect(tmp_dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = record_batch!(("id", Int32, (0..512).collect::<Vec<_>>())).unwrap();
|
||||
let table = conn.create_table("t", batch).execute().await.unwrap();
|
||||
|
||||
let job = Arc::new(
|
||||
table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
// Cancel before yielding, so the build cannot have started and the
|
||||
// outcome is always the cancellation.
|
||||
job.cancel().await.unwrap();
|
||||
|
||||
let waiters = (0..2)
|
||||
.map(|_| {
|
||||
let job = job.clone();
|
||||
tokio::spawn(async move { job.wait().await })
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for waiter in waiters {
|
||||
match waiter.await.unwrap() {
|
||||
Err(crate::Error::JobCancelled { .. }) => {}
|
||||
other => panic!("expected the cancellation, got {other:?}"),
|
||||
}
|
||||
}
|
||||
match job.wait().await {
|
||||
Err(crate::Error::JobCancelled { .. }) => {}
|
||||
other => panic!("expected the cancellation, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_execute_async_job_cancel_stops_local_build() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let conn = connect(tmp_dir.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let batch = record_batch!(("id", Int32, (0..512).collect::<Vec<_>>())).unwrap();
|
||||
let table = conn.create_table("t", batch).execute().await.unwrap();
|
||||
|
||||
let job = table
|
||||
.create_index(&["id"], Index::BTree(BTreeIndexBuilder::default()))
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap();
|
||||
job.cancel().await.unwrap();
|
||||
match job.wait().await {
|
||||
Err(crate::Error::JobCancelled { .. }) => {}
|
||||
// The build may finish before the abort lands.
|
||||
Ok(()) => {}
|
||||
other => panic!("unexpected job outcome: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_ivf_pq_uses_default_partition_size_for_num_partitions() {
|
||||
use crate::index::vector::IvfPqIndexBuilder;
|
||||
|
||||
@@ -193,12 +193,10 @@ mod tests {
|
||||
|
||||
// Add a computed column
|
||||
let result = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"doubled".into(),
|
||||
"id * 2".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("doubled".into(), "id * 2".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -253,12 +251,13 @@ mod tests {
|
||||
|
||||
// Add multiple columns at once
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![
|
||||
("y".into(), "x + 1".into()),
|
||||
("z".into(), "x * x".into()),
|
||||
]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![
|
||||
("y".into(), "x + 1".into()),
|
||||
("z".into(), "x * x".into()),
|
||||
]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -284,12 +283,10 @@ mod tests {
|
||||
|
||||
// Add a column with a constant value
|
||||
table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"constant".into(),
|
||||
"42".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("constant".into(), "42".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -662,12 +659,10 @@ mod tests {
|
||||
|
||||
// Add column increments version
|
||||
let add_result = table
|
||||
.add_columns()
|
||||
.transform(NewColumnTransform::SqlExpressions(vec![(
|
||||
"c".into(),
|
||||
"a + b".into(),
|
||||
)]))
|
||||
.execute()
|
||||
.add_columns(
|
||||
NewColumnTransform::SqlExpressions(vec![("c".into(), "a + b".into())]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(add_result.version > v1);
|
||||
|
||||
@@ -9,17 +9,14 @@ use arrow_array::{
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
use futures::TryStreamExt;
|
||||
use lance::Dataset;
|
||||
use lance_file::version::LanceFileVersion;
|
||||
use lance_encoding::version::LanceFileVersion;
|
||||
use lancedb::{
|
||||
Connection, Error, Result, Table,
|
||||
blob::{BlobRangeRequest, blob},
|
||||
connect, connect_namespace,
|
||||
database::listing::{
|
||||
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
},
|
||||
database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
query::{ExecutableQuery, QueryBase},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction, OptimizeStats},
|
||||
table::{AddDataMode, CompactionOptions, OptimizeAction},
|
||||
};
|
||||
use tempfile::tempdir;
|
||||
|
||||
@@ -650,7 +647,7 @@ async fn fetch_blob_ranges_validates_requests() -> Result<()> {
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
||||
assert!(err.to_string().contains("row IDs"));
|
||||
assert!(err.to_string().contains("row ids"));
|
||||
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 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(())
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
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();
|
||||
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 err = table
|
||||
@@ -709,7 +706,7 @@ async fn fetch_blob_apis_reject_mixed_valid_and_missing_row_ids() -> Result<()>
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(&err, Error::InvalidInput { .. }), "got {err:?}");
|
||||
assert!(err.to_string().contains("row IDs"));
|
||||
assert!(err.to_string().contains("row ids"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1078,252 +1075,3 @@ async fn fetch_blob_files_aligns_across_fragments_with_nulls_and_dups() -> Resul
|
||||
}
|
||||
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