mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-09 06:42:30 +00:00
Compare commits
33
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1f8a790004 | ||
|
|
7b29fb2f51 | ||
|
|
0111a72dc3 | ||
|
|
a487d4033e | ||
|
|
02ea0dda9f | ||
|
|
c7980dbc40 | ||
|
|
1b0f9329ea | ||
|
|
e5cc7a4d66 | ||
|
|
21f11b4463 | ||
|
|
8c9c5c5a5f | ||
|
|
aab23eb39e | ||
|
|
e639b1b650 | ||
|
|
2779b75d0d | ||
|
|
0d19a6c546 | ||
|
|
c0f33f8627 | ||
|
|
904bd975e5 | ||
|
|
d2ca0ce0ab | ||
|
|
9a1ffb9e02 | ||
|
|
f2eb4a245d | ||
|
|
e6867f7d04 | ||
|
|
193c5e3458 | ||
|
|
7ebd3c222d | ||
|
|
d118ef168b | ||
|
|
19232f9c50 | ||
|
|
5cbd979455 | ||
|
|
e773d1e093 | ||
|
|
c8fd3e97d1 | ||
|
|
c196d033e9 | ||
|
|
16753b805a | ||
|
|
840e1d7313 | ||
|
|
1a9414c47c | ||
|
|
c4ee8ae670 | ||
|
|
57b8d3bf05 |
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.38.0-beta.15"
|
||||
current_version = "0.39.0-beta.4"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -232,7 +232,10 @@ jobs:
|
||||
ALL_FEATURES=`cargo metadata --format-version=1 --no-deps \
|
||||
| jq -r '.packages[] | .features | keys | .[]' \
|
||||
| grep -v s3-test | sort | uniq | paste -s -d "," -`
|
||||
cargo test --profile ci --features $ALL_FEATURES --locked
|
||||
# Run doctests before test binaries fill the runner disk. Examples are
|
||||
# already built by the Linux job, so avoid retaining them here.
|
||||
cargo test --profile ci --features $ALL_FEATURES --locked --doc
|
||||
cargo test --profile ci --features $ALL_FEATURES --locked --lib --tests
|
||||
|
||||
windows:
|
||||
strategy:
|
||||
|
||||
Generated
+164
-74
@@ -332,6 +332,34 @@ dependencies = [
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "arrow-flight"
|
||||
version = "58.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b2dbe34824c639e43136af8f106992792ab456540d54b880bc320a3192502d2e"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
"arrow-cast",
|
||||
"arrow-data",
|
||||
"arrow-ipc",
|
||||
"arrow-ord",
|
||||
"arrow-row",
|
||||
"arrow-schema",
|
||||
"arrow-select",
|
||||
"arrow-string",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures",
|
||||
"once_cell",
|
||||
"paste",
|
||||
"prost",
|
||||
"prost-types",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "arrow-ipc"
|
||||
version = "58.4.0"
|
||||
@@ -535,9 +563,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.91"
|
||||
version = "0.1.92"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec"
|
||||
checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -1129,7 +1157,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum-core",
|
||||
"axum-core 0.4.5",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http 1.5.0",
|
||||
@@ -1138,7 +1166,7 @@ dependencies = [
|
||||
"hyper 1.9.0",
|
||||
"hyper-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"matchit 0.7.3",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
@@ -1156,6 +1184,31 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90"
|
||||
dependencies = [
|
||||
"axum-core 0.5.6",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http 1.5.0",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"itoa",
|
||||
"matchit 0.8.4",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"serde_core",
|
||||
"sync_wrapper",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.4.5"
|
||||
@@ -1177,6 +1230,24 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.5.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"http 1.5.0",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"sync_wrapper",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "backoff"
|
||||
version = "0.4.0"
|
||||
@@ -1443,9 +1514,9 @@ checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e"
|
||||
|
||||
[[package]]
|
||||
name = "bytemuck"
|
||||
version = "1.25.0"
|
||||
version = "1.25.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec"
|
||||
checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797"
|
||||
dependencies = [
|
||||
"bytemuck_derive",
|
||||
]
|
||||
@@ -3455,8 +3526,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "fsst"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"rand 0.9.5",
|
||||
@@ -4815,8 +4886,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
|
||||
|
||||
[[package]]
|
||||
name = "lance"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -4888,8 +4959,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-arrow"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4911,7 +4982,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4925,7 +4996,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -4934,8 +5005,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-bitpacking"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrayref",
|
||||
"crunchy",
|
||||
@@ -4945,8 +5016,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-core"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4983,8 +5054,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datafusion"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5000,6 +5071,7 @@ dependencies = [
|
||||
"datafusion-functions",
|
||||
"datafusion-physical-expr",
|
||||
"futures",
|
||||
"half",
|
||||
"jsonb",
|
||||
"lance-arrow",
|
||||
"lance-core",
|
||||
@@ -5013,8 +5085,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datagen"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5031,8 +5103,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-derive"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5041,8 +5113,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-encoding"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5075,8 +5147,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-file"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5107,8 +5179,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5172,8 +5244,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index-core"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5195,8 +5267,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-io"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5236,8 +5308,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-linalg"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5251,27 +5323,29 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
"bytes",
|
||||
"lance-core",
|
||||
"lance-namespace-reqwest-client",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"snafu 0.9.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-impls"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
"arrow-schema",
|
||||
"async-trait",
|
||||
"axum",
|
||||
"axum 0.7.9",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"chrono",
|
||||
@@ -5304,9 +5378,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-reqwest-client"
|
||||
version = "0.11.0"
|
||||
version = "0.12.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0a030196da1c994b63a96a4f0bf5b0cfa459fe6dadc9e962320246ca328da22a"
|
||||
checksum = "d8d23e54b1634d5bbb434f8dd33dc3c05f6e58d876a9a27b3b4aef58ddbe11af"
|
||||
dependencies = [
|
||||
"reqwest 0.12.28",
|
||||
"serde",
|
||||
@@ -5318,8 +5392,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-select"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5333,8 +5407,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-table"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5374,8 +5448,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-testing"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5388,8 +5462,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-tokenizer"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
version = "12.0.0-beta.14"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.14#a8a101774a1c9647065cc60137094feadbe55296"
|
||||
dependencies = [
|
||||
"frostem",
|
||||
"icu_segmenter",
|
||||
@@ -5402,7 +5476,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.14"
|
||||
version = "0.39.0-beta.4"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"anyhow",
|
||||
@@ -5411,6 +5485,7 @@ dependencies = [
|
||||
"arrow-buffer",
|
||||
"arrow-cast",
|
||||
"arrow-data",
|
||||
"arrow-flight",
|
||||
"arrow-ipc",
|
||||
"arrow-ord",
|
||||
"arrow-schema",
|
||||
@@ -5466,6 +5541,7 @@ dependencies = [
|
||||
"polars",
|
||||
"polars-arrow",
|
||||
"pprof 0.14.1",
|
||||
"prost",
|
||||
"rand 0.9.5",
|
||||
"random_word",
|
||||
"regex",
|
||||
@@ -5482,6 +5558,7 @@ dependencies = [
|
||||
"test-log",
|
||||
"tokenizers",
|
||||
"tokio",
|
||||
"tonic",
|
||||
"url",
|
||||
"urlencoding",
|
||||
"uuid",
|
||||
@@ -5490,7 +5567,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-nodejs"
|
||||
version = "0.38.0-beta.14"
|
||||
version = "0.39.0-beta.4"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5515,7 +5592,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.14"
|
||||
version = "0.39.0-beta.4"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5539,6 +5616,7 @@ dependencies = [
|
||||
"serde_json",
|
||||
"snafu 0.8.9",
|
||||
"tokio",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -5748,9 +5826,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "log"
|
||||
version = "0.4.33"
|
||||
version = "0.4.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
|
||||
checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6"
|
||||
|
||||
[[package]]
|
||||
name = "loom"
|
||||
@@ -5859,6 +5937,12 @@ version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94"
|
||||
|
||||
[[package]]
|
||||
name = "matchit"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "matrixmultiply"
|
||||
version = "0.3.10"
|
||||
@@ -6001,9 +6085,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "moka"
|
||||
version = "0.12.15"
|
||||
version = "0.12.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046"
|
||||
checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9"
|
||||
dependencies = [
|
||||
"async-lock",
|
||||
"crossbeam-channel",
|
||||
@@ -6097,14 +6181,15 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi"
|
||||
version = "3.11.0"
|
||||
version = "3.12.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "de33522036981030a75c231829566bc63414e08101a6f5ff4ac6cef19c8e0941"
|
||||
checksum = "58c5f4d5375213fdb7be2655e152386e82f026f9a5ba36a75556e11359aafe09"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"chrono",
|
||||
"ctor 1.0.12",
|
||||
"futures",
|
||||
"libc",
|
||||
"napi-build",
|
||||
"napi-sys",
|
||||
"nohash-hasher",
|
||||
@@ -6116,15 +6201,15 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi-build"
|
||||
version = "2.4.0"
|
||||
version = "2.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5282704fbe8d49b0cf8b08e3f33233416a528658f205c7e5ace63b582de0b11c"
|
||||
checksum = "60fdf9b392c50e7c4170fa633bd909490ed7835cea4c046776d1a4dd8d2ae0ab"
|
||||
|
||||
[[package]]
|
||||
name = "napi-derive"
|
||||
version = "3.6.1"
|
||||
version = "3.6.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4d5c9c02556ea6dc99dffd36c1ce60141411657438501a125b675776d011ce92"
|
||||
checksum = "0fa55ea69990c90b888e9e77044410e304ce7f35de599dc6d0b5c1923d2e59af"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"ctor 1.0.12",
|
||||
@@ -6136,9 +6221,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi-derive-backend"
|
||||
version = "6.1.1"
|
||||
version = "6.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d60b5d773ad46c698c8cc2cd9fde0b283d39cbb7f71c04bee633c7bdba4423bd"
|
||||
checksum = "df4056ac7c18e4438ccf0edaed4340ca0d269278c8ec19284f7b23cb039fd0ae"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"proc-macro2",
|
||||
@@ -7733,6 +7818,7 @@ dependencies = [
|
||||
"pyo3-build-config",
|
||||
"pyo3-ffi",
|
||||
"pyo3-macros",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8601,9 +8687,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "roaring"
|
||||
version = "0.11.4"
|
||||
version = "0.11.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1dedc5658c6ecb3bdb5ef5f3295bb9253f42dcf3fd1402c03f6b1f7659c3c4a9"
|
||||
checksum = "18bd8a37d17a58532776dcdf6041ce64929adca78e8489d5cacbafe99229d3e1"
|
||||
dependencies = [
|
||||
"bytemuck",
|
||||
"byteorder",
|
||||
@@ -9063,9 +9149,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "serde_with"
|
||||
version = "3.21.0"
|
||||
version = "3.22.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "76a5c54c7310e7b8b9577c286d7e399ddd876c3e12b3ed917a8aabc4b96e9e8c"
|
||||
checksum = "ee78f1fbe43ac4a0e47aadb3dbd357b69eb0d3793e948624cd03dd2750ab1c0a"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bs58",
|
||||
@@ -9073,6 +9159,7 @@ dependencies = [
|
||||
"hex",
|
||||
"indexmap 1.9.3",
|
||||
"indexmap 2.14.0",
|
||||
"jiff",
|
||||
"schemars 0.9.0",
|
||||
"schemars 1.2.1",
|
||||
"serde_core",
|
||||
@@ -9083,9 +9170,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "serde_with_macros"
|
||||
version = "3.21.0"
|
||||
version = "3.22.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660"
|
||||
checksum = "8705578779c2b6bd90d84d66eb2e206b708b1a4d7b9f17641b293545bf1c7e46"
|
||||
dependencies = [
|
||||
"darling 0.23.0",
|
||||
"proc-macro2",
|
||||
@@ -10084,6 +10171,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum 0.8.9",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"h2 0.4.16",
|
||||
@@ -10095,9 +10183,11 @@ dependencies = [
|
||||
"hyper-util",
|
||||
"percent-encoding",
|
||||
"pin-project",
|
||||
"rustls-native-certs",
|
||||
"socket2 0.6.3",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls 0.26.4",
|
||||
"tokio-stream",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
@@ -10452,9 +10542,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.24.0"
|
||||
version = "1.26.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239"
|
||||
checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812"
|
||||
dependencies = [
|
||||
"getrandom 0.4.2",
|
||||
"js-sys",
|
||||
|
||||
+17
-15
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=12.0.0-beta.5", default-features = false, "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=12.0.0-beta.5", default-features = false, "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=12.0.0-beta.5", default-features = false, "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=12.0.0-beta.5", "tag" = "v12.0.0-beta.5", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=12.0.0-beta.14", default-features = false, "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=12.0.0-beta.14", default-features = false, "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=12.0.0-beta.14", default-features = false, "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=12.0.0-beta.14", "tag" = "v12.0.0-beta.14", "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
|
||||
@@ -39,6 +39,7 @@ arrow-ord = "58.0.0"
|
||||
arrow-schema = "58.0.0"
|
||||
arrow-select = "58.0.0"
|
||||
arrow-cast = "58.0.0"
|
||||
arrow-flight = { version = "58.0.0", features = ["flight-sql-experimental"] }
|
||||
async-trait = "0"
|
||||
bytes = "1"
|
||||
datafusion = { version = "54.0.0", default-features = false }
|
||||
@@ -71,7 +72,8 @@ serde = "1"
|
||||
serde_json = "1"
|
||||
tempfile = "3.5.0"
|
||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||
uuid = { version = "1.7.0", features = ["v4"] }
|
||||
tonic = { version = "0.14", features = ["tls-native-roots", "tls-ring"] }
|
||||
uuid = { version = "1.7.0", features = ["v4", "v7"] }
|
||||
chrono = { version = "0.4", default-features = false, features = ["clock"] }
|
||||
|
||||
[profile.ci]
|
||||
|
||||
@@ -446,6 +446,15 @@ paths:
|
||||
properties:
|
||||
column:
|
||||
type: string
|
||||
name:
|
||||
type: string
|
||||
description: Optional name for the created index.
|
||||
replace:
|
||||
type: boolean
|
||||
default: true
|
||||
description: |
|
||||
Whether to replace an existing index with the same resolved
|
||||
name. Defaults to true.
|
||||
metric_type:
|
||||
type: string
|
||||
nullable: false
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.38.0-beta.15</version>
|
||||
<version>0.39.0-beta.4</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -448,26 +448,6 @@ on the returned job to know when cleanup has finished.
|
||||
|
||||
***
|
||||
|
||||
### getJob()
|
||||
|
||||
```ts
|
||||
abstract getJob(jobId): Promise<null | JobDescription>
|
||||
```
|
||||
|
||||
Describe a single server-side job by id.
|
||||
|
||||
Resolves to `null` when the server has no such job.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`null` \| [`JobDescription`](../interfaces/JobDescription.md)>
|
||||
|
||||
***
|
||||
|
||||
### isOpen()
|
||||
|
||||
```ts
|
||||
@@ -482,48 +462,6 @@ Return true if the connection has not been closed
|
||||
|
||||
***
|
||||
|
||||
### job()
|
||||
|
||||
```ts
|
||||
abstract job(jobId): Job
|
||||
```
|
||||
|
||||
A [Job](Job.md) handle for a server-side job by id.
|
||||
|
||||
The handle is constructed without a server round trip; an unknown id
|
||||
surfaces when the handle is used. Dropping the handle has no effect on
|
||||
the job itself.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
[`Job`](Job.md)
|
||||
|
||||
***
|
||||
|
||||
### jobHistory()
|
||||
|
||||
```ts
|
||||
abstract jobHistory(jobId?): Promise<Table<any>>
|
||||
```
|
||||
|
||||
The lifecycle event history of a server-side job, as an Arrow table.
|
||||
|
||||
Lists history across all jobs when `jobId` is omitted.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId?**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Table`<`any`>>
|
||||
|
||||
***
|
||||
|
||||
### listJobs()
|
||||
|
||||
```ts
|
||||
@@ -648,6 +586,30 @@ A page of table names and an
|
||||
|
||||
***
|
||||
|
||||
### openJob()
|
||||
|
||||
```ts
|
||||
abstract openJob(jobId): Promise<Job>
|
||||
```
|
||||
|
||||
Open a server-side job by id, returning a handle with its record already
|
||||
populated. Rejects when the server has no such job, the way
|
||||
[Connection.openTable](Connection.md#opentable) does for a missing table.
|
||||
|
||||
The returned [Job](Job.md) answers for its own state, specification,
|
||||
result, failure and event history, so there is no separate
|
||||
connection-level call for any of them.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **jobId**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`Job`](Job.md)>
|
||||
|
||||
***
|
||||
|
||||
### openMaterializedView()
|
||||
|
||||
```ts
|
||||
|
||||
+159
-12
@@ -8,19 +8,46 @@
|
||||
|
||||
A handle to an operation that may still be running.
|
||||
|
||||
## Constructors
|
||||
The operation may already be complete when the handle is created.
|
||||
|
||||
### new Job()
|
||||
The detail getters read what the handle last observed. Submitting an
|
||||
operation returns only a job id, so populating them eagerly would cost an
|
||||
extra round trip on every call:
|
||||
|
||||
- [Job.refresh](Job.md#refresh) and [Job.status](Job.md#status) fetch the whole record.
|
||||
- [Job.wait](Job.md#wait) records the terminal state it establishes, but not the
|
||||
rest of the record.
|
||||
- Everything is null until one of those runs.
|
||||
|
||||
## Accessors
|
||||
|
||||
### creationMs
|
||||
|
||||
```ts
|
||||
new Job(): Job
|
||||
get creationMs(): null | number
|
||||
```
|
||||
|
||||
When the job was created, in milliseconds since the epoch.
|
||||
|
||||
#### Returns
|
||||
|
||||
[`Job`](Job.md)
|
||||
`null` \| `number`
|
||||
|
||||
## Accessors
|
||||
***
|
||||
|
||||
### failure
|
||||
|
||||
```ts
|
||||
get failure(): null | JobFailureInfo
|
||||
```
|
||||
|
||||
Why the job failed, when it failed and the server reports a reason.
|
||||
|
||||
#### Returns
|
||||
|
||||
`null` \| [`JobFailureInfo`](../interfaces/JobFailureInfo.md)
|
||||
|
||||
***
|
||||
|
||||
### id
|
||||
|
||||
@@ -28,8 +55,69 @@ new Job(): Job
|
||||
get id(): null | string
|
||||
```
|
||||
|
||||
Identifies the operation on the server that is running it. Operations
|
||||
that run in this process have no server id. The value is opaque.
|
||||
Identifies the operation on the server that is running it.
|
||||
|
||||
Operations that run in this process have no server id. The value is
|
||||
opaque: parsing it or storing it to resume the job later is not supported.
|
||||
|
||||
#### Returns
|
||||
|
||||
`null` \| `string`
|
||||
|
||||
***
|
||||
|
||||
### jobType
|
||||
|
||||
```ts
|
||||
get jobType(): null | string
|
||||
```
|
||||
|
||||
The job's type, as the server names it. Null for an in-process job, which
|
||||
has no server-side record.
|
||||
|
||||
#### Returns
|
||||
|
||||
`null` \| `string`
|
||||
|
||||
***
|
||||
|
||||
### result
|
||||
|
||||
```ts
|
||||
get result(): any
|
||||
```
|
||||
|
||||
The job-type-specific terminal result. Null until the job succeeds, so a
|
||||
job that never terminates reports its progress through [Job.events](Job.md#events)
|
||||
instead.
|
||||
|
||||
#### Returns
|
||||
|
||||
`any`
|
||||
|
||||
***
|
||||
|
||||
### spec
|
||||
|
||||
```ts
|
||||
get spec(): any
|
||||
```
|
||||
|
||||
The job-type-specific specification it was submitted with.
|
||||
|
||||
#### Returns
|
||||
|
||||
`any`
|
||||
|
||||
***
|
||||
|
||||
### state
|
||||
|
||||
```ts
|
||||
get state(): null | string
|
||||
```
|
||||
|
||||
The last observed lifecycle state, without contacting the backend.
|
||||
|
||||
#### Returns
|
||||
|
||||
@@ -51,18 +139,61 @@ Request cancellation. Cancelling a finished operation is a no-op.
|
||||
|
||||
***
|
||||
|
||||
### events()
|
||||
|
||||
```ts
|
||||
events(options?): Promise<Table<any>>
|
||||
```
|
||||
|
||||
This job's recorded lifecycle events.
|
||||
|
||||
Where the getters above report a terminal result only once the job reaches
|
||||
one, events are written as the job runs and outlive the workers that
|
||||
produced them. A distributed job records a `claim`/`claim_complete` pair
|
||||
per unit of work, each carrying `rows_processed`, so a job that never
|
||||
finishes still accounts for what it did.
|
||||
|
||||
The server caps results at 1000 rows by default and 10,000 at most, and
|
||||
truncates without saying so, so pass `limit` for a job that emits an event
|
||||
per fragment. `filter` is a SQL-like expression over the `state`,
|
||||
`updated_by`, `emitted_from`, `emitted_by`, and `claim_entity` columns.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **options?**: [`JobEventsOptions`](../interfaces/JobEventsOptions.md)
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Table`<`any`>>
|
||||
|
||||
***
|
||||
|
||||
### refresh()
|
||||
|
||||
```ts
|
||||
refresh(): Promise<void>
|
||||
```
|
||||
|
||||
Ask the backend for this job's current state, and for a server-side job
|
||||
its full record, then cache it for the getters above.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`void`>
|
||||
|
||||
***
|
||||
|
||||
### status()
|
||||
|
||||
```ts
|
||||
status(): Promise<string>
|
||||
```
|
||||
|
||||
The operation's current lifecycle state: "running", "finished",
|
||||
"failed", or "cancelled".
|
||||
The operation's current lifecycle state: "running", "finished", "failed",
|
||||
or "cancelled".
|
||||
|
||||
A point snapshot; unlike [Job.wait](Job.md#wait) it does not block or reject
|
||||
on a terminal failure state. States a newer server reports that this
|
||||
client version does not know pass through as-is.
|
||||
A point snapshot; unlike [Job.wait](Job.md#wait) it does not block or reject on a
|
||||
terminal failure state. Also refreshes the getters above.
|
||||
|
||||
#### Returns
|
||||
|
||||
@@ -70,6 +201,22 @@ client version does not know pass through as-is.
|
||||
|
||||
***
|
||||
|
||||
### toString()
|
||||
|
||||
```ts
|
||||
toString(): string
|
||||
```
|
||||
|
||||
Every field the handle currently knows, one per line, with the JSON
|
||||
payloads indented -- a refresh job's spec and result are the point of
|
||||
printing it.
|
||||
|
||||
#### Returns
|
||||
|
||||
`string`
|
||||
|
||||
***
|
||||
|
||||
### wait()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -676,9 +676,17 @@ List all the versions of the table
|
||||
abstract mergeInsert(on): MergeInsertBuilder
|
||||
```
|
||||
|
||||
Create a [MergeInsertBuilder](MergeInsertBuilder.md), which combines new data with the
|
||||
existing table in a single transaction — inserting, updating and deleting
|
||||
rows depending on how they match.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **on**: `string` \| `string`[]
|
||||
The column, or columns, to match source rows against target
|
||||
rows on. Typically a key or id column. Several columns match on the
|
||||
composite key: a source row updates a target row only when it agrees on
|
||||
every one of them.
|
||||
|
||||
#### Returns
|
||||
|
||||
|
||||
@@ -96,7 +96,7 @@
|
||||
- [IvfFlatOptions](interfaces/IvfFlatOptions.md)
|
||||
- [IvfPqOptions](interfaces/IvfPqOptions.md)
|
||||
- [IvfRqOptions](interfaces/IvfRqOptions.md)
|
||||
- [JobDescription](interfaces/JobDescription.md)
|
||||
- [JobEventsOptions](interfaces/JobEventsOptions.md)
|
||||
- [JobFailureInfo](interfaces/JobFailureInfo.md)
|
||||
- [JobInfo](interfaces/JobInfo.md)
|
||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||
|
||||
@@ -1,66 +0,0 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / JobDescription
|
||||
|
||||
# Interface: JobDescription
|
||||
|
||||
A described job from `Connection.getJob`.
|
||||
|
||||
## Properties
|
||||
|
||||
### creationMs
|
||||
|
||||
```ts
|
||||
creationMs: number;
|
||||
```
|
||||
|
||||
When the job was created, in milliseconds since the epoch.
|
||||
|
||||
***
|
||||
|
||||
### failure?
|
||||
|
||||
```ts
|
||||
optional failure: JobFailureInfo;
|
||||
```
|
||||
|
||||
Why the job failed, when the job is failed and the server reports a
|
||||
reason.
|
||||
|
||||
***
|
||||
|
||||
### jobId
|
||||
|
||||
```ts
|
||||
jobId: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### jobType
|
||||
|
||||
```ts
|
||||
jobType: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### specJson?
|
||||
|
||||
```ts
|
||||
optional specJson: string;
|
||||
```
|
||||
|
||||
The job-type-specific specification as a JSON string, when present.
|
||||
|
||||
***
|
||||
|
||||
### state
|
||||
|
||||
```ts
|
||||
state: string;
|
||||
```
|
||||
|
||||
Lifecycle state: "running", "finished", "failed", or "cancelled".
|
||||
@@ -0,0 +1,29 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / JobEventsOptions
|
||||
|
||||
# Interface: JobEventsOptions
|
||||
|
||||
Which of a job's events [Job.events](../classes/Job.md#events) returns.
|
||||
|
||||
## Properties
|
||||
|
||||
### filter?
|
||||
|
||||
```ts
|
||||
optional filter: string;
|
||||
```
|
||||
|
||||
SQL-like filter over the event columns.
|
||||
|
||||
***
|
||||
|
||||
### limit?
|
||||
|
||||
```ts
|
||||
optional limit: number;
|
||||
```
|
||||
|
||||
Maximum event rows to return, up to the server maximum of 10,000.
|
||||
@@ -26,7 +26,7 @@ When the job was created, in milliseconds since the epoch.
|
||||
jobId: string;
|
||||
```
|
||||
|
||||
The job id -- what `Connection.getJob` and `Connection.cancelJob`
|
||||
The job id -- what `Connection.openJob` and `Connection.cancelJob`
|
||||
accept.
|
||||
|
||||
***
|
||||
|
||||
@@ -50,6 +50,16 @@ projections: [string, string][];
|
||||
|
||||
***
|
||||
|
||||
### sourceNamespace
|
||||
|
||||
```ts
|
||||
sourceNamespace: string[];
|
||||
```
|
||||
|
||||
Namespace holding the source table; empty is the root namespace.
|
||||
|
||||
***
|
||||
|
||||
### sourceTable
|
||||
|
||||
```ts
|
||||
|
||||
@@ -28,6 +28,59 @@ is also an [asynchronous API client](#connections-asynchronous).
|
||||
|
||||
::: lancedb.Session
|
||||
|
||||
## Remote SQL
|
||||
|
||||
Submit SQL against a remote LanceDB database through the connection.
|
||||
The connected database and `default_namespace_path=["public"]` are used for
|
||||
unqualified tables. Fully qualified references can still query other databases
|
||||
and namespaces available to the same deployment. `execute_query` returns a
|
||||
reader as soon as its initial result stream is available. `execute_query_async`
|
||||
returns a query handle immediately; use it to inspect progress, open a reader,
|
||||
or cancel the query. The SQL client is initialized by the first query and
|
||||
retained for the lifetime of the remote connection. Query ids are random,
|
||||
connection-scoped references rather than encoded SQL or durable resume tokens:
|
||||
|
||||
```python
|
||||
import lancedb
|
||||
|
||||
db = lancedb.connect(
|
||||
"db://analytics",
|
||||
api_key="ldb_...",
|
||||
host_override="https://api.example.com",
|
||||
sql_host_override="grpc+tls://sql.example.com:10026",
|
||||
)
|
||||
reader = db.execute_query(
|
||||
"""
|
||||
SELECT events.id, accounts.name
|
||||
FROM analytics.public.events AS events
|
||||
JOIN users.public.accounts AS accounts ON events.user_id = accounts.id
|
||||
""",
|
||||
default_namespace_path=["public"],
|
||||
)
|
||||
for batch in reader:
|
||||
print(batch.num_rows)
|
||||
|
||||
query = db.execute_query_async("SELECT * FROM events")
|
||||
print(query.id)
|
||||
print(query.describe().status)
|
||||
for batch in query.reader():
|
||||
print(batch.num_rows)
|
||||
|
||||
# The async connection exposes the same lifecycle without blocking:
|
||||
# async_db = await lancedb.connect_async(
|
||||
# "db://analytics",
|
||||
# api_key="ldb_...",
|
||||
# host_override="https://api.example.com",
|
||||
# sql_host_override="grpc+tls://sql.example.com:10026",
|
||||
# )
|
||||
# reader = await async_db.execute_query("SELECT * FROM events")
|
||||
# query = await async_db.execute_query_async("SELECT * FROM events")
|
||||
# description = await async_db.describe_query(query.id)
|
||||
# async for batch in await query.reader():
|
||||
# print(batch.num_rows)
|
||||
# await query.cancel()
|
||||
```
|
||||
|
||||
## Namespaces (Synchronous)
|
||||
|
||||
A namespace-backed connection resolves tables through a
|
||||
@@ -72,6 +125,10 @@ listing a storage directory.
|
||||
|
||||
::: lancedb.functions.UdfDefinition
|
||||
|
||||
::: lancedb.secrets.EnvVarSecret
|
||||
|
||||
::: lancedb.secrets.SecretInfo
|
||||
|
||||
::: lancedb.functions.FunctionRegistrationRequest
|
||||
|
||||
::: lancedb.functions.FunctionArtifactRequest
|
||||
@@ -94,6 +151,8 @@ listing a storage directory.
|
||||
|
||||
::: lancedb.functions.OutputMapping
|
||||
|
||||
::: lancedb.functions.AssignmentMapping
|
||||
|
||||
::: lancedb.functions.FunctionBinding
|
||||
|
||||
::: lancedb.functions.RefreshColumnResult
|
||||
@@ -102,6 +161,18 @@ listing a storage directory.
|
||||
|
||||
::: lancedb.job.AsyncJob
|
||||
|
||||
::: lancedb.job.JobInfo
|
||||
|
||||
::: lancedb.job.JobDescription
|
||||
|
||||
::: lancedb.job.JobFailureInfo
|
||||
|
||||
::: lancedb.sql.Query
|
||||
|
||||
::: lancedb.sql.AsyncQuery
|
||||
|
||||
::: lancedb.sql.QueryDescription
|
||||
|
||||
## Materialized Views (Synchronous)
|
||||
|
||||
::: lancedb.materialized_view.MaterializedView
|
||||
@@ -249,6 +320,12 @@ still work. Queries return descriptors. Call
|
||||
|
||||
::: lancedb.exceptions.MissingColumnError
|
||||
|
||||
::: lancedb.exceptions.JobNotFoundError
|
||||
|
||||
::: lancedb.exceptions.JobFailedError
|
||||
|
||||
::: lancedb.exceptions.JobCancelledError
|
||||
|
||||
## Integrations
|
||||
|
||||
## Pydantic
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.15</version>
|
||||
<version>0.39.0-beta.4</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.15</version>
|
||||
<version>0.39.0-beta.4</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>${project.artifactId}</name>
|
||||
<description>LanceDB Java SDK Parent POM</description>
|
||||
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>12.0.0-beta.5</lance-core.version>
|
||||
<lance-core.version>12.0.0-beta.14</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>
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.38.0-beta.15"
|
||||
version = "0.39.0-beta.4"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
|
||||
@@ -48,6 +48,28 @@ describe("materialized views", () => {
|
||||
expect(definitionFromMetadata(safe, "v").limit).toBe(42);
|
||||
});
|
||||
|
||||
it("reads the namespaced select kind and refuses unknown kinds", () => {
|
||||
// "namespaced_select" is the namespaced form of "select": same shape, a
|
||||
// separate kind so readers that predate it refuse instead of resolving
|
||||
// the source at the root.
|
||||
const namespaced = new Map([
|
||||
[
|
||||
DEFINITION_META_KEY,
|
||||
'{"kind":"namespaced_select","source_table":"people","source_namespace":["ns"]}',
|
||||
],
|
||||
]);
|
||||
const definition = definitionFromMetadata(namespaced, "v");
|
||||
expect(definition.sourceTable).toBe("people");
|
||||
expect(definition.sourceNamespace).toEqual(["ns"]);
|
||||
|
||||
const unknown = new Map([
|
||||
[DEFINITION_META_KEY, '{"kind":"select_v3","source_table":"people"}'],
|
||||
]);
|
||||
expect(() => definitionFromMetadata(unknown, "v")).toThrow(
|
||||
/cannot refresh/,
|
||||
);
|
||||
});
|
||||
|
||||
it("creates, refreshes and queries a view", async () => {
|
||||
const view = await db.createMaterializedView("adults", "people", {
|
||||
select: ["name", ["shout", "upper(name)"]],
|
||||
|
||||
@@ -939,6 +939,7 @@ describe("remote connection jobs surface", () => {
|
||||
const { tableFromArrays, tableToIPC } = await import("apache-arrow");
|
||||
const eventsTable = tableFromArrays({ state: ["created", "succeeded"] });
|
||||
const eventsBody = Buffer.from(tableToIPC(eventsTable, "stream"));
|
||||
const queryEventsPayloads: Record<string, unknown>[] = [];
|
||||
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
@@ -967,6 +968,16 @@ describe("remote connection jobs surface", () => {
|
||||
);
|
||||
}
|
||||
} else if (req.url === "/v1/jobs/describe") {
|
||||
if (payload["job_id"] === "job-2") {
|
||||
res
|
||||
.writeHead(200, { "Content-Type": "application/json" })
|
||||
.end(
|
||||
'{"job_id": "job-2", "job_type": "refresh_column", ' +
|
||||
'"job_state": "DONE", "creation_ms": 2000, ' +
|
||||
'"result": {"rows_assigned": 1000000}}',
|
||||
);
|
||||
return;
|
||||
}
|
||||
if (payload["job_id"] !== "job-1") {
|
||||
res.writeHead(404).end("no such job");
|
||||
return;
|
||||
@@ -988,6 +999,7 @@ describe("remote connection jobs surface", () => {
|
||||
.writeHead(200, { "Content-Type": "application/json" })
|
||||
.end('{"job_id": "job-1"}');
|
||||
} else if (req.url === "/v1/jobs/query_events") {
|
||||
queryEventsPayloads.push(payload);
|
||||
res
|
||||
.writeHead(200, {
|
||||
"Content-Type": "application/vnd.apache.arrow.stream",
|
||||
@@ -1004,22 +1016,65 @@ describe("remote connection jobs surface", () => {
|
||||
expect(jobs[0].state).toEqual("running");
|
||||
expect(jobs[1].state).toEqual("finished");
|
||||
|
||||
const description = await db.getJob("job-1");
|
||||
expect(description?.state).toEqual("failed");
|
||||
expect(JSON.parse(description?.specJson ?? "")).toEqual({
|
||||
column: "vec",
|
||||
});
|
||||
expect(description?.failure?.message).toEqual("worker died");
|
||||
expect(await db.getJob("missing")).toBeNull();
|
||||
|
||||
expect(await db.cancelJob("job-1")).toBe(true);
|
||||
expect(await db.cancelJob("missing")).toBe(false);
|
||||
|
||||
const history = await db.jobHistory("job-1");
|
||||
expect(history.numRows).toEqual(2);
|
||||
// Opening a job hands back a populated handle; a missing one rejects.
|
||||
await expect(db.openJob("missing")).rejects.toThrow("not found");
|
||||
const finished = await db.openJob("job-2");
|
||||
expect(finished.state).toEqual("finished");
|
||||
expect(finished.result).toEqual({
|
||||
// biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
|
||||
rows_assigned: 1000000,
|
||||
});
|
||||
|
||||
const job = db.job("job-1");
|
||||
const job = await db.openJob("job-1");
|
||||
expect(job.id).toEqual("job-1");
|
||||
|
||||
// openJob already populated the handle; refresh() re-reads it.
|
||||
expect(job.state).toEqual("failed");
|
||||
await job.refresh();
|
||||
expect(job.state).toEqual("failed");
|
||||
expect(job.jobType).toEqual("create_index");
|
||||
expect(job.creationMs).toEqual(1000);
|
||||
expect(job.spec).toEqual({ column: "vec" });
|
||||
expect(job.result).toBeNull();
|
||||
expect(job.failure?.message).toEqual("worker died");
|
||||
|
||||
// The handle reaches its own events, supplying its job id.
|
||||
const jobEvents = await job.events({
|
||||
limit: 500,
|
||||
filter: "state = 'claim_complete'",
|
||||
});
|
||||
expect(jobEvents.numRows).toEqual(2);
|
||||
expect(queryEventsPayloads.pop()).toEqual({
|
||||
// biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
|
||||
job_id: "job-1",
|
||||
limit: 500,
|
||||
filter: "state = 'claim_complete'",
|
||||
});
|
||||
|
||||
// Printing lays every known field out on its own line, with the JSON
|
||||
// payloads indented rather than crammed onto one line.
|
||||
expect(`${job}`).toEqual(
|
||||
[
|
||||
"Job(",
|
||||
' id="job-1",',
|
||||
' state="failed",',
|
||||
' jobType="create_index",',
|
||||
" creationMs=1000,",
|
||||
" spec={",
|
||||
' "column": "vec"',
|
||||
" },",
|
||||
" failure={",
|
||||
' "phase": "execute",',
|
||||
' "message": "worker died",',
|
||||
' "retryable": true',
|
||||
" },",
|
||||
")",
|
||||
].join("\n"),
|
||||
);
|
||||
|
||||
expect(await job.status()).toEqual("failed");
|
||||
await expect(job.wait()).rejects.toThrow("worker died");
|
||||
},
|
||||
|
||||
@@ -737,11 +737,12 @@ it("should query documents with LangChain PDF metadata", async () => {
|
||||
|
||||
describe("merge insert", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let conn: Connection;
|
||||
let table: Table;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const conn = await connect(tmpDir.name);
|
||||
conn = await connect(tmpDir.name);
|
||||
|
||||
table = await conn.createTable("some_table", [
|
||||
{ a: 1, b: "a" },
|
||||
@@ -779,6 +780,38 @@ describe("merge insert", () => {
|
||||
|
||||
expect(result.map((row) => ({ ...row }))).toEqual(expected);
|
||||
});
|
||||
test("upsert on a composite key", async () => {
|
||||
const composite = await conn.createTable("composite", [
|
||||
{ shard: "a", id: 1, val: "x" },
|
||||
{ shard: "a", id: 2, val: "y" },
|
||||
{ shard: "b", id: 1, val: "z" },
|
||||
]);
|
||||
|
||||
// ("a", 1) matches an existing row and updates it. ("b", 2) agrees with an
|
||||
// existing row on each key column separately but on neither pair, so it is
|
||||
// an insert.
|
||||
const mergeInsertRes = await composite
|
||||
.mergeInsert(["shard", "id"])
|
||||
.whenMatchedUpdateAll()
|
||||
.whenNotMatchedInsertAll()
|
||||
.execute([
|
||||
{ shard: "a", id: 1, val: "X" },
|
||||
{ shard: "b", id: 2, val: "W" },
|
||||
]);
|
||||
expect(mergeInsertRes.numUpdatedRows).toBe(1);
|
||||
expect(mergeInsertRes.numInsertedRows).toBe(1);
|
||||
|
||||
const result = (await composite.toArrow())
|
||||
.toArray()
|
||||
.sort((a, b) => a.shard.localeCompare(b.shard) || a.id - b.id);
|
||||
|
||||
expect(result.map((row) => ({ ...row }))).toEqual([
|
||||
{ shard: "a", id: 1, val: "X" },
|
||||
{ shard: "a", id: 2, val: "y" },
|
||||
{ shard: "b", id: 1, val: "z" },
|
||||
{ shard: "b", id: 2, val: "W" },
|
||||
]);
|
||||
});
|
||||
test("conditional update", async () => {
|
||||
const newData = [
|
||||
{ a: 2, b: "x" },
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { tableFromIPC } from "apache-arrow";
|
||||
import {
|
||||
Data,
|
||||
SchemaLike,
|
||||
@@ -16,6 +15,7 @@ import {
|
||||
makeEmptyTable,
|
||||
} from "./arrow";
|
||||
import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
|
||||
import { Job } from "./job";
|
||||
import {
|
||||
MaterializedView,
|
||||
MaterializedViewSelect,
|
||||
@@ -27,8 +27,6 @@ import type {
|
||||
CreateNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
Job,
|
||||
JobDescription,
|
||||
JobInfo,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
@@ -557,24 +555,19 @@ export abstract class Connection {
|
||||
): Promise<void>;
|
||||
|
||||
/**
|
||||
* A {@link Job} handle for a server-side job by id.
|
||||
* Open a server-side job by id, returning a handle with its record already
|
||||
* populated. Rejects when the server has no such job, the way
|
||||
* {@link Connection.openTable} does for a missing table.
|
||||
*
|
||||
* The handle is constructed without a server round trip; an unknown id
|
||||
* surfaces when the handle is used. Dropping the handle has no effect on
|
||||
* the job itself.
|
||||
* The returned {@link Job} answers for its own state, specification,
|
||||
* result, failure and event history, so there is no separate
|
||||
* connection-level call for any of them.
|
||||
*/
|
||||
abstract job(jobId: string): Job;
|
||||
abstract openJob(jobId: string): Promise<Job>;
|
||||
|
||||
/** List server-side jobs across the database's tables. */
|
||||
abstract listJobs(): Promise<JobInfo[]>;
|
||||
|
||||
/**
|
||||
* Describe a single server-side job by id.
|
||||
*
|
||||
* Resolves to `null` when the server has no such job.
|
||||
*/
|
||||
abstract getJob(jobId: string): Promise<JobDescription | null>;
|
||||
|
||||
/**
|
||||
* Request cancellation of a server-side job by id.
|
||||
*
|
||||
@@ -582,13 +575,6 @@ export abstract class Connection {
|
||||
* such job exists. Cancelling an already-terminal job is a no-op success.
|
||||
*/
|
||||
abstract cancelJob(jobId: string): Promise<boolean>;
|
||||
|
||||
/**
|
||||
* The lifecycle event history of a server-side job, as an Arrow table.
|
||||
*
|
||||
* Lists history across all jobs when `jobId` is omitted.
|
||||
*/
|
||||
abstract jobHistory(jobId?: string): Promise<ArrowTable>;
|
||||
}
|
||||
|
||||
/** @hideconstructor */
|
||||
@@ -869,7 +855,7 @@ export class LocalConnection extends Connection {
|
||||
}
|
||||
|
||||
async dropTableAsync(name: string, namespacePath?: string[]): Promise<Job> {
|
||||
return this.inner.dropTableAsync(name, namespacePath ?? []);
|
||||
return new Job(await this.inner.dropTableAsync(name, namespacePath ?? []));
|
||||
}
|
||||
|
||||
async dropAllTables(namespacePath?: string[]): Promise<void> {
|
||||
@@ -928,29 +914,17 @@ export class LocalConnection extends Connection {
|
||||
);
|
||||
}
|
||||
|
||||
job(jobId: string): Job {
|
||||
return this.inner.job(jobId);
|
||||
async openJob(jobId: string): Promise<Job> {
|
||||
return new Job(await this.inner.openJob(jobId));
|
||||
}
|
||||
|
||||
async listJobs(): Promise<JobInfo[]> {
|
||||
return this.inner.listJobs();
|
||||
}
|
||||
|
||||
async getJob(jobId: string): Promise<JobDescription | null> {
|
||||
return this.inner.getJob(jobId);
|
||||
}
|
||||
|
||||
async cancelJob(jobId: string): Promise<boolean> {
|
||||
return this.inner.cancelJob(jobId);
|
||||
}
|
||||
|
||||
async jobHistory(jobId?: string): Promise<ArrowTable> {
|
||||
const buf = await this.inner.jobHistory(jobId);
|
||||
if (buf.length === 0) {
|
||||
return new ArrowTable();
|
||||
}
|
||||
return tableFromIPC(buf);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -94,13 +94,9 @@ export {
|
||||
RenameTableOptions,
|
||||
} from "./connection";
|
||||
|
||||
export {
|
||||
Job,
|
||||
JobDescription,
|
||||
JobFailureInfo,
|
||||
JobInfo,
|
||||
Session,
|
||||
} from "./native.js";
|
||||
export { JobFailureInfo, JobInfo, Session } from "./native.js";
|
||||
|
||||
export { Job, JobEventsOptions } from "./job";
|
||||
|
||||
export {
|
||||
AutoQuery,
|
||||
|
||||
@@ -0,0 +1,188 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { Table as ArrowTable, tableFromIPC } from "apache-arrow";
|
||||
import { JobFailureInfo, Job as NativeJob } from "./native";
|
||||
|
||||
/** Which of a job's events {@link Job.events} returns. */
|
||||
export interface JobEventsOptions {
|
||||
/** Maximum event rows to return, up to the server maximum of 10,000. */
|
||||
limit?: number;
|
||||
/** SQL-like filter over the event columns. */
|
||||
filter?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* A handle to an operation that may still be running.
|
||||
*
|
||||
* The operation may already be complete when the handle is created.
|
||||
*
|
||||
* The detail getters read what the handle last observed. Submitting an
|
||||
* operation returns only a job id, so populating them eagerly would cost an
|
||||
* extra round trip on every call:
|
||||
*
|
||||
* - {@link Job.refresh} and {@link Job.status} fetch the whole record.
|
||||
* - {@link Job.wait} records the terminal state it establishes, but not the
|
||||
* rest of the record.
|
||||
* - Everything is null until one of those runs.
|
||||
*
|
||||
* @hideconstructor
|
||||
*/
|
||||
export class Job {
|
||||
private readonly inner: NativeJob;
|
||||
|
||||
constructor(inner: NativeJob) {
|
||||
this.inner = inner;
|
||||
}
|
||||
|
||||
/**
|
||||
* Identifies the operation on the server that is running it.
|
||||
*
|
||||
* Operations that run in this process have no server id. The value is
|
||||
* opaque: parsing it or storing it to resume the job later is not supported.
|
||||
*/
|
||||
get id(): string | null {
|
||||
return this.inner.id ?? null;
|
||||
}
|
||||
|
||||
/** The last observed lifecycle state, without contacting the backend. */
|
||||
get state(): string | null {
|
||||
return this.inner.state ?? null;
|
||||
}
|
||||
|
||||
/**
|
||||
* The job's type, as the server names it. Null for an in-process job, which
|
||||
* has no server-side record.
|
||||
*/
|
||||
get jobType(): string | null {
|
||||
return this.inner.jobType ?? null;
|
||||
}
|
||||
|
||||
/** When the job was created, in milliseconds since the epoch. */
|
||||
get creationMs(): number | null {
|
||||
return this.inner.creationMs ?? null;
|
||||
}
|
||||
|
||||
/** The job-type-specific specification it was submitted with. */
|
||||
// biome-ignore lint/suspicious/noExplicitAny: shape varies by job type
|
||||
get spec(): any | null {
|
||||
return parseJson(this.inner.specJson);
|
||||
}
|
||||
|
||||
/**
|
||||
* The job-type-specific terminal result. Null until the job succeeds, so a
|
||||
* job that never terminates reports its progress through {@link Job.events}
|
||||
* instead.
|
||||
*/
|
||||
// biome-ignore lint/suspicious/noExplicitAny: shape varies by job type
|
||||
get result(): any | null {
|
||||
return parseJson(this.inner.resultJson);
|
||||
}
|
||||
|
||||
/** Why the job failed, when it failed and the server reports a reason. */
|
||||
get failure(): JobFailureInfo | null {
|
||||
return this.inner.failure ?? null;
|
||||
}
|
||||
|
||||
/**
|
||||
* The operation's current lifecycle state: "running", "finished", "failed",
|
||||
* or "cancelled".
|
||||
*
|
||||
* A point snapshot; unlike {@link Job.wait} it does not block or reject on a
|
||||
* terminal failure state. Also refreshes the getters above.
|
||||
*/
|
||||
async status(): Promise<string> {
|
||||
return this.inner.status();
|
||||
}
|
||||
|
||||
/** Wait until the operation reaches a terminal state. */
|
||||
async wait(): Promise<void> {
|
||||
return this.inner.wait();
|
||||
}
|
||||
|
||||
/** Request cancellation. Cancelling a finished operation is a no-op. */
|
||||
async cancel(): Promise<void> {
|
||||
return this.inner.cancel();
|
||||
}
|
||||
|
||||
/**
|
||||
* Ask the backend for this job's current state, and for a server-side job
|
||||
* its full record, then cache it for the getters above.
|
||||
*/
|
||||
async refresh(): Promise<void> {
|
||||
return this.inner.refresh();
|
||||
}
|
||||
|
||||
/**
|
||||
* This job's recorded lifecycle events.
|
||||
*
|
||||
* Where the getters above report a terminal result only once the job reaches
|
||||
* one, events are written as the job runs and outlive the workers that
|
||||
* produced them. A distributed job records a `claim`/`claim_complete` pair
|
||||
* per unit of work, each carrying `rows_processed`, so a job that never
|
||||
* finishes still accounts for what it did.
|
||||
*
|
||||
* The server caps results at 1000 rows by default and 10,000 at most, and
|
||||
* truncates without saying so, so pass `limit` for a job that emits an event
|
||||
* per fragment. `filter` is a SQL-like expression over the `state`,
|
||||
* `updated_by`, `emitted_from`, `emitted_by`, and `claim_entity` columns.
|
||||
*/
|
||||
async events(options?: JobEventsOptions): Promise<ArrowTable> {
|
||||
const buf = await this.inner.events(options?.limit, options?.filter);
|
||||
if (buf.length === 0) {
|
||||
return new ArrowTable();
|
||||
}
|
||||
return tableFromIPC(buf);
|
||||
}
|
||||
|
||||
/**
|
||||
* Every field the handle currently knows, one per line, with the JSON
|
||||
* payloads indented -- a refresh job's spec and result are the point of
|
||||
* printing it.
|
||||
*/
|
||||
toString(): string {
|
||||
if (this.state === null) {
|
||||
const known = this.id === null ? "" : `id=${JSON.stringify(this.id)}, `;
|
||||
return `Job(${known}not refreshed)`;
|
||||
}
|
||||
const fields: string[] = [];
|
||||
if (this.id !== null) {
|
||||
fields.push(`id=${JSON.stringify(this.id)}`);
|
||||
}
|
||||
fields.push(`state=${JSON.stringify(this.state)}`);
|
||||
if (this.jobType !== null) {
|
||||
fields.push(`jobType=${JSON.stringify(this.jobType)}`);
|
||||
}
|
||||
if (this.creationMs !== null) {
|
||||
fields.push(`creationMs=${this.creationMs}`);
|
||||
}
|
||||
for (const [name, value] of [
|
||||
["spec", this.spec],
|
||||
["result", this.result],
|
||||
] as const) {
|
||||
if (value !== null) {
|
||||
fields.push(`${name}=${indentJson(value)}`);
|
||||
}
|
||||
}
|
||||
if (this.failure !== null) {
|
||||
fields.push(`failure=${indentJson(this.failure)}`);
|
||||
}
|
||||
return `Job(${fields.map((field) => `\n${REPR_INDENT}${field},`).join("")}\n)`;
|
||||
}
|
||||
|
||||
[Symbol.for("nodejs.util.inspect.custom")](): string {
|
||||
return this.toString();
|
||||
}
|
||||
}
|
||||
|
||||
const REPR_INDENT = " ";
|
||||
|
||||
// biome-ignore lint/suspicious/noExplicitAny: shape varies by job type
|
||||
function indentJson(value: any): string {
|
||||
return JSON.stringify(value, null, 4).replace(/\n/g, `\n${REPR_INDENT}`);
|
||||
}
|
||||
|
||||
// biome-ignore lint/suspicious/noExplicitAny: shape varies by job type
|
||||
function parseJson(raw: string | null | undefined): any | null {
|
||||
return raw === null || raw === undefined ? null : JSON.parse(raw);
|
||||
}
|
||||
@@ -19,6 +19,8 @@ export interface MaterializedViewDefinition {
|
||||
limit?: number;
|
||||
/** Source columns the projections and filter read. */
|
||||
inputs: string[];
|
||||
/** Namespace holding the source table; empty is the root namespace. */
|
||||
sourceNamespace: string[];
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -78,7 +80,8 @@ export function definitionFromMetadata(
|
||||
}
|
||||
// biome-ignore lint/suspicious/noExplicitAny: raw JSON
|
||||
const value: any = JSON.parse(raw);
|
||||
if (value.kind !== "select") {
|
||||
// "namespaced_select" keeps older readers from resolving the source at root.
|
||||
if (value.kind !== "select" && value.kind !== "namespaced_select") {
|
||||
throw new Error(
|
||||
`materialized view '${name}' is defined by '${value.kind}', which this ` +
|
||||
"version of lancedb cannot refresh",
|
||||
@@ -103,6 +106,7 @@ export function definitionFromMetadata(
|
||||
filter: value.filter ?? undefined,
|
||||
limit,
|
||||
inputs: value.inputs ?? [],
|
||||
sourceNamespace: value.source_namespace ?? [],
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
+21
-9
@@ -19,6 +19,7 @@ import {
|
||||
|
||||
import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
|
||||
import { IndexOptions } from "./indices";
|
||||
import { Job } from "./job";
|
||||
import { MergeInsertBuilder } from "./merge";
|
||||
import {
|
||||
AddColumnsResult,
|
||||
@@ -30,7 +31,6 @@ import {
|
||||
DropColumnsResult,
|
||||
IndexConfig,
|
||||
IndexStatistics,
|
||||
Job,
|
||||
LsmStats,
|
||||
Branches as NativeBranches,
|
||||
OptimizeStats,
|
||||
@@ -919,6 +919,16 @@ export abstract class Table {
|
||||
/** Return the table as an arrow table */
|
||||
abstract toArrow(): Promise<ArrowTable>;
|
||||
|
||||
/**
|
||||
* Create a {@link MergeInsertBuilder}, which combines new data with the
|
||||
* existing table in a single transaction — inserting, updating and deleting
|
||||
* rows depending on how they match.
|
||||
*
|
||||
* @param on - The column, or columns, to match source rows against target
|
||||
* rows on. Typically a key or id column. Several columns match on the
|
||||
* composite key: a source row updates a target row only when it agrees on
|
||||
* every one of them.
|
||||
*/
|
||||
abstract mergeInsert(on: string | string[]): MergeInsertBuilder;
|
||||
|
||||
/** List all the stats of a specified index
|
||||
@@ -1114,13 +1124,15 @@ export class LocalTable extends Table {
|
||||
): Promise<Job> {
|
||||
// biome-ignore lint/suspicious/noExplicitAny: skip
|
||||
const nativeIndex = (options?.config as any)?.inner;
|
||||
return await this.inner.createIndexAsync(
|
||||
nativeIndex,
|
||||
column,
|
||||
options?.replace,
|
||||
options?.waitTimeoutSeconds,
|
||||
options?.name,
|
||||
options?.train,
|
||||
return new Job(
|
||||
await this.inner.createIndexAsync(
|
||||
nativeIndex,
|
||||
column,
|
||||
options?.replace,
|
||||
options?.waitTimeoutSeconds,
|
||||
options?.name,
|
||||
options?.train,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1303,7 +1315,7 @@ export class LocalTable extends Table {
|
||||
}
|
||||
|
||||
async refreshColumnAsync(column: string): Promise<Job> {
|
||||
return await this.inner.refreshColumnAsync(column);
|
||||
return new Job(await this.inner.refreshColumnAsync(column));
|
||||
}
|
||||
|
||||
async refreshMaterializedView(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
+3
-3
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.38.0-beta.15",
|
||||
"version": "0.39.0-beta.4",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
@@ -44,7 +44,7 @@
|
||||
"@biomejs/biome": "^1.7.3",
|
||||
"@jest/globals": "^29.7.0",
|
||||
"@napi-rs/cli": "3.7.0",
|
||||
"@opentelemetry/sdk-metrics": "^1.30.0",
|
||||
"@opentelemetry/sdk-metrics": "^2.10.0",
|
||||
"@types/axios": "^0.14.0",
|
||||
"@types/jest": "^29.1.2",
|
||||
"@types/node": "22.7.4",
|
||||
@@ -56,7 +56,7 @@
|
||||
"eslint": "^8.57.0",
|
||||
"jest": "^29.7.0",
|
||||
"shx": "^0.3.4",
|
||||
"tmp": "^0.2.3",
|
||||
"tmp": "^0.2.7",
|
||||
"ts-jest": "^29.1.2",
|
||||
"typedoc": "0.26.4",
|
||||
"typedoc-plugin-markdown": "4.2.1",
|
||||
|
||||
Generated
+600
-444
File diff suppressed because it is too large
Load Diff
@@ -16,3 +16,41 @@ allowBuilds:
|
||||
onnxruntime-node: true
|
||||
protobufjs: true
|
||||
sharp: true
|
||||
|
||||
minimumReleaseAgeExclude:
|
||||
- protobufjs@7.5.8
|
||||
- tmp@0.2.6
|
||||
- form-data@4.0.6
|
||||
- tar@7.5.16
|
||||
- markdown-it@14.1.2
|
||||
- linkify-it@5.0.1
|
||||
- js-yaml@3.15.0
|
||||
- js-yaml@4.1.2
|
||||
- protobufjs@7.6.1
|
||||
- protobufjs@7.6.3
|
||||
- '@babel/core@7.29.1'
|
||||
- axios@1.18.0
|
||||
- brace-expansion@2.1.2
|
||||
- brace-expansion@1.1.16
|
||||
- js-yaml@4.3.0
|
||||
- tar@7.5.18
|
||||
- tar@7.5.19
|
||||
- tar@7.5.17
|
||||
- protobufjs@7.6.5
|
||||
- linkify-it@5.0.2
|
||||
- sharp@0.35.0
|
||||
- brace-expansion@1.1.17
|
||||
- brace-expansion@2.1.3
|
||||
- brace-expansion@2.1.4
|
||||
- brace-expansion@1.1.18
|
||||
- js-yaml@3.15.1
|
||||
- js-yaml@4.3.1
|
||||
- tar@7.5.21
|
||||
- '@opentelemetry/core@2.8.0'
|
||||
|
||||
# @huggingface/transformers pins sharp ^0.33.5 and no released version has moved
|
||||
# past ^0.34.5, all of which inherit the libvips CVEs in GHSA-f88m-g3jw-g9cj.
|
||||
# Force the patched line. sharp is only reached by transformers' image pipeline,
|
||||
# which LanceDB's text embedding function never uses.
|
||||
overrides:
|
||||
sharp: ^0.35.4
|
||||
|
||||
@@ -442,13 +442,15 @@ impl Connection {
|
||||
self.get_inner()?.drop_all_tables(&ns).await.default_error()
|
||||
}
|
||||
|
||||
/// A `Job` handle for a server-side job by id.
|
||||
/// Open a server-side job by id, returning a handle with its record
|
||||
/// already populated. Rejects when the server has no such job.
|
||||
///
|
||||
/// The handle is constructed without a server round trip; an unknown id
|
||||
/// surfaces when the handle is used.
|
||||
#[napi]
|
||||
pub fn job(&self, job_id: String) -> napi::Result<crate::job::Job> {
|
||||
let job = self.get_inner()?.job(job_id).default_error()?;
|
||||
/// The returned handle answers for its own state, specification, result,
|
||||
/// failure and event history, so there is no separate connection-level
|
||||
/// call for any of them.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn open_job(&self, job_id: String) -> napi::Result<crate::job::Job> {
|
||||
let job = self.get_inner()?.open_job(&job_id).await.default_error()?;
|
||||
Ok(crate::job::Job::new(job))
|
||||
}
|
||||
|
||||
@@ -459,17 +461,6 @@ impl Connection {
|
||||
Ok(jobs.into_iter().map(Into::into).collect())
|
||||
}
|
||||
|
||||
/// Describe a single server-side job by id. `null` when the server has
|
||||
/// no such job.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn get_job(
|
||||
&self,
|
||||
job_id: String,
|
||||
) -> napi::Result<Option<crate::job::JobDescription>> {
|
||||
let description = self.get_inner()?.get_job(&job_id).await.default_error()?;
|
||||
Ok(description.map(Into::into))
|
||||
}
|
||||
|
||||
/// Request cancellation of a server-side job by id. Returns true if the
|
||||
/// server accepted the cancellation, false if no such job exists.
|
||||
#[napi(catch_unwind)]
|
||||
@@ -477,34 +468,6 @@ impl Connection {
|
||||
self.get_inner()?.cancel_job(&job_id).await.default_error()
|
||||
}
|
||||
|
||||
/// The lifecycle event history of a server-side job (all jobs when
|
||||
/// `job_id` is null), as an Arrow IPC stream buffer. Empty when there is
|
||||
/// no history.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn job_history(&self, job_id: Option<String>) -> napi::Result<Buffer> {
|
||||
let batches = self
|
||||
.get_inner()?
|
||||
.job_history(job_id.as_deref())
|
||||
.await
|
||||
.default_error()?;
|
||||
let Some(first) = batches.first() else {
|
||||
return Ok(Buffer::from(Vec::<u8>::new()));
|
||||
};
|
||||
let mut out = Vec::new();
|
||||
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut out, &first.schema())
|
||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
||||
for batch in &batches {
|
||||
writer
|
||||
.write(batch)
|
||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
||||
}
|
||||
writer
|
||||
.finish()
|
||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
||||
drop(writer);
|
||||
Ok(Buffer::from(out))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
/// Describe a namespace and return its properties.
|
||||
pub async fn describe_namespace(
|
||||
|
||||
+90
-34
@@ -3,6 +3,9 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_array::RecordBatch;
|
||||
use lancedb::job::JobEventsRequest;
|
||||
use napi::bindgen_prelude::Buffer;
|
||||
use napi_derive::napi;
|
||||
|
||||
use crate::error::NapiErrorExt;
|
||||
@@ -55,12 +58,98 @@ impl Job {
|
||||
pub async fn cancel(&self) -> napi::Result<()> {
|
||||
self.inner.cancel().await.default_error()
|
||||
}
|
||||
|
||||
/// Ask the backend for this job's current state, and for a server-side job
|
||||
/// its full record, then cache it for the getters below.
|
||||
///
|
||||
/// They are all null until this runs, because submitting an operation
|
||||
/// returns only a job id. {@link Job.status} fetches the whole record too;
|
||||
/// {@link Job.wait} records only the terminal state it establishes.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn refresh(&self) -> napi::Result<()> {
|
||||
self.inner.refresh().await.default_error()
|
||||
}
|
||||
|
||||
/// The last observed lifecycle state, without contacting the backend.
|
||||
#[napi(getter)]
|
||||
pub fn state(&self) -> Option<String> {
|
||||
self.inner.state()
|
||||
}
|
||||
|
||||
/// The job's type, as the server names it. Null for an in-process job,
|
||||
/// which has no server-side record.
|
||||
#[napi(getter)]
|
||||
pub fn job_type(&self) -> Option<String> {
|
||||
self.inner.job_type()
|
||||
}
|
||||
|
||||
/// When the job was created, in milliseconds since the epoch.
|
||||
#[napi(getter)]
|
||||
pub fn creation_ms(&self) -> Option<i64> {
|
||||
self.inner.creation_ms()
|
||||
}
|
||||
|
||||
/// The job-type-specific specification as a JSON string, when present.
|
||||
#[napi(getter)]
|
||||
pub fn spec_json(&self) -> Option<String> {
|
||||
self.inner.spec().map(|spec| spec.to_string())
|
||||
}
|
||||
|
||||
/// The job-type-specific terminal result as a JSON string. Null until the
|
||||
/// job succeeds, so a job that never terminates reports its progress
|
||||
/// through {@link Job.events} instead.
|
||||
#[napi(getter)]
|
||||
pub fn result_json(&self) -> Option<String> {
|
||||
self.inner.result().map(|result| result.to_string())
|
||||
}
|
||||
|
||||
/// Why the job failed, when it failed and the server reports a reason.
|
||||
#[napi(getter)]
|
||||
pub fn failure(&self) -> Option<JobFailureInfo> {
|
||||
self.inner.failure().map(|failure| JobFailureInfo {
|
||||
phase: failure.phase,
|
||||
message: failure.message,
|
||||
retryable: failure.retryable,
|
||||
})
|
||||
}
|
||||
|
||||
/// This job's recorded lifecycle events, as an Arrow IPC stream buffer.
|
||||
/// The TypeScript wrapper turns it into an Arrow table.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn events(&self, limit: Option<u32>, filter: Option<String>) -> napi::Result<Buffer> {
|
||||
let batches = self
|
||||
.inner
|
||||
.events(JobEventsRequest { limit, filter })
|
||||
.await
|
||||
.default_error()?;
|
||||
batches_to_ipc_buffer(&batches)
|
||||
}
|
||||
}
|
||||
|
||||
/// Serialise Arrow batches as a single IPC stream for the TypeScript layer.
|
||||
fn batches_to_ipc_buffer(batches: &[RecordBatch]) -> napi::Result<Buffer> {
|
||||
let Some(first) = batches.first() else {
|
||||
return Ok(Buffer::from(Vec::<u8>::new()));
|
||||
};
|
||||
let mut out = Vec::new();
|
||||
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut out, &first.schema())
|
||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
||||
for batch in batches {
|
||||
writer
|
||||
.write(batch)
|
||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
||||
}
|
||||
writer
|
||||
.finish()
|
||||
.map_err(|e| napi::Error::from_reason(e.to_string()))?;
|
||||
drop(writer);
|
||||
Ok(Buffer::from(out))
|
||||
}
|
||||
|
||||
/// A row from `Connection.listJobs`: one server-side job.
|
||||
#[napi(object)]
|
||||
pub struct JobInfo {
|
||||
/// The job id -- what `Connection.getJob` and `Connection.cancelJob`
|
||||
/// The job id -- what `Connection.openJob` and `Connection.cancelJob`
|
||||
/// accept.
|
||||
pub job_id: String,
|
||||
/// The table the job runs against, without URI or namespace.
|
||||
@@ -91,36 +180,3 @@ pub struct JobFailureInfo {
|
||||
pub message: Option<String>,
|
||||
pub retryable: Option<bool>,
|
||||
}
|
||||
|
||||
/// A described job from `Connection.getJob`.
|
||||
#[napi(object)]
|
||||
pub struct JobDescription {
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
/// Lifecycle state: "running", "finished", "failed", or "cancelled".
|
||||
pub state: String,
|
||||
/// When the job was created, in milliseconds since the epoch.
|
||||
pub creation_ms: i64,
|
||||
/// The job-type-specific specification as a JSON string, when present.
|
||||
pub spec_json: Option<String>,
|
||||
/// Why the job failed, when the job is failed and the server reports a
|
||||
/// reason.
|
||||
pub failure: Option<JobFailureInfo>,
|
||||
}
|
||||
|
||||
impl From<lancedb::database::JobDescription> for JobDescription {
|
||||
fn from(description: lancedb::database::JobDescription) -> Self {
|
||||
Self {
|
||||
job_id: description.job_id,
|
||||
job_type: description.job_type,
|
||||
state: description.state,
|
||||
creation_ms: description.creation_ms,
|
||||
spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
|
||||
failure: description.failure.map(|failure| JobFailureInfo {
|
||||
phase: failure.phase,
|
||||
message: failure.message,
|
||||
retryable: failure.retryable,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+5
-1
@@ -664,7 +664,11 @@ impl JsFullTextQuery {
|
||||
}
|
||||
|
||||
fn parse_fts_query(query: Object) -> napi::Result<FullTextSearchQuery> {
|
||||
if let Ok(Some(query)) = query.get::<&JsFullTextQuery>("query") {
|
||||
// `&JsFullTextQuery` recovers a native class reference through napi's borrow-tracked
|
||||
// path, which is only usable from generated `#[napi]` argument conversion. This is a
|
||||
// manual lookup on a nested `Object` property instead, so use `ClassInstance`, which
|
||||
// unwraps the class without requiring a borrow scope.
|
||||
if let Ok(Some(query)) = query.get::<ClassInstance<JsFullTextQuery>>("query") {
|
||||
Ok(FullTextSearchQuery::new_query(query.inner.clone()))
|
||||
} else if let Ok(Some(query_text)) = query.get::<String>("query") {
|
||||
let mut query_text = query_text;
|
||||
|
||||
+3
-2
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.15"
|
||||
version = "0.39.0-beta.4"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
@@ -28,7 +28,7 @@ env_logger.workspace = true
|
||||
log.workspace = true
|
||||
# Maturin enables extension-module mode for Python builds. Keeping it out of
|
||||
# Cargo features lets Rust unit tests link against libpython.
|
||||
pyo3 = { version = "0.28", features = ["abi3-py310", "chrono"] }
|
||||
pyo3 = { version = "0.28", features = ["abi3-py310", "chrono", "uuid"] }
|
||||
chrono.workspace = true
|
||||
pyo3-async-runtimes = { version = "0.28", features = [
|
||||
"attributes",
|
||||
@@ -40,6 +40,7 @@ serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
snafu.workspace = true
|
||||
tokio.workspace = true
|
||||
uuid.workspace = true
|
||||
libc = "0.2"
|
||||
|
||||
[build-dependencies]
|
||||
|
||||
@@ -139,6 +139,7 @@ include = [
|
||||
"python/lancedb/exceptions.py",
|
||||
"python/lancedb/background_loop.py",
|
||||
"python/lancedb/schema.py",
|
||||
"python/lancedb/sql.py",
|
||||
"python/lancedb/remote/__init__.py",
|
||||
"python/lancedb/remote/errors.py",
|
||||
"python/lancedb/embeddings/__init__.py",
|
||||
|
||||
@@ -22,7 +22,11 @@ from .remote.db import RemoteDBConnection
|
||||
from .expr import Expr, col, lit, func
|
||||
from .schema import blob, vector
|
||||
from .job import AsyncJob, Job
|
||||
from .sql import AsyncQuery as AsyncSqlQuery
|
||||
from .sql import Query as SqlQuery
|
||||
from .sql import QueryDescription
|
||||
from .functions import (
|
||||
AssignmentMapping as AssignmentMapping,
|
||||
FunctionArtifactRequest as FunctionArtifactRequest,
|
||||
FunctionApplication as FunctionApplication,
|
||||
FunctionBinding as FunctionBinding,
|
||||
@@ -33,6 +37,8 @@ from .functions import (
|
||||
UdfDefinition as UdfDefinition,
|
||||
udf as udf,
|
||||
)
|
||||
from .secrets import EnvVarSecret as EnvVarSecret
|
||||
from .secrets import SecretInfo as SecretInfo
|
||||
from .materialized_view import (
|
||||
AsyncMaterializedView,
|
||||
MaterializedView,
|
||||
@@ -101,6 +107,7 @@ def connect(
|
||||
api_key: Optional[str] = None,
|
||||
region: str = "us-east-1",
|
||||
host_override: Optional[str] = None,
|
||||
sql_host_override: Optional[str] = None,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
request_thread_pool: Optional[Union[int, ThreadPoolExecutor]] = None,
|
||||
client_config: Union[ClientConfig, Dict[str, Any], None] = None,
|
||||
@@ -129,6 +136,9 @@ def connect(
|
||||
The region to use for LanceDB Cloud.
|
||||
host_override: str, optional
|
||||
The override url for LanceDB Cloud.
|
||||
sql_host_override: str, optional
|
||||
The remote SQL service endpoint override. The client connects lazily when SQL
|
||||
is first executed and retains that connection.
|
||||
read_consistency_interval: timedelta, default None
|
||||
The interval at which to check for updates to the table from other
|
||||
processes. If None, then consistency is not checked. For performance
|
||||
@@ -270,6 +280,7 @@ def connect(
|
||||
api_key,
|
||||
region,
|
||||
host_override,
|
||||
sql_host_override=sql_host_override,
|
||||
# TODO: remove this (deprecation warning downstream)
|
||||
request_thread_pool=request_thread_pool,
|
||||
client_config=client_config,
|
||||
@@ -412,6 +423,7 @@ def deserialize_conn(
|
||||
parsed["api_key"],
|
||||
parsed.get("region", "us-east-1"),
|
||||
host_override=parsed.get("host_override"),
|
||||
sql_host_override=parsed.get("sql_host_override"),
|
||||
client_config=parsed.get("client_config"),
|
||||
storage_options=storage_options,
|
||||
)
|
||||
@@ -425,6 +437,7 @@ async def connect_async(
|
||||
api_key: Optional[str] = None,
|
||||
region: str = "us-east-1",
|
||||
host_override: Optional[str] = None,
|
||||
sql_host_override: Optional[str] = None,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
client_config: Optional[Union[ClientConfig, Dict[str, Any]]] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
@@ -447,6 +460,9 @@ async def connect_async(
|
||||
The region to use for LanceDB Cloud.
|
||||
host_override: str, optional
|
||||
The override url for LanceDB Cloud.
|
||||
sql_host_override: str, optional
|
||||
The remote SQL service endpoint override. The client connects lazily when SQL
|
||||
is first executed and retains that connection.
|
||||
read_consistency_interval: timedelta, default None
|
||||
The interval at which to check for updates to the table from other
|
||||
processes. If None, then consistency is not checked. For performance
|
||||
@@ -534,6 +550,7 @@ async def connect_async(
|
||||
api_key,
|
||||
region,
|
||||
host_override,
|
||||
sql_host_override,
|
||||
read_consistency_interval_secs,
|
||||
client_config,
|
||||
storage_options,
|
||||
@@ -556,6 +573,7 @@ __all__ = [
|
||||
"connect_namespace_async",
|
||||
"AsyncConnection",
|
||||
"AsyncJob",
|
||||
"AsyncSqlQuery",
|
||||
"AsyncLanceNamespaceDBConnection",
|
||||
"AsyncTable",
|
||||
"FtsToken",
|
||||
@@ -570,6 +588,8 @@ __all__ = [
|
||||
"vector",
|
||||
"DBConnection",
|
||||
"Job",
|
||||
"QueryDescription",
|
||||
"SqlQuery",
|
||||
"LanceDBConnection",
|
||||
"LanceNamespaceDBConnection",
|
||||
"LsmWriteSpec",
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
||||
from uuid import UUID
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
@@ -147,15 +148,31 @@ class Connection(object):
|
||||
start_after: Optional[str],
|
||||
limit: Optional[int],
|
||||
) -> list[str]: ... # Deprecated: Use list_tables instead
|
||||
def job(self, job_id: str) -> Job: ...
|
||||
async def open_job(self, job_id: str) -> Job: ...
|
||||
async def create_function_async(self, request_json: str) -> Job: ...
|
||||
async def get_function(self, name: str, version: str) -> str: ...
|
||||
async def list_functions(self) -> List[str]: ...
|
||||
async def drop_function(self, name: str, version: str) -> bool: ...
|
||||
async def create_secret(
|
||||
self, name: str, value: str, namespace_path: List[str]
|
||||
) -> None: ...
|
||||
async def alter_secret(
|
||||
self, name: str, value: str, namespace_path: List[str]
|
||||
) -> None: ...
|
||||
async def list_secrets(self, namespace_path: List[str]) -> List[str]: ...
|
||||
async def drop_secret(self, name: str, namespace_path: List[str]) -> None: ...
|
||||
async def describe_secret(
|
||||
self, name: str, namespace_path: List[str]
|
||||
) -> Dict[str, str]: ...
|
||||
async def list_jobs(self) -> List[JobInfo]: ...
|
||||
async def get_job(self, job_id: str) -> Optional[JobDescription]: ...
|
||||
async def cancel_job(self, job_id: str) -> bool: ...
|
||||
async def job_history(
|
||||
self, job_id: Optional[str] = None
|
||||
) -> List[pa.RecordBatch]: ...
|
||||
async def execute_query_async(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> SqlQuery: ...
|
||||
async def describe_query(self, query_id: UUID) -> QueryDescription: ...
|
||||
async def create_table(
|
||||
self,
|
||||
name: str,
|
||||
@@ -234,9 +251,20 @@ class BlobFile:
|
||||
class Job:
|
||||
@property
|
||||
def id(self) -> Optional[str]: ...
|
||||
@property
|
||||
def _state(self) -> Optional[str]: ...
|
||||
@property
|
||||
def _description(self) -> Optional[JobDescription]: ...
|
||||
async def status(self) -> str: ...
|
||||
async def wait(self) -> Optional[str]: ...
|
||||
async def cancel(self) -> None: ...
|
||||
async def refresh(self) -> None: ...
|
||||
async def events(
|
||||
self,
|
||||
*,
|
||||
limit: Optional[int] = None,
|
||||
filter: Optional[str] = None,
|
||||
) -> pa.Table: ...
|
||||
|
||||
class JobInfo:
|
||||
@property
|
||||
@@ -268,10 +296,33 @@ class JobDescription:
|
||||
@property
|
||||
def creation_ms(self) -> int: ...
|
||||
@property
|
||||
def spec_json(self) -> Optional[str]: ...
|
||||
def _spec_json(self) -> Optional[str]: ...
|
||||
@property
|
||||
def _result_json(self) -> Optional[str]: ...
|
||||
@property
|
||||
def spec(self) -> Optional[Any]: ...
|
||||
@property
|
||||
def result(self) -> Optional[Any]: ...
|
||||
@property
|
||||
def failure(self) -> Optional[JobFailureInfo]: ...
|
||||
|
||||
class SqlQuery:
|
||||
@property
|
||||
def id(self) -> UUID: ...
|
||||
async def describe(self) -> QueryDescription: ...
|
||||
async def reader(self) -> RecordBatchStream: ...
|
||||
async def cancel(self) -> None: ...
|
||||
|
||||
class QueryDescription:
|
||||
@property
|
||||
def id(self) -> UUID: ...
|
||||
@property
|
||||
def status(self) -> str: ...
|
||||
@property
|
||||
def progress(self) -> Optional[float]: ...
|
||||
@property
|
||||
def expires_at(self) -> Optional[datetime]: ...
|
||||
|
||||
class Table:
|
||||
def name(self) -> str: ...
|
||||
def __repr__(self) -> str: ...
|
||||
@@ -450,6 +501,7 @@ async def connect(
|
||||
api_key: Optional[str],
|
||||
region: Optional[str],
|
||||
host_override: Optional[str],
|
||||
sql_host_override: Optional[str],
|
||||
read_consistency_interval: Optional[float],
|
||||
client_config: Optional[Union[ClientConfig, Dict[str, Any]]],
|
||||
storage_options: Optional[Dict[str, str]],
|
||||
@@ -606,6 +658,7 @@ class FullTextQuery:
|
||||
class PyQueryRequest:
|
||||
limit: Optional[int]
|
||||
offset: Optional[int]
|
||||
take_offsets: Optional[List[int]]
|
||||
filter: Optional[Union[str, bytes]]
|
||||
full_text_search: Optional[FullTextQuery]
|
||||
select: Optional[Union[str, List[str]]]
|
||||
|
||||
+351
-73
@@ -17,8 +17,10 @@ from typing import (
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
)
|
||||
from uuid import UUID
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override
|
||||
@@ -47,12 +49,21 @@ from . import __version__
|
||||
from ._lancedb import connect as lancedb_connect # type: ignore
|
||||
from .functions import FunctionVersion, UdfDefinition
|
||||
from .job import AsyncJob, Job, _typed_job
|
||||
from .sql import AsyncQuery as AsyncSqlQuery
|
||||
from .sql import Query as SqlQuery
|
||||
from .sql import QueryDescription
|
||||
from .materialized_view import (
|
||||
AsyncMaterializedView,
|
||||
MaterializedView,
|
||||
SelectArg,
|
||||
normalize_select,
|
||||
)
|
||||
from .secrets import (
|
||||
EnvVarSecret,
|
||||
SecretInfo,
|
||||
validate_namespace_path,
|
||||
validate_secret_name,
|
||||
)
|
||||
from .table import (
|
||||
AsyncTable,
|
||||
LanceTable,
|
||||
@@ -68,10 +79,11 @@ import deprecation
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pyarrow as pa
|
||||
from .arrow import AsyncRecordBatchReader
|
||||
from .pydantic import LanceModel
|
||||
|
||||
from ._lancedb import Connection as LanceDbConnection
|
||||
from ._lancedb import JobDescription, JobInfo
|
||||
from ._lancedb import JobInfo
|
||||
from .common import DATA, URI
|
||||
from .embeddings import EmbeddingFunctionConfig
|
||||
from ._lancedb import Session
|
||||
@@ -687,15 +699,47 @@ class DBConnection(EnforceOverrides):
|
||||
"""
|
||||
raise NotImplementedError("serialize is not supported for this connection type")
|
||||
|
||||
def create_function(self, definition: UdfDefinition) -> FunctionVersion:
|
||||
def create_function(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> FunctionVersion:
|
||||
"""Register a scalar Python UDF and wait for its immutable version.
|
||||
|
||||
This is the blocking counterpart of :meth:`create_function_async`.
|
||||
Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
return self.create_function_async(definition).wait()
|
||||
|
||||
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
|
||||
Parameters
|
||||
----------
|
||||
definition : UdfDefinition
|
||||
A callable decorated with [udf][lancedb.udf].
|
||||
secrets : sequence of EnvVarSecret, optional
|
||||
One [EnvVarSecret][lancedb.secrets.EnvVarSecret] per credential the
|
||||
Function needs, each naming a Secret and the environment variable
|
||||
its value arrives in. The Function's source is unchanged by this;
|
||||
it reads the variable the way it already did.
|
||||
|
||||
Examples
|
||||
--------
|
||||
```python
|
||||
db.create_secret("openai-prod", os.environ["OPENAI_API_KEY"])
|
||||
db.create_function(
|
||||
analyze_caption,
|
||||
secrets=[
|
||||
EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")
|
||||
],
|
||||
)
|
||||
```
|
||||
"""
|
||||
return self.create_function_async(definition, secrets=secrets).wait()
|
||||
|
||||
def create_function_async(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
"""Register a scalar Python UDF through the remote Function catalog.
|
||||
|
||||
Submission returns a typed job. The immutable Function version becomes
|
||||
@@ -712,26 +756,115 @@ class DBConnection(EnforceOverrides):
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def job(self, job_id: str) -> Job:
|
||||
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
|
||||
def list_functions(self) -> List[FunctionVersion]:
|
||||
"""List every published immutable Function version.
|
||||
|
||||
The handle is constructed without a server round trip; an unknown id
|
||||
surfaces when the handle is used. Dropping the handle has no effect
|
||||
on the job itself.
|
||||
Results are ordered by Function name then version. Local connections
|
||||
raise ``NotImplementedError``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
List the identities available to use in Function-backed columns:
|
||||
|
||||
```python
|
||||
[(function.name, function.version) for function in db.list_functions()]
|
||||
```
|
||||
"""
|
||||
raise NotImplementedError("job is not supported for this connection type")
|
||||
raise NotImplementedError(
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
"""Drop one exact immutable Function version from the remote catalog.
|
||||
|
||||
Returns True when the version changed to Dropped and False for an
|
||||
idempotent replay. Local connections raise NotImplementedError.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Function catalog operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Create a named Secret in this database.
|
||||
|
||||
Fails if the name is taken, so a create never silently becomes a
|
||||
rotation. Nothing reads the value back: it is bound to a Function by
|
||||
name and resolved by the service when that Function runs. Local
|
||||
connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Replace the credential behind an existing Secret.
|
||||
|
||||
Fails if it does not exist. Every Function bound to the Secret uses the
|
||||
new value from its next job, and no new Function version is created --
|
||||
which is how a rotation reaches columns pinned to a version registered
|
||||
before it. Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
"""The names of every Secret in this database.
|
||||
|
||||
Names only. No method returns a stored credential, by construction
|
||||
rather than by policy. Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Drop a Secret.
|
||||
|
||||
Functions bound to it fail at their next job, naming the Secret; that
|
||||
is the revocation path. The name becomes free to reuse, and a new
|
||||
Secret under it is picked up by everything still bound to that name.
|
||||
Local connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
"""What this database records about a Secret: name and timestamps.
|
||||
|
||||
Never the value -- there is no code path that could return one. Local
|
||||
connections raise ``NotImplementedError``.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Secret operations are not supported for this connection type"
|
||||
)
|
||||
|
||||
def open_job(self, job_id: str) -> Job:
|
||||
"""Open a server-side job by id, returning a handle with its record
|
||||
already populated.
|
||||
|
||||
The returned [Job][lancedb.job.Job] answers for its own state,
|
||||
specification, result, failure and event history, so there is no
|
||||
separate connection-level call for any of them.
|
||||
|
||||
Raises `JobNotFoundError` when the server has no such job, the way
|
||||
`open_table` does for a missing table.
|
||||
"""
|
||||
raise NotImplementedError("open_job is not supported for this connection type")
|
||||
|
||||
def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
raise NotImplementedError("list_jobs is not supported for this connection type")
|
||||
|
||||
def get_job(self, job_id: str) -> Optional[JobDescription]:
|
||||
"""Describe a single server-side job by id.
|
||||
|
||||
Returns None when the server has no such job.
|
||||
"""
|
||||
raise NotImplementedError("get_job is not supported for this connection type")
|
||||
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""Request cancellation of a server-side job by id.
|
||||
|
||||
@@ -743,14 +876,38 @@ class DBConnection(EnforceOverrides):
|
||||
"cancel_job is not supported for this connection type"
|
||||
)
|
||||
|
||||
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
||||
def execute_query(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> pa.RecordBatchReader:
|
||||
"""Execute SQL and return a blocking Arrow reader.
|
||||
|
||||
Lists history across all jobs when `job_id` is None.
|
||||
This submits through :meth:`execute_query_async` and waits until the
|
||||
initial result stream is readable. It does not wait for the full query
|
||||
to finish.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"job_history is not supported for this connection type"
|
||||
)
|
||||
return self.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
).reader()
|
||||
|
||||
def execute_query_async(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> SqlQuery:
|
||||
"""Start executing SQL and return its query handle.
|
||||
|
||||
Local connections do not support SQL.
|
||||
"""
|
||||
raise NotImplementedError("SQL is not supported for this connection type")
|
||||
|
||||
def describe_query(self, query_id: UUID) -> QueryDescription:
|
||||
"""Describe a submitted SQL query by its connection-scoped id."""
|
||||
raise NotImplementedError("SQL is not supported for this connection type")
|
||||
|
||||
|
||||
class LanceDBConnection(DBConnection):
|
||||
@@ -847,6 +1004,7 @@ class LanceDBConnection(DBConnection):
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
read_consistency_interval_secs,
|
||||
None,
|
||||
storage_options,
|
||||
@@ -1395,37 +1553,67 @@ class LanceDBConnection(DBConnection):
|
||||
)
|
||||
|
||||
@override
|
||||
def job(self, job_id: str) -> Job:
|
||||
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
|
||||
|
||||
The handle is constructed without a server round trip; an unknown id
|
||||
surfaces when the handle is used. Dropping the handle has no effect
|
||||
on the job itself.
|
||||
def open_job(self, job_id: str) -> Job:
|
||||
"""Open a server-side job by id. See
|
||||
[DBConnection.open_job][lancedb.db.DBConnection.open_job].
|
||||
"""
|
||||
return Job(self._conn.job(job_id))
|
||||
return Job(LOOP.run(self._conn.open_job(job_id)))
|
||||
|
||||
@override
|
||||
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition))
|
||||
def create_function_async(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
return LOOP.run(self._conn.get_function(name, version=version))
|
||||
|
||||
@override
|
||||
def list_functions(self) -> List[FunctionVersion]:
|
||||
return LOOP.run(self._conn.list_functions())
|
||||
|
||||
@override
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
|
||||
@override
|
||||
def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.drop_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return LOOP.run(self._conn.list_jobs())
|
||||
|
||||
@override
|
||||
def get_job(self, job_id: str) -> Optional[JobDescription]:
|
||||
"""Describe a single server-side job by id.
|
||||
|
||||
Returns None when the server has no such job.
|
||||
"""
|
||||
return LOOP.run(self._conn.get_job(job_id))
|
||||
|
||||
@override
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""Request cancellation of a server-side job by id.
|
||||
@@ -1436,14 +1624,6 @@ class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
return LOOP.run(self._conn.cancel_job(job_id))
|
||||
|
||||
@override
|
||||
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
||||
|
||||
Lists history across all jobs when `job_id` is None.
|
||||
"""
|
||||
return LOOP.run(self._conn.job_history(job_id))
|
||||
|
||||
@override
|
||||
def namespace_client(self) -> LanceNamespace:
|
||||
"""Get the equivalent namespace client for this connection.
|
||||
@@ -2214,46 +2394,110 @@ class AsyncConnection(object):
|
||||
namespace_path = []
|
||||
await self._inner.drop_all_tables(namespace_path=namespace_path)
|
||||
|
||||
def job(self, job_id: str) -> AsyncJob:
|
||||
"""An [AsyncJob][lancedb.job.AsyncJob] handle for a server-side job
|
||||
by id.
|
||||
|
||||
The handle is constructed without a server round trip; an unknown id
|
||||
surfaces when the handle is used. Dropping the handle has no effect
|
||||
on the job itself.
|
||||
async def open_job(self, job_id: str) -> AsyncJob:
|
||||
"""Open a server-side job by id. See
|
||||
[DBConnection.open_job][lancedb.db.DBConnection.open_job].
|
||||
"""
|
||||
return AsyncJob(self._inner.job(job_id))
|
||||
return AsyncJob(await self._inner.open_job(job_id))
|
||||
|
||||
async def create_function_async(
|
||||
self, definition: UdfDefinition
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> AsyncJob[FunctionVersion]:
|
||||
"""Register a scalar Python UDF through the remote Function catalog.
|
||||
|
||||
The returned typed job resolves to the immutable Function version.
|
||||
Local connections raise ``NotImplementedError``.
|
||||
``secrets`` is a sequence of
|
||||
[EnvVarSecret][lancedb.secrets.EnvVarSecret], each naming a Secret and
|
||||
the environment variable its value arrives in. Local connections raise
|
||||
``NotImplementedError``.
|
||||
"""
|
||||
if not isinstance(definition, UdfDefinition):
|
||||
raise TypeError("create_function_async requires a @udf definition")
|
||||
inner = await self._inner.create_function_async(
|
||||
definition.registration_request.to_canonical_json()
|
||||
)
|
||||
request = definition.bind_secrets(secrets)
|
||||
inner = await self._inner.create_function_async(request.to_canonical_json())
|
||||
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."""
|
||||
return FunctionVersion.from_json(await self._inner.get_function(name, version))
|
||||
|
||||
async def list_functions(self) -> List[FunctionVersion]:
|
||||
"""List every published immutable Function version.
|
||||
|
||||
Results are ordered by Function name then version. Local connections
|
||||
raise ``NotImplementedError``.
|
||||
"""
|
||||
return [
|
||||
FunctionVersion.from_json(value)
|
||||
for value in await self._inner.list_functions()
|
||||
]
|
||||
|
||||
async def drop_function(self, name: str, *, version: str) -> bool:
|
||||
"""Drop one exact immutable Function version from the remote catalog."""
|
||||
return await self._inner.drop_function(name, version)
|
||||
|
||||
async def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Create a named Secret in this database.
|
||||
|
||||
Fails if the name is taken, so a create never silently becomes a
|
||||
rotation. Nothing reads the value back.
|
||||
"""
|
||||
await self._inner.create_secret(
|
||||
validate_secret_name(name),
|
||||
value,
|
||||
list(validate_namespace_path(namespace_path)),
|
||||
)
|
||||
|
||||
async def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Replace the credential behind an existing Secret.
|
||||
|
||||
Fails if it does not exist. Bound Functions use the new value from
|
||||
their next job, with no new Function version.
|
||||
"""
|
||||
await self._inner.alter_secret(
|
||||
validate_secret_name(name),
|
||||
value,
|
||||
list(validate_namespace_path(namespace_path)),
|
||||
)
|
||||
|
||||
async def list_secrets(
|
||||
self, *, namespace_path: Optional[List[str]] = None
|
||||
) -> List[str]:
|
||||
"""The names of every Secret in this database. Names only."""
|
||||
return await self._inner.list_secrets(
|
||||
list(validate_namespace_path(namespace_path))
|
||||
)
|
||||
|
||||
async def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
"""Drop a Secret. Bound Functions fail at their next job."""
|
||||
await self._inner.drop_secret(
|
||||
validate_secret_name(name), list(validate_namespace_path(namespace_path))
|
||||
)
|
||||
|
||||
async def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
"""What this database records about a Secret. Never the value."""
|
||||
return SecretInfo.from_json(
|
||||
await self._inner.describe_secret(
|
||||
validate_secret_name(name),
|
||||
list(validate_namespace_path(namespace_path)),
|
||||
)
|
||||
)
|
||||
|
||||
async def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return await self._inner.list_jobs()
|
||||
|
||||
async def get_job(self, job_id: str) -> Optional[JobDescription]:
|
||||
"""Describe a single server-side job by id.
|
||||
|
||||
Returns None when the server has no such job.
|
||||
"""
|
||||
return await self._inner.get_job(job_id)
|
||||
|
||||
async def cancel_job(self, job_id: str) -> bool:
|
||||
"""Request cancellation of a server-side job by id.
|
||||
|
||||
@@ -2263,12 +2507,46 @@ class AsyncConnection(object):
|
||||
"""
|
||||
return await self._inner.cancel_job(job_id)
|
||||
|
||||
async def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
||||
async def execute_query(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> AsyncRecordBatchReader:
|
||||
"""Execute SQL and return an asynchronous Arrow reader.
|
||||
|
||||
Lists history across all jobs when `job_id` is None.
|
||||
This submits through :meth:`execute_query_async` and waits until the
|
||||
initial result stream is readable. It does not wait for the full query
|
||||
to finish.
|
||||
"""
|
||||
return await self._inner.job_history(job_id)
|
||||
submitted = await self.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
)
|
||||
return await submitted.reader()
|
||||
|
||||
async def execute_query_async(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> AsyncSqlQuery:
|
||||
"""Start executing SQL and return its query handle.
|
||||
|
||||
The database from ``connect_async`` is used for unqualified database
|
||||
references. The namespace defaults to ``["public"]``. Local
|
||||
connections raise ``NotImplementedError``.
|
||||
"""
|
||||
return AsyncSqlQuery(
|
||||
await self._inner.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
)
|
||||
)
|
||||
|
||||
async def describe_query(self, query_id: UUID) -> QueryDescription:
|
||||
"""Describe a submitted SQL query by its connection-scoped id."""
|
||||
return await self._inner.describe_query(query_id)
|
||||
|
||||
async def namespace_client(self) -> LanceNamespace:
|
||||
"""Get the equivalent namespace client for this connection.
|
||||
|
||||
@@ -35,3 +35,9 @@ class JobCancelledError(RuntimeError):
|
||||
"""Exception raised when an asynchronous job was cancelled."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class JobNotFoundError(ValueError):
|
||||
"""Exception raised when opening a job the server does not have."""
|
||||
|
||||
pass
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
"""Canonical Function values exchanged with LanceDB Enterprise services.
|
||||
|
||||
These immutable models contain client/wire state only. Catalog persistence,
|
||||
environment bake, and execution are owned by Sophon.
|
||||
environment bake, secret resolution, and execution are owned by Sophon.
|
||||
``RefreshColumnResult`` is also the backend-neutral result of a local
|
||||
expression-backed refresh job.
|
||||
"""
|
||||
@@ -25,7 +25,7 @@ import re
|
||||
import sys
|
||||
import textwrap
|
||||
import types
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import date, datetime
|
||||
from typing import (
|
||||
Annotated,
|
||||
@@ -50,6 +50,7 @@ from pydantic import (
|
||||
)
|
||||
|
||||
from .schema import is_blob_v2_field as _is_blob_v2_field
|
||||
from .secrets import EnvVarSecret
|
||||
|
||||
_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
|
||||
_UInt32 = conint(strict=True, ge=0, le=2**32 - 1)
|
||||
@@ -309,6 +310,7 @@ class FunctionVersion(_RemoteValue):
|
||||
runtime: PythonRuntimeSpec
|
||||
runtime_digest: str
|
||||
environment_digest: str
|
||||
secret_bindings: Mapping[str, str] = {}
|
||||
created_at: str
|
||||
|
||||
def __call__(self, **inputs: Any) -> FunctionApplication:
|
||||
@@ -370,12 +372,18 @@ class FunctionVersion(_RemoteValue):
|
||||
|
||||
|
||||
class FunctionRegistrationRequest(_RemoteValue):
|
||||
"""Stable remote registration envelope produced by :func:`udf`."""
|
||||
"""Stable remote registration envelope produced by :func:`udf`.
|
||||
|
||||
Credential values deliberately have no field here. The only secret-shaped
|
||||
thing a client sends is ``secret_bindings``: the name of a Secret the
|
||||
database already holds, which the remote service resolves at execution.
|
||||
"""
|
||||
|
||||
name: str
|
||||
artifact: FunctionArtifactRequest
|
||||
signature: FunctionSignature
|
||||
runtime: PythonRuntimeSpec
|
||||
secret_bindings: Mapping[str, str] = {}
|
||||
|
||||
|
||||
class FunctionVersionRef(_OpenRemoteValue):
|
||||
@@ -470,11 +478,7 @@ class InputBinding(_RemoteValue):
|
||||
|
||||
|
||||
class OutputMapping(_RemoteValue):
|
||||
"""One stable result-field mapping.
|
||||
|
||||
Assignment state is outside the Slice 1 client contract. During the NULL
|
||||
transition Lance exposes no public cell-flag identifier to persist here.
|
||||
"""
|
||||
"""One stable result-field mapping."""
|
||||
|
||||
result_field: str
|
||||
output_name: str
|
||||
@@ -484,6 +488,13 @@ class OutputMapping(_RemoteValue):
|
||||
nullable: bool
|
||||
|
||||
|
||||
class AssignmentMapping(_RemoteValue):
|
||||
"""Internal physical column preserving flattened struct validity."""
|
||||
|
||||
output_name: str
|
||||
output_field_id: _Int32
|
||||
|
||||
|
||||
class FunctionBinding(_RemoteValue):
|
||||
"""Immutable Function binding persisted by the Enterprise table service."""
|
||||
|
||||
@@ -491,6 +502,7 @@ class FunctionBinding(_RemoteValue):
|
||||
function: FunctionVersionRef
|
||||
inputs: tuple[InputBinding, ...]
|
||||
outputs: tuple[OutputMapping, ...]
|
||||
assignment: Optional[AssignmentMapping] = None
|
||||
input_schema: Optional[Mapping[str, Any]] = None
|
||||
output_schema: Optional[Mapping[str, Any]] = None
|
||||
|
||||
@@ -520,7 +532,21 @@ class RefreshColumnResult(_RemoteValue):
|
||||
|
||||
|
||||
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
|
||||
_DECLARED_SECRET = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
|
||||
MAX_FUNCTION_SECRET_BINDINGS = 16
|
||||
"""A Function binds at most this many Secrets.
|
||||
|
||||
Each bound Secret is one extra read on the launch path of every fragment, so
|
||||
the count needs a bound for the same reason a credential needs a size limit.
|
||||
"""
|
||||
_FUNCTION_BLOB_V2_TYPE = "blob_v2"
|
||||
_ARROW_EXTENSION_NAME_KEY = "ARROW:extension:name"
|
||||
_BLOB_V2_EXTENSION_NAME = "lance.blob.v2"
|
||||
_NESTED_BLOB_COLLECTION_ERROR = (
|
||||
"unsupported Arrow type for Function signature: Blob v2 fields nested under "
|
||||
"collection types are not supported"
|
||||
)
|
||||
|
||||
|
||||
_GRAMMAR_PRIMITIVES = (
|
||||
@@ -591,6 +617,19 @@ def _validate_exact_arrow_field(field: pa.Field) -> None:
|
||||
"unsupported Arrow type for Function signature: lance.blob.v2 "
|
||||
f"requires a supported Blob storage layout, got {field}"
|
||||
)
|
||||
metadata = {
|
||||
(key.decode() if isinstance(key, bytes) else key): (
|
||||
value.decode() if isinstance(value, bytes) else value
|
||||
)
|
||||
for key, value in (field.metadata or {}).items()
|
||||
}
|
||||
if metadata and metadata != {
|
||||
_ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME
|
||||
}:
|
||||
raise TypeError(
|
||||
"unsupported Arrow type for Function signature: lance.blob.v2 "
|
||||
"field metadata must contain only its canonical extension marker"
|
||||
)
|
||||
elif field.metadata:
|
||||
raise TypeError(
|
||||
"unsupported Arrow type for Function signature: field metadata "
|
||||
@@ -655,23 +694,84 @@ def _canonical_arrow_field(field: pa.Field) -> str:
|
||||
return _canonical_arrow_type(field.type)
|
||||
|
||||
|
||||
def _exact_arrow_field(field: pa.Field) -> dict[str, Any]:
|
||||
def _blob_storage_type(field: pa.Field) -> pa.DataType:
|
||||
data_type = field.type
|
||||
if isinstance(data_type, pa.ExtensionType):
|
||||
return data_type.storage_type
|
||||
return data_type
|
||||
|
||||
|
||||
def _exact_blob_storage_type(field: pa.Field) -> dict[str, Any]:
|
||||
storage = _blob_storage_type(field)
|
||||
if not pa.types.is_struct(storage):
|
||||
raise TypeError(
|
||||
"unsupported Arrow type for Function signature: lance.blob.v2 "
|
||||
"requires struct storage"
|
||||
)
|
||||
return {
|
||||
"type": "struct",
|
||||
"fields": [
|
||||
{
|
||||
"name": child.name,
|
||||
"nullable": child.nullable,
|
||||
"type": (
|
||||
{"type": "large_binary"}
|
||||
if pa.types.is_large_binary(child.type)
|
||||
else _exact_arrow_type(child.type)
|
||||
),
|
||||
}
|
||||
for child in storage
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _data_type_has_blob_v2(data_type: pa.DataType) -> bool:
|
||||
if pa.types.is_struct(data_type):
|
||||
return any(
|
||||
_is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
|
||||
for field in data_type
|
||||
)
|
||||
if (
|
||||
pa.types.is_list(data_type)
|
||||
or pa.types.is_large_list(data_type)
|
||||
or pa.types.is_fixed_size_list(data_type)
|
||||
):
|
||||
field = data_type.value_field
|
||||
return _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
|
||||
if pa.types.is_map(data_type):
|
||||
return any(
|
||||
_is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
|
||||
for field in (data_type.key_field, data_type.item_field)
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _exact_arrow_field(
|
||||
field: pa.Field, *, inside_collection: bool = False
|
||||
) -> dict[str, Any]:
|
||||
_validate_exact_arrow_field(field)
|
||||
if _is_blob_v2_field(field):
|
||||
raise TypeError(
|
||||
"unsupported Arrow type for Function signature: nested Blob v2 "
|
||||
"fields are not supported; declare Blob parameters or named result "
|
||||
"fields directly"
|
||||
)
|
||||
if inside_collection:
|
||||
raise TypeError(_NESTED_BLOB_COLLECTION_ERROR)
|
||||
return {
|
||||
"name": field.name,
|
||||
"nullable": field.nullable,
|
||||
"type": _exact_blob_storage_type(field),
|
||||
"metadata": {
|
||||
_ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME,
|
||||
},
|
||||
}
|
||||
value = {
|
||||
"name": field.name,
|
||||
"nullable": field.nullable,
|
||||
"type": _exact_arrow_type(field.type),
|
||||
"type": _exact_arrow_type(field.type, inside_collection=inside_collection),
|
||||
}
|
||||
return value
|
||||
|
||||
|
||||
def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]:
|
||||
def _exact_arrow_type(
|
||||
data_type: pa.DataType, *, inside_collection: bool = False
|
||||
) -> dict[str, Any]:
|
||||
for candidate, name in _GRAMMAR_PRIMITIVES:
|
||||
if data_type == candidate:
|
||||
return {"type": name}
|
||||
@@ -685,7 +785,10 @@ def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]:
|
||||
)
|
||||
return {
|
||||
"type": "struct",
|
||||
"fields": [_exact_arrow_field(field) for field in fields],
|
||||
"fields": [
|
||||
_exact_arrow_field(field, inside_collection=inside_collection)
|
||||
for field in fields
|
||||
],
|
||||
}
|
||||
if (
|
||||
pa.types.is_list(data_type)
|
||||
@@ -710,11 +813,15 @@ def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]:
|
||||
if pa.types.is_large_list(data_type)
|
||||
else "fixed_size_list"
|
||||
),
|
||||
"fields": [_exact_arrow_field(data_type.value_field)],
|
||||
"fields": [
|
||||
_exact_arrow_field(data_type.value_field, inside_collection=True)
|
||||
],
|
||||
}
|
||||
if pa.types.is_fixed_size_list(data_type):
|
||||
value["length"] = data_type.list_size
|
||||
return value
|
||||
if pa.types.is_map(data_type) and _data_type_has_blob_v2(data_type):
|
||||
raise TypeError(_NESTED_BLOB_COLLECTION_ERROR)
|
||||
raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
|
||||
|
||||
|
||||
@@ -824,8 +931,6 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp
|
||||
|
||||
if not fields:
|
||||
raise ValueError("named-struct Function output must contain at least one field")
|
||||
if any(field.nullable for field in fields):
|
||||
raise ValueError("Function output fields must be non-nullable")
|
||||
for field in fields:
|
||||
_validate_exact_arrow_field(field)
|
||||
names = [field.name for field in fields]
|
||||
@@ -837,7 +942,7 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp
|
||||
FunctionResultField(
|
||||
name=field.name,
|
||||
arrow_type=_canonical_arrow_field(field),
|
||||
nullable=False,
|
||||
nullable=field.nullable,
|
||||
)
|
||||
for field in fields
|
||||
),
|
||||
@@ -1176,9 +1281,63 @@ class UdfDefinition:
|
||||
|
||||
@property
|
||||
def registration_request(self) -> FunctionRegistrationRequest:
|
||||
"""The immutable request sent by ``create_function_async``."""
|
||||
"""The immutable request sent by ``create_function_async``.
|
||||
|
||||
Carries no secret bindings. Binding is a registration-time decision,
|
||||
so a Function bound to Secrets is registered through :meth:`bind_secrets`,
|
||||
which is what ``create_function`` calls.
|
||||
"""
|
||||
return self._request
|
||||
|
||||
def bind_secrets(
|
||||
self, secrets: Optional[Sequence[EnvVarSecret]]
|
||||
) -> FunctionRegistrationRequest:
|
||||
"""The registration request for this definition bound to ``secrets``.
|
||||
|
||||
Binding does not change the Function's source: each
|
||||
[EnvVarSecret][lancedb.secrets.EnvVarSecret] names a Secret and the
|
||||
environment variable its value should arrive in, and the Function reads
|
||||
that variable the way it already did. Whether the named Secrets exist is
|
||||
the server's answer, not this one.
|
||||
"""
|
||||
bindings = () if secrets is None else tuple(secrets)
|
||||
wrong_type = [
|
||||
binding for binding in bindings if not isinstance(binding, EnvVarSecret)
|
||||
]
|
||||
if wrong_type:
|
||||
kinds = sorted({type(binding).__name__ for binding in wrong_type})
|
||||
raise TypeError(
|
||||
f"Function secrets must be EnvVarSecret values, not {kinds!r}; a "
|
||||
"credential value is never sent to this API"
|
||||
)
|
||||
if len(bindings) > MAX_FUNCTION_SECRET_BINDINGS:
|
||||
raise ValueError(
|
||||
f"a Function binds at most {MAX_FUNCTION_SECRET_BINDINGS} secrets, "
|
||||
f"not {len(bindings)}"
|
||||
)
|
||||
variables = [binding.env_variable for binding in bindings]
|
||||
duplicates = sorted({name for name in variables if variables.count(name) > 1})
|
||||
if duplicates:
|
||||
raise ValueError(
|
||||
"a Function binds each environment variable once; duplicated: "
|
||||
f"{duplicates!r}"
|
||||
)
|
||||
# `env` is ordinary configuration carried in the definition, so a name in
|
||||
# both would have a value visible in the Function's record and a value
|
||||
# that is not. Refuse rather than pick.
|
||||
environment = self._request.runtime.env or {}
|
||||
overlap = sorted(set(environment) & set(variables))
|
||||
if overlap:
|
||||
raise ValueError(
|
||||
f"Function env and secret bindings must be disjoint: {overlap!r}"
|
||||
)
|
||||
if not bindings:
|
||||
return self._request
|
||||
# The binding records the full id -- path plus name -- because that is
|
||||
# what the service resolves. At the root it is the bare name.
|
||||
resolved = {binding.env_variable: binding.secret_id for binding in bindings}
|
||||
return self._request._copy(update={"secret_bindings": resolved})
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self._function(*args, **kwargs)
|
||||
|
||||
@@ -1220,8 +1379,9 @@ def udf(
|
||||
|
||||
Input and output signatures are inferred from supported annotations. For
|
||||
Arrow types annotations cannot express precisely, pass ``input_schema``
|
||||
and ``output_schema`` together. Nullable outputs are rejected because V1
|
||||
uses physical NULL to represent unassigned computed-column rows.
|
||||
and ``output_schema`` together. Scalar outputs must be non-nullable. Every
|
||||
named-struct field may be nullable; Enterprise preserves the struct's
|
||||
validity when the result is expanded into sibling columns.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -1233,8 +1393,8 @@ def udf(
|
||||
Explicit input fields in the exact order of the callable parameters.
|
||||
Must be provided together with ``output_schema``.
|
||||
output_schema : pyarrow.DataType, pyarrow.Field, or pyarrow.Schema, optional
|
||||
Explicit scalar or named-struct output. Must be non-nullable and be
|
||||
provided together with ``input_schema``.
|
||||
Explicit scalar or named-struct output. Scalar outputs must be
|
||||
non-nullable. Must be provided together with ``input_schema``.
|
||||
pip : sequence of str, optional
|
||||
Pip requirements for the remote environment.
|
||||
conda : sequence of str, optional
|
||||
@@ -1242,7 +1402,9 @@ def udf(
|
||||
conda_channels : sequence of str, optional
|
||||
Conda channels in priority order; requires ``conda``.
|
||||
env : mapping of str to str, optional
|
||||
Environment variables included in the Function definition.
|
||||
Environment variables included in the Function definition. Not for
|
||||
credentials -- these are ordinary configuration, stored with the
|
||||
Function and visible wherever it is.
|
||||
python_version : str, optional
|
||||
Remote Python major/minor version. Defaults to the client version.
|
||||
gpu : bool, default False
|
||||
@@ -1300,6 +1462,7 @@ def udf(
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AssignmentMapping",
|
||||
"ApplicationInput",
|
||||
"FunctionApplication",
|
||||
"FunctionArtifact",
|
||||
|
||||
@@ -4,15 +4,27 @@
|
||||
"""Handles to operations a server may run asynchronously."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import timedelta
|
||||
from typing import Any, Callable, Generic, Optional, TypeVar, cast
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from lancedb.background_loop import LOOP
|
||||
|
||||
from . import _lancedb
|
||||
from ._lancedb import JobDescription, JobFailureInfo, JobInfo
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
__all__ = [
|
||||
"AsyncJob",
|
||||
"Job",
|
||||
"JobDescription",
|
||||
"JobFailureInfo",
|
||||
"JobInfo",
|
||||
]
|
||||
|
||||
|
||||
class AsyncJob(Generic[T]):
|
||||
"""A handle to an operation that may still be running.
|
||||
@@ -78,6 +90,149 @@ class AsyncJob(Generic[T]):
|
||||
return
|
||||
await self._inner.cancel()
|
||||
|
||||
async def refresh(self) -> None:
|
||||
"""Ask the backend for this job's current state, and for a server-side
|
||||
job its full record, then cache it for the properties below.
|
||||
|
||||
The properties are all `None` until this runs, because submitting an
|
||||
operation returns only a job id. `status` fetches the whole record too;
|
||||
`wait` records only the terminal state it establishes.
|
||||
"""
|
||||
if self._inner is None:
|
||||
return
|
||||
await self._inner.refresh()
|
||||
|
||||
@property
|
||||
def state(self) -> Optional[str]:
|
||||
"""The last observed lifecycle state, without contacting the backend.
|
||||
|
||||
`None` until the handle has talked to it. See :meth:`AsyncJob.refresh`.
|
||||
"""
|
||||
if self._inner is None:
|
||||
return "finished"
|
||||
return self._inner._state
|
||||
|
||||
@property
|
||||
def job_type(self) -> Optional[str]:
|
||||
"""The job's type, as the server names it.
|
||||
|
||||
`None` for an in-process job, which has no server-side record.
|
||||
"""
|
||||
return self._field("job_type")
|
||||
|
||||
@property
|
||||
def creation_ms(self) -> Optional[int]:
|
||||
"""When the job was created, in milliseconds since the epoch."""
|
||||
return self._field("creation_ms")
|
||||
|
||||
@property
|
||||
def spec(self) -> Optional[Any]:
|
||||
"""The job-type-specific specification it was submitted with."""
|
||||
return self._field("spec")
|
||||
|
||||
@property
|
||||
def result(self) -> Optional[Any]:
|
||||
"""The job-type-specific terminal result, as reported data rather than
|
||||
the typed model :meth:`AsyncJob.wait` returns.
|
||||
|
||||
`None` until the job succeeds, so a job that never terminates reports
|
||||
its progress through :meth:`AsyncJob.events` instead.
|
||||
"""
|
||||
return self._field("result")
|
||||
|
||||
@property
|
||||
def failure(self) -> Optional[JobFailureInfo]:
|
||||
"""Why the job failed, when it failed and the server reports a reason."""
|
||||
return self._field("failure")
|
||||
|
||||
@property
|
||||
def _spec_json(self) -> Optional[str]:
|
||||
return self._field("_spec_json")
|
||||
|
||||
@property
|
||||
def _result_json(self) -> Optional[str]:
|
||||
return self._field("_result_json")
|
||||
|
||||
def _field(self, name: str) -> Optional[Any]:
|
||||
description = self._inner._description if self._inner is not None else None
|
||||
return getattr(description, name) if description is not None else None
|
||||
|
||||
async def events(
|
||||
self,
|
||||
*,
|
||||
limit: Optional[int] = None,
|
||||
filter: Optional[str] = None,
|
||||
) -> "pa.Table":
|
||||
"""This job's recorded lifecycle events.
|
||||
|
||||
Where the properties above report a terminal result only once the job
|
||||
reaches one, events are written as the job runs and outlive the workers
|
||||
that produced them. A distributed job records a `claim`/`claim_complete`
|
||||
pair per unit of work, each carrying `rows_processed`, so a job that
|
||||
never finishes still accounts for what it did.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
limit: int, optional
|
||||
Maximum event rows to return. The server caps results at 1000 by
|
||||
default and 10,000 at most, and truncates without saying so, so
|
||||
pass this for a job that emits an event per fragment.
|
||||
filter: str, optional
|
||||
SQL-like expression over the `state`, `updated_by`, `emitted_from`,
|
||||
`emitted_by`, and `claim_entity` columns, such as
|
||||
``state = 'claim_complete'``.
|
||||
"""
|
||||
if self._inner is None:
|
||||
raise NotImplementedError(
|
||||
"job event history is only available for server-side jobs"
|
||||
)
|
||||
return await self._inner.events(limit=limit, filter=filter)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return _job_repr("AsyncJob", self)
|
||||
|
||||
|
||||
_REPR_INDENT = " " * 4
|
||||
|
||||
|
||||
def _repr_payload(value: Any) -> str:
|
||||
"""Render a job payload as indented JSON, aligned under its field."""
|
||||
try:
|
||||
rendered = json.dumps(value, indent=4)
|
||||
except TypeError:
|
||||
return repr(value)
|
||||
return rendered.replace("\n", "\n" + _REPR_INDENT)
|
||||
|
||||
|
||||
def _job_repr(kind: str, job: Any) -> str:
|
||||
"""Render every field the handle currently knows, omitting the rest.
|
||||
|
||||
One field per line, with the JSON payloads indented, because a refresh
|
||||
job's spec and result are the point of printing it.
|
||||
"""
|
||||
state = job.state
|
||||
if state is None:
|
||||
# Nothing has been fetched yet, so there is nothing to lay out.
|
||||
known = f"id={job.id!r}, " if job.id is not None else ""
|
||||
return f"{kind}({known}not refreshed)"
|
||||
|
||||
fields = []
|
||||
if job.id is not None:
|
||||
fields.append(f"id={job.id!r}")
|
||||
fields.append(f"state={state!r}")
|
||||
for name in ("job_type", "creation_ms"):
|
||||
value = getattr(job, name)
|
||||
if value is not None:
|
||||
fields.append(f"{name}={value!r}")
|
||||
for name in ("spec", "result"):
|
||||
value = getattr(job, name)
|
||||
if value is not None:
|
||||
fields.append(f"{name}={_repr_payload(value)}")
|
||||
if job.failure is not None:
|
||||
fields.append(f"failure={job.failure!r}")
|
||||
body = "".join(f"\n{_REPR_INDENT}{field}," for field in fields)
|
||||
return f"{kind}({body}\n)"
|
||||
|
||||
|
||||
class Job(Generic[T]):
|
||||
"""Synchronous counterpart of `AsyncJob` with the same result type."""
|
||||
@@ -122,6 +277,75 @@ class Job(Generic[T]):
|
||||
return
|
||||
LOOP.run(self._inner.cancel())
|
||||
|
||||
def refresh(self) -> None:
|
||||
"""Ask the backend for this job's current state and record.
|
||||
|
||||
See :meth:`AsyncJob.refresh`.
|
||||
"""
|
||||
if self._inner is None:
|
||||
return
|
||||
LOOP.run(self._inner.refresh())
|
||||
|
||||
@property
|
||||
def state(self) -> Optional[str]:
|
||||
"""The last observed lifecycle state. See :attr:`AsyncJob.state`."""
|
||||
return self._inner.state if self._inner is not None else "finished"
|
||||
|
||||
@property
|
||||
def job_type(self) -> Optional[str]:
|
||||
"""The job's type. See :attr:`AsyncJob.job_type`."""
|
||||
return self._field("job_type")
|
||||
|
||||
@property
|
||||
def creation_ms(self) -> Optional[int]:
|
||||
"""When the job was created. See :attr:`AsyncJob.creation_ms`."""
|
||||
return self._field("creation_ms")
|
||||
|
||||
@property
|
||||
def spec(self) -> Optional[Any]:
|
||||
"""The job's specification. See :attr:`AsyncJob.spec`."""
|
||||
return self._field("spec")
|
||||
|
||||
@property
|
||||
def result(self) -> Optional[Any]:
|
||||
"""The job's terminal result. See :attr:`AsyncJob.result`."""
|
||||
return self._field("result")
|
||||
|
||||
@property
|
||||
def failure(self) -> Optional[JobFailureInfo]:
|
||||
"""Why the job failed. See :attr:`AsyncJob.failure`."""
|
||||
return self._field("failure")
|
||||
|
||||
@property
|
||||
def _spec_json(self) -> Optional[str]:
|
||||
return self._field("_spec_json")
|
||||
|
||||
@property
|
||||
def _result_json(self) -> Optional[str]:
|
||||
return self._field("_result_json")
|
||||
|
||||
def _field(self, name: str) -> Optional[Any]:
|
||||
return getattr(self._inner, name) if self._inner is not None else None
|
||||
|
||||
def events(
|
||||
self,
|
||||
*,
|
||||
limit: Optional[int] = None,
|
||||
filter: Optional[str] = None,
|
||||
) -> "pa.Table":
|
||||
"""This job's recorded lifecycle events.
|
||||
|
||||
See :meth:`AsyncJob.events`.
|
||||
"""
|
||||
if self._inner is None:
|
||||
raise NotImplementedError(
|
||||
"job event history is only available for server-side jobs"
|
||||
)
|
||||
return LOOP.run(self._inner.events(limit=limit, filter=filter))
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return _job_repr("Job", self)
|
||||
|
||||
|
||||
def _typed_job(
|
||||
inner: "_lancedb.Job", result_decoder: Callable[[str], T]
|
||||
|
||||
@@ -42,6 +42,8 @@ class MaterializedViewDefinition:
|
||||
"""Cap on the number of rows the view holds."""
|
||||
inputs: List[str] = field(default_factory=list)
|
||||
"""Source columns the projections and filter read."""
|
||||
source_namespace: List[str] = field(default_factory=list)
|
||||
"""Namespace holding the source table; empty is the root namespace."""
|
||||
|
||||
|
||||
def _definition_from_schema(
|
||||
@@ -53,7 +55,8 @@ def _definition_from_schema(
|
||||
raise ValueError(f"Table '{name}' is not a materialized view")
|
||||
value = json.loads(raw)
|
||||
kind = value.get("kind")
|
||||
if kind != "select":
|
||||
# "namespaced_select" keeps older readers from resolving the source at root.
|
||||
if kind not in ("select", "namespaced_select"):
|
||||
raise NotImplementedError(
|
||||
f"materialized view '{name}' is defined by '{kind}', which this "
|
||||
"version of lancedb cannot refresh"
|
||||
@@ -66,6 +69,7 @@ def _definition_from_schema(
|
||||
filter=value.get("filter"),
|
||||
limit=value.get("limit"),
|
||||
inputs=value.get("inputs", []),
|
||||
source_namespace=value.get("source_namespace", []),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
|
||||
from uuid import UUID
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override
|
||||
@@ -48,8 +49,11 @@ from lancedb._lancedb import (
|
||||
connect_namespace_client as _connect_namespace_client,
|
||||
)
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.arrow import AsyncRecordBatchReader
|
||||
from lancedb.db import AsyncConnection, DBConnection
|
||||
from lancedb.job import AsyncJob, Job
|
||||
from lancedb.sql import AsyncQuery as AsyncSqlQuery
|
||||
from lancedb.sql import QueryDescription
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
connect as namespace_connect,
|
||||
@@ -1447,6 +1451,37 @@ class AsyncLanceNamespaceDBConnection:
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
async def execute_query(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> AsyncRecordBatchReader:
|
||||
"""Execute SQL when supported by the underlying connection."""
|
||||
return await self._inner.execute_query(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
)
|
||||
|
||||
async def execute_query_async(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> AsyncSqlQuery:
|
||||
"""Start executing SQL when supported by the underlying connection.
|
||||
|
||||
Namespace-backed local connections do not support SQL.
|
||||
"""
|
||||
return await self._inner.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
)
|
||||
|
||||
async def describe_query(self, query_id: UUID) -> QueryDescription:
|
||||
"""Describe a submitted SQL query when supported."""
|
||||
return await self._inner.describe_query(query_id)
|
||||
|
||||
async def namespace_client(self) -> LanceNamespace:
|
||||
"""Get the namespace client for this connection.
|
||||
|
||||
|
||||
@@ -109,6 +109,7 @@ def _query_is_plain_scan(query: Query) -> bool:
|
||||
return (
|
||||
query.vector is None
|
||||
and query.full_text_query is None
|
||||
and query.take_offsets is None
|
||||
and not query.postfilter
|
||||
and not query.order_by
|
||||
)
|
||||
@@ -804,6 +805,10 @@ class Query(pydantic.BaseModel):
|
||||
# offset to start fetching results from
|
||||
offset: Optional[int] = None
|
||||
|
||||
# Dataset offsets whose duplicate occurrences must be restored after lookup.
|
||||
# This is populated when a take query is converted to this serializable form.
|
||||
take_offsets: Optional[List[int]] = None
|
||||
|
||||
# if true, will only search the indexed data
|
||||
fast_search: Optional[bool] = None
|
||||
|
||||
@@ -825,6 +830,7 @@ class Query(pydantic.BaseModel):
|
||||
query = cls()
|
||||
query.limit = req.limit
|
||||
query.offset = req.offset
|
||||
query.take_offsets = req.take_offsets
|
||||
query.filter = req.filter
|
||||
query.full_text_query = req.full_text_search
|
||||
query.columns = req.select
|
||||
|
||||
@@ -7,8 +7,18 @@ import json
|
||||
import logging
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Dict,
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
)
|
||||
from urllib.parse import urlparse
|
||||
from uuid import UUID
|
||||
import warnings
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
@@ -25,10 +35,13 @@ from ..common import DATA
|
||||
from ..db import DBConnection, LOOP
|
||||
from ..functions import FunctionVersion, UdfDefinition
|
||||
from ..job import AsyncJob, Job
|
||||
from ..sql import Query as SqlQuery
|
||||
from ..sql import QueryDescription
|
||||
from ..materialized_view import MaterializedView, SelectArg
|
||||
from ..secrets import EnvVarSecret, SecretInfo
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .._lancedb import JobDescription, JobInfo
|
||||
from .._lancedb import JobInfo
|
||||
from ..embeddings import EmbeddingFunctionConfig
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
@@ -116,6 +129,7 @@ class RemoteDBConnection(DBConnection):
|
||||
read_timeout: Optional[float] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
sql_host_override: Optional[str] = None,
|
||||
):
|
||||
"""Connect to a remote LanceDB database."""
|
||||
if isinstance(client_config, dict):
|
||||
@@ -161,6 +175,7 @@ class RemoteDBConnection(DBConnection):
|
||||
self.api_key = api_key
|
||||
self.region = region
|
||||
self.host_override = host_override
|
||||
self.sql_host_override = sql_host_override
|
||||
self.storage_options = storage_options
|
||||
self.db_name = parsed.netloc
|
||||
|
||||
@@ -175,6 +190,7 @@ class RemoteDBConnection(DBConnection):
|
||||
api_key=api_key,
|
||||
region=region,
|
||||
host_override=host_override,
|
||||
sql_host_override=sql_host_override,
|
||||
client_config=client_config,
|
||||
storage_options=storage_options,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
@@ -193,6 +209,7 @@ class RemoteDBConnection(DBConnection):
|
||||
"api_key": self.api_key,
|
||||
"region": self.region,
|
||||
"host_override": self.host_override,
|
||||
"sql_host_override": self.sql_host_override,
|
||||
"client_config": _client_config_to_dict(self.client_config),
|
||||
"storage_options": self.storage_options,
|
||||
}
|
||||
@@ -732,36 +749,67 @@ class RemoteDBConnection(DBConnection):
|
||||
)
|
||||
|
||||
@override
|
||||
def job(self, job_id: str) -> Job:
|
||||
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
|
||||
|
||||
The handle is constructed without a server round trip; an unknown id
|
||||
surfaces when the handle is used. Dropping the handle has no effect
|
||||
on the job itself.
|
||||
def open_job(self, job_id: str) -> Job:
|
||||
"""Open a server-side job by id. See
|
||||
[DBConnection.open_job][lancedb.db.DBConnection.open_job].
|
||||
"""
|
||||
return Job(self._conn.job(job_id))
|
||||
return Job(LOOP.run(self._conn.open_job(job_id)))
|
||||
|
||||
@override
|
||||
def create_function_async(self, definition: UdfDefinition) -> Job[FunctionVersion]:
|
||||
return Job(LOOP.run(self._conn.create_function_async(definition)))
|
||||
def create_function_async(
|
||||
self,
|
||||
definition: UdfDefinition,
|
||||
*,
|
||||
secrets: Optional[Sequence[EnvVarSecret]] = None,
|
||||
) -> Job[FunctionVersion]:
|
||||
job = LOOP.run(self._conn.create_function_async(definition, secrets=secrets))
|
||||
return Job(job)
|
||||
|
||||
@override
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
return LOOP.run(self._conn.get_function(name, version=version))
|
||||
|
||||
@override
|
||||
def list_functions(self) -> List[FunctionVersion]:
|
||||
return LOOP.run(self._conn.list_functions())
|
||||
|
||||
@override
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
|
||||
@override
|
||||
def create_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def alter_secret(
|
||||
self, name: str, value: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.alter_secret(name, value, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def describe_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> SecretInfo:
|
||||
return LOOP.run(self._conn.describe_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_secrets(self, *, namespace_path: Optional[List[str]] = None) -> List[str]:
|
||||
return LOOP.run(self._conn.list_secrets(namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def drop_secret(
|
||||
self, name: str, *, namespace_path: Optional[List[str]] = None
|
||||
) -> None:
|
||||
LOOP.run(self._conn.drop_secret(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List["JobInfo"]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return LOOP.run(self._conn.list_jobs())
|
||||
|
||||
@override
|
||||
def get_job(self, job_id: str) -> Optional["JobDescription"]:
|
||||
"""Describe a single server-side job by id.
|
||||
|
||||
Returns None when the server has no such job.
|
||||
"""
|
||||
return LOOP.run(self._conn.get_job(job_id))
|
||||
|
||||
@override
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""Request cancellation of a server-side job by id.
|
||||
@@ -773,12 +821,35 @@ class RemoteDBConnection(DBConnection):
|
||||
return LOOP.run(self._conn.cancel_job(job_id))
|
||||
|
||||
@override
|
||||
def job_history(self, job_id: Optional[str] = None) -> List[pa.RecordBatch]:
|
||||
"""The lifecycle event history of a server-side job, as Arrow batches.
|
||||
def execute_query_async(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
default_namespace_path: Optional[List[str]] = None,
|
||||
) -> SqlQuery:
|
||||
"""Start executing SQL through this remote connection.
|
||||
|
||||
Lists history across all jobs when `job_id` is None.
|
||||
Unqualified tables use this connection's database and the
|
||||
``["public"]`` namespace by default. Fully qualified table names may
|
||||
reference other databases available to the same deployment.
|
||||
"""
|
||||
return LOOP.run(self._conn.job_history(job_id))
|
||||
return SqlQuery(
|
||||
LOOP.run(
|
||||
self._conn.execute_query_async(
|
||||
query,
|
||||
default_namespace_path=default_namespace_path,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def describe_query(self, query_id: UUID) -> QueryDescription:
|
||||
"""Describe a submitted SQL query by its connection-scoped id."""
|
||||
return LOOP.run(
|
||||
self._conn.describe_query(
|
||||
query_id,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
def namespace_client(self) -> LanceNamespace:
|
||||
|
||||
@@ -177,4 +177,7 @@ class OAuthProvider(HeaderProvider):
|
||||
if not self._current_token:
|
||||
raise RuntimeError("Failed to obtain OAuth token")
|
||||
|
||||
return {"Authorization": f"Bearer {self._current_token}"}
|
||||
return {
|
||||
"Authorization": f"Bearer {self._current_token}",
|
||||
"x-lancedb-credential-type": "oidc",
|
||||
}
|
||||
|
||||
@@ -548,6 +548,7 @@ class RemoteTable(Table):
|
||||
LOOP.run(
|
||||
self._table.create_index(
|
||||
column,
|
||||
replace=replace,
|
||||
config=config,
|
||||
wait_timeout=wait_timeout,
|
||||
name=name,
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Named Secrets, and the bindings that deliver them to Functions.
|
||||
|
||||
A Secret is a database-scoped named credential. Nothing in this module holds a
|
||||
value: :class:`EnvVarSecret` names one and says which environment variable it
|
||||
should arrive in, and the value is resolved by the remote service when a
|
||||
Function bound to it runs. No API returns a stored credential, by construction
|
||||
rather than by policy -- there is no code path that could.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
_SECRET_NAME = re.compile(r"^[A-Za-z0-9_-]{1,255}$")
|
||||
_ENV_VARIABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
|
||||
|
||||
def validate_secret_name(name: str) -> str:
|
||||
"""Check a Secret name locally and return it unchanged."""
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(f"Secret name must be a string, not {type(name).__name__}")
|
||||
if not _SECRET_NAME.fullmatch(name):
|
||||
raise ValueError(f"invalid Secret name: {name!r}")
|
||||
return name
|
||||
|
||||
|
||||
def validate_namespace_path(namespace_path=None):
|
||||
"""Check a namespace path locally and return it as a tuple.
|
||||
|
||||
``None`` and ``[]`` both mean the root namespace. Segments follow the same
|
||||
rule as Secret names, and necessarily so: an id is the path plus the name
|
||||
joined with ``$``, so the join has to read the same from either side.
|
||||
"""
|
||||
if namespace_path is None:
|
||||
return ()
|
||||
if isinstance(namespace_path, str):
|
||||
raise TypeError(
|
||||
"namespace_path must be a list of segments, not a string; "
|
||||
f"did you mean [{namespace_path!r}]?"
|
||||
)
|
||||
segments = tuple(namespace_path)
|
||||
for segment in segments:
|
||||
if not isinstance(segment, str):
|
||||
raise TypeError(
|
||||
f"namespace path segment must be a string, not {type(segment).__name__}"
|
||||
)
|
||||
if not _SECRET_NAME.fullmatch(segment):
|
||||
raise ValueError(f"invalid namespace path segment: {segment!r}")
|
||||
return segments
|
||||
|
||||
|
||||
def secret_id(secret: str, namespace_path=()) -> str:
|
||||
"""The id a Secret is named by: its namespace path plus its name.
|
||||
|
||||
Joined with ``$``, which is why neither a name nor a segment may contain
|
||||
one. At the root this is the bare name, so a root id is unchanged from
|
||||
before namespaces existed.
|
||||
"""
|
||||
return "$".join((*namespace_path, secret))
|
||||
|
||||
|
||||
def validate_env_variable(name: str) -> str:
|
||||
"""Check an environment variable name locally and return it unchanged."""
|
||||
if not isinstance(name, str):
|
||||
raise TypeError(
|
||||
f"environment variable name must be a string, not {type(name).__name__}"
|
||||
)
|
||||
if not _ENV_VARIABLE.fullmatch(name):
|
||||
raise ValueError(f"invalid environment variable name: {name!r}")
|
||||
return name
|
||||
|
||||
|
||||
class EnvVarSecret:
|
||||
"""A Secret bound to the environment variable a Function's library reads.
|
||||
|
||||
Pass these in the ``secrets`` sequence of
|
||||
[DBConnection.create_function][lancedb.db.DBConnection.create_function]. The
|
||||
Function's source is unchanged by binding: it reads ``OPENAI_API_KEY`` the
|
||||
way it always did, and the binding is what puts a value there.
|
||||
|
||||
This is a local value. Constructing it contacts no server, so it always
|
||||
succeeds and says nothing about whether the Secret exists; that is checked
|
||||
at registration, where a mistyped Secret name surfaces as a clear "does not
|
||||
exist" naming both the Secret and the variable bound to it. A mistyped
|
||||
*variable* name cannot be caught anywhere -- nothing knows which variables a
|
||||
Function reads -- so it surfaces on the first rows instead.
|
||||
|
||||
The type exists so a credential cannot be passed by accident. A bare string
|
||||
in the same position is a plausible-looking mistake with the opposite
|
||||
meaning, and it reads identically in a diff.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
secret : str
|
||||
The Secret's database-scoped name.
|
||||
env_variable : str
|
||||
The environment variable the Function reads it from.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> from lancedb import EnvVarSecret
|
||||
>>> binding = EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")
|
||||
>>> binding.secret, binding.env_variable
|
||||
('openai-prod', 'OPENAI_API_KEY')
|
||||
"""
|
||||
|
||||
__slots__ = ("_secret", "_env_variable", "_namespace_path")
|
||||
|
||||
def __init__(self, secret: str, env_variable: str, *, namespace_path=None):
|
||||
self._secret = validate_secret_name(secret)
|
||||
self._env_variable = validate_env_variable(env_variable)
|
||||
self._namespace_path = validate_namespace_path(namespace_path)
|
||||
|
||||
@property
|
||||
def secret(self) -> str:
|
||||
"""The Secret's database-scoped name."""
|
||||
return self._secret
|
||||
|
||||
@property
|
||||
def env_variable(self) -> str:
|
||||
"""The environment variable the value is delivered in."""
|
||||
return self._env_variable
|
||||
|
||||
@property
|
||||
def namespace_path(self):
|
||||
"""The namespace path the Secret is addressed within, root when empty."""
|
||||
return list(self._namespace_path)
|
||||
|
||||
@property
|
||||
def secret_id(self) -> str:
|
||||
"""The full id the binding records: the path plus the name.
|
||||
|
||||
Pinned at construction, so a worker resolves the id it was handed and
|
||||
never re-resolves against its own default namespace -- the same
|
||||
Function resolves the same Secret wherever it runs.
|
||||
"""
|
||||
return secret_id(self._secret, self._namespace_path)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
path = (
|
||||
f", namespace_path={list(self._namespace_path)!r}"
|
||||
if self._namespace_path
|
||||
else ""
|
||||
)
|
||||
return (
|
||||
f"EnvVarSecret(secret={self._secret!r}, "
|
||||
f"env_variable={self._env_variable!r}{path})"
|
||||
)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return (
|
||||
isinstance(other, EnvVarSecret)
|
||||
and other._secret == self._secret
|
||||
and other._env_variable == self._env_variable
|
||||
and other._namespace_path == self._namespace_path
|
||||
)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(
|
||||
(EnvVarSecret, self._secret, self._env_variable, self._namespace_path)
|
||||
)
|
||||
|
||||
|
||||
class SecretInfo:
|
||||
"""What a database records about a Secret. Never its value.
|
||||
|
||||
Returned by
|
||||
[DBConnection.describe_secret][lancedb.db.DBConnection.describe_secret].
|
||||
"""
|
||||
|
||||
__slots__ = ("_name", "_created_at", "_updated_at")
|
||||
|
||||
def __init__(self, name: str, created_at: str, updated_at: str):
|
||||
self._name = name
|
||||
self._created_at = created_at
|
||||
self._updated_at = updated_at
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""The Secret's database-scoped name."""
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def created_at(self) -> str:
|
||||
"""When the Secret was created, as an RFC 3339 timestamp."""
|
||||
return self._created_at
|
||||
|
||||
@property
|
||||
def updated_at(self) -> str:
|
||||
"""When the Secret's value was last rotated, as an RFC 3339 timestamp."""
|
||||
return self._updated_at
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, value: dict) -> "SecretInfo":
|
||||
return cls(
|
||||
name=value["name"],
|
||||
created_at=value["created_at"],
|
||||
updated_at=value["updated_at"],
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"SecretInfo(name={self._name!r}, created_at={self._created_at!r}, "
|
||||
f"updated_at={self._updated_at!r})"
|
||||
)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return (
|
||||
isinstance(other, SecretInfo)
|
||||
and other._name == self._name
|
||||
and other._created_at == self._created_at
|
||||
and other._updated_at == self._updated_at
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"EnvVarSecret",
|
||||
"SecretInfo",
|
||||
"validate_env_variable",
|
||||
"validate_secret_name",
|
||||
]
|
||||
@@ -0,0 +1,88 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Handles to SQL queries running on a remote database."""
|
||||
|
||||
from uuid import UUID
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from lancedb.background_loop import LOOP
|
||||
|
||||
from . import _lancedb
|
||||
from .arrow import AsyncRecordBatchReader
|
||||
|
||||
QueryDescription = _lancedb.QueryDescription
|
||||
|
||||
|
||||
class AsyncQuery:
|
||||
"""A handle to a submitted SQL query on an asynchronous connection."""
|
||||
|
||||
def __init__(self, inner: "_lancedb.SqlQuery"):
|
||||
self._inner = inner
|
||||
|
||||
@property
|
||||
def id(self) -> UUID:
|
||||
"""The stable identifier scoped to the connection that submitted it."""
|
||||
return self._inner.id
|
||||
|
||||
async def describe(self) -> QueryDescription:
|
||||
"""Get a point-in-time description of the query."""
|
||||
return await self._inner.describe()
|
||||
|
||||
async def reader(self) -> AsyncRecordBatchReader:
|
||||
"""Wait for the initial result stream and return its Arrow reader.
|
||||
|
||||
Results are single-consumer. Calling this method more than once on the
|
||||
same query raises an error. Later batches are streamed as they become
|
||||
available without waiting for the full query to finish.
|
||||
"""
|
||||
return AsyncRecordBatchReader(await self._inner.reader())
|
||||
|
||||
async def cancel(self) -> None:
|
||||
"""Request cancellation of the query."""
|
||||
await self._inner.cancel()
|
||||
|
||||
|
||||
class Query:
|
||||
"""Synchronous counterpart of :class:`AsyncQuery`."""
|
||||
|
||||
def __init__(self, inner: AsyncQuery):
|
||||
self._inner = inner
|
||||
|
||||
@property
|
||||
def id(self) -> UUID:
|
||||
"""The stable identifier scoped to the connection that submitted it."""
|
||||
return self._inner.id
|
||||
|
||||
def describe(self) -> QueryDescription:
|
||||
"""Get a point-in-time description of the query."""
|
||||
return LOOP.run(self._inner.describe())
|
||||
|
||||
def reader(self) -> pa.RecordBatchReader:
|
||||
"""Wait for the initial result stream and return a blocking reader.
|
||||
|
||||
Results are single-consumer. Calling this method more than once on the
|
||||
same query raises an error. Later batches block only until they become
|
||||
available, without waiting for the full query to finish.
|
||||
"""
|
||||
reader = LOOP.run(self._inner.reader())
|
||||
|
||||
def next_batch():
|
||||
try:
|
||||
return LOOP.run(reader.__anext__())
|
||||
except StopAsyncIteration:
|
||||
return None
|
||||
|
||||
def batches():
|
||||
while (batch := next_batch()) is not None:
|
||||
yield batch
|
||||
|
||||
return pa.RecordBatchReader.from_batches(reader.schema, batches())
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Request cancellation of the query."""
|
||||
LOOP.run(self._inner.cancel())
|
||||
|
||||
|
||||
__all__ = ["AsyncQuery", "Query", "QueryDescription"]
|
||||
@@ -1547,7 +1547,9 @@ class Table(ABC):
|
||||
on: Union[str, Iterable[str]]
|
||||
A column (or columns) to join on. This is how records from the
|
||||
source table and target table are matched. Typically this is some
|
||||
kind of key or id column.
|
||||
kind of key or id column. Passing several columns matches on the
|
||||
composite key: a source row updates a target row only when it
|
||||
agrees on every one of them.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -1678,9 +1680,9 @@ class Table(ABC):
|
||||
Offsets are mostly useful for sampling as the set of all valid offsets is easily
|
||||
known in advance to be [0, len(table)).
|
||||
|
||||
No guarantees are made regarding the order in which results are returned. If
|
||||
you desire an output order that matches the order of the given offsets, you will
|
||||
need to add the row offset column to the output and align it yourself.
|
||||
No guarantees are made regarding the order in which results are returned.
|
||||
Repeated offsets produce repeated rows, which makes this method suitable for
|
||||
sampling with replacement.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -1791,6 +1793,9 @@ class Table(ABC):
|
||||
The result has the same length and order as ``row_ids``. Null blobs
|
||||
produce null slots; valid empty blobs produce ``b""``.
|
||||
|
||||
``_rowid`` values stay valid after compaction when the table has stable
|
||||
row ids.
|
||||
|
||||
Convenience for small payloads. For large values use
|
||||
:meth:`fetch_blob_files`.
|
||||
"""
|
||||
@@ -1808,6 +1813,9 @@ class Table(ABC):
|
||||
The result has the same length and order as ``requests``; null blobs
|
||||
produce null slots and empty ranges on non-null blobs produce ``b""``.
|
||||
|
||||
``_rowid`` values stay valid after compaction when the table has stable
|
||||
row ids.
|
||||
|
||||
Row IDs can be obtained from a query with ``with_row_id(True)``. This
|
||||
API is currently supported only by local tables.
|
||||
"""
|
||||
@@ -1823,6 +1831,9 @@ class Table(ABC):
|
||||
``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null
|
||||
rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or
|
||||
newer.
|
||||
|
||||
``_rowid`` values stay valid after compaction when the table has stable
|
||||
row ids.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -4090,6 +4101,7 @@ class LanceTable(Table):
|
||||
)
|
||||
and not self._route_pushdown_to_rust
|
||||
and self.current_branch() is None
|
||||
and query.take_offsets is None
|
||||
):
|
||||
from lancedb.namespace import _execute_server_side_query
|
||||
|
||||
@@ -5700,7 +5712,9 @@ class AsyncTable:
|
||||
on: Union[str, Iterable[str]]
|
||||
A column (or columns) to join on. This is how records from the
|
||||
source table and target table are matched. Typically this is some
|
||||
kind of key or id column.
|
||||
kind of key or id column. Passing several columns matches on the
|
||||
composite key: a source row updates a target row only when it
|
||||
agrees on every one of them.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -5983,7 +5997,23 @@ class AsyncTable:
|
||||
|
||||
def _sync_query_to_async(
|
||||
self, query: Query
|
||||
) -> AsyncHybridQuery | AsyncFTSQuery | AsyncVectorQuery | AsyncQuery:
|
||||
) -> (
|
||||
AsyncHybridQuery
|
||||
| AsyncFTSQuery
|
||||
| AsyncVectorQuery
|
||||
| AsyncQuery
|
||||
| AsyncTakeQuery
|
||||
):
|
||||
if query.take_offsets is not None:
|
||||
take_query = self.take_offsets(query.take_offsets)
|
||||
if query.columns:
|
||||
take_query = take_query.select(query.columns)
|
||||
if query.use_lsm is not None:
|
||||
take_query = take_query.use_lsm(query.use_lsm)
|
||||
if query.with_row_id:
|
||||
take_query = take_query.with_row_id()
|
||||
return take_query
|
||||
|
||||
async_query = self.query()
|
||||
if query.limit is not None:
|
||||
async_query = async_query.limit(query.limit)
|
||||
@@ -6048,6 +6078,7 @@ class AsyncTable:
|
||||
self._namespace_client, self._pushdown_operations
|
||||
)
|
||||
and not self._route_pushdown_to_rust
|
||||
and query.take_offsets is None
|
||||
):
|
||||
from lancedb.namespace import _execute_server_side_query
|
||||
|
||||
@@ -6545,6 +6576,9 @@ class AsyncTable:
|
||||
Offsets are mostly useful for sampling as the set of all valid offsets is easily
|
||||
known in advance to be [0, len(table)).
|
||||
|
||||
No guarantees are made regarding the order in which results are returned.
|
||||
Repeated offsets produce repeated rows.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
offsets: list[int]
|
||||
|
||||
@@ -66,6 +66,25 @@ def _row_ids_by_id(table):
|
||||
return dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
|
||||
|
||||
def _assert_missing_blob_row_ids(exc_info):
|
||||
message = str(exc_info.value)
|
||||
assert "row ids" in message
|
||||
assert "rowaddr" not in message
|
||||
assert "fragment" not in message
|
||||
|
||||
|
||||
def _assert_fetch_apis_reject_missing_row_ids(table, row_ids):
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
table.fetch_blobs("image", row_ids)
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
table.fetch_blob_files("image", row_ids)
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
table.fetch_blob_ranges("image", [(row_id, 0, 1) for row_id in row_ids])
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
|
||||
|
||||
def test_blob_factory_declares_v2_field():
|
||||
field = lancedb.blob("image")
|
||||
assert isinstance(field.type, pa.ExtensionType)
|
||||
@@ -691,6 +710,25 @@ def test_fetch_blobs_accepts_query_result():
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"}
|
||||
|
||||
|
||||
def test_fetch_blobs_after_compact_with_stable_row_ids(tmp_path):
|
||||
db = lancedb.connect(
|
||||
tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"}
|
||||
)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("t", schema=schema)
|
||||
table.add([{"id": 1, "image": b"frag-one"}])
|
||||
table.add([{"id": 2, "image": b"frag-two"}])
|
||||
by_id = _row_ids_by_id(table)
|
||||
ids = [by_id[1], by_id[2]]
|
||||
|
||||
table.optimize()
|
||||
|
||||
blobs = table.fetch_blobs("image", ids)
|
||||
assert blobs.to_pylist() == [b"frag-one", b"frag-two"]
|
||||
ranges = table.fetch_blob_ranges("image", [(ids[0], 5, 3), (ids[1], 5, 3)])
|
||||
assert ranges.to_pylist() == [b"one", b"two"]
|
||||
|
||||
|
||||
def test_fetch_blobs_preserves_null_and_empty_values():
|
||||
table = _blob_table(
|
||||
"nulls",
|
||||
@@ -739,8 +777,25 @@ def test_fetch_blob_ranges_validates_requests():
|
||||
with pytest.raises(ValueError, match="offset \\+ length overflowed"):
|
||||
table.fetch_blob_ranges("image", [(row_id, 2**64 - 1, 1)])
|
||||
|
||||
with pytest.raises(ValueError, match="row IDs"):
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)])
|
||||
_assert_missing_blob_row_ids(exc_info)
|
||||
|
||||
|
||||
def test_fetch_blob_apis_reject_missing_fragment_row_addr():
|
||||
table = _blob_table("missing_frag", [{"id": 1, "image": b"x"}])
|
||||
live = _row_ids_by_id(table)[1]
|
||||
_assert_fetch_apis_reject_missing_row_ids(table, [1 << 32, live])
|
||||
|
||||
|
||||
def test_fetch_blob_apis_reject_deleted_row_ids():
|
||||
table = _blob_table(
|
||||
"deleted_rows",
|
||||
[{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
table.delete("id = 2")
|
||||
_assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]])
|
||||
|
||||
|
||||
def test_fetch_blob_ranges_empty_requests_returns_empty_array():
|
||||
|
||||
@@ -37,6 +37,22 @@ def job_result(name: str) -> dict:
|
||||
return json.loads(fixture(name))["result"]
|
||||
|
||||
|
||||
def assert_no_secret_values(value):
|
||||
"""No client value models a resolved credential, at any nesting depth."""
|
||||
if isinstance(value, dict):
|
||||
for key, child in value.items():
|
||||
assert key not in {
|
||||
"secret_value",
|
||||
"secret_values",
|
||||
"resolved_secret",
|
||||
"resolved_secrets",
|
||||
}
|
||||
assert_no_secret_values(child)
|
||||
elif isinstance(value, list):
|
||||
for child in value:
|
||||
assert_no_secret_values(child)
|
||||
|
||||
|
||||
def test_public_function_values_are_in_api_reference():
|
||||
docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md"
|
||||
rendered = docs.read_text()
|
||||
@@ -94,6 +110,7 @@ def test_function_version_identity_is_immutable_and_exact():
|
||||
version = FunctionVersion.from_json(json.dumps(value))
|
||||
assert version.name == "embed"
|
||||
assert version.version == "fv_01K3EXACT"
|
||||
assert dict(version.secret_bindings) == {"HF_TOKEN": "hf-prod"}
|
||||
|
||||
with pytest.raises((TypeError, ValueError)):
|
||||
version.version = "fv_changed"
|
||||
@@ -276,6 +293,25 @@ def test_refresh_result_rejects_non_u64_values(field):
|
||||
RefreshColumnResult.from_json(json.dumps(value))
|
||||
|
||||
|
||||
def test_canonical_client_values_carry_bindings_and_no_credentials():
|
||||
"""A binding names a Secret; the credential behind it has no client field."""
|
||||
version = FunctionVersion.from_json(
|
||||
json.dumps(job_result("remote_function_job.json"))
|
||||
)
|
||||
canonical = json.loads(version.to_canonical_json())
|
||||
assert canonical["secret_bindings"] == {"HF_TOKEN": "hf-prod"}
|
||||
assert_no_secret_values(canonical)
|
||||
|
||||
|
||||
def test_a_version_without_bindings_keeps_the_original_wire_shape():
|
||||
"""Every Function registered before Secrets existed serializes unchanged."""
|
||||
value = job_result("remote_function_job.json")
|
||||
del value["secret_bindings"]
|
||||
version = FunctionVersion.from_json(json.dumps(value))
|
||||
assert dict(version.secret_bindings) == {}
|
||||
assert "secret_bindings" not in json.loads(version.to_canonical_json())
|
||||
|
||||
|
||||
class _FunctionDeclarationInner:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
@@ -12,6 +12,8 @@ from datetime import date
|
||||
import http.server
|
||||
import json
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
@@ -26,6 +28,7 @@ from lancedb.functions import (
|
||||
_GRAMMAR_PRIMITIVES,
|
||||
udf,
|
||||
)
|
||||
from lancedb.secrets import EnvVarSecret
|
||||
|
||||
THRESHOLD = 20
|
||||
_CACHE = None
|
||||
@@ -51,6 +54,15 @@ def normalize_score(value: float) -> float:
|
||||
return value / 100.0
|
||||
|
||||
|
||||
@udf(
|
||||
pip=["openai==3.7.0"],
|
||||
env={"MODE": "test"},
|
||||
python_version="3.12",
|
||||
)
|
||||
def analyze_caption(caption: str) -> str:
|
||||
return caption.strip()
|
||||
|
||||
|
||||
def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
assert isinstance(normalize_score, UdfDefinition)
|
||||
assert normalize_score(25.0) == 0.25
|
||||
@@ -67,6 +79,209 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
}
|
||||
|
||||
|
||||
def test_secret_bound_udf_matches_its_shared_registration_golden():
|
||||
assert analyze_caption(" hello ") == "hello"
|
||||
bound = analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")]
|
||||
)
|
||||
assert (
|
||||
bound.to_canonical_json()
|
||||
== (FIXTURES / "remote_function_secret_registration_request.canonical.json")
|
||||
.read_text()
|
||||
.strip()
|
||||
)
|
||||
|
||||
|
||||
def test_a_namespaced_binding_records_the_full_id():
|
||||
"""A binding records the id the service resolves: path plus name.
|
||||
|
||||
At the root that is the bare name, so a root binding is unchanged from
|
||||
before namespaces existed -- which is what keeps the wire shape stable.
|
||||
"""
|
||||
root = EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")
|
||||
assert root.secret_id == "openai-prod"
|
||||
assert root.namespace_path == []
|
||||
|
||||
nested = EnvVarSecret(
|
||||
secret="openai-prod",
|
||||
env_variable="OPENAI_API_KEY",
|
||||
namespace_path=["prod", "vision"],
|
||||
)
|
||||
assert nested.secret_id == "prod$vision$openai-prod"
|
||||
assert nested.namespace_path == ["prod", "vision"]
|
||||
assert nested != root
|
||||
|
||||
bound = analyze_caption.bind_secrets([nested])
|
||||
assert bound.secret_bindings == {"OPENAI_API_KEY": "prod$vision$openai-prod"}
|
||||
|
||||
|
||||
def test_a_namespace_path_is_validated_locally():
|
||||
# The delimiter is why the charset is closed: a segment carrying one would
|
||||
# make the same id parse two ways.
|
||||
with pytest.raises(ValueError):
|
||||
EnvVarSecret(
|
||||
secret="openai-prod", env_variable="K", namespace_path=["with$delim"]
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
EnvVarSecret(secret="openai-prod", env_variable="K", namespace_path=["a/b"])
|
||||
# A bare string is a plausible mistake with the wrong meaning.
|
||||
with pytest.raises(TypeError):
|
||||
EnvVarSecret(secret="openai-prod", env_variable="K", namespace_path="prod")
|
||||
|
||||
|
||||
def test_an_unbound_request_carries_no_binding_at_all():
|
||||
"""Binding is a registration-time decision, so the definition holds none.
|
||||
|
||||
The decorator declares nothing about secrets, which is what makes the PRD's
|
||||
claim true: a Function's source and its registration request are identical
|
||||
whether or not a credential is later bound to it.
|
||||
"""
|
||||
unbound = json.loads(analyze_caption.registration_request.to_canonical_json())
|
||||
assert "secret_bindings" not in unbound
|
||||
assert "OPENAI_API_KEY" not in json.dumps(unbound)
|
||||
|
||||
|
||||
def test_binding_a_secret_leaves_the_packaged_artifact_untouched():
|
||||
"""The artifact is source bytes and nothing else, with or without secrets."""
|
||||
bound = analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")]
|
||||
)
|
||||
assert bound.artifact == analyze_caption.registration_request.artifact
|
||||
assert bound.artifact.digest == analyze_caption.registration_request.artifact.digest
|
||||
|
||||
|
||||
def test_a_function_declaring_no_secret_is_registered_exactly_as_before():
|
||||
"""The compatibility claim: nothing about the no-secret path moves."""
|
||||
assert (
|
||||
normalize_score.bind_secrets(None).to_canonical_json()
|
||||
== normalize_score.registration_request.to_canonical_json()
|
||||
)
|
||||
assert (
|
||||
"secret_bindings"
|
||||
not in normalize_score.registration_request.to_canonical_json()
|
||||
)
|
||||
|
||||
|
||||
def test_a_function_binds_each_variable_once():
|
||||
with pytest.raises(ValueError, match="binds each environment variable once"):
|
||||
analyze_caption.bind_secrets(
|
||||
[
|
||||
EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY"),
|
||||
EnvVarSecret(secret="openai-staging", env_variable="OPENAI_API_KEY"),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_bindings_may_not_collide_with_plain_configuration():
|
||||
"""`env` is stored with the Function; a Secret is not. Refuse, do not pick."""
|
||||
with pytest.raises(ValueError, match="must be disjoint"):
|
||||
analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret="mode-prod", env_variable="MODE")]
|
||||
)
|
||||
|
||||
|
||||
def test_a_function_binds_at_most_sixteen_secrets():
|
||||
bindings = [
|
||||
EnvVarSecret(secret=f"secret-{index}", env_variable=f"TOKEN_{index}")
|
||||
for index in range(17)
|
||||
]
|
||||
with pytest.raises(ValueError, match="at most 16 secrets"):
|
||||
normalize_score.bind_secrets(bindings)
|
||||
|
||||
|
||||
def test_a_credential_value_is_rejected_in_the_binding_position():
|
||||
"""The one mistake the typed binding exists to stop."""
|
||||
with pytest.raises(TypeError, match="EnvVarSecret"):
|
||||
analyze_caption.bind_secrets(["sk-live-0001"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("secret", "variable", "message"),
|
||||
[
|
||||
("openai-prod", "not-a-var", "invalid environment variable name"),
|
||||
("openai-prod", "API-TOKEN", "invalid environment variable name"),
|
||||
("not a name", "API_TOKEN", "invalid Secret name"),
|
||||
],
|
||||
)
|
||||
def test_a_binding_validates_both_names_locally(secret, variable, message):
|
||||
with pytest.raises(ValueError, match=message):
|
||||
EnvVarSecret(secret=secret, env_variable=variable)
|
||||
|
||||
|
||||
def _main_udf_source(
|
||||
*, threshold: int = 20, input_annotation: str = "int", comparison: str = ">="
|
||||
) -> str:
|
||||
return (
|
||||
"from __future__ import annotations\n"
|
||||
"from lancedb.functions import udf\n"
|
||||
f"THRESHOLD = {threshold}\n"
|
||||
"\n"
|
||||
"@udf\n"
|
||||
f"def label(value: {input_annotation}) -> str:\n"
|
||||
f" return 'big' if value {comparison} THRESHOLD else 'small'\n"
|
||||
"\n"
|
||||
"assert label.__module__ == '__main__'\n"
|
||||
"print(label.registration_request.to_canonical_json())\n"
|
||||
)
|
||||
|
||||
|
||||
def _run_main_udf(path: Path, source: str) -> dict:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(source)
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(path)],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
return json.loads(result.stdout)
|
||||
|
||||
|
||||
def test_main_udf_registration_identity_is_stable_across_processes_and_paths(
|
||||
tmp_path,
|
||||
):
|
||||
source = _main_udf_source()
|
||||
original_path = tmp_path / "original" / "job.py"
|
||||
moved_path = tmp_path / "moved" / "renamed_job.py"
|
||||
|
||||
original_runs = [_run_main_udf(original_path, source) for _ in range(2)]
|
||||
moved_run = _run_main_udf(moved_path, source)
|
||||
|
||||
assert len({run["artifact"]["digest"] for run in [*original_runs, moved_run]}) == 1
|
||||
assert all(
|
||||
run["signature"] == original_runs[0]["signature"]
|
||||
for run in [original_runs[1], moved_run]
|
||||
)
|
||||
assert original_runs[0] == original_runs[1] == moved_run
|
||||
|
||||
body_change = _run_main_udf(
|
||||
tmp_path / "changes" / "body.py", _main_udf_source(comparison=">")
|
||||
)
|
||||
global_change = _run_main_udf(
|
||||
tmp_path / "changes" / "global.py", _main_udf_source(threshold=21)
|
||||
)
|
||||
annotation_change = _run_main_udf(
|
||||
tmp_path / "changes" / "annotation.py",
|
||||
_main_udf_source(input_annotation="float"),
|
||||
)
|
||||
|
||||
baseline = original_runs[0]
|
||||
assert baseline["signature"] == body_change["signature"]
|
||||
assert baseline["signature"] == global_change["signature"]
|
||||
assert baseline["signature"] != annotation_change["signature"]
|
||||
assert (
|
||||
len(
|
||||
{
|
||||
baseline["artifact"]["digest"],
|
||||
body_change["artifact"]["digest"],
|
||||
global_change["artifact"]["digest"],
|
||||
annotation_change["artifact"]["digest"],
|
||||
}
|
||||
)
|
||||
== 4
|
||||
)
|
||||
|
||||
|
||||
def _run_packaged(definition, *args):
|
||||
"""Execute the shipped artifact in a fresh namespace, as a worker would."""
|
||||
source = base64.b64decode(definition.registration_request.artifact.content.data)
|
||||
@@ -592,6 +807,167 @@ def test_blob_fields_use_the_scalar_function_semantic_type():
|
||||
assert signature.output.arrow_type == "blob_v2"
|
||||
|
||||
|
||||
def test_whole_named_struct_function_can_include_a_blob_result_field():
|
||||
@udf(
|
||||
input_schema=pa.schema([lancedb.blob("image", nullable=False)]),
|
||||
output_schema=pa.field(
|
||||
"payload",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("mime_type", pa.string(), nullable=False),
|
||||
lancedb.blob("image", nullable=False),
|
||||
]
|
||||
),
|
||||
nullable=False,
|
||||
),
|
||||
)
|
||||
def inspect_blob(image):
|
||||
return {"mime_type": "image/png", "image": image}
|
||||
|
||||
output = inspect_blob.registration_request.signature.output
|
||||
assert output.kind == "named_struct"
|
||||
assert [(field.name, field.arrow_type) for field in output.fields] == [
|
||||
("mime_type", "utf8"),
|
||||
("image", "blob_v2"),
|
||||
]
|
||||
|
||||
|
||||
def test_struct_blob_signature_fields_preserve_exact_metadata_and_nullability():
|
||||
nested_input = pa.field(
|
||||
"payload",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("mime_type", pa.string(), nullable=False),
|
||||
pa.field(
|
||||
"nested",
|
||||
pa.struct([lancedb.blob("image", nullable=True)]),
|
||||
nullable=True,
|
||||
),
|
||||
]
|
||||
),
|
||||
nullable=True,
|
||||
)
|
||||
nested_output = pa.field(
|
||||
"result",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("mime_type", pa.string(), nullable=False),
|
||||
pa.field(
|
||||
"nested",
|
||||
pa.struct([lancedb.blob("image", nullable=True)]),
|
||||
nullable=False,
|
||||
),
|
||||
]
|
||||
),
|
||||
nullable=False,
|
||||
)
|
||||
|
||||
@udf(input_schema=pa.schema([nested_input]), output_schema=nested_output)
|
||||
def copy_payload(payload):
|
||||
return payload
|
||||
|
||||
signature = copy_payload.registration_request.signature
|
||||
input_type = json.loads(signature.inputs[0].arrow_type)
|
||||
assert input_type["fields"][1]["nullable"] is True
|
||||
input_blob = input_type["fields"][1]["type"]["fields"][0]
|
||||
assert input_blob["nullable"] is True
|
||||
assert input_blob["metadata"] == {"ARROW:extension:name": "lance.blob.v2"}
|
||||
|
||||
assert signature.output.kind == "named_struct"
|
||||
nested_result = next(
|
||||
field for field in signature.output.fields if field.name == "nested"
|
||||
)
|
||||
output_type = json.loads(nested_result.arrow_type)
|
||||
output_blob = output_type["fields"][0]
|
||||
assert output_blob["nullable"] is True
|
||||
assert output_blob["metadata"] == {"ARROW:extension:name": "lance.blob.v2"}
|
||||
|
||||
|
||||
def test_struct_blob_signature_supports_multiple_struct_levels():
|
||||
recursive = pa.field(
|
||||
"value",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field(
|
||||
"level_1",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field(
|
||||
"level_2",
|
||||
pa.struct([lancedb.blob("image", nullable=False)]),
|
||||
nullable=False,
|
||||
)
|
||||
]
|
||||
),
|
||||
nullable=False,
|
||||
)
|
||||
]
|
||||
),
|
||||
nullable=False,
|
||||
)
|
||||
|
||||
@udf(
|
||||
input_schema=pa.schema([recursive]),
|
||||
output_schema=pa.field("size", pa.int64(), nullable=False),
|
||||
)
|
||||
def blob_size(value):
|
||||
return len(value["level_1"]["level_2"]["image"])
|
||||
|
||||
encoded = json.loads(blob_size.registration_request.signature.inputs[0].arrow_type)
|
||||
blob = encoded["fields"][0]["type"]["fields"][0]["type"]["fields"][0]
|
||||
assert blob["metadata"]["ARROW:extension:name"] == "lance.blob.v2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_type",
|
||||
[
|
||||
pa.list_(lancedb.blob("item", nullable=False)),
|
||||
pa.large_list(lancedb.blob("item", nullable=False)),
|
||||
pa.list_(lancedb.blob("item", nullable=False), 2),
|
||||
pa.map_(pa.string(), lancedb.blob("value", nullable=False).type),
|
||||
],
|
||||
)
|
||||
def test_blob_signature_rejects_collection_ancestors(data_type):
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match="Blob v2 fields nested under collection types are not supported",
|
||||
):
|
||||
|
||||
@udf(
|
||||
input_schema=pa.schema([pa.field("value", data_type, nullable=False)]),
|
||||
output_schema=pa.field("size", pa.int64(), nullable=False),
|
||||
)
|
||||
def blob_size(value):
|
||||
return len(value)
|
||||
|
||||
|
||||
def test_blob_signature_rejects_collection_below_a_struct():
|
||||
nested = pa.field(
|
||||
"value",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field(
|
||||
"images",
|
||||
pa.list_(lancedb.blob("item", nullable=False)),
|
||||
nullable=False,
|
||||
)
|
||||
]
|
||||
),
|
||||
nullable=False,
|
||||
)
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match="Blob v2 fields nested under collection types are not supported",
|
||||
):
|
||||
|
||||
@udf(
|
||||
input_schema=pa.schema([nested]),
|
||||
output_schema=pa.field("size", pa.int64(), nullable=False),
|
||||
)
|
||||
def blob_size(value):
|
||||
return len(value["images"])
|
||||
|
||||
|
||||
def test_named_struct_function_can_include_a_blob_result_field():
|
||||
@udf(
|
||||
input_schema=pa.schema([lancedb.blob("image", nullable=False)]),
|
||||
@@ -613,6 +989,43 @@ def test_named_struct_function_can_include_a_blob_result_field():
|
||||
]
|
||||
|
||||
|
||||
def test_named_struct_function_preserves_nullable_result_fields():
|
||||
@udf(
|
||||
input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]),
|
||||
output_schema=pa.schema(
|
||||
[
|
||||
pa.field("result", pa.int64(), nullable=True),
|
||||
pa.field("failure_code", pa.int32(), nullable=False),
|
||||
]
|
||||
),
|
||||
)
|
||||
def nullable_result(value):
|
||||
return {"result": value, "failure_code": 0}
|
||||
|
||||
output = nullable_result.registration_request.signature.output
|
||||
assert [(field.name, field.nullable) for field in output.fields] == [
|
||||
("result", True),
|
||||
("failure_code", False),
|
||||
]
|
||||
|
||||
@udf(
|
||||
input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]),
|
||||
output_schema=pa.schema(
|
||||
[
|
||||
pa.field("result", pa.int64(), nullable=True),
|
||||
pa.field("failure_code", pa.int32(), nullable=True),
|
||||
]
|
||||
),
|
||||
)
|
||||
def all_nullable(value):
|
||||
return {"result": value, "failure_code": None}
|
||||
|
||||
assert all(
|
||||
field.nullable
|
||||
for field in all_nullable.registration_request.signature.output.fields
|
||||
)
|
||||
|
||||
|
||||
def test_metadata_marked_blob_field_uses_the_semantic_type():
|
||||
extension = lancedb.blob("image", nullable=False).type
|
||||
storage = (
|
||||
@@ -653,22 +1066,6 @@ def test_blob_marker_rejects_invalid_storage_layout():
|
||||
return len(image)
|
||||
|
||||
|
||||
def test_nested_blob_signature_field_has_a_clear_error():
|
||||
nested = pa.field(
|
||||
"value",
|
||||
pa.struct([lancedb.blob("image", nullable=False)]),
|
||||
nullable=False,
|
||||
)
|
||||
with pytest.raises(TypeError, match="nested Blob v2 fields are not supported"):
|
||||
|
||||
@udf(
|
||||
input_schema=pa.schema([nested]),
|
||||
output_schema=pa.field("size", pa.int64(), nullable=False),
|
||||
)
|
||||
def blob_size(value):
|
||||
return len(value["image"])
|
||||
|
||||
|
||||
def test_nested_non_blob_extension_is_not_silently_unwrapped():
|
||||
class TestExtension(pa.ExtensionType):
|
||||
def __init__(self):
|
||||
@@ -941,6 +1338,10 @@ def test_local_function_catalog_operations_are_not_supported(tmp_path):
|
||||
db.create_function_async(normalize_score)
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.get_function("normalize_score", version="fv_exact")
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.list_functions()
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
db.drop_function("normalize_score", version="fv_exact")
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
@@ -968,6 +1369,7 @@ def _mock_remote_function_catalog():
|
||||
"runtime": body["runtime"],
|
||||
"runtime_digest": "sha256:runtime",
|
||||
"environment_digest": "sha256:environment",
|
||||
"secret_bindings": body.get("secret_bindings", {}),
|
||||
"created_at": "2026-08-21T00:00:00Z",
|
||||
}
|
||||
response = {"job_id": "job-register"}
|
||||
@@ -986,6 +1388,43 @@ def _mock_remote_function_catalog():
|
||||
"version": "fv_exact",
|
||||
}
|
||||
response = state["version"]
|
||||
elif self.path == "/v1/functions/list":
|
||||
assert body["include_definition"] is True
|
||||
if "page_token" not in body:
|
||||
response = {
|
||||
"functions": [
|
||||
{
|
||||
"name": "normalize_score",
|
||||
"version": "fv_exact",
|
||||
"definition": state["version"],
|
||||
}
|
||||
],
|
||||
"page_token": "next",
|
||||
}
|
||||
else:
|
||||
assert body["page_token"] == "next"
|
||||
response = {"functions": []}
|
||||
elif self.path == "/v1/functions/drop":
|
||||
assert body == {
|
||||
"name": "normalize_score",
|
||||
"version": "fv_exact",
|
||||
}
|
||||
response = {"dropped": True}
|
||||
elif self.path in ("/v1/secrets/create", "/v1/secrets/alter"):
|
||||
assert set(body) == {"name", "value"}
|
||||
response = {}
|
||||
elif self.path == "/v1/secrets/list":
|
||||
if "page_token" not in body:
|
||||
response = {
|
||||
"secrets": [{"name": "openai-prod"}],
|
||||
"page_token": "next",
|
||||
}
|
||||
else:
|
||||
assert body["page_token"] == "next"
|
||||
response = {"secrets": [{"name": "hf-prod"}]}
|
||||
elif self.path == "/v1/secrets/drop":
|
||||
assert body == {"name": "openai-prod"}
|
||||
response = {}
|
||||
else:
|
||||
status = 404
|
||||
response = {"error": "not found"}
|
||||
@@ -1028,6 +1467,71 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
|
||||
)
|
||||
|
||||
|
||||
def test_remote_registration_sends_bindings_and_never_a_credential():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
created = db.create_function(
|
||||
analyze_caption,
|
||||
secrets=[EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")],
|
||||
)
|
||||
|
||||
assert dict(created.secret_bindings) == {"OPENAI_API_KEY": "openai-prod"}
|
||||
path, create_request = state["requests"][0]
|
||||
assert path == "/v1/functions/create"
|
||||
assert create_request["secret_bindings"] == {"OPENAI_API_KEY": "openai-prod"}
|
||||
# The request names a Secret and carries nothing that could be one.
|
||||
assert create_request == json.loads(
|
||||
analyze_caption.bind_secrets(
|
||||
[EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")]
|
||||
).to_canonical_json()
|
||||
)
|
||||
|
||||
|
||||
def test_remote_secret_verbs_round_trip():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert db.create_secret("openai-prod", "sk-live-0001") is None
|
||||
assert db.alter_secret("openai-prod", "sk-live-0002") is None
|
||||
assert db.list_secrets() == ["openai-prod", "hf-prod"]
|
||||
assert db.drop_secret("openai-prod") is None
|
||||
|
||||
routes = [path for path, _ in state["requests"]]
|
||||
assert routes == [
|
||||
"/v1/secrets/create",
|
||||
"/v1/secrets/alter",
|
||||
"/v1/secrets/list",
|
||||
"/v1/secrets/list",
|
||||
"/v1/secrets/drop",
|
||||
]
|
||||
assert state["requests"][0][1] == {"name": "openai-prod", "value": "sk-live-0001"}
|
||||
# The listing returns names, and the client has no way to ask for more.
|
||||
assert state["requests"][2][1] == {}
|
||||
|
||||
|
||||
def test_building_a_binding_contacts_no_server():
|
||||
"""A binding is a local value: it says nothing about whether the Secret exists.
|
||||
|
||||
Existence is the server's answer at registration, where a mistyped name is a
|
||||
clear error rather than a client-side check that was already stale.
|
||||
"""
|
||||
with _mock_remote_function_catalog() as (_host, state):
|
||||
binding = EnvVarSecret(secret="openai-prod", env_variable="OPENAI_API_KEY")
|
||||
assert binding.secret == "openai-prod"
|
||||
assert binding.env_variable == "OPENAI_API_KEY"
|
||||
|
||||
assert state["requests"] == []
|
||||
|
||||
|
||||
def test_blocking_remote_registration_returns_function_version():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
@@ -1044,3 +1548,83 @@ def test_blocking_remote_registration_returns_function_version():
|
||||
"/v1/functions/create",
|
||||
"/v1/jobs/describe",
|
||||
]
|
||||
|
||||
|
||||
def test_remote_list_functions_paginates_and_returns_typed_versions():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
created = db.create_function(normalize_score)
|
||||
state["requests"].clear()
|
||||
functions = db.list_functions()
|
||||
|
||||
assert functions == [created]
|
||||
assert state["requests"] == [
|
||||
("/v1/functions/list", {"include_definition": True}),
|
||||
(
|
||||
"/v1/functions/list",
|
||||
{"include_definition": True, "page_token": "next"},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_list_functions_returns_typed_versions():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = await lancedb.connect_async(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
registration = await db.create_function_async(normalize_score)
|
||||
created = await registration.wait()
|
||||
state["requests"].clear()
|
||||
functions = await db.list_functions()
|
||||
|
||||
assert functions == [created]
|
||||
assert [path for path, _ in state["requests"]] == [
|
||||
"/v1/functions/list",
|
||||
"/v1/functions/list",
|
||||
]
|
||||
|
||||
|
||||
def test_remote_drop_function_sends_exact_version():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = lancedb.connect(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert db.drop_function("normalize_score", version="fv_exact") is True
|
||||
|
||||
assert state["requests"] == [
|
||||
(
|
||||
"/v1/functions/drop",
|
||||
{"name": "normalize_score", "version": "fv_exact"},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_drop_function_sends_exact_version():
|
||||
with _mock_remote_function_catalog() as (host, state):
|
||||
db = await lancedb.connect_async(
|
||||
"db://dev",
|
||||
api_key="fake",
|
||||
host_override=host,
|
||||
client_config={"retry_config": {"retries": 0}},
|
||||
)
|
||||
assert await db.drop_function("normalize_score", version="fv_exact") is True
|
||||
|
||||
assert state["requests"] == [
|
||||
(
|
||||
"/v1/functions/drop",
|
||||
{"name": "normalize_score", "version": "fv_exact"},
|
||||
)
|
||||
]
|
||||
|
||||
@@ -54,7 +54,10 @@ class TestOAuthProvider:
|
||||
provider = OAuthProvider(fetcher)
|
||||
headers = provider.get_headers()
|
||||
|
||||
assert headers == {"Authorization": "Bearer token123"}
|
||||
assert headers == {
|
||||
"Authorization": "Bearer token123",
|
||||
"x-lancedb-credential-type": "oidc",
|
||||
}
|
||||
assert provider._current_token == "token123"
|
||||
assert provider._token_expires_at is not None
|
||||
|
||||
@@ -73,14 +76,20 @@ class TestOAuthProvider:
|
||||
|
||||
# First call
|
||||
headers1 = provider.get_headers()
|
||||
assert headers1 == {"Authorization": "Bearer token1"}
|
||||
assert headers1 == {
|
||||
"Authorization": "Bearer token1",
|
||||
"x-lancedb-credential-type": "oidc",
|
||||
}
|
||||
|
||||
# Wait for token to expire
|
||||
time.sleep(1.1)
|
||||
|
||||
# Second call should refresh
|
||||
headers2 = provider.get_headers()
|
||||
assert headers2 == {"Authorization": "Bearer token2"}
|
||||
assert headers2 == {
|
||||
"Authorization": "Bearer token2",
|
||||
"x-lancedb-credential-type": "oidc",
|
||||
}
|
||||
assert call_count == 2
|
||||
|
||||
def test_no_expiry_info(self):
|
||||
@@ -92,12 +101,18 @@ class TestOAuthProvider:
|
||||
provider = OAuthProvider(fetcher)
|
||||
headers = provider.get_headers()
|
||||
|
||||
assert headers == {"Authorization": "Bearer permanent_token"}
|
||||
assert headers == {
|
||||
"Authorization": "Bearer permanent_token",
|
||||
"x-lancedb-credential-type": "oidc",
|
||||
}
|
||||
assert provider._token_expires_at is None
|
||||
|
||||
# Should not refresh on second call
|
||||
headers2 = provider.get_headers()
|
||||
assert headers2 == {"Authorization": "Bearer permanent_token"}
|
||||
assert headers2 == {
|
||||
"Authorization": "Bearer permanent_token",
|
||||
"x-lancedb-credential-type": "oidc",
|
||||
}
|
||||
|
||||
def test_missing_access_token(self):
|
||||
"""Test error handling when access_token is missing."""
|
||||
@@ -121,7 +136,10 @@ class TestOAuthProvider:
|
||||
provider = OAuthProvider(fetcher)
|
||||
headers = provider.get_headers()
|
||||
|
||||
assert headers == {"Authorization": "Bearer sync_token"}
|
||||
assert headers == {
|
||||
"Authorization": "Bearer sync_token",
|
||||
"x-lancedb-credential-type": "oidc",
|
||||
}
|
||||
|
||||
|
||||
class TestClientConfigIntegration:
|
||||
|
||||
@@ -266,3 +266,38 @@ async def test_async_namespace_connection_materialized_views(tmp_path):
|
||||
handle._route_pushdown_to_rust == through_namespace._route_pushdown_to_rust
|
||||
)
|
||||
assert handle._namespace_path == through_namespace._namespace_path
|
||||
|
||||
|
||||
def test_namespaced_select_kind_is_read_and_unknown_kinds_are_refused():
|
||||
import json
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from lancedb.materialized_view import _definition_from_schema
|
||||
|
||||
def schema_with(definition: dict) -> pa.Schema:
|
||||
return pa.schema([pa.field("id", pa.int32())]).with_metadata(
|
||||
{b"mv.definition": json.dumps(definition).encode()}
|
||||
)
|
||||
|
||||
# "namespaced_select" is the namespaced form of "select": same shape,
|
||||
# a separate kind so readers that predate it refuse instead of
|
||||
# resolving the source at the root.
|
||||
definition = _definition_from_schema(
|
||||
schema_with(
|
||||
{
|
||||
"kind": "namespaced_select",
|
||||
"source_table": "people",
|
||||
"source_namespace": ["ns"],
|
||||
"projections": [{"output": "name", "expression": "name"}],
|
||||
}
|
||||
),
|
||||
"v",
|
||||
)
|
||||
assert definition.source_table == "people"
|
||||
assert definition.source_namespace == ["ns"]
|
||||
|
||||
with pytest.raises(NotImplementedError, match="cannot refresh"):
|
||||
_definition_from_schema(
|
||||
schema_with({"kind": "select_v3", "source_table": "people"}), "v"
|
||||
)
|
||||
|
||||
@@ -1923,6 +1923,21 @@ def test_take_queries(tmp_path):
|
||||
17,
|
||||
]
|
||||
|
||||
# Duplicate offsets are occurrences, not set members. Ordering is unspecified.
|
||||
assert sorted(table.take_offsets([5, 2, 5, 17]).to_pandas()["idx"].to_list()) == [
|
||||
2,
|
||||
5,
|
||||
5,
|
||||
17,
|
||||
]
|
||||
|
||||
# Converting a take builder to its serializable query representation must
|
||||
# retain occurrence metadata and execute with the same multiplicity.
|
||||
query = table.take_offsets([5, 2, 5, 17]).select(["idx"]).to_query_object()
|
||||
assert query.take_offsets == [5, 2, 5, 17]
|
||||
converted = table._execute_query(query).read_all()
|
||||
assert sorted(converted["idx"].to_pylist()) == [2, 5, 5, 17]
|
||||
|
||||
# Take by row id
|
||||
assert list(
|
||||
sorted(table.take_row_ids([5, 2, 17]).to_pandas()["idx"].to_list())
|
||||
|
||||
@@ -479,24 +479,49 @@ def test_remote_permutation_is_picklable():
|
||||
match = re.search(
|
||||
r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE
|
||||
)
|
||||
offsets = [int(o.strip()) for o in match.group(1).split(",")]
|
||||
offsets = list(
|
||||
dict.fromkeys(int(o.strip()) for o in match.group(1).split(","))
|
||||
)
|
||||
else:
|
||||
offsets = list(range(len(rows)))
|
||||
table = pa.table({"a": [rows[offset] for offset in offsets]})
|
||||
columns = body.get("columns") or ["a"]
|
||||
table = pa.table(
|
||||
{
|
||||
column: (
|
||||
[rows[offset] for offset in offsets]
|
||||
if column == "a"
|
||||
else offsets
|
||||
)
|
||||
for column in columns
|
||||
}
|
||||
)
|
||||
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/vnd.apache.arrow.file")
|
||||
request.end_headers()
|
||||
with pa.ipc.new_file(request.wfile, schema=table.schema) as writer:
|
||||
writer.write_table(table)
|
||||
writer.write_table(table, max_chunksize=2)
|
||||
else:
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
permutation = Permutation.identity(db.open_table("test"))
|
||||
table = db.open_table("test")
|
||||
assert table.take_offsets([0, 2, 0, 4]).to_list() == [
|
||||
{"a": 0},
|
||||
{"a": 0},
|
||||
{"a": 2},
|
||||
{"a": 4},
|
||||
]
|
||||
|
||||
permutation = Permutation.identity(table)
|
||||
restored = pickle.loads(pickle.dumps(permutation))
|
||||
assert restored.__getitems__([0, 2, 4]) == [{"a": 0}, {"a": 2}, {"a": 4}]
|
||||
assert restored.__getitems__([0, 2, 0, 4]) == [
|
||||
{"a": 0},
|
||||
{"a": 2},
|
||||
{"a": 0},
|
||||
{"a": 4},
|
||||
]
|
||||
|
||||
|
||||
def test_create_table_exist_ok():
|
||||
@@ -795,11 +820,13 @@ def test_table_create_indices():
|
||||
scalar_req = received_requests[0]
|
||||
assert "name" in scalar_req
|
||||
assert scalar_req["name"] == "custom_scalar_idx"
|
||||
assert scalar_req["replace"] is False
|
||||
|
||||
# Check FTS index request has custom name
|
||||
fts_req = received_requests[1]
|
||||
assert "name" in fts_req
|
||||
assert fts_req["name"] == "custom_fts_idx"
|
||||
assert fts_req["replace"] is False
|
||||
assert fts_req["block_size"] == 256
|
||||
assert fts_req["custom_stop_words"] == ["cloud"]
|
||||
|
||||
@@ -807,6 +834,7 @@ def test_table_create_indices():
|
||||
vector_req = received_requests[2]
|
||||
assert "name" in vector_req
|
||||
assert vector_req["name"] == "custom_vector_idx"
|
||||
assert "replace" not in vector_req
|
||||
|
||||
table.wait_for_index(["custom_scalar_idx"], timedelta(seconds=2))
|
||||
table.wait_for_index(
|
||||
@@ -1079,6 +1107,9 @@ def test_remote_create_index_new_api():
|
||||
table.create_index("text", config=FTS(block_size=256))
|
||||
# IvfRq via new API
|
||||
table.create_index("vector", config=IvfRq(distance_type="l2"))
|
||||
table.create_index(
|
||||
"vector", config=IvfPq(distance_type="l2"), replace=False
|
||||
)
|
||||
|
||||
# Legacy index_type="IVF_RQ" routes to IvfRq config under the hood.
|
||||
with pytest.warns(DeprecationWarning, match="create_index"):
|
||||
@@ -1088,15 +1119,17 @@ def test_remote_create_index_new_api():
|
||||
num_partitions=8,
|
||||
)
|
||||
|
||||
assert len(received_requests) == 5
|
||||
assert len(received_requests) == 6
|
||||
assert [req["column"] for req in received_requests] == [
|
||||
"vector",
|
||||
"category",
|
||||
"text",
|
||||
"vector",
|
||||
"vector",
|
||||
"vector",
|
||||
]
|
||||
assert received_requests[2]["block_size"] == 256
|
||||
assert received_requests[4]["replace"] is False
|
||||
|
||||
|
||||
def test_table_wait_for_index_timeout():
|
||||
@@ -2434,7 +2467,7 @@ def test_remote_blob_byte_apis_not_supported_on_old_server():
|
||||
|
||||
|
||||
def test_remote_connection_jobs_surface():
|
||||
from lancedb.exceptions import JobFailedError
|
||||
from lancedb.exceptions import JobFailedError, JobNotFoundError
|
||||
|
||||
schema = pa.schema([("state", pa.string())])
|
||||
batch = pa.record_batch([pa.array(["created", "done"])], schema=schema)
|
||||
@@ -2442,6 +2475,7 @@ def test_remote_connection_jobs_surface():
|
||||
with pa.ipc.new_stream(sink, schema) as writer:
|
||||
writer.write_batch(batch)
|
||||
events_body = sink.getvalue().to_pybytes()
|
||||
query_events_payloads = []
|
||||
|
||||
def handler(request):
|
||||
content_len = int(request.headers.get("Content-Length", 0))
|
||||
@@ -2479,6 +2513,22 @@ def test_remote_connection_jobs_surface():
|
||||
request.end_headers()
|
||||
request.wfile.write(json.dumps(rsp).encode())
|
||||
elif request.path == "/v1/jobs/describe":
|
||||
if payload["job_id"] == "job-2":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
json.dumps(
|
||||
dict(
|
||||
job_id="job-2",
|
||||
job_type="refresh_column",
|
||||
job_state="DONE",
|
||||
creation_ms=2000,
|
||||
result=dict(rows_assigned=1000000, rows_failed=0),
|
||||
)
|
||||
).encode()
|
||||
)
|
||||
return
|
||||
if payload["job_id"] != "job-1":
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
@@ -2510,7 +2560,7 @@ def test_remote_connection_jobs_surface():
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"job_id": "job-1"}')
|
||||
elif request.path == "/v1/jobs/query_events":
|
||||
assert payload["job_id"] == "job-1"
|
||||
query_events_payloads.append(payload)
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
|
||||
request.end_headers()
|
||||
@@ -2526,24 +2576,109 @@ def test_remote_connection_jobs_surface():
|
||||
assert jobs[0].table == "t1"
|
||||
assert jobs[1].state == "finished"
|
||||
|
||||
description = db.get_job("job-1")
|
||||
assert description.job_type == "create_index"
|
||||
assert description.state == "failed"
|
||||
assert json.loads(description.spec_json) == {"column": "vec"}
|
||||
assert description.failure.message == "worker died"
|
||||
assert description.failure.retryable is True
|
||||
assert db.get_job("missing") is None
|
||||
|
||||
assert db.cancel_job("job-1") is True
|
||||
assert db.cancel_job("missing") is False
|
||||
|
||||
batches = db.job_history("job-1")
|
||||
assert len(batches) == 1
|
||||
assert batches[0].num_rows == 2
|
||||
assert batches[0].column("state").to_pylist() == ["created", "done"]
|
||||
# Opening a job hands back a populated handle; a missing one fails.
|
||||
with pytest.raises(JobNotFoundError, match="missing"):
|
||||
db.open_job("missing")
|
||||
finished = db.open_job("job-2")
|
||||
assert finished.state == "finished"
|
||||
assert finished.result == {"rows_assigned": 1000000, "rows_failed": 0}
|
||||
|
||||
job = db.job("job-1")
|
||||
job = db.open_job("job-1")
|
||||
assert job.id == "job-1"
|
||||
# Opening already populated the handle.
|
||||
assert job.state == "failed"
|
||||
assert job.spec == {"column": "vec"}
|
||||
assert job.failure.message == "worker died"
|
||||
assert job.status() == "failed"
|
||||
with pytest.raises(JobFailedError, match="worker died"):
|
||||
job.wait(timeout=timedelta(seconds=5))
|
||||
|
||||
|
||||
def test_remote_job_handle_reports_its_own_detail():
|
||||
schema = pa.schema([("state", pa.string())])
|
||||
batch = pa.record_batch([pa.array(["claim_complete"])], schema=schema)
|
||||
sink = pa.BufferOutputStream()
|
||||
with pa.ipc.new_stream(sink, schema) as writer:
|
||||
writer.write_batch(batch)
|
||||
events_body = sink.getvalue().to_pybytes()
|
||||
event_payloads = []
|
||||
|
||||
def handler(request):
|
||||
content_len = int(request.headers.get("Content-Length", 0))
|
||||
body = request.rfile.read(content_len) if content_len > 0 else b""
|
||||
payload = json.loads(body) if body else {}
|
||||
if request.path == "/v1/jobs/describe":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
json.dumps(
|
||||
dict(
|
||||
job_id="job-1",
|
||||
job_type="refresh_column",
|
||||
job_state="DONE",
|
||||
creation_ms=2000,
|
||||
spec=dict(column="vec"),
|
||||
result=dict(rows_assigned=1000000),
|
||||
)
|
||||
).encode()
|
||||
)
|
||||
elif request.path == "/v1/jobs/query_events":
|
||||
event_payloads.append(payload)
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/vnd.apache.arrow.stream")
|
||||
request.end_headers()
|
||||
request.wfile.write(events_body)
|
||||
else:
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
job = db.open_job("job-1")
|
||||
|
||||
# Opening populates the handle in the same round trip.
|
||||
assert job.state == "finished"
|
||||
job.refresh()
|
||||
assert job.job_type == "refresh_column"
|
||||
assert job.creation_ms == 2000
|
||||
assert job.spec == {"column": "vec"}
|
||||
assert job.result == {"rows_assigned": 1000000}
|
||||
assert job.failure is None
|
||||
# The JSON payloads stay reachable, but as internal APIs.
|
||||
assert json.loads(job._spec_json) == {"column": "vec"}
|
||||
assert json.loads(job._result_json) == {"rows_assigned": 1000000}
|
||||
|
||||
# print() shows everything the handle knows and nothing it does not.
|
||||
# print() lays every known field out on its own line, with the JSON
|
||||
# payloads indented rather than crammed onto one line.
|
||||
assert repr(job) == "\n".join(
|
||||
[
|
||||
"Job(",
|
||||
" id='job-1',",
|
||||
" state='finished',",
|
||||
" job_type='refresh_column',",
|
||||
" creation_ms=2000,",
|
||||
" spec={",
|
||||
' "column": "vec"',
|
||||
" },",
|
||||
" result={",
|
||||
' "rows_assigned": 1000000',
|
||||
" },",
|
||||
")",
|
||||
]
|
||||
)
|
||||
# Nothing it does not know shows up.
|
||||
assert "failure" not in repr(job)
|
||||
|
||||
events = job.events(filter="state = 'claim_complete'", limit=500)
|
||||
assert isinstance(events, pa.Table)
|
||||
assert events.column("state").to_pylist() == ["claim_complete"]
|
||||
# The handle supplies job_id; the caller only narrows the query.
|
||||
assert event_payloads[-1] == {
|
||||
"job_id": "job-1",
|
||||
"limit": 500,
|
||||
"filter": "state = 'claim_complete'",
|
||||
}
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
import pyarrow as pa
|
||||
|
||||
import lancedb
|
||||
from lancedb import _lancedb
|
||||
from lancedb.arrow import AsyncRecordBatchReader
|
||||
from lancedb.db import AsyncConnection
|
||||
from lancedb.remote.db import RemoteDBConnection
|
||||
from lancedb.sql import AsyncQuery, Query
|
||||
|
||||
NIL_QUERY_ID = UUID(int=0)
|
||||
|
||||
|
||||
class FakeNativeQuery:
|
||||
id = UUID("0198f1b2-c3d4-7e5f-8123-456789abcdef")
|
||||
|
||||
async def reader(self):
|
||||
return pa.table({"value": [1, 2]})
|
||||
|
||||
|
||||
class FakeNativeConnection:
|
||||
async def execute_query_async(self, query, *, default_namespace_path=None):
|
||||
return FakeNativeQuery()
|
||||
|
||||
|
||||
class FakeAsyncConnection:
|
||||
async def execute_query_async(self, query, *, default_namespace_path=None):
|
||||
return AsyncQuery(FakeNativeQuery())
|
||||
|
||||
|
||||
def remote_connection(sql_host_override=None):
|
||||
return lancedb.connect(
|
||||
"db://analytics",
|
||||
api_key="test-key",
|
||||
host_override="http://localhost:10024",
|
||||
sql_host_override=sql_host_override,
|
||||
)
|
||||
|
||||
|
||||
def test_sql_is_connection_scoped():
|
||||
assert hasattr(lancedb, "sql")
|
||||
assert not callable(lancedb.sql)
|
||||
assert not hasattr(_lancedb, "sql")
|
||||
assert not hasattr(remote_connection(), "sql")
|
||||
assert hasattr(remote_connection(), "execute_query")
|
||||
assert hasattr(remote_connection(), "execute_query_async")
|
||||
assert hasattr(remote_connection(), "describe_query")
|
||||
|
||||
|
||||
def test_query_id_is_uuid():
|
||||
query = AsyncQuery(FakeNativeQuery())
|
||||
assert isinstance(query.id, UUID)
|
||||
assert Query(query).id == query.id
|
||||
|
||||
|
||||
def test_connection_serializes_sql_host_override():
|
||||
endpoint = "grpc+tls://sql.example.com:10026"
|
||||
restored = lancedb.deserialize_conn(
|
||||
remote_connection(sql_host_override=endpoint).serialize()
|
||||
)
|
||||
assert restored.sql_host_override == endpoint
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sql_reader_is_record_batch_stream():
|
||||
reader = await AsyncQuery(FakeNativeQuery()).reader()
|
||||
assert isinstance(reader, AsyncRecordBatchReader)
|
||||
assert (await reader.read_all())[0].column(0).to_pylist() == [1, 2]
|
||||
|
||||
|
||||
def test_sync_sql_reader_is_record_batch_reader():
|
||||
reader = Query(AsyncQuery(FakeNativeQuery())).reader()
|
||||
assert isinstance(reader, pa.RecordBatchReader)
|
||||
assert reader.read_all().column(0).to_pylist() == [1, 2]
|
||||
|
||||
|
||||
def test_execute_query_returns_blocking_reader():
|
||||
connection = RemoteDBConnection.__new__(RemoteDBConnection)
|
||||
connection._conn = FakeAsyncConnection()
|
||||
reader = connection.execute_query("SELECT 1")
|
||||
assert isinstance(reader, pa.RecordBatchReader)
|
||||
assert reader.read_all().column(0).to_pylist() == [1, 2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_execute_query_returns_async_reader():
|
||||
connection = AsyncConnection(FakeNativeConnection())
|
||||
reader = await connection.execute_query("SELECT 1")
|
||||
assert isinstance(reader, AsyncRecordBatchReader)
|
||||
assert (await reader.read_all())[0].column(0).to_pylist() == [1, 2]
|
||||
|
||||
|
||||
def test_local_connection_rejects_sql(tmp_path):
|
||||
connection = lancedb.connect(tmp_path)
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
connection.execute_query("SELECT 1")
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
connection.execute_query_async("SELECT 1")
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
connection.describe_query(NIL_QUERY_ID)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_async_connection_rejects_sql(tmp_path):
|
||||
connection = await lancedb.connect_async(tmp_path)
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
await connection.execute_query("SELECT 1")
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
await connection.execute_query_async("SELECT 1")
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
await connection.describe_query(NIL_QUERY_ID)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_namespace_connection_rejects_sql(tmp_path):
|
||||
connection = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
await connection.execute_query("SELECT 1")
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
await connection.execute_query_async("SELECT 1")
|
||||
with pytest.raises(NotImplementedError, match="SQL"):
|
||||
await connection.describe_query(NIL_QUERY_ID)
|
||||
|
||||
|
||||
def test_describe_query_requires_uuid():
|
||||
with pytest.raises(TypeError, match="UUID"):
|
||||
remote_connection().describe_query(str(NIL_QUERY_ID))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"default_namespace_path",
|
||||
["public", ("public",), [1]],
|
||||
)
|
||||
def test_execute_query_async_requires_namespace_path_list(default_namespace_path):
|
||||
with pytest.raises(ValueError, match="default_namespace_path"):
|
||||
remote_connection().execute_query_async(
|
||||
"SELECT 1", default_namespace_path=default_namespace_path
|
||||
)
|
||||
|
||||
|
||||
def test_execute_query_async_rejects_invalid_endpoint():
|
||||
connection = remote_connection(sql_host_override="invalid://localhost")
|
||||
with pytest.raises(ValueError, match="sql_host_override"):
|
||||
connection.execute_query_async("SELECT 1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"default_namespace_path",
|
||||
[[""], ["café"], ["pub\tlic"], ["events$raw"]],
|
||||
)
|
||||
def test_execute_query_async_rejects_invalid_namespace_components(
|
||||
default_namespace_path,
|
||||
):
|
||||
with pytest.raises(ValueError, match="default_namespace_path"):
|
||||
remote_connection().execute_query_async(
|
||||
"SELECT 1", default_namespace_path=default_namespace_path
|
||||
)
|
||||
@@ -2682,6 +2682,43 @@ def test_merge_insert(mem_db: DBConnection):
|
||||
)
|
||||
|
||||
|
||||
def test_merge_insert_composite_key(mem_db: DBConnection):
|
||||
table = mem_db.create_table(
|
||||
"my_table",
|
||||
data=pa.table(
|
||||
{
|
||||
"shard": ["a", "a", "b"],
|
||||
"id": [1, 2, 1],
|
||||
"val": ["x", "y", "z"],
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
# ("a", 1) matches an existing row and updates it. ("b", 2) agrees with an
|
||||
# existing row on each key column separately but on neither pair, so it is
|
||||
# an insert.
|
||||
new_data = pa.table({"shard": ["a", "b"], "id": [1, 2], "val": ["X", "W"]})
|
||||
res = (
|
||||
table.merge_insert(["shard", "id"])
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(new_data)
|
||||
)
|
||||
assert res.num_updated_rows == 1
|
||||
assert res.num_inserted_rows == 1
|
||||
|
||||
expected = pa.table(
|
||||
{
|
||||
"shard": ["a", "a", "b", "b"],
|
||||
"id": [1, 2, 1, 2],
|
||||
"val": ["X", "y", "z", "W"],
|
||||
}
|
||||
)
|
||||
assert table.to_arrow().sort_by([("shard", "ascending"), ("id", "ascending")]) == (
|
||||
expected
|
||||
)
|
||||
|
||||
|
||||
def test_merge_insert_nullable_pandas_into_pydantic_schema(mem_db: DBConnection):
|
||||
# Regression test for https://github.com/lancedb/lancedb/issues/2366
|
||||
pd = pytest.importorskip("pandas")
|
||||
|
||||
+164
-36
@@ -13,11 +13,7 @@ use crate::{
|
||||
runtime::future_into_py,
|
||||
table::Table,
|
||||
};
|
||||
use arrow::{
|
||||
datatypes::Schema,
|
||||
ffi_stream::ArrowArrayStreamReader,
|
||||
pyarrow::{FromPyArrow, ToPyArrow},
|
||||
};
|
||||
use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow};
|
||||
use lancedb::{
|
||||
connection::Connection as LanceConnection,
|
||||
connection::NamespaceClientPushdownOperation,
|
||||
@@ -28,7 +24,7 @@ use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
pyclass, pyfunction, pymethods,
|
||||
types::{PyDict, PyDictMethods, PyList, PyListMethods},
|
||||
types::{PyAnyMethods, PyDict, PyDictMethods, PyList},
|
||||
};
|
||||
|
||||
#[pyclass]
|
||||
@@ -86,6 +82,24 @@ impl Connection {
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_default_namespace_path(path: Option<Bound<'_, PyAny>>) -> PyResult<Vec<String>> {
|
||||
match path {
|
||||
Some(path) => {
|
||||
if !path.is_instance_of::<PyList>() {
|
||||
return Err(PyValueError::new_err(
|
||||
"Connection.execute_query_async default_namespace_path must be a list",
|
||||
));
|
||||
}
|
||||
path.extract::<Vec<String>>().map_err(|_| {
|
||||
PyValueError::new_err(
|
||||
"Connection.execute_query_async default_namespace_path components must be strings",
|
||||
)
|
||||
})
|
||||
}
|
||||
None => Ok(vec!["public".to_string()]),
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Connection {
|
||||
fn __repr__(&self) -> String {
|
||||
@@ -108,6 +122,40 @@ impl Connection {
|
||||
self.get_inner().map(|inner| inner.uri().to_string())
|
||||
}
|
||||
|
||||
#[pyo3(signature = (query, *, default_namespace_path=None))]
|
||||
pub fn execute_query_async<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
query: String,
|
||||
default_namespace_path: Option<Bound<'_, PyAny>>,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
let default_namespace_path = parse_default_namespace_path(default_namespace_path)?;
|
||||
future_into_py(self_.py(), async move {
|
||||
let operation = inner
|
||||
.execute_query_async(query)
|
||||
.default_namespace_path(default_namespace_path);
|
||||
operation
|
||||
.execute()
|
||||
.await
|
||||
.map(crate::sql::Query::new)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn describe_query<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
query_id: uuid::Uuid,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.describe_query(query_id)
|
||||
.await
|
||||
.map(crate::sql::QueryDescription::from)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = ())]
|
||||
pub fn get_read_consistency_interval(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
@@ -592,9 +640,12 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn job(&self, job_id: String) -> PyResult<crate::job::Job> {
|
||||
let inner = self.get_inner()?.clone();
|
||||
Ok(crate::job::Job::new(inner.job(job_id).infer_error()?))
|
||||
pub fn open_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let job = inner.open_job(&job_id).await.infer_error()?;
|
||||
Ok(crate::job::Job::new(job))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn create_function_async(
|
||||
@@ -629,6 +680,102 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.list_functions()
|
||||
.await
|
||||
.infer_error()?
|
||||
.into_iter()
|
||||
.map(|function| function.to_canonical_json().infer_error())
|
||||
.collect::<PyResult<Vec<_>>>()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_function(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
version: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_function(name, version).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn create_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
value: String,
|
||||
namespace_path: Vec<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_secret(name, value, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn alter_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
value: String,
|
||||
namespace_path: Vec<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.alter_secret(name, value, &namespace_path)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_secrets(
|
||||
self_: PyRef<'_, Self>,
|
||||
namespace_path: Vec<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.list_secrets(&namespace_path).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
namespace_path: Vec<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_secret(name, &namespace_path).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Name and timestamps as a plain mapping. `SecretInfo` carries no value,
|
||||
/// so there is none to filter out here.
|
||||
pub fn describe_secret(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
namespace_path: Vec<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let info = inner
|
||||
.describe_secret(name, &namespace_path)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(HashMap::from([
|
||||
("name".to_string(), info.name),
|
||||
("created_at".to_string(), info.created_at),
|
||||
("updated_at".to_string(), info.updated_at),
|
||||
]))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -640,42 +787,16 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let description = inner.get_job(&job_id).await.infer_error()?;
|
||||
Ok(description.map(crate::job::JobDescription::from))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.cancel_job(&job_id).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None))]
|
||||
pub fn job_history(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let batches = inner.job_history(job_id.as_deref()).await.infer_error()?;
|
||||
Python::attach(|py| {
|
||||
let list = PyList::empty(py);
|
||||
for batch in batches {
|
||||
list.append(batch.to_pyarrow(py)?)?;
|
||||
}
|
||||
Ok(list.unbind())
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None, oauth_config=None))]
|
||||
#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, sql_host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None, oauth_config=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn connect(
|
||||
py: Python<'_>,
|
||||
@@ -683,6 +804,7 @@ pub fn connect(
|
||||
api_key: Option<String>,
|
||||
region: Option<String>,
|
||||
host_override: Option<String>,
|
||||
sql_host_override: Option<String>,
|
||||
read_consistency_interval: Option<f64>,
|
||||
client_config: Option<PyClientConfig>,
|
||||
storage_options: Option<HashMap<String, String>>,
|
||||
@@ -702,6 +824,12 @@ pub fn connect(
|
||||
if let Some(host_override) = host_override {
|
||||
builder = builder.host_override(&host_override);
|
||||
}
|
||||
#[cfg(feature = "remote")]
|
||||
if let Some(sql_host_override) = sql_host_override {
|
||||
builder = builder.sql_host_override(&sql_host_override);
|
||||
}
|
||||
#[cfg(not(feature = "remote"))]
|
||||
let _ = sql_host_override;
|
||||
if let Some(read_consistency_interval) = read_consistency_interval {
|
||||
let read_consistency_interval = Duration::from_secs_f64(read_consistency_interval);
|
||||
builder = builder.read_consistency_interval(read_consistency_interval);
|
||||
|
||||
@@ -114,6 +114,12 @@ impl<T> PythonErrorExt<T> for std::result::Result<T, LanceError> {
|
||||
.getattr(intern!(py, "JobCancelledError"))?;
|
||||
Err(PyErr::from_value(cls.call1((err.to_string(),))?))
|
||||
}),
|
||||
LanceError::JobNotFound { .. } => Python::attach(|py| {
|
||||
let cls = py
|
||||
.import(intern!(py, "lancedb.exceptions"))?
|
||||
.getattr(intern!(py, "JobNotFoundError"))?;
|
||||
Err(PyErr::from_value(cls.call1((err.to_string(),))?))
|
||||
}),
|
||||
_ => self.runtime_error(),
|
||||
},
|
||||
}
|
||||
|
||||
+126
-9
@@ -4,11 +4,50 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::runtime::future_into_py;
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
|
||||
use arrow::{
|
||||
datatypes::Schema,
|
||||
pyarrow::{IntoPyArrow, Table as PyArrowTable},
|
||||
};
|
||||
use lancedb::job::JobEventsRequest;
|
||||
use pyo3::{
|
||||
Bound, PyAny, PyRef, PyResult, Python,
|
||||
exceptions::PyValueError,
|
||||
pyclass, pymethods,
|
||||
types::{PyAnyMethods, PyDict, PyDictMethods},
|
||||
};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::error::PythonErrorExt;
|
||||
|
||||
const REPR_INDENT: &str = " ";
|
||||
|
||||
/// Parse a stored JSON payload into Python data. The bindings carry these as
|
||||
/// strings because that is what crosses the boundary cheaply; the public
|
||||
/// Python surface is the parsed form.
|
||||
fn parse_json_payload<'py>(
|
||||
py: Python<'py>,
|
||||
raw: Option<&str>,
|
||||
) -> PyResult<Option<Bound<'py, PyAny>>> {
|
||||
match raw {
|
||||
None => Ok(None),
|
||||
Some(raw) => Ok(Some(py.import("json")?.call_method1("loads", (raw,))?)),
|
||||
}
|
||||
}
|
||||
|
||||
/// A payload rendered as indented JSON, aligned under the field that holds it.
|
||||
fn pretty_json_payload(py: Python<'_>, raw: Option<&str>) -> PyResult<Option<String>> {
|
||||
let Some(parsed) = parse_json_payload(py, raw)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("indent", 4)?;
|
||||
let rendered: String = py
|
||||
.import("json")?
|
||||
.call_method("dumps", (parsed,), Some(&kwargs))?
|
||||
.extract()?;
|
||||
Ok(Some(rendered.replace('\n', &format!("\n{REPR_INDENT}"))))
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct Job {
|
||||
inner: Arc<lancedb::Job<std::result::Result<Option<String>, String>>>,
|
||||
@@ -67,6 +106,48 @@ impl Job {
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn refresh(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.refresh().await.infer_error()?;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
/// The last observed lifecycle state, without contacting the backend.
|
||||
#[getter]
|
||||
pub fn _state(&self) -> Option<String> {
|
||||
self.inner.state()
|
||||
}
|
||||
|
||||
/// The last observed server-side record. `None` for an in-process job.
|
||||
#[getter]
|
||||
pub fn _description(&self) -> Option<JobDescription> {
|
||||
self.inner.description().map(JobDescription::from)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (*, limit=None, filter=None))]
|
||||
pub fn events(
|
||||
self_: PyRef<'_, Self>,
|
||||
limit: Option<u32>,
|
||||
filter: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let request = JobEventsRequest { limit, filter };
|
||||
future_into_py(self_.py(), async move {
|
||||
let batches = inner.events(request).await.infer_error()?;
|
||||
Python::attach(|py| {
|
||||
let schema = batches
|
||||
.first()
|
||||
.map(|batch| batch.schema())
|
||||
.unwrap_or_else(|| Arc::new(Schema::empty()));
|
||||
let table = PyArrowTable::try_new(batches, schema)
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
table.into_pyarrow(py).map(|table| table.unbind())
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A row from `Connection.list_jobs`: one server-side job.
|
||||
@@ -121,7 +202,7 @@ impl JobFailureInfo {
|
||||
}
|
||||
}
|
||||
|
||||
/// A described job from `Connection.get_job`.
|
||||
/// The server-side record behind a `Job` handle.
|
||||
#[pyclass(get_all, skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobDescription {
|
||||
@@ -129,17 +210,49 @@ pub struct JobDescription {
|
||||
job_type: String,
|
||||
state: String,
|
||||
creation_ms: i64,
|
||||
spec_json: Option<String>,
|
||||
/// Internal: the wire form behind the `spec` property.
|
||||
_spec_json: Option<String>,
|
||||
/// Internal: the wire form behind the `result` property.
|
||||
_result_json: Option<String>,
|
||||
failure: Option<JobFailureInfo>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl JobDescription {
|
||||
fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"JobDescription(job_id={:?}, job_type={:?}, state={:?}, creation_ms={})",
|
||||
self.job_id, self.job_type, self.state, self.creation_ms
|
||||
)
|
||||
/// The job-type-specific specification it was submitted with.
|
||||
#[getter]
|
||||
fn spec<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
|
||||
parse_json_payload(py, self._spec_json.as_deref())
|
||||
}
|
||||
|
||||
/// The job-type-specific terminal result. `None` until the job succeeds.
|
||||
#[getter]
|
||||
fn result<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyAny>>> {
|
||||
parse_json_payload(py, self._result_json.as_deref())
|
||||
}
|
||||
|
||||
fn __repr__(&self, py: Python<'_>) -> PyResult<String> {
|
||||
let mut fields = vec![
|
||||
format!("job_id={:?}", self.job_id),
|
||||
format!("job_type={:?}", self.job_type),
|
||||
format!("state={:?}", self.state),
|
||||
format!("creation_ms={}", self.creation_ms),
|
||||
];
|
||||
// Lay the payloads out as indented JSON, the same way the `Job` repr
|
||||
// does, so the two agree on how the same data looks.
|
||||
for (name, payload) in [("spec", &self._spec_json), ("result", &self._result_json)] {
|
||||
if let Some(rendered) = pretty_json_payload(py, payload.as_deref())? {
|
||||
fields.push(format!("{name}={rendered}"));
|
||||
}
|
||||
}
|
||||
if let Some(failure) = &self.failure {
|
||||
fields.push(format!("failure={}", failure.__repr__()));
|
||||
}
|
||||
let body = fields
|
||||
.iter()
|
||||
.map(|field| format!("\n{REPR_INDENT}{field},"))
|
||||
.collect::<String>();
|
||||
Ok(format!("JobDescription({body}\n)"))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -150,7 +263,11 @@ impl From<lancedb::database::JobDescription> for JobDescription {
|
||||
job_type: description.job_type,
|
||||
state: description.state,
|
||||
creation_ms: description.creation_ms,
|
||||
spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
|
||||
_spec_json: (!description.spec.is_null()).then(|| description.spec.to_string()),
|
||||
_result_json: description
|
||||
.result
|
||||
.filter(|result| !result.is_null())
|
||||
.map(|result| result.to_string()),
|
||||
failure: description.failure.map(|failure| JobFailureInfo {
|
||||
phase: failure.phase,
|
||||
message: failure.message,
|
||||
|
||||
@@ -34,6 +34,7 @@ pub mod permutation;
|
||||
pub mod query;
|
||||
pub mod runtime;
|
||||
pub mod session;
|
||||
pub mod sql;
|
||||
pub mod table;
|
||||
pub mod util;
|
||||
|
||||
@@ -50,6 +51,8 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<crate::job::JobInfo>()?;
|
||||
m.add_class::<crate::job::JobDescription>()?;
|
||||
m.add_class::<crate::job::JobFailureInfo>()?;
|
||||
m.add_class::<crate::sql::Query>()?;
|
||||
m.add_class::<crate::sql::QueryDescription>()?;
|
||||
m.add_class::<PyBlobFile>()?;
|
||||
m.add_class::<IndexConfig>()?;
|
||||
m.add_class::<Query>()?;
|
||||
|
||||
@@ -323,6 +323,7 @@ impl<'py> IntoPyObject<'py> for PyQueryVectors {
|
||||
pub struct PyQueryRequest {
|
||||
pub limit: Option<usize>,
|
||||
pub offset: Option<usize>,
|
||||
pub take_offsets: Option<Vec<u64>>,
|
||||
pub filter: Option<PyQueryFilter>,
|
||||
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
|
||||
pub select: PySelect,
|
||||
@@ -353,6 +354,7 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
AnyQuery::Query(query_request) => Self {
|
||||
limit: query_request.limit,
|
||||
offset: query_request.offset,
|
||||
take_offsets: query_request.take_offsets,
|
||||
filter: query_request.filter.map(PyQueryFilter),
|
||||
full_text_search: query_request
|
||||
.full_text_search
|
||||
@@ -381,6 +383,7 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
AnyQuery::VectorQuery(vector_query) => Self {
|
||||
limit: vector_query.base.limit,
|
||||
offset: vector_query.base.offset,
|
||||
take_offsets: vector_query.base.take_offsets,
|
||||
filter: vector_query.base.filter.map(PyQueryFilter),
|
||||
full_text_search: None,
|
||||
select_source_columns: PySelect::source_columns(&vector_query.base.select),
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::arrow::RecordBatchStream;
|
||||
use crate::error::PythonErrorExt;
|
||||
use crate::runtime::future_into_py;
|
||||
|
||||
#[pyclass(name = "SqlQuery")]
|
||||
pub struct Query {
|
||||
inner: Arc<lancedb::sql::Query>,
|
||||
}
|
||||
|
||||
impl Query {
|
||||
pub(crate) fn new(inner: lancedb::sql::Query) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(inner),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl Query {
|
||||
#[getter]
|
||||
pub fn id(&self) -> Uuid {
|
||||
self.inner.id()
|
||||
}
|
||||
|
||||
pub fn describe(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.describe()
|
||||
.await
|
||||
.map(QueryDescription::from)
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn reader(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let stream = inner.reader().await.infer_error()?;
|
||||
Ok(RecordBatchStream::new(stream))
|
||||
})
|
||||
}
|
||||
|
||||
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(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(get_all, skip_from_py_object)]
|
||||
#[derive(Clone)]
|
||||
pub struct QueryDescription {
|
||||
id: Uuid,
|
||||
status: String,
|
||||
progress: Option<f64>,
|
||||
expires_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl QueryDescription {
|
||||
fn __repr__(&self) -> String {
|
||||
format!(
|
||||
"QueryDescription(id={:?}, status={:?}, progress={:?}, expires_at={:?})",
|
||||
self.id, self.status, self.progress, self.expires_at
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::sql::QueryDescription> for QueryDescription {
|
||||
fn from(description: lancedb::sql::QueryDescription) -> Self {
|
||||
Self {
|
||||
id: description.id,
|
||||
status: description.status.to_string(),
|
||||
progress: description.progress,
|
||||
expires_at: description.expires_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
Generated
+98
-95
@@ -10,6 +10,9 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
|
||||
[options]
|
||||
prerelease-mode = "allow"
|
||||
|
||||
[[package]]
|
||||
name = "accelerate"
|
||||
version = "1.14.0"
|
||||
@@ -799,7 +802,7 @@ name = "cuda-bindings"
|
||||
version = "13.3.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cuda-pathfinder", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "cuda-pathfinder" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" },
|
||||
@@ -834,37 +837,37 @@ wheels = [
|
||||
|
||||
[package.optional-dependencies]
|
||||
cublas = [
|
||||
{ name = "nvidia-cublas", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cublas" },
|
||||
]
|
||||
cudart = [
|
||||
{ name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-runtime" },
|
||||
]
|
||||
cufft = [
|
||||
{ name = "nvidia-cufft", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cufft" },
|
||||
]
|
||||
cufile = [
|
||||
{ name = "nvidia-cufile", marker = "sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cufile" },
|
||||
]
|
||||
cupti = [
|
||||
{ name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-cupti" },
|
||||
]
|
||||
curand = [
|
||||
{ name = "nvidia-curand", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-curand" },
|
||||
]
|
||||
cusolver = [
|
||||
{ name = "nvidia-cusolver", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cusolver" },
|
||||
]
|
||||
cusparse = [
|
||||
{ name = "nvidia-cusparse", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cusparse" },
|
||||
]
|
||||
nvjitlink = [
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nvjitlink" },
|
||||
]
|
||||
nvrtc = [
|
||||
{ name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-nvrtc" },
|
||||
]
|
||||
nvtx = [
|
||||
{ name = "nvidia-nvtx", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nvtx" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1023,7 +1026,7 @@ name = "exceptiongroup"
|
||||
version = "1.3.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.11'" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" }
|
||||
wheels = [
|
||||
@@ -1440,16 +1443,16 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "cachetools", marker = "python_full_version < '3.11'" },
|
||||
{ name = "certifi", marker = "python_full_version < '3.11'" },
|
||||
{ name = "httpx", marker = "python_full_version < '3.11'" },
|
||||
{ name = "ibm-cos-sdk", marker = "python_full_version < '3.11'" },
|
||||
{ name = "lomond", marker = "python_full_version < '3.11'" },
|
||||
{ name = "packaging", marker = "python_full_version < '3.11'" },
|
||||
{ name = "pandas", version = "2.2.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "requests", marker = "python_full_version < '3.11'" },
|
||||
{ name = "tabulate", marker = "python_full_version < '3.11'" },
|
||||
{ name = "urllib3", marker = "python_full_version < '3.11'" },
|
||||
{ name = "cachetools" },
|
||||
{ name = "certifi" },
|
||||
{ name = "httpx" },
|
||||
{ name = "ibm-cos-sdk" },
|
||||
{ name = "lomond" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pandas", version = "2.2.3", source = { registry = "https://pypi.org/simple" } },
|
||||
{ name = "requests" },
|
||||
{ name = "tabulate" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/56/2e3df38a1f13062095d7bde23c87a92f3898982993a15186b1bfecbd206f/ibm_watsonx_ai-1.3.42.tar.gz", hash = "sha256:ee5be59009004245d957ce97d1227355516df95a2640189749487614fef674ff", size = 688651, upload-time = "2025-10-01T13:35:41.527Z" }
|
||||
wheels = [
|
||||
@@ -1468,17 +1471,17 @@ resolution-markers = [
|
||||
"python_full_version == '3.11.*'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "cachetools", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "certifi", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "httpx", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "ibm-cos-sdk", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "lomond", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "packaging", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" },
|
||||
{ name = "cachetools" },
|
||||
{ name = "certifi" },
|
||||
{ name = "httpx" },
|
||||
{ name = "ibm-cos-sdk" },
|
||||
{ name = "lomond" },
|
||||
{ name = "packaging" },
|
||||
{ name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" },
|
||||
{ name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" },
|
||||
{ name = "requests", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "tabulate", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "urllib3", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "requests" },
|
||||
{ name = "tabulate" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/29/a3/c756b534696ab2f3f29882fdb7ca7198b7a5c94e10c0a3a327853d6d6b79/ibm_watsonx_ai-1.5.14.tar.gz", hash = "sha256:a756488bd57e87c0fc51be42dcba871143cfe0ac1e805c497c5047e1e4f13e9d", size = 735804, upload-time = "2026-06-22T12:32:43.85Z" }
|
||||
wheels = [
|
||||
@@ -1554,17 +1557,17 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" },
|
||||
{ name = "decorator", marker = "python_full_version < '3.11'" },
|
||||
{ name = "exceptiongroup", marker = "python_full_version < '3.11'" },
|
||||
{ name = "jedi", marker = "python_full_version < '3.11'" },
|
||||
{ name = "matplotlib-inline", marker = "python_full_version < '3.11'" },
|
||||
{ name = "pexpect", marker = "python_full_version < '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" },
|
||||
{ name = "prompt-toolkit", marker = "python_full_version < '3.11'" },
|
||||
{ name = "pygments", marker = "python_full_version < '3.11'" },
|
||||
{ name = "stack-data", marker = "python_full_version < '3.11'" },
|
||||
{ name = "traitlets", marker = "python_full_version < '3.11'" },
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.11'" },
|
||||
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||
{ name = "decorator" },
|
||||
{ name = "exceptiongroup" },
|
||||
{ name = "jedi" },
|
||||
{ name = "matplotlib-inline" },
|
||||
{ name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
|
||||
{ name = "prompt-toolkit" },
|
||||
{ name = "pygments" },
|
||||
{ name = "stack-data" },
|
||||
{ name = "traitlets" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/40/18/f8598d287006885e7136451fdea0755af4ebcbfe342836f24deefaed1164/ipython-8.39.0.tar.gz", hash = "sha256:4110ae96012c379b8b6db898a07e186c40a2a1ef5d57a7fa83166047d9da7624", size = 5513971, upload-time = "2026-03-27T10:02:13.94Z" }
|
||||
wheels = [
|
||||
@@ -1583,18 +1586,18 @@ resolution-markers = [
|
||||
"python_full_version == '3.11.*'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "colorama", marker = "python_full_version >= '3.11' and sys_platform == 'win32'" },
|
||||
{ name = "decorator", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "ipython-pygments-lexers", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "jedi", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "matplotlib-inline", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "pexpect", marker = "python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" },
|
||||
{ name = "prompt-toolkit", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "psutil", marker = "python_full_version >= '3.11' and sys_platform != 'cygwin' and sys_platform != 'emscripten'" },
|
||||
{ name = "pygments", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "stack-data", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "traitlets", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "typing-extensions", marker = "python_full_version == '3.11.*'" },
|
||||
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||
{ name = "decorator" },
|
||||
{ name = "ipython-pygments-lexers" },
|
||||
{ name = "jedi" },
|
||||
{ name = "matplotlib-inline" },
|
||||
{ name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
|
||||
{ name = "prompt-toolkit" },
|
||||
{ name = "psutil", marker = "sys_platform != 'cygwin' and sys_platform != 'emscripten'" },
|
||||
{ name = "pygments" },
|
||||
{ name = "stack-data" },
|
||||
{ name = "traitlets" },
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.12'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/53/59/165d3b4d75cc34add3122c4417ecb229085140ac573103c223cd01dde96f/ipython-9.15.0.tar.gz", hash = "sha256:da2819ce2aa83135257df830660b1176d986c3d2876db24df01974fa955b2756", size = 4442580, upload-time = "2026-06-26T11:03:35.913Z" }
|
||||
wheels = [
|
||||
@@ -1606,7 +1609,7 @@ name = "ipython-pygments-lexers"
|
||||
version = "1.1.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pygments", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "pygments" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" }
|
||||
wheels = [
|
||||
@@ -2858,7 +2861,7 @@ name = "nvidia-cudnn-cu13"
|
||||
version = "9.19.0.56"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-cublas" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" },
|
||||
@@ -2870,7 +2873,7 @@ name = "nvidia-cufft"
|
||||
version = "12.0.0.61"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-nvjitlink" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
|
||||
@@ -2900,9 +2903,9 @@ name = "nvidia-cusolver"
|
||||
version = "12.0.4.66"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-cusparse", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-cublas" },
|
||||
{ name = "nvidia-cusparse" },
|
||||
{ name = "nvidia-nvjitlink" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
|
||||
@@ -2914,7 +2917,7 @@ name = "nvidia-cusparse"
|
||||
version = "12.6.3.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-nvjitlink" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
|
||||
@@ -3091,10 +3094,10 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "python-dateutil", marker = "python_full_version < '3.11'" },
|
||||
{ name = "pytz", marker = "python_full_version < '3.11'" },
|
||||
{ name = "tzdata", marker = "python_full_version < '3.11'" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
|
||||
{ name = "python-dateutil" },
|
||||
{ name = "pytz" },
|
||||
{ name = "tzdata" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/9c/d6/9f8431bacc2e19dca897724cd097b1bb224a6ad5433784a44b587c7c13af/pandas-2.2.3.tar.gz", hash = "sha256:4f18ba62b61d7e192368b84517265a99b4d7ee8912f8708660fb4a366cc82667", size = 4399213, upload-time = "2024-09-20T13:10:04.827Z" }
|
||||
wheels = [
|
||||
@@ -3143,11 +3146,11 @@ resolution-markers = [
|
||||
"python_full_version == '3.11.*'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
|
||||
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or python_full_version >= '3.14'" },
|
||||
{ name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' and python_full_version < '3.14'" },
|
||||
{ name = "python-dateutil", marker = "python_full_version >= '3.11' and python_full_version < '3.14'" },
|
||||
{ name = "pytz", marker = "python_full_version >= '3.11' and python_full_version < '3.14'" },
|
||||
{ name = "tzdata", marker = "python_full_version >= '3.11' and python_full_version < '3.14'" },
|
||||
{ name = "python-dateutil" },
|
||||
{ name = "pytz" },
|
||||
{ name = "tzdata" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" }
|
||||
wheels = [
|
||||
@@ -3210,9 +3213,9 @@ resolution-markers = [
|
||||
"python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" },
|
||||
{ name = "python-dateutil", marker = "python_full_version >= '3.14'" },
|
||||
{ name = "tzdata", marker = "(python_full_version >= '3.14' and sys_platform == 'emscripten') or (python_full_version >= '3.14' and sys_platform == 'win32')" },
|
||||
{ name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } },
|
||||
{ name = "python-dateutil" },
|
||||
{ name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/f8/87/4341c6252d1c47b08768c3d25ac487362bf403f0313ddae4a2a26c9b1b4c/pandas-3.0.3.tar.gz", hash = "sha256:696a4a00a2a2a35d4e5deb3fc946641b96c944f02230e4f76137fe35d806c4fc", size = 4651414, upload-time = "2026-05-11T18:54:29.21Z" }
|
||||
wheels = [
|
||||
@@ -3320,7 +3323,7 @@ name = "pexpect"
|
||||
version = "4.9.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "ptyprocess", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "ptyprocess" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" }
|
||||
wheels = [
|
||||
@@ -3912,8 +3915,8 @@ crypto = [
|
||||
|
||||
[[package]]
|
||||
name = "pylance"
|
||||
version = "7.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
version = "9.0.0rc1"
|
||||
source = { registry = "https://pypi.fury.io/lance-format/" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
@@ -3922,12 +3925,12 @@ dependencies = [
|
||||
{ name = "pyarrow" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ac/ad/2f64921bf346e7075aef24a72595db44821724a3d89a9a92dd24e79632aa/pylance-7.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:98422021975be76e72b1572f41b8c9abb3bee5bdc9bfa5e9ce731110a65ed4d1", size = 62134146, upload-time = "2026-05-27T21:59:37.459Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/73/1c/c5a01bee0160b55d9a98895cbd33091d038f0a0995b121ab72e629008d02/pylance-7.0.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4bec86ee5b6fbd8bfc493e653f0a1fba0303cfe5492b9b46fc25ab908edc7183", size = 65373684, upload-time = "2026-05-27T22:04:01.584Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/eb/da/1fe8b8f7dbfe734d76af76acc994fc360a0d0c79a4874ef69f5a72a58fe3/pylance-7.0.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:881491432c53184e52f8d1db8d5f872f39a03f36fb104bec77b33d379519d8b5", size = 69458555, upload-time = "2026-05-27T22:16:50.567Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/76/f0/dd505cf3fd0226ab9d94759acd713125af1d3bfacfd80bbd52e3b9f89509/pylance-7.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:18453999e7fff4f76b16d6b7882c9df0628bd142ff95e2461bd7dd5ee3fe0af3", size = 65394430, upload-time = "2026-05-27T22:05:30.923Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/17/ba/2357b81034f28eb00790e258ed140289a6a887a7468ca9df6349fd186b27/pylance-7.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:04a58051d408c60fe76d41a220dcaf8fea8fb6d1aa0ca78a709b60bc3cc8d19a", size = 69473470, upload-time = "2026-05-27T22:17:18.935Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1f/ec/5c00b6303a67d787f9475141832cbdc513d674ac3dcaeef8a7b169905e65/pylance-7.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:467d4864af047eaab4e1370e2f1e88e2c6f507c079874421116cb41d78bc3629", size = 74792863, upload-time = "2026-05-27T22:19:23.875Z" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_vEHBE/pylance-9.0.0rc1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:f0b6b02a1808bb3072ee7fe4e36614cae6f86302513e73ec7f55b2234a963b24" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1Jipm4/pylance-9.0.0rc1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:30f0ebf0d88034301819eb964f9236ce555aaa58e7ab89c5975a3e2250bbb405" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_IvKxo/pylance-9.0.0rc1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:44609ea2615ea6e684b85478d1694af2026458f61cf7895ecc75e238bfd17aa8" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2hidj1/pylance-9.0.0rc1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:182167a8dba9eeabffbffd53bd5b8548613d4d459b7cd7b34a840dd00cbb806f" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1dFx3r/pylance-9.0.0rc1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8a63b11e814b7eab758bcaf0d6f97eb05ea86203d9fb0af718c462c24c7d6c9c" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2a8dSh/pylance-9.0.0rc1-cp310-abi3-win_amd64.whl", hash = "sha256:2ff8b953ae2b0550490c1a7efd210aa91bc223d200ffac28849056cfd7436d97" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4683,10 +4686,10 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "joblib", marker = "python_full_version < '3.11'" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "threadpoolctl", marker = "python_full_version < '3.11'" },
|
||||
{ name = "joblib" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
|
||||
{ name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" } },
|
||||
{ name = "threadpoolctl" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" }
|
||||
wheels = [
|
||||
@@ -4734,13 +4737,13 @@ resolution-markers = [
|
||||
"python_full_version == '3.11.*'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "joblib", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "narwhals", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
|
||||
{ name = "joblib" },
|
||||
{ name = "narwhals" },
|
||||
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
|
||||
{ name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
|
||||
{ name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
|
||||
{ name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
|
||||
{ name = "scipy", version = "1.18.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
|
||||
{ name = "threadpoolctl", marker = "python_full_version >= '3.11'" },
|
||||
{ name = "threadpoolctl" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/fa/6f/37092bdb25f712817231799fc5674d8e704066a8a70c1d2d40517e18b4ab/scikit_learn-1.9.0.tar.gz", hash = "sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557", size = 7750767, upload-time = "2026-06-02T11:54:32.706Z" }
|
||||
wheels = [
|
||||
@@ -4784,7 +4787,7 @@ resolution-markers = [
|
||||
"python_full_version < '3.11'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" }
|
||||
wheels = [
|
||||
@@ -4843,7 +4846,7 @@ resolution-markers = [
|
||||
"python_full_version == '3.11.*'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" },
|
||||
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" } },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" }
|
||||
wheels = [
|
||||
@@ -4920,7 +4923,7 @@ resolution-markers = [
|
||||
"python_full_version >= '3.12' and python_full_version < '3.14'",
|
||||
]
|
||||
dependencies = [
|
||||
{ name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
|
||||
{ name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" } },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a7/25/c2700dfaf6442b4effaa91af24ebce5dc9d31bb4a69706313aae70d72cd0/scipy-1.18.0.tar.gz", hash = "sha256:67b2ad2ad54c72ca6d04975a9b2df8c3638c34ddd5b28738e94fc2b57929d378", size = 30774447, upload-time = "2026-06-19T15:01:43.456Z" }
|
||||
wheels = [
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.15"
|
||||
version = "0.39.0-beta.4"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
@@ -21,6 +21,8 @@ arrow-select = { workspace = true }
|
||||
arrow-ord = { workspace = true }
|
||||
arrow-cast = { workspace = true }
|
||||
arrow-ipc.workspace = true
|
||||
arrow-flight = { workspace = true, optional = true }
|
||||
prost = { version = "0.14", optional = true }
|
||||
chrono = { workspace = true }
|
||||
datafusion-catalog.workspace = true
|
||||
datafusion-common.workspace = true
|
||||
@@ -77,6 +79,7 @@ reqwest = { version = "0.12.0", default-features = false, features = [
|
||||
"rustls-tls-native-roots",
|
||||
"stream",
|
||||
], optional = true }
|
||||
tonic = { workspace = true, optional = true }
|
||||
http = { version = "1", optional = true } # Matching what is in reqwest
|
||||
urlencoding = { version = "2", optional = true }
|
||||
uuid = { workspace = true, features = ["v5"] }
|
||||
@@ -145,8 +148,11 @@ huggingface = [
|
||||
]
|
||||
dynamodb = ["lance/dynamodb", "aws"]
|
||||
remote = [
|
||||
"dep:arrow-flight",
|
||||
"dep:prost",
|
||||
"dep:reqwest",
|
||||
"dep:http",
|
||||
"dep:tonic",
|
||||
"dep:urlencoding",
|
||||
"lance-namespace-impls/rest",
|
||||
"lance-namespace-impls/rest-adapter",
|
||||
|
||||
@@ -7,7 +7,9 @@
|
||||
//! raw `Binary` / `LargeBinary` into the blob struct layout. Queries return
|
||||
//! small descriptors, not bytes.
|
||||
//!
|
||||
//! Blob tables require Lance file format >= 2.2 and stable row ids at create.
|
||||
//! Blob tables require Lance file format >= 2.2. `_rowid` values stay valid
|
||||
//! after compaction when the table has stable row ids. Overwrite is a new
|
||||
//! create and does not keep the previous table's stable row id setting.
|
||||
|
||||
use std::ops::Range;
|
||||
use std::sync::Arc;
|
||||
@@ -324,6 +326,7 @@ pub(crate) fn blob_column_names(schema: &Schema) -> Vec<String> {
|
||||
}
|
||||
|
||||
/// Bumps storage format to at least [`LanceFileVersion::V2_2`] for blob schemas.
|
||||
/// Leaves `enable_stable_row_ids` unchanged.
|
||||
pub(crate) fn ensure_blob_storage_version(schema: &Schema, params: &mut WriteParams) {
|
||||
if !has_blob_columns(schema) {
|
||||
return;
|
||||
@@ -385,6 +388,30 @@ fn ensure_all_row_ids_resolved(column: &str, requested: usize, resolved: usize)
|
||||
}
|
||||
}
|
||||
|
||||
/// Lance take reports a missing physical row address as NotSupported or InvalidInput.
|
||||
fn map_blob_take_error(column: &str, requested: usize, err: lance::Error) -> Error {
|
||||
let missing_row_addr = match &err {
|
||||
lance::Error::NotSupported { source, .. } => {
|
||||
source.to_string().contains("must not target deleted rows")
|
||||
}
|
||||
lance::Error::InvalidInput { source, .. } => source
|
||||
.to_string()
|
||||
.contains("belongs to non-existent fragment"),
|
||||
_ => false,
|
||||
};
|
||||
|
||||
if missing_row_addr {
|
||||
Error::InvalidInput {
|
||||
message: format!(
|
||||
"blob read for column '{column}' requested {requested} row ids but some \
|
||||
do not exist in the table; pass row ids collected from this table"
|
||||
),
|
||||
}
|
||||
} else {
|
||||
err.into()
|
||||
}
|
||||
}
|
||||
|
||||
/// Materialize blob-local ranges (same length and order as `requests`, nulls preserved).
|
||||
pub(crate) async fn take_blob_ranges_aligned(
|
||||
dataset: &Arc<Dataset>,
|
||||
@@ -405,7 +432,8 @@ pub(crate) async fn take_blob_ranges_aligned(
|
||||
.with_row_ids(lance_requests)
|
||||
.preserve_order(true)
|
||||
.execute()
|
||||
.await?;
|
||||
.await
|
||||
.map_err(|err| map_blob_take_error(column, requests.len(), err))?;
|
||||
ensure_all_row_ids_resolved(column, requests.len(), payloads.len())?;
|
||||
|
||||
let mut builder = LargeBinaryBuilder::new();
|
||||
@@ -434,7 +462,8 @@ pub(crate) async fn take_blobs_aligned(
|
||||
.with_row_ids(row_ids.to_vec())
|
||||
.preserve_order(true)
|
||||
.execute()
|
||||
.await?;
|
||||
.await
|
||||
.map_err(|err| map_blob_take_error(column, row_ids.len(), err))?;
|
||||
ensure_all_row_ids_resolved(column, row_ids.len(), payloads.len())?;
|
||||
|
||||
let mut builder = LargeBinaryBuilder::new();
|
||||
@@ -458,7 +487,10 @@ pub(crate) async fn take_blob_files_aligned(
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let handles = dataset.take_blobs(row_ids, column).await?;
|
||||
let handles = dataset
|
||||
.take_blobs(row_ids, column)
|
||||
.await
|
||||
.map_err(|err| map_blob_take_error(column, row_ids.len(), err))?;
|
||||
ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?;
|
||||
Ok(handles
|
||||
.into_iter()
|
||||
@@ -504,6 +536,21 @@ mod tests {
|
||||
params.data_storage_version.unwrap().resolve(),
|
||||
ConcreteFileVersion::V2_2
|
||||
);
|
||||
assert!(!params.enable_stable_row_ids);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn storage_version_leaves_stable_row_ids_enabled() {
|
||||
let mut params = WriteParams {
|
||||
enable_stable_row_ids: true,
|
||||
..Default::default()
|
||||
};
|
||||
ensure_blob_storage_version(&blob_schema(), &mut params);
|
||||
assert!(params.enable_stable_row_ids);
|
||||
assert_eq!(
|
||||
params.data_storage_version.unwrap().resolve(),
|
||||
ConcreteFileVersion::V2_2
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -576,5 +623,6 @@ mod tests {
|
||||
let mut params = WriteParams::default();
|
||||
ensure_blob_storage_version(&schema, &mut params);
|
||||
assert!(params.data_storage_version.is_none());
|
||||
assert!(!params.enable_stable_row_ids);
|
||||
}
|
||||
}
|
||||
|
||||
+263
-23
@@ -23,15 +23,18 @@ use crate::connection::create_table::CreateTableBuilder;
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::database::listing::ListingDatabase;
|
||||
use crate::database::{
|
||||
CloneTableRequest, Database, DatabaseOptions, JobDescription, JobInfo, OpenTableRequest,
|
||||
ReadConsistency, TableNamesRequest,
|
||||
CloneTableRequest, Database, DatabaseOptions, JobInfo, OpenTableRequest, ReadConsistency,
|
||||
SecretInfo, TableNamesRequest,
|
||||
};
|
||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
#[cfg(feature = "remote")]
|
||||
use crate::remote::{
|
||||
client::ClientConfig,
|
||||
db::{OPT_REMOTE_API_KEY, OPT_REMOTE_HOST_OVERRIDE, OPT_REMOTE_REGION},
|
||||
db::{
|
||||
OPT_REMOTE_API_KEY, OPT_REMOTE_HOST_OVERRIDE, OPT_REMOTE_REGION,
|
||||
OPT_REMOTE_SQL_HOST_OVERRIDE,
|
||||
},
|
||||
};
|
||||
use lance::io::ObjectStoreParams;
|
||||
pub use lance_file::version::LanceFileVersion;
|
||||
@@ -322,6 +325,43 @@ pub struct CloneTableBuilder {
|
||||
request: CloneTableRequest,
|
||||
}
|
||||
|
||||
/// Builder for asynchronously executing a SQL statement on a remote database.
|
||||
pub struct ExecuteQueryAsyncBuilder {
|
||||
parent: Arc<dyn Database>,
|
||||
query: String,
|
||||
default_namespace_path: Vec<String>,
|
||||
}
|
||||
|
||||
impl ExecuteQueryAsyncBuilder {
|
||||
fn new(parent: Arc<dyn Database>, query: String) -> Self {
|
||||
Self {
|
||||
parent,
|
||||
query,
|
||||
default_namespace_path: vec!["public".to_string()],
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the namespace used for unqualified table names.
|
||||
///
|
||||
/// An empty path is treated as `public`, which is the SQL name for the
|
||||
/// root Lance namespace.
|
||||
pub fn default_namespace_path<I, S>(mut self, path: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = S>,
|
||||
S: Into<String>,
|
||||
{
|
||||
self.default_namespace_path = path.into_iter().map(Into::into).collect();
|
||||
self
|
||||
}
|
||||
|
||||
/// Start the statement and return its asynchronous query handle.
|
||||
pub async fn execute(self) -> Result<crate::sql::Query> {
|
||||
self.parent
|
||||
.execute_query_async(&self.query, &self.default_namespace_path)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
impl CloneTableBuilder {
|
||||
fn new(parent: Arc<dyn Database>, target_table_name: String, source_uri: String) -> Self {
|
||||
Self {
|
||||
@@ -405,6 +445,51 @@ impl Connection {
|
||||
&self.internal
|
||||
}
|
||||
|
||||
/// Start executing SQL on a remote LanceDB database.
|
||||
///
|
||||
/// The query can reference tables in other databases with SQL dot notation.
|
||||
/// Use [`ExecuteQueryAsyncBuilder::default_namespace_path`] to avoid qualifying
|
||||
/// tables in the default namespace. Local connections return
|
||||
/// [`Error::NotSupported`].
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```no_run
|
||||
/// # async fn query(db: &lancedb::Connection) -> lancedb::Result<()> {
|
||||
/// use futures::TryStreamExt;
|
||||
///
|
||||
/// let query = db
|
||||
/// .execute_query_async("SELECT * FROM events LIMIT 10")
|
||||
/// .default_namespace_path(["public"])
|
||||
/// .execute()
|
||||
/// .await?;
|
||||
/// println!("query id: {}", query.id());
|
||||
/// let mut batches = query.reader().await?;
|
||||
/// while let Some(batch) = batches.try_next().await? {
|
||||
/// println!("received {} rows", batch.num_rows());
|
||||
/// }
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub fn execute_query_async(&self, query: impl Into<String>) -> ExecuteQueryAsyncBuilder {
|
||||
ExecuteQueryAsyncBuilder::new(self.internal.clone(), query.into())
|
||||
}
|
||||
|
||||
/// Describe a submitted SQL query by its connection-scoped id.
|
||||
///
|
||||
/// This performs one bounded status poll using state retained by this
|
||||
/// connection. Running state with a live query handle is not evicted;
|
||||
/// abandoned state has bounded retention, and server expiration is
|
||||
/// honored. Terminal state is retained briefly.
|
||||
/// Query ids are not portable to another connection. Local connections
|
||||
/// return [`Error::NotSupported`].
|
||||
pub async fn describe_query(
|
||||
&self,
|
||||
query_id: uuid::Uuid,
|
||||
) -> Result<crate::sql::QueryDescription> {
|
||||
self.internal.describe_query(query_id).await
|
||||
}
|
||||
|
||||
/// Get the names of all tables in the database
|
||||
///
|
||||
/// The names will be returned in lexicographical order (ascending)
|
||||
@@ -523,6 +608,118 @@ impl Connection {
|
||||
.await
|
||||
}
|
||||
|
||||
/// List every published immutable Function version in the remote catalog.
|
||||
///
|
||||
/// Results are ordered by Function name then version. The client walks all
|
||||
/// server pages before returning. Local databases return
|
||||
/// [`Error::NotSupported`].
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```no_run
|
||||
/// # async fn list_functions(
|
||||
/// # connection: &lancedb::Connection,
|
||||
/// # ) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// for function in connection.list_functions().await? {
|
||||
/// println!("{} {}", function.name(), function.version());
|
||||
/// }
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn list_functions(&self) -> Result<Vec<crate::function::FunctionVersion>> {
|
||||
self.internal.list_functions().await
|
||||
}
|
||||
|
||||
/// Drop one exact immutable Function version from the remote catalog.
|
||||
///
|
||||
/// Returns `true` when the server appended a Dropped transition and
|
||||
/// `false` for an idempotent replay. Local databases return
|
||||
/// [`Error::NotSupported`].
|
||||
pub async fn drop_function(
|
||||
&self,
|
||||
name: impl AsRef<str>,
|
||||
version: impl AsRef<str>,
|
||||
) -> Result<bool> {
|
||||
self.internal
|
||||
.drop_function(name.as_ref(), version.as_ref())
|
||||
.await
|
||||
}
|
||||
|
||||
/// Create a named Secret in this database.
|
||||
///
|
||||
/// Fails if the name is taken, so a create can never silently become a
|
||||
/// rotation. There is no API that reads a stored credential back; the only
|
||||
/// consumer is a Function that binds the Secret by name. Local databases
|
||||
/// return [`Error::NotSupported`].
|
||||
pub async fn create_secret(
|
||||
&self,
|
||||
name: impl AsRef<str>,
|
||||
value: impl AsRef<str>,
|
||||
namespace_path: &[String],
|
||||
) -> Result<()> {
|
||||
self.internal
|
||||
.create_secret(name.as_ref(), value.as_ref(), namespace_path)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Replace the credential behind an existing Secret.
|
||||
///
|
||||
/// Fails if it does not exist. Every Function bound to the Secret resolves
|
||||
/// the new value from its next execution, and no new Function version is
|
||||
/// minted -- which is what lets a rotation reach columns pinned to a
|
||||
/// version registered before it. Local databases return
|
||||
/// [`Error::NotSupported`].
|
||||
pub async fn alter_secret(
|
||||
&self,
|
||||
name: impl AsRef<str>,
|
||||
value: impl AsRef<str>,
|
||||
namespace_path: &[String],
|
||||
) -> Result<()> {
|
||||
self.internal
|
||||
.alter_secret(name.as_ref(), value.as_ref(), namespace_path)
|
||||
.await
|
||||
}
|
||||
|
||||
/// The names of every Secret in this database.
|
||||
///
|
||||
/// Names only. No path in this API returns a stored credential, by
|
||||
/// construction rather than by policy. Local databases return
|
||||
/// [`Error::NotSupported`].
|
||||
pub async fn list_secrets(&self, namespace_path: &[String]) -> Result<Vec<String>> {
|
||||
self.internal.list_secrets(namespace_path).await
|
||||
}
|
||||
|
||||
/// Drop a Secret.
|
||||
///
|
||||
/// Functions bound to it fail at their next job, naming the Secret; that
|
||||
/// is the revocation path. The name becomes free to reuse, and a new
|
||||
/// Secret under it is picked up by everything still bound to that name.
|
||||
/// Local databases return [`Error::NotSupported`].
|
||||
pub async fn drop_secret(
|
||||
&self,
|
||||
name: impl AsRef<str>,
|
||||
namespace_path: &[String],
|
||||
) -> Result<()> {
|
||||
self.internal
|
||||
.drop_secret(name.as_ref(), namespace_path)
|
||||
.await
|
||||
}
|
||||
|
||||
/// What this database records about one Secret: its name and timestamps.
|
||||
///
|
||||
/// Never the value. The type it returns has no field for one, so this is a
|
||||
/// property of the API rather than of what the caller chooses to read.
|
||||
/// Local databases return [`Error::NotSupported`].
|
||||
pub async fn describe_secret(
|
||||
&self,
|
||||
name: impl AsRef<str>,
|
||||
namespace_path: &[String],
|
||||
) -> Result<SecretInfo> {
|
||||
self.internal
|
||||
.describe_secret(name.as_ref(), namespace_path)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Rename a table in the database.
|
||||
///
|
||||
/// This is only supported in LanceDB Cloud.
|
||||
@@ -548,14 +745,34 @@ impl Connection {
|
||||
self.internal.read_consistency().await
|
||||
}
|
||||
|
||||
/// A [`crate::job::Job`] handle for a server-side job by id, suitable for
|
||||
/// waiting on or cancelling the job.
|
||||
/// Open a server-side job by id, returning a handle with its record
|
||||
/// already populated. Fails with [`crate::Error::JobNotFound`] when the
|
||||
/// server has no such job, the way [`Connection::open_table`] does for a
|
||||
/// missing table.
|
||||
///
|
||||
/// The handle is constructed without a server round trip; an unknown id
|
||||
/// surfaces when the handle is used. Only server-backed databases support
|
||||
/// job handles by id.
|
||||
pub fn job(&self, job_id: impl AsRef<str>) -> Result<crate::job::Job> {
|
||||
self.internal.job(job_id.as_ref())
|
||||
/// This is the one way in: the returned [`crate::job::Job`] answers for
|
||||
/// its own state, specification, result, failure and event history, so
|
||||
/// there is no separate connection-level call for any of them.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```no_run
|
||||
/// # use lancedb::job::JobEventsRequest;
|
||||
/// # async fn open_job(
|
||||
/// # connection: &lancedb::Connection,
|
||||
/// # job_id: &str,
|
||||
/// # ) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// let job = connection.open_job(job_id).await?;
|
||||
/// println!("{:?} {:?}", job.state(), job.result());
|
||||
/// let done = job
|
||||
/// .events(JobEventsRequest::default().filter("state = 'claim_complete'"))
|
||||
/// .await?;
|
||||
/// println!("{} completions", done.iter().map(|b| b.num_rows()).sum::<usize>());
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn open_job(&self, job_id: impl AsRef<str>) -> Result<crate::job::Job> {
|
||||
self.internal.open_job(job_id.as_ref()).await
|
||||
}
|
||||
|
||||
/// List server-side jobs across the database's tables.
|
||||
@@ -563,24 +780,12 @@ impl Connection {
|
||||
self.internal.list_jobs().await
|
||||
}
|
||||
|
||||
/// Describe a single server-side job by id. `None` when the server has no
|
||||
/// such job.
|
||||
pub async fn get_job(&self, job_id: impl AsRef<str>) -> Result<Option<JobDescription>> {
|
||||
self.internal.get_job(job_id.as_ref()).await
|
||||
}
|
||||
|
||||
/// Request cancellation of a server-side job by id. Returns true if the
|
||||
/// server accepted the cancellation, false if no such job exists.
|
||||
pub async fn cancel_job(&self, job_id: impl AsRef<str>) -> Result<bool> {
|
||||
self.internal.cancel_job(job_id.as_ref()).await
|
||||
}
|
||||
|
||||
/// The lifecycle event history of a server-side job (all jobs when
|
||||
/// `job_id` is `None`), as recorded Arrow batches.
|
||||
pub async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<RecordBatch>> {
|
||||
self.internal.job_history(job_id).await
|
||||
}
|
||||
|
||||
/// Drop a table in the database.
|
||||
///
|
||||
/// # Arguments
|
||||
@@ -827,6 +1032,19 @@ impl ConnectBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the SQL service host override for a remote connection.
|
||||
///
|
||||
/// The SQL client is initialized lazily when the connection first executes
|
||||
/// SQL and is retained for the connection's lifetime.
|
||||
#[cfg(feature = "remote")]
|
||||
pub fn sql_host_override(mut self, sql_host_override: &str) -> Self {
|
||||
self.request.options.insert(
|
||||
OPT_REMOTE_SQL_HOST_OVERRIDE.to_string(),
|
||||
sql_host_override.to_string(),
|
||||
);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the database specific options
|
||||
///
|
||||
/// See [crate::database::listing::ListingDatabaseOptions] for the options available for
|
||||
@@ -1016,6 +1234,7 @@ impl ConnectBuilder {
|
||||
|
||||
let mut merged_options = self.request.options.clone();
|
||||
Self::apply_env_defaults(&ENV_VARS_TO_STORAGE_OPTS, &mut merged_options);
|
||||
let sql_host_override = merged_options.get(OPT_REMOTE_SQL_HOST_OVERRIDE).cloned();
|
||||
let options = RemoteDatabaseOptions::parse_from_map(&merged_options)?;
|
||||
|
||||
let region = options.region.ok_or_else(|| Error::InvalidInput {
|
||||
@@ -1057,11 +1276,15 @@ impl ConnectBuilder {
|
||||
}
|
||||
|
||||
let storage_options = StorageOptions(options.storage_options.clone());
|
||||
let host_overrides = crate::remote::db::RemoteHostOverrides {
|
||||
rest: options.host_override,
|
||||
sql: sql_host_override,
|
||||
};
|
||||
let internal = Arc::new(crate::remote::db::RemoteDatabase::try_new(
|
||||
&self.request.uri,
|
||||
&api_key,
|
||||
®ion,
|
||||
options.host_override,
|
||||
host_overrides,
|
||||
client_config,
|
||||
storage_options.into(),
|
||||
self.request.read_consistency_interval,
|
||||
@@ -1355,6 +1578,23 @@ mod tests {
|
||||
assert_eq!(tc.connection.uri(), tc.uri);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_local_connection_rejects_sql_queries() {
|
||||
let directory = tempdir().unwrap();
|
||||
let connection = connect(directory.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
connection.execute_query_async("SELECT 1").execute().await,
|
||||
Err(Error::NotSupported { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
connection.describe_query(uuid::Uuid::nil()).await,
|
||||
Err(Error::NotSupported { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[cfg(feature = "remote")]
|
||||
#[test]
|
||||
fn test_apply_env_defaults() {
|
||||
|
||||
@@ -18,8 +18,6 @@ use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use arrow_array::RecordBatch;
|
||||
|
||||
use lance::dataset::ReadParams;
|
||||
use lance_namespace::LanceNamespace;
|
||||
use lance_namespace::models::{
|
||||
@@ -206,8 +204,8 @@ pub enum ReadConsistency {
|
||||
/// compaction, column refresh, ...).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobInfo {
|
||||
/// The job id -- what [`Database::get_job`] and [`Database::cancel_job`]
|
||||
/// accept.
|
||||
/// The job id -- what [`Database::open_job`] and
|
||||
/// [`Database::cancel_job`] accept.
|
||||
pub job_id: String,
|
||||
/// The table the job runs against, without URI or namespace.
|
||||
pub table: String,
|
||||
@@ -218,8 +216,8 @@ pub struct JobInfo {
|
||||
pub created_at_millis: i64,
|
||||
}
|
||||
|
||||
/// A described job from [`Database::get_job`]: lifecycle state plus the
|
||||
/// job-type-specific specification.
|
||||
/// The server-side record behind a [`crate::job::Job`] handle: lifecycle
|
||||
/// state plus the job-type-specific specification and result.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobDescription {
|
||||
pub job_id: String,
|
||||
@@ -230,6 +228,10 @@ pub struct JobDescription {
|
||||
pub creation_ms: i64,
|
||||
/// The job-type-specific specification. Null when the server omits it.
|
||||
pub spec: serde_json::Value,
|
||||
/// The job-type-specific terminal result, for job types that define one.
|
||||
/// `None` until the job succeeds, so a job that never terminates reports
|
||||
/// its progress through [`crate::job::Job::events`] instead.
|
||||
pub result: Option<serde_json::Value>,
|
||||
/// Why the job failed, when the job is failed and the server reports a
|
||||
/// reason.
|
||||
pub failure: Option<crate::error::JobFailure>,
|
||||
@@ -247,9 +249,29 @@ fn function_catalog_not_supported<T>() -> Result<T> {
|
||||
})
|
||||
}
|
||||
|
||||
fn secret_catalog_not_supported<T>() -> Result<T> {
|
||||
Err(crate::error::Error::NotSupported {
|
||||
message: "Secret operations are not supported by this database".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// The `Database` trait defines the interface for database implementations.
|
||||
///
|
||||
/// A database is responsible for managing tables and their metadata.
|
||||
/// What a database records about a Secret. Never its value.
|
||||
///
|
||||
/// Returned by [`crate::connection::Connection::describe_secret`]. There is no
|
||||
/// field for the credential and no method that could produce one.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)]
|
||||
pub struct SecretInfo {
|
||||
/// The Secret's database-scoped name.
|
||||
pub name: String,
|
||||
/// When the Secret was created, as an RFC 3339 timestamp.
|
||||
pub created_at: String,
|
||||
/// When the Secret's value was last rotated, as an RFC 3339 timestamp.
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
pub trait Database:
|
||||
Send + Sync + std::any::Any + std::fmt::Debug + std::fmt::Display + 'static
|
||||
@@ -307,30 +329,83 @@ pub trait Database:
|
||||
) -> Result<crate::function::FunctionVersion> {
|
||||
function_catalog_not_supported()
|
||||
}
|
||||
/// A [`crate::job::Job`] handle for a server-side job by id, suitable for
|
||||
/// waiting on or cancelling the job. The handle is constructed without a
|
||||
/// server round trip; an unknown id surfaces when the handle is used.
|
||||
fn job(&self, _job_id: &str) -> Result<crate::job::Job> {
|
||||
job_op_not_supported("job")
|
||||
/// List every published immutable Function version in the remote catalog.
|
||||
async fn list_functions(&self) -> Result<Vec<crate::function::FunctionVersion>> {
|
||||
function_catalog_not_supported()
|
||||
}
|
||||
/// Drop one exact immutable Function version from the remote catalog.
|
||||
async fn drop_function(&self, _name: &str, _version: &str) -> Result<bool> {
|
||||
function_catalog_not_supported()
|
||||
}
|
||||
/// Create a named Secret in this database. Fails if the name is taken, so
|
||||
/// a create can never silently become a rotation.
|
||||
async fn create_secret(
|
||||
&self,
|
||||
_name: &str,
|
||||
_value: &str,
|
||||
_namespace_path: &[String],
|
||||
) -> Result<()> {
|
||||
secret_catalog_not_supported()
|
||||
}
|
||||
/// Replace the credential behind an existing Secret. Fails if it does not
|
||||
/// exist. Every Function bound to it resolves the new value from its next
|
||||
/// execution, with no new Function version.
|
||||
async fn alter_secret(
|
||||
&self,
|
||||
_name: &str,
|
||||
_value: &str,
|
||||
_namespace_path: &[String],
|
||||
) -> Result<()> {
|
||||
secret_catalog_not_supported()
|
||||
}
|
||||
/// The names of every Secret in this database.
|
||||
///
|
||||
/// Names only. No API path returns a stored credential, by construction
|
||||
/// rather than by policy.
|
||||
async fn list_secrets(&self, _namespace_path: &[String]) -> Result<Vec<String>> {
|
||||
secret_catalog_not_supported()
|
||||
}
|
||||
/// Drop a Secret. Functions bound to it fail at their next job, which is
|
||||
/// the revocation path.
|
||||
async fn drop_secret(&self, _name: &str, _namespace_path: &[String]) -> Result<()> {
|
||||
secret_catalog_not_supported()
|
||||
}
|
||||
/// What the database records about one Secret: its name and timestamps,
|
||||
/// never its value.
|
||||
async fn describe_secret(&self, _name: &str, _namespace_path: &[String]) -> Result<SecretInfo> {
|
||||
secret_catalog_not_supported()
|
||||
}
|
||||
/// Open a job by id, returning a handle with its record already
|
||||
/// populated. Fails with [`crate::Error::JobNotFound`] when the server has
|
||||
/// no such job.
|
||||
async fn open_job(&self, _job_id: &str) -> Result<crate::job::Job> {
|
||||
job_op_not_supported("open_job")
|
||||
}
|
||||
/// List server-side jobs across the database's tables.
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
job_op_not_supported("list_jobs")
|
||||
}
|
||||
/// Describe a single job by id. `None` when the server has no such job.
|
||||
async fn get_job(&self, _job_id: &str) -> Result<Option<JobDescription>> {
|
||||
job_op_not_supported("get_job")
|
||||
}
|
||||
/// Request cancellation of a job by id. Returns true if the server
|
||||
/// accepted the cancellation, false if no such job exists. Cancelling an
|
||||
/// already-terminal job is a no-op success.
|
||||
async fn cancel_job(&self, _job_id: &str) -> Result<bool> {
|
||||
job_op_not_supported("cancel_job")
|
||||
}
|
||||
/// The lifecycle event history of a job (all jobs when `job_id` is
|
||||
/// `None`), as recorded Arrow batches.
|
||||
async fn job_history(&self, _job_id: Option<&str>) -> Result<Vec<RecordBatch>> {
|
||||
job_op_not_supported("job_history")
|
||||
/// Start executing a SQL statement on a remote database.
|
||||
async fn execute_query_async(
|
||||
&self,
|
||||
_query: &str,
|
||||
_default_namespace_path: &[String],
|
||||
) -> Result<crate::sql::Query> {
|
||||
Err(crate::error::Error::NotSupported {
|
||||
message: "SQL is not supported by this database".to_string(),
|
||||
})
|
||||
}
|
||||
/// Describe a submitted SQL query by its connection-scoped id.
|
||||
async fn describe_query(&self, _query_id: uuid::Uuid) -> Result<crate::sql::QueryDescription> {
|
||||
Err(crate::error::Error::NotSupported {
|
||||
message: "SQL is not supported by this database".to_string(),
|
||||
})
|
||||
}
|
||||
/// Open a table in the database
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
|
||||
@@ -18,7 +18,7 @@ use lance_table::io::commit::commit_handler_from_url;
|
||||
use object_store::local::LocalFileSystem;
|
||||
use snafu::ResultExt;
|
||||
|
||||
use crate::blob::{ensure_blob_storage_version, has_blob_columns};
|
||||
use crate::blob::ensure_blob_storage_version;
|
||||
use crate::connection::ConnectRequest;
|
||||
use crate::database::ReadConsistency;
|
||||
use crate::database::namespace::LanceNamespaceDatabase;
|
||||
@@ -827,7 +827,6 @@ impl ListingDatabase {
|
||||
if let Some(enable_stable_row_ids) = overrides
|
||||
.enable_stable_row_ids
|
||||
.or(self.new_table_config.enable_stable_row_ids)
|
||||
.or(has_blob_columns(&data_schema).then_some(true))
|
||||
{
|
||||
write_params.enable_stable_row_ids = enable_stable_row_ids;
|
||||
}
|
||||
|
||||
@@ -23,7 +23,7 @@ use lance_namespace_impls::ConnectBuilder;
|
||||
use lance_table::io::commit::CommitHandler;
|
||||
use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
|
||||
|
||||
use crate::blob::{ensure_blob_storage_version, has_blob_columns};
|
||||
use crate::blob::ensure_blob_storage_version;
|
||||
use crate::connection::NamespaceClientPushdownOperation;
|
||||
use crate::database::ReadConsistency;
|
||||
use crate::database::listing::{NewTableConfig, take_request_creation_overrides};
|
||||
@@ -217,7 +217,6 @@ impl LanceNamespaceDatabase {
|
||||
if let Some(enable_stable_row_ids) = overrides
|
||||
.enable_stable_row_ids
|
||||
.or(self.new_table_config.enable_stable_row_ids)
|
||||
.or(has_blob_columns(data_schema.as_ref()).then_some(true))
|
||||
{
|
||||
params.enable_stable_row_ids = enable_stable_row_ids;
|
||||
}
|
||||
@@ -539,9 +538,7 @@ impl Database for LanceNamespaceDatabase {
|
||||
self.namespace
|
||||
.drop_table(drop_request)
|
||||
.await
|
||||
.map_err(|e| Error::Runtime {
|
||||
message: format!("Failed to drop table: {}", e),
|
||||
})?;
|
||||
.map_err(|e| map_namespace_lance_error(e, name))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1495,6 +1492,15 @@ mod tests {
|
||||
.expect("Failed to list tables");
|
||||
assert!(!table_names_after.contains(&"drop_test".to_string()));
|
||||
|
||||
let error = conn
|
||||
.drop_table("drop_test", &["test_ns".into()])
|
||||
.await
|
||||
.expect_err("dropping a missing table should fail");
|
||||
assert!(
|
||||
matches!(error, Error::TableNotFound { ref name, .. } if name == "drop_test"),
|
||||
"expected TableNotFound, got: {error:?}"
|
||||
);
|
||||
|
||||
// Verify: Cannot open dropped table
|
||||
let open_result = conn.open_table("drop_test").execute().await;
|
||||
assert!(open_result.is_err());
|
||||
|
||||
@@ -31,7 +31,7 @@ use lance::io::RecordBatchStream;
|
||||
use lance_arrow::RecordBatchExt;
|
||||
use lance_core::ROW_ID;
|
||||
use lance_core::error::LanceOptionExt;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Reads a permutation of a source table based on row IDs stored in a separate table
|
||||
@@ -234,7 +234,14 @@ impl PermutationReader {
|
||||
.expect_ok()?
|
||||
.values();
|
||||
|
||||
let in_list: Vec<Expr> = row_ids.iter().map(|id| lit(*id)).collect();
|
||||
let mut unique_row_ids = HashSet::with_capacity(num_rows);
|
||||
let in_list: Vec<Expr> = row_ids
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|row_id| unique_row_ids.insert(*row_id))
|
||||
.map(lit)
|
||||
.collect();
|
||||
let num_unique_row_ids = unique_row_ids.len();
|
||||
|
||||
let base_query = QueryRequest {
|
||||
filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))),
|
||||
@@ -247,7 +254,7 @@ impl PermutationReader {
|
||||
.query(
|
||||
&AnyQuery::Query(base_query),
|
||||
QueryExecutionOptions {
|
||||
max_batch_length: num_rows as u32,
|
||||
max_batch_length: num_unique_row_ids as u32,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
@@ -262,9 +269,9 @@ impl PermutationReader {
|
||||
});
|
||||
}
|
||||
|
||||
if batches.iter().map(|b| b.num_rows()).sum::<usize>() != num_rows {
|
||||
if batches.iter().map(|b| b.num_rows()).sum::<usize>() != num_unique_row_ids {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "Base table returned different number of rows than the number of row IDs"
|
||||
message: "Base table returned a different number of rows than the number of unique row IDs"
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
@@ -504,6 +511,7 @@ impl PermutationReader {
|
||||
let table = Table::from(self.base_table.clone());
|
||||
let batches = table
|
||||
.take_offsets(offsets.to_vec())
|
||||
.preserve_order()
|
||||
.select(selection.clone())
|
||||
.execute()
|
||||
.await?
|
||||
@@ -803,10 +811,10 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
// Take offsets in reverse order and verify returned rows match that order
|
||||
let offsets = vec![5, 3, 1, 0];
|
||||
let offsets = vec![5, 3, 5, 1, 0];
|
||||
let batch = reader.take_offsets(&offsets, Select::All).await.unwrap();
|
||||
|
||||
assert_eq!(batch.num_rows(), 4);
|
||||
assert_eq!(batch.num_rows(), 5);
|
||||
|
||||
let idx_values = batch
|
||||
.column(0)
|
||||
@@ -820,6 +828,52 @@ mod tests {
|
||||
assert_eq!(idx_values, expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_preserves_repeated_rows_in_permutation() {
|
||||
let base_table = lance_datagen::gen_batch()
|
||||
.col("idx", lance_datagen::array::step::<Int32Type>())
|
||||
.into_mem_table("tbl", RowCount::from(5), BatchCount::from(1))
|
||||
.await;
|
||||
let base_row_ids = collect_column::<UInt64Type>(&base_table, "_rowid").await;
|
||||
let permutation_row_ids = vec![
|
||||
base_row_ids[3],
|
||||
base_row_ids[1],
|
||||
base_row_ids[3],
|
||||
base_row_ids[2],
|
||||
];
|
||||
let permutation_batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("row_id", DataType::UInt64, false),
|
||||
Field::new(SPLIT_ID_COLUMN, DataType::UInt64, false),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(UInt64Array::from(permutation_row_ids)),
|
||||
Arc::new(UInt64Array::from(vec![0; 4])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let permutation_table = virtual_table("row_ids", &permutation_batch).await;
|
||||
let reader = PermutationReader::try_from_tables(
|
||||
base_table.base_table().clone(),
|
||||
permutation_table.base_table().clone(),
|
||||
0,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let batch = reader
|
||||
.take_offsets(&[0, 1, 2, 3], Select::All)
|
||||
.await
|
||||
.unwrap();
|
||||
let idx_values = batch
|
||||
.column(0)
|
||||
.as_primitive::<Int32Type>()
|
||||
.values()
|
||||
.to_vec();
|
||||
|
||||
assert_eq!(idx_values, vec![3, 1, 3, 2]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_with_column_selection() {
|
||||
let (base_table, row_ids_table, row_ids) = setup_permutation_tables(10).await;
|
||||
@@ -883,17 +937,17 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
// With no permutation table, take_offsets uses the base table directly
|
||||
let offsets = vec![0, 2, 4, 6];
|
||||
let offsets = vec![0, 2, 0, 4, 6];
|
||||
let batch = reader.take_offsets(&offsets, Select::All).await.unwrap();
|
||||
|
||||
assert_eq!(batch.num_rows(), 4);
|
||||
assert_eq!(batch.num_rows(), 5);
|
||||
|
||||
let idx_values = batch
|
||||
.column(0)
|
||||
.as_primitive::<Int32Type>()
|
||||
.values()
|
||||
.to_vec();
|
||||
assert_eq!(idx_values, vec![0, 2, 4, 6]);
|
||||
assert_eq!(idx_values, vec![0, 2, 0, 4, 6]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -102,6 +102,8 @@ pub enum Error {
|
||||
},
|
||||
#[snafu(display("Job{} was cancelled", job_id.as_ref().map(|id| format!(" {id}")).unwrap_or_default()))]
|
||||
JobCancelled { job_id: Option<String> },
|
||||
#[snafu(display("Job '{job_id}' was not found"))]
|
||||
JobNotFound { job_id: String },
|
||||
|
||||
// 3rd party / external errors
|
||||
#[snafu(display("object_store error: {source}"))]
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
//! backend-neutral terminal result of a computed-column refresh.
|
||||
//!
|
||||
//! This module contains client/wire values only. Catalog persistence,
|
||||
//! environment bake, and execution are owned by Sophon.
|
||||
//! environment bake, secret resolution, and execution are owned by Sophon.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
@@ -409,6 +409,8 @@ pub struct FunctionVersion {
|
||||
runtime: PythonRuntimeSpec,
|
||||
runtime_digest: String,
|
||||
environment_digest: String,
|
||||
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
|
||||
secret_bindings: BTreeMap<String, String>,
|
||||
created_at: String,
|
||||
}
|
||||
|
||||
@@ -441,6 +443,16 @@ impl FunctionVersion {
|
||||
&self.environment_digest
|
||||
}
|
||||
|
||||
/// Declared environment variable name to the Secret each one resolves.
|
||||
///
|
||||
/// Bindings are part of this version's identity; the credentials behind
|
||||
/// them are not, and resolve at execution. Rotating a bound Secret
|
||||
/// therefore changes what the same version runs with, and no value has a
|
||||
/// field in this model.
|
||||
pub fn secret_bindings(&self) -> &BTreeMap<String, String> {
|
||||
&self.secret_bindings
|
||||
}
|
||||
|
||||
pub fn created_at(&self) -> &str {
|
||||
&self.created_at
|
||||
}
|
||||
@@ -482,12 +494,20 @@ pub struct FunctionArtifactRequest {
|
||||
}
|
||||
|
||||
/// Stable request envelope for remote immutable Function registration.
|
||||
///
|
||||
/// Credential values deliberately have no field here. The only secret-shaped
|
||||
/// thing a client sends is `secret_bindings`: the name of a Secret the
|
||||
/// database already holds, which Sophon resolves inside the remote runtime.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct FunctionRegistrationRequest {
|
||||
pub name: String,
|
||||
pub artifact: FunctionArtifactRequest,
|
||||
pub signature: FunctionSignature,
|
||||
pub runtime: PythonRuntimeSpec,
|
||||
/// Declared environment variable name to the Secret it binds. Every bound
|
||||
/// Secret must already exist; registration fails otherwise.
|
||||
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
|
||||
pub secret_bindings: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
impl_json!(FunctionRegistrationRequest);
|
||||
@@ -582,8 +602,8 @@ pub struct InputBinding {
|
||||
|
||||
/// Ordered result-field to table-field mapping for a Function binding.
|
||||
///
|
||||
/// Assignment state is not part of the Slice 1 client contract. During the
|
||||
/// NULL transition there is no public Lance cell-flag identifier to persist.
|
||||
/// `nullable` describes the logical Function result. Physical computed-column
|
||||
/// fields remain nullable while unassigned.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct OutputMapping {
|
||||
pub result_field: String,
|
||||
@@ -594,6 +614,14 @@ pub struct OutputMapping {
|
||||
pub nullable: bool,
|
||||
}
|
||||
|
||||
/// Internal physical column preserving the parent validity of a flattened
|
||||
/// named-struct result.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AssignmentMapping {
|
||||
pub output_name: String,
|
||||
pub output_field_id: i32,
|
||||
}
|
||||
|
||||
/// Immutable Function binding persisted by the Enterprise table service.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct FunctionBinding {
|
||||
@@ -601,6 +629,8 @@ pub struct FunctionBinding {
|
||||
function: FunctionVersionRef,
|
||||
inputs: Vec<InputBinding>,
|
||||
outputs: Vec<OutputMapping>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
assignment: Option<AssignmentMapping>,
|
||||
/// Exact Arrow schema presented to the Function, encoded with the Lance
|
||||
/// Namespace Arrow JSON representation.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
@@ -627,6 +657,10 @@ impl FunctionBinding {
|
||||
&self.outputs
|
||||
}
|
||||
|
||||
pub fn assignment(&self) -> Option<&AssignmentMapping> {
|
||||
self.assignment.as_ref()
|
||||
}
|
||||
|
||||
pub fn input_schema(&self) -> Option<&Value> {
|
||||
self.input_schema.as_ref()
|
||||
}
|
||||
|
||||
+247
-12
@@ -3,16 +3,53 @@
|
||||
|
||||
//! Handles to operations a server may run asynchronously.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
|
||||
|
||||
use arrow_array::RecordBatch;
|
||||
use async_trait::async_trait;
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::Value;
|
||||
use tokio::sync::watch;
|
||||
use tokio::task::{AbortHandle, JoinHandle};
|
||||
|
||||
use crate::database::JobDescription;
|
||||
use crate::error::{Error, JobFailure, Result};
|
||||
|
||||
/// Which of a job's events [`Job::events`] returns.
|
||||
///
|
||||
/// The handle already knows which job to ask about, so this narrows the
|
||||
/// query rather than naming one.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct JobEventsRequest {
|
||||
/// Maximum event rows to return. The server applies its own default
|
||||
/// (1000 rows) and maximum (10,000 rows) when this is `None`, and
|
||||
/// truncates without saying so, which matters for a job with one event
|
||||
/// per fragment.
|
||||
pub limit: Option<u32>,
|
||||
/// SQL-like filter over the event columns `state`, `updated_by`,
|
||||
/// `emitted_from`, `emitted_by`, and `claim_entity`. For example
|
||||
/// `state = 'claim_complete'` selects only per-claim completions.
|
||||
pub filter: Option<String>,
|
||||
}
|
||||
|
||||
impl JobEventsRequest {
|
||||
pub fn limit(mut self, limit: u32) -> Self {
|
||||
self.limit = Some(limit);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn filter(mut self, filter: impl Into<String>) -> Self {
|
||||
self.filter = Some(filter.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
fn job_detail_not_supported<T>(what: &str) -> Result<T> {
|
||||
Err(Error::NotSupported {
|
||||
message: format!("{what} is only available for server-side jobs"),
|
||||
})
|
||||
}
|
||||
|
||||
/// Backend-specific tracking for an asynchronous operation.
|
||||
#[async_trait]
|
||||
pub(crate) trait JobHandle: Send + Sync {
|
||||
@@ -23,6 +60,15 @@ pub(crate) trait JobHandle: Send + Sync {
|
||||
async fn status(&self) -> Result<String>;
|
||||
async fn wait(&self) -> Result<TerminalResult>;
|
||||
async fn cancel(&self) -> Result<()>;
|
||||
/// The job's full server-side record. Backends that run the operation in
|
||||
/// this process have none and keep the default.
|
||||
async fn describe(&self) -> Result<JobDescription> {
|
||||
job_detail_not_supported("describing a job")
|
||||
}
|
||||
/// The job's recorded lifecycle events.
|
||||
async fn events(&self, _request: JobEventsRequest) -> Result<Vec<RecordBatch>> {
|
||||
job_detail_not_supported("job event history")
|
||||
}
|
||||
}
|
||||
|
||||
/// A backend-neutral successful terminal result.
|
||||
@@ -85,16 +131,34 @@ enum JobInner<T> {
|
||||
Completed(T),
|
||||
}
|
||||
|
||||
/// What a handle last learned about its job. `state` is separate because an
|
||||
/// in-process job can report one but has no server-side record behind it.
|
||||
#[derive(Default)]
|
||||
struct JobCache {
|
||||
state: Option<String>,
|
||||
description: Option<JobDescription>,
|
||||
}
|
||||
|
||||
/// A handle to an operation that may still be running.
|
||||
///
|
||||
/// 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<()>`.
|
||||
///
|
||||
/// The detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the
|
||||
/// handle last observed. Submitting an operation returns only a job id, so
|
||||
/// populating them eagerly would cost an extra round trip on every call:
|
||||
///
|
||||
/// - [`Job::refresh`] and [`Job::status`] fetch the whole record.
|
||||
/// - [`Job::wait`] records the terminal state it establishes, but not the rest
|
||||
/// of the record; call [`Job::refresh`] for that.
|
||||
/// - Everything is `None` until one of those runs.
|
||||
pub struct Job<T = ()>
|
||||
where
|
||||
T: Clone + Send + Sync + 'static,
|
||||
{
|
||||
inner: JobInner<T>,
|
||||
cache: RwLock<JobCache>,
|
||||
}
|
||||
|
||||
impl<T> std::fmt::Debug for Job<T>
|
||||
@@ -102,18 +166,40 @@ where
|
||||
T: Clone + Send + Sync + 'static,
|
||||
{
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("Job")
|
||||
.field("id", &self.id())
|
||||
.field("done", &matches!(self.inner, JobInner::Completed(_)))
|
||||
.finish()
|
||||
let cache = self.cache_read();
|
||||
let mut out = f.debug_struct("Job");
|
||||
out.field("id", &self.id())
|
||||
.field("done", &matches!(self.inner, JobInner::Completed(_)));
|
||||
if let Some(state) = &cache.state {
|
||||
out.field("state", state);
|
||||
}
|
||||
if let Some(description) = &cache.description {
|
||||
out.field("job_type", &description.job_type)
|
||||
.field("creation_ms", &description.creation_ms);
|
||||
if !description.spec.is_null() {
|
||||
out.field("spec", &description.spec);
|
||||
}
|
||||
if let Some(result) = &description.result {
|
||||
out.field("result", result);
|
||||
}
|
||||
if let Some(failure) = &description.failure {
|
||||
out.field("failure", failure);
|
||||
}
|
||||
}
|
||||
out.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Job<()> {
|
||||
/// A job whose operation finished before the handle was created.
|
||||
/// A job whose operation finished before the handle was created. Its
|
||||
/// state is known without asking anyone, so the cache starts populated.
|
||||
pub(crate) fn new_done() -> Self {
|
||||
Self {
|
||||
inner: JobInner::Completed(()),
|
||||
cache: RwLock::new(JobCache {
|
||||
state: Some("finished".to_string()),
|
||||
description: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -123,8 +209,21 @@ impl Job<()> {
|
||||
handle,
|
||||
decode: Arc::new(|_| Ok(())),
|
||||
},
|
||||
cache: RwLock::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// A handle whose record the caller has already fetched, so the detail
|
||||
/// accessors answer without a second round trip.
|
||||
pub(crate) fn opened(handle: Box<dyn JobHandle>, description: JobDescription) -> Self {
|
||||
let job = Self::new(handle);
|
||||
{
|
||||
let mut cache = job.cache_write();
|
||||
cache.state = Some(description.state.clone());
|
||||
cache.description = Some(description);
|
||||
}
|
||||
job
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Job<T>
|
||||
@@ -138,6 +237,7 @@ where
|
||||
handle,
|
||||
decode: Arc::new(TerminalResult::decode::<T>),
|
||||
},
|
||||
cache: RwLock::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -169,16 +269,124 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
fn cache_read(&self) -> RwLockReadGuard<'_, JobCache> {
|
||||
self.cache.read().unwrap_or_else(|err| err.into_inner())
|
||||
}
|
||||
|
||||
fn cache_write(&self) -> RwLockWriteGuard<'_, JobCache> {
|
||||
self.cache.write().unwrap_or_else(|err| err.into_inner())
|
||||
}
|
||||
|
||||
/// Asks the backend for this job's current state, and for a server-side
|
||||
/// job its full record, then caches the answer for the detail accessors.
|
||||
///
|
||||
/// In-process operations have no server-side record, so only
|
||||
/// [`Job::state`] is populated for them.
|
||||
pub async fn refresh(&self) -> Result<()> {
|
||||
self.refresh_state().await.map(|_| ())
|
||||
}
|
||||
|
||||
/// Refreshes and reports the state, which every backend can answer.
|
||||
async fn refresh_state(&self) -> Result<String> {
|
||||
let JobInner::Handle { handle, .. } = &self.inner else {
|
||||
let state = "finished".to_string();
|
||||
self.cache_write().state = Some(state.clone());
|
||||
return Ok(state);
|
||||
};
|
||||
match handle.describe().await {
|
||||
Ok(description) => {
|
||||
let state = description.state.clone();
|
||||
let mut cache = self.cache_write();
|
||||
cache.state = Some(state.clone());
|
||||
cache.description = Some(description);
|
||||
Ok(state)
|
||||
}
|
||||
// An in-process job knows its own state and nothing more.
|
||||
Err(Error::NotSupported { .. }) => {
|
||||
let state = handle.status().await?;
|
||||
self.cache_write().state = Some(state.clone());
|
||||
Ok(state)
|
||||
}
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
/// The operation's current lifecycle state: "running", "finished",
|
||||
/// "failed", or "cancelled".
|
||||
///
|
||||
/// A point snapshot; unlike [`Job::wait`] it does not block, raise on a
|
||||
/// terminal failure state, or retry. States a newer server reports that
|
||||
/// this client version does not know pass through as-is.
|
||||
/// this client version does not know pass through as-is. Also refreshes
|
||||
/// the detail accessors.
|
||||
pub async fn status(&self) -> Result<String> {
|
||||
self.refresh_state().await
|
||||
}
|
||||
|
||||
/// The last lifecycle state this handle observed, without contacting the
|
||||
/// backend. `None` until the handle has.
|
||||
pub fn state(&self) -> Option<String> {
|
||||
self.cache_read().state.clone()
|
||||
}
|
||||
|
||||
/// The whole server-side record this handle last observed. The accessors
|
||||
/// below read individual fields out of it. `None` for an in-process job,
|
||||
/// which has no such record.
|
||||
pub fn description(&self) -> Option<JobDescription> {
|
||||
self.cache_read().description.clone()
|
||||
}
|
||||
|
||||
/// The job's type, as the server names it. `None` for an in-process job.
|
||||
pub fn job_type(&self) -> Option<String> {
|
||||
self.with_description(|description| description.job_type.clone())
|
||||
}
|
||||
|
||||
/// When the job was created, in milliseconds since the epoch. `None` for
|
||||
/// an in-process job.
|
||||
pub fn creation_ms(&self) -> Option<i64> {
|
||||
self.with_description(|description| description.creation_ms)
|
||||
}
|
||||
|
||||
/// The job-type-specific specification it was submitted with.
|
||||
pub fn spec(&self) -> Option<Value> {
|
||||
self.with_description(|description| description.spec.clone())
|
||||
.filter(|spec| !spec.is_null())
|
||||
}
|
||||
|
||||
/// The job-type-specific terminal result, as reported data rather than the
|
||||
/// typed model [`Job::wait`] returns. `None` until the job succeeds.
|
||||
pub fn result(&self) -> Option<Value> {
|
||||
self.with_description(|description| description.result.clone())
|
||||
.flatten()
|
||||
}
|
||||
|
||||
/// Why the job failed, when it failed and the server reports a reason.
|
||||
pub fn failure(&self) -> Option<JobFailure> {
|
||||
self.with_description(|description| description.failure.clone())
|
||||
.flatten()
|
||||
}
|
||||
|
||||
fn with_description<R>(&self, read: impl FnOnce(&JobDescription) -> R) -> Option<R> {
|
||||
self.cache_read().description.as_ref().map(read)
|
||||
}
|
||||
|
||||
/// This job's recorded lifecycle events.
|
||||
///
|
||||
/// Unlike the detail accessors, which report a terminal result only once
|
||||
/// the job reaches one, events are written as the job runs and outlive the
|
||||
/// workers that produced them. A distributed job records a
|
||||
/// `claim`/`claim_complete` pair per unit of work, each carrying
|
||||
/// `rows_processed`, so a job that never finishes still accounts for what
|
||||
/// it did. In-process operations keep no event history.
|
||||
pub async fn events(&self, request: JobEventsRequest) -> Result<Vec<RecordBatch>> {
|
||||
match &self.inner {
|
||||
JobInner::Handle { handle, .. } => handle.status().await,
|
||||
JobInner::Completed(_) => Ok("finished".to_string()),
|
||||
JobInner::Handle { handle, .. } => handle.events(request).await,
|
||||
// The operation finished before the handle existed, so there is no
|
||||
// id to query with even when a server ran it.
|
||||
JobInner::Completed(_) => Err(Error::NotSupported {
|
||||
message: "this operation completed before its handle was created, so it \
|
||||
carries no job id to query events with"
|
||||
.to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,8 +398,19 @@ where
|
||||
/// [`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::Completed(result) => Ok(result.clone()),
|
||||
JobInner::Handle { handle, decode } => {
|
||||
let settled = handle.wait().await;
|
||||
// Waiting already established a terminal state; record it so
|
||||
// the detail accessors do not need another round trip for it.
|
||||
if let Some(state) = terminal_state(&settled) {
|
||||
self.cache_write().state = Some(state.to_string());
|
||||
}
|
||||
(decode)(settled?)
|
||||
}
|
||||
JobInner::Completed(result) => {
|
||||
self.cache_write().state = Some("finished".to_string());
|
||||
Ok(result.clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -224,20 +443,36 @@ where
|
||||
U: Clone + Send + Sync + 'static,
|
||||
F: Fn(T) -> U + Send + Sync + 'static,
|
||||
{
|
||||
match self.inner {
|
||||
// The mapped handle tracks the same job, so it inherits what this one
|
||||
// has already learned about it.
|
||||
let Self { inner, cache } = self;
|
||||
match inner {
|
||||
JobInner::Handle { handle, decode } => Job {
|
||||
inner: JobInner::Handle {
|
||||
handle,
|
||||
decode: Arc::new(move |result| Ok(map((decode)(result)?))),
|
||||
},
|
||||
cache,
|
||||
},
|
||||
JobInner::Completed(result) => Job {
|
||||
inner: JobInner::Completed(map(result)),
|
||||
cache,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The lifecycle state a settled [`JobHandle::wait`] implies.
|
||||
fn terminal_state(settled: &Result<TerminalResult>) -> Option<&'static str> {
|
||||
match settled {
|
||||
Ok(_) => Some("finished"),
|
||||
Err(Error::JobFailed { .. }) => Some("failed"),
|
||||
Err(Error::JobCancelled { .. }) => Some("cancelled"),
|
||||
// Anything else is a transport failure, not a verdict on the job.
|
||||
Err(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 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)]
|
||||
|
||||
@@ -195,6 +195,7 @@ pub mod query;
|
||||
#[cfg(feature = "remote")]
|
||||
pub mod remote;
|
||||
pub mod rerankers;
|
||||
pub mod sql;
|
||||
pub mod table;
|
||||
#[cfg(test)]
|
||||
pub mod test_utils;
|
||||
|
||||
@@ -74,8 +74,15 @@ const EMBEDDING_FUNCTIONS_META_KEY: &str = "embedding_functions";
|
||||
const COLUMN_DEFINITIONS_META_KEY: &str = "lancedb::column_definitions";
|
||||
|
||||
/// Value of the definition's `kind` tag for the projected `select` form.
|
||||
/// Reserved for root-namespace sources; see [`NAMESPACED_SELECT_KIND`].
|
||||
pub const SELECT_KIND: &str = "select";
|
||||
|
||||
/// The `select` form over a namespaced source: its own kind, because released
|
||||
/// readers drop unknown fields and resolve a `select` source at the root, so
|
||||
/// this routes them to the [`MaterializedViewKind::Unrecognized`] refusal
|
||||
/// instead of a wrong-table refresh.
|
||||
pub const NAMESPACED_SELECT_KIND: &str = "namespaced_select";
|
||||
|
||||
/// Which view outputs each source column is projected to directly. A column
|
||||
/// may be projected more than once, so each carries every name the view gives
|
||||
/// it, in projection order.
|
||||
@@ -95,6 +102,10 @@ pub struct ViewProjection {
|
||||
pub struct MaterializedViewDefinition {
|
||||
/// Name of the source table, in the same database as the view.
|
||||
pub source_table: String,
|
||||
/// Namespace path holding the source table; empty is the root namespace.
|
||||
/// A definition written before namespaced sources reads as root.
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub source_namespace: Vec<String>,
|
||||
/// The projected output columns, in view schema order.
|
||||
pub projections: Vec<ViewProjection>,
|
||||
/// SQL predicate selecting the source rows the view holds.
|
||||
@@ -129,7 +140,12 @@ pub(crate) fn definition_to_metadata(definition: &MaterializedViewDefinition) ->
|
||||
let mut value = serde_json::to_value(definition).map_err(|e| Error::Runtime {
|
||||
message: format!("failed to serialize view definition: {e}"),
|
||||
})?;
|
||||
value["kind"] = serde_json::Value::String(SELECT_KIND.to_string());
|
||||
let kind = if definition.source_namespace.is_empty() {
|
||||
SELECT_KIND
|
||||
} else {
|
||||
NAMESPACED_SELECT_KIND
|
||||
};
|
||||
value["kind"] = serde_json::Value::String(kind.to_string());
|
||||
Ok(value.to_string())
|
||||
}
|
||||
|
||||
@@ -150,12 +166,21 @@ pub fn materialized_view_kind(
|
||||
.get("kind")
|
||||
.and_then(|k| k.as_str())
|
||||
.ok_or_else(|| unreadable(&"missing kind tag"))?;
|
||||
if kind != SELECT_KIND {
|
||||
if kind != SELECT_KIND && kind != NAMESPACED_SELECT_KIND {
|
||||
return Ok(Some(MaterializedViewKind::Unrecognized {
|
||||
kind: kind.to_string(),
|
||||
}));
|
||||
}
|
||||
let definition = serde_json::from_value(value).map_err(|e| unreadable(&e))?;
|
||||
let kind = kind.to_string();
|
||||
let definition: MaterializedViewDefinition =
|
||||
serde_json::from_value(value).map_err(|e| unreadable(&e))?;
|
||||
// No correct writer produces a kind that disagrees with its namespace.
|
||||
if (kind == SELECT_KIND) != definition.source_namespace.is_empty() {
|
||||
return Err(unreadable(&format!(
|
||||
"kind '{kind}' does not match its source namespace {:?}",
|
||||
definition.source_namespace
|
||||
)));
|
||||
}
|
||||
Ok(Some(MaterializedViewKind::Select(definition)))
|
||||
}
|
||||
|
||||
@@ -166,6 +191,7 @@ pub fn materialized_view_kind(
|
||||
pub(crate) fn plan(
|
||||
source_schema: SchemaRef,
|
||||
source_table: &str,
|
||||
source_namespace: &[String],
|
||||
projections: &[(String, String)],
|
||||
filter: Option<&str>,
|
||||
limit: Option<u64>,
|
||||
@@ -319,6 +345,7 @@ pub(crate) fn plan(
|
||||
|
||||
let definition = MaterializedViewDefinition {
|
||||
source_table: source_table.to_string(),
|
||||
source_namespace: source_namespace.to_vec(),
|
||||
projections: projections
|
||||
.into_iter()
|
||||
.map(|(output, expression)| ViewProjection { output, expression })
|
||||
@@ -602,7 +629,7 @@ pub struct PreparedDeclaration {
|
||||
definition: MaterializedViewDefinition,
|
||||
/// The source's own database: the only place
|
||||
/// [`PreparedDeclaration::create`] will put the view, because refresh
|
||||
/// resolves the recorded source name through the view's database.
|
||||
/// resolves the recorded source coordinate through the view's database.
|
||||
database: Arc<dyn Database>,
|
||||
}
|
||||
|
||||
@@ -622,10 +649,21 @@ impl PreparedDeclaration {
|
||||
|
||||
/// Create the view table and verify it, consuming the declaration.
|
||||
///
|
||||
/// The view goes in the source's own database, where refresh resolves the
|
||||
/// recorded source name. Stable row ids are requested at both levels and
|
||||
/// verified rather than trusted; nothing is rolled back on failure.
|
||||
/// The view goes at the root of the source's own database, where refresh
|
||||
/// resolves the recorded source coordinate. Stable row ids are requested
|
||||
/// at both levels and verified rather than trusted; nothing is rolled
|
||||
/// back on failure.
|
||||
pub async fn create(self, name: &str) -> Result<MaterializedView> {
|
||||
self.create_in(&[], name).await
|
||||
}
|
||||
|
||||
/// Create the view in `namespace_path`, empty for the root namespace.
|
||||
/// Otherwise [`PreparedDeclaration::create`].
|
||||
pub async fn create_in(
|
||||
self,
|
||||
namespace_path: &[String],
|
||||
name: &str,
|
||||
) -> Result<MaterializedView> {
|
||||
let empty: Vec<std::result::Result<arrow_array::RecordBatch, arrow_schema::ArrowError>> =
|
||||
vec![];
|
||||
// Minted here, not at preparation: a declaration can be cloned and
|
||||
@@ -640,6 +678,7 @@ impl PreparedDeclaration {
|
||||
let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
|
||||
Box::new(arrow_array::RecordBatchIterator::new(empty, schema));
|
||||
let mut request = CreateTableRequest::new(name.to_string(), Box::new(reader));
|
||||
request.namespace_path = namespace_path.to_vec();
|
||||
let write_params = request
|
||||
.write_options
|
||||
.lance_write_params
|
||||
@@ -680,8 +719,8 @@ impl PreparedDeclaration {
|
||||
|
||||
/// Validate a view declaration against its live source and hold what its
|
||||
/// creation needs. The declaration is canonicalized through the coordinate a
|
||||
/// refresh will resolve, so a handle that does not resolve back to itself is
|
||||
/// rejected, as is a namespaced source. Same creation-time checks as
|
||||
/// refresh will resolve -- name and namespace both -- so a handle that does
|
||||
/// not resolve back to itself is rejected. Same creation-time checks as
|
||||
/// [`Connection::create_materialized_view`].
|
||||
///
|
||||
/// ```no_run
|
||||
@@ -710,17 +749,9 @@ pub async fn prepare_declaration(
|
||||
message: "materialized views are supported only on local databases".into(),
|
||||
});
|
||||
};
|
||||
// The definition records the source by bare name; any other source
|
||||
// form would be recorded as a name its refresh cannot resolve.
|
||||
if !source.namespace().is_empty() {
|
||||
return Err(Error::NotSupported {
|
||||
message: format!(
|
||||
"a namespaced source cannot be recorded in a view definition; \
|
||||
'{}' must be a root-namespace table",
|
||||
source.name()
|
||||
),
|
||||
});
|
||||
}
|
||||
// Refresh resolves the source at exactly this coordinate, so the
|
||||
// definition records the namespace alongside the name.
|
||||
let source_namespace = source.namespace().to_vec();
|
||||
let database = source
|
||||
.database_opt()
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
@@ -734,7 +765,7 @@ pub async fn prepare_declaration(
|
||||
let resolved = database
|
||||
.open_table(OpenTableRequest {
|
||||
name: source.name().to_string(),
|
||||
namespace_path: vec![],
|
||||
namespace_path: source_namespace.clone(),
|
||||
index_cache_size: None,
|
||||
lance_read_params: None,
|
||||
location: None,
|
||||
@@ -780,6 +811,7 @@ pub async fn prepare_declaration(
|
||||
let (definition, mut fields, lineage) = plan(
|
||||
source_schema.clone(),
|
||||
resolved.name(),
|
||||
&source_namespace,
|
||||
projections,
|
||||
filter,
|
||||
limit,
|
||||
@@ -839,7 +871,9 @@ fn ensure_local(connection: &Connection) -> Result<()> {
|
||||
pub struct CreateMaterializedViewBuilder {
|
||||
connection: Connection,
|
||||
name: String,
|
||||
namespace: Vec<String>,
|
||||
source: String,
|
||||
source_namespace: Vec<String>,
|
||||
projections: Vec<(String, String)>,
|
||||
filter: Option<String>,
|
||||
limit: Option<u64>,
|
||||
@@ -850,13 +884,28 @@ impl CreateMaterializedViewBuilder {
|
||||
Self {
|
||||
connection,
|
||||
name,
|
||||
namespace: Vec::new(),
|
||||
source,
|
||||
source_namespace: Vec::new(),
|
||||
projections: Vec::new(),
|
||||
filter: None,
|
||||
limit: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// The namespace to create the view in. Defaults to the root namespace.
|
||||
pub fn namespace(mut self, namespace_path: Vec<String>) -> Self {
|
||||
self.namespace = namespace_path;
|
||||
self
|
||||
}
|
||||
|
||||
/// The namespace holding the source table; recorded in the definition
|
||||
/// for refresh to resolve. Defaults to the root namespace.
|
||||
pub fn source_namespace(mut self, namespace_path: Vec<String>) -> Self {
|
||||
self.source_namespace = namespace_path;
|
||||
self
|
||||
}
|
||||
|
||||
/// The view's columns, as `(name, SQL expression)` pairs. Not calling
|
||||
/// this selects every source column, expanded at creation time.
|
||||
pub fn select(
|
||||
@@ -887,7 +936,12 @@ impl CreateMaterializedViewBuilder {
|
||||
/// provenance across compaction, and cannot be enabled later.
|
||||
pub async fn execute(self) -> Result<MaterializedView> {
|
||||
ensure_local(&self.connection)?;
|
||||
let source = self.connection.open_table(&self.source).execute().await?;
|
||||
let source = self
|
||||
.connection
|
||||
.open_table(&self.source)
|
||||
.namespace(self.source_namespace.clone())
|
||||
.execute()
|
||||
.await?;
|
||||
let prepared = prepare_declaration(
|
||||
&source,
|
||||
&self.projections,
|
||||
@@ -895,7 +949,7 @@ impl CreateMaterializedViewBuilder {
|
||||
self.limit,
|
||||
)
|
||||
.await?;
|
||||
prepared.create(&self.name).await
|
||||
prepared.create_in(&self.namespace, &self.name).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1152,6 +1206,7 @@ mod tests {
|
||||
view.definition(),
|
||||
&MaterializedViewDefinition {
|
||||
source_table: "people".into(),
|
||||
source_namespace: Vec::new(),
|
||||
projections: vec![
|
||||
ViewProjection {
|
||||
output: "name".into(),
|
||||
@@ -2083,33 +2138,138 @@ mod tests {
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("custom_loc"), "{err}");
|
||||
}
|
||||
|
||||
// A namespaced source cannot be recorded in the definition: the
|
||||
// bare name refresh resolves would reach a different table or none.
|
||||
let namespaced = crate::table::NativeTable::create(
|
||||
"memory://ns_src",
|
||||
"ns_src",
|
||||
vec!["ns".to_string()],
|
||||
Box::new(arrow_array::RecordBatchIterator::new(
|
||||
vec![],
|
||||
std::sync::Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new(
|
||||
"id",
|
||||
arrow_schema::DataType::Int32,
|
||||
true,
|
||||
)])),
|
||||
)) as Box<dyn arrow_array::RecordBatchReader + Send>,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
std::collections::HashSet::new(),
|
||||
)
|
||||
/// A view declared over a namespaced source records that namespace, and
|
||||
/// refresh resolves the source through it -- the coordinate round-trips.
|
||||
#[tokio::test]
|
||||
async fn a_namespaced_source_round_trips_through_refresh() {
|
||||
use lance_namespace::models::CreateNamespaceRequest;
|
||||
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let mut properties = std::collections::HashMap::new();
|
||||
properties.insert("root".to_string(), tmp.path().to_str().unwrap().to_string());
|
||||
let conn = crate::connect_namespace("dir", properties)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
conn.create_namespace(CreateNamespaceRequest {
|
||||
id: Some(vec!["ns".into()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let namespaced = Table::new(std::sync::Arc::new(namespaced), conn.database().clone());
|
||||
let err = prepare_declaration(&namespaced, &[], None, None)
|
||||
|
||||
let batch = record_batch!(
|
||||
("name", Utf8, ["ada", "grace", "alan"]),
|
||||
("age", Int32, [36, 85, 41])
|
||||
)
|
||||
.unwrap();
|
||||
conn.create_table("people", batch)
|
||||
.namespace(vec!["ns".to_string()])
|
||||
.write_options(stable_row_ids())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.to_string().contains("namespaced source"), "{err}");
|
||||
.unwrap();
|
||||
|
||||
// A decoy of the same name at the root: resolving the source at the
|
||||
// wrong namespace materializes one row here instead of three.
|
||||
let decoy = record_batch!(("name", Utf8, ["mallory"]), ("age", Int32, [42])).unwrap();
|
||||
conn.create_table("people", decoy)
|
||||
.write_options(stable_row_ids())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let view = conn
|
||||
.create_materialized_view("adults", "people")
|
||||
.namespace(vec!["ns".to_string()])
|
||||
.source_namespace(vec!["ns".to_string()])
|
||||
.select([("name", "name")])
|
||||
.only_if("age >= 18")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(view.definition().source_table, "people");
|
||||
assert_eq!(view.definition().source_namespace, vec!["ns".to_string()]);
|
||||
assert_eq!(view.table().namespace(), &["ns"]);
|
||||
|
||||
// Refresh resolves the source at the recorded namespace, not at root.
|
||||
let result = view.refresh().execute().await.unwrap();
|
||||
assert_eq!(result.rows_written, 3);
|
||||
}
|
||||
|
||||
/// A definition stored before namespaced sources existed carries no
|
||||
/// namespace key and must read as the root namespace.
|
||||
#[test]
|
||||
fn a_definition_without_a_namespace_reads_as_root() {
|
||||
let stored =
|
||||
r#"{"source_table":"people","projections":[{"output":"name","expression":"name"}]}"#;
|
||||
let definition: MaterializedViewDefinition = serde_json::from_str(stored).unwrap();
|
||||
assert!(definition.source_namespace.is_empty());
|
||||
}
|
||||
|
||||
fn definition(source_namespace: Vec<String>) -> MaterializedViewDefinition {
|
||||
MaterializedViewDefinition {
|
||||
source_table: "people".to_string(),
|
||||
source_namespace,
|
||||
projections: vec![ViewProjection {
|
||||
output: "name".to_string(),
|
||||
expression: "name".to_string(),
|
||||
}],
|
||||
filter: None,
|
||||
limit: None,
|
||||
inputs: vec!["name".to_string()],
|
||||
}
|
||||
}
|
||||
|
||||
/// A root definition keeps the pre-namespace `select` form byte-stably;
|
||||
/// a namespaced one moves off `select`, which sends pre-namespace readers
|
||||
/// to the `Unrecognized` refusal instead of a root resolve.
|
||||
#[test]
|
||||
fn a_namespaced_definition_is_refused_by_the_pre_namespace_reader() {
|
||||
let root = definition_to_metadata(&definition(Vec::new())).unwrap();
|
||||
let root: serde_json::Value = serde_json::from_str(&root).unwrap();
|
||||
assert_eq!(root["kind"], "select");
|
||||
assert!(
|
||||
root.get("source_namespace").is_none(),
|
||||
"a root definition must not grow new keys: {root}"
|
||||
);
|
||||
|
||||
let stored = definition_to_metadata(&definition(vec!["ns".to_string()])).unwrap();
|
||||
let value: serde_json::Value = serde_json::from_str(&stored).unwrap();
|
||||
// The pre-namespace discriminator is `kind == "select"`; anything
|
||||
// else lands in its Unrecognized refusal rather than in a root open.
|
||||
assert_eq!(value["kind"], "namespaced_select");
|
||||
|
||||
// The current reader round-trips the coordinate.
|
||||
let metadata = HashMap::from([(DEFINITION_META_KEY.to_string(), stored)]);
|
||||
match materialized_view_kind(&metadata).unwrap() {
|
||||
Some(MaterializedViewKind::Select(read)) => {
|
||||
assert_eq!(read.source_namespace, vec!["ns".to_string()])
|
||||
}
|
||||
other => panic!("expected the namespaced select form, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// A kind that disagrees with its namespace is an error, not a view:
|
||||
/// under `select` it is the shape old readers would resolve at the root.
|
||||
#[test]
|
||||
fn a_kind_namespace_mismatch_is_refused() {
|
||||
for (kind, namespace) in [
|
||||
(SELECT_KIND, vec!["ns".to_string()]),
|
||||
(NAMESPACED_SELECT_KIND, Vec::new()),
|
||||
] {
|
||||
let mut value = serde_json::to_value(definition(namespace)).unwrap();
|
||||
value["kind"] = serde_json::Value::String(kind.to_string());
|
||||
let metadata = HashMap::from([(DEFINITION_META_KEY.to_string(), value.to_string())]);
|
||||
let err = materialized_view_kind(&metadata).unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("does not match its source namespace"),
|
||||
"kind '{kind}': {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,6 +170,7 @@ pub(crate) async fn execute_refresh(
|
||||
let (replanned, mut planned_fields, _renames) = super::plan(
|
||||
source_schema,
|
||||
&definition.source_table,
|
||||
&definition.source_namespace,
|
||||
&projections,
|
||||
definition.filter.as_deref(),
|
||||
definition.limit,
|
||||
@@ -590,7 +591,7 @@ async fn open_source(view: &Table, definition: &MaterializedViewDefinition) -> R
|
||||
let source = database
|
||||
.open_table(OpenTableRequest {
|
||||
name: definition.source_table.clone(),
|
||||
namespace_path: Vec::new(),
|
||||
namespace_path: definition.source_namespace.clone(),
|
||||
index_cache_size: None,
|
||||
lance_read_params: None,
|
||||
location: None,
|
||||
@@ -2919,6 +2920,7 @@ mod tests {
|
||||
|
||||
let replacement = crate::materialized_view::MaterializedViewDefinition {
|
||||
source_table: "src".into(),
|
||||
source_namespace: Vec::new(),
|
||||
projections: vec![
|
||||
crate::materialized_view::ViewProjection {
|
||||
output: "x".into(),
|
||||
@@ -2958,6 +2960,7 @@ mod tests {
|
||||
|
||||
let narrower = crate::materialized_view::MaterializedViewDefinition {
|
||||
source_table: "src".into(),
|
||||
source_namespace: Vec::new(),
|
||||
projections: vec![crate::materialized_view::ViewProjection {
|
||||
output: "x".into(),
|
||||
expression: "x".into(),
|
||||
|
||||
+835
-5
@@ -1,21 +1,37 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::{future::Future, time::Duration};
|
||||
|
||||
use arrow::compute::concat_batches;
|
||||
use arrow_array::{Array, Float16Array, Float32Array, Float64Array, RecordBatch, make_array};
|
||||
use arrow_array::{
|
||||
Array, Float16Array, Float32Array, Float64Array, RecordBatch, UInt64Array,
|
||||
cast::AsArray,
|
||||
make_array,
|
||||
types::{Int64Type, UInt64Type},
|
||||
};
|
||||
use arrow_schema::{DataType, SchemaRef};
|
||||
use datafusion_common::{DataFusionError, Result as DataFusionResult};
|
||||
use datafusion_execution::TaskContext;
|
||||
use datafusion_expr::{Expr, col, lit};
|
||||
use datafusion_physical_plan::ExecutionPlan;
|
||||
use futures::{FutureExt, TryFutureExt, TryStreamExt, stream, try_join};
|
||||
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
|
||||
use datafusion_physical_plan::{
|
||||
DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
|
||||
coalesce_partitions::CoalescePartitionsExec,
|
||||
execution_plan::{Boundedness, EmissionType},
|
||||
limit::GlobalLimitExec,
|
||||
stream::RecordBatchStreamAdapter,
|
||||
};
|
||||
use futures::{FutureExt, StreamExt, TryFutureExt, TryStreamExt, stream, try_join};
|
||||
use half::f16;
|
||||
/// Re-export Lance ColumnOrdering type for use in query ordering
|
||||
pub use lance::dataset::scanner::ColumnOrdering;
|
||||
use lance::dataset::{ROW_ID, scanner::DatasetRecordBatchStream};
|
||||
use lance_arrow::RecordBatchExt;
|
||||
use lance_datafusion::exec::execute_plan;
|
||||
use lance_datafusion::exec::{execute_plan, format_plan as format_analyzed_plan};
|
||||
use lance_index::scalar::FullTextSearchQuery;
|
||||
use lance_index::scalar::inverted::SCORE_COL;
|
||||
use lance_index::vector::DIST_COL;
|
||||
@@ -825,6 +841,14 @@ pub struct QueryRequest {
|
||||
/// Offset of the query.
|
||||
pub offset: Option<usize>,
|
||||
|
||||
/// Dataset offsets whose occurrence multiplicity must be restored after
|
||||
/// executing the physical lookup represented by this request.
|
||||
///
|
||||
/// This is client-side execution metadata used when a [`TakeQuery`] is
|
||||
/// converted into a request. It is not sent to remote services.
|
||||
#[doc(hidden)]
|
||||
pub take_offsets: Option<Vec<u64>>,
|
||||
|
||||
/// Apply filter to the returned rows.
|
||||
pub filter: Option<QueryFilter>,
|
||||
|
||||
@@ -893,6 +917,7 @@ impl Default for QueryRequest {
|
||||
Self {
|
||||
limit: None,
|
||||
offset: None,
|
||||
take_offsets: None,
|
||||
filter: None,
|
||||
filter_error: None,
|
||||
full_text_search: None,
|
||||
@@ -1529,6 +1554,302 @@ impl HasQuery for VectorQuery {
|
||||
}
|
||||
}
|
||||
|
||||
fn take_occurrences(offsets: &[u64]) -> HashMap<u64, usize> {
|
||||
let mut occurrences = HashMap::with_capacity(offsets.len());
|
||||
for offset in offsets {
|
||||
*occurrences.entry(*offset).or_insert(0) += 1;
|
||||
}
|
||||
occurrences
|
||||
}
|
||||
|
||||
fn restore_take_batch_with_occurrences(
|
||||
batch: RecordBatch,
|
||||
offsets: &[u64],
|
||||
occurrences: &HashMap<u64, usize>,
|
||||
ordering_column: &str,
|
||||
drop_ordering_column: bool,
|
||||
preserve_order: bool,
|
||||
) -> Result<RecordBatch> {
|
||||
let actual_offsets = batch
|
||||
.column_by_name(ordering_column)
|
||||
.ok_or_else(|| Error::Schema {
|
||||
message: format!(
|
||||
"take query result did not include ordering column '{ordering_column}'"
|
||||
),
|
||||
})?;
|
||||
let actual_offsets = match actual_offsets.data_type() {
|
||||
DataType::UInt64 => actual_offsets
|
||||
.as_primitive::<UInt64Type>()
|
||||
.values()
|
||||
.to_vec(),
|
||||
DataType::Int64 => actual_offsets
|
||||
.as_primitive::<Int64Type>()
|
||||
.values()
|
||||
.iter()
|
||||
.map(|offset| {
|
||||
u64::try_from(*offset).map_err(|_| Error::Schema {
|
||||
message: format!(
|
||||
"take query ordering column '{ordering_column}' contained a negative offset"
|
||||
),
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?,
|
||||
data_type => {
|
||||
return Err(Error::Schema {
|
||||
message: format!(
|
||||
"take query ordering column '{ordering_column}' had unsupported type {data_type}"
|
||||
),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
let mut desired_order = Vec::with_capacity(offsets.len());
|
||||
if preserve_order {
|
||||
let ordering = actual_offsets
|
||||
.iter()
|
||||
.copied()
|
||||
.enumerate()
|
||||
.map(|(index, offset)| (offset, index as u64))
|
||||
.collect::<HashMap<_, _>>();
|
||||
// Missing offsets retain the filter-based behavior of returning no row.
|
||||
desired_order.extend(
|
||||
offsets
|
||||
.iter()
|
||||
.filter_map(|offset| ordering.get(offset).copied()),
|
||||
);
|
||||
} else {
|
||||
// Public take queries do not guarantee output order. Preserve the lookup's
|
||||
// existing order and only restore the multiplicity of each matching row.
|
||||
for (index, offset) in actual_offsets.iter().enumerate() {
|
||||
if let Some(count) = occurrences.get(offset) {
|
||||
desired_order.extend(std::iter::repeat_n(index as u64, *count));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut ordered_batch = if desired_order.len() == batch.num_rows()
|
||||
&& desired_order
|
||||
.iter()
|
||||
.enumerate()
|
||||
.all(|(index, desired)| *desired == index as u64)
|
||||
{
|
||||
batch
|
||||
} else {
|
||||
arrow_select::take::take_record_batch(&batch, &UInt64Array::from(desired_order))?
|
||||
};
|
||||
|
||||
if drop_ordering_column {
|
||||
ordered_batch = ordered_batch.drop_column(ordering_column)?;
|
||||
}
|
||||
|
||||
Ok(ordered_batch)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn restore_take_batch(
|
||||
batch: RecordBatch,
|
||||
offsets: &[u64],
|
||||
ordering_column: &str,
|
||||
drop_ordering_column: bool,
|
||||
preserve_order: bool,
|
||||
) -> Result<RecordBatch> {
|
||||
restore_take_batch_with_occurrences(
|
||||
batch,
|
||||
offsets,
|
||||
&take_occurrences(offsets),
|
||||
ordering_column,
|
||||
drop_ordering_column,
|
||||
preserve_order,
|
||||
)
|
||||
}
|
||||
|
||||
/// Restores the logical offset occurrence sequence above the physical lookup plan.
|
||||
///
|
||||
/// The lookup plan returns each matching row at most once. For ordinary unordered
|
||||
/// takes this operator expands each input batch incrementally and preserves the
|
||||
/// lookup's partitioning. The explicitly ordered reader path collects one coalesced
|
||||
/// input before restoring requested order. Pagination must remain above this operator
|
||||
/// so it applies to occurrences.
|
||||
#[derive(Debug)]
|
||||
struct TakeRestoreExec {
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
offsets: Vec<u64>,
|
||||
occurrences: Arc<HashMap<u64, usize>>,
|
||||
ordering_column: String,
|
||||
drop_ordering_column: bool,
|
||||
preserve_order: bool,
|
||||
schema: SchemaRef,
|
||||
properties: Arc<PlanProperties>,
|
||||
}
|
||||
|
||||
impl TakeRestoreExec {
|
||||
fn try_new(
|
||||
input: Arc<dyn ExecutionPlan>,
|
||||
offsets: Vec<u64>,
|
||||
ordering_column: String,
|
||||
drop_ordering_column: bool,
|
||||
preserve_order: bool,
|
||||
) -> Result<Self> {
|
||||
let schema = if drop_ordering_column {
|
||||
RecordBatch::new_empty(input.schema())
|
||||
.drop_column(&ordering_column)?
|
||||
.schema()
|
||||
} else {
|
||||
input.schema()
|
||||
};
|
||||
let partition_count = if preserve_order {
|
||||
1
|
||||
} else {
|
||||
input.output_partitioning().partition_count()
|
||||
};
|
||||
let emission_type = if preserve_order {
|
||||
EmissionType::Final
|
||||
} else {
|
||||
EmissionType::Incremental
|
||||
};
|
||||
let properties = Arc::new(PlanProperties::new(
|
||||
EquivalenceProperties::new(schema.clone()),
|
||||
Partitioning::UnknownPartitioning(partition_count),
|
||||
emission_type,
|
||||
Boundedness::Bounded,
|
||||
));
|
||||
|
||||
Ok(Self {
|
||||
input,
|
||||
occurrences: Arc::new(take_occurrences(&offsets)),
|
||||
offsets,
|
||||
ordering_column,
|
||||
drop_ordering_column,
|
||||
preserve_order,
|
||||
schema,
|
||||
properties,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl DisplayAs for TakeRestoreExec {
|
||||
fn fmt_as(
|
||||
&self,
|
||||
_display_type: DisplayFormatType,
|
||||
formatter: &mut std::fmt::Formatter<'_>,
|
||||
) -> std::fmt::Result {
|
||||
write!(
|
||||
formatter,
|
||||
"TakeRestoreExec: occurrences={}",
|
||||
self.offsets.len()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionPlan for TakeRestoreExec {
|
||||
fn name(&self) -> &str {
|
||||
"TakeRestoreExec"
|
||||
}
|
||||
|
||||
fn properties(&self) -> &Arc<PlanProperties> {
|
||||
&self.properties
|
||||
}
|
||||
|
||||
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
||||
vec![&self.input]
|
||||
}
|
||||
|
||||
fn maintains_input_order(&self) -> Vec<bool> {
|
||||
vec![!self.preserve_order]
|
||||
}
|
||||
|
||||
fn benefits_from_input_partitioning(&self) -> Vec<bool> {
|
||||
vec![false]
|
||||
}
|
||||
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
children: Vec<Arc<dyn ExecutionPlan>>,
|
||||
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
|
||||
if children.len() != 1 {
|
||||
return Err(DataFusionError::Internal(format!(
|
||||
"TakeRestoreExec expected one child, got {}",
|
||||
children.len()
|
||||
)));
|
||||
}
|
||||
let child = children.into_iter().next().unwrap();
|
||||
let plan = Self::try_new(
|
||||
child,
|
||||
self.offsets.clone(),
|
||||
self.ordering_column.clone(),
|
||||
self.drop_ordering_column,
|
||||
self.preserve_order,
|
||||
)
|
||||
.map_err(|error| DataFusionError::External(Box::new(error)))?;
|
||||
Ok(Arc::new(plan))
|
||||
}
|
||||
|
||||
fn execute(
|
||||
&self,
|
||||
partition: usize,
|
||||
context: Arc<TaskContext>,
|
||||
) -> DataFusionResult<datafusion_physical_plan::SendableRecordBatchStream> {
|
||||
let partition_count = self.input.output_partitioning().partition_count();
|
||||
if partition >= partition_count || (self.preserve_order && partition != 0) {
|
||||
return Err(DataFusionError::Internal(format!(
|
||||
"TakeRestoreExec cannot execute partition {partition}; input has {partition_count} partitions"
|
||||
)));
|
||||
}
|
||||
|
||||
let input = self.input.execute(partition, context)?;
|
||||
let output_schema = self.schema.clone();
|
||||
let offsets = self.offsets.clone();
|
||||
let occurrences = self.occurrences.clone();
|
||||
let ordering_column = self.ordering_column.clone();
|
||||
let drop_ordering_column = self.drop_ordering_column;
|
||||
let preserve_order = self.preserve_order;
|
||||
let stream: Pin<Box<dyn futures::Stream<Item = DataFusionResult<RecordBatch>> + Send>> =
|
||||
if preserve_order {
|
||||
let input_schema = input.schema();
|
||||
Box::pin(stream::once(async move {
|
||||
let batches = input.try_collect::<Vec<_>>().await?;
|
||||
let batch = if batches.is_empty() {
|
||||
RecordBatch::new_empty(input_schema.clone())
|
||||
} else {
|
||||
concat_batches(&input_schema, &batches)?
|
||||
};
|
||||
restore_take_batch_with_occurrences(
|
||||
batch,
|
||||
&offsets,
|
||||
&occurrences,
|
||||
&ordering_column,
|
||||
drop_ordering_column,
|
||||
true,
|
||||
)
|
||||
.map_err(|error| DataFusionError::External(Box::new(error)))
|
||||
}))
|
||||
} else {
|
||||
Box::pin(input.map(move |batch| {
|
||||
batch.and_then(|batch| {
|
||||
restore_take_batch_with_occurrences(
|
||||
batch,
|
||||
&offsets,
|
||||
&occurrences,
|
||||
&ordering_column,
|
||||
drop_ordering_column,
|
||||
false,
|
||||
)
|
||||
.map_err(|error| DataFusionError::External(Box::new(error)))
|
||||
})
|
||||
}))
|
||||
};
|
||||
|
||||
Ok(Box::pin(RecordBatchStreamAdapter::new(
|
||||
output_schema,
|
||||
stream,
|
||||
)))
|
||||
}
|
||||
|
||||
fn supports_limit_pushdown(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// A builder for LanceDB take queries.
|
||||
///
|
||||
/// See [`crate::Table::query`] for more details on queries
|
||||
@@ -1545,6 +1866,8 @@ impl HasQuery for VectorQuery {
|
||||
pub struct TakeQuery {
|
||||
parent: Arc<dyn BaseTable>,
|
||||
request: QueryRequest,
|
||||
offsets: Option<Vec<u64>>,
|
||||
preserve_order: bool,
|
||||
}
|
||||
|
||||
impl TakeQuery {
|
||||
@@ -1552,15 +1875,24 @@ impl TakeQuery {
|
||||
///
|
||||
/// See [`crate::Table::take_offsets`] for more details.
|
||||
pub fn from_offsets(parent: Arc<dyn BaseTable>, offsets: Vec<u64>) -> Self {
|
||||
let in_list: Vec<Expr> = offsets.iter().map(|o| lit(*o)).collect();
|
||||
let mut seen = HashSet::with_capacity(offsets.len());
|
||||
let in_list: Vec<Expr> = offsets
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|offset| seen.insert(*offset))
|
||||
.map(lit)
|
||||
.collect();
|
||||
Self {
|
||||
parent,
|
||||
request: QueryRequest {
|
||||
filter: Some(QueryFilter::Datafusion(
|
||||
col("_rowoffset").in_list(in_list, false),
|
||||
)),
|
||||
take_offsets: Some(offsets.clone()),
|
||||
..Default::default()
|
||||
},
|
||||
offsets: Some(offsets),
|
||||
preserve_order: false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1575,9 +1907,181 @@ impl TakeQuery {
|
||||
filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))),
|
||||
..Default::default()
|
||||
},
|
||||
offsets: None,
|
||||
preserve_order: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Preserve the requested offset order when restoring duplicate occurrences.
|
||||
///
|
||||
/// This is reserved for readers whose API explicitly guarantees ordering.
|
||||
pub(crate) fn preserve_order(mut self) -> Self {
|
||||
debug_assert!(self.offsets.is_some());
|
||||
self.preserve_order = true;
|
||||
self
|
||||
}
|
||||
|
||||
async fn request_with_row_offset(
|
||||
parent: &dyn BaseTable,
|
||||
request: &QueryRequest,
|
||||
) -> Result<(QueryRequest, String, bool)> {
|
||||
const ROW_OFFSET: &str = "_rowoffset";
|
||||
const INTERNAL_ROW_OFFSET: &str = "__lancedb_take_row_offset";
|
||||
|
||||
let mut request = request.clone();
|
||||
// The physical lookup must not recursively restore occurrences. The
|
||||
// wrapper above this request owns that logical operation.
|
||||
request.take_offsets = None;
|
||||
let (ordering_column, drop_ordering_column) = match &mut request.select {
|
||||
Select::All => {
|
||||
let mut columns = parent
|
||||
.schema()
|
||||
.await?
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|field| field.name().clone())
|
||||
.collect::<Vec<_>>();
|
||||
columns.push(ROW_OFFSET.to_string());
|
||||
request.select = Select::Columns(columns);
|
||||
(ROW_OFFSET.to_string(), true)
|
||||
}
|
||||
Select::Columns(columns) => {
|
||||
if columns.iter().any(|column| column == ROW_OFFSET) {
|
||||
(ROW_OFFSET.to_string(), false)
|
||||
} else {
|
||||
columns.push(ROW_OFFSET.to_string());
|
||||
(ROW_OFFSET.to_string(), true)
|
||||
}
|
||||
}
|
||||
Select::Dynamic(columns) => {
|
||||
let mut ordering_column = INTERNAL_ROW_OFFSET.to_string();
|
||||
while columns.iter().any(|(name, _)| name == &ordering_column) {
|
||||
ordering_column.push('_');
|
||||
}
|
||||
columns.push((ordering_column.clone(), ROW_OFFSET.to_string()));
|
||||
(ordering_column, true)
|
||||
}
|
||||
Select::Expr(columns) => {
|
||||
let mut ordering_column = INTERNAL_ROW_OFFSET.to_string();
|
||||
while columns.iter().any(|(name, _)| name == &ordering_column) {
|
||||
ordering_column.push('_');
|
||||
}
|
||||
columns.push((ordering_column.clone(), col(ROW_OFFSET)));
|
||||
(ordering_column, true)
|
||||
}
|
||||
};
|
||||
|
||||
Ok((request, ordering_column, drop_ordering_column))
|
||||
}
|
||||
|
||||
async fn prepare_offsets_lookup(
|
||||
parent: &dyn BaseTable,
|
||||
request: &QueryRequest,
|
||||
) -> Result<(QueryRequest, String, bool, usize, Option<usize>)> {
|
||||
let (mut request, ordering_column, drop_ordering_column) =
|
||||
Self::request_with_row_offset(parent, request).await?;
|
||||
// The lookup operates on distinct physical rows. Pagination is a logical
|
||||
// operation over occurrences and must be applied only after restoration.
|
||||
let output_offset = request.offset.take().unwrap_or_default();
|
||||
let output_limit = request.limit.take();
|
||||
|
||||
Ok((
|
||||
request,
|
||||
ordering_column,
|
||||
drop_ordering_column,
|
||||
output_offset,
|
||||
output_limit,
|
||||
))
|
||||
}
|
||||
|
||||
fn wrap_offsets_plan(
|
||||
lookup: Arc<dyn ExecutionPlan>,
|
||||
offsets: &[u64],
|
||||
ordering_column: String,
|
||||
drop_ordering_column: bool,
|
||||
output_offset: usize,
|
||||
output_limit: Option<usize>,
|
||||
preserve_order: bool,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
let lookup = if preserve_order {
|
||||
Arc::new(CoalescePartitionsExec::new(lookup)) as Arc<dyn ExecutionPlan>
|
||||
} else {
|
||||
lookup
|
||||
};
|
||||
let restored: Arc<dyn ExecutionPlan> = Arc::new(TakeRestoreExec::try_new(
|
||||
lookup,
|
||||
offsets.to_vec(),
|
||||
ordering_column,
|
||||
drop_ordering_column,
|
||||
preserve_order,
|
||||
)?);
|
||||
|
||||
if output_offset > 0 || output_limit.is_some() {
|
||||
Ok(Arc::new(GlobalLimitExec::new(
|
||||
restored,
|
||||
output_offset,
|
||||
output_limit,
|
||||
)))
|
||||
} else {
|
||||
Ok(restored)
|
||||
}
|
||||
}
|
||||
|
||||
fn wrap_offsets_explanation(
|
||||
lookup: &str,
|
||||
occurrence_count: usize,
|
||||
output_offset: usize,
|
||||
output_limit: Option<usize>,
|
||||
preserve_order: bool,
|
||||
) -> String {
|
||||
fn indent(plan: &str, spaces: usize) -> String {
|
||||
let indentation = " ".repeat(spaces);
|
||||
plan.lines()
|
||||
.map(|line| format!("{indentation}{line}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
|
||||
let restored = if preserve_order {
|
||||
format!(
|
||||
"TakeRestoreExec: occurrences={occurrence_count}\n CoalescePartitionsExec\n{}",
|
||||
indent(lookup, 4)
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"TakeRestoreExec: occurrences={occurrence_count}\n{}",
|
||||
indent(lookup, 2)
|
||||
)
|
||||
};
|
||||
|
||||
if output_offset > 0 || output_limit.is_some() {
|
||||
let fetch = output_limit
|
||||
.map(|limit| limit.to_string())
|
||||
.unwrap_or_else(|| "None".to_string());
|
||||
format!(
|
||||
"GlobalLimitExec: skip={output_offset}, fetch={fetch}\n{}",
|
||||
indent(&restored, 2)
|
||||
)
|
||||
} else {
|
||||
restored
|
||||
}
|
||||
}
|
||||
|
||||
async fn create_offsets_plan(
|
||||
&self,
|
||||
offsets: &[u64],
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
create_take_offsets_plan(
|
||||
self.parent.as_ref(),
|
||||
&self.request,
|
||||
offsets,
|
||||
options,
|
||||
self.preserve_order,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Convert the `TakeQuery` into a `QueryRequest`.
|
||||
pub fn into_request(self) -> QueryRequest {
|
||||
self.request
|
||||
@@ -1622,6 +2126,63 @@ impl TakeQuery {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_take_offsets_plan(
|
||||
parent: &dyn BaseTable,
|
||||
request: &QueryRequest,
|
||||
offsets: &[u64],
|
||||
options: QueryExecutionOptions,
|
||||
preserve_order: bool,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
let (request, ordering_column, drop_ordering_column, output_offset, output_limit) =
|
||||
TakeQuery::prepare_offsets_lookup(parent, request).await?;
|
||||
let lookup_options = if preserve_order {
|
||||
options.without_output_batch_length_limit()
|
||||
} else {
|
||||
options
|
||||
};
|
||||
let lookup = parent
|
||||
.create_plan(&AnyQuery::Query(request), lookup_options)
|
||||
.await?;
|
||||
|
||||
TakeQuery::wrap_offsets_plan(
|
||||
lookup,
|
||||
offsets,
|
||||
ordering_column,
|
||||
drop_ordering_column,
|
||||
output_offset,
|
||||
output_limit,
|
||||
preserve_order,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn explain_take_offsets_plan(
|
||||
parent: &dyn BaseTable,
|
||||
request: &QueryRequest,
|
||||
offsets: &[u64],
|
||||
verbose: bool,
|
||||
) -> Result<String> {
|
||||
let (request, _, _, output_offset, output_limit) =
|
||||
TakeQuery::prepare_offsets_lookup(parent, request).await?;
|
||||
let lookup = parent
|
||||
.explain_plan(&AnyQuery::Query(request), verbose)
|
||||
.await?;
|
||||
Ok(TakeQuery::wrap_offsets_explanation(
|
||||
&lookup,
|
||||
offsets.len(),
|
||||
output_offset,
|
||||
output_limit,
|
||||
false,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) async fn prepare_take_offsets_request(
|
||||
parent: &dyn BaseTable,
|
||||
request: &QueryRequest,
|
||||
) -> Result<QueryRequest> {
|
||||
let (request, _, _, _, _) = TakeQuery::prepare_offsets_lookup(parent, request).await?;
|
||||
Ok(request)
|
||||
}
|
||||
|
||||
impl HasQuery for TakeQuery {
|
||||
fn mut_query(&mut self) -> &mut QueryRequest {
|
||||
&mut self.request
|
||||
@@ -1630,6 +2191,10 @@ impl HasQuery for TakeQuery {
|
||||
|
||||
impl ExecutableQuery for TakeQuery {
|
||||
async fn create_plan(&self, options: QueryExecutionOptions) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
if let Some(offsets) = &self.offsets {
|
||||
return self.create_offsets_plan(offsets, options).await;
|
||||
}
|
||||
|
||||
let req = AnyQuery::Query(self.request.clone());
|
||||
self.parent.clone().create_plan(&req, options).await
|
||||
}
|
||||
@@ -1638,6 +2203,18 @@ impl ExecutableQuery for TakeQuery {
|
||||
&self,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<SendableRecordBatchStream> {
|
||||
if self.offsets.is_some() {
|
||||
let plan = self.create_plan(options.clone()).await?;
|
||||
let inner = execute_plan(plan, Default::default())?;
|
||||
let inner = MaxBatchLengthStream::new_boxed(inner, options.max_batch_length as usize);
|
||||
let inner = if let Some(timeout) = options.timeout {
|
||||
TimeoutStream::new_boxed(inner, timeout)
|
||||
} else {
|
||||
inner
|
||||
};
|
||||
return Ok(DatasetRecordBatchStream::new(inner).into());
|
||||
}
|
||||
|
||||
let query = AnyQuery::Query(self.request.clone());
|
||||
Ok(SendableRecordBatchStream::from(
|
||||
self.parent.clone().query(&query, options).await?,
|
||||
@@ -1645,11 +2222,51 @@ impl ExecutableQuery for TakeQuery {
|
||||
}
|
||||
|
||||
async fn explain_plan(&self, verbose: bool) -> Result<String> {
|
||||
if let Some(offsets) = &self.offsets {
|
||||
let (request, _, _, output_offset, output_limit) =
|
||||
Self::prepare_offsets_lookup(self.parent.as_ref(), &self.request).await?;
|
||||
// Ask the backend to explain only the distinct-row lookup. This keeps
|
||||
// remote explanation non-executing while still showing the client-side
|
||||
// operators that create_plan and execution place above that lookup.
|
||||
let lookup = self
|
||||
.parent
|
||||
.explain_plan(&AnyQuery::Query(request), verbose)
|
||||
.await?;
|
||||
return Ok(Self::wrap_offsets_explanation(
|
||||
&lookup,
|
||||
offsets.len(),
|
||||
output_offset,
|
||||
output_limit,
|
||||
self.preserve_order,
|
||||
));
|
||||
}
|
||||
|
||||
let query = AnyQuery::Query(self.request.clone());
|
||||
self.parent.explain_plan(&query, verbose).await
|
||||
}
|
||||
|
||||
async fn analyze_plan_with_options(&self, options: QueryExecutionOptions) -> Result<String> {
|
||||
if self.offsets.is_some() {
|
||||
if self.parent.analyze_plan_is_remote() {
|
||||
let (request, _, _, _, _) =
|
||||
Self::prepare_offsets_lookup(self.parent.as_ref(), &self.request).await?;
|
||||
// Remote analysis is owned by the service. The current wire
|
||||
// request represents only the distinct-row lookup, so return
|
||||
// the service report unchanged instead of fabricating metrics
|
||||
// for client-side restoration operators.
|
||||
return self
|
||||
.parent
|
||||
.analyze_plan(&AnyQuery::Query(request), options)
|
||||
.await;
|
||||
}
|
||||
|
||||
let plan = self.create_plan(options).await?;
|
||||
execute_plan(plan.clone(), Default::default())?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
return Ok(format_analyzed_plan(plan));
|
||||
}
|
||||
|
||||
let query = AnyQuery::Query(self.request.clone());
|
||||
self.parent.analyze_plan(&query, options).await
|
||||
}
|
||||
@@ -1670,6 +2287,7 @@ mod tests {
|
||||
StringArray, cast::AsArray, types::Float32Type,
|
||||
};
|
||||
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
|
||||
use datafusion_physical_plan::display::DisplayableExecutionPlan;
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
use lance_testing::datagen::{BatchGenerator, IncrementingInt32, RandomVector};
|
||||
use rand::seq::IndexedRandom;
|
||||
@@ -2924,6 +3542,218 @@ mod tests {
|
||||
assert_eq!(results[0].num_columns(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_preserves_duplicate_multiplicity() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
|
||||
let results = table
|
||||
.take_offsets(vec![5, 1, 5, 17])
|
||||
.select(Select::Columns(vec!["id".to_string()]))
|
||||
.execute_with_options(QueryExecutionOptions {
|
||||
max_batch_length: 2,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(results.len(), 2);
|
||||
assert!(results.iter().all(|batch| batch.num_columns() == 1));
|
||||
let mut ids = results
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_primitive::<Int32Type>()
|
||||
.values()
|
||||
.to_vec()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
ids.sort_unstable();
|
||||
assert_eq!(ids, vec![1, 5, 5, 17]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_plan_is_incremental() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
|
||||
let plan = table
|
||||
.take_offsets(vec![5, 1, 17])
|
||||
.create_plan(QueryExecutionOptions {
|
||||
max_batch_length: 1,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(plan.properties().emission_type, EmissionType::Incremental);
|
||||
let displayed = DisplayableExecutionPlan::new(plan.as_ref())
|
||||
.indent(false)
|
||||
.to_string();
|
||||
assert!(displayed.contains("TakeRestoreExec"));
|
||||
assert!(!displayed.contains("CoalescePartitionsExec"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_into_request_preserves_duplicate_multiplicity() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
let request = table.take_offsets(vec![5, 5]).into_request();
|
||||
assert_eq!(request.take_offsets, Some(vec![5, 5]));
|
||||
|
||||
let batches = table
|
||||
.base_table()
|
||||
.query(&AnyQuery::Query(request), QueryExecutionOptions::default())
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_restore_take_batch_only_reorders_when_requested() {
|
||||
let batch = RecordBatch::try_from_iter([
|
||||
(
|
||||
"id",
|
||||
Arc::new(Int32Array::from(vec![17, 5, 1])) as Arc<dyn Array>,
|
||||
),
|
||||
(
|
||||
"_rowoffset",
|
||||
Arc::new(UInt64Array::from(vec![17, 5, 1])) as Arc<dyn Array>,
|
||||
),
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let restored =
|
||||
restore_take_batch(batch.clone(), &[5, 1, 5, 17], "_rowoffset", true, false).unwrap();
|
||||
assert_eq!(
|
||||
restored
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_primitive::<Int32Type>()
|
||||
.values(),
|
||||
&[17, 5, 5, 1]
|
||||
);
|
||||
|
||||
let ordered = restore_take_batch(batch, &[5, 1, 5, 17], "_rowoffset", true, true).unwrap();
|
||||
assert_eq!(
|
||||
ordered
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_primitive::<Int32Type>()
|
||||
.values(),
|
||||
&[5, 1, 5, 17]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_applies_pagination_after_restoration() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
|
||||
let limited = table
|
||||
.take_offsets(vec![0, 1, 0, 2])
|
||||
.select(Select::Columns(vec!["id".to_string()]))
|
||||
.limit(3)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let limited = concat_batches(&limited[0].schema(), &limited).unwrap();
|
||||
assert_eq!(limited.num_rows(), 3);
|
||||
assert!(
|
||||
limited
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_primitive::<Int32Type>()
|
||||
.values()
|
||||
.iter()
|
||||
.all(|id| [0, 1, 2].contains(id))
|
||||
);
|
||||
|
||||
let offset = table
|
||||
.take_offsets(vec![5, 1, 5, 17])
|
||||
.select(Select::Columns(vec!["id".to_string()]))
|
||||
.offset(1)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let offset = concat_batches(&offset[0].schema(), &offset).unwrap();
|
||||
assert_eq!(offset.num_rows(), 3);
|
||||
assert!(
|
||||
offset
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_primitive::<Int32Type>()
|
||||
.values()
|
||||
.iter()
|
||||
.all(|id| [1, 5, 17].contains(id))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_create_plan_restores_occurrences() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
let take = table
|
||||
.take_offsets(vec![5, 1, 5, 17])
|
||||
.select(Select::Columns(vec!["id".to_string()]));
|
||||
|
||||
let plan = take
|
||||
.create_plan(QueryExecutionOptions::default())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(plan.schema().fields().len(), 1);
|
||||
assert_eq!(plan.schema().field(0).name(), "id");
|
||||
let planned = execute_plan(plan, Default::default())
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let planned = concat_batches(&planned[0].schema(), &planned).unwrap();
|
||||
let mut ids = planned
|
||||
.column_by_name("id")
|
||||
.unwrap()
|
||||
.as_primitive::<Int32Type>()
|
||||
.values()
|
||||
.to_vec();
|
||||
ids.sort_unstable();
|
||||
assert_eq!(ids, vec![1, 5, 5, 17]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_plan_introspection_shows_restoration() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let table = make_test_table(&tmp_dir).await;
|
||||
let take = table
|
||||
.take_offsets(vec![0, 1, 0, 2])
|
||||
.select(Select::Columns(vec!["id".to_string()]))
|
||||
.limit(3);
|
||||
|
||||
let explained = take.explain_plan(false).await.unwrap();
|
||||
assert!(explained.contains("GlobalLimitExec"));
|
||||
assert!(explained.contains("TakeRestoreExec"));
|
||||
assert!(!explained.contains("CoalescePartitionsExec"));
|
||||
|
||||
let analyzed = take.analyze_plan().await.unwrap();
|
||||
assert!(analyzed.contains("GlobalLimitExec"));
|
||||
assert!(analyzed.contains("TakeRestoreExec"));
|
||||
assert!(!analyzed.contains("CoalescePartitionsExec"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_row_ids() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
|
||||
@@ -11,6 +11,7 @@ pub(crate) mod db;
|
||||
pub(crate) mod job;
|
||||
pub mod oauth;
|
||||
mod retry;
|
||||
pub(crate) mod sql;
|
||||
pub(crate) mod table;
|
||||
pub(crate) mod util;
|
||||
|
||||
|
||||
+596
-88
@@ -1,7 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -20,19 +20,20 @@ use lance_namespace::models::{
|
||||
|
||||
use crate::Error;
|
||||
use crate::database::{
|
||||
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
|
||||
JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest,
|
||||
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions, JobInfo,
|
||||
OpenTableRequest, ReadConsistency, SecretInfo, TableNamesRequest,
|
||||
};
|
||||
use crate::error::Result;
|
||||
use crate::function::{FunctionRegistrationRequest, FunctionVersion};
|
||||
use crate::job::Job;
|
||||
use crate::remote::job::{DescribeJobResponse, RemoteJob, job_state_to_client};
|
||||
use crate::remote::job::{RemoteJob, job_state_to_client};
|
||||
use crate::remote::util::stream_as_body;
|
||||
use crate::table::BaseTable;
|
||||
|
||||
use super::client::{
|
||||
ClientConfig, HeaderProvider, HttpSend, RequestResultExt, RestfulLanceDbClient, Sender,
|
||||
};
|
||||
use super::sql::SqlClient;
|
||||
use super::table::RemoteTable;
|
||||
use super::util::parse_server_version;
|
||||
use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id};
|
||||
@@ -97,6 +98,7 @@ pub const OPT_REMOTE_PREFIX: &str = "remote_database_";
|
||||
pub const OPT_REMOTE_API_KEY: &str = "remote_database_api_key";
|
||||
pub const OPT_REMOTE_REGION: &str = "remote_database_region";
|
||||
pub const OPT_REMOTE_HOST_OVERRIDE: &str = "remote_database_host_override";
|
||||
pub const OPT_REMOTE_SQL_HOST_OVERRIDE: &str = "remote_database_sql_host_override";
|
||||
// TODO: add support for configuring client config via key/value options
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
@@ -212,6 +214,7 @@ pub struct RemoteDatabase<S: HttpSend = Sender> {
|
||||
namespace_context_provider: Option<Arc<dyn DynamicContextProvider>>,
|
||||
/// TLS configuration for mTLS support
|
||||
tls_config: Option<super::client::TlsConfig>,
|
||||
sql_client: Option<SqlClient>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -269,22 +272,51 @@ impl DynamicContextProvider for NamespaceHeaderProviderContext {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct RemoteHostOverrides {
|
||||
pub rest: Option<String>,
|
||||
pub sql: Option<String>,
|
||||
}
|
||||
|
||||
/// Attach a namespace path to a Secret request body.
|
||||
///
|
||||
/// A root path is omitted rather than sent empty, so a root request is byte
|
||||
/// identical to one from a client that predates namespace addressing.
|
||||
fn add_namespace_path(body: &mut serde_json::Value, namespace_path: &[String]) {
|
||||
if namespace_path.is_empty() {
|
||||
return;
|
||||
}
|
||||
body["namespace_path"] = serde_json::Value::Array(
|
||||
namespace_path
|
||||
.iter()
|
||||
.map(|segment| serde_json::Value::String(segment.clone()))
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
|
||||
impl RemoteDatabase {
|
||||
pub fn try_new(
|
||||
pub(crate) fn try_new(
|
||||
uri: &str,
|
||||
api_key: &str,
|
||||
region: &str,
|
||||
host_override: Option<String>,
|
||||
host_overrides: RemoteHostOverrides,
|
||||
client_config: ClientConfig,
|
||||
options: RemoteOptions,
|
||||
read_consistency_interval: Option<std::time::Duration>,
|
||||
) -> Result<Self> {
|
||||
let parsed = super::client::parse_db_url(uri)?;
|
||||
let sql_client = SqlClient::new(
|
||||
parsed.db_name.clone(),
|
||||
parsed.db_prefix.clone(),
|
||||
api_key.to_string(),
|
||||
host_overrides.rest.clone(),
|
||||
host_overrides.sql,
|
||||
client_config.clone(),
|
||||
);
|
||||
let header_map = RestfulLanceDbClient::<Sender>::default_headers(
|
||||
api_key,
|
||||
region,
|
||||
&parsed.db_name,
|
||||
host_override.is_some(),
|
||||
host_overrides.rest.is_some(),
|
||||
&options,
|
||||
parsed.db_prefix.as_deref(),
|
||||
&client_config,
|
||||
@@ -312,7 +344,7 @@ impl RemoteDatabase {
|
||||
let client = RestfulLanceDbClient::try_new(
|
||||
&parsed,
|
||||
region,
|
||||
host_override,
|
||||
host_overrides.rest,
|
||||
header_map,
|
||||
client_config.clone(),
|
||||
read_consistency_interval,
|
||||
@@ -330,11 +362,34 @@ impl RemoteDatabase {
|
||||
namespace_headers,
|
||||
namespace_context_provider,
|
||||
tls_config: client_config.tls_config,
|
||||
sql_client: Some(sql_client),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: HttpSend> RemoteDatabase<S> {
|
||||
/// `create` and `alter` differ only in which name state the server
|
||||
/// requires, so they share one request shape. The value is a request field
|
||||
/// and never a path segment or query parameter, which keeps it out of
|
||||
/// access logs and proxy traces.
|
||||
async fn write_secret(
|
||||
&self,
|
||||
route: &str,
|
||||
name: &str,
|
||||
value: &str,
|
||||
namespace_path: &[String],
|
||||
) -> Result<()> {
|
||||
let mut body = serde_json::json!({
|
||||
"name": name,
|
||||
"value": value,
|
||||
});
|
||||
add_namespace_path(&mut body, namespace_path);
|
||||
let req = self.client.post(route).json(&body);
|
||||
let (request_id, response) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, response).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn submit_drop_table(
|
||||
&self,
|
||||
name: &str,
|
||||
@@ -427,6 +482,7 @@ mod test_utils {
|
||||
namespace_headers: HashMap::new(),
|
||||
namespace_context_provider: None,
|
||||
tls_config: None,
|
||||
sql_client: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -449,6 +505,7 @@ mod test_utils {
|
||||
namespace_headers: config.extra_headers.clone(),
|
||||
namespace_context_provider,
|
||||
tls_config: config.tls_config.clone(),
|
||||
sql_client: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -533,6 +590,39 @@ struct RemoteListJobsResponse {
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListedFunctionVersion {
|
||||
definition: FunctionVersion,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListFunctionsResponse {
|
||||
#[serde(default)]
|
||||
functions: Vec<RemoteListedFunctionVersion>,
|
||||
#[serde(default)]
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteDropFunctionResponse {
|
||||
dropped: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListSecretsResponse {
|
||||
#[serde(default)]
|
||||
secrets: Vec<RemoteListedSecret>,
|
||||
#[serde(default)]
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
/// An object rather than a bare name so a later listing can carry a Secret's
|
||||
/// type or last-updated time without breaking this one.
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListedSecret {
|
||||
name: String,
|
||||
}
|
||||
|
||||
/// Bound on `list_jobs` page walking; a warning is logged when the listing
|
||||
/// is truncated at this many pages.
|
||||
const MAX_LIST_JOBS_PAGES: usize = 100;
|
||||
@@ -583,11 +673,135 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
response.json().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
fn job(&self, job_id: &str) -> Result<crate::job::Job> {
|
||||
Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new(
|
||||
self.client.clone(),
|
||||
job_id.to_string(),
|
||||
))))
|
||||
async fn list_functions(&self) -> Result<Vec<FunctionVersion>> {
|
||||
let mut functions = Vec::new();
|
||||
let mut page_token: Option<String> = None;
|
||||
let mut seen_page_tokens = HashSet::new();
|
||||
loop {
|
||||
let mut body = serde_json::json!({ "include_definition": true });
|
||||
if let Some(token) = &page_token {
|
||||
body["page_token"] = serde_json::Value::String(token.clone());
|
||||
}
|
||||
let req = self.client.post("/v1/functions/list").json(&body);
|
||||
let (request_id, response) = self.client.send(req).await?;
|
||||
let response = self.client.check_response(&request_id, response).await?;
|
||||
let status = response.status();
|
||||
let response: RemoteListFunctionsResponse =
|
||||
response.json().await.err_to_http(request_id.clone())?;
|
||||
functions.extend(
|
||||
response
|
||||
.functions
|
||||
.into_iter()
|
||||
.map(|listed| listed.definition),
|
||||
);
|
||||
let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty())
|
||||
else {
|
||||
break;
|
||||
};
|
||||
if !seen_page_tokens.insert(next_page_token.clone()) {
|
||||
return Err(Error::Http {
|
||||
source: "Function listing response repeated a page_token".into(),
|
||||
request_id,
|
||||
status_code: Some(status),
|
||||
});
|
||||
}
|
||||
page_token = Some(next_page_token);
|
||||
}
|
||||
Ok(functions)
|
||||
}
|
||||
|
||||
async fn drop_function(&self, name: &str, version: &str) -> Result<bool> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/functions/drop")
|
||||
.json(&serde_json::json!({
|
||||
"name": name,
|
||||
"version": version,
|
||||
}));
|
||||
let (request_id, response) = self.client.send(req).await?;
|
||||
let response = self.client.check_response(&request_id, response).await?;
|
||||
let response: RemoteDropFunctionResponse = response.json().await.err_to_http(request_id)?;
|
||||
Ok(response.dropped)
|
||||
}
|
||||
|
||||
async fn create_secret(
|
||||
&self,
|
||||
name: &str,
|
||||
value: &str,
|
||||
namespace_path: &[String],
|
||||
) -> Result<()> {
|
||||
self.write_secret("/v1/secrets/create", name, value, namespace_path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn alter_secret(&self, name: &str, value: &str, namespace_path: &[String]) -> Result<()> {
|
||||
self.write_secret("/v1/secrets/alter", name, value, namespace_path)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_secrets(&self, namespace_path: &[String]) -> Result<Vec<String>> {
|
||||
let mut names = Vec::new();
|
||||
let mut page_token: Option<String> = None;
|
||||
let mut seen_page_tokens = HashSet::new();
|
||||
loop {
|
||||
let mut body = serde_json::json!({});
|
||||
if let Some(token) = &page_token {
|
||||
body["page_token"] = serde_json::Value::String(token.clone());
|
||||
}
|
||||
add_namespace_path(&mut body, namespace_path);
|
||||
let req = self.client.post("/v1/secrets/list").json(&body);
|
||||
let (request_id, response) = self.client.send(req).await?;
|
||||
let response = self.client.check_response(&request_id, response).await?;
|
||||
let status = response.status();
|
||||
let response: RemoteListSecretsResponse =
|
||||
response.json().await.err_to_http(request_id.clone())?;
|
||||
names.extend(response.secrets.into_iter().map(|secret| secret.name));
|
||||
let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty())
|
||||
else {
|
||||
break;
|
||||
};
|
||||
if !seen_page_tokens.insert(next_page_token.clone()) {
|
||||
return Err(Error::Http {
|
||||
source: "Secret listing response repeated a page_token".into(),
|
||||
request_id,
|
||||
status_code: Some(status),
|
||||
});
|
||||
}
|
||||
page_token = Some(next_page_token);
|
||||
}
|
||||
Ok(names)
|
||||
}
|
||||
|
||||
async fn drop_secret(&self, name: &str, namespace_path: &[String]) -> Result<()> {
|
||||
let mut body = serde_json::json!({ "name": name });
|
||||
add_namespace_path(&mut body, namespace_path);
|
||||
let req = self.client.post("/v1/secrets/drop").json(&body);
|
||||
let (request_id, response) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, response).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn describe_secret(&self, name: &str, namespace_path: &[String]) -> Result<SecretInfo> {
|
||||
let mut body = serde_json::json!({ "name": name });
|
||||
add_namespace_path(&mut body, namespace_path);
|
||||
let req = self.client.post("/v1/secrets/describe").json(&body);
|
||||
let (request_id, response) = self.client.send(req).await?;
|
||||
let response = self.client.check_response(&request_id, response).await?;
|
||||
response.json().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
async fn open_job(&self, job_id: &str) -> Result<Job> {
|
||||
let handle = super::job::RemoteJob::new(self.client.clone(), job_id.to_string());
|
||||
match crate::job::JobHandle::describe(&handle).await {
|
||||
Ok(description) => Ok(Job::opened(Box::new(handle), description)),
|
||||
Err(Error::Http {
|
||||
status_code: Some(StatusCode::NOT_FOUND),
|
||||
..
|
||||
}) => Err(Error::JobNotFound {
|
||||
job_id: job_id.to_string(),
|
||||
}),
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
@@ -624,31 +838,6 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
async fn get_job(&self, job_id: &str) -> Result<Option<JobDescription>> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/describe")
|
||||
.json(&serde_json::json!({ "job_id": job_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = match self.client.check_response(&request_id, rsp).await {
|
||||
Ok(rsp) => rsp,
|
||||
Err(Error::Http {
|
||||
status_code: Some(StatusCode::NOT_FOUND),
|
||||
..
|
||||
}) => return Ok(None),
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
let body: DescribeJobResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(Some(JobDescription {
|
||||
job_id: body.job_id,
|
||||
job_type: body.job_type,
|
||||
state: job_state_to_client(&body.job_state),
|
||||
creation_ms: body.creation_ms,
|
||||
spec: body.spec,
|
||||
failure: body.failure.map(|reported| reported.into_job_failure()),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
||||
let req = self
|
||||
.client
|
||||
@@ -665,19 +854,28 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<arrow_array::RecordBatch>> {
|
||||
let mut body = serde_json::json!({});
|
||||
if let Some(job_id) = job_id {
|
||||
body["job_id"] = serde_json::Value::String(job_id.to_string());
|
||||
}
|
||||
let req = self.client.post("/v1/jobs/query_events").json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let bytes = rsp.bytes().await.err_to_http(request_id)?;
|
||||
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None)?;
|
||||
reader
|
||||
.collect::<std::result::Result<Vec<_>, _>>()
|
||||
.map_err(Into::into)
|
||||
async fn execute_query_async(
|
||||
&self,
|
||||
query: &str,
|
||||
default_namespace_path: &[String],
|
||||
) -> Result<crate::sql::Query> {
|
||||
let client = self
|
||||
.sql_client
|
||||
.as_ref()
|
||||
.ok_or_else(|| Error::NotSupported {
|
||||
message: "SQL is unavailable for this remote database client".to_string(),
|
||||
})?;
|
||||
client.submit(query, default_namespace_path).await
|
||||
}
|
||||
|
||||
async fn describe_query(&self, query_id: uuid::Uuid) -> Result<crate::sql::QueryDescription> {
|
||||
let client = self
|
||||
.sql_client
|
||||
.as_ref()
|
||||
.ok_or_else(|| Error::NotSupported {
|
||||
message: "SQL is unavailable for this remote database client".to_string(),
|
||||
})?;
|
||||
client.describe(query_id).await
|
||||
}
|
||||
|
||||
async fn table_names(&self, request: TableNamesRequest) -> Result<Vec<String>> {
|
||||
@@ -1195,6 +1393,7 @@ mod tests {
|
||||
use crate::{
|
||||
Connection, Error,
|
||||
database::CreateTableMode,
|
||||
job::JobEventsRequest,
|
||||
remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE},
|
||||
};
|
||||
|
||||
@@ -2543,7 +2742,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_job() {
|
||||
async fn test_open_job() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/jobs/describe");
|
||||
@@ -2557,51 +2756,55 @@ mod tests {
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let job = conn.get_job("job-1").await.unwrap().unwrap();
|
||||
assert_eq!(job.job_id, "job-1");
|
||||
assert_eq!(job.job_type, "create_index");
|
||||
assert_eq!(job.state, "failed");
|
||||
assert_eq!(job.creation_ms, 1000);
|
||||
assert_eq!(job.spec["column"], "vec");
|
||||
let failure = job.failure.unwrap();
|
||||
// Opening populates the handle, so the accessors answer without a
|
||||
// second round trip.
|
||||
let job = conn.open_job("job-1").await.unwrap();
|
||||
assert_eq!(job.id(), Some("job-1"));
|
||||
assert_eq!(job.job_type().as_deref(), Some("create_index"));
|
||||
assert_eq!(job.state().as_deref(), Some("failed"));
|
||||
assert_eq!(job.creation_ms(), Some(1000));
|
||||
assert_eq!(job.spec().unwrap()["column"], "vec");
|
||||
assert!(job.result().is_none());
|
||||
let failure = job.failure().unwrap();
|
||||
assert_eq!(failure.phase.as_deref(), Some("execute"));
|
||||
assert_eq!(failure.message.as_deref(), Some("worker died"));
|
||||
assert_eq!(failure.retryable, Some(true));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_job_missing_is_none() {
|
||||
async fn test_open_job_reports_the_terminal_result() {
|
||||
let conn = Connection::new_with_handler(|_| {
|
||||
http::Response::builder()
|
||||
.status(404)
|
||||
.body("no such job")
|
||||
.unwrap()
|
||||
});
|
||||
assert!(conn.get_job("nope").await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cancel_job() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/cancel");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id": "job-1"}"#)
|
||||
.body(
|
||||
r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "DONE", "creation_ms": 1000, "result": {"rows_assigned": 1000000, "rows_failed": 0}}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
assert!(conn.cancel_job("job-1").await.unwrap());
|
||||
let job = conn.open_job("job-1").await.unwrap();
|
||||
assert_eq!(job.state().as_deref(), Some("finished"));
|
||||
let result = job.result().unwrap();
|
||||
assert_eq!(result["rows_assigned"], 1_000_000);
|
||||
assert_eq!(result["rows_failed"], 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_open_job_missing_fails() {
|
||||
let conn = Connection::new_with_handler(|_| {
|
||||
http::Response::builder()
|
||||
.status(404)
|
||||
.body("no such job")
|
||||
.unwrap()
|
||||
});
|
||||
assert!(!conn.cancel_job("nope").await.unwrap());
|
||||
let err = conn.open_job("nope").await.unwrap_err();
|
||||
assert!(
|
||||
matches!(&err, Error::JobNotFound { job_id } if job_id == "nope"),
|
||||
"{err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_history_parses_arrow_stream() {
|
||||
async fn test_job_events_scope_to_that_job() {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"state",
|
||||
DataType::Utf8,
|
||||
@@ -2610,29 +2813,198 @@ mod tests {
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![Arc::new(arrow_array::StringArray::from(vec![
|
||||
"created", "done",
|
||||
"claim_complete",
|
||||
]))],
|
||||
)
|
||||
.unwrap();
|
||||
let mut body = Vec::new();
|
||||
let mut events = Vec::new();
|
||||
{
|
||||
let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, &schema).unwrap();
|
||||
let mut writer =
|
||||
arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap();
|
||||
writer.write(&batch).unwrap();
|
||||
writer.finish().unwrap();
|
||||
}
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/query_events");
|
||||
let req_body: serde_json::Value =
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(req_body["job_id"], "job-1");
|
||||
if request.url().path() == "/v1/jobs/describe" {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"#
|
||||
.as_bytes()
|
||||
.to_vec(),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(request.url().path(), "/v1/jobs/query_events");
|
||||
// The handle supplies job_id; the caller only narrows the query.
|
||||
assert_eq!(body["job_id"], "job-1");
|
||||
assert_eq!(body["limit"], 500);
|
||||
assert_eq!(body["filter"], "state = 'claim_complete'");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(body.clone())
|
||||
.body(events.clone())
|
||||
.unwrap()
|
||||
});
|
||||
let batches = conn.job_history(Some("job-1")).await.unwrap();
|
||||
let job = conn.open_job("job-1").await.unwrap();
|
||||
let batches = job
|
||||
.events(
|
||||
JobEventsRequest::default()
|
||||
.limit(500)
|
||||
.filter("state = 'claim_complete'"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches.len(), 1);
|
||||
assert_eq!(batches[0].num_rows(), 2);
|
||||
assert_eq!(batches[0].num_rows(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_events_keep_the_schema_when_nothing_matches() {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"state",
|
||||
DataType::Utf8,
|
||||
false,
|
||||
)]));
|
||||
let mut events = Vec::new();
|
||||
{
|
||||
let mut writer =
|
||||
arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap();
|
||||
writer.finish().unwrap();
|
||||
}
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
if request.url().path() == "/v1/jobs/describe" {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"job_id": "job-1", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 1}"#
|
||||
.as_bytes()
|
||||
.to_vec(),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
// Only the job id when the caller narrows nothing.
|
||||
assert_eq!(body, serde_json::json!({ "job_id": "job-1" }));
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(events.clone())
|
||||
.unwrap()
|
||||
});
|
||||
let job = conn.open_job("job-1").await.unwrap();
|
||||
let batches = job.events(JobEventsRequest::default()).await.unwrap();
|
||||
assert_eq!(batches.len(), 1);
|
||||
assert_eq!(batches[0].num_rows(), 0);
|
||||
assert_eq!(batches[0].schema(), schema);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_and_alter_secret_send_the_value_in_the_request_body() {
|
||||
for (route, call) in [("/v1/secrets/create", true), ("/v1/secrets/alter", false)] {
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), route);
|
||||
// Never a path segment or query parameter, which is what keeps
|
||||
// it out of access logs and proxy traces.
|
||||
assert!(request.url().query().is_none(), "{:?}", request.url());
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["name"], "openai-prod");
|
||||
assert_eq!(body["value"], "sk-live-0001");
|
||||
http::Response::builder().status(200).body("{}").unwrap()
|
||||
});
|
||||
if call {
|
||||
conn.create_secret("openai-prod", "sk-live-0001", &[])
|
||||
.await
|
||||
.unwrap();
|
||||
} else {
|
||||
conn.alter_secret("openai-prod", "sk-live-0001", &[])
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_secrets_walks_pages_and_returns_names_only() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/secrets/list");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
let page = body.get("page_token").and_then(|token| token.as_str());
|
||||
let body = match page {
|
||||
None => r#"{"secrets":[{"name":"openai-prod"}],"page_token":"p2"}"#,
|
||||
Some("p2") => r#"{"secrets":[{"name":"hf-prod"}]}"#,
|
||||
Some(other) => panic!("unexpected page token: {other}"),
|
||||
};
|
||||
http::Response::builder().status(200).body(body).unwrap()
|
||||
});
|
||||
assert_eq!(
|
||||
conn.list_secrets(&[]).await.unwrap(),
|
||||
vec!["openai-prod".to_string(), "hf-prod".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
/// A server that keeps handing back the same token would otherwise spin
|
||||
/// forever.
|
||||
#[tokio::test]
|
||||
async fn test_list_secrets_rejects_a_repeated_page_token() {
|
||||
let conn = Connection::new_with_handler(|_| {
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"secrets":[{"name":"openai-prod"}],"page_token":"same"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
let error = conn.list_secrets(&[]).await.unwrap_err();
|
||||
assert!(
|
||||
error.to_string().contains("repeated a page_token"),
|
||||
"{error}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_drop_secret_posts_the_name_alone() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/secrets/drop");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body, serde_json::json!({"name": "openai-prod"}));
|
||||
http::Response::builder().status(200).body("{}").unwrap()
|
||||
});
|
||||
conn.drop_secret("openai-prod", &[]).await.unwrap();
|
||||
}
|
||||
|
||||
/// A namespace path is sent when there is one and omitted when there is
|
||||
/// not, so a root request stays byte identical to one from a client that
|
||||
/// predates namespace addressing -- which is what lets the parameter ship
|
||||
/// before every server implements it.
|
||||
#[tokio::test]
|
||||
async fn test_a_namespace_path_is_sent_only_when_it_is_not_root() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
body,
|
||||
serde_json::json!({
|
||||
"name": "openai-prod",
|
||||
"namespace_path": ["prod", "vision"],
|
||||
})
|
||||
);
|
||||
http::Response::builder().status(200).body("{}").unwrap()
|
||||
});
|
||||
conn.drop_secret("openai-prod", &["prod".to_string(), "vision".to_string()])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert!(body.get("namespace_path").is_none(), "{body}");
|
||||
http::Response::builder().status(200).body("{}").unwrap()
|
||||
});
|
||||
conn.drop_secret("openai-prod", &[]).await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -2689,13 +3061,147 @@ mod tests {
|
||||
assert_eq!(version.version(), "fv_01K3EXACT");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_functions_requests_definitions_and_paginates() {
|
||||
const VERSION: &str = include_str!(
|
||||
"../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json"
|
||||
);
|
||||
let version: serde_json::Value = serde_json::from_str(VERSION).unwrap();
|
||||
let page = Arc::new(AtomicUsize::new(0));
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/functions/list");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["include_definition"], true);
|
||||
match page.fetch_add(1, Ordering::SeqCst) {
|
||||
0 => {
|
||||
assert!(body.get("page_token").is_none());
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"functions": [], "page_token": "next"}"#.to_string())
|
||||
.unwrap()
|
||||
}
|
||||
_ => {
|
||||
assert_eq!(body["page_token"], "next");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
serde_json::json!({
|
||||
"functions": [{
|
||||
"name": "embed",
|
||||
"version": "fv_01K3EXACT",
|
||||
"definition": version.clone(),
|
||||
}],
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
});
|
||||
let functions = conn.list_functions().await.unwrap();
|
||||
assert_eq!(functions.len(), 1);
|
||||
assert_eq!(functions[0].name(), "embed");
|
||||
assert_eq!(functions[0].version(), "fv_01K3EXACT");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_functions_stops_on_an_empty_page_token() {
|
||||
let requests = Arc::new(AtomicUsize::new(0));
|
||||
let seen = requests.clone();
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
seen.fetch_add(1, Ordering::SeqCst);
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert!(body.get("page_token").is_none());
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"functions": [], "page_token": ""}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let functions = conn.list_functions().await.unwrap();
|
||||
assert!(functions.is_empty());
|
||||
assert_eq!(requests.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_functions_rejects_a_page_token_cycle() {
|
||||
let page = Arc::new(AtomicUsize::new(0));
|
||||
let requests = page.clone();
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
let next_page_token = match page.fetch_add(1, Ordering::SeqCst) {
|
||||
0 => {
|
||||
assert!(body.get("page_token").is_none());
|
||||
"one"
|
||||
}
|
||||
1 => {
|
||||
assert_eq!(body["page_token"], "one");
|
||||
"two"
|
||||
}
|
||||
2 => {
|
||||
assert_eq!(body["page_token"], "two");
|
||||
"one"
|
||||
}
|
||||
page => panic!("unexpected page: {page}"),
|
||||
};
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
serde_json::json!({
|
||||
"functions": [],
|
||||
"page_token": next_page_token,
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let error = conn.list_functions().await.unwrap_err();
|
||||
assert!(
|
||||
matches!(
|
||||
&error,
|
||||
Error::Http {
|
||||
status_code: Some(http::StatusCode::OK),
|
||||
..
|
||||
}
|
||||
),
|
||||
"got {error:?}"
|
||||
);
|
||||
assert_eq!(requests.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_drop_function_sends_exact_version_and_decodes_replay() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/functions/drop");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
body,
|
||||
serde_json::json!({"name": "embed", "version": "fv_01K3EXACT"})
|
||||
);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"dropped":false}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert!(!conn.drop_function("embed", "fv_01K3EXACT").await.unwrap());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_conn_job_waits_to_done() {
|
||||
let polls = Arc::new(AtomicUsize::new(0));
|
||||
let polls_ref = polls.clone();
|
||||
let conn = Connection::new_with_handler(move |request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/describe");
|
||||
let state = if polls_ref.fetch_add(1, Ordering::SeqCst) == 0 {
|
||||
// Two in-progress answers: one for the load, one for the first
|
||||
// status poll.
|
||||
let state = if polls_ref.fetch_add(1, Ordering::SeqCst) < 2 {
|
||||
"IN_PROGRESS"
|
||||
} else {
|
||||
"DONE"
|
||||
@@ -2708,11 +3214,13 @@ mod tests {
|
||||
))
|
||||
.unwrap()
|
||||
});
|
||||
let job = conn.job("job-1").unwrap();
|
||||
let job = conn.open_job("job-1").await.unwrap();
|
||||
assert_eq!(job.id(), Some("job-1"));
|
||||
// Opening already answered the state; no extra call needed for it.
|
||||
assert_eq!(job.state().as_deref(), Some("running"));
|
||||
assert_eq!(job.status().await.unwrap(), "running");
|
||||
job.wait().await.unwrap();
|
||||
assert_eq!(job.status().await.unwrap(), "finished");
|
||||
assert!(polls.load(Ordering::SeqCst) >= 3);
|
||||
assert!(polls.load(Ordering::SeqCst) >= 4);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,13 +5,15 @@
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use arrow_array::RecordBatch;
|
||||
use async_trait::async_trait;
|
||||
use tokio::time::sleep;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::database::JobDescription;
|
||||
use crate::error::{Error, JobFailure, Result};
|
||||
use crate::job::{JobHandle, TerminalResult};
|
||||
use crate::job::{JobEventsRequest, JobHandle, TerminalResult};
|
||||
use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient};
|
||||
|
||||
/// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`].
|
||||
@@ -86,7 +88,7 @@ pub(super) struct DescribeJobResponse {
|
||||
#[serde(default)]
|
||||
pub(super) spec: serde_json::Value,
|
||||
#[serde(default)]
|
||||
result: Option<serde_json::Value>,
|
||||
pub(super) result: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub(super) failure: Option<ReportedFailure>,
|
||||
}
|
||||
@@ -110,6 +112,39 @@ impl DescribeJobResponse {
|
||||
fn into_terminal_result(self, request_id: String) -> TerminalResult {
|
||||
TerminalResult::remote(self.result, request_id)
|
||||
}
|
||||
|
||||
/// The public description this wire envelope stands for.
|
||||
pub(super) fn into_description(self) -> JobDescription {
|
||||
JobDescription {
|
||||
job_id: self.job_id,
|
||||
job_type: self.job_type,
|
||||
state: JobState::from(self.job_state.as_str()).client_label(),
|
||||
creation_ms: self.creation_ms,
|
||||
spec: self.spec,
|
||||
result: self.result,
|
||||
failure: self.failure.map(ReportedFailure::into_job_failure),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One `/v1/jobs/query_events` round trip.
|
||||
pub(super) async fn fetch_job_events<S: HttpSend>(
|
||||
client: &RestfulLanceDbClient<S>,
|
||||
body: serde_json::Value,
|
||||
) -> Result<Vec<RecordBatch>> {
|
||||
let request = client.post("/v1/jobs/query_events").json(&body);
|
||||
let (request_id, response) = client.send(request).await?;
|
||||
let response = client.check_response(&request_id, response).await?;
|
||||
let bytes = response.bytes().await.err_to_http(request_id)?;
|
||||
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None)?;
|
||||
let schema = reader.schema();
|
||||
let mut batches = reader.collect::<std::result::Result<Vec<_>, _>>()?;
|
||||
// A query that matched nothing still describes the event columns.
|
||||
// Keep that schema so callers can build a typed empty result.
|
||||
if batches.is_empty() {
|
||||
batches.push(RecordBatch::new_empty(schema));
|
||||
}
|
||||
Ok(batches)
|
||||
}
|
||||
|
||||
pub struct RemoteJob<S: HttpSend> {
|
||||
@@ -123,7 +158,7 @@ impl<S: HttpSend> RemoteJob<S> {
|
||||
}
|
||||
|
||||
/// One `/v1/jobs/describe` round trip.
|
||||
async fn describe(&self) -> Result<(String, DescribeJobResponse)> {
|
||||
async fn fetch_description(&self) -> Result<(String, DescribeJobResponse)> {
|
||||
let request = self
|
||||
.client
|
||||
.post("/v1/jobs/describe")
|
||||
@@ -148,13 +183,28 @@ impl<S: HttpSend> JobHandle for RemoteJob<S> {
|
||||
}
|
||||
|
||||
async fn status(&self) -> Result<String> {
|
||||
Ok(self.describe().await?.1.state().client_label())
|
||||
Ok(self.fetch_description().await?.1.state().client_label())
|
||||
}
|
||||
|
||||
async fn describe(&self) -> Result<JobDescription> {
|
||||
Ok(self.fetch_description().await?.1.into_description())
|
||||
}
|
||||
|
||||
async fn events(&self, request: JobEventsRequest) -> Result<Vec<RecordBatch>> {
|
||||
let mut body = serde_json::json!({ "job_id": self.job_id });
|
||||
if let Some(limit) = request.limit {
|
||||
body["limit"] = serde_json::Value::from(limit);
|
||||
}
|
||||
if let Some(filter) = request.filter {
|
||||
body["filter"] = serde_json::Value::String(filter);
|
||||
}
|
||||
fetch_job_events(&self.client, body).await
|
||||
}
|
||||
|
||||
async fn wait(&self) -> Result<TerminalResult> {
|
||||
let mut interval = INITIAL_POLL_INTERVAL;
|
||||
loop {
|
||||
let (request_id, description) = self.describe().await?;
|
||||
let (request_id, description) = self.fetch_description().await?;
|
||||
match description.state() {
|
||||
JobState::Done => return Ok(description.into_terminal_result(request_id)),
|
||||
JobState::Failed => {
|
||||
|
||||
@@ -466,7 +466,9 @@ impl TokenSource for AzureImdsSource {
|
||||
/// OAuth header provider that manages the full token lifecycle.
|
||||
///
|
||||
/// Implements [`HeaderProvider`] to inject `Authorization: Bearer <token>`
|
||||
/// headers into every LanceDB request, with automatic token refresh.
|
||||
/// headers into every LanceDB request, with automatic token refresh. It also
|
||||
/// identifies the bearer credential as OIDC so LanceDB's SQL service selects
|
||||
/// OIDC validation instead of API-key validation.
|
||||
pub struct OAuthHeaderProvider {
|
||||
token_source: Box<dyn TokenSource>,
|
||||
token_state: Arc<RwLock<TokenState>>,
|
||||
@@ -554,10 +556,10 @@ impl OAuthHeaderProvider {
|
||||
impl HeaderProvider for OAuthHeaderProvider {
|
||||
async fn get_headers(&self) -> Result<HashMap<String, String>> {
|
||||
let token = self.get_valid_token().await?;
|
||||
Ok(HashMap::from([(
|
||||
"authorization".to_string(),
|
||||
format!("Bearer {token}"),
|
||||
)]))
|
||||
Ok(HashMap::from([
|
||||
("authorization".to_string(), format!("Bearer {token}")),
|
||||
("x-lancedb-credential-type".to_string(), "oidc".to_string()),
|
||||
]))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -40,8 +40,8 @@ use crate::table::{
|
||||
use crate::table::{AnyQuery, Filter, Predicate, PreprocessingOutput, TableStatistics};
|
||||
use crate::utils::background_cache::BackgroundCache;
|
||||
use crate::utils::{
|
||||
resolve_arrow_field_path, resolve_arrow_fts_field_path, supported_btree_data_type,
|
||||
supported_vector_data_type,
|
||||
MaxBatchLengthStream, TimeoutStream, resolve_arrow_field_path, resolve_arrow_fts_field_path,
|
||||
supported_btree_data_type, supported_vector_data_type,
|
||||
};
|
||||
use crate::{DistanceType, Error};
|
||||
use crate::{
|
||||
@@ -72,7 +72,7 @@ use lance_datafusion::exec::{OneShotExec, execute_plan};
|
||||
use reqwest::{RequestBuilder, Response};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Number;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::io::Cursor;
|
||||
use std::pin::Pin;
|
||||
use std::sync::{Arc, Mutex};
|
||||
@@ -264,6 +264,17 @@ impl<S: HttpSend> crate::job::JobHandle for FreshnessJob<S> {
|
||||
crate::job::JobHandle::status(&self.inner).await
|
||||
}
|
||||
|
||||
async fn describe(&self) -> Result<crate::database::JobDescription> {
|
||||
crate::job::JobHandle::describe(&self.inner).await
|
||||
}
|
||||
|
||||
async fn events(
|
||||
&self,
|
||||
request: crate::job::JobEventsRequest,
|
||||
) -> Result<Vec<arrow_array::RecordBatch>> {
|
||||
crate::job::JobHandle::events(&self.inner, request).await
|
||||
}
|
||||
|
||||
async fn wait(&self) -> Result<crate::job::TerminalResult> {
|
||||
let result = crate::job::JobHandle::wait(&self.inner).await?;
|
||||
let version = self.version.read().await;
|
||||
@@ -527,6 +538,10 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
"column": canonical_column
|
||||
});
|
||||
|
||||
if !index.replace {
|
||||
body["replace"] = false.into();
|
||||
}
|
||||
|
||||
// Add name parameter if provided (for backwards compatibility, only include if Some)
|
||||
if let Some(ref name) = index.name {
|
||||
body["name"] = serde_json::Value::String(name.clone());
|
||||
@@ -2022,6 +2037,9 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
fn as_any(&self) -> &dyn std::any::Any {
|
||||
self
|
||||
}
|
||||
fn analyze_plan_is_remote(&self) -> bool {
|
||||
true
|
||||
}
|
||||
fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
@@ -2594,6 +2612,13 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
if let AnyQuery::Query(request) = query
|
||||
&& let Some(offsets) = &request.take_offsets
|
||||
{
|
||||
return crate::query::create_take_offsets_plan(self, request, offsets, options, false)
|
||||
.await;
|
||||
}
|
||||
|
||||
let streams = self.execute_query(query, &options).await?;
|
||||
if streams.len() == 1 {
|
||||
let stream = streams.into_iter().next().unwrap();
|
||||
@@ -2612,6 +2637,27 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<DatasetRecordBatchStream> {
|
||||
if let AnyQuery::Query(request) = query
|
||||
&& let Some(offsets) = &request.take_offsets
|
||||
{
|
||||
let plan = crate::query::create_take_offsets_plan(
|
||||
self,
|
||||
request,
|
||||
offsets,
|
||||
options.clone(),
|
||||
false,
|
||||
)
|
||||
.await?;
|
||||
let inner = execute_plan(plan, Default::default())?;
|
||||
let inner = MaxBatchLengthStream::new_boxed(inner, options.max_batch_length as usize);
|
||||
let inner = if let Some(timeout) = options.timeout {
|
||||
TimeoutStream::new_boxed(inner, timeout)
|
||||
} else {
|
||||
inner
|
||||
};
|
||||
return Ok(DatasetRecordBatchStream::new(inner));
|
||||
}
|
||||
|
||||
let streams = self.execute_query(query, &options).await?;
|
||||
|
||||
if streams.len() == 1 {
|
||||
@@ -2649,6 +2695,12 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
}
|
||||
|
||||
async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result<String> {
|
||||
if let AnyQuery::Query(request) = query
|
||||
&& let Some(offsets) = &request.take_offsets
|
||||
{
|
||||
return crate::query::explain_take_offsets_plan(self, request, offsets, verbose).await;
|
||||
}
|
||||
|
||||
let base_request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/explain_plan/", self.identifier));
|
||||
@@ -2701,6 +2753,17 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<String> {
|
||||
let prepared_query = if let AnyQuery::Query(request) = query
|
||||
&& request.take_offsets.is_some()
|
||||
{
|
||||
Some(AnyQuery::Query(
|
||||
crate::query::prepare_take_offsets_request(self, request).await?,
|
||||
))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let query = prepared_query.as_ref().unwrap_or(query);
|
||||
|
||||
let mut request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/analyze_plan/", self.identifier));
|
||||
@@ -3599,7 +3662,12 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
#[derive(Serialize, Clone, Debug)]
|
||||
pub struct MergeInsertRequest {
|
||||
on: String,
|
||||
// Sent as one repeated `on` query parameter per column, which is how the
|
||||
// namespace spec encodes an array-valued `on`. serde_urlencoded (which
|
||||
// reqwest's `query()` uses) cannot serialize a sequence nested in a struct,
|
||||
// so this field is emitted separately by [`Self::on_query_params`].
|
||||
#[serde(skip_serializing)]
|
||||
on: Vec<String>,
|
||||
when_matched_update_all: bool,
|
||||
when_matched_update_all_filt: Option<String>,
|
||||
when_not_matched_insert_all: bool,
|
||||
@@ -3615,6 +3683,17 @@ pub struct MergeInsertRequest {
|
||||
use_lsm: Option<bool>,
|
||||
}
|
||||
|
||||
impl MergeInsertRequest {
|
||||
/// The `on` columns as repeated query parameters: `?on=a&on=b`.
|
||||
///
|
||||
/// A single column serializes to `?on=a`, exactly what clients sent before
|
||||
/// `on` became a list, so a server that predates composite keys sees no
|
||||
/// change from a single-column caller.
|
||||
pub(crate) fn on_query_params(&self) -> Vec<(&str, &str)> {
|
||||
self.on.iter().map(|col| ("on", col.as_str())).collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn is_true(b: &bool) -> bool {
|
||||
*b
|
||||
}
|
||||
@@ -3627,12 +3706,15 @@ impl TryFrom<MergeInsertBuilder> for MergeInsertRequest {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "MergeInsertBuilder missing required 'on' field".into(),
|
||||
});
|
||||
} else if value.on.len() > 1 {
|
||||
return Err(Error::NotSupported {
|
||||
message: "MergeInsertBuilder only supports a single 'on' column".into(),
|
||||
}
|
||||
// The server rejects a repeated column with a 400; catching it here
|
||||
// names the offending column and costs no round trip.
|
||||
let mut seen = HashSet::with_capacity(value.on.len());
|
||||
if let Some(dup) = value.on.iter().find(|col| !seen.insert(*col)) {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("MergeInsertBuilder 'on' column '{dup}' is repeated"),
|
||||
});
|
||||
}
|
||||
let on = value.on[0].clone();
|
||||
|
||||
let when_matched_update_all_filt = match value.when_matched_update_all_filt {
|
||||
Some(MergeFilter::Sql(sql)) => Some(sql),
|
||||
@@ -3656,7 +3738,7 @@ impl TryFrom<MergeInsertBuilder> for MergeInsertRequest {
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
on,
|
||||
on: value.on,
|
||||
when_matched_update_all: value.when_matched_update_all,
|
||||
when_matched_update_all_filt,
|
||||
when_not_matched_insert_all: value.when_not_matched_insert_all,
|
||||
@@ -3690,7 +3772,7 @@ mod tests {
|
||||
};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::{StreamExt, TryFutureExt, future::BoxFuture};
|
||||
use futures::{StreamExt, TryFutureExt, TryStreamExt, future::BoxFuture};
|
||||
use lance_index::scalar::inverted::{DocumentGranularity, query::MatchQuery};
|
||||
use lance_index::scalar::{FullTextSearchQuery, InvertedIndexParams};
|
||||
use reqwest::Body;
|
||||
@@ -4500,6 +4582,76 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_merge_insert_composite_key() {
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
|
||||
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
|
||||
)
|
||||
.unwrap();
|
||||
let data: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
|
||||
[Ok(batch.clone())],
|
||||
batch.schema(),
|
||||
));
|
||||
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
|
||||
|
||||
// One repeated `on` per column, in the order the caller gave them.
|
||||
let on = request
|
||||
.url()
|
||||
.query_pairs()
|
||||
.filter(|(key, _)| key == "on")
|
||||
.map(|(_, value)| value.into_owned())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(on, vec!["shard_key".to_string(), "id".to_string()]);
|
||||
|
||||
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
|
||||
assert_eq!(params["when_matched_update_all"], "true");
|
||||
assert_eq!(params["when_not_matched_insert_all"], "true");
|
||||
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 43, "num_deleted_rows": 0, "num_inserted_rows": 3, "num_updated_rows": 0}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let mut merge = table.merge_insert(&["shard_key", "id"]);
|
||||
merge.when_matched_update_all(None);
|
||||
merge.when_not_matched_insert_all();
|
||||
let result = table.base_table().merge_insert(merge, data).await.unwrap();
|
||||
|
||||
assert_eq!(result.num_inserted_rows, 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_merge_insert_rejects_repeated_on_column() {
|
||||
let batch = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
|
||||
vec![Arc::new(Int32Array::from(vec![1]))],
|
||||
)
|
||||
.unwrap();
|
||||
let data: Box<dyn RecordBatchReader + Send> = Box::new(RecordBatchIterator::new(
|
||||
[Ok(batch.clone())],
|
||||
batch.schema(),
|
||||
));
|
||||
|
||||
let table = Table::new_with_handler::<&str>("my_table", |request| {
|
||||
panic!("Unexpected request: {}", request.url());
|
||||
});
|
||||
|
||||
let merge = table.merge_insert(&["id", "id"]);
|
||||
let err = table
|
||||
.base_table()
|
||||
.merge_insert(merge, data)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
matches!(&err, Error::InvalidInput { message } if message.contains("'id' is repeated")),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_merge_insert_retries_on_409() {
|
||||
let batch = RecordBatch::try_new(
|
||||
@@ -5611,6 +5763,114 @@ mod tests {
|
||||
assert_eq!(result, "analyzed plan");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_explain_plan_does_not_execute_query() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/explain_plan/");
|
||||
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#""RemoteLookupExec""#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let explained = table
|
||||
.take_offsets(vec![0, 1, 0, 2])
|
||||
.select(crate::query::Select::columns(&["id"]))
|
||||
.limit(3)
|
||||
.explain_plan(false)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(explained.contains("GlobalLimitExec"));
|
||||
assert!(explained.contains("TakeRestoreExec"));
|
||||
assert!(!explained.contains("CoalescePartitionsExec"));
|
||||
assert!(explained.contains("RemoteLookupExec"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_converted_take_request_restores_remote_occurrences() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/query/");
|
||||
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["columns"], json!(["id", "_rowoffset"]));
|
||||
|
||||
let data = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int32, false),
|
||||
Field::new("_rowoffset", DataType::UInt64, false),
|
||||
])),
|
||||
vec![
|
||||
Arc::new(Int32Array::from(vec![5])),
|
||||
Arc::new(arrow_array::UInt64Array::from(vec![5])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE)
|
||||
.body(write_ipc_file(&data))
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let request = table
|
||||
.take_offsets(vec![5, 5])
|
||||
.select(crate::query::Select::columns(&["id"]))
|
||||
.into_request();
|
||||
let batches = table
|
||||
.base_table()
|
||||
.query(&AnyQuery::Query(request), QueryExecutionOptions::default())
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
|
||||
assert!(
|
||||
batches
|
||||
.iter()
|
||||
.all(|batch| batch.schema().fields().len() == 1)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_take_offsets_analyze_plan_delegates_to_remote() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/analyze_plan/");
|
||||
assert_eq!(
|
||||
request
|
||||
.url()
|
||||
.query_pairs()
|
||||
.find(|(key, _)| key == "distributed_metrics"),
|
||||
Some(("distributed_metrics".into(), "per_worker".into()))
|
||||
);
|
||||
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#""Remote analyzed plan: worker metrics""#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let analyzed = table
|
||||
.take_offsets(vec![0, 1, 0, 2])
|
||||
.select(crate::query::Select::columns(&["id"]))
|
||||
.limit(3)
|
||||
.analyze_plan_with_options(QueryExecutionOptions {
|
||||
analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::PerWorker,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(analyzed, "Remote analyzed plan: worker metrics");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_structured_fts() {
|
||||
let table =
|
||||
@@ -6077,6 +6337,40 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_forwards_replace_false_on_existing_route() {
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => {
|
||||
let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&schema))
|
||||
.unwrap()
|
||||
}
|
||||
"/v1/table/my_table/create_index/" => {
|
||||
let body = request.body().unwrap().as_bytes().unwrap();
|
||||
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
|
||||
assert_eq!(body["replace"], json!(false));
|
||||
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body("{}".to_string())
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("Unexpected path: {}", path),
|
||||
}
|
||||
});
|
||||
|
||||
table
|
||||
.create_index(&["a"], Index::BTree(Default::default()))
|
||||
.replace(false)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_returns_job() {
|
||||
let describe_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
@@ -7699,6 +7993,72 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// The refresh handle is wrapped for read-freshness tracking, so it has to
|
||||
/// forward the detail APIs too -- this is the job an operator is holding
|
||||
/// when a backfill goes quiet.
|
||||
#[tokio::test]
|
||||
async fn test_refresh_job_handle_reports_detail_and_events() {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
"state",
|
||||
DataType::Utf8,
|
||||
false,
|
||||
)]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![Arc::new(arrow_array::StringArray::from(vec![
|
||||
"claim_complete",
|
||||
]))],
|
||||
)
|
||||
.unwrap();
|
||||
let mut events = Vec::new();
|
||||
{
|
||||
let mut writer =
|
||||
arrow_ipc::writer::StreamWriter::try_new(&mut events, &schema).unwrap();
|
||||
writer.write(&batch).unwrap();
|
||||
writer.finish().unwrap();
|
||||
}
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
match request.url().path() {
|
||||
"/v1/table/my_table/backfill_column" => http::Response::builder()
|
||||
.status(202)
|
||||
.body(br#"{"job_id": "j-42"}"#.to_vec())
|
||||
.unwrap(),
|
||||
"/v1/jobs/describe" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"job_id": "j-42", "job_type": "refresh_column", "job_state": "IN_PROGRESS", "creation_ms": 7, "spec": {"column": "doubled"}}"#
|
||||
.as_bytes()
|
||||
.to_vec(),
|
||||
)
|
||||
.unwrap(),
|
||||
"/v1/jobs/query_events" => {
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(body["job_id"], "j-42");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(events.clone())
|
||||
.unwrap()
|
||||
}
|
||||
other => panic!("unexpected path {other}"),
|
||||
}
|
||||
});
|
||||
|
||||
let job = table.refresh_column_async("doubled").await.unwrap();
|
||||
job.refresh().await.unwrap();
|
||||
assert_eq!(job.state().as_deref(), Some("running"));
|
||||
assert_eq!(job.job_type().as_deref(), Some("refresh_column"));
|
||||
assert_eq!(job.creation_ms(), Some(7));
|
||||
assert_eq!(job.spec().unwrap()["column"], "doubled");
|
||||
|
||||
let batches = job
|
||||
.events(crate::job::JobEventsRequest::default())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_submission_uses_add_columns_version_fence() {
|
||||
let table = Table::new_with_handler("my_table", |request| match request.url().path() {
|
||||
|
||||
@@ -734,6 +734,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
|
||||
WriteOp::MergeInsert { query, timeout } => {
|
||||
let mut request = client
|
||||
.post(&format!("/v1/table/{}/merge_insert/", identifier))
|
||||
.query(&query.on_query_params())
|
||||
.query(query)
|
||||
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE);
|
||||
if let Some(timeout) = timeout {
|
||||
@@ -1489,7 +1490,7 @@ mod tests {
|
||||
});
|
||||
|
||||
let query = MergeInsertRequest {
|
||||
on: "id".to_string(),
|
||||
on: vec!["id".to_string()],
|
||||
when_matched_update_all: false,
|
||||
when_matched_update_all_filt: None,
|
||||
when_not_matched_insert_all: false,
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Handles to SQL queries running on a remote database.
|
||||
|
||||
use std::{fmt, sync::Arc};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{Result, arrow::SendableRecordBatchStream};
|
||||
|
||||
/// The externally visible lifecycle state of a submitted SQL query.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum QueryStatus {
|
||||
/// The server is still executing the query.
|
||||
Running,
|
||||
/// The server has made the complete result available.
|
||||
Finished,
|
||||
/// The server accepted cancellation but has not confirmed it yet.
|
||||
Cancelling,
|
||||
/// The server confirmed cancellation.
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl fmt::Display for QueryStatus {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str(match self {
|
||||
Self::Running => "running",
|
||||
Self::Finished => "finished",
|
||||
Self::Cancelling => "cancelling",
|
||||
Self::Cancelled => "cancelled",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A point-in-time description of a submitted SQL query.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct QueryDescription {
|
||||
/// The stable, connection-scoped identifier assigned when the query was submitted.
|
||||
pub id: Uuid,
|
||||
/// The server-visible lifecycle state.
|
||||
pub status: QueryStatus,
|
||||
/// Server-reported completion progress, when known. Values are in `[0.0, 1.0]`,
|
||||
/// with `1.0` meaning complete.
|
||||
pub progress: Option<f64>,
|
||||
/// When the server may stop accepting this query's continuation token.
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub(crate) trait QueryHandle: Send + Sync {
|
||||
fn id(&self) -> Uuid;
|
||||
async fn describe(&self) -> Result<QueryDescription>;
|
||||
async fn reader(&self) -> Result<SendableRecordBatchStream>;
|
||||
async fn cancel(&self) -> Result<()>;
|
||||
}
|
||||
|
||||
/// A handle to a submitted SQL query.
|
||||
///
|
||||
/// The handle can be inspected, opened as an Arrow reader, or cancelled.
|
||||
/// Dropping it does not cancel the server-side query.
|
||||
/// Identifier lookup is scoped to the connection that submitted the query and
|
||||
/// is not a durable resume mechanism.
|
||||
pub struct Query {
|
||||
handle: Arc<dyn QueryHandle>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Query {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("Query")
|
||||
.field("id", &self.id())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Query {
|
||||
#[cfg(feature = "remote")]
|
||||
pub(crate) fn new(handle: Arc<dyn QueryHandle>) -> Self {
|
||||
Self { handle }
|
||||
}
|
||||
|
||||
/// Return the stable, connection-scoped identifier for this query.
|
||||
pub fn id(&self) -> Uuid {
|
||||
self.handle.id()
|
||||
}
|
||||
|
||||
/// Get a point-in-time description of the query.
|
||||
pub async fn describe(&self) -> Result<QueryDescription> {
|
||||
self.handle.describe().await
|
||||
}
|
||||
|
||||
/// Wait for the initial result stream and return its Arrow record batches.
|
||||
///
|
||||
/// The stream can begin yielding partial results before query execution is
|
||||
/// complete. It continues polling for newly available result endpoints
|
||||
/// until the query finishes and all endpoints have been consumed.
|
||||
///
|
||||
/// Results are single-consumer. Calling this method more than once on the
|
||||
/// same handle returns an error.
|
||||
pub async fn reader(&self) -> Result<SendableRecordBatchStream> {
|
||||
self.handle.reader().await
|
||||
}
|
||||
|
||||
/// Request cancellation of the query.
|
||||
pub async fn cancel(&self) -> Result<()> {
|
||||
self.handle.cancel().await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::QueryStatus;
|
||||
|
||||
#[test]
|
||||
fn query_status_display_is_stable() {
|
||||
assert_eq!(QueryStatus::Running.to_string(), "running");
|
||||
assert_eq!(QueryStatus::Finished.to_string(), "finished");
|
||||
assert_eq!(QueryStatus::Cancelling.to_string(), "cancelling");
|
||||
assert_eq!(QueryStatus::Cancelled.to_string(), "cancelled");
|
||||
}
|
||||
}
|
||||
@@ -595,6 +595,14 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
query: &AnyQuery,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<String>;
|
||||
/// Whether [`BaseTable::analyze_plan`] is provided by a remote service.
|
||||
///
|
||||
/// Client-side query wrappers use this to preserve backend metrics and
|
||||
/// distributed-analysis options instead of replacing them with a local plan.
|
||||
#[doc(hidden)]
|
||||
fn analyze_plan_is_remote(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Add new records to the table.
|
||||
async fn add(&self, add: AddDataBuilder) -> Result<AddResult>;
|
||||
@@ -1184,6 +1192,9 @@ impl Table {
|
||||
/// valid empty blobs contain empty byte strings. Prefer
|
||||
/// [`Self::fetch_blob_files`] for large selections.
|
||||
///
|
||||
/// `_rowid` values stay valid after compaction when the table has stable
|
||||
/// row ids.
|
||||
///
|
||||
/// ```
|
||||
/// use arrow_array::UInt64Array;
|
||||
/// use futures::TryStreamExt;
|
||||
@@ -1225,6 +1236,9 @@ impl Table {
|
||||
/// the requests. Null blobs produce null output slots; empty ranges on
|
||||
/// non-null blobs produce empty byte strings.
|
||||
///
|
||||
/// `_rowid` values stay valid after compaction when the table has stable
|
||||
/// row ids.
|
||||
///
|
||||
/// ```
|
||||
/// use lancedb::blob::BlobRangeRequest;
|
||||
///
|
||||
@@ -1263,6 +1277,9 @@ impl Table {
|
||||
/// Same length and order as `row_ids`. Null rows are `None`. Bytes are not
|
||||
/// read from disk until a call to [`BlobFile::read`].
|
||||
///
|
||||
/// `_rowid` values stay valid after compaction when the table has stable
|
||||
/// row ids.
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::Table;
|
||||
/// # async fn lazy_read(table: &Table, row_ids: &[u64]) -> Result<(), Box<dyn std::error::Error>> {
|
||||
@@ -1498,7 +1515,9 @@ impl Table {
|
||||
///
|
||||
/// * `on` One or more columns to join on. This is how records from the
|
||||
/// source table and target table are matched. Typically this is some
|
||||
/// kind of key or id column.
|
||||
/// kind of key or id column. Several columns match on the composite
|
||||
/// key: a source row updates a target row only when it agrees on every
|
||||
/// one of them.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
@@ -1652,9 +1671,9 @@ impl Table {
|
||||
/// Offsets are useful for sampling as the set of all valid offsets is easily
|
||||
/// known in advance to be [0, len(table)).
|
||||
///
|
||||
/// No guarantees are made regarding the order in which results are returned. If you
|
||||
/// desire an output order that matches the order of the given offsets, you will need
|
||||
/// to add the row offset column to the output and align it yourself.
|
||||
/// No guarantees are made regarding the order in which results are returned.
|
||||
/// Repeated offsets produce repeated rows, which makes this method suitable for
|
||||
/// sampling with replacement.
|
||||
///
|
||||
/// Parameters
|
||||
/// ----------
|
||||
|
||||
@@ -60,6 +60,10 @@ pub const FUNCTION_BINDING_ID_META_KEY: &str = "computed_column.function.binding
|
||||
/// Field metadata key holding this sibling's ordered Function output ordinal.
|
||||
pub const FUNCTION_OUTPUT_ORDINAL_META_KEY: &str = "computed_column.function.output_ordinal";
|
||||
|
||||
/// Reserved Function output ordinal for an internal flattened-result
|
||||
/// assignment column.
|
||||
pub const FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL: u32 = u32::MAX;
|
||||
|
||||
/// Schema metadata key holding all immutable Function bindings.
|
||||
pub const FUNCTION_BINDINGS_META_KEY: &str = "lancedb::function_bindings";
|
||||
|
||||
@@ -312,22 +316,29 @@ pub(crate) fn ensure_supported_function_metadata(schema: &ArrowSchema) -> Result
|
||||
binding_id
|
||||
),
|
||||
})?;
|
||||
let output = binding
|
||||
.outputs()
|
||||
.get(output_ordinal as usize)
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: format!(
|
||||
"Function output '{}' has invalid ordinal {}",
|
||||
field.name(),
|
||||
output_ordinal
|
||||
),
|
||||
})?;
|
||||
if output.output_name != field.name().as_str() {
|
||||
let destination = if output_ordinal == FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL {
|
||||
binding
|
||||
.assignment()
|
||||
.map(|assignment| assignment.output_name.as_str())
|
||||
} else {
|
||||
binding
|
||||
.outputs()
|
||||
.get(output_ordinal as usize)
|
||||
.map(|output| output.output_name.as_str())
|
||||
}
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: format!(
|
||||
"Function output '{}' has invalid ordinal {}",
|
||||
field.name(),
|
||||
output_ordinal
|
||||
),
|
||||
})?;
|
||||
if destination != field.name().as_str() {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"Function output '{}' does not match binding destination '{}'",
|
||||
field.name(),
|
||||
output.output_name
|
||||
destination
|
||||
),
|
||||
});
|
||||
}
|
||||
@@ -498,6 +509,7 @@ fn ensure_known_binding_shape(value: &Value) -> Result<()> {
|
||||
"function",
|
||||
"inputs",
|
||||
"outputs",
|
||||
"assignment",
|
||||
"input_schema",
|
||||
"output_schema",
|
||||
],
|
||||
@@ -546,6 +558,13 @@ fn ensure_known_binding_shape(value: &Value) -> Result<()> {
|
||||
"output mapping",
|
||||
)?;
|
||||
}
|
||||
if let Some(assignment) = object.get("assignment") {
|
||||
reject_unknown_object_fields(
|
||||
assignment,
|
||||
&["output_name", "output_field_id"],
|
||||
"assignment mapping",
|
||||
)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -589,16 +608,13 @@ fn canonical_input_arrow_type(field: &JsonArrowField) -> Result<String> {
|
||||
.and_then(|metadata| metadata.get(ARROW_EXT_NAME_KEY))
|
||||
.map(String::as_str)
|
||||
== Some(BLOB_V2_EXT_NAME);
|
||||
if is_blob_v2 {
|
||||
if is_blob_v2 || field.r#type.fields.is_some() {
|
||||
let arrow_field = lance_namespace::schema::convert_json_arrow_field(field)
|
||||
.map_err(|e| invalid_function(format!("invalid Function input field: {e}")))?;
|
||||
if !has_supported_blob_v2_layout(&arrow_field) {
|
||||
return Err(invalid_function(format!(
|
||||
"Function input '{}' has an invalid Blob v2 storage layout",
|
||||
arrow_field.name()
|
||||
)));
|
||||
validate_function_blob_nesting(&arrow_field, false)?;
|
||||
if is_blob_v2 {
|
||||
return Ok(FUNCTION_BLOB_V2_TYPE.to_string());
|
||||
}
|
||||
return Ok(FUNCTION_BLOB_V2_TYPE.to_string());
|
||||
}
|
||||
if field.r#type.fields.is_none() && field.r#type.length.is_none() {
|
||||
Ok(field.r#type.r#type.clone())
|
||||
@@ -617,6 +633,34 @@ fn has_supported_blob_v2_layout(field: &ArrowField) -> bool {
|
||||
)
|
||||
}
|
||||
|
||||
fn validate_function_blob_nesting(field: &ArrowField, inside_collection: bool) -> Result<()> {
|
||||
if field.is_blob_v2() {
|
||||
if inside_collection {
|
||||
return Err(invalid_function(format!(
|
||||
"Function field '{}' nests Blob v2 under a collection, which Function signatures do not support",
|
||||
field.name()
|
||||
)));
|
||||
}
|
||||
if !has_supported_blob_v2_layout(field) {
|
||||
return Err(invalid_function(format!(
|
||||
"Function field '{}' has an invalid Blob v2 storage layout",
|
||||
field.name()
|
||||
)));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
match field.data_type() {
|
||||
DataType::Struct(fields) => fields
|
||||
.iter()
|
||||
.try_for_each(|field| validate_function_blob_nesting(field, inside_collection)),
|
||||
DataType::List(field)
|
||||
| DataType::LargeList(field)
|
||||
| DataType::FixedSizeList(field, _)
|
||||
| DataType::Map(field, _) => validate_function_blob_nesting(field, true),
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
/// `fixed_size_list<item, size>` -> (`item`, `size`); the comma must sit outside
|
||||
/// any nested `<...>`.
|
||||
fn split_fixed_size_list(raw: &str) -> Option<(&str, i32)> {
|
||||
@@ -697,31 +741,51 @@ fn parse_output_arrow_type(raw: &str) -> Result<JsonArrowDataType> {
|
||||
}
|
||||
|
||||
fn function_output_field(name: &str, nullable: bool, raw: &str) -> Result<JsonArrowField> {
|
||||
if raw == FUNCTION_BLOB_V2_TYPE {
|
||||
return lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![
|
||||
crate::blob(name, nullable),
|
||||
]))
|
||||
let field = if raw == FUNCTION_BLOB_V2_TYPE {
|
||||
lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![crate::blob(
|
||||
name, nullable,
|
||||
)]))
|
||||
.map_err(|e| invalid_function(format!("could not encode Blob v2 output field: {e}")))?
|
||||
.fields
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| invalid_function("Blob v2 output field is missing"));
|
||||
}
|
||||
Ok(JsonArrowField::new(
|
||||
name.to_string(),
|
||||
nullable,
|
||||
parse_output_arrow_type(raw)?,
|
||||
))
|
||||
.ok_or_else(|| invalid_function("Blob v2 output field is missing"))?
|
||||
} else {
|
||||
JsonArrowField::new(name.to_string(), nullable, parse_output_arrow_type(raw)?)
|
||||
};
|
||||
let arrow_field = lance_namespace::schema::convert_json_arrow_field(&field)
|
||||
.map_err(|e| invalid_function(format!("invalid Function output field: {e}")))?;
|
||||
validate_function_blob_nesting(&arrow_field, false)?;
|
||||
Ok(field)
|
||||
}
|
||||
|
||||
fn function_output_field_matches(expected: &ArrowField, actual: &ArrowField) -> bool {
|
||||
expected.name() == actual.name()
|
||||
&& expected.is_nullable() == actual.is_nullable()
|
||||
&& if expected.is_blob_v2() {
|
||||
/// Whether two fields describe the same Function output.
|
||||
///
|
||||
/// `compare_identity` covers the field's own name and nullability. Struct
|
||||
/// children carry both as part of the declaration and compare with it on. List
|
||||
/// children do not: Lance rewrites a list item's name and nullability when it
|
||||
/// writes, so a stored `fixed_size_list<item: float not null>` comes back as
|
||||
/// `fixed_size_list<item: float>` and never matches the declaration again.
|
||||
/// Comparing those by type alone keeps this agreeing with the server, which
|
||||
/// draws the same distinction and is what accepted the column when it was
|
||||
/// declared.
|
||||
fn function_output_field_matches(
|
||||
expected: &ArrowField,
|
||||
actual: &ArrowField,
|
||||
compare_identity: bool,
|
||||
) -> bool {
|
||||
if compare_identity
|
||||
&& (expected.name() != actual.name() || expected.is_nullable() != actual.is_nullable())
|
||||
{
|
||||
return false;
|
||||
}
|
||||
match (expected.is_blob_v2(), actual.is_blob_v2()) {
|
||||
(false, false) => function_output_type_matches(expected.data_type(), actual.data_type()),
|
||||
(true, true) => {
|
||||
has_supported_blob_v2_layout(expected) && has_supported_blob_v2_layout(actual)
|
||||
} else {
|
||||
function_output_type_matches(expected.data_type(), actual.data_type())
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn function_output_type_matches(expected: &DataType, actual: &DataType) -> bool {
|
||||
@@ -734,33 +798,19 @@ fn function_output_type_matches(expected: &DataType, actual: &DataType) -> bool
|
||||
&& expected
|
||||
.iter()
|
||||
.zip(actual)
|
||||
.all(|(expected, actual)| function_output_field_matches(expected, actual))
|
||||
.all(|(expected, actual)| function_output_field_matches(expected, actual, true))
|
||||
}
|
||||
(DataType::List(expected), DataType::List(actual))
|
||||
| (DataType::LargeList(expected), DataType::LargeList(actual)) => {
|
||||
function_output_field_matches(expected, actual)
|
||||
function_output_field_matches(expected, actual, false)
|
||||
}
|
||||
(
|
||||
DataType::FixedSizeList(expected, expected_size),
|
||||
DataType::FixedSizeList(actual, actual_size),
|
||||
) => expected_size == actual_size && function_output_field_matches(expected, actual),
|
||||
) => expected_size == actual_size && function_output_field_matches(expected, actual, false),
|
||||
(DataType::Map(expected, expected_sorted), DataType::Map(actual, actual_sorted)) => {
|
||||
expected_sorted == actual_sorted && function_output_field_matches(expected, actual)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn function_output_type_has_blob(data_type: &DataType) -> bool {
|
||||
match data_type {
|
||||
DataType::Struct(fields) => fields
|
||||
.iter()
|
||||
.any(|field| field.is_blob_v2() || function_output_type_has_blob(field.data_type())),
|
||||
DataType::List(field)
|
||||
| DataType::LargeList(field)
|
||||
| DataType::FixedSizeList(field, _)
|
||||
| DataType::Map(field, _) => {
|
||||
field.is_blob_v2() || function_output_type_has_blob(field.data_type())
|
||||
expected_sorted == actual_sorted
|
||||
&& function_output_field_matches(expected, actual, true)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
@@ -839,23 +889,20 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding
|
||||
output.output_name
|
||||
))
|
||||
})?;
|
||||
if field.name() != &output.output_name || !field.is_nullable() || output.nullable {
|
||||
if field.name() != &output.output_name || !field.is_nullable() {
|
||||
return Err(invalid_function(format!(
|
||||
"Function output '{}' no longer matches binding '{}'",
|
||||
output.output_name,
|
||||
binding.binding_id()
|
||||
)));
|
||||
}
|
||||
let (type_matches, has_semantic_blob) = if output.arrow_type == FUNCTION_BLOB_V2_TYPE {
|
||||
(has_supported_blob_v2_layout(field), true)
|
||||
let type_matches = if output.arrow_type == FUNCTION_BLOB_V2_TYPE {
|
||||
has_supported_blob_v2_layout(field)
|
||||
} else {
|
||||
let expected_type = parse_output_arrow_type(&output.arrow_type)?;
|
||||
let expected_type = lance_namespace::schema::convert_json_arrow_type(&expected_type)
|
||||
.map_err(|e| invalid_function(format!("invalid Function output type: {e}")))?;
|
||||
(
|
||||
function_output_type_matches(&expected_type, field.data_type()),
|
||||
function_output_type_has_blob(&expected_type),
|
||||
)
|
||||
function_output_type_matches(&expected_type, field.data_type())
|
||||
};
|
||||
if !type_matches {
|
||||
return Err(invalid_function(format!(
|
||||
@@ -886,19 +933,71 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding
|
||||
binding.binding_id()
|
||||
)));
|
||||
}
|
||||
if has_semantic_blob {
|
||||
output_fields.push(function_output_field(
|
||||
field.name(),
|
||||
true,
|
||||
&output.arrow_type,
|
||||
)?);
|
||||
} else {
|
||||
let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![
|
||||
ArrowField::new(field.name().clone(), field.data_type().clone(), true),
|
||||
]))
|
||||
.map_err(|e| invalid_function(format!("invalid Function output schema: {e}")))?;
|
||||
output_fields.push(json.fields.into_iter().next().unwrap());
|
||||
// Rebuild from the declaration rather than from the stored field. The
|
||||
// stored field carries Lance's write-time normalization, which would
|
||||
// never round-trip back to the schema the binding recorded -- the same
|
||||
// reason list children compare by type above. Whether the column on
|
||||
// disk still matches is settled by that comparison, not here.
|
||||
output_fields.push(function_output_field(
|
||||
field.name(),
|
||||
true,
|
||||
&output.arrow_type,
|
||||
)?);
|
||||
}
|
||||
if let Some(assignment) = binding.assignment() {
|
||||
if binding
|
||||
.outputs()
|
||||
.iter()
|
||||
.any(|output| output.result_field == WHOLE_RESULT_FIELD)
|
||||
{
|
||||
return Err(invalid_function(format!(
|
||||
"Function binding '{}' cannot attach an assignment column to a whole result",
|
||||
binding.binding_id()
|
||||
)));
|
||||
}
|
||||
let field = schema
|
||||
.field_with_name(&assignment.output_name)
|
||||
.map_err(|_| {
|
||||
invalid_function(format!(
|
||||
"Function binding '{}' assignment column '{}' is missing",
|
||||
binding.binding_id(),
|
||||
assignment.output_name
|
||||
))
|
||||
})?;
|
||||
let metadata = field.metadata();
|
||||
if field.data_type() != &DataType::Boolean
|
||||
|| !field.is_nullable()
|
||||
|| metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str) != Some("true")
|
||||
|| metadata.get(KIND_META_KEY).map(String::as_str) != Some(FUNCTION_KIND)
|
||||
|| metadata
|
||||
.get(FUNCTION_BINDING_ID_META_KEY)
|
||||
.map(String::as_str)
|
||||
!= Some(binding.binding_id())
|
||||
|| metadata
|
||||
.get(FUNCTION_OUTPUT_ORDINAL_META_KEY)
|
||||
.and_then(|value| value.parse::<u32>().ok())
|
||||
!= Some(FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL)
|
||||
{
|
||||
return Err(invalid_function(format!(
|
||||
"Function binding '{}' assignment column no longer matches its declaration",
|
||||
binding.binding_id()
|
||||
)));
|
||||
}
|
||||
let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![
|
||||
ArrowField::new(assignment.output_name.clone(), DataType::Boolean, true),
|
||||
]))
|
||||
.map_err(|e| invalid_function(format!("invalid Function assignment schema: {e}")))?;
|
||||
output_fields.push(json.fields.into_iter().next().unwrap());
|
||||
} else if binding.outputs().iter().all(|output| output.nullable)
|
||||
&& binding
|
||||
.outputs()
|
||||
.iter()
|
||||
.all(|output| output.result_field != WHOLE_RESULT_FIELD)
|
||||
{
|
||||
return Err(invalid_function(format!(
|
||||
"Function binding '{}' has no flattened-result assignment column",
|
||||
binding.binding_id()
|
||||
)));
|
||||
}
|
||||
let output_schema = JsonArrowSchema::new(output_fields);
|
||||
let output_schema = serde_json::to_value(output_schema).map_err(|e| {
|
||||
@@ -1055,11 +1154,6 @@ pub(crate) fn plan_function_application(
|
||||
"named-struct Function result field names must be unique",
|
||||
));
|
||||
}
|
||||
if output.fields.iter().any(|field| field.nullable) {
|
||||
return Err(invalid_function(
|
||||
"Function logical outputs must be non-nullable during NULL assignment",
|
||||
));
|
||||
}
|
||||
let unknown = application
|
||||
.columns()
|
||||
.keys()
|
||||
@@ -1081,7 +1175,9 @@ pub(crate) fn plan_function_application(
|
||||
let fields = output
|
||||
.fields
|
||||
.iter()
|
||||
.map(|field| function_output_field(&field.name, false, &field.arrow_type))
|
||||
.map(|field| {
|
||||
function_output_field(&field.name, field.nullable, &field.arrow_type)
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let mut data_type = JsonArrowDataType::new("struct".to_string());
|
||||
data_type.fields = Some(fields);
|
||||
@@ -1718,6 +1814,69 @@ mod tests {
|
||||
assert!(super::validate_declarations(schema, &declarations).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_children_match_by_type_but_struct_children_by_identity() {
|
||||
use arrow_schema::Field as F;
|
||||
|
||||
// Lance rewrites a list item's name and nullability on write, so the
|
||||
// stored field is no longer identical to what was declared. Comparing
|
||||
// those by type keeps a table with a vector output usable.
|
||||
let declared =
|
||||
DataType::FixedSizeList(Arc::new(F::new("item", DataType::Float32, false)), 4);
|
||||
let stored = DataType::FixedSizeList(Arc::new(F::new("item", DataType::Float32, true)), 4);
|
||||
assert!(super::function_output_type_matches(&declared, &stored));
|
||||
|
||||
let renamed =
|
||||
DataType::FixedSizeList(Arc::new(F::new("element", DataType::Float32, true)), 4);
|
||||
assert!(super::function_output_type_matches(&declared, &renamed));
|
||||
|
||||
// The dimension is still part of the declaration.
|
||||
let resized = DataType::FixedSizeList(Arc::new(F::new("item", DataType::Float32, true)), 8);
|
||||
assert!(!super::function_output_type_matches(&declared, &resized));
|
||||
|
||||
// Struct children keep comparing by name and nullability.
|
||||
let struct_declared =
|
||||
DataType::Struct(vec![F::new("changed", DataType::Boolean, false)].into());
|
||||
let struct_nullable =
|
||||
DataType::Struct(vec![F::new("changed", DataType::Boolean, true)].into());
|
||||
let struct_renamed =
|
||||
DataType::Struct(vec![F::new("altered", DataType::Boolean, false)].into());
|
||||
assert!(super::function_output_type_matches(
|
||||
&struct_declared,
|
||||
&struct_declared
|
||||
));
|
||||
assert!(!super::function_output_type_matches(
|
||||
&struct_declared,
|
||||
&struct_nullable
|
||||
));
|
||||
assert!(!super::function_output_type_matches(
|
||||
&struct_declared,
|
||||
&struct_renamed
|
||||
));
|
||||
|
||||
// A list nested inside a struct gets the list rule.
|
||||
let nested_declared = DataType::Struct(
|
||||
vec![F::new(
|
||||
"tokens",
|
||||
DataType::List(Arc::new(F::new("item", DataType::Utf8, false))),
|
||||
true,
|
||||
)]
|
||||
.into(),
|
||||
);
|
||||
let nested_stored = DataType::Struct(
|
||||
vec![F::new(
|
||||
"tokens",
|
||||
DataType::List(Arc::new(F::new("item", DataType::Utf8, true))),
|
||||
true,
|
||||
)]
|
||||
.into(),
|
||||
);
|
||||
assert!(super::function_output_type_matches(
|
||||
&nested_declared,
|
||||
&nested_stored
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn output_arrow_type_grammar_matches_the_shared_golden() {
|
||||
let golden: serde_json::Value = serde_json::from_str(include_str!(
|
||||
@@ -2719,6 +2878,28 @@ mod tests {
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn exact_arrow_type(field: ArrowField) -> String {
|
||||
let json =
|
||||
lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![field])).unwrap();
|
||||
serde_json::to_string(json.fields[0].r#type.as_ref()).unwrap()
|
||||
}
|
||||
|
||||
fn single_input_application(path: &str) -> FunctionApplication {
|
||||
FunctionApplication::from_json(
|
||||
&serde_json::json!({
|
||||
"function": {"name": "inspect", "version": "fv_nested_blob"},
|
||||
"inputs": [{
|
||||
"parameter": "value",
|
||||
"kind": "column",
|
||||
"value": {"path": path}
|
||||
}],
|
||||
"output": {"kind": "scalar", "arrow_type": "int64", "nullable": false}
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn binding_from_plan(plan: &FunctionDeclarationPlan) -> FunctionBinding {
|
||||
let inputs = plan
|
||||
.input_bindings
|
||||
@@ -2810,6 +2991,17 @@ mod tests {
|
||||
&inputs,
|
||||
));
|
||||
}
|
||||
if let Some(assignment) = binding.assignment() {
|
||||
fields.push(
|
||||
ArrowField::new(&assignment.output_name, DataType::Boolean, true).with_metadata(
|
||||
function_computed_column_metadata(
|
||||
binding.binding_id(),
|
||||
FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL,
|
||||
&inputs,
|
||||
),
|
||||
),
|
||||
);
|
||||
}
|
||||
ArrowSchema::new(fields)
|
||||
}
|
||||
|
||||
@@ -2827,6 +3019,74 @@ mod tests {
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_binding_preserves_all_nullable_outputs_with_an_assignment_column() {
|
||||
let mut raw_binding: Value = serde_json::from_str(include_str!(
|
||||
"../../tests/fixtures/first_class_functions/v1/remote_function_binding.json"
|
||||
))
|
||||
.unwrap();
|
||||
raw_binding["outputs"][0]["nullable"] = Value::Bool(true);
|
||||
raw_binding["outputs"][1]["nullable"] = Value::Bool(true);
|
||||
let without_assignment: FunctionBinding =
|
||||
serde_json::from_value(raw_binding.clone()).unwrap();
|
||||
let error = ensure_binding_matches_schema(
|
||||
&valid_function_binding_schema(true, true, &without_assignment),
|
||||
&without_assignment,
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
error
|
||||
.to_string()
|
||||
.contains("flattened-result assignment column")
|
||||
);
|
||||
|
||||
raw_binding["assignment"] = serde_json::json!({
|
||||
"output_name": "__function_assignment_fb_01K3TEXT",
|
||||
"output_field_id": -1,
|
||||
});
|
||||
raw_binding["output_schema"]["fields"]
|
||||
.as_array_mut()
|
||||
.unwrap()
|
||||
.push(serde_json::json!({
|
||||
"name": "__function_assignment_fb_01K3TEXT",
|
||||
"nullable": true,
|
||||
"type": {"type": "bool"},
|
||||
}));
|
||||
let binding: FunctionBinding = serde_json::from_value(raw_binding).unwrap();
|
||||
ensure_binding_matches_schema(
|
||||
&valid_function_binding_schema(true, true, &binding),
|
||||
&binding,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let schema = ArrowSchema::new_with_metadata(
|
||||
valid_function_binding_schema(true, true, &binding)
|
||||
.fields()
|
||||
.to_vec(),
|
||||
HashMap::from([(
|
||||
FUNCTION_BINDINGS_META_KEY.to_string(),
|
||||
function_bindings_metadata(std::slice::from_ref(&binding)).unwrap(),
|
||||
)]),
|
||||
);
|
||||
ensure_supported_function_metadata(&schema).unwrap();
|
||||
|
||||
let mut metadata: Value =
|
||||
serde_json::from_str(schema.metadata().get(FUNCTION_BINDINGS_META_KEY).unwrap())
|
||||
.unwrap();
|
||||
metadata["bindings"][0]["assignment"]["future"] = Value::Bool(true);
|
||||
let future_schema = ArrowSchema::new_with_metadata(
|
||||
schema.fields().to_vec(),
|
||||
HashMap::from([(
|
||||
FUNCTION_BINDINGS_META_KEY.to_string(),
|
||||
serde_json::to_string(&metadata).unwrap(),
|
||||
)]),
|
||||
);
|
||||
assert!(matches!(
|
||||
ensure_supported_function_metadata(&future_schema),
|
||||
Err(Error::NotSupported { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nullable_function_input_cannot_bind_to_non_nullable_parameter() {
|
||||
let mut raw_binding: Value = serde_json::from_str(include_str!(
|
||||
@@ -3046,6 +3306,35 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_named_struct_plan_preserves_nullable_result_fields() {
|
||||
let mut value = serde_json::to_value(named_struct_application("{}")).unwrap();
|
||||
value["output"]["fields"][0]["nullable"] = Value::Bool(true);
|
||||
value["output"]["fields"][1]["nullable"] = Value::Bool(true);
|
||||
let application = FunctionApplication::from_json(&value.to_string()).unwrap();
|
||||
|
||||
let expanded =
|
||||
plan_function_application(&function_input_schema(), &application, None).unwrap();
|
||||
assert!(
|
||||
expanded
|
||||
.output_schema
|
||||
.fields
|
||||
.iter()
|
||||
.all(|field| field.nullable)
|
||||
);
|
||||
|
||||
let whole =
|
||||
plan_function_application(&function_input_schema(), &application, Some("features"))
|
||||
.unwrap();
|
||||
let fields = whole.output_schema.fields[0]
|
||||
.r#type
|
||||
.fields
|
||||
.as_ref()
|
||||
.unwrap();
|
||||
assert!(fields[0].nullable);
|
||||
assert!(fields[1].nullable);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blob_function_plans_semantic_input_and_scalar_output() {
|
||||
let schema = ArrowSchema::new(vec![crate::blob("image", false)]);
|
||||
@@ -3062,6 +3351,70 @@ mod tests {
|
||||
assert!(output_schema.field(0).is_blob_v2());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binding_accepts_a_lance_normalized_list_child() {
|
||||
// The whole guard, not just the type helper: this also reaches the
|
||||
// output-schema comparison at the end of ensure_binding_matches_schema,
|
||||
// which used to rebuild the schema from the stored field and so failed
|
||||
// on exactly the same normalization.
|
||||
let input = ArrowField::new("value", DataType::Int64, false);
|
||||
let application = FunctionApplication::from_json(
|
||||
&serde_json::json!({
|
||||
"function": {"name": "embed", "version": "fv_embed"},
|
||||
"inputs": [{
|
||||
"parameter": "value",
|
||||
"kind": "column",
|
||||
"value": {"path": "value"}
|
||||
}],
|
||||
"output": {
|
||||
"kind": "scalar",
|
||||
"arrow_type": "fixed_size_list<float32, 4>",
|
||||
"nullable": false
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
let plan = plan_function_application(
|
||||
&ArrowSchema::new(vec![input.clone()]),
|
||||
&application,
|
||||
Some("embedding"),
|
||||
)
|
||||
.unwrap();
|
||||
let binding = binding_from_plan(&plan);
|
||||
|
||||
// The declaration says the item is non-nullable; Lance rewrites it to
|
||||
// nullable on write, so this is what the column looks like on disk.
|
||||
let stored = DataType::FixedSizeList(
|
||||
Arc::new(ArrowField::new("item", DataType::Float32, true)),
|
||||
4,
|
||||
);
|
||||
let output = ArrowField::new("embedding", stored, true).with_metadata(
|
||||
function_computed_column_metadata(binding.binding_id(), 0, &["value".into()]),
|
||||
);
|
||||
|
||||
ensure_binding_matches_schema(&ArrowSchema::new(vec![input.clone(), output]), &binding)
|
||||
.unwrap();
|
||||
|
||||
// A different element type is still a mismatch.
|
||||
let wrong = ArrowField::new(
|
||||
"embedding",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(ArrowField::new("item", DataType::Float64, true)),
|
||||
4,
|
||||
),
|
||||
true,
|
||||
)
|
||||
.with_metadata(function_computed_column_metadata(
|
||||
binding.binding_id(),
|
||||
0,
|
||||
&["value".into()],
|
||||
));
|
||||
assert!(
|
||||
ensure_binding_matches_schema(&ArrowSchema::new(vec![input, wrong]), &binding).is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blob_scalar_binding_accepts_full_logical_layout() {
|
||||
let input = crate::blob("image", false);
|
||||
@@ -3158,6 +3511,128 @@ mod tests {
|
||||
assert_eq!(fields[1].data_type(), &DataType::Int32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_struct_blob_input_preserves_exact_schema_and_nullability() {
|
||||
let payload = ArrowField::new(
|
||||
"payload",
|
||||
DataType::Struct(Fields::from(vec![
|
||||
ArrowField::new("mime_type", DataType::Utf8, false),
|
||||
ArrowField::new(
|
||||
"nested",
|
||||
DataType::Struct(Fields::from(vec![crate::blob("image", true)])),
|
||||
true,
|
||||
),
|
||||
])),
|
||||
true,
|
||||
);
|
||||
let plan = plan_function_application(
|
||||
&ArrowSchema::new(vec![payload]),
|
||||
&single_input_application("payload"),
|
||||
Some("size"),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let declared: JsonArrowDataType =
|
||||
serde_json::from_str(&plan.input_bindings[0].arrow_type).unwrap();
|
||||
let DataType::Struct(fields) =
|
||||
lance_namespace::schema::convert_json_arrow_type(&declared).unwrap()
|
||||
else {
|
||||
panic!("expected a struct Function input")
|
||||
};
|
||||
assert!(fields[1].is_nullable());
|
||||
let DataType::Struct(nested) = fields[1].data_type() else {
|
||||
panic!("expected a recursive struct Function input")
|
||||
};
|
||||
assert!(nested[0].is_blob_v2());
|
||||
assert!(nested[0].is_nullable());
|
||||
|
||||
let exact = lance_namespace::schema::convert_json_arrow_schema(&plan.input_schema).unwrap();
|
||||
let DataType::Struct(fields) = exact.field(0).data_type() else {
|
||||
panic!("expected exact input schema to retain the struct")
|
||||
};
|
||||
let DataType::Struct(nested) = fields[1].data_type() else {
|
||||
panic!("expected exact input schema to retain the nested struct")
|
||||
};
|
||||
assert!(nested[0].is_blob_v2());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recursive_blob_result_plans_one_whole_named_struct_column() {
|
||||
let details_type = exact_arrow_type(ArrowField::new(
|
||||
"details",
|
||||
DataType::Struct(Fields::from(vec![crate::blob("image", true)])),
|
||||
false,
|
||||
));
|
||||
let application = FunctionApplication::from_json(
|
||||
&serde_json::json!({
|
||||
"function": {"name": "inspect", "version": "fv_nested_blob"},
|
||||
"inputs": [],
|
||||
"output": {
|
||||
"kind": "named_struct",
|
||||
"fields": [
|
||||
{"name": "mime_type", "arrow_type": "utf8", "nullable": false},
|
||||
{"name": "details", "arrow_type": details_type, "nullable": false}
|
||||
]
|
||||
}
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.unwrap();
|
||||
let plan = plan_function_application(&ArrowSchema::empty(), &application, Some("payload"))
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(plan.outputs.len(), 1);
|
||||
assert_eq!(plan.outputs[0].result_field, WHOLE_RESULT_FIELD);
|
||||
let schema =
|
||||
lance_namespace::schema::convert_json_arrow_schema(&plan.output_schema).unwrap();
|
||||
assert_eq!(schema.field(0).name(), "payload");
|
||||
let DataType::Struct(fields) = schema.field(0).data_type() else {
|
||||
panic!("whole named result must be one struct column")
|
||||
};
|
||||
assert_eq!(
|
||||
fields.iter().map(|field| field.name()).collect::<Vec<_>>(),
|
||||
["mime_type", "details"]
|
||||
);
|
||||
let DataType::Struct(details) = fields[1].data_type() else {
|
||||
panic!("expected recursive result struct")
|
||||
};
|
||||
assert!(details[0].is_blob_v2());
|
||||
assert!(!fields.iter().any(|field| field.name() == "payload"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blob_children_under_collections_are_rejected() {
|
||||
let collections = vec![
|
||||
DataType::List(Arc::new(crate::blob("item", false))),
|
||||
DataType::LargeList(Arc::new(crate::blob("item", false))),
|
||||
DataType::FixedSizeList(Arc::new(crate::blob("item", false)), 2),
|
||||
DataType::Map(
|
||||
Arc::new(ArrowField::new(
|
||||
"entries",
|
||||
DataType::Struct(Fields::from(vec![
|
||||
ArrowField::new("key", DataType::Utf8, false),
|
||||
crate::blob("value", false),
|
||||
])),
|
||||
false,
|
||||
)),
|
||||
false,
|
||||
),
|
||||
];
|
||||
for data_type in collections {
|
||||
let schema = ArrowSchema::new(vec![ArrowField::new("value", data_type, false)]);
|
||||
let error = plan_function_application(
|
||||
&schema,
|
||||
&single_input_application("value"),
|
||||
Some("size"),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
error.to_string().contains("under a collection"),
|
||||
"got: {error}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_blob_whole_struct_binding_accepts_full_logical_layout() {
|
||||
let input = crate::blob("image", false);
|
||||
|
||||
@@ -133,7 +133,7 @@ impl NativeTable {
|
||||
),
|
||||
});
|
||||
}
|
||||
(resolved.canonical_path, resolved.field)
|
||||
(resolved.canonical_path, resolved.terminal_field)
|
||||
} else {
|
||||
Self::resolve_index_field(dataset.schema(), &opts.columns[0])?
|
||||
};
|
||||
@@ -439,7 +439,8 @@ mod tests {
|
||||
use arrow_array::record_batch;
|
||||
use arrow_array::{
|
||||
Array, ArrayRef, BinaryArray, BooleanArray, FixedSizeListArray, Float32Array, Int32Array,
|
||||
LargeBinaryArray, LargeStringArray, RecordBatch, StringArray, StructArray,
|
||||
LargeBinaryArray, LargeStringArray, ListArray, RecordBatch, StringArray, StructArray,
|
||||
UInt32Array,
|
||||
};
|
||||
use arrow_data::ArrayDataBuilder;
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
@@ -458,6 +459,7 @@ mod tests {
|
||||
use crate::query::{ExecutableQuery, QueryBase};
|
||||
use crate::table::optimize::{CompactionOptions, OptimizeAction};
|
||||
use lance_index::scalar::FullTextSearchQuery;
|
||||
use lance_index::scalar::inverted::query::{FtsQuery, MatchQuery};
|
||||
|
||||
fn create_fixed_size_list<T: Array>(
|
||||
values: T,
|
||||
@@ -599,6 +601,80 @@ mod tests {
|
||||
assert!(invalid_granularity.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_nested_list_fts_uses_deepest_document_coordinates() {
|
||||
let conn = connect("memory://").execute().await.unwrap();
|
||||
let mut docs = ListBuilder::new(ListBuilder::new(StringBuilder::new()));
|
||||
|
||||
docs.values().values().append_value("alpha");
|
||||
docs.values().values().append_value("beta");
|
||||
docs.values().append(true);
|
||||
docs.values().values().append_value("gamma");
|
||||
docs.values().values().append_value("alpha delta");
|
||||
docs.values().append(true);
|
||||
docs.append(true);
|
||||
|
||||
docs.values().append(true);
|
||||
docs.values().values().append_value("alpha");
|
||||
docs.values().append(true);
|
||||
docs.append(true);
|
||||
|
||||
let batch = RecordBatch::try_from_iter(vec![
|
||||
("id", Arc::new(Int32Array::from(vec![0, 1])) as ArrayRef),
|
||||
("docs", Arc::new(docs.finish()) as ArrayRef),
|
||||
])
|
||||
.unwrap();
|
||||
let table = conn.create_table("nested", batch).execute().await.unwrap();
|
||||
|
||||
let job = table
|
||||
.create_index(
|
||||
&["docs"],
|
||||
Index::FTS(
|
||||
FtsIndexBuilder::default()
|
||||
.document_granularity(DocumentGranularity::ListElement),
|
||||
),
|
||||
)
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap();
|
||||
job.wait().await.unwrap();
|
||||
|
||||
let query = FullTextSearchQuery::new_query(FtsQuery::Match(
|
||||
MatchQuery::new("alpha".to_string())
|
||||
.with_column(Some("docs".to_string()))
|
||||
.with_document_granularity(DocumentGranularity::ListElement),
|
||||
));
|
||||
let batches = table
|
||||
.query()
|
||||
.full_text_search(query)
|
||||
.limit(10)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut hits = Vec::new();
|
||||
for batch in batches {
|
||||
let ids = batch["id"].as_any().downcast_ref::<Int32Array>().unwrap();
|
||||
let coordinates = batch["_doc_index"]
|
||||
.as_any()
|
||||
.downcast_ref::<ListArray>()
|
||||
.unwrap();
|
||||
for row in 0..batch.num_rows() {
|
||||
let coordinate = coordinates.value(row);
|
||||
let coordinate = coordinate.as_any().downcast_ref::<UInt32Array>().unwrap();
|
||||
hits.push((ids.value(row), coordinate.values().to_vec()));
|
||||
}
|
||||
}
|
||||
hits.sort_unstable();
|
||||
assert_eq!(
|
||||
hits,
|
||||
vec![(0, vec![0, 0]), (0, vec![1, 1]), (1, vec![1, 0])]
|
||||
);
|
||||
}
|
||||
|
||||
/// Concurrent waiters, and a wait issued after the job settled, all
|
||||
/// succeed once the build does.
|
||||
#[tokio::test]
|
||||
|
||||
@@ -5,12 +5,17 @@
|
||||
//!
|
||||
//! [`super::cast::cast_to_table_schema`] calls [`coerce_blob_expr`].
|
||||
|
||||
use std::fmt;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::{DataType, Field, FieldRef, Fields};
|
||||
use arrow_array::{Array, BooleanArray, RecordBatch};
|
||||
use arrow_schema::{DataType, Field, FieldRef, Fields, Schema};
|
||||
use arrow_select::nullif::nullif;
|
||||
use datafusion::functions::core::{get_field, named_struct};
|
||||
use datafusion_common::ScalarValue;
|
||||
use datafusion_common::config::ConfigOptions;
|
||||
use datafusion_expr::ColumnarValue;
|
||||
use datafusion_physical_expr::ScalarFunctionExpr;
|
||||
use datafusion_physical_expr::expressions::{CastExpr, Literal};
|
||||
use datafusion_physical_plan::PhysicalExpr;
|
||||
@@ -133,16 +138,102 @@ pub(super) fn coerce_blob_expr(
|
||||
ns_args.push(value);
|
||||
}
|
||||
|
||||
let expr: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::new(
|
||||
let built: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::new(
|
||||
&format!("named_struct({})", table_field.name()),
|
||||
named_struct(),
|
||||
ns_args,
|
||||
table_field.clone(),
|
||||
config.clone(),
|
||||
));
|
||||
|
||||
// `named_struct` always yields a valid struct, so a null input would land
|
||||
// as a row that set neither `data` nor `uri` -- not an absent blob but a
|
||||
// malformed one, which Lance rejects on write.
|
||||
let expr: Arc<dyn PhysicalExpr> = Arc::new(AbsentBlobIsNull {
|
||||
source: input_expr,
|
||||
built,
|
||||
field: table_field.clone(),
|
||||
});
|
||||
Ok((expr, table_field.clone()))
|
||||
}
|
||||
|
||||
/// Carries the source column's nullity onto the struct built for it.
|
||||
///
|
||||
/// This is its own expression rather than a `CASE` because the projection
|
||||
/// takes its output field from `return_field`, and the generic implementation
|
||||
/// rebuilds a bare field -- which would drop the `lance.blob.v2` extension
|
||||
/// metadata and stop the column being recognised as a blob at all.
|
||||
#[derive(Debug, Clone)]
|
||||
struct AbsentBlobIsNull {
|
||||
source: Arc<dyn PhysicalExpr>,
|
||||
built: Arc<dyn PhysicalExpr>,
|
||||
field: FieldRef,
|
||||
}
|
||||
|
||||
impl fmt::Display for AbsentBlobIsNull {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "absent_blob_is_null({}, {})", self.source, self.built)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for AbsentBlobIsNull {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
self.source.eq(&other.source) && self.built.eq(&other.built) && self.field == other.field
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for AbsentBlobIsNull {}
|
||||
|
||||
impl Hash for AbsentBlobIsNull {
|
||||
fn hash<H: Hasher>(&self, state: &mut H) {
|
||||
self.source.hash(state);
|
||||
self.built.hash(state);
|
||||
self.field.hash(state);
|
||||
}
|
||||
}
|
||||
|
||||
impl PhysicalExpr for AbsentBlobIsNull {
|
||||
fn return_field(&self, _input_schema: &Schema) -> datafusion_common::Result<FieldRef> {
|
||||
Ok(self.field.clone())
|
||||
}
|
||||
|
||||
fn nullable(&self, _input_schema: &Schema) -> datafusion_common::Result<bool> {
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn evaluate(&self, batch: &RecordBatch) -> datafusion_common::Result<ColumnarValue> {
|
||||
let rows = batch.num_rows();
|
||||
let built = self.built.evaluate(batch)?.into_array(rows)?;
|
||||
let source = self.source.evaluate(batch)?.into_array(rows)?;
|
||||
let Some(nulls) = source.logical_nulls() else {
|
||||
return Ok(ColumnarValue::Array(built));
|
||||
};
|
||||
// `nullif` nulls the rows the mask marks true, which is where the
|
||||
// source had no value.
|
||||
let absent = BooleanArray::new(!nulls.inner(), None);
|
||||
Ok(ColumnarValue::Array(nullif(built.as_ref(), &absent)?))
|
||||
}
|
||||
|
||||
fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
|
||||
vec![&self.source, &self.built]
|
||||
}
|
||||
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
children: Vec<Arc<dyn PhysicalExpr>>,
|
||||
) -> datafusion_common::Result<Arc<dyn PhysicalExpr>> {
|
||||
Ok(Arc::new(Self {
|
||||
source: children[0].clone(),
|
||||
built: children[1].clone(),
|
||||
field: self.field.clone(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{self}")
|
||||
}
|
||||
}
|
||||
|
||||
enum BlobInputShape<'a> {
|
||||
Bytes,
|
||||
String,
|
||||
@@ -313,6 +404,11 @@ mod tests {
|
||||
let data = image.column_by_name("data").unwrap();
|
||||
assert!(!data.is_null(0));
|
||||
assert!(data.is_null(1));
|
||||
// The row itself has to be null, not merely a struct whose children
|
||||
// are. A present-but-empty struct set neither `data` nor `uri`, which
|
||||
// Lance rejects as malformed rather than reading as an absent blob.
|
||||
assert!(!image.is_null(0));
|
||||
assert!(image.is_null(1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -110,7 +110,7 @@ fn requires_local_namespace_execution(query: &AnyQuery) -> bool {
|
||||
// pushing these down would silently ignore the user's setting. For use_lsm that
|
||||
// is worse than a tuning miss: MemWAL read routing lives only in `create_plan`,
|
||||
// so a pushed-down query would return stale base-only data with no error.
|
||||
if query.base().use_lsm.is_some() {
|
||||
if query.base().use_lsm.is_some() || query.base().take_offsets.is_some() {
|
||||
return true;
|
||||
}
|
||||
matches!(
|
||||
@@ -154,6 +154,13 @@ pub async fn create_plan(
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<Arc<dyn ExecutionPlan>> {
|
||||
let query = query.canonicalized()?;
|
||||
if let AnyQuery::Query(request) = &query
|
||||
&& let Some(offsets) = &request.take_offsets
|
||||
{
|
||||
return crate::query::create_take_offsets_plan(table, request, offsets, options, false)
|
||||
.await;
|
||||
}
|
||||
|
||||
let query = match query {
|
||||
AnyQuery::VectorQuery(query) => query,
|
||||
AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query),
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user