Compare commits

..
Author SHA1 Message Date
lancedb automation 456dae3e62 chore: update lance dependency to v11.0.0-rc.1 2026-08-24 17:55:03 +00:00
Xuanwo b0dae5eb0b feat: return typed refresh job results (#4013)
## Problem

`refresh_column_async` returned a unit-result job even though durable
refresh jobs carry a canonical terminal result. Python callers could not
obtain row counts or source and published versions through the public
`Job` API, and local and remote refresh jobs exposed different result
semantics.

## Behavior

`refresh_column_async` now returns `Job[RefreshColumnResult]` for local
and remote tables. The general typed-job bridge binds each endpoint to
its public result model while preserving unit-result jobs and existing
status, wait, cancel, and timeout behavior. A local no-op refresh
reports no published version.

The Node.js API continues to resolve `wait()` as `void`; its binding
erases the Rust result type internally to preserve the existing public
contract.

## Ownership and integration boundary

LanceDB owns the language-neutral `Job<T>` contract and language-binding
decode. Sophon owns production and durable persistence of terminal
payloads. Sophon #7348 and #7378 now publish the canonical refresh
result for Function-backed and expression-backed refresh jobs,
respectively. The remote client fixture matches the merged server
schema; live deployment and end-to-end demo acceptance remain separate
rollout checks.
2026-08-25 00:01:32 +08:00
Ayush ChaurasiaandOpenAI Codex 242ade8017 feat(python): support sequence packing in streaming dataset (#3920)
## How packing works

  Consider four tokenized documents:

  
  [1]
  [2]
  [10, 11, 12, 13, 14, 15, 16, 17]
  [20]

  With:
```
  StreamingDataset(
      table,
      shuffle=False,
      columns=["tokens"],
      num_splits=2,
      pack_sequences=5,
      eos_id=9,
      pad_id=0,
      blocks_per_epoch=6,
  )
```
the documents are assigned to two fixed logical splits. Each split
maintains an independent token buffer, appends eos_id after every
document, and emits blocks of five tokens.

  Because blocks_per_epoch=6, each split emits exactly three blocks:

  Cycl/e 1:
    Split 0: [1, 9, 2, 9, 0] # 9 is eos, 0 is padding
    Split 1: [10, 11, 12, 13, 14]

  Cycle 2:
    Split 0: [0, 0, 0, 0, 0]
    Split 1: [15, 16, 17, 9, 20]

  Cycle 3:
    Split 0: [0, 0, 0, 0, 0]
    Split 1: [9, 0, 0, 0, 0]

If a split runs out of tokens early, it emits padded blocks through the
fixed budget. This prevents one rank from finishing before another.

Logical splits are independent of rank and worker ownership. A
checkpoint records each split’s consumed-document count, emitted-block
count, remaining tokens, and document boundaries. Merging
those per-split states allows the same packed stream to resume after the
topology changes.

doc_ids identifies document segments, including continuations across
block boundaries. It is not a padding mask: padding retains the
preceding document ID, so callers must mask padding using a
  reserved pad_id.

blocks_per_epoch="auto" is also available. It estimates the budget from
a deterministic bounded sample and warns that the result is approximate.

WIP pre-training tests:
```
  ┌────────────────────────────────────┬────────────────────────┬─────────────────────────────────────┐
  │                                    │       GPT-2 124M       │          GPT-2 medium 354M          │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Corpus                             │ 2.4M docs / 12GB table │ 9.67M docs / 45GB table             │
  │ Tokens (Chinchilla)                │ 2.43B                  │ 7.0B                                │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Data prep (ingest→curate→tokenize) │ ~12 min                │ ~51 min                             │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Training wall time                 │ ~50 min                │ 3h 06m                              │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Throughput / MFU                   │ 1.60M tok/s / 35%      │ 684k tok/s / 42.0%, │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Final val loss                     │ 3.230                  │ 2.840                               │
  └────────────────────────────────────┴────────────────────────┴─────────────────────────────────────┘
```

---------

Co-authored-by: OpenAI Codex <codex@openai.com>
2026-08-24 15:56:36 +08:00
Lance Release 40d4d012e7 Bump version: 0.38.0-beta.5 → 0.38.0-beta.6 2026-08-23 17:33:19 +00:00
29 changed files with 1129 additions and 310 deletions
Generated
+45 -45
View File
@@ -3455,8 +3455,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "fsst"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"rand 0.9.5",
@@ -4815,8 +4815,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
[[package]]
name = "lance"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arc-swap",
"arrow",
@@ -4888,8 +4888,8 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4911,7 +4911,7 @@ dependencies = [
[[package]]
name = "lance-arrow-scalar"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4925,7 +4925,7 @@ dependencies = [
[[package]]
name = "lance-arrow-stats"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -4934,8 +4934,8 @@ dependencies = [
[[package]]
name = "lance-bitpacking"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrayref",
"crunchy",
@@ -4945,8 +4945,8 @@ dependencies = [
[[package]]
name = "lance-core"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4983,8 +4983,8 @@ dependencies = [
[[package]]
name = "lance-datafusion"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow",
"arrow-array",
@@ -5013,8 +5013,8 @@ dependencies = [
[[package]]
name = "lance-datagen"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow",
"arrow-array",
@@ -5031,8 +5031,8 @@ dependencies = [
[[package]]
name = "lance-derive"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"proc-macro2",
"quote",
@@ -5041,8 +5041,8 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5075,8 +5075,8 @@ dependencies = [
[[package]]
name = "lance-file"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5107,8 +5107,8 @@ dependencies = [
[[package]]
name = "lance-index"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arc-swap",
"arrow",
@@ -5172,8 +5172,8 @@ dependencies = [
[[package]]
name = "lance-index-core"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5195,8 +5195,8 @@ dependencies = [
[[package]]
name = "lance-io"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow",
"arrow-array",
@@ -5236,8 +5236,8 @@ dependencies = [
[[package]]
name = "lance-linalg"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5251,8 +5251,8 @@ dependencies = [
[[package]]
name = "lance-namespace"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow",
"async-trait",
@@ -5264,8 +5264,8 @@ dependencies = [
[[package]]
name = "lance-namespace-impls"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow",
"arrow-ipc",
@@ -5318,8 +5318,8 @@ dependencies = [
[[package]]
name = "lance-select"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5333,8 +5333,8 @@ dependencies = [
[[package]]
name = "lance-table"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow",
"arrow-array",
@@ -5374,8 +5374,8 @@ dependencies = [
[[package]]
name = "lance-testing"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5388,8 +5388,8 @@ dependencies = [
[[package]]
name = "lance-tokenizer"
version = "11.0.0-beta.22"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.22#ea3cb4d799c468232735e9bcb43959487aca5c20"
version = "11.0.0-rc.1"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
dependencies = [
"frostem",
"icu_segmenter",
@@ -5402,7 +5402,7 @@ dependencies = [
[[package]]
name = "lancedb"
version = "0.38.0-beta.5"
version = "0.38.0-beta.6"
dependencies = [
"ahash",
"anyhow",
@@ -5490,7 +5490,7 @@ dependencies = [
[[package]]
name = "lancedb-nodejs"
version = "0.38.0-beta.5"
version = "0.38.0-beta.6"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5515,7 +5515,7 @@ dependencies = [
[[package]]
name = "lancedb-python"
version = "0.38.0-beta.5"
version = "0.38.0-beta.6"
dependencies = [
"arrow",
"async-trait",
+14 -14
View File
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
rust-version = "1.91.0"
[workspace.dependencies]
lance = { "version" = "=11.0.0-beta.22", default-features = false, "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=11.0.0-beta.22", default-features = false, "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=11.0.0-beta.22", default-features = false, "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=11.0.0-beta.22", "tag" = "v11.0.0-beta.22", "git" = "https://github.com/lance-format/lance.git" }
lance = { "version" = "=11.0.0-rc.1", default-features = false, "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=11.0.0-rc.1", default-features = false, "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=11.0.0-rc.1", default-features = false, "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
lancedb = { path = "rust/lancedb", default-features = false }
ahash = "0.8"
# Note that this one does not include pyarrow
+1 -1
View File
@@ -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.22</lance-core.version>
<lance-core.version>11.0.0-rc.1</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>
+2 -2
View File
@@ -1,12 +1,12 @@
{
"name": "@lancedb/lancedb",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.6",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "@lancedb/lancedb",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.6",
"cpu": [
"x64",
"arm64"
+5 -2
View File
@@ -14,9 +14,12 @@ pub struct Job {
}
impl Job {
pub(crate) fn new(inner: lancedb::Job) -> Self {
pub(crate) fn new<T>(inner: lancedb::Job<T>) -> Self
where
T: Clone + Send + Sync + 'static,
{
Self {
inner: Arc::new(inner),
inner: Arc::new(inner.map(|_| ())),
}
}
}
+1
View File
@@ -29,6 +29,7 @@ from .functions import (
FunctionRegistrationRequest as FunctionRegistrationRequest,
FunctionVersion as FunctionVersion,
PythonRuntimeSpec as PythonRuntimeSpec,
RefreshColumnResult as RefreshColumnResult,
UdfDefinition as UdfDefinition,
udf as udf,
)
+2 -9
View File
@@ -147,7 +147,7 @@ class Connection(object):
limit: Optional[int],
) -> list[str]: ... # Deprecated: Use list_tables instead
def job(self, job_id: str) -> Job: ...
async def create_function_async(self, request_json: str) -> FunctionJob: ...
async def create_function_async(self, request_json: str) -> Job: ...
async def get_function(self, name: str, version: str) -> str: ...
async def list_jobs(self) -> List[JobInfo]: ...
async def get_job(self, job_id: str) -> Optional[JobDescription]: ...
@@ -234,14 +234,7 @@ class Job:
@property
def id(self) -> Optional[str]: ...
async def status(self) -> str: ...
async def wait(self) -> None: ...
async def cancel(self) -> None: ...
class FunctionJob:
@property
def id(self) -> Optional[str]: ...
async def status(self) -> str: ...
async def wait(self) -> str: ...
async def wait(self) -> Optional[str]: ...
async def cancel(self) -> None: ...
class JobInfo:
+2 -2
View File
@@ -46,7 +46,7 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
from . import __version__
from ._lancedb import connect as lancedb_connect # type: ignore
from .functions import FunctionVersion, UdfDefinition
from .job import AsyncJob, Job, _function_job
from .job import AsyncJob, Job, _typed_job
from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
@@ -2237,7 +2237,7 @@ class AsyncConnection(object):
inner = await self._inner.create_function_async(
definition.registration_request.to_canonical_json()
)
return _function_job(inner)
return _typed_job(inner, FunctionVersion.from_json)
async def get_function(self, name: str, *, version: str) -> FunctionVersion:
"""Open one exact immutable Function version from the remote catalog."""
+8 -2
View File
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Canonical values exchanged with LanceDB Enterprise Function services.
"""Canonical Function values exchanged with LanceDB Enterprise services.
These immutable models contain client/wire state only. Catalog persistence,
environment bake, secret resolution, and execution are owned by Sophon.
``RefreshColumnResult`` is also the backend-neutral result of a local
expression-backed refresh job.
"""
from __future__ import annotations
@@ -460,7 +462,11 @@ class FunctionBinding(_RemoteValue):
class RefreshColumnResult(_RemoteValue):
"""Terminal result of a remote Function-column refresh Job."""
"""Terminal result of an expression-backed or Function-backed refresh Job.
Local jobs produce this value in process. LanceDB Cloud and Enterprise
decode the same value from the durable server-job terminal payload.
"""
rows_assigned: _UInt64
rows_failed: _UInt64
+28 -30
View File
@@ -5,12 +5,11 @@
import asyncio
from datetime import timedelta
from typing import Any, Generic, Optional, TypeVar, cast
from typing import Any, Callable, Generic, Optional, TypeVar, cast
from lancedb.background_loop import LOOP
from . import _lancedb
from .functions import FunctionVersion
T = TypeVar("T")
@@ -18,11 +17,18 @@ T = TypeVar("T")
class AsyncJob(Generic[T]):
"""A handle to an operation that may still be running.
The operation may already be complete when the handle is created.
The operation may already be complete when the handle is created. ``T``
is the endpoint's terminal result type; unit-result jobs resolve to
``None``.
"""
def __init__(self, inner: Optional[Any]):
def __init__(
self,
inner: Optional[Any],
result_decoder: Optional[Callable[[Any], T]] = None,
):
self._inner = inner
self._result_decoder = result_decoder
@property
def id(self) -> Optional[str]:
@@ -50,17 +56,21 @@ class AsyncJob(Generic[T]):
async def wait(self, timeout: Optional[timedelta] = None) -> T:
"""Wait until the operation reaches a terminal state.
Returns the endpoint's typed result, or ``None`` for a unit-result
job.
Raises `JobFailedError` if the operation failed, `JobCancelledError`
if it was cancelled, and `TimeoutError` if `timeout` elapses first.
"""
if self._inner is None:
return cast(T, None)
if timeout is None:
return cast(T, await self._inner.wait())
return cast(
T,
await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()),
)
result = await self._inner.wait()
else:
result = await asyncio.wait_for(self._inner.wait(), timeout.total_seconds())
if self._result_decoder is not None:
return self._result_decoder(result)
return cast(T, result)
async def cancel(self):
"""Request cancellation. Cancelling a finished operation is a no-op."""
@@ -70,7 +80,7 @@ class AsyncJob(Generic[T]):
class Job(Generic[T]):
"""Synchronous counterpart of `AsyncJob`."""
"""Synchronous counterpart of `AsyncJob` with the same result type."""
def __init__(self, inner: Optional[AsyncJob[T]]):
self._inner = inner
@@ -96,6 +106,9 @@ class Job(Generic[T]):
def wait(self, timeout: Optional[timedelta] = None) -> T:
"""Block until the operation reaches a terminal state.
Returns the endpoint's typed result, or ``None`` for a unit-result
job.
Raises `JobFailedError` if the operation failed, `JobCancelledError`
if it was cancelled, and `TimeoutError` if `timeout` elapses first.
"""
@@ -110,23 +123,8 @@ class Job(Generic[T]):
LOOP.run(self._inner.cancel())
class _FunctionJobAdapter:
def __init__(self, inner: "_lancedb.FunctionJob"):
self._inner = inner
@property
def id(self) -> Optional[str]:
return self._inner.id
async def status(self) -> str:
return await self._inner.status()
async def wait(self) -> FunctionVersion:
return FunctionVersion.from_json(await self._inner.wait())
async def cancel(self):
await self._inner.cancel()
def _function_job(inner: "_lancedb.FunctionJob") -> AsyncJob[FunctionVersion]:
return AsyncJob(_FunctionJobAdapter(inner))
def _typed_job(
inner: "_lancedb.Job", result_decoder: Callable[[str], T]
) -> AsyncJob[T]:
"""Bind an internal JSON-producing job to its public result model."""
return AsyncJob(inner, result_decoder)
+2 -2
View File
@@ -49,7 +49,7 @@ from lancedb.index import (
LabelList,
)
from lancedb.job import Job
from lancedb.functions import FunctionApplication
from lancedb.functions import FunctionApplication, RefreshColumnResult
from lancedb.remote.db import LOOP
from lancedb.table import IndexConfigType, KNOWN_METRICS
import pyarrow as pa
@@ -972,7 +972,7 @@ class RemoteTable(Table):
def refresh_column(self, column: str):
return LOOP.run(self._table.refresh_column(column))
def refresh_column_async(self, column: str) -> Job:
def refresh_column_async(self, column: str) -> Job[RefreshColumnResult]:
return Job(LOOP.run(self._table.refresh_column_async(column)))
def alter_columns(
+430 -61
View File
@@ -24,11 +24,16 @@ import os
import random
import threading
import time
import warnings
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from multiprocessing import RawArray
from typing import Any, Callable, Iterator, Optional, Union
from typing import Any, Callable, cast, Iterator, Literal, Optional, Union
import pyarrow as pa
import pyarrow.compute as pc
import torch
from torch.utils.data import IterableDataset, get_worker_info
from .permutation import (
@@ -132,6 +137,39 @@ class StreamingDataset(IterableDataset):
Maximum number of transforms to run concurrently. Must be greater
than zero. When ``None`` (the default), uses ``os.cpu_count()`` or 1
when the CPU count is unavailable.
pack_sequences:
Sequence-packing mode: token lists from consecutive documents are
joined with ``eos_id`` and sliced into blocks of this many tokens.
Each item is then a dict of two ``(pack_sequences,)`` LongTensors —
``input_ids`` and ``doc_ids`` (per-position document index within
the block, for block-diagonal masks or position-id resets).
* Packing happens independently per owned split and preserves per-split
resume state.
* When a split cannot fill a real block for a cycle but an owned sibling
still can, or when only a short tail remains at epoch end, the short buffer
is padded to ``pack_sequences`` with ``pad_id`` so every local cycle emits
one block per owned split.
* ``eos_id``, ``pad_id``, and ``columns`` naming a single integer-list
column are required; incompatible with ``transform``.
eos_id:
Separator token id between packed documents. Required with
``pack_sequences``, ignored otherwise.
pad_id:
Padding token id used to complete blocks when a split runs out of
real tokens mid-cycle or at epoch end. Required with
``pack_sequences``, ignored otherwise. It must be reserved for padding:
padding positions retain the preceding document's ``doc_id`` (or zero
in an all-padding block), so callers must mask them separately using
``input_ids == pad_id``.
blocks_per_epoch:
Total number of packed blocks emitted globally per epoch. Required with
``pack_sequences``. An integer must be divisible by ``num_splits``.
Every logical split emits exactly ``blocks_per_epoch / num_splits``
blocks: exhausted splits emit padding, while tokens beyond the budget
are left out of the epoch. This fixed per-split budget keeps packed
iteration and checkpoints independent of rank topology.
Pass ``"auto"`` to estimate a corpus-level budget from a bounded sample
of token lists. The estimate may be inaccurate.
on_transform_error:
What to do when the transform raises an exception:
@@ -210,6 +248,10 @@ class StreamingDataset(IterableDataset):
filter: Optional[str] = None,
transform: Optional[Callable] = None,
transform_parallelism: Optional[int] = None,
pack_sequences: Optional[int] = None,
eos_id: Optional[int] = None,
pad_id: Optional[int] = None,
blocks_per_epoch: Optional[Union[int, Literal["auto"]]] = None,
on_transform_error: Union[str, Callable[[Exception], bool]] = "raise",
transform_queue_depth: Optional[int] = None,
connection_factory: Optional[Callable[[str], Any]] = None,
@@ -237,6 +279,55 @@ class StreamingDataset(IterableDataset):
raise ValueError("io_queue_depth must be greater than 0")
if transform_parallelism is not None and transform_parallelism <= 0:
raise ValueError("transform_parallelism must be greater than 0")
if pack_sequences is not None:
if pack_sequences <= 0:
raise ValueError("pack_sequences must be greater than 0")
if eos_id is None:
raise ValueError("eos_id is required when pack_sequences is set")
if pad_id is None:
raise ValueError("pad_id is required when pack_sequences is set")
if blocks_per_epoch is None:
raise ValueError(
"blocks_per_epoch is required when pack_sequences is set"
)
if blocks_per_epoch != "auto":
if not isinstance(blocks_per_epoch, int) or isinstance(
blocks_per_epoch, bool
):
raise ValueError(
"blocks_per_epoch must be a positive integer or 'auto'"
)
if blocks_per_epoch <= 0:
raise ValueError("blocks_per_epoch must be greater than 0")
if blocks_per_epoch % num_splits != 0:
raise ValueError(
f"blocks_per_epoch ({blocks_per_epoch}) must be divisible by "
f"num_splits ({num_splits})"
)
if transform is not None:
raise ValueError("transform cannot be combined with pack_sequences")
if columns is None or len(columns) != 1:
raise ValueError(
"pack_sequences requires columns to name exactly one "
"list-typed column of token ids"
)
field = table.schema.field(columns[0])
if not (
pa.types.is_list(field.type)
or pa.types.is_large_list(field.type)
or pa.types.is_fixed_size_list(field.type)
):
raise ValueError(
f"pack_sequences requires a list-typed token column; "
f"{columns[0]} has type {field.type}"
)
if not pa.types.is_integer(field.type.value_type):
raise ValueError(
"pack_sequences requires a token column with integer values; "
f"{columns[0]} has value type {field.type.value_type}"
)
elif blocks_per_epoch is not None:
raise ValueError("blocks_per_epoch requires pack_sequences")
if on_transform_error not in ("raise", "skip", "warn") and not callable(
on_transform_error
):
@@ -261,11 +352,20 @@ class StreamingDataset(IterableDataset):
self._filter = filter
self._transform = transform
self._transform_parallelism = transform_parallelism
self._pack_sequences = pack_sequences
self._eos_id = eos_id
self._pad_id = pad_id
self._blocks_per_epoch = blocks_per_epoch
self._on_transform_error = on_transform_error
self._transform_queue_depth = transform_queue_depth
self._connection_factory = connection_factory
self._worker_info_override = worker_info_override
# Packing resume state: permutation positions and partial-block buffers.
self._pack_consumed: list[int] = [0] * num_splits
self._pack_buffers: dict[int, dict[str, list[int]]] = {}
self._pack_blocks_emitted: list[int] = [0] * num_splits
# Live references to pipeline state, set only while __iter__ is running
# in the same process. Used by the observability properties when the
# DataLoader runs with num_workers=0.
@@ -315,6 +415,9 @@ class StreamingDataset(IterableDataset):
else:
self._perm_table = builder.split_sequential(fixed=num_splits).execute()
if self._blocks_per_epoch == "auto":
self._blocks_per_epoch = self._estimate_blocks_per_epoch()
# Contiguous block of global split indices assigned to this rank.
splits_per_rank = num_splits // world_size
rank_start = rank * splits_per_rank
@@ -322,6 +425,71 @@ class StreamingDataset(IterableDataset):
range(rank_start, rank_start + splits_per_rank)
)
def _estimate_blocks_per_epoch(self) -> int:
"""Estimate a fixed packed-block budget from a bounded token sample."""
# TODO: Replace this fallback with Lance's dedicated exact token-count
# estimation API once it is available.
if self._pack_sequences is None or not self._columns:
raise RuntimeError(
"packing must be configured before estimating its budget"
)
pack_len = self._pack_sequences
token_column = self._columns[0]
sample_cap_per_split = max(1, 100_000 // self._num_splits)
sampled_tokens = 0
total_sampled = 0
total_rows = 0
rng = random.Random(self._shuffle_seed)
warnings.warn(
"blocks_per_epoch='auto' uses an approximate token-count sample; "
"pass an explicit value for exact epoch sizing",
)
for split in range(self._num_splits):
permutation = Permutation.from_tables(
self._table, self._perm_table, split=split
)
permutation = permutation.select_columns([token_column])
permutation = permutation.with_transform(Transforms.arrow2arrow)
split_rows = permutation.num_rows
if split_rows == 0:
raise ValueError(
"blocks_per_epoch='auto' cannot estimate an empty dataset"
)
# Sample roughly 1% from each logical split, with at least one row
# per split and a global target cap of 100,000 rows.
sample_rows = min(
split_rows,
max(1, min((split_rows + 99) // 100, sample_cap_per_split)),
)
sample_offsets = sorted(rng.sample(range(split_rows), sample_rows))
sample_batch_size = max(1, self._read_batch_size)
for start in range(0, sample_rows, sample_batch_size):
batch = permutation.__getitems__(
sample_offsets[start : start + sample_batch_size]
)
lengths = pc.list_value_length(batch.column(0))
if lengths.null_count:
raise ValueError("pack_sequences does not support null token lists")
sampled_tokens += int(pc.sum(lengths).as_py())
total_sampled += sample_rows
total_rows += split_rows
# Pool the samples into one global average. Each document contributes
# one EOS token.
estimated_tokens = (
(sampled_tokens + total_sampled) * total_rows // total_sampled
)
blocks = estimated_tokens // pack_len
return max(
self._num_splits,
blocks - blocks % self._num_splits,
)
def _resolve_my_splits(self) -> list[int]:
"""Return the split indices this instance should read in __iter__."""
torch_worker_info = get_worker_info()
@@ -372,8 +540,14 @@ class StreamingDataset(IterableDataset):
)
if self._columns is not None:
perm = perm.select_columns(self._columns)
perm = perm.with_transform(lambda batch: batch)
start_pos = self._resume_positions.get(split_idx, self._resume_offset)
perm = perm.with_transform(Transforms.arrow2arrow)
# Both modes resume from absolute permutation positions. Packing
# stores them separately because it also checkpoints partial blocks.
start_pos = (
self._pack_consumed[split_idx]
if self._pack_sequences is not None
else self._resume_positions.get(split_idx, self._resume_offset)
)
if start_pos > 0:
perm = perm.with_skip(start_pos)
initial_positions.append(start_pos)
@@ -395,9 +569,24 @@ class StreamingDataset(IterableDataset):
if self._transform_parallelism is not None
else (os.cpu_count() or 1)
)
final_transform = (
self._transform if self._transform is not None else Transforms.arrow2python
)
final_transform: Callable[[pa.RecordBatch], Any]
if self._pack_sequences is not None:
# Packing consumes raw token lists, one per document.
def arrow_tokens(batch: pa.RecordBatch) -> list[list[int]]:
token_column = batch.column(0)
if token_column.null_count or token_column.flatten().null_count:
raise ValueError(
"pack_sequences does not support null token lists or values"
)
return cast(list[list[int]], token_column.to_pylist())
final_transform = arrow_tokens
else:
final_transform = (
self._transform
if self._transform is not None
else Transforms.arrow2python
)
# None means no limit; otherwise cap rows per split to
# transform_queue_depth batches worth (including in-flight transforms).
max_cooked_rows = (
@@ -574,6 +763,82 @@ class StreamingDataset(IterableDataset):
else:
break # split exhausted
def _update_stats(*, idle: bool = False) -> None:
"""Refresh pipeline statistics visible to the parent process."""
ws = self._worker_stats
ws[0] = sum(split_sizes[j] - fetch_head[j] for j in range(n))
ws[1] = (
0
if idle
else sum(batch.num_rows for q in raw_batches for _, batch in q)
)
ws[2] = 0 if idle else sum(len(q) for q in cooked)
ws[3] = sum(local_consumed)
ws[4] = self._bytes_loaded
ws[5] = int(self._fetch_time * 1_000_000)
ws[6] = int(self._transform_time * 1_000_000)
ws[7] = self._rows_skipped
# Sequence-packing helpers
pack_len = cast(int, self._pack_sequences)
eos_id = cast(int, self._eos_id)
pad_id = cast(int, self._pad_id)
blocks_per_split = (
cast(int, self._blocks_per_epoch) // self._num_splits
if self._pack_sequences is not None
else 0
)
pack_consumed = list(self._pack_consumed)
pack_buffers = deepcopy(self._pack_buffers)
pack_blocks_emitted = list(self._pack_blocks_emitted)
def _pack_buffer(i: int) -> dict[str, list[int]]:
return pack_buffers.setdefault(my_splits[i], {"tokens": [], "starts": []})
def _fill_block(i: int) -> None:
"""Fill split i's buffer to one block or exhaust the split."""
buf = _pack_buffer(i)
while len(buf["tokens"]) < pack_len:
_ensure_cooked(i)
if not cooked[i]:
return
buf["starts"].append(len(buf["tokens"]))
pos, tokens = cooked[i].popleft()
buf["tokens"].extend(tokens)
buf["tokens"].append(eos_id)
pack_consumed[my_splits[i]] = pos + 1
local_consumed[i] += 1
_advance(i)
def _emit_block(i: int) -> dict[str, Any]:
buf = _pack_buffer(i)
tokens, starts = buf["tokens"], buf["starts"]
# doc_ids label document segments within the block; 0 also covers
# the continuation of a document begun in a prior block.
doc_ids = torch.zeros(pack_len, dtype=torch.int64)
doc_starts = [s for s in starts if 0 < s < pack_len]
doc_ids[doc_starts] = 1
doc_ids.cumsum_(dim=0) # cumulative sum marks document boundaries
block = {
"input_ids": torch.tensor(tokens[:pack_len], dtype=torch.int64),
"doc_ids": doc_ids,
}
del tokens[:pack_len]
# Shift start boundaries for the next call.
buf["starts"] = [s - pack_len for s in starts if s >= pack_len]
return block
def _commit_pack_state() -> None:
self._pack_consumed = list(pack_consumed)
self._pack_buffers = {
split: {
"tokens": list(buffer["tokens"]),
"starts": list(buffer["starts"]),
}
for split, buffer in pack_buffers.items()
}
self._pack_blocks_emitted = list(pack_blocks_emitted)
# ── Main loop ─────────────────────────────────────────────────────────
with ThreadPoolExecutor(max_workers=n * io_queue_depth) as io_pool:
@@ -583,10 +848,42 @@ class StreamingDataset(IterableDataset):
self._fetch_head_ref = fetch_head
self._split_sizes_ref = split_sizes
self._local_consumed_ref = local_consumed
try:
for i in range(n):
_fill_io(i)
if self._pack_sequences is not None:
first_count = pack_blocks_emitted[my_splits[0]]
if any(
pack_blocks_emitted[split] != first_count
for split in my_splits[1:]
):
raise ValueError(
"Packed checkpoint is not aligned across the splits "
"owned by this iterator; merge every rank "
"state with merge_state_dicts before resuming on a "
"different topology"
)
while pack_blocks_emitted[my_splits[0]] < blocks_per_split:
# Each logical split gets one block per cycle. Exhausted
# splits are padded through the fixed global budget.
for i in range(n):
_fill_block(i)
for i in range(n):
tokens = _pack_buffer(i)["tokens"]
if len(tokens) < pack_len:
tokens.extend([pad_id] * (pack_len - len(tokens)))
block = _emit_block(i)
pack_blocks_emitted[my_splits[i]] += 1
if i == n - 1:
_commit_pack_state()
_update_stats()
yield block
return
while True:
# A cycle only runs if every split can still produce a
# row. Without skips all splits exhaust simultaneously
@@ -620,21 +917,7 @@ class StreamingDataset(IterableDataset):
self._resume_offset = initial_offset + local_consumed[i]
for j, split_idx in enumerate(my_splits):
self._resume_positions[split_idx] = pos_consumed[j]
ws = self._worker_stats
ws[0] = sum(
split_sizes[j] - fetch_head[j] for j in range(n)
)
ws[1] = sum(
batch.num_rows
for q in raw_batches
for _, batch in q
)
ws[2] = sum(len(q) for q in cooked)
ws[3] = sum(local_consumed)
ws[4] = self._bytes_loaded
ws[5] = int(self._fetch_time * 1_000_000)
ws[6] = int(self._transform_time * 1_000_000)
ws[7] = self._rows_skipped
_update_stats()
yield row
finally:
@@ -642,15 +925,7 @@ class StreamingDataset(IterableDataset):
# when iteration ends mid-cycle (e.g. a split whose rows
# were all skipped before completing a single cycle), so
# counters like rows_skipped would otherwise be stale.
ws = self._worker_stats
ws[0] = sum(split_sizes[j] - fetch_head[j] for j in range(n))
ws[1] = 0 # queue-depth properties document 0 when idle
ws[2] = 0
ws[3] = sum(local_consumed)
ws[4] = self._bytes_loaded
ws[5] = int(self._fetch_time * 1_000_000)
ws[6] = int(self._transform_time * 1_000_000)
ws[7] = self._rows_skipped
_update_stats(idle=True)
self._raw_batches_ref = None
self._cooked_ref = None
self._fetch_head_ref = None
@@ -808,21 +1083,31 @@ class StreamingDataset(IterableDataset):
def state_dict(self) -> dict:
"""Snapshot the dataset's consumption state.
The returned dict is topology-independent: at global step boundaries
every split has been consumed the same number of times (by the
round-robin design), so the per-split count is a single uniform value
that is identical across all ranks and DataLoader workers.
``positions_consumed_per_split`` records how far into each split's
permutation iteration has advanced. It only differs from
``samples_consumed_per_split`` when ``on_transform_error`` skipped
rows, in which case entries are exact for the splits this instance
iterated and a lower bound (the sample count) for splits owned by
other ranks or workers. Combine the state dicts from all ranks with
In row mode, the returned dict is topology-independent at global step
boundaries. ``positions_consumed_per_split`` records how far each
split's permutation has advanced, which can differ from the sample
count when ``on_transform_error`` skips rows. Combine state dicts from
every rank with
[merge_state_dicts][lancedb.streaming.StreamingDataset.merge_state_dicts]
to recover the exact value for every split before resuming on a
different topology.
before resuming on a different topology.
Packed state includes partial token buffers and emitted block counts
for every logical split. When packing is sharded, merge every rank
state with ``merge_state_dicts`` before loading it.
"""
if self._pack_sequences is not None:
return {
"shuffle_seed": self._shuffle_seed,
"num_splits": self._num_splits,
"epoch": self._epoch,
"pack_sequences": self._pack_sequences,
"eos_id": self._eos_id,
"pad_id": self._pad_id,
"blocks_per_epoch": self._blocks_per_epoch,
"samples_consumed_per_split": list(self._pack_consumed),
"blocks_emitted_per_split": list(self._pack_blocks_emitted),
"pack_buffers": deepcopy(self._pack_buffers),
}
positions = [
self._resume_positions.get(split, self._resume_offset)
for split in range(self._num_splits)
@@ -840,7 +1125,9 @@ class StreamingDataset(IterableDataset):
Raises ``ValueError`` if ``num_splits`` or ``shuffle_seed`` differ
from the checkpoint, since a different split structure or shuffle order
makes mid-epoch resumption meaningless.
makes mid-epoch resumption meaningless. Packed checkpoints
pin ``pack_sequences``, ``eos_id``, ``pad_id``,
``blocks_per_epoch``, and ``epoch``.
"""
if state["num_splits"] != self._num_splits:
raise ValueError(
@@ -852,6 +1139,31 @@ class StreamingDataset(IterableDataset):
f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, "
f"current dataset has {self._shuffle_seed}"
)
if "pack_buffers" in state or self._pack_sequences is not None:
for key in (
"pack_sequences",
"eos_id",
"pad_id",
"blocks_per_epoch",
"epoch",
):
ours = getattr(self, f"_{key}")
if state.get(key) != ours:
raise ValueError(
f"{key} mismatch: checkpoint has {state.get(key)}, "
f"current dataset has {ours}"
)
self._pack_consumed = [int(c) for c in state["samples_consumed_per_split"]]
self._pack_blocks_emitted = [
int(c) for c in state["blocks_emitted_per_split"]
]
self._pack_buffers = {
int(g): {"tokens": list(b["tokens"]), "starts": list(b["starts"])}
for g, b in state["pack_buffers"].items()
}
return
consumed = state["samples_consumed_per_split"]
# All entries are equal at step boundaries; use the first.
if isinstance(consumed, list):
@@ -873,25 +1185,22 @@ class StreamingDataset(IterableDataset):
def merge_state_dicts(states: list[dict]) -> dict:
"""Merge state dicts saved by different ranks into one exact state.
Only needed when ``on_transform_error`` skips rows in multi-rank
training: each rank then knows the exact permutation position only for
its own splits, and records a lower bound for the rest. Because
exactly one rank owns each split, the elementwise maximum across all
ranks' ``positions_consumed_per_split`` recovers the exact position of
every split. Without skipped rows every rank's state is already
identical and merging is a no-op.
For row mode, the elementwise maximum of permutation positions recovers
splits advanced by different ranks after transform failures. For packed
mode, the state that emitted the most blocks for each logical split
supplies that split's permutation position and partial token buffer. Packed
states must cover every rank at the same global step.
Raises ``ValueError`` if the states are empty or were not produced by
the same run (mismatched seed, split count, epoch, or sample counts).
Raises ``ValueError`` if the states are empty, were not produced by
the same run, or do not represent the same global step.
The merge is always all-to-all and topology-agnostic: collect the
``state_dict()`` from every rank of the *previous* run into one list,
merge that whole list, and hand the identical merged result to every
rank of the *next* run — regardless of whether the rank count grew,
shrank, or stayed the same. There is no pairwise or subset merging
step, because each split's exact position is only known to whichever
rank owned that split, and the elementwise maximum needs every rank's
contribution to be correct.
``state_dict()`` from every rank of the *previous* run into
one list, merge that whole list, and hand the identical merged result
to every rank of the *next* run — regardless of whether the
topology grew, shrank, or stayed the same. There is no pairwise or
subset merging step, because each split's exact state is only known to
whichever iterator owned that split.
For example, checkpointing 8 ranks and resuming on 4 (the same
pattern applies when growing, e.g. 4 ranks resuming on 8)::
@@ -924,13 +1233,73 @@ class StreamingDataset(IterableDataset):
if not states:
raise ValueError("merge_state_dicts requires at least one state dict")
first = states[0]
packed = "pack_buffers" in first
config_keys = ["shuffle_seed", "num_splits", "epoch"]
if packed:
config_keys.extend(
["pack_sequences", "eos_id", "pad_id", "blocks_per_epoch"]
)
for state in states[1:]:
for key in ("shuffle_seed", "num_splits", "epoch"):
if ("pack_buffers" in state) != packed:
raise ValueError("cannot merge packed and unpacked state dicts")
for key in config_keys:
if state[key] != first[key]:
raise ValueError(
f"{key} mismatch across state dicts: "
f"{state[key]} != {first[key]}"
)
if packed:
num_splits = first["num_splits"]
for state in states:
for key in (
"samples_consumed_per_split",
"blocks_emitted_per_split",
):
if len(state[key]) != num_splits:
raise ValueError(
f"{key} must contain one entry per logical split"
)
merged_consumed = []
merged_emitted = []
merged_buffers = {}
for split in range(num_splits):
owner = states[0]
owner_progress = (
owner["blocks_emitted_per_split"][split],
owner["samples_consumed_per_split"][split],
)
for state in states[1:]:
progress = (
state["blocks_emitted_per_split"][split],
state["samples_consumed_per_split"][split],
)
if progress > owner_progress:
owner = state
owner_progress = progress
merged_consumed.append(owner["samples_consumed_per_split"][split])
merged_emitted.append(owner["blocks_emitted_per_split"][split])
buffer = owner["pack_buffers"].get(
split, owner["pack_buffers"].get(str(split))
)
if buffer is not None:
merged_buffers[split] = deepcopy(buffer)
if len(set(merged_emitted)) > 1:
raise ValueError(
"packed state dicts were not captured at the same global "
"step or do not cover every rank"
)
merged = dict(first)
merged["samples_consumed_per_split"] = merged_consumed
merged["blocks_emitted_per_split"] = merged_emitted
merged["pack_buffers"] = merged_buffers
return merged
for state in states[1:]:
if (
state["samples_consumed_per_split"]
!= first["samples_consumed_per_split"]
+31 -8
View File
@@ -40,7 +40,7 @@ from ._blob import (
from .types import BlobMode
from lancedb.arrow import peek_reader
from lancedb.background_loop import LOOP, embedding_executor
from lancedb.job import AsyncJob, Job
from lancedb.job import AsyncJob, Job, _typed_job
from .dependencies import (
_check_for_hugging_face,
_check_for_lance,
@@ -72,7 +72,10 @@ from .index import (
FTS,
)
from .expr import Expr
from .functions import FunctionApplication
from .functions import (
FunctionApplication,
RefreshColumnResult as RefreshColumnJobResult,
)
from .merge import LanceMergeInsertBuilder
from .pydantic import LanceModel, model_to_dict
from .query import (
@@ -2039,7 +2042,7 @@ class Table(ABC):
"""
@abstractmethod
def refresh_column_async(self, column: str) -> Job:
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
@@ -2050,6 +2053,12 @@ class Table(ABC):
than failing the job. On local tables the job runs in-process; on
LanceDB Cloud and Enterprise it is the server's backfill job.
Returns
-------
Job[RefreshColumnResult]
A job whose successful ``wait`` returns row counts plus the source
and published table versions.
Examples
--------
>>> import lancedb
@@ -2058,7 +2067,9 @@ class Table(ABC):
>>> table.add_columns(computed={"doubled": "x * 2"})
AddColumnsResult(version=2)
>>> job = table.refresh_column_async("doubled")
>>> job.wait()
>>> result = job.wait()
>>> result.rows_assigned
2
>>> job.status()
'finished'
"""
@@ -4082,7 +4093,7 @@ class LanceTable(Table):
[`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
return LOOP.run(self._table.refresh_column(column))
def refresh_column_async(self, column: str) -> Job:
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
"""Fill a computed column's unfilled rows, returning a handle to the
refresh job. See
[`Table.refresh_column_async`][lancedb.table.Table.refresh_column_async].
@@ -6122,7 +6133,9 @@ class AsyncTable:
"""
return await self._inner.refresh_column(column)
async def refresh_column_async(self, column: str) -> AsyncJob:
async def refresh_column_async(
self, column: str
) -> AsyncJob[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
@@ -6134,6 +6147,12 @@ class AsyncTable:
in-process; on LanceDB Cloud and Enterprise it is the server's
backfill job.
Returns
-------
AsyncJob[RefreshColumnResult]
A job whose successful ``wait`` returns row counts plus the source
and published table versions.
Examples
--------
>>> import asyncio
@@ -6143,12 +6162,16 @@ class AsyncTable:
... table = await db.create_table("computed_job_async_demo", [{"x": 1}])
... await table.add_columns(computed={"doubled": "x * 2"})
... job = await table.refresh_column_async("doubled")
... await job.wait()
... result = await job.wait()
... assert result.rows_assigned == 1
... return await job.status()
>>> asyncio.run(refresh_in_background())
'finished'
"""
return AsyncJob(await self._inner.refresh_column_async(column))
return _typed_job(
await self._inner.refresh_column_async(column),
RefreshColumnJobResult.from_json,
)
async def alter_columns(
self, *alterations: Iterable[dict[str, Any]]
+2 -2
View File
@@ -774,7 +774,7 @@ def test_drop_table_async(tmp_db: lancedb.DBConnection):
job = tmp_db.drop_table_async("test")
assert job.id is None
assert job.status() == "finished"
job.wait()
assert job.wait() is None
assert tmp_db.table_names() == []
tmp_db.create_table("test", data=data)
@@ -790,7 +790,7 @@ async def test_drop_table_async_connection(tmp_db_async: lancedb.AsyncConnection
job = await tmp_db_async.drop_table_async("test")
assert job.id is None
assert await job.status() == "finished"
await job.wait()
assert await job.wait() is None
assert await tmp_db_async.table_names() == []
@@ -2113,6 +2113,214 @@ def test_shuffle_seed_none_generates_stable_seed(lance_table):
assert first == second, "Same resolved seed must produce the same ordering"
# Sequence packing tests
def _create_token_table(tmp_path, documents):
db = lancedb.connect(tmp_path)
tokens = pa.array(documents, type=pa.list_(pa.int64()))
return db.create_table("tokens", pa.table({"tokens": tokens}))
def _packed_dataset(table, pack_sequences, *, blocks_per_epoch, pad_id=0, **kwargs):
return StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
pack_sequences=pack_sequences,
eos_id=9,
pad_id=pad_id,
blocks_per_epoch=blocks_per_epoch,
**kwargs,
)
def test_pack_sequences_emits_blocks_and_pads_final_tail(tmp_path):
table = _create_token_table(tmp_path, [[1, 2], [3, 4], [5]])
dataset = _packed_dataset(table, 6, blocks_per_epoch=2)
blocks = list(dataset)
assert len(blocks) == 2
assert blocks[0]["input_ids"].tolist() == [1, 2, 9, 3, 4, 9]
assert blocks[0]["doc_ids"].tolist() == [0, 0, 0, 1, 1, 1]
assert blocks[1]["input_ids"].tolist() == [5, 9, 0, 0, 0, 0]
assert blocks[1]["doc_ids"].tolist() == [0, 0, 0, 0, 0, 0]
assert blocks[0]["input_ids"].dtype == torch.int64
assert blocks[0]["doc_ids"].dtype == torch.int64
def test_pack_sequences_pads_lagging_splits(tmp_path):
table = _create_token_table(
tmp_path,
[[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]],
)
dataset = _packed_dataset(table, 5, blocks_per_epoch=6, num_splits=2)
input_ids = [block["input_ids"].tolist() for block in dataset]
# Split 0 has four real tokens including EOS markers, while split 1 has
# eleven. Packing must emit three complete two-split cycles.
assert input_ids == [
[1, 9, 2, 9, 0],
[10, 11, 12, 13, 14],
[0, 0, 0, 0, 0],
[15, 16, 17, 9, 20],
[0, 0, 0, 0, 0],
[9, 0, 0, 0, 0],
]
per_rank = []
for rank in range(2):
rank_dataset = _packed_dataset(
table,
5,
blocks_per_epoch=6,
num_splits=2,
world_size=2,
rank=rank,
)
per_rank.append([block["input_ids"].tolist() for block in rank_dataset])
assert [len(blocks) for blocks in per_rank] == [3, 3]
sharded = [block for cycle in zip(*per_rank) for block in cycle]
assert sharded == input_ids
def test_pack_sequences_auto_estimates_filtered_token_column(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table(
"tokens",
pa.table(
{
"tokens": pa.array([[1] * 4, [2] * 9], type=pa.list_(pa.int64())),
"keep": [True, False],
}
),
)
table.add(
pa.table(
{
"tokens": pa.array([[3] * 4, [4] * 9], type=pa.list_(pa.int64())),
"keep": [True, False],
}
)
)
with pytest.warns(UserWarning, match="approximate token-count sample"):
dataset = _packed_dataset(
table,
5,
blocks_per_epoch="auto",
num_splits=2,
filter="keep",
)
# Two kept documents contain 8 tokens plus 2 EOS tokens: two blocks.
assert dataset.state_dict()["blocks_per_epoch"] == 2
def test_pack_sequences_checkpoint_resumes_on_new_topology(tmp_path):
table = _create_token_table(
tmp_path,
[[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]],
)
kwargs = dict(pack_sequences=5, blocks_per_epoch=6, num_splits=2)
reference = list(_packed_dataset(table, **kwargs))
datasets = [
_packed_dataset(table, world_size=2, rank=rank, **kwargs) for rank in range(2)
]
iterators = [iter(dataset) for dataset in datasets]
first_cycle = [next(iterator) for iterator in iterators]
checkpoint = StreamingDataset.merge_state_dicts(
[dataset.state_dict() for dataset in datasets]
)
for iterator in iterators:
iterator.close()
resumed = _packed_dataset(table, **kwargs)
resumed.load_state_dict(checkpoint)
actual_remaining = list(resumed)
assert [block["input_ids"].tolist() for block in first_cycle] == [
[1, 9, 2, 9, 0],
[10, 11, 12, 13, 14],
]
assert checkpoint["blocks_emitted_per_split"] == [1, 1]
assert [block["input_ids"].tolist() for block in actual_remaining] == [
block["input_ids"].tolist() for block in reference[2:]
]
assert [block["doc_ids"].tolist() for block in actual_remaining] == [
block["doc_ids"].tolist() for block in reference[2:]
]
def test_pack_sequences_validates_configuration_and_tokens(tmp_path):
table = _create_token_table(tmp_path, [[1, 2]])
with pytest.raises(ValueError, match="pad_id is required"):
StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
pack_sequences=4,
eos_id=9,
)
with pytest.raises(ValueError, match="blocks_per_epoch is required"):
StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
pack_sequences=4,
eos_id=9,
pad_id=0,
)
with pytest.raises(ValueError, match="must be divisible"):
_packed_dataset(table, 4, blocks_per_epoch=3, num_splits=2)
with pytest.raises(ValueError, match="positive integer or 'auto'"):
_packed_dataset(table, 4, blocks_per_epoch="estimate")
checkpoint = _packed_dataset(table, 4, blocks_per_epoch=1).state_dict()
resumed = _packed_dataset(table, 4, blocks_per_epoch=1, pad_id=8)
with pytest.raises(ValueError, match="pad_id mismatch"):
resumed.load_state_dict(checkpoint)
float_db = lancedb.connect(tmp_path / "float")
float_table = float_db.create_table(
"tokens",
pa.table({"tokens": pa.array([[1.5, 2.5]], type=pa.list_(pa.float64()))}),
)
with pytest.raises(ValueError, match="token column with integer values"):
_packed_dataset(float_table, 4, blocks_per_epoch=1)
null_db = lancedb.connect(tmp_path / "null")
null_table = null_db.create_table(
"tokens",
pa.table({"tokens": pa.array([None], type=pa.list_(pa.int64()))}),
)
with pytest.raises(ValueError, match="does not support null token lists"):
list(_packed_dataset(null_table, 4, blocks_per_epoch=1))
null_value_db = lancedb.connect(tmp_path / "null_value")
null_value_table = null_value_db.create_table(
"tokens",
pa.table(
{"tokens": pa.array([[1], [2, None], [3]], type=pa.list_(pa.int64()))}
),
)
blocks = list(
_packed_dataset(
null_value_table,
2,
blocks_per_epoch=2,
on_transform_error="skip",
)
)
assert [block["input_ids"].tolist() for block in blocks] == [[1, 9], [3, 9]]
# ---------------------------------------------------------------------------
# Doc examples — each test mirrors the code snippet in index.mdx so that
# broken doc examples are caught before they ship.
@@ -187,7 +187,7 @@ def _mock_remote_function_catalog():
"job_state": "DONE",
"result": state["version"],
}
elif self.path == "/v1/functions/get":
elif self.path == "/v1/functions/describe":
assert body == {
"name": "normalize_score",
"version": "fv_exact",
+1 -1
View File
@@ -88,7 +88,7 @@ async def binary_table(db_async):
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 await job.wait() is None
assert len(await some_table.list_indices()) == 1
await job.cancel()
+75 -1
View File
@@ -875,11 +875,85 @@ def test_remote_create_index_async_returns_job():
table = db.create_table("test", [{"id": 1}])
job = table.create_index_async("id", config=BTree())
assert job.id == "job-1"
job.wait(timeout=timedelta(seconds=30))
assert job.wait(timeout=timedelta(seconds=30)) is None
assert len(describe_calls) == 2
job.cancel()
def test_remote_refresh_async_returns_typed_terminal_result():
terminal_result = {
"rows_assigned": 12,
"rows_failed": 0,
"rows_remaining": 0,
"source_version": 7,
"published_version": 8,
}
def handler(request):
content_len = int(request.headers.get("Content-Length", 0))
body = request.rfile.read(content_len) if content_len > 0 else b""
if request.path == "/v1/table/test/backfill_column":
assert json.loads(body)["column"] == "derived"
request.send_response(202)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"job_id": "refresh-1"}')
elif request.path == "/v1/jobs/describe":
assert json.loads(body)["job_id"] == "refresh-1"
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
{
"job_id": "refresh-1",
"job_type": "function_refresh",
"job_state": "DONE",
"result": terminal_result,
}
).encode()
)
elif request.path == "/v1/table/test/create/?mode=create":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b"{}")
elif request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
{
"version": 1,
"schema": {
"fields": [
{
"name": "id",
"type": {"type": "int64"},
"nullable": False,
}
]
},
}
).encode()
)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
table = db.create_table("test", [{"id": 1}])
job = table.refresh_column_async("derived")
assert job.id == "refresh-1"
result = job.wait(timeout=timedelta(seconds=30))
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.model_dump() == terminal_result
assert result.rows_filled == 12
assert result.version == 8
def test_remote_job_wait_raises_on_failure():
from lancedb.exceptions import JobFailedError
from lancedb.index import BTree
+18 -3
View File
@@ -1467,7 +1467,7 @@ 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 job.wait() is None
assert len(table.list_indices()) == 1
job.cancel()
@@ -3947,10 +3947,21 @@ def test_refresh_column_async_returns_job(tmp_path):
job = table.refresh_column_async("doubled")
assert job.id is None # in-process jobs have no server id
assert job.wait() is None
result = job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 2
assert result.rows_failed == 0
assert result.rows_remaining == 0
assert result.source_version == 2
assert result.published_version == 3
assert job.status() == "finished"
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4]
no_op = table.refresh_column_async("doubled").wait()
assert no_op.rows_assigned == 0
assert no_op.source_version == 3
assert no_op.published_version is None
# Bad input raises at the call, not through the job.
with pytest.raises(Exception, match="not a computed column"):
table.refresh_column_async("x")
@@ -3963,6 +3974,10 @@ async def test_refresh_column_async_job_async_table(tmp_path):
await table.add_columns(computed={"tripled": "x * 3"})
job = await table.refresh_column_async("tripled")
assert await job.wait() is None
result = await job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 1
assert result.source_version == 2
assert result.published_version == 3
assert await job.status() == "finished"
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
+1 -1
View File
@@ -609,7 +609,7 @@ impl Connection {
.create_function_async(request)
.await
.infer_error()
.map(crate::job::FunctionJob::new)
.map(crate::job::Job::new_typed)
})
}
+18 -55
View File
@@ -5,72 +5,33 @@ use std::sync::Arc;
use crate::runtime::future_into_py;
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
use serde::Serialize;
use crate::error::PythonErrorExt;
#[pyclass]
pub struct Job {
inner: Arc<lancedb::Job>,
}
/// Python bridge for a typed remote Function registration job.
///
/// The public Python layer decodes the canonical JSON returned by `wait`
/// into its immutable `FunctionVersion` model.
#[pyclass]
pub struct FunctionJob {
inner: Arc<lancedb::Job<lancedb::function::FunctionVersion>>,
}
impl FunctionJob {
pub(crate) fn new(inner: lancedb::Job<lancedb::function::FunctionVersion>) -> Self {
Self {
inner: Arc::new(inner),
}
}
inner: Arc<lancedb::Job<std::result::Result<Option<String>, String>>>,
}
impl Job {
pub(crate) fn new(inner: lancedb::Job) -> Self {
Self {
inner: Arc::new(inner),
inner: Arc::new(inner.map(|()| Ok(None))),
}
}
}
#[pymethods]
impl FunctionJob {
#[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()?
.to_canonical_json()
.infer_error()
})
}
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(())
})
pub(crate) fn new_typed<T>(inner: lancedb::Job<T>) -> Self
where
T: Clone + Serialize + Send + Sync + 'static,
{
Self {
inner: Arc::new(inner.map(|result| {
serde_json::to_string(&result)
.map(Some)
.map_err(|error| format!("failed to serialize typed job result: {error}"))
})),
}
}
}
@@ -92,8 +53,10 @@ impl Job {
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(None::<()>)
let result = inner.wait().await.infer_error()?;
result
.map_err(|message| lancedb::Error::Runtime { message })
.infer_error()
})
}
-1
View File
@@ -47,7 +47,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<Session>()?;
m.add_class::<Table>()?;
m.add_class::<crate::job::Job>()?;
m.add_class::<crate::job::FunctionJob>()?;
m.add_class::<crate::job::JobInfo>()?;
m.add_class::<crate::job::JobDescription>()?;
m.add_class::<crate::job::JobFailureInfo>()?;
+1 -1
View File
@@ -1619,7 +1619,7 @@ impl Table {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let job = inner.refresh_column_async(column).await.infer_error()?;
Ok(crate::job::Job::new(job))
Ok(crate::job::Job::new_typed(job))
})
}
+12 -2
View File
@@ -1,7 +1,8 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Canonical values exchanged with the Enterprise Function service.
//! Canonical Function values exchanged with the Enterprise service, plus the
//! backend-neutral terminal result of a computed-column refresh.
//!
//! This module contains client/wire values only. Catalog persistence,
//! environment bake, secret resolution, and execution are owned by Sophon.
@@ -580,13 +581,22 @@ impl FunctionBinding {
impl_json!(FunctionBinding);
/// Stable terminal result of a remote Function-column refresh Job.
/// Stable terminal result of an expression-backed or Function-backed column
/// refresh [`crate::Job`].
///
/// Local refresh jobs produce this value in process. LanceDB Cloud and
/// Enterprise decode the same value from the durable job's terminal payload.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RefreshColumnResult {
/// Rows assigned a value by this refresh.
pub rows_assigned: u64,
/// Rows whose computation failed.
pub rows_failed: u64,
/// Rows that still need a value when the job completes.
pub rows_remaining: u64,
/// Exact table version the refresh read.
pub source_version: u64,
/// Table version made visible by the refresh, when one was published.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub published_version: Option<u64>,
}
+115 -37
View File
@@ -6,7 +6,7 @@
use std::sync::Arc;
use async_trait::async_trait;
use serde::de::DeserializeOwned;
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
use tokio::sync::watch;
use tokio::task::{AbortHandle, JoinHandle};
@@ -26,20 +26,16 @@ pub(crate) trait JobHandle: Send + Sync {
}
/// A backend-neutral successful terminal result.
///
/// Local operations do not carry a value. Remote operations may carry JSON
/// that the public [`Job`] decodes according to its result type.
#[derive(Clone)]
pub(crate) struct TerminalResult {
#[allow(dead_code)] // Typed remote submit endpoints consume this after Slice 1.
value: Option<Value>,
#[allow(dead_code)] // Preserved so typed decode errors retain request correlation.
request_id: Option<String>,
}
impl TerminalResult {
pub(crate) fn local() -> Self {
fn local(value: Value) -> Self {
Self {
value: None,
value: Some(value),
request_id: None,
}
}
@@ -51,23 +47,31 @@ impl TerminalResult {
}
}
#[allow(dead_code)] // Exercised by the remote typed-result fixtures in Slice 1.
fn decode<T: DeserializeOwned>(self) -> Result<T> {
let request_id = self.request_id.unwrap_or_default();
let value = self.value.ok_or_else(|| Error::Http {
source: "successful typed job response did not contain a result".into(),
request_id: request_id.clone(),
status_code: None,
let value = self.value.ok_or_else(|| match &self.request_id {
Some(request_id) => Error::Http {
source: "successful typed job response did not contain a result".into(),
request_id: request_id.clone(),
status_code: None,
},
None => Error::Runtime {
message: "successful typed job did not contain a result".to_string(),
},
})?;
serde_json::from_value(value).map_err(|error| Error::Http {
source: format!("failed to parse typed job result: {error}").into(),
request_id,
status_code: None,
serde_json::from_value(value).map_err(|error| match self.request_id {
Some(request_id) => Error::Http {
source: format!("failed to parse typed job result: {error}").into(),
request_id,
status_code: None,
},
None => Error::Runtime {
message: format!("failed to parse typed job result: {error}"),
},
})
}
}
type ResultDecoder<T> = fn(TerminalResult) -> Result<T>;
type ResultDecoder<T> = Arc<dyn Fn(TerminalResult) -> Result<T> + Send + Sync>;
enum JobInner<T> {
Handle {
@@ -79,7 +83,9 @@ enum JobInner<T> {
/// A handle to an operation that may still be running.
///
/// The operation may already be complete when the handle is created.
/// The operation may already be complete when the handle is created. `T` is
/// the endpoint's successful terminal result; unit-result operations use the
/// default `Job<()>`.
pub struct Job<T = ()>
where
T: Clone + Send + Sync + 'static,
@@ -111,15 +117,10 @@ impl Job<()> {
Self {
inner: JobInner::Handle {
handle,
decode: |_| Ok(()),
decode: Arc::new(|_| Ok(())),
},
}
}
/// A unit-result job running as a task in this process.
pub(crate) fn spawned(task: JoinHandle<Result<()>>) -> Self {
Self::new(Box::new(SpawnedJob::new(task)))
}
}
impl<T> Job<T>
@@ -131,12 +132,22 @@ where
Self {
inner: JobInner::Handle {
handle,
decode: TerminalResult::decode::<T>,
decode: Arc::new(TerminalResult::decode::<T>),
},
}
}
}
impl<T> Job<T>
where
T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static,
{
/// A typed job running as a task in this process.
pub(crate) fn spawned(task: JoinHandle<Result<T>>) -> Self {
Self::new_typed(Box::new(SpawnedJob::new(task)))
}
}
impl<T> Job<T>
where
T: Clone + Send + Sync + 'static,
@@ -169,11 +180,13 @@ where
/// Waits until the operation reaches a terminal state.
///
/// Returns the endpoint's typed result. Unit-result jobs return `()`.
///
/// Returns [`crate::Error::JobFailed`] if the operation failed and
/// [`crate::Error::JobCancelled`] if it was cancelled.
pub async fn wait(&self) -> Result<T> {
match &self.inner {
JobInner::Handle { handle, decode } => decode(handle.wait().await?),
JobInner::Handle { handle, decode } => (decode)(handle.wait().await?),
JobInner::Completed(result) => Ok(result.clone()),
}
}
@@ -187,21 +200,53 @@ where
JobInner::Completed(_) => Ok(()),
}
}
/// Maps a successful terminal result without changing the job lifecycle.
/// The mapping may run once for each call to [`Job::wait`], so it should
/// be deterministic and free of externally visible side effects.
///
/// ```
/// use lancedb::{Job, function::RefreshColumnResult};
///
/// # async fn rows_assigned(
/// # job: Job<RefreshColumnResult>,
/// # ) -> lancedb::Result<u64> {
/// let job = job.map(|result| result.rows_assigned);
/// job.wait().await
/// # }
/// ```
pub fn map<U, F>(self, map: F) -> Job<U>
where
U: Clone + Send + Sync + 'static,
F: Fn(T) -> U + Send + Sync + 'static,
{
match self.inner {
JobInner::Handle { handle, decode } => Job {
inner: JobInner::Handle {
handle,
decode: Arc::new(move |result| Ok(map((decode)(result)?))),
},
},
JobInner::Completed(result) => Job {
inner: JobInner::Completed(map(result)),
},
}
}
}
/// 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,
Succeeded(TerminalResult),
Failed(Arc<Error>),
Cancelled,
}
impl Outcome {
fn into_result(self) -> Result<()> {
fn into_result(self) -> Result<TerminalResult> {
match self {
Self::Succeeded => Ok(()),
Self::Succeeded(result) => Ok(result),
Self::Failed(source) => Err(Error::JobFailed {
job_id: None,
failure: JobFailure::from_source(source),
@@ -220,12 +265,20 @@ struct SpawnedJob {
}
impl SpawnedJob {
fn new(task: JoinHandle<Result<()>>) -> Self {
fn new<T>(task: JoinHandle<Result<T>>) -> Self
where
T: Serialize + Send + 'static,
{
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(Ok(result)) => match serde_json::to_value(result) {
Ok(value) => Outcome::Succeeded(TerminalResult::local(value)),
Err(err) => Outcome::Failed(Arc::new(Error::Runtime {
message: format!("failed to serialize job result: {err}"),
})),
},
Ok(Err(err)) => Outcome::Failed(Arc::new(err)),
Err(err) if err.is_cancelled() => Outcome::Cancelled,
Err(err) => Outcome::Failed(Arc::new(Error::Runtime {
@@ -243,7 +296,7 @@ impl JobHandle for SpawnedJob {
async fn status(&self) -> Result<String> {
let label = match &*self.outcome.borrow() {
None => "running",
Some(Outcome::Succeeded) => "finished",
Some(Outcome::Succeeded(_)) => "finished",
Some(Outcome::Failed(_)) => "failed",
Some(Outcome::Cancelled) => "cancelled",
};
@@ -256,12 +309,11 @@ impl JobHandle for SpawnedJob {
.wait_for(|outcome| outcome.is_some())
.await
.map_err(|_| Error::Runtime {
message: "index job outcome was dropped before it completed".to_string(),
message: "job outcome was dropped before it completed".to_string(),
})?
.clone()
.expect("wait_for returns once an outcome is set");
settled.into_result()?;
Ok(TerminalResult::local())
settled.into_result()
}
async fn cancel(&self) -> Result<()> {
@@ -269,3 +321,29 @@ impl JobHandle for SpawnedJob {
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::future::pending;
use super::*;
#[tokio::test]
async fn mapped_spawned_job_reuses_outcome() {
let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1);
assert_eq!(job.wait().await.unwrap(), 42);
assert_eq!(job.wait().await.unwrap(), 42);
assert_eq!(job.status().await.unwrap(), "finished");
}
#[tokio::test]
async fn mapped_spawned_job_preserves_cancellation() {
let job = Job::spawned(tokio::spawn(async { pending::<Result<u64>>().await }))
.map(|value| value.to_string());
job.cancel().await.unwrap();
assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. })));
assert_eq!(job.status().await.unwrap(), "cancelled");
}
}
+2 -2
View File
@@ -513,7 +513,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
async fn get_function(&self, name: &str, version: &str) -> Result<FunctionVersion> {
let req = self
.client
.post("/v1/functions/get")
.post("/v1/functions/describe")
.json(&serde_json::json!({
"name": name,
"version": version,
@@ -2520,7 +2520,7 @@ mod tests {
);
let conn = Connection::new_with_handler(|request| {
assert_eq!(request.method(), &reqwest::Method::POST);
assert_eq!(request.url().path(), "/v1/functions/get");
assert_eq!(request.url().path(), "/v1/functions/describe");
let body: serde_json::Value =
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(
+28 -8
View File
@@ -2864,7 +2864,10 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
})
}
async fn refresh_column_async(&self, column: &str) -> Result<Job> {
async fn refresh_column_async(
&self,
column: &str,
) -> Result<Job<crate::function::RefreshColumnResult>> {
self.check_mutable().await?;
let mut body = serde_json::json!({ "column": column });
self.apply_branch_body(&mut body);
@@ -2885,7 +2888,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
status_code: None,
})?;
Ok(Job::new(Box::new(FreshnessJob {
Ok(Job::new_typed(Box::new(FreshnessJob {
inner: RemoteJob::new(self.client.clone(), response.job_id),
freshness: self.freshness.clone(),
version: self.version.clone(),
@@ -3280,6 +3283,21 @@ mod tests {
},
};
fn refresh_done(job_id: &str) -> String {
json!({
"job_id": job_id,
"job_state": "DONE",
"result": {
"rows_assigned": 12,
"rows_failed": 0,
"rows_remaining": 0,
"source_version": 7,
"published_version": 8,
}
})
.to_string()
}
#[tokio::test]
async fn test_not_found() {
let table = Table::new_with_handler("my_table", |_| {
@@ -6892,7 +6910,7 @@ mod tests {
.unwrap(),
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(r#"{"job_id": "j-7", "job_state": "DONE"}"#.to_string())
.body(refresh_done("j-7"))
.unwrap(),
"/v1/table/my_table/count_rows/" => {
saw.store(
@@ -6908,7 +6926,9 @@ mod tests {
});
let job = table.refresh_column_async("doubled").await.unwrap();
job.wait().await.unwrap();
let result = job.wait().await.unwrap();
assert_eq!(result.rows_assigned, 12);
assert_eq!(result.published_version, Some(8));
table.count_rows(None).await.unwrap();
assert!(
saw_min_timestamp.load(std::sync::atomic::Ordering::SeqCst),
@@ -6930,7 +6950,7 @@ mod tests {
.unwrap(),
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(r#"{"job_id": "j-8", "job_state": "DONE"}"#.to_string())
.body(refresh_done("j-8"))
.unwrap(),
"/v1/table/my_table/describe/" => {
let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
@@ -6976,7 +6996,7 @@ mod tests {
.unwrap(),
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(r#"{"job_id": "j-9", "job_state": "DONE"}"#.to_string())
.body(refresh_done("j-9"))
.unwrap(),
"/v1/table/my_table/tags/version/" => http::Response::builder()
.status(200)
@@ -7040,7 +7060,7 @@ mod tests {
}
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(r#"{"job_id": "j-10", "job_state": "DONE"}"#.to_string())
.body(refresh_done("j-10"))
.unwrap(),
"/v1/table/my_table/describe/" => {
let schema = Schema::new(vec![Field::new("x", DataType::Int32, true)]);
@@ -7113,7 +7133,7 @@ mod tests {
}
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(r#"{"job_id": "j-11", "job_state": "DONE"}"#.to_string())
.body(refresh_done("j-11"))
.unwrap(),
"/v1/table/my_table/count_rows/" => {
*saw.lock().unwrap() = request
+17 -5
View File
@@ -771,7 +771,10 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
}
/// Fill a computed column's unfilled rows, returning a [`Job`] tracking
/// the operation.
async fn refresh_column_async(&self, _column: &str) -> Result<Job> {
async fn refresh_column_async(
&self,
_column: &str,
) -> Result<Job<crate::function::RefreshColumnResult>> {
Err(Error::NotSupported {
message: "computed columns are supported only on local tables".into(),
})
@@ -1708,7 +1711,9 @@ impl Table {
/// operation instead of blocking until it completes.
///
/// The job may already be complete when returned, and callers must not
/// assume the column is filled until [`Job::wait`] returns. Invalid input
/// assume the column is filled until [`Job::wait`] returns. A successful
/// wait returns the durable [`crate::function::RefreshColumnResult`] for
/// both expression-backed and Function-backed columns. Invalid input
/// -- an unknown column, or one that is not computed -- is reported by
/// this call rather than by the job. On local tables the job runs as an
/// in-process task; on LanceDB Cloud and Enterprise it is the server's
@@ -1719,11 +1724,15 @@ impl Table {
/// # async fn refresh_in_background(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
/// let job = table.refresh_column_async("doubled").await?;
/// println!("refresh running: {:?}", job.status().await?);
/// job.wait().await?;
/// let result = job.wait().await?;
/// println!("assigned {} rows", result.rows_assigned);
/// # Ok(())
/// # }
/// ```
pub async fn refresh_column_async(&self, column: impl AsRef<str>) -> Result<Job> {
pub async fn refresh_column_async(
&self,
column: impl AsRef<str>,
) -> Result<Job<crate::function::RefreshColumnResult>> {
self.inner.refresh_column_async(column.as_ref()).await
}
@@ -3425,7 +3434,10 @@ impl BaseTable for NativeTable {
Ok(result)
}
async fn refresh_column_async(&self, column: &str) -> Result<Job> {
async fn refresh_column_async(
&self,
column: &str,
) -> Result<Job<crate::function::RefreshColumnResult>> {
refresh::execute_refresh_column_async(self, column).await
}
+59 -12
View File
@@ -49,11 +49,25 @@ pub struct RefreshColumnResult {
pub version: u64,
}
struct RefreshExecution {
result: RefreshColumnResult,
source_version: u64,
}
/// Internal implementation of the refresh logic.
pub(crate) async fn execute_refresh_column(
table: &NativeTable,
column: &str,
) -> Result<RefreshColumnResult> {
Ok(execute_refresh_column_with_source(table, column)
.await?
.result)
}
async fn execute_refresh_column_with_source(
table: &NativeTable,
column: &str,
) -> Result<RefreshExecution> {
table.dataset.ensure_mutable()?;
ensure_no_lsm_write_spec(table).await?;
let dataset = table.dataset.get().await?;
@@ -87,9 +101,13 @@ pub(crate) async fn execute_refresh_column(
}
if replacements.is_empty() {
return Ok(RefreshColumnResult {
rows_filled: 0,
version: dataset.version().version,
let source_version = dataset.version().version;
return Ok(RefreshExecution {
result: RefreshColumnResult {
rows_filled: 0,
version: source_version,
},
source_version,
});
}
@@ -110,14 +128,20 @@ pub(crate) async fn execute_refresh_column(
let version = new_dataset.version().version;
table.dataset.update(new_dataset);
Ok(RefreshColumnResult {
rows_filled,
version,
Ok(RefreshExecution {
result: RefreshColumnResult {
rows_filled,
version,
},
source_version: read_version,
})
}
/// Run the refresh as a [`Job`] in this process.
pub(crate) async fn execute_refresh_column_async(table: &NativeTable, column: &str) -> Result<Job> {
pub(crate) async fn execute_refresh_column_async(
table: &NativeTable,
column: &str,
) -> Result<Job<crate::function::RefreshColumnResult>> {
// Validate before spawning so bad input is reported by this call rather
// than only by the job.
table.dataset.ensure_mutable()?;
@@ -129,9 +153,16 @@ pub(crate) async fn execute_refresh_column_async(table: &NativeTable, column: &s
let table = table.clone();
let column = column.to_string();
Ok(Job::spawned(tokio::spawn(async move {
execute_refresh_column(&table, &column).await?;
let execution = execute_refresh_column_with_source(&table, &column).await?;
table.bump_freshness();
Ok(())
Ok(crate::function::RefreshColumnResult {
rows_assigned: execution.result.rows_filled,
rows_failed: 0,
rows_remaining: 0,
source_version: execution.source_version,
published_version: (execution.result.rows_filled > 0)
.then_some(execution.result.version),
})
})))
}
@@ -396,6 +427,17 @@ mod tests {
read(&table, "doubled").await,
vec![Some(2), Some(4), Some(6)]
);
let no_op = table
.refresh_column_async("doubled")
.await
.unwrap()
.wait()
.await
.unwrap();
assert_eq!(no_op.rows_assigned, 0);
assert_eq!(no_op.source_version, 3);
assert_eq!(no_op.published_version, None);
}
/// Values written after the last refresh must be reachable by another one.
@@ -646,7 +688,12 @@ mod tests {
let job = table.refresh_column_async("doubled").await.unwrap();
assert!(job.id().is_none(), "in-process jobs have no server id");
job.wait().await.unwrap();
let result = job.wait().await.unwrap();
assert_eq!(result.rows_assigned, 3);
assert_eq!(result.rows_failed, 0);
assert_eq!(result.rows_remaining, 0);
assert_eq!(result.source_version, 2);
assert_eq!(result.published_version, Some(3));
assert_eq!(job.status().await.unwrap(), "finished");
assert_eq!(
read(&table, "doubled").await,
@@ -672,9 +719,9 @@ mod tests {
declare_doubled(&table).await.unwrap();
let job = table.refresh_column_async("doubled").await.unwrap();
job.wait().await.unwrap();
let first = job.wait().await.unwrap();
// A second wait after completion observes the same outcome.
job.wait().await.unwrap();
assert_eq!(job.wait().await.unwrap(), first);
assert_eq!(job.status().await.unwrap(), "finished");
}