diff --git a/.bumpversion.toml b/.bumpversion.toml index f14c9b1cb..12d7f0804 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.1" +current_version = "0.41.0-beta.6" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/.github/actions/setup-protoc/action.yml b/.github/actions/setup-protoc/action.yml new file mode 100644 index 000000000..a99cf3320 --- /dev/null +++ b/.github/actions/setup-protoc/action.yml @@ -0,0 +1,80 @@ +name: Setup protoc +description: Install the protoc version pinned for wheel builds, and expose it for builds that install protoc themselves + +inputs: + install: + description: "Install protoc on the runner. Set to false when the build runs in a container (e.g. manylinux) and installs protoc there." + required: false + default: "true" + +outputs: + version: + description: "The pinned protoc version, e.g. 36.2" + value: ${{ steps.pin.outputs.version }} + linux-x86_64-sha256: + description: "SHA-256 of the protoc linux-x86_64 zip" + value: ${{ steps.pin.outputs.linux-x86_64-sha256 }} + linux-aarch_64-sha256: + description: "SHA-256 of the protoc linux-aarch_64 zip" + value: ${{ steps.pin.outputs.linux-aarch_64-sha256 }} + +runs: + using: composite + steps: + # protobuf does not publish checksum files for protoc. These are the digests + # GitHub records for the release assets; when bumping, copy them from + # `gh api repos/protocolbuffers/protobuf/releases/tags/v`. + - name: Set protoc pin + id: pin + shell: bash + run: | + { + echo "version=36.2" + echo "linux-x86_64-sha256=121f6c7afe1d4d0e3ea6aab9432038599250134cbf4474cb1167d2c7decd4278" + echo "linux-aarch_64-sha256=8b8f18bd2b30346efbc698dd5a73dd7c805f3ef8380f6dfc95c768f3f1852f6a" + echo "osx-universal_binary-sha256=83d0dc80672f53486f54299bca6177159f9cad1fd54d4d3143eeea53c4423e1b" + echo "win64-sha256=f0c128dc0d8492eceece83bb459a4c0e316764b929ffbf1aa416357fd644edd3" + } >> "$GITHUB_OUTPUT" + - name: Install protoc + if: ${{ inputs.install == 'true' }} + shell: bash + env: + PROTOC_VERSION: ${{ steps.pin.outputs.version }} + RUNNER_OS_NAME: ${{ runner.os }} + RUNNER_ARCH_NAME: ${{ runner.arch }} + LINUX_X86_64_SHA256: ${{ steps.pin.outputs.linux-x86_64-sha256 }} + LINUX_AARCH_64_SHA256: ${{ steps.pin.outputs.linux-aarch_64-sha256 }} + OSX_SHA256: ${{ steps.pin.outputs.osx-universal_binary-sha256 }} + WIN64_SHA256: ${{ steps.pin.outputs.win64-sha256 }} + run: | + set -euo pipefail + case "$RUNNER_OS_NAME/$RUNNER_ARCH_NAME" in + Linux/X64) platform="linux-x86_64"; sha256="$LINUX_X86_64_SHA256" ;; + Linux/ARM64) platform="linux-aarch_64"; sha256="$LINUX_AARCH_64_SHA256" ;; + macOS/*) platform="osx-universal_binary"; sha256="$OSX_SHA256" ;; + Windows/X64) platform="win64"; sha256="$WIN64_SHA256" ;; + *) echo "Unsupported runner: $RUNNER_OS_NAME/$RUNNER_ARCH_NAME" >&2; exit 1 ;; + esac + zip="$RUNNER_TEMP/protoc.zip" + dest="$RUNNER_TEMP/protoc" + curl -fsSL "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-${platform}.zip" -o "$zip" + # Hash stdin rather than a path: Windows paths contain backslashes, + # which make sha256sum escape its output. + if [ "$RUNNER_OS_NAME" = "macOS" ]; then + actual=$(shasum -a 256 < "$zip" | cut -d' ' -f1) + else + actual=$(sha256sum < "$zip" | cut -d' ' -f1) + fi + if [ "$actual" != "$sha256" ]; then + echo "protoc checksum mismatch: expected $sha256, got $actual" >&2 + exit 1 + fi + if [ "$RUNNER_OS_NAME" = "Windows" ]; then + # Git Bash on Windows runners does not reliably ship unzip. + 7z x -y -o"$dest" "$zip" > /dev/null + else + unzip -q -o "$zip" -d "$dest" + fi + rm "$zip" + echo "$dest/bin" >> "$GITHUB_PATH" + "$dest/bin/protoc" --version diff --git a/.github/workflows/build_linux_wheel/action.yml b/.github/workflows/build_linux_wheel/action.yml index 8fa5b1f79..44ce82ee5 100644 --- a/.github/workflows/build_linux_wheel/action.yml +++ b/.github/workflows/build_linux_wheel/action.yml @@ -29,6 +29,11 @@ inputs: runs: using: "composite" steps: + - name: Get protoc version + id: protoc + uses: ./.github/actions/setup-protoc + with: + install: "false" - name: CONFIRM ARM BUILD shell: bash env: @@ -54,13 +59,14 @@ runs: maturin-version: "1.12.4" command: build working-directory: python - docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc ${{ inputs.rustflags != '' && format('-e RUSTFLAGS={0}', inputs.rustflags) || '' }}" + docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc -e PROTOC_VERSION=${{ steps.protoc.outputs.version }} -e PROTOC_SHA256=${{ steps.protoc.outputs.linux-x86_64-sha256 }} ${{ inputs.rustflags != '' && format('-e RUSTFLAGS={0}', inputs.rustflags) || '' }}" target: x86_64-unknown-linux-gnu manylinux: ${{ inputs.manylinux }} args: ${{ inputs.args }} before-script-linux: | set -e - curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-x86_64.zip -o /tmp/protoc.zip + curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-x86_64.zip -o /tmp/protoc.zip + echo "${PROTOC_SHA256} /tmp/protoc.zip" | sha256sum -c - unzip /tmp/protoc.zip -d /usr/local rm /tmp/protoc.zip /usr/local/bin/protoc --version @@ -71,14 +77,15 @@ runs: maturin-version: "1.12.4" command: build working-directory: python - docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc ${{ inputs.rustflags != '' && format('-e RUSTFLAGS={0}', inputs.rustflags) || '' }}" + docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc -e PROTOC_VERSION=${{ steps.protoc.outputs.version }} -e PROTOC_SHA256=${{ steps.protoc.outputs.linux-aarch_64-sha256 }} ${{ inputs.rustflags != '' && format('-e RUSTFLAGS={0}', inputs.rustflags) || '' }}" target: aarch64-unknown-linux-gnu manylinux: ${{ inputs.manylinux }} args: ${{ inputs.args }} before-script-linux: | set -e yum install -y clang - curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-aarch_64.zip -o /tmp/protoc.zip + curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-aarch_64.zip -o /tmp/protoc.zip + echo "${PROTOC_SHA256} /tmp/protoc.zip" | sha256sum -c - unzip /tmp/protoc.zip -d /usr/local rm /tmp/protoc.zip /usr/local/bin/protoc --version diff --git a/.github/workflows/build_mac_wheel/action.yml b/.github/workflows/build_mac_wheel/action.yml index 75d942c96..727cd190c 100644 --- a/.github/workflows/build_mac_wheel/action.yml +++ b/.github/workflows/build_mac_wheel/action.yml @@ -12,10 +12,7 @@ inputs: runs: using: "composite" steps: - - name: Install macos dependency - shell: bash - run: | - brew install protobuf + - uses: ./.github/actions/setup-protoc - name: Build wheel uses: PyO3/maturin-action@v1 with: diff --git a/.github/workflows/build_windows_wheel/action.yml b/.github/workflows/build_windows_wheel/action.yml index 18f39818f..954ee7d35 100644 --- a/.github/workflows/build_windows_wheel/action.yml +++ b/.github/workflows/build_windows_wheel/action.yml @@ -12,15 +12,7 @@ inputs: runs: using: "composite" steps: - - name: Install Protoc v21.12 - working-directory: C:\ - run: | - New-Item -Path 'C:\protoc' -ItemType Directory - Set-Location C:\protoc - Invoke-WebRequest https://github.com/protocolbuffers/protobuf/releases/download/v21.12/protoc-21.12-win64.zip -OutFile C:\protoc\protoc.zip - 7z x protoc.zip - Add-Content $env:GITHUB_PATH "C:\protoc\bin" - shell: powershell + - uses: ./.github/actions/setup-protoc - name: Build wheel uses: PyO3/maturin-action@v1 with: diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml index 4f883aec1..d16203979 100644 --- a/.github/workflows/pypi-publish.yml +++ b/.github/workflows/pypi-publish.yml @@ -11,6 +11,7 @@ on: - .github/workflows/build_linux_wheel/action.yml - .github/workflows/build_mac_wheel/action.yml - .github/workflows/build_windows_wheel/action.yml + - .github/actions/setup-protoc/action.yml - Cargo.toml # Change in dependency frequently breaks builds - Cargo.lock diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index 3594bd826..3608c24ca 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -15,6 +15,7 @@ on: - .github/workflows/build_linux_wheel/** - .github/workflows/build_mac_wheel/** - .github/workflows/build_windows_wheel/** + - .github/actions/setup-protoc/** - .github/workflows/run_tests/** permissions: diff --git a/Cargo.lock b/Cargo.lock index a3d865be8..0515319b2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -986,40 +986,36 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.1.12" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a2f165a7feee6f263028b899d0a181987f4fa7179a6411a32a439fba7c5f769" +checksum = "5c3ef8931ad1c98aa6a55b4256f847f3116090819844e0dd41ea682cac5dd2d3" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", "aws-smithy-types", - "h2 0.3.27", - "h2 0.4.16", - "http 0.2.12", + "h2", "http 1.5.0", - "http-body 0.4.6", - "hyper 0.14.32", - "hyper 1.9.0", - "hyper-rustls 0.24.2", - "hyper-rustls 0.27.9", + "hyper", + "hyper-rustls", "hyper-util", "pin-project-lite", - "rustls 0.21.12", - "rustls 0.23.45", + "rustls", "rustls-native-certs", "rustls-pki-types", "tokio", - "tokio-rustls 0.26.4", + "tokio-rustls", "tower", "tracing", ] [[package]] name = "aws-smithy-json" -version = "0.62.5" +version = "0.62.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9648b0bb82a2eedd844052c6ad2a1a822d1f8e3adee5fbf668366717e428856a" +checksum = "701a947f4797e52a911e114a898667c746c39feea467bbd1abd7b3721f702ffa" dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", ] @@ -1070,9 +1066,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.12.1" +version = "1.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc117c179ecf39a62a0a3f49f600e9ac26a7ad7dd172177999f83933af776c32" +checksum = "9db177daa6ba8afb9ee1aefcf548c907abcf52065e394ee11a92780057fe0e8c" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", @@ -1170,7 +1166,7 @@ dependencies = [ "http 1.5.0", "http-body 1.1.0", "http-body-util", - "hyper 1.9.0", + "hyper", "hyper-util", "itoa", "matchit 0.7.3", @@ -1819,7 +1815,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.52.0", ] [[package]] @@ -3315,7 +3311,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -3588,8 +3584,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "rand 0.9.5", @@ -4000,11 +3996,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if 1.0.4", - "js-sys", "libc", "r-efi 5.3.0", "wasip2", - "wasm-bindgen", ] [[package]] @@ -4110,25 +4104,6 @@ dependencies = [ "subtle", ] -[[package]] -name = "h2" -version = "0.3.27" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" -dependencies = [ - "bytes", - "fnv", - "futures-core", - "futures-sink", - "futures-util", - "http 0.2.12", - "indexmap 2.14.0", - "slab", - "tokio", - "tokio-util", - "tracing", -] - [[package]] name = "h2" version = "0.4.16" @@ -4430,30 +4405,6 @@ dependencies = [ "typenum", ] -[[package]] -name = "hyper" -version = "0.14.32" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" -dependencies = [ - "bytes", - "futures-channel", - "futures-core", - "futures-util", - "h2 0.3.27", - "http 0.2.12", - "http-body 0.4.6", - "httparse", - "httpdate", - "itoa", - "pin-project-lite", - "socket2 0.5.10", - "tokio", - "tower-service", - "tracing", - "want", -] - [[package]] name = "hyper" version = "1.9.0" @@ -4464,7 +4415,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "h2 0.4.16", + "h2", "http 1.5.0", "http-body 1.1.0", "httparse", @@ -4476,21 +4427,6 @@ dependencies = [ "want", ] -[[package]] -name = "hyper-rustls" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" -dependencies = [ - "futures-util", - "http 0.2.12", - "hyper 0.14.32", - "log", - "rustls 0.21.12", - "tokio", - "tokio-rustls 0.24.1", -] - [[package]] name = "hyper-rustls" version = "0.27.9" @@ -4498,12 +4434,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ "http 1.5.0", - "hyper 1.9.0", + "hyper", "hyper-util", - "rustls 0.23.45", + "rustls", "rustls-native-certs", "tokio", - "tokio-rustls 0.26.4", + "tokio-rustls", "tower-service", "webpki-roots 1.0.7", ] @@ -4514,7 +4450,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" dependencies = [ - "hyper 1.9.0", + "hyper", "hyper-util", "pin-project-lite", "tokio", @@ -4533,12 +4469,12 @@ dependencies = [ "futures-util", "http 1.5.0", "http-body 1.1.0", - "hyper 1.9.0", + "hyper", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.3", + "socket2", "system-configuration", "tokio", "tower-service", @@ -4880,7 +4816,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -5095,26 +5031,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "jsonb" -version = "0.5.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb98fb29636087c40ad0d1274d9a30c0c1e83e03ae93f6e7e89247b37fcc6953" -dependencies = [ - "byteorder", - "ethnum", - "fast-float2", - "itoa", - "jiff", - "nom 8.0.0", - "num-traits", - "ordered-float 5.3.0", - "rand 0.9.5", - "serde", - "serde_json", - "zmij", -] - [[package]] name = "kanaria" version = "0.2.0" @@ -5143,8 +5059,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arc-swap", "arrow", @@ -5217,8 +5133,8 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -5232,7 +5148,7 @@ dependencies = [ "futures", "getrandom 0.2.17", "half", - "jsonb", + "lance-jsonb", "num-traits", "rand 0.9.5", ] @@ -5240,7 +5156,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -5254,7 +5170,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-schema", @@ -5263,8 +5179,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrayref", "crunchy", @@ -5274,8 +5190,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -5312,8 +5228,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow", "arrow-array", @@ -5330,10 +5246,10 @@ dependencies = [ "datafusion-physical-expr", "futures", "half", - "jsonb", "lance-arrow", "lance-core", "lance-geo", + "lance-jsonb", "log", "pin-project", "prost", @@ -5344,8 +5260,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow", "arrow-array", @@ -5362,8 +5278,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "proc-macro2", "quote", @@ -5372,8 +5288,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-arith", "arrow-array", @@ -5406,8 +5322,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-arith", "arrow-array", @@ -5439,8 +5355,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "datafusion", "geo-traits", @@ -5454,8 +5370,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arc-swap", "arrow", @@ -5486,7 +5402,6 @@ dependencies = [ "hex", "itertools 0.14.0", "jieba-rs", - "jsonb", "lance-arrow", "lance-arrow-stats", "lance-bitpacking", @@ -5497,6 +5412,7 @@ dependencies = [ "lance-geo", "lance-index-core", "lance-io", + "lance-jsonb", "lance-linalg", "lance-select", "lance-table", @@ -5526,8 +5442,8 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-schema", @@ -5550,8 +5466,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow", "arrow-array", @@ -5591,10 +5507,28 @@ dependencies = [ "uuid", ] +[[package]] +name = "lance-jsonb" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" +dependencies = [ + "byteorder", + "ethnum", + "fast-float2", + "itoa", + "jiff", + "nom 8.0.0", + "num-traits", + "ordered-float 5.3.0", + "serde", + "serde_json", + "zmij", +] + [[package]] name = "lance-linalg" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-schema", @@ -5608,8 +5542,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow", "async-trait", @@ -5623,8 +5557,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arc-swap", "arrow", @@ -5678,8 +5612,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -5693,8 +5627,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow", "arrow-array", @@ -5734,8 +5668,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "arrow-array", "arrow-schema", @@ -5748,8 +5682,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "14.0.0-beta.7" -source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.7#0fd26effc79c62c244a5baf92cdda06e301baf21" +version = "14.0.0-beta.15" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.15#765641dc0d7100f032fdebc00d73eef49711c4d4" dependencies = [ "frostem", "icu_segmenter", @@ -5762,7 +5696,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.1" +version = "0.41.0-beta.6" dependencies = [ "ahash", "anyhow", @@ -5835,7 +5769,7 @@ dependencies = [ "rand 0.9.5", "random_word", "regex", - "reqwest 0.12.28", + "reqwest 0.13.4", "roaring", "rstest", "semver", @@ -5845,6 +5779,7 @@ dependencies = [ "serial_test", "sha2 0.10.9", "snafu 0.8.9", + "strum_macros 0.28.0", "tempfile", "test-log", "tokenizers", @@ -5859,7 +5794,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.1" +version = "0.41.0-beta.6" dependencies = [ "arrow-array", "arrow-buffer", @@ -5884,7 +5819,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.1" +version = "0.41.0-beta.6" dependencies = [ "arc-swap", "arrow", @@ -6901,7 +6836,7 @@ dependencies = [ "http-body-util", "httparse", "humantime", - "hyper 1.9.0", + "hyper", "itertools 0.15.0", "md-5", "nix 0.31.3", @@ -8381,8 +8316,8 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls 0.23.45", - "socket2 0.6.3", + "rustls", + "socket2", "thiserror 2.0.18", "tokio", "tracing", @@ -8391,18 +8326,19 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "0e750cca55fe4f0439a15d0bb529da9651e79993e8e72c61a899a36d462befbe" dependencies = [ "aws-lc-rs", "bytes", - "getrandom 0.3.4", + "getrandom 0.4.2", "lru-slab", - "rand 0.9.5", + "rand 0.10.1", + "rand_pcg", "ring", "rustc-hash", - "rustls 0.23.45", + "rustls", "rustls-pki-types", "slab", "thiserror 2.0.18", @@ -8420,9 +8356,9 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.3", + "socket2", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -8567,6 +8503,15 @@ dependencies = [ "rand 0.9.5", ] +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", +] + [[package]] name = "rand_xoshiro" version = "0.7.0" @@ -8931,12 +8876,12 @@ dependencies = [ "encoding_rs", "futures-core", "futures-util", - "h2 0.4.16", + "h2", "http 1.5.0", "http-body 1.1.0", "http-body-util", - "hyper 1.9.0", - "hyper-rustls 0.27.9", + "hyper", + "hyper-rustls", "hyper-util", "js-sys", "log", @@ -8945,7 +8890,7 @@ dependencies = [ "percent-encoding", "pin-project-lite", "quinn", - "rustls 0.23.45", + "rustls", "rustls-native-certs", "rustls-pki-types", "serde", @@ -8953,7 +8898,7 @@ dependencies = [ "serde_urlencoded", "sync_wrapper", "tokio", - "tokio-rustls 0.26.4", + "tokio-rustls", "tokio-util", "tower", "tower-http 0.6.10", @@ -8974,28 +8919,31 @@ checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" dependencies = [ "base64 0.22.1", "bytes", + "encoding_rs", "futures-core", "futures-util", - "h2 0.4.16", + "h2", "http 1.5.0", "http-body 1.1.0", "http-body-util", - "hyper 1.9.0", - "hyper-rustls 0.27.9", + "hyper", + "hyper-rustls", "hyper-util", "js-sys", "log", + "mime", "percent-encoding", "pin-project-lite", "quinn", - "rustls 0.23.45", + "rustls", "rustls-pki-types", "rustls-platform-verifier", "serde", "serde_json", + "serde_urlencoded", "sync_wrapper", "tokio", - "tokio-rustls 0.26.4", + "tokio-rustls", "tokio-util", "tower", "tower-http 0.6.10", @@ -9220,19 +9168,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", -] - -[[package]] -name = "rustls" -version = "0.21.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" -dependencies = [ - "log", - "ring", - "rustls-webpki 0.101.7", - "sct", + "windows-sys 0.52.0", ] [[package]] @@ -9246,7 +9182,7 @@ dependencies = [ "once_cell", "ring", "rustls-pki-types", - "rustls-webpki 0.103.15", + "rustls-webpki", "subtle", "zeroize", ] @@ -9284,14 +9220,14 @@ dependencies = [ "jni", "log", "once_cell", - "rustls 0.23.45", + "rustls", "rustls-native-certs", "rustls-platform-verifier-android", - "rustls-webpki 0.103.15", + "rustls-webpki", "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -9300,16 +9236,6 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" -[[package]] -name = "rustls-webpki" -version = "0.101.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" -dependencies = [ - "ring", - "untrusted 0.9.0", -] - [[package]] name = "rustls-webpki" version = "0.103.15" @@ -9425,16 +9351,6 @@ dependencies = [ "sha2 0.10.9", ] -[[package]] -name = "sct" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" -dependencies = [ - "ring", - "untrusted 0.9.0", -] - [[package]] name = "sec1" version = "0.3.0" @@ -9530,6 +9446,7 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ + "indexmap 2.14.0", "itoa", "memchr", "serde", @@ -9856,16 +9773,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "socket2" -version = "0.5.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" -dependencies = [ - "libc", - "windows-sys 0.52.0", -] - [[package]] name = "socket2" version = "0.6.3" @@ -10284,7 +10191,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -10497,7 +10404,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2 0.6.3", + "socket2", "tokio-macros", "windows-sys 0.61.2", ] @@ -10524,23 +10431,13 @@ dependencies = [ "tokio", ] -[[package]] -name = "tokio-rustls" -version = "0.24.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" -dependencies = [ - "rustls 0.21.12", - "tokio", -] - [[package]] name = "tokio-rustls" version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls 0.23.45", + "rustls", "tokio", ] @@ -10633,20 +10530,20 @@ dependencies = [ "axum 0.8.9", "base64 0.22.1", "bytes", - "h2 0.4.16", + "h2", "http 1.5.0", "http-body 1.1.0", "http-body-util", - "hyper 1.9.0", + "hyper", "hyper-timeout", "hyper-util", "percent-encoding", "pin-project", "rustls-native-certs", - "socket2 0.6.3", + "socket2", "sync_wrapper", "tokio", - "tokio-rustls 0.26.4", + "tokio-rustls", "tokio-stream", "tower", "tower-layer", @@ -10954,7 +10851,7 @@ dependencies = [ "flate2", "log", "once_cell", - "rustls 0.23.45", + "rustls", "rustls-pki-types", "serde", "serde_json", @@ -11305,7 +11202,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.52.0", ] [[package]] @@ -11775,7 +11672,7 @@ dependencies = [ "crc32fast", "futures", "http 1.5.0", - "hyper 1.9.0", + "hyper", "more-asserts", "rand 0.10.1", "redb", @@ -11911,9 +11808,9 @@ checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "550a2b930b62486a393c52d5c3b84bff264b28aa437ed64694d31e93b1757af7" [[package]] name = "yoke" diff --git a/Cargo.toml b/Cargo.toml index ad7c00477..5033b1529 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.94.0" [workspace.dependencies] -lance = { "version" = "=14.0.0-beta.7", default-features = false, "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=14.0.0-beta.7", default-features = false, "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=14.0.0-beta.7", default-features = false, "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=14.0.0-beta.7", "tag" = "v14.0.0-beta.7", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=14.0.0-beta.15", default-features = false, "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=14.0.0-beta.15", default-features = false, "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=14.0.0-beta.15", default-features = false, "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=14.0.0-beta.15", "tag" = "v14.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow @@ -72,10 +72,10 @@ oauth2 = { version = "5.0", default-features = false } regex = "1.10" semver = "1.0.25" serde = "1" -serde_json = "1" +serde_json = { version = "1", features = ["preserve_order"] } tempfile = "3.5.0" tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] } -tonic = { version = "0.14", features = ["tls-native-roots", "tls-ring"] } +tonic = { version = "0.14", features = ["tls-native-roots", "tls-aws-lc"] } uuid = { version = "1.7.0", features = ["v4", "v7"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } diff --git a/docs/README.md b/docs/README.md index bce1ec668..c5a99943d 100644 --- a/docs/README.md +++ b/docs/README.md @@ -15,18 +15,21 @@ unreleased features. ### Preview the docs +Run these commands from the LanceDB repository root. The subshell changes into +`docs` to load `mkdocs.yml`, while leaving your shell in the repository root. + ```shell -cd docs -mkdocs serve +(cd docs && PYTHONPATH=. mkdocs serve) ``` If you want to just generate the HTML files: ```shell -PYTHONPATH=. mkdocs build -f docs/mkdocs.yml +(cd docs && PYTHONPATH=. mkdocs build) ``` -If successful, you should see a `docs/site` directory that you can verify locally. +If successful, you should see a `docs/site` directory in the repository root that +you can verify locally. ## Adding examples diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 20db07d56..e2d46921e 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.41.0-beta.1 + 0.41.0-beta.6 ``` diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 3ac1dd031..59c3209c8 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -143,7 +143,8 @@ containing the new version number of the table after altering the columns. abstract blobColumns(): Promise ``` -Blob v2 columns, including nested dotted paths. +Blob v2 columns addressable by row ID, including struct-nested dotted paths. +Blobs inside lists are omitted because a row ID does not identify an element. #### Returns @@ -527,6 +528,8 @@ table's current checkout. Preserves input order, duplicates, and nulls. Use this for large payloads. See [Table.fetchBlobs](Table.md#fetchblobs) for row-ID validity across versions. +One missing or deleted row ID rejects the entire batch, including on remote +tables; no partial handles are returned. #### Parameters @@ -556,6 +559,10 @@ Bytes for `column` at row IDs from [Query.withRowId](Query.md#withrowid). Reads the table's current checkout. IDs from another version can fail after compaction unless stable row ids are enabled. Results keep input order and duplicates. Null blobs are `null`. Empty blobs are empty buffers. +One missing or deleted row ID rejects the entire batch, including on remote +tables; no partial results are returned. +Blobs inside lists cannot be fetched by row ID. + Remote servers limit each request to 1024 row IDs and 64 MiB of blob bytes. The client splits requests automatically and reads an individual larger blob through the Range route. This method still materializes all bytes in diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index a0a77fd7c..4809ebefd 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -460,17 +460,16 @@ exhaustiely searched to find matches. This parameter controls how many partitions should be searched. Increasing this value will increase the recall of your query but will -also increase the latency of your query. The default value is 20. This -default is good for many cases but the best value to use will depend on -your data and the recall that you need to achieve. +also increase the latency of your query. If this method is not called, +Lance's adaptive probe defaults are used. For best results we recommend tuning this parameter with a benchmark against your actual data to find the smallest possible value that will still give you the desired recall. -For more fine grained control over behavior when you have a very narrow filter -you can use `minimumNprobes` and `maximumNprobes`. This method sets both -the minimum and maximum to the same value. +The value is retained as `nprobes` through client and server request +construction. Lance sets both probe bounds to this value. Explicit minimum +or maximum settings can override their respective bounds. #### Parameters @@ -647,6 +646,38 @@ rerank(reranker): VectorQuery *** +### searchEffort() + +```ts +searchEffort(searchEffort): VectorQuery +``` + +Set the initial IVF search effort to a finite value in [0, 1]. + +The default, null, uses Auto (internally 0.5). Zero starts at the caller +minimum; one searches all partitions within the caller maximum. Later +probing may continue when too few candidates are found. +Any explicit effort, including 0.5, overrides nprobes in either call order. +Explicit minimum and maximum bounds still apply. This is not a recall +guarantee and does not change refinement or HNSW ef. The MemWAL LSM +scanner does not yet support this option. + +#### Parameters + +* **searchEffort**: `null` \| `number` + +#### Returns + +[`VectorQuery`](VectorQuery.md) + +#### Example + +```ts +table.vectorSearch([0.1, 0.2]).searchEffort(0.75).limit(10); +``` + +*** + ### select() ```ts diff --git a/docs/src/js/functions/blob.md b/docs/src/js/functions/blob.md index 20a734cd4..1a65b5133 100644 --- a/docs/src/js/functions/blob.md +++ b/docs/src/js/functions/blob.md @@ -14,6 +14,8 @@ Declares a `lance.blob.v2` column. Query results are descriptors, not payload bytes. Use [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles) to read bytes. +Blobs inside lists can be written and queried as descriptors, but cannot +currently be fetched because a row ID does not identify a list element. ## Parameters diff --git a/docs/src/js/interfaces/ClientConfig.md b/docs/src/js/interfaces/ClientConfig.md index 2b4984b8d..d5ce32117 100644 --- a/docs/src/js/interfaces/ClientConfig.md +++ b/docs/src/js/interfaces/ClientConfig.md @@ -8,6 +8,18 @@ ## Properties +### blobRequestConcurrency? + +```ts +optional blobRequestConcurrency: number; +``` + +Maximum number of concurrent blob HTTP requests across all tables and +blob handles on this connection. Defaults to 8. Can also be set with +`LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY`. Must be greater than zero. + +*** + ### extraHeaders? ```ts diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 3776283dc..24372a83d 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -48,6 +48,169 @@ are ordinary connections. Dropping a database requires it to be empty. ::: lancedb.catalog.DatabaseNames +## Authorization + +Some catalogs support the authorization API. This is an API that can be used to +bind access control lists (ACLs) to individual catalog objects. + +Each ACL contains a subject, an object, and privilege. + +### Subjects + +There are three kinds of subjects: principals, groups, and roles. + +A principal uniquely identifies each user accessing the system. Each LanceDB +Enterprise API key corresponds to a unique principal, and each OpenID user has +a unique principal. + +All principals have an immutable ID which never changes, and is safe to include +in log files. For example, the API key +`sk_AAAQEAYEAUDAOCAJBIFQYDIOB4IBCEQTCQKRMFYYDENBWHA5DYPQ====` maps to the +principal ID `pk_aaaqeayeaudaocajbifqydiob4`. APIs such as `list_acls` will +return the principal ID, never the API key. + +OpenID principals also typically have a human-readable name associated with +them. For example, a principal might be named `bob@example.com` and might have +a principal ID of `8a004400-9443-4fab-9997-b9372a7ca8fd`. Whenever +possible, APIs will return the principal name for convenience. However, +principal names are mutable and can be changed by the identity provider. + +Here is an example of getting information about an API key: +```python +from lancedb.authz import Subject + +catalog = lancedb.connect_catalog("https://catalog.example.com", api_key=api_key) +principal = catalog.authz.get_principal(Subject.principal_api_key(api_key)) +print(f"Found principal {principal.principal_id}") +``` + +OpenID principals can be members of groups. Like principals, groups also have +IDs and human-readable names. It is not possible to modify group membership via +this API; it is managed by the OpenID identity provider. + +Here is an example of getting information about a group based on its name: +```python +from lancedb.authz import Subject + +catalog = lancedb.connect_catalog("https://catalog.example.com", api_key=api_key) +group = catalog.authz.get_group(Subject.group_name("Engineering")) +print(f"Found group {group.group_id} with name {group.group_name}") +``` + +Roles can contain either principals or groups. Unlike groups, they are managed +by the catalog, not the identity provider. Roles can also contain API key +principals. + +Here is an example of managing the `reader` role: +```python +from lancedb.authz import Subject + +catalog = lancedb.connect_catalog("https://catalog.example.com", api_key=api_key) +catalog.authz.add_role_binding(subject=Subject.principal_name("Alice"), role="reader") +``` +It may take a few milliseconds for changes to role bindings to become visible +to all clients. + +More subject types may be added in the future. + +### Objects + +An object is any resource in the system that we can attach authorization rules to. +The Python builders and Rust client share resource validation and wire encoding. +Methods also accept literal strings such as `table:tenant/db:public$events`; +malformed names for known resource types raise `ValueError` locally before a +request is sent. Unknown resource types are preserved for compatibility with +newer servers. + +Object types include: + +The System object, a singleton which represents the entire system. You can +refer to it as `Object.system()`. + +Database objects which represents specific databases. For example, +`Object.database("analytics")` represents the "analytics" database. + +Table objects which represent specific tables. For example, +`Object.table(database="analytics", name="events")` represents the "events" +table inside the "analytics" database. Tables are in the default namespace (which +is named "public") unless a specific namespace is given. + +Similarly, View objects represent specific views. For example, +`Object.view(database="analytics", name="recent_events")` represents the +"recent_events" view inside the "analytics" database. Views are in the default +namespace (which is named "public") unless a specific namespace is given. + +Secret objects identify named credentials, for example, +`Object.secret(database="analytics", name="api-key")`. Function objects identify +registered functions, for example, +`Object.function(database="analytics", name="caption")`. Both default to the +`public` namespace. Pass `namespace_path=["platform", "ml"]` to select a nested +namespace. These objects identify resources for ACL operations; a secret object +contains the credential's name, not its value. + +More object types may be added in the future. + +### Access control lists + +An access control list attaches a specific set of privileges to a specific +(subject, object) pair. So, for example, we might have an access control list +that gives `bob@example.com` the ability to read the foo table, that looks like +this: + +```python +from lancedb.authz import Object, Subject, Privilege + +catalog = lancedb.connect_catalog("https://catalog.example.com", api_key=api_key) +catalog.authz.add_acl( + object=Object.table(database="analytics", name="foo"), + subject=Subject.principal_name("bob@example.com"), + privilege=Privilege.SELECT, +) +``` + +Reading a table generally requires USAGE on its database and namespace, plus +SELECT on the table. Each object has at most one owner. + +OWNERSHIP is a special privilege which grants all the other privileges. By +default, the principal who created an object owns it. Only one subject can own +an object; granting OWNERSHIP to a new subject removes it from the previous +owner. OWNERSHIP cannot be deleted. + +To delegate secret or function creation, grant `CREATE_SECRET` or +`CREATE_FUNCTION` on the containing namespace. This preserves namespace ownership: + +```python +for privilege in [Privilege.CREATE_SECRET, Privilege.CREATE_FUNCTION]: + catalog.authz.add_acl( + object=Object.namespace(database="analytics", namespace_path=["public"]), + subject=Subject.role("developer"), + privilege=privilege, + ) +``` + +Both privileges also accept their string names in synchronous and asynchronous +ACL methods. + +By default, the `list_acls` function lists all the access control entries in +the system. By adding a function parameter, the output is filtered to just ACLs +matching that parameter. For example, this code lists all the ACLs for the +"reader" role: + +```python +from lancedb.authz import Subject + +catalog = lancedb.connect_catalog("https://catalog.example.com", api_key=api_key) +for entry in catalog.authz.list_acls(subject=Subject.role("reader")): + print(entry) +``` + +Or, when using the async API: +```python +catalog = await lancedb.connect_catalog_async(endpoint, api_key=api_key) +async for entry in catalog.authz.list_acls(subject=Subject.role("reader")): + print(entry) +``` + ## Remote SQL Submit SQL against a remote LanceDB database through the connection. @@ -342,6 +505,12 @@ pylance, LanceDB uses a matching `lance.blob.v2` extension type so blob columns still work. Queries return descriptors. Call [`fetch_blob_files`][lancedb.table.Table.fetch_blob_files] for lazy reads or [`fetch_blobs`][lancedb.table.Table.fetch_blobs] for eager bytes. +One missing or deleted row ID rejects the entire blob-read batch; the APIs do +not return partial results. + +Blobs inside lists can be written and queried as descriptors. They are omitted +from `blob_columns()` because row ids do not identify individual list elements, +so the fetch APIs cannot read their bytes yet. ::: lancedb.blob @@ -359,6 +528,8 @@ still work. Queries return descriptors. Call ## Exceptions +::: lancedb.remote.errors.HttpError + ::: lancedb.exceptions.MissingValueError ::: lancedb.exceptions.MissingColumnError diff --git a/docs/web-tests/py/test_multimodal.py b/docs/web-tests/py/test_multimodal.py index cb4a647c3..bf7ec1bc8 100644 --- a/docs/web-tests/py/test_multimodal.py +++ b/docs/web-tests/py/test_multimodal.py @@ -215,18 +215,28 @@ async def test_query_to_pandas_kwargs(db_path_factory): .to_pandas(blob_mode="bytes") ) - # Vector / FTS / hybrid queries can't materialize blob columns, - # so omit them from the projection + # Vector / FTS / hybrid queries return "bytes" or "descriptions". + # Lazy BlobFile handles need a plain scan. df_vec = ( tbl.search(query_vector) .limit(10) - .select(["id", "vector"]) - .to_pandas(split_blocks=True, self_destruct=True) + .select(["id", "video"]) + .to_pandas(blob_mode="bytes", split_blocks=True, self_destruct=True) ) + + # to_arrow and to_list take the same blob_mode + hits = tbl.search(query_vector).limit(2).select(["id", "video"]) + arrow_bytes = hits.to_arrow(blob_mode="bytes") + rows = hits.to_list(blob_mode="bytes") # --8<-- [end:query_to_pandas_kwargs] assert len(df_lazy) == 1 assert len(df_bytes) == 1 assert df_bytes["video"].iloc[0] == b"fake_video_bytes_1" assert len(df_vec) == 10 - assert "video" not in df_vec.columns + assert all( + video == f"fake_video_bytes_{i}".encode() + for i, video in zip(df_vec["id"], df_vec["video"]) + ) + assert arrow_bytes["video"].to_pylist() == [row["video"] for row in rows] + assert all(isinstance(row["video"], bytes) for row in rows) diff --git a/docs/web-tests/rs/Cargo.lock b/docs/web-tests/rs/Cargo.lock index 0dff8ddfc..599633c5e 100644 --- a/docs/web-tests/rs/Cargo.lock +++ b/docs/web-tests/rs/Cargo.lock @@ -556,6 +556,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bs58" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4" +dependencies = [ + "tinyvec", +] + [[package]] name = "bumpalo" version = "3.19.0" @@ -1065,12 +1074,12 @@ dependencies = [ [[package]] name = "darling" -version = "0.21.3" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" dependencies = [ - "darling_core 0.21.3", - "darling_macro 0.21.3", + "darling_core 0.23.0", + "darling_macro 0.23.0", ] [[package]] @@ -1089,11 +1098,10 @@ dependencies = [ [[package]] name = "darling_core" -version = "0.21.3" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" dependencies = [ - "fnv", "ident_case", "proc-macro2", "quote", @@ -1114,11 +1122,11 @@ dependencies = [ [[package]] name = "darling_macro" -version = "0.21.3" +version = "0.23.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ - "darling_core 0.21.3", + "darling_core 0.23.0", "quote", "syn 2.0.117", ] @@ -1177,7 +1185,7 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.2", + "rand 0.9.5", "regex", "sqlparser 0.61.0", "tempfile", @@ -1292,7 +1300,7 @@ dependencies = [ "itertools 0.14.0", "log", "object_store", - "rand 0.9.2", + "rand 0.9.5", "tokio", "url", ] @@ -1392,7 +1400,7 @@ dependencies = [ "log", "object_store", "parking_lot", - "rand 0.9.2", + "rand 0.9.5", "tempfile", "url", ] @@ -1457,7 +1465,7 @@ dependencies = [ "md-5", "memchr", "num-traits", - "rand 0.9.2", + "rand 0.9.5", "regex", "sha2", "unicode-segmentation", @@ -2078,7 +2086,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7af6e24ed12cf382082d5a7f365df2b8e3d6b1518f615d54c85165e9b9718250" dependencies = [ "arrow-array", - "rand 0.9.2", + "rand 0.9.5", ] [[package]] @@ -2498,7 +2506,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core 0.62.2", + "windows-core 0.61.2", ] [[package]] @@ -2877,7 +2885,7 @@ dependencies = [ "nom", "num-traits", "ordered-float", - "rand 0.9.2", + "rand 0.9.5", "ryu", "serde", "serde_json", @@ -2950,7 +2958,7 @@ dependencies = [ "prost", "prost-build", "prost-types", - "rand 0.9.2", + "rand 0.9.5", "rayon", "roaring", "rustc-hash", @@ -2985,7 +2993,7 @@ dependencies = [ "half", "jsonb", "num-traits", - "rand 0.9.2", + "rand 0.9.5", ] [[package]] @@ -3007,7 +3015,7 @@ dependencies = [ "half", "jsonb", "num-traits", - "rand 0.9.2", + "rand 0.9.5", ] [[package]] @@ -3071,7 +3079,7 @@ dependencies = [ "object_store", "pin-project", "prost", - "rand 0.9.2", + "rand 0.9.5", "roaring", "serde_json", "snafu 0.9.1", @@ -3110,7 +3118,7 @@ dependencies = [ "object_store", "pin-project", "prost", - "rand 0.9.2", + "rand 0.9.5", "roaring", "serde_json", "snafu 0.9.1", @@ -3169,7 +3177,7 @@ dependencies = [ "futures", "half", "hex", - "rand 0.9.2", + "rand 0.9.5", "rand_distr 0.5.1", "rand_xoshiro", ] @@ -3214,7 +3222,7 @@ dependencies = [ "num-traits", "prost", "prost-build", - "rand 0.9.2", + "rand 0.9.5", "strum 0.26.3", "tokio", "tracing", @@ -3306,7 +3314,7 @@ dependencies = [ "prost", "prost-build", "prost-types", - "rand 0.9.2", + "rand 0.9.5", "rand_distr 0.5.1", "rangemap", "rayon", @@ -3352,7 +3360,7 @@ dependencies = [ "path_abs", "pin-project", "prost", - "rand 0.9.2", + "rand 0.9.5", "serde", "tempfile", "tokio", @@ -3374,7 +3382,7 @@ dependencies = [ "lance-arrow 8.0.0", "lance-core 8.0.0", "num-traits", - "rand 0.9.2", + "rand 0.9.5", ] [[package]] @@ -3428,7 +3436,7 @@ dependencies = [ "lance-table", "log", "object_store", - "rand 0.9.2", + "rand 0.9.5", "roaring", "serde_json", "time", @@ -3508,7 +3516,7 @@ dependencies = [ "prost", "prost-build", "prost-types", - "rand 0.9.2", + "rand 0.9.5", "rangemap", "roaring", "semver", @@ -3533,7 +3541,7 @@ dependencies = [ "lance-arrow 8.0.0", "num-traits", "pprof", - "rand 0.9.2", + "rand 0.9.5", ] [[package]] @@ -3603,7 +3611,7 @@ dependencies = [ "pin-project", "polars", "polars-arrow", - "rand 0.9.2", + "rand 0.9.5", "regex", "semver", "serde", @@ -4366,7 +4374,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" dependencies = [ "phf_shared 0.11.3", - "rand 0.8.5", + "rand 0.8.8", ] [[package]] @@ -4584,7 +4592,7 @@ dependencies = [ "polars-error", "polars-row", "polars-utils", - "rand 0.8.5", + "rand 0.8.8", "rand_distr 0.4.3", "rayon", "regex", @@ -4783,7 +4791,7 @@ dependencies = [ "polars-error", "polars-lazy", "polars-plan", - "rand 0.8.5", + "rand 0.8.8", "serde", "serde_json", "sqlparser 0.39.0", @@ -5096,9 +5104,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.8.5" +version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" dependencies = [ "libc", "rand_chacha 0.3.1", @@ -5107,9 +5115,9 @@ dependencies = [ [[package]] name = "rand" -version = "0.9.2" +version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ "rand_chacha 0.9.0", "rand_core 0.9.3", @@ -5177,7 +5185,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" dependencies = [ "num-traits", - "rand 0.8.5", + "rand 0.8.8", ] [[package]] @@ -5187,7 +5195,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" dependencies = [ "num-traits", - "rand 0.9.2", + "rand 0.9.5", ] [[package]] @@ -5533,9 +5541,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.35" +version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "533f54bc6a7d4f647e46ad909549eda97bf5afc1585190ef692b4286b198bd8f" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "once_cell", "ring", @@ -5569,9 +5577,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ "ring", "rustls-pki-types", @@ -5747,15 +5755,17 @@ dependencies = [ [[package]] name = "serde_with" -version = "3.16.1" +version = "3.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fa237f2807440d238e0364a218270b98f767a00d3dada77b1c53ae88940e2e7" +checksum = "ee78f1fbe43ac4a0e47aadb3dbd357b69eb0d3793e948624cd03dd2750ab1c0a" dependencies = [ "base64 0.22.1", + "bs58", "chrono", "hex", "indexmap 1.9.3", "indexmap 2.14.0", + "jiff", "schemars 0.9.0", "schemars 1.1.0", "serde_core", @@ -5766,11 +5776,11 @@ dependencies = [ [[package]] name = "serde_with_macros" -version = "3.16.1" +version = "3.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52a8e3ca0ca629121f70ab50f95249e5a6f925cc0f6ffe8256c45b728875706c" +checksum = "8705578779c2b6bd90d84d66eb2e206b708b1a4d7b9f17641b293545bf1c7e46" dependencies = [ - "darling 0.21.3", + "darling 0.23.0", "proc-macro2", "quote", "syn 2.0.117", @@ -6544,7 +6554,7 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" dependencies = [ - "rand 0.9.2", + "rand 0.9.5", ] [[package]] @@ -6917,21 +6927,8 @@ dependencies = [ "windows-implement", "windows-interface", "windows-link 0.1.3", - "windows-result 0.3.4", - "windows-strings 0.4.2", -] - -[[package]] -name = "windows-core" -version = "0.62.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" -dependencies = [ - "windows-implement", - "windows-interface", - "windows-link 0.2.1", - "windows-result 0.4.1", - "windows-strings 0.5.1", + "windows-result", + "windows-strings", ] [[package]] @@ -6998,15 +6995,6 @@ dependencies = [ "windows-link 0.1.3", ] -[[package]] -name = "windows-result" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" -dependencies = [ - "windows-link 0.2.1", -] - [[package]] name = "windows-strings" version = "0.4.2" @@ -7016,15 +7004,6 @@ dependencies = [ "windows-link 0.1.3", ] -[[package]] -name = "windows-strings" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" -dependencies = [ - "windows-link 0.2.1", -] - [[package]] name = "windows-sys" version = "0.52.0" @@ -7310,9 +7289,9 @@ dependencies = [ [[package]] name = "xxhash-rust" -version = "0.8.15" +version = "0.8.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +checksum = "550a2b930b62486a393c52d5c3b84bff264b28aa437ed64694d31e93b1757af7" [[package]] name = "yoke" diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index 3cd3a970a..b10da3f6f 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -139,9 +139,9 @@ Learn how to configure, build, and search LanceDB vector indexes, including buil | Parameter | Description | | :--- | :--- | | `limit` | Number of results to return (the `k` in `k-ANN`). | -| `nprobes` | Shorthand that sets both `minimum_nprobes` and `maximum_nprobes` to the same value. LanceDB defaults this to 20 which is a good default for many cases. At very high scales you will likely need to experiment to find your preferred recall / latency balance. | -| `minimum_nprobes` | Minimum number of partitions scanned. | -| `maximum_nprobes` | Maximum number of partitions scanned. Only scans more than `minimum_nprobes` if an initial pass does not return enough results — useful for narrow filters. Set to `0` to remove the cap. | +| `nprobes` | Sets both `minimum_nprobes` and `maximum_nprobes` to the same value in Lance. Leave unset to use Lance's adaptive defaults. At very high scales you will likely need to experiment to find your preferred recall / latency balance. | +| `minimum_nprobes` | Minimum number of partitions scanned. Leave unset to use Lance's default. | +| `maximum_nprobes` | Maximum number of partitions scanned. Only scans more than `minimum_nprobes` if an initial pass does not return enough results — useful for narrow filters. Leave unset, or set to `0`, to remove the cap. | | `ef` | HNSW search-time exploration factor. Start around `1.5 * k` and increase up to `10 * k` for higher recall. | | `refine_factor` | Reads and reranks additional candidates in memory to recover recall lost to quantization. Start between `5` and `50` (a factor of `25` reads and reranks the top `25 * k` candidates) and measure the latency cost, typically a few milliseconds, against the recall gained. | | `distance_range(lower_bound, upper_bound)` | Return only rows whose distance falls within `[lower_bound, upper_bound)`. Either bound is optional. Useful for near-duplicate detection or "close-enough" matching. | @@ -165,8 +165,8 @@ older client. | Index type | Guidance | | :--- | :--- | -| `IVF_HNSW_FLAT`, `IVF_HNSW_SQ` | Keep the default `nprobes`, then tune `ef` first. Expect higher latency variance under filtered search. | -| `IVF_RQ`, `IVF_PQ` | Keep the default `nprobes` and raise when recall is insufficient. If using a selective filter you may need to increase `maximum_nprobes` if you do not received enough results. | +| `IVF_HNSW_FLAT`, `IVF_HNSW_SQ` | Leave the probe settings unset, then tune `ef` first. Expect higher latency variance under filtered search. | +| `IVF_RQ`, `IVF_PQ` | Leave the probe settings unset and configure them only when recall is insufficient. If using a selective filter you may need to increase `maximum_nprobes` if you do not receive enough results. | diff --git a/docs/web/management/table.mdx b/docs/web/management/table.mdx index 7f84b9131..447b59029 100644 --- a/docs/web/management/table.mdx +++ b/docs/web/management/table.mdx @@ -456,8 +456,6 @@ When you call `to_pandas()` on a local LanceDB table that contains Blob API colu - `"bytes"`: eagerly materializes each blob as `bytes`. Use this when you need the raw payload in the DataFrame, for example to decode an image or audio clip in-process. - `"descriptions"`: returns blob descriptors (offsets, sizes, and positions) instead of the data itself. Use this when you want to plan I/O without paying the cost of loading every blob. -`"bytes"` and `"descriptions"` require a filesystem-backed Lance dataset and are not supported on in-memory tables. - Extra keyword arguments are forwarded to the underlying PyArrow / Lance pandas conversion, so you can also pass options like `split_blocks` or `self_destruct`: @@ -466,10 +464,13 @@ Extra keyword arguments are forwarded to the underlying PyArrow / Lance pandas c -Query builders also accept `blob_mode` on their `to_pandas()` method: +Query builders also accept `blob_mode` on `to_pandas()`, `to_arrow()`, and `to_list()`: - Plain scans support `"lazy"`, `"bytes"`, and `"descriptions"` with filters, projections, aliases, `limit`, and `offset`. -- Vector, FTS, hybrid, and ordered queries can't materialize blob columns through `to_pandas()`; omit blob columns from the projection for those query shapes. +- Vector, FTS, hybrid, and take queries support `"bytes"` and `"descriptions"`. Lazy `BlobFile` handles need a plain scan. +- In `to_arrow()` and `to_list()`, `"lazy"` returns blob descriptors that carry the row id `fetch_blobs` reads, and `"descriptions"` returns them without it. +- A plain scan's `to_pandas()` without `timeout` reads the local dataset through Lance's scanner, so it also returns `"bytes"` for legacy blob columns and under namespace query pushdown. In-memory tables have no local dataset, so they take the `to_arrow()` route. +- Every other `"bytes"` route raises `ValueError` for legacy blob columns and blobs inside lists, and `NotImplementedError` on remote tables and namespace query pushdown, where the blob fetch could read a different table version than the query. - This works on both sync and async query builders. Extra PyArrow kwargs like `split_blocks` and `self_destruct` are still forwarded. diff --git a/docs/web/snippets/multimodal.mdx b/docs/web/snippets/multimodal.mdx index f41af6629..90e275db7 100644 --- a/docs/web/snippets/multimodal.mdx +++ b/docs/web/snippets/multimodal.mdx @@ -16,7 +16,7 @@ export const PyMultimodalImports = "import lancedb\nimport pyarrow as pa\nimport export const PyProcessResults = "# Convert back to PIL Image\nfor _, row in results.iterrows():\n image_bytes = row['image_blob']\n image = Image.open(io.BytesIO(image_bytes))\n print(f\"Retrieved image: {row['filename']}, Size: {image.size}\")\n # You can now use 'image' with other libraries or display it\n"; -export const PyQueryToPandasKwargs = "# Plain scan query: blob_mode is supported end to end\ndf_lazy = (\n tbl.search()\n .where(\"id = 1\")\n .select([\"id\", \"video\"])\n .to_pandas(blob_mode=\"lazy\")\n)\n\n# Same call shape works on async query builders\ndf_bytes = await (\n tbl_async.query()\n .where(\"id = 1\")\n .select([\"id\", \"video\"])\n .to_pandas(blob_mode=\"bytes\")\n)\n\n# Vector / FTS / hybrid queries can't materialize blob columns,\n# so omit them from the projection\ndf_vec = (\n tbl.search(query_vector)\n .limit(10)\n .select([\"id\", \"vector\"])\n .to_pandas(split_blocks=True, self_destruct=True)\n)\n"; +export const PyQueryToPandasKwargs = "# Plain scan query: blob_mode is supported end to end\ndf_lazy = (\n tbl.search()\n .where(\"id = 1\")\n .select([\"id\", \"video\"])\n .to_pandas(blob_mode=\"lazy\")\n)\n\n# Same call shape works on async query builders\ndf_bytes = await (\n tbl_async.query()\n .where(\"id = 1\")\n .select([\"id\", \"video\"])\n .to_pandas(blob_mode=\"bytes\")\n)\n\n# Vector / FTS / hybrid queries return \"bytes\" or \"descriptions\".\n# Lazy BlobFile handles need a plain scan.\ndf_vec = (\n tbl.search(query_vector)\n .limit(10)\n .select([\"id\", \"video\"])\n .to_pandas(blob_mode=\"bytes\", split_blocks=True, self_destruct=True)\n)\n\n# to_arrow and to_list take the same blob_mode\nhits = tbl.search(query_vector).limit(2).select([\"id\", \"video\"])\narrow_bytes = hits.to_arrow(blob_mode=\"bytes\")\nrows = hits.to_list(blob_mode=\"bytes\")\n"; export const PySearchData = "# Search for similar images\nquery_vector = np.random.rand(128).astype(np.float32)\nresults = tbl.search(query_vector).limit(1).to_pandas()\n"; diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index c5f30e13d..b312ffebf 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.1 + 0.41.0-beta.6 ../pom.xml @@ -44,7 +44,30 @@ com.fasterxml.jackson.core jackson-databind - 2.17.1 + + + + + com.fasterxml.jackson.core + jackson-core + + + + com.fasterxml.jackson.core + jackson-annotations + + + + com.fasterxml.jackson.datatype + jackson-datatype-jsr310 + + + + com.fasterxml.jackson.jaxrs + jackson-jaxrs-json-provider diff --git a/java/pom.xml b/java/pom.xml index d01e30465..33446040b 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.1 + 0.41.0-beta.6 pom ${project.artifactId} LanceDB Java SDK Parent POM @@ -28,7 +28,8 @@ UTF-8 15.0.0 - 14.0.0-beta.7 + 14.0.0-beta.15 + 2.18.11 false 2.30.0 1.7 @@ -62,6 +63,15 @@ + + + com.fasterxml.jackson + jackson-bom + ${jackson.version} + pom + import + org.lance lance-core diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 20363ec85..9404a0913 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.1" +version = "0.41.0-beta.6" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/__test__/catalog.test.ts b/nodejs/__test__/catalog.test.ts index 945956b5b..f8b14ca6f 100644 --- a/nodejs/__test__/catalog.test.ts +++ b/nodejs/__test__/catalog.test.ts @@ -266,8 +266,41 @@ describe("remote catalog", () => { "limit", ); } - await expect(catalog.connectDatabase("a$b")).rejects.toThrow( - "Invalid database name", - ); }); + + it.each([ + "", + ".", + "..", + "/db", + "db/", + "a//b", + "a/./b", + "a/../b", + "a b", + "a:b", + "a$b", + "a%b", + "a?b", + "a#b", + "a\\b", + "café", + "a\nb", + ])( + "rejects invalid database name %j before sending requests", + async (name) => { + await withCatalog([], async (catalog, requests) => { + for (const method of [ + "createDatabase", + "connectDatabase", + "dropDatabase", + ] as const) { + await expect(catalog[method](name)).rejects.toThrow( + "Invalid database name", + ); + } + expect(requests).toEqual([]); + }); + }, + ); }); diff --git a/nodejs/__test__/query.test.ts b/nodejs/__test__/query.test.ts index 16a595002..786f04015 100644 --- a/nodejs/__test__/query.test.ts +++ b/nodejs/__test__/query.test.ts @@ -527,3 +527,106 @@ describe("Query orderBy", () => { ]); }); }); + +describe("IVF search effort", () => { + let directory: tmp.DirResult; + let table: Table; + const vector = Array.from({ length: 16 }, (_, i) => (i === 0 ? 1 : 0)); + + beforeEach(async () => { + directory = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(directory.name); + const rows = Array.from({ length: 512 }, (_, id) => ({ + id, + vector: Array.from( + { length: 16 }, + (_, j) => + (Math.floor(id / 32) === j ? 1 : 0) + Math.sin(id * 17 + j) * 0.001, + ), + })); + table = await db.createTable("effort", rows.slice(0, 256)); + await table.add(rows.slice(256)); + await table.createIndex("vector", { + config: Index.ivfFlat({ numPartitions: 16 }), + }); + }); + + afterEach(() => { + directory.removeCallback(); + }); + + it.each([ + [null, 6], + [0, 1], + [0.25, 3], + [0.5, 6], + [0.75, 10], + [1, 16], + ])( + "preserves effort %s through execution and probe precedence", + async (effort, expected) => { + const partitions = async (query: ReturnType) => { + const plan = await query.limit(8).analyzePlan(); + const matches = [...plan.matchAll(/partitions_searched=(\d+)/g)]; + expect(matches.length).toBeGreaterThan(0); + return matches.reduce((total, match) => total + Number(match[1]), 0); + }; + expect( + await partitions(table.vectorSearch(vector).searchEffort(effort)), + ).toBe(expected); + expect( + await partitions( + table.vectorSearch(vector).nprobes(1).searchEffort(effort), + ), + ).toBe(effort === null ? 1 : expected); + expect( + await partitions( + table.vectorSearch(vector).searchEffort(effort).nprobes(1), + ), + ).toBe(effort === null ? 1 : expected); + expect( + await partitions( + table + .vectorSearch(vector) + .searchEffort(effort) + .nprobes(1) + .searchEffort(null), + ), + ).toBe(1); + expect( + await partitions( + table.vectorSearch(vector).searchEffort(effort).maximumNprobes(2), + ), + ).toBeLessThanOrEqual(2); + const actual = await table + .vectorSearch(vector) + .searchEffort(effort) + .limit(8) + .toArray(); + const exact = await table + .vectorSearch(vector) + .bypassVectorIndex() + .limit(8) + .toArray(); + const ids = new Set(exact.map((row) => row.id)); + expect( + actual.filter((row) => ids.has(row.id)).length / 8, + ).toBeGreaterThanOrEqual(0.5); + }, + ); + + it.each([ + -0.1, + 1.1, + Number.NaN, + Number.POSITIVE_INFINITY, + Number.NEGATIVE_INFINITY, + ])("rejects invalid effort %s before execution", (effort) => { + expect(() => table.vectorSearch(vector).searchEffort(effort)).toThrow( + "search_effort must be finite and in [0, 1]", + ); + expect(() => + table.vectorSearch(vector).fullTextSearch("dog").searchEffort(effort), + ).toThrow("search_effort must be finite and in [0, 1]"); + }); +}); diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 6911b24b9..4f70ac1aa 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -480,6 +480,15 @@ describe("remote connection", () => { ); }); + it("rejects a zero blob request concurrency limit", async () => { + await expect( + connect("db://dev", { + apiKey: "fake", + clientConfig: { blobRequestConcurrency: 0 }, + }), + ).rejects.toThrow("blob_request_concurrency"); + }); + it("shows the full error messages on retry errors", async () => { await withMockDatabase( (_req, res) => { diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index eee40fb9c..6c770494b 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -1231,7 +1231,9 @@ describe("When creating an index", () => { expect(() => tbl.query().nearestTo(queryVec).minimumNprobes(0)).toThrow( "Invalid input, minimum_nprobes must be greater than 0", ); - expect(() => tbl.query().nearestTo(queryVec).maximumNprobes(5)).toThrow( + expect(() => + tbl.query().nearestTo(queryVec).minimumNprobes(10).maximumNprobes(5), + ).toThrow( "Invalid input, maximum_nprobes must be greater than or equal to minimum_nprobes", ); @@ -2815,6 +2817,18 @@ describe("when dealing with blob columns", () => { expect(Buffer.from(await files[1]!.read()).toString()).toBe("beta"); }); + it("rejects a whole blob read batch when one row ID was deleted", async () => { + const { table, rowIds } = await openBlobTable(); + await table.delete("id = 2"); + + await expect(table.fetchBlobs("image", rowIds)).rejects.toThrow( + `first missing row ids: [${rowIds[1]}]`, + ); + await expect(table.fetchBlobFiles("image", rowIds)).rejects.toThrow( + `first missing row ids: [${rowIds[1]}]`, + ); + }); + it("adds a LargeBinary Arrow column to a blob table", async () => { const db = await connect(tmpDir.name); const schema = new Schema([new Field("id", new Int64()), blob("image")]); @@ -3027,7 +3041,7 @@ describe("when dealing with blob columns", () => { expect(bytes[0]!.equals(payload)).toBe(true); }); - it("creates and adds list blob columns", async () => { + it("creates list blobs without advertising unfetchable elements", async () => { const db = await connect(tmpDir.name); const schema = new Schema([ new Field("id", new Int64(), true), @@ -3046,7 +3060,16 @@ describe("when dealing with blob columns", () => { { id: 3n, images: [gamma, null] }, { id: 4n, images: [] }, ]); - expect(await table.blobColumns()).toEqual(["images.image"]); + expect(await table.blobColumns()).toEqual([]); + const [first] = await table.query().withRowId().toArray(); + for (const column of ["images", "images.image"]) { + await expect( + table.fetchBlobs(column, [first._rowid as bigint]), + ).rejects.toThrow(/blobs inside lists cannot be fetched/); + await expect( + table.fetchBlobFiles(column, [first._rowid as bigint]), + ).rejects.toThrow(/blobs inside lists cannot be fetched/); + } const rows = await table.query().toArray(); const byId = new Map(rows.map((row) => [Number(row.id), row])); expect(descriptorSizes(byId.get(1)!.images)).toEqual([ @@ -3092,6 +3115,7 @@ describe("when dealing with blob columns", () => { ], }, ]); + expect(await table.blobColumns()).toEqual([]); const rows = await table.query().toArray(); const byId = new Map(rows.map((row) => [Number(row.id), row])); expect( diff --git a/nodejs/lancedb/blob.ts b/nodejs/lancedb/blob.ts index f3d648f48..7b0276f14 100644 --- a/nodejs/lancedb/blob.ts +++ b/nodejs/lancedb/blob.ts @@ -81,6 +81,8 @@ export type BlobOptions = { * * Query results are descriptors, not payload bytes. Use {@link Table.fetchBlobs} * or {@link Table.fetchBlobFiles} to read bytes. + * Blobs inside lists can be written and queried as descriptors, but cannot + * currently be fetched because a row ID does not identify a list element. * * @example * ```ts diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index f370e1b31..cce505c97 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -532,17 +532,16 @@ export class VectorQuery extends StandardQueryBase { * partitions should be searched. * * Increasing this value will increase the recall of your query but will - * also increase the latency of your query. The default value is 20. This - * default is good for many cases but the best value to use will depend on - * your data and the recall that you need to achieve. + * also increase the latency of your query. If this method is not called, + * Lance's adaptive probe defaults are used. * * For best results we recommend tuning this parameter with a benchmark against * your actual data to find the smallest possible value that will still give * you the desired recall. * - * For more fine grained control over behavior when you have a very narrow filter - * you can use `minimumNprobes` and `maximumNprobes`. This method sets both - * the minimum and maximum to the same value. + * The value is retained as `nprobes` through client and server request + * construction. Lance sets both probe bounds to this value. Explicit minimum + * or maximum settings can override their respective bounds. */ nprobes(nprobes: number): VectorQuery { this.doVectorCall((inner) => inner.nprobes(nprobes)); @@ -550,6 +549,27 @@ export class VectorQuery extends StandardQueryBase { return this; } + /** + * Set the initial IVF search effort to a finite value in [0, 1]. + * + * The default, null, uses Auto (internally 0.5). Zero starts at the caller + * minimum; one searches all partitions within the caller maximum. Later + * probing may continue when too few candidates are found. + * Any explicit effort, including 0.5, overrides nprobes in either call order. + * Explicit minimum and maximum bounds still apply. This is not a recall + * guarantee and does not change refinement or HNSW ef. The MemWAL LSM + * scanner does not yet support this option. + * + * @example + * ```ts + * table.vectorSearch([0.1, 0.2]).searchEffort(0.75).limit(10); + * ``` + */ + searchEffort(searchEffort: number | null): VectorQuery { + this.doVectorCall((inner) => inner.searchEffort(searchEffort)); + return this; + } + /** * Set the minimum number of probes used. * diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index eae718168..630d06f29 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -526,7 +526,8 @@ export abstract class Table { abstract takeRowIds(rowIds: readonly (bigint | number)[]): TakeQuery; /** - * Blob v2 columns, including nested dotted paths. + * Blob v2 columns addressable by row ID, including struct-nested dotted paths. + * Blobs inside lists are omitted because a row ID does not identify an element. */ abstract blobColumns(): Promise; @@ -536,6 +537,10 @@ export abstract class Table { * Reads the table's current checkout. IDs from another version can fail after * compaction unless stable row ids are enabled. Results keep input order and * duplicates. Null blobs are `null`. Empty blobs are empty buffers. + * One missing or deleted row ID rejects the entire batch, including on remote + * tables; no partial results are returned. + * Blobs inside lists cannot be fetched by row ID. + * * Remote servers limit each request to 1024 row IDs and 64 MiB of blob bytes. * The client splits requests automatically and reads an individual larger * blob through the Range route. This method still materializes all bytes in @@ -553,6 +558,8 @@ export abstract class Table { * * Preserves input order, duplicates, and nulls. Use this for large payloads. * See {@link Table.fetchBlobs} for row-ID validity across versions. + * One missing or deleted row ID rejects the entire batch, including on remote + * tables; no partial handles are returned. */ abstract fetchBlobFiles( column: string, diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 9ff0cd5e8..08e5b3a41 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.41.0-beta.1", + "version": "0.41.0-beta.6", "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 a04377631..5aaa172ac 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.41.0-beta.1", + "version": "0.41.0-beta.6", "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 da117d8a5..cb1881aaf 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.41.0-beta.1", + "version": "0.41.0-beta.6", "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 46189293d..9881b85d1 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.41.0-beta.1", + "version": "0.41.0-beta.6", "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 c576aa720..29f3e7986 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.41.0-beta.1", + "version": "0.41.0-beta.6", "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 b73302811..f70f4eacd 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.41.0-beta.1", + "version": "0.41.0-beta.6", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index bb686b5bb..c06f89a36 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.41.0-beta.1", + "version": "0.41.0-beta.6", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 71232b770..4c9be04a7 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.1", + "version": "0.41.0-beta.6", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index 712569b2f..480426411 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -298,6 +298,16 @@ impl VectorQuery { Ok(()) } + #[napi] + pub fn search_effort(&mut self, search_effort: Option) -> napi::Result<()> { + self.inner = self + .inner + .clone() + .search_effort(search_effort) + .default_error()?; + Ok(()) + } + #[napi] pub fn minimum_nprobes(&mut self, minimum_nprobe: u32) -> napi::Result<()> { self.inner = self diff --git a/nodejs/src/remote.rs b/nodejs/src/remote.rs index c2d671735..d7e2a21cd 100644 --- a/nodejs/src/remote.rs +++ b/nodejs/src/remote.rs @@ -92,6 +92,10 @@ pub struct ClientConfig { pub user_agent: Option, pub retry_config: Option, pub timeout_config: Option, + /// Maximum number of concurrent blob HTTP requests across all tables and + /// blob handles on this connection. Defaults to 8. Can also be set with + /// `LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY`. Must be greater than zero. + pub blob_request_concurrency: Option, pub extra_headers: Option>, /// The delimiter joining a namespace path and a name into one object /// identifier. `"$"` is the only supported value, and leaving this unset is @@ -428,6 +432,7 @@ impl From for lancedb::remote::ClientConfig { .unwrap_or(concat!("LanceDB-Node-Client/", env!("CARGO_PKG_VERSION")).to_string()), retry_config: config.retry_config.map(Into::into).unwrap_or_default(), timeout_config: config.timeout_config.map(Into::into).unwrap_or_default(), + blob_request_concurrency: config.blob_request_concurrency.map(|value| value as usize), extra_headers: config.extra_headers.unwrap_or_default(), id_delimiter: config.id_delimiter, tls_config: config.tls_config.map(Into::into), diff --git a/python/Cargo.toml b/python/Cargo.toml index 94fe5fddf..40b74520e 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.1" +version = "0.41.0-beta.6" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/python/pyproject.toml b/python/pyproject.toml index 2d753c139..ed54ae536 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "pyarrow>=16", "pydantic>=2.7.4,<3", "tqdm>=4.27.0", - "lance-namespace>=0.3.2" + "lance-namespace>=0.13.0" ] description = "lancedb" authors = [{ name = "LanceDB Devs", email = "dev@lancedb.com" }] @@ -63,7 +63,7 @@ tests = [ "polars>=0.19, <=1.32.3", "pyarrow<25", "pyarrow-stubs>=16.0", - "pylance==14.0.0b6", + "pylance==14.0.0b12", "requests>=2.31.0", "datafusion>=54,<55", "opentelemetry-sdk>=1.30.0", @@ -73,6 +73,7 @@ dev = [ "pre-commit>=3.5.0", "pyright>=1.1.350", 'typing-extensions>=4.0.0; python_version < "3.11"', + "maturin>=1.10", ] docs = ["mkdocs", "mkdocs-jupyter", "mkdocs-material", "mkdocstrings-python"] clip = ["torch", "pillow>=12.1.1", "open-clip-torch"] diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index dc4ed37df..1e9110dd0 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -6,13 +6,14 @@ from __future__ import annotations import io +import operator from collections.abc import Awaitable, Callable, Iterable from typing import TYPE_CHECKING, Optional, Union import pyarrow as pa from .expr import Expr -from .schema import row_addressable_blob_v2_paths +from .schema import blob_column_paths, row_addressable_blob_v2_paths from .types import BlobMode, QueryProjection, QueryProjectionSpec if TYPE_CHECKING: @@ -26,21 +27,24 @@ BLOB_MODE_TO_HANDLING = { ROW_ID_FIELD_NAME = "_lance_row_id" -FetchBlobsSync = Callable[[str, pa.Table], pa.Array | pa.ChunkedArray] -FetchBlobsAsync = Callable[[str, pa.Table], Awaitable[pa.Array | pa.ChunkedArray]] +FetchBlobsSync = Callable[[str, list[int]], pa.Array | pa.ChunkedArray] +FetchBlobsAsync = Callable[[str, list[int]], Awaitable[pa.Array | pa.ChunkedArray]] class BlobFile(io.RawIOBase): """Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`. - Bytes load on ``read`` or ``read_range``, not when the handle is opened. - Use :meth:`aread` from async code. + Bytes load on ``read``, ``read1``, ``read_range``, or ``read_ranges``, not + when the handle is opened. Range reads leave the cursor unchanged. Use + :meth:`aread` from async code. """ def __init__(self, inner) -> None: self._inner = inner async def aread(self) -> bytes: + if self.closed: + raise ValueError("read of closed file") return await self._inner.read() def close(self) -> None: @@ -58,13 +62,16 @@ class BlobFile(io.RawIOBase): def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: if whence == io.SEEK_SET: - self._inner.seek(offset) + position = offset elif whence == io.SEEK_CUR: - self._inner.seek(self._inner.tell() + offset) + position = self._inner.tell() + offset elif whence == io.SEEK_END: - self._inner.seek(self._inner.size() + offset) + position = self._inner.size() + offset else: raise ValueError(f"invalid whence: {whence}") + if position < 0: + raise ValueError(f"negative seek value {position}") + self._inner.seek(position) return self._inner.tell() def tell(self) -> int: @@ -74,17 +81,39 @@ class BlobFile(io.RawIOBase): return self._inner.size() def readall(self) -> bytes: + if self.closed: + raise ValueError("read of closed file") return self._inner.read_bytes() - def read(self, size: int = -1) -> bytes: - if size == -1: + def read(self, size: Optional[int] = -1) -> bytes: + if self.closed: + raise ValueError("read of closed file") + if size is None or size < 0: return self._inner.read_bytes() return super().read(size) + def read1(self, size: Optional[int] = -1) -> bytes: + """Read up to ``size`` bytes with one underlying read.""" + if self.closed: + raise ValueError("read of closed file") + if size is None or size < 0: + size = io.DEFAULT_BUFFER_SIZE + return self._inner.read_up_to(size) + def read_range(self, offset: int, length: int) -> bytes: + if self.closed: + raise ValueError("read of closed file") return self._inner.read_range(offset, length) + def read_ranges(self, ranges: list[tuple[int, int]]) -> list[bytes]: + """Read ``(offset, length)`` pairs without moving the cursor.""" + if self.closed: + raise ValueError("read of closed file") + return self._inner.read_ranges(ranges) + def readinto(self, b: WriteableBuffer) -> int: + if self.closed: + raise ValueError("readinto of closed file") view = memoryview(b).cast("B") chunk = self._inner.read_up_to(len(view)) view[: len(chunk)] = chunk @@ -107,10 +136,12 @@ def projection_includes_blob_column( columns = set(blob_columns) if not columns: return False - if projection is None: + if not projection: return True for output, source in _iter_projection_pairs(projection): - if output in columns or source in columns: + if output in columns or any( + column == source or column.startswith(f"{source}.") for column in columns + ): return True return False @@ -119,17 +150,69 @@ def blob_v2_projection_sources( schema: pa.Schema, projection: QueryProjection, ) -> dict[str, str]: - blob_columns = row_addressable_blob_v2_paths(schema) - if not blob_columns: - return {} - columns = set(blob_columns) - if projection is None: + return _project_blob_paths(row_addressable_blob_v2_paths(schema), projection) + + +def reject_unfetchable_blob_columns( + schema: pa.Schema, projection: QueryProjection, blob_sources: dict[str, str] +) -> None: + # Rerankers can drop field metadata, so validate against the source schema. + unfetchable = [ + output + for output in _project_blob_paths(blob_column_paths(schema), projection) + if output not in blob_sources + ] + if unfetchable: + raise ValueError( + f"blob_mode='bytes' cannot fetch blob columns {unfetchable}. Bytes are " + "fetched by row id, which covers blob v2 columns with one blob per " + "row but not legacy blob columns or blobs inside lists. Use " + "blob_mode='descriptions' or remove these columns from the projection." + ) + + +def reject_unmatched_blob_outputs( + output_schema: pa.Schema, blob_sources: dict[str, str] +) -> None: + # Another spelling of a blob column, such as a quoted or recased name, is a + # blob in the query's output schema but maps to no source by name. + unmatched = [ + path for path in blob_column_paths(output_schema) if path not in blob_sources + ] + if unmatched: + raise ValueError( + f"blob_mode='bytes' could not match result columns {unmatched} to a blob " + "column of the table. Select blob columns by their plain names, or " + "use blob_mode='descriptions'." + ) + + +def _project_blob_paths( + blob_columns: list[str], projection: QueryProjection +) -> dict[str, str]: + """Map each projected output path to the blob column path it reads.""" + # An empty projection selects all columns. + if not projection: return {column: column for column in blob_columns} - return { - output: source - for output, source in _iter_projection_pairs(projection) - if source in columns - } + sources: dict[str, str] = {} + entries: dict[str, str] = {} + for output, source in _iter_projection_pairs(projection): + for column in blob_columns: + if column == source: + path = output + elif column.startswith(f"{source}."): + path = f"{output}{column[len(source) :]}" + else: + continue + if path in sources: + raise ValueError( + f"projection entries {entries[path]!r} and {output!r} both " + f"produce blob column '{path}'. Rename one of them or drop " + "the redundant selection." + ) + sources[path] = column + entries[path] = output + return sources def v2_projection_needs_row_id( @@ -179,11 +262,14 @@ async def replace_v2_blob_columns_with_bytes( blob_sources: dict[str, str], fetch_blobs: FetchBlobsAsync, ) -> pa.Table: - for output_name, source_name in blob_sources.items(): - if output_name not in tbl.column_names: - continue - blobs = await fetch_blobs(source_name, tbl) - tbl = _set_blob_column(tbl, output_name, blobs) + present = _present_blob_sources(tbl, blob_sources) + if not present: + return tbl + row_ids = _hit_row_ids(tbl, present) + for output_name, source_name in present.items(): + tbl = _set_blob_column( + tbl, output_name, await fetch_blobs(source_name, row_ids) + ) return tbl @@ -192,19 +278,48 @@ def replace_v2_blob_columns_with_bytes_sync( blob_sources: dict[str, str], fetch_blobs: FetchBlobsSync, ) -> pa.Table: - for output_name, source_name in blob_sources.items(): - if output_name not in tbl.column_names: - continue - blobs = fetch_blobs(source_name, tbl) - tbl = _set_blob_column(tbl, output_name, blobs) + present = _present_blob_sources(tbl, blob_sources) + if not present: + return tbl + row_ids = _hit_row_ids(tbl, present) + for output_name, source_name in present.items(): + tbl = _set_blob_column(tbl, output_name, fetch_blobs(source_name, row_ids)) return tbl +def _present_blob_sources( + tbl: pa.Table, blob_sources: dict[str, str] +) -> dict[str, str]: + return { + output: source + for output, source in blob_sources.items() + if _has_column_path(tbl.schema, output) + } + + +def _hit_row_ids(tbl: pa.Table, blob_sources: dict[str, str]) -> list[int]: + """Read row ids before replacing columns, since a swapped alias can overwrite + the lookup source.""" + if "_rowid" in tbl.column_names: + return tbl["_rowid"].to_pylist() + for output in blob_sources: + descriptor = _leaf_struct_column(tbl, output) + if pa.types.is_struct(descriptor.type) and ( + ROW_ID_FIELD_NAME in descriptor.type.names + ): + return descriptor.field(ROW_ID_FIELD_NAME).to_pylist() + raise ValueError( + f"query result has no '_rowid' column and no '{ROW_ID_FIELD_NAME}' " + f"field on blob columns {list(blob_sources)}. Call .with_row_id(True) " + "to fetch blob bytes for this query." + ) + + def stash_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table: if "_rowid" not in tbl.column_names: raise ValueError("query result has no '_rowid' column to hide") - present_paths = [p for p in blob_paths if p.split(".")[0] in tbl.column_names] + present_paths = [p for p in blob_paths if _has_column_path(tbl.schema, p)] if not present_paths: raise ValueError("query result has no blob v2 column to carry a row id") @@ -224,7 +339,7 @@ def read_row_ids_from_hits(hits: pa.Table, blob_column: str) -> list[int]: try: leaf = _leaf_struct_column(hits, blob_column) - if ROW_ID_FIELD_NAME in leaf.type.names: + if pa.types.is_struct(leaf.type) and ROW_ID_FIELD_NAME in leaf.type.names: return leaf.field(ROW_ID_FIELD_NAME).to_pylist() except KeyError: pass @@ -287,9 +402,37 @@ def _iter_projection_pairs( yield name, source if source is not None else expr.to_sql() -def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table: - index = tbl.schema.get_field_index(output_name) - return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs]) +def _has_column_path(schema: pa.Schema, path: str) -> bool: + top_name, rest = _split_column_path(schema, path) + index = schema.get_field_index(top_name) + return ( + index >= 0 and _walk_struct_fields(schema.field(index).type, rest) is not None + ) + + +def _set_blob_column( + tbl: pa.Table, output_name: str, blobs: pa.Array | pa.ChunkedArray +) -> pa.Table: + top_name, rest = _split_column_path(tbl.schema, output_name) + if not rest: + index = tbl.schema.get_field_index(top_name) + return tbl.set_column(index, pa.field(top_name, blobs.type), [blobs]) + + leaf_name = rest[-1] + if isinstance(blobs, pa.ChunkedArray): + blobs = blobs.combine_chunks() + + def replace_blob(children: list, child_fields: list) -> None: + for i, field in enumerate(child_fields): + if field.name == leaf_name: + children[i] = blobs + child_fields[i] = pa.field( + leaf_name, blobs.type, nullable=field.nullable + ) + return + raise KeyError(f"Blob field '{leaf_name}' not found in '{output_name}'") + + return _transform_struct_column(tbl, top_name, rest[:-1], replace_blob) def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.Table: @@ -297,7 +440,9 @@ def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.T children.append(row_ids) child_fields.append(pa.field(ROW_ID_FIELD_NAME, pa.uint64(), nullable=False)) - return _transform_struct_column(tbl, path, add_row_id) + return _transform_struct_column( + tbl, *_split_column_path(tbl.schema, path), add_row_id + ) def strip_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table: @@ -314,16 +459,20 @@ def strip_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table: return for path in blob_paths: - if path.split(".")[0] not in tbl.column_names: + if not _has_column_path(tbl.schema, path): continue - tbl = _transform_struct_column(tbl, path, drop_row_id) + tbl = _transform_struct_column( + tbl, *_split_column_path(tbl.schema, path), drop_row_id + ) return tbl def _transform_struct_column( - tbl: pa.Table, path: str, leaf_transform: Callable[[list, list], None] + tbl: pa.Table, + top_name: str, + rest: list[str], + leaf_transform: Callable[[list, list], None], ) -> pa.Table: - top_name, *rest = path.split(".") top_index = tbl.schema.get_field_index(top_name) top_field = tbl.schema.field(top_index) top_array = tbl.column(top_name) @@ -381,26 +530,74 @@ def _rebuild_struct( def _leaf_struct_column(tbl: pa.Table, path: str) -> pa.StructArray: - parts = path.split(".") - column = tbl.column(parts[0]) + top_name, rest = _split_column_path(tbl.schema, path) + column = tbl.column(top_name) if isinstance(column, pa.ChunkedArray): column = column.combine_chunks() - for part in parts[1:]: + for part in rest: + if not pa.types.is_struct(column.type): + raise KeyError(f"blob path '{path}' crosses non-struct column") column = column.field(part) return column +def _split_column_path(schema: pa.Schema, path: str) -> tuple[str, list[str]]: + """Prefer a path ending in a struct, since a flat alias can shadow a nested + blob path.""" + parts = path.split(".") + fallback = None + for cut in range(len(parts), 0, -1): + top_name, rest = ".".join(parts[:cut]), parts[cut:] + index = schema.get_field_index(top_name) + if index < 0: + continue + leaf_type = _walk_struct_fields(schema.field(index).type, rest) + if leaf_type is None: + continue + if pa.types.is_struct(leaf_type): + return top_name, rest + fallback = fallback or (top_name, rest) + return fallback or (parts[0], parts[1:]) + + +def _walk_struct_fields( + data_type: pa.DataType, names: list[str] +) -> Optional[pa.DataType]: + for name in names: + if not pa.types.is_struct(data_type): + return None + index = data_type.get_field_index(name) + if index < 0: + return None + data_type = data_type.field(index).type + return data_type + + def _normalize_blob_row_ids( - row_ids: Union[list[int], pa.Table], blob_column: str + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], blob_column: str ) -> list[int]: if isinstance(row_ids, pa.Table): - return read_row_ids_from_hits(row_ids, blob_column) - if isinstance(row_ids, (pa.Array, pa.ChunkedArray)): - raise ValueError( - "pass a query table with _rowid, not a column array " - "(use fetch_blobs('image', hits), not fetch_blobs('image', hits['image']))" - ) - return list(row_ids) + values = read_row_ids_from_hits(row_ids, blob_column) + elif isinstance(row_ids, (pa.Array, pa.ChunkedArray)): + if not pa.types.is_integer(row_ids.type): + raise ValueError( + "row id arrays must have an integer type; pass a query table with " + "_rowid when using a blob descriptor column" + ) + values = row_ids.to_pylist() + else: + values = list(row_ids) + + normalized = [] + for value in values: + try: + row_id = operator.index(value) + except TypeError as exc: + raise ValueError("row ids must be integers without nulls") from exc + if row_id < 0: + raise ValueError("row ids must be non-negative") + normalized.append(row_id) + return normalized def _wrap_blob_files(handles: Iterable[object]) -> list[Optional[BlobFile]]: diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index a8f02e458..80554baf0 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -1,6 +1,6 @@ from datetime import date, datetime, timedelta from decimal import Decimal -from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal +from typing import Dict, List, Optional, Sequence, Tuple, Any, TypedDict, Union, Literal from uuid import UUID import pyarrow as pa @@ -311,6 +311,7 @@ class BlobFile: def tell(self) -> int: ... def size(self) -> int: ... def read_range(self, offset: int, length: int) -> bytes: ... + def read_ranges(self, ranges: list[tuple[int, int]]) -> list[bytes]: ... def read_up_to(self, length: int) -> bytes: ... class Job: @@ -563,7 +564,7 @@ class Table: async def fetch_blob_ranges( self, column: str, - requests: List[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: ... async def fetch_blob_files( self, column: str, row_ids: list[int] @@ -727,6 +728,7 @@ class VectorQuery: def use_lsm(self, enable: bool): ... def refine_factor(self, refine_factor: int): ... def nprobes(self, nprobes: int): ... + def search_effort(self, search_effort: Optional[float]): ... def minimum_nprobes(self, minimum_nprobes: int): ... def maximum_nprobes(self, maximum_nprobes: int): ... def bypass_vector_index(self): ... @@ -752,6 +754,7 @@ class HybridQuery: def distance_type(self, distance_type: str): ... def refine_factor(self, refine_factor: int): ... def nprobes(self, nprobes: int): ... + def search_effort(self, search_effort: Optional[float]): ... def minimum_nprobes(self, minimum_nprobes: int): ... def maximum_nprobes(self, maximum_nprobes: int): ... def bypass_vector_index(self): ... @@ -778,8 +781,10 @@ class PyQueryRequest: use_lsm: Optional[bool] column: Optional[str] query_vector: Optional[List[pa.Array]] + nprobes: Optional[int] minimum_nprobes: Optional[int] maximum_nprobes: Optional[int] + search_effort: Optional[float] lower_bound: Optional[float] upper_bound: Optional[float] ef: Optional[int] @@ -951,6 +956,8 @@ class PermutationReader: def __init__(self, base_table: Table, permutation_table: Table): ... class Catalog: + @property + def authz(self) -> AuthorizationClient: ... @property def uri(self) -> str: ... async def create_database( @@ -997,3 +1004,60 @@ class JobListing: def page_token(self) -> Optional[str]: ... def __aiter__(self) -> JobListing: ... async def __anext__(self) -> JobInfo: ... + +class AuthorizationClient: + async def list_acls(self, request: str) -> str: ... + async def add_acl(self, request: str) -> None: ... + async def delete_acl(self, request: str) -> None: ... + async def list_role_bindings(self, request: str) -> str: ... + async def add_role_binding(self, request: str) -> None: ... + async def delete_role_binding(self, request: str) -> None: ... + async def get_principal(self, request: str) -> str: ... + async def get_group(self, request: str) -> str: ... + +class AuthzSubject: + def __init__(self, kind: str, value: str): ... + @staticmethod + def from_canonical(value: str) -> AuthzSubject: ... + @property + def kind(self) -> str: ... + @property + def value(self) -> str: ... + @property + def display_value(self) -> str: ... + def to_wire(self, context: Optional[str] = None) -> str: ... + def __str__(self) -> str: ... + def __repr__(self) -> str: ... + +class AuthzPrivilege: + def __init__(self, value: str): ... + @staticmethod + def known() -> list[str]: ... + def __str__(self) -> str: ... + +class AuthzObject: + def __init__(self, value: str): ... + @staticmethod + def system() -> AuthzObject: ... + @staticmethod + def database(name: str) -> AuthzObject: ... + @staticmethod + def namespace(database: str, namespace_path: Sequence[str]) -> AuthzObject: ... + @staticmethod + def table( + database: str, namespace_path: Sequence[str], name: str + ) -> AuthzObject: ... + @staticmethod + def view( + database: str, namespace_path: Sequence[str], name: str + ) -> AuthzObject: ... + @staticmethod + def secret( + database: str, namespace_path: Sequence[str], name: str + ) -> AuthzObject: ... + @staticmethod + def function( + database: str, namespace_path: Sequence[str], name: str + ) -> AuthzObject: ... + def __str__(self) -> str: ... + def __repr__(self) -> str: ... diff --git a/python/python/lancedb/authz.py b/python/python/lancedb/authz.py new file mode 100644 index 000000000..e62a276cf --- /dev/null +++ b/python/python/lancedb/authz.py @@ -0,0 +1,588 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Deployment-wide authorization through ``catalog.authz``. + +ACLs grant one privilege on an object to a principal, group, or role. Role +bindings grant a role to a principal or group. Identity-provider group membership +is read-only. Name resolution and permission checks are performed by the server. +""" + +from collections.abc import AsyncIterator, Iterator, Sequence +from dataclasses import dataclass, field +from enum import Enum +import json +from typing import Literal, Optional, TypeVar, Union, cast + +from . import _lancedb +from .background_loop import LOOP + +__all__ = [ + "Authorization", + "AsyncAuthorization", + "Subject", + "Object", + "Privilege", + "AccessControlEntry", + "RoleBinding", + "IdentityInfo", + "PrincipalInfo", + "GroupInfo", +] + +_PAGE_SIZE = 100 +_T = TypeVar("_T") +_SubjectKind = Literal[ + "principal_id", + "principal_name", + "group_id", + "group_name", + "role", + "principal_api_key", + "unknown", +] + + +@dataclass(frozen=True) +class Subject: + """An immutable identity selector. Prefer the named factory methods. + + IDs are stable; names must resolve unambiguously on the server. API-key + selectors identify the subject of an operation, not its authenticated caller. + Secret values are omitted from representations. Unknown canonical subject + kinds returned by newer servers use ``kind="unknown"`` and preserve the full + wire string in ``value``; they can be passed back to ACL methods. + """ + + kind: _SubjectKind + _value: str = field(repr=False) + + _inner: _lancedb.AuthzSubject = field(init=False, repr=False, compare=False) + + def __post_init__(self): + object.__setattr__( + self, "_inner", _lancedb.AuthzSubject(self.kind, self._value) + ) + + @classmethod + def principal_id(cls, id: str) -> "Subject": + """Select a principal by canonical identity-provider ID.""" + return cls("principal_id", id) + + @classmethod + def principal_name(cls, name: str) -> "Subject": + """Select a principal by an unambiguous provider name.""" + return cls("principal_name", name) + + @classmethod + def group_id(cls, id: str) -> "Subject": + """Select a group by canonical identity-provider ID.""" + return cls("group_id", id) + + @classmethod + def group_name(cls, name: str) -> "Subject": + """Select a group by an unambiguous provider name.""" + return cls("group_name", name) + + @classmethod + def role(cls, name: str) -> "Subject": + """Select a LanceDB-managed role, without an ``r:`` prefix.""" + return cls("role", name) + + @classmethod + def principal_api_key(cls, api_key: str) -> "Subject": + """Select the principal holding this secret API key.""" + return cls("principal_api_key", api_key) + + def __repr__(self) -> str: + value = self._inner.display_value + return f"Subject.{self.kind}({value!r})" + + @property + def value(self) -> str: + """The selected value, or full wire string for unknown subject kinds. + + API-key values are secrets. + """ + return self._value + + def __reduce__(self): + return type(self), (self.kind, self._value) + + def _wire(self) -> str: + return self._inner.to_wire() + + @classmethod + def _canonical(cls, value: str) -> "Subject": + inner = _lancedb.AuthzSubject.from_canonical(value) + return cls(cast(_SubjectKind, inner.kind), inner.value) + + +@dataclass(frozen=True) +class Object: + """An immutable resource name in the authorization syntax. + + Use the builders to avoid interpreting separators as part of a name. + Methods also accept raw object strings, validated by the same Rust parser. + Unknown resource types are preserved for compatibility with newer servers. + Omitted or empty namespace paths select the database's public namespace. + """ + + _value: str + _inner: _lancedb.AuthzObject = field(init=False, repr=False, compare=False) + + def __post_init__(self): + object.__setattr__(self, "_inner", _lancedb.AuthzObject(self._value)) + + @classmethod + def system(cls) -> "Object": + """The deployment-wide system object.""" + return cls(str(_lancedb.AuthzObject.system())) + + @classmethod + def database(cls, name: str) -> "Object": + """A database; slash-containing names remain one logical database.""" + return cls(str(_lancedb.AuthzObject.database(name))) + + @classmethod + def namespace( + cls, *, database: str, namespace_path: Sequence[str] = ("public",) + ) -> "Object": + """A namespace path, defaulting to the database's public namespace.""" + return cls(str(_lancedb.AuthzObject.namespace(database, namespace_path))) + + @classmethod + def table( + cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",) + ) -> "Object": + """A table, defaulting to the database's public namespace.""" + return cls(str(_lancedb.AuthzObject.table(database, namespace_path, name))) + + @classmethod + def view( + cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",) + ) -> "Object": + """A view, defaulting to the database's public namespace.""" + return cls(str(_lancedb.AuthzObject.view(database, namespace_path, name))) + + @classmethod + def secret( + cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",) + ) -> "Object": + """A secret, defaulting to the database's public namespace.""" + return cls(str(_lancedb.AuthzObject.secret(database, namespace_path, name))) + + @classmethod + def function( + cls, *, database: str, name: str, namespace_path: Sequence[str] = ("public",) + ) -> "Object": + """A function, defaulting to the database's public namespace.""" + return cls(str(_lancedb.AuthzObject.function(database, namespace_path, name))) + + def __str__(self) -> str: + return str(self._inner) + + def __reduce__(self): + return type(self), (self._value,) + + +class Privilege(str, Enum): + """Known privileges, validated and serialized by Rust. + + ACL methods also accept strings such as ``"SELECT"`` and future privilege + names such as ``"PRIVILEGE20"``. Future names are preserved verbatim; + their meaning and validity are determined by the server. + + Grant ``CREATE_SECRET`` or ``CREATE_FUNCTION`` on a namespace to delegate + creation of those resources without transferring namespace ownership. + """ + + def __new__(cls, value: str): + canonical = str(_lancedb.AuthzPrivilege(value)) + member = str.__new__(cls, canonical) + member._value_ = canonical + return member + + OWNERSHIP = "OWNERSHIP" + USAGE = "USAGE" + SELECT = "SELECT" + INSERT = "INSERT" + UPDATE = "UPDATE" + DELETE = "DELETE" + CREATE_DATABASE = "CREATE_DATABASE" + CREATE_NAMESPACE = "CREATE_NAMESPACE" + CREATE_TABLE = "CREATE_TABLE" + CREATE_MATERIALIZED_VIEW = "CREATE_MATERIALIZED_VIEW" + CREATE_VIEW = "CREATE_VIEW" + CREATE_SECRET = "CREATE_SECRET" + CREATE_FUNCTION = "CREATE_FUNCTION" + OPERATE = "OPERATE" + + +@dataclass(frozen=True) +class AccessControlEntry: + """One object/subject/privilege grant; subject contains a canonical ID or role.""" + + object: str + subject: Subject + subject_name: Optional[str] + privilege: str + + +@dataclass(frozen=True) +class RoleBinding: + """A direct principal/group role binding, not an expanded effective role.""" + + subject: Subject + subject_name: Optional[str] + role: str + + +@dataclass(frozen=True) +class IdentityInfo: + """A canonical identity-provider ID and its optional display name.""" + + id: str + name: Optional[str] = None + + +@dataclass(frozen=True) +class PrincipalInfo: + """A principal's groups and effective roles, including group-inherited roles.""" + + principal_id: str + principal_name: Optional[str] + groups: list[IdentityInfo] + roles: list[str] + + +@dataclass(frozen=True) +class GroupInfo: + """A group's provider-managed members and LanceDB-managed direct roles.""" + + group_id: str + group_name: Optional[str] + principals: list[IdentityInfo] + roles: list[str] + + +def _subject(subject: Subject, context: Optional[str] = None) -> str: + if not isinstance(subject, Subject): + raise ValueError("subject must be a Subject selector") + return subject._inner.to_wire(context) + + +def _object(value: Union[str, Object]) -> str: + if not isinstance(value, (str, Object)) or not str(value): + raise ValueError("object must be a nonempty string or Object") + return str(value) if isinstance(value, Object) else str(_lancedb.AuthzObject(value)) + + +def _limit(limit: Optional[int]) -> None: + if limit is not None and ( + isinstance(limit, bool) or not isinstance(limit, int) or limit < 0 + ): + raise ValueError("limit must be a nonnegative integer or None") + + +def _privilege(value: Union[str, Privilege]) -> str: + if not isinstance(value, str): + raise ValueError("privilege must be a string or Privilege") + return str(_lancedb.AuthzPrivilege(value)) + + +def _entry(value: dict) -> AccessControlEntry: + return AccessControlEntry( + value["object"], + Subject._canonical(value["subject_id"]), + value.get("subject_name"), + value["privilege"], + ) + + +def _binding(value: dict) -> RoleBinding: + return RoleBinding( + Subject._canonical(value["subject_id"]), + value.get("subject_name"), + value["role"], + ) + + +class AsyncAuthorization: + """Async authorization returned by ``AsyncCatalog.authz``. + + HTTP failures raise [HttpError][lancedb.remote.errors.HttpError], preserving + status and request ID. Invalid local arguments raise ValueError. Lists are + lazy async iterators: failures can occur during iteration after some results + have been yielded. Iteration does not provide a consistent snapshot. + """ + + def __init__(self, inner: _lancedb.AuthorizationClient): + self._inner = inner + + async def _pages(self, method, request: dict, key: str, limit: Optional[int]): + yielded = 0 + while limit is None or yielded < limit: + request["max_results"] = ( + _PAGE_SIZE if limit is None else min(_PAGE_SIZE, limit - yielded) + ) + response = json.loads(await method(json.dumps(request))) + for value in response[key]: + yield value + yielded += 1 + if limit is not None and yielded >= limit: + return + token = response.get("continuation_token") + if not token: + return + if token == request.get("continuation_token"): + raise RuntimeError( + "Authorization listing returned a repeated continuation token" + ) + request["continuation_token"] = token + + def list_acls( + self, + *, + object: Optional[Union[str, Object]] = None, + subject: Optional[Subject] = None, + limit: Optional[int] = None, + ) -> AsyncIterator[AccessControlEntry]: + """Iterate grants, optionally filtered by exact object and/or subject. + + Requires USAGE on the supplied object, or on system if object is omitted. + ``limit`` bounds yielded entries across all pages; None lists all and zero + makes no requests. Missing matches produce an empty iterator. + """ + _limit(limit) + request = {} + if object is not None: + request["object"] = _object(object) + if subject is not None: + request["subject"] = _subject(subject) + + async def entries(): + async for value in self._pages( + self._inner.list_acls, request, "entries", limit + ): + yield _entry(value) + + return entries() + + async def add_acl( + self, + *, + object: Union[str, Object], + subject: Subject, + privilege: Union[str, Privilege], + ) -> None: + """Grant one privilege; requires OWNERSHIP on the object. + + A duplicate grant raises HTTP 409. Granting OWNERSHIP transfers ownership. + Writes are not automatically retried after ambiguous failures. + """ + await self._inner.add_acl( + json.dumps( + { + "object": _object(object), + "subject": _subject(subject), + "privilege": _privilege(privilege), + } + ) + ) + + async def delete_acl( + self, + *, + object: Union[str, Object], + subject: Subject, + privilege: Union[str, Privilege], + ) -> None: + """Revoke one privilege; requires OWNERSHIP on the object. + + A missing grant raises HTTP 404. OWNERSHIP itself cannot be deleted. + """ + await self._inner.delete_acl( + json.dumps( + { + "object": _object(object), + "subject": _subject(subject), + "privilege": _privilege(privilege), + } + ) + ) + + def list_role_bindings( + self, + *, + subject: Optional[Subject] = None, + role: Optional[str] = None, + limit: Optional[int] = None, + ) -> AsyncIterator[RoleBinding]: + """Iterate direct bindings, optionally filtered by principal/group and role. + + Requires USAGE on system. Limit semantics match ``list_acls``. Principals' + roles inherited through groups are available through ``get_principal``. + """ + _limit(limit) + request = {} + if subject is not None: + request["subject"] = _subject(subject, "role_member") + if role is not None: + request["role"] = role + + async def bindings(): + async for value in self._pages( + self._inner.list_role_bindings, request, "bindings", limit + ): + yield _binding(value) + + return bindings() + + async def add_role_binding(self, *, subject: Subject, role: str) -> None: + """Grant a role idempotently; requires OPERATE on system.""" + await self._inner.add_role_binding( + json.dumps({"subject": _subject(subject, "role_member"), "role": role}) + ) + + async def delete_role_binding(self, *, subject: Subject, role: str) -> None: + """Idempotently revoke a role; requires OPERATE on system.""" + await self._inner.delete_role_binding( + json.dumps({"subject": _subject(subject, "role_member"), "role": role}) + ) + + async def get_principal(self, subject: Subject) -> PrincipalInfo: + """Look up groups and effective roles; requires USAGE on system. + + Select by ID, name, or API key. Missing/ambiguous identities raise HTTP + 404/409; unavailable or unconfigured resolution raises HTTP 503/501. + """ + value = json.loads( + await self._inner.get_principal( + json.dumps({"principal": _subject(subject, "principal")}) + ) + ) + return PrincipalInfo( + value["principal_id"], + value.get("principal_name"), + [IdentityInfo(**group) for group in value["groups"]], + value["roles"], + ) + + async def get_group(self, subject: Subject) -> GroupInfo: + """Look up a group's members and roles by ID/name; requires USAGE on system. + + Membership is read-only and comes from the identity provider. Lookup + errors have the same HTTP semantics as ``get_principal``. + """ + value = json.loads( + await self._inner.get_group( + json.dumps({"group": _subject(subject, "group")}) + ) + ) + return GroupInfo( + value["group_id"], + value.get("group_name"), + [IdentityInfo(**principal) for principal in value["principals"]], + value["roles"], + ) + + +async def _next(iterator: AsyncIterator[_T]) -> _T: + return await iterator.__anext__() + + +def _sync_iter(iterator: AsyncIterator[_T]) -> Iterator[_T]: + try: + while True: + yield LOOP.run(_next(iterator)) + except StopAsyncIteration: + return + finally: + close = getattr(iterator, "aclose", None) + if close is not None: + LOOP.run(close()) + + +class Authorization: + """Synchronous authorization returned by ``Catalog.authz``. + + Shares the async API's permissions, errors, and lazy pagination behavior. + Group membership comes from the identity provider and cannot be edited here. + """ + + def __init__(self, inner: AsyncAuthorization): + self._inner = inner + + def list_acls( + self, + *, + object: Optional[Union[str, Object]] = None, + subject: Optional[Subject] = None, + limit: Optional[int] = None, + ) -> Iterator[AccessControlEntry]: + """Iterate grants. + + See [list_acls][lancedb.authz.AsyncAuthorization.list_acls] for permissions. + """ + return _sync_iter( + self._inner.list_acls(object=object, subject=subject, limit=limit) + ) + + def add_acl( + self, + *, + object: Union[str, Object], + subject: Subject, + privilege: Union[str, Privilege], + ) -> None: + """Grant a privilege; duplicates raise HTTP 409. + + OWNERSHIP transfers ownership. + """ + LOOP.run( + self._inner.add_acl(object=object, subject=subject, privilege=privilege) + ) + + def delete_acl( + self, + *, + object: Union[str, Object], + subject: Subject, + privilege: Union[str, Privilege], + ) -> None: + """Revoke a privilege. Missing grants raise HTTP 404. + + OWNERSHIP cannot be deleted. + """ + LOOP.run( + self._inner.delete_acl(object=object, subject=subject, privilege=privilege) + ) + + def list_role_bindings( + self, + *, + subject: Optional[Subject] = None, + role: Optional[str] = None, + limit: Optional[int] = None, + ) -> Iterator[RoleBinding]: + """Iterate direct role bindings, using the same limit semantics as list_acls.""" + return _sync_iter( + self._inner.list_role_bindings(subject=subject, role=role, limit=limit) + ) + + def add_role_binding(self, *, subject: Subject, role: str) -> None: + """Grant a role idempotently; requires OPERATE on system.""" + LOOP.run(self._inner.add_role_binding(subject=subject, role=role)) + + def delete_role_binding(self, *, subject: Subject, role: str) -> None: + """Idempotently revoke a role; requires OPERATE on system.""" + LOOP.run(self._inner.delete_role_binding(subject=subject, role=role)) + + def get_principal(self, subject: Subject) -> PrincipalInfo: + """Look up groups and effective roles; requires USAGE on system.""" + return LOOP.run(self._inner.get_principal(subject)) + + def get_group(self, subject: Subject) -> GroupInfo: + """Look up a group's members and roles; requires USAGE on system.""" + return LOOP.run(self._inner.get_group(subject)) diff --git a/python/python/lancedb/catalog.py b/python/python/lancedb/catalog.py index 70aab1bbb..666d414a4 100644 --- a/python/python/lancedb/catalog.py +++ b/python/python/lancedb/catalog.py @@ -3,11 +3,13 @@ """Remote catalogs manage databases through a server's root namespace.""" +from functools import cached_property from datetime import timedelta from typing import Any, AsyncIterator, Iterator, Optional, Union from . import _lancedb from .background_loop import LOOP +from .authz import AsyncAuthorization, Authorization from .db import AsyncConnection, DBConnection from .remote import ClientConfig, OAuthConfig from .remote.db import RemoteDBConnection @@ -86,6 +88,11 @@ class AsyncCatalog: def __init__(self, inner: _lancedb.Catalog): self._inner = inner + @cached_property + def authz(self) -> AsyncAuthorization: + """Authorization sharing this catalog's scope and credentials.""" + return AsyncAuthorization(self._inner.authz) + @property def uri(self) -> str: """The catalog's root namespace endpoint.""" @@ -157,6 +164,11 @@ class Catalog: self._sql_host_override = sql_host_override self._oauth_config = oauth_config + @cached_property + def authz(self) -> Authorization: + """Authorization sharing this catalog's scope and credentials.""" + return Authorization(self._inner.authz) + @property def uri(self) -> str: """The catalog's root namespace endpoint.""" diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index c170c7375..445d31027 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -372,7 +372,7 @@ class FunctionVersion(_RemoteValue): image: FunctionImage signature: FunctionSignature secret_bindings: tuple[SecretBinding, ...] = () - created_at: str + created_at_millis: int metadata: Mapping[str, str] disabled: bool diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index 2a7731be3..cf3bf8e4f 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -112,9 +112,42 @@ class Fm: class FTS: """Describe a FTS index configuration. - `FTS` is a full-text search index that can be used on `String` columns + `FTS` supports string columns, lists of strings, and JSON columns created + from PyArrow ``pa.json_()``. JSON columns use Lance's JSON-aware tokenizer + and are stored as ``LargeBinary`` with the ``lance.json`` logical type. + Arbitrary binary columns are not supported. - For example, it works with `title`, `description`, `content`, etc. + JSON queries use ``path,str,value`` triples. Paths are relative to the JSON + column and use dots for nested object keys. For example, + ``meta.tag,str,nature`` matches the string value at ``$.meta.tag``; + it does not search other paths or JSON key names. The value is processed + with the configured text tokenizer. Use ``str``, not ``string``. + + Create and query a JSON index through the modern API:: + + import pyarrow as pa + from lancedb.index import FTS + + table = db.create_table( + "documents", + pa.table({"doc": pa.array( + ['{"meta":{"tag":"nature"}}'], type=pa.json_() + )}), + ) + table.create_index("doc", config=FTS()) + table.search( + "meta.tag,str,nature", query_type="fts", fts_columns="doc" + ).to_list() + + Create the index on the JSON column itself, not a JSON subpath such as + ``doc.meta.tag``. JSON-path scalar indexes do not provide FTS path scoping. + Plain text such as ``nature`` is not a valid JSON FTS query. The current + Lance tokenizer can panic on malformed triples; see + https://github.com/lance-format/lance/issues/8812. + + Remote tables use the same ``FTS()`` configuration and require a server + that supports JSON FTS. Wait for ``table.wait_for_index(["doc_idx"])`` + before searching a newly created remote index with its default name. Examples -------- diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 4edc33843..bd3435292 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -105,6 +105,11 @@ def _query_to_namespace_request( if not isinstance(query, Query): raise TypeError(f"Expected Query, got {type(query)}") + if query.search_effort is not None: + raise ValueError( + "namespace QueryTable does not support search_effort; " + "execute the query locally" + ) vector_request = None if query.vector is not None: @@ -190,8 +195,12 @@ def _query_to_namespace_request( kwargs["with_row_id"] = query.with_row_id if query.ef is not None: kwargs["ef"] = query.ef + if query.nprobes is not None: + kwargs["nprobes"] = query.nprobes if query.minimum_nprobes is not None: - kwargs["nprobes"] = query.minimum_nprobes + kwargs["minimum_nprobes"] = query.minimum_nprobes + if query.maximum_nprobes is not None: + kwargs["maximum_nprobes"] = query.maximum_nprobes if query.refine_factor is not None: kwargs["refine_factor"] = query.refine_factor if query.lower_bound is not None: @@ -205,6 +214,13 @@ def _query_to_namespace_request( if query.distance_type is not None: kwargs["distance_type"] = query.distance_type + if kwargs.get("maximum_nprobes") == 0: + # The generated 0.13 model rejects LanceDB's unbounded sentinel; + # model_copy preserves direct wire pass-through to the Lance receiver. + maximum_nprobes = kwargs.pop("maximum_nprobes") + return QueryTableRequest(**kwargs).model_copy( + update={"maximum_nprobes": maximum_nprobes} + ) return QueryTableRequest(**kwargs) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 6f68692a2..15a52c5e3 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -12,6 +12,7 @@ from typing import ( TYPE_CHECKING, Any, Dict, + Iterable, List, Literal, Optional, @@ -40,7 +41,7 @@ from .rerankers.base import Reranker from .rerankers.rrf import RRFReranker from .rerankers.util import check_reranker_result from .schema import blob_column_paths, is_blob_like_field, schema_has_blob_field -from .util import _validate_query_vector, flatten_columns +from .util import _validate_query_vector, flatten_columns, get_uri_scheme from . import _wal_hybrid # WAL-PK-FUSION: delete. from ._blob import ( BLOB_MODE_TO_HANDLING, @@ -49,6 +50,8 @@ from ._blob import ( blob_auto_row_id_for_scan, blob_v2_projection_sources, finalize_blob_query_table, + reject_unfetchable_blob_columns, + reject_unmatched_blob_outputs, replace_v2_blob_columns_with_bytes, replace_v2_blob_columns_with_bytes_sync, strip_auto_row_ids, @@ -99,15 +102,45 @@ class _LanceScanner(Protocol): def to_reader(self): ... -def _blob_mode_requires_native_pandas(blob_mode: BlobMode, schema: pa.Schema) -> bool: - return blob_mode in BLOB_MODE_TO_HANDLING and schema_has_blob_field(schema) - - def _unsupported_blob_pandas_error(reason: str) -> RuntimeError: return RuntimeError( "blob columns require Lance native scanner conversion for query " - f"to_pandas(), but {reason}. Use a plain scan query or remove blob " - "columns from the projection." + f"to_pandas(), but {reason}. Use a plain scan query, pass " + "blob_mode='bytes' or blob_mode='descriptions', or remove blob columns " + "from the projection." + ) + + +def _supports_local_blob_fetch( + async_table: "AsyncTable", pushdown_operations: Iterable[str] +) -> bool: + # Remote tables and QueryTable pushdown answer queries on a server, so a + # fetch through this handle could read another version. + return async_table._inner._is_native() and "QueryTable" not in pushdown_operations + + +def _blob_bytes_without_snapshot_error() -> NotImplementedError: + return NotImplementedError( + "blob_mode='bytes' cannot safely materialize blob columns without a " + "stable table snapshot. Remote tables and namespace QueryTable pushdown " + "run the query on a server, so the blob fetch could read another table " + "version. Use blob_mode='descriptions'." + ) + + +def _query_needs_lance_scanner(builder: Any) -> bool: + """Whether the query sets options that to_arrow ignores.""" + to_query_object = getattr(builder, "to_query_object", None) + if to_query_object is None: + return False + try: + query = to_query_object() + except NotImplementedError: + return False + return ( + query.fragments is not None + or query.fragment_ids is not None + or bool(query.with_row_address) ) @@ -151,7 +184,8 @@ def _combine_where( def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]: - if columns is None: + # An empty projection selects all columns. + if not columns: return {} if isinstance(columns, list): if all(isinstance(column, str) for column in columns): @@ -175,7 +209,8 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]: def _query_request_projection(req: "PyQueryRequest") -> QueryProjection: - if req.select_source_columns is not None: + # Treat empty map as a calculated select + if req.select_source_columns: return req.select_source_columns return req.select @@ -475,8 +510,8 @@ class MatchQuery(FullTextQuery): ---------- query : str The query string to match against. - column : str - The name of the column to match against. + column : str, optional + The column to match against. None leaves column selection to query planning. boost : float, default 1.0 The boost factor for the query. The score of each matching document is multiplied by this value. @@ -505,7 +540,7 @@ class MatchQuery(FullTextQuery): """ query: str - column: str + column: Optional[str] = None boost: float = pydantic.Field(1.0, kw_only=True) fuzziness: Optional[int] = pydantic.Field(0, kw_only=True) max_expansions: int = pydantic.Field(50, kw_only=True) @@ -693,8 +728,9 @@ class Query(pydantic.BaseModel): if True then apply the filter after vector / FTS search. This is ignored for plain SQL filtering. nprobes : Optional[int] - The number of IVF partitions to search. If this is None then a default - number of partitions will be used. + The legacy number of IVF partitions to search. Lance sets both probe + bounds to this value. If this is None then Lance's default probe settings + will be used. - A higher number makes search more accurate but also slower. @@ -773,22 +809,28 @@ class Query(pydantic.BaseModel): # distance type to use for vector search distance_type: Optional[str] = None + # legacy number of IVF partitions to search + # + # Lance sets both probe bounds to this value. Explicit bounds can override it. + nprobes: Optional[int] = None + # which columns to return in the results (dict values may be str or Expr) columns: QueryProjection = None # minimum number of IVF partitions to search # - # If None then a default value (20) will be used. + # If None then Lance's default will be used. minimum_nprobes: Optional[int] = None # maximum number of IVF partitions to search # - # If None then a default value (20) will be used. - # - # If 0 then no limit will be applied and all partitions could be searched - # if needed to satisfy the limit. + # If None then Lance's default will be used. If 0 then no limit will be applied + # and all partitions could be searched if needed to satisfy the limit. maximum_nprobes: Optional[int] = None + # Optional initial IVF effort; None preserves Auto and explicit nprobes. + search_effort: Optional[float] = None + # lower bound for distance search lower_bound: Optional[float] = None @@ -847,8 +889,10 @@ class Query(pydantic.BaseModel): query.vector_column = req.column query.vector = req.query_vector query.distance_type = req.distance_type + query.nprobes = req.nprobes query.minimum_nprobes = req.minimum_nprobes query.maximum_nprobes = req.maximum_nprobes + query.search_effort = req.search_effort query.lower_bound = req.lower_bound query.upper_bound = req.upper_bound query.ef = req.ef @@ -1063,44 +1107,42 @@ class LanceQueryBuilder(ABC): The maximum time to wait for the query to complete. If None, wait indefinitely. blob_mode: str, default "lazy" - Controls how blob columns are returned for plain scan queries. - Vector, FTS, hybrid, and other non-native query shapes keep the - existing Arrow conversion path and only support blob descriptions. + Controls how blob columns are returned. ``"lazy"`` returns + :class:`~lancedb._blob.BlobFile` handles and needs a plain scan. + ``"descriptions"`` and ``"bytes"`` work for every query shape. A + plain scan without ``timeout`` reads the local dataset through + Lance's scanner, which also returns legacy blob columns as bytes. + Other shapes use :meth:`to_arrow` with the same ``blob_mode``. **kwargs Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. """ validate_blob_mode(blob_mode) output_schema = getattr(self, "output_schema", None) - if output_schema is not None: - schema = output_schema() - if _blob_mode_requires_native_pandas(blob_mode, schema): - native_error = None - if (flatten is None or blob_mode == "descriptions") and timeout is None: - try: - df = self._plain_scan_to_pandas( - blob_mode, flatten=flatten, **kwargs - ) - if df is not None: - return df - except Exception as err: - native_error = err - reason = ( + if output_schema is not None and schema_has_blob_field(output_schema()): + if (flatten is None or blob_mode != "lazy") and timeout is None: + df = self._plain_scan_to_pandas(blob_mode, flatten=flatten, **kwargs) + if df is not None: + return df + if blob_mode == "lazy" or _query_needs_lance_scanner(self): + raise _unsupported_blob_pandas_error( "this query shape cannot use Lance native pandas conversion" - if native_error is None - else str(native_error) ) - raise _unsupported_blob_pandas_error(reason) from native_error - tbl = flatten_columns(self.to_arrow(timeout=timeout), flatten) - if _blob_mode_requires_native_pandas(blob_mode, tbl.schema): + tbl = flatten_columns(self._to_arrow_in_mode(timeout, blob_mode), flatten) + if blob_mode == "lazy" and schema_has_blob_field(tbl.schema): raise _unsupported_blob_pandas_error( "this query shape cannot use Lance native pandas conversion" ) return tbl.to_pandas(**kwargs) @abstractmethod - def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table: + def to_arrow( + self, + *, + timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", + ) -> pa.Table: """ Execute the query and return the results as an [Apache Arrow Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table). @@ -1114,6 +1156,11 @@ class LanceQueryBuilder(ABC): timeout: Optional[timedelta] The maximum time to wait for the query to complete. If None, wait indefinitely. + blob_mode: str, default "lazy" + ``"lazy"`` returns blob v2 descriptors that carry the row id + ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. + ``"bytes"`` returns payload bytes and raises for legacy blob columns + and blobs inside lists. """ raise NotImplementedError @@ -1144,7 +1191,12 @@ class LanceQueryBuilder(ABC): """ raise NotImplementedError - def to_list(self, *, timeout: Optional[timedelta] = None) -> List[dict]: + def to_list( + self, + *, + timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", + ) -> List[dict]: """ Execute the query and return the results as a list of dictionaries. @@ -1157,8 +1209,19 @@ class LanceQueryBuilder(ABC): timeout: Optional[timedelta] The maximum time to wait for the query to complete. If None, wait indefinitely. + blob_mode: str, default "lazy" + Passed through to :meth:`to_arrow`. """ - return self.to_arrow(timeout=timeout).to_pylist() + validate_blob_mode(blob_mode) + return self._to_arrow_in_mode(timeout, blob_mode).to_pylist() + + def _to_arrow_in_mode( + self, timeout: Optional[timedelta], blob_mode: BlobMode + ) -> pa.Table: + # Older subclasses may not accept blob_mode. Omit it when it has no effect. + if blob_mode == "lazy" or not schema_has_blob_field(self._table.schema): + return self.to_arrow(timeout=timeout) + return self.to_arrow(timeout=timeout, blob_mode=blob_mode) def to_pydantic( self, model: type[T], *, timeout: Optional[timedelta] = None @@ -1331,19 +1394,62 @@ class LanceQueryBuilder(ABC): query.with_row_id = True return query - def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table: - blob_auto_row_id = self._blob_auto_row_id_enabled() - blob_paths = ( - blob_v2_projection_sources(self._table.schema, self._columns).keys() - if blob_auto_row_id - else () + def _validated_blob_sources(self, blob_mode: BlobMode) -> dict[str, str]: + """Map projected blob v2 outputs to their source columns. + + Bytes mode is checked here, before the query runs. + """ + validate_blob_mode(blob_mode) + schema = self._table.schema + blob_sources = blob_v2_projection_sources(schema, self._columns) + if blob_mode != "bytes": + return blob_sources + reject_unfetchable_blob_columns(schema, self._columns, blob_sources) + if blob_sources and not _supports_local_blob_fetch( + self._table._table, getattr(self._table, "_pushdown_operations", ()) + ): + raise _blob_bytes_without_snapshot_error() + reject_unmatched_blob_outputs(self._blob_output_schema(), blob_sources) + return blob_sources + + def _blob_output_schema(self) -> pa.Schema: + return self.output_schema() + + def _apply_arrow_blob_mode( + self, tbl: pa.Table, blob_mode: BlobMode, blob_sources: dict[str, str] + ) -> pa.Table: + if blob_mode == "lazy": + return tbl + if blob_mode == "descriptions": + return strip_auto_row_ids(tbl, blob_sources) + return replace_v2_blob_columns_with_bytes_sync( + tbl, blob_sources, self._table.fetch_blobs ) - return finalize_blob_query_table( + + def _finalize_blob_query_table( + self, + tbl: pa.Table, + *, + blob_mode: BlobMode, + blob_sources: dict[str, str], + ) -> pa.Table: + blob_auto_row_id = self._blob_auto_row_id_enabled() + finalized = finalize_blob_query_table( tbl, user_requested_row_id=self._user_requested_row_id(), blob_auto_row_id=blob_auto_row_id, - blob_paths=blob_paths, + blob_paths=blob_sources.keys() if blob_auto_row_id else (), ) + finalized = self._apply_arrow_blob_mode(finalized, blob_mode, blob_sources) + # with_row_id(False) does not reach the scan, which adds _rowid anyway. + # Lazy keeps it because fetch_blobs(hits) reads it there. + if ( + blob_mode != "lazy" + and self._with_row_id is False + and "_rowid" in finalized.column_names + ): + finalized = finalized.drop_columns(["_rowid"]) + return finalized def with_row_address(self, with_row_address: bool = True) -> Self: """Set whether to return row addresses. @@ -1556,6 +1662,11 @@ class LanceQueryBuilder(ABC): if not _query_is_plain_scan(query): return None + # Remote and in-memory tables have no Lance dataset to scan. + if not self._table._table._inner._is_native() or ( + get_uri_scheme(self._table._dataset_path) == "memory" + ): + return None dataset = self._table.to_lance() blob_auto_row_id = self._blob_auto_row_id_enabled() blob_sources = ( @@ -1563,6 +1674,10 @@ class LanceQueryBuilder(ABC): if blob_mode == "bytes" else {} ) + if blob_sources: + reject_unfetchable_blob_columns( + self._table.schema, query.columns, blob_sources + ) scanner = dataset.scanner( **_scanner_kwargs_for_query( query, @@ -1576,7 +1691,7 @@ class LanceQueryBuilder(ABC): blob_mode=blob_mode, blob_sources=blob_sources, fetch_blobs=self._table.fetch_blobs, - strip_auto_row_id=blob_auto_row_id, + strip_auto_row_id=not self._user_requested_row_id(), flatten=flatten, **kwargs, ) @@ -1627,8 +1742,10 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): super().__init__(table) self._query = query self._distance_type = None + self._nprobes = None self._minimum_nprobes = None self._maximum_nprobes = None + self._search_effort = None self._lower_bound = None self._upper_bound = None self._refine_factor = None @@ -1691,9 +1808,9 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): See discussion in [Querying an ANN Index](https://lancedb.com/docs/indexing/) for tuning advice. - This method sets both the minimum and maximum number of probes to the same - value. See `minimum_nprobes` and `maximum_nprobes` for more fine-grained - control. + The value is retained as `nprobes` through client and server request + construction. Lance sets both probe bounds to this value. Explicit minimum + or maximum settings can override their respective bounds. Parameters ---------- @@ -1707,8 +1824,28 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): """ if nprobes <= 0: raise ValueError("Invalid input, nprobes must be greater than 0") - self._minimum_nprobes = nprobes - self._maximum_nprobes = nprobes + self._nprobes = nprobes + return self + + def search_effort(self, search_effort: Optional[float]) -> LanceVectorQueryBuilder: + """Set the initial IVF search effort to a finite value in [0, 1]. + + None (the default) uses Auto, internally 0.5. Zero starts at the caller + minimum; one searches all partitions within the caller maximum. Later + probing may continue when too few candidates are found. Any explicit + effort, including 0.5, overrides nprobes in either call order. Explicit + minimum and maximum bounds still apply. This is not a recall guarantee + and does not change refinement or HNSW ef. MemWAL LSM queries do not + yet support this option. + + Examples + -------- + >>> import lancedb + >>> db = lancedb.connect("memory://") + >>> table = db.create_table("effort_vectors", [{"vector": [0.1, 0.2]}]) + >>> query = table.search([0.1, 0.2]).search_effort(0.75) + """ + self._search_effort = search_effort return self def minimum_nprobes(self, minimum_nprobes: int) -> LanceVectorQueryBuilder: @@ -1815,7 +1952,12 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): """ return self._table._output_schema(self.to_query_object()) - def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table: + def to_arrow( + self, + *, + timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", + ) -> pa.Table: """ Execute the query and return the results as an [Apache Arrow Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table). @@ -1829,9 +1971,17 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): timeout: Optional[timedelta] The maximum time to wait for the query to complete. If None, wait indefinitely. + blob_mode: str, default "lazy" + ``"lazy"`` returns blob v2 descriptors that carry the row id + ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. + ``"bytes"`` returns payload bytes and raises for legacy blob columns + and blobs inside lists. """ + blob_sources = self._validated_blob_sources(blob_mode) return self._finalize_blob_query_table( - self.to_batches(timeout=timeout).read_all() + self.to_batches(timeout=timeout).read_all(), + blob_mode=blob_mode, + blob_sources=blob_sources, ) def to_query_object(self) -> Query: @@ -1850,8 +2000,10 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): limit=self._limit, distance_type=self._distance_type, columns=self._columns, + nprobes=self._nprobes, minimum_nprobes=self._minimum_nprobes, maximum_nprobes=self._maximum_nprobes, + search_effort=self._search_effort, lower_bound=self._lower_bound, upper_bound=self._upper_bound, refine_factor=self._refine_factor, @@ -2102,16 +2254,23 @@ class LanceFtsQueryBuilder(LanceQueryBuilder): """ return self._table._output_schema(self.to_query_object()) - def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table: + def to_arrow( + self, + *, + timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", + ) -> pa.Table: + blob_sources = self._validated_blob_sources(blob_mode) self._table._ensure_no_legacy_fts_index() query = self._query_for_scan() - results = self._table._execute_query(query, timeout=timeout) - results = results.read_all() + results = self._table._execute_query(query, timeout=timeout).read_all() if self._reranker is not None: results = self._reranker.rerank_fts(self._query, results) check_reranker_result(results) - return self._finalize_blob_query_table(results) + return self._finalize_blob_query_table( + results, blob_mode=blob_mode, blob_sources=blob_sources + ) def to_batches( self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None @@ -2138,9 +2297,17 @@ class LanceFtsQueryBuilder(LanceQueryBuilder): class LanceEmptyQueryBuilder(LanceQueryBuilder): - def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table: + def to_arrow( + self, + *, + timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", + ) -> pa.Table: + blob_sources = self._validated_blob_sources(blob_mode) return self._finalize_blob_query_table( - self.to_batches(timeout=timeout).read_all() + self.to_batches(timeout=timeout).read_all(), + blob_mode=blob_mode, + blob_sources=blob_sources, ) def to_query_object(self) -> Query: @@ -2208,8 +2375,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): self._fts_columns = fts_columns self._norm = None self._reranker = None + self._nprobes = None self._minimum_nprobes = None self._maximum_nprobes = None + self._search_effort = None self._refine_factor = None self._distance_type = None self._phrase_query = None @@ -2271,16 +2440,30 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): def to_query_object(self) -> Query: raise NotImplementedError("to_query_object not yet supported on a hybrid query") - def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table: + def to_arrow( + self, + *, + timeout: Optional[timedelta] = None, + blob_mode: BlobMode = "lazy", + ) -> pa.Table: + blob_sources = self._validated_blob_sources(blob_mode) # WAL-PK-FUSION: without the fallback, this body is `_run_hybrid`'s. return _wal_hybrid.with_pk_fallback( - lambda: self._run_hybrid(timeout=timeout), + lambda: self._run_hybrid( + timeout=timeout, blob_mode=blob_mode, blob_sources=blob_sources + ), self._table, fused=lambda: self._pk_fusion is not None, caller_requested_row_id=self._caller_requested_row_id, ) - def _run_hybrid(self, *, timeout: Optional[timedelta] = None) -> pa.Table: + def _run_hybrid( + self, + *, + timeout: Optional[timedelta] = None, + blob_mode: BlobMode, + blob_sources: dict[str, str], + ) -> pa.Table: self._create_query_builders() fts_query, vector_query = self._fts_query, self._vector_query # WAL-PK-FUSION: without the fallback, both legs always ask for row ids. @@ -2307,18 +2490,33 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): with_row_ids=True, offset=self._offset, ) - return self._finish_hybrid_results(results) + return self._finish_hybrid_results( + results, blob_mode=blob_mode, blob_sources=blob_sources + ) - def _finish_hybrid_results(self, results: pa.Table) -> pa.Table: + def _finish_hybrid_results( + self, + results: pa.Table, + *, + blob_mode: BlobMode, + blob_sources: dict[str, str], + ) -> pa.Table: if self._pk_fusion is not None: # WAL-PK-FUSION: delete. - return self._pk_fusion.strip(results) + return self._apply_arrow_blob_mode( + self._pk_fusion.strip(results), blob_mode, blob_sources + ) if self._user_requested_row_id(): - return results + return self._apply_arrow_blob_mode(results, blob_mode, blob_sources) if self._blob_auto_row_id_enabled(): - return self._finalize_blob_query_table(results) + return self._finalize_blob_query_table( + results, blob_mode=blob_mode, blob_sources=blob_sources + ) if "_rowid" in results.column_names: - return results.drop(["_rowid"]) - return results + # Fetch bytes before dropping _rowid. + return self._apply_arrow_blob_mode(results, blob_mode, blob_sources).drop( + ["_rowid"] + ) + return self._apply_arrow_blob_mode(results, blob_mode, blob_sources) @staticmethod def _combine_hybrid_results( @@ -2499,6 +2697,9 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): Higher values will yield better recall (more likely to find vectors if they exist) at the expense of latency. + The value is retained as `nprobes` until Lance sets both probe bounds to it. + Explicit minimum or maximum settings can override their respective bounds. + Parameters ---------- nprobes: int @@ -2511,8 +2712,33 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): """ if nprobes <= 0: raise ValueError("Invalid input, nprobes must be greater than 0") - self._minimum_nprobes = nprobes - self._maximum_nprobes = nprobes + self._nprobes = nprobes + return self + + def search_effort(self, search_effort: Optional[float]) -> LanceHybridQueryBuilder: + """Set the initial IVF search effort to a finite value in [0, 1]. + + None (the default) uses Auto, internally 0.5. Zero starts at the caller + minimum; one searches all partitions within the caller maximum. Later + probing may continue when too few candidates are found. Any explicit + effort, including 0.5, overrides nprobes in either call order. Explicit + minimum and maximum bounds still apply. This is not a recall guarantee + and does not change refinement or HNSW ef. MemWAL LSM queries do not + yet support this option. + + Examples + -------- + >>> import lancedb + >>> db = lancedb.connect("memory://") + >>> table = db.create_table( + ... "effort_hybrid", [{"vector": [0.1, 0.2], "text": "dog"}] + ... ) + >>> query = ( + ... table.search(query_type="hybrid").vector([0.1, 0.2]) + ... .text("dog").search_effort(0.75) + ... ) + """ + self._search_effort = search_effort return self def minimum_nprobes(self, minimum_nprobes: int) -> LanceHybridQueryBuilder: @@ -2741,6 +2967,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): ) return "\n".join(results) + def _blob_output_schema(self) -> pa.Schema: + self._create_query_builders() + return self._vector_query.output_schema() + def _create_query_builders(self): """Set up and configure the vector and FTS query builders.""" vector_query, fts_query = self._validate_query( @@ -2791,6 +3021,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): self._fts_query.phrase_query(True) if self._distance_type: self._vector_query.metric(self._distance_type) + if self._nprobes is not None: + self._vector_query.nprobes(self._nprobes) + if self._search_effort is not None: + self._vector_query.search_effort(self._search_effort) if self._minimum_nprobes is not None: self._vector_query.minimum_nprobes(self._minimum_nprobes) if self._maximum_nprobes is not None: @@ -2997,7 +3231,48 @@ class AsyncQueryBase(object): """ return await self._inner.output_schema() - async def to_arrow(self, timeout: Optional[timedelta] = None) -> pa.Table: + async def _validated_blob_sources(self, blob_mode: BlobMode) -> dict[str, str]: + """Map projected blob v2 outputs to their source columns. + + Bytes mode is checked here, before the query runs. + """ + validate_blob_mode(blob_mode) + # Lazy needs no mapping, and skipping it saves remote tables a request. + if blob_mode == "lazy" or self._table is None: + return {} + schema = await self._table.schema() + projection = _query_request_projection(self._inner.to_query_request()) + blob_sources = blob_v2_projection_sources(schema, projection) + if blob_mode != "bytes": + return blob_sources + reject_unfetchable_blob_columns(schema, projection, blob_sources) + if blob_sources and not _supports_local_blob_fetch( + self._table, self._table._pushdown_operations + ): + raise _blob_bytes_without_snapshot_error() + reject_unmatched_blob_outputs(await self._blob_output_schema(), blob_sources) + return blob_sources + + async def _blob_output_schema(self) -> pa.Schema: + return await self.output_schema() + + async def _apply_arrow_blob_mode( + self, tbl: pa.Table, blob_mode: BlobMode, blob_sources: dict[str, str] + ) -> pa.Table: + if blob_mode == "lazy": + return tbl + if blob_mode == "descriptions": + return strip_auto_row_ids(tbl, blob_sources) + return await replace_v2_blob_columns_with_bytes( + tbl, blob_sources, self._table.fetch_blobs + ) + + async def to_arrow( + self, + timeout: Optional[timedelta] = None, + *, + blob_mode: BlobMode = "lazy", + ) -> pa.Table: """ Execute the query and collect the results into an Apache Arrow Table. @@ -3011,13 +3286,26 @@ class AsyncQueryBase(object): The maximum time to wait for the query to complete. If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. + blob_mode: str, default "lazy" + ``"lazy"`` returns blob v2 descriptors that carry the row id + ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. + ``"bytes"`` returns payload bytes and raises for legacy blob columns + and blobs inside lists. + Remote tables and namespace query pushdown refuse ``"bytes"``. """ + blob_sources = await self._validated_blob_sources(blob_mode) batch_iter = await self.to_batches(timeout=timeout) - return self._finalize_blob_query_table( + tbl = self._finalize_blob_query_table( pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema) ) + return await self._apply_arrow_blob_mode(tbl, blob_mode, blob_sources) - async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]: + async def to_list( + self, + timeout: Optional[timedelta] = None, + *, + blob_mode: BlobMode = "lazy", + ) -> List[dict]: """ Execute the query and return the results as a list of dictionaries. @@ -3031,8 +3319,10 @@ class AsyncQueryBase(object): The maximum time to wait for the query to complete. If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. + blob_mode: str, default "lazy" + Passed through to :meth:`to_arrow`. """ - return (await self.to_arrow(timeout=timeout)).to_pylist() + return (await self.to_arrow(timeout=timeout, blob_mode=blob_mode)).to_pylist() async def to_pandas( self, @@ -3074,40 +3364,42 @@ class AsyncQueryBase(object): If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. blob_mode: str, default "lazy" - Controls how blob columns are returned. Remote queries support - "descriptions"; "bytes" and "lazy" are not yet supported. + Controls how blob columns are returned. ``"lazy"`` returns + :class:`~lancedb._blob.BlobFile` handles and needs a plain scan. + ``"descriptions"`` and ``"bytes"`` work for every query shape. A + plain scan without ``timeout`` reads the local dataset through + Lance's scanner, which also returns legacy blob columns as bytes. + Other shapes use :meth:`to_arrow` with the same ``blob_mode``. + Remote queries support ``"descriptions"`` only. **kwargs Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. """ validate_blob_mode(blob_mode) if self._table is not None and not self._table._inner._is_native(): + await self._validated_blob_sources(blob_mode) return await self._remote_to_pandas( flatten=flatten, timeout=timeout, blob_mode=blob_mode, **kwargs ) - if hasattr(self._inner, "output_schema"): - schema = await self.output_schema() - if _blob_mode_requires_native_pandas(blob_mode, schema): - native_error = None - if (flatten is None or blob_mode == "descriptions") and timeout is None: - try: - df = await self._plain_scan_to_pandas( - blob_mode, flatten=flatten, **kwargs - ) - if df is not None: - return df - except Exception as err: - native_error = err - reason = ( + if hasattr(self._inner, "output_schema") and schema_has_blob_field( + await self.output_schema() + ): + if (flatten is None or blob_mode != "lazy") and timeout is None: + df = await self._plain_scan_to_pandas( + blob_mode, flatten=flatten, **kwargs + ) + if df is not None: + return df + if blob_mode == "lazy" or _query_needs_lance_scanner(self): + raise _unsupported_blob_pandas_error( "this query shape cannot use Lance native pandas conversion" - if native_error is None - else str(native_error) ) - raise _unsupported_blob_pandas_error(reason) from native_error - tbl = flatten_columns(await self.to_arrow(timeout=timeout), flatten) - if _blob_mode_requires_native_pandas(blob_mode, tbl.schema): + tbl = flatten_columns( + await self.to_arrow(timeout=timeout, blob_mode=blob_mode), flatten + ) + if blob_mode == "lazy" and schema_has_blob_field(tbl.schema): raise _unsupported_blob_pandas_error( "this query shape cannot use Lance native pandas conversion" ) @@ -3162,6 +3454,13 @@ class AsyncQueryBase(object): if blob_mode == "bytes" else {} ) + if blob_sources: + reject_unfetchable_blob_columns(schema, query.columns, blob_sources) + # Remote and in-memory tables have no Lance dataset to scan. + if not self._table._inner._is_native() or ( + get_uri_scheme(await self._table.uri()) == "memory" + ): + return None dataset = await self._table.to_lance() scanner = dataset.scanner( **_scanner_kwargs_for_query( @@ -3176,7 +3475,7 @@ class AsyncQueryBase(object): blob_mode=blob_mode, blob_sources=blob_sources, fetch_blobs=self._table.fetch_blobs, - strip_auto_row_id=blob_auto_row_id, + strip_auto_row_id=self._with_row_id is not True, flatten=flatten, **kwargs, ) @@ -3690,7 +3989,7 @@ class AsyncVectorQueryBase: def nprobes(self, nprobes: int) -> Self: """ - Set the number of partitions to search (probe) + Set the legacy IVF probe parameter The number of probes must be greater than 0. @@ -3702,20 +4001,49 @@ class AsyncVectorQueryBase: The partition whose centroids are closest to the query vector will be exhaustiely searched to find matches. This parameter controls how many - partitions should be searched. + partitions will be searched. Increasing this value will increase the recall of your query but will - also increase the latency of your query. The default value is 20. This - default is good for many cases but the best value to use will depend on - your data and the recall that you need to achieve. + also increase the latency of your query. If this method is not called, + Lance's adaptive probe defaults are used. For best results we recommend tuning this parameter with a benchmark against your actual data to find the smallest possible value that will still give you the desired recall. + + LanceDB retains this as `nprobes` through local and remote request + construction. Lance sets both probe bounds to this value. Explicit minimum + or maximum settings can override their respective bounds. """ self._inner.nprobes(nprobes) return self + def search_effort(self, search_effort: Optional[float]) -> Self: + """Set the initial IVF search effort to a finite value in [0, 1]. + + None preserves Auto (internally 0.5) and fixed nprobes. Any explicit + effort, including 0.5, overrides nprobes in either call order; explicit + minimum and maximum bounds remain effective. Zero starts at the caller + minimum and one searches all partitions within the maximum. Later + probing may continue to obtain enough candidates. This is not a recall + guarantee and does not change refinement or HNSW ef. MemWAL LSM queries + do not yet support this option. + + Examples + -------- + >>> import asyncio + >>> import lancedb + >>> async def build_query(): + ... db = await lancedb.connect_async("memory://") + ... table = await db.create_table( + ... "effort_async", [{"vector": [0.1, 0.2]}] + ... ) + ... return table.vector_search([0.1, 0.2]).search_effort(0.75) + >>> query = asyncio.run(build_query()) + """ + self._inner.search_effort(search_effort) + return self + def minimum_nprobes(self, minimum_nprobes: int) -> Self: """Set the minimum number of probes to use. @@ -3987,6 +4315,10 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): return self + async def _blob_output_schema(self) -> pa.Schema: + vector_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table) + return await vector_query.output_schema() + async def _create_child_queries( self, ) -> Tuple["AsyncFTSQuery", "AsyncVectorQuery", int, int]: @@ -4348,7 +4680,12 @@ class BaseQueryBuilder(object): return pa.RecordBatchReader.from_batches(async_reader.schema, iter_sync()) - def to_arrow(self, timeout: Optional[timedelta] = None) -> pa.Table: + def to_arrow( + self, + timeout: Optional[timedelta] = None, + *, + blob_mode: BlobMode = "lazy", + ) -> pa.Table: """ Execute the query and collect the results into an Apache Arrow Table. @@ -4362,10 +4699,20 @@ class BaseQueryBuilder(object): The maximum time to wait for the query to complete. If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. + blob_mode: str, default "lazy" + ``"lazy"`` returns blob v2 descriptors that carry the row id + ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. + ``"bytes"`` returns payload bytes and raises for legacy blob columns + and blobs inside lists. """ - return LOOP.run(self._inner.to_arrow(timeout)) + return LOOP.run(self._inner.to_arrow(timeout, blob_mode=blob_mode)) - def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]: + def to_list( + self, + timeout: Optional[timedelta] = None, + *, + blob_mode: BlobMode = "lazy", + ) -> List[dict]: """ Execute the query and return the results as a list of dictionaries. @@ -4379,8 +4726,10 @@ class BaseQueryBuilder(object): The maximum time to wait for the query to complete. If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. + blob_mode: str, default "lazy" + Passed through to :meth:`to_arrow`. """ - return LOOP.run(self._inner.to_list(timeout)) + return LOOP.run(self._inner.to_list(timeout, blob_mode=blob_mode)) def to_pandas( self, @@ -4422,7 +4771,7 @@ class BaseQueryBuilder(object): If not specified, no timeout is applied. If the query does not complete within the specified time, an error will be raised. blob_mode: str, default "lazy" - Controls how blob columns are returned for plain scan queries. + Passed through to :meth:`AsyncQueryBase.to_pandas`. **kwargs Forwarded to pyarrow.Table.to_pandas after query execution and optional flattening. diff --git a/python/python/lancedb/remote/__init__.py b/python/python/lancedb/remote/__init__.py index dbcaa1b64..a1c8ebee5 100644 --- a/python/python/lancedb/remote/__init__.py +++ b/python/python/lancedb/remote/__init__.py @@ -170,6 +170,10 @@ class ClientConfig: Configuration for retrying failed requests. timeout_config: Optional[TimeoutConfig] Configuration for request timeouts. + blob_request_concurrency: Optional[int] + Maximum concurrent blob HTTP requests across all handles on this + connection. Defaults to 8. Set a positive integer here or via + `LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY`. extra_headers: Optional[dict] Additional headers to include in requests. id_delimiter: Optional[str] @@ -193,6 +197,7 @@ class ClientConfig: user_agent: str = f"LanceDB-Python-Client/{__version__}" retry_config: RetryConfig = field(default_factory=RetryConfig) timeout_config: Optional[TimeoutConfig] = field(default_factory=TimeoutConfig) + blob_request_concurrency: Optional[int] = None extra_headers: Optional[dict] = None id_delimiter: Optional[str] = None tls_config: Optional[TlsConfig] = None diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 0f3d35fc7..ac6b8b99d 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -112,6 +112,7 @@ def _client_config_to_dict(config: ClientConfig) -> dict[str, Any]: "user_agent": config.user_agent, "retry_config": _retry_config_to_dict(config.retry_config), "timeout_config": _timeout_config_to_dict(config.timeout_config), + "blob_request_concurrency": config.blob_request_concurrency, "extra_headers": config.extra_headers, "id_delimiter": config.id_delimiter, "tls_config": _tls_config_to_dict(config.tls_config), diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 7f03f8fa9..53b6c402a 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -57,6 +57,7 @@ from lancedb.table import IndexConfigType, KNOWN_METRICS import pyarrow as pa from lancedb.common import DATA, VEC, VECTOR_COLUMN_NAME +from lancedb._blob import validate_blob_mode from lancedb.merge import LanceMergeInsertBuilder from lancedb.embeddings import EmbeddingFunctionRegistry from lancedb.table import _normalize_progress @@ -239,12 +240,14 @@ class RemoteTable(Table): """List all versions of the table""" return LOOP.run(self._table.list_versions()) - def to_arrow(self) -> pa.Table: + def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """to_arrow() is not supported for remote tables.""" + validate_blob_mode(blob_mode) raise NotImplementedError("to_arrow() is not supported for remote tables.") def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs): """to_pandas() is not supported for remote tables.""" + validate_blob_mode(blob_mode) raise NotImplementedError("to_pandas() is not supported for remote tables.") def checkout(self, version: Union[int, str]): @@ -1154,7 +1157,9 @@ class RemoteTable(Table): return LOOP.run(self._table.blob_columns()) def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: return LOOP.run(self._table.fetch_blobs(column, row_ids)) @@ -1162,7 +1167,9 @@ class RemoteTable(Table): return LOOP.run(self._table.fetch_blob_ranges(column, list(requests))) def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) diff --git a/python/python/lancedb/schema.py b/python/python/lancedb/schema.py index 4dce09f0f..5edda5610 100644 --- a/python/python/lancedb/schema.py +++ b/python/python/lancedb/schema.py @@ -5,7 +5,7 @@ """Schema helpers for Lance blob columns.""" import importlib -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pyarrow as pa import pyarrow.ipc @@ -181,13 +181,60 @@ def _resolve_blob_type(): return _resolved_blob_type -def blob(name: str, nullable: bool = True) -> pa.Field: +def blob( + name: str, + nullable: bool = True, + *, + inline_size_threshold: Optional[int] = None, + dedicated_size_threshold: Optional[int] = None, + pack_file_size_threshold: Optional[int] = None, +) -> pa.Field: """Create a Lance blob v2 column field. When pylance is installed this is ``lance.blob.BlobType``. + + Parameters + ---------- + name: + Field name. + nullable: + Whether the column accepts null values. + inline_size_threshold: + Maximum payload size in bytes to keep inline in the data file + before using packed blob storage. Zero is allowed. + dedicated_size_threshold: + Maximum payload size in bytes to store in packed blob storage + before using dedicated blob storage. Must be positive. + pack_file_size_threshold: + Maximum size in bytes of a single packed blob sidecar file before + a new one is started. Must be positive. + + The thresholds are forwarded to ``lance.blob_field``, which validates + them and stores them as field metadata, so pylance is required whenever + any threshold is set. When all thresholds are omitted the behavior is + unchanged. """ - blob_type = _resolve_blob_type() - return pa.field(name, blob_type(), nullable=nullable) + if ( + inline_size_threshold is None + and dedicated_size_threshold is None + and pack_file_size_threshold is None + ): + blob_type = _resolve_blob_type() + return pa.field(name, blob_type(), nullable=nullable) + try: + lance = importlib.import_module("lance") + except ModuleNotFoundError as err: + raise ValueError( + "Setting blob tier thresholds requires pylance to be installed" + ) from err + _resolve_blob_type() + return lance.blob_field( + name, + nullable=nullable, + inline_size_threshold=inline_size_threshold, + dedicated_size_threshold=dedicated_size_threshold, + pack_file_size_threshold=pack_file_size_threshold, + ) def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType: diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index f20ca1a50..f137c2257 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -578,6 +578,55 @@ def _cast_to_target_schema( return pa.RecordBatchReader.from_batches(reordered_schema, gen()) +def _preserve_blob_v2_fields( + inferred_schema: pa.Schema, existing_schema: pa.Schema +) -> pa.Schema: + """Keep existing blob fields when an overwrite infers a replacement schema.""" + existing_fields = {field.name: field for field in existing_schema} + fields = [ + _preserve_blob_v2_field(field, existing_fields[field.name]) + if field.name in existing_fields + else field + for field in inferred_schema + ] + return pa.schema(fields, metadata=inferred_schema.metadata) + + +def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field: + if is_blob_v2_field(existing): + return existing + + inferred_type, existing_type = inferred.type, existing.type + if pa.types.is_struct(inferred_type) and pa.types.is_struct(existing_type): + existing_children = {field.name: field for field in existing_type} + children = [ + _preserve_blob_v2_field(field, existing_children[field.name]) + if field.name in existing_children + else field + for field in inferred_type + ] + return inferred.with_type(pa.struct(children)) + + if _is_list_like(inferred_type) and _is_list_like(existing_type): + value_field = _preserve_blob_v2_field( + inferred_type.value_field, existing_type.value_field + ) + if value_field == inferred_type.value_field: + return inferred + value_field = value_field.with_name(existing_type.value_field.name) + if pa.types.is_fixed_size_list(inferred_type): + value_type = pa.list_(value_field, inferred_type.list_size) + elif pa.types.is_large_list(inferred_type): + value_type = pa.large_list(value_field) + else: + value_type = pa.list_(value_field) + return inferred.with_type(value_type) + + if blob_v2_column_paths(pa.schema([existing])): + return existing + return inferred + + def _coerce_blob_write_columns( batch: pa.RecordBatch, target_schema: pa.Schema ) -> pa.RecordBatch: @@ -693,12 +742,7 @@ def _coerce_blob_list_values( def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array: - if pa.types.is_null(values.type): - data = pa.nulls(len(values), type=pa.large_binary()) - elif pa.types.is_large_binary(values.type): - data = values - else: - data = values.cast(pa.large_binary()) + value_field = "uri" if _is_string_like(values.type) else "data" length = len(values) storage_type = target_field.type if isinstance(storage_type, pa.ExtensionType): @@ -706,8 +750,11 @@ def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array: storage_fields = list(storage_type) children = [] for storage_field in storage_fields: - if storage_field.name == "data": - children.append(data) + if storage_field.name == value_field and not pa.types.is_null(values.type): + if values.type == storage_field.type: + children.append(values) + else: + children.append(values.cast(storage_field.type)) else: children.append(pa.nulls(length, type=storage_field.type)) storage = pa.StructArray.from_arrays( @@ -727,7 +774,19 @@ def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]: def _can_coerce_to_blob(data_type: pa.DataType) -> bool: - return _is_binary_like(data_type) or pa.types.is_null(data_type) + return ( + _is_binary_like(data_type) + or _is_string_like(data_type) + or pa.types.is_null(data_type) + ) + + +def _is_string_like(data_type: pa.DataType) -> bool: + return ( + pa.types.is_string(data_type) + or pa.types.is_large_string(data_type) + or pa.types.is_string_view(data_type) + ) def _is_binary_like(data_type: pa.DataType) -> bool: @@ -889,6 +948,24 @@ def _arrow_json_storage_type(input_type: pa.DataType) -> Optional[pa.DataType]: def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field: + if is_blob_v2_field(target_field): + source_type = field.type + if pa.types.is_struct(source_type): + if not any(child.name in ("data", "uri") for child in source_type): + raise ValueError( + f"blob struct input for column '{field.name}' must contain " + "a 'data' or 'uri' child" + ) + elif not _can_coerce_to_blob(source_type) and not isinstance( + source_type, pa.ExtensionType + ): + raise ValueError( + f"cannot coerce column '{field.name}' with type {source_type} " + "into a blob v2 struct. expected binary bytes " + "(Binary, LargeBinary, BinaryView), strings " + "(Utf8, LargeUtf8, Utf8View), or a Struct with a " + "'data' or 'uri' child" + ) # LanceDB exposes stored JSON columns as lance.json (JSONB-backed LargeBinary), but # casting the input to that storage type here merely relabels the raw JSON bytes as # JSONB. Lance must see arrow.json so it can perform the JSONB encoding. @@ -1324,9 +1401,17 @@ class Table(ABC): return self.to_arrow().to_pandas(**kwargs) @abstractmethod - def to_arrow(self) -> pa.Table: + def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """Return the table as a pyarrow Table. + Parameters + ---------- + blob_mode: str, default "lazy" + ``"lazy"`` returns blob v2 descriptors that carry the row id + ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. + ``"bytes"`` returns payload bytes and raises for legacy blob columns + and blobs inside lists. + Returns ------- pa.Table @@ -1740,7 +1825,8 @@ class Table(ABC): - pyarrow.Table or pyarrow.RecordBatch mode: str The mode to use when writing the data. Valid values are - "append" and "overwrite". + "append" and "overwrite". Overwrite infers a replacement schema + but preserves existing blob v2 fields present in the new data. on_bad_vectors: str, default "error" What to do if any of the vectors are not the same size or contains NaNs. One of "error", "drop", "fill". @@ -2056,19 +2142,30 @@ class Table(ABC): @abstractmethod def blob_columns(self) -> list[str]: - """Names of the blob v2 columns declared on this table.""" + """Names of blob v2 columns that can be fetched by row id. + + Struct-nested blobs use dotted paths. Blobs inside lists are omitted + because a row id does not identify an individual list element. + """ @abstractmethod def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: """Materialize full blob bytes for ``column`` at the given rows. The result has the same length and order as ``row_ids``. Null blobs produce null slots; valid empty blobs produce ``b""``. - ``_rowid`` values stay valid after compaction when the table has stable - row ids. + ``row_ids`` may be a list, an integer Arrow array or chunked array, + or a query table with row ids. ``_rowid`` values stay valid after + compaction when the table has stable row ids. + One missing or deleted row ID rejects the entire batch; no partial + results are returned. + + Blobs inside lists cannot be fetched by row id yet. Remote servers limit each request to 1024 row IDs and 64 MiB of blob bytes. The client splits requests automatically and reads an individual @@ -2080,17 +2177,20 @@ class Table(ABC): def fetch_blob_ranges( self, column: str, - requests: Sequence[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: """Materialize row-specific byte ranges from a blob v2 column. - Each request is a ``(row_id, offset, length)`` tuple. Requests may be - repeated or reordered, including multiple ranges for the same blob. + Each request is a three-item ``(row_id, offset, length)`` sequence. + Requests may be repeated or reordered, including multiple ranges for + the same blob. The result has the same length and order as ``requests``; null blobs produce null slots and empty ranges on non-null blobs produce ``b""``. ``_rowid`` values stay valid after compaction when the table has stable row ids. + One missing or deleted row ID rejects the entire batch; no partial + results are returned. Row IDs can be obtained from a query with ``with_row_id(True)``. This API is currently supported only by local tables. @@ -2098,18 +2198,22 @@ class Table(ABC): @abstractmethod def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": """Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles. Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is - a ``list[int]`` or a query ``pyarrow.Table`` carrying row identity via - ``_rowid`` or a ``_lance_row_id`` field on the blob descriptor. Null - rows are ``None``. Remote tables require LanceDB Cloud server 0.5.0 or - newer. + a ``list[int]``, an integer Arrow array or chunked array, or a query + ``pyarrow.Table`` carrying row identity via ``_rowid`` or a + ``_lance_row_id`` field on the blob descriptor. Null blob values return + ``None``. Remote tables require LanceDB Cloud server 0.5.0 or newer. ``_rowid`` values stay valid after compaction when the table has stable row ids. + One missing or deleted row ID rejects the entire batch; no partial + handles are returned. """ @abstractmethod @@ -2887,6 +2991,8 @@ class LanceTable(Table): managed_versioning=managed_versioning, ) ) + # Take queries read the pushdown set from the inner table. + self._table._pushdown_operations = self._pushdown_operations self._initialize_reopen_state(name) def _initialize_reopen_state(self, name: str) -> None: @@ -3232,19 +3338,23 @@ class LanceTable(Table): return LOOP.run(self._table.blob_columns()) def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: return LOOP.run(self._table.fetch_blobs(column, row_ids)) def fetch_blob_ranges( self, column: str, - requests: Sequence[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: return LOOP.run(self._table.fetch_blob_ranges(column, list(requests))) def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) @@ -3527,12 +3637,23 @@ class LanceTable(Table): return self.to_lance().to_pandas(blob_mode=blob_mode, **kwargs) - def to_arrow(self) -> pa.Table: + def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """Return the table as a pyarrow Table. + Parameters + ---------- + blob_mode: str, default "lazy" + ``"lazy"`` returns blob v2 descriptors that carry the row id + ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. + ``"bytes"`` returns payload bytes and raises for legacy blob columns + and blobs inside lists. + Returns ------- pa.Table""" + validate_blob_mode(blob_mode) + if blob_mode != "lazy" and schema_has_blob_field(self.schema): + return self.search().to_arrow(blob_mode=blob_mode) if ( _should_push_down_query_table( self._namespace_client, self._pushdown_operations @@ -4268,7 +4389,8 @@ class LanceTable(Table): The data to insert into the table. mode: str The mode to use when writing the data. Valid values are - "append" and "overwrite". + "append" and "overwrite". Overwrite infers a replacement schema + but preserves existing blob v2 fields present in the new data. on_bad_vectors: str, default "error" What to do if any of the vectors are not the same size or contains NaNs. One of "error", "drop", "fill", "null". @@ -4732,6 +4854,7 @@ class LanceTable(Table): and not self._route_pushdown_to_rust and self.current_branch() is None and query.take_offsets is None + and query.search_effort is None ): from lancedb.namespace import _execute_server_side_query @@ -5978,22 +6101,37 @@ class AsyncTable: return await self.query().to_pandas(blob_mode=blob_mode, **kwargs) return (await self.to_lance()).to_pandas(blob_mode=blob_mode, **kwargs) - async def to_arrow(self) -> pa.Table: + async def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """Return the table as a pyarrow Table. + Parameters + ---------- + blob_mode: str, default "lazy" + ``"lazy"`` returns blob v2 descriptors that carry the row id + ``fetch_blobs`` reads. ``"descriptions"`` returns them without it. + ``"bytes"`` returns payload bytes and raises for legacy blob columns + and blobs inside lists. + Returns ------- pa.Table """ + validate_blob_mode(blob_mode) + schema = await self.schema() + fetches_bytes = blob_mode == "bytes" and schema_has_blob_field(schema) if ( - _should_push_down_query_table( + not fetches_bytes + and _should_push_down_query_table( self._namespace_client, self._pushdown_operations ) and not self._route_pushdown_to_rust ): - return (await self._execute_query(Query())).read_all() - - return await self.query().to_arrow() + tbl = (await self._execute_query(Query())).read_all() + if blob_mode == "descriptions": + tbl = strip_auto_row_ids(tbl, row_addressable_blob_v2_paths(schema)) + return tbl + # The query builder refuses bytes under pushdown before running the query. + return await self.query().to_arrow(blob_mode=blob_mode) async def create_index( self, @@ -6289,7 +6427,8 @@ class AsyncTable: - pyarrow.Table or pyarrow.RecordBatch mode: str The mode to use when writing the data. Valid values are - "append" and "overwrite". + "append" and "overwrite". Overwrite infers a replacement schema + but preserves existing blob v2 fields present in the new data. on_bad_vectors: str, default "error" What to do if any of the vectors are not the same size or contains NaNs. One of "error", "drop", "fill", "null". @@ -6327,6 +6466,11 @@ class AsyncTable: data, _ = sanitize_create_table( data, None, on_bad_vectors=on_bad_vectors, fill_value=fill_value ) + if blob_v2_column_paths(schema): + # Plain bytes cannot be inferred as blob v2. Keep matching blob + # fields while allowing other columns to use the inferred schema. + target_schema = _preserve_blob_v2_fields(data.schema, schema) + data = _cast_to_target_schema(data, target_schema) elif on_bad_vectors != "error" or ( schema.metadata is not None and b"embedding_functions" in schema.metadata ): @@ -6731,9 +6875,12 @@ class AsyncTable: ) if query.distance_type is not None: async_query = async_query.distance_type(query.distance_type) + if query.nprobes is not None: + async_query = async_query.nprobes(query.nprobes) + if query.search_effort is not None: + async_query = async_query.search_effort(query.search_effort) if query.minimum_nprobes is not None and query.maximum_nprobes is not None: - # Set both to the minimum first to avoid min > max error. - async_query = async_query.nprobes( + async_query = async_query.minimum_nprobes( query.minimum_nprobes ).maximum_nprobes(query.maximum_nprobes) elif query.minimum_nprobes is not None: @@ -6772,6 +6919,7 @@ class AsyncTable: ) and not self._route_pushdown_to_rust and query.take_offsets is None + and query.search_effort is None ): from lancedb.namespace import _execute_server_side_query @@ -7366,7 +7514,9 @@ class AsyncTable: return await self._inner.blob_columns() async def fetch_blobs( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> pa.LargeBinaryArray: return await self._inner.fetch_blobs( column, _normalize_blob_row_ids(row_ids, column) @@ -7375,12 +7525,14 @@ class AsyncTable: async def fetch_blob_ranges( self, column: str, - requests: Sequence[Tuple[int, int, int]], + requests: Sequence[Sequence[int]], ) -> pa.LargeBinaryArray: return await self._inner.fetch_blob_ranges(column, list(requests)) async def fetch_blob_files( - self, column: str, row_ids: Union[list[int], pa.Table] + self, + column: str, + row_ids: Union[list[int], pa.Array, pa.ChunkedArray, pa.Table], ) -> "list[Optional[BlobFile]]": handles = await self._inner.fetch_blob_files( column, _normalize_blob_row_ids(row_ids, column) diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index 3a6a8a102..fefc234ce 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -10,6 +10,7 @@ import os import pathlib import warnings from datetime import date, datetime +from decimal import Decimal from functools import singledispatch from typing import Tuple, Union, Optional, Any, List from urllib.parse import urlparse @@ -373,6 +374,14 @@ def _(value: float): return str(value) +@value_to_sql.register(Decimal) +def _(value: Decimal): + if not value.is_finite(): + raise ValueError("Non-finite Decimal values cannot be converted to SQL") + # Quote exact decimal text so the target-column cast avoids float rounding. + return value_to_sql(format(value, "f")) + + @value_to_sql.register(bool) def _(value: bool): return str(value).upper() diff --git a/python/python/tests/test_authz.py b/python/python/tests/test_authz.py new file mode 100644 index 000000000..b4e145317 --- /dev/null +++ b/python/python/tests/test_authz.py @@ -0,0 +1,821 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import asyncio +from dataclasses import FrozenInstanceError +import json +from unittest.mock import AsyncMock + +import pytest + +import lancedb +from lancedb.authz import ( + AccessControlEntry, + AsyncAuthorization, + Authorization, + GroupInfo, + IdentityInfo, + Object, + PrincipalInfo, + Privilege, + RoleBinding, + Subject, +) +from lancedb.remote import HeaderProvider +from lancedb.remote.errors import HttpError +from test_catalog import catalog_server as _catalog_server + +catalog_server = _catalog_server + + +class DynamicHeaders(HeaderProvider): + def get_headers(self): + return { + "Authorization": "Bearer rotating-token", + "X-LanceDB-Database": "wrong-dynamic-db", + "X-LanceDB-Database-Prefix": "wrong-prefix", + } + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_authz_operations_and_scope(catalog_server, asynchronous): + endpoint, requests, responses = catalog_server + responses.extend( + [ + (200, {}), + (200, {}), + (200, {}), + (200, {}), + ( + 200, + { + "entries": [ + { + "subject_id": "g:team-id", + "object": "table:tenant/db:a$b$events", + "privilege": "PRIVILEGE20", + } + ] + }, + ), + ( + 200, + { + "bindings": [ + { + "subject_id": "p:alice-id", + "subject_name": "Alice", + "role": "reader", + } + ] + }, + ), + ( + 200, + { + "principal_id": "alice-id", + "groups": [{"id": "team-id"}], + "roles": ["reader"], + }, + ), + ( + 200, + { + "group_id": "team-id", + "group_name": "Engineering", + "principals": [{"id": "alice-id", "name": "Alice"}], + "roles": ["reader"], + }, + ), + ] + ) + config = { + "header_provider": DynamicHeaders(), + "extra_headers": {"x-lancedb-database": "wrong-static-db"}, + } + obj = Object.table(database="tenant/db", namespace_path=["a", "b"], name="events") + + async def exercise(): + if asynchronous: + catalog = await lancedb.connect_catalog_async( + endpoint, client_config=config + ) + else: + catalog = lancedb.connect_catalog(endpoint, client_config=config) + authz = catalog.authz + assert authz is catalog.authz + assert requests == [] # Constructing handles does not probe a database. + + async def call(method, *args, **kwargs): + value = method(*args, **kwargs) + return await value if asynchronous else value + + await call( + authz.add_acl, + object=obj, + subject=Subject.group_name("Engineering"), + privilege=Privilege.SELECT, + ) + await call( + authz.delete_acl, + object=obj, + subject=Subject.group_id("team-id"), + privilege="SELECT", + ) + await call( + authz.add_role_binding, + subject=Subject.principal_name("Alice"), + role="reader", + ) + await call( + authz.delete_role_binding, + subject=Subject.principal_api_key("subject-secret"), + role="reader", + ) + iterator = authz.list_acls(object=obj, subject=Subject.role("reader")) + entries = ( + [entry async for entry in iterator] if asynchronous else list(iterator) + ) + assert entries == [ + AccessControlEntry( + str(obj), Subject.group_id("team-id"), None, "PRIVILEGE20" + ) + ] + iterator = authz.list_role_bindings( + subject=Subject.group_id("team-id"), role="reader" + ) + bindings = ( + [binding async for binding in iterator] if asynchronous else list(iterator) + ) + assert bindings == [ + RoleBinding(Subject.principal_id("alice-id"), "Alice", "reader") + ] + assert await call( + authz.get_principal, Subject.principal_id("alice-id") + ) == PrincipalInfo("alice-id", None, [IdentityInfo("team-id")], ["reader"]) + assert await call( + authz.get_group, Subject.group_name("Engineering") + ) == GroupInfo( + "team-id", "Engineering", [IdentityInfo("alice-id", "Alice")], ["reader"] + ) + + asyncio.run(exercise()) + assert [path for path, _, _ in requests] == [ + "/admin/authz/acl/add", + "/admin/authz/acl/delete", + "/admin/authz/role_binding/add", + "/admin/authz/role_binding/delete", + "/admin/authz/acl/list", + "/admin/authz/role_binding/list", + "/admin/authz/principal", + "/admin/authz/group", + ] + assert [body for _, _, body in requests] == [ + {"object": str(obj), "subject": "G:Engineering", "privilege": "SELECT"}, + {"object": str(obj), "subject": "g:team-id", "privilege": "SELECT"}, + {"subject": "P:Alice", "role": "reader"}, + {"subject": "a:subject-secret", "role": "reader"}, + {"object": str(obj), "subject": "r:reader", "max_results": 100}, + {"subject": "g:team-id", "role": "reader", "max_results": 100}, + {"principal": "p:alice-id"}, + {"group": "G:Engineering"}, + ] + for _, headers, _ in requests: + headers = {key.lower(): value for key, value in headers.items()} + assert "x-lancedb-database" not in headers + assert "x-lancedb-database-prefix" not in headers + assert headers["authorization"] == "Bearer rotating-token" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("limit", [None, 0, 1, 3]) +def test_lazy_expanded_acl_pagination(catalog_server, asynchronous, limit): + endpoint, requests, responses = catalog_server + + def entry(privilege): + return { + "object": "database:db", + "subject_id": "r:reader", + "privilege": privilege, + } + + responses.extend( + [ + ( + 200, + { + "entries": [entry("SELECT"), entry("INSERT")], + "continuation_token": "opaque/first", + }, + ), + (200, {"entries": [], "continuation_token": "opaque/second"}), + (200, {"entries": [entry("UPDATE"), entry("DELETE")]}), + ] + ) + + async def exercise(): + catalog = ( + await lancedb.connect_catalog_async(endpoint, api_key="caller-secret") + if asynchronous + else lancedb.connect_catalog(endpoint, api_key="caller-secret") + ) + iterator = catalog.authz.list_acls(subject=Subject.role("reader"), limit=limit) + assert requests == [] + return [entry async for entry in iterator] if asynchronous else list(iterator) + + values = asyncio.run(exercise()) + expected = ["SELECT", "INSERT", "UPDATE", "DELETE"] + assert [entry.privilege for entry in values] == ( + expected if limit is None else expected[:limit] + ) + assert len(requests) == (0 if limit == 0 else 1 if limit == 1 else 3) + for _, headers, body in requests: + assert headers["x-api-key"] == "caller-secret" + assert body["subject"] == "r:reader" + if len(requests) == 3: + assert requests[1][2]["continuation_token"] == "opaque/first" + assert requests[2][2]["continuation_token"] == "opaque/second" + + +@pytest.mark.parametrize("status", [400, 401, 403, 404, 409, 501, 503]) +def test_http_errors_preserve_status_without_replaying_writes(catalog_server, status): + endpoint, requests, responses = catalog_server + responses.append((status, {"error": "authorization failure"})) + authz = lancedb.connect_catalog(endpoint).authz + with pytest.raises(HttpError) as error: + authz.add_acl( + object="system", subject=Subject.role("operator"), privilege="OPERATE" + ) + assert error.value.status_code == status + assert error.value.request_id + assert len(requests) == 1 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_role_pagination_and_late_failure(catalog_server, asynchronous): + endpoint, requests, responses = catalog_server + responses.extend( + [ + ( + 200, + { + "bindings": [{"subject_id": "g:team", "role": "reader"}], + "continuation_token": "next", + }, + ), + (403, {"error": "denied"}), + ] + ) + + async def exercise(): + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + iterator = catalog.authz.list_role_bindings(role="reader") + first = await anext(iterator) if asynchronous else next(iterator) + assert first.subject == Subject.group_id("team") + assert len(requests) == 1 + with pytest.raises(HttpError) as error: + if asynchronous: + await anext(iterator) + else: + next(iterator) + assert error.value.status_code == 403 + + asyncio.run(exercise()) + assert requests[1][2] == { + "role": "reader", + "max_results": 100, + "continuation_token": "next", + } + + +def test_early_close_does_not_fetch_next_page(catalog_server): + endpoint, requests, responses = catalog_server + responses.append( + ( + 200, + { + "bindings": [{"subject_id": "g:team", "role": "reader"}], + "continuation_token": "next", + }, + ) + ) + iterator = lancedb.connect_catalog(endpoint).authz.list_role_bindings() + next(iterator) + iterator.close() + assert len(requests) == 1 + + +def test_selectors_objects_and_redaction(): + secret = Subject.principal_api_key("do-not-print-this") + assert "do-not-print-this" not in repr(secret) + assert "do-not-print-this" not in str(secret) + with pytest.raises(FrozenInstanceError): + secret.kind = "role" + assert str(Object.system()) == "system" + assert str(Object.database("tenant/db")) == "database:tenant/db" + assert ( + str(Object.namespace(database="db", namespace_path=["a", "b"])) + == "namespace:db:a$b" + ) + assert str(Object.table(database="db", name="events")) == "table:db:public$events" + assert str(Object.view(database="db", name="recent")) == "view:db:public$recent" + assert Subject.principal_id("p:id")._wire() == "p:p:id" + assert Subject.group_id("team-id").value == "team-id" + for namespace in (["a$b"], [""], "public", [".."]): + with pytest.raises(ValueError): + Object.table(database="db", namespace_path=namespace, name="events") + for database in ("", "a//b", "a/../b", "a:b", "a%2Fb", "a$b"): + with pytest.raises(ValueError): + Object.database(database) + + +@pytest.mark.asyncio +async def test_validation_before_io_and_empty_lists(): + inner = AsyncMock() + inner.list_acls.return_value = json.dumps({"entries": []}) + inner.list_role_bindings.return_value = json.dumps({"bindings": []}) + authz = AsyncAuthorization(inner) + for limit in (-1, 1.5, True): + with pytest.raises(ValueError): + authz.list_acls(limit=limit) + with pytest.raises(ValueError): + Authorization(authz).list_role_bindings(limit=limit) + with pytest.raises(ValueError): + await authz.get_group(Subject.principal_id("alice")) + with pytest.raises(ValueError): + await authz.get_principal(Subject.group_id("team")) + for method in (authz.add_role_binding, authz.delete_role_binding): + with pytest.raises(ValueError): + await method(subject=Subject.role("reader"), role="writer") + with pytest.raises(ValueError): + authz.list_role_bindings(subject=Subject.role("reader")) + assert [entry async for entry in authz.list_acls(limit=0)] == [] + assert inner.mock_calls == [] + assert [entry async for entry in authz.list_acls()] == [] + assert [entry async for entry in authz.list_role_bindings()] == [] + + +def test_api_key_subject_is_not_logged(catalog_server): + import os + import subprocess + import sys + + endpoint, requests, responses = catalog_server + responses.append((200, {"principal_id": "public-id", "groups": [], "roles": []})) + result = subprocess.run( + [ + sys.executable, + "-c", + "import lancedb, sys; from lancedb.authz import Subject; " + "lancedb.connect_catalog(sys.argv[1]).authz.get_principal(" + "Subject.principal_api_key('test-subject-secret-never-log'))", + endpoint, + ], + env={**os.environ, "LANCEDB_LOG": "debug"}, + capture_output=True, + text=True, + check=True, + timeout=30, + ) + assert requests[0][2] == {"principal": "a:test-subject-secret-never-log"} + assert "test-subject-secret-never-log" not in result.stdout + result.stderr + + +@pytest.mark.parametrize( + "builder, kwargs, wire", + [ + ("system", {}, "system"), + ("secret", {"database": "db", "name": "api-key"}, "secret:db:public$api-key"), + ( + "function", + {"database": "db", "name": "caption"}, + "function:db:public$caption", + ), + ( + "secret", + {"database": "tenant/db", "namespace_path": ["a", "b"], "name": "api-key"}, + "secret:tenant/db:a$b$api-key", + ), + ( + "function", + {"database": "tenant/db", "namespace_path": ["a", "b"], "name": "caption"}, + "function:tenant/db:a$b$caption", + ), + ("database", {"name": "tenant/db"}, "database:tenant/db"), + ( + "namespace", + {"database": "tenant/db", "namespace_path": ["a", "b"]}, + "namespace:tenant/db:a$b", + ), + ( + "table", + {"database": "tenant/db", "namespace_path": ["a", "b"], "name": "events"}, + "table:tenant/db:a$b$events", + ), + ( + "view", + {"database": "db", "namespace_path": ["public"], "name": "recent"}, + "view:db:public$recent", + ), + ], +) +def test_rust_backed_object_builders(builder, kwargs, wire): + import pickle + from lancedb import _lancedb + + value = getattr(Object, builder)(**kwargs) + assert str(value) == wire + assert str(value._inner) == wire + assert str(_lancedb.AuthzObject(wire)) == wire + assert value == Object(wire) + assert hash(value) == hash(Object(wire)) + assert pickle.loads(pickle.dumps(value)) == value + + +@pytest.mark.parametrize( + "kind, prefix", + [ + ("principal_id", "p:"), + ("principal_name", "P:"), + ("group_id", "g:"), + ("group_name", "G:"), + ("role", "r:"), + ("principal_api_key", "a:"), + ], +) +def test_rust_backed_subject_selectors(kind, prefix): + import pickle + from lancedb import _lancedb + + value = "p:literal@example.com / 名前" + subject = getattr(Subject, kind)(value) + assert subject.kind == kind + assert subject.value == value + assert subject._wire() == prefix + value + assert _lancedb.AuthzSubject(kind, value).to_wire() == prefix + value + assert pickle.loads(pickle.dumps(subject)) == subject + assert hash(subject) == hash(getattr(Subject, kind)(value)) + with pytest.raises(ValueError): + getattr(Subject, kind)("") + if kind == "principal_api_key": + assert value not in repr(subject._inner) + assert value not in str(subject._inner) + if kind in ("principal_id", "group_id", "role"): + assert Subject._canonical(prefix + value) == subject + else: + with pytest.raises(ValueError) as error: + Subject._canonical(prefix + value) + assert value not in str(error.value) + + +@pytest.mark.parametrize( + "wire", + [ + "", + "database:", + "database:a//b", + "database:a/../b", + "database:a%2Fb", + "namespace:db:a$$b", + "table:db:events", + "table:db:public$", + "table:db:public$..", + "table:db:a/b$events", + "view:db:public$café", + "secret:db:public$", + "function:db:public$..", + "secret:db:a$$key", + "function:db:caption", + "secret:db:public$bad/name", + "function:db:public$bad:name", + "system:db", + ], +) +def test_invalid_raw_objects_fail_before_io(catalog_server, wire): + endpoint, requests, _ = catalog_server + authz = lancedb.connect_catalog(endpoint).authz + with pytest.raises(ValueError): + Object(wire) + with pytest.raises(ValueError): + authz.list_acls(object=wire) + with pytest.raises(ValueError): + authz.add_acl(object=wire, subject=Subject.role("reader"), privilege="SELECT") + assert requests == [] + + +@pytest.mark.parametrize( + "wire", + [ + "future:db:resource", + "table_extension:db:resource", + "secret:tenant/db:a$b$key", + "function:tenant/db:a$b$caption", + ], +) +def test_objects_round_trip_through_acl_requests(catalog_server, wire): + endpoint, requests, responses = catalog_server + responses.extend( + [ + (200, {}), + ( + 200, + { + "entries": [ + { + "object": wire, + "subject_id": "r:reader", + "privilege": "SELECT", + } + ] + }, + ), + (200, {}), + ] + ) + authz = lancedb.connect_catalog(endpoint).authz + obj = Object(wire) + assert str(obj) == wire + authz.add_acl(object=obj, subject=Subject.role("reader"), privilege="SELECT") + entries = list(authz.list_acls(object=wire)) + assert entries[0].object == wire + authz.delete_acl( + object=entries[0].object, subject=Subject.role("reader"), privilege="SELECT" + ) + assert len(requests) == 3 + for _, _, body in requests: + assert body["object"] == wire + + +@pytest.mark.asyncio +async def test_native_requests_validate_selector_context_before_io(catalog_server): + endpoint, requests, _ = catalog_server + catalog = await lancedb.connect_catalog_async(endpoint) + native = catalog.authz._inner + for method, request in [ + (native.list_role_bindings, {"subject": "r:reader"}), + (native.add_role_binding, {"subject": "r:reader", "role": "writer"}), + (native.delete_role_binding, {"subject": "r:reader", "role": "writer"}), + (native.get_principal, {"principal": "g:group"}), + (native.get_group, {"group": "a:never-log-this-key"}), + ]: + with pytest.raises(ValueError) as error: + await method(json.dumps(request)) + assert "never-log-this-key" not in str(error.value) + assert requests == [] + + +def test_privilege_values_use_rust_serialization(): + import pickle + from lancedb import _lancedb + + names = [value.value for value in Privilege] + assert names == _lancedb.AuthzPrivilege.known() + for privilege in Privilege: + assert json.loads(json.dumps(privilege)) == privilege.value + assert str(_lancedb.AuthzPrivilege(privilege)) == privilege.value + assert pickle.loads(pickle.dumps(privilege)) is privilege + for index in range(64): + alias = f"PRIVILEGE{index}" + expected = names[index] if index < len(names) else alias + assert str(_lancedb.AuthzPrivilege(alias)) == expected + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "creation_privilege", [None, "CREATE_SECRET", "CREATE_FUNCTION"] +) +def test_privilege_names_and_future_values_round_trip( + catalog_server, asynchronous, creation_privilege +): + endpoint, requests, responses = catalog_server + inputs = [ + Privilege.SELECT, + "SELECT", + "PRIVILEGE2", + "PRIVILEGE20", + "PRIVILEGE63", + "PRIVILEGE64", + "PRIVILEGE256", + "PRIVILEGE18446744073709551616", + ] + expected = ["SELECT", "SELECT", "SELECT", *inputs[3:]] + object_id = "table:db:public$t" + if creation_privilege is not None: + inputs = [Privilege[creation_privilege], creation_privilege] + expected = [creation_privilege, creation_privilege] + object_id = "namespace:db:public" + for privilege in inputs: + responses.extend( + [ + (200, {}), + ( + 200, + { + "entries": [ + { + "object": object_id, + "subject_id": "r:reader", + "privilege": privilege, + } + ] + }, + ), + (200, {}), + ] + ) + + async def exercise(): + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + authz = catalog.authz + for privilege, canonical in zip(inputs, expected): + added = authz.add_acl( + object=object_id, + subject=Subject.role("reader"), + privilege=privilege, + ) + if asynchronous: + await added + iterator = authz.list_acls() + entries = ( + [entry async for entry in iterator] if asynchronous else list(iterator) + ) + assert entries[0].privilege == canonical + deleted = authz.delete_acl( + object=entries[0].object, + subject=entries[0].subject, + privilege=entries[0].privilege, + ) + if asynchronous: + await deleted + + asyncio.run(exercise()) + assert len(requests) == 3 * len(inputs) + for index, canonical in enumerate(expected): + assert requests[index * 3][2]["object"] == object_id + assert requests[index * 3][2]["privilege"] == canonical + assert requests[index * 3 + 2][2]["object"] == object_id + assert requests[index * 3 + 2][2]["privilege"] == canonical + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_invalid_privileges_fail_before_io(catalog_server, asynchronous): + endpoint, requests, _ = catalog_server + invalid = [ + "", + "select", + "SELECT ", + "UNKNOWN", + "PRIVILEGE", + "PRIVILEGE-1", + "PRIVILEGE+1", + "PRIVILEGE1x", + None, + 2, + True, + ["SELECT"], + ] + + async def exercise(): + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + async_authz = catalog.authz if asynchronous else catalog.authz._inner + native = async_authz._inner + for privilege in invalid: + for method in [catalog.authz.add_acl, catalog.authz.delete_acl]: + with pytest.raises(ValueError, match="privilege"): + result = method( + object="system", + subject=Subject.role("reader"), + privilege=privilege, + ) + if asynchronous: + await result + for method in [ + native.add_acl, + native.delete_acl, + ]: + with pytest.raises(ValueError): + await method( + json.dumps( + { + "object": "system", + "subject": "r:reader", + "privilege": privilege, + } + ) + ) + + asyncio.run(exercise()) + assert requests == [] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_acl_page_preserves_unknown_subjects(catalog_server, asynchronous): + import pickle + + endpoint, requests, responses = catalog_server + responses.extend( + [ + ( + 200, + { + "entries": [ + { + "subject_id": "p:known", + "object": "system", + "privilege": "USAGE", + }, + { + "subject_id": "x:future-subject", + "object": "system", + "privilege": "USAGE", + }, + ] + }, + ), + (200, {}), + ] + ) + + async def exercise(): + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + iterator = catalog.authz.list_acls() + entries = ( + [entry async for entry in iterator] if asynchronous else list(iterator) + ) + assert len(entries) == 2 + assert entries[0].subject == Subject.principal_id("known") + subject = entries[1].subject + assert subject.kind == "unknown" + assert subject.value == "x:future-subject" + assert subject._wire() == "x:future-subject" + assert pickle.loads(pickle.dumps(subject)) == subject + result = catalog.authz.delete_acl( + object=entries[1].object, subject=subject, privilege=entries[1].privilege + ) + if asynchronous: + await result + + asyncio.run(exercise()) + assert len(requests) == 2 + assert requests[1][2]["subject"] == "x:future-subject" + + +@pytest.mark.parametrize("builder", ["secret", "function"]) +def test_secret_and_function_builders_validate_components(builder): + from lancedb import _lancedb + + for database, namespace, name in [ + ("a//b", ["public"], "valid"), + ("db", ["a$b"], "valid"), + ("db", "public", "valid"), + ("db", ["public"], "bad$name"), + ("db", ["public"], ".."), + ]: + with pytest.raises(ValueError): + getattr(Object, builder)( + database=database, namespace_path=namespace, name=name + ) + with pytest.raises(ValueError): + getattr(_lancedb.AuthzObject, builder)(database, namespace, name) + + +@pytest.mark.parametrize( + "builder", ["namespace", "table", "view", "secret", "function"] +) +def test_empty_and_omitted_namespace_paths_select_public(builder): + from lancedb import _lancedb + + kwargs = {"database": "db"} + native_args = () + wire = "namespace:db:public" + if builder != "namespace": + kwargs["name"] = "item" + native_args = ("item",) + wire = f"{builder}:db:public$item" + + build = getattr(Object, builder) + expected = build(**kwargs, namespace_path=["public"]) + assert str(expected) == wire + assert build(**kwargs) == expected + for namespace in [[], ()]: + assert build(**kwargs, namespace_path=namespace) == expected + assert ( + str(getattr(_lancedb.AuthzObject, builder)("db", namespace, *native_args)) + == wire + ) diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 72107364c..71e78643c 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -15,12 +15,14 @@ from lance.blob import Blob, BlobType as LanceBlobType import lancedb from lancedb._blob import ( + _set_blob_column, blob_v2_projection_sources, read_row_ids_from_hits, stash_auto_row_ids, ) from lancedb.expr import col from lancedb.index import FTS +from lancedb.rerankers import Reranker from lancedb.schema import blob_column_paths, blob_v2_column_paths @@ -69,21 +71,25 @@ def _row_ids_by_id(table): def _assert_missing_blob_row_ids(exc_info): message = str(exc_info.value) - assert "row ids" in message + assert "row ids but some do not exist in the table" in message assert "rowaddr" not in message assert "fragment" not in message def _assert_fetch_apis_reject_missing_row_ids(table, row_ids): + expected = f"first missing row ids: [{row_ids[0]}]" with pytest.raises(ValueError) as exc_info: table.fetch_blobs("image", row_ids) _assert_missing_blob_row_ids(exc_info) + assert expected in str(exc_info.value) with pytest.raises(ValueError) as exc_info: table.fetch_blob_files("image", row_ids) _assert_missing_blob_row_ids(exc_info) + assert expected in str(exc_info.value) with pytest.raises(ValueError) as exc_info: table.fetch_blob_ranges("image", [(row_id, 0, 1) for row_id in row_ids]) _assert_missing_blob_row_ids(exc_info) + assert expected in str(exc_info.value) def test_blob_factory_declares_v2_field(): @@ -94,6 +100,37 @@ def test_blob_factory_declares_v2_field(): assert type(field.type) is LanceBlobType +def test_blob_thresholds_land_in_field_metadata(): + field = lancedb.blob( + "payload", + inline_size_threshold=8 * 1024, + dedicated_size_threshold=1 << 20, + pack_file_size_threshold=1 << 26, + ) + assert field.type.extension_name == "lance.blob.v2" + metadata = field.metadata + assert metadata[b"lance-encoding:blob-inline-size-threshold"] == b"8192" + assert metadata[b"lance-encoding:blob-dedicated-size-threshold"] == b"1048576" + assert metadata[b"lance-encoding:blob-pack-file-size-threshold"] == b"67108864" + + +def test_blob_without_thresholds_has_no_threshold_metadata(): + field = lancedb.blob("image") + metadata = field.metadata or {} + assert not any(b"threshold" in key for key in metadata) + + +def test_blob_thresholds_are_validated(): + with pytest.raises(ValueError, match="inline_size_threshold must be non-negative"): + lancedb.blob("payload", inline_size_threshold=-1) + with pytest.raises(ValueError, match="dedicated_size_threshold must be positive"): + lancedb.blob("payload", dedicated_size_threshold=0) + with pytest.raises(ValueError, match="pack_file_size_threshold must be positive"): + lancedb.blob("payload", pack_file_size_threshold=0) + with pytest.raises(TypeError, match="inline_size_threshold must be an int"): + lancedb.blob("payload", inline_size_threshold="8192") + + def test_blob_type_works_without_pylance(): script = _HIDE_LANCE_BLOB + textwrap.dedent( """\ @@ -297,6 +334,273 @@ def test_blob_v2_projection_sources_use_typed_column_name(): } +def test_blob_v2_projection_sources_include_selected_struct(): + schema = pa.schema( + [ + pa.field( + "info", pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]) + ) + ] + ) + + assert blob_v2_projection_sources(schema, ["info"]) == {"info.blob": "info.blob"} + assert blob_v2_projection_sources(schema, {"renamed": "info"}) == { + "renamed.blob": "info.blob" + } + + +@pytest.mark.parametrize("parent_path", ["info", "outer.info"]) +def test_set_blob_column_missing_nested_field_raises(parent_path): + row = {"name": "first"} + for name in reversed(parent_path.split(".")): + row = {name: row} + table = pa.Table.from_pylist([row]) + blobs = pa.array([b"hello"], type=pa.large_binary()) + + with pytest.raises(KeyError) as exc_info: + _set_blob_column(table, f"{parent_path}.blob", blobs) + + assert parent_path in str(exc_info.value) + assert "'blob'" in str(exc_info.value) + + +def _nested_blob_table(db, name): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field( + "info", pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]) + ), + ] + ) + table = db.create_table(name, schema=schema) + info = pa.StructArray.from_arrays( + [ + pa.array(["first", "second", "third"]), + pa.array([b"hello", None, b"hidden"], type=pa.large_binary()), + ], + names=["name", "blob"], + mask=pa.array([False, False, True]), + ) + table.add(pa.table({"id": pa.array([0, 1, 2]), "info": info})) + return table + + +def test_nested_blob_to_pandas_bytes(tmp_path): + table = _nested_blob_table(lancedb.connect(tmp_path), "nested_pandas_bytes") + expected = [ + {"name": "first", "blob": b"hello"}, + {"name": "second", "blob": None}, + None, + ] + + assert table.to_pandas(blob_mode="bytes")["info"].tolist() == expected + assert table.search().to_pandas(blob_mode="bytes")["info"].tolist() == expected + assert ( + table.search().select(["info"]).to_pandas(blob_mode="bytes")["info"].tolist() + == expected + ) + assert ( + table.search() + .select({"renamed": "info"}) + .to_pandas(blob_mode="bytes")["renamed"] + .tolist() + == expected + ) + assert table.search().select({"b": "info.blob"}).to_pandas(blob_mode="bytes")[ + "b" + ].tolist() == [b"hello", None, None] + + +def test_to_arrow_and_to_list_blob_mode(tmp_path): + table = _nested_blob_table(lancedb.connect(tmp_path), "arrow_blob_mode") + expected = [ + {"name": "first", "blob": b"hello"}, + {"name": "second", "blob": None}, + None, + ] + + assert table.to_arrow(blob_mode="bytes")["info"].to_pylist() == expected + assert table.search().to_arrow(blob_mode="bytes")["info"].to_pylist() == expected + assert table.search().to_list(blob_mode="bytes") == [ + {"id": 0, "info": expected[0]}, + {"id": 1, "info": expected[1]}, + {"id": 2, "info": expected[2]}, + ] + assert table.search().select({"b": "info.blob"}).to_list(blob_mode="bytes") == [ + {"b": b"hello"}, + {"b": None}, + {"b": None}, + ] + assert table.search().select(["id", "info.blob"]).to_list(blob_mode="bytes") == [ + {"id": 0, "info.blob": b"hello"}, + {"id": 1, "info.blob": None}, + {"id": 2, "info.blob": None}, + ] + assert table.search().select({"x.y": "info"}).to_list(blob_mode="bytes") == [ + {"x.y": expected[0]}, + {"x.y": expected[1]}, + {"x.y": expected[2]}, + ] + assert [ + row["info"] for row in table.search().select([]).to_list(blob_mode="bytes") + ] == expected + assert table.search().select([]).to_pandas(blob_mode="bytes")["info"].tolist() == ( + table.search().select([]).to_arrow(blob_mode="bytes")["info"].to_pylist() + ) + assert "_rowid" not in table.search().select([]).to_arrow().column_names + # A flat non-blob column named like the nested blob path keeps its values. + shadowed = table.search().select({"info": "info", "info.blob": "id"}) + assert [row["info.blob"] for row in shadowed.to_list()] == [0, 1, 2] + assert [row["info"] for row in shadowed.to_list(blob_mode="bytes")] == expected + with pytest.raises(ValueError, match="both produce blob column 'info.blob'"): + table.search().select(["info", "info.blob"]).to_list(blob_mode="bytes") + + default_blob = table.search().select({"b": "info.blob"}).to_list()[0]["b"] + assert isinstance(default_blob, dict) + lazy = table.to_arrow(blob_mode="lazy")["info"].to_pylist() + assert table.to_arrow()["info"].to_pylist() == lazy + lazy_descriptor = lazy[0]["blob"] + assert "_lance_row_id" in lazy_descriptor + lazy_descriptor.pop("_lance_row_id") + descriptions = table.to_arrow(blob_mode="descriptions")["info"].to_pylist() + assert descriptions[0]["blob"] == lazy_descriptor + assert ( + table.search().to_list(blob_mode="descriptions")[0]["info"]["blob"] + == lazy_descriptor + ) + + with pytest.raises(ValueError, match="blob_mode must be one of"): + table.to_arrow(blob_mode="handles") + with pytest.raises(ValueError, match="blob_mode must be one of"): + table.search().to_list(blob_mode="handles") + + +@pytest.mark.asyncio +async def test_async_to_arrow_and_to_list_blob_mode(tmp_path): + db = await lancedb.connect_async(tmp_path) + table = await db.create_table( + "async_arrow_blob_mode", + schema=pa.schema( + [ + pa.field("id", pa.int64()), + pa.field( + "info", + pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]), + ), + ] + ), + ) + await table.add( + pa.table( + { + "id": pa.array([0, 1]), + "info": pa.StructArray.from_arrays( + [ + pa.array(["first", "second"]), + pa.array([b"hello", None], type=pa.large_binary()), + ], + names=["name", "blob"], + ), + } + ) + ) + expected = [ + {"name": "first", "blob": b"hello"}, + {"name": "second", "blob": None}, + ] + + assert (await table.to_arrow(blob_mode="bytes"))["info"].to_pylist() == expected + assert (await table.query().to_arrow(blob_mode="bytes"))[ + "info" + ].to_pylist() == expected + assert await table.query().to_list(blob_mode="bytes") == [ + {"id": 0, "info": expected[0]}, + {"id": 1, "info": expected[1]}, + ] + default_blob = (await table.query().select({"b": "info.blob"}).to_list())[0]["b"] + assert isinstance(default_blob, dict) + descriptions = await table.to_arrow(blob_mode="descriptions") + assert ( + "_lance_row_id" + not in descriptions.schema.field("info").type.field("blob").type.names + ) + + +def test_vector_query_to_arrow_blob_mode_bytes(tmp_path): + db = lancedb.connect(tmp_path) + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("vector", pa.list_(pa.float32(), 2)), + lancedb.blob("blob"), + ] + ) + table = db.create_table("vector_arrow_blob", schema=schema) + table.add( + [ + {"id": 1, "vector": [1.0, 0.0], "blob": b"one"}, + {"id": 2, "vector": [0.0, 1.0], "blob": b"two"}, + ] + ) + table.add([{"id": 3, "vector": [0.9, 0.1], "blob": b"three"}]) + table.add([{"id": 4, "vector": [0.1, 0.9], "blob": None}]) + + hits = table.search([0.0, 1.0]).limit(4).to_arrow(blob_mode="bytes").to_pylist() + assert [(hit["id"], hit["blob"]) for hit in hits] == [ + (2, b"two"), + (4, None), + (3, b"three"), + (1, b"one"), + ] + arrow = table.search([1.0, 0.0]).limit(1).to_arrow(blob_mode="bytes") + assert arrow["blob"].to_pylist() == [b"one"] + assert table.search([1.0, 0.0]).limit(1).to_list(blob_mode="bytes")[0]["blob"] == ( + b"one" + ) + assert isinstance( + table.search([1.0, 0.0]).limit(1).to_list()[0]["blob"], + dict, + ) + + +@pytest.mark.asyncio +async def test_nested_blob_to_pandas_bytes_async(tmp_path): + db = await lancedb.connect_async(tmp_path) + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field( + "info", pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]) + ), + ] + ) + table = await db.create_table("nested_pandas_bytes_async", schema=schema) + await table.add( + pa.table( + { + "id": pa.array([0, 1]), + "info": pa.StructArray.from_arrays( + [ + pa.array(["first", "second"]), + pa.array([b"hello", None], type=pa.large_binary()), + ], + names=["name", "blob"], + ), + } + ) + ) + expected = [ + {"name": "first", "blob": b"hello"}, + {"name": "second", "blob": None}, + ] + + assert (await table.to_pandas(blob_mode="bytes"))["info"].tolist() == expected + assert (await table.query().to_pandas(blob_mode="bytes"))[ + "info" + ].tolist() == expected + + def _legacy_v1_table(name): # Legacy v1 blob columns are only writable at file version <= 2.1. db = lancedb.connect( @@ -396,6 +700,19 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id(): assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"} +@pytest.mark.asyncio +async def test_async_calculated_select_does_not_read_every_blob_column(): + db = await lancedb.connect_async("memory:///calculated_select") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = await db.create_table("calculated_select", schema=schema) + await table.add([{"id": 1, "image": b"hello"}]) + + hits = await table.query().select({"next_id": col("id") + 1}).to_arrow() + + assert hits.column_names == ["next_id"] + assert hits["next_id"].to_pylist() == [2] + + @pytest.mark.asyncio async def test_async_typed_blob_projection_preserves_source_column(): db = await lancedb.connect_async("memory:///typed_blob_projection") @@ -410,6 +727,185 @@ async def test_async_typed_blob_projection_preserves_source_column(): assert blobs.to_pylist() == [b"alpha"] +@pytest.mark.asyncio +@pytest.mark.parametrize("shape", ["scan", "hybrid"]) +async def test_async_bytes_mode_rejects_blob_spelling_it_cannot_match(shape): + db = await lancedb.connect_async("memory:///") + table = await _hybrid_blob_table_async(db) + query = table.query() + if shape == "hybrid": + query = query.nearest_to([1.0, 0.0]).nearest_to_text("hello") + + with pytest.raises(ValueError, match=r"could not match result columns \['pic'\]"): + await query.select({"id": "id", "pic": "`image`"}).to_arrow(blob_mode="bytes") + + +def test_bytes_mode_keeps_struct_with_blob_descriptor_field_names(): + descriptor_like = pa.struct( + [ + pa.field("kind", pa.uint8()), + pa.field("position", pa.uint64()), + pa.field("size", pa.uint64()), + pa.field("blob_id", pa.uint32()), + pa.field("blob_uri", pa.utf8()), + ] + ) + meta = {"kind": 1, "position": 2, "size": 3, "blob_id": 4, "blob_uri": "u"} + table = lancedb.connect("memory:///").create_table( + "descriptor_like", + schema=pa.schema( + [ + pa.field("id", pa.int64()), + lancedb.blob("image"), + pa.field("meta", descriptor_like), + ] + ), + ) + table.add([{"id": 1, "image": b"one", "meta": meta}]) + + assert table.search().to_list(blob_mode="bytes") == [ + {"id": 1, "image": b"one", "meta": meta} + ] + + +@pytest.mark.parametrize("source", ["`image`", "IMAGE"]) +def test_bytes_mode_rejects_blob_spelling_it_cannot_match(source): + table = _blob_table("blob_spelling", [{"id": 1, "image": b"one"}]) + + with pytest.raises(ValueError, match="could not match result columns"): + table.search().select({"id": "id", "pic": source}).to_arrow(blob_mode="bytes") + + +def test_to_pandas_does_not_drop_fragment_selection_on_fallback(): + table = _blob_table("blob_fragments", [{"id": 1, "image": b"one"}]) + table.add([{"id": 2, "image": b"two"}]) + + query = table.search().fragment_ids([0]) + with pytest.raises(RuntimeError, match="native pandas conversion"): + query.to_pandas(blob_mode="bytes") + with pytest.raises(RuntimeError, match="native pandas conversion"): + table.search().with_row_address(True).to_pandas(blob_mode="descriptions") + + +@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes"]) +def test_to_pandas_reports_lance_dataset_failures(tmp_path, blob_mode): + table = lancedb.connect(tmp_path).create_table( + "dataset_failure", + schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), + ) + table.add([{"id": 1, "image": b"one"}]) + + def reject_credentials(): + raise PermissionError("storage credentials rejected") + + table.to_lance = reject_credentials + with pytest.raises(PermissionError, match="storage credentials rejected"): + table.search().to_pandas(blob_mode=blob_mode) + + +def test_to_pandas_keeps_to_arrow_overrides_without_blob_mode(): + from lancedb.query import LanceEmptyQueryBuilder + + class ToArrowWithoutBlobMode(LanceEmptyQueryBuilder): + def to_arrow(self, *, timeout=None): + return super().to_arrow(timeout=timeout) + + table = lancedb.connect("memory:///").create_table("no_blob", data=[{"id": 1}]) + query = ToArrowWithoutBlobMode(table) + assert query.to_list(blob_mode="bytes") == [{"id": 1}] + assert query.to_pandas(blob_mode="descriptions")["id"].tolist() == [1] + with pytest.raises(ValueError, match="blob_mode must be one of"): + query.to_list(blob_mode="handles") + with pytest.raises(ValueError, match="blob_mode must be one of"): + table.search().to_list(blob_mode="handles") + + +def test_take_queries_honor_blob_mode(): + table = _blob_table( + "take_blob_mode", [{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}] + ) + + hits = table.take_offsets([1, 0]).to_arrow(blob_mode="bytes") + assert dict(zip(hits["id"].to_pylist(), hits["image"].to_pylist())) == { + 1: b"one", + 2: b"two", + } + assert [ + row["image"] for row in table.take_row_ids([0]).to_list(blob_mode="bytes") + ] == [b"one"] + assert table.take_offsets([0]).to_pandas(blob_mode="bytes")["image"].tolist() == [ + b"one" + ] + description = table.take_row_ids([0]).to_list(blob_mode="descriptions")[0]["image"] + assert "_lance_row_id" not in description + + +def test_bytes_mode_follows_swapped_blob_aliases(): + db = lancedb.connect("memory:///") + table = db.create_table( + "swapped_aliases", + schema=pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("caption", pa.utf8()), + lancedb.blob("image"), + lancedb.blob("thumb"), + ] + ), + ) + table.add([{"id": 0, "caption": "cap", "image": b"img", "thumb": b"th"}]) + + swapped = table.search().select({"id": "id", "image": "thumb", "thumb": "image"}) + assert swapped.to_list(blob_mode="bytes") == [ + {"id": 0, "image": b"th", "thumb": b"img"} + ] + over_text = table.search().select({"id": "id", "image": "caption", "pic": "image"}) + assert over_text.to_list(blob_mode="bytes") == [ + {"id": 0, "image": "cap", "pic": b"img"} + ] + + +def test_bytes_mode_refused_under_namespace_query_pushdown(tmp_path): + db = lancedb.connect_namespace( + "dir", + {"root": str(tmp_path)}, + namespace_client_pushdown_operations=["QueryTable"], + ) + table = db.create_table( + "pushdown_blob", + schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), + ) + table.add([{"id": 1, "image": b"one"}]) + + with pytest.raises(NotImplementedError, match="stable table snapshot"): + table.to_arrow(blob_mode="bytes") + with pytest.raises(NotImplementedError, match="stable table snapshot"): + table.search().to_list(blob_mode="bytes") + with pytest.raises(NotImplementedError, match="stable table snapshot"): + table.take_row_ids([0]).to_list(blob_mode="bytes") + descriptor = table.to_arrow(blob_mode="descriptions")["image"].to_pylist()[0] + assert "_lance_row_id" not in descriptor + + +@pytest.mark.asyncio +async def test_async_bytes_mode_refused_under_namespace_query_pushdown(tmp_path): + db = lancedb.connect_namespace_async( + "dir", + {"root": str(tmp_path)}, + namespace_client_pushdown_operations=["QueryTable"], + ) + table = await db.create_table( + "pushdown_blob", + schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), + ) + await table.add([{"id": 1, "image": b"one"}]) + + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.to_arrow(blob_mode="bytes") + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.query().to_list(blob_mode="bytes") + + def test_fetch_blobs_round_trip(): table = _blob_table( "round_trip", @@ -420,6 +916,73 @@ def test_fetch_blobs_round_trip(): assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"] +def test_add_overwrite_preserves_blob_v2(tmp_path): + db = lancedb.connect(tmp_path) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("overwrite_blob", schema=schema) + table.add([{"id": 1, "image": b"old"}]) + + payload = b"x" * 200_000 + table.add([{"id": 2, "image": payload}], mode="overwrite") + + assert table.count_rows() == 1 + assert blob_v2_column_paths(table.schema) == ["image"] + assert table.blob_columns() == ["image"] + assert table.fetch_blobs("image", [_row_ids_by_id(table)[2]]).to_pylist() == [ + payload + ] + assert list(tmp_path.rglob("*.blob")) + + +def test_add_overwrite_with_blob_still_infers_vector_schema(): + db = lancedb.connect("memory:///overwrite_blob_vector") + schema = pa.schema( + [ + pa.field("vector", pa.list_(pa.float32(), 2)), + lancedb.blob("image"), + ] + ) + table = db.create_table("overwrite_blob_vector", schema=schema) + + table.add([{"vector": [1.0, 2.0, 3.0], "image": b"data"}], mode="overwrite") + + assert table.schema.field("vector").type == pa.list_(pa.float32(), 3) + assert table.blob_columns() == ["image"] + + +@pytest.mark.asyncio +async def test_add_overwrite_preserves_blob_v2_async(): + db = await lancedb.connect_async("memory:///overwrite_blob_async") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = await db.create_table("overwrite_blob_async", schema=schema) + + await table.add([{"id": 1, "image": b"updated"}], mode="overwrite") + + assert blob_v2_column_paths(await table.schema()) == ["image"] + assert await table.blob_columns() == ["image"] + hits = await table.query().with_row_id().to_arrow() + assert (await table.fetch_blobs("image", hits)).to_pylist() == [b"updated"] + + +def test_add_overwrite_preserves_nested_blob_v2(): + db = lancedb.connect("memory:///overwrite_nested_blob") + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("info", pa.struct([lancedb.blob("image")])), + ] + ) + table = db.create_table("overwrite_nested_blob", schema=schema) + + table.add([{"id": 1, "info": {"image": b"nested"}}], mode="overwrite") + + assert blob_v2_column_paths(table.schema) == ["info.image"] + assert table.blob_columns() == ["info.image"] + assert table.fetch_blobs("info.image", [_row_ids_by_id(table)[1]]).to_pylist() == [ + b"nested" + ] + + def test_merge_insert_writes_python_bytes(): table = _blob_table("merge_bytes", [{"id": 1, "image": b"before"}]) result = ( @@ -435,6 +998,50 @@ def test_merge_insert_writes_python_bytes(): assert blobs.to_pylist() == [b"updated", b"inserted"] +@pytest.mark.parametrize("write", ["merge_insert", "create_table"]) +def test_uri_string_blob_write_reaches_external_base_validation(tmp_path, write): + blob_path = tmp_path / "payload.bin" + blob_path.write_bytes(b"payload") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + rows = pa.table( + {"id": pa.array([1], pa.int64()), "image": pa.array([blob_path.as_uri()])} + ) + db = lancedb.connect(tmp_path) + + # These APIs do not expose add's outside-base opt-in. The URI should reach + # the blob policy check instead of failing at PyArrow's string-to-struct cast. + with pytest.raises(ValueError, match="outside registered external bases"): + if write == "merge_insert": + table = db.create_table("merge_uri", schema=schema) + table.merge_insert("id").when_not_matched_insert_all().execute(rows) + else: + db.create_table("create_uri", data=rows, schema=schema) + + +@pytest.mark.parametrize("write", ["merge_insert", "create_table"]) +@pytest.mark.parametrize( + "source_type, value", + [ + (pa.binary(4), b"abcd"), + (pa.list_(pa.uint8()), [1, 2]), + (pa.int64(), 42), + ], +) +def test_blob_write_rejects_unsupported_input(write, source_type, value): + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + rows = pa.table( + {"id": pa.array([1], pa.int64()), "image": pa.array([value], source_type)} + ) + db = lancedb.connect("memory:///") + + with pytest.raises(ValueError, match="cannot coerce column 'image' with type"): + if write == "merge_insert": + table = db.create_table("merge_invalid", schema=schema) + table.merge_insert("id").when_not_matched_insert_all().execute(rows) + else: + db.create_table("create_invalid", data=rows, schema=schema) + + def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path): db = lancedb.connect(tmp_path) schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) @@ -661,12 +1268,29 @@ def test_merge_insert_list_blob_dicts(): def test_list_blob_column_queries_as_raw_descriptors(): table = _list_blob_table("list_query") + assert table.blob_columns() == [] hits = table.search().limit(10).to_arrow() element = hits.schema.field("images").type.value_type assert pa.types.is_struct(element) assert "_lance_row_id" not in element.names - with pytest.raises(ValueError, match="expected struct before segment"): - table.fetch_blobs("images.image", [0]) + for column in ("images", "images.image"): + with pytest.raises(ValueError, match="blobs inside lists cannot be fetched"): + table.fetch_blobs(column, [0]) + with pytest.raises(ValueError, match="blobs inside lists cannot be fetched"): + table.fetch_blob_files(column, [0]) + with pytest.raises(ValueError, match="blobs inside lists cannot be fetched"): + table.fetch_blob_ranges(column, [(0, 0, 1)]) + + +def test_bytes_mode_rejects_blobs_inside_lists(): + table = _list_blob_table("list_blob_bytes") + + with pytest.raises(ValueError, match=r"\['images.image'\].*blobs inside lists"): + table.to_arrow(blob_mode="bytes") + with pytest.raises(ValueError, match="blobs inside lists"): + table.search().select({"pics": "images"}).to_list(blob_mode="bytes") + descriptions = table.to_arrow(blob_mode="descriptions")["images"].to_pylist() + assert descriptions[0][0]["size"] == len(b"before") def test_row_addressable_paths_exclude_list_children(): @@ -714,6 +1338,38 @@ def test_fetch_blobs_accepts_query_result(): assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"} +def test_fetch_blob_apis_accept_integer_arrow_row_ids(): + table = _blob_table( + "arrow_row_ids", + [{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}], + ) + by_id = _row_ids_by_id(table) + ids = [by_id[2], by_id[1]] + + for row_ids in ( + pa.array(ids, type=pa.uint64()), + pa.chunked_array([ids[:1], ids[1:]], type=pa.int64()), + ): + assert table.fetch_blobs("image", row_ids).to_pylist() == [b"beta", b"alpha"] + files = table.fetch_blob_files("image", row_ids) + assert [file.read() for file in files] == [b"beta", b"alpha"] + + +def test_fetch_blob_apis_validate_arrow_and_negative_row_ids(): + table = _blob_table("invalid_arrow_row_ids", [{"id": 1, "image": b"alpha"}]) + hits = table.search().with_row_id(True).to_arrow() + + for fetch in (table.fetch_blobs, table.fetch_blob_files): + with pytest.raises(ValueError, match="non-negative"): + fetch("image", [-1]) + with pytest.raises(ValueError, match="non-negative"): + fetch("image", pa.array([-1], type=pa.int64())) + with pytest.raises(ValueError, match="without nulls"): + fetch("image", pa.array([None], type=pa.int64())) + with pytest.raises(ValueError, match="integer type"): + fetch("image", hits["image"]) + + def test_fetch_blobs_after_compact_with_stable_row_ids(tmp_path): db = lancedb.connect( tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"} @@ -816,13 +1472,14 @@ def test_add_list_of_dicts_to_blob_list_column(large_list): table = db.create_table(f"blob_{large_list}_list_add", data=seed) table.add([{"id": 1, "items": [None]}]) + table.add([{"id": 2, "items": [b"a", None, b""]}]) table.add( - [{"id": 2, "items": [b"a", None]}], + [{"id": 3, "items": [b"b", None]}], on_bad_vectors="fill", ) ids = table.search().select(["id"]).to_arrow()["id"].to_pylist() - assert sorted(ids) == [0, 1, 2] + assert sorted(ids) == [0, 1, 2, 3] assert pa.types.is_large_list(table.schema.field("items").type) is large_list @@ -845,6 +1502,15 @@ def test_fetch_blob_ranges_aligns_repeated_ranges_and_nulls(): assert ranges.to_pylist() == [b"cde", None, b"ab", b"cde", b""] +def test_fetch_blob_ranges_accepts_three_item_sequences(): + table = _blob_table("range_sequences", [{"id": 1, "image": b"abcdef"}]) + row_id = _row_ids_by_id(table)[1] + + ranges = table.fetch_blob_ranges("image", [[row_id, 1, 2], (row_id, 4, 2)]) + + assert ranges.to_pylist() == [b"bc", b"ef"] + + def test_fetch_blob_ranges_validates_requests(): table = _blob_table("range_validation", [{"id": 1, "image": b"abc"}]) row_id = _row_ids_by_id(table)[1] @@ -859,6 +1525,30 @@ def test_fetch_blob_ranges_validates_requests(): table.fetch_blob_ranges("image", [(2**64 - 1, 0, 1)]) _assert_missing_blob_row_ids(exc_info) + for index, name in enumerate(("row id", "offset", "length")): + request = [row_id, 0, 1] + request[index] = -1 + with pytest.raises( + ValueError, match=f"blob range {name} must be a non-negative" + ): + table.fetch_blob_ranges("image", [request]) + + with pytest.raises(ValueError, match="three values"): + table.fetch_blob_ranges("image", [[row_id, 0]]) + + +def test_fetch_blob_apis_reject_row_past_fragment_end(): + table = _blob_table( + "past_fragment_end", + [ + {"id": 1, "image": b"a"}, + {"id": 2, "image": b"b"}, + {"id": 3, "image": b"c"}, + ], + ) + _assert_fetch_apis_reject_missing_row_ids(table, [3]) + _assert_fetch_apis_reject_missing_row_ids(table, [3, 0]) + def test_fetch_blob_apis_reject_missing_fragment_row_addr(): table = _blob_table("missing_frag", [{"id": 1, "image": b"x"}]) @@ -876,6 +1566,19 @@ def test_fetch_blob_apis_reject_deleted_row_ids(): _assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]]) +def test_fetch_blob_apis_reject_deleted_stable_row_ids(tmp_path): + db = lancedb.connect( + tmp_path, storage_options={"new_table_enable_stable_row_ids": "true"} + ) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("stable_deleted", schema=schema) + table.add([{"id": 1, "image": b"one"}, {"id": 2, "image": b"two"}]) + by_id = _row_ids_by_id(table) + table.delete("id = 2") + + _assert_fetch_apis_reject_missing_row_ids(table, [by_id[2], by_id[1]]) + + def test_fetch_blob_ranges_empty_requests_returns_empty_array(): table = _blob_table("range_empty", [{"id": 1, "image": b"x"}]) assert table.fetch_blob_ranges("image", []).to_pylist() == [] @@ -914,6 +1617,11 @@ def test_fetch_blobs_nested_path(): blobs = table.fetch_blobs("info.blob", [by_id[1], by_id[2]]) assert [blobs[0].as_py(), blobs[1].as_py()] == [b"nested-alpha", b"nested-beta"] + # The source path's parent name now holds a string column in the hits. + hits = table.search().select({"id": "id", "info": "info.name", "b": "info.blob"}) + blobs = table.fetch_blobs("info.blob", hits.to_arrow()) + assert blobs.to_pylist() == [b"nested-alpha", b"nested-beta"] + def test_fetch_blob_files_lazy_read(): payload = b"lazy-read" * 100 @@ -1032,6 +1740,9 @@ async def test_async_fetch_blob_files_lazy_read(): handles = await table.fetch_blob_files("image", hits) assert len(handles) == 1 assert await handles[0].aread() == payload + handles[0].close() + with pytest.raises(ValueError, match="read of closed file"): + await handles[0].aread() def test_fetch_blobs_from_query_result_without_row_id_raises(): @@ -1088,6 +1799,122 @@ def test_blob_v2_hybrid_fetch_blobs(): assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"} +def _blob_query_shape(table, shape): + if shape == "scan": + return table.search().where("id < 3") + if shape == "vector": + return table.search([1.0, 0.0]) + if shape == "fts": + return table.search("hello", query_type="fts") + return table.search(query_type="hybrid").vector([1.0, 0.0]).text("hello") + + +@pytest.mark.parametrize("with_row_id", [None, True, False]) +@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"]) +@pytest.mark.parametrize("shape", ["scan", "vector", "fts", "hybrid"]) +def test_blob_mode_for_every_query_shape(shape, blob_mode, with_row_id): + table = _hybrid_blob_table(lancedb.connect("memory:///")) + query = _blob_query_shape(table, shape).select(["id", "image"]).limit(2) + if with_row_id is not None: + query = query.with_row_id(with_row_id) + + hits = query.to_arrow(blob_mode=blob_mode) + + assert ("_rowid" in hits.column_names) == (with_row_id is True) + images = {row["id"]: row["image"] for row in hits.to_pylist()} + if blob_mode == "bytes": + assert images == {1: b"alpha", 2: b"beta"} + else: + assert sorted(images) == [1, 2] + assert all("_lance_row_id" not in image for image in images.values()) + df = query.to_pandas(blob_mode=blob_mode) + assert dict(zip(df["id"], df["image"])) == images + + +@pytest.mark.asyncio +@pytest.mark.parametrize("shape", ["scan", "vector", "fts", "hybrid"]) +async def test_async_blob_bytes_for_every_query_shape(shape): + db = await lancedb.connect_async("memory:///") + table = await _hybrid_blob_table_async(db) + query = table.query() + if shape in ("vector", "hybrid"): + query = query.nearest_to([1.0, 0.0]) + if shape in ("fts", "hybrid"): + query = query.nearest_to_text("hello") + if shape == "scan": + query = query.where("id < 3") + query = query.select(["id", "image"]).limit(2) + + hits = await query.to_arrow(blob_mode="bytes") + + assert {row["id"]: row["image"] for row in hits.to_pylist()} == { + 1: b"alpha", + 2: b"beta", + } + df = await query.to_pandas(blob_mode="bytes") + assert dict(zip(df["id"], df["image"])) == {1: b"alpha", 2: b"beta"} + + +class _KeepAllScoresReranker(Reranker): + """Rebuilds the merged table like model rerankers with return_score="all".""" + + def __init__(self): + super().__init__(return_score="all") + + def rerank_hybrid(self, query, vector_results, fts_results): + combined = self._merge_and_keep_scores(vector_results, fts_results) + scores = pa.array([1.0] * len(combined), type=pa.float32()) + return combined.append_column("_relevance_score", scores) + + +def test_bytes_mode_rejects_legacy_blob_after_reranker_drops_metadata(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table( + "legacy_rerank", + pa.table( + { + "id": pa.array([1, 2], pa.int64()), + "text": ["hello alpha", "hello beta"], + "vector": pa.array( + [[1.0, 0.0], [0.9, 0.1]], type=pa.list_(pa.float32(), 2) + ), + "blob": pa.array([b"one", b"two"], pa.large_binary()), + }, + schema=_HYBRID_BLOB_SCHEMA.remove(3).append( + pa.field( + "blob", pa.large_binary(), metadata={"lance-encoding:blob": "true"} + ) + ), + ), + storage_options={"new_table_data_storage_version": "2.1"}, + ) + table.create_index("text", config=FTS(with_position=False)) + query = ( + table.search(query_type="hybrid") + .vector([1.0, 0.0]) + .text("hello") + .rerank(_KeepAllScoresReranker()) + ) + + with pytest.raises(ValueError, match=r"\['blob'\].*legacy blob columns"): + query.to_arrow(blob_mode="bytes") + + +@pytest.mark.parametrize("source", ["`image`", "IMAGE"]) +def test_bytes_mode_rejects_blob_spelling_after_reranker_drops_metadata(source): + table = _hybrid_blob_table(lancedb.connect("memory:///")) + query = ( + table.search(query_type="hybrid") + .vector([1.0, 0.0]) + .text("hello") + .select({"id": "id", "pic": source}) + .rerank(_KeepAllScoresReranker()) + ) + + with pytest.raises(ValueError, match=r"could not match result columns \['pic'\]"): + query.to_arrow(blob_mode="bytes") + + @pytest.mark.asyncio async def test_blob_v2_hybrid_fetch_blobs_async(): db = await lancedb.connect_async("memory:///hybrid_blob_fetch_async") @@ -1167,6 +1994,49 @@ def test_blob_file_seek_read_and_read_range(): handle.seek(0, 99) +def test_blob_file_matches_file_error_and_range_contract(): + payload = b"abcdef" + table = _blob_table("file_contract", [{"id": 1, "image": payload}]) + row_id = _row_ids_by_id(table)[1] + handle = table.fetch_blob_files("image", [row_id])[0] + + handle.seek(2) + for offset, whence in [(-1, io.SEEK_SET), (-3, io.SEEK_CUR), (-7, io.SEEK_END)]: + with pytest.raises(ValueError, match="negative seek value -1"): + handle.seek(offset, whence) + assert handle.tell() == 2 + + assert handle.read(None) == payload[2:] + handle.seek(2) + assert handle.read1(2) == payload[2:4] + assert handle.read1() == payload[4:] + handle.seek(2) + assert handle.read_ranges([(0, 2), (4, 2), (2, 0)]) == [b"ab", b"ef", b""] + assert handle.tell() == 2 + + with pytest.raises(ValueError): + handle.read_range(5, 2) + with pytest.raises(ValueError): + handle.read_ranges([(0, 1), (5, 2)]) + assert handle.tell() == 2 + + handle.close() + with pytest.raises(ValueError, match="read of closed file"): + handle.read(1) + with pytest.raises(ValueError, match="read of closed file"): + handle.read(None) + with pytest.raises(ValueError, match="read of closed file"): + handle.read1(1) + with pytest.raises(ValueError, match="read of closed file"): + handle.readall() + with pytest.raises(ValueError, match="readinto of closed file"): + handle.readinto(bytearray(1)) + with pytest.raises(ValueError): + handle.read_range(0, 1) + with pytest.raises(ValueError): + handle.read_ranges([]) + + def test_fetch_blob_files_from_query_partial_read(): payload = _identifiable_payload(65536) table = _blob_table("query_partial", [{"id": 1, "image": payload}]) diff --git a/python/python/tests/test_catalog.py b/python/python/tests/test_catalog.py index e7b7b1ec2..b2036b7a6 100644 --- a/python/python/tests/test_catalog.py +++ b/python/python/tests/test_catalog.py @@ -263,3 +263,45 @@ async def test_database_names_invalid_limit(catalog_server, asynchronous, page_l next(names) assert names.num_page_results() == 0 assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_catalog_rejects_invalid_database_names_before_io( + catalog_server, asynchronous +): + endpoint, requests, _ = catalog_server + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + for name in [ + "", + ".", + "..", + "/db", + "db/", + "a//b", + "a/./b", + "a/../b", + "a b", + "a:b", + "a$b", + "a%b", + "a?b", + "a#b", + "a\\b", + "café", + "a\nb", + ]: + for method in [ + catalog.create_database, + catalog.connect_database, + catalog.drop_database, + ]: + with pytest.raises(ValueError, match="Invalid database name"): + result = method(name) + if asynchronous: + await result + assert requests == [] diff --git a/python/python/tests/test_e2e_remote_db.py b/python/python/tests/test_e2e_remote_db.py index d4364bc43..74d5cd7de 100644 --- a/python/python/tests/test_e2e_remote_db.py +++ b/python/python/tests/test_e2e_remote_db.py @@ -2,13 +2,78 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import os +from datetime import timedelta +from uuid import uuid4 + +import lancedb import numpy as np +import pyarrow as pa import pytest from lancedb import LanceDBConnection +from lancedb.index import FTS # TODO: setup integ test mark and script +@pytest.mark.skipif( + not os.environ.get("LANCEDB_URI"), reason="Set LANCEDB_URI to a JSON FTS server" +) +def test_json_fts_against_server(): + """Requires a writable database and capacity to run the distributed indexer. + + Configure LANCEDB_URI, LANCEDB_API_KEY, and optionally LANCEDB_REGION and + LANCEDB_HOST_OVERRIDE before running this test. + """ + uri = os.environ["LANCEDB_URI"] + assert uri.startswith("db://"), "LANCEDB_URI must identify a remote database" + conn = lancedb.connect( + uri, + api_key=os.environ.get("LANCEDB_API_KEY"), + region=os.environ.get("LANCEDB_REGION", "us-east-1"), + host_override=os.environ.get("LANCEDB_HOST_OVERRIDE"), + ) + name = f"json_fts_{uuid4().hex}" + table = conn.create_table( + name, + pa.table( + { + "id": [0, 1, 2], + "doc": pa.array( + [ + '{"meta":{"tag":"nature"},"author":"alice"}', + '{"meta":{"tag":"city"},"author":"nature"}', + '{"meta":{"tag":"nature"},"author":"bob"}', + ], + type=pa.json_(), + ), + "raw": pa.array([b"nature"] * 3, type=pa.large_binary()), + } + ), + ) + try: + table.create_index("doc", config=FTS(), name="doc_fts") + table.wait_for_index(["doc_fts"], timeout=timedelta(minutes=5)) + stats = table.index_stats("doc_fts") + assert stats.num_indexed_rows == 3 + assert stats.num_unindexed_rows == 0 + assert stats.index_type == "FTS" + for query, expected in [ + ("meta.tag,str,nature", [0, 2]), + ("meta.tag,str,alice", []), + ("meta.tag,str,tag", []), + ]: + search = table.search(query, query_type="fts", fts_columns="doc") + assert sorted(row["id"] for row in search.to_list()) == expected + plan = search.explain_plan(True) + assert "MatchQuery" in plan, plan + assert "FlatMatchQuery" not in plan, plan + with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"): + table.create_index("raw", config=FTS()) + finally: + conn.drop_table(name) + + @pytest.mark.skip(reason="Need to set up a local server") def test_against_local_server(): conn = LanceDBConnection("lancedb+http://localhost:10024") diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index c18893f1b..975310ed4 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1555,7 +1555,7 @@ def _mock_remote_function_catalog(): )["image"], "signature": body["signature"], "secret_bindings": body.get("secret_bindings", []), - "created_at": "2026-08-21T00:00:00Z", + "created_at_millis": 1787270400000, # 2026-08-21T00:00:00Z } response = {"job_id": "job-register"} status = 202 diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index e3963559d..d7fa2e0ec 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -45,6 +45,75 @@ from utils import exception_output TEST_LANGUAGE_MODEL_HOME = Path(__file__).parent / "models" +@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type") +@pytest.mark.parametrize("block_size", [128, 256]) +def test_json_fts(tmp_path, block_size): + db = ldb.connect(tmp_path) + data = pa.table( + { + "id": [0, 1, 2, 3], + "doc": pa.array( + [ + '{"title":"brown fox","meta":{"tag":"nature"},"author":"alice"}', + '{"title":"blue car","meta":{"tag":"city"}}', + '{"title":"brown bear","meta":{"tag":"nature"}}', + '{"title":"red fox","author":"alice"}', + ], + type=pa.json_(), + ), + } + ) + table = db.create_table("json_fts", data) + table.create_index("doc", config=FTS(block_size=block_size), name="doc_fts") + table = db.open_table("json_fts") + index = table.list_indices()[0] + assert index.name == "doc_fts" + assert index.index_type == "FTS" + assert index.columns == ["doc"] + stats = table.index_stats("doc_fts") + assert stats.num_indexed_rows == 4 + assert stats.num_unindexed_rows == 0 + + for query, expected in [ + ("title,str,brown", [0, 2]), + ("meta.tag,str,nature", [0, 2]), + ("author,str,alice", [0, 3]), + ("title,str,alice", []), + ("title,str,title", []), + ]: + search = table.search(query, query_type="fts", fts_columns="doc") + assert sorted(row["id"] for row in search.to_list()) == expected + plan = search.explain_plan(True) + assert "MatchQuery" in plan, plan + assert "FlatMatchQuery" not in plan, plan + + +@pytest.mark.skipif(not hasattr(pa, "json_"), reason="requires PyArrow JSON type") +@pytest.mark.asyncio +async def test_json_fts_async(tmp_path): + db = await ldb.connect_async(tmp_path) + table = await db.create_table( + "json_fts", + pa.table({"doc": pa.array(['{"meta":{"tag":"nature"}}'], type=pa.json_())}), + ) + await table.create_index("doc", config=FTS()) + search = table.query().nearest_to_text(MatchQuery("meta.tag,str,nature", "doc")) + assert len(await search.to_list()) == 1 + plan = await search.explain_plan(True) + assert "MatchQuery" in plan, plan + assert "FlatMatchQuery" not in plan, plan + + +@pytest.mark.parametrize("dtype", [pa.binary(), pa.large_binary()]) +def test_fts_rejects_raw_binary(tmp_path, dtype): + table = ldb.connect(tmp_path).create_table( + "binary_fts", pa.table({"doc": pa.array([b'{"title":"brown"}'], type=dtype)}) + ) + with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"): + table.create_index("doc", config=FTS()) + assert table.list_indices() == [] + + @pytest.fixture def table(tmp_path) -> ldb.table.LanceTable: # Use local random state to avoid affecting other tests diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index c5f58308d..8d670a136 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -355,6 +355,20 @@ def test_hybrid_query_minimum_nprobes_zero_raises(sync_table: Table): ) +def test_hybrid_query_nprobes_remains_independent(sync_table: Table): + query = sync_table.search(query_type="hybrid").nprobes(20) + + assert query._nprobes == 20 + assert query._minimum_nprobes is None + assert query._maximum_nprobes is None + + query = query.minimum_nprobes(5).nprobes(20) + + assert query._nprobes == 20 + assert query._minimum_nprobes == 5 + assert query._maximum_nprobes is None + + def test_hybrid_query_distance_range(sync_table: Table): reranker = RRFReranker(return_score="all") result = ( diff --git a/python/python/tests/test_namespace.py b/python/python/tests/test_namespace.py index 4a67ed31e..0ec509fbf 100644 --- a/python/python/tests/test_namespace.py +++ b/python/python/tests/test_namespace.py @@ -19,7 +19,8 @@ import lancedb from lance_namespace import connect as namespace_connect from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError from lancedb.listing import AsyncListing, Listing -from lancedb.namespace import _MAX_QUERY_K +from lancedb.namespace import _MAX_QUERY_K, _query_to_namespace_request +from lancedb.query import Query from lancedb.table import AsyncTable, LanceTable @@ -1213,6 +1214,29 @@ class TestPushdownOperations: ] assert all(r.k <= 2**31 - 1 for r in namespace_client.requests) + def test_probe_fields_pass_through_independently(self): + request = _query_to_namespace_request( + ["geneva", "hist"], + Query( + vector=[1.0, 2.0], + nprobes=20, + minimum_nprobes=3, + maximum_nprobes=10, + ), + ) + + assert request.nprobes == 20 + assert request.minimum_nprobes == 3 + assert request.maximum_nprobes == 10 + + request = _query_to_namespace_request( + ["geneva", "hist"], + Query(vector=[1.0, 2.0], nprobes=5, maximum_nprobes=0), + ) + assert request.nprobes == 5 + assert request.maximum_nprobes == 0 + assert request.to_dict()["maximum_nprobes"] == 0 + @pytest.mark.asyncio class TestAsyncPushdownOperations: @@ -1287,3 +1311,15 @@ def test_local_table_to_arrow_and_to_pandas_are_unchanged(tmp_path): assert table.to_arrow().column("id").to_pylist() == [1, 2] assert table.to_pandas()["id"].tolist() == [1, 2] + + +@pytest.mark.parametrize("effort", [0.0, 0.5, 1.0]) +def test_namespace_request_rejects_search_effort(effort): + from lancedb.query import Query + + with pytest.raises( + ValueError, match="namespace QueryTable does not support search_effort" + ): + _query_to_namespace_request( + ["test"], Query(vector=[1.0, 0.0], search_effort=effort) + ) diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index d3fedbcc8..a05247c5e 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -509,18 +509,30 @@ def test_vector_query_to_pandas_blob_mode_requires_native_path(tmp_db): ) -def test_vector_query_to_pandas_blob_descriptions_requires_plain_scan(tmp_db): +def test_vector_query_to_pandas_legacy_blob_descriptions_and_bytes(tmp_db): pytest.importorskip("lance") table = tmp_db.create_table( "test_vector_query_blob_descriptions", _blob_query_data(), storage_options=LEGACY_BLOB_STORAGE_OPTIONS, ) + query = table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1) - with pytest.raises(RuntimeError, match="plain scan query"): - table.search([1.0, 0.0]).select(["blob", "vector"]).limit(1).to_pandas( - blob_mode="descriptions" - ) + df = query.to_pandas(blob_mode="descriptions") + assert df["blob"].tolist() == [{"position": 0, "size": 3}] + + with pytest.raises(ValueError, match="legacy blob columns"): + query.to_pandas(blob_mode="bytes") + with pytest.raises(ValueError, match="legacy blob columns"): + query.to_arrow(blob_mode="bytes") + with pytest.raises(ValueError, match="legacy blob columns"): + table.to_arrow(blob_mode="bytes") + assert table.to_pandas(blob_mode="bytes")["blob"].tolist() == [ + b"one", + b"two", + b"three", + b"four", + ] def test_order_by_plain_query(mem_db): @@ -995,12 +1007,22 @@ def test_invalid_nprobes_sync(table): ValueError, match="maximum_nprobes must be greater than or equal to minimum_nprobes", ): - LanceVectorQueryBuilder(table, [0, 0], "vector").maximum_nprobes(5).to_list() + ( + LanceVectorQueryBuilder(table, [0, 0], "vector") + .minimum_nprobes(10) + .maximum_nprobes(5) + .to_list() + ) with pytest.raises( ValueError, - match="minimum_nprobes must be less than or equal to maximum_nprobes", + match="maximum_nprobes must be greater than or equal to minimum_nprobes", ): - LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(100).to_list() + ( + LanceVectorQueryBuilder(table, [0, 0], "vector") + .maximum_nprobes(5) + .minimum_nprobes(100) + .to_list() + ) def test_nprobes_works_sync(table): @@ -1039,12 +1061,22 @@ async def test_invalid_nprobes_async(table_async: AsyncTable): ValueError, match="maximum_nprobes must be greater than or equal to minimum_nprobes", ): - await table_async.vector_search([0, 0]).maximum_nprobes(5).to_list() + await ( + table_async.vector_search([0, 0]) + .minimum_nprobes(10) + .maximum_nprobes(5) + .to_list() + ) with pytest.raises( ValueError, match="minimum_nprobes must be less than or equal to maximum_nprobes", ): - await table_async.vector_search([0, 0]).minimum_nprobes(100).to_list() + await ( + table_async.vector_search([0, 0]) + .maximum_nprobes(5) + .minimum_nprobes(100) + .to_list() + ) def test_query_builder_with_prefilter(table): @@ -1635,8 +1667,9 @@ def test_query_serialization_sync(table: lancedb.table.Table): q, vector_column="vector", vector=[5.0, 6.0], - minimum_nprobes=10, - maximum_nprobes=10, + nprobes=10, + minimum_nprobes=None, + maximum_nprobes=None, refine_factor=5, ) @@ -1654,8 +1687,19 @@ def test_query_serialization_sync(table: lancedb.table.Table): q, vector_column="vector", vector=[5.0, 6.0], - minimum_nprobes=50, - maximum_nprobes=50, + nprobes=50, + minimum_nprobes=None, + maximum_nprobes=None, + ) + + q = table.search([5.0, 6.0]).minimum_nprobes(5).nprobes(50).to_query_object() + check_set_props( + q, + vector_column="vector", + vector=[5.0, 6.0], + nprobes=50, + minimum_nprobes=5, + maximum_nprobes=None, ) q = table.search([5.0, 6.0]).maximum_nprobes(10).to_query_object() @@ -1667,6 +1711,15 @@ def test_query_serialization_sync(table: lancedb.table.Table): minimum_nprobes=None, ) + q = table.search([5.0, 6.0]).nprobes(5).maximum_nprobes(0).to_query_object() + check_set_props( + q, + vector_column="vector", + vector=[5.0, 6.0], + nprobes=5, + maximum_nprobes=0, + ) + q = table.search([5.0, 6.0]).distance_range(0.0, 1.0).to_query_object() check_set_props( q, vector_column="vector", vector=[5.0, 6.0], lower_bound=0.0, upper_bound=1.0 @@ -1716,8 +1769,8 @@ async def test_query_serialization_async(table_async: AsyncTable): limit=10, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, ) @@ -1727,8 +1780,8 @@ async def test_query_serialization_async(table_async: AsyncTable): q, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1739,8 +1792,9 @@ async def test_query_serialization_async(table_async: AsyncTable): q, vector=sample_vector, postfilter=False, - minimum_nprobes=50, - maximum_nprobes=50, + nprobes=50, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1759,8 +1813,8 @@ async def test_query_serialization_async(table_async: AsyncTable): filter="id = 1", postfilter=True, vector=sample_vector, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, ) @@ -1774,8 +1828,9 @@ async def test_query_serialization_async(table_async: AsyncTable): check_set_props( q, vector=sample_vector, - minimum_nprobes=10, - maximum_nprobes=10, + nprobes=10, + minimum_nprobes=None, + maximum_nprobes=None, refine_factor=5, postfilter=False, with_row_id=False, @@ -1788,7 +1843,24 @@ async def test_query_serialization_async(table_async: AsyncTable): q, vector=sample_vector, minimum_nprobes=5, - maximum_nprobes=20, + maximum_nprobes=None, + postfilter=False, + with_row_id=False, + bypass_vector_index=False, + limit=10, + ) + + q = ( + (await table_async.search([5.0, 6.0])) + .nprobes(5) + .maximum_nprobes(0) + .to_query_object() + ) + check_set_props( + q, + vector=sample_vector, + nprobes=5, + maximum_nprobes=0, postfilter=False, with_row_id=False, bypass_vector_index=False, @@ -1806,8 +1878,8 @@ async def test_query_serialization_async(table_async: AsyncTable): lower_bound=0.0, upper_bound=1.0, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1819,8 +1891,8 @@ async def test_query_serialization_async(table_async: AsyncTable): distance_type="cosine", vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1832,8 +1904,8 @@ async def test_query_serialization_async(table_async: AsyncTable): ef=7, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, bypass_vector_index=False, limit=10, @@ -1845,8 +1917,8 @@ async def test_query_serialization_async(table_async: AsyncTable): bypass_vector_index=True, vector=sample_vector, postfilter=False, - minimum_nprobes=20, - maximum_nprobes=20, + minimum_nprobes=None, + maximum_nprobes=None, with_row_id=False, limit=10, ) @@ -2205,3 +2277,127 @@ def test_blob_v2_to_batches_row_id(tmp_db): assert "_rowid" in hits.column_names blobs = table.fetch_blobs("blob", hits) assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"] + + +@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"]) +@pytest.mark.parametrize( + "effort,probes", [(None, 4), (0, 1), (0.25, 2), (0.5, 4), (0.75, 8), (1, 16)] +) +def test_search_effort_probe_budget(tmp_path, monkeypatch, metric, effort, probes): + import re + from lancedb.index import IvfFlat + + monkeypatch.setenv("LANCE_AUTO_PROBE_MARGIN", "0") + monkeypatch.setenv("LANCE_AUTO_MIN_INITIAL_NPROBES", "4") + monkeypatch.setenv("LANCE_AUTO_MAX_INITIAL_NPROBES", "4") + rng = np.random.default_rng(42) + vectors = np.repeat(np.eye(16, dtype=np.float32), 32, axis=0) + vectors += rng.normal(0, 0.001, vectors.shape).astype(np.float32) + data = pa.table( + { + "id": np.arange(512), + "text": ["dog"] * 512, + "vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 16), + } + ) + table = lancedb.connect(tmp_path).create_table("effort", data.slice(0, 256)) + table.add(data.slice(256)) + table.create_index( + "vector", config=IvfFlat(distance_type=metric, num_partitions=16) + ) + vector = np.eye(16, dtype=np.float32)[0] + + def search(value=effort): + return table.search(vector).distance_type(metric).search_effort(value).limit(8) + + def partitions(query): + vector_plan = query.analyze_plan().split("FTS Search Plan:", 1)[0] + values = re.findall(r"partitions_searched=(\d+)", vector_plan) + assert values + return sum(map(int, values)) + + assert partitions(search()) == probes + exact = set( + table.search(vector) + .distance_type(metric) + .bypass_vector_index() + .limit(8) + .to_arrow()["id"] + .to_pylist() + ) + actual = set(search().to_arrow()["id"].to_pylist()) + assert len(actual & exact) / 8 >= 0.5 + expected = 1 if effort is None else probes + assert partitions(search().nprobes(1)) == expected + assert ( + partitions( + table.search(vector) + .distance_type(metric) + .nprobes(1) + .search_effort(effort) + .limit(8) + ) + == expected + ) + assert partitions(search().nprobes(1).search_effort(None)) == 1 + assert partitions(search().maximum_nprobes(2)) <= 2 + assert partitions(search().minimum_nprobes(3).maximum_nprobes(3)) == 3 + assert partitions(search(0).nprobes(16).maximum_nprobes(0)) == 1 + if effort is None: + assert search().to_arrow() == search(0.5).to_arrow() + filtered = search(0).maximum_nprobes(2).where("id % 8 = 0", prefilter=True) + assert partitions(filtered) <= 2 + assert all(row["id"] % 8 == 0 for row in filtered.to_list()) + + if metric == "l2" and effort == 0.5: + table.create_index("text", config=FTS()) + hybrid = ( + table.search(query_type="hybrid") + .vector(vector) + .text("dog") + .nprobes(1) + .search_effort(0.5) + .limit(8) + ) + assert partitions(hybrid) == 4 + assert partitions(hybrid.search_effort(None)) == 1 + + +@pytest.mark.parametrize( + "effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")] +) +def test_search_effort_invalid_sync(table, effort): + with pytest.raises(ValueError, match="search_effort must be finite and in"): + table.search([0, 0]).search_effort(effort).to_list() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hybrid", [False, True]) +@pytest.mark.parametrize("effort", [None, 0, 0.25, 0.5, 0.75, 1]) +async def test_search_effort_async_roundtrip(table_async, hybrid, effort): + query = table_async.vector_search([0, 0]) + if hybrid: + query = query.nearest_to_text("dog") + query = query.nprobes(5).search_effort(effort).minimum_nprobes(2).maximum_nprobes(3) + request = query.to_query_object() + assert request.search_effort == effort + assert request.nprobes == 5 + assert request.minimum_nprobes == 2 + assert request.maximum_nprobes == 3 + assert query.search_effort(None).to_query_object().search_effort is None + query.search_effort(effort) + if not hybrid: + assert len(await query.limit(1).to_list()) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hybrid", [False, True]) +@pytest.mark.parametrize( + "effort", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")] +) +async def test_search_effort_invalid_async(table_async, hybrid, effort): + query = table_async.vector_search([0, 0]) + if hybrid: + query = query.nearest_to_text("dog") + with pytest.raises(ValueError, match="search_effort must be finite and in"): + query.search_effort(effort) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 00e575e46..f43972a22 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -4,6 +4,7 @@ from concurrent.futures import ThreadPoolExecutor import contextlib from datetime import timedelta import http.server +import io import json import multiprocessing as mp import pickle @@ -1410,6 +1411,72 @@ def test_remote_job_wait_raises_on_failure(): job.wait() +@pytest.mark.parametrize("server_version", ["0.2.0", "0.6.0"]) +def test_remote_json_fts_config(server_version): + from lancedb.index import FTS + + received_requests = [] + + def handler(request): + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.send_header("phalanx-version", server_version) + request.end_headers() + if request.path == "/v1/table/test/describe/": + fields = [ + { + "name": "doc", + "type": {"type": "large_binary"}, + "nullable": True, + "metadata": {"ARROW:extension:name": "lance.json"}, + }, + { + "name": "raw", + "type": {"type": "large_binary"}, + "nullable": True, + }, + ] + request.wfile.write( + json.dumps({"version": 1, "schema": {"fields": fields}}).encode() + ) + elif request.path == "/v1/table/test/create_index/": + received_requests.append(read_json_body(request)) + request.wfile.write(b"{}") + else: + raise AssertionError(f"Unexpected request: {request.path}") + + with mock_lancedb_connection(handler) as db: + table = db.open_table("test") + table.create_index( + "doc", + config=FTS( + with_position=True, + base_tokenizer="whitespace", + stem=False, + remove_stop_words=False, + lower_case=False, + block_size=256, + ), + name="doc_fts", + replace=False, + ) + with pytest.raises(ValueError, match="FTS.*lance.json.*raw binary"): + table.create_index("raw", config=FTS()) + + assert len(received_requests) == 1 + request = received_requests[0] + assert request["column"] == "doc" + assert request["index_type"] == "FTS" + assert request["name"] == "doc_fts" + assert request["replace"] is False + assert request["with_position"] is True + assert request["base_tokenizer"] == "whitespace" + assert request["stem"] is False + assert request["remove_stop_words"] is False + assert request["lower_case"] is False + assert request["block_size"] == 256 + + def test_remote_create_index_new_api(caplog): received_requests = [] @@ -2098,9 +2165,6 @@ def test_query_sync_minimal(): "ef": None, "vector": [1.0, 2.0, 3.0], "vector_column": "vector", - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "version": None, } @@ -2272,8 +2336,6 @@ def test_query_sync_maximal(): "refine_factor": 10, "vector": [1.0, 2.0, 3.0], "nprobes": 5, - "minimum_nprobes": 5, - "maximum_nprobes": 5, "lower_bound": None, "upper_bound": None, "ef": None, @@ -2369,7 +2431,7 @@ async def test_query_async_nprobes_zero(hybrid): assert requests == ["/v1/table/test/describe/"] -def test_query_sync_nprobes(): +def test_query_sync_probe_bounds(): def handler(body): assert body == { "k": 10, @@ -2381,7 +2443,6 @@ def test_query_sync_nprobes(): "upper_bound": None, "ef": None, "vector": [1.0, 2.0, 3.0], - "nprobes": 5, "minimum_nprobes": 5, "maximum_nprobes": 15, "version": None, @@ -2398,7 +2459,28 @@ def test_query_sync_nprobes(): ) -def test_query_sync_no_max_nprobes(): +def test_query_sync_nprobes_passthrough(): + def handler(body): + assert body["nprobes"] == 20 + assert "minimum_nprobes" not in body + assert "maximum_nprobes" not in body + return pa.table({"id": [1]}) + + with query_test_table(handler) as table: + table.search([1, 2, 3]).nprobes(20).to_list() + + +def test_query_sync_nprobes_with_unbounded_maximum(): + def handler(body): + assert body["nprobes"] == 20 + assert body["maximum_nprobes"] == 0 + return pa.table({"id": [1]}) + + with query_test_table(handler) as table: + table.search([1, 2, 3]).nprobes(20).maximum_nprobes(0).to_list() + + +def test_query_sync_minimum_nprobes_only(): def handler(body): assert body == { "k": 10, @@ -2410,9 +2492,7 @@ def test_query_sync_no_max_nprobes(): "upper_bound": None, "ef": None, "vector": [1.0, 2.0, 3.0], - "nprobes": 5, "minimum_nprobes": 5, - "maximum_nprobes": 0, "version": None, } @@ -2422,7 +2502,6 @@ def test_query_sync_no_max_nprobes(): ( table.search([1, 2, 3], vector_column_name="vector2", fast_search=True) .minimum_nprobes(5) - .maximum_nprobes(0) .to_list() ) @@ -2577,9 +2656,6 @@ def test_query_sync_hybrid(): "refine_factor": None, "vector": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], "vector_column": "vector", - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "lower_bound": None, "upper_bound": None, "ef": None, @@ -2653,6 +2729,11 @@ def test_create_client(): assert isinstance(db.client_config, ClientConfig) assert db.client_config.retry_config.retries == 42 + db = lancedb.connect( + **mandatory_args, client_config={"blob_request_concurrency": 3} + ) + assert db.client_config.blob_request_concurrency == 3 + with pytest.warns(DeprecationWarning): db = lancedb.connect(**mandatory_args, connection_timeout=42) assert db.client_config.timeout_config.connect_timeout == timedelta(seconds=42) @@ -3388,6 +3469,45 @@ async def test_async_remote_blob_pandas_rejects_unpinned_fetch(monkeypatch, blob await table.to_pandas(**pandas_kwargs) +@pytest.mark.parametrize("method", ["to_arrow", "to_list"]) +def test_remote_search_bytes_refused_before_query(monkeypatch, method): + with blob_remote_table() as table: + query = table.search().select(["id", "image"]) + descriptions = getattr(query, method)(blob_mode="descriptions") + images = ( + descriptions["image"].to_pylist() + if method == "to_arrow" + else [row["image"] for row in descriptions] + ) + assert "_lance_row_id" not in images[0] + + def fail_query(*args, **kwargs): + raise AssertionError("bytes must be refused before the query runs") + + monkeypatch.setattr(table, "_execute_query", fail_query) + with pytest.raises(NotImplementedError, match="stable table snapshot"): + getattr(query, method)(blob_mode="bytes") + with pytest.raises(ValueError, match="blob_mode must be one of"): + table.to_pandas(blob_mode="handles") + + +@pytest.mark.asyncio +async def test_async_remote_to_arrow_bytes_refused_before_query(monkeypatch): + with blob_remote_table() as remote_table: + table = remote_table._table + + async def fail_fetch(*args, **kwargs): + raise AssertionError("blob fetch must not use an unpinned remote table") + + monkeypatch.setattr(table, "fetch_blobs", fail_fetch) + query = table.query().select(["id", "image"]) + monkeypatch.setattr(query, "to_batches", fail_fetch) + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await query.to_arrow(blob_mode="bytes") + with pytest.raises(NotImplementedError, match="stable table snapshot"): + await table.to_arrow(blob_mode="bytes") + + @pytest.mark.asyncio @pytest.mark.parametrize("blob_mode", ["bytes", "lazy"]) async def test_async_remote_nonblob_projection_to_pandas(blob_mode): @@ -3423,6 +3543,28 @@ def test_remote_blob_files_are_lazy_seekable_handles(): alpha.read_range(0, 1) +def test_remote_blob_file_matches_file_error_and_range_contract(): + with blob_remote_table() as table: + handle = table.fetch_blob_files("image", [10])[0] + handle.seek(2) + + with pytest.raises(ValueError, match="negative seek value -1"): + handle.seek(-3, io.SEEK_CUR) + assert handle.tell() == 2 + assert handle.read(None) == b"pha" + + handle.seek(2) + assert handle.read1(2) == b"ph" + assert handle.read_ranges([(0, 2), (3, 2)]) == [b"al", b"ha"] + assert handle.tell() == 4 + with pytest.raises(ValueError): + handle.read_range(4, 2) + + handle.close() + with pytest.raises(ValueError, match="read of closed file"): + handle.read(1) + + def test_remote_blob_fetch_accepts_query_table(): hits = pa.table({"_rowid": pa.array([10, 20, 30], type=pa.uint64())}) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index bae6e1347..99cf2af54 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -13,7 +13,7 @@ import warnings import weakref from concurrent.futures import CancelledError, ThreadPoolExecutor from datetime import date, datetime, timedelta -from decimal import Decimal +from decimal import Decimal, localcontext from time import sleep from typing import List from unittest.mock import patch @@ -2867,6 +2867,51 @@ def test_update_with_arrow_scalar(mem_db: DBConnection): assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]] +@pytest.mark.parametrize("use_arrow_scalar", [False, True]) +@pytest.mark.parametrize( + "decimal_type, value", + [ + (pa.decimal128(10, 2), "9.99"), + (pa.decimal128(10, 2), "-9.99"), + (pa.decimal128(10, 2), "0.00"), + (pa.decimal128(10, 2), "1.23E+5"), + (pa.decimal128(38, 18), "1E-18"), + (pa.decimal128(38, 18), "12345678901234567.123456789012345678"), + (pa.decimal128(38, 18), "12345678901234567890.123456789012345678"), + (pa.decimal128(38, 18), "-12345678901234567890.123456789012345678"), + ( + pa.decimal256(76, 38), + "12345678901234567890123456789012345678." + "12345678901234567890123456789012345678", + ), + ( + pa.decimal256(76, 38), + "-12345678901234567890123456789012345678." + "12345678901234567890123456789012345678", + ), + ], +) +def test_update_decimal(mem_db: DBConnection, decimal_type, value, use_arrow_scalar): + expected = Decimal(value) + table = mem_db.create_table( + "my_table", + pa.table( + { + "id": [1, 2], + "price": pa.array([Decimal("1"), expected], type=decimal_type), + } + ), + ) + update = table.to_arrow()["price"][1] if use_arrow_scalar else expected + with localcontext() as context: + context.prec = 6 + result = table.update(where="id = 1", values={"price": update}) + assert result.rows_updated == 1 + actual = table.to_arrow() + assert actual.schema.field("price").type == decimal_type + assert actual["price"].to_pylist() == [expected, expected] + + def test_update_types(mem_db: DBConnection): table = mem_db.create_table( "my_table", diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index d9c574eb9..9deb0c073 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -4,6 +4,7 @@ import os import pathlib +from decimal import Decimal, localcontext from typing import Optional import lance @@ -157,6 +158,41 @@ def test_value_to_sql_string(tmp_path): assert table.to_pandas().query("search == @value")["replace"].item() == value +@pytest.mark.parametrize( + "value, expected", + [ + ("9.99", "'9.99'"), + ("-9.99", "'-9.99'"), + ("0", "'0'"), + ("-0.00", "'-0.00'"), + ("1.23E+5", "'123000'"), + ("1E-18", "'0.000000000000000001'"), + ( + "12345678901234567890.123456789012345678", + "'12345678901234567890.123456789012345678'", + ), + ( + "12345678901234567890123456789012345678." + "12345678901234567890123456789012345678", + "'12345678901234567890123456789012345678." + "12345678901234567890123456789012345678'", + ), + ], +) +def test_value_to_sql_decimal(value, expected): + with localcontext() as context: + context.prec = 6 + assert value_to_sql(Decimal(value)) == expected + + +@pytest.mark.parametrize("value", ["NaN", "sNaN", "NaN123456", "Infinity", "-Infinity"]) +def test_value_to_sql_decimal_non_finite(value): + with pytest.raises( + ValueError, match="^Non-finite Decimal values cannot be converted to SQL$" + ): + value_to_sql(Decimal(value)) + + def test_value_to_sql_dict(): # Simple flat struct assert value_to_sql({"a": 1, "b": "hello"}) == "named_struct('a', 1, 'b', 'hello')" @@ -951,6 +987,39 @@ def test_cast_to_target_schema_coerces_binary_to_blob_v2(): ] +@pytest.mark.parametrize( + "string_type", [pa.string(), pa.large_string(), pa.string_view()] +) +def test_cast_to_target_schema_coerces_string_to_blob_v2_uri(string_type): + data = pa.table({"image": pa.array(["file:///payload", None], type=string_type)}) + target = pa.schema([lancedb.blob("image")]) + + output = _cast_to_target_schema(data.to_reader(), target).read_all() + + image = output["image"].chunk(0) + assert type(image.type) is lancedb.BlobType + assert image.storage.to_pylist() == [ + {"data": None, "uri": "file:///payload", "position": None, "size": None}, + None, + ] + + +@pytest.mark.parametrize( + "source_type, value", + [ + (pa.binary(4), b"abcd"), + (pa.list_(pa.uint8()), [1, 2]), + (pa.int64(), 42), + ], +) +def test_cast_to_target_schema_rejects_unsupported_blob_input(source_type, value): + data = pa.table({"image": pa.array([value], type=source_type)}) + target = pa.schema([lancedb.blob("image")]) + + with pytest.raises(ValueError, match="cannot coerce column 'image' with type"): + _cast_to_target_schema(data.to_reader(), target) + + def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct(): storage = lancedb.blob("image").type.storage_type target = pa.schema( diff --git a/python/python/tests/test_wal_hybrid.py b/python/python/tests/test_wal_hybrid.py index 11a9755e0..9c7efe151 100644 --- a/python/python/tests/test_wal_hybrid.py +++ b/python/python/tests/test_wal_hybrid.py @@ -119,7 +119,9 @@ def test_wal_hybrid_fuses_on_the_key_and_drops_it_again(): limit=10, with_row_ids=True, ) - results = builder._finish_hybrid_results(combined) + results = builder._finish_hybrid_results( + combined, blob_mode="lazy", blob_sources={} + ) assert "_rowid" not in results.column_names, "the surrogate is internal" assert "id" not in results.column_names, "the caller selected only `text`" diff --git a/python/src/authz.rs b/python/src/authz.rs new file mode 100644 index 000000000..6ad3e69b1 --- /dev/null +++ b/python/src/authz.rs @@ -0,0 +1,280 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use crate::error::PythonErrorExt; +use crate::runtime::future_into_py; +use lancedb::authz::*; +use pyo3::exceptions::{PyRuntimeError, PyValueError}; +use pyo3::types::{PyAnyMethods, PyBytes, PyString}; +use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods}; +use std::sync::Arc; + +#[pyclass] +pub struct AuthorizationClient { + pub(crate) inner: Arc, +} + +#[pymethods] +impl AuthorizationClient { + fn list_acls<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult> { + let request: ListAccessControlEntriesRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let result = inner.list_acls(request).await.infer_error()?; + serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string())) + }) + } + fn add_acl<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult> { + let request: AccessControlEntryRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.add_acl(request).await.infer_error() + }) + } + fn delete_acl<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult> { + let request: AccessControlEntryRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.delete_acl(request).await.infer_error() + }) + } + fn list_role_bindings<'py>( + self_: PyRef<'py, Self>, + request: &str, + ) -> PyResult> { + let request: ListRoleBindingsRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + request.validate().infer_error()?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let result = inner.list_role_bindings(request).await.infer_error()?; + serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string())) + }) + } + fn add_role_binding<'py>( + self_: PyRef<'py, Self>, + request: &str, + ) -> PyResult> { + let request: RoleBindingRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + request.validate().infer_error()?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.add_role_binding(request).await.infer_error() + }) + } + fn delete_role_binding<'py>( + self_: PyRef<'py, Self>, + request: &str, + ) -> PyResult> { + let request: RoleBindingRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + request.validate().infer_error()?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + inner.delete_role_binding(request).await.infer_error() + }) + } + fn get_principal<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult> { + let request: GetPrincipalRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + request.validate().infer_error()?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let result = inner.get_principal(request).await.infer_error()?; + serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string())) + }) + } + fn get_group<'py>(self_: PyRef<'py, Self>, request: &str) -> PyResult> { + let request: GetGroupRequest = serde_json::from_str(request) + .map_err(|_| PyValueError::new_err("Invalid authorization request"))?; + request.validate().infer_error()?; + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let result = inner.get_group(request).await.infer_error()?; + serde_json::to_string(&result).map_err(|err| PyRuntimeError::new_err(err.to_string())) + }) + } +} + +/// The shared privilege parser and canonical wire names. +#[pyclass(frozen)] +pub struct AuthzPrivilege { + inner: Privilege, +} + +#[pymethods] +impl AuthzPrivilege { + #[new] + fn new(value: &str) -> PyResult { + Ok(Self { + inner: value.parse().map_err(PyValueError::new_err)?, + }) + } + + #[staticmethod] + fn known() -> Vec { + Privilege::KNOWN.iter().map(ToString::to_string).collect() + } + + fn __str__(&self) -> String { + self.inner.to_string() + } +} + +#[pyclass(frozen)] +pub struct AuthzSubject { + inner: Subject, +} + +#[pymethods] +impl AuthzSubject { + #[new] + fn new(kind: &str, value: &Bound<'_, PyAny>) -> PyResult { + let value: String = value + .extract() + .map_err(|_| PyValueError::new_err("Subject value must be a nonempty string"))?; + Ok(Self { + inner: Subject::new(kind.parse().infer_error()?, value).infer_error()?, + }) + } + + #[staticmethod] + fn from_canonical(value: &str) -> PyResult { + Ok(Self { + inner: Subject::from_canonical(value).infer_error()?, + }) + } + + #[getter] + fn kind(&self) -> &str { + self.inner.kind().as_str() + } + #[getter] + fn value(&self) -> &str { + self.inner.value() + } + #[getter] + fn display_value(&self) -> &str { + self.inner.display_value() + } + + #[pyo3(signature = (context=None))] + fn to_wire(&self, context: Option<&str>) -> PyResult { + match context { + None => {} + Some("principal") => self.inner.validate_principal().infer_error()?, + Some("group") => self.inner.validate_group().infer_error()?, + Some("role_member") => self.inner.validate_role_member().infer_error()?, + Some(_) => return Err(PyValueError::new_err("Invalid subject context")), + } + Ok(self.inner.to_wire()) + } + + fn __repr__(&self) -> String { + format!("{:?}", self.inner) + } + fn __str__(&self) -> String { + self.inner.to_string() + } +} + +#[pyclass(frozen)] +pub struct AuthzObject { + inner: Object, +} + +impl AuthzObject { + fn validated(inner: Object) -> PyResult { + inner.validate().map_err(PyValueError::new_err)?; + Ok(Self { inner }) + } +} + +fn namespace_path(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_instance_of::() || value.is_instance_of::() { + return Err(PyValueError::new_err( + "namespace must be a sequence of components", + )); + } + let components = value + .extract() + .map_err(|_| PyValueError::new_err("namespace must be a sequence of string components"))?; + NamespacePath::new(components).map_err(PyValueError::new_err) +} + +#[pymethods] +impl AuthzObject { + #[new] + fn new(value: &str) -> PyResult { + Ok(Self { + inner: value.parse().infer_error()?, + }) + } + + #[staticmethod] + fn system() -> Self { + Self { + inner: Object::System, + } + } + + #[staticmethod] + fn database(name: &str) -> PyResult { + Self::validated(Object::Database(DatabaseObject::new(name.into()))) + } + + #[staticmethod] + fn namespace(database: &str, namespace: &Bound<'_, PyAny>) -> PyResult { + Self::validated(Object::Namespace(NamespaceObject { + database: database.into(), + namespace: namespace_path(namespace)?, + })) + } + + #[staticmethod] + fn table(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult { + Self::validated(Object::Table(TableObject { + database: database.into(), + namespace: namespace_path(namespace)?, + table: name.into(), + })) + } + + #[staticmethod] + fn view(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult { + Self::validated(Object::View(ViewObject { + database: database.into(), + namespace: namespace_path(namespace)?, + view: name.into(), + })) + } + + #[staticmethod] + fn secret(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult { + Self::validated(Object::Secret(SecretObject { + database: database.into(), + namespace: namespace_path(namespace)?, + secret: name.into(), + })) + } + + #[staticmethod] + fn function(database: &str, namespace: &Bound<'_, PyAny>, name: &str) -> PyResult { + Self::validated(Object::Function(FunctionObject { + database: database.into(), + namespace: namespace_path(namespace)?, + function: name.into(), + })) + } + + fn __str__(&self) -> String { + self.inner.to_string() + } + fn __repr__(&self) -> String { + format!("{:?}", self.inner) + } +} diff --git a/python/src/catalog.rs b/python/src/catalog.rs index e7cc56c49..9e375b314 100644 --- a/python/src/catalog.rs +++ b/python/src/catalog.rs @@ -21,6 +21,13 @@ pub struct Catalog { #[pymethods] impl Catalog { + #[getter] + fn authz(&self) -> PyResult { + Ok(crate::authz::AuthorizationClient { + inner: self.inner.authz().infer_error()?, + }) + } + #[getter] fn uri(&self) -> &str { self.inner.uri() diff --git a/python/src/connection.rs b/python/src/connection.rs index ec0768eb2..9ac081b9d 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -1220,6 +1220,7 @@ pub struct PyClientConfig { user_agent: String, retry_config: Option, timeout_config: Option, + blob_request_concurrency: Option, extra_headers: Option>, id_delimiter: Option, tls_config: Option, @@ -1305,6 +1306,7 @@ impl From for lancedb::remote::ClientConfig { user_agent: value.user_agent, retry_config: value.retry_config.map(Into::into).unwrap_or_default(), timeout_config: value.timeout_config.map(Into::into).unwrap_or_default(), + blob_request_concurrency: value.blob_request_concurrency, extra_headers: value.extra_headers.unwrap_or_default(), id_delimiter: value.id_delimiter, tls_config: value.tls_config.map(Into::into), diff --git a/python/src/lib.rs b/python/src/lib.rs index 9f88030ca..066e85927 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -22,6 +22,7 @@ use table::{ }; pub mod arrow; +pub mod authz; pub mod catalog; pub mod connection; pub mod error; @@ -70,6 +71,10 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/query.rs b/python/src/query.rs index 1079fb314..e5547881a 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -71,7 +71,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB { match class_name(&ob)?.as_str() { "MatchQuery" => { let query = ob.getattr("query")?.extract()?; - let column = ob.getattr("column")?.extract()?; + let column = ob.getattr("column")?.extract::>()?; let boost = ob.getattr("boost")?.extract()?; let fuzziness = ob.getattr("fuzziness")?.extract()?; let max_expansions = ob.getattr("max_expansions")?.extract()?; @@ -88,7 +88,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB { let mut query = MatchQuery::new(query) - .with_column(Some(column)) + .with_column(column) .with_boost(boost) .with_fuzziness(fuzziness) .with_max_expansions(max_expansions) @@ -199,7 +199,7 @@ impl<'py> IntoPyObject<'py> for PyLanceDB { } namespace .getattr(intern!(py, "MatchQuery"))? - .call((query.terms, query.column.unwrap()), Some(&kwargs)) + .call((query.terms, query.column), Some(&kwargs)) } FtsQuery::Phrase(query) => { let kwargs = PyDict::new(py); @@ -336,10 +336,10 @@ pub struct PyQueryRequest { pub use_lsm: Option, pub column: Option, pub query_vector: Option, + pub nprobes: Option, pub minimum_nprobes: Option, - // None means user did not set it and default should be used (currently 20) - // Some(0) means user set it to None and there is no limit pub maximum_nprobes: Option, + pub search_effort: Option, pub lower_bound: Option, pub upper_bound: Option, pub ef: Option, @@ -369,8 +369,10 @@ impl From for PyQueryRequest { use_lsm: query_request.use_lsm, column: None, query_vector: None, + nprobes: None, minimum_nprobes: None, maximum_nprobes: None, + search_effort: None, lower_bound: None, upper_bound: None, ef: None, @@ -396,11 +398,10 @@ impl From for PyQueryRequest { use_lsm: vector_query.base.use_lsm, column: vector_query.column, query_vector: Some(PyQueryVectors(vector_query.query_vector)), - minimum_nprobes: Some(vector_query.minimum_nprobes), - maximum_nprobes: match vector_query.maximum_nprobes { - None => Some(0), - Some(value) => Some(value), - }, + nprobes: vector_query.nprobes, + minimum_nprobes: vector_query.minimum_nprobes, + maximum_nprobes: vector_query.maximum_nprobes, + search_effort: vector_query.search_effort, lower_bound: vector_query.lower_bound, upper_bound: vector_query.upper_bound, ef: vector_query.ef, @@ -1000,6 +1001,15 @@ impl VectorQuery { Ok(()) } + pub fn search_effort(&mut self, search_effort: Option) -> PyResult<()> { + self.inner = self + .inner + .clone() + .search_effort(search_effort) + .infer_error()?; + Ok(()) + } + pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> { self.inner = self .inner @@ -1193,6 +1203,18 @@ impl HybridQuery { self.inner_vec.nprobes(nprobe) } + pub fn search_effort(&mut self, search_effort: Option) -> PyResult<()> { + self.inner_vec.search_effort(search_effort) + } + + pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> { + self.inner_vec.minimum_nprobes(minimum_nprobes) + } + + pub fn maximum_nprobes(&mut self, maximum_nprobes: u32) -> PyResult<()> { + self.inner_vec.maximum_nprobes(maximum_nprobes) + } + pub fn ef(&mut self, ef: u32) { self.inner_vec.ef(ef); } @@ -1237,6 +1259,10 @@ impl HybridQuery { req.query_vector = vec_req.query_vector; req.column = vec_req.column; req.distance_type = vec_req.distance_type; + req.nprobes = vec_req.nprobes; + req.minimum_nprobes = vec_req.minimum_nprobes; + req.maximum_nprobes = vec_req.maximum_nprobes; + req.search_effort = vec_req.search_effort; req.ef = vec_req.ef; req.refine_factor = vec_req.refine_factor; req.lower_bound = vec_req.lower_bound; diff --git a/python/src/table.rs b/python/src/table.rs index d82ebbaaf..23b48e3fd 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -697,6 +697,14 @@ impl PyBlobFile { } } +fn blob_file_error(operation: &str, error: lance_core::Error) -> pyo3::PyErr { + let message = format!("blob {operation} failed: {error}"); + match error { + lance_core::Error::InvalidInput { .. } => PyValueError::new_err(message), + _ => PyRuntimeError::new_err(message), + } +} + #[pymethods] impl PyBlobFile { fn read_bytes(self_: PyRef<'_, Self>) -> PyResult> { @@ -705,7 +713,7 @@ impl PyBlobFile { let py = self_.py(); let bytes = py .detach(move || block_on(async move { inner.read().await })) - .map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?; + .map_err(|e| blob_file_error("read", e))?; Ok(PyBytes::new(py, bytes.as_ref()).unbind()) } @@ -713,10 +721,7 @@ impl PyBlobFile { self_.ensure_open()?; let inner = self_.inner.clone(); future_into_py(self_.py(), async move { - let bytes = inner - .read() - .await - .map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?; + let bytes = inner.read().await.map_err(|e| blob_file_error("read", e))?; Python::attach(|py| Ok(PyBytes::new(py, bytes.as_ref()).unbind())) }) } @@ -739,7 +744,7 @@ impl PyBlobFile { self_ .py() .detach(move || block_on(async move { inner.close().await })) - .map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}"))) + .map_err(|e| blob_file_error("close", e)) } fn is_closed(self_: PyRef<'_, Self>) -> bool { @@ -752,7 +757,7 @@ impl PyBlobFile { self_ .py() .detach(move || block_on(async move { inner.seek(position).await })) - .map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}"))) + .map_err(|e| blob_file_error("seek", e)) } fn tell(self_: PyRef<'_, Self>) -> PyResult { @@ -761,7 +766,7 @@ impl PyBlobFile { self_ .py() .detach(move || block_on(async move { inner.tell().await })) - .map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}"))) + .map_err(|e| blob_file_error("tell", e)) } fn size(self_: PyRef<'_, Self>) -> u64 { @@ -778,17 +783,39 @@ impl PyBlobFile { let py = self_.py(); let bytes = py .detach(move || block_on(async move { inner.read_range(offset..end).await })) - .map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?; + .map_err(|e| blob_file_error("read_range", e))?; Ok(PyBytes::new(py, bytes.as_ref()).unbind()) } + /// Read blob-local byte ranges without moving the cursor. + fn read_ranges(self_: PyRef<'_, Self>, ranges: Vec<(u64, u64)>) -> PyResult>> { + let ranges = ranges + .into_iter() + .map(|(offset, length)| { + offset + .checked_add(length) + .map(|end| offset..end) + .ok_or_else(|| PyValueError::new_err("offset + length overflowed")) + }) + .collect::>>()?; + let inner = self_.inner.clone(); + let py = self_.py(); + let bytes = py + .detach(move || block_on(async move { inner.read_ranges(&ranges).await })) + .map_err(|e| blob_file_error("read_ranges", e))?; + Ok(bytes + .iter() + .map(|bytes| PyBytes::new(py, bytes.as_ref()).unbind()) + .collect()) + } + 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 .detach(move || block_on(async move { inner.read_up_to(length).await })) - .map_err(|e| PyRuntimeError::new_err(format!("blob read_up_to failed: {e}")))?; + .map_err(|e| blob_file_error("read_up_to", e))?; Ok(PyBytes::new(py, bytes.as_ref()).unbind()) } } @@ -1462,14 +1489,30 @@ impl Table { pub fn fetch_blob_ranges( self_: PyRef<'_, Self>, column: String, - requests: Vec<(u64, u64, u64)>, + requests: Vec>, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); + let requests = requests + .into_iter() + .map(|values| { + let [row_id, offset, length] = <[i128; 3]>::try_from(values).map_err(|_| { + PyValueError::new_err("each blob range request must have three values") + })?; + let to_u64 = |value, name| { + u64::try_from(value).map_err(|_| { + PyValueError::new_err(format!( + "blob range {name} must be a non-negative 64-bit integer" + )) + }) + }; + Ok(BlobRangeRequest::new( + to_u64(row_id, "row id")?, + to_u64(offset, "offset")?, + to_u64(length, "length")?, + )) + }) + .collect::>>()?; future_into_py(self_.py(), async move { - let requests = requests - .into_iter() - .map(|(row_id, offset, length)| BlobRangeRequest::new(row_id, offset, length)) - .collect::>(); let blobs: LargeBinaryArray = inner .fetch_blob_ranges(column, requests) .await diff --git a/python/uv.lock b/python/uv.lock index f80541711..ace8b98b9 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -1849,19 +1849,19 @@ wheels = [ [[package]] name = "lance-namespace" -version = "0.11.1" +version = "0.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "lance-namespace-urllib3-client" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bf/93/da5f7fcac690db9b282a3439ed9e34960c147619a0d6e1f4eb8cd240e7a5/lance_namespace-0.11.1.tar.gz", hash = "sha256:f67cfbbe0647b7cb42f23b673e7edf8a75b7d8a047265a916492f8d247ee1bc2", size = 11631, upload-time = "2026-08-18T17:40:06.294Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c0/e7/d5d46594678ee479c0eda830c47b2f5c46133bd100e84e9aa6e01306eca7/lance_namespace-0.13.0.tar.gz", hash = "sha256:24554a0997bdb39595c6e4cb3ac6722069f6cd3bd1a74e7acbba7bfaf774a40d", size = 11636, upload-time = "2026-09-14T02:47:03.871523Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/bc/601f2b3cc4cfa0070d858a33223bc823fffdd7981a25c45984a5216ca952/lance_namespace-0.11.1-py3-none-any.whl", hash = "sha256:07643fce9a42ad4d58cc8bf91e3f592bc7f4cbd8d0ad5233223506debf67551c", size = 13507, upload-time = "2026-08-18T17:40:03.561Z" }, + { url = "https://files.pythonhosted.org/packages/e7/47/cfa33cca1ba7c749fd2918cfc5c8ded788f378cc6424d23e4fead5a14125/lance_namespace-0.13.0-py3-none-any.whl", hash = "sha256:438c7b17aef421c21c138196e715f2510d62f07e865372047f87c1e75e618c7a", size = 13507, upload-time = "2026-09-14T02:47:01.765655Z" }, ] [[package]] name = "lance-namespace-urllib3-client" -version = "0.11.1" +version = "0.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pydantic" }, @@ -1869,9 +1869,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e4/c5/2bdd0ff98b469894c8a73be809d26ffdad5402517b0e5f9e758026cba29e/lance_namespace_urllib3_client-0.11.1.tar.gz", hash = "sha256:145a9e9424d7597487249b5b95ee274423bf2910e1a9160b6a07b676b61ea46a", size = 237345, upload-time = "2026-08-18T17:40:07.308Z" } +sdist = { url = "https://files.pythonhosted.org/packages/69/25/9aaa4a5e8999693fb0f227c2c0c4b97bd8f0539066408cdff0b89d41b2d5/lance_namespace_urllib3_client-0.13.0.tar.gz", hash = "sha256:1e8a79c6e4e6277033597fd76aa0e2f33d909ca1436c935936e9e569221f43ef", size = 238348, upload-time = "2026-09-14T02:47:04.521878Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/2a/eaaefd55d1190291207049fedc6b3eb22b506e57d6de91bae46bbaaa9c60/lance_namespace_urllib3_client-0.11.1-py3-none-any.whl", hash = "sha256:36537f529294da6d884ba0fe783704483f0a75463497c7705fd083a4d0257990", size = 406311, upload-time = "2026-08-18T17:40:04.842Z" }, + { url = "https://files.pythonhosted.org/packages/60/4c/7b8f0712a7fe1b8655b711342bc4989696635a79e7027ee56ba3cae23e99/lance_namespace_urllib3_client-0.13.0-py3-none-any.whl", hash = "sha256:fb361eb4f6c7f2d1f9e92809609657b73e38241393e9c4516e12d6b6ee643a0a", size = 406848, upload-time = "2026-09-14T02:47:02.715954Z" }, ] [[package]] @@ -1900,6 +1900,7 @@ clip = [ { name = "torch" }, ] dev = [ + { name = "maturin" }, { name = "pre-commit" }, { name = "pyright" }, { name = "ruff" }, @@ -1980,7 +1981,8 @@ requires-dist = [ { name = "huggingface-hub", marker = "extra == 'embeddings'", specifier = ">=0.19.0" }, { name = "ibm-watsonx-ai", marker = "python_full_version >= '3.10' and extra == 'embeddings'", specifier = ">=1.1.2" }, { name = "instructorembedding", marker = "extra == 'embeddings'", specifier = ">=1.0.1" }, - { name = "lance-namespace", specifier = ">=0.3.2" }, + { name = "lance-namespace", specifier = ">=0.13.0" }, + { name = "maturin", marker = "extra == 'dev'", specifier = ">=1.10" }, { name = "mkdocs", marker = "extra == 'docs'" }, { name = "mkdocs-jupyter", marker = "extra == 'docs'" }, { name = "mkdocs-material", marker = "extra == 'docs'" }, @@ -2005,7 +2007,7 @@ requires-dist = [ { name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" }, { name = "pydantic", specifier = ">=2.7.4,<3" }, { name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" }, - { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b6" }, + { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" }, { name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" }, { name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" }, @@ -2156,6 +2158,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl", hash = "sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6", size = 9534, upload-time = "2026-05-08T17:33:32.055Z" }, ] +[[package]] +name = "maturin" +version = "1.15.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b9/c8/22e5e21b2679c9bce6415ca578034ca2cc9316be0642ae21e051a2d5198c/maturin-1.15.0.tar.gz", hash = "sha256:94b26cc8e8aba61a5f2099715fe640e18c5f678e9a500408b38761263954228a", size = 385504, upload-time = "2026-08-24T12:11:22.665Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/14/69/5c01b461044eb1f45ddcce006706eb88110c793cdb11c7ae0b5e08492e94/maturin-1.15.0-py3-none-linux_armv6l.whl", hash = "sha256:6bf6dc62e22d4dcfd5a51244ff0d58975fa4979c48209fe84159617648956d82", size = 10206220, upload-time = "2026-08-24T12:10:53.327Z" }, + { url = "https://files.pythonhosted.org/packages/eb/1f/2b431554e11687cdb1077e0cdadcc118c53f611086b3af00c8545a67c6a5/maturin-1.15.0-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:cd35772633f489841132bc8e71d6fc7f842df30b9c05cd5cdf1ee1ddcb744cc7", size = 19416513, upload-time = "2026-08-24T12:10:56.126Z" }, + { url = "https://files.pythonhosted.org/packages/51/36/e23a21cb34a648b711036b9b2fe1d4f3f4ee24f8db54215d73f1a9a3a3ec/maturin-1.15.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:c40b4eae7bf5ef1f4b1af8d623fe4105016f93578fb15b764e741d08ec3b92dd", size = 10014962, upload-time = "2026-08-24T12:10:58.486Z" }, + { url = "https://files.pythonhosted.org/packages/8c/80/33b15cb2d8f30f12c807955e8f2fd775027692904e30ec0784744ce8cd83/maturin-1.15.0-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:7eb066372f541f8eb4909c79c5d9bd0b9e8125980bdf1ec9e8aba23c6c8d6c55", size = 10196223, upload-time = "2026-08-24T12:11:00.696Z" }, + { url = "https://files.pythonhosted.org/packages/fe/91/b495e19e2f5c503b540452b2039115e7b2363867e8c5ad4179eb752fa92c/maturin-1.15.0-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:653020a63525bb224e5ab0adf02e17a2e08bc86dbea7fc1399c9a56d7529b99e", size = 10541186, upload-time = "2026-08-24T12:11:02.857Z" }, + { url = "https://files.pythonhosted.org/packages/5d/2b/2abff58037188d852b124871b1f0d720e1c2bfb3d4f1b03d87c52cd66488/maturin-1.15.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:0ebf9767892725083138e671c34482c660317a2f3d6a29fc0e0f34e9d8c99136", size = 10083468, upload-time = "2026-08-24T12:11:05.012Z" }, + { url = "https://files.pythonhosted.org/packages/3f/07/b7e9f8be99a6627849e81ac7b6694876bce8f50a92995fe17e3cf2610f0a/maturin-1.15.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:7ab7eebffd7b8debca2265985de4eaeb332141276d24b9560b5ad484d4b3add1", size = 10047786, upload-time = "2026-08-24T12:11:07.042Z" }, + { url = "https://files.pythonhosted.org/packages/ca/ab/167e3cb7accee11b507dbe53e0e87aeccb376d44ae66284c96ee4df3a9fd/maturin-1.15.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:126e12e618b4db42f68c779a56d41f82a390145ba36ac3f621d057eb34f5ad9d", size = 13315332, upload-time = "2026-08-24T12:11:09.433Z" }, + { url = "https://files.pythonhosted.org/packages/14/4d/801379f646cbc6b00998e5289b0630a886be3a4ee4c75b6bc9b87478a7f1/maturin-1.15.0-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4f9d33e6c3f9615c8caceecbbbd440f8eb25a3ddeb687077682cd5eca2e9ae15", size = 10807183, upload-time = "2026-08-24T12:11:11.73Z" }, + { url = "https://files.pythonhosted.org/packages/89/27/2e612e1cbd1580e9e94d4722c227b5180dca27b32b955a34b79918aa1292/maturin-1.15.0-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:bf29beddd0c6708f112db51d5275fc28b28b9e9c9c5faae387eaef662918b176", size = 10413274, upload-time = "2026-08-24T12:11:14.04Z" }, + { url = "https://files.pythonhosted.org/packages/70/d8/202a7b4d75a51f20f84ec9ce3b7345b12b822207072164e2b1c6ef665125/maturin-1.15.0-py3-none-win32.whl", hash = "sha256:da649988be98e87e009e51b1bf0d301b6a301bc0cecbdd60d40d8ba60748d1ca", size = 8928744, upload-time = "2026-08-24T12:11:16.269Z" }, + { url = "https://files.pythonhosted.org/packages/40/dc/4e90da594986ba78dd3bc8a5921ecdcdb11085b22b02a412caab3b225601/maturin-1.15.0-py3-none-win_amd64.whl", hash = "sha256:552c2be4afd43fe8d5c9f3ec8d4c4756d973b8dcbe94c14084390301f50243e1", size = 10335085, upload-time = "2026-08-24T12:11:18.326Z" }, + { url = "https://files.pythonhosted.org/packages/8b/10/15d4314edf130955edf2dc237aa393a8a7c10f2b9b57b89fa2f61f915659/maturin-1.15.0-py3-none-win_arm64.whl", hash = "sha256:c7dc0c66c78d3debdd9c5aa807e861fbcbf07f3505d34b125df74c03986b0f48", size = 9713795, upload-time = "2026-08-24T12:11:20.83Z" }, +] + [[package]] name = "mdit-py-plugins" version = "0.6.1" @@ -3912,7 +3938,7 @@ crypto = [ [[package]] name = "pylance" -version = "14.0.0b6" +version = "14.0.0b12" source = { registry = "https://pypi.fury.io/lance-format/" } dependencies = [ { name = "lance-namespace" }, @@ -3922,12 +3948,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://pypi.fury.io/lance-format/-/ver_DuQEn/pylance-14.0.0b6-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3eade86fa30d9ff3248ee7ffb0b27ce29495394d12327da4d22b5c54bc2b59fe" }, - { url = "https://pypi.fury.io/lance-format/-/ver_K5sqj/pylance-14.0.0b6-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:10e291974ee738a1d5d7e9dad17a634394a91cb60080b5ceac209077fbaec43c" }, - { url = "https://pypi.fury.io/lance-format/-/ver_EnbkT/pylance-14.0.0b6-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:61bac10f5081c1210edd017f81c8c784d800b89446bd0ed05d4a24e7001965bc" }, - { url = "https://pypi.fury.io/lance-format/-/ver_19ffYg/pylance-14.0.0b6-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:3c627ad4a9870611becf84409d8e5ecddc6be64e1dd53bee4bd5fb10ab561e5c" }, - { url = "https://pypi.fury.io/lance-format/-/ver_2eLLY5/pylance-14.0.0b6-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:bc603a7dd690233f1fe000600171e1c092f153bc9060cdf198825c4bbe1418aa" }, - { url = "https://pypi.fury.io/lance-format/-/ver_t8ggc/pylance-14.0.0b6-cp310-abi3-win_amd64.whl", hash = "sha256:1d32df32a3d157611450def07eae8ace3a7401286eac68cad9a24e0d5deed2db" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" }, + { url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" }, + { url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" }, + { url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" }, + { url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" }, ] [[package]] diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 2d7299a0e..ec78fd103 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.1" +version = "0.41.0-beta.6" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true @@ -66,19 +66,23 @@ serde = { workspace = true } serde_json = { workspace = true } async-openai = { version = "0.20.0", optional = true } serde_with = { version = "3.8.1" } +strum_macros = "0.28" tempfile = { workspace = true } aws-smithy-types = { workspace = true, optional = true } -aws-sdk-bedrockruntime = { version = "1.27.0", optional = true } +# Default features include the legacy `rustls` feature, which pulls in rustls 0.21. +aws-sdk-bedrockruntime = { version = "1.77.0", optional = true, default-features = false, features = ["default-https-client", "rt-tokio"] } # For remote feature -reqwest = { version = "0.12.0", default-features = false, features = [ +reqwest = { version = "0.13.0", default-features = false, features = [ "charset", "gzip", "http2", "json", - "macos-system-configuration", + "query", # Avoid linking OpenSSL into Python wheels, which breaks on FIPS hosts. - "rustls-tls-native-roots", + # This uses the aws-lc-rs crypto provider. + "rustls", "stream", + "system-proxy", ], optional = true } tonic = { workspace = true, optional = true } http = { version = "1", optional = true } # Matching what is in reqwest @@ -114,9 +118,10 @@ random_word = { version = "0.4.3", features = ["en"] } tokio = { workspace = true, features = ["io-util", "macros", "net", "test-util"] } uuid = { workspace = true } walkdir = "2" -aws-sdk-dynamodb = { version = "1.55.0" } -aws-sdk-s3 = { version = "1.55.0" } -aws-sdk-kms = { version = "1.48.0" } +# Default features include the legacy `rustls` feature, which pulls in rustls 0.21. +aws-sdk-dynamodb = { version = "1.68.0", default-features = false, features = ["default-https-client", "rt-tokio"] } +aws-sdk-s3 = { version = "1.79.0", default-features = false, features = ["default-https-client", "rt-tokio", "sigv4a", "http-1x"] } +aws-sdk-kms = { version = "1.63.0", default-features = false, features = ["default-https-client", "rt-tokio"] } aws-config = { version = "1.5.10" } aws-smithy-runtime = { version = "1.9.1" } aws-smithy-types.workspace = true diff --git a/rust/lancedb/src/authz/api.rs b/rust/lancedb/src/authz/api.rs new file mode 100644 index 000000000..38e7d381b --- /dev/null +++ b/rust/lancedb/src/authz/api.rs @@ -0,0 +1,239 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Authorization requests, responses, and the deployment-wide API. + +use super::{Object, Privilege, Subject}; +use crate::Result; +use serde::{Deserialize, Serialize}; + +/// Filters and pagination for a page of ACL grants. Omitted filters match all. +#[derive(Clone, Default, Serialize, Deserialize)] +#[non_exhaustive] +pub struct ListAccessControlEntriesRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub object: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subject: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub continuation_token: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + /// Maximum stored subject/object pairs, not expanded privilege entries. + /// The server caps this at 100; None uses its default. Must be positive. + pub max_results: Option, +} + +/// One privilege granted to a canonical subject on an object. +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub struct AccessControlEntry { + #[serde(deserialize_with = "Subject::deserialize_canonical")] + pub subject_id: Subject, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subject_name: Option, + pub object: Object, + pub privilege: Privilege, +} + +/// One page of grants. A stored subject/object pair may expand to several entries. +#[derive(Clone, Default, Serialize, Deserialize, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub struct ListAccessControlEntriesResponse { + pub entries: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub continuation_token: Option, +} + +/// One grant to add or delete. The subject may be selected by ID, name, or API key. +#[derive(Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct AccessControlEntryRequest { + pub subject: Subject, + pub object: Object, + /// A privilege such as [`Privilege::Select`] or [`Privilege::Usage`]. + /// [`Privilege::Unknown`] preserves future privilege names. + pub privilege: Privilege, +} + +/// Filters and pagination for direct principal/group role bindings. +#[derive(Clone, Default, Serialize, Deserialize)] +#[non_exhaustive] +pub struct ListRoleBindingsRequest { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subject: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub role: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub continuation_token: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + /// Maximum bindings on this page; the server caps this at 100. + pub max_results: Option, +} + +/// A direct principal or group membership in a LanceDB-managed role. +#[derive(Clone, Serialize, Deserialize, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub struct RoleBinding { + #[serde(deserialize_with = "Subject::deserialize_canonical")] + pub subject_id: Subject, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subject_name: Option, + pub role: String, +} + +/// One page of direct bindings, without expanding group membership. +#[derive(Clone, Default, Serialize, Deserialize, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub struct ListRoleBindingsResponse { + pub bindings: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub continuation_token: Option, +} + +/// A role to grant to or revoke from a principal or group. Roles cannot nest. +#[derive(Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct RoleBindingRequest { + pub subject: Subject, + pub role: String, +} + +/// Select a principal by canonical ID, provider name, or API key. +#[derive(Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct GetPrincipalRequest { + pub principal: Subject, +} + +/// Select a group by canonical ID or provider name. +#[derive(Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct GetGroupRequest { + pub group: Subject, +} + +/// A canonical identity-provider ID and optional display name. +#[derive(Clone, Default, Serialize, Deserialize, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub struct IdentityInfo { + pub id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub name: Option, +} + +/// A principal, its provider-managed groups, and roles held directly or via groups. +#[derive(Clone, Default, Serialize, Deserialize, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub struct PrincipalInfo { + pub principal_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub principal_name: Option, + pub groups: Vec, + pub roles: Vec, +} + +/// A group, its provider-managed members, and its directly granted roles. +#[derive(Clone, Default, Serialize, Deserialize, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub struct GroupInfo { + pub group_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub group_name: Option, + pub principals: Vec, + pub roles: Vec, +} + +/// Authorization operations on one deployment, independent of database selection. +/// +/// ACL writes preserve server conflict/not-found errors and must not be replayed +/// after ambiguous transport failures. Role-binding writes are idempotent. +#[async_trait::async_trait] +pub trait Authorization: Send + Sync + std::fmt::Debug + 'static { + /// List a page of ACLs. Requires USAGE on the object, or system if unfiltered. + async fn list_acls( + &self, + request: ListAccessControlEntriesRequest, + ) -> Result; + /// Grant a privilege; requires OWNERSHIP. Duplicate grants return HTTP 409. + /// Granting OWNERSHIP transfers it from the previous owner. + async fn add_acl(&self, request: AccessControlEntryRequest) -> Result<()>; + /// Revoke a privilege; requires OWNERSHIP. Missing grants return HTTP 404. + /// OWNERSHIP cannot be deleted. + async fn delete_acl(&self, request: AccessControlEntryRequest) -> Result<()>; + /// List a page of direct bindings; requires USAGE on system. + async fn list_role_bindings( + &self, + request: ListRoleBindingsRequest, + ) -> Result; + /// Idempotently grant a role; requires OPERATE on system. + async fn add_role_binding(&self, request: RoleBindingRequest) -> Result<()>; + /// Idempotently revoke a role; requires OPERATE on system. + async fn delete_role_binding(&self, request: RoleBindingRequest) -> Result<()>; + /// Inspect a principal and effective roles; requires USAGE on system. + async fn get_principal(&self, request: GetPrincipalRequest) -> Result; + /// Inspect a group and direct roles; requires USAGE on system. + async fn get_group(&self, request: GetGroupRequest) -> Result; +} + +impl AccessControlEntryRequest { + /// Build one grant from validated resource and identity selectors. + pub fn new(object: Object, subject: Subject, privilege: Privilege) -> Self { + Self { + object, + subject, + privilege, + } + } +} + +impl ListRoleBindingsRequest { + /// Reject role subjects, since roles cannot contain other roles. + pub fn validate(&self) -> Result<()> { + if let Some(subject) = &self.subject { + subject.validate_role_member()?; + } + Ok(()) + } +} + +impl RoleBindingRequest { + /// Build a role grant/revocation for a principal or group. + pub fn new(subject: Subject, role: impl Into) -> Result { + subject.validate_role_member()?; + Ok(Self { + subject, + role: role.into(), + }) + } + + /// Validate the member selector, including deserialized requests. + pub fn validate(&self) -> Result<()> { + self.subject.validate_role_member() + } +} + +impl GetPrincipalRequest { + /// Select a principal by ID, name, or API key. + pub fn new(principal: Subject) -> Result { + principal.validate_principal()?; + Ok(Self { principal }) + } + + /// Validate the lookup selector, including deserialized requests. + pub fn validate(&self) -> Result<()> { + self.principal.validate_principal() + } +} + +impl GetGroupRequest { + /// Select a group by ID or name. + pub fn new(group: Subject) -> Result { + group.validate_group()?; + Ok(Self { group }) + } + + /// Validate the lookup selector, including deserialized requests. + pub fn validate(&self) -> Result<()> { + self.group.validate_group() + } +} diff --git a/rust/lancedb/src/authz/mod.rs b/rust/lancedb/src/authz/mod.rs new file mode 100644 index 000000000..3191627e0 --- /dev/null +++ b/rust/lancedb/src/authz/mod.rs @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Deployment-wide authorization, accessed through a catalog. +//! +//! Subjects use `p:ID`, `P:NAME`, `g:ID`, `G:NAME`, `r:ROLE`, or `a:API_KEY`. +//! Objects use literal names such as `table:tenant/db:public$events`. +//! Names are resolved by the server. Request types intentionally do not implement +//! Debug because their subject selectors may contain secret API keys. +//! +//! ```no_run +//! # #[cfg(feature = "remote")] +//! # async fn example() -> lancedb::Result<()> { +//! let catalog = lancedb::connect_catalog("https://example.com") +//! .api_key("secret") +//! .execute().await?; +//! let authz = catalog.authz()?; +//! use lancedb::authz::{AccessControlEntryRequest, Object, TableObject, NamespacePath, Subject, Privilege}; +//! let grant = AccessControlEntryRequest::new( +//! Object::Table(TableObject { +//! database: "analytics".into(), +//! namespace: NamespacePath::from_component("public"), +//! table: "foo".into(), +//! }), +//! Subject::role("reader")?, +//! Privilege::Select, +//! ); +//! authz.add_acl(grant).await?; +//! let page = authz.list_acls(Default::default()).await?; +//! for entry in page.entries { +//! println!("{} {} {}", entry.object, entry.subject_id, entry.privilege); +//! } +//! # Ok(()) +//! # } +//! ``` + +pub mod api; +pub mod object; +pub mod privilege; +pub mod subject; + +pub use api::*; +pub use object::{ + DatabaseObject, FunctionObject, NamespaceObject, NamespacePath, Object, SecretObject, + TableObject, ViewObject, +}; +pub use privilege::*; +pub use subject::{Subject, SubjectKind}; diff --git a/rust/lancedb/src/authz/object.rs b/rust/lancedb/src/authz/object.rs new file mode 100644 index 000000000..fd324b748 --- /dev/null +++ b/rust/lancedb/src/authz/object.rs @@ -0,0 +1,992 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use crate::utils::{ + validate_database_name, validate_namespace_component, validate_secret_name, + validate_table_name, validate_view_name, +}; +use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _, ser::Error as _}; + +/// The default namespace for authorization resources. +pub const DEFAULT_NAMESPACE: &str = "public"; + +pub const SYSTEM_TYPE_STRING: &str = "system"; +pub const DATABASE_TYPE_STRING: &str = "database"; +pub const NAMESPACE_TYPE_STRING: &str = "namespace"; +pub const TABLE_TYPE_STRING: &str = "table"; +pub const VIEW_TYPE_STRING: &str = "view"; +pub const SECRET_TYPE_STRING: &str = "secret"; +pub const FUNCTION_TYPE_STRING: &str = "function"; +pub const UNKNOWN_TYPE_STRING: &str = "unknown"; + +/// A non-empty path through the namespace hierarchy. +/// +/// The `$` separator is only the canonical wire encoding. Authorization code +/// should use the components and prefix helpers instead of interpreting the +/// encoded form. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct NamespacePath(Vec); + +impl NamespacePath { + /// Construct a path, validating each component without splitting separators. + /// An empty path selects the default `public` namespace. + /// + /// ``` + /// use lancedb::authz::NamespacePath; + /// + /// let path = NamespacePath::new(vec![]).unwrap(); + /// assert_eq!(path.components(), &["public"]); + /// assert_eq!(path, NamespacePath::default()); + /// ``` + pub fn new(components: Vec) -> Result { + if components.is_empty() { + return Ok(Self::default()); + } + for component in &components { + validate_namespace_component(component).map_err(|error| error.to_string())?; + } + Ok(Self(components)) + } + + /// Construct a single-component path. + /// + /// # Panics + /// Panics if the component is invalid. Use [`Self::new`] for untrusted input. + pub fn from_component(component: impl Into) -> Self { + Self::new(vec![component.into()]).expect("namespace components must be valid") + } + + pub fn components(&self) -> &[String] { + &self.0 + } + + pub fn parent(&self) -> Option { + (self.0.len() > 1).then(|| Self(self.0[..self.0.len() - 1].to_vec())) + } + + pub fn prefixes(&self) -> impl Iterator + '_ { + (1..=self.0.len()).map(|length| Self(self.0[..length].to_vec())) + } + + pub fn is_descendant_of(&self, ancestor: &Self) -> bool { + self.0.len() > ancestor.0.len() && self.0.starts_with(&ancestor.0) + } + + pub fn parse(path: &str) -> Result { + Self::new(path.split('$').map(str::to_string).collect()) + } +} + +impl Default for NamespacePath { + /// Select the default `public` namespace. + fn default() -> Self { + Self(vec![DEFAULT_NAMESPACE.to_string()]) + } +} + +impl std::fmt::Display for NamespacePath { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0.join("$")) + } +} + +/// An authorization resource, serialized using its canonical wire name. +/// +/// Unknown resource types are preserved so clients can read rules created by +/// newer servers. Known resource types validate names during parsing and serialization. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub enum Object { + System, + Database(DatabaseObject), + Namespace(NamespaceObject), + Table(TableObject), + View(ViewObject), + Secret(SecretObject), + Function(FunctionObject), + Unknown(String), +} + +impl Object { + /// Validate resource names, including objects recovered from storage keys. + pub fn validate(&self) -> Result<(), String> { + match self { + Self::System => Ok(()), + Self::Database(object) => object.validate(), + Self::Namespace(object) => object.validate(), + Self::Table(object) => object.validate(), + Self::View(object) => object.validate(), + Self::Secret(object) => object.validate(), + Self::Function(object) => object.validate(), + Self::Unknown(value) if value.is_empty() => Err("object must not be empty".into()), + Self::Unknown(_) => Ok(()), + } + } + + /// The database an object belongs to, or None. + pub fn database(&self) -> Option<&str> { + match self { + Self::System => None, + Self::Database(object) => Some(&object.database), + Self::Namespace(object) => Some(&object.database), + Self::Table(object) => Some(&object.database), + Self::View(object) => Some(&object.database), + Self::Secret(object) => Some(&object.database), + Self::Function(object) => Some(&object.database), + Self::Unknown(_) => None, + } + } + + pub fn type_name(&self) -> &'static str { + match self { + Self::System => SYSTEM_TYPE_STRING, + Self::Database(_) => DATABASE_TYPE_STRING, + Self::Namespace(_) => NAMESPACE_TYPE_STRING, + Self::Table(_) => TABLE_TYPE_STRING, + Self::View(_) => VIEW_TYPE_STRING, + Self::Secret(_) => SECRET_TYPE_STRING, + Self::Function(_) => FUNCTION_TYPE_STRING, + Self::Unknown(_) => UNKNOWN_TYPE_STRING, + } + } +} + +impl std::str::FromStr for Object { + type Err = crate::Error; + + fn from_str(value: &str) -> crate::Result { + Self::deserialize(serde::de::value::StrDeserializer::::new(value)) + .map_err(|error| crate::Error::InvalidInput { + message: error.to_string(), + }) + } +} + +impl std::fmt::Display for Object { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::System => f.write_str(SYSTEM_TYPE_STRING), + Self::Database(object) => object.fmt(f), + Self::Namespace(object) => object.fmt(f), + Self::Table(object) => object.fmt(f), + Self::View(object) => object.fmt(f), + Self::Secret(object) => object.fmt(f), + Self::Function(object) => object.fmt(f), + Self::Unknown(s) => f.write_str(s), + } + } +} + +impl Serialize for Object { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for Object { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let kind = s.split(':').next().unwrap_or_default(); + if s == SYSTEM_TYPE_STRING { + Ok(Self::System) + } else if s.is_empty() || kind == SYSTEM_TYPE_STRING { + Err(D::Error::custom(format!("invalid Object {s:?}"))) + } else if kind == DATABASE_TYPE_STRING { + let database = DatabaseObject::deserialize(serde::de::value::StringDeserializer::< + D::Error, + >::new(s))?; + Ok(Self::Database(database)) + } else if kind == NAMESPACE_TYPE_STRING { + let namespace = NamespaceObject::deserialize(serde::de::value::StringDeserializer::< + D::Error, + >::new(s))?; + Ok(Self::Namespace(namespace)) + } else if kind == TABLE_TYPE_STRING { + let table = + TableObject::deserialize(serde::de::value::StringDeserializer::::new(s))?; + Ok(Self::Table(table)) + } else if kind == VIEW_TYPE_STRING { + let view = + ViewObject::deserialize(serde::de::value::StringDeserializer::::new(s))?; + Ok(Self::View(view)) + } else if kind == SECRET_TYPE_STRING { + let secret = SecretObject::deserialize( + serde::de::value::StringDeserializer::::new(s), + )?; + Ok(Self::Secret(secret)) + } else if kind == FUNCTION_TYPE_STRING { + let function = FunctionObject::deserialize(serde::de::value::StringDeserializer::< + D::Error, + >::new(s))?; + Ok(Self::Function(function)) + } else { + Ok(Self::Unknown(s)) + } + } +} + +/// A database resource, encoded as `database:NAME`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct DatabaseObject { + pub database: String, +} + +impl DatabaseObject { + pub fn new(database: String) -> Self { + Self { database } + } + + fn validate(&self) -> Result<(), String> { + validate_database_name(&self.database).map_err(|error| error.to_string())?; + Ok(()) + } +} + +impl std::fmt::Display for DatabaseObject { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "database:{}", self.database) + } +} + +impl Serialize for DatabaseObject { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for DatabaseObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let database = s + .strip_prefix("database:") + .ok_or_else(|| D::Error::custom(format!("invalid DatabaseObject {s:?}")))?; + let object = Self { + database: database.to_string(), + }; + object.validate().map_err(D::Error::custom)?; + Ok(object) + } +} + +/// A namespace resource, encoded as `namespace:DATABASE:PATH`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct NamespaceObject { + pub database: String, + pub namespace: NamespacePath, +} + +impl NamespaceObject { + fn validate(&self) -> Result<(), String> { + validate_database_name(&self.database).map_err(|error| error.to_string())?; + Ok(()) + } +} + +impl std::fmt::Display for NamespaceObject { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "namespace:{}:{}", self.database, self.namespace) + } +} + +impl Serialize for NamespaceObject { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for NamespaceObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let invalid = || D::Error::custom(format!("invalid NamespaceObject {s:?}")); + let rest = s.strip_prefix("namespace:").ok_or_else(invalid)?; + let (database, namespace) = rest.split_once(':').ok_or_else(invalid)?; + let object = Self { + database: database.to_string(), + namespace: NamespacePath::parse(namespace).map_err(D::Error::custom)?, + }; + object.validate().map_err(D::Error::custom)?; + Ok(object) + } +} + +/// A table resource, encoded as `table:DATABASE:NAMESPACE$TABLE`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct TableObject { + pub database: String, + pub namespace: NamespacePath, + pub table: String, +} + +impl TableObject { + fn validate(&self) -> Result<(), String> { + validate_database_name(&self.database).map_err(|error| error.to_string())?; + validate_table_name(&self.table).map_err(|error| error.to_string())?; + Ok(()) + } +} + +impl std::fmt::Display for TableObject { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "table:{}:{}${}", + self.database, self.namespace, self.table + ) + } +} + +impl Serialize for TableObject { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for TableObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let invalid = || D::Error::custom(format!("invalid TableObject {s:?}")); + let rest = s.strip_prefix("table:").ok_or_else(invalid)?; + let (database, path) = rest.split_once(':').ok_or_else(invalid)?; + let (namespace, table) = path.rsplit_once('$').ok_or_else(invalid)?; + let object = Self { + database: database.to_string(), + namespace: NamespacePath::parse(namespace).map_err(D::Error::custom)?, + table: table.to_string(), + }; + object.validate().map_err(D::Error::custom)?; + Ok(object) + } +} + +/// A view resource, encoded as `view:DATABASE:NAMESPACE$VIEW`. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct ViewObject { + pub database: String, + pub namespace: NamespacePath, + pub view: String, +} + +impl ViewObject { + fn validate(&self) -> Result<(), String> { + validate_database_name(&self.database).map_err(|error| error.to_string())?; + validate_view_name(&self.view).map_err(|error| error.to_string())?; + Ok(()) + } +} + +impl std::fmt::Display for ViewObject { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "view:{}:{}${}", self.database, self.namespace, self.view) + } +} + +impl Serialize for ViewObject { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for ViewObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let invalid = || D::Error::custom(format!("invalid ViewObject {s:?}")); + let rest = s.strip_prefix("view:").ok_or_else(invalid)?; + let (database, path) = rest.split_once(':').ok_or_else(invalid)?; + let (namespace, view) = path.rsplit_once('$').ok_or_else(invalid)?; + let object = Self { + database: database.to_string(), + namespace: NamespacePath::parse(namespace).map_err(D::Error::custom)?, + view: view.to_string(), + }; + object.validate().map_err(D::Error::custom)?; + Ok(object) + } +} +/// A secret resource, encoded as `secret:DATABASE:NAMESPACE$SECRET`. +/// +/// ``` +/// use lancedb::authz::{SecretObject, NamespacePath, Object}; +/// +/// let object = Object::Secret(SecretObject { +/// database: "analytics".into(), +/// namespace: NamespacePath::from_component("public"), +/// secret: "example".into(), +/// }); +/// assert_eq!(object.to_string(), "secret:analytics:public$example"); +/// ``` +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct SecretObject { + pub database: String, + pub namespace: NamespacePath, + pub secret: String, +} + +impl SecretObject { + fn validate(&self) -> Result<(), String> { + validate_database_name(&self.database).map_err(|error| error.to_string())?; + validate_secret_name(&self.secret).map_err(|error| error.to_string())?; + Ok(()) + } +} + +impl std::fmt::Display for SecretObject { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "secret:{}:{}${}", + self.database, self.namespace, self.secret + ) + } +} + +impl Serialize for SecretObject { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for SecretObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let invalid = || D::Error::custom(format!("invalid SecretObject {s:?}")); + let rest = s.strip_prefix("secret:").ok_or_else(invalid)?; + let (database, path) = rest.split_once(':').ok_or_else(invalid)?; + let (namespace, secret) = path.rsplit_once('$').ok_or_else(invalid)?; + let object = Self { + database: database.to_string(), + namespace: NamespacePath::parse(namespace).map_err(D::Error::custom)?, + secret: secret.to_string(), + }; + object.validate().map_err(D::Error::custom)?; + Ok(object) + } +} + +/// A function resource, encoded as `function:DATABASE:NAMESPACE$FUNCTION`. +/// +/// ``` +/// use lancedb::authz::{FunctionObject, NamespacePath, Object}; +/// +/// let object = Object::Function(FunctionObject { +/// database: "analytics".into(), +/// namespace: NamespacePath::from_component("public"), +/// function: "example".into(), +/// }); +/// assert_eq!(object.to_string(), "function:analytics:public$example"); +/// ``` +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub struct FunctionObject { + pub database: String, + pub namespace: NamespacePath, + pub function: String, +} + +impl FunctionObject { + fn validate(&self) -> Result<(), String> { + validate_database_name(&self.database).map_err(|error| error.to_string())?; + validate_table_name(&self.function).map_err(|error| error.to_string())?; + Ok(()) + } +} + +impl std::fmt::Display for FunctionObject { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "function:{}:{}${}", + self.database, self.namespace, self.function + ) + } +} + +impl Serialize for FunctionObject { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.validate().map_err(S::Error::custom)?; + self.to_string().serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for FunctionObject { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + let invalid = || D::Error::custom(format!("invalid FunctionObject {s:?}")); + let rest = s.strip_prefix("function:").ok_or_else(invalid)?; + let (database, path) = rest.split_once(':').ok_or_else(invalid)?; + let (namespace, function) = path.rsplit_once('$').ok_or_else(invalid)?; + let object = Self { + database: database.to_string(), + namespace: NamespacePath::parse(namespace).map_err(D::Error::custom)?, + function: function.to_string(), + }; + object.validate().map_err(D::Error::custom)?; + Ok(object) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn system_object_round_trip() { + let encoded = serde_json::to_string(&Object::System).unwrap(); + assert_eq!(encoded, r#""system""#); + assert_eq!( + serde_json::from_str::(&encoded).unwrap(), + Object::System + ); + } + + #[test] + fn object_database_round_trip() { + let object = Object::Database(DatabaseObject { + database: "db".to_string(), + }); + let encoded = serde_json::to_string(&object).unwrap(); + assert_eq!(encoded, r#""database:db""#); + assert_eq!(serde_json::from_str::(&encoded).unwrap(), object); + } + + #[test] + fn object_namespace_round_trip() { + let object = Object::Namespace(NamespaceObject { + database: "db".to_string(), + namespace: NamespacePath::from_component("ns"), + }); + let encoded = serde_json::to_string(&object).unwrap(); + assert_eq!(encoded, r#""namespace:db:ns""#); + assert_eq!(serde_json::from_str::(&encoded).unwrap(), object); + } + + #[test] + fn object_table_round_trip() { + let object = Object::Table(TableObject { + database: "db".to_string(), + namespace: NamespacePath::from_component("ns"), + table: "foo".to_string(), + }); + let encoded = serde_json::to_string(&object).unwrap(); + assert_eq!(encoded, r#""table:db:ns$foo""#); + assert_eq!(serde_json::from_str::(&encoded).unwrap(), object); + } + + #[test] + fn object_view_round_trip() { + let object = Object::View(ViewObject { + database: "db".to_string(), + namespace: NamespacePath::from_component("ns"), + view: "foo".to_string(), + }); + let encoded = serde_json::to_string(&object).unwrap(); + assert_eq!(encoded, r#""view:db:ns$foo""#); + assert_eq!(serde_json::from_str::(&encoded).unwrap(), object); + } + + #[test] + fn secret_objects_round_trip_and_validate() { + let resource = SecretObject { + database: "tenant/db".into(), + namespace: NamespacePath::new(vec!["a".into(), "b".into()]).unwrap(), + secret: "example".into(), + }; + let wire = "secret:tenant/db:a$b$example"; + assert_eq!(serde_json::to_value(&resource).unwrap(), wire); + assert_eq!( + serde_json::from_value::(serde_json::json!(wire)).unwrap(), + resource + ); + let object = Object::Secret(resource.clone()); + assert_eq!(object.database(), Some("tenant/db")); + assert_eq!(object.type_name(), "secret"); + assert_eq!(wire.parse::().unwrap(), object); + assert_eq!(serde_json::to_value(object).unwrap(), wire); + for name in ["", ".", "..", "bad$name", "bad/name", "bad:name"] { + let invalid = SecretObject { + secret: name.into(), + ..resource.clone() + }; + assert!(serde_json::to_value(&invalid).is_err()); + assert!(serde_json::to_value(Object::Secret(invalid)).is_err()); + } + } + + #[test] + fn function_objects_round_trip_and_validate() { + let resource = FunctionObject { + database: "tenant/db".into(), + namespace: NamespacePath::new(vec!["a".into(), "b".into()]).unwrap(), + function: "example".into(), + }; + let wire = "function:tenant/db:a$b$example"; + assert_eq!(serde_json::to_value(&resource).unwrap(), wire); + assert_eq!( + serde_json::from_value::(serde_json::json!(wire)).unwrap(), + resource + ); + let object = Object::Function(resource.clone()); + assert_eq!(object.database(), Some("tenant/db")); + assert_eq!(object.type_name(), "function"); + assert_eq!(wire.parse::().unwrap(), object); + assert_eq!(serde_json::to_value(object).unwrap(), wire); + for name in ["", ".", "..", "bad$name", "bad/name", "bad:name"] { + let invalid = FunctionObject { + function: name.into(), + ..resource.clone() + }; + assert!(serde_json::to_value(&invalid).is_err()); + assert!(serde_json::to_value(Object::Function(invalid)).is_err()); + } + } + + #[test] + fn unknown_objects_round_trip() { + for wire in [ + "future:db:resource", + "table_extension:db:resource", + "garbage", + ] { + let object: Object = wire.parse().unwrap(); + assert_eq!(object, Object::Unknown(wire.into())); + assert_eq!(object.database(), None); + assert_eq!(object.type_name(), UNKNOWN_TYPE_STRING); + assert_eq!(object.to_string(), wire); + assert_eq!(serde_json::to_value(&object).unwrap(), wire); + assert_eq!( + serde_json::from_value::(serde_json::json!(wire)).unwrap(), + object + ); + } + } + + #[test] + fn database_object_round_trip() { + let object = DatabaseObject { + database: "db".to_string(), + }; + let encoded = serde_json::to_string(&object).unwrap(); + assert_eq!(encoded, r#""database:db""#); + assert_eq!( + serde_json::from_str::(&encoded).unwrap(), + object + ); + } + + #[test] + fn namespace_object_round_trip() { + let object = NamespaceObject { + database: "db".to_string(), + namespace: NamespacePath::from_component("ns"), + }; + let encoded = serde_json::to_string(&object).unwrap(); + assert_eq!(encoded, r#""namespace:db:ns""#); + assert_eq!( + serde_json::from_str::(&encoded).unwrap(), + object + ); + } + + #[test] + fn namespace_object_rejects_missing_namespace() { + assert!(serde_json::from_str::(r#""namespace:db""#).is_err()); + } + + #[test] + fn table_object_round_trip() { + let object = TableObject { + database: "db".to_string(), + namespace: NamespacePath::from_component("ns"), + table: "foo".to_string(), + }; + let encoded = serde_json::to_string(&object).unwrap(); + assert_eq!(encoded, r#""table:db:ns$foo""#); + assert_eq!( + serde_json::from_str::(&encoded).unwrap(), + object + ); + } + + #[test] + fn table_object_rejects_missing_prefix() { + assert!(serde_json::from_str::(r#""foo""#).is_err()); + } + + #[test] + fn table_object_rejects_missing_table() { + assert!(serde_json::from_str::(r#""table:db:ns""#).is_err()); + } + + #[test] + fn objects_order_by_variant_then_path_components() { + let database = Object::Database(DatabaseObject::new("z".to_string())); + let namespace = Object::Namespace(NamespaceObject { + database: "a".to_string(), + namespace: NamespacePath::from_component("a"), + }); + assert!(Object::System < database); + assert!(database < namespace); + + let first_table = Object::Table(TableObject { + database: "db".to_string(), + namespace: NamespacePath::from_component("a"), + table: "z".to_string(), + }); + let second_table = Object::Table(TableObject { + database: "db".to_string(), + namespace: NamespacePath::from_component("b"), + table: "a".to_string(), + }); + assert!(first_table < second_table); + } + + #[test] + fn empty_namespace_paths_select_public() { + let namespace = NamespacePath::new(vec![]).unwrap(); + assert_eq!(namespace, NamespacePath::default()); + assert_eq!(namespace.components(), &["public"]); + assert!(namespace.parent().is_none()); + assert_eq!( + namespace.prefixes().collect::>(), + vec![namespace.clone()] + ); + + let database = "db".to_string(); + let objects = [ + Object::Namespace(NamespaceObject { + database: database.clone(), + namespace: namespace.clone(), + }), + Object::Table(TableObject { + database: database.clone(), + namespace: namespace.clone(), + table: "item".into(), + }), + Object::View(ViewObject { + database: database.clone(), + namespace: namespace.clone(), + view: "item".into(), + }), + Object::Secret(SecretObject { + database: database.clone(), + namespace: namespace.clone(), + secret: "item".into(), + }), + Object::Function(FunctionObject { + database, + namespace, + function: "item".into(), + }), + ]; + for (object, wire) in objects.into_iter().zip([ + "namespace:db:public", + "table:db:public$item", + "view:db:public$item", + "secret:db:public$item", + "function:db:public$item", + ]) { + assert_eq!(serde_json::to_value(&object).unwrap(), wire); + assert_eq!(wire.parse::().unwrap(), object); + } + } + + #[test] + fn namespace_path_returns_cumulative_prefixes() { + let path = NamespacePath::new(vec!["ns1".into(), "ns2".into(), "ns3".into()]).unwrap(); + assert_eq!( + path.prefixes() + .map(|prefix| prefix.to_string()) + .collect::>(), + vec!["ns1", "ns1$ns2", "ns1$ns2$ns3"] + ); + assert_eq!(path.parent().unwrap().to_string(), "ns1$ns2"); + } + + #[test] + fn nested_objects_round_trip_without_escaping() { + for wire in [ + "database:tenant/db", + "namespace:tenant/db:a$b", + "table:tenant/db:a$b$c", + "view:tenant/db:a$b$c", + "secret:tenant/db:a$b$c", + "function:tenant/db:a$b$c", + ] { + let object: Object = serde_json::from_value(serde_json::json!(wire)).unwrap(); + assert_eq!(wire.parse::().unwrap(), object); + assert_eq!(object.database(), Some("tenant/db")); + assert_eq!(object.to_string(), wire); + assert_eq!(serde_json::to_value(object).unwrap(), wire); + } + } + + #[test] + fn objects_reject_invalid_names_and_legacy_forms() { + for database in [ + "", + ".", + "..", + "db:bad", + "db%2Fchild", + "db$child", + "/db", + "db/", + "db//child", + "db/../child", + "db name", + ] { + for wire in [ + format!("database:{database}"), + format!("namespace:{database}:ns"), + format!("table:{database}:ns$t"), + format!("view:{database}:ns$v"), + format!("secret:{database}:ns$s"), + format!("function:{database}:ns$f"), + ] { + assert!( + serde_json::from_value::(serde_json::json!(wire)).is_err(), + "{wire}" + ); + } + } + for name in ["", "bad:name", "bad/name", "bad%24name", "bad name"] { + for wire in [ + format!("namespace:db:{name}"), + format!("table:db:{name}$t"), + format!("view:db:{name}$v"), + format!("table:db:ns${name}"), + format!("view:db:ns${name}"), + format!("secret:db:{name}$s"), + format!("function:db:{name}$f"), + format!("secret:db:ns${name}"), + format!("function:db:ns${name}"), + ] { + assert!( + serde_json::from_value::(serde_json::json!(wire)).is_err(), + "{wire}" + ); + } + } + for wire in [ + "", + "system:db", + "database", + "namespace", + "table", + "view", + "namespace:db/ns", + "table:db/ns/t", + "view:db/ns/v", + "table:db:ns", + "view:db:ns", + "namespace:db:a$$b", + "table:db:ns$.", + "view:db:ns$..", + "table:db:a$$t", + "view:db:$v", + "secret", + "function", + "secret:db/ns/s", + "function:db/ns/f", + "secret:db:ns", + "function:db:ns", + "secret:db:a$$s", + "function:db:$f", + "secret:db:ns$.", + "function:db:ns$..", + ] { + assert!( + serde_json::from_value::(serde_json::json!(wire)).is_err(), + "{wire}" + ); + assert!(matches!( + wire.parse::(), + Err(crate::Error::InvalidInput { .. }) + )); + } + } + + #[test] + fn namespace_components_use_namespace_validation() { + for component in ["a$b", "c/d", "e%f", "a:b", "", ".", ".."] { + assert!(NamespacePath::new(vec![component.into()]).is_err()); + } + assert!(NamespacePath::new(vec!["a.b".into(), "c..d".into()]).is_ok()); + } + + #[test] + fn serialization_rejects_invalid_directly_constructed_objects() { + let database = DatabaseObject::new("db:bad".into()); + let namespace = NamespaceObject { + database: "db%2Fbad".into(), + namespace: NamespacePath::from_component("ns"), + }; + let table = TableObject { + database: "db".into(), + namespace: NamespacePath::from_component("ns"), + table: "bad$name".into(), + }; + let view = ViewObject { + database: "db".into(), + namespace: NamespacePath::from_component("ns"), + view: "..".into(), + }; + assert!(serde_json::to_string(&database).is_err()); + assert!(serde_json::to_string(&namespace).is_err()); + assert!(serde_json::to_string(&table).is_err()); + assert!(serde_json::to_string(&view).is_err()); + for object in [ + Object::Database(database), + Object::Namespace(namespace), + Object::Table(table), + Object::View(view), + ] { + assert!(serde_json::to_string(&object).is_err()); + } + } +} diff --git a/rust/lancedb/src/authz/privilege.rs b/rust/lancedb/src/authz/privilege.rs new file mode 100644 index 000000000..66b07643a --- /dev/null +++ b/rust/lancedb/src/authz/privilege.rs @@ -0,0 +1,514 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::{fmt, str::FromStr}; + +use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _}; + +/// A LanceDB privilege, serialized by name (for example, `SELECT`). +/// Numeric aliases use `PRIVILEGE` followed by decimal digits. Numbers 0 through +/// 13 map to the known privileges in declaration order. Unassigned numbers are +/// preserved verbatim; their meaning and validity are determined by the server. +/// +/// ``` +/// use lancedb::authz::Privilege; +/// +/// assert_eq!("PRIVILEGE0".parse::().unwrap(), Privilege::Ownership); +/// assert_eq!("PRIVILEGE2".parse::().unwrap(), Privilege::Select); +/// +/// let privilege: Privilege = "PRIVILEGE256".parse().unwrap(); +/// assert_eq!(privilege.to_string(), "PRIVILEGE256"); +/// +/// // Delegate creation on a namespace without transferring its ownership. +/// for privilege in [Privilege::CreateSecret, Privilege::CreateFunction] { +/// assert!(privilege.valid_for_namespace()); +/// } +/// ``` +#[derive(Debug, Clone, PartialEq, Eq, strum_macros::Display)] +#[strum(serialize_all = "SCREAMING_SNAKE_CASE")] +#[non_exhaustive] +pub enum Privilege { + /// A special privilege that conveys ownership. + Ownership, + + /// On a database: allows accessing the namespaces inside. + /// On a namespace: allows accessing the tables inside. + /// + /// When looking for Usage on a Database, if there are no ACLs + /// found for that database, we assume it doesn't exist yet and + /// check for CreateDatabase on System. + /// + /// When looking for Usage on a public (aka default) Namespace, + /// if there are no ACLs found for that Namespace, we assume it + /// doesn't exist yet and check for CreateNamespace on the + /// containing database instead. + Usage, + + /// On a table: allows reading rows. + Select, + + /// On a table: allows appending rows. + Insert, + + /// On a table: allows modifying existing rows. + Update, + + /// On a table: allows deleting rows. + Delete, + + /// On the system object: can create a database. + CreateDatabase, + + /// On a database or namespace: can create immediate child namespaces. + CreateNamespace, + + /// On a namespace: can create a table in that namespace. + CreateTable, + + /// On a namespace: can create a materialized view in that namespace. + CreateMaterializedView, + + /// On a namespace: can create a view in that namespace. + CreateView, + + /// On a namespace: can create a secret without taking ownership of the namespace. + CreateSecret, + + /// On a namespace: can create a function without taking ownership of the namespace. + CreateFunction, + + /// On the system object: + /// - can read server configuration. + Operate, + + /// A future privilege name, preserved verbatim. + #[strum(to_string = "{0}")] + Unknown(UnknownPrivilege), +} + +/// An unknown privilege. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UnknownPrivilege(String); + +impl fmt::Display for UnknownPrivilege { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl Privilege { + pub const KNOWN: [Self; 14] = [ + Self::Ownership, + Self::Usage, + Self::Select, + Self::Insert, + Self::Update, + Self::Delete, + Self::CreateDatabase, + Self::CreateNamespace, + Self::CreateTable, + Self::CreateMaterializedView, + Self::CreateView, + Self::CreateSecret, + Self::CreateFunction, + Self::Operate, + ]; + + pub fn valid_for_system(&self) -> bool { + matches!( + self, + Self::Operate | Self::CreateDatabase | Self::Usage | Self::Unknown(_) + ) + } + + pub fn valid_for_database(&self) -> bool { + matches!( + self, + Self::Ownership | Self::Usage | Self::CreateNamespace | Self::Unknown(_) + ) + } + + pub fn valid_for_namespace(&self) -> bool { + matches!( + self, + Self::Ownership + | Self::Usage + | Self::CreateNamespace + | Self::CreateTable + | Self::CreateMaterializedView + | Self::CreateView + | Self::CreateSecret + | Self::CreateFunction + | Self::Unknown(_) + ) + } + + pub fn valid_for_table(&self) -> bool { + matches!( + self, + Self::Ownership + | Self::Select + | Self::Insert + | Self::Update + | Self::Delete + | Self::Unknown(_) + ) + } + + pub fn valid_for_view(&self) -> bool { + matches!(self, Self::Ownership | Self::Select | Self::Unknown(_)) + } + + pub fn valid_for_secret(&self) -> bool { + matches!( + self, + Self::Ownership | Self::Usage | Self::Update | Self::Unknown(_) + ) + } + + pub fn valid_for_function(&self) -> bool { + matches!(self, Self::Ownership | Self::Usage | Self::Unknown(_)) + } + + /// Whether this privilege can mutate database contents or schema. + /// Unknown privileges are conservatively treated as writes. + pub fn is_write(&self) -> bool { + !matches!(self, Self::Usage | Self::Select | Self::Operate) + } +} + +impl FromStr for Privilege { + type Err = String; + + fn from_str(name: &str) -> Result { + if let Some(privilege) = Self::KNOWN.into_iter().find(|p| p.to_string() == name) { + return Ok(privilege); + } + if let Some(digits) = name.strip_prefix("PRIVILEGE") + && !digits.is_empty() + && digits.bytes().all(|b| b.is_ascii_digit()) + { + return Ok(match digits.parse::() { + Ok(0) => Self::Ownership, + Ok(1) => Self::Usage, + Ok(2) => Self::Select, + Ok(3) => Self::Insert, + Ok(4) => Self::Update, + Ok(5) => Self::Delete, + Ok(6) => Self::CreateDatabase, + Ok(7) => Self::CreateNamespace, + Ok(8) => Self::CreateTable, + Ok(9) => Self::CreateMaterializedView, + Ok(10) => Self::CreateView, + Ok(11) => Self::CreateSecret, + Ok(12) => Self::CreateFunction, + Ok(13) => Self::Operate, + _ => Self::Unknown(UnknownPrivilege(name.to_owned())), + }); + } + Err(format!("invalid privilege {name:?}")) + } +} + +/// Converts a privilege to its numeric identifier. +/// +/// Unknown privileges return an error if their number exceeds [`u8::MAX`]. +/// +/// ``` +/// use lancedb::authz::Privilege; +/// +/// assert_eq!(u8::try_from(&Privilege::Select).unwrap(), 2); +/// let unknown: Privilege = "PRIVILEGE255".parse().unwrap(); +/// assert_eq!(u8::try_from(&unknown).unwrap(), 255); +/// let too_large: Privilege = "PRIVILEGE256".parse().unwrap(); +/// assert!(u8::try_from(&too_large).is_err()); +/// ``` +impl TryFrom<&Privilege> for u8 { + type Error = std::num::ParseIntError; + + fn try_from(privilege: &Privilege) -> Result { + Ok(match privilege { + Privilege::Ownership => 0, + Privilege::Usage => 1, + Privilege::Select => 2, + Privilege::Insert => 3, + Privilege::Update => 4, + Privilege::Delete => 5, + Privilege::CreateDatabase => 6, + Privilege::CreateNamespace => 7, + Privilege::CreateTable => 8, + Privilege::CreateMaterializedView => 9, + Privilege::CreateView => 10, + Privilege::CreateSecret => 11, + Privilege::CreateFunction => 12, + Privilege::Operate => 13, + Privilege::Unknown(unknown) => unknown.0["PRIVILEGE".len()..].parse()?, + }) + } +} + +/// Converts an owned privilege to its numeric identifier. +/// +/// Like the borrowed conversion, this returns an error for numbers above 255. +/// +/// ``` +/// use lancedb::authz::Privilege; +/// +/// assert_eq!(u8::try_from(Privilege::Ownership).unwrap(), 0); +/// ``` +impl TryFrom for u8 { + type Error = std::num::ParseIntError; + + fn try_from(privilege: Privilege) -> Result { + Self::try_from(&privilege) + } +} + +impl Serialize for Privilege { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.collect_str(self) + } +} + +impl<'de> Deserialize<'de> for Privilege { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + String::deserialize(deserializer)? + .parse() + .map_err(D::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::Privilege; + + #[test] + fn display_matches_serde() { + for privilege in Privilege::KNOWN { + assert_eq!( + serde_json::to_string(&privilege).unwrap(), + format!("\"{privilege}\""), + ); + } + } + + #[test] + fn secret_and_function_creation_privileges() { + for (name, privilege) in [ + ("CREATE_SECRET", Privilege::CreateSecret), + ("CREATE_FUNCTION", Privilege::CreateFunction), + ] { + assert_eq!(name.parse::().unwrap(), privilege); + assert_eq!(privilege.to_string(), name); + assert_eq!(serde_json::to_value(&privilege).unwrap(), name); + assert_eq!( + serde_json::from_value::(serde_json::json!(name)).unwrap(), + privilege + ); + assert!(privilege.valid_for_namespace()); + assert!(!privilege.valid_for_system()); + assert!(!privilege.valid_for_database()); + assert!(!privilege.valid_for_table()); + assert!(!privilege.valid_for_view()); + assert!(privilege.is_write()); + } + } + + #[test] + fn numeric_aliases_resolve_to_known_privileges() { + for (number, privilege) in [ + (0, Privilege::Ownership), + (1, Privilege::Usage), + (2, Privilege::Select), + (3, Privilege::Insert), + (4, Privilege::Update), + (5, Privilege::Delete), + (6, Privilege::CreateDatabase), + (7, Privilege::CreateNamespace), + (8, Privilege::CreateTable), + (9, Privilege::CreateMaterializedView), + (10, Privilege::CreateView), + (11, Privilege::CreateSecret), + (12, Privilege::CreateFunction), + (13, Privilege::Operate), + ] { + assert_eq!(u8::try_from(&privilege).unwrap(), number); + assert_eq!(u8::try_from(privilege.clone()).unwrap(), number); + for name in [ + format!("PRIVILEGE{number}"), + format!("PRIVILEGE{number:03}"), + ] { + let parsed = name.parse::().unwrap(); + assert_eq!(parsed, privilege, "{name}"); + assert_eq!( + serde_json::from_value::(serde_json::json!(name)).unwrap(), + privilege + ); + assert_eq!( + serde_json::to_value(&parsed).unwrap(), + privilege.to_string() + ); + assert_eq!( + privilege.to_string().parse::().unwrap(), + privilege + ); + } + } + } + + #[test] + fn unknown_privileges_convert_to_u8() { + for number in 14..=u8::MAX { + for name in [ + format!("PRIVILEGE{number}"), + format!("PRIVILEGE{number:03}"), + ] { + let privilege: Privilege = name.parse().unwrap(); + assert!(matches!(privilege, Privilege::Unknown(_))); + assert_eq!(u8::try_from(&privilege).unwrap(), number); + assert_eq!(u8::try_from(privilege).unwrap(), number); + } + } + } + + #[test] + fn privileges_outside_u8_range_fail_conversion() { + for name in [ + "PRIVILEGE256", + "PRIVILEGE0256", + "PRIVILEGE18446744073709551616", + ] { + let privilege: Privilege = name.parse().unwrap(); + assert_eq!( + u8::try_from(&privilege).unwrap_err().kind(), + &std::num::IntErrorKind::PosOverflow + ); + assert_eq!( + u8::try_from(privilege).unwrap_err().kind(), + &std::num::IntErrorKind::PosOverflow + ); + } + } + + #[test] + fn future_privilege_names_round_trip() { + for name in [ + "PRIVILEGE14", + "PRIVILEGE20", + "PRIVILEGE64", + "PRIVILEGE256", + "PRIVILEGE18446744073709551616", + "PRIVILEGE020", + ] { + let privilege: Privilege = name.parse().unwrap(); + assert!(matches!(privilege, Privilege::Unknown(_))); + assert_eq!(privilege.to_string(), name); + assert_eq!(serde_json::to_value(&privilege).unwrap(), name); + assert_eq!( + serde_json::from_value::(serde_json::json!(name)).unwrap(), + privilege + ); + } + } + + #[test] + fn rejects_invalid_privileges() { + for name in [ + "", + "select", + " SELECT", + "SELECT ", + "PRIVILEGE", + "PRIVILEGE-1", + "PRIVILEGE+1", + "PRIVILEGE 1", + "PRIVILEGE1x", + "privilege12", + "UNKNOWN", + ] { + assert!(name.parse::().is_err(), "{name}"); + assert!( + serde_json::from_value::(serde_json::json!(name)).is_err(), + "{name}" + ); + } + } + + #[test] + fn unknown_privileges_are_valid_for_all_objects_and_treated_as_writes() { + for name in ["PRIVILEGE14", "PRIVILEGE20", "PRIVILEGE64", "PRIVILEGE256"] { + let privilege: Privilege = name.parse().unwrap(); + assert!(privilege.valid_for_system()); + assert!(privilege.valid_for_database()); + assert!(privilege.valid_for_namespace()); + assert!(privilege.valid_for_table()); + assert!(privilege.valid_for_view()); + assert!(privilege.valid_for_secret()); + assert!(privilege.valid_for_function()); + assert!(privilege.is_write()); + } + assert!(!Privilege::Ownership.valid_for_system()); + assert!(!Privilege::Select.valid_for_database()); + } + + #[test] + fn test_privilege_is_write() { + for privilege in Privilege::KNOWN { + assert_eq!( + privilege.is_write(), + matches!( + privilege, + Privilege::Ownership + | Privilege::Insert + | Privilege::Update + | Privilege::Delete + | Privilege::CreateDatabase + | Privilege::CreateNamespace + | Privilege::CreateTable + | Privilege::CreateMaterializedView + | Privilege::CreateView + | Privilege::CreateSecret + | Privilege::CreateFunction + ), + "unexpected write classification for {privilege}" + ); + } + } + + #[test] + fn create_secret_and_create_function_are_namespace_scoped() { + for privilege in [Privilege::CreateSecret, Privilege::CreateFunction] { + assert!(privilege.valid_for_namespace()); + assert!(!privilege.valid_for_system()); + assert!(!privilege.valid_for_database()); + assert!(!privilege.valid_for_table()); + assert!(!privilege.valid_for_view()); + assert!(!privilege.valid_for_secret()); + assert!(!privilege.valid_for_function()); + } + } + + #[test] + fn secret_and_function_privileges() { + assert_eq!( + Privilege::KNOWN + .into_iter() + .filter(Privilege::valid_for_secret) + .collect::>(), + vec![Privilege::Ownership, Privilege::Usage, Privilege::Update] + ); + assert_eq!( + Privilege::KNOWN + .into_iter() + .filter(Privilege::valid_for_function) + .collect::>(), + vec![Privilege::Ownership, Privilege::Usage] + ); + } +} diff --git a/rust/lancedb/src/authz/subject.rs b/rust/lancedb/src/authz/subject.rs new file mode 100644 index 000000000..2cc607fc4 --- /dev/null +++ b/rust/lancedb/src/authz/subject.rs @@ -0,0 +1,378 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Identity selectors. Provider names resolve on the server; IDs remain literal. + +use std::{fmt, str::FromStr}; + +use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _}; + +use crate::{Error, Result}; + +/// The kind of identifier supplied to an authorization operation. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +#[non_exhaustive] +pub enum SubjectKind { + PrincipalId, + PrincipalName, + GroupId, + GroupName, + Role, + PrincipalApiKey, + /// An unrecognized canonical subject. Its full wire string is stored as the value. + Unknown, +} + +impl SubjectKind { + /// Stable name used by language bindings. + pub fn as_str(self) -> &'static str { + match self { + Self::PrincipalId => "principal_id", + Self::PrincipalName => "principal_name", + Self::GroupId => "group_id", + Self::GroupName => "group_name", + Self::Role => "role", + Self::PrincipalApiKey => "principal_api_key", + Self::Unknown => "unknown", + } + } + + fn prefix(self) -> &'static str { + match self { + Self::PrincipalId => "p:", + Self::PrincipalName => "P:", + Self::GroupId => "g:", + Self::GroupName => "G:", + Self::Role => "r:", + Self::PrincipalApiKey => "a:", + Self::Unknown => "", + } + } +} + +const KINDS: [SubjectKind; 6] = [ + SubjectKind::PrincipalId, + SubjectKind::PrincipalName, + SubjectKind::GroupId, + SubjectKind::GroupName, + SubjectKind::Role, + SubjectKind::PrincipalApiKey, +]; + +impl FromStr for SubjectKind { + type Err = Error; + + fn from_str(value: &str) -> Result { + if value == "unknown" { + return Ok(Self::Unknown); + } + KINDS + .into_iter() + .find(|kind| kind.as_str() == value) + .ok_or_else(|| invalid("Invalid subject kind")) + } +} + +/// An immutable validated principal, group, or role selector. +/// +/// Unrecognized canonical prefixes are preserved as [`SubjectKind::Unknown`], +/// with the complete wire string in [`Self::value`], so future subject kinds can +/// be listed and sent back to the server. +/// +/// ``` +/// use lancedb::authz::{Subject, SubjectKind}; +/// let subject = Subject::from_canonical("x:future-subject")?; +/// assert_eq!(subject.kind(), SubjectKind::Unknown); +/// assert_eq!(subject.to_wire(), "x:future-subject"); +/// # Ok::<(), lancedb::Error>(()) +/// ``` +/// +/// Debug and Display redact API keys. [`Self::to_wire`] and serialization +/// intentionally expose the key for transport and must not be logged. +/// +/// ``` +/// use lancedb::authz::Subject; +/// let subject = Subject::group_name("Engineering")?; +/// assert_eq!(subject.to_wire(), "G:Engineering"); +/// assert_eq!(subject, "G:Engineering".parse::()?); +/// # Ok::<(), lancedb::Error>(()) +/// ``` +#[derive(Clone, PartialEq, Eq, Hash)] +pub struct Subject { + kind: SubjectKind, + value: String, +} + +fn invalid(message: &str) -> Error { + // Never include the selector: it may contain an API key. + Error::InvalidInput { + message: message.into(), + } +} + +impl Subject { + /// Construct a selector without resolving it against an identity provider. + pub fn new(kind: SubjectKind, value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() { + return Err(invalid("Subject value must be a nonempty string")); + } + if kind == SubjectKind::Unknown + && (!value + .split_once(':') + .is_some_and(|(prefix, value)| !prefix.is_empty() && !value.is_empty()) + || KINDS.iter().any(|kind| value.starts_with(kind.prefix()))) + { + return Err(invalid("Invalid unknown subject selector")); + } + Ok(Self { kind, value }) + } + + /// Select a principal by canonical ID, without stripping any prefix in the ID. + pub fn principal_id(id: impl Into) -> Result { + Self::new(SubjectKind::PrincipalId, id) + } + /// Select a principal by an unambiguous provider name. + pub fn principal_name(name: impl Into) -> Result { + Self::new(SubjectKind::PrincipalName, name) + } + /// Select a group by canonical ID. + pub fn group_id(id: impl Into) -> Result { + Self::new(SubjectKind::GroupId, id) + } + /// Select a group by an unambiguous provider name. + pub fn group_name(name: impl Into) -> Result { + Self::new(SubjectKind::GroupName, name) + } + /// Select a LanceDB-managed role. + pub fn role(name: impl Into) -> Result { + Self::new(SubjectKind::Role, name) + } + /// Select the principal holding this secret API key, not the request's caller. + pub fn principal_api_key(key: impl Into) -> Result { + Self::new(SubjectKind::PrincipalApiKey, key) + } + + /// The selector kind. + pub fn kind(&self) -> SubjectKind { + self.kind + } + /// The selected ID, name, role, or API key. API-key values are secrets. + /// For unknown kinds, this is the full canonical wire string. + pub fn value(&self) -> &str { + &self.value + } + /// The value safe for display; API keys are replaced with ``. + pub fn display_value(&self) -> &str { + if self.kind == SubjectKind::PrincipalApiKey { + "" + } else { + &self.value + } + } + /// Encode the literal wire selector. The result can contain a secret API key. + pub fn to_wire(&self) -> String { + format!("{}{}", self.kind.prefix(), self.value) + } + + /// Parse a returned canonical identity, rejecting names and secret selectors. + pub fn from_canonical(value: &str) -> Result { + let subject: Self = value.parse()?; + if !matches!( + subject.kind, + SubjectKind::PrincipalId + | SubjectKind::GroupId + | SubjectKind::Role + | SubjectKind::Unknown + ) { + return Err(invalid("Server returned a non-canonical subject")); + } + Ok(subject) + } + + /// Deserialize a canonical response subject, rejecting secrets and name selectors. + pub fn deserialize_canonical<'de, D: Deserializer<'de>>( + deserializer: D, + ) -> std::result::Result { + Self::from_canonical(&String::deserialize(deserializer)?).map_err(D::Error::custom) + } + + /// Require a principal selector for a principal lookup. + pub fn validate_principal(&self) -> Result<()> { + if matches!( + self.kind, + SubjectKind::PrincipalId | SubjectKind::PrincipalName | SubjectKind::PrincipalApiKey + ) { + Ok(()) + } else { + Err(invalid("Subject kind is not valid for a principal lookup")) + } + } + /// Require a group selector for a group lookup. + pub fn validate_group(&self) -> Result<()> { + if matches!(self.kind, SubjectKind::GroupId | SubjectKind::GroupName) { + Ok(()) + } else { + Err(invalid("Subject kind is not valid for a group lookup")) + } + } + /// Require a principal/group selector. Roles cannot contain other roles. + pub fn validate_role_member(&self) -> Result<()> { + if self.kind == SubjectKind::Role { + Err(invalid("Roles cannot be members of other roles")) + } else { + Ok(()) + } + } +} + +impl FromStr for Subject { + type Err = Error; + fn from_str(value: &str) -> Result { + for kind in KINDS { + if let Some(value) = value.strip_prefix(kind.prefix()) { + return Self::new(kind, value); + } + } + Self::new(SubjectKind::Unknown, value) + } +} + +impl fmt::Debug for Subject { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Subject") + .field("kind", &self.kind) + .field("value", &self.display_value()) + .finish() + } +} +impl fmt::Display for Subject { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}{}", self.kind.prefix(), self.display_value()) + } +} +impl Serialize for Subject { + fn serialize(&self, serializer: S) -> std::result::Result { + serializer.serialize_str(&self.to_wire()) + } +} +impl<'de> Deserialize<'de> for Subject { + fn deserialize>(deserializer: D) -> std::result::Result { + String::deserialize(deserializer)? + .parse() + .map_err(D::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn selectors_roundtrip_without_normalizing_values() { + for kind in KINDS { + for value in [ + "alice@example.com", + "p:literal-id", + "Engineering / R&D", + "名前", + ] { + let subject = Subject::new(kind, value).unwrap(); + assert_eq!(subject.to_wire().parse::().unwrap(), subject); + let json = serde_json::to_string(&subject).unwrap(); + assert_eq!(serde_json::from_str::(&json).unwrap(), subject); + assert_eq!(subject.value(), value); + assert!(Subject::new(kind, "").is_err()); + } + } + } + + #[test] + fn contexts_and_canonical_ids() { + for kind in KINDS { + let subject = Subject::new(kind, "value").unwrap(); + assert_eq!( + subject.validate_principal().is_ok(), + matches!( + kind, + SubjectKind::PrincipalId + | SubjectKind::PrincipalName + | SubjectKind::PrincipalApiKey + ) + ); + assert_eq!( + subject.validate_group().is_ok(), + matches!(kind, SubjectKind::GroupId | SubjectKind::GroupName) + ); + assert_eq!( + subject.validate_role_member().is_ok(), + kind != SubjectKind::Role + ); + assert_eq!( + Subject::from_canonical(&subject.to_wire()).is_ok(), + matches!( + kind, + SubjectKind::PrincipalId | SubjectKind::GroupId | SubjectKind::Role + ) + ); + } + } + + #[test] + fn unknown_subjects_preserve_canonical_wire_strings() { + for wire in ["x:future-subject", "future:literal:identifier"] { + let subject = Subject::from_canonical(wire).unwrap(); + assert_eq!(subject.kind(), SubjectKind::Unknown); + assert_eq!(subject.value(), wire); + assert_eq!(subject.to_wire(), wire); + assert_eq!(serde_json::to_value(&subject).unwrap(), wire); + assert_eq!( + serde_json::from_value::(serde_json::json!(wire)).unwrap(), + subject + ); + assert!(subject.validate_principal().is_err()); + assert!(subject.validate_group().is_err()); + } + for wire in [ + "", + "no-prefix", + ":value", + "x:", + "p:value", + "P:value", + "g:value", + "G:value", + "r:value", + "a:secret", + ] { + assert!(Subject::new(SubjectKind::Unknown, wire).is_err()); + } + } + + #[test] + fn acl_page_preserves_known_and_unknown_subjects() { + let payload = serde_json::json!({"entries": [ + {"subject_id": "p:known", "object": "system", "privilege": "USAGE"}, + {"subject_id": "x:future-subject", "object": "system", "privilege": "USAGE"} + ]}); + let page: crate::authz::ListAccessControlEntriesResponse = + serde_json::from_value(payload.clone()).unwrap(); + assert_eq!(page.entries.len(), 2); + assert_eq!(page.entries[0].subject_id.kind(), SubjectKind::PrincipalId); + assert_eq!(page.entries[1].subject_id.kind(), SubjectKind::Unknown); + assert_eq!(serde_json::to_value(page).unwrap(), payload); + } + + #[test] + fn secrets_are_only_exposed_by_explicit_serialization() { + let secret = "never-log-this-key"; + let subject = Subject::principal_api_key(secret).unwrap(); + assert!(!format!("{subject:?} {subject}").contains(secret)); + assert!(serde_json::to_string(&subject).unwrap().contains(secret)); + let error = Subject::from_canonical(&subject.to_wire()).unwrap_err(); + assert!(!error.to_string().contains(secret)); + assert!(!format!("{error:?}").contains(secret)); + let error = secret.parse::().unwrap_err(); + assert!(!error.to_string().contains(secret)); + } +} diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index cbc0724c6..ca68164e4 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -247,8 +247,10 @@ impl BlobRangeRequest { /// `Struct` with the `lance.blob.v2` marker. Same layout Lance /// expects on write. /// -/// A blob column may be top-level or nested inside a struct or list. Nested -/// blobs are addressed by a dotted path (e.g. `info.blob`) in the read APIs. +/// A blob column may be top-level or nested inside a struct or list. Blobs +/// nested inside structs use dotted paths (e.g. `info.blob`) in the read APIs. +/// Blobs inside lists can be written and queried as descriptors, but cannot +/// currently be fetched by row id. /// /// ``` /// use arrow_schema::{DataType, Field, Schema}; @@ -262,6 +264,37 @@ pub fn blob(name: impl AsRef, nullable: bool) -> Field { lance::blob::blob_field(name.as_ref(), nullable) } +/// Options for creating a blob v2 column field with tier thresholds. +/// +/// Re-exported from lance so the thresholds are built exactly the way +/// [`lance::blob::blob_field_with_options`] builds them. Note that lance does +/// not currently expose a pack file size threshold option, so only the inline +/// and dedicated size thresholds are supported here. +pub use lance::blob::BlobFieldOptions as BlobOptions; + +/// Creates an Arrow field for a Lance blob v2 column with tier thresholds. +/// +/// Same as [`blob`], but forwards `options` to +/// [`lance::blob::blob_field_with_options`], storing the inline and dedicated +/// size thresholds as field metadata. With [`BlobOptions::default()`] this is +/// identical to [`blob`]. +/// +/// ``` +/// use std::num::NonZeroUsize; +/// +/// use lancedb::blob::BlobOptions; +/// +/// let options = BlobOptions { +/// inline_size_threshold: Some(8 * 1024), +/// dedicated_size_threshold: NonZeroUsize::new(1 << 20), +/// }; +/// let field = lancedb::blob::blob_with_options("payload", true, options); +/// assert!(lancedb::blob::is_blob(&field)); +/// ``` +pub fn blob_with_options(name: impl AsRef, nullable: bool, options: BlobOptions) -> Field { + lance::blob::blob_field_with_options(name.as_ref(), nullable, options) +} + /// Returns true if `field` is a blob v2 column. /// /// ``` @@ -286,7 +319,7 @@ fn field_tree_has_blob_v2(field: &Field) -> bool { } } -/// Collects the dotted paths of blob v2 columns under `field`, into `paths`. +/// Collects the dotted paths of row-addressable blob v2 columns under `field`. fn collect_blob_paths(field: &Field, prefix: &str, paths: &mut Vec) { let path = if prefix.is_empty() { field.name().clone() @@ -297,26 +330,58 @@ fn collect_blob_paths(field: &Field, prefix: &str, paths: &mut Vec) { paths.push(path); return; } - match field.data_type() { - DataType::Struct(children) => { - for child in children { - collect_blob_paths(child, &path, paths); - } + if let DataType::Struct(children) = field.data_type() { + for child in children { + collect_blob_paths(child, &path, paths); } - DataType::List(child) | DataType::LargeList(child) | DataType::FixedSizeList(child, _) => { - collect_blob_paths(child, &path, paths) - } - _ => {} } } +fn is_list_blob_path(field: &Field, prefix: &str, column: &str, inside_list: bool) -> bool { + let path = if prefix.is_empty() { + field.name().clone() + } else { + format!("{prefix}.{}", field.name()) + }; + if field.is_blob_v2() { + return inside_list && column == path; + } + match field.data_type() { + DataType::Struct(children) => children + .iter() + .any(|child| is_list_blob_path(child, &path, column, inside_list)), + DataType::List(child) | DataType::LargeList(child) | DataType::FixedSizeList(child, _) => { + (column == path && field_tree_has_blob_v2(child)) + || is_list_blob_path(child, &path, column, true) + } + _ => false, + } +} + +/// Give list blob paths a useful error before they reach the row-based reader. +pub(crate) fn reject_list_blob_column(schema: &Schema, column: &str) -> Result<()> { + if schema + .fields() + .iter() + .any(|field| is_list_blob_path(field, "", column, false)) + { + return Err(Error::InvalidInput { + message: format!( + "column '{column}' contains a blob inside a list; blobs inside lists cannot be fetched by row id yet" + ), + }); + } + Ok(()) +} + /// Returns true if `schema` declares any blob v2 column, including nested ones. pub(crate) fn has_blob_columns(schema: &Schema) -> bool { schema.fields().iter().any(|f| field_tree_has_blob_v2(f)) } -/// Blob v2 column paths in `schema`, declaration order preserved. Nested blobs -/// are dotted paths (e.g. `info.blob`). +/// Row-addressable blob v2 column paths in `schema`, in declaration order. +/// Struct-nested blobs use dotted paths (e.g. `info.blob`); list-nested blobs +/// are omitted because a row id cannot identify an individual list element. pub(crate) fn blob_column_names(schema: &Schema) -> Vec { let mut paths = Vec::new(); for field in schema.fields() { @@ -351,6 +416,11 @@ pub(crate) fn ensure_blob_v2_column( schema: &lance_core::datatypes::Schema, column: &str, ) -> Result<()> { + // A top-level blob cannot have a list ancestor. Keep the common read path + // from converting the entire schema just to validate that condition. + if column.contains('.') || schema.field(column).is_none_or(|field| !field.is_blob_v2()) { + reject_list_blob_column(&Schema::from(schema), column)?; + } match schema.field(column) { Some(field) if field.is_blob_v2() => Ok(()), Some(field) if field.is_blob() => Err(Error::InvalidInput { @@ -368,45 +438,115 @@ pub(crate) fn ensure_blob_v2_column( } } -fn ensure_all_row_ids_resolved(column: &str, requested: usize, resolved: usize) -> Result<()> { - if requested == resolved { +fn missing_blob_row_ids_error(column: &str, requested: usize, missing: &[u64]) -> Error { + let detail = if missing.is_empty() { + String::new() + } else { + format!("; first missing row ids: {missing:?}") + }; + Error::InvalidInput { + message: format!( + "blob read for column '{column}' requested {requested} row ids but some \ + do not exist in the table{detail}; pass row ids collected from this table" + ), + } +} + +/// Diagnose only after a failed read. Stable row IDs use Lance's deletion-aware +/// index; physical row IDs are fragment/offset pairs. Limit the error to three +/// distinct IDs so a large failed batch stays readable. +async fn find_missing_blob_row_ids(dataset: &Dataset, row_ids: &[u64]) -> Result> { + let mut missing = Vec::new(); + if let Some(index) = lance::dataset::rowids::get_row_id_index(dataset).await? { + for (row_id, address) in row_ids.iter().zip(index.get_many(row_ids)?) { + if address.is_none() && !missing.contains(row_id) { + missing.push(*row_id); + if missing.len() == 3 { + break; + } + } + } + } else { + let mut fragments = std::collections::HashMap::new(); + for &row_id in row_ids { + let fragment_id = row_id >> 32; + let offset = row_id as u32; + if let std::collections::hash_map::Entry::Vacant(entry) = fragments.entry(fragment_id) { + let fragment = dataset.get_fragment(fragment_id as usize); + let state = if let Some(fragment) = fragment { + Some(( + fragment.physical_rows().await?, + fragment.get_deletion_vector().await?, + )) + } else { + None + }; + entry.insert(state); + } + let is_missing = match fragments.get(&fragment_id) { + Some(Some((physical_rows, deletions))) => { + offset as usize >= *physical_rows + || deletions + .as_ref() + .is_some_and(|deletions| deletions.contains(offset)) + } + Some(None) => true, + None => unreachable!(), + }; + if is_missing && !missing.contains(&row_id) { + missing.push(row_id); + if missing.len() == 3 { + break; + } + } + } + } + Ok(missing) +} + +async fn ensure_all_row_ids_resolved( + dataset: &Dataset, + column: &str, + row_ids: &[u64], + resolved: usize, +) -> Result<()> { + if row_ids.len() == resolved { return Ok(()); } - if resolved < requested { - Err(Error::InvalidInput { - message: format!( - "blob read for column '{column}' requested {requested} row ids but only {resolved} \ - exist in the table; pass row ids collected from this table" - ), - }) + if resolved < row_ids.len() { + let missing = find_missing_blob_row_ids(dataset, row_ids) + .await + .unwrap_or_default(); + Err(missing_blob_row_ids_error(column, row_ids.len(), &missing)) } else { Err(Error::Runtime { message: format!( - "blob read for column '{column}' returned {resolved} results for {requested} row ids" + "blob read for column '{column}' returned {resolved} results for {} row ids", + row_ids.len() ), }) } } -/// Lance take reports a missing physical row address as NotSupported or InvalidInput. -fn map_blob_take_error(column: &str, requested: usize, err: lance::Error) -> Error { - let missing_row_addr = match &err { - lance::Error::NotSupported { source, .. } => { - source.to_string().contains("must not target deleted rows") - } - lance::Error::InvalidInput { source, .. } => source - .to_string() - .contains("belongs to non-existent fragment"), - _ => false, - }; +/// Lance reports missing physical and stable row ids through several error paths +/// without a shared typed error. +async fn map_blob_take_error( + dataset: &Dataset, + column: &str, + row_ids: &[u64], + err: lance::Error, +) -> Error { + let message = err.to_string(); + let missing_row_id = message.contains("must not target deleted rows") + || message.contains("belongs to non-existent fragment") + || (message.contains("Invalid read params ") && message.contains("addressable rows")) + || message.contains("Could not resolve all requested row IDs"); - if missing_row_addr { - Error::InvalidInput { - message: format!( - "blob read for column '{column}' requested {requested} row ids but some \ - do not exist in the table; pass row ids collected from this table" - ), - } + if missing_row_id { + let missing = find_missing_blob_row_ids(dataset, row_ids) + .await + .unwrap_or_default(); + missing_blob_row_ids_error(column, row_ids.len(), &missing) } else { err.into() } @@ -423,18 +563,25 @@ pub(crate) async fn take_blob_ranges_aligned( return Ok(LargeBinaryBuilder::new().finish()); } + let row_ids = requests + .iter() + .map(|request| request.row_id) + .collect::>(); let lance_requests = requests .iter() .map(|request| LanceBlobRangeRequest::new(request.row_id, request.offset, request.length)) .collect::>(); - let payloads = dataset + let payloads = match dataset .read_blob_ranges(column)? .with_row_ids(lance_requests) .preserve_order(true) .execute() .await - .map_err(|err| map_blob_take_error(column, requests.len(), err))?; - ensure_all_row_ids_resolved(column, requests.len(), payloads.len())?; + { + Ok(payloads) => payloads, + Err(err) => return Err(map_blob_take_error(dataset, column, &row_ids, err).await), + }; + ensure_all_row_ids_resolved(dataset, column, &row_ids, payloads.len()).await?; let mut builder = LargeBinaryBuilder::new(); for payload in payloads { @@ -457,14 +604,17 @@ pub(crate) async fn take_blobs_aligned( return Ok(LargeBinaryBuilder::new().finish()); } - let payloads = dataset + let payloads = match dataset .read_blobs(column)? .with_row_ids(row_ids.to_vec()) .preserve_order(true) .execute() .await - .map_err(|err| map_blob_take_error(column, row_ids.len(), err))?; - ensure_all_row_ids_resolved(column, row_ids.len(), payloads.len())?; + { + Ok(payloads) => payloads, + Err(err) => return Err(map_blob_take_error(dataset, column, row_ids, err).await), + }; + ensure_all_row_ids_resolved(dataset, column, row_ids, payloads.len()).await?; let mut builder = LargeBinaryBuilder::new(); for payload in payloads { @@ -487,11 +637,11 @@ pub(crate) async fn take_blob_files_aligned( return Ok(Vec::new()); } - let handles = dataset - .take_blobs(row_ids, column) - .await - .map_err(|err| map_blob_take_error(column, row_ids.len(), err))?; - ensure_all_row_ids_resolved(column, row_ids.len(), handles.len())?; + let handles = match dataset.take_blobs(row_ids, column).await { + Ok(handles) => handles, + Err(err) => return Err(map_blob_take_error(dataset, column, row_ids, err).await), + }; + ensure_all_row_ids_resolved(dataset, column, row_ids, handles.len()).await?; Ok(handles .into_iter() .map(|handle| handle.map(Into::into)) @@ -501,8 +651,13 @@ pub(crate) async fn take_blob_files_aligned( #[cfg(test)] mod tests { use super::*; + use std::num::NonZeroUsize; + use arrow_schema::DataType; - use lance_arrow::ARROW_EXT_NAME_KEY; + use lance_arrow::{ + ARROW_EXT_NAME_KEY, BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, + BLOB_INLINE_SIZE_THRESHOLD_META_KEY, + }; fn blob_schema() -> Schema { Schema::new(vec![ @@ -521,6 +676,38 @@ mod tests { assert!(matches!(field.data_type(), DataType::Struct(_))); } + #[test] + fn blob_with_options_sets_threshold_metadata() { + let options = BlobOptions { + inline_size_threshold: Some(8 * 1024), + dedicated_size_threshold: NonZeroUsize::new(1 << 20), + }; + let field = blob_with_options("payload", true, options); + assert!(is_blob(&field)); + assert_eq!( + field + .metadata() + .get(BLOB_INLINE_SIZE_THRESHOLD_META_KEY) + .map(String::as_str), + Some("8192") + ); + assert_eq!( + field + .metadata() + .get(BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY) + .map(String::as_str), + Some("1048576") + ); + } + + #[test] + fn blob_with_default_options_matches_plain_blob() { + let field = blob_with_options("payload", true, BlobOptions::default()); + let plain = blob("payload", true); + assert_eq!(field.metadata(), plain.metadata()); + assert_eq!(field.data_type(), plain.data_type()); + } + #[test] fn has_blob_columns_detects_blob_fields() { assert!(has_blob_columns(&blob_schema())); @@ -619,6 +806,54 @@ mod tests { assert_eq!(blob_column_names(&schema), vec!["info.blob"]); } + #[test] + fn list_blob_paths_are_not_row_addressable() { + let schema = Schema::new(vec![ + blob("image", true), + Field::new( + "info", + DataType::Struct(vec![blob("thumbnail", true)].into()), + true, + ), + Field::new("images", DataType::List(Arc::new(blob("item", true))), true), + Field::new( + "groups", + DataType::LargeList(Arc::new(Field::new( + "group", + DataType::Struct(vec![blob("photo", true)].into()), + true, + ))), + true, + ), + Field::new( + "frames", + DataType::FixedSizeList(Arc::new(blob("frame", true)), 2), + true, + ), + ]); + assert!(has_blob_columns(&schema)); + assert_eq!(blob_column_names(&schema), vec!["image", "info.thumbnail"]); + + for column in [ + "images", + "images.item", + "groups", + "groups.group.photo", + "frames", + "frames.frame", + ] { + let err = reject_list_blob_column(&schema, column).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + assert!( + err.to_string() + .contains("blobs inside lists cannot be fetched") + ); + } + for column in ["image", "info.thumbnail", "info"] { + reject_list_blob_column(&schema, column).unwrap(); + } + } + #[test] fn storage_version_noop_without_blob_columns() { let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]); diff --git a/rust/lancedb/src/catalog.rs b/rust/lancedb/src/catalog.rs index 576eb8e88..ad8674b77 100644 --- a/rust/lancedb/src/catalog.rs +++ b/rust/lancedb/src/catalog.rs @@ -193,6 +193,13 @@ impl Stream for DatabaseNames { /// A backend that manages databases. Implementations own database lifecycle semantics. #[async_trait::async_trait] pub trait Catalog: Send + Sync + std::fmt::Debug + 'static { + /// Obtain deployment-wide authorization, if supported by this backend. + fn authz(&self) -> Result> { + Err(crate::Error::NotSupported { + message: "Authorization administration is not supported by this catalog".into(), + }) + } + /// Catalog endpoint or location. fn uri(&self) -> &str; /// Create a database, or open it when `exist_ok` permits. @@ -239,6 +246,11 @@ impl CatalogConnection { pub fn uri(&self) -> &str { self.catalog.uri() } + /// Obtain a handle sharing this catalog's deployment scope and credentials. + pub fn authz(&self) -> Result> { + self.catalog.authz() + } + /// Access the underlying backend. pub fn catalog(&self) -> &Arc { &self.catalog diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index ff01f4109..8a0e489c2 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -38,10 +38,7 @@ use crate::remote::{ }, }; use crate::secrets::SecretInfo; -use crate::utils::{ - validate_namespace, validate_secret_component, validate_secret_reference, - validate_view_reference, -}; +use crate::utils::{validate_namespace, validate_secret_reference, validate_view_reference}; use crate::view::ViewDescription; use lance::io::ObjectStoreParams; pub use lance_file::version::LanceFileVersion; @@ -811,9 +808,7 @@ impl Connection { let database = database.clone(); let namespace_path = namespace_path.clone(); async move { - for segment in &namespace_path { - validate_secret_component("Secret namespace path segment", segment)?; - } + validate_namespace(&namespace_path)?; database.list_secrets(&namespace_path, options).await } }) diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index f6a2df9a6..6052b67d8 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -127,7 +127,7 @@ pub enum Error { "Hit retry limit for request_id={request_id} (\ request_failures={request_failures}/{max_request_failures}, \ connect_failures={connect_failures}/{max_connect_failures}, \ - read_failures={read_failures}/{max_read_failures})" + read_failures={read_failures}/{max_read_failures}): {source}" ))] Retry { request_id: String, diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index f05633281..98d915118 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -453,7 +453,8 @@ pub struct FunctionVersion { signature: FunctionSignature, #[serde(default, skip_serializing_if = "Vec::is_empty")] secret_bindings: Vec, - created_at: String, + /// When this version was created, in milliseconds since the epoch. + created_at_millis: i64, metadata: BTreeMap, disabled: bool, } @@ -507,8 +508,9 @@ impl FunctionVersion { &self.secret_bindings } - pub fn created_at(&self) -> &str { - &self.created_at + /// When this version was created, in milliseconds since the epoch. + pub fn created_at_millis(&self) -> i64 { + self.created_at_millis } } diff --git a/rust/lancedb/src/index.rs b/rust/lancedb/src/index.rs index cdc973044..82f8ad845 100644 --- a/rust/lancedb/src/index.rs +++ b/rust/lancedb/src/index.rs @@ -75,6 +75,17 @@ pub enum Index { /// Full text search index using BM25. /// + /// Supports strings, lists of strings, and `lance.json` logical fields + /// stored as `LargeBinary`. Raw binary fields are not supported. Lance + /// selects its JSON-aware tokenizer from the field metadata. + /// + /// JSON queries use `path,str,value` triples, where the dotted path is + /// relative to the indexed column. For example, `meta.tag,str,nature` + /// searches the string value at `$.meta.tag`, not other paths or key names. + /// Create the index on the JSON column, not on a JSON subpath. Plain text + /// queries are invalid for this tokenizer; malformed triples can panic in + /// the current Lance dependency (see [lance#8812](https://github.com/lance-format/lance/issues/8812)). + /// /// The posting block size defaults to 128. Supported values are 128 and 256; /// a value of 256 uses the experimental FTS V3 format and may introduce /// breaking changes. diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index e215f3c6b..7750f66b5 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -173,6 +173,7 @@ #![recursion_limit = "256"] pub mod arrow; +pub mod authz; pub mod blob; pub mod catalog; pub mod connection; diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 6b20fb123..07db69bb3 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1114,10 +1114,24 @@ pub struct VectorQueryRequest { pub column: Option, /// The vector(s) to search for pub query_vector: Vec>, - /// The minimum number of partitions to search - pub minimum_nprobes: usize, - /// The maximum number of partitions to search + /// The legacy number of partitions to search. + /// + /// Lance sets both probe bounds to this value. Explicit bounds can override it. + pub nprobes: Option, + /// The minimum number of partitions to search. + /// + /// If unset, Lance's default is used. + pub minimum_nprobes: Option, + /// The maximum number of partitions to search. + /// + /// If unset, Lance's default is used. Zero represents an explicitly unbounded + /// maximum on request and wire representations. pub maximum_nprobes: Option, + /// Optional initial IVF search effort in [0, 1]. + /// + /// None preserves Auto and fixed `nprobes`; any explicit value overrides + /// `nprobes` while retaining explicit minimum and maximum bounds. + pub search_effort: Option, /// The lower bound (inclusive) of the distance to search for. pub lower_bound: Option, /// The upper bound (exclusive) of the distance to search for. @@ -1141,8 +1155,10 @@ impl Default for VectorQueryRequest { base: QueryRequest::default(), column: None, query_vector: Vec::new(), - minimum_nprobes: 20, - maximum_nprobes: Some(20), + nprobes: None, + minimum_nprobes: None, + maximum_nprobes: None, + search_effort: None, lower_bound: None, upper_bound: None, ef: None, @@ -1155,6 +1171,17 @@ impl Default for VectorQueryRequest { } impl VectorQueryRequest { + pub(crate) fn validate_search_effort(&self) -> Result<()> { + if let Some(effort) = self.search_effort + && (!effort.is_finite() || !(0.0..=1.0).contains(&effort)) + { + return Err(Error::InvalidInput { + message: format!("search_effort must be finite and in [0, 1], got {effort}"), + }); + } + Ok(()) + } + fn add_query_vector(&mut self, vector: impl IntoQueryVector) -> Result<()> { let vector = vector.to_query_vector(&DataType::Float32, "default")?; let array = @@ -1248,7 +1275,7 @@ impl VectorQuery { Ok(self) } - /// Set the number of partitions to search (probe) + /// Set the legacy IVF probe parameter /// /// This argument is only used when the vector column has an IVF PQ index. /// If there is no index then this value is ignored. @@ -1258,19 +1285,20 @@ impl VectorQuery { /// /// The partition whose centroids are closest to the query vector will be /// exhaustiely searched to find matches. This parameter controls how many - /// partitions should be searched. + /// partitions may be searched. /// /// Increasing this value will increase the recall of your query but will - /// also increase the latency of your query. The default value is 20. This - /// default is good for many cases but the best value to use will depend on - /// your data and the recall that you need to achieve. + /// also increase the latency of your query. If this method is not called, + /// Lance's adaptive probe defaults are used. /// /// For best results we recommend tuning this parameter with a benchmark against /// your actual data to find the smallest possible value that will still give /// you the desired recall. /// - /// This method sets both the minimum and maximum number of partitions to search. - /// For more fine-grained control see [`VectorQuery::minimum_nprobes`] and + /// LanceDB retains this as `nprobes` through local and remote request + /// construction. Lance sets both probe bounds to this value only when + /// [`Self::search_effort`] is unset. For more + /// fine-grained control see [`VectorQuery::minimum_nprobes`] and /// [`VectorQuery::maximum_nprobes`]. /// /// Returns an error if `nprobes` is not greater than 0. @@ -1280,8 +1308,32 @@ impl VectorQuery { message: "nprobes must be greater than 0".to_string(), }); } - self.request.minimum_nprobes = nprobes; - self.request.maximum_nprobes = Some(nprobes); + self.request.nprobes = Some(nprobes); + Ok(self) + } + + /// Set the initial IVF search effort to a finite value in [0, 1]. + /// + /// None (the default) resolves to 0.5, preserving Auto probing. At 0 the + /// search starts at the caller minimum; at 1 it searches all partitions + /// within the caller maximum. Intermediate values interpolate around Auto. + /// Later probing may continue when too few candidates are found. + /// Any explicit value, including 0.5, overrides [`Self::nprobes`] regardless + /// of setter order. Explicit minimum and maximum bounds still apply. + /// This does not change refinement or HNSW ef and is not a recall guarantee. + /// Lance's MemWAL LSM scanner does not yet support this option. + /// + /// ``` + /// # use lancedb::{Table, query::{QueryBase, ExecutableQuery}}; + /// # async fn example(table: &Table) -> lancedb::Result<()> { + /// let results = table.query().nearest_to([0.1, 0.2].as_slice())? + /// .search_effort(0.75)?.execute().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn search_effort(mut self, search_effort: impl Into>) -> Result { + self.request.search_effort = search_effort.into(); + self.request.validate_search_effort()?; Ok(self) } @@ -1303,6 +1355,7 @@ impl VectorQuery { }); } if let Some(maximum_nprobes) = self.request.maximum_nprobes + && maximum_nprobes != 0 && minimum_nprobes > maximum_nprobes { return Err(Error::InvalidInput { @@ -1310,7 +1363,7 @@ impl VectorQuery { .to_string(), }); } - self.request.minimum_nprobes = minimum_nprobes; + self.request.minimum_nprobes = Some(minimum_nprobes); Ok(self) } @@ -1323,11 +1376,13 @@ impl VectorQuery { /// /// If this value is greater than minimum_nprobes then the excess partitions will /// only be searched if the initial search does not return enough results. + /// The bound applies per index segment to initial and later probing without + /// disabling Auto probing. /// /// This can be useful when there is a narrow filter to allow these queries to /// spend more time searching and avoid potential false negatives. /// - /// Set to None to search all partitions, if needed, to satisfy the limit + /// Set to None to use Lance's default maximum probe behavior. pub fn maximum_nprobes(mut self, maximum_nprobes: Option) -> Result { if let Some(maximum_nprobes) = maximum_nprobes { if maximum_nprobes == 0 { @@ -1335,14 +1390,18 @@ impl VectorQuery { message: "maximum_nprobes must be greater than 0".to_string(), }); } - if maximum_nprobes < self.request.minimum_nprobes { + if let Some(minimum_nprobes) = self.request.minimum_nprobes + && maximum_nprobes < minimum_nprobes + { return Err(Error::InvalidInput { message: "maximum_nprobes must be greater than or equal to minimum_nprobes" .to_string(), }); } } - self.request.maximum_nprobes = maximum_nprobes; + // Preserve an explicit reset separately from an untouched maximum. The + // request protocols already use zero to represent an unbounded maximum. + self.request.maximum_nprobes = Some(maximum_nprobes.unwrap_or(0)); Ok(self) } @@ -2431,6 +2490,29 @@ mod tests { let vector = Float32Array::from_iter_values([0.1, 0.2]); let query = table.query().nearest_to(&[0.1, 0.2]).unwrap(); + assert_eq!(query.request.nprobes, None); + assert_eq!(query.request.minimum_nprobes, None); + assert_eq!(query.request.maximum_nprobes, None); + assert_eq!(query.request.search_effort, None); + + let effort_query = query + .clone() + .nprobes(20) + .unwrap() + .search_effort(0.5) + .unwrap(); + assert_eq!(effort_query.request.search_effort, Some(0.5)); + assert_eq!(effort_query.request.nprobes, Some(20)); + assert_eq!(effort_query.request.minimum_nprobes, None); + assert_eq!(effort_query.request.maximum_nprobes, None); + assert_eq!( + effort_query + .search_effort(None) + .unwrap() + .request + .search_effort, + None + ); assert_eq!( *query .request @@ -2442,6 +2524,37 @@ mod tests { vector ); + let minimum_only_query = query.clone().minimum_nprobes(20).unwrap(); + assert_eq!(minimum_only_query.request.minimum_nprobes, Some(20)); + assert_eq!(minimum_only_query.request.maximum_nprobes, None); + + let maximum_only_query = query.maximum_nprobes(Some(20)).unwrap(); + assert_eq!(maximum_only_query.request.minimum_nprobes, None); + assert_eq!(maximum_only_query.request.maximum_nprobes, Some(20)); + + let explicitly_unbounded_query = table + .query() + .nearest_to(&[0.1, 0.2]) + .unwrap() + .nprobes(20) + .unwrap() + .maximum_nprobes(None) + .unwrap(); + assert_eq!(explicitly_unbounded_query.request.nprobes, Some(20)); + assert_eq!(explicitly_unbounded_query.request.maximum_nprobes, Some(0)); + + let bounded_query = table + .query() + .nearest_to(&[0.1, 0.2]) + .unwrap() + .minimum_nprobes(5) + .unwrap() + .nprobes(20) + .unwrap(); + assert_eq!(bounded_query.request.nprobes, Some(20)); + assert_eq!(bounded_query.request.minimum_nprobes, Some(5)); + assert_eq!(bounded_query.request.maximum_nprobes, None); + let new_vector = Float32Array::from_iter_values([9.8, 8.7]); let query = table @@ -2450,7 +2563,7 @@ mod tests { .offset(1) .nearest_to(&[9.8, 8.7]) .unwrap() - .nprobes(1000) + .nprobes(20) .unwrap() .postfilter() .distance_type(DistanceType::Cosine) @@ -2469,8 +2582,9 @@ mod tests { ); assert_eq!(query.request.base.limit.unwrap(), 100); assert_eq!(query.request.base.offset.unwrap(), 1); - assert_eq!(query.request.minimum_nprobes, 1000); - assert_eq!(query.request.maximum_nprobes, Some(1000)); + assert_eq!(query.request.nprobes, Some(20)); + assert_eq!(query.request.minimum_nprobes, None); + assert_eq!(query.request.maximum_nprobes, None); assert!(query.request.use_index); assert_eq!(query.request.distance_type, Some(DistanceType::Cosine)); assert_eq!(query.request.approx_mode, Some(ApproxMode::Accurate)); @@ -2492,13 +2606,13 @@ mod tests { Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0" )); - // Fixed probe counts replace both bounds, even when moving below or - // above the previous range. + // Legacy probe counts remain independent until Lance resolves them. let mut query = query; for nprobes in [30, 1, 50] { query = query.nprobes(nprobes).unwrap(); - assert_eq!(query.request.minimum_nprobes, nprobes); - assert_eq!(query.request.maximum_nprobes, Some(nprobes)); + assert_eq!(query.request.nprobes, Some(nprobes)); + assert_eq!(query.request.minimum_nprobes, None); + assert_eq!(query.request.maximum_nprobes, None); } } diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index bc534c0f4..a6b2f0e62 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -6,6 +6,7 @@ //! building client/server applications with LanceDB or as a client for some //! other custom LanceDB service. +mod authz; pub mod catalog; pub(crate) mod client; pub(crate) mod db; diff --git a/rust/lancedb/src/remote/authz.rs b/rust/lancedb/src/remote/authz.rs new file mode 100644 index 000000000..88c2841ce --- /dev/null +++ b/rust/lancedb/src/remote/authz.rs @@ -0,0 +1,254 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::client::{HttpSend, RequestResultExt, RestfulLanceDbClient, Sender}; +use crate::Result; +use crate::authz::*; +use serde::{Serialize, de::DeserializeOwned}; + +pub struct RemoteAuthorization { + pub(crate) client: RestfulLanceDbClient, +} + +impl std::fmt::Debug for RemoteAuthorization { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RemoteAuthorization") + .finish_non_exhaustive() + } +} + +impl RemoteAuthorization { + async fn post( + &self, + route: &str, + request: &T, + ) -> Result<(String, reqwest::Response)> { + // Structured objects can be constructed with invalid names. Serialize + // before building the request so validation errors remain local errors. + let body = serde_json::to_value(request).map_err(|error| crate::Error::InvalidInput { + message: error.to_string(), + })?; + let request = self + .client + .post(&format!("/admin/authz/{route}")) + .json(&body); + // Every subject selector can contain a secret. Never log these bodies, + // and never replay an ACL mutation after an ambiguous response. + let (id, response) = self.client.send_suppressing_body(request).await?; + let response = self.client.check_response(&id, response).await?; + Ok((id, response)) + } + + async fn read( + &self, + route: &str, + request: &T, + ) -> Result { + let (id, response) = self.post(route, request).await?; + response.json().await.err_to_http(id) + } +} + +#[async_trait::async_trait] +impl Authorization for RemoteAuthorization { + async fn list_acls( + &self, + request: ListAccessControlEntriesRequest, + ) -> Result { + self.read("acl/list", &request).await + } + async fn add_acl(&self, request: AccessControlEntryRequest) -> Result<()> { + self.post("acl/add", &request).await?; + Ok(()) + } + async fn delete_acl(&self, request: AccessControlEntryRequest) -> Result<()> { + self.post("acl/delete", &request).await?; + Ok(()) + } + async fn list_role_bindings( + &self, + request: ListRoleBindingsRequest, + ) -> Result { + request.validate()?; + self.read("role_binding/list", &request).await + } + async fn add_role_binding(&self, request: RoleBindingRequest) -> Result<()> { + request.validate()?; + self.post("role_binding/add", &request).await?; + Ok(()) + } + async fn delete_role_binding(&self, request: RoleBindingRequest) -> Result<()> { + request.validate()?; + self.post("role_binding/delete", &request).await?; + Ok(()) + } + async fn get_principal(&self, request: GetPrincipalRequest) -> Result { + request.validate()?; + self.read("principal", &request).await + } + async fn get_group(&self, request: GetGroupRequest) -> Result { + request.validate()?; + self.read("group", &request).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::remote::client::test_utils::client_with_handler; + use serde_json::json; + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + #[tokio::test] + async fn list_preserves_wire_identity_and_continuation() { + let authz = RemoteAuthorization { + client: client_with_handler(|request| { + assert_eq!(request.url().path(), "/admin/authz/acl/list"); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + body, + json!({"object": "table:tenant/db:a$b$events", "subject": "G:Engineering", "max_results": 100, "continuation_token": "opaque/token"}) + ); + http::Response::builder().status(200).body(json!({ + "entries": [{"object": "table:tenant/db:a$b$events", "subject_id": "g:team-id", "privilege": "PRIVILEGE20"}], + "continuation_token": "next/token" + }).to_string()).unwrap() + }), + }; + let result = authz + .list_acls(ListAccessControlEntriesRequest { + object: Some("table:tenant/db:a$b$events".parse().unwrap()), + subject: Some(Subject::group_name("Engineering").unwrap()), + max_results: Some(100), + continuation_token: Some("opaque/token".into()), + }) + .await + .unwrap(); + assert_eq!(result.continuation_token.as_deref(), Some("next/token")); + assert_eq!( + result.entries[0].subject_id, + Subject::group_id("team-id").unwrap() + ); + assert_eq!(result.entries[0].subject_name, None); + assert_eq!( + result.entries[0].privilege, + "PRIVILEGE20".parse::().unwrap() + ); + } + + #[tokio::test] + async fn acl_failure_is_not_retried_or_reclassified() { + for status in [404, 409, 503] { + let calls = Arc::new(AtomicUsize::new(0)); + let count = calls.clone(); + let authz = RemoteAuthorization { + client: client_with_handler(move |request| { + count.fetch_add(1, Ordering::SeqCst); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["privilege"], "USAGE"); + http::Response::builder() + .status(status) + .body("failure") + .unwrap() + }), + }; + let result = authz + .add_acl(AccessControlEntryRequest { + object: Object::Database(DatabaseObject::new("db".into())), + subject: Subject::principal_api_key("secret").unwrap(), + privilege: Privilege::Usage, + }) + .await; + let Err(crate::Error::Http { + status_code, + request_id, + .. + }) = result + else { + panic!("expected Http error") + }; + assert_eq!(status_code.unwrap().as_u16(), status); + assert!(!request_id.is_empty()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + } + + #[tokio::test] + async fn rejects_invalid_objects_before_sending() { + let authz = RemoteAuthorization { + client: client_with_handler(|_| -> http::Response { + panic!("invalid objects must not send a request") + }), + }; + let request = AccessControlEntryRequest::new( + Object::Database(DatabaseObject::new("bad:database".into())), + Subject::principal_api_key("secret").unwrap(), + Privilege::Usage, + ); + for result in [ + authz.add_acl(request.clone()).await, + authz.delete_acl(request.clone()).await, + authz + .list_acls(ListAccessControlEntriesRequest { + object: Some(request.object), + ..Default::default() + }) + .await + .map(|_| ()), + ] { + assert!(matches!(result, Err(crate::Error::InvalidInput { .. }))); + assert!(!result.unwrap_err().to_string().contains("secret")); + } + } + + #[tokio::test] + async fn rejects_incompatible_subjects_before_sending() { + let authz = RemoteAuthorization { + client: client_with_handler(|_| -> http::Response { + panic!("invalid selectors must not send a request") + }), + }; + let request = RoleBindingRequest { + subject: Subject::role("reader").unwrap(), + role: "writer".into(), + }; + assert!(matches!( + authz.add_role_binding(request.clone()).await, + Err(crate::Error::InvalidInput { .. }) + )); + assert!(matches!( + authz.delete_role_binding(request).await, + Err(crate::Error::InvalidInput { .. }) + )); + let request = ListRoleBindingsRequest { + subject: Some(Subject::role("reader").unwrap()), + ..Default::default() + }; + assert!(matches!( + authz.list_role_bindings(request).await, + Err(crate::Error::InvalidInput { .. }) + )); + assert!(matches!( + authz + .get_principal(GetPrincipalRequest { + principal: Subject::group_id("group").unwrap() + }) + .await, + Err(crate::Error::InvalidInput { .. }) + )); + assert!(matches!( + authz + .get_group(GetGroupRequest { + group: Subject::principal_api_key("secret").unwrap() + }) + .await, + Err(crate::Error::InvalidInput { .. }) + )); + } +} diff --git a/rust/lancedb/src/remote/catalog.rs b/rust/lancedb/src/remote/catalog.rs index 72b77a87e..df3782368 100644 --- a/rust/lancedb/src/remote/catalog.rs +++ b/rust/lancedb/src/remote/catalog.rs @@ -19,6 +19,7 @@ use crate::catalog::{ ListDatabasesResponse, }; use crate::database::Database; +use crate::utils::validate_database_name; use crate::{Error, Result}; /// Authentication and client settings shared by a catalog and its databases. @@ -154,29 +155,6 @@ impl RemoteCatalog { }) } - fn validate_name(&self, name: &str) -> Result<()> { - let delimiter = self - .options - .client_config - .id_delimiter - .as_deref() - .unwrap_or("$"); - if name.is_empty() - || name.trim() != name - || !name.is_ascii() - || name.chars().any(char::is_control) - || name.contains(delimiter) - || matches!(name, "." | "..") - { - return Err(Error::InvalidInput { - message: format!( - "Invalid database name '{name}': expected a nonempty ASCII name without surrounding whitespace, control characters, or namespace delimiter '{delimiter}'" - ), - }); - } - Ok(()) - } - fn database(&self, name: &str) -> Result> { Ok(Arc::new(RemoteDatabase::for_catalog( &self.endpoint, @@ -198,12 +176,16 @@ impl RemoteCatalog { #[async_trait] impl Catalog for RemoteCatalog { + fn authz(&self) -> Result> { + Ok(self.root.authz()) + } + fn uri(&self) -> &str { &self.endpoint } async fn create_database(&self, request: CreateDatabaseRequest) -> Result> { - self.validate_name(&request.name)?; + validate_database_name(&request.name)?; self.root .create_namespace(CreateNamespaceRequest { id: Some(vec![request.name.clone()]), @@ -231,7 +213,7 @@ impl Catalog for RemoteCatalog { } async fn drop_database(&self, request: DropDatabaseRequest) -> Result<()> { - self.validate_name(&request.name)?; + validate_database_name(&request.name)?; let result = self .root .drop_namespace(DropNamespaceRequest { @@ -286,7 +268,7 @@ impl Catalog for RemoteCatalog { } async fn connect_database(&self, name: &str) -> Result> { - self.validate_name(name)?; + validate_database_name(name)?; self.root .describe_namespace(DescribeNamespaceRequest { id: Some(vec![name.into()]), @@ -676,7 +658,24 @@ mod tests { } let catalog = RemoteCatalog::try_new("http://127.0.0.1:1", RemoteCatalogOptions::default()).unwrap(); - for name in ["", "a$b", "\r\ninjected", "..", "café", " padded "] { + for name in [ + "", + "/", + "//", + "a$b", + "\r\ninjected", + "..", + "café", + " padded ", + "a b", + "a?b", + "/db", + "db/", + "a//b", + "a/../b", + "a/./b", + &"a".repeat(65), + ] { assert!(matches!( catalog.create_database(name.into()).await, Err(Error::InvalidInput { .. }) diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 705506705..783d2ef84 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -8,12 +8,15 @@ use reqwest::{ header::{HeaderMap, HeaderValue}, }; use std::{collections::HashMap, future::Future, str::FromStr, sync::Arc, time::Duration}; +use tokio::sync::Semaphore; use crate::error::{Error, Result}; use crate::remote::db::RemoteOptions; use crate::remote::retry::{ResolvedRetryConfig, RetryCounter}; +use crate::utils::validate_database_name; const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id"); +const DEFAULT_BLOB_REQUEST_CONCURRENCY: usize = 8; pub fn redact_sensitive_headers(headers: &mut HeaderMap) { const SENSITIVE_HEADERS: [&str; 5] = [ @@ -85,6 +88,11 @@ const DEFAULT_MAX_REQUEST_DURATION_DIVISOR: u32 = 2; pub struct ClientConfig { pub timeout_config: TimeoutConfig, pub retry_config: RetryConfig, + /// Maximum concurrent blob HTTP requests for this connection, shared by all + /// tables and blob handles. Defaults to 8. The environment variable + /// `LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY` can also set this limit. + /// The value must be greater than zero. + pub blob_request_concurrency: Option, /// User agent to use for requests. The default provides the library /// name and version. pub user_agent: String, @@ -139,6 +147,7 @@ impl std::fmt::Debug for ClientConfig { f.debug_struct("ClientConfig") .field("timeout_config", &self.timeout_config) .field("retry_config", &self.retry_config) + .field("blob_request_concurrency", &self.blob_request_concurrency) .field("user_agent", &self.user_agent) .field("extra_headers", &self.extra_headers) .field("id_delimiter", &self.id_delimiter) @@ -159,6 +168,7 @@ impl Default for ClientConfig { Self { timeout_config: TimeoutConfig::default(), retry_config: RetryConfig::default(), + blob_request_concurrency: None, user_agent: concat!("LanceDB-Rust-Client/", env!("CARGO_PKG_VERSION")).into(), extra_headers: HashMap::new(), id_delimiter: None, @@ -307,6 +317,7 @@ pub struct RestfulLanceDbClient { client: reqwest::Client, host: String, pub(crate) retry_config: ResolvedRetryConfig, + pub(crate) blob_request_semaphore: Arc, pub(crate) sender: S, pub(crate) header_provider: Option>, /// Connection-level read consistency interval. Drives the @@ -383,6 +394,7 @@ pub fn parse_db_url(db_url: &str) -> Result { message: format!("Invalid encoded database name: {err}"), })? .into_owned(); + validate_database_name(&db_name)?; let db_prefix = { let prefix = parsed_url.path().trim_start_matches('/'); if prefix.is_empty() { @@ -467,6 +479,11 @@ impl ClientConfig { /// mistake is reported where it was made rather than as a confusing /// response later. Public so a caller can ask without connecting. pub fn validate(&self) -> Result<()> { + if self.blob_request_concurrency == Some(0) { + return Err(Error::InvalidInput { + message: "blob_request_concurrency must be greater than zero".into(), + }); + } if let Some(delimiter) = &self.id_delimiter { validate_id_delimiter(delimiter)?; } @@ -502,6 +519,8 @@ impl RestfulLanceDbClient { // Before anything is built from it, so the error names the caller's // configuration rather than a request. client_config.validate()?; + let blob_request_concurrency = + Self::resolve_blob_request_concurrency(client_config.blob_request_concurrency)?; // Get the timeouts let timeout = @@ -601,6 +620,7 @@ impl RestfulLanceDbClient { client, host, retry_config, + blob_request_semaphore: Arc::new(Semaphore::new(blob_request_concurrency)), sender: Sender, header_provider: client_config.header_provider, read_consistency_interval, @@ -609,6 +629,26 @@ impl RestfulLanceDbClient { }) } + fn resolve_blob_request_concurrency(passed: Option) -> Result { + let value = if let Some(value) = passed { + value + } else if let Ok(env) = std::env::var("LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY") { + env.parse::().map_err(|_| Error::InvalidInput { + message: format!( + "LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY must be a positive integer, got '{env}'" + ), + })? + } else { + DEFAULT_BLOB_REQUEST_CONCURRENCY + }; + if value == 0 { + return Err(Error::InvalidInput { + message: "blob request concurrency must be greater than zero".into(), + }); + } + Ok(value) + } + /// Resolve the max bytes per insert request from config, environment, or the /// default. A value of `0` (from either source) disables request splitting. fn resolve_max_bytes_per_request(passed: Option) -> Result> { @@ -1104,6 +1144,7 @@ pub mod test_utils { client: reqwest::Client::new(), host: "http://localhost".to_string(), retry_config: RetryConfig::default().try_into().unwrap(), + blob_request_semaphore: Arc::new(Semaphore::new(DEFAULT_BLOB_REQUEST_CONCURRENCY)), sender: MockSender { f: Arc::new(wrapper), }, @@ -1130,6 +1171,11 @@ pub mod test_utils { client: reqwest::Client::new(), host: "http://localhost".to_string(), retry_config: config.retry_config.try_into().unwrap(), + blob_request_semaphore: Arc::new(Semaphore::new( + config + .blob_request_concurrency + .unwrap_or(DEFAULT_BLOB_REQUEST_CONCURRENCY), + )), sender: MockSender { f: Arc::new(wrapper), }, @@ -1181,6 +1227,21 @@ mod tests { .expect_err("a configured delimiter other than the supported one must be refused"); } + #[test] + fn test_blob_request_concurrency_must_be_positive() { + let error = super::ClientConfig { + blob_request_concurrency: Some(0), + ..Default::default() + } + .validate() + .unwrap_err(); + assert!(error.to_string().contains("blob_request_concurrency")); + assert_eq!( + RestfulLanceDbClient::::resolve_blob_request_concurrency(Some(3)).unwrap(), + 3 + ); + } + use super::*; use serial_test::serial; use std::time::Duration; @@ -1204,6 +1265,29 @@ mod tests { assert_eq!(parsed.db_prefix.as_deref(), Some("prefix")); } + #[test] + fn test_parse_database_uri_validates_decoded_name() { + for name in [ + "a b", + "a$b", + "a?b", + "café", + "/db", + "db/", + "a//b", + "a/../b", + "a/./b", + &"a".repeat(65), + ] { + let uri = format!("db://{}", urlencoding::encode(name)); + assert!( + matches!(parse_db_url(&uri), Err(Error::InvalidInput { .. })), + "accepted {name:?}" + ); + } + assert!(parse_db_url(&format!("db://{}", "a".repeat(64))).is_ok()); + } + #[test] fn test_timeout_config_default() { let config = TimeoutConfig::default(); @@ -1490,6 +1574,7 @@ mod tests { client: reqwest::Client::new(), host: "https://example.com".to_string(), retry_config: RetryConfig::default().try_into().unwrap(), + blob_request_semaphore: Arc::new(Semaphore::new(DEFAULT_BLOB_REQUEST_CONCURRENCY)), sender: Sender, header_provider: Some(Arc::new(provider) as Arc), read_consistency_interval: None, @@ -1528,6 +1613,7 @@ mod tests { client: reqwest::Client::new(), host: "https://example.com".to_string(), retry_config: RetryConfig::default().try_into().unwrap(), + blob_request_semaphore: Arc::new(Semaphore::new(DEFAULT_BLOB_REQUEST_CONCURRENCY)), sender: Sender, header_provider: Some(Arc::new(provider) as Arc), read_consistency_interval: None, @@ -1592,6 +1678,7 @@ mod tests { client: reqwest::Client::new(), host: "https://example.com".to_string(), retry_config: RetryConfig::default().try_into().unwrap(), + blob_request_semaphore: Arc::new(Semaphore::new(DEFAULT_BLOB_REQUEST_CONCURRENCY)), sender: Sender, header_provider: Some(Arc::new(provider) as Arc), read_consistency_interval: None, diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index c6462a834..a0f4aa953 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -36,7 +36,10 @@ use crate::remote::util::stream_as_body; use crate::secrets::SecretBinding; use crate::secrets::SecretInfo; use crate::table::BaseTable; -use crate::utils::{reject_relative_segment, validate_table_name}; +use crate::utils::{ + validate_database_name, validate_function_name, validate_namespace, validate_secret_name, + validate_table_name, validate_view_name, +}; use crate::view::ViewDescription; use super::client::{ @@ -230,6 +233,7 @@ pub struct RemoteDatabase { /// TLS configuration for mTLS support tls_config: Option, sql_client: Option, + is_catalog_root: bool, } #[derive(Clone)] @@ -293,6 +297,12 @@ pub struct RemoteHostOverrides { } impl RemoteDatabase { + pub(crate) fn authz(&self) -> Arc { + Arc::new(super::authz::RemoteAuthorization { + client: self.client.clone(), + }) + } + pub(crate) fn try_new( uri: &str, api_key: &str, @@ -320,6 +330,9 @@ impl RemoteDatabase { name: Option<&str>, options: &super::catalog::RemoteCatalogOptions, ) -> Result { + if let Some(name) = name { + validate_database_name(name)?; + } let scope = super::catalog::ScopedHeaderProvider { provider: options.client_config.header_provider.clone(), database: name.map(str::to_string), @@ -348,6 +361,7 @@ impl RemoteDatabase { )?; if name.is_none() { db.sql_client = None; + db.is_catalog_root = true; } Ok(db) } @@ -422,11 +436,38 @@ impl RemoteDatabase { namespace_context_provider, tls_config: client_config.tls_config, sql_client: Some(sql_client), + is_catalog_root: false, }) } } impl RemoteDatabase { + // A catalog's root namespaces are database names, which may contain slashes. + fn namespace_identifier(&self, namespace: &[String]) -> Result { + match namespace.first() { + None => Ok(ID_DELIMITER.to_string()), + Some(first) => { + if self.is_catalog_root { + // When dealing with the catalog root, the first component is a + // database name. Those need to be validated differently because + // they can contain slashes. Also, we need to use percent encoding + // for the slashes. + validate_database_name(first)?; + validate_namespace(&namespace[1..])?; + Ok(namespace + .iter() + .map(|component| urlencoding::encode(component).into_owned()) + .collect::>() + .join(ID_DELIMITER)) + } else { + // If we're not dealing with the catalog root, build the namespace identifier in the + // usual way. + build_namespace_identifier(namespace) + } + } + } + } + /// Post a request whose body carries a credential. /// /// Shared by the create and alter verbs, which declare their own request @@ -537,6 +578,7 @@ mod test_utils { namespace_context_provider: None, tls_config: None, sql_client: None, + is_catalog_root: false, } } @@ -560,6 +602,7 @@ mod test_utils { namespace_context_provider, tls_config: config.tls_config.clone(), sql_client: None, + is_catalog_root: false, } } } @@ -581,93 +624,60 @@ impl From<&CreateTableMode> for &'static str { } } -/// The path segment addressing one object: its namespace path and its name. -/// -/// One builder for tables, Secrets, Functions and materialized views: the -/// identifier grammar belongs to the namespace spec, not to an object type. An -/// empty path addresses an object with no namespace. -/// -/// Components are checked for addressability, not a character set. The name's -/// grammar is the caller's, so a table reports [`Error::InvalidTableName`], a -/// Function admits names a table may not, and a catalog database carries the -/// `/` that [`RemoteCatalog`] allows. -/// -/// [`RemoteCatalog`]: super::catalog::RemoteCatalog -fn build_object_identifier(what: &str, name: &str, namespace: &[String]) -> Result { - for segment in namespace { - reject_unaddressable_component("namespace segment", segment)?; +fn build_compound_identifier(name: &str, namespace: &[String]) -> Result { + validate_namespace(namespace)?; + let mut identifier = namespace.join("$"); + if !namespace.is_empty() { + identifier.push('$'); } - reject_unaddressable_component(what, name)?; - Ok(join_identifier( - namespace.iter().map(String::as_str).chain([name]), - )) + identifier.push_str(name); + Ok(identifier) } -/// What a component may not be if the join is to survive being split back -/// apart: empty, a segment URL parsing resolves away, or the delimiter itself. -/// -/// Each erases a boundary no encoding of the joined form recovers. `["prod", -/// ""]` joins to `prod$`, which reads back as `["prod"]`, so a drop reaches the -/// parent of the namespace the caller named. -/// -/// Not a character set: per-component percent-encoding makes the wider set -/// safe, since a `/` in a name reaches the service as `%2F`, still one -/// segment. -fn reject_unaddressable_component(what: &str, value: &str) -> Result<()> { - if value.is_empty() { - return Err(Error::InvalidInput { - message: format!( - "{what} must not be empty: the identifier would carry two delimiters in a row, \ - and splitting it back apart would name a different object" - ), - }); - } - reject_relative_segment(what, value)?; - if value.contains(ID_DELIMITER) { - return Err(Error::InvalidInput { - message: format!( - "{what} '{value}' contains the identifier delimiter '{ID_DELIMITER}', so the \ - namespace path and the name it joins could not be told apart" - ), - }); - } - Ok(()) -} - -/// The path segment addressing one table. A wrapper for the error type: -/// callers match on [`Error::InvalidTableName`]. +/// Build a full path to a table. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. fn build_table_identifier(name: &str, namespace: &[String]) -> Result { validate_table_name(name)?; - build_object_identifier("table name", name, namespace) + build_compound_identifier(name, namespace) } -/// Join components into the `{id}` a route addresses: each percent-encoded, -/// then joined by the delimiter. -/// -/// Per component rather than over the joined string, so the delimiter stays a -/// delimiter and nothing inside a component can end the path segment. -/// -/// A second line, not the first: a component from the object charset is all -/// unreserved and encodes to itself, so the route reads as the caller wrote it. -/// It does not cover `.` and `..`, which are unreserved too and resolve away -/// after decoding -- [`build_object_identifier`] refuses those. -fn join_identifier<'a>(components: impl Iterator) -> String { - components - .map(|component| urlencoding::encode(component).into_owned()) - .collect::>() - .join(ID_DELIMITER) -} - -/// The path segment addressing one namespace. +/// Translate an array of namespace components into a namespace string. +/// Note that we do not need to percent-encode this when using it in a URL +/// because validate_namespace ensures that only URL-safe characters are +/// present. fn build_namespace_identifier(namespace: &[String]) -> Result { - for segment in namespace { - reject_unaddressable_component("namespace segment", segment)?; - } + validate_namespace(namespace)?; if namespace.is_empty() { // According to the namespace spec, use delimiter to represent root namespace - return Ok(ID_DELIMITER.to_string()); + Ok(ID_DELIMITER.to_string()) + } else { + Ok(namespace.join(ID_DELIMITER)) } - Ok(join_identifier(namespace.iter().map(String::as_str))) +} + +/// Build a full path to a function. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. +fn build_function_identifier(name: &str, namespace: &[String]) -> Result { + validate_function_name(name)?; + build_compound_identifier(name, namespace) +} + +/// Build a full path to a secret. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. +fn build_secret_identifier(name: &str, namespace: &[String]) -> Result { + validate_secret_name(name)?; + build_compound_identifier(name, namespace) +} + +/// Build a full path to a view. Note that we do not need to percent-encode +/// this when using it in a URL because the validation functions ensure that +/// only URL-safe characters are present. +fn build_view_identifier(name: &str, namespace: &[String]) -> Result { + validate_view_name(name)?; + build_compound_identifier(name, namespace) } /// Build a secure cache key using length prefixes. @@ -945,7 +955,7 @@ impl Database for RemoteDatabase { request: FunctionRegistrationRequest, namespace_path: &[String], ) -> Result> { - let function_id = build_object_identifier("Function name", &request.name, namespace_path)?; + let function_id = build_function_identifier(&request.name, namespace_path)?; let req = self .client .post(&format!("/v1/function/{function_id}/create")) @@ -976,7 +986,7 @@ impl Database for RemoteDatabase { version: &str, namespace_path: &[String], ) -> Result { - let function_id = build_object_identifier("Function name", name, namespace_path)?; + let function_id = build_function_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/function/{function_id}/describe")) @@ -1036,7 +1046,7 @@ impl Database for RemoteDatabase { version: &str, namespace_path: &[String], ) -> Result<(bool, Job)> { - let function_id = build_object_identifier("Function name", name, namespace_path)?; + let function_id = build_function_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/function/{function_id}/drop")) @@ -1081,7 +1091,7 @@ impl Database for RemoteDatabase { value: &str, namespace_path: &[String], ) -> Result<()> { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; self.post_secret_write( &format!("/v1/secret/{secret_id}/create"), &RemoteCreateSecretRequest { value }, @@ -1090,7 +1100,7 @@ impl Database for RemoteDatabase { } async fn alter_secret(&self, name: &str, value: &str, namespace_path: &[String]) -> Result<()> { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; self.post_secret_write( &format!("/v1/secret/{secret_id}/alter"), &RemoteAlterSecretRequest { value }, @@ -1127,7 +1137,7 @@ impl Database for RemoteDatabase { } async fn drop_secret(&self, name: &str, namespace_path: &[String]) -> Result<()> { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; let req = self.client.post(&format!("/v1/secret/{secret_id}/drop")); let (request_id, response) = self.client.send(req).await?; self.client.check_response(&request_id, response).await?; @@ -1135,7 +1145,7 @@ impl Database for RemoteDatabase { } async fn describe_secret(&self, name: &str, namespace_path: &[String]) -> Result { - let secret_id = build_object_identifier("Secret name", name, namespace_path)?; + let secret_id = build_secret_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/secret/{secret_id}/describe")); @@ -1150,7 +1160,7 @@ impl Database for RemoteDatabase { query: &str, namespace_path: &[String], ) -> Result { - let view_id = build_object_identifier("View name", name, namespace_path)?; + let view_id = build_view_identifier(name, namespace_path)?; let req = self .client .post(&format!("/v1/view/{view_id}/create")) @@ -1167,7 +1177,7 @@ impl Database for RemoteDatabase { name: &str, namespace_path: &[String], ) -> Result { - let view_id = build_object_identifier("View name", name, namespace_path)?; + let view_id = build_view_identifier(name, namespace_path)?; let req = self.client.post(&format!("/v1/view/{view_id}/describe")); let (request_id, response) = self.client.send(req).await?; let response = self.client.check_response(&request_id, response).await?; @@ -1184,7 +1194,7 @@ impl Database for RemoteDatabase { } async fn drop_view_async(&self, name: &str, namespace_path: &[String]) -> Result { - let view_id = build_object_identifier("View name", name, namespace_path)?; + let view_id = build_view_identifier(name, namespace_path)?; let req = self.client.post(&format!("/v1/view/{view_id}/drop")); let (request_id, response) = self.client.send(req).await?; let response = self.client.check_response(&request_id, response).await?; @@ -1407,7 +1417,7 @@ impl Database for RemoteDatabase { async fn list_tables(&self, request: ListTablesRequest) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .get(&format!("/v1/namespace/{}/table/list", namespace_id)); @@ -1674,7 +1684,7 @@ impl Database for RemoteDatabase { request: ListNamespacesRequest, ) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .get(&format!("/v1/namespace/{}/list", namespace_id)); @@ -1696,7 +1706,7 @@ impl Database for RemoteDatabase { request: CreateNamespaceRequest, ) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .post(&format!("/v1/namespace/{}/create", namespace_id)); @@ -1727,7 +1737,7 @@ impl Database for RemoteDatabase { async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let mut req = self .client .post(&format!("/v1/namespace/{}/drop", namespace_id)); @@ -1761,7 +1771,7 @@ impl Database for RemoteDatabase { request: DescribeNamespaceRequest, ) -> Result { let namespace_parts = request.id.as_deref().unwrap_or(&[]); - let namespace_id = build_namespace_identifier(namespace_parts)?; + let namespace_id = self.namespace_identifier(namespace_parts)?; let req = self .client .post(&format!("/v1/namespace/{}/describe", namespace_id)) @@ -1887,6 +1897,40 @@ mod tests { remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE}, }; + #[tokio::test] + async fn test_catalog_namespace_identifier_validates_all_namespace_components() { + let mut db = super::RemoteDatabase::new_mock(|_| -> http::Response { + panic!("identifier validation must not send a request") + }); + db.is_catalog_root = true; + + for (path, expected) in [ + (vec![], "$"), + (vec!["team/search"], "team%2Fsearch"), + (vec!["team/search", "ns"], "team%2Fsearch$ns"), + (vec!["team/search", "ns", "child"], "team%2Fsearch$ns$child"), + ] { + let path = path.into_iter().map(String::from).collect::>(); + assert_eq!(db.namespace_identifier(&path).unwrap(), expected); + } + + for path in [ + vec!["../search"], + vec!["team/search", "bad/name"], + vec!["team/search", "ns", "bad/name"], + vec!["team/search", "public", "child"], + ] { + let path = path.into_iter().map(String::from).collect::>(); + assert!( + matches!( + db.namespace_identifier(&path), + Err(Error::InvalidInput { .. }) + ), + "accepted {path:?}" + ); + } + } + #[test] fn test_cache_key_security() { // Test that cache keys are unique regardless of delimiter manipulation @@ -4045,7 +4089,7 @@ mod tests { .expect_err("an illegal component must be refused"); assert!(!*reached.lock().unwrap(), "{name:?} reached the transport"); assert!( - error.to_string().contains("Secret name"), + error.to_string().contains("secret name"), "{name:?}: {error}" ); } @@ -4076,7 +4120,7 @@ mod tests { .await .expect_err("a dot-only name must be refused"); assert!( - by_name.to_string().contains("relative path segments"), + by_name.to_string().contains("Illegal secret name"), "{by_name}" ); @@ -4086,7 +4130,9 @@ mod tests { .await .expect_err("a dot-only namespace segment must be refused"); assert!( - by_segment.to_string().contains("relative path segments"), + by_segment + .to_string() + .contains("Illegal namespace component"), "{by_segment}" ); assert!( @@ -4165,28 +4211,20 @@ mod tests { } } - /// A segment outside the table charset still addresses one segment: the - /// service decides whether it may exist, and percent-encoding is what keeps - /// the question reaching the right route. A catalog database is named this - /// way. + /// Slashes belong to catalog database names, not table namespace segments. #[tokio::test] - async fn test_a_namespace_segment_outside_the_charset_is_encoded_not_refused() { - use std::sync::{Arc, Mutex}; - let seen = Arc::new(Mutex::new(String::new())); - let path = seen.clone(); - let conn = Connection::new_with_handler(move |request| { - *path.lock().unwrap() = request.url().path().to_string(); - http::Response::builder().status(200).body("{}").unwrap() + async fn test_a_namespace_segment_outside_the_charset_is_refused() { + let conn = Connection::new_with_handler(|_| -> http::Response { + panic!("invalid namespace reached the transport") }); - conn.drop_table("t", &["team/search".to_string()]) - .await - .unwrap(); - assert_eq!(*seen.lock().unwrap(), "/v1/table/team%2Fsearch$t/drop/"); + assert!(matches!( + conn.drop_table("t", &["team/search".to_string()]).await, + Err(Error::InvalidInput { .. }) + )); } - /// A Function name is percent-encoded, which covers everything but the - /// relative segment: `..` is unreserved, so it survives encoding and is - /// then resolved away, posting a registration body to `/v1/create`. + /// Reject relative segments and their encoded forms before sending a request, + /// so a function name cannot change the route through URL normalization. #[tokio::test] async fn test_a_relative_segment_function_name_is_refused() { use std::sync::{Arc, Mutex}; @@ -4201,10 +4239,7 @@ mod tests { .drop_function(name, "fv_1", &[]) .await .expect_err("a dot-only Function name must be refused"); - assert!( - error.to_string().contains("relative path segments"), - "{error}" - ); + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); assert!(!*reached.lock().unwrap(), "{name:?} reached the transport"); } } @@ -4585,7 +4620,7 @@ mod tests { }; // A view is not a table, so the refusal says so. assert!( - error.contains("view name") || error.contains("view namespace path segment"), + error.contains("view name") || error.contains("namespace component"), "{error}" ); } diff --git a/rust/lancedb/src/remote/retry.rs b/rust/lancedb/src/remote/retry.rs index 1a354eebb..19aa301e4 100644 --- a/rust/lancedb/src/remote/retry.rs +++ b/rust/lancedb/src/remote/retry.rs @@ -198,6 +198,29 @@ mod tests { assert_eq!(counter.connect_failures, 0); } + #[test] + fn retry_limit_message_includes_last_status_error() { + let config = test_config(); + let mut counter = RetryCounter::new(&config, "test".to_string()); + for _ in 0..2 { + counter + .increment_request_failures(crate::Error::Http { + source: "502 Bad Gateway".into(), + request_id: "test".to_string(), + status_code: Some(reqwest::StatusCode::BAD_GATEWAY), + }) + .unwrap(); + } + let error = counter + .increment_request_failures(crate::Error::Http { + source: "last response: 502 Bad Gateway".into(), + request_id: "test".to_string(), + status_code: Some(reqwest::StatusCode::BAD_GATEWAY), + }) + .unwrap_err(); + assert!(error.to_string().contains("last response: 502 Bad Gateway")); + } + #[tokio::test] async fn test_increment_from_error_respects_global_limits() { // If request_failures is already at max, a connect error should still diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index a17aadc9f..7d7a6301a 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -47,6 +47,7 @@ use crate::utils::background_cache::BackgroundCache; use crate::utils::{ MaxBatchLengthStream, TimeoutStream, public_fts_field_path_by_id, resolve_arrow_field_path, resolve_arrow_fts_field_path, supported_btree_data_type, supported_vector_data_type, + validate_fts_field, }; use crate::{DistanceType, Error}; use crate::{ @@ -76,7 +77,6 @@ use lance::dataset::{ColumnAlteration, NewColumnTransform, Version}; use lance_datafusion::exec::{OneShotExec, execute_plan}; use reqwest::{RequestBuilder, Response}; use serde::{Deserialize, Serialize}; -use serde_json::Number; use std::collections::{HashMap, HashSet}; use std::io::Cursor; use std::pin::Pin; @@ -597,6 +597,7 @@ impl RemoteTable { Index::BloomFilter(p) => ("BLOOM_FILTER", Some(to_json(p)?)), Index::RTree(p) => ("RTREE", Some(to_json(p)?)), Index::FTS(p) => { + validate_fts_field(&field)?; let mut params = to_json(p)?; if p.get_document_granularity().is_list_element() { params["document_granularity"] = "list_element".into(); @@ -1056,6 +1057,8 @@ impl RemoteTable { ); } Select::Dynamic(pairs) => { + // Keep the caller's alias order in the JSON object sent to the server. + // serde_json's preserve_order feature is required here. let alias_map = serde_json::Map::from_iter(pairs.iter().map(|(name, expr)| { (name.clone(), serde_json::Value::String(expr.clone())) @@ -1136,6 +1139,7 @@ impl RemoteTable { mut body: serde_json::Value, query: &VectorQueryRequest, ) -> Result> { + query.validate_search_effort()?; self.apply_query_params(&mut body, &query.base)?; // Apply general parameters, before we dispatch based on number of query vectors. @@ -1145,17 +1149,17 @@ impl RemoteTable { if let Some(approx_mode) = query.approx_mode { body["approx_mode"] = serde_json::json!(approx_mode); } - // In 0.23.1 we migrated from `nprobes` to `minimum_nprobes` and `maximum_nprobes`. - // Old client / new server: since minimum_nprobes is missing, fallback to nprobes - // New client / old server: old server will only see nprobes, make sure to set both - // nprobes and minimum_nprobes - // New client / new server: since minimum_nprobes is present, server can ignore nprobes - body["nprobes"] = query.minimum_nprobes.into(); - body["minimum_nprobes"] = query.minimum_nprobes.into(); + if let Some(nprobes) = query.nprobes { + body["nprobes"] = nprobes.into(); + } + if let Some(minimum_nprobes) = query.minimum_nprobes { + body["minimum_nprobes"] = minimum_nprobes.into(); + } if let Some(maximum_nprobes) = query.maximum_nprobes { body["maximum_nprobes"] = maximum_nprobes.into(); - } else { - body["maximum_nprobes"] = serde_json::Value::Number(Number::from_u128(0).unwrap()) + } + if let Some(search_effort) = query.search_effort { + body["search_effort"] = search_effort.into(); } body["lower_bound"] = query.lower_bound.into(); body["upper_bound"] = query.upper_bound.into(); @@ -1400,24 +1404,29 @@ impl RemoteTable { }); let streams = futures::future::try_join_all(futures); - if let Some(timeout) = options.timeout { + let streams = if let Some(timeout) = options.timeout { let timeout_future = tokio::time::sleep(timeout); tokio::pin!(timeout_future); tokio::pin!(streams); tokio::select! { _ = &mut timeout_future => { - Err(Error::Other { + return Err(Error::Other { message: format!("Query timeout after {} ms", timeout.as_millis()), source: None, }) } result = &mut streams => { - Ok(result?) + result? } } } else { - Ok(streams.await?) - } + streams.await? + }; + + Ok(streams + .into_iter() + .map(|stream| restore_projection_order(stream, &query.base().select)) + .collect()) } fn prepare_query_bodies( @@ -1459,6 +1468,52 @@ impl RemoteTable { } } +fn restore_projection_order( + stream: Pin>, + select: &Select, +) -> Pin> { + let pairs = match select { + Select::Dynamic(pairs) => pairs.iter().map(|(name, _)| name).collect::>(), + Select::Expr(pairs) => pairs.iter().map(|(name, _)| name).collect::>(), + _ => return stream, + }; + + let schema = stream.schema(); + let mut remaining = (0..schema.fields().len()).collect::>(); + let mut indices = Vec::with_capacity(remaining.len()); + for name in pairs { + if let Some(position) = remaining + .iter() + .position(|&index| schema.field(index).name() == name) + { + indices.push(remaining.remove(position)); + } + } + indices.extend(remaining); + if indices + .iter() + .enumerate() + .all(|(index, &field)| index == field) + { + return stream; + } + + // Older servers may sort alias maps even when the request preserves their + // order. Reorder the schema and every batch while keeping server-added fields. + let fields = indices + .iter() + .map(|&index| schema.field(index).clone()) + .collect::>(); + let schema = Arc::new(arrow_schema::Schema::new_with_metadata( + fields, + schema.metadata().clone(), + )); + let batches = stream.map(move |batch| { + batch.and_then(|batch| batch.project(&indices).map_err(DataFusionError::from)) + }); + Box::pin(RecordBatchStreamAdapter::new(schema, batches)) +} + #[derive(Deserialize)] struct TableDescription { version: u64, @@ -5249,6 +5304,14 @@ mod tests { assert_eq!(data[0].as_ref().unwrap(), &expected_data); } + fn schema_describe_response(schema: &Schema) -> http::Response { + let json_schema = JsonSchema::try_from(schema).unwrap(); + http::Response::builder() + .status(200) + .body(serde_json::json!({ "version": 1, "schema": json_schema }).to_string()) + .unwrap() + } + fn blob_describe_response() -> http::Response { let schema = Schema::new(vec![ Field::new("id", DataType::Int64, false), @@ -5256,11 +5319,7 @@ mod tests { Field::new("caption", DataType::Utf8, true), crate::blob("thumbnail", true), ]); - let json_schema = JsonSchema::try_from(&schema).unwrap(); - http::Response::builder() - .status(200) - .body(serde_json::json!({ "version": 1, "schema": json_schema }).to_string()) - .unwrap() + schema_describe_response(&schema) } #[rstest] @@ -5279,6 +5338,95 @@ mod tests { assert_eq!(columns, vec!["image".to_string(), "thumbnail".to_string()]); } + #[tokio::test] + async fn test_list_blob_paths_are_not_advertised_or_fetchable() { + let schema = Schema::new(vec![ + crate::blob("thumbnail", true), + Field::new( + "imgs", + DataType::List(Arc::new(crate::blob("item", true))), + true, + ), + ]); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| match request.url().path() { + "/v1/table/my_table/describe/" => schema_describe_response(&schema), + "/v1/table/my_table/fetch_blobs/" | "/v1/table/my_table/query/" => { + http::Response::builder() + .status(400) + .body("Invalid input, column 'imgs' is not a blob column".to_string()) + .unwrap() + } + path if path.starts_with("/v1/table/my_table/blob/") => http::Response::builder() + .status(400) + .body("Invalid input, column 'imgs' is not a blob column".to_string()) + .unwrap(), + path => panic!("unexpected request: {path}"), + }, + ); + + assert_eq!(table.blob_columns().await.unwrap(), vec!["thumbnail"]); + for column in ["imgs", "imgs.item"] { + let err = table.fetch_blobs(column, &[10]).await.unwrap_err(); + assert!( + err.to_string() + .contains("blobs inside lists cannot be fetched") + ); + let err = table.fetch_blob_files(column, &[10]).await.unwrap_err(); + assert!( + err.to_string() + .contains("blobs inside lists cannot be fetched") + ); + } + } + + #[tokio::test] + async fn test_fetch_blob_files_explains_uncached_list_descriptors() { + let list_field = Field::new( + "imgs", + DataType::List(Arc::new(crate::blob("item", true))), + true, + ); + let schema = Schema::new(vec![list_field.clone()]); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + list_field.clone(), + Field::new(lance::dataset::ROW_ID, DataType::UInt64, false), + ])), + vec![ + arrow_array::new_null_array(list_field.data_type(), 1), + Arc::new(arrow_array::UInt64Array::from(vec![10])), + ], + ) + .unwrap(); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + schema_describe_response(&schema).map(|body| body.into_bytes()) + } + "/v1/table/my_table/query/" => http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .header(VERSION_HEADER, "1") + .body(write_ipc_stream_uncompressed(&batch)) + .unwrap(), + path => panic!("unexpected request: {path}"), + }, + ); + + let error = table.fetch_blob_files("imgs", &[10]).await.unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("blobs inside lists cannot be fetched") + ); + } + #[tokio::test] async fn test_fetch_blobs_decodes_null_aligned_bytes() { let mut builder = LargeBinaryBuilder::new(); @@ -5512,7 +5660,7 @@ mod tests { .header(VERSION_HEADER, "42") .body(b"0".to_vec()) .unwrap(), - "bytes=0-" => http::Response::builder() + "bytes=0-9" => http::Response::builder() .status(206) .header(reqwest::header::CONTENT_RANGE, "bytes 0-9/10") .body(b"0123456789".to_vec()) @@ -5534,21 +5682,35 @@ mod tests { assert_eq!(blobs.value(4), b"bbb"); let requests = requests.lock().unwrap(); assert!(requests.contains(&"GET bytes=0-0".to_string())); - assert!(requests.contains(&"GET bytes=0-".to_string())); + assert!(requests.contains(&"GET bytes=0-9".to_string())); } #[tokio::test] async fn test_fetch_blobs_does_not_split_unrelated_bad_requests() { let requests = Arc::new(AtomicUsize::new(0)); let seen = requests.clone(); - let table = - Table::new_with_handler_version("my_table", semver::Version::new(0, 5, 0), move |_| { - seen.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(400) - .body(b"unknown blob column".to_vec()) - .unwrap() - }); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 5, 0), + move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + schema_describe_response(&Schema::new(vec![Field::new( + "id", + DataType::Int64, + false, + )])) + .map(|body| body.into_bytes()) + } + "/v1/table/my_table/fetch_blobs/" => { + seen.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(400) + .body(b"unknown blob column".to_vec()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + }, + ); assert_fetch_blobs_http_error( table.fetch_blobs("missing", &[10, 20]).await.unwrap_err(), @@ -5711,10 +5873,18 @@ mod tests { #[tokio::test] async fn test_fetch_blobs_sends_a_nested_column_as_a_dotted_path() { let ipc = one_row_blob_ipc_stream("info.blob"); + let schema = Schema::new(vec![Field::new( + "info", + DataType::Struct(vec![crate::blob("blob", true)].into()), + true, + )]); let table = Table::new_with_handler_version( "my_table", semver::Version::new(0, 5, 0), move |request| { + if request.url().path() == "/v1/table/my_table/describe/" { + return schema_describe_response(&schema).map(|body| body.into_bytes()); + } let body = request_body_json(&request); assert_eq!(body["column"], "info.blob"); http::Response::builder() @@ -6088,9 +6258,6 @@ mod tests { let body: serde_json::Value = serde_json::from_slice(body).unwrap(); let mut expected_body = serde_json::json!({ "prefilter": true, - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "lower_bound": Option::::None, "upper_bound": Option::::None, "k": 10, @@ -6122,7 +6289,7 @@ mod tests { } #[tokio::test] - async fn test_query_vector_approx_mode_sent_when_set() { + async fn test_query_vector_minimum_only_and_approx_mode() { let expected_data = RecordBatch::try_new( Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], @@ -6142,9 +6309,7 @@ mod tests { let body: serde_json::Value = serde_json::from_slice(body).unwrap(); let mut expected_body = serde_json::json!({ "prefilter": true, - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, + "minimum_nprobes": 5, "approx_mode": "accurate", "lower_bound": Option::::None, "upper_bound": Option::::None, @@ -6168,6 +6333,8 @@ mod tests { .query() .nearest_to(vec![0.1, 0.2, 0.3]) .unwrap() + .minimum_nprobes(5) + .unwrap() .approx_mode(crate::ApproxMode::Accurate) .execute() .await; @@ -6176,6 +6343,39 @@ mod tests { assert_eq!(data[0].as_ref().unwrap(), &expected_data); } + #[tokio::test] + async fn test_query_vector_explicitly_unbounded_maximum() { + let table = Table::new_with_handler("my_table", |request| { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["nprobes"], 5); + assert_eq!(body["maximum_nprobes"], 0); + + let data = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(write_ipc_file(&data)) + .unwrap() + }); + + table + .query() + .nearest_to(vec![0.1, 0.2, 0.3]) + .unwrap() + .nprobes(5) + .unwrap() + .maximum_nprobes(None) + .unwrap() + .execute() + .await + .unwrap(); + } + #[tokio::test] async fn test_query_fts_default_values() { let expected_data = RecordBatch::try_new( @@ -6241,6 +6441,77 @@ mod tests { )); } + #[rstest] + #[case::unset(None)] + #[case::minimum(Some(0.0))] + #[case::low(Some(0.25))] + #[case::auto(Some(0.5))] + #[case::high(Some(0.75))] + #[case::maximum(Some(1.0))] + #[tokio::test] + async fn test_query_vector_search_effort_wire(#[case] effort: Option) { + let table = Table::new_with_handler("my_table", move |request| { + let body = request_body_json(&request); + assert_eq!(body["nprobes"], 7); + assert_eq!(body["minimum_nprobes"], 2); + assert_eq!(body["maximum_nprobes"], 8); + if let Some(effort) = effort { + assert_eq!(body["search_effort"], effort); + } else { + assert!(body.get("search_effort").is_none()); + } + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(write_ipc_file(&batch)) + .unwrap() + }); + let batches = table + .query() + .nearest_to(vec![0.1, 0.2]) + .unwrap() + .search_effort(effort) + .unwrap() + .nprobes(7) + .unwrap() + .minimum_nprobes(2) + .unwrap() + .maximum_nprobes(Some(8)) + .unwrap() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches[0].num_rows(), 1); + } + + #[rstest] + #[case::negative(-0.1)] + #[case::too_large(1.1)] + #[case::nan(f64::NAN)] + #[case::infinity(f64::INFINITY)] + #[case::negative_infinity(f64::NEG_INFINITY)] + #[tokio::test] + async fn test_query_vector_search_effort_invalid(#[case] effort: f64) { + let table = Table::new_with_handler::<&str>("my_table", |_| { + panic!("invalid search_effort must not send a request") + }); + let result = table + .query() + .nearest_to(vec![0.1, 0.2]) + .unwrap() + .search_effort(effort); + assert!(matches!(result, Err(Error::InvalidInput { message }) + if message.contains("search_effort must be finite and in [0, 1]"))); + } + #[tokio::test] async fn test_query_vector_all_params() { let table = Table::new_with_handler("my_table", |request| { @@ -6274,8 +6545,8 @@ mod tests { } ], "nprobes": 12, - "minimum_nprobes": 12, - "maximum_nprobes": 12, + "minimum_nprobes": 3, + "maximum_nprobes": 10, "lower_bound": Option::::None, "upper_bound": Option::::None, "ef": Option::::None, @@ -6315,6 +6586,10 @@ mod tests { .distance_type(crate::DistanceType::Cosine) .nprobes(12) .unwrap() + .minimum_nprobes(3) + .unwrap() + .maximum_nprobes(Some(10)) + .unwrap() .refine_factor(2) .bypass_vector_index() .execute() @@ -6322,6 +6597,72 @@ mod tests { .unwrap(); } + #[rstest] + #[case::sql(Select::dynamic(&[("p", "x"), ("id", "id")]))] + #[case::typed(Select::expr_projection(&[ + ("p", crate::expr::col("x")), + ("id", crate::expr::col("id")), + ]))] + #[tokio::test] + async fn test_query_dynamic_projection_order(#[case] select: Select) { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/query/"); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + let aliases = body["columns"] + .as_object() + .unwrap() + .keys() + .map(String::as_str) + .collect::>(); + assert_eq!(aliases, ["p", "id"]); + + // Simulate a server that returns aliases in alphabetical order. + let data = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("p", DataType::Int32, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![7])), + Arc::new(Int32Array::from(vec![42])), + ], + ) + .unwrap(); + http::Response::builder() + .status(200) + .header(CONTENT_TYPE, ARROW_FILE_CONTENT_TYPE) + .body(write_ipc_file(&data)) + .unwrap() + }); + + let batches = table + .query() + .select(select) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches.len(), 1); + let schema = batches[0].schema(); + assert_eq!( + schema + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + ["p", "id"] + ); + let projected = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(projected.value(0), 42); + } + #[tokio::test] async fn test_query_vector_nested_field_path() { let expected_data = RecordBatch::try_new( @@ -6345,9 +6686,6 @@ mod tests { "vector_column": "image.embedding", "prefilter": true, "k": 10, - "nprobes": 20, - "minimum_nprobes": 20, - "maximum_nprobes": 20, "lower_bound": Option::::None, "upper_bound": Option::::None, "ef": Option::::None, @@ -7083,19 +7421,24 @@ mod tests { ]; for (index_type, expected_body, index) in cases { + let data_type = match &index { + Index::FTS(params) if params.get_document_granularity().is_list_element() => { + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))) + } + Index::FTS(_) => DataType::Utf8, + _ => DataType::Int32, + }; + let schema = Schema::new(vec![Field::new("a", data_type, false)]); let table = Table::new_with_handler_version( "my_table", semver::Version::new(0, 6, 0), move |request| { assert_eq!(request.method(), "POST"); match request.url().path() { - "/v1/table/my_table/describe/" => { - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - http::Response::builder() - .status(200) - .body(describe_response(&schema)) - .unwrap() - } + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), "/v1/table/my_table/create_index/" => { assert_eq!( request.headers().get("Content-Type").unwrap(), diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index b8506adb3..901616fb2 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -11,13 +11,13 @@ use std::time::Duration; use arrow_array::{Array, LargeBinaryArray, StructArray, UInt8Array, UInt64Array}; use arrow_schema::DataType; -use bytes::{Bytes, BytesMut}; +use bytes::Bytes; use datafusion_expr::{col, lit}; use futures::{StreamExt, TryStreamExt}; use lance::dataset::ROW_ID; use lance_core::datatypes::BlobKind; use reqwest::{Response, StatusCode, header}; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, OwnedSemaphorePermit}; use crate::Error; use crate::blob::BlobFile; @@ -57,7 +57,26 @@ trait BlobRangeRequester: Send + Sync + std::fmt::Debug { range_header: &str, mode: RangeRequestMode, version: Option, - ) -> Result<(String, Response)>; + ) -> Result; +} + +/// Keep the connection slot until the response body has been consumed or dropped. +#[derive(Debug)] +struct BlobRangeResponse { + request_id: String, + response: Response, + _permit: Option, +} + +#[cfg(test)] +impl From<(String, Response)> for BlobRangeResponse { + fn from((request_id, response): (String, Response)) -> Self { + Self { + request_id, + response, + _permit: None, + } + } } #[derive(Debug)] @@ -79,7 +98,14 @@ impl BlobRangeRequester for TableBlobRangeRequester { range_header: &str, mode: RangeRequestMode, version: Option, - ) -> Result<(String, Response)> { + ) -> Result { + let permit = self + .client + .blob_request_semaphore + .clone() + .acquire_owned() + .await + .expect("blob request semaphore is never closed"); let freshness_request = freshness_headers_snapshot(&self.freshness, self.read_consistency_interval); let mut request = freshness_request @@ -96,32 +122,32 @@ impl BlobRangeRequester for TableBlobRangeRequester { if response.status() == StatusCode::RANGE_NOT_SATISFIABLE && matches!(mode, RangeRequestMode::SizeProbe) { - return Ok((request_id, response)); + return Ok(BlobRangeResponse { + request_id, + response, + _permit: Some(permit), + }); } let response = self.client.check_response(&request_id, response).await?; freshness_request.observe_headers(&self.freshness, response.headers()); self.parent_freshness_request .observe_headers(&self.parent_freshness, response.headers()); - Ok((request_id, response)) + Ok(BlobRangeResponse { + request_id, + response, + _permit: Some(permit), + }) } } -#[derive(Debug)] -struct SequentialResponse { - response: Response, - request_id: String, - buffered: Bytes, -} - #[derive(Debug, Default)] struct RemoteBlobState { cursor: u64, - sequential_response: Option, } /// Seekable Cloud blob handle over HTTP Range. /// -/// Nonempty handles read the table version returned by their size probe. +/// Handles read the table version used to fetch their descriptor or probe their size. #[derive(Debug)] pub struct RemoteBlobFile { requester: Arc, @@ -145,11 +171,6 @@ impl RemoteBlobFile { /// Close the handle without waiting for an in-flight read. pub(crate) async fn close(&self) -> lance_core::Result<()> { self.closed.store(true, Ordering::Release); - // Drop a retained response when the state lock is immediately available. - // A reader holding the lock drops it instead once it observes the flag. - if let Ok(mut state) = self.state.try_lock() { - state.sequential_response = None; - } Ok(()) } @@ -185,7 +206,11 @@ impl RemoteBlobFile { return Ok(Bytes::new()); } let range_header = format!("bytes={}-{}", range.start, range.end - 1); - let (request_id, response) = self + let BlobRangeResponse { + request_id, + response, + _permit, + } = self .requester .request_range(&range_header, RangeRequestMode::DataRead, self.version) .await @@ -252,62 +277,20 @@ impl RemoteBlobFile { return Ok(Bytes::new()); } - // Remove the retained response from shared state before awaiting. Failed or - // cancelled reads leave the committed cursor unchanged and force the next - // read to open a fresh response. - let mut sequential_response = state.sequential_response.take(); - let mut cursor = state.cursor; - let mut output = BytesMut::with_capacity(target_len); - while output.len() < target_len { - if sequential_response.is_none() { - let range_header = format!("bytes={cursor}-"); - let (request_id, response) = self - .requester - .request_range(&range_header, RangeRequestMode::DataRead, self.version) - .await - .map_err(remote_blob_error)?; - self.ensure_open()?; - validate_partial_response(&response, cursor..self.size, self.size)?; - sequential_response = Some(SequentialResponse { - response, - request_id, - buffered: Bytes::new(), - }); - } - - let needed = target_len - output.len(); - let active = sequential_response.as_mut().unwrap(); - if !active.buffered.is_empty() { - let take = needed.min(active.buffered.len()); - output.extend_from_slice(&active.buffered.split_to(take)); - cursor += take as u64; - continue; - } - let chunk = active - .response - .chunk() - .await - .err_to_http(active.request_id.clone()) - .map_err(remote_blob_error)?; - self.ensure_open()?; - let chunk = chunk.ok_or_else(|| { - remote_blob_error("response ended before the requested blob range") - })?; - active.buffered = chunk; - } + // A partial read must release its HTTP response before returning. Keeping + // it open across calls would hold a connection slot indefinitely and + // prevent other handles from making progress. + let end = state.cursor + target_len as u64; + let bytes = self.read_range(state.cursor..end).await?; self.ensure_open()?; - state.cursor = cursor; - if state.cursor < self.size { - state.sequential_response = sequential_response; - } - Ok(output.freeze()) + state.cursor = end; + Ok(bytes) } pub(crate) async fn seek(&self, new_cursor: u64) -> lance_core::Result<()> { self.ensure_open()?; let mut state = self.state.lock().await; self.ensure_open()?; - state.sequential_response = None; state.cursor = new_cursor; Ok(()) } @@ -375,6 +358,34 @@ fn validate_partial_response( } impl RemoteTable { + async fn reject_list_blob_path_if_known(&self, column: &str) -> Result<()> { + // A dotted path needs client-side validation because the server may + // report an internal type mismatch when the path crosses a list. Use + // an already cached schema for top-level list fields as well. + if column.contains('.') || self.schema_cache.try_get().is_some() { + let schema = self.schema().await?; + crate::blob::reject_list_blob_column(schema.as_ref(), column)?; + } + Ok(()) + } + + async fn explain_list_blob_bad_request(&self, column: &str, error: Error) -> Error { + let is_bad_request = match &error { + Error::Http { + status_code: Some(code), + .. + } => *code == StatusCode::BAD_REQUEST || *code == StatusCode::UNPROCESSABLE_ENTITY, + _ => false, + }; + if is_bad_request + && let Ok(schema) = self.schema().await + && let Err(list_error) = crate::blob::reject_list_blob_column(schema.as_ref(), column) + { + return list_error; + } + error + } + /// Blob v2 columns are marked in field metadata, which `describe` returns. Reading /// them from the cached schema needs no route of its own and no version gate. pub(super) async fn blob_columns_impl(&self) -> Result> { @@ -387,6 +398,7 @@ impl RemoteTable { column: &str, row_ids: &[u64], ) -> Result { + self.reject_list_blob_path_if_known(column).await?; // Empty requests do not require blob-route support. if row_ids.is_empty() { return Ok(LargeBinaryArray::from(Vec::>::new())); @@ -466,7 +478,10 @@ impl RemoteTable { let (request_id, response) = self .send_with_freshness(request, true, read_snapshot.freshness) .await?; - let mut stream = self.read_arrow_response(&request_id, response).await?; + let mut stream = match self.read_arrow_response(&request_id, response).await { + Ok(stream) => stream, + Err(error) => return Err(self.explain_list_blob_bad_request(column, error).await), + }; let mut blob_chunks: Vec> = Vec::new(); while let Some(batch) = stream.try_next().await? { @@ -535,6 +550,7 @@ impl RemoteTable { column: &str, row_ids: &[u64], ) -> Result>> { + self.reject_list_blob_path_if_known(column).await?; // Empty requests do not require blob-route support. if row_ids.is_empty() { return Ok(Vec::new()); @@ -542,9 +558,13 @@ impl RemoteTable { self.ensure_blob_files_supported()?; let mut read_snapshot = self.snapshot_read_state().await; - let sizes = self + let sizes = match self .take_blob_descriptor_sizes(column, row_ids, &mut read_snapshot) - .await?; + .await + { + Ok(sizes) => sizes, + Err(error) => return Err(self.explain_list_blob_bad_request(column, error).await), + }; let version = read_snapshot.version; let mut files = Vec::with_capacity(row_ids.len()); @@ -568,7 +588,10 @@ impl RemoteTable { } } } - let probed = probe_blob_files(probe_requesters).await?; + let probed = match probe_blob_files(probe_requesters).await { + Ok(files) => files, + Err(error) => return Err(self.explain_list_blob_bad_request(column, error).await), + }; for (index, file) in probe_indices.into_iter().zip(probed) { files[index] = file; } @@ -589,7 +612,10 @@ impl RemoteTable { .iter() .map(|row_id| self.blob_range_requester(column, *row_id, &read_snapshot)) .collect(); - probe_blob_files(requesters).await + match probe_blob_files(requesters).await { + Ok(files) => Ok(files), + Err(error) => Err(self.explain_list_blob_bad_request(column, error).await), + } } /// Read the blob descriptor of every unique row in `row_ids`. @@ -696,6 +722,15 @@ impl RemoteTable { let values = batch.column_by_name(column).ok_or_else(|| { invalid_response(format!("blob descriptor take is missing '{column}'")) })?; + // An uncached top-level list can reach the descriptor query. Explain + // that unsupported path before the generic descriptor type check. + if matches!( + values.data_type(), + DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + ) && let Ok(schema) = self.schema().await + { + crate::blob::reject_list_blob_column(schema.as_ref(), column)?; + } // Only blob v2 columns read back as descriptors. let descriptors = values .as_any() @@ -817,7 +852,11 @@ fn probe_blob_version(response: &Response, request_id: &str) -> Result { /// /// `204` represents null. `416` with `bytes */0` represents an empty blob. async fn probe_blob_file(requester: Arc) -> Result> { - let (request_id, response) = requester + let BlobRangeResponse { + request_id, + response, + _permit, + } = requester .request_range("bytes=0-0", RangeRequestMode::SizeProbe, None) .await?; match response.status() { @@ -1372,6 +1411,114 @@ mod tests { assert!(requests.lock().unwrap().contains(&"bytes=5-11".to_string())); } + #[tokio::test(flavor = "multi_thread", worker_threads = 8)] + async fn blob_requests_are_bounded_across_handles() { + let in_flight = Arc::new(AtomicUsize::new(0)); + let max_in_flight = Arc::new(AtomicUsize::new(0)); + let table = RemoteTable::new_mock_with_version_and_config( + "my_table".to_string(), + { + let in_flight = in_flight.clone(); + let max_in_flight = max_in_flight.clone(); + move |request| { + let range = request.headers().get(header::RANGE).unwrap(); + if range != "bytes=0-0" { + let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1; + max_in_flight.fetch_max(now, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(20)); + in_flight.fetch_sub(1, Ordering::SeqCst); + } + range_response(&request, PAYLOAD) + } + }, + Some(Version::new(0, 5, 0)), + crate::remote::ClientConfig { + blob_request_concurrency: Some(3), + ..Default::default() + }, + ); + let files = table + .probe_files("image", &(0..16).collect::>()) + .await + .unwrap(); + let start = Arc::new(tokio::sync::Barrier::new(files.len() + 1)); + let reads = files.into_iter().map(|file| { + let start = start.clone(); + tokio::spawn(async move { + start.wait().await; + file.unwrap().read_range(1..2).await + }) + }); + let reads = reads.collect::>(); + start.wait().await; + for result in futures::future::join_all(reads).await { + assert_eq!(result.unwrap().unwrap(), &PAYLOAD[1..2]); + } + let max = max_in_flight.load(Ordering::SeqCst); + assert!(max > 1, "reads never overlapped"); + assert!( + max <= 3, + "{max} blob requests exceeded the configured limit" + ); + } + + #[tokio::test] + async fn partial_reads_release_slots_for_other_handles() { + let table = RemoteTable::new_mock_with_version_and_config( + "my_table".to_string(), + |request| range_response(&request, PAYLOAD), + Some(Version::new(0, 5, 0)), + crate::remote::ClientConfig { + blob_request_concurrency: Some(1), + ..Default::default() + }, + ); + let files = table.probe_files("image", &[10, 20]).await.unwrap(); + let reads = files + .iter() + .map(|file| file.as_ref().unwrap().read_up_to(1)); + let results = + tokio::time::timeout(Duration::from_secs(1), futures::future::join_all(reads)) + .await + .expect("an idle partial response held the only connection slot"); + for result in results { + assert_eq!(result.unwrap(), &PAYLOAD[..1]); + } + } + + #[tokio::test] + async fn blob_retry_error_includes_last_status() { + let table = RemoteTable::new_mock_with_version_and_config( + "my_table".to_string(), + |request| { + if request.headers().get(header::RANGE).unwrap() == "bytes=0-0" { + range_response(&request, PAYLOAD) + } else { + http::Response::builder() + .status(StatusCode::SERVICE_UNAVAILABLE) + .body(Vec::new()) + .unwrap() + } + }, + Some(Version::new(0, 5, 0)), + crate::remote::ClientConfig { + retry_config: crate::remote::RetryConfig { + retries: Some(1), + ..Default::default() + }, + ..Default::default() + }, + ); + let file = table.probe_files("image", &[10]).await.unwrap(); + let error = file[0] + .as_ref() + .unwrap() + .read_range(1..2) + .await + .unwrap_err(); + assert!(error.to_string().contains("503"), "got: {error}"); + } + #[tokio::test] async fn remote_blob_file_reads_the_probed_version_after_table_changes() { let changed = Arc::new(AtomicBool::new(false)); @@ -1428,8 +1575,11 @@ mod tests { [ ("bytes=0-0".to_string(), None), ("bytes=5-8".to_string(), Some("5".to_string())), - ("bytes=0-".to_string(), Some("5".to_string())), - ("bytes=10-".to_string(), Some("5".to_string())), + ("bytes=0-3".to_string(), Some("5".to_string())), + ( + format!("bytes=10-{}", PAYLOAD.len() - 1), + Some("5".to_string()) + ), ] ); } @@ -1475,7 +1625,7 @@ mod tests { } #[tokio::test] - async fn remote_blob_file_reuses_sequential_response_until_seek() { + async fn remote_blob_file_releases_response_after_each_partial_read() { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); let file = table @@ -1493,7 +1643,7 @@ mod tests { assert_eq!(file.tell().await.unwrap(), 24); assert_eq!( requests.lock().unwrap().as_slice(), - ["bytes=0-0", "bytes=0-", "bytes=20-"] + ["bytes=0-0", "bytes=0-3", "bytes=4-6", "bytes=20-23"] ); } @@ -1689,11 +1839,11 @@ mod tests { #[tokio::test] async fn remote_blob_file_failed_read_preserves_cursor_and_retries_fresh() { - let sequential_requests = Arc::new(AtomicUsize::new(0)); + let read_requests = Arc::new(AtomicUsize::new(0)); let table = RemoteTable::new_mock( "my_table".to_string(), { - let sequential_requests = sequential_requests.clone(); + let read_requests = read_requests.clone(); move |request| { let range = request .headers() @@ -1705,15 +1855,14 @@ mod tests { if range == "bytes=0-0" { return range_response(&request, PAYLOAD); } - let attempt = sequential_requests.fetch_add(1, Ordering::SeqCst); + let attempt = read_requests.fetch_add(1, Ordering::SeqCst); if attempt == 0 { - // End the response five bytes early to simulate a truncated - // sequential read. + // Return too few bytes for the requested range. return http::Response::builder() .status(StatusCode::PARTIAL_CONTENT) .header( header::CONTENT_RANGE, - format!("bytes 0-{}/{}", PAYLOAD.len() - 1, PAYLOAD.len()), + format!("bytes 0-9/{}", PAYLOAD.len()), ) .body(PAYLOAD[..5].to_vec()) .unwrap(); @@ -1733,9 +1882,7 @@ mod tests { let error = file.read_up_to(10).await.unwrap_err(); assert!( - error - .to_string() - .contains("response ended before the requested blob range"), + error.to_string().contains("returned 5 bytes, expected 10"), "got: {error}" ); assert_eq!(file.tell().await.unwrap(), 0); @@ -1743,7 +1890,7 @@ mod tests { // Retry from the last committed cursor with a fresh request. let retried = file.read_up_to(4).await.unwrap(); assert_eq!(retried, b"0123".as_slice()); - assert_eq!(sequential_requests.load(Ordering::SeqCst), 2); + assert_eq!(read_requests.load(Ordering::SeqCst), 2); } #[tokio::test] @@ -1776,7 +1923,7 @@ mod tests { range_header: &str, _mode: RangeRequestMode, _version: Option, - ) -> Result<(String, Response)> { + ) -> Result { assert_eq!(range_header, "bytes=0-0"); let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; self.max_in_flight.fetch_max(now, Ordering::SeqCst); @@ -1793,7 +1940,7 @@ mod tests { .header(VERSION_HEADER, "5") .body(vec![0u8]) .unwrap(); - Ok((format!("probe-{}", self.index), Response::from(response))) + Ok((format!("probe-{}", self.index), Response::from(response)).into()) } } @@ -1957,7 +2104,7 @@ mod tests { range_header: &str, _mode: RangeRequestMode, _version: Option, - ) -> Result<(String, Response)> { + ) -> Result { let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; self.max_in_flight.fetch_max(now, Ordering::SeqCst); tokio::time::sleep(Duration::from_millis(5)).await; @@ -1977,7 +2124,7 @@ mod tests { ) .body(PAYLOAD[start..=end].to_vec()) .unwrap(); - Ok(("range".to_string(), Response::from(response))) + Ok(("range".to_string(), Response::from(response)).into()) } } @@ -2078,8 +2225,8 @@ mod tests { range_header: &str, _mode: RangeRequestMode, _version: Option, - ) -> Result<(String, Response)> { - if range_header.ends_with('-') { + ) -> Result { + if range_header == "bytes=0-3" { self.started.notify_one(); self.release.wait().await; } @@ -2091,7 +2238,7 @@ mod tests { ) .body(PAYLOAD.to_vec()) .unwrap(); - Ok(("hung".to_string(), Response::from(response))) + Ok(("hung".to_string(), Response::from(response)).into()) } } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index b296d135d..ca58b54de 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -19,6 +19,9 @@ pub use lance::dataset::ReadParams; pub use lance::dataset::Version; use lance::dataset::WriteMode; use lance::dataset::builder::DatasetBuilder; +pub use lance::dataset::mem_wal::MemIndexRegistry; +/// The traits a memtable index plugin implements; see [`MemIndexRegistry`]. +pub use lance::dataset::mem_wal::index as mem_index; use lance::dataset::{InsertBuilder, WriteParams}; use lance::index::DatasetIndexExt; use lance::index::scalar::load_segment_params; @@ -424,7 +427,8 @@ impl LsmWriteSpec { /// changes, so an index created later is maintained too; one of a kind the /// MemWAL cannot mirror is skipped rather than failing the table. A list is /// verbatim: each name must already exist and be maintainable, and an empty - /// list maintains nothing. + /// list maintains nothing. The maintainable kinds are Lance's built-ins plus + /// any added with [`NativeTable::set_mem_index_registry`]. /// /// ``` /// # use lancedb::table::LsmWriteSpec; @@ -1237,10 +1241,11 @@ impl Table { self.inner.count_rows(filter).await } - /// Names of the blob v2 columns in this table, in declaration order. + /// Names of the row-addressable blob v2 columns, in declaration order. /// - /// Nested blobs use dotted paths (e.g. `info.blob`). Returns - /// [`Error::NotSupported`] on table types without blob support. + /// Blobs nested in structs use dotted paths (e.g. `info.blob`). Blobs + /// inside lists are omitted because row ids cannot identify their elements. + /// Returns [`Error::NotSupported`] on table types without blob support. pub async fn blob_columns(&self) -> Result> { self.inner.blob_columns().await } @@ -1255,7 +1260,8 @@ impl Table { /// [`Self::fetch_blob_files`] for large selections. /// /// `_rowid` values stay valid after compaction when the table has stable - /// row ids. + /// row ids. One missing or deleted row id rejects the whole request; no + /// partial results are returned. /// /// ``` /// use arrow_array::UInt64Array; @@ -1279,9 +1285,10 @@ impl Table { /// # } /// ``` /// - /// Returns [`Error::InvalidInput`] when the column does not exist or is - /// not a blob v2 column, and [`Error::NotSupported`] on table types - /// without blob support. + /// Blobs inside lists cannot be fetched by row id. Returns + /// [`Error::InvalidInput`] for those paths, when the column does not exist, + /// or when it is not a blob v2 column. Returns [`Error::NotSupported`] on + /// table types without blob support. pub async fn fetch_blobs( &self, column: impl AsRef, @@ -1299,7 +1306,8 @@ impl Table { /// non-null blobs produce empty byte strings. /// /// `_rowid` values stay valid after compaction when the table has stable - /// row ids. + /// row ids. One missing or deleted row id rejects the whole request; no + /// partial results are returned. /// /// ``` /// use lancedb::blob::BlobRangeRequest; @@ -1338,9 +1346,13 @@ impl Table { /// /// Same length and order as `row_ids`. Null rows are `None`. Bytes are not /// read from disk until a call to [`BlobFile::read`]. + /// Remote blob requests share a connection-wide concurrency limit. Set + /// `ClientConfig::blob_request_concurrency` or the + /// `LANCE_CLIENT_BLOB_REQUEST_CONCURRENCY` environment variable to tune it. /// /// `_rowid` values stay valid after compaction when the table has stable - /// row ids. + /// row ids. One missing or deleted row id rejects the whole request; no + /// partial handles are returned. /// /// ``` /// # use lancedb::Table; @@ -2590,6 +2602,9 @@ pub struct NativeTable { pub(crate) pushdown_operations: HashSet, // Read-freshness baseline; `Some` only for namespace-backed tables. freshness: Option, + // Memtable index kinds this handle's LSM writes maintain, shared with its + // clones and branch handles. + mem_index_registry: Arc>, } impl std::fmt::Debug for NativeTable { @@ -2756,9 +2771,51 @@ impl NativeTable { namespace_client, pushdown_operations, freshness: None, + mem_index_registry: Default::default(), }) } + /// Maintain the memtable index kinds in `registry` when this table writes + /// through its [`LsmWriteSpec`], and accept them when + /// [`LsmWriteSpec::with_maintained_indexes`] names one. + /// + /// The registry belongs to this handle, its clones and its branch handles, + /// and is not stored with the table: a handle opened separately starts from + /// Lance's built-ins. Set it before the handle's first LSM write; writers + /// already open keep the registry they were opened with. Local tables only, + /// reached through [`Table::as_native`]. + /// + /// ``` + /// # use std::sync::Arc; + /// # use lancedb::Table; + /// # use lancedb::table::{LsmWriteSpec, MemIndexRegistry, mem_index::MemIndexPlugin}; + /// # async fn example( + /// # table: &Table, + /// # plugin: Arc, + /// # ) -> Result<(), Box> { + /// let registry = MemIndexRegistry::default().with_plugin(plugin)?; + /// if let Some(native) = table.as_native() { + /// native.set_mem_index_registry(registry); + /// } + /// let spec = LsmWriteSpec::unsharded().with_maintained_indexes(vec!["region_bitmap".to_string()]); + /// table.set_lsm_write_spec(spec).await?; + /// # Ok(()) + /// # } + /// ``` + pub fn set_mem_index_registry(&self, registry: MemIndexRegistry) { + *self + .mem_index_registry + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = registry; + } + + pub(crate) fn mem_index_registry(&self) -> MemIndexRegistry { + self.mem_index_registry + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + /// Set the namespace client for server-side query execution. /// /// When set, queries will be executed on the namespace server instead of locally. @@ -2786,6 +2843,7 @@ impl NativeTable { namespace_client: self.namespace_client.clone(), pushdown_operations: self.pushdown_operations.clone(), freshness: self.freshness.clone(), + mem_index_registry: self.mem_index_registry.clone(), } } @@ -2893,6 +2951,7 @@ impl NativeTable { namespace_client: stored_namespace_client, pushdown_operations, freshness: None, + mem_index_registry: Default::default(), }) } @@ -2984,6 +3043,7 @@ impl NativeTable { namespace_client, pushdown_operations, freshness: None, + mem_index_registry: Default::default(), }) } @@ -3117,6 +3177,7 @@ impl NativeTable { namespace_client: stored_namespace_client, pushdown_operations, freshness: None, + mem_index_registry: Default::default(), }) } diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 916b28ed1..3ac759bdd 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -29,8 +29,8 @@ use crate::index::Index; use crate::index::vector::{VectorIndex, suggested_num_sub_vectors}; use crate::utils::{ resolve_lance_fts_field_path, supported_bitmap_data_type, supported_btree_data_type, - supported_fm_data_type, supported_fts_data_type, supported_label_list_data_type, - supported_vector_data_type, supported_zonemap_data_type, + supported_fm_data_type, supported_label_list_data_type, supported_vector_data_type, + supported_zonemap_data_type, validate_fts_field, }; use super::NativeTable; @@ -292,7 +292,7 @@ impl NativeTable { )) } Index::FTS(fts_opts) => { - Self::validate_index_type(field, "FTS", supported_fts_data_type)?; + validate_fts_field(field)?; Ok(Box::new(fts_opts)) } Index::IvfFlat(index) => { diff --git a/rust/lancedb/src/table/datafusion.rs b/rust/lancedb/src/table/datafusion.rs index ea5222b15..0b820494f 100644 --- a/rust/lancedb/src/table/datafusion.rs +++ b/rust/lancedb/src/table/datafusion.rs @@ -166,6 +166,7 @@ pub struct BaseTableAdapter { table: Arc, schema: Arc, fts_query: Option, + use_lsm: Option, } impl BaseTableAdapter { @@ -181,9 +182,16 @@ impl BaseTableAdapter { table, schema: Arc::new(schema), fts_query: None, + use_lsm: None, }) } + /// Set the MemWAL read routing of every scan; see [`QueryRequest::use_lsm`]. + pub fn with_use_lsm(mut self, use_lsm: Option) -> Self { + self.use_lsm = use_lsm; + self + } + /// Create a new adapter with an FTS query applied. pub fn with_fts_query(&self, fts_query: FullTextSearchQuery) -> Self { // Add _score column to the schema @@ -196,6 +204,7 @@ impl BaseTableAdapter { table: self.table.clone(), schema, fts_query: Some(fts_query), + use_lsm: self.use_lsm, } } } @@ -223,6 +232,7 @@ impl TableProvider for BaseTableAdapter { let mut query = QueryRequest { full_text_search: self.fts_query.clone(), disable_scoring_autoprojection: disable_scoring, + use_lsm: self.use_lsm, ..Default::default() }; @@ -699,6 +709,44 @@ pub mod tests { ); } + #[tokio::test] + async fn test_topk_late_materialization_over_table_adapter() { + use lance::io::exec::topk_late_materialization::TopKLateMaterialization; + + let fixture = TestFixture::new().await; + let run = |rule: Option| { + let adapter = fixture.adapter2.clone(); + async move { + let mut state = + datafusion::execution::SessionStateBuilder::new().with_default_features(); + if let Some(rule) = rule { + state = state.with_physical_optimizer_rule(Arc::new(rule)); + } + let ctx = SessionContext::new_with_state(state.build()); + ctx.register_table("tbl2", adapter).unwrap(); + let df = ctx + .sql("SELECT ints, strings, jsons FROM tbl2 WHERE floats > 10 ORDER BY ints DESC LIMIT 5") + .await + .unwrap(); + let plan = df.clone().create_physical_plan().await.unwrap(); + let plan = datafusion_physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + (plan, df.collect().await.unwrap()) + } + }; + + let (plan, expected) = run(None).await; + assert!(!plan.contains("Take"), "{plan}"); + + // The sort lands below `MetadataEraserExec`, so the rule needs no + // certification for it. + let (plan, batches) = run(Some(TopKLateMaterialization::new())).await; + assert!(plan.contains("Take"), "{plan}"); + assert_eq!(batches, expected); + assert!(batches[0].schema().metadata().is_empty()); + } + #[tokio::test] async fn test_metadata_eraser_with_new_children_takes_child_schema() { let fixture = TestFixture::new().await; diff --git a/rust/lancedb/src/table/datafusion/blob_coerce.rs b/rust/lancedb/src/table/datafusion/blob_coerce.rs index 3b79ec022..80b6cd6fc 100644 --- a/rust/lancedb/src/table/datafusion/blob_coerce.rs +++ b/rust/lancedb/src/table/datafusion/blob_coerce.rs @@ -9,16 +9,22 @@ use std::fmt; use std::hash::{Hash, Hasher}; use std::sync::Arc; -use arrow_array::{Array, BooleanArray, RecordBatch}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, FixedSizeListArray, LargeListArray, ListArray, RecordBatch, +}; +use arrow_buffer::OffsetBuffer; +use arrow_cast::cast; use arrow_schema::{DataType, Field, FieldRef, Fields, Schema}; use arrow_select::nullif::nullif; use datafusion::functions::core::{get_field, named_struct}; -use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; +use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_expr::ColumnarValue; use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::expressions::{CastExpr, Literal}; use datafusion_physical_plan::PhysicalExpr; +use datafusion_physical_plan::expressions::Column; +use lance_arrow::FieldExt; use crate::error::{Error, Result}; @@ -157,6 +163,198 @@ pub(super) fn coerce_blob_expr( Ok((expr, table_field.clone())) } +/// Apply the normal blob coercion to list values while retaining the list's +/// offsets and validity. A child batch lets the existing blob expression handle +/// binary, null, and descriptor inputs exactly as it does for top-level fields. +pub(super) fn coerce_blob_list_expr( + input_expr: Arc, + input_field: &Field, + table_field: &FieldRef, + config: &Arc, +) -> Result, FieldRef)>> { + let (DataType::List(input_item) + | DataType::LargeList(input_item) + | DataType::FixedSizeList(input_item, _)) = input_field.data_type() + else { + return Ok(None); + }; + let (DataType::List(table_item) + | DataType::LargeList(table_item) + | DataType::FixedSizeList(table_item, _)) = table_field.data_type() + else { + return Ok(None); + }; + if !table_item.is_blob_v2() || input_item == table_item { + return Ok(None); + } + + let (item_expr, _) = coerce_blob_expr( + Arc::new(Column::new(input_item.name(), 0)), + input_item, + table_item, + config, + )?; + let expr: Arc = Arc::new(CoerceBlobList { + source: input_expr, + input_item: input_item.clone(), + item_expr, + field: table_field.clone(), + }); + Ok(Some((expr, table_field.clone()))) +} + +#[derive(Debug, Clone)] +struct CoerceBlobList { + source: Arc, + input_item: FieldRef, + item_expr: Arc, + field: FieldRef, +} + +impl fmt::Display for CoerceBlobList { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "coerce_blob_list({})", self.source) + } +} + +impl PartialEq for CoerceBlobList { + fn eq(&self, other: &Self) -> bool { + self.source.eq(&other.source) + && self.input_item == other.input_item + && self.item_expr.eq(&other.item_expr) + && self.field == other.field + } +} + +impl Eq for CoerceBlobList {} + +impl Hash for CoerceBlobList { + fn hash(&self, state: &mut H) { + self.source.hash(state); + self.input_item.hash(state); + self.item_expr.hash(state); + self.field.hash(state); + } +} + +impl PhysicalExpr for CoerceBlobList { + fn return_field(&self, _input_schema: &Schema) -> datafusion_common::Result { + Ok(self.field.clone()) + } + + fn nullable(&self, _input_schema: &Schema) -> datafusion_common::Result { + Ok(self.field.is_nullable()) + } + + fn evaluate(&self, batch: &RecordBatch) -> datafusion_common::Result { + let source = self.source.evaluate(batch)?.into_array(batch.num_rows())?; + // First convert the outer list kind, if needed, while keeping its binary + // item type. PyArrow commonly infers List even for a LargeList column. + let intermediate_type = match self.field.data_type() { + DataType::List(_) => DataType::List(self.input_item.clone()), + DataType::LargeList(_) => DataType::LargeList(self.input_item.clone()), + DataType::FixedSizeList(_, size) => { + DataType::FixedSizeList(self.input_item.clone(), *size) + } + _ => unreachable!("validated list type when building the expression"), + }; + let source = if source.data_type() == &intermediate_type { + source + } else { + cast(source.as_ref(), &intermediate_type)? + }; + let values = match source.data_type() { + DataType::List(_) => { + let list = source.as_any().downcast_ref::().unwrap(); + let first = list.offsets()[0] as usize; + let last = *list.offsets().last().unwrap() as usize; + list.values().slice(first, last - first) + } + DataType::LargeList(_) => { + let list = source.as_any().downcast_ref::().unwrap(); + let first = list.offsets()[0] as usize; + let last = *list.offsets().last().unwrap() as usize; + list.values().slice(first, last - first) + } + DataType::FixedSizeList(_, _) => source + .as_any() + .downcast_ref::() + .unwrap() + .values() + .clone(), + other => { + return Err(DataFusionError::Internal(format!( + "coerce_blob_list expected a list, got {other}" + ))); + } + }; + let item_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![self.input_item.clone()])), + vec![values.clone()], + )?; + let coerced = self + .item_expr + .evaluate(&item_batch)? + .into_array(values.len())?; + + let rebuilt: ArrayRef = match (source.data_type(), self.field.data_type()) { + (DataType::List(_), DataType::List(item)) => { + let list = source.as_any().downcast_ref::().unwrap(); + Arc::new(ListArray::try_new( + item.clone(), + OffsetBuffer::from_lengths(list.offsets().lengths()), + coerced, + list.nulls().cloned(), + )?) + } + (DataType::LargeList(_), DataType::LargeList(item)) => { + let list = source.as_any().downcast_ref::().unwrap(); + Arc::new(LargeListArray::try_new( + item.clone(), + OffsetBuffer::from_lengths(list.offsets().lengths()), + coerced, + list.nulls().cloned(), + )?) + } + (DataType::FixedSizeList(_, _), DataType::FixedSizeList(item, size)) => { + let list = source + .as_any() + .downcast_ref::() + .unwrap(); + Arc::new(FixedSizeListArray::try_new_with_length( + item.clone(), + *size, + coerced, + list.nulls().cloned(), + list.len(), + )?) + } + _ => unreachable!("validated list types when building the expression"), + }; + Ok(ColumnarValue::Array(rebuilt)) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.source] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion_common::Result> { + Ok(Arc::new(Self { + source: children[0].clone(), + input_item: self.input_item.clone(), + item_expr: self.item_expr.clone(), + field: self.field.clone(), + })) + } + + fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +} + /// Carries the source column's nullity onto the struct built for it. /// /// This is its own expression rather than a `CASE` because the projection @@ -253,9 +451,11 @@ mod tests { use super::*; use crate::blob::blob; use arrow_array::{ - Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, Int64Array, LargeBinaryArray, - NullArray, RecordBatch, StringArray, StringViewArray, StructArray, UInt8Array, UInt64Array, + Array, ArrayRef, BinaryArray, BinaryViewArray, FixedSizeListArray, Int32Array, Int64Array, + LargeBinaryArray, LargeListArray, ListArray, NullArray, RecordBatch, StringArray, + StringViewArray, StructArray, UInt8Array, UInt64Array, }; + use arrow_buffer::{NullBuffer, OffsetBuffer}; use arrow_schema::Schema; use datafusion::prelude::SessionContext; use datafusion_catalog::MemTable; @@ -366,6 +566,207 @@ mod tests { ); } + #[tokio::test] + async fn binary_list_items_coerce_to_blobs_without_losing_list_shape() { + for kind in [ + "list", + "large_list", + "large_list_from_list", + "fixed_size_list", + ] { + let input_item = Arc::new(Field::new("item", DataType::Binary, true)); + let table_item = Arc::new(wide_blob_field("item")); + let values: ArrayRef = Arc::new(BinaryArray::from_iter(vec![ + Some(b"a".as_slice()), + None, + Some(b"".as_slice()), + Some(b"ignored".as_slice()), + None, + None, + ])); + let (input_type, table_type, list): (DataType, DataType, ArrayRef) = match kind { + "list" => ( + DataType::List(input_item.clone()), + DataType::List(table_item.clone()), + Arc::new(ListArray::new( + input_item, + OffsetBuffer::new(vec![0, 3, 3].into()), + values, + Some(NullBuffer::from(vec![true, false])), + )), + ), + "large_list" => ( + DataType::LargeList(input_item.clone()), + DataType::LargeList(table_item.clone()), + Arc::new(LargeListArray::new( + input_item, + OffsetBuffer::new(vec![0_i64, 3, 3].into()), + values, + Some(NullBuffer::from(vec![true, false])), + )), + ), + "large_list_from_list" => ( + DataType::List(input_item.clone()), + DataType::LargeList(table_item.clone()), + Arc::new(ListArray::new( + input_item, + OffsetBuffer::new(vec![0, 3, 3].into()), + values, + Some(NullBuffer::from(vec![true, false])), + )), + ), + "fixed_size_list" => ( + DataType::FixedSizeList(input_item.clone(), 3), + DataType::FixedSizeList(table_item.clone(), 3), + Arc::new(FixedSizeListArray::new( + input_item, + 3, + values, + Some(NullBuffer::from(vec![true, false])), + )), + ), + _ => unreachable!(), + }; + let batch = batch_with_image(Field::new("image", input_type, true), list); + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("image", table_type, true), + ]); + let coerced = coerce(batch, &schema).await; + let image = coerced.column_by_name("image").unwrap(); + let output_schema = coerced.schema(); + let (DataType::List(item_field) + | DataType::LargeList(item_field) + | DataType::FixedSizeList(item_field, _)) = + output_schema.field_with_name("image").unwrap().data_type() + else { + unreachable!() + }; + assert!(item_field.is_blob_v2(), "{kind} lost blob metadata"); + let (offsets, items): (Vec, &StructArray) = match kind { + "list" => { + let array = image.as_any().downcast_ref::().unwrap(); + assert!(array.is_null(1)); + ( + array.offsets().iter().map(|v| i64::from(*v)).collect(), + array.values().as_any().downcast_ref().unwrap(), + ) + } + "large_list" | "large_list_from_list" => { + let array = image.as_any().downcast_ref::().unwrap(); + assert!(array.is_null(1)); + ( + array.offsets().iter().copied().collect(), + array.values().as_any().downcast_ref().unwrap(), + ) + } + "fixed_size_list" => { + let array = image.as_any().downcast_ref::().unwrap(); + assert!(array.is_null(1)); + ( + vec![0, 3, 6], + array.values().as_any().downcast_ref().unwrap(), + ) + } + _ => unreachable!(), + }; + assert_eq!( + offsets, + if kind == "fixed_size_list" { + vec![0, 3, 6] + } else { + vec![0, 3, 3] + } + ); + assert_eq!(items.len(), if kind == "fixed_size_list" { 6 } else { 3 }); + let data: &LargeBinaryArray = items + .column_by_name("data") + .unwrap() + .as_any() + .downcast_ref() + .unwrap(); + assert_eq!(data.value(0), b"a"); + assert!(items.is_null(1)); + assert_eq!(data.value(2), b""); + assert!(!items.is_null(2)); + } + } + + #[tokio::test] + async fn sliced_binary_lists_only_coerce_visible_items() { + for large in [false, true] { + let input_item = Arc::new(Field::new("item", DataType::Binary, true)); + let table_item = Arc::new(wide_blob_field("item")); + let values: ArrayRef = Arc::new(BinaryArray::from_iter(vec![ + Some(b"before".as_slice()), + Some(b"a".as_slice()), + None, + Some(b"after".as_slice()), + Some(b"tail".as_slice()), + ])); + let (input_type, table_type, list): (DataType, DataType, ArrayRef) = if large { + ( + DataType::LargeList(input_item.clone()), + DataType::LargeList(table_item), + Arc::new( + LargeListArray::new( + input_item, + OffsetBuffer::new(vec![0_i64, 1, 3, 4, 5].into()), + values, + None, + ) + .slice(1, 2), + ), + ) + } else { + ( + DataType::List(input_item.clone()), + DataType::List(table_item), + Arc::new( + ListArray::new( + input_item, + OffsetBuffer::new(vec![0, 1, 3, 4, 5].into()), + values, + None, + ) + .slice(1, 2), + ), + ) + }; + let batch = batch_with_image(Field::new("image", input_type, true), list); + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("image", table_type, true), + ]); + let coerced = coerce(batch, &schema).await; + let image = coerced.column_by_name("image").unwrap(); + let (offsets, items): (Vec, &StructArray) = if large { + let list = image.as_any().downcast_ref::().unwrap(); + ( + list.offsets().iter().copied().collect(), + list.values().as_any().downcast_ref().unwrap(), + ) + } else { + let list = image.as_any().downcast_ref::().unwrap(); + ( + list.offsets().iter().map(|v| i64::from(*v)).collect(), + list.values().as_any().downcast_ref().unwrap(), + ) + }; + assert_eq!(offsets, [0, 2, 3]); + assert_eq!(items.len(), 3); + let data: &LargeBinaryArray = items + .column_by_name("data") + .unwrap() + .as_any() + .downcast_ref() + .unwrap(); + assert_eq!(data.value(0), b"a"); + assert!(items.is_null(1)); + assert_eq!(data.value(2), b"after"); + } + } + #[tokio::test] async fn binary_view_coerces_to_declared_blob_struct() { let batch = batch_with_image( diff --git a/rust/lancedb/src/table/datafusion/cast.rs b/rust/lancedb/src/table/datafusion/cast.rs index 2f3f8b868..063572cbe 100644 --- a/rust/lancedb/src/table/datafusion/cast.rs +++ b/rust/lancedb/src/table/datafusion/cast.rs @@ -20,7 +20,7 @@ use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use lance_arrow::json::{ARROW_JSON_EXT_NAME, has_json_fields, is_arrow_json_field, is_json_field}; use lance_arrow::{ARROW_EXT_NAME_KEY, FieldExt}; -use super::blob_coerce::coerce_blob_expr; +use super::blob_coerce::{coerce_blob_expr, coerce_blob_list_expr}; use crate::{Error, Result}; pub fn cast_to_table_schema( @@ -113,6 +113,14 @@ fn build_field_exprs( continue; } + // Blob values in lists need the same coercion as top-level blob columns. + if let Some(coerced) = + coerce_blob_list_expr(input_expr.clone(), input_field, table_field, &config)? + { + result.push(coerced); + continue; + } + // A column whose values are all null infers as `Null` (pyarrow does this for a list of // dicts), so there is no input type to cast from. Emit typed nulls carrying the table // field verbatim: a plain cast would drop the field metadata, and extension columns diff --git a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs index 000addbe1..543a59589 100644 --- a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs +++ b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs @@ -411,9 +411,7 @@ impl ExecutionPlan for VectorDedupExec { .map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?; let mut scanner = dataset.scan(); scanner - .with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter( - removed, - ))) + .with_row_id_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter(removed))) .blob_handling(BlobHandling::AllBinary) .batch_size(MATERIALIZE_BATCH_ROWS) .strict_batch_size(true); diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 354eef5a5..fbe82a198 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -1296,6 +1296,21 @@ mod lsm_tests { rows } + async fn collect_ids(stream: SendableRecordBatchStream) -> Vec { + let batches: Vec<_> = stream.try_collect().await.unwrap(); + batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect() + } + /// Upsert `ids` (value = 0..n) through the LSM `merge_insert` path. async fn lsm_upsert(table: &Table, ids: Vec) { let mut builder = table.merge_insert(&[]); @@ -1334,6 +1349,41 @@ mod lsm_tests { ); } + #[tokio::test] + async fn base_table_adapter_use_lsm_false_scans_the_base_only() { + use crate::table::datafusion::BaseTableAdapter; + use datafusion::prelude::SessionContext; + + let dir = tempdir().unwrap(); + let table = id_value_table(&dir).await; // base: ids 1,2,3 + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + // Ids 4,5 land in the active memtable, not the base table. + lsm_upsert(&table, vec![4, 5]).await; + + let adapter = BaseTableAdapter::try_new(table.base_table().clone()) + .await + .unwrap() + .with_use_lsm(Some(false)); + let ctx = SessionContext::new(); + ctx.register_table("base_only", Arc::new(adapter)).unwrap(); + let batches = ctx + .sql("SELECT id FROM base_only") + .await + .unwrap() + .collect() + .await + .unwrap(); + let mut ids: Vec = batches + .iter() + .flat_map(|b| b.column(0).as_primitive::().values().to_vec()) + .collect(); + ids.sort(); + assert_eq!(ids, vec![1, 2, 3]); + } + #[tokio::test] async fn query_snapshot_preserves_lsm_read_semantics() { let dir = tempdir().unwrap(); @@ -1767,12 +1817,16 @@ mod lsm_tests { assert_eq!(rows, 1, "LSM FTS must still see base rows"); } - #[tokio::test] - async fn lsm_read_vector_search() { + /// A table whose vector index the LSM maintains: base rows filled with + /// their id, and one memtable row filled with 1000, nearest to `[1000; 8]`. + /// `registry` replaces the table's memtable kinds. + async fn lsm_vector_table( + dir: &std::path::Path, + registry: Option, + ) -> Table { use crate::index::Index; use crate::index::vector::IvfPqIndexBuilder; use arrow::array::{FixedSizeListBuilder, Float32Builder}; - use arrow::datatypes::Int64Type; const DIM: i32 = 8; const N: i64 = 256; @@ -1801,12 +1855,7 @@ mod lsm_tests { .unwrap() }; - let dir = tempdir().unwrap(); - let conn = connect(dir.path().to_str().unwrap()) - .execute() - .await - .unwrap(); - // Base rows fill each vector with its own id (0..256); all far from 1000. + let conn = connect(dir.to_str().unwrap()).execute().await.unwrap(); let base = make_batch((0..N).map(|i| (i, i as f32)).collect()); let base_reader: Box = Box::new(RecordBatchIterator::new(vec![Ok(base)], schema.clone())); @@ -1817,20 +1866,22 @@ mod lsm_tests { &["vec"], Index::IvfPq( IvfPqIndexBuilder::default() - .num_partitions(1) + .num_partitions(8) .num_sub_vectors(2), ), ) .execute() .await .unwrap(); + if let Some(registry) = registry { + table.as_native().unwrap().set_mem_index_registry(registry); + } let vec_index = table.list_indices().await.unwrap()[0].name.clone(); table .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes(vec![vec_index])) .await .unwrap(); - // Insert a vector (filled with 1000) that is nearest to the query. let mut builder = table.merge_insert(&[]); builder .when_matched_update_all(None) @@ -1840,32 +1891,457 @@ mod lsm_tests { schema.clone(), )); builder.execute(insert_reader).await.unwrap(); + table + } - // KNN near [1000; DIM]: the default (auto-routed) read surfaces the - // memtable row. - let stream = table - .query() - .nearest_to(&[1000.0_f32; 8]) - .unwrap() - .limit(1) - .execute() - .await - .unwrap(); - let batches: Vec<_> = stream.try_collect().await.unwrap(); - let ids: Vec = batches - .iter() - .flat_map(|b| { - b.column_by_name("id") - .unwrap() - .as_primitive::() - .values() - .to_vec() - }) - .collect(); + /// The ids of the nearest row to `[1000; 8]` through the default read. + async fn nearest_ids(table: &Table) -> Vec { + let query = table.query().nearest_to(&[1000.0_f32; 8]).unwrap().limit(1); + collect_ids(query.execute().await.unwrap()).await + } + + #[tokio::test] + async fn lsm_read_vector_search() { + let dir = tempdir().unwrap(); + let table = lsm_vector_table(dir.path(), None).await; assert_eq!( - ids, + nearest_ids(&table).await, vec![9999], "LSM vector search must rank the memtable row first" ); + + let base_query = table + .query() + .nearest_to(&[0.0_f32; 8]) + .unwrap() + .only_if("id = 255") + .limit(1); + for query in [ + base_query.clone(), + base_query.clone().nprobes(8).unwrap(), + base_query.clone().minimum_nprobes(8).unwrap(), + base_query.clone().maximum_nprobes(Some(8)).unwrap(), + base_query + .nprobes(1) + .unwrap() + .maximum_nprobes(None) + .unwrap(), + ] { + let ids = collect_ids(query.execute().await.unwrap()).await; + assert_eq!(ids, vec![255]); + } + } + + /// The built-in vector plugin, except that its index answers no search. + #[derive(Debug)] + struct AnswersNothing; + + #[derive(Debug)] + struct AnswersNothingIndex(Arc); + + fn vector_plugin() -> Arc { + Arc::new(lance::dataset::mem_wal::index::HnswMemIndexPlugin) + } + + #[async_trait::async_trait] + impl lance::dataset::mem_wal::index::MemIndexPlugin for AnswersNothing { + fn name(&self) -> &str { + "AnswersNothing" + } + fn details_message(&self) -> &str { + "VectorIndexDetails" + } + fn flush_index_type(&self) -> lance_index::IndexType { + vector_plugin().flush_index_type() + } + fn training_criteria(&self) -> lance_index::scalar::registry::TrainingCriteria { + vector_plugin().training_criteria() + } + async fn resolve( + &self, + ctx: &lance::dataset::mem_wal::index::ResolveContext<'_>, + ) -> lance_core::Result { + vector_plugin().resolve(ctx).await + } + fn validate( + &self, + ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>, + ) -> lance_core::Result<()> { + vector_plugin().validate(ctx) + } + fn create( + &self, + ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>, + ) -> lance_core::Result> { + Ok(Arc::new(AnswersNothingIndex(vector_plugin().create(ctx)?))) + } + } + + #[async_trait::async_trait] + impl lance::dataset::mem_wal::index::MemIndex for AnswersNothingIndex { + fn columns(&self) -> &[String] { + self.0.columns() + } + fn can_answer(&self, _query: &dyn lance::dataset::mem_wal::index::MemQuery) -> bool { + false + } + fn insert(&self, batch: &RecordBatch, row_offset: u64) -> lance_core::Result<()> { + self.0.insert(batch, row_offset) + } + fn insert_batches( + &self, + batches: &[lance::dataset::mem_wal::write::StoredBatch], + ) -> lance_core::Result<()> { + self.0.insert_batches(batches) + } + fn resident_bytes(&self) -> usize { + self.0.resident_bytes() + } + fn search( + &self, + _query: &dyn lance::dataset::mem_wal::index::MemQuery, + _ctx: &lance::dataset::mem_wal::index::SearchContext, + ) -> lance_core::Result> { + Ok(None) + } + async fn flush( + &self, + ctx: &lance::dataset::mem_wal::index::FlushContext<'_>, + ) -> lance_core::Result { + self.0.flush(ctx).await + } + } + + /// A resident memtable whose vector index answers nothing still serves an + /// LSM vector search: Lance compares every row instead. + #[tokio::test] + async fn lsm_vector_search_compares_every_row_when_no_resident_index_answers() { + let dir = tempdir().unwrap(); + let mut registry = lance::dataset::mem_wal::MemIndexRegistry::default(); + registry.replace_plugin(Arc::new(AnswersNothing)).unwrap(); + let table = lsm_vector_table(dir.path(), Some(registry)).await; + assert_eq!(nearest_ids(&table).await, vec![9999]); + + // The second page starts after the memtable row. + let page = table + .query() + .nearest_to(&[1000.0_f32; 8]) + .unwrap() + .limit(2) + .offset(1); + assert_eq!( + collect_ids(page.execute().await.unwrap()).await, + vec![255, 254] + ); + } + + /// Claims the base table's bitmap index and keeps it as a B-tree, whose + /// flush rows are in the shape the bitmap trainer takes. + #[derive(Debug)] + struct BitmapAsBTree; + + fn btree_plugin() -> Arc { + lance::dataset::mem_wal::MemIndexRegistry::default() + .plugin_for_details_url("/lance.table.BTreeIndexDetails") + .unwrap() + .clone() + } + + #[async_trait::async_trait] + impl lance::dataset::mem_wal::index::MemIndexPlugin for BitmapAsBTree { + fn name(&self) -> &str { + "BitmapAsBTree" + } + fn details_message(&self) -> &str { + "BitmapIndexDetails" + } + fn flush_index_type(&self) -> lance_index::IndexType { + lance_index::IndexType::Bitmap + } + fn training_criteria(&self) -> lance_index::scalar::registry::TrainingCriteria { + btree_plugin().training_criteria() + } + fn validate( + &self, + ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>, + ) -> lance_core::Result<()> { + btree_plugin().validate(ctx) + } + fn create( + &self, + ctx: &lance::dataset::mem_wal::index::MemIndexBuildContext<'_>, + ) -> lance_core::Result> { + btree_plugin().create(ctx) + } + } + + /// A kind Lance has no memtable index for is refused when named until the + /// table's registry maintains it; then it is maintained, flushed with each + /// generation, and answers searches on the flushed generations. + #[tokio::test] + async fn lsm_maintains_an_index_kind_from_the_tables_registry() { + let dir = tempdir().unwrap(); + let conn = connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let table = conn + .create_table("t", id_region_reader(vec![(1, "us"), (2, "eu"), (3, "us")])) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .create_index(&["region"], crate::index::Index::Bitmap(Default::default())) + .name("region_bitmap".to_string()) + .execute() + .await + .unwrap(); + let spec = + || LsmWriteSpec::unsharded().with_maintained_indexes(vec!["region_bitmap".to_string()]); + + let refused = table.set_lsm_write_spec(spec()).await.unwrap_err(); + assert!( + refused.to_string().contains("region_bitmap") + && refused + .to_string() + .contains("no registered plugin maintains"), + "{refused:?}" + ); + + let registry = lance::dataset::mem_wal::MemIndexRegistry::default() + .with_plugin(Arc::new(BitmapAsBTree)) + .unwrap(); + table.as_native().unwrap().set_mem_index_registry(registry); + table.set_lsm_write_spec(spec()).await.unwrap(); + + // The registry is not stored with the table. + let other = conn.open_table("t").execute().await.unwrap(); + let mut builder = other.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let refused = builder + .execute(id_region_reader(vec![(20, "us")])) + .await + .unwrap_err(); + assert!( + refused + .to_string() + .contains("no registered plugin maintains"), + "{refused:?}" + ); + + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder + .execute(id_region_reader(vec![ + (10, "fresh"), + (11, "fresh"), + (12, "eu"), + ])) + .await + .unwrap(); + + let fresh = || async { + let query = table.query().only_if("region = 'fresh'"); + let mut ids = collect_ids(query.execute().await.unwrap()).await; + ids.sort_unstable(); + ids + }; + assert_eq!(fresh().await, vec![10, 11]); + let (_, _, memtables) = table + .as_native() + .unwrap() + .dataset + .shard_writer() + .read_snapshot() + .await + .unwrap() + .expect("the merge opened a writer"); + let active = memtables.expect("the writer keeps memtables").active; + assert!( + active + .index_store + .index_names() + .contains(&"region_bitmap".to_string()), + "the memtable maintains the plugin's index: {:?}", + active.index_store.index_names() + ); + + table.close_lsm_writers().await.unwrap(); + assert_eq!(fresh().await, vec![10, 11]); + assert_generations_answer_through_region_bitmap(&table, vec![10, 11]).await; + } + + /// The generations each shard's manifest publishes. + async fn published_generations(table: &Table) -> Vec { + use lance::dataset::mem_wal::ShardManifestStore; + + let uri = table.uri().await.unwrap(); + let (store, base) = lance_io::object_store::ObjectStore::from_uri(&uri) + .await + .unwrap(); + let mem_wal = std::path::Path::new(&uri).join("_mem_wal"); + let mut generations = Vec::new(); + for shard in std::fs::read_dir(&mem_wal).unwrap() { + let Ok(shard_id) = shard.unwrap().file_name().to_string_lossy().parse() else { + continue; + }; + let manifest = ShardManifestStore::new(store.clone(), &base, shard_id, 2) + .latest() + .await + .unwrap() + .expect("a shard directory has a manifest"); + for sstable in manifest.sstables { + let path = mem_wal.join(shard_id.to_string()).join(&sstable.path); + generations.push(lance::Dataset::open(path.to_str().unwrap()).await.unwrap()); + } + } + generations + } + + /// Every published generation carries `region_bitmap` as a bitmap index on + /// `region`, and a search for `region = 'fresh'` on it uses that index. + async fn assert_generations_answer_through_region_bitmap(table: &Table, fresh: Vec) { + use lance::index::DatasetIndexExt; + + let generations = published_generations(table).await; + assert!( + !generations.is_empty(), + "closing the writer flushed a generation" + ); + let mut found = Vec::new(); + for generation in generations { + let region = generation.schema().field("region").unwrap().id; + let indices = generation.load_indices().await.unwrap(); + let index = indices + .iter() + .find(|index| index.name == "region_bitmap") + .expect("the generation carries the index the plugin trained"); + assert_eq!(index.fields, vec![region]); + assert!( + index + .index_details + .as_ref() + .is_some_and(|details| details.type_url.ends_with("BitmapIndexDetails")), + "{:?}", + index.index_details + ); + + let mut scan = generation.scan(); + scan.filter("region = 'fresh'").unwrap(); + let plan = scan.explain_plan(false).await.unwrap(); + assert!(plan.contains("ScalarIndexQuery"), "{plan}"); + let batch = scan.try_into_batch().await.unwrap(); + found.extend( + batch["id"] + .as_primitive::() + .values() + .iter() + .copied(), + ); + } + found.sort_unstable(); + assert_eq!(found, fresh); + } + + /// Under the default set, a kind only the registry maintains is picked up + /// when its index is created after the writer opened, and the table's + /// clones share the registry. + #[tokio::test] + async fn lsm_default_set_maintains_a_registry_kind_created_while_writing() { + let dir = tempdir().unwrap(); + let conn = connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let table = conn + .create_table("t", id_region_reader(vec![(1, "us"), (2, "eu")])) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); + let registry = lance::dataset::mem_wal::MemIndexRegistry::default() + .with_plugin(Arc::new(BitmapAsBTree)) + .unwrap(); + table.as_native().unwrap().set_mem_index_registry(registry); + let clone = table.clone(); + assert!( + clone + .as_native() + .unwrap() + .mem_index_registry() + .plugin_for_details_url("/lance.table.BitmapIndexDetails") + .is_some(), + "a clone shares the handle's registry" + ); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + + let upsert = |rows: Vec<(i64, &'static str)>| { + let table = table.clone(); + async move { + let mut builder = table.merge_insert(&["id"]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder.execute(id_region_reader(rows)).await.unwrap(); + } + }; + // Opens the writer before the table has a bitmap index. + upsert(vec![(3, "us")]).await; + table + .create_index(&["region"], crate::index::Index::Bitmap(Default::default())) + .name("region_bitmap".to_string()) + .execute() + .await + .unwrap(); + upsert(vec![(10, "fresh"), (11, "fresh")]).await; + + let (_, _, memtables) = table + .as_native() + .unwrap() + .dataset + .shard_writer() + .read_snapshot() + .await + .unwrap() + .expect("the merge opened a writer"); + let active = memtables.expect("the writer keeps memtables").active; + assert!( + active + .index_store + .index_names() + .contains(&"region_bitmap".to_string()), + "the open writer picked up the new index: {:?}", + active.index_store.index_names() + ); + + table.close_lsm_writers().await.unwrap(); + // Only the generation written after the index exists carries it. + let generations = published_generations(&table).await; + let mut carrying = 0; + for generation in &generations { + use lance::index::DatasetIndexExt; + if generation + .load_indices() + .await + .unwrap() + .iter() + .any(|index| index.name == "region_bitmap") + { + carrying += 1; + } + } + assert!( + carrying >= 1, + "{} generations, none carry the index", + generations.len() + ); } } diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index 0f86fdcac..33bf54535 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -29,7 +29,7 @@ use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef}; use lance::Dataset; use lance::dataset::mem_wal::{ DatasetMemWalExt, ShardWriter, ShardWriterConfig, evaluate_sharding_spec, - validate_maintained_indexes, + validate_maintained_indexes_with, }; use lance::index::DatasetIndexExt; use lance_core::datatypes::Schema as LanceSchema; @@ -103,7 +103,8 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec) } } let dataset = table.dataset.get().await?; - validate_maintained_indexes(&dataset, requested).await?; + validate_maintained_indexes_with(&dataset, requested, &table.mem_index_registry()) + .await?; Some(requested.to_vec()) } None => None, @@ -126,7 +127,9 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec) .into(), }); } - let mut builder = dataset.initialize_mem_wal(); + let mut builder = dataset + .initialize_mem_wal() + .mem_index_registry(table.mem_index_registry()); let writer_config_defaults = match spec { LsmWriteSpec::Bucket { column, @@ -725,7 +728,8 @@ pub(crate) async fn execute_lsm_merge_insert( return Ok(lsm_merge_result(0)); }; - let config = shard_writer_config_from_defaults(&plan.writer_config_defaults); + let mut config = shard_writer_config_from_defaults(&plan.writer_config_defaults); + config.mem_index_registry = table.mem_index_registry(); let writer = table .dataset .shard_writer() diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index d611fe8a5..3ab54ae07 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -106,19 +106,16 @@ async fn can_execute_namespace_query(table: &NativeTable, query: &AnyQuery) -> R } fn requires_local_namespace_execution(query: &AnyQuery) -> bool { - // The namespace QueryTable request has no approx_mode or use_lsm field yet, so - // pushing these down would silently ignore the user's setting. For use_lsm that - // is worse than a tuning miss: MemWAL read routing lives only in `create_plan`, - // so a pushed-down query would return stale base-only data with no error. + // The namespace QueryTable request cannot represent approx_mode, search_effort or use_lsm. + // For use_lsm that is worse than a tuning miss: MemWAL read routing lives only + // in `create_plan`, so a pushed-down query would return stale base-only data. if query.base().use_lsm.is_some() || query.base().take_offsets.is_some() { return true; } matches!( query, - AnyQuery::VectorQuery(VectorQueryRequest { - approx_mode: Some(_), - .. - }) + AnyQuery::VectorQuery(vector_query) + if vector_query.approx_mode.is_some() || vector_query.search_effort.is_some() ) } @@ -166,6 +163,7 @@ pub async fn create_plan( AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query), }; query.base.check_filter()?; + query.validate_search_effort()?; let ds_ref = table.dataset.get().await?; @@ -296,10 +294,24 @@ pub async fn create_plan( scanner.approx_mode(approx_mode.into()); } - scanner.minimum_nprobes(query.minimum_nprobes); - if let Some(maximum_nprobes) = query.maximum_nprobes { + if query.search_effort.is_none() + && let Some(nprobes) = query.nprobes + { + if query.maximum_nprobes == Some(0) { + scanner.minimum_nprobes(nprobes); + } else { + scanner.nprobes(nprobes); + } + } + if let Some(minimum_nprobes) = query.minimum_nprobes { + scanner.minimum_nprobes(minimum_nprobes); + } + if let Some(maximum_nprobes) = query.maximum_nprobes.filter(|nprobes| *nprobes != 0) { scanner.maximum_nprobes(maximum_nprobes); } + if let Some(search_effort) = query.search_effort { + scanner.search_effort(search_effort)?; + } } // For a batch query, `nearest` already applies k to each query vector. @@ -598,6 +610,11 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result { }; // Check for unsupported features + if vq.search_effort.is_some() { + return Err(Error::NotSupported { + message: "namespace QueryTable does not support search_effort; execute the query locally".to_string(), + }); + } if vq.base.reranker.is_some() { return Err(Error::NotSupported { message: "Reranker is not supported for server-side queries".to_string(), @@ -630,7 +647,9 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result { columns, offset: vq.base.offset.map(|o| o as i32), distance_type: vq.distance_type.map(|dt| dt.to_string()), - nprobes: Some(vq.minimum_nprobes as i32), + nprobes: vq.nprobes.map(|nprobes| nprobes as i32), + minimum_nprobes: vq.minimum_nprobes.map(|nprobes| nprobes as i32), + maximum_nprobes: vq.maximum_nprobes.map(|nprobes| nprobes as i32), ef: vq.ef.map(|e| e as i32), refine_factor: vq.refine_factor.map(|r| r as i32), lower_bound: vq.lower_bound, @@ -919,7 +938,6 @@ mod tests { column: Some("vector".to_string()), // We cast here to satisfy the struct definition query_vector: vec![query_vector as Arc], - minimum_nprobes: 20, distance_type: Some(crate::DistanceType::L2), ..Default::default() }; @@ -940,6 +958,7 @@ mod tests { ); assert_eq!(ns_request.vector_column, Some("vector".to_string())); assert_eq!(ns_request.distance_type, Some("l2".to_string())); + assert_eq!(ns_request.nprobes, None); // Verify the vector data was extracted correctly assert!(ns_request.vector.single_vector.is_some()); @@ -949,6 +968,79 @@ mod tests { ); } + #[test] + fn test_namespace_query_probe_pushdown_compatibility() { + let default_query = AnyQuery::VectorQuery(VectorQueryRequest::default()); + assert!(!requires_local_namespace_execution(&default_query)); + + for effort in [0.0, 0.5, 1.0] { + let effort_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + search_effort: Some(effort), + ..Default::default() + }); + assert!(requires_local_namespace_execution(&effort_query)); + assert!(matches!( + convert_to_namespace_query(&effort_query), + Err(Error::NotSupported { message }) if message.contains("search_effort") + )); + } + + let nprobes_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + nprobes: Some(20), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&nprobes_query)); + assert_eq!( + convert_to_namespace_query(&nprobes_query).unwrap().nprobes, + Some(20) + ); + + let minimum_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + minimum_nprobes: Some(5), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&minimum_query)); + let request = convert_to_namespace_query(&minimum_query).unwrap(); + assert_eq!(request.minimum_nprobes, Some(5)); + assert_eq!(request.maximum_nprobes, None); + + let maximum_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + maximum_nprobes: Some(10), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&maximum_query)); + let request = convert_to_namespace_query(&maximum_query).unwrap(); + assert_eq!(request.minimum_nprobes, None); + assert_eq!(request.maximum_nprobes, Some(10)); + + let exact_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + nprobes: Some(8), + minimum_nprobes: Some(20), + maximum_nprobes: Some(20), + ..Default::default() + }); + assert!(!requires_local_namespace_execution(&exact_query)); + let request = convert_to_namespace_query(&exact_query).unwrap(); + assert_eq!(request.nprobes, Some(8)); + assert_eq!(request.minimum_nprobes, Some(20)); + assert_eq!(request.maximum_nprobes, Some(20)); + + let explicitly_unbounded_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![Arc::new(Float32Array::from(vec![1.0, 2.0]))], + nprobes: Some(8), + maximum_nprobes: Some(0), + ..Default::default() + }); + let request = convert_to_namespace_query(&explicitly_unbounded_query).unwrap(); + assert_eq!(request.nprobes, Some(8)); + assert_eq!(request.maximum_nprobes, Some(0)); + } + #[test] fn test_convert_to_namespace_query_plain_query() { let q = QueryRequest { diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs index f3e40fc7b..4798da431 100644 --- a/rust/lancedb/src/table/query/lsm.rs +++ b/rust/lancedb/src/table/query/lsm.rs @@ -170,10 +170,12 @@ fn reject_unsupported(query: &VectorQueryRequest) -> Result<()> { // Vector-only knobs the LSM scanner cannot honor. Both change results rather // than just recall, so error instead of silently ignoring them: distance_range // would return rows outside the bound, and use_index(false) asks for a - // brute-force search the index-only base arm can't do. (ef / approx_mode / - // maximum_nprobes are recall/speed knobs and are left to no-op — and - // maximum_nprobes defaults to Some, so it cannot be rejected on presence.) + // brute-force search the index-only base arm can't do. (ef / approx_mode are + // recall/speed knobs and are left to no-op.) if !query.query_vector.is_empty() { + if query.search_effort.is_some() { + return unsupported("search_effort on vector search"); + } if query.lower_bound.is_some() || query.upper_bound.is_some() { return unsupported("distance_range on vector search"); } @@ -475,7 +477,7 @@ async fn fts_plan( // omits un-compacted documents. Reject rather than mislead. if !index_maintained(dataset, column, details, "InvertedIndexDetails").await? || !resident_memtables_carry(&in_memory, |memtable| { - memtable.index_store.get_fts_by_column(column).is_some() + !memtable.index_store.fts_granularities_on(column).is_empty() }) { return Err(Error::NotSupported { @@ -708,17 +710,14 @@ async fn vector_plan( // The base arm relies on the column's vector index (`fast_search`). Unmaintained, // its catch-up is untracked and exclusion falls back to the compaction watermark, - // dropping compacted SSTables the lagging base index has not re-indexed; and a - // resident MemTable built before the index joined the set carries none either. + // dropping compacted SSTables the lagging base index has not re-indexed. // Reject rather than silently omit rows, mirroring the FTS arm. - if !index_maintained(dataset, &column, details, "VectorIndexDetails").await? - || !resident_memtables_carry(&in_memory, |memtable| { - memtable.index_store.get_hnsw_by_column(&column).is_some() - }) - { + // A resident MemTable needs no index of its own: one that cannot answer + // compares every row. + if !index_maintained(dataset, &column, details, "VectorIndexDetails").await? { return Err(Error::NotSupported { message: format!( - "the MemWAL LSM scanner requires the vector index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes) and carried by every resident MemTable; otherwise compacted rows not yet re-indexed are omitted. set use_lsm(false) to read the base table only" + "the MemWAL LSM scanner requires the vector index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes); otherwise compacted rows not yet re-indexed are omitted. set use_lsm(false) to read the base table only" ), }); } @@ -740,8 +739,20 @@ async fn vector_plan( let mut scanner = base_scanner(dataset, query, pk_columns, snapshots, in_memory)? .with_overfetch_factor(LSM_OVERFETCH_FACTOR) .nearest(&column, query_vector.as_ref(), k)? - .nprobes(query.minimum_nprobes) .distance_metric(distance_type.into()); + if let Some(nprobes) = query.nprobes { + scanner = if query.maximum_nprobes == Some(0) { + scanner.minimum_nprobes(nprobes) + } else { + scanner.nprobes(nprobes) + }; + } + if let Some(minimum_nprobes) = query.minimum_nprobes { + scanner = scanner.minimum_nprobes(minimum_nprobes); + } + if let Some(maximum_nprobes) = query.maximum_nprobes.filter(|nprobes| *nprobes != 0) { + scanner = scanner.maximum_nprobes(maximum_nprobes); + } if let Some(refine_factor) = query.refine_factor { scanner = scanner.refine(refine_factor); } @@ -793,6 +804,53 @@ mod tests { use super::*; use lance_index::mem_wal::{CompactedSsTable, IndexCatchupProgress}; + #[test] + fn search_effort_is_not_silently_ignored() { + let query = VectorQueryRequest { + query_vector: vec![Arc::new(arrow_array::Float32Array::from(vec![1.0, 0.0]))], + search_effort: Some(0.5), + ..Default::default() + }; + assert!( + matches!(reject_unsupported(&query), Err(Error::NotSupported { message }) + if message.contains("search_effort") && message.contains("un-compacted MemWAL")) + ); + } + + #[test] + fn resident_fts_check_distinguishes_columns() { + use arrow_schema::Field; + use lance::dataset::mem_wal::index::{IndexStore, MemIndexSpec}; + + let schema = ArrowSchema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("text", DataType::Utf8, true), + Field::new( + "vector", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), + true, + ), + ]); + let schema = lance_core::datatypes::Schema::try_from(&schema).unwrap(); + let specs = vec![ + MemIndexSpec::btree("id_idx", 0, "id"), + MemIndexSpec::fts("text_idx", 1, "text"), + MemIndexSpec::hnsw( + "vector_idx", + 2, + "vector", + lance_linalg::distance::DistanceType::Cosine, + ), + ]; + let store = IndexStore::from_specs(&specs, &schema, 1_000, 16).unwrap(); + assert!(!store.fts_granularities_on("text").is_empty()); + for column in ["id", "vector", "missing"] { + assert!(store.fts_granularities_on(column).is_empty()); + } + let empty = IndexStore::from_specs(&[], &schema, 1_000, 16).unwrap(); + assert!(empty.fts_granularities_on("text").is_empty()); + } + #[test] fn exclusion_watermark_gates_on_lagging_index_catchup() { let shard = Uuid::from_u128(1); diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 884648f54..0119ba36d 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -1847,8 +1847,8 @@ mod tests { assert_eq!( layouts, vec![ - (1, Some(BlobKind::Managed)), - (2, Some(BlobKind::Managed)), + (1, Some(BlobKind::Dedicated)), + (2, Some(BlobKind::Packed)), (3, Some(BlobKind::Inline)), (4, None), ] diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index 52b7bd5e8..20f209f2e 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -19,15 +19,27 @@ use std::pin::Pin; use crate::error::{Error, Result}; use datafusion_physical_plan::SendableRecordBatchStream; -/// The characters any object name may contain: a table, a namespace segment, a -/// Secret, a materialized view. -/// -/// No positional rule on top of it -- a name may begin with `_`, `-` or `.`, -/// as LanceDB namespaces already do. `.` and `..` are excluded separately, by -/// [`reject_relative_segment`]: that is a property of where a name sits in a -/// URL, not of the name. Length is the service's to bound. -static OBJECT_NAME_REGEX: std::sync::LazyLock = - std::sync::LazyLock::new(|| regex::Regex::new(r"^[A-Za-z0-9_.\-]+$").unwrap()); +/// The maximum length of a table name. Table names are used in URLs and in the keys used with the +/// object store. +const MAX_TABLE_NAME_CHARS: usize = 128; + +/// The maximum length of a component of a namespace. +const MAX_NAMESPACE_NAME_CHARS: usize = 64; + +/// The maximum number of namespace components that are allowed. +const MAX_NAMESPACE_COMPONENTS: usize = 10; + +/// The maximum length of a view name. +const MAX_VIEW_NAME_CHARS: usize = 128; + +/// The maximum length of a secret name. +const MAX_SECRET_NAME_CHARS: usize = 128; + +/// The maximum length of a function name. +const MAX_FUNCTION_NAME_CHARS: usize = 128; + +/// The maximum length of a database name. +const MAX_DB_NAME_CHARS: usize = 64; pub trait PatchStoreParam { fn patch_with_store_wrapper( @@ -86,129 +98,190 @@ impl PatchReadParam for ReadParams { } } -/// The reason `.` and `..` are refused wherever a name becomes a path segment. -const RELATIVE_SEGMENT_REASON: &str = - "'.' and '..' are read as relative path segments and cannot address an object"; - -/// Whether URL parsing would resolve this component away rather than keep it. -/// -/// Exactly `.` and `..`, and their percent-encoded spellings -- resolution -/// happens after decoding, so `%2E%2E` collapses as surely as `..` does, and -/// `drop_table("..")` would reach `/v1/drop/`. No wider than that: `...` is an -/// ordinary segment that addresses fine. -fn is_relative_segment(value: &str) -> bool { - let decoded = value.replace("%2e", ".").replace("%2E", "."); - decoded == "." || decoded == ".." +/// Translate a string into something that can be safely logged. This involves removing control +/// characters that could be used in an exploit like CVE-2003-0063. +fn sanitize(input: &str) -> String { + input.chars().filter(|c| !c.is_control()).collect() } -/// Refuse a path component that URL parsing resolves as a relative segment. -/// -/// Reachable on its own for an identifier with no other validator: a Function -/// name has no client-side grammar, so this is the only rule that applies. -pub(crate) fn reject_relative_segment(what: &str, value: &str) -> Result<()> { - if is_relative_segment(value) { - return Err(Error::InvalidInput { - message: format!("invalid {what} '{value}': {RELATIVE_SEGMENT_REASON}"), +pub fn validate_table_name(name: &str) -> Result<()> { + if name.is_empty() || name == "." || name == ".." { + return Err(Error::InvalidTableName { + name: name.to_string(), + reason: "Illegal table name.".to_string(), + }); + } + if name.len() > MAX_TABLE_NAME_CHARS { + return Err(Error::InvalidTableName { + name: sanitize(name), + reason: format!( + "Table name cannot be more than {MAX_TABLE_NAME_CHARS} characters long." + ), + }); + } + if !name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.') + { + return Err(Error::InvalidTableName { + name: sanitize(name), + reason: "Only alphanumeric characters, dashes, underscores, and\ + periods are allowed in table names." + .to_string(), }); } Ok(()) } -/// Every rule an object name obeys: non-empty, inside [`OBJECT_NAME_REGEX`], -/// and addressable as a path segment. -/// -/// Returns the reason rather than an [`Error`], because the error type is each -/// API's own -- a table reports [`Error::InvalidTableName`], the rest -/// [`Error::InvalidInput`]. Sharing the rules but not the error keeps a table, -/// a namespace segment and a Secret from drifting apart. -fn check_object_name(name: &str) -> std::result::Result<(), &'static str> { - if name.is_empty() { - return Err("it must not be empty"); +/// Validate that an input string contains alphanumerics, or dashes, underscores, and periods. +fn validate_view_or_namespace_or_secret(input: &str, kind: &str, max_len: usize) -> Result<()> { + if input.is_empty() || input == "." || input == ".." { + return Err(Error::InvalidInput { + message: format!("Illegal {} '{}'.", kind, input), + }); } - if !OBJECT_NAME_REGEX.is_match(name) { - return Err( - "it may contain only alphanumeric characters, underscores, hyphens and periods", - ); + if input.len() > max_len { + return Err(Error::InvalidInput { + message: format!( + "Invalid {} '{}' is more than {} characters long.", + kind, + sanitize(input), + max_len + ), + }); } - if is_relative_segment(name) { - return Err(RELATIVE_SEGMENT_REASON); + if !input + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.') + { + return Err(Error::InvalidInput { + message: format!( + "Invalid {} '{}': only alphanumeric characters, \ + dashes, underscores, and periods are allowed.", + kind, + sanitize(input) + ), + }); } Ok(()) } -/// Validate a table name. -pub fn validate_table_name(name: &str) -> Result<()> { - check_object_name(name).map_err(|reason| Error::InvalidTableName { - name: name.to_string(), - reason: reason.to_string(), - }) -} - /// Validate one component of a namespace path -- a single segment, not the /// whole path. [`validate_namespace`] covers a path. -pub fn validate_namespace_name(name: &str) -> Result<()> { - check_object_name(name).map_err(|reason| Error::InvalidInput { - message: format!("invalid namespace name '{name}': {reason}"), - }) +pub fn validate_namespace_component(name: &str) -> Result<()> { + validate_view_or_namespace_or_secret(name, "namespace component", MAX_NAMESPACE_NAME_CHARS) } -/// Validate one component of a Secret identifier: a Secret name, or one segment -/// of the namespace path holding it. -/// -/// The join decides identity, and the service only sees what the split -/// produced. `"a$b"` is not a name the service accepts, but joined and split it -/// reads as the namespace `a` and the name `b` -- a different Secret that may -/// already exist. This is not a second opinion on the name; it is what lets the -/// service have one. -pub fn validate_secret_component(what: &str, value: &str) -> Result<()> { - check_object_name(value).map_err(|reason| Error::InvalidInput { - message: format!("invalid {what} '{value}': {reason}"), - }) +/// Validate all components of a namespace +pub fn validate_namespace(namespace: &[String]) -> Result<()> { + if namespace.len() > 1 && namespace.first().map(String::as_str) == Some("public") { + return Err(Error::InvalidInput { + message: "Namespaces that begin with public cannot have multiple namespace components." + .to_string(), + }); + } + if namespace.len() > MAX_NAMESPACE_COMPONENTS { + return Err(Error::InvalidInput { + message: format!( + "Cannot have a namespace with more than {MAX_NAMESPACE_COMPONENTS} components." + ), + }); + } + for component in namespace { + validate_namespace_component(component)?; + } + Ok(()) +} + +/// Validate a secret name. +pub fn validate_secret_name(name: &str) -> Result<()> { + validate_view_or_namespace_or_secret(name, "secret name", MAX_SECRET_NAME_CHARS) } /// Validate a Secret name and every segment of the namespace path holding it. pub fn validate_secret_reference(name: &str, namespace_path: &[String]) -> Result<()> { - for segment in namespace_path { - validate_secret_component("Secret namespace path segment", segment)?; - } - validate_secret_component("Secret name", name) + validate_namespace(namespace_path)?; + validate_secret_name(name) +} + +pub fn validate_view_name(name: &str) -> Result<()> { + validate_view_or_namespace_or_secret(name, "view name", MAX_VIEW_NAME_CHARS) } /// Validate a view name and every segment of the namespace path holding it. -/// -/// Worded for a view rather than deferring to [`validate_table_name`]: a view -/// is not a table, and a caller who mistypes one should not be told their -/// table name is invalid. The rule is the same one every object name follows, -/// which is what lets the `$`-joined identifier split back apart. pub fn validate_view_reference(name: &str, namespace_path: &[String]) -> Result<()> { - for segment in namespace_path { - validate_view_component("view namespace path segment", segment)?; + validate_namespace(namespace_path)?; + validate_view_name(name) +} + +/// Validate a function name. These do not allow dashes or periods. +pub fn validate_function_name(name: &str) -> Result<()> { + if name.is_empty() { + return Err(Error::InvalidInput { + message: format!("Invalid function name '{}'.", name), + }); } - validate_view_component("view name", name) + if name.len() > MAX_FUNCTION_NAME_CHARS { + return Err(Error::InvalidInput { + message: format!( + "Invalid function name '{}' is more than {} characters long.", + sanitize(name), + MAX_FUNCTION_NAME_CHARS + ), + }); + } + if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { + return Err(Error::InvalidInput { + message: format!( + "Invalid function name '{}': only alphanumeric \ + characters and underscores are allowed.", + sanitize(name) + ), + }); + } + Ok(()) } -/// Validate one component of a view identifier: a view name, or one segment of -/// the namespace path holding it. -pub fn validate_view_component(what: &str, value: &str) -> Result<()> { - check_object_name(value).map_err(|reason| Error::InvalidInput { - message: format!("invalid {what} '{value}': {reason}"), - }) -} - -/// Validate all components of a namespace -/// -/// Iterates through all namespace components and validates each one. -/// Returns an error if any component is invalid. -/// -/// # Arguments -/// * `namespace` - The namespace components to validate -/// -/// # Returns -/// * `Ok(())` if all namespace components are valid -/// * `Err(Error)` if any component is invalid -pub fn validate_namespace(namespace: &[String]) -> Result<()> { - for component in namespace { - validate_namespace_name(component)?; +/// Validate a database name, including slash-separated components. +pub fn validate_database_name(name: &str) -> Result<()> { + if name.len() > MAX_DB_NAME_CHARS { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}' is more than {} characters long.", + sanitize(name), + MAX_DB_NAME_CHARS + ), + }); + } + if !name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' || c == '/') + { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}': only alphanumeric characters, dashes, \ + underscores, periods, and slashes are allowed.", + sanitize(name) + ), + }); + } + for component in name.split('/') { + if component.is_empty() { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}': invalid empty path component.", + sanitize(name) + ), + }); + } + if component == "." || component == ".." { + return Err(Error::InvalidInput { + message: format!( + "Invalid database name '{}': illegal path component.", + sanitize(name) + ), + }); + } } Ok(()) } @@ -494,6 +567,26 @@ pub fn supported_fts_data_type(dtype: &DataType) -> bool { supported_fts_data_type_impl(dtype, false) } +/// Validate FTS input without losing the logical type in field metadata. +pub(crate) fn validate_fts_field(field: &Field) -> Result<()> { + if lance_arrow::json::is_json_field(field) || supported_fts_data_type(field.data_type()) { + return Ok(()); + } + let logical_type = field + .metadata() + .get("ARROW:extension:name") + .map(|name| format!(" (logical type {name})")) + .unwrap_or_default(); + Err(Error::Schema { + message: format!( + "A FTS index cannot be created on the field `{}` which has data type {}{}. FTS supports strings, lists of strings, and lance.json fields stored as LargeBinary; raw binary is not supported", + field.name(), + field.data_type(), + logical_type, + ), + }) +} + fn supported_fts_data_type_impl(dtype: &DataType, in_list: bool) -> bool { match (dtype, in_list) { (DataType::Utf8 | DataType::LargeUtf8, _) => true, @@ -707,6 +800,43 @@ mod tests { use super::*; + #[test] + fn test_fts_field_logical_type_validation() { + let json = lance_arrow::json::json_field("doc", true); + assert!(validate_fts_field(&json).is_ok()); + + for dtype in [DataType::Binary, DataType::LargeBinary] { + let raw = Field::new("doc", dtype.clone(), true); + let err = validate_fts_field(&raw).unwrap_err().to_string(); + assert!(err.contains("raw binary is not supported"), "{err}"); + assert!(err.contains("lance.json"), "{err}"); + + let other_extension = raw.with_metadata(std::collections::HashMap::from([( + "ARROW:extension:name".into(), + "other.json".into(), + )])); + let err = validate_fts_field(&other_extension) + .unwrap_err() + .to_string(); + assert!(err.contains("logical type other.json"), "{err}"); + } + + // A JSON marker on the wrong binary storage type is insufficient. + let wrong_storage = + Field::new("doc", DataType::Binary, true).with_metadata(json.metadata().clone()); + assert!(validate_fts_field(&wrong_storage).is_err()); + for dtype in [DataType::Utf8, DataType::LargeUtf8] { + let text = Field::new("text", dtype, true); + assert!(validate_fts_field(&text).is_ok()); + for list_type in [ + DataType::List(Arc::new(text.clone())), + DataType::LargeList(Arc::new(text.clone())), + ] { + assert!(validate_fts_field(&Field::new("texts", list_type, true)).is_ok()); + } + } + } + #[test] fn test_public_fts_field_path_prefers_exact_case() { let text_list = || { @@ -870,6 +1000,37 @@ mod tests { assert!(err.contains("text.embedding")); } + #[test] + fn test_validate_database_name() { + for name in [ + "db", + "team/search", + "Team_1/db-2.3", + &"a".repeat(MAX_DB_NAME_CHARS), + ] { + assert!(validate_database_name(name).is_ok(), "rejected {name:?}"); + } + for name in [ + "", + "/", + "//", + ".", + "..", + "/db", + "db/", + "a//b", + "a/./b", + "a/../b", + "a$b", + "a b", + "a?b", + "café", + &"a".repeat(MAX_DB_NAME_CHARS + 1), + ] { + assert!(validate_database_name(name).is_err(), "accepted {name:?}"); + } + } + #[test] fn test_validate_table_name() { assert!(validate_table_name("my_table").is_ok()); @@ -890,27 +1051,27 @@ mod tests { } #[test] - fn test_validate_namespace_name() { + fn test_validate_namespace_component() { // Valid namespace names - assert!(validate_namespace_name("ns1").is_ok()); - assert!(validate_namespace_name("namespace_123").is_ok()); - assert!(validate_namespace_name("my-namespace").is_ok()); - assert!(validate_namespace_name("my.namespace").is_ok()); - assert!(validate_namespace_name("NS_1.2.3").is_ok()); - assert!(validate_namespace_name("a").is_ok()); - assert!(validate_namespace_name("123").is_ok()); - assert!(validate_namespace_name("_underscore").is_ok()); - assert!(validate_namespace_name("-hyphen").is_ok()); - assert!(validate_namespace_name(".period").is_ok()); + assert!(validate_namespace_component("ns1").is_ok()); + assert!(validate_namespace_component("namespace_123").is_ok()); + assert!(validate_namespace_component("my-namespace").is_ok()); + assert!(validate_namespace_component("my.namespace").is_ok()); + assert!(validate_namespace_component("NS_1.2.3").is_ok()); + assert!(validate_namespace_component("a").is_ok()); + assert!(validate_namespace_component("123").is_ok()); + assert!(validate_namespace_component("_underscore").is_ok()); + assert!(validate_namespace_component("-hyphen").is_ok()); + assert!(validate_namespace_component(".period").is_ok()); // Invalid namespace names - assert!(validate_namespace_name("").is_err()); - assert!(validate_namespace_name("namespace with spaces").is_err()); - assert!(validate_namespace_name("namespace/with/slashes").is_err()); - assert!(validate_namespace_name("namespace\\with\\backslashes").is_err()); - assert!(validate_namespace_name("namespace$with$delimiter").is_err()); - assert!(validate_namespace_name("namespace@special").is_err()); - assert!(validate_namespace_name("namespace#hash").is_err()); + assert!(validate_namespace_component("").is_err()); + assert!(validate_namespace_component("namespace with spaces").is_err()); + assert!(validate_namespace_component("namespace/with/slashes").is_err()); + assert!(validate_namespace_component("namespace\\with\\backslashes").is_err()); + assert!(validate_namespace_component("namespace$with$delimiter").is_err()); + assert!(validate_namespace_component("namespace@special").is_err()); + assert!(validate_namespace_component("namespace#hash").is_err()); } #[test] @@ -926,6 +1087,9 @@ mod tests { // Empty namespace (root) is valid assert!(validate_namespace(&[]).is_ok()); + // Valid public namespace (which is equivalent to the root namespace.) + assert!(validate_namespace(&["public".to_string()]).is_ok()); + // Invalid: contains empty component assert!(validate_namespace(&["ns1".to_string(), "".to_string()]).is_err()); @@ -937,6 +1101,9 @@ mod tests { assert!(validate_namespace(&["ns1".to_string(), "ns/2".to_string()]).is_err()); assert!(validate_namespace(&["ns1".to_string(), "ns$2".to_string()]).is_err()); + // Invalid: starts with public but contains other components. + assert!(validate_namespace(&["public".to_string(), "foo".to_string()]).is_err()); + // Valid: underscores, hyphens, and periods are allowed assert!( validate_namespace(&["ns_1".to_string(), "ns-2".to_string(), "ns.3".to_string()]) diff --git a/rust/lancedb/tests/blob_integration.rs b/rust/lancedb/tests/blob_integration.rs index dd4051657..063690129 100644 --- a/rust/lancedb/tests/blob_integration.rs +++ b/rust/lancedb/tests/blob_integration.rs @@ -501,17 +501,23 @@ async fn collect_id_rowid(table: &Table) -> Result> { fn assert_missing_blob_row_ids(err: &Error) { assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); let message = err.to_string(); - assert!(message.contains("row ids"), "{message}"); + assert!( + message.contains("row ids but some do not exist in the table"), + "{message}" + ); assert!(!message.contains("rowaddr"), "{message}"); assert!(!message.contains("fragment"), "{message}"); } async fn assert_fetch_apis_reject_missing_row_ids(table: &Table, row_ids: &[u64]) -> Result<()> { + let expected = format!("first missing row ids: [{}]", row_ids[0]); let err = table.fetch_blobs("image", row_ids).await.unwrap_err(); assert_missing_blob_row_ids(&err); + assert!(err.to_string().contains(&expected), "{err}"); let err = table.fetch_blob_files("image", row_ids).await.unwrap_err(); assert_missing_blob_row_ids(&err); + assert!(err.to_string().contains(&expected), "{err}"); let err = table .fetch_blob_ranges( @@ -524,6 +530,7 @@ async fn assert_fetch_apis_reject_missing_row_ids(table: &Table, row_ids: &[u64] .await .unwrap_err(); assert_missing_blob_row_ids(&err); + assert!(err.to_string().contains(&expected), "{err}"); Ok(()) } @@ -611,6 +618,47 @@ async fn blob_columns_lists_nested_dotted_paths() -> Result<()> { Ok(()) } +#[tokio::test] +async fn blob_columns_excludes_list_elements_and_fetches_explain_limitation() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + blob("thumbnail", true), + Field::new("imgs", DataType::List(Arc::new(blob("item", true))), true), + ])); + let table = db + .create_empty_table("list_blobs", schema) + .execute() + .await?; + + assert_eq!(table.blob_columns().await?, vec!["thumbnail"]); + for column in ["imgs", "imgs.item"] { + let err = table.fetch_blobs(column, &[0]).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + assert!( + err.to_string() + .contains("blobs inside lists cannot be fetched") + ); + + let err = table.fetch_blob_files(column, &[0]).await.unwrap_err(); + assert!( + err.to_string() + .contains("blobs inside lists cannot be fetched") + ); + + let err = table + .fetch_blob_ranges(column, [BlobRangeRequest::new(0, 0, 1)]) + .await + .unwrap_err(); + assert!( + err.to_string() + .contains("blobs inside lists cannot be fetched") + ); + } + Ok(()) +} + #[tokio::test] async fn blob_columns_lists_blob_fields_in_order() -> Result<()> { let tmp = tempdir().unwrap(); @@ -790,6 +838,30 @@ async fn fetch_blobs_out_of_range_id_errors_without_panic() -> Result<()> { Ok(()) } +#[tokio::test] +async fn fetch_blob_apis_reject_row_past_fragment_end() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = + create_inline_blob_table(&db, "t", &[1, 2, 3], &[Some(b"a"), Some(b"b"), Some(b"c")]) + .await?; + + assert_fetch_apis_reject_missing_row_ids(&table, &[3]).await?; + assert_fetch_apis_reject_missing_row_ids(&table, &[3, 0]).await +} + +#[tokio::test] +async fn fetch_blobs_identifies_multiple_missing_ids_without_partial_results() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()).execute().await?; + let table = create_inline_blob_table(&db, "t", &[1, 2], &[Some(b"a"), Some(b"b")]).await?; + + let err = table.fetch_blobs("image", &[0, 2, 3]).await.unwrap_err(); + assert_missing_blob_row_ids(&err); + assert!(err.to_string().contains("first missing row ids: [2, 3]")); + Ok(()) +} + #[tokio::test] async fn fetch_blob_files_rejects_missing_fragment_row_addr() -> Result<()> { let tmp = tempdir().unwrap(); @@ -831,6 +903,23 @@ async fn fetch_blob_apis_reject_deleted_row_ids() -> Result<()> { assert_fetch_apis_reject_missing_row_ids(&table, &[deleted_row_addr, live_row_addr]).await } +#[tokio::test] +async fn fetch_blob_apis_reject_deleted_stable_row_ids() -> Result<()> { + let tmp = tempdir().unwrap(); + let db = connect(tmp.path().to_str().unwrap()) + .storage_option(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS, "true") + .execute() + .await?; + let table = create_inline_blob_table(&db, "t", &[1, 2], &[Some(b"one"), Some(b"two")]).await?; + let pairs = collect_id_rowid(&table).await?; + let deleted_row_id = pairs.iter().find(|(id, _)| *id == 2).unwrap().1; + let live_row_id = pairs.iter().find(|(id, _)| *id == 1).unwrap().1; + + table.delete("id = 2").await?; + + assert_fetch_apis_reject_missing_row_ids(&table, &[deleted_row_id, live_row_id]).await +} + #[tokio::test] async fn fetch_blobs_rejects_non_blob_column() -> Result<()> { let tmp = tempdir().unwrap(); diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json index c6bbd83e3..e07c29f98 100644 --- a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_job.json @@ -23,7 +23,7 @@ "nullable": false } }, - "created_at": "2026-08-21T00:00:00Z", + "created_at_millis": 1787270400000, "image": { "manifest_digest": "sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6", "descriptor": { diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json index 63f3c824e..44a84b17c 100644 --- a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json @@ -1 +1 @@ -{"created_at":"2026-08-21T00:00:00Z","disabled":false,"image":{"descriptor":{"behavior":{"result_stability":"input_and_context","side_effects":"non_idempotent"},"entrypoint":"app.function:create","format_version":1,"interface":{"type":"lance.scalar","version":1},"python":{"abi_tag":"cp312","executable":"/usr/local/bin/python3","implementation":"cpython","import_paths":["/opt/function/code"],"version":"3.12.14"},"python_api":1,"requires":{"capabilities":[],"kernel_min":"4.18.0"},"schemas":{"initialization":"/opt/function/schemas/initialization.arrow","input":"/opt/function/schemas/input.arrow","output":"/opt/function/schemas/output.arrow"}},"manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6","source":false},"location":"memory:///fixture","metadata":{},"name":"embed","object_id":"fixture","secret_bindings":[{"kind":"env","secret_ref":{"name":"hf-prod"},"variable":"HF_TOKEN"}],"signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list","kind":"scalar","nullable":false}},"version":"1"} +{"created_at_millis":1787270400000,"disabled":false,"image":{"descriptor":{"behavior":{"result_stability":"input_and_context","side_effects":"non_idempotent"},"entrypoint":"app.function:create","format_version":1,"interface":{"type":"lance.scalar","version":1},"python":{"abi_tag":"cp312","executable":"/usr/local/bin/python3","implementation":"cpython","import_paths":["/opt/function/code"],"version":"3.12.14"},"python_api":1,"requires":{"capabilities":[],"kernel_min":"4.18.0"},"schemas":{"initialization":"/opt/function/schemas/initialization.arrow","input":"/opt/function/schemas/input.arrow","output":"/opt/function/schemas/output.arrow"}},"manifest_digest":"sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6","source":false},"location":"memory:///fixture","metadata":{},"name":"embed","object_id":"fixture","secret_bindings":[{"kind":"env","secret_ref":{"name":"hf-prod"},"variable":"HF_TOKEN"}],"signature":{"inputs":[{"arrow_type":"utf8","name":"text","nullable":true}],"output":{"arrow_type":"list","kind":"scalar","nullable":false}},"version":"1"}