mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-01 19:18:38 +00:00
Compare commits
8 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| d82fe8520f | |||
| 19232f9c50 | |||
| 5cbd979455 | |||
| e773d1e093 | |||
| c8fd3e97d1 | |||
| c196d033e9 | |||
| 16753b805a | |||
| 840e1d7313 |
Generated
+71
-92
@@ -535,9 +535,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",
|
||||
@@ -1443,9 +1443,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,9 +3455,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "fsst"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f727719438dfdb74f358a347c91ff81b6e7084a6421f34de3e473ce271f10caa"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"rand 0.9.5",
|
||||
@@ -4816,9 +4815,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
|
||||
|
||||
[[package]]
|
||||
name = "lance"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "be781f40c7a75f9eae2188a2f71174acb7a360dca97163db40b041d0828dea48"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -4890,9 +4888,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-arrow"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fb97fd9875f3036d7c2561aa5b16eb87b80ccabaa4eeb5e6099b19cc662f1cd8"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4914,8 +4911,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "771f68b04b47f3addf781116f65061808de94b05e1e9411c23c18f32d14ebe79"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4929,20 +4925,17 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dd47ec33c90bf29f688fd02118e37d3a5ad5c339caa3163f89e417dc0867001f"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
"half",
|
||||
"lance-arrow-scalar",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-bitpacking"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6f45658c5b2dc9aada41b66ee44b83af3fa888b7385ae414bae951b12a9f1cd3"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrayref",
|
||||
"crunchy",
|
||||
@@ -4952,9 +4945,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-core"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "27af3df3a7d08897efccd04461df31cedf0880c4b86a055ddce48e423d27f967"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4991,9 +4983,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datafusion"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c377f837df5296e92f9fad724c83c1bef4e74d5af6e5a9312e9307e1dead8614"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5022,9 +5013,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datagen"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "778e1a5065fa4bc184e36e32681f10f8f4680ad8cedc9377b4c088dce8c5b8da"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5041,9 +5031,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-derive"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "13e5e95e0fd3d74f7938f4bee623041421b323b5c61f242c8622a1f48a202527"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5052,9 +5041,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-encoding"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1625653c55c65f3426e281f6e29b54c603f38a40bd4cebd707bd4f3ea48be6c5"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5087,9 +5075,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-file"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b7e13c9266b478fc98f36ee19347c4658f7a6613fed77778b1a455fe1b88552e"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5120,9 +5107,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0e0cb95f2c4f341c4dd04ac60f6a89ea26a6f75e09570225cbda4854c8b088e"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5186,9 +5172,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index-core"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "79ccd371977c1f7168da259d66ad37154f23146f093d46136bc7f79559f00f2c"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5210,9 +5195,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-io"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "414d50997391b1ac83dc183c1612fdff88f58b959806078dc4c5e465154566de"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5252,9 +5236,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-linalg"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6ff55b152ef23a56d7ba7e4d1b2c9cf0cc79aef6ee607c115597557ea4059f41"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5268,9 +5251,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09991c13ab282b731e323619613914e08da9cc82b312f904e58c128b23f2f0e3"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5282,9 +5264,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-impls"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ec0bc005f6bb8f120774eb4a9ba02e10463d8338167a46cbb1391d46680a174"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -5337,9 +5318,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-select"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e8f676a2a1837cc85b77feb5326d3296827da964e40d67144f646563302a6ce9"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5353,9 +5333,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-table"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bd33054347395048b1d842dfb85a13f7801392c2da5f39425db62f00a481744b"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5395,9 +5374,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-testing"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ecc9ad9ae24f045dfddd538a39e55a28fa7e1ca6ad9f23e20d4087eaf2bb66f7"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5410,9 +5388,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-tokenizer"
|
||||
version = "11.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3bfa6f0164c8b7056150f5682ce4d415a335b59b04c479873fda04b200117d27"
|
||||
version = "12.0.0-beta.5"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.5#556637791d0048c2b4f1342dd84b67c8bbd65259"
|
||||
dependencies = [
|
||||
"frostem",
|
||||
"icu_segmenter",
|
||||
@@ -5425,7 +5402,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.15"
|
||||
version = "0.38.0"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"anyhow",
|
||||
@@ -5513,7 +5490,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-nodejs"
|
||||
version = "0.38.0-beta.15"
|
||||
version = "0.38.0"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5538,7 +5515,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.15"
|
||||
version = "0.38.0"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5771,9 +5748,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"
|
||||
@@ -6024,9 +6001,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",
|
||||
@@ -6120,14 +6097,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",
|
||||
@@ -6139,15 +6117,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",
|
||||
@@ -6159,9 +6137,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",
|
||||
@@ -8624,9 +8602,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",
|
||||
@@ -9086,9 +9064,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",
|
||||
@@ -9096,6 +9074,7 @@ dependencies = [
|
||||
"hex",
|
||||
"indexmap 1.9.3",
|
||||
"indexmap 2.14.0",
|
||||
"jiff",
|
||||
"schemars 0.9.0",
|
||||
"schemars 1.2.1",
|
||||
"serde_core",
|
||||
@@ -9106,9 +9085,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",
|
||||
@@ -10475,9 +10454,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",
|
||||
|
||||
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=11.0.0", default-features = false }
|
||||
lance-core = "=11.0.0"
|
||||
lance-datagen = "=11.0.0"
|
||||
lance-file = "=11.0.0"
|
||||
lance-io = { "version" = "=11.0.0", default-features = false }
|
||||
lance-index = "=11.0.0"
|
||||
lance-linalg = "=11.0.0"
|
||||
lance-namespace = "=11.0.0"
|
||||
lance-namespace-impls = { "version" = "=11.0.0", default-features = false }
|
||||
lance-table = "=11.0.0"
|
||||
lance-testing = "=11.0.0"
|
||||
lance-datafusion = "=11.0.0"
|
||||
lance-encoding = "=11.0.0"
|
||||
lance-arrow = "=11.0.0"
|
||||
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" }
|
||||
lancedb = { path = "rust/lancedb", default-features = false }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
|
||||
+1
-1
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>11.0.0</lance-core.version>
|
||||
<lance-core.version>12.0.0-beta.5</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
+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;
|
||||
|
||||
@@ -150,9 +150,12 @@ class Connection(object):
|
||||
def 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 drop_function(self, name: str, version: str) -> bool: ...
|
||||
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 pause_job(self, job_id: str) -> str: ...
|
||||
async def resume_job(self, job_id: str) -> str: ...
|
||||
async def job_history(
|
||||
self, job_id: Optional[str] = None
|
||||
) -> List[pa.RecordBatch]: ...
|
||||
@@ -606,6 +609,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]]]
|
||||
|
||||
@@ -712,6 +712,16 @@ class DBConnection(EnforceOverrides):
|
||||
"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 job(self, job_id: str) -> Job:
|
||||
"""A [Job][lancedb.job.Job] handle for a server-side job by id.
|
||||
|
||||
@@ -743,6 +753,26 @@ class DBConnection(EnforceOverrides):
|
||||
"cancel_job is not supported for this connection type"
|
||||
)
|
||||
|
||||
def pause_job(self, job_id: str) -> str:
|
||||
"""Pause a server-side job by id.
|
||||
|
||||
The job's workers drain and it stays parked until resumed. Returns
|
||||
"pausing", "already_paused", or "committing" -- a job finalizing its
|
||||
results cannot be parked; retry shortly.
|
||||
"""
|
||||
raise NotImplementedError("pause_job is not supported for this connection type")
|
||||
|
||||
def resume_job(self, job_id: str) -> str:
|
||||
"""Resume a paused server-side job by id.
|
||||
|
||||
Its workers pick their work back up from checkpoints. Returns
|
||||
"resumed", "still_pausing" -- the pause's worker drain is not
|
||||
confirmed yet; retry shortly -- or "not_paused".
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"resume_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.
|
||||
|
||||
@@ -1413,6 +1443,10 @@ class LanceDBConnection(DBConnection):
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
return LOOP.run(self._conn.get_function(name, version=version))
|
||||
|
||||
@override
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
@@ -1436,6 +1470,22 @@ class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
return LOOP.run(self._conn.cancel_job(job_id))
|
||||
|
||||
@override
|
||||
def pause_job(self, job_id: str) -> str:
|
||||
"""Pause a server-side job by id.
|
||||
|
||||
Returns "pausing", "already_paused", or "committing".
|
||||
"""
|
||||
return LOOP.run(self._conn.pause_job(job_id))
|
||||
|
||||
@override
|
||||
def resume_job(self, job_id: str) -> str:
|
||||
"""Resume a paused server-side job by id.
|
||||
|
||||
Returns "resumed", "still_pausing", or "not_paused".
|
||||
"""
|
||||
return LOOP.run(self._conn.resume_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.
|
||||
@@ -2243,6 +2293,10 @@ class AsyncConnection(object):
|
||||
"""Open one exact immutable Function version from the remote catalog."""
|
||||
return FunctionVersion.from_json(await self._inner.get_function(name, version))
|
||||
|
||||
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 list_jobs(self) -> List[JobInfo]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
return await self._inner.list_jobs()
|
||||
@@ -2263,6 +2317,23 @@ class AsyncConnection(object):
|
||||
"""
|
||||
return await self._inner.cancel_job(job_id)
|
||||
|
||||
async def pause_job(self, job_id: str) -> str:
|
||||
"""Pause a server-side job by id.
|
||||
|
||||
The job's workers drain and it stays parked until resumed. Returns
|
||||
"pausing", "already_paused", or "committing" -- a job finalizing its
|
||||
results cannot be parked; retry shortly.
|
||||
"""
|
||||
return await self._inner.pause_job(job_id)
|
||||
|
||||
async def resume_job(self, job_id: str) -> str:
|
||||
"""Resume a paused server-side job by id.
|
||||
|
||||
Its workers pick their work back up from checkpoints. Returns
|
||||
"resumed", "still_pausing" -- retry shortly -- or "not_paused".
|
||||
"""
|
||||
return await self._inner.resume_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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -749,6 +749,10 @@ class RemoteDBConnection(DBConnection):
|
||||
def get_function(self, name: str, *, version: str) -> FunctionVersion:
|
||||
return LOOP.run(self._conn.get_function(name, version=version))
|
||||
|
||||
@override
|
||||
def drop_function(self, name: str, *, version: str) -> bool:
|
||||
return LOOP.run(self._conn.drop_function(name, version=version))
|
||||
|
||||
@override
|
||||
def list_jobs(self) -> List["JobInfo"]:
|
||||
"""List server-side jobs across the database's tables."""
|
||||
@@ -772,6 +776,22 @@ class RemoteDBConnection(DBConnection):
|
||||
"""
|
||||
return LOOP.run(self._conn.cancel_job(job_id))
|
||||
|
||||
@override
|
||||
def pause_job(self, job_id: str) -> str:
|
||||
"""Pause a server-side job by id.
|
||||
|
||||
Returns "pausing", "already_paused", or "committing".
|
||||
"""
|
||||
return LOOP.run(self._conn.pause_job(job_id))
|
||||
|
||||
@override
|
||||
def resume_job(self, job_id: str) -> str:
|
||||
"""Resume a paused server-side job by id.
|
||||
|
||||
Returns "resumed", "still_pausing", or "not_paused".
|
||||
"""
|
||||
return LOOP.run(self._conn.resume_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.
|
||||
|
||||
@@ -1678,9 +1678,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
|
||||
----------
|
||||
@@ -4090,6 +4090,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
|
||||
|
||||
@@ -5983,7 +5984,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 +6065,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 +6563,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]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -67,6 +69,80 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
}
|
||||
|
||||
|
||||
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)
|
||||
@@ -941,6 +1017,8 @@ 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.drop_function("normalize_score", version="fv_exact")
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
@@ -986,6 +1064,12 @@ def _mock_remote_function_catalog():
|
||||
"version": "fv_exact",
|
||||
}
|
||||
response = state["version"]
|
||||
elif self.path == "/v1/functions/drop":
|
||||
assert body == {
|
||||
"name": "normalize_score",
|
||||
"version": "fv_exact",
|
||||
}
|
||||
response = {"dropped": True}
|
||||
else:
|
||||
status = 404
|
||||
response = {"error": "not found"}
|
||||
@@ -1044,3 +1128,40 @@ def test_blocking_remote_registration_returns_function_version():
|
||||
"/v1/functions/create",
|
||||
"/v1/jobs/describe",
|
||||
]
|
||||
|
||||
|
||||
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"},
|
||||
)
|
||||
]
|
||||
|
||||
@@ -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():
|
||||
@@ -2509,6 +2534,26 @@ def test_remote_connection_jobs_surface():
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"job_id": "job-1"}')
|
||||
elif request.path == "/v1/jobs/pause":
|
||||
if payload["job_id"] != "job-1":
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
return
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"job_id": "job-1", "paused": true}')
|
||||
elif request.path == "/v1/jobs/resume":
|
||||
if payload["job_id"] != "job-1":
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
return
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(
|
||||
b'{"job_id": "job-1", "resumed": false, "still_pausing": true}'
|
||||
)
|
||||
elif request.path == "/v1/jobs/query_events":
|
||||
assert payload["job_id"] == "job-1"
|
||||
request.send_response(200)
|
||||
@@ -2537,6 +2582,9 @@ def test_remote_connection_jobs_surface():
|
||||
assert db.cancel_job("job-1") is True
|
||||
assert db.cancel_job("missing") is False
|
||||
|
||||
assert db.pause_job("job-1") == "pausing"
|
||||
assert db.resume_job("job-1") == "still_pausing"
|
||||
|
||||
batches = db.job_history("job-1")
|
||||
assert len(batches) == 1
|
||||
assert batches[0].num_rows == 2
|
||||
|
||||
@@ -629,6 +629,17 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
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 list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -655,6 +666,30 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn pause_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let status = inner.pause_job(&job_id).await.infer_error()?;
|
||||
Ok(match status {
|
||||
lancedb::database::PauseJobStatus::Pausing => "pausing",
|
||||
lancedb::database::PauseJobStatus::AlreadyPaused => "already_paused",
|
||||
lancedb::database::PauseJobStatus::Committing => "committing",
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resume_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let status = inner.resume_job(&job_id).await.infer_error()?;
|
||||
Ok(match status {
|
||||
lancedb::database::ResumeJobStatus::Resumed => "resumed",
|
||||
lancedb::database::ResumeJobStatus::StillPausing => "still_pausing",
|
||||
lancedb::database::ResumeJobStatus::NotPaused => "not_paused",
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None))]
|
||||
pub fn job_history(
|
||||
self_: PyRef<'_, Self>,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -24,7 +24,7 @@ use crate::data::scannable::Scannable;
|
||||
use crate::database::listing::ListingDatabase;
|
||||
use crate::database::{
|
||||
CloneTableRequest, Database, DatabaseOptions, JobDescription, JobInfo, OpenTableRequest,
|
||||
ReadConsistency, TableNamesRequest,
|
||||
PauseJobStatus, ReadConsistency, ResumeJobStatus, TableNamesRequest,
|
||||
};
|
||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
@@ -523,6 +523,21 @@ impl Connection {
|
||||
.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
|
||||
}
|
||||
|
||||
/// Rename a table in the database.
|
||||
///
|
||||
/// This is only supported in LanceDB Cloud.
|
||||
@@ -575,6 +590,18 @@ impl Connection {
|
||||
self.internal.cancel_job(job_id.as_ref()).await
|
||||
}
|
||||
|
||||
/// Pause a server-side job by id. Its workers drain and it stays parked
|
||||
/// until resumed; see [`PauseJobStatus`] for the outcomes.
|
||||
pub async fn pause_job(&self, job_id: impl AsRef<str>) -> Result<PauseJobStatus> {
|
||||
self.internal.pause_job(job_id.as_ref()).await
|
||||
}
|
||||
|
||||
/// Resume a paused server-side job by id. Its workers pick their work
|
||||
/// back up from checkpoints; see [`ResumeJobStatus`] for the outcomes.
|
||||
pub async fn resume_job(&self, job_id: impl AsRef<str>) -> Result<ResumeJobStatus> {
|
||||
self.internal.resume_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>> {
|
||||
|
||||
@@ -235,6 +235,29 @@ pub struct JobDescription {
|
||||
pub failure: Option<crate::error::JobFailure>,
|
||||
}
|
||||
|
||||
/// The server's answer to a pause request.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PauseJobStatus {
|
||||
/// The pause was accepted; workers drain and the job stays parked.
|
||||
Pausing,
|
||||
/// The job was already paused, so a repeated pause changed nothing.
|
||||
AlreadyPaused,
|
||||
/// The job is finalizing its results and cannot be parked right now.
|
||||
/// The commit is the short tail of a long job; retry shortly.
|
||||
Committing,
|
||||
}
|
||||
|
||||
/// The server's answer to a resume request.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ResumeJobStatus {
|
||||
/// The job re-entered the queue and will run again.
|
||||
Resumed,
|
||||
/// The pause's worker drain is not confirmed yet; retry shortly.
|
||||
StillPausing,
|
||||
/// The job was not paused, so there was nothing to resume.
|
||||
NotPaused,
|
||||
}
|
||||
|
||||
fn job_op_not_supported<T>(what: &str) -> Result<T> {
|
||||
Err(crate::error::Error::NotSupported {
|
||||
message: format!("{} is not supported by this database", what),
|
||||
@@ -307,6 +330,10 @@ pub trait Database:
|
||||
) -> Result<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()
|
||||
}
|
||||
/// 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.
|
||||
@@ -327,6 +354,16 @@ pub trait Database:
|
||||
async fn cancel_job(&self, _job_id: &str) -> Result<bool> {
|
||||
job_op_not_supported("cancel_job")
|
||||
}
|
||||
/// Pause a job by id. The job's workers drain and it stays parked until
|
||||
/// resumed; see [`PauseJobStatus`] for the outcomes.
|
||||
async fn pause_job(&self, _job_id: &str) -> Result<PauseJobStatus> {
|
||||
job_op_not_supported("pause_job")
|
||||
}
|
||||
/// Resume a paused job by id. It re-enters the queue and its workers pick
|
||||
/// their work back up from checkpoints; see [`ResumeJobStatus`].
|
||||
async fn resume_job(&self, _job_id: &str) -> Result<ResumeJobStatus> {
|
||||
job_op_not_supported("resume_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>> {
|
||||
|
||||
@@ -13,7 +13,7 @@ use lance::dataset::{ReadParams, WriteMode, builder::DatasetBuilder};
|
||||
use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore};
|
||||
use lance_datafusion::utils::StreamingWriteSource;
|
||||
use lance_file::version::LanceFileVersion;
|
||||
use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider};
|
||||
use lance_io::object_store::{ReadDirOptions, StorageOptionsAccessor, StorageOptionsProvider};
|
||||
use lance_table::io::commit::commit_handler_from_url;
|
||||
use object_store::local::LocalFileSystem;
|
||||
use snafu::ResultExt;
|
||||
@@ -282,11 +282,14 @@ impl std::fmt::Display for ListingDatabase {
|
||||
|
||||
const LANCE_EXTENSION: &str = "lance";
|
||||
|
||||
/// The table a listed child directory holds, or `None` if it is not a table at all.
|
||||
/// The table a listed child of the database names, or `None` if the child is not a table.
|
||||
///
|
||||
/// A table is the directory `<name>.lance`; a loose file or any other directory under the
|
||||
/// database prefix belongs to something else. `dir_suffix` is `.lance`, built once by the
|
||||
/// caller rather than per child.
|
||||
/// The table a listed child directory holds, or `None` if it is not a table at all.
|
||||
///
|
||||
/// Only directories are considered, so a loose object named like a table is not one.
|
||||
fn table_name(location: &object_store::path::Path, dir_suffix: &str) -> Option<String> {
|
||||
location
|
||||
.filename()?
|
||||
@@ -294,75 +297,6 @@ fn table_name(location: &object_store::path::Path, dir_suffix: &str) -> Option<S
|
||||
.map(String::from)
|
||||
.filter(|name| !name.is_empty())
|
||||
}
|
||||
|
||||
/// One page of the table directories under the database directory, in key order.
|
||||
struct DirPage {
|
||||
/// The table directories the page holds, as the store lists them.
|
||||
common_prefixes: Vec<object_store::path::Path>,
|
||||
/// Resumes after this page, or `None` when the page reached the end of the level.
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
/// Where a listed location sits inside the database directory — the space page tokens live
|
||||
/// in — or `None` if it is not a child of that directory at all. Matching both halves of the
|
||||
/// prefix drops a location that merely starts with the directory's name (`dbx/y` against
|
||||
/// `db/`) as well as the marker object some stores keep for the directory itself.
|
||||
fn relative_key<'a>(prefix: Option<&str>, location: &'a str) -> Option<&'a str> {
|
||||
let relative = match prefix {
|
||||
Some(prefix) => location.strip_prefix(prefix)?,
|
||||
None => location,
|
||||
};
|
||||
(!relative.is_empty()).then_some(relative)
|
||||
}
|
||||
|
||||
/// One page of the table directories under `base_path`, one directory level deep.
|
||||
///
|
||||
/// Lance 11 exposes no paginated directory listing, so the level is listed in full and paged
|
||||
/// locally: table directories go into key order (a directory's key keeps its trailing `/`,
|
||||
/// so a token is never a table name), the page is the smallest `limit` of them past
|
||||
/// `page_token`, and the token handed back is the key of the last directory the page took —
|
||||
/// so a page that took nothing ends the listing rather than resuming from a position no page
|
||||
/// ever reached. Only `<name>.lance/` directories enter the page: loose objects, other
|
||||
/// directories, and a bare `.lance/` never take a page slot or name a token, which keeps a
|
||||
/// page to exactly one listing of the level. Correct on every store, at the cost of that one
|
||||
/// full-level listing per page.
|
||||
async fn read_dir_page(
|
||||
object_store: &ObjectStore,
|
||||
base_path: &object_store::path::Path,
|
||||
page_token: Option<String>,
|
||||
limit: Option<usize>,
|
||||
) -> Result<DirPage> {
|
||||
let listed = object_store.list_with_delimiter(Some(base_path)).await?;
|
||||
let prefix = {
|
||||
let base = base_path.as_ref();
|
||||
(!base.is_empty()).then(|| format!("{base}/"))
|
||||
};
|
||||
let table_dir_suffix = format!(".{LANCE_EXTENSION}/");
|
||||
let mut children: Vec<(String, object_store::path::Path)> = listed
|
||||
.common_prefixes
|
||||
.into_iter()
|
||||
.filter_map(|location| {
|
||||
let key = format!("{}/", relative_key(prefix.as_deref(), location.as_ref())?);
|
||||
(key.len() > table_dir_suffix.len() && key.ends_with(&table_dir_suffix))
|
||||
.then_some((key, location))
|
||||
})
|
||||
.collect();
|
||||
children.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
|
||||
if let Some(resume) = &page_token {
|
||||
children.retain(|(key, _)| key > resume);
|
||||
}
|
||||
let total = children.len();
|
||||
children.truncate(limit.unwrap_or(total).min(total));
|
||||
let page_token = match children.last() {
|
||||
Some((last, _)) if children.len() < total => Some(last.clone()),
|
||||
_ => None,
|
||||
};
|
||||
Ok(DirPage {
|
||||
common_prefixes: children.into_iter().map(|(_, location)| location).collect(),
|
||||
page_token,
|
||||
})
|
||||
}
|
||||
|
||||
const ENGINE: &str = "engine";
|
||||
const MIRRORED_STORE: &str = "mirroredStore";
|
||||
|
||||
@@ -1048,7 +982,8 @@ impl Database for ListingDatabase {
|
||||
let mut tables = Vec::new();
|
||||
let mut page_token = request.page_token.filter(|token| !token.is_empty());
|
||||
|
||||
// A page of nothing: no table was handed over for a token to resume after.
|
||||
// A page of nothing: the store rejects a limit of zero, and no table was handed over
|
||||
// for a token to resume after.
|
||||
if limit == Some(0) {
|
||||
return Ok(ListTablesResponse {
|
||||
context: None,
|
||||
@@ -1057,21 +992,35 @@ impl Database for ListingDatabase {
|
||||
});
|
||||
}
|
||||
|
||||
// The page holds only table directories, so one call — and the one full-level
|
||||
// listing behind it — fills it.
|
||||
let page = read_dir_page(
|
||||
&self.object_store,
|
||||
&self.base_path,
|
||||
page_token.take(),
|
||||
limit,
|
||||
)
|
||||
.await?;
|
||||
page_token = page.page_token;
|
||||
tables.extend(
|
||||
page.common_prefixes
|
||||
.iter()
|
||||
.filter_map(|location| table_name(location, &dir_suffix)),
|
||||
);
|
||||
loop {
|
||||
// Ask only for what the page still has room for, so a database holding more
|
||||
// than one page costs one request per page rather than one per table.
|
||||
let listing = self
|
||||
.object_store
|
||||
.read_dir_page(
|
||||
self.base_path.clone(),
|
||||
ReadDirOptions {
|
||||
page_token: page_token.take(),
|
||||
limit: limit.map(|limit| limit - tables.len()),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
page_token = listing.page_token;
|
||||
// Only child directories can be tables, and the store already separates them
|
||||
// out, so the objects in the page are not looked at.
|
||||
tables.extend(
|
||||
listing
|
||||
.result
|
||||
.common_prefixes
|
||||
.iter()
|
||||
.filter_map(|location| table_name(location, &dir_suffix)),
|
||||
);
|
||||
// Children that are not tables leave the page short of the limit, so keep
|
||||
// going until the page is full or the database runs out.
|
||||
if page_token.is_none() || limit.is_none_or(|limit| tables.len() >= limit) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ListTablesResponse {
|
||||
context: None,
|
||||
@@ -1717,8 +1666,8 @@ mod tests {
|
||||
}
|
||||
|
||||
/// Only directories named `<name>.lance` are tables; loose files and other directories
|
||||
/// under the database prefix are not. They never take a page slot, so even a `limit`
|
||||
/// smaller than the clutter ahead of the first table returns that table.
|
||||
/// under the database prefix are not. A page spent on them is filled from the next one,
|
||||
/// so a page holding only non-tables does not read as an empty database.
|
||||
#[tokio::test]
|
||||
async fn test_listing_ignores_non_table_children() {
|
||||
let (tempdir, db) = setup_database().await;
|
||||
@@ -1737,37 +1686,6 @@ mod tests {
|
||||
assert_eq!(page.tables, vec!["real"]);
|
||||
}
|
||||
|
||||
/// The Lance 11 fallback pages locally over one full-level listing, so a bounded page
|
||||
/// costs exactly one listing call — clutter ahead of the first table must not buy extra
|
||||
/// round trips.
|
||||
#[tokio::test]
|
||||
async fn test_one_full_listing_per_public_page() {
|
||||
use crate::io::object_store::io_tracking::IoStatsHolder;
|
||||
use lance_io::object_store::WrappingObjectStore;
|
||||
|
||||
let (tempdir, mut db) = setup_database().await;
|
||||
create_tables(&db, &["real"]).await;
|
||||
std::fs::write(tempdir.path().join("aaa-loose.lance"), b"not a table").unwrap();
|
||||
create_dir_all(tempdir.path().join("aaa-scratch")).unwrap();
|
||||
|
||||
let io_stats = IoStatsHolder::default();
|
||||
let mut tracked_store = (*db.object_store).clone();
|
||||
tracked_store.inner =
|
||||
io_stats.wrap(&tracked_store.store_prefix, tracked_store.inner.clone());
|
||||
db.object_store = Arc::new(tracked_store);
|
||||
|
||||
let page = db
|
||||
.list_tables(ListTablesRequest {
|
||||
limit: Some(1),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(page.tables, vec!["real"]);
|
||||
assert_eq!(io_stats.incremental_stats().read_iops, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn listing_ignores_empty_table_name() {
|
||||
let (tempdir, db) = setup_database().await;
|
||||
|
||||
@@ -539,9 +539,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 +1493,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]
|
||||
|
||||
@@ -10,7 +10,7 @@ use lance::io::WrappingObjectStore;
|
||||
use object_store::{
|
||||
CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
|
||||
ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result,
|
||||
UploadPart, path::Path,
|
||||
UploadPart, list::PaginatedListStore, path::Path,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -187,6 +187,14 @@ impl WrappingObjectStore for MirroringObjectStoreWrapper {
|
||||
secondary: self.secondary.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
original: Arc<dyn PaginatedListStore>,
|
||||
) -> Option<Arc<dyn PaginatedListStore>> {
|
||||
Some(original)
|
||||
}
|
||||
}
|
||||
|
||||
// windows pathing can't be simply concatenated
|
||||
|
||||
@@ -12,7 +12,7 @@ use lance::io::WrappingObjectStore;
|
||||
use object_store::{
|
||||
CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
|
||||
PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult,
|
||||
UploadPart, path::Path,
|
||||
UploadPart, list::PaginatedListStore, path::Path,
|
||||
};
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
@@ -57,6 +57,14 @@ impl WrappingObjectStore for IoStatsHolder {
|
||||
stats: self.0.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
original: Arc<dyn PaginatedListStore>,
|
||||
) -> Option<Arc<dyn PaginatedListStore>> {
|
||||
Some(original)
|
||||
}
|
||||
}
|
||||
|
||||
impl IoTrackingStore {
|
||||
|
||||
+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();
|
||||
|
||||
@@ -26,7 +26,9 @@ use crate::database::{
|
||||
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::{
|
||||
DescribeJobResponse, PauseJobResponse, RemoteJob, ResumeJobResponse, job_state_to_client,
|
||||
};
|
||||
use crate::remote::util::stream_as_body;
|
||||
use crate::table::BaseTable;
|
||||
|
||||
@@ -533,6 +535,11 @@ struct RemoteListJobsResponse {
|
||||
page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteDropFunctionResponse {
|
||||
dropped: bool,
|
||||
}
|
||||
|
||||
/// 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,6 +590,20 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
response.json().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
fn job(&self, job_id: &str) -> Result<crate::job::Job> {
|
||||
Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new(
|
||||
self.client.clone(),
|
||||
@@ -665,6 +686,40 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn pause_job(&self, job_id: &str) -> Result<crate::database::PauseJobStatus> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/pause")
|
||||
.json(&serde_json::json!({ "job_id": job_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: PauseJobResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(if body.paused {
|
||||
crate::database::PauseJobStatus::Pausing
|
||||
} else if body.committing {
|
||||
crate::database::PauseJobStatus::Committing
|
||||
} else {
|
||||
crate::database::PauseJobStatus::AlreadyPaused
|
||||
})
|
||||
}
|
||||
|
||||
async fn resume_job(&self, job_id: &str) -> Result<crate::database::ResumeJobStatus> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/resume")
|
||||
.json(&serde_json::json!({ "job_id": job_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: ResumeJobResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(if body.resumed {
|
||||
crate::database::ResumeJobStatus::Resumed
|
||||
} else if body.still_pausing {
|
||||
crate::database::ResumeJobStatus::StillPausing
|
||||
} else {
|
||||
crate::database::ResumeJobStatus::NotPaused
|
||||
})
|
||||
}
|
||||
|
||||
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 {
|
||||
@@ -2600,6 +2655,45 @@ mod tests {
|
||||
assert!(!conn.cancel_job("nope").await.unwrap());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pause_and_resume_job() {
|
||||
use crate::database::{PauseJobStatus, ResumeJobStatus};
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/pause");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id": "job-1", "paused": true}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert_eq!(
|
||||
conn.pause_job("job-1").await.unwrap(),
|
||||
PauseJobStatus::Pausing
|
||||
);
|
||||
|
||||
let conn = Connection::new_with_handler(|_| {
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id": "job-1", "paused": false, "committing": true}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert_eq!(
|
||||
conn.pause_job("job-1").await.unwrap(),
|
||||
PauseJobStatus::Committing
|
||||
);
|
||||
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/resume");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id": "job-1", "resumed": false, "still_pausing": true}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert_eq!(
|
||||
conn.resume_job("job-1").await.unwrap(),
|
||||
ResumeJobStatus::StillPausing
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_job_history_parses_arrow_stream() {
|
||||
let schema = Arc::new(Schema::new(vec![Field::new(
|
||||
@@ -2689,6 +2783,25 @@ mod tests {
|
||||
assert_eq!(version.version(), "fv_01K3EXACT");
|
||||
}
|
||||
|
||||
#[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));
|
||||
|
||||
@@ -73,6 +73,28 @@ pub(super) struct ReportedFailure {
|
||||
retryable: Option<bool>,
|
||||
}
|
||||
|
||||
/// Forward-compatible `/v1/jobs/pause` wire envelope.
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct PauseJobResponse {
|
||||
/// False when the job was already paused, so a repeated pause changed
|
||||
/// nothing.
|
||||
#[serde(default)]
|
||||
pub(super) paused: bool,
|
||||
/// The job is finalizing its results and cannot be parked right now.
|
||||
#[serde(default)]
|
||||
pub(super) committing: bool,
|
||||
}
|
||||
|
||||
/// Forward-compatible `/v1/jobs/resume` wire envelope.
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct ResumeJobResponse {
|
||||
#[serde(default)]
|
||||
pub(super) resumed: bool,
|
||||
/// The pause's worker drain is not confirmed yet.
|
||||
#[serde(default)]
|
||||
pub(super) still_pausing: bool,
|
||||
}
|
||||
|
||||
/// Forward-compatible `/v1/jobs/describe` wire envelope.
|
||||
#[derive(Deserialize)]
|
||||
pub(super) struct DescribeJobResponse {
|
||||
|
||||
@@ -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::{
|
||||
@@ -2022,6 +2022,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 +2597,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 +2622,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 +2680,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 +2738,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));
|
||||
@@ -3690,7 +3738,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;
|
||||
@@ -5611,6 +5659,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 =
|
||||
|
||||
@@ -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>;
|
||||
@@ -1652,9 +1660,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
|
||||
/// ----------
|
||||
@@ -4183,6 +4191,14 @@ mod tests {
|
||||
parent_list_calls: self.parent_list_calls.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
_original: Arc<dyn object_store::list::PaginatedListStore>,
|
||||
) -> Option<Arc<dyn object_store::list::PaginatedListStore>> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -4286,6 +4302,14 @@ mod tests {
|
||||
self.called.store(true, Ordering::Relaxed);
|
||||
original
|
||||
}
|
||||
|
||||
fn wrap_paginated(
|
||||
&self,
|
||||
_store_prefix: &str,
|
||||
original: Arc<dyn object_store::list::PaginatedListStore>,
|
||||
) -> Option<Arc<dyn object_store::list::PaginatedListStore>> {
|
||||
Some(original)
|
||||
}
|
||||
}
|
||||
|
||||
#[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),
|
||||
|
||||
@@ -300,7 +300,7 @@ async fn build_read_context(
|
||||
for shard_id in shard_ids {
|
||||
let manifest_store =
|
||||
ShardManifestStore::new(store.clone(), &base_path, shard_id, scan_batch_size);
|
||||
if let Some(manifest) = manifest_store.read_latest().await? {
|
||||
if let Some(manifest) = manifest_store.latest().await? {
|
||||
snapshots.push(snapshot_from_manifest(shard_id, &manifest, &exclude));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -45,7 +45,11 @@ async fn local_function_catalog_operations_return_stable_not_supported() {
|
||||
.get_function("normalize_score", "fv_exact")
|
||||
.await
|
||||
.unwrap_err();
|
||||
for error in [create_error, lookup_error] {
|
||||
let drop_error = connection
|
||||
.drop_function("normalize_score", "fv_exact")
|
||||
.await
|
||||
.unwrap_err();
|
||||
for error in [create_error, lookup_error, drop_error] {
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::NotSupported { message }
|
||||
|
||||
Reference in New Issue
Block a user