mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-08 06:18:59 +00:00
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
456dae3e62 | ||
|
|
b0dae5eb0b | ||
|
|
242ade8017 | ||
|
|
40d4d012e7 |
Generated
+45
-45
@@ -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
@@ -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
@@ -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>
|
||||
|
||||
Generated
+2
-2
@@ -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
@@ -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(|_| ())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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]]
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
@@ -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()
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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))
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user