diff --git a/.bumpversion.toml b/.bumpversion.toml index bbad8583d..17e99cfb8 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.40.0-beta.7" +current_version = "0.40.0-beta.8" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml index 5f45ba68c..4f883aec1 100644 --- a/.github/workflows/pypi-publish.yml +++ b/.github/workflows/pypi-publish.yml @@ -21,7 +21,7 @@ permissions: contents: read # Without this, a force-push to a PR leaves the previous run going -- including -# a ~74 minute Windows job and a billed arm64 wheel build. +# a long fat-LTO Windows job and a billed arm64 wheel build. concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true @@ -123,8 +123,11 @@ jobs: path: target/wheels/lancedb-*.whl if-no-files-found: error windows: - timeout-minutes: 90 - runs-on: windows-latest + # The fat-LTO build of the `lancedb` crate does not finish within 90 minutes + # on the 4-core standard runner, so use the 8-core larger runner with extra + # headroom. + timeout-minutes: 150 + runs-on: windows-2025-8x-x64 env: # link.exe is single-threaded and the long pole on Windows builds. Use # rustc's bundled lld-link instead. @@ -145,7 +148,7 @@ jobs: # or the default branch -- so with no run on main there is nothing that # can populate an entry the release build would be allowed to read. Fixing # this needs a main/nightly trigger (which would also catch wheel-build - # breakage before a release); the ~74 minutes here is otherwise dominated + # breakage before a release); the build time here is otherwise dominated # by the fat-LTO link, which no cache avoids. - uses: ./.github/workflows/build_windows_wheel with: diff --git a/Cargo.lock b/Cargo.lock index f6775b0e8..c6ecd450c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -141,7 +141,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -152,7 +152,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1819,7 +1819,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3087,7 +3087,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3315,7 +3315,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "rand 0.9.5", @@ -4538,7 +4538,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.3", "system-configuration", "tokio", "tower-service", @@ -4880,7 +4880,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5143,8 +5143,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arc-swap", "arrow", @@ -5216,7 +5216,7 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "13.0.0-beta.8" +version = "13.0.0-beta.12" dependencies = [ "arrow-array", "arrow-buffer", @@ -5238,7 +5238,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "arrow-buffer", @@ -5252,7 +5252,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "arrow-schema", @@ -5261,8 +5261,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrayref", "crunchy", @@ -5272,8 +5272,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "arrow-buffer", @@ -5310,8 +5310,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow", "arrow-array", @@ -5342,8 +5342,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow", "arrow-array", @@ -5360,8 +5360,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "proc-macro2", "quote", @@ -5370,8 +5370,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-arith", "arrow-array", @@ -5404,8 +5404,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-arith", "arrow-array", @@ -5437,8 +5437,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "datafusion", "geo-traits", @@ -5452,8 +5452,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arc-swap", "arrow", @@ -5480,6 +5480,7 @@ dependencies = [ "geoarrow-array", "geoarrow-schema", "half", + "hex", "itertools 0.14.0", "jieba-rs", "jsonb", @@ -5517,12 +5518,13 @@ dependencies = [ "tempfile", "tokio", "tracing", + "twox-hash", ] [[package]] name = "lance-index-core" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "arrow-schema", @@ -5544,8 +5546,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow", "arrow-array", @@ -5587,8 +5589,8 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "arrow-schema", @@ -5602,8 +5604,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow", "async-trait", @@ -5617,8 +5619,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow", "arrow-ipc", @@ -5671,8 +5673,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "arrow-buffer", @@ -5686,8 +5688,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow", "arrow-array", @@ -5727,8 +5729,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "arrow-array", "arrow-schema", @@ -5741,8 +5743,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "13.0.0-beta.8" -source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.8#3c20b7875f621c929fdcecabe74dbca507ec0f0d" +version = "13.0.0-beta.12" +source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.12#5d2946fa452c61cb018144f3200892b6dd4a8b4f" dependencies = [ "frostem", "icu_segmenter", @@ -5755,7 +5757,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.40.0-beta.7" +version = "0.40.0-beta.8" dependencies = [ "ahash", "anyhow", @@ -5852,7 +5854,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.40.0-beta.7" +version = "0.40.0-beta.8" dependencies = [ "arrow-array", "arrow-buffer", @@ -5877,7 +5879,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.40.0-beta.7" +version = "0.40.0-beta.8" dependencies = [ "arc-swap", "arrow", @@ -6631,7 +6633,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -8134,8 +8136,8 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ - "heck 0.4.1", - "itertools 0.11.0", + "heck 0.5.0", + "itertools 0.12.1", "log", "multimap", "petgraph", @@ -8154,7 +8156,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.12.1", "proc-macro2", "quote", "syn 2.0.117", @@ -8376,7 +8378,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.45", - "socket2 0.5.10", + "socket2 0.6.3", "thiserror 2.0.18", "tokio", "tracing", @@ -8414,7 +8416,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.3", "tracing", "windows-sys 0.60.2", ] @@ -9214,7 +9216,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9285,7 +9287,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9832,7 +9834,7 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9844,7 +9846,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9867,7 +9869,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -9986,7 +9988,7 @@ dependencies = [ "cfg-if 1.0.4", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -10278,7 +10280,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -11299,7 +11301,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index f94710c72..55980dcab 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.91.0" [workspace.dependencies] -lance = { "version" = "=13.0.0-beta.8", default-features = false, "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=13.0.0-beta.8", default-features = false, "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=13.0.0-beta.8", default-features = false, "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=13.0.0-beta.8", "tag" = "v13.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { version = "=13.0.0-beta.8", path = "python/vendor/lance-arrow" } +lance = { "version" = "=13.0.0-beta.12", default-features = false, "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=13.0.0-beta.12", default-features = false, "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=13.0.0-beta.12", default-features = false, "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=13.0.0-beta.12", "tag" = "v13.0.0-beta.12", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { version = "=13.0.0-beta.12", path = "python/vendor/lance-arrow" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 073b57e16..d15242673 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.40.0-beta.7 + 0.40.0-beta.8 ``` diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index f61dd8633..ef947dc66 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.40.0-beta.7 + 0.40.0-beta.8 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 9e1d76f9a..832df9054 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.40.0-beta.7 + 0.40.0-beta.8 pom ${project.artifactId} LanceDB Java SDK Parent POM @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 13.0.0-beta.8 + 13.0.0-beta.12 false 2.30.0 1.7 diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 6e7b6e245..931fea960 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.40.0-beta.7" +version = "0.40.0-beta.8" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/__test__/arrow.test.ts b/nodejs/__test__/arrow.test.ts index 6c1d5ecba..df02b4936 100644 --- a/nodejs/__test__/arrow.test.ts +++ b/nodejs/__test__/arrow.test.ts @@ -601,6 +601,103 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( ); }); + it.each([ + ["ArrayBuffer", new Uint8Array([104]).buffer], + ["URL", new URL("https://example.com/")], + ["empty object", {}], + [ + "class instance with enumerable fields", + Object.assign(new Blob(["x"]), { position: 0 }), + ], + ])("rejects %s values without a schema", (_name, value) => { + expect(() => makeArrowTable([{ id: 1, value }])).toThrow( + /field value at row 0/, + ); + expect(() => makeArrowTable([{ value }])).toThrow( + /field value at row 0/, + ); + }); + + it.each([ + ["Binary", new Binary(), new Uint8Array([104]).buffer], + ["Utf8", new Utf8(), new URL("https://example.com/")], + ])("rejects %s values with a schema", (_name, type, value) => { + const schema = new Schema([new Field("value", type, true)]); + expect(() => makeArrowTable([{ value }], { schema })).toThrow( + /field value at row 0/, + ); + }); + + it("rejects unsupported values in nested struct fields", () => { + const schema = new Schema([ + new Field( + "metadata", + new Struct([new Field("bytes", new Binary(), true)]), + true, + ), + ]); + expect(() => + makeArrowTable( + [{ metadata: { bytes: new Uint8Array([104]).buffer } }], + { schema }, + ), + ).toThrow(/field metadata\.bytes at row 0/); + }); + + it("rejects unsupported objects in list elements", () => { + const schema = new Schema([ + new Field( + "images", + new List(new Field("item", new Binary(), true)), + true, + ), + ]); + expect(() => + makeArrowTable([{ images: [new Blob(["abc"])] }], { schema }), + ).toThrow(/field images\[0\] at row 0/); + }); + + it("rejects unsupported objects in nested lists", () => { + const schema = new Schema([ + new Field( + "images", + new List( + new Field( + "item", + new List(new Field("item", new Binary(), true)), + true, + ), + ), + true, + ), + ]); + expect(() => + makeArrowTable([{ images: [[new Blob(["abc"])]] }], { schema }), + ).toThrow(/field images\[0\]\[0\] at row 0/); + }); + + it("rejects unsupported nested struct fields in lists", () => { + const schema = new Schema([ + new Field( + "items", + new List( + new Field( + "item", + new Struct([new Field("bytes", new Binary(), true)]), + true, + ), + ), + true, + ), + ]); + expect(() => + makeArrowTable( + [{ items: [{ bytes: new Uint8Array([104]).buffer }] }], + { schema }, + ), + ).toThrow(/field items\[0\]\.bytes at row 0/); + }); + it("will ignore generated dictionary IDs when comparing inferred types", function () { const table = makeArrowTable([{ str: "a" }, { str: "b" }], { dictionaryEncodeStrings: true, diff --git a/nodejs/__test__/blob.test.ts b/nodejs/__test__/blob.test.ts index e8e9357e3..5c01897b2 100644 --- a/nodejs/__test__/blob.test.ts +++ b/nodejs/__test__/blob.test.ts @@ -102,6 +102,27 @@ describe("BlobFile", () => { }); describe("makeArrowTable blob columns", () => { + it.each([ + ["ArrayBuffer", new Uint8Array([104]).buffer], + ["Blob", new Blob(["hello"])], + ["File", new File(["hello"], "hello.txt")], + ["URL", new URL("https://example.com/blob")], + ["ReadableStream", new ReadableStream()], + ["data struct with ArrayBuffer", { data: new Uint8Array([104]).buffer }], + ["uri struct with URL", { uri: new URL("https://example.com/blob") }], + ])("rejects %s input rather than writing a null blob", (_name, image) => { + const schema = new Schema([blob("image")]); + expect(() => makeArrowTable([{ image }], { schema })).toThrow(/image/); + }); + + it("rejects Blob instances with enumerable fields as one invalid value", () => { + const schema = new Schema([blob("image")]); + const image = Object.assign(new Blob(["hello"]), { position: 0 }); + expect(() => makeArrowTable([{ image }], { schema })).toThrow( + "Unsupported object value for field image at row 0.", + ); + }); + it("coerces Buffer input onto a blob field", () => { const schema = new Schema([ new Field("id", new Int64(), true), diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 4142af696..ac8904c16 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2412,6 +2412,23 @@ describe("when dealing with blob columns", () => { tmpDir.removeCallback(); }); + it("rejects unsupported blob values in createTable and add", async () => { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + blob("image"), + ]); + const row = { id: 1n, image: new Uint8Array([104]).buffer }; + + await expect(db.createTable("invalid", [row], { schema })).rejects.toThrow( + /field image at row 0/, + ); + + const table = await db.createEmptyTable("empty", schema); + await expect(table.add([row])).rejects.toThrow(/field image at row 0/); + await expect(table.countRows()).resolves.toBe(0); + }); + it("discovers blob columns", async () => { const { table } = await openBlobTable(); expect(await table.blobColumns()).toEqual(["image"]); diff --git a/nodejs/lancedb/schema.ts b/nodejs/lancedb/schema.ts index 3a3ee9316..724a81fc9 100644 --- a/nodejs/lancedb/schema.ts +++ b/nodejs/lancedb/schema.ts @@ -62,6 +62,8 @@ class SchemaInferrer { } private observe(path: string[], value: unknown, row: number): void { + assertSupportedValue(value, path.join("."), row); + const current = this.fields.get(path); if (current === undefined) { this.addField(path, value, row); @@ -430,19 +432,57 @@ function* recordPathsAndValues( } function isRecord(value: unknown): value is Record { + if (typeof value !== "object" || value === null) { + return false; + } + const prototype = Object.getPrototypeOf(value); return ( - typeof value === "object" && - value !== null && - !Array.isArray(value) && - !(value instanceof RegExp) && - !(value instanceof Date) && - !(value instanceof Set) && - !(value instanceof Map) && - !(value instanceof Buffer) && - !ArrayBuffer.isView(value) + (prototype === Object.prototype || prototype === null) && + Object.keys(value).length > 0 ); } +function isUnsupportedObject(value: unknown): boolean { + if ( + typeof value !== "object" || + value === null || + Array.isArray(value) || + isRecord(value) + ) { + return false; + } + // Preserve the object types already passed through to Arrow. Other objects + // must not disappear while walking struct fields. + return !( + value instanceof RegExp || + value instanceof Date || + value instanceof Set || + value instanceof Map || + Buffer.isBuffer(value) || + ArrayBuffer.isView(value) + ); +} + +function assertSupportedValue( + value: unknown, + field: string, + row: number, +): void { + if (Array.isArray(value)) { + for (const [index, item] of value.entries()) { + assertSupportedValue(item, `${field}[${index}]`, row); + } + } else if (isRecord(value)) { + for (const [name, child] of Object.entries(value)) { + assertSupportedValue(child, `${field}.${name}`, row); + } + } else if (isUnsupportedObject(value)) { + throw new Error( + `Unsupported object value for field ${field} at row ${row}.`, + ); + } +} + function fieldAtPath(schema: Schema, path: string[]): Field | undefined { let fields = schema.fields; let field: Field | undefined; diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 88739355c..6235b56b6 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index a095b0cf4..21748b27d 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index edda19630..f82fa2587 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index 2d7e4b3a3..9103dea1c 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index c3060c5f4..2cc815310 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index c75b22291..7ae321f4a 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 95c8b3838..10a83fd86 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 9da9a073f..e6b525654 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.40.0-beta.7", + "version": "0.40.0-beta.8", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index 8f0ca3ec0..f1abd51be 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.40.0-beta.7" +version = "0.40.0-beta.8" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 21c8a31e0..380625f57 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -20,6 +20,7 @@ import importlib import inspect import symtable import json +import linecache import math import re import sys @@ -367,14 +368,20 @@ class FunctionVersion(_RemoteValue): metadata: Mapping[str, str] disabled: bool - def __call__(self, **inputs: Any) -> FunctionApplication: + def __call__(self, **arguments: Any) -> FunctionApplication: """Bind this exact version to named table columns. Every input must be a direct [lancedb.col][lancedb.expr.col] - reference. The returned application is immutable and retains a - named-struct output as one binding, so every row's sibling values - come from one logical Function evaluation. Map result fields to table - columns with + reference. A Function that declares initialization fields takes their + values as further keyword arguments; they are constants of this + binding, fixed for every row, and every instance the binding runs is + created with them. A non-nullable field without a default must be + given; an omitted nullable field is null, and the Function's own + default applies to a null value. + + The returned application is immutable and retains a named-struct + output as one binding, so every row's sibling values come from one + logical Function evaluation. Map result fields to table columns with [FunctionApplication.rename][lancedb.functions.FunctionApplication.rename], then pass the application to [Table.add_columns][lancedb.table.Table.add_columns]. @@ -390,23 +397,34 @@ class FunctionVersion(_RemoteValue): ... "token_count": "search_token_count", ... }) >>> table.add_columns(application) # doctest: +SKIP + >>> table.add_columns({ # doctest: +SKIP + ... "embedding": embed(text=col("body"), model="small", dimensions=512) + ... }) """ from lancedb.expr import Expr parameters = tuple(parameter.name for parameter in self.signature.inputs) - missing = [parameter for parameter in parameters if parameter not in inputs] - unknown = sorted(set(inputs) - set(parameters)) - if missing or unknown: + fields = {field.name: field for field in self.signature.initialization} + missing = [parameter for parameter in parameters if parameter not in arguments] + missing_initialization = [ + name + for name, field in fields.items() + if not field.nullable and name not in arguments + ] + unknown = sorted(set(arguments) - set(parameters) - set(fields)) + if missing or missing_initialization or unknown: details = [] if missing: details.append(f"missing inputs: {missing!r}") + if missing_initialization: + details.append(f"missing initialization: {missing_initialization!r}") if unknown: - details.append(f"unknown inputs: {unknown!r}") - raise TypeError("invalid Function inputs (" + "; ".join(details) + ")") + details.append(f"unknown arguments: {unknown!r}") + raise TypeError("invalid Function arguments (" + "; ".join(details) + ")") bindings = [] for parameter in parameters: - value = inputs[parameter] + value = arguments[parameter] if not isinstance(value, Expr) or value._column_path is None: raise TypeError( f"Function input {parameter!r} must be a direct col(...) reference" @@ -418,6 +436,10 @@ class FunctionVersion(_RemoteValue): value={"path": value._column_path}, ) ) + initialization = {} + for name in fields: + if name in arguments: + initialization[name] = _initialization_value(name, arguments[name]) return FunctionApplication( function=FunctionVersionRef( name=self.name, @@ -428,9 +450,40 @@ class FunctionVersion(_RemoteValue): ), inputs=tuple(bindings), output=self.signature.output, + initialization=initialization, ) +def _initialization_value(name: str, value: Any) -> Any: + """One initialization value as plain JSON; the service checks its type.""" + from lancedb.expr import Expr + + def plain(value): + if isinstance(value, Expr): + raise TypeError( + f"initialization field {name!r} takes a constant, not a column " + "expression; every row of a binding shares one initialization" + ) + if value is None or type(value) in (bool, int, str): + return value + if type(value) is float: + if not math.isfinite(value): + raise ValueError( + f"initialization field {name!r} must be finite, got {value!r}" + ) + return value + if isinstance(value, (list, tuple)): + return [plain(child) for child in value] + if isinstance(value, Mapping) and all(isinstance(key, str) for key in value): + return {key: plain(child) for key, child in value.items()} + raise TypeError( + f"initialization field {name!r} takes booleans, numbers, strings, " + f"lists, and string-keyed mappings, got {type(value).__name__}" + ) + + return plain(value) + + class FunctionRegistrationRequest(_RemoteValue): """Stable remote registration envelope produced by :func:`udf`. @@ -919,7 +972,12 @@ def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]: raise TypeError( "Annotated Function types require exactly one PyArrow DataType" ) - _, base_nullable = _annotation_type(base) + # A struct value is written as a dict, which has no Arrow type of its + # own; the Annotated metadata supplies it. + if base is dict or get_origin(base) is dict: + base_nullable = False + else: + _, base_nullable = _annotation_type(base) return arrow_types[0], nullable or base_nullable if isinstance(annotation, pa.DataType): @@ -1022,8 +1080,12 @@ def _infer_signature( function: Callable[..., Any], input_schema: Optional[pa.Schema], output_schema: Optional[pa.DataType | pa.Field | pa.Schema], + *, + method: bool = False, ) -> FunctionSignature: parameters = _callable_parameters(function) + if method: + parameters = parameters[1:] if (input_schema is None) != (output_schema is None): raise ValueError("input_schema and output_schema must be provided together") @@ -1135,10 +1197,11 @@ def _namespace_acquisition( def _module_references(module_source: str) -> set[str]: - """Names any scope in `module_source` binds or loads at module scope. - Python's own scope analysis on the exact text that ships: free variables - belong to an enclosing scope inside the function, and postponed - annotations are not runtime loads.""" + """Names `module_source` loads from module scope, by Python's own scope + analysis on the exact text that ships: loads inside any definition, plus + what the definition statement itself evaluates at module scope (bases, + decorators, defaults). Free variables belong to an enclosing scope inside + the definition, and postponed annotations are not runtime loads.""" def visit(table: symtable.SymbolTable, found: set[str]) -> None: for symbol in table.get_symbols(): @@ -1150,138 +1213,522 @@ def _module_references(module_source: str) -> set[str]: visit(child, found) found: set[str] = set() - for table in symtable.symtable(module_source, "", "exec").get_children(): + module = symtable.symtable(module_source, "", "exec") + for symbol in module.get_symbols(): + if ( + symbol.is_referenced() + and not symbol.is_assigned() + and not symbol.is_imported() + ): + found.add(symbol.get_name()) + for table in module.get_children(): visit(table, found) return found -def _global_source(name: str, value: Any) -> str: - """One module-level line that rebinds `name` to `value` in the artifact: - an import for modules and importable classes/functions, a literal otherwise.""" - if isinstance(value, types.ModuleType): - if value.__name__.split(".")[0] in _NAMESPACE_MODULES: - raise ValueError( - f"@udf cannot package dynamic namespace access: {value.__name__!r}" - ) +_SOURCE_HEADER = "from __future__ import annotations" +_SOURCE_ENTRY_FILE = "user_function.py" +_PYTHON_CALLABLE_ARTIFACT = "python_callable" +_PYTHON_BUNDLE_ARTIFACT = "python_bundle" +# Modules the Function image places beside user code. +_RESERVED_SOURCE_MODULES = frozenset( + {"user_function", "lance_source", "lance_callable"} +) +# The service accepts bundle paths made of these components only. +_SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +_REQUIREMENT_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") + + +def _class_source(cls: type) -> str: + """Source of a module-level class. + + `inspect.getsource` finds a class through its module's file, which a + notebook cell or a doctest does not have. A method's code object records + where the class body really is, so that is consulted first. + """ + for member in vars(cls).values(): + if isinstance(member, (staticmethod, classmethod)): + member = member.__func__ + code = ( + getattr(inspect.unwrap(member), "__code__", None) + if callable(member) + else None + ) + if code is None: + continue + lines = linecache.getlines(code.co_filename) try: - imported = importlib.import_module(value.__name__) - except ImportError: - imported = None - if imported is not value: - raise TypeError( - f"Function source references module {name!r} that does not import " - f"as {value.__name__!r}" - ) - return f"import {value.__name__} as {name}" - module_name = getattr(value, "__module__", None) - qualname = getattr(value, "__qualname__", None) - if ( - isinstance(module_name, str) - and isinstance(qualname, str) - and module_name != "__main__" - and "." not in qualname - and "<" not in qualname - ): - try: - imported = getattr(importlib.import_module(module_name), qualname) - except (ImportError, AttributeError): - imported = None - if imported is value: - return f"from {module_name} import {qualname} as {name}" - return f"{name} = {_literal_source(value)}" + module = ast.parse("".join(lines)) + except SyntaxError: + continue + for node in module.body: + if ( + isinstance(node, ast.ClassDef) + and node.name == cls.__name__ + and node.lineno <= code.co_firstlineno <= node.end_lineno + ): + start = min( + [node.lineno] + + [decorator.lineno for decorator in node.decorator_list] + ) + return "".join(lines[start - 1 : node.end_lineno]) + return inspect.getsource(cls) -def _is_recursive_reference(function: Callable[..., Any], name: str) -> bool: - """`name` inside the body means the function itself unless the module has - since bound it to something else.""" - if name != function.__name__: - return False - bound = function.__globals__.get(name, function) - if bound is function: - return True - # The decorator's own result is the one wrapper known to call `function` - # unchanged; any other binding may behave differently from a self-call. - return type(bound) is UdfDefinition and bound._function is function - - -def _package_source(function: Callable[..., Any]) -> bytes: - if not inspect.isfunction(function) or inspect.iscoroutinefunction(function): - raise TypeError("@udf requires a synchronous Python function") +def _definition_node( + target: Any, *, entry: bool +) -> Union[ast.FunctionDef, ast.ClassDef]: try: - source = textwrap.dedent(inspect.getsource(function)) + source = textwrap.dedent( + _class_source(target) + if inspect.isclass(target) + else inspect.getsource(target) + ) except (OSError, TypeError) as error: - raise ValueError("@udf requires inspectable Python source") from error + raise ValueError( + f"@udf requires inspectable Python source for {target.__qualname__!r}" + ) from error module = ast.parse(source) + is_class = inspect.isclass(target) definitions = [ node for node in module.body - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) - and node.name == function.__name__ + if isinstance( + node, + (ast.ClassDef,) if is_class else (ast.FunctionDef, ast.AsyncFunctionDef), + ) + and node.name == target.__name__ ] - if len(definitions) != 1 or not isinstance(definitions[0], ast.FunctionDef): - raise ValueError("@udf source must contain exactly one synchronous function") + if len(definitions) != 1 or isinstance(definitions[0], ast.AsyncFunctionDef): + raise ValueError( + "@udf source must contain exactly one " + + ("class" if is_class else "synchronous function") + ) definition = definitions[0] - if any(not _is_udf_decorator(decorator) for decorator in definition.decorator_list): - raise ValueError("@udf cannot package additional Python decorators") - definition.decorator_list = [] + if entry: + if any( + not _is_udf_decorator(decorator) for decorator in definition.decorator_list + ): + raise ValueError("@udf cannot package additional Python decorators") + definition.decorator_list = [] + return definition - closure = inspect.getclosurevars(function) - if closure.nonlocals: - raise ValueError("@udf cannot package functions that capture closure values") - function_source = ast.unparse(definition) - module_header = "from __future__ import annotations" - references = _module_references(f"{module_header}\n\n{function_source}\n") - dynamic = _namespace_acquisition(definition, references) - if dynamic: - raise ValueError(f"@udf cannot package dynamic namespace access: {dynamic!r}") - # Resolve every module-scope reference the way the interpreter would: the - # function's own globals first (a module global may shadow a builtin, and - # nested scopes are not visible to getclosurevars), then its builtins. + +def _definition_namespace(target: Any) -> Mapping[str, Any]: + """The module namespace a function or class resolves its globals in.""" + if inspect.isclass(target): + module = sys.modules.get(target.__module__) + if module is None: + raise ValueError( + f"@udf cannot resolve the module of {target.__qualname__!r}" + ) + namespace = vars(module) + environment = namespace.get("__builtins__", builtins) + environment = ( + vars(environment) + if isinstance(environment, types.ModuleType) + else environment + ) + else: + namespace = target.__globals__ + environment = target.__builtins__ # The artifact runs under the standard builtins; only the exact mapping is # provably equivalent (a subclass or copy can change lookups and hooks). - if function.__builtins__ is not vars(builtins): + if environment is not vars(builtins): raise ValueError("@udf cannot package a non-standard builtins environment") - globals_source = [] - unresolved = [] - for name in sorted(references): - if name == function.__name__: - if not _is_recursive_reference(function, name): - raise ValueError( - f"@udf cannot package {name!r}: the module binds that name to " - "another value, which the artifact's own definition would shadow" - ) - continue - if name in function.__globals__: - globals_source.append(_global_source(name, function.__globals__[name])) - elif hasattr(builtins, name): - pass - else: - unresolved.append(name) - if unresolved: + return namespace + + +def _is_self_reference(target: Any, namespace: Mapping[str, Any], name: str) -> bool: + """`name` inside the definition means the definition itself unless the + module has since bound it to something else.""" + if name != target.__name__: + return False + bound = namespace.get(name, target) + if bound is target: + return True + # The decorator's own result is the one wrapper known to behave like + # `target`; any other binding may differ from the packaged definition. + return type(bound) is UdfDefinition and bound._function is target + + +def _packaged_by_source(value: Any) -> bool: + """Functions and classes defined in `__main__` -- notebook cells and + scripts -- have no module a worker can import, so their source travels.""" + if not (inspect.isfunction(value) or inspect.isclass(value)): + return False + if getattr(value, "__module__", None) != "__main__": + return False + qualname = value.__qualname__ + if "<" in qualname or "." in qualname: raise ValueError( - f"@udf source contains unresolved global names: {unresolved!r}" + f"@udf cannot package {qualname!r}: only module-level functions and " + "classes travel with a Function, not lambdas, closures, or nested " + "definitions" + ) + return True + + +def _requirement_name(requirement: str) -> str: + requirement = requirement.rpartition("::")[2] + match = _REQUIREMENT_NAME.match(requirement.strip()) + return _normalized_distribution(match.group(0) if match else requirement) + + +def _normalized_distribution(name: str) -> str: + return re.sub(r"[-_.]+", "_", name).lower() + + +def _installed_roots() -> tuple[str, ...]: + import os + import site + import sysconfig + + roots = set(site.getsitepackages()) + roots.add(site.getusersitepackages()) + paths = sysconfig.get_paths() + for key in ("purelib", "platlib", "stdlib", "platstdlib"): + if key in paths: + roots.add(paths[key]) + return tuple(os.path.realpath(root) + os.sep for root in roots) + + +def _module_locations(module: types.ModuleType) -> tuple[str, ...]: + import os + + file = getattr(module, "__file__", None) + if file: + return (os.path.realpath(file),) + return tuple(os.path.realpath(path) for path in getattr(module, "__path__", ())) + + +class _SourcePackager: + """Builds the entry module of a Function artifact. + + The module holds the definition plus exactly the module-level names it + references: modules as imports, importable classes and functions as + imports, literals inline, and functions and classes defined in + ``__main__`` by source, recursively, dependencies first. + """ + + def __init__(self, code_modules: frozenset[str], requirements: frozenset[str]): + self._code_modules = code_modules + self._requirements = requirements + self._imports: dict[str, str] = {} + self._definitions: list[str] = [] + self._packaged: set[int] = set() + self._names: dict[str, Any] = {} + + def package(self, target: Any) -> str: + entry = self._definition(target, entry=True) + parts = [_SOURCE_HEADER] + if self._imports: + parts.extend(["", *(self._imports[name] for name in sorted(self._imports))]) + for definition in self._definitions: + parts.extend(["", definition]) + parts.extend(["", entry, ""]) + return "\n".join(parts) + + def _definition(self, target: Any, *, entry: bool) -> str: + definition = _definition_node(target, entry=entry) + if not inspect.isclass(target) and inspect.getclosurevars(target).nonlocals: + raise ValueError( + "@udf cannot package functions that capture closure values" + ) + source = ast.unparse(definition) + references = _module_references(f"{_SOURCE_HEADER}\n\n{source}\n") + dynamic = _namespace_acquisition(definition, references) + if dynamic: + raise ValueError( + f"@udf cannot package dynamic namespace access: {dynamic!r}" + ) + namespace = _definition_namespace(target) + self._packaged.add(id(target)) + self._claim(target.__name__, ("definition", id(target))) + unresolved = [] + for name in sorted(references): + if name == target.__name__: + if not _is_self_reference(target, namespace, name): + raise ValueError( + f"@udf cannot package {name!r}: the module binds that name " + "to another value, which the artifact's own definition " + "would shadow" + ) + continue + if name in namespace: + self._bind(name, namespace[name]) + elif not hasattr(builtins, name): + unresolved.append(name) + if unresolved: + raise ValueError( + f"@udf source contains unresolved global names: {unresolved!r}" + ) + return source + + def _claim(self, name: str, meaning: Any) -> None: + claimed = self._names.setdefault(name, meaning) + if claimed != meaning: + raise ValueError( + f"@udf cannot package {name!r}: packaged definitions bind that " + "name to different values" + ) + + def _bind(self, name: str, value: Any) -> None: + if _packaged_by_source(value): + if value.__name__ != name: + raise ValueError( + f"@udf packages {value.__qualname__!r} by source, so it must be " + f"referenced by its own name, not {name!r}" + ) + if id(value) not in self._packaged: + self._definitions.append(self._definition(value, entry=False)) + self._claim(name, ("definition", id(value))) + return + line = self._global_source(name, value) + self._claim(name, line) + self._imports[name] = line + + def _global_source(self, name: str, value: Any) -> str: + """One module-level line that rebinds `name` to `value` in the artifact: + an import for modules and importable classes/functions, a literal + otherwise.""" + if isinstance(value, types.ModuleType): + if value.__name__.split(".")[0] in _NAMESPACE_MODULES: + raise ValueError( + f"@udf cannot package dynamic namespace access: {value.__name__!r}" + ) + try: + imported = importlib.import_module(value.__name__) + except ImportError: + imported = None + if imported is not value: + raise TypeError( + f"Function source references module {name!r} that does not " + f"import as {value.__name__!r}" + ) + self._require_importable(value.__name__) + return f"import {value.__name__} as {name}" + module_name = getattr(value, "__module__", None) + qualname = getattr(value, "__qualname__", None) + if ( + isinstance(module_name, str) + and isinstance(qualname, str) + and module_name != "__main__" + and "." not in qualname + and "<" not in qualname + ): + try: + imported = getattr(importlib.import_module(module_name), qualname) + except (ImportError, AttributeError): + imported = None + if imported is value: + self._require_importable(module_name) + return f"from {module_name} import {qualname} as {name}" + return f"{name} = {_literal_source(value)}" + + def _require_importable(self, module_name: str) -> None: + """Refuse an import the Function's environment has no way to satisfy: + a module that lives in a local source tree, is not shipped with + ``code=``, and names no declared package.""" + top = module_name.partition(".")[0] + if ( + top in self._code_modules + or top in sys.stdlib_module_names + or top in sys.builtin_module_names + or _normalized_distribution(top) in self._requirements + ): + return + module = sys.modules.get(top) + locations = () if module is None else _module_locations(module) + roots = _installed_roots() + if not locations or any(location.startswith(roots) for location in locations): + return + raise ValueError( + f"@udf source imports {module_name!r} from {locations[0]}, a local " + "module the Function's environment cannot import; ship it with " + f"code=[{top}] or declare the package that provides it in pip/conda" ) - parts = [module_header] - if globals_source: - parts.extend(["", *globals_source]) - parts.extend(["", function_source, ""]) - packaged = "\n".join(parts) - return packaged.encode("utf-8") + +def _code_files(modules: Sequence[Any]) -> tuple[dict[str, str], frozenset[str]]: + """The Python source files of the modules and packages in ``code=``, + keyed by their path on the Function's import path.""" + import os + from pathlib import Path + + files: dict[str, str] = {} + names: set[str] = set() + for module in modules: + if not isinstance(module, types.ModuleType): + raise TypeError( + f"code= takes imported modules or packages, got {type(module).__name__}" + ) + name = module.__name__ + if "." in name: + raise ValueError( + f"code= takes top-level modules or packages; pass " + f"{name.partition('.')[0]!r} instead of {name!r}" + ) + if ( + name in _RESERVED_SOURCE_MODULES + or name == "__main__" + or name in sys.stdlib_module_names + or not _SOURCE_PATH_COMPONENT.fullmatch(name) + ): + raise ValueError(f"code= cannot ship module {name!r}") + if name in names: + raise ValueError(f"code= lists module {name!r} more than once") + names.add(name) + paths = list(getattr(module, "__path__", ())) + if paths: + if len(paths) != 1 or not getattr(module, "__file__", None): + raise ValueError( + f"code= cannot ship namespace package {name!r}; give it an " + "__init__.py" + ) + root = Path(paths[0]) + sources = [] + for directory, subdirectories, filenames in os.walk(root): + # Only importable modules travel: data files, caches, and + # directories that are not identifiers cannot be imported. + subdirectories[:] = sorted( + child + for child in subdirectories + if _SOURCE_PATH_COMPONENT.fullmatch(child) + and child != "__pycache__" + ) + for filename in filenames: + stem, suffix = os.path.splitext(filename) + if suffix == ".py" and _SOURCE_PATH_COMPONENT.fullmatch(stem): + sources.append(Path(directory) / filename) + for path in sources: + relative = path.relative_to(root).as_posix() + files[f"{name}/{relative}"] = _source_text(path) + else: + file = getattr(module, "__file__", None) + if not file or not file.endswith(".py"): + raise ValueError(f"code= ships Python source; module {name!r} has none") + files[f"{name}.py"] = _source_text(Path(file)) + return files, frozenset(names) + + +def _source_text(path) -> str: + try: + return path.read_text(encoding="utf-8") + except UnicodeDecodeError as error: + raise ValueError(f"code= ships UTF-8 source; {path} is not UTF-8") from error + + +def _initialization_type(data_type: pa.DataType) -> bool: + """Types with one unambiguous JSON and SQL literal spelling.""" + if data_type in ( + pa.bool_(), + pa.int8(), + pa.int16(), + pa.int32(), + pa.int64(), + pa.uint8(), + pa.uint16(), + pa.uint32(), + pa.uint64(), + pa.float32(), + pa.float64(), + pa.string(), + pa.large_string(), + ): + return True + if ( + pa.types.is_list(data_type) + or pa.types.is_large_list(data_type) + or pa.types.is_fixed_size_list(data_type) + ): + return _initialization_type(data_type.value_type) + if pa.types.is_struct(data_type): + return all(_initialization_type(field.type) for field in data_type) + return False + + +def _initialization_signature(cls: type) -> tuple[FunctionParameter, ...]: + """The initialization fields a class Function's ``__init__`` declares.""" + if cls.__init__ is object.__init__: + return () + parameters = tuple(inspect.signature(cls.__init__).parameters.values())[1:] + for parameter in parameters: + if parameter.kind not in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ): + raise TypeError( + "Function initialization parameters must be named and " + f"non-variadic: {parameter.name!r}" + ) + try: + annotations = get_type_hints(cls.__init__, include_extras=True) + except Exception as error: + raise TypeError( + f"failed to resolve Function initialization annotations: {error}" + ) from error + missing = [ + parameter.name for parameter in parameters if parameter.name not in annotations + ] + if missing: + raise TypeError(f"missing Function initialization annotations: {missing!r}") + fields = [] + for parameter in parameters: + data_type, nullable = _annotation_type(annotations[parameter.name]) + if not _initialization_type(data_type): + raise TypeError( + f"Function initialization parameter {parameter.name!r} has type " + f"{data_type}; initialization takes booleans, integers, floats, " + "strings, and lists or structs of them" + ) + # A parameter with a default may be omitted from a binding; a null + # value then takes the default. + if parameter.default is not inspect.Parameter.empty: + nullable = True + fields.append( + FunctionParameter( + name=parameter.name, + arrow_type=_canonical_arrow_field( + pa.field(parameter.name, data_type, nullable=nullable) + ), + nullable=nullable, + ) + ) + return tuple(fields) + + +def _class_member(cls: type, name: str) -> Any: + for klass in cls.__mro__: + if klass is not object and name in vars(klass): + return vars(klass)[name] + return None + + +def _class_call(cls: type) -> Callable[..., Any]: + call = _class_member(cls, "__call__") + if not inspect.isfunction(call) or inspect.iscoroutinefunction(call): + raise TypeError("a class Function must define a synchronous __call__ method") + close = _class_member(cls, "close") + if close is not None and ( + not inspect.isfunction(close) or len(inspect.signature(close).parameters) != 1 + ): + raise TypeError("a class Function's close must be a method taking no arguments") + return call class UdfDefinition: - """A scalar Python callable prepared for remote Function registration. + """A Python callable prepared for remote Function registration. Instances are created with :func:`udf`. Calling an instance executes the - original scalar Python function, which keeps local unit testing ordinary. - Remote execution adapts that scalar callable to the internal Arrow batch - ABI described by the registration artifact. + original Python function, or constructs the original class, which keeps + local unit testing ordinary. Remote execution adapts the callable to the + internal Arrow batch ABI described by the registration artifact. """ def __init__( self, - function: Callable[..., Any], + function: Any, *, name: Optional[str], input_schema: Optional[pa.Schema], @@ -1292,7 +1739,13 @@ class UdfDefinition: gpu: bool = False, conda: tuple[str, ...] = (), conda_channels: tuple[str, ...] = (), + code: tuple[types.ModuleType, ...] = (), ): + is_class = inspect.isclass(function) + if not is_class and ( + not inspect.isfunction(function) or inspect.iscoroutinefunction(function) + ): + raise TypeError("@udf requires a synchronous Python function or a class") function_name = name or function.__name__ if not _FUNCTION_NAME.fullmatch(function_name): raise ValueError(f"invalid Function name: {function_name!r}") @@ -1315,8 +1768,43 @@ class UdfDefinition: for key, value in environment.items() ): raise TypeError("Function env keys and values must be strings") - signature = _infer_signature(function, input_schema, output_schema) - source = _package_source(function) + if is_class: + if "<" in function.__qualname__ or "." in function.__qualname__: + raise ValueError( + "@udf classes must be defined at module level; nested classes " + "would capture their enclosing scope" + ) + signature = _infer_signature( + _class_call(function), input_schema, output_schema, method=True + ) + initialization = _initialization_signature(function) + shared = sorted( + {field.name for field in initialization} + & {parameter.name for parameter in signature.inputs} + ) + if shared: + raise ValueError( + "Function initialization parameters and inputs share names " + f"{shared!r}; a binding passes both as keyword arguments" + ) + signature = signature._copy(update={"initialization": initialization}) + else: + signature = _infer_signature(function, input_schema, output_schema) + code_files, code_modules = _code_files(code) + entry_source = _SourcePackager( + code_modules, frozenset(_requirement_name(package) for package in packages) + ).package(function) + if code_files: + kind = _PYTHON_BUNDLE_ARTIFACT + source = json.dumps( + {"files": {**code_files, _SOURCE_ENTRY_FILE: entry_source}}, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + else: + kind = _PYTHON_CALLABLE_ARTIFACT + source = entry_source.encode("utf-8") digest = f"sha256:{hashlib.sha256(source).hexdigest()}" gpu_marker = _normalize_gpu_marker(gpu) runtime = PythonRuntimeSpec( @@ -1331,7 +1819,7 @@ class UdfDefinition: self._request = FunctionRegistrationRequest( name=function_name, artifact=FunctionArtifactRequest( - kind="python_callable", + kind=kind, digest=digest, entrypoint=function.__name__, content=FunctionArtifactContent( @@ -1346,7 +1834,11 @@ class UdfDefinition: signature=signature, runtime=runtime, ) - functools.update_wrapper(self, function) + # A class's own __dict__ holds its methods; copying it would shadow + # this wrapper's attributes. + functools.update_wrapper( + self, function, updated=() if is_class else functools.WRAPPER_UPDATES + ) @property def registration_request(self) -> FunctionRegistrationRequest: @@ -1438,6 +1930,7 @@ def udf( gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), + code: Sequence[types.ModuleType] = (), ) -> Callable[[Callable[..., Any]], UdfDefinition]: ... @@ -1453,8 +1946,9 @@ def udf( gpu: bool = False, conda: tuple[str, ...] | list[str] = (), conda_channels: tuple[str, ...] | list[str] = (), + code: Sequence[types.ModuleType] = (), ): - """Prepare a scalar Python callable for remote Function registration. + """Prepare a Python function or class for remote Function registration. Input and output signatures are inferred from supported annotations. For Arrow types annotations cannot express precisely, pass ``input_schema`` @@ -1462,10 +1956,54 @@ def udf( named-struct field may be nullable; Enterprise preserves the struct's validity when the result is expanded into sibling columns. + **Functions and classes.** A decorated function is called once per row, + or once per batch when its parameters are annotated ``pyarrow.Array`` or + it takes one ``pyarrow.RecordBatch``. A decorated class is the same + callable with state: each remote instance calls ``__init__`` once, calls + ``__call__`` for every row or batch of every input it processes, and calls + ``close()``, if defined, once when it retires. Load models, open clients, + and build rate limiters in ``__init__``; every batch of the instance then + reuses them. An instance serves many batches of one binding, so state is + shared across them, but not across instances or workers. + + **Initialization.** The annotated parameters of a class's ``__init__`` + are the Function's initialization fields. Their values belong to each + binding, not to the Function: ``function(text=col("body"), model="small")`` + binds column ``body`` and initializes every instance with ``model="small"``, + and another column can bind the same Function version with other values. + Initialization fields take booleans, integers, floats, strings, and lists + or structs of them (``Annotated[dict, pa.struct(...)]``). A parameter with + a default may be omitted from a binding; a null value then takes the + default. Use Secrets, not initialization, for credentials. + + **What travels with the Function.** The artifact is a snapshot of the + decorated source plus exactly the module-level names it references: + + - modules, and classes and functions importable from an installed + package, become imports; the package must be declared in ``pip`` or + ``conda``; + - literals (``None``, booleans, numbers, strings, bytes, tuples of them) + are inlined; + - functions and classes defined in ``__main__``, such as notebook cells, + are packaged by source, together with what they reference; + - modules and packages listed in ``code`` are shipped as Python source and + imported normally, so helpers shared by several Functions live in one + place. + + A reference to a module in a local source tree that is neither shipped + with ``code`` nor provided by a declared package is rejected at + registration. Closures, lambdas, nested definitions, mutable module-level + objects, code that reaches the module namespace another way + (``globals()``/``eval``, ``sys.modules``, ``builtins``), and a + non-standard ``__builtins__`` are rejected where they can be seen and + otherwise unsupported. Changing any packaged source, including a module in + ``code``, produces a new Function version; the environment built from + ``pip`` or ``conda`` is reused. + Parameters ---------- - function : Callable, optional - The synchronous scalar callable to package. + function : function or class, optional + The synchronous callable, or the class whose instances are called. name : str, optional The remote Function name. Defaults to the callable name. input_schema : pyarrow.Schema, optional @@ -1483,20 +2021,17 @@ def udf( env : mapping of str to str, optional Environment variables included in the Function definition. Not for credentials -- these are ordinary configuration, stored with the - Function and visible wherever it is. + Function and visible wherever it is. Prefer initialization for + values that differ between bindings. python_version : str, optional Remote Python major/minor version. Defaults to the client version. gpu : bool, default False Whether every remote execution requires a GPU. The execution platform selects one compatible GPU for each worker. The requirement is part of the immutable Function version. - - The packaged artifact is a snapshot: the function source plus exactly - the module-level names it references (modules as imports, importable - classes and functions as imports, literals inline). Code that reaches the - module namespace another way -- ``globals()``/``eval``, ``sys.modules``, - ``builtins`` -- is rejected where it can be seen and otherwise - unsupported; closures and a non-standard ``__builtins__`` are rejected. + code : sequence of module, optional + Top-level modules or packages whose importable ``.py`` files ship + with the Function and are importable by their own names. Returns ------- @@ -1519,9 +2054,19 @@ def udf( ... return value * 2 >>> gpu_score.registration_request.runtime.gpu True + >>> @udf + ... class scale: + ... def __init__(self, factor: float, offset: float = 0.0): + ... self.factor, self.offset = factor, offset + ... def __call__(self, value: float) -> float: + ... return value * self.factor + self.offset + >>> scale(factor=2.0)(1.5) + 3.0 + >>> [field.name for field in scale.registration_request.signature.initialization] + ['factor', 'offset'] """ - def decorate(target: Callable[..., Any]) -> UdfDefinition: + def decorate(target: Any) -> UdfDefinition: return UdfDefinition( target, name=name, @@ -1533,6 +2078,7 @@ def udf( gpu=gpu, conda=tuple(conda), conda_channels=tuple(conda_channels), + code=tuple(code), ) if function is None: diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 297b4237f..d87012bcf 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -771,12 +771,13 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field: if json_storage is not None: # Labelled through metadata rather than pa.json_(), which only exists on # newer PyArrow; Lance reads the extension name off the field either way. - return pa.field( - field.name, - json_storage, - field.nullable, - {"ARROW:extension:name": "arrow.json"}, - ) + # The other metadata keys mirror the table field's: Lance swaps the name + # back to lance.json on write and then requires the field to match the + # stored one exactly, including the empty ``ARROW:extension:metadata`` + # that pyarrow records for a ``pa.json_()`` column. + metadata = dict(target_field.metadata or {}) + metadata[b"ARROW:extension:name"] = b"arrow.json" + return pa.field(field.name, json_storage, field.nullable, metadata) if pa.types.is_struct(target_field.type): if pa.types.is_struct(field.type): new_type = pa.struct( diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 0d3a89c21..cada6e83e 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -921,6 +921,70 @@ def test_fetch_blob_files_lazy_read(): handles = table.fetch_blob_files("image", [by_id[1]]) assert len(handles) == 1 assert handles[0].read() == payload + handles[0].close() + assert handles[0].closed + with pytest.raises(RuntimeError, match="already closed"): + handles[0].read() + + +def test_blob_file_finalizer_and_close_on_runtime_worker(): + # A tokio panic in IOBase.__del__ is printed but does not fail the add call. + # Run in a child process so the assertion catches that panic on stderr. + script = textwrap.dedent( + """\ + import gc + import threading + + import lancedb + import pyarrow as pa + + db = lancedb.connect("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("finalize", schema=schema) + table.add([{"id": 1, "image": b"first"}]) + hits = table.search().with_row_id(True).limit(1).to_arrow() + row_id = hits["_rowid"][0].as_py() + gc_handle, explicit_handle = table.fetch_blob_files("image", [row_id, row_id]) + + gc.disable() + cycle = [gc_handle] + cycle.append(cycle) + del gc_handle, cycle + + observed = [] + def on_progress(_): + if not observed: + observed.append(( + threading.current_thread() is not threading.main_thread(), + explicit_handle.closed, + gc.collect(), + )) + explicit_handle.close() + + table.add([{"id": 2, "image": b"second"}], progress=on_progress) + gc.enable() + + assert len(observed) == 1, observed + assert observed[0][0] is True, observed + assert observed[0][1] is False, observed + assert observed[0][2] > 0, observed + assert explicit_handle.closed + explicit_handle.close() # Close remains idempotent after the worker call. + try: + explicit_handle.read() + except RuntimeError as error: + assert "already closed" in str(error) + else: + raise AssertionError("a closed blob file was readable") + print("worker close completed") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, check=False + ) + assert result.returncode == 0, result.stderr + assert "Cannot start a runtime from within a runtime" not in result.stderr + assert "worker close completed" in result.stdout def test_fetch_blob_files_null_alignment(): diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py index 849923fda..1abd42e56 100644 --- a/python/python/tests/test_first_class_function_slice1.py +++ b/python/python/tests/test_first_class_function_slice1.py @@ -169,7 +169,7 @@ def test_function_version_binding_validates_names_and_direct_columns(): with pytest.raises(TypeError, match=r"missing inputs: \['text'\]"): version() - with pytest.raises(TypeError, match=r"unknown inputs: \['body'\]"): + with pytest.raises(TypeError, match=r"unknown arguments: \['body'\]"): version(text=col("text"), body=col("body")) with pytest.raises(TypeError, match="direct col"): version(text=col("text").lower()) diff --git a/python/python/tests/test_function_packaging.py b/python/python/tests/test_function_packaging.py new file mode 100644 index 000000000..239d07be1 --- /dev/null +++ b/python/python/tests/test_function_packaging.py @@ -0,0 +1,521 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Class Functions, initialization, and the source that travels with a Function.""" + +from __future__ import annotations + +import base64 +import importlib +import json +import subprocess +import sys +import textwrap +from pathlib import Path +from typing import Annotated, Optional + +import pyarrow as pa +import pytest + +from lancedb import col +from lancedb.functions import FunctionVersion, udf + +FIXTURES = ( + Path(__file__).parents[3] + / "rust" + / "lancedb" + / "tests" + / "fixtures" + / "first_class_functions" + / "v1" +) + + +def artifact(definition) -> bytes: + return base64.b64decode(definition.registration_request.artifact.content.data) + + +def bundle_files(definition) -> dict[str, str]: + assert definition.registration_request.artifact.kind == "python_bundle" + return json.loads(artifact(definition))["files"] + + +@pytest.fixture +def helper_package(tmp_path, monkeypatch): + """A local package on sys.path, the way a project's own helpers are.""" + + def write(name: str, files: dict[str, str]): + root = tmp_path / name + for path, source in files.items(): + target = root / path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(textwrap.dedent(source)) + monkeypatch.syspath_prepend(str(tmp_path)) + for module in [ + module for module in sys.modules if module.split(".")[0] == name + ]: + del sys.modules[module] + importlib.invalidate_caches() + return importlib.import_module(name) + + return write + + +@udf +class scale: + def __init__( + self, + factor: float, + label: str, + offset: float = 0.0, + tags: Optional[list[str]] = None, + retry: Annotated[dict, pa.struct([("attempts", pa.int32())])] = None, + ): + self.factor, self.offset = factor, offset + + def __call__(self, value: float) -> float: + return value * self.factor + self.offset + + +class NoCall: + pass + + +class BadClose: + def __call__(self, value: int) -> int: + return value + + def close(self, force): + pass + + +def test_class_function_declares_inputs_output_and_initialization(): + signature = scale.registration_request.signature + assert [(p.name, p.arrow_type) for p in signature.inputs] == [("value", "float64")] + assert [ + (field.name, field.arrow_type, field.nullable) + for field in signature.initialization + ] == [ + ("factor", "float64", False), + ("label", "utf8", False), + ("offset", "float64", True), + ("tags", "list", True), + ( + "retry", + '{"fields":[{"name":"attempts","nullable":true,"type":{"type":"int32"}}],' + '"type":"struct"}', + True, + ), + ] + # Calling the definition constructs the class, so local tests are ordinary. + assert scale(factor=2.0, label="x")(1.5) == 3.0 + assert scale.registration_request.artifact.entrypoint == "scale" + source = artifact(scale).decode() + assert source.startswith("from __future__ import annotations\n") + assert "class scale:" in source and "@udf" not in source + + +def test_function_without_initialization_keeps_its_wire_form(): + @udf + def double(value: int) -> int: + return value * 2 + + request = json.loads(double.registration_request.to_canonical_json()) + assert "initialization" not in request["signature"] + assert request["artifact"]["kind"] == "python_callable" + + +@pytest.mark.parametrize( + ("init", "message"), + [ + ("def __init__(self, blob: bytes): pass", "initialization takes booleans"), + ("def __init__(self, *args: int): pass", "named and non-variadic"), + ("def __init__(self, factor): pass", "missing Function initialization"), + ("def __init__(self, value: int): pass", "share names \\['value'\\]"), + ], +) +def test_class_function_rejects_unsupported_initialization(init, message): + namespace = {} + exec( + textwrap.dedent( + f""" + class bad: + {init} + def __call__(self, value: int) -> int: + return value + """ + ), + namespace, + ) + with pytest.raises((TypeError, ValueError), match=message): + udf(namespace["bad"]) + + +def test_class_function_requires_call_and_a_plain_close(): + with pytest.raises(TypeError, match="synchronous __call__"): + udf(NoCall) + with pytest.raises(TypeError, match="close must be a method taking no"): + udf(BadClose) + + +def test_nested_class_and_closure_are_rejected_with_the_reason(): + def build(): + class nested: + def __call__(self, value: int) -> int: + return value + + return nested + + with pytest.raises(ValueError, match="defined at module level"): + udf(build()) + + factor = 3 + + with pytest.raises(ValueError, match="capture closure values"): + + @udf + def closes_over(value: int) -> int: + return value * factor + + +def test_binding_passes_initialization_as_constants(): + version = FunctionVersion.from_json( + json.dumps( + { + **json.loads( + (FIXTURES / "remote_function_version.canonical.json").read_text() + ), + "signature": { + "inputs": [ + {"name": "text", "arrow_type": "utf8", "nullable": True} + ], + "output": { + "kind": "scalar", + "arrow_type": "list", + "nullable": False, + }, + "initialization": [ + {"name": "model", "arrow_type": "utf8", "nullable": False}, + {"name": "dimensions", "arrow_type": "int32", "nullable": True}, + ], + }, + } + ) + ) + application = version(text=col("body"), model="small", dimensions=(512)) + assert dict(application.initialization) == {"model": "small", "dimensions": 512} + assert json.loads(application.to_canonical_json())["initialization"] == { + "dimensions": 512, + "model": "small", + } + assert version(text=col("body"), model="small").initialization == {"model": "small"} + + with pytest.raises(TypeError, match=r"missing initialization: \['model'\]"): + version(text=col("body")) + with pytest.raises(TypeError, match="takes a constant"): + version(text=col("body"), model=col("model")) + with pytest.raises(ValueError, match="must be finite"): + version(text=col("body"), model="small", dimensions=float("nan")) + with pytest.raises(TypeError, match="takes booleans"): + version(text=col("body"), model=object()) + + +def test_code_packages_ship_as_a_bundle_and_version_with_their_source( + tmp_path, helper_package +): + files = { + "__init__.py": "", + "rate.py": """ + class TokenBucket: + def __init__(self, tokens_per_minute): + self.tokens_per_minute = tokens_per_minute + """, + "__pycache__/rate.cpython-312.pyc": "", + "fixtures/data.json": "{}", + "not-a-module.py": "", + } + helper_package("fx_helpers", files) + (tmp_path / "fx_helpers_functions.py").write_text( + textwrap.dedent( + """ + import fx_helpers + from fx_helpers.rate import TokenBucket + from lancedb import udf + + + @udf(code=[fx_helpers]) + class limited: + def __init__(self, tpm: int): + self.limiter = TokenBucket(tpm) + + def __call__(self, value: int) -> int: + return value + """ + ) + ) + + def define(): + sys.modules.pop("fx_helpers_functions", None) + return importlib.import_module("fx_helpers_functions").limited + + first = define() + shipped = bundle_files(first) + assert sorted(shipped) == [ + "fx_helpers/__init__.py", + "fx_helpers/rate.py", + "user_function.py", + ] + assert ( + "from fx_helpers.rate import TokenBucket as TokenBucket" + in (shipped["user_function.py"]) + ) + + # Only the helper changes: a new artifact digest, the same environment. + helper_package( + "fx_helpers", {**files, "rate.py": textwrap.dedent(files["rate.py"]) + "# v2\n"} + ) + second = define() + assert ( + first.registration_request.artifact.digest + != second.registration_request.artifact.digest + ) + assert first.registration_request.runtime == second.registration_request.runtime + + +def test_code_rejects_submodules_duplicates_and_non_modules(helper_package): + helpers = helper_package("fx_rejects", {"__init__.py": "", "sub.py": ""}) + import fx_rejects.sub + + for code, message in [ + ([fx_rejects.sub], "top-level modules"), + ([helpers, helpers], "more than once"), + (["fx_rejects"], "imported modules"), + ([json], "cannot ship module 'json'"), + ]: + with pytest.raises((TypeError, ValueError), match=message): + + @udf(code=code) + def f(value: int) -> int: + return value + + +def test_local_module_outside_code_is_rejected_at_registration( + tmp_path, helper_package +): + helper_package( + "fx_local", {"__init__.py": "def normalize(value):\n return value\n"} + ) + + def define(name, options): + (tmp_path / f"{name}.py").write_text( + textwrap.dedent( + f""" + from fx_local import normalize + from lancedb import udf + + + @udf({options}) + def {name}(value: int) -> int: + return normalize(value) + """ + ) + ) + return getattr(importlib.import_module(name), name) + + with pytest.raises(ValueError, match=r"local module.*code=\[fx_local\]"): + define("fx_uses_local", "") + + # A declared package that provides the module is trusted to install it. + installed = define("fx_uses_installed", "pip=['fx-local==1.0']") + assert "from fx_local import normalize" in artifact(installed).decode() + + +MMLB_HELPERS = { + "__init__.py": "", + "rate.py": """ + class TokenBucket: + constructed = 0 + + def __init__(self, tokens_per_minute): + TokenBucket.constructed += 1 + self.tokens_per_minute = tokens_per_minute + self.acquired = 0 + + def acquire(self, tokens): + self.acquired += tokens + """, + "embedding.py": """ + from fx_mmlb.rate import TokenBucket + + + class OpenAIEmbedder: + def __init__(self, model: str, tpm_limit: int, max_attempts: int = 6): + self.model = model + self.limiter = TokenBucket(tpm_limit) + self.max_attempts = max_attempts + + def embed(self, texts): + self.limiter.acquire(sum(len(text) for text in texts)) + return [[float(len(text))] for text in texts] + """, +} + +MMLB_FUNCTIONS = """ +import pyarrow as pa + +import fx_mmlb +from fx_mmlb.embedding import OpenAIEmbedder +from lancedb import udf + +TEXT = pa.schema([pa.field("text", pa.string())]) +EMBEDDING = pa.list_(pa.field("item", pa.float64(), nullable=False)) + + +@udf(code=[fx_mmlb], input_schema=TEXT, output_schema=EMBEDDING) +class openai_embedding_batch(OpenAIEmbedder): + def __call__(self, text: pa.Array) -> pa.Array: + item = pa.field("item", pa.float64(), nullable=False) + return pa.array(self.embed(text.to_pylist()), pa.list_(item)) + + +@udf(code=[fx_mmlb]) +class openai_embedding(OpenAIEmbedder): + def __call__(self, text: str) -> list[float]: + return self.embed([text])[0] +""" + + +def test_shared_helpers_construct_state_on_every_path(tmp_path, helper_package): + """The MMLB rate limiter, shared by a row and a batch Function: both are + initialized through the same helper, so neither path can skip it.""" + helper_package("fx_mmlb", MMLB_HELPERS) + (tmp_path / "fx_mmlb_functions.py").write_text(MMLB_FUNCTIONS) + functions = importlib.import_module("fx_mmlb_functions") + + signatures = [ + functions.openai_embedding_batch.registration_request.signature, + functions.openai_embedding.registration_request.signature, + ] + for signature in signatures: + assert [field.name for field in signature.initialization] == [ + "model", + "tpm_limit", + "max_attempts", + ] + + # Each bundle imports on its own, in a process that has never seen the + # client's modules, the way the worker's image code layer does. + for definition, call in [ + ( + functions.openai_embedding_batch, + "instance(pa.array(['ab', 'abc'])).to_pylist()", + ), + (functions.openai_embedding, "[instance('ab'), instance('abc')]"), + ]: + worker = tmp_path / definition.__name__ + for path, source in bundle_files(definition).items(): + (worker / path).parent.mkdir(parents=True, exist_ok=True) + (worker / path).write_text(source) + probe = textwrap.dedent( + f""" + import sys + sys.path.insert(0, {str(worker)!r}) + import pyarrow as pa + import user_function + from fx_mmlb.rate import TokenBucket + cls = getattr(user_function, {definition.__name__!r}) + instance = cls(model="small", tpm_limit=1000) + print(repr(({call}, TokenBucket.constructed, instance.limiter.acquired))) + """ + ) + result = subprocess.run( + [sys.executable, "-I", "-c", probe], + capture_output=True, + text=True, + cwd=worker, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "([[2.0], [3.0]], 1, 5)" + + +def test_main_definitions_travel_by_source(tmp_path): + """Notebook cells and scripts define helpers in __main__, which a worker + cannot import; they are packaged by source, dependencies first.""" + script = tmp_path / "notebook.py" + script.write_text( + textwrap.dedent( + """ + import json + import math + from dataclasses import dataclass + + from lancedb import udf + + SCALE = 2 + + + @dataclass + class Settings: + factor: int + + + def helper(value): + return math.floor(value) * SCALE + + + class Base: + def run(self, value): + return helper(value) * self.settings.factor + + + @udf + class scaled(Base): + def __init__(self, factor: int): + self.settings = Settings(factor) + + def __call__(self, value: float) -> int: + return self.run(value) + + + print(json.dumps(scaled.registration_request.to_canonical_json())) + """ + ) + ) + result = subprocess.run( + [sys.executable, str(script)], capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr + request = json.loads(json.loads(result.stdout)) + assert request["artifact"]["kind"] == "python_callable" + source = base64.b64decode(request["artifact"]["content"]["data"]).decode() + assert source.index("def helper") < source.index("class Base") + assert source.index("class Base") < source.index("class scaled") + assert source.index("class Settings") < source.index("class scaled") + assert "@dataclass" in source + assert "SCALE = 2" in source + + namespace = {} + exec(compile(source, "user_function.py", "exec"), namespace) + assert namespace["scaled"](factor=3)(2.7) == 12 + + +def test_main_lambdas_are_rejected(tmp_path): + script = tmp_path / "lambda.py" + script.write_text( + textwrap.dedent( + """ + from lancedb import udf + + double = lambda value: value * 2 + + @udf + def uses_lambda(value: int) -> int: + return double(value) + """ + ) + ) + result = subprocess.run( + [sys.executable, str(script)], capture_output=True, text=True + ) + assert result.returncode != 0 + assert "not lambdas, closures, or nested definitions" in result.stderr diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index c03862a49..0a357839a 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -2470,6 +2470,10 @@ def test_remote_blob_files_are_lazy_seekable_handles(): assert alpha.read_range(1, 3) == b"lph" gamma.seek(2) assert gamma.read() == b"mma" + alpha.close() + assert alpha.closed + with pytest.raises(RuntimeError, match="already closed"): + alpha.read_range(0, 1) def test_remote_blob_fetch_accepts_query_table(): diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index acdef8eac..29cfd7f0b 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -948,6 +948,55 @@ def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct(): ] +def test_cast_to_target_schema_labels_json_text_with_table_extension_metadata(): + # A table created from pa.json_() stores the empty ARROW:extension:metadata that + # pyarrow exports next to the name. Lance keeps that key when it turns the + # arrow.json label back into lance.json on write and then requires the field to + # match the stored one exactly, so the label has to mirror the table's metadata. + target = pa.schema( + [ + pa.field( + "payload", + pa.large_binary(), + metadata={ + b"ARROW:extension:name": b"lance.json", + b"ARROW:extension:metadata": b"", + }, + ), + pa.field( + "docs", + pa.list_( + pa.field( + "item", + pa.large_binary(), + metadata={b"ARROW:extension:name": b"lance.json"}, + ) + ), + ), + ] + ) + data = pa.table( + { + "payload": pa.array(['{"a": 1}', None], type=pa.string()), + "docs": pa.array([['{"b": 2}'], []], type=pa.list_(pa.string())), + } + ) + + output = _cast_to_target_schema(data.to_reader(), target).read_all() + + payload = output.schema.field("payload") + assert payload.type == pa.string() + assert payload.metadata == { + b"ARROW:extension:name": b"arrow.json", + b"ARROW:extension:metadata": b"", + } + docs_item = output.schema.field("docs").type.value_field + assert docs_item.type == pa.string() + assert docs_item.metadata == {b"ARROW:extension:name": b"arrow.json"} + assert output["payload"].to_pylist() == ['{"a": 1}', None] + assert output["docs"].to_pylist() == [['{"b": 2}'], []] + + def test_cast_to_target_schema_coerces_nested_binary_blob(): data = pa.table( { diff --git a/python/src/runtime.rs b/python/src/runtime.rs index 170f7e592..33ffd8bf9 100644 --- a/python/src/runtime.rs +++ b/python/src/runtime.rs @@ -197,6 +197,27 @@ pub fn block_on(fut: F) -> F::Output { get_runtime().block_on(fut) } +/// Run a detached task, including when the caller is already on a runtime worker. +/// Keep it visible to [`shutdown`] until it finishes. +pub fn spawn_background(fut: F) +where + F: Future + Send + 'static, +{ + let guard = OutstandingGuard::new(); + let task = async move { + let _guard = guard; + fut.await; + }; + match runtime::Handle::try_current() { + Ok(handle) => { + let _ = handle.spawn(task); + } + Err(_) => { + let _ = get_runtime().spawn(task); + } + } +} + /// Gracefully quiesce the shared runtime, meant to run at normal process exit. /// /// Waits (bounded by `timeout`) for [`OUTSTANDING`] to reach zero -- i.e. diff --git a/python/src/table.rs b/python/src/table.rs index a3dad3d86..bf5250732 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1,8 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::HashMap, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; -use crate::runtime::{block_on, future_into_py}; +use crate::runtime::{block_on, future_into_py, spawn_background}; use crate::{ connection::Connection, error::PythonErrorExt, @@ -680,11 +686,23 @@ impl From for DropColumnsResult { #[pyclass(name = "BlobFile")] pub struct PyBlobFile { inner: Arc, + closed: AtomicBool, +} + +impl PyBlobFile { + fn ensure_open(&self) -> PyResult<()> { + if self.closed.load(Ordering::Acquire) { + Err(PyRuntimeError::new_err("blob file is already closed")) + } else { + Ok(()) + } + } } #[pymethods] impl PyBlobFile { fn read_bytes(self_: PyRef<'_, Self>) -> PyResult> { + self_.ensure_open()?; let inner = self_.inner.clone(); let py = self_.py(); let bytes = py @@ -694,6 +712,7 @@ impl PyBlobFile { } pub fn read(self_: PyRef<'_, Self>) -> PyResult> { + self_.ensure_open()?; let inner = self_.inner.clone(); future_into_py(self_.py(), async move { let bytes = inner @@ -705,7 +724,20 @@ impl PyBlobFile { } fn close(self_: PyRef<'_, Self>) -> PyResult<()> { + if self_.closed.swap(true, Ordering::AcqRel) { + return Ok(()); + } let inner = self_.inner.clone(); + if tokio::runtime::Handle::try_current().is_ok() { + // IOBase.__del__ can call close while cyclic GC runs on a worker. + // The status changes immediately; release Lance's async state there. + spawn_background(async move { + if let Err(error) = inner.close().await { + log::warn!("blob close failed: {error}"); + } + }); + return Ok(()); + } self_ .py() .detach(move || block_on(async move { inner.close().await })) @@ -713,13 +745,11 @@ impl PyBlobFile { } fn is_closed(self_: PyRef<'_, Self>) -> bool { - let inner = self_.inner.clone(); - self_ - .py() - .detach(move || block_on(async move { inner.is_closed().await })) + self_.closed.load(Ordering::Acquire) } fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> { + self_.ensure_open()?; let inner = self_.inner.clone(); self_ .py() @@ -728,6 +758,7 @@ impl PyBlobFile { } fn tell(self_: PyRef<'_, Self>) -> PyResult { + self_.ensure_open()?; let inner = self_.inner.clone(); self_ .py() @@ -741,6 +772,7 @@ impl PyBlobFile { /// Read a blob-local byte range without moving the cursor. fn read_range(self_: PyRef<'_, Self>, offset: u64, length: usize) -> PyResult> { + self_.ensure_open()?; let end = offset .checked_add(length as u64) .ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?; @@ -753,6 +785,7 @@ impl PyBlobFile { } fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult> { + self_.ensure_open()?; let inner = self_.inner.clone(); let py = self_.py(); let bytes = py @@ -1455,6 +1488,7 @@ impl Table { .map(|handle| { handle.map(|file| PyBlobFile { inner: Arc::new(file), + closed: AtomicBool::new(false), }) }) .collect::>()) diff --git a/python/vendor/lance-arrow/Cargo.toml b/python/vendor/lance-arrow/Cargo.toml index 0e6e0ba91..fcd5d05e5 100644 --- a/python/vendor/lance-arrow/Cargo.toml +++ b/python/vendor/lance-arrow/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lance-arrow" -version = "13.0.0-beta.8" +version = "13.0.0-beta.12" edition = "2024" authors = ["Lance Devs "] license = "Apache-2.0" diff --git a/python/vendor/lance-arrow/README.lancedb.md b/python/vendor/lance-arrow/README.lancedb.md index e0f84c4a2..8fb76ce58 100644 --- a/python/vendor/lance-arrow/README.lancedb.md +++ b/python/vendor/lance-arrow/README.lancedb.md @@ -1,5 +1,5 @@ -This directory contains `rust/lance-arrow` from Lance tag `v13.0.0-beta.8` -(commit `3c20b7875f621c929fdcecabe74dbca507ec0f0d`), licensed under Apache-2.0. +This directory contains `rust/lance-arrow` from Lance tag `v13.0.0-beta.12` +(commit `5d2946fa452c61cb018144f3200892b6dd4a8b4f`), licensed under Apache-2.0. LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps until a Lance release includes the fix. The patch is selected in the root diff --git a/python/vendor/lance-arrow/src/json.rs b/python/vendor/lance-arrow/src/json.rs index b43777ba2..9fad06b34 100644 --- a/python/vendor/lance-arrow/src/json.rs +++ b/python/vendor/lance-arrow/src/json.rs @@ -7,13 +7,15 @@ use std::convert::TryFrom; use std::sync::Arc; -use arrow_array::builder::{LargeBinaryBuilder, StringBuilder}; +use arrow_array::builder::{GenericStringBuilder, LargeBinaryBuilder}; use arrow_array::cast::AsArray; use arrow_array::{ - Array, ArrayRef, FixedSizeListArray, LargeBinaryArray, LargeListArray, LargeStringArray, - ListArray, MapArray, RecordBatch, StringArray, StructArray, + Array, ArrayRef, FixedSizeListArray, GenericStringArray, LargeBinaryArray, LargeListArray, + LargeStringArray, ListArray, MapArray, OffsetSizeTrait, RecordBatch, StringArray, StructArray, }; use arrow_schema::{ArrowError, DataType, Field as ArrowField, Fields, Schema}; +use jsonb::OwnedJsonb; +use jsonb::jsonpath::{JsonPath, Selector}; use crate::ARROW_EXT_NAME_KEY; @@ -76,6 +78,72 @@ pub fn has_arrow_json_fields(field: &ArrowField) -> bool { } } +/// The Arrow representation holding a JSON field's values. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JsonEncoding { + /// JSONB bytes: [`JSON_EXT_NAME`] over `LargeBinary`. Lance stores JSON + /// this way. + Jsonb, + /// JSON text: [`ARROW_JSON_EXT_NAME`] over `Utf8` or `LargeUtf8`. Callers + /// write JSON this way, and reads return it this way by default. + Text, +} + +impl JsonEncoding { + /// The encoding of `field`'s values, or `None` when `field` does not hold + /// JSON. + pub fn of_field(field: &ArrowField) -> Option { + if is_json_field(field) { + Some(Self::Jsonb) + } else if is_arrow_json_field(field) { + Some(Self::Text) + } else { + None + } + } +} + +/// The values of a JSON column, readable as JSONB or as text whichever +/// encoding they arrived in. +/// +/// Code that computes on JSON, such as tokenizers and indices, reads values +/// through this type, so its behavior does not depend on whether they come +/// from storage or from a caller. +#[derive(Debug, Clone, Copy)] +pub struct JsonValues<'a> { + encoding: JsonEncoding, + array: &'a ArrayRef, +} + +impl<'a> JsonValues<'a> { + /// Read `array` as the values of `field`, or `None` when `field` does not + /// hold JSON. + pub fn try_new(field: &ArrowField, array: &'a ArrayRef) -> Option { + JsonEncoding::of_field(field).map(|encoding| Self { encoding, array }) + } + + pub fn encoding(&self) -> JsonEncoding { + self.encoding + } + + /// The values as JSONB. Text values are parsed, so invalid JSON fails. + pub fn to_jsonb(&self) -> Result { + match self.encoding { + JsonEncoding::Jsonb => Ok(self.array.as_binary::().clone()), + JsonEncoding::Text => Ok(JsonArray::try_from(self.array.clone())?.into_inner()), + } + } + + /// The values as JSON text: the input array when it already holds text, + /// otherwise `LargeUtf8` decoded from JSONB. + pub fn to_text(&self) -> ArrayRef { + match self.encoding { + JsonEncoding::Jsonb => Arc::new(decode_jsonb_array::(self.array.as_binary())), + JsonEncoding::Text => self.array.clone(), + } + } +} + /// Create a JSON field with the appropriate extension metadata pub fn json_field(name: &str, nullable: bool) -> ArrowField { let mut field = ArrowField::new(name, DataType::LargeBinary, nullable); @@ -150,28 +218,13 @@ impl JsonArray { return Ok(None); } - let jsonb_bytes = self.inner.value(i); - get_json_path(jsonb_bytes, path).map_err(|e| { - ArrowError::InvalidArgumentError(format!("Failed to extract JSONPath: {}", e)) - }) + let path = parse_json_path(path)?; + Ok(select_json_path(self.inner.value(i), &path)?.map(|value| value.to_string())) } /// Convert to Arrow string array (JSON as UTF-8) pub fn to_arrow_json(&self) -> ArrayRef { - let mut builder = arrow_array::builder::StringBuilder::new(); - - for i in 0..self.inner.len() { - if self.inner.is_null(i) { - builder.append_null(); - } else { - let jsonb_bytes = self.inner.value(i); - let json_str = decode_json(jsonb_bytes); - builder.append_value(&json_str); - } - } - - // Return as UTF-8 string array (Arrow represents JSON as strings) - Arc::new(builder.finish()) + Arc::new(decode_jsonb_array::(&self.inner)) } pub fn len(&self) -> usize { @@ -200,23 +253,7 @@ impl TryFrom<&StringArray> for JsonArray { type Error = ArrowError; fn try_from(array: &StringArray) -> Result { - let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len()); - - for i in 0..array.len() { - if array.is_null(i) { - builder.append_null(); - } else { - let json_str = array.value(i); - let encoded = encode_json(json_str).map_err(|e| { - ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e)) - })?; - builder.append_value(&encoded); - } - } - - Ok(Self { - inner: builder.finish(), - }) + encode_json_text_array(array) } } @@ -232,24 +269,40 @@ impl TryFrom<&LargeStringArray> for JsonArray { type Error = ArrowError; fn try_from(array: &LargeStringArray) -> Result { - let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len()); + encode_json_text_array(array) + } +} - for i in 0..array.len() { - if array.is_null(i) { - builder.append_null(); - } else { - let json_str = array.value(i); +fn encode_json_text_array( + array: &GenericStringArray, +) -> Result { + let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len()); + for value in array { + match value { + Some(json_str) => { let encoded = encode_json(json_str).map_err(|e| { ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e)) })?; builder.append_value(&encoded); } + None => builder.append_null(), } - - Ok(Self { - inner: builder.finish(), - }) } + Ok(JsonArray { + inner: builder.finish(), + }) +} + +fn decode_jsonb_array(array: &LargeBinaryArray) -> GenericStringArray { + let mut builder = + GenericStringBuilder::::with_capacity(array.len(), array.value_data().len()); + for value in array { + match value { + Some(jsonb_bytes) => builder.append_value(decode_json(jsonb_bytes)), + None => builder.append_null(), + } + } + builder.finish() } impl TryFrom for JsonArray { @@ -293,21 +346,46 @@ pub fn decode_json(jsonb_bytes: &[u8]) -> String { raw_jsonb.to_string() } -/// Extract JSONPath value from JSONB -fn get_json_path( - jsonb_bytes: &[u8], - path: &str, -) -> Result, Box> { - let json_path = jsonb::jsonpath::parse_json_path(path.as_bytes())?; - let raw_jsonb = jsonb::RawJsonb::new(jsonb_bytes); - let mut selector = jsonb::jsonpath::Selector::new(raw_jsonb); +/// Parse a JSONPath expression. +pub fn parse_json_path(path: &str) -> Result, ArrowError> { + jsonb::jsonpath::parse_json_path(path.as_bytes()) + .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid JSONPath '{path}': {e}"))) +} - let values = selector.select_values(&json_path)?; - if values.is_empty() { - Ok(None) - } else { - Ok(Some(values[0].to_string())) - } +/// Select the value at `path` in a JSONB value. +/// +/// A path matching several values selects them together as one JSON array, +/// so no match is silently dropped. +pub fn select_json_path( + jsonb_bytes: &[u8], + path: &JsonPath<'_>, +) -> Result, ArrowError> { + Selector::new(jsonb::RawJsonb::new(jsonb_bytes)) + .select_value(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) + }) +} + +/// Whether `path` matches any value in a JSONB value. +pub fn json_path_exists(jsonb_bytes: &[u8], path: &JsonPath<'_>) -> Result { + Selector::new(jsonb::RawJsonb::new(jsonb_bytes)) + .exists(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to match JSONPath {path}: {e}")) + }) +} + +/// Select every value matched by `path` in a JSONB value. +pub fn select_json_path_values( + jsonb_bytes: &[u8], + path: &JsonPath<'_>, +) -> Result, ArrowError> { + Selector::new(jsonb::RawJsonb::new(jsonb_bytes)) + .select_values(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) + }) } /// Convert an Arrow JSON field to Lance JSON field (with JSONB storage) @@ -532,25 +610,9 @@ fn convert_lance_json_array( ) -> Result<(ArrowField, ArrayRef, bool), ArrowError> { convert_json_array(field, array, &|field, array| { if is_json_field(field) { - let binary_array = array - .as_any() - .downcast_ref::() - .expect("Lance JSON field must be LargeBinaryArray"); - let mut builder = StringBuilder::new(); - - for i in 0..binary_array.len() { - if binary_array.is_null(i) { - builder.append_null(); - } else { - let jsonb_bytes = binary_array.value(i); - let json_str = decode_json(jsonb_bytes); - builder.append_value(&json_str); - } - } - Ok(Some(( lance_json_to_arrow_json(field), - Arc::new(builder.finish()) as ArrayRef, + Arc::new(decode_jsonb_array::(array.as_binary())) as ArrayRef, ))) } else { Ok(None) @@ -677,6 +739,41 @@ mod tests { let age = json_array.json_path(1, "$.user.age").unwrap(); assert_eq!(age, None); + + // Several matches are selected together, as the SQL JSON functions do. + let json_array = + JsonArray::try_from_iter(vec![Some(r#"{"a": [{"b": 1}, {"b": 2}]}"#)]).unwrap(); + let bs = json_array.json_path(0, "$.a[*].b").unwrap(); + assert_eq!(bs, Some("[1,2]".to_string())); + } + + #[test] + fn test_json_values_read_either_encoding() { + let docs = vec![Some(r#"{"a":1}"#), None]; + let text: ArrayRef = Arc::new(StringArray::from(docs.clone())); + let jsonb: ArrayRef = Arc::new(JsonArray::try_from_iter(docs).unwrap().into_inner()); + let mut text_field = ArrowField::new("j", DataType::Utf8, true); + text_field.set_metadata(std::collections::HashMap::from([( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + )])); + let jsonb_field = json_field("j", true); + + let from_text = JsonValues::try_new(&text_field, &text).unwrap(); + let from_jsonb = JsonValues::try_new(&jsonb_field, &jsonb).unwrap(); + assert_eq!(from_text.encoding(), JsonEncoding::Text); + assert_eq!(from_jsonb.encoding(), JsonEncoding::Jsonb); + + assert_eq!(&from_text.to_jsonb().unwrap(), jsonb.as_binary::()); + assert_eq!(&from_jsonb.to_jsonb().unwrap(), jsonb.as_binary::()); + assert_eq!(&from_text.to_text(), &text); + let decoded = from_jsonb.to_text(); + let decoded = decoded.as_string::(); + assert_eq!(decoded.value(0), r#"{"a":1}"#); + assert!(decoded.is_null(1)); + + let plain = ArrowField::new("j", DataType::Utf8, true); + assert!(JsonValues::try_new(&plain, &text).is_none()); } #[test] @@ -1301,12 +1398,7 @@ mod tests { // Invalid JSONPath syntax should return error let result = json_array.json_path(0, "invalid path without $"); assert!(result.is_err()); - assert!( - result - .unwrap_err() - .to_string() - .contains("Failed to extract JSONPath") - ); + assert!(result.unwrap_err().to_string().contains("Invalid JSONPath")); } #[test] diff --git a/python/vendor/lance-arrow/src/lib.rs b/python/vendor/lance-arrow/src/lib.rs index e2e470cd2..3c111a19e 100644 --- a/python/vendor/lance-arrow/src/lib.rs +++ b/python/vendor/lance-arrow/src/lib.rs @@ -472,7 +472,11 @@ pub fn iter_str_array(arr: &dyn Array) -> Box> match arr.data_type() { DataType::Utf8 => Box::new(arr.as_string::().iter()), DataType::LargeUtf8 => Box::new(arr.as_string::().iter()), - _ => panic!("Expecting Utf8 or LargeUtf8, found {:?}", arr.data_type()), + DataType::Utf8View => Box::new(arr.as_string_view().iter()), + _ => panic!( + "Expecting Utf8, LargeUtf8 or Utf8View, found {:?}", + arr.data_type() + ), } } diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 179b602dd..bf3b76981 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.40.0-beta.7" +version = "0.40.0-beta.8" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 893c584e1..ad4a1781b 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -216,7 +216,8 @@ impl RemoteDatabaseOptionsBuilder { #[derive(Debug)] pub struct RemoteDatabase { client: RestfulLanceDbClient, - table_cache: Cache>>, + // Cache existence and server capabilities, not mutable per-handle table state. + table_cache: Cache, uri: String, /// Headers to pass to the namespace client for authentication namespace_headers: HashMap, @@ -1371,16 +1372,9 @@ impl Database for RemoteDatabase { }; for table in &tables { - let table_identifier = build_table_identifier(table, &request.namespace_path)?; + build_table_identifier(table, &request.namespace_path)?; let cache_key = build_cache_key(table, &request.namespace_path); - let remote_table = Arc::new(RemoteTable::new( - self.client.clone(), - table.clone(), - request.namespace_path.clone(), - table_identifier.clone(), - version.clone(), - )); - self.table_cache.insert(cache_key, remote_table).await; + self.table_cache.insert(cache_key, version.clone()).await; } Ok(tables) } @@ -1407,16 +1401,9 @@ impl Database for RemoteDatabase { // Cache the tables for future use let namespace_vec = namespace_parts.to_vec(); for table in &response.tables { - let table_identifier = build_table_identifier(table, &namespace_vec)?; + build_table_identifier(table, &namespace_vec)?; let cache_key = build_cache_key(table, &namespace_vec); - let remote_table = Arc::new(RemoteTable::new( - self.client.clone(), - table.clone(), - namespace_vec.clone(), - table_identifier.clone(), - version.clone(), - )); - self.table_cache.insert(cache_key, remote_table).await; + self.table_cache.insert(cache_key, version.clone()).await; } Ok(response) @@ -1484,9 +1471,9 @@ impl Database for RemoteDatabase { request.name.clone(), request.namespace_path.clone(), table_identifier, - version, + version.clone(), )); - self.table_cache.insert(cache_key, table.clone()).await; + self.table_cache.insert(cache_key, version).await; Ok(table) } @@ -1526,9 +1513,9 @@ impl Database for RemoteDatabase { request.target_table_name.clone(), request.target_namespace_path.clone(), table_identifier, - version, + version.clone(), )); - self.table_cache.insert(cache_key, table.clone()).await; + self.table_cache.insert(cache_key, version).await; Ok(table) } @@ -1537,10 +1524,17 @@ impl Database for RemoteDatabase { let identifier = build_table_identifier(&request.name, &request.namespace_path)?; let cache_key = build_cache_key(&request.name, &request.namespace_path); - // We describe the table to confirm it exists before moving on. - if let Some(table) = self.table_cache.get(&cache_key).await { - Ok(table.clone()) + // Every open gets its own checkout, schema cache, and freshness state. + if let Some(version) = self.table_cache.get(&cache_key).await { + Ok(Arc::new(RemoteTable::new( + self.client.clone(), + request.name, + request.namespace_path, + identifier, + version, + ))) } else { + // Describe the table to confirm it exists before moving on. let req = self .client .post(&format!("/v1/table/{}/describe/", identifier)); @@ -1550,13 +1544,12 @@ impl Database for RemoteDatabase { let rsp = self.client.check_response(&request_id, rsp).await?; let version = parse_server_version(&request_id, &rsp)?; let describe_body = rsp.text().await.ok(); - let table_identifier = build_table_identifier(&request.name, &request.namespace_path)?; let table = Arc::new(RemoteTable::new( self.client.clone(), request.name.clone(), request.namespace_path.clone(), - table_identifier, - version, + identifier, + version.clone(), )); // This describe already carries the schema, so hand it to the table // instead of making the first schema read fetch it again. A version or @@ -1564,8 +1557,7 @@ impl Database for RemoteDatabase { if let Some(body) = &describe_body { table.seed_schema(body); } - let cache_key = build_cache_key(&request.name, &request.namespace_path); - self.table_cache.insert(cache_key, table.clone()).await; + self.table_cache.insert(cache_key, version).await; Ok(table) } } @@ -2271,6 +2263,59 @@ mod tests { assert_eq!(table.name(), "table1"); } + #[tokio::test] + async fn test_open_table_checkout_is_independent_per_handle() { + let latest = Arc::new(AtomicUsize::new(2)); + let current = latest.clone(); + let mut db = super::RemoteDatabase::new_mock(move |request| { + let body = match request.url().path() { + "/v1/table/table1/describe/" => { + let requested_version = request + .body() + .and_then(|body| body.as_bytes()) + .and_then(|body| serde_json::from_slice::(body).ok()) + .and_then(|body| body["version"].as_u64()); + let version = + requested_version.unwrap_or_else(|| current.load(Ordering::SeqCst) as u64); + serde_json::json!({ + "version": version, + "schema": {"fields": [ + {"name": "id", "type": {"type": "int32"}, "nullable": false} + ]} + }) + .to_string() + } + "/v1/table/table1/insert/" => { + let version = current.fetch_add(1, Ordering::SeqCst) + 1; + serde_json::json!({"version": version}).to_string() + } + path => panic!("unexpected request path: {path}"), + }; + http::Response::builder().status(200).body(body).unwrap() + }); + db.table_cache = moka::future::Cache::new(10); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + + let writer = conn.open_table("table1").execute().await.unwrap(); + let reader = conn.open_table("table1").execute().await.unwrap(); + reader.checkout(1).await.unwrap(); + + assert_eq!(reader.version().await.unwrap(), 1); + assert_eq!(writer.version().await.unwrap(), 2); + + let data = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![3]))], + ) + .unwrap(); + assert_eq!(writer.add(data).execute().await.unwrap().version, 3); + assert_eq!(writer.version().await.unwrap(), 3); + assert_eq!(reader.version().await.unwrap(), 1); + } + #[tokio::test] async fn test_open_table_seeds_the_schema_from_its_describe() { let describe_calls = Arc::new(AtomicUsize::new(0)); diff --git a/rust/lancedb/src/table/datafusion/cast.rs b/rust/lancedb/src/table/datafusion/cast.rs index 0df1cb4ba..2f3f8b868 100644 --- a/rust/lancedb/src/table/datafusion/cast.rs +++ b/rust/lancedb/src/table/datafusion/cast.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::collections::HashMap; use std::sync::{Arc, LazyLock}; use arrow_array::StructArray; @@ -73,8 +72,12 @@ fn build_field_exprs( let input_expr = get_input_expr(input_idx); // PyArrow's pa.json_() is already labelled arrow.json, which is what lance-core wants - // to see, so pass it straight through. - if is_json_field(table_field) && is_arrow_json_field(input_field) { + // to see, so pass it straight through once its extension metadata matches the table's + // (see `arrow_json_field`); otherwise it takes the relabelling cast below. + if is_json_field(table_field) + && is_arrow_json_field(input_field) + && input_field.as_ref() == &arrow_json_field(table_field, input_field.data_type()) + { result.push((input_expr, Arc::clone(input_field) as FieldRef)); continue; } @@ -299,11 +302,25 @@ fn arrow_json_storage_type(input: &DataType) -> Option { } } -fn arrow_json_field(name: &str, storage: DataType, nullable: bool) -> Field { - Field::new(name, storage, nullable).with_metadata(HashMap::from([( +/// The arrow.json field lance-core should see for `table_field`, a lance.json leaf. +/// +/// Lance-core converts the leaf to lance.json by swapping the extension name and keeping the +/// rest of the field metadata, and the file writer then requires the result to match the stored +/// field exactly, `ARROW:extension:metadata` included. PyArrow exports every extension field +/// with that key (empty for JSON), while a schema built in Rust omits it, so the relabelled +/// field mirrors the table field's metadata rather than carrying a fixed set of keys. +fn arrow_json_field(table_field: &Field, storage: &DataType) -> Field { + let mut metadata = table_field.metadata().clone(); + metadata.insert( ARROW_EXT_NAME_KEY.to_string(), ARROW_JSON_EXT_NAME.to_string(), - )])) + ); + Field::new( + table_field.name(), + storage.clone(), + table_field.is_nullable(), + ) + .with_metadata(metadata) } /// Rewrite `table_field` so that every lance.json leaf the input supplies as text becomes an @@ -320,11 +337,7 @@ fn json_write_target(input_field: &Field, table_field: &Field) -> Option } else { arrow_json_storage_type(input_field.data_type())? }; - return Some(arrow_json_field( - table_field.name(), - storage, - table_field.is_nullable(), - )); + return Some(arrow_json_field(table_field, &storage)); } if !has_json_fields(table_field) { @@ -401,6 +414,7 @@ fn null_literal(field: &FieldRef) -> Result> { #[cfg(test)] mod tests { + use std::collections::HashMap; use std::sync::Arc; use arrow::buffer::OffsetBuffer; @@ -408,7 +422,7 @@ mod tests { Array, Float32Array, Float64Array, Int32Array, Int64Array, ListArray, RecordBatch, StringArray, StructArray, UInt32Array, UInt64Array, }; - use arrow_schema::{DataType, Field, Fields, Schema}; + use arrow_schema::{DataType, Field, FieldRef, Fields, Schema}; use datafusion::prelude::SessionContext; use datafusion_catalog::MemTable; use futures::TryStreamExt; @@ -1007,7 +1021,7 @@ mod tests { )]); let input_item = if input_labelled { - Arc::new(arrow_json_field("item", item_type, true)) + Arc::new(arrow_json_field(&json_field("item", true), &item_type)) } else { Arc::new(Field::new("item", item_type, true)) }; @@ -1318,4 +1332,150 @@ mod tests { assert_eq!(result.num_rows(), 2); assert_eq!(result.column(0).null_count(), 2); } + + /// PyArrow records `ARROW:extension:metadata` (empty) alongside the extension name when it + /// exports a `pa.json_()` column, so a table created from one stores that key. Lance-core + /// swaps the name of a relabelled leaf back to lance.json but keeps the other metadata, and + /// the file writer then requires the field to match the stored one key for key. The + /// relabelled field must therefore mirror the table field's metadata. + #[rstest::rstest] + #[case::unlabelled_text(false)] + #[case::arrow_json_without_metadata_key(true)] + #[tokio::test] + async fn test_relabelled_json_mirrors_table_extension_metadata(#[case] input_labelled: bool) { + use lance_arrow::json::{ARROW_JSON_EXT_NAME, JSON_EXT_NAME}; + use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY}; + + let table_schema = Schema::new(vec![ + Field::new("data", DataType::LargeBinary, true).with_metadata(HashMap::from([ + (ARROW_EXT_NAME_KEY.to_string(), JSON_EXT_NAME.to_string()), + (ARROW_EXT_META_KEY.to_string(), String::new()), + ])), + ]); + + let mut input_field = Field::new("data", DataType::Utf8, true); + if input_labelled { + input_field = input_field.with_metadata(HashMap::from([( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + )])); + } + let input_schema = Arc::new(Schema::new(vec![input_field])); + let values = StringArray::from(vec![Some(r#"{"x": 1}"#), None]); + let input_batch = RecordBatch::try_new(input_schema, vec![Arc::new(values)]).unwrap(); + + let plan = plan_from_batch(input_batch).await; + let projected = cast_to_table_schema(plan, &table_schema).unwrap(); + + let out_field = projected.schema().field_with_name("data").unwrap().clone(); + assert_eq!(out_field.data_type(), &DataType::Utf8); + assert_eq!( + out_field.metadata(), + &HashMap::from([ + ( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string() + ), + (ARROW_EXT_META_KEY.to_string(), String::new()), + ]), + "the arrow.json label must carry the table field's other extension metadata" + ); + + let result = collect(projected).await; + assert_eq!(result.num_rows(), 2); + assert_eq!(result.column(0).null_count(), 1); + } + + /// The inverse: a table built in Rust records only the extension name, so an input that + /// carries the empty metadata key (what pyarrow exports) must drop it, or lance-core sees a + /// field the file schema does not describe. + #[tokio::test] + async fn test_arrow_json_input_drops_metadata_key_the_table_lacks() { + use lance_arrow::json::{ARROW_JSON_EXT_NAME, json_field}; + use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY}; + + let table_schema = Schema::new(vec![json_field("data", true)]); + + let input_field = Field::new("data", DataType::Utf8, true).with_metadata(HashMap::from([ + ( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ), + (ARROW_EXT_META_KEY.to_string(), String::new()), + ])); + let input_schema = Arc::new(Schema::new(vec![input_field])); + let values = StringArray::from(vec![Some(r#"{"x": 1}"#)]); + let input_batch = RecordBatch::try_new(input_schema, vec![Arc::new(values)]).unwrap(); + + let plan = plan_from_batch(input_batch).await; + let projected = cast_to_table_schema(plan, &table_schema).unwrap(); + + let out_field = projected.schema().field_with_name("data").unwrap().clone(); + assert_eq!(out_field.data_type(), &DataType::Utf8); + assert_eq!( + out_field.metadata(), + &HashMap::from([( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string() + )]) + ); + assert_eq!(collect(projected).await.num_rows(), 1); + } + + /// Nested leaves mirror their table leaf's metadata too, since lance-core checks each + /// child of a list, map or struct against the stored child. + #[tokio::test] + async fn test_nested_relabelled_json_mirrors_table_extension_metadata() { + use lance_arrow::json::{ARROW_JSON_EXT_NAME, JSON_EXT_NAME}; + use lance_arrow::{ARROW_EXT_META_KEY, ARROW_EXT_NAME_KEY}; + + let table_item = + Field::new("item", DataType::LargeBinary, true).with_metadata(HashMap::from([ + (ARROW_EXT_NAME_KEY.to_string(), JSON_EXT_NAME.to_string()), + (ARROW_EXT_META_KEY.to_string(), String::new()), + ])); + let table_schema = Schema::new(vec![Field::new( + "docs", + DataType::List(Arc::new(table_item)), + true, + )]); + + let input_item: FieldRef = Arc::new(Field::new("item", DataType::Utf8, true)); + let input_schema = Arc::new(Schema::new(vec![Field::new( + "docs", + DataType::List(input_item.clone()), + true, + )])); + let values = StringArray::from(vec![Some(r#"{"k": 1}"#), Some(r#"{"k": 2}"#)]); + let input_batch = RecordBatch::try_new( + input_schema, + vec![Arc::new(ListArray::new( + input_item, + OffsetBuffer::new(vec![0, 1, 2].into()), + Arc::new(values), + None, + ))], + ) + .unwrap(); + + let plan = plan_from_batch(input_batch).await; + let projected = cast_to_table_schema(plan, &table_schema).unwrap(); + + let out_field = projected.schema().field_with_name("docs").unwrap().clone(); + let DataType::List(out_item) = out_field.data_type() else { + panic!("expected a list, got {}", out_field.data_type()); + }; + assert_eq!(out_item.data_type(), &DataType::Utf8); + assert_eq!( + out_item.metadata(), + &HashMap::from([ + ( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string() + ), + (ARROW_EXT_META_KEY.to_string(), String::new()), + ]) + ); + assert_eq!(collect(projected).await.num_rows(), 2); + } }