diff --git a/.bumpversion.toml b/.bumpversion.toml index a6912bed5..353755206 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.40.0-beta.12" +current_version = "0.40.0-beta.13" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/.github/release_notes.json b/.github/release_notes.json index a97860698..6593ef66a 100644 --- a/.github/release_notes.json +++ b/.github/release_notes.json @@ -6,6 +6,10 @@ "title": "## 🏆 Highlights", "labels": ["highlight"] }, + { + "title": "## 🚨 Critical Fixes", + "labels": ["critical-fix"] + }, { "title": "## 🛠 Breaking Changes", "labels": ["breaking-change"] @@ -38,4 +42,4 @@ "labels": ["ci"] } ] -} \ No newline at end of file +} diff --git a/.github/workflows/nodejs.yml b/.github/workflows/nodejs.yml index 98acccf02..2e290166e 100644 --- a/.github/workflows/nodejs.yml +++ b/.github/workflows/nodejs.yml @@ -164,6 +164,71 @@ jobs: python ci/mock_openai.py & cd nodejs/examples node_modules/.bin/jest --testEnvironment jest-environment-node-single-context --verbose + arrow-compat: + # The main jobs resolve `apache-arrow` to the version pinned in + # devDependencies (the newest supported release). This job re-runs the + # suite with the library itself running on each older Arrow release the + # peer range still accepts, which is what a consumer on that version gets. + name: Linux (Arrow ${{ matrix.arrow-version }}) + timeout-minutes: 30 + strategy: + matrix: + arrow-version: [ "19.0.1", "20.0.0" ] + runs-on: "ubuntu-22.04" + defaults: + run: + shell: bash + working-directory: nodejs + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + lfs: true + - uses: pnpm/action-setup@v6 + with: + version: 11.1.1 + - uses: actions/setup-node@v6 + with: + node-version: 24 + cache: 'pnpm' + cache-dependency-path: nodejs/pnpm-lock.yaml + - uses: Swatinem/rust-cache@v2 + with: + # Same crate, profile and target as the `linux` job above, so reuse + # its cache: rust-cache keys on the job id by default, and this job + # would otherwise start cold on every run until it had saved its own. + shared-key: linux + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Install dependencies + run: | + sudo apt update + sudo apt install -y protobuf-compiler libssl-dev + - name: Build + run: | + pnpm install --frozen-lockfile + # No `--` separator: pnpm forwards it literally, which would + # make napi-rs treat `--profile ci` as a cargo passthrough arg. + pnpm build:debug --profile ci + - name: Resolve apache-arrow to ${{ matrix.arrow-version }} + run: | + # The older releases are already installed as the `apache-arrow-NN` + # test aliases. With pnpm's hoisted linker each alias is a complete, + # self-contained directory, so swap it in for the library's own + # `apache-arrow` instead of running a second install, which has been + # seen to hang after printing "Done" in CI. A copy, not a symlink: the + # cross-version tests rely on the alias being a separate module + # instance from the one the library loads. + major="$(echo '${{ matrix.arrow-version }}' | cut -d. -f1)" + mv node_modules/apache-arrow node_modules/apache-arrow.pinned + cp -R "node_modules/apache-arrow-${major}" node_modules/apache-arrow + # `require("apache-arrow/package.json")` is blocked by the package's exports map. + node -p "'apache-arrow resolved to ' + JSON.parse(require('fs').readFileSync('node_modules/apache-arrow/package.json')).version" + pnpm tsc + - name: Test + env: + # Newer @smithy/core uses dynamic ESM imports. + NODE_OPTIONS: "--experimental-vm-modules" + run: node_modules/.bin/jest --verbose macos: timeout-minutes: 30 # macos-15 ships a newer linker; the older macos-14 linker fails to insert diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 585abf5ee..cbfefc83d 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -285,7 +285,7 @@ jobs: runs-on: ubuntu-24.04 strategy: matrix: - msrv: ["1.91.0"] # This should match up with rust-version in Cargo.toml + msrv: ["1.94.0"] # This should match up with rust-version in Cargo.toml env: # Need up-to-date compilers for kernels CC: clang-18 @@ -308,40 +308,9 @@ jobs: # unreadable outside their own branch anyway, since GitHub scopes # caches to the creating ref. save-if: ${{ github.ref == 'refs/heads/main' }} - - name: Downgrade dependencies - # These packages have newer requirements for MSRV - run: | - cargo update -p aws-sdk-bedrockruntime --precise 1.77.0 - cargo update -p aws-sdk-dynamodb --precise 1.68.0 - cargo update -p aws-config --precise 1.6.0 - cargo update -p aws-sdk-kms --precise 1.63.0 - cargo update -p aws-sdk-s3 --precise 1.79.0 - cargo update -p aws-sdk-sso --precise 1.62.0 - cargo update -p aws-sdk-ssooidc --precise 1.63.0 - cargo update -p aws-sdk-sts --precise 1.63.0 - # aws-runtime/sigv4/credential-types/types and the aws-smithy-* - # crates bumped their MSRV to 1.91.1 in late 2026; pin to the last - # 1.91.0-compatible versions. The order matters — each downgrade - # only succeeds once everything that still pins it at a higher - # version has itself been downgraded. - cargo update -p aws-runtime --precise 1.5.12 - cargo update -p aws-types --precise 1.3.9 - cargo update -p aws-sigv4 --precise 1.3.5 - cargo update -p aws-credential-types --precise 1.2.8 - # aws-smithy-checksums must stay at or above 0.63.13: OpenDAL's S3 - # service needs crc-fast ~1.9, and older releases pin it to ~1.3. - cargo update -p aws-smithy-checksums --precise 0.63.13 - cargo update -p aws-smithy-runtime --precise 1.9.3 - cargo update -p aws-smithy-http --precise 0.62.6 - cargo update -p aws-smithy-eventstream --precise 0.60.14 - cargo update -p aws-smithy-http-client --precise 1.1.3 - cargo update -p aws-smithy-observability --precise 0.1.4 - cargo update -p aws-smithy-query --precise 0.60.8 - cargo update -p aws-smithy-runtime-api --precise 1.9.3 - cargo update -p aws-smithy-async --precise 1.2.7 - cargo update -p aws-smithy-types --precise 1.3.6 - cargo update -p aws-smithy-xml --precise 0.60.11 - cargo update -p home --precise 0.5.9 + - name: Downgrade dependencies + # aws-smithy-types 1.6.1+ requires rustc 1.94.1 + run: cargo update -p aws-smithy-types --precise 1.5.0 - name: cargo +${{ matrix.msrv }} check env: RUSTUP_TOOLCHAIN: ${{ matrix.msrv }} diff --git a/Cargo.lock b/Cargo.lock index bb986e05a..905f587af 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -141,7 +141,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -152,7 +152,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -181,9 +181,9 @@ dependencies = [ [[package]] name = "arc-swap" -version = "1.9.1" +version = "1.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" dependencies = [ "rustversion", ] @@ -1110,9 +1110,9 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.6.3" +version = "1.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f94d16e797ec62cd999fc9d5942b48fa7050c3093ddadff48e4d7528d16fcb9" +checksum = "0b791f3ac597193fe1d08b82366986eb1f5bc31f2ac6c194c0855276116c76cd" dependencies = [ "base64-simd", "bytes", @@ -1819,7 +1819,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3087,7 +3087,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3315,7 +3315,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "rand 0.9.5", @@ -4538,7 +4538,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.3", "system-configuration", "tokio", "tower-service", @@ -4880,7 +4880,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5143,8 +5143,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arc-swap", "arrow", @@ -5193,7 +5193,7 @@ dependencies = [ "lance-tokenizer", "log", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "permutation", "pin-project", "prost", @@ -5217,7 +5217,7 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "13.1.0-beta.1" +version = "14.0.0-beta.1" dependencies = [ "arrow-array", "arrow-buffer", @@ -5239,7 +5239,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "arrow-buffer", @@ -5253,7 +5253,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "arrow-schema", @@ -5262,8 +5262,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrayref", "crunchy", @@ -5273,8 +5273,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "arrow-buffer", @@ -5293,7 +5293,7 @@ dependencies = [ "log", "moka", "num_cpus", - "object_store 0.14.1", + "object_store 0.14.2", "pin-project", "prost", "quick_cache", @@ -5311,8 +5311,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow", "arrow-array", @@ -5343,8 +5343,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow", "arrow-array", @@ -5361,8 +5361,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "proc-macro2", "quote", @@ -5371,8 +5371,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-arith", "arrow-array", @@ -5405,8 +5405,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-arith", "arrow-array", @@ -5427,7 +5427,7 @@ dependencies = [ "lance-io", "log", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -5438,8 +5438,8 @@ dependencies = [ [[package]] name = "lance-geo" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "datafusion", "geo-traits", @@ -5453,8 +5453,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arc-swap", "arrow", @@ -5504,7 +5504,7 @@ dependencies = [ "log", "ndarray", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -5525,8 +5525,8 @@ dependencies = [ [[package]] name = "lance-index-core" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "arrow-schema", @@ -5549,8 +5549,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow", "arrow-array", @@ -5570,7 +5570,7 @@ dependencies = [ "log", "metrics", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "object_store_opendal", "opendal", "opendal-service-goosefs", @@ -5592,8 +5592,8 @@ dependencies = [ [[package]] name = "lance-linalg" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "arrow-schema", @@ -5607,8 +5607,8 @@ dependencies = [ [[package]] name = "lance-namespace" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow", "async-trait", @@ -5622,8 +5622,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arc-swap", "arrow", @@ -5646,7 +5646,7 @@ dependencies = [ "lance-namespace", "lance-table", "log", - "object_store 0.14.1", + "object_store 0.14.2", "quick-xml 0.40.1", "rand 0.9.5", "reqwest 0.12.28", @@ -5677,8 +5677,8 @@ dependencies = [ [[package]] name = "lance-select" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "arrow-buffer", @@ -5692,8 +5692,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow", "arrow-array", @@ -5714,7 +5714,7 @@ dependencies = [ "lance-io", "lance-select", "log", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -5733,8 +5733,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "arrow-array", "arrow-schema", @@ -5747,8 +5747,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "13.1.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d" +version = "14.0.0-beta.1" +source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3" dependencies = [ "frostem", "icu_segmenter", @@ -5761,7 +5761,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.40.0-beta.12" +version = "0.40.0-beta.13" dependencies = [ "ahash", "anyhow", @@ -5825,7 +5825,7 @@ dependencies = [ "moka", "num-traits", "oauth2", - "object_store 0.14.1", + "object_store 0.14.2", "pin-project", "polars", "polars-arrow", @@ -5858,7 +5858,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.40.0-beta.12" +version = "0.40.0-beta.13" dependencies = [ "arrow-array", "arrow-buffer", @@ -5883,7 +5883,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.40.0-beta.12" +version = "0.40.0-beta.13" dependencies = [ "arc-swap", "arrow", @@ -6475,9 +6475,9 @@ dependencies = [ [[package]] name = "napi" -version = "3.12.2" +version = "3.12.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58c5f4d5375213fdb7be2655e152386e82f026f9a5ba36a75556e11359aafe09" +checksum = "f0c007d4a8ead952a81887661d41fd56b7e7a70d3d4d921f445fb37a8f6efa1e" dependencies = [ "bitflags 2.11.1", "chrono", @@ -6495,15 +6495,15 @@ dependencies = [ [[package]] name = "napi-build" -version = "2.4.2" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "860e7c40864f95cfb83cde99f9ebadd88ef3d9bdccd7dd2cee0cc96a2dd4ffa7" +checksum = "0941983559cd62bfa7f04646a16428967a847365adeef27711e4e3927046bbed" [[package]] name = "napi-derive" -version = "3.6.5" +version = "3.6.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "350057056a30368aa76c11a0d406b0aa61321710be2c36656ab4f6efe0b785a2" +checksum = "6b1d232c24dc6754979b27efdf2047a5849efc7c92c1245de87e6bedbb9f5c3c" dependencies = [ "convert_case", "ctor", @@ -6515,9 +6515,9 @@ dependencies = [ [[package]] name = "napi-derive-backend" -version = "6.1.3" +version = "6.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c1c87a71568f3fe5c736b10878ff55064b250bb4ea4b48c8055713e07b9e463" +checksum = "fd6856fe1686867da90ed0dc292d7c0c87a695eb7c578b913360f33619ee01e3" dependencies = [ "convert_case", "proc-macro2", @@ -6528,9 +6528,9 @@ dependencies = [ [[package]] name = "napi-sys" -version = "3.3.0" +version = "3.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85fbf1fa9f1babfe396d74bbbf52b3643770243e8f5b0b46715d4caf7f0dfc9a" +checksum = "e22a4f25c16a5c5411d987cd6fbd48313522ae1789ade2d6dd3efdc6d40a0fc8" dependencies = [ "libloading", ] @@ -6637,7 +6637,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -6766,7 +6766,7 @@ version = "5.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" dependencies = [ - "base64 0.21.7", + "base64 0.22.1", "chrono", "getrandom 0.2.17", "http 1.5.0", @@ -6883,13 +6883,13 @@ dependencies = [ [[package]] name = "object_store" -version = "0.14.1" +version = "0.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d354792e39fa5f0009e47623cf8b15b099bf9a652fa55c6f817fe28ac84fea50" +checksum = "f1796bc93603f78c5760a69f2d58badc9618d22adade0a95385bb2adbae4eb94" dependencies = [ "async-trait", "aws-lc-rs", - "base64 0.22.1", + "base64 0.23.1", "bytes", "chrono", "crc-fast", @@ -6935,7 +6935,7 @@ dependencies = [ "bytes", "chrono", "futures", - "object_store 0.14.1", + "object_store 0.14.2", "opendal", "pin-project", "tokio", @@ -8140,8 +8140,8 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ - "heck 0.4.1", - "itertools 0.11.0", + "heck 0.5.0", + "itertools 0.14.0", "log", "multimap", "petgraph", @@ -8160,7 +8160,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools 0.11.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -8382,7 +8382,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.45", - "socket2 0.5.10", + "socket2 0.6.3", "thiserror 2.0.18", "tokio", "tracing", @@ -8420,7 +8420,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.3", "tracing", "windows-sys 0.60.2", ] @@ -9220,7 +9220,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9291,7 +9291,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9582,9 +9582,9 @@ dependencies = [ [[package]] name = "serde_with" -version = "3.23.0" +version = "3.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "935177bb8c0cd8ca1a4e6d1a2ac8988bea69cab4f9d3a31311e012ad27868ea4" +checksum = "df9adc193c780ef8f159aee8b61e2d5801aaa555e6eb0947fe45530ec506296f" dependencies = [ "base64 0.23.1", "bs58", @@ -9603,9 +9603,9 @@ dependencies = [ [[package]] name = "serde_with_macros" -version = "3.23.0" +version = "3.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d607aa01a3cb0ad757d6fd216136910db3c97b102fe686585689615a02dbcdc" +checksum = "3e17bbc68e28663bbbb90df47e058aa7eda4fb445b89fe70457bb94fbccf6e49" dependencies = [ "darling 0.24.1", "proc-macro2", @@ -9838,7 +9838,7 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9850,7 +9850,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9873,7 +9873,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -9992,7 +9992,7 @@ dependencies = [ "cfg-if 1.0.4", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -10284,7 +10284,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -11305,7 +11305,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 1d072e654..b42be7a0d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,23 +10,23 @@ repository = "https://github.com/lancedb/lancedb" description = "Serverless, low-latency vector database for AI applications" keywords = ["lancedb", "lance", "database", "vector", "search"] categories = ["database-implementations"] -rust-version = "1.91.0" +rust-version = "1.94.0" [workspace.dependencies] -lance = { "version" = "=13.1.0-beta.1", default-features = false, "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=13.1.0-beta.1", default-features = false, "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=13.1.0-beta.1", default-features = false, "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { version = "=13.1.0-beta.1", path = "python/vendor/lance-arrow" } +lance = { "version" = "=14.0.0-beta.1", default-features = false, "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=14.0.0-beta.1", default-features = false, "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=14.0.0-beta.1", default-features = false, "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { version = "=14.0.0-beta.1", path = "python/vendor/lance-arrow" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow diff --git a/README.md b/README.md index c8a70780d..7bf5d39fc 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,6 @@ - - LanceDB Cloud Public Beta -
+[![Reverie](docs/src/assets/reverie-banner.png)](https://www.reveriesummit.com/?utm_medium=social&utm_source=github&utm_campaign=reverie-2026) [![LanceDB](docs/src/assets/hero-header.png)](https://lancedb.com) [![Website](https://img.shields.io/badge/-Website-100000?style=for-the-badge&labelColor=645cfb&color=645cfb)](https://lancedb.com/) [![Blog](https://img.shields.io/badge/Blog-100000?style=for-the-badge&labelColor=645cfb&color=645cfb)](https://blog.lancedb.com/) @@ -13,7 +11,7 @@ LanceDB -# **The Multimodal AI Lakehouse** +# **The Multimodal Lakehouse for AI** [**How to Install** ](#how-to-install) ✦ [**Detailed Documentation**](https://docs.lancedb.com) ✦ [**Tutorials and Recipes**](https://github.com/lancedb/vectordb-recipes/tree/main) ✦ [**Contributors**](#contributors) diff --git a/docs/src/assets/hero-header.png b/docs/src/assets/hero-header.png index 588053722..aa474b06b 100644 Binary files a/docs/src/assets/hero-header.png and b/docs/src/assets/hero-header.png differ diff --git a/docs/src/assets/reverie-banner.png b/docs/src/assets/reverie-banner.png new file mode 100644 index 000000000..d7d49cb3b Binary files /dev/null and b/docs/src/assets/reverie-banner.png differ diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 0565156dc..8b31b07b2 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.40.0-beta.12 + 0.40.0-beta.13 ``` diff --git a/docs/src/js/classes/AutoQuery.md b/docs/src/js/classes/AutoQuery.md index 1d7ea6952..12f517dbd 100644 --- a/docs/src/js/classes/AutoQuery.md +++ b/docs/src/js/classes/AutoQuery.md @@ -350,7 +350,7 @@ input to this method would be: #### Parameters -* **columns**: `string` \| `string`[] \| `Record`<`string`, `string`> \| `Map`<`string`, `string`> +* **columns**: `string` \| `string`[] \| `Map`<`string`, `string`> \| `Record`<`string`, `string`> #### Returns @@ -417,6 +417,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/docs/src/js/classes/Catalog.md b/docs/src/js/classes/Catalog.md index deac5e3ce..8351380e2 100644 --- a/docs/src/js/classes/Catalog.md +++ b/docs/src/js/classes/Catalog.md @@ -89,19 +89,30 @@ Drop an empty database. The server rejects nonempty databases. ### listDatabases() ```ts -listDatabases(options): Promise +listDatabases(options): DatabaseNames ``` -List one page of databases; pass pageToken from a response for the next page. +Iterate lazily over all database names using `for await...of`. +pageToken resumes from a saved token; omitted starts at the beginning. +pageLimit limits each REST response, not the total; omitted uses the server default. +Request errors are raised during iteration and terminate the iterator. #### Parameters * **options** = `{}` -* **options.limit?**: `number` +* **options.pageLimit?**: `number` * **options.pageToken?**: `string` #### Returns -`Promise`<[`ListDatabasesResponse`](../interfaces/ListDatabasesResponse.md)> +[`DatabaseNames`](DatabaseNames.md) + +#### Example + +```ts +for await (const name of catalog.listDatabases({ pageLimit: 20 })) { + console.log(name); +} +``` diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index f1aace87b..4bb769a0e 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -183,10 +183,6 @@ Define a materialized view named `name` over the table `source`. The view is populated before creation returns. Set `withNoData` to create only its definition and empty backing table. The view is a normal table: it can be queried, indexed and searched, and it appears in `tableNames`. -The source table must have stable row ids (create it with -the `newTableEnableStableRowIds` storage option); they keep the view's -provenance valid across source compactions and cannot be enabled after -a table exists. #### Parameters diff --git a/docs/src/js/classes/DatabaseNames.md b/docs/src/js/classes/DatabaseNames.md new file mode 100644 index 000000000..6b78f93b8 --- /dev/null +++ b/docs/src/js/classes/DatabaseNames.md @@ -0,0 +1,79 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / DatabaseNames + +# Class: DatabaseNames + +A lazy async iterator of database names with pagination state. +Returned by [Catalog.listDatabases](Catalog.md#listdatabases). + +## Implements + +- `AsyncIterableIterator`<`string`> + +## Methods + +### \[asyncIterator\]() + +```ts +asyncIterator: AsyncIterableIterator +``` + +#### Returns + +`AsyncIterableIterator`<`string`> + +#### Implementation of + +`AsyncIterableIterator.[asyncIterator]` + +*** + +### next() + +```ts +next(): Promise> +``` + +Fetch the next name, requesting another page only when needed. + +#### Returns + +`Promise`<`IteratorResult`<`string`, `any`>> + +#### Implementation of + +`AsyncIterableIterator.next` + +*** + +### numPageResults() + +```ts +numPageResults(): number +``` + +Number of names available without another REST request. + +#### Returns + +`number` + +*** + +### pageToken() + +```ts +pageToken(): undefined | string +``` + +Token for the next REST request. Initially this is the supplied starting token. +Returns undefined after the final page is fetched, even if names remain cached. +Drain the cache before saving a token to avoid skipping names when resuming. +A failed request terminates iteration and retains its token for a new iterator. + +#### Returns + +`undefined` \| `string` diff --git a/docs/src/js/classes/MaterializedView.md b/docs/src/js/classes/MaterializedView.md index 58eaa3c13..0e2413328 100644 --- a/docs/src/js/classes/MaterializedView.md +++ b/docs/src/js/classes/MaterializedView.md @@ -65,21 +65,14 @@ refresh(options?): Promise Recompute the view from its source. -The refresh is incremental when the source's changes can be reconciled -into the view -- rows added, changed or removed since the last one -- -and otherwise rebuilds. `full` forces a rebuild; `sourceVersion` -refreshes to that source version instead of the latest. - -Concurrent refreshes of one view do not duplicate its rows. Two that -plan the same source rows conflict on commit, and the loser throws -rather than writing them a second time. +Every refresh rebuilds the complete view. For local views, +`sourceVersion` refreshes to that source version instead of the latest; +remote SQL refreshes do not support source-version pinning. #### Parameters * **options?** -* **options.full?**: `boolean` - * **options.sourceVersion?**: `number` #### Returns diff --git a/docs/src/js/classes/Query.md b/docs/src/js/classes/Query.md index 6ebaebd75..3038bf2b3 100644 --- a/docs/src/js/classes/Query.md +++ b/docs/src/js/classes/Query.md @@ -424,7 +424,7 @@ input to this method would be: #### Parameters -* **columns**: `string` \| `string`[] \| `Record`<`string`, `string`> \| `Map`<`string`, `string`> +* **columns**: `string` \| `string`[] \| `Map`<`string`, `string`> \| `Record`<`string`, `string`> #### Returns @@ -491,6 +491,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/docs/src/js/classes/QueryBase.md b/docs/src/js/classes/QueryBase.md index 35c071525..0ef4fa286 100644 --- a/docs/src/js/classes/QueryBase.md +++ b/docs/src/js/classes/QueryBase.md @@ -193,7 +193,7 @@ input to this method would be: #### Parameters -* **columns**: `string` \| `string`[] \| `Record`<`string`, `string`> \| `Map`<`string`, `string`> +* **columns**: `string` \| `string`[] \| `Map`<`string`, `string`> \| `Record`<`string`, `string`> #### Returns @@ -252,6 +252,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + *** ### withRowId() diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 0d813ff1e..ac16b8a81 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -1381,7 +1381,7 @@ repeatedly calling this method. ##### Parameters -* **updates**: `Record`<`string`, `string`> \| `Map`<`string`, `string`> +* **updates**: `Map`<`string`, `string`> \| `Record`<`string`, `string`> the columns to update diff --git a/docs/src/js/classes/TakeQuery.md b/docs/src/js/classes/TakeQuery.md index 6ae2f9c8c..f3d59a2a3 100644 --- a/docs/src/js/classes/TakeQuery.md +++ b/docs/src/js/classes/TakeQuery.md @@ -200,7 +200,7 @@ input to this method would be: #### Parameters -* **columns**: `string` \| `string`[] \| `Record`<`string`, `string`> \| `Map`<`string`, `string`> +* **columns**: `string` \| `string`[] \| `Map`<`string`, `string`> \| `Record`<`string`, `string`> #### Returns @@ -267,6 +267,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from [`QueryBase`](QueryBase.md).[`toArrow`](QueryBase.md#toarrow) diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index f9412c76d..a0a77fd7c 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -447,6 +447,8 @@ nprobes(nprobes): VectorQuery Set the number of partitions to search (probe) +The number of probes must be greater than 0. + This argument is only used when the vector column has an IVF PQ index. If there is no index then this value is ignored. @@ -673,7 +675,7 @@ input to this method would be: #### Parameters -* **columns**: `string` \| `string`[] \| `Record`<`string`, `string`> \| `Map`<`string`, `string`> +* **columns**: `string` \| `string`[] \| `Map`<`string`, `string`> \| `Record`<`string`, `string`> #### Returns @@ -740,6 +742,8 @@ Collect the results as an Arrow ArrowTable. +Empty results retain the query's output schema. + #### Inherited from `StandardQueryBase.toArrow` diff --git a/docs/src/js/functions/connectCatalog.md b/docs/src/js/functions/connectCatalog.md index 2bdb3182f..ca51c0ed4 100644 --- a/docs/src/js/functions/connectCatalog.md +++ b/docs/src/js/functions/connectCatalog.md @@ -28,5 +28,7 @@ headers; opened database connections inherit authentication and client options. ```ts const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" }); const db = await catalog.createDatabase("analytics", { existOk: true }); -const page = await catalog.listDatabases({ limit: 20 }); +for await (const name of catalog.listDatabases({ pageLimit: 20 })) { + console.log(name); +} ``` diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index a82d9e334..1a939ae28 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -27,6 +27,7 @@ - [Branches](classes/Branches.md) - [Catalog](classes/Catalog.md) - [Connection](classes/Connection.md) +- [DatabaseNames](classes/DatabaseNames.md) - [HeaderProvider](classes/HeaderProvider.md) - [Index](classes/Index.md) - [Job](classes/Job.md) @@ -108,7 +109,6 @@ - [JobEventsOptions](interfaces/JobEventsOptions.md) - [JobFailureInfo](interfaces/JobFailureInfo.md) - [JobInfo](interfaces/JobInfo.md) -- [ListDatabasesResponse](interfaces/ListDatabasesResponse.md) - [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) - [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) - [ListTablesOptions](interfaces/ListTablesOptions.md) @@ -157,11 +157,15 @@ - [AnalyzePlanDistributedMetrics](type-aliases/AnalyzePlanDistributedMetrics.md) - [BaseTokenizer](type-aliases/BaseTokenizer.md) +- [BlobData](type-aliases/BlobData.md) +- [BlobInput](type-aliases/BlobInput.md) - [BlobOptions](type-aliases/BlobOptions.md) - [BlobRange](type-aliases/BlobRange.md) - [BlobReadOptions](type-aliases/BlobReadOptions.md) +- [BlobUri](type-aliases/BlobUri.md) - [Data](type-aliases/Data.md) - [DataLike](type-aliases/DataLike.md) +- [DataTypeLike](type-aliases/DataTypeLike.md) - [FieldLike](type-aliases/FieldLike.md) - [IntoSql](type-aliases/IntoSql.md) - [IntoVector](type-aliases/IntoVector.md) diff --git a/docs/src/js/interfaces/CherryPickResult.md b/docs/src/js/interfaces/CherryPickResult.md index c837ab143..fdec8eedd 100644 --- a/docs/src/js/interfaces/CherryPickResult.md +++ b/docs/src/js/interfaces/CherryPickResult.md @@ -38,9 +38,9 @@ preview: CherryPickPreview; ```ts status: - | "failed" - | "unknown" | "ready" + | "failed" | "notImplemented" - | "cherryPicked"; + | "cherryPicked" + | "unknown"; ``` diff --git a/docs/src/js/interfaces/DropNamespaceOptions.md b/docs/src/js/interfaces/DropNamespaceOptions.md index 1863bb0cd..6e5f01d34 100644 --- a/docs/src/js/interfaces/DropNamespaceOptions.md +++ b/docs/src/js/interfaces/DropNamespaceOptions.md @@ -21,7 +21,7 @@ Refuse to drop if non-empty (restrict) or drop recursively (cascade). ### mode? ```ts -optional mode: "fail" | "skip"; +optional mode: "skip" | "fail"; ``` Whether to skip if the namespace doesn't exist, or fail. diff --git a/docs/src/js/interfaces/ListDatabasesResponse.md b/docs/src/js/interfaces/ListDatabasesResponse.md deleted file mode 100644 index a1a247746..000000000 --- a/docs/src/js/interfaces/ListDatabasesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListDatabasesResponse - -# Interface: ListDatabasesResponse - -## Properties - -### databases - -```ts -databases: string[]; -``` - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` diff --git a/docs/src/js/interfaces/RefreshMaterializedViewResult.md b/docs/src/js/interfaces/RefreshMaterializedViewResult.md index cb7100cd8..1c8b8bc21 100644 --- a/docs/src/js/interfaces/RefreshMaterializedViewResult.md +++ b/docs/src/js/interfaces/RefreshMaterializedViewResult.md @@ -14,7 +14,7 @@ mode: string; ``` -How the view was brought up to date: "rebuild", "incremental" or "no_op". +How the view was brought up to date: "rebuild". *** diff --git a/docs/src/js/namespaces/embedding/classes/EmbeddingFunction.md b/docs/src/js/namespaces/embedding/classes/EmbeddingFunction.md index 574a0d71f..b6291e0be 100644 --- a/docs/src/js/namespaces/embedding/classes/EmbeddingFunction.md +++ b/docs/src/js/namespaces/embedding/classes/EmbeddingFunction.md @@ -52,7 +52,7 @@ new EmbeddingFunction(): EmbeddingFunction ### computeQueryEmbeddings() ```ts -computeQueryEmbeddings(data): Promise +computeQueryEmbeddings(data): Promise ``` Compute the embeddings for a single query @@ -63,7 +63,7 @@ Compute the embeddings for a single query #### Returns -`Promise`<`number`[] \| `Uint8Array` \| `Float32Array` \| `Float64Array`> +`Promise`<`Float32Array` \| `Float64Array` \| `Uint8Array` \| `number`[]> *** diff --git a/docs/src/js/namespaces/embedding/classes/TextEmbeddingFunction.md b/docs/src/js/namespaces/embedding/classes/TextEmbeddingFunction.md index 444c4c3f0..0a0cc619d 100644 --- a/docs/src/js/namespaces/embedding/classes/TextEmbeddingFunction.md +++ b/docs/src/js/namespaces/embedding/classes/TextEmbeddingFunction.md @@ -37,7 +37,7 @@ new TextEmbeddingFunction(): TextEmbeddingFunction ### computeQueryEmbeddings() ```ts -computeQueryEmbeddings(data): Promise +computeQueryEmbeddings(data): Promise ``` Compute the embeddings for a single query @@ -48,7 +48,7 @@ Compute the embeddings for a single query #### Returns -`Promise`<`number`[] \| `Uint8Array` \| `Float32Array` \| `Float64Array`> +`Promise`<`Float32Array` \| `Float64Array` \| `Uint8Array` \| `number`[]> #### Overrides diff --git a/docs/src/js/type-aliases/BlobData.md b/docs/src/js/type-aliases/BlobData.md new file mode 100644 index 000000000..2af0a3659 --- /dev/null +++ b/docs/src/js/type-aliases/BlobData.md @@ -0,0 +1,16 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BlobData + +# Type Alias: BlobData + +```ts +type BlobData: Buffer | Uint8Array | ArrayBuffer | Blob; +``` + +Bytes accepted for a blob value. `ArrayBuffer` is wrapped without copying. +`Blob` and `File` are read with `arrayBuffer()`, which only the async write +paths ([Connection.createTable](../classes/Connection.md#createtable), [Table.add](../classes/Table.md#add), +[Table.mergeInsert](../classes/Table.md#mergeinsert)) can do. diff --git a/docs/src/js/type-aliases/BlobInput.md b/docs/src/js/type-aliases/BlobInput.md new file mode 100644 index 000000000..d1738a760 --- /dev/null +++ b/docs/src/js/type-aliases/BlobInput.md @@ -0,0 +1,36 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BlobInput + +# Type Alias: BlobInput + +```ts +type BlobInput: + | BlobData + | BlobUri + | object + | object + | null + | undefined; +``` + +A value accepted for a `lance.blob.v2` column (see [blob](../functions/blob.md)). + +Either inline bytes, a URI pointing at external bytes, or a struct that sets +exactly one of `data` or `uri`. `null` and `undefined` write a null blob. + +## Example + +```ts +import { pathToFileURL } from "node:url"; +import type { BlobInput } from "@lancedb/lancedb"; + +const rows: { id: bigint; image: BlobInput }[] = [ + { id: 1n, image: await (await fetch(url)).arrayBuffer() }, + { id: 2n, image: pathToFileURL("/data/cat.png") }, + { id: 3n, image: { data: new Blob(["hello"]) } }, +]; +await table.add(rows); +``` diff --git a/docs/src/js/type-aliases/BlobUri.md b/docs/src/js/type-aliases/BlobUri.md new file mode 100644 index 000000000..a90b395fc --- /dev/null +++ b/docs/src/js/type-aliases/BlobUri.md @@ -0,0 +1,13 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BlobUri + +# Type Alias: BlobUri + +```ts +type BlobUri: string | URL; +``` + +A URI accepted for a blob value. A `URL` is stored as its `href`. diff --git a/docs/src/js/type-aliases/DataTypeLike.md b/docs/src/js/type-aliases/DataTypeLike.md new file mode 100644 index 000000000..1f0ee6b81 --- /dev/null +++ b/docs/src/js/type-aliases/DataTypeLike.md @@ -0,0 +1,18 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / DataTypeLike + +# Type Alias: DataTypeLike + +```ts +type DataTypeLike: DataType | object; +``` + +A `DataType` from any copy or version of apache-arrow. + +Arrow 21 brands its classes with `unique symbol` properties, so a type +object from a second copy of the library no longer satisfies the `DataType` +type of this one even though it is structurally identical. Inputs that only +need to be sanitized accept this looser shape instead. diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 66fda4378..814bcb12a 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -37,7 +37,8 @@ are ordinary connections. Dropping a database requires it to be empty. ::: lancedb.catalog.Catalog -::: lancedb.catalog.ListDatabasesResponse +::: lancedb.catalog.DatabaseNames + ## Remote SQL @@ -198,33 +199,6 @@ listing a storage directory. ::: lancedb.materialized_view.MaterializedViewDefinition -::: lancedb.vector_dedup - -::: lancedb.materialized_view.VectorDedupSource - -For an existing indexed source, declare and materialize a dedup result without -writing SQL. Creation captures the source version once and waits for the complete -result; remote connections use the service's MV jobs. - -```python -from lancedb import vector_dedup - -view = db.create_materialized_view( - "images_clean", - vector_dedup("images", column="phash", distance_threshold=4), -) -cleaned = view.table -``` - -To submit without waiting, call `db.create_materialized_view_async(...)` with the -same source expression and wait on the returned job. To declare only, pass -`with_no_data=True`, then use `view.refresh_async().wait()` when ready. Pass -`dataset_version=...` to `vector_dedup` for an explicit snapshot. The default -policy retains direct representatives (A-B-C without A-C keeps A and C), preserves -the source, and materializes its original columns. Only within-segment/partition -pairs are covered, using the index's distance representation. - - ## Views ::: lancedb.view.ViewDescription @@ -440,6 +414,8 @@ latency on your workload before opting in. ::: lancedb.catalog.AsyncCatalog +::: lancedb.catalog.AsyncDatabaseNames + Connections represent a connection to a LanceDb database and can be used to create, list, or open tables. diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 77635c74e..1458d79f0 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.40.0-beta.12 + 0.40.0-beta.13 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 9d7249f92..52c6f84c8 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.40.0-beta.12 + 0.40.0-beta.13 pom ${project.artifactId} LanceDB Java SDK Parent POM @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 13.1.0-beta.1 + 14.0.0-beta.1 false 2.30.0 1.7 diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 051e16474..cb44cf34a 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.40.0-beta.12" +version = "0.40.0-beta.13" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/NODEJS_THIRD_PARTY_LICENSES.md b/nodejs/NODEJS_THIRD_PARTY_LICENSES.md index c0881f61e..550d2888c 100644 --- a/nodejs/NODEJS_THIRD_PARTY_LICENSES.md +++ b/nodejs/NODEJS_THIRD_PARTY_LICENSES.md @@ -275,6 +275,7 @@ [@types/node@20.16.10](https://github.com/DefinitelyTyped/DefinitelyTyped) - MIT [@types/node@20.17.9](https://github.com/DefinitelyTyped/DefinitelyTyped) - MIT [@types/node@22.7.4](https://github.com/DefinitelyTyped/DefinitelyTyped) - MIT +[@types/node@25.9.8](https://github.com/DefinitelyTyped/DefinitelyTyped) - MIT [@types/semver@7.5.6](https://github.com/DefinitelyTyped/DefinitelyTyped) - MIT [@types/stack-utils@2.0.3](https://github.com/DefinitelyTyped/DefinitelyTyped) - MIT [@types/tmp@0.2.6](https://github.com/DefinitelyTyped/DefinitelyTyped) - MIT @@ -306,6 +307,9 @@ [apache-arrow@16.0.0](https://github.com/apache/arrow) - Apache-2.0 [apache-arrow@17.0.0](https://github.com/apache/arrow) - Apache-2.0 [apache-arrow@18.0.0](https://github.com/apache/arrow) - Apache-2.0 +[apache-arrow@19.0.1](https://github.com/apache/arrow) - Apache-2.0 +[apache-arrow@20.0.0](https://github.com/apache/arrow) - Apache-2.0 +[apache-arrow@21.2.0](https://github.com/apache/arrow-js) - Apache-2.0 [argparse@1.0.10](https://github.com/nodeca/argparse) - MIT [argparse@2.0.1](https://github.com/nodeca/argparse) - Python-2.0 [array-back@3.1.0](https://github.com/75lb/array-back) - MIT @@ -348,6 +352,7 @@ [color@4.2.3](https://github.com/Qix-/color) - MIT [combined-stream@1.0.8](https://github.com/felixge/node-combined-stream) - MIT [command-line-args@5.2.1](https://github.com/75lb/command-line-args) - MIT +[command-line-args@6.0.2](https://github.com/75lb/command-line-args) - MIT [command-line-usage@7.0.1](https://github.com/75lb/command-line-usage) - MIT [concat-map@0.0.1](https://github.com/substack/node-concat-map) - MIT [convert-source-map@2.0.0](https://github.com/thlorenz/convert-source-map) - MIT @@ -398,12 +403,14 @@ [file-entry-cache@6.0.1](https://github.com/royriojas/file-entry-cache) - MIT [fill-range@7.1.1](https://github.com/jonschlinkert/fill-range) - MIT [find-replace@3.0.0](https://github.com/75lb/find-replace) - MIT +[find-replace@5.0.2](https://github.com/75lb/find-replace) - MIT [find-up@4.1.0](https://github.com/sindresorhus/find-up) - MIT [find-up@5.0.0](https://github.com/sindresorhus/find-up) - MIT [flat-cache@3.2.0](https://github.com/jaredwray/flat-cache) - MIT [flatbuffers@1.12.0](https://github.com/google/flatbuffers) - Apache* [flatbuffers@23.5.26](https://github.com/google/flatbuffers) - Apache* [flatbuffers@24.3.25](https://github.com/google/flatbuffers) - Apache-2.0 +[flatbuffers@25.9.23](https://github.com/google/flatbuffers) - Apache-2.0 [flatted@3.2.9](https://github.com/WebReflection/flatted) - ISC [follow-redirects@1.15.6](https://github.com/follow-redirects/follow-redirects) - MIT [foreground-child@3.3.0](https://github.com/tapjs/foreground-child) - ISC @@ -644,6 +651,7 @@ [uc.micro@2.1.0](https://github.com/markdown-it/uc.micro) - MIT [undici-types@5.26.5](https://github.com/nodejs/undici) - MIT [undici-types@6.19.8](https://github.com/nodejs/undici) - MIT +[undici-types@7.24.6](https://github.com/nodejs/undici) - MIT [update-browserslist-db@1.0.13](https://github.com/browserslist/update-db) - MIT [uri-js@4.4.1](https://github.com/garycourt/uri-js) - BSD-2-Clause [uuid@9.0.1](https://github.com/uuidjs/uuid) - MIT diff --git a/nodejs/__test__/arrow.test.ts b/nodejs/__test__/arrow.test.ts index df02b4936..376533961 100644 --- a/nodejs/__test__/arrow.test.ts +++ b/nodejs/__test__/arrow.test.ts @@ -2,10 +2,8 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors import * as fs from "node:fs"; import * as vm from "node:vm"; -import * as arrow15 from "apache-arrow-15"; -import * as arrow16 from "apache-arrow-16"; -import * as arrow17 from "apache-arrow-17"; -import * as arrow18 from "apache-arrow-18"; + +import { type ApacheArrow, arrow15, arrowVersions } from "./arrow_versions"; import { Field as CurrentField, @@ -110,1541 +108,1539 @@ it("preserves field metadata from a provided schema", async function () { expect(roundTripped.schema.fields[0].metadata).toEqual(jsonMetadata); }); -describe.each([arrow15, arrow16, arrow17, arrow18])( - "Arrow", - ( - arrow: typeof arrow15 | typeof arrow16 | typeof arrow17 | typeof arrow18, - ) => { - type ApacheArrow = - | typeof arrow15 - | typeof arrow16 - | typeof arrow17 - | typeof arrow18; - const { - Schema, - Field, - Binary, - Bool, - Utf8, - Float64, - Struct, - List, - Map_, - Int16, - Int32, - Int64, - Float, - Float16, - Float32, - FixedSizeList, - Precision, - tableFromIPC, - DataType, - Dictionary, - RecordBatch: ArrowRecordBatch, - Table: ArrowTable, - Uint8: ArrowUint8, - makeData: arrowMakeData, - vectorFromArray, +describe.each(arrowVersions)("Arrow %#", (arrow: ApacheArrow) => { + const { + Schema, + Field, + Binary, + Bool, + Utf8, + Float64, + Struct, + List, + Map_, + Int16, + Int32, + Int64, + Float, + Float16, + Float32, + FixedSizeList, + Precision, + tableFromIPC, + DataType, + Dictionary, + RecordBatch: ArrowRecordBatch, + Table: ArrowTable, + Uint8: ArrowUint8, + makeData: arrowMakeData, + vectorFromArray, + // biome-ignore lint/suspicious/noExplicitAny: + } = arrow; + type Schema = ApacheArrow["Schema"]; + type Table = ApacheArrow["Table"]; + + function expectValidMapField( + // biome-ignore lint/suspicious/noExplicitAny: Arrow Field types vary across supported versions + field: any, + ): void { + expect(DataType.isMap(field.type)).toBe(true); + expect(field.type.keysSorted).toBe(true); + expect(field.type.children).toHaveLength(1); + + const entries = field.type.children[0]; + expect(entries.name).toBe("entries"); + expect(entries.nullable).toBe(false); + expect(DataType.isStruct(entries.type)).toBe(true); + expect(entries.type.children).toHaveLength(2); + + const [key, value] = entries.type.children; + expect([key.name, value.name]).toEqual(["key", "value"]); + expect(key.nullable).toBe(false); + expect(DataType.isUtf8(key.type)).toBe(true); + expect(value.nullable).toBe(true); + expect(DataType.isInt(value.type)).toBe(true); + expect(value.type.bitWidth).toBe(32); + expect(value.type.isSigned).toBe(true); + } + + // Helper method to verify various ways to create a table + async function checkTableCreation( + tableCreationMethod: ( + records: Record[], + recordsReversed: Record[], + schema: Schema, + ) => Promise, + infersTypes: boolean, + ): Promise { + const records = sampleRecords(); + const recordsReversed = [ + { + list: ["anime", "action", "comedy"], + struct: { x: 0, y: 0 }, + string: "hello", + number: 7, + boolean: false, + binary: Buffer.alloc(5), + }, + ]; + const schema = new Schema([ + new Field("binary", new Binary(), false), + new Field("boolean", new Bool(), false), + new Field("number", new Float64(), false), + new Field("string", new Utf8(), false), + new Field( + "struct", + new Struct([ + new Field("x", new Float64(), false), + new Field("y", new Float64(), false), + ]), + ), + new Field("list", new List(new Field("item", new Utf8(), false)), false), + ]); + const table = (await tableCreationMethod( + records, + recordsReversed, + schema, // biome-ignore lint/suspicious/noExplicitAny: - } = arrow; - type Schema = ApacheArrow["Schema"]; - type Table = ApacheArrow["Table"]; + )) as any; - function expectValidMapField( - // biome-ignore lint/suspicious/noExplicitAny: Arrow Field types vary across supported versions - field: any, - ): void { - expect(DataType.isMap(field.type)).toBe(true); - expect(field.type.keysSorted).toBe(true); - expect(field.type.children).toHaveLength(1); + // We expect deterministic ordering of the fields + expect(table.schema.names).toEqual(schema.names); - const entries = field.type.children[0]; - expect(entries.name).toBe("entries"); - expect(entries.nullable).toBe(false); - expect(DataType.isStruct(entries.type)).toBe(true); - expect(entries.type.children).toHaveLength(2); + schema.fields.forEach( + ( + // biome-ignore lint/suspicious/noExplicitAny: + field: { name: any; type: { toString: () => any } }, + idx: string | number, + ) => { + const actualField = table.schema.fields[idx]; + // Type inference always assumes nullable=true + if (infersTypes) { + expect(actualField.nullable).toBe(true); + } else { + expect(actualField.nullable).toBe(false); + } + expect(table.getChild(field.name)?.type.toString()).toEqual( + field.type.toString(), + ); + expect(table.getChildAt(idx)?.type.toString()).toEqual( + field.type.toString(), + ); + }, + ); + } - const [key, value] = entries.type.children; - expect([key.name, value.name]).toEqual(["key", "value"]); - expect(key.nullable).toBe(false); - expect(DataType.isUtf8(key.type)).toBe(true); - expect(value.nullable).toBe(true); - expect(DataType.isInt(value.type)).toBe(true); - expect(value.type.bitWidth).toBe(32); - expect(value.type.isSigned).toBe(true); - } - - // Helper method to verify various ways to create a table - async function checkTableCreation( - tableCreationMethod: ( - records: Record[], - recordsReversed: Record[], - schema: Schema, - ) => Promise
, - infersTypes: boolean, - ): Promise { - const records = sampleRecords(); - const recordsReversed = [ - { - list: ["anime", "action", "comedy"], - struct: { x: 0, y: 0 }, - string: "hello", - number: 7, - boolean: false, - binary: Buffer.alloc(5), - }, + describe("The function makeArrowTable", function () { + it("accepts snake_case embedding metadata like camelCase", function () { + const spellings = [ + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + { source_column: "text", vector_column: "vector" }, + { sourceColumn: "text", vectorColumn: "vector" }, ]; - const schema = new Schema([ - new Field("binary", new Binary(), false), - new Field("boolean", new Bool(), false), - new Field("number", new Float64(), false), - new Field("string", new Utf8(), false), - new Field( - "struct", - new Struct([ - new Field("x", new Float64(), false), - new Field("y", new Float64(), false), + for (const columns of spellings) { + const schema = new Schema( + [ + new Field("text", new Utf8(), false), + new Field( + "vector", + new FixedSizeList(3, new Field("item", new Float32(), true)), + false, + ), + ], + new Map([ + [ + "embedding_functions", + JSON.stringify([{ name: "mock", model: {}, ...columns }]), + ], ]), - ), + ); + // The vector field is non-nullable and absent from the data; only a + // recognized embedding config makes that acceptable. + const table = makeArrowTable([{ text: "hello" }], { schema }); + expect(table.numRows).toBe(1); + } + }); + + it("will use data types from a provided schema instead of inference", async function () { + const schema = new Schema([ + new Field("a", new Int32(), false), + new Field("b", new Float32(), true), + new Field("c", new FixedSizeList(3, new Field("item", new Float16()))), + new Field("d", new Int64(), true), + ]); + const table = makeArrowTable( + [ + { a: 1, b: 2, c: [1, 2, 3], d: 9 }, + { a: 4, b: 5, c: [4, 5, 6], d: 10 }, + { a: 7, b: 8, c: [7, 8, 9], d: null }, + ], + { schema }, + ); + + const buf = await fromTableToBuffer(table); + expect(buf.byteLength).toBeGreaterThan(0); + + const actual = tableFromIPC(buf); + expect(actual.numRows).toBe(3); + const actualSchema = actual.schema; + expect(actualSchema).toEqual(schema); + expect(table.getChild("a")?.toJSON()).toEqual([1, 4, 7]); + expect(table.getChild("b")?.toJSON()).toEqual([2, 5, 8]); + expect(table.getChild("d")?.toJSON()).toEqual([9n, 10n, null]); + }); + + it("will use a provided FixedSizeList schema with typed array values", function () { + const schema = new Schema([ + new Field("text", new Utf8(), false), new Field( - "list", - new List(new Field("item", new Utf8(), false)), + "vector", + new FixedSizeList(3, new Field("item", new Float32(), false)), false, ), ]); - const table = (await tableCreationMethod( - records, - recordsReversed, - schema, - // biome-ignore lint/suspicious/noExplicitAny: - )) as any; - // We expect deterministic ordering of the fields - expect(table.schema.names).toEqual(schema.names); + const table = makeArrowTable( + [ + { + text: "foo", + vector: new Float32Array([1, 2, 3]), + }, + ], + { schema }, + ); - schema.fields.forEach( - ( - // biome-ignore lint/suspicious/noExplicitAny: - field: { name: any; type: { toString: () => any } }, - idx: string | number, - ) => { - const actualField = table.schema.fields[idx]; - // Type inference always assumes nullable=true - if (infersTypes) { - expect(actualField.nullable).toBe(true); - } else { - expect(actualField.nullable).toBe(false); - } - expect(table.getChild(field.name)?.type.toString()).toEqual( - field.type.toString(), - ); - expect(table.getChildAt(idx)?.type.toString()).toEqual( - field.type.toString(), - ); + expect(table.getChild("text")?.toJSON()).toEqual(["foo"]); + expect( + table + .getChild("vector") + ?.toJSON() + .map((value) => value.toJSON()), + ).toEqual([[1, 2, 3]]); + }); + + it("will assume the column `vector` is FixedSizeList by default", async function () { + const schema = new Schema([ + new Field("a", new Float(Precision.DOUBLE), true), + new Field("b", new Int64(), true), + new Field( + "vector", + new FixedSizeList( + 3, + new Field("item", new Float(Precision.SINGLE), true), + ), + true, + ), + ]); + const table = makeArrowTable([ + { a: 1, b: 2n, vector: [1, 2, 3] }, + { a: 4, b: 5n, vector: [4, 5, 6] }, + { a: 7, b: 8n, vector: [7, 8, 9] }, + ]); + + const buf = await fromTableToBuffer(table); + expect(buf.byteLength).toBeGreaterThan(0); + + const actual = tableFromIPC(buf); + expect(actual.numRows).toBe(3); + const actualSchema = actual.schema; + expect(actualSchema).toEqual(schema); + + expect(table.getChild("a")?.toJSON()).toEqual([1, 4, 7]); + expect(table.getChild("b")?.toJSON()).toEqual([2n, 5n, 8n]); + expect( + table + .getChild("vector") + ?.toJSON() + .map((v) => v.toJSON()), + ).toEqual([ + [1, 2, 3], + [4, 5, 6], + [7, 8, 9], + ]); + }); + + it("can support multiple vector columns", async function () { + const schema = new Schema([ + new Field("a", new Float(Precision.DOUBLE), true), + new Field("b", new Float(Precision.DOUBLE), true), + new Field( + "vec1", + new FixedSizeList(3, new Field("item", new Float16(), true)), + true, + ), + new Field( + "vec2", + new FixedSizeList(3, new Field("item", new Float64(), true)), + true, + ), + ]); + const table = makeArrowTable( + [ + { a: 1, b: 2, vec1: [1, 2, 3], vec2: [2, 4, 6] }, + { a: 4, b: 5, vec1: [4, 5, 6], vec2: [8, 10, 12] }, + { a: 7, b: 8, vec1: [7, 8, 9], vec2: [14, 16, 18] }, + ], + { + vectorColumns: { + vec1: { type: new Float16() }, + vec2: { type: new Float64() }, + }, }, ); - } - describe("The function makeArrowTable", function () { - it("accepts snake_case embedding metadata like camelCase", function () { - const spellings = [ - // biome-ignore lint/style/useNamingConvention: the Python wire spelling - { source_column: "text", vector_column: "vector" }, - { sourceColumn: "text", vectorColumn: "vector" }, - ]; - for (const columns of spellings) { - const schema = new Schema( - [ - new Field("text", new Utf8(), false), - new Field( - "vector", - new FixedSizeList(3, new Field("item", new Float32(), true)), - false, - ), - ], - new Map([ - [ - "embedding_functions", - JSON.stringify([{ name: "mock", model: {}, ...columns }]), - ], - ]), - ); - // The vector field is non-nullable and absent from the data; only a - // recognized embedding config makes that acceptable. - const table = makeArrowTable([{ text: "hello" }], { schema }); - expect(table.numRows).toBe(1); - } - }); + const buf = await fromTableToBuffer(table); + expect(buf.byteLength).toBeGreaterThan(0); - it("will use data types from a provided schema instead of inference", async function () { - const schema = new Schema([ - new Field("a", new Int32(), false), - new Field("b", new Float32(), true), - new Field( - "c", - new FixedSizeList(3, new Field("item", new Float16())), - ), - new Field("d", new Int64(), true), - ]); - const table = makeArrowTable( - [ - { a: 1, b: 2, c: [1, 2, 3], d: 9 }, - { a: 4, b: 5, c: [4, 5, 6], d: 10 }, - { a: 7, b: 8, c: [7, 8, 9], d: null }, - ], - { schema }, - ); + const actual = tableFromIPC(buf); + expect(actual.numRows).toBe(3); + const actualSchema = actual.schema; + expect(actualSchema).toEqual(schema); + }); - const buf = await fromTableToBuffer(table); - expect(buf.byteLength).toBeGreaterThan(0); + it("will detect vector columns when name contains 'vector' or 'embedding'", async function () { + // Test various naming patterns that should be detected as vector columns + const floatVectorTable = makeArrowTable([ + { + // Float vectors (use decimal values to ensure they're treated as floats) + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + user_vector: [1.1, 2.2], + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + text_embedding: [3.3, 4.4], + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + doc_embeddings: [5.5, 6.6], + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + my_vector_field: [7.7, 8.8], + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + embedding_model: [9.9, 10.1], + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + VECTOR_COL: [11.1, 12.2], // uppercase + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + Vector_Mixed: [13.3, 14.4], // mixed case + }, + ]); - const actual = tableFromIPC(buf); - expect(actual.numRows).toBe(3); - const actualSchema = actual.schema; - expect(actualSchema).toEqual(schema); - expect(table.getChild("a")?.toJSON()).toEqual([1, 4, 7]); - expect(table.getChild("b")?.toJSON()).toEqual([2, 5, 8]); - expect(table.getChild("d")?.toJSON()).toEqual([9n, 10n, null]); - }); + // Check that columns with 'vector' or 'embedding' in name are converted to FixedSizeList + const floatVectorColumns = [ + "user_vector", + "text_embedding", + "doc_embeddings", + "my_vector_field", + "embedding_model", + "VECTOR_COL", + "Vector_Mixed", + ]; - it("will use a provided FixedSizeList schema with typed array values", function () { - const schema = new Schema([ - new Field("text", new Utf8(), false), - new Field( - "vector", - new FixedSizeList(3, new Field("item", new Float32(), false)), - false, - ), - ]); - - const table = makeArrowTable( - [ - { - text: "foo", - vector: new Float32Array([1, 2, 3]), - }, - ], - { schema }, - ); - - expect(table.getChild("text")?.toJSON()).toEqual(["foo"]); + for (const columnName of floatVectorColumns) { expect( - table - .getChild("vector") - ?.toJSON() - .map((value) => value.toJSON()), - ).toEqual([[1, 2, 3]]); - }); - - it("will assume the column `vector` is FixedSizeList by default", async function () { - const schema = new Schema([ - new Field("a", new Float(Precision.DOUBLE), true), - new Field("b", new Int64(), true), - new Field( - "vector", - new FixedSizeList( - 3, - new Field("item", new Float(Precision.SINGLE), true), - ), - true, - ), - ]); - const table = makeArrowTable([ - { a: 1, b: 2n, vector: [1, 2, 3] }, - { a: 4, b: 5n, vector: [4, 5, 6] }, - { a: 7, b: 8n, vector: [7, 8, 9] }, - ]); - - const buf = await fromTableToBuffer(table); - expect(buf.byteLength).toBeGreaterThan(0); - - const actual = tableFromIPC(buf); - expect(actual.numRows).toBe(3); - const actualSchema = actual.schema; - expect(actualSchema).toEqual(schema); - - expect(table.getChild("a")?.toJSON()).toEqual([1, 4, 7]); - expect(table.getChild("b")?.toJSON()).toEqual([2n, 5n, 8n]); + DataType.isFixedSizeList(floatVectorTable.getChild(columnName)?.type), + ).toBe(true); + // Check that float vectors use Float32 by default expect( - table - .getChild("vector") - ?.toJSON() - .map((v) => v.toJSON()), - ).toEqual([ - [1, 2, 3], - [4, 5, 6], - [7, 8, 9], - ]); - }); + floatVectorTable + .getChild(columnName) + ?.type.children[0].type.toString(), + ).toEqual(new Float32().toString()); + } - it("can support multiple vector columns", async function () { - const schema = new Schema([ - new Field("a", new Float(Precision.DOUBLE), true), - new Field("b", new Float(Precision.DOUBLE), true), - new Field( - "vec1", - new FixedSizeList(3, new Field("item", new Float16(), true)), - true, - ), - new Field( - "vec2", - new FixedSizeList(3, new Field("item", new Float64(), true)), - true, - ), - ]); - const table = makeArrowTable( - [ - { a: 1, b: 2, vec1: [1, 2, 3], vec2: [2, 4, 6] }, - { a: 4, b: 5, vec1: [4, 5, 6], vec2: [8, 10, 12] }, - { a: 7, b: 8, vec1: [7, 8, 9], vec2: [14, 16, 18] }, - ], - { - vectorColumns: { - vec1: { type: new Float16() }, - vec2: { type: new Float64() }, - }, - }, - ); + // Test that regular integer arrays still get treated as float vectors + // (since JavaScript doesn't distinguish integers from floats at runtime) + const integerArrayTable = makeArrowTable([ + { + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + vector_int: [1, 2], // Regular array with integers - should be Float32 + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + embedding_int: [3, 4], // Regular array with integers - should be Float32 + }, + ]); - const buf = await fromTableToBuffer(table); - expect(buf.byteLength).toBeGreaterThan(0); - - const actual = tableFromIPC(buf); - expect(actual.numRows).toBe(3); - const actualSchema = actual.schema; - expect(actualSchema).toEqual(schema); - }); - - it("will detect vector columns when name contains 'vector' or 'embedding'", async function () { - // Test various naming patterns that should be detected as vector columns - const floatVectorTable = makeArrowTable([ - { - // Float vectors (use decimal values to ensure they're treated as floats) - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - user_vector: [1.1, 2.2], - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - text_embedding: [3.3, 4.4], - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - doc_embeddings: [5.5, 6.6], - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - my_vector_field: [7.7, 8.8], - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - embedding_model: [9.9, 10.1], - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - VECTOR_COL: [11.1, 12.2], // uppercase - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - Vector_Mixed: [13.3, 14.4], // mixed case - }, - ]); - - // Check that columns with 'vector' or 'embedding' in name are converted to FixedSizeList - const floatVectorColumns = [ - "user_vector", - "text_embedding", - "doc_embeddings", - "my_vector_field", - "embedding_model", - "VECTOR_COL", - "Vector_Mixed", - ]; - - for (const columnName of floatVectorColumns) { - expect( - DataType.isFixedSizeList( - floatVectorTable.getChild(columnName)?.type, - ), - ).toBe(true); - // Check that float vectors use Float32 by default - expect( - floatVectorTable - .getChild(columnName) - ?.type.children[0].type.toString(), - ).toEqual(new Float32().toString()); - } - - // Test that regular integer arrays still get treated as float vectors - // (since JavaScript doesn't distinguish integers from floats at runtime) - const integerArrayTable = makeArrowTable([ - { - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - vector_int: [1, 2], // Regular array with integers - should be Float32 - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - embedding_int: [3, 4], // Regular array with integers - should be Float32 - }, - ]); - - const integerArrayColumns = ["vector_int", "embedding_int"]; - - for (const columnName of integerArrayColumns) { - expect( - DataType.isFixedSizeList( - integerArrayTable.getChild(columnName)?.type, - ), - ).toBe(true); - // Regular integer arrays should use Float32 (avoiding false positives) - expect( - integerArrayTable - .getChild(columnName) - ?.type.children[0].type.toString(), - ).toEqual(new Float32().toString()); - } - - // Test normal list should NOT be converted to FixedSizeList - const normalListTable = makeArrowTable([ - { - // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns - normal_list: [15.5, 16.6], // should NOT be detected as vector - }, - ]); + const integerArrayColumns = ["vector_int", "embedding_int"]; + for (const columnName of integerArrayColumns) { expect( DataType.isFixedSizeList( - normalListTable.getChild("normal_list")?.type, + integerArrayTable.getChild(columnName)?.type, ), - ).toBe(false); - expect( - DataType.isList(normalListTable.getChild("normal_list")?.type), ).toBe(true); - }); - - it("will allow different vector column types", async function () { - const table = makeArrowTable([{ fp16: [1], fp32: [1], fp64: [1] }], { - vectorColumns: { - fp16: { type: new Float16() }, - fp32: { type: new Float32() }, - fp64: { type: new Float64() }, - }, - }); - + // Regular integer arrays should use Float32 (avoiding false positives) expect( - table.getChild("fp16")?.type.children[0].type.toString(), - ).toEqual(new Float16().toString()); - expect( - table.getChild("fp32")?.type.children[0].type.toString(), + integerArrayTable + .getChild(columnName) + ?.type.children[0].type.toString(), ).toEqual(new Float32().toString()); - expect( - table.getChild("fp64")?.type.children[0].type.toString(), - ).toEqual(new Float64().toString()); + } + + // Test normal list should NOT be converted to FixedSizeList + const normalListTable = makeArrowTable([ + { + // biome-ignore lint/style/useNamingConvention: Testing vector column detection patterns + normal_list: [15.5, 16.6], // should NOT be detected as vector + }, + ]); + + expect( + DataType.isFixedSizeList(normalListTable.getChild("normal_list")?.type), + ).toBe(false); + expect( + DataType.isList(normalListTable.getChild("normal_list")?.type), + ).toBe(true); + }); + + it("will allow different vector column types", async function () { + const table = makeArrowTable([{ fp16: [1], fp32: [1], fp64: [1] }], { + vectorColumns: { + fp16: { type: new Float16() }, + fp32: { type: new Float32() }, + fp64: { type: new Float64() }, + }, }); - it("will infer FixedSizeList from Float32Array values", async function () { - const table = makeArrowTable([ - { id: "a", vector: new Float32Array([0.1, 0.2, 0.3]) }, - { id: "b", vector: new Float32Array([0.4, 0.5, 0.6]) }, - ]); + expect(table.getChild("fp16")?.type.children[0].type.toString()).toEqual( + new Float16().toString(), + ); + expect(table.getChild("fp32")?.type.children[0].type.toString()).toEqual( + new Float32().toString(), + ); + expect(table.getChild("fp64")?.type.children[0].type.toString()).toEqual( + new Float64().toString(), + ); + }); - expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( + it("will infer FixedSizeList from Float32Array values", async function () { + const table = makeArrowTable([ + { id: "a", vector: new Float32Array([0.1, 0.2, 0.3]) }, + { id: "b", vector: new Float32Array([0.4, 0.5, 0.6]) }, + ]); + + expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( + true, + ); + const vectorType = table.getChild("vector")?.type; + expect(vectorType.listSize).toBe(3); + expect(vectorType.children[0].type.toString()).toEqual( + new Float32().toString(), + ); + }); + + it("will infer FixedSizeList from Uint8Array values", async function () { + const table = makeArrowTable([ + { id: "a", vector: new Uint8Array([1, 2, 3]) }, + { id: "b", vector: new Uint8Array([4, 5, 6]) }, + ]); + + expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( + true, + ); + const vectorType = table.getChild("vector")?.type; + expect(vectorType.listSize).toBe(3); + expect(vectorType.children[0].type.toString()).toEqual( + new ArrowUint8().toString(), + ); + }); + + it("will use dictionary encoded strings if asked", async function () { + const table = makeArrowTable([{ str: "hello" }]); + expect(DataType.isUtf8(table.getChild("str")?.type)).toBe(true); + + const tableWithDict = makeArrowTable([{ str: "hello" }], { + dictionaryEncodeStrings: true, + }); + expect(DataType.isDictionary(tableWithDict.getChild("str")?.type)).toBe( + true, + ); + + const schema = new Schema([ + new Field("str", new Dictionary(new Utf8(), new Int32())), + ]); + + const tableWithDict2 = makeArrowTable([{ str: "hello" }], { schema }); + expect(DataType.isDictionary(tableWithDict2.getChild("str")?.type)).toBe( + true, + ); + }); + + it("will infer data types correctly", async function () { + await checkTableCreation( + // biome-ignore lint/suspicious/noExplicitAny: + async (records) => (makeArrowTable)(records), + true, + ); + }); + + it("will allow matching inferred types across records", function () { + expect(() => makeArrowTable([{ value: 1 }, { value: 2 }])).not.toThrow(); + }); + + it("will reject mismatched inferred types across records", function () { + expect(() => makeArrowTable([{ value: 1 }, { value: "two" }])).toThrow( + "Failed to infer schema for data. Previously inferred type Float64 but found Utf8 for field value at row 1. Consider providing an explicit schema.", + ); + }); + + it.each([ + ["ArrayBuffer", new Uint8Array([104]).buffer], + ["URL", new URL("https://example.com/")], + ["empty object", {}], + [ + "class instance with enumerable fields", + Object.assign(new Blob(["x"]), { position: 0 }), + ], + ])("rejects %s values without a schema", (_name, value) => { + expect(() => makeArrowTable([{ id: 1, value }])).toThrow( + /field value at row 0/, + ); + expect(() => makeArrowTable([{ value }])).toThrow(/field value at row 0/); + }); + + it.each([ + ["Binary", new Binary(), new Uint8Array([104]).buffer], + ["Utf8", new Utf8(), new URL("https://example.com/")], + ])("rejects %s values with a schema", (_name, type, value) => { + const schema = new Schema([new Field("value", type, true)]); + expect(() => makeArrowTable([{ value }], { schema })).toThrow( + /field value at row 0/, + ); + }); + + it("rejects unsupported values in nested struct fields", () => { + const schema = new Schema([ + new Field( + "metadata", + new Struct([new Field("bytes", new Binary(), true)]), true, - ); - const vectorType = table.getChild("vector")?.type; - expect(vectorType.listSize).toBe(3); - expect(vectorType.children[0].type.toString()).toEqual( - new Float32().toString(), - ); - }); + ), + ]); + expect(() => + makeArrowTable( + [{ metadata: { bytes: new Uint8Array([104]).buffer } }], + { schema }, + ), + ).toThrow(/field metadata\.bytes at row 0/); + }); - it("will infer FixedSizeList from Uint8Array values", async function () { - const table = makeArrowTable([ - { id: "a", vector: new Uint8Array([1, 2, 3]) }, - { id: "b", vector: new Uint8Array([4, 5, 6]) }, - ]); - - expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( + it("rejects unsupported objects in list elements", () => { + const schema = new Schema([ + new Field( + "images", + new List(new Field("item", new Binary(), true)), true, - ); - const vectorType = table.getChild("vector")?.type; - expect(vectorType.listSize).toBe(3); - expect(vectorType.children[0].type.toString()).toEqual( - new ArrowUint8().toString(), - ); - }); + ), + ]); + expect(() => + makeArrowTable([{ images: [new Blob(["abc"])] }], { schema }), + ).toThrow(/field images\[0\] at row 0/); + }); - it("will use dictionary encoded strings if asked", async function () { - const table = makeArrowTable([{ str: "hello" }]); - expect(DataType.isUtf8(table.getChild("str")?.type)).toBe(true); - - const tableWithDict = makeArrowTable([{ str: "hello" }], { - dictionaryEncodeStrings: true, - }); - expect(DataType.isDictionary(tableWithDict.getChild("str")?.type)).toBe( - true, - ); - - const schema = new Schema([ - new Field("str", new Dictionary(new Utf8(), new Int32())), - ]); - - const tableWithDict2 = makeArrowTable([{ str: "hello" }], { schema }); - expect( - DataType.isDictionary(tableWithDict2.getChild("str")?.type), - ).toBe(true); - }); - - it("will infer data types correctly", async function () { - await checkTableCreation( - // biome-ignore lint/suspicious/noExplicitAny: - async (records) => (makeArrowTable)(records), - true, - ); - }); - - it("will allow matching inferred types across records", function () { - expect(() => - makeArrowTable([{ value: 1 }, { value: 2 }]), - ).not.toThrow(); - }); - - it("will reject mismatched inferred types across records", function () { - expect(() => makeArrowTable([{ value: 1 }, { value: "two" }])).toThrow( - "Failed to infer schema for data. Previously inferred type Float64 but found Utf8 for field value at row 1. Consider providing an explicit schema.", - ); - }); - - it.each([ - ["ArrayBuffer", new Uint8Array([104]).buffer], - ["URL", new URL("https://example.com/")], - ["empty object", {}], - [ - "class instance with enumerable fields", - Object.assign(new Blob(["x"]), { position: 0 }), - ], - ])("rejects %s values without a schema", (_name, value) => { - expect(() => makeArrowTable([{ id: 1, value }])).toThrow( - /field value at row 0/, - ); - expect(() => makeArrowTable([{ value }])).toThrow( - /field value at row 0/, - ); - }); - - it.each([ - ["Binary", new Binary(), new Uint8Array([104]).buffer], - ["Utf8", new Utf8(), new URL("https://example.com/")], - ])("rejects %s values with a schema", (_name, type, value) => { - const schema = new Schema([new Field("value", type, true)]); - expect(() => makeArrowTable([{ value }], { schema })).toThrow( - /field value at row 0/, - ); - }); - - it("rejects unsupported values in nested struct fields", () => { - const schema = new Schema([ - new Field( - "metadata", - new Struct([new Field("bytes", new Binary(), true)]), - true, - ), - ]); - expect(() => - makeArrowTable( - [{ metadata: { bytes: new Uint8Array([104]).buffer } }], - { schema }, - ), - ).toThrow(/field metadata\.bytes at row 0/); - }); - - it("rejects unsupported objects in list elements", () => { - const schema = new Schema([ - new Field( - "images", - new List(new Field("item", new Binary(), true)), - true, - ), - ]); - expect(() => - makeArrowTable([{ images: [new Blob(["abc"])] }], { schema }), - ).toThrow(/field images\[0\] at row 0/); - }); - - it("rejects unsupported objects in nested lists", () => { - const schema = new Schema([ - new Field( - "images", - new List( - new Field( - "item", - new List(new Field("item", new Binary(), true)), - true, - ), + it("rejects unsupported objects in nested lists", () => { + const schema = new Schema([ + new Field( + "images", + new List( + new Field( + "item", + new List(new Field("item", new Binary(), true)), + true, ), - true, ), - ]); - expect(() => - makeArrowTable([{ images: [[new Blob(["abc"])]] }], { schema }), - ).toThrow(/field images\[0\]\[0\] at row 0/); - }); + true, + ), + ]); + expect(() => + makeArrowTable([{ images: [[new Blob(["abc"])]] }], { schema }), + ).toThrow(/field images\[0\]\[0\] at row 0/); + }); - it("rejects unsupported nested struct fields in lists", () => { - const schema = new Schema([ - new Field( - "items", - new List( - new Field( - "item", - new Struct([new Field("bytes", new Binary(), true)]), - true, - ), + it("rejects unsupported nested struct fields in lists", () => { + const schema = new Schema([ + new Field( + "items", + new List( + new Field( + "item", + new Struct([new Field("bytes", new Binary(), true)]), + true, ), - true, ), - ]); - expect(() => - makeArrowTable( - [{ items: [{ bytes: new Uint8Array([104]).buffer }] }], - { schema }, - ), - ).toThrow(/field items\[0\]\.bytes at row 0/); + true, + ), + ]); + expect(() => + makeArrowTable([{ items: [{ bytes: new Uint8Array([104]).buffer }] }], { + schema, + }), + ).toThrow(/field items\[0\]\.bytes at row 0/); + }); + + it("will ignore generated dictionary IDs when comparing inferred types", function () { + const table = makeArrowTable([{ str: "a" }, { str: "b" }], { + dictionaryEncodeStrings: true, }); - it("will ignore generated dictionary IDs when comparing inferred types", function () { - const table = makeArrowTable([{ str: "a" }, { str: "b" }], { - dictionaryEncodeStrings: true, - }); + expect(table.getChild("str")?.toJSON()).toEqual(["a", "b"]); + }); - expect(table.getChild("str")?.toJSON()).toEqual(["a", "b"]); - }); + it("will preserve null values without treating them as type mismatches", function () { + for (const records of [ + [{ vector: [1, 2, 3] }, { vector: null }], + [{ vector: null }, { vector: [1, 2, 3] }], + ]) { + const table = makeArrowTable(records); - it("will preserve null values without treating them as type mismatches", function () { - for (const records of [ - [{ vector: [1, 2, 3] }, { vector: null }], - [{ vector: null }, { vector: [1, 2, 3] }], - ]) { - const table = makeArrowTable(records); + expect(table.numRows).toBe(2); + expect(table.getChild("vector")?.nullCount).toBe(1); + } + }); - expect(table.numRows).toBe(2); - expect(table.getChild("vector")?.nullCount).toBe(1); - } - }); - - it("will preserve empty variable-size lists", function () { - for (const records of [ - [{ items: [1] }, { items: [] }], - [{ items: [] }, { items: [1] }], - ]) { - const table = makeArrowTable(records); - expect( - table - .getChild("items") - ?.toJSON() - .map((value) => value.toJSON()), - ).toEqual(records.map((record) => record.items)); - } - }); - - it("will propagate deferred evidence through nested lists", function () { - for (const records of [ - [{ items: [1] }, { items: [null] }], - [{ items: [null] }, { items: [1] }], - [{ items: [null, 1] }, { items: [2, null] }], - ]) { - const table = makeArrowTable(records); - expect( - table - .getChild("items") - ?.toJSON() - .map((value) => value.toJSON()), - ).toEqual(records.map((record) => record.items)); - } - - const nestedRecords = [{ items: [[1]] }, { items: [[null]] }]; - const nestedTable = makeArrowTable(nestedRecords); + it("will preserve empty variable-size lists", function () { + for (const records of [ + [{ items: [1] }, { items: [] }], + [{ items: [] }, { items: [1] }], + ]) { + const table = makeArrowTable(records); expect( - nestedTable + table .getChild("items") ?.toJSON() - .map((value) => - value - .toJSON() - .map((nestedValue: { toJSON: () => unknown[] }) => - nestedValue.toJSON(), - ), - ), - ).toEqual(nestedRecords.map((record) => record.items)); - }); + .map((value) => value.toJSON()), + ).toEqual(records.map((record) => record.items)); + } + }); - it("will reject incompatible deferred evidence within a list", function () { - for (const items of [ - [[], 1], - [1, []], - [[null], 1], - [1, [null]], - ]) { - expect(() => makeArrowTable([{ items }])).toThrow( - "Failed to infer data type for field items at row 0.", - ); - } - }); + it("will propagate deferred evidence through nested lists", function () { + for (const records of [ + [{ items: [1] }, { items: [null] }], + [{ items: [null] }, { items: [1] }], + [{ items: [null, 1] }, { items: [2, null] }], + ]) { + const table = makeArrowTable(records); + expect( + table + .getChild("items") + ?.toJSON() + .map((value) => value.toJSON()), + ).toEqual(records.map((record) => record.items)); + } - it("will reject empty fixed-size lists", function () { - expect(() => - makeArrowTable([{ vector: [1, 2, 3] }, { vector: [] }]), - ).toThrow( - "Failed to infer schema for data. Previously inferred type FixedSizeList[3] but found List[0] for field vector at row 1.", + const nestedRecords = [{ items: [[1]] }, { items: [[null]] }]; + const nestedTable = makeArrowTable(nestedRecords); + expect( + nestedTable + .getChild("items") + ?.toJSON() + .map((value) => + value + .toJSON() + .map((nestedValue: { toJSON: () => unknown[] }) => + nestedValue.toJSON(), + ), + ), + ).toEqual(nestedRecords.map((record) => record.items)); + }); + + it("will reject incompatible deferred evidence within a list", function () { + for (const items of [ + [[], 1], + [1, []], + [[null], 1], + [1, [null]], + ]) { + expect(() => makeArrowTable([{ items }])).toThrow( + "Failed to infer data type for field items at row 0.", ); - }); + } + }); - it("will reject inferred leaf and branch shape changes", function () { - expect(() => - makeArrowTable([{ value: 1 }, { value: { nested: 2 } }]), - ).toThrow( - "Failed to infer schema for data. Previously inferred type Float64 but found Struct for field value at row 1.", + it("will reject empty fixed-size lists", function () { + expect(() => + makeArrowTable([{ vector: [1, 2, 3] }, { vector: [] }]), + ).toThrow( + "Failed to infer schema for data. Previously inferred type FixedSizeList[3] but found List[0] for field vector at row 1.", + ); + }); + + it("will reject inferred leaf and branch shape changes", function () { + expect(() => + makeArrowTable([{ value: 1 }, { value: { nested: 2 } }]), + ).toThrow( + "Failed to infer schema for data. Previously inferred type Float64 but found Struct for field value at row 1.", + ); + expect(() => + makeArrowTable([{ value: { nested: 1 } }, { value: 2 }]), + ).toThrow( + "Failed to infer schema for data. Previously inferred type Struct but found Float64 for field value at row 1.", + ); + }); + + it("will allow null values around inferred struct values", function () { + for (const { records, nullIndex } of [ + { + records: [{ value: null }, { value: { nested: 2 } }], + nullIndex: 0, + }, + { + records: [{ value: { nested: 1 } }, { value: null }], + nullIndex: 1, + }, + ]) { + const table = makeArrowTable(records); + const values = table.getChild("value"); + + expect(values?.nullCount).toBe(1); + expect(values?.get(nullIndex)).toBeNull(); + } + }); + + it("will allow null values around struct values with a schema", function () { + const schema = new Schema([ + new Field("value", new Struct([new Field("nested", new Int32())])), + ]); + for (const { records, nullIndex } of [ + { + records: [{ value: null }, { value: { nested: 2 } }], + nullIndex: 0, + }, + { + records: [{ value: { nested: 1 } }, { value: null }], + nullIndex: 1, + }, + ]) { + const table = makeArrowTable(records, { schema }); + const values = table.getChild("value"); + + expect(values?.nullCount).toBe(1); + expect(values?.get(nullIndex)).toBeNull(); + } + }); + + it("will reject non-object values for struct fields in a schema", function () { + const schema = new Schema([ + new Field("value", new Struct([new Field("nested", new Int32())])), + ]); + for (const records of [ + [{ value: 5 }], + [{ value: 5 }, { value: { nested: 1 } }], + [{ value: { nested: 1 } }, { value: "oops" }], + ]) { + expect(() => makeArrowTable(records, { schema })).toThrow( + /Expected a struct value for field value at row \d, got/, ); - expect(() => - makeArrowTable([{ value: { nested: 1 } }, { value: 2 }]), - ).toThrow( - "Failed to infer schema for data. Previously inferred type Struct but found Float64 for field value at row 1.", - ); - }); + } + }); - it("will allow null values around inferred struct values", function () { - for (const { records, nullIndex } of [ - { - records: [{ value: null }, { value: { nested: 2 } }], - nullIndex: 0, - }, - { - records: [{ value: { nested: 1 } }, { value: null }], - nullIndex: 1, - }, - ]) { - const table = makeArrowTable(records); - const values = table.getChild("value"); - - expect(values?.nullCount).toBe(1); - expect(values?.get(nullIndex)).toBeNull(); - } - }); - - it("will allow a schema to be provided", async function () { - await checkTableCreation( - async (records, _, schema) => - // biome-ignore lint/suspicious/noExplicitAny: - (makeArrowTable)(records, { schema }), - false, - ); - }); - - it("will use the field order of any provided schema", async function () { - await checkTableCreation( - async (_, recordsReversed, schema) => - // biome-ignore lint/suspicious/noExplicitAny: - (makeArrowTable)(recordsReversed, { schema }), - false, - ); - }); - - it("will make an empty table", async function () { - await checkTableCreation( + it("will allow a schema to be provided", async function () { + await checkTableCreation( + async (records, _, schema) => // biome-ignore lint/suspicious/noExplicitAny: - async (_, __, schema) => (makeArrowTable)([], { schema }), - false, - ); - }); + (makeArrowTable)(records, { schema }), + false, + ); + }); - it("will allow subsets of columns if nullable", async function () { - const schema = new Schema([ - new Field("a", new Int64(), true), + it("will use the field order of any provided schema", async function () { + await checkTableCreation( + async (_, recordsReversed, schema) => + // biome-ignore lint/suspicious/noExplicitAny: + (makeArrowTable)(recordsReversed, { schema }), + false, + ); + }); + + it("will make an empty table", async function () { + await checkTableCreation( + // biome-ignore lint/suspicious/noExplicitAny: + async (_, __, schema) => (makeArrowTable)([], { schema }), + false, + ); + }); + + it("will allow subsets of columns if nullable", async function () { + const schema = new Schema([ + new Field("a", new Int64(), true), + new Field( + "s", + new Struct([ + new Field("x", new Int32(), true), + new Field("y", new Int32(), true), + ]), + true, + ), + new Field("d", new Int16(), true), + ]); + + const table = makeArrowTable([{ a: 1n }], { schema }); + expect(table.numCols).toBe(1); + expect(table.numRows).toBe(1); + + const table2 = makeArrowTable([{ a: 1n, d: 2 }], { schema }); + expect(table2.numCols).toBe(2); + + const table3 = makeArrowTable([{ s: { y: 3 } }], { schema }); + expect(table3.numCols).toBe(1); + const expectedSchema = new Schema([ + new Field("s", new Struct([new Field("y", new Int32(), true)]), true), + ]); + expect(table3.schema).toEqual(expectedSchema); + }); + + it("will work even if columns are sparsely provided", async function () { + const sparseRecords = [{ a: 1n }, { b: 2n }, { c: 3n }, { d: 4n }]; + const table = makeArrowTable(sparseRecords); + expect(table.numCols).toBe(4); + expect(table.numRows).toBe(4); + + const schema = new Schema([ + new Field("a", new Int64(), true), + new Field("b", new Int32(), true), + new Field("c", new Int64(), true), + new Field("d", new Int16(), true), + ]); + const table2 = makeArrowTable(sparseRecords, { schema }); + expect(table2.numCols).toBe(4); + expect(table2.numRows).toBe(4); + expect(table2.schema).toEqual(schema); + }); + + it("will handle missing columns in schema alignment when using embeddings", async function () { + const schema = new Schema( + [ + new Field("domain", new Utf8(), true), + new Field("name", new Utf8(), true), + new Field("description", new Utf8(), true), + ], + new Map([["embedding_functions", JSON.stringify([])]]), + ); + + const data = [ + { domain: "google.com", name: "Google" }, + { domain: "facebook.com", name: "Facebook" }, + ]; + + const table = await convertToTable(data, undefined, { schema }); + + expect(table.numCols).toBe(3); + expect(table.numRows).toBe(2); + + const descriptionColumn = table.getChild("description"); + expect(descriptionColumn).toBeDefined(); + expect(descriptionColumn?.nullCount).toBe(2); + expect(descriptionColumn?.toArray()).toEqual([null, null]); + + expect(table.getChild("domain")?.toArray()).toEqual([ + "google.com", + "facebook.com", + ]); + expect(table.getChild("name")?.toArray()).toEqual(["Google", "Facebook"]); + }); + + it("will handle completely missing nested struct columns", async function () { + const schema = new Schema( + [ + new Field("id", new Utf8(), true), + new Field("name", new Utf8(), true), new Field( - "s", + "metadata", new Struct([ - new Field("x", new Int32(), true), - new Field("y", new Int32(), true), + new Field("version", new Int32(), true), + new Field("author", new Utf8(), true), + new Field( + "tags", + new List(new Field("item", new Utf8(), true)), + true, + ), ]), true, ), - new Field("d", new Int16(), true), - ]); - - const table = makeArrowTable([{ a: 1n }], { schema }); - expect(table.numCols).toBe(1); - expect(table.numRows).toBe(1); - - const table2 = makeArrowTable([{ a: 1n, d: 2 }], { schema }); - expect(table2.numCols).toBe(2); - - const table3 = makeArrowTable([{ s: { y: 3 } }], { schema }); - expect(table3.numCols).toBe(1); - const expectedSchema = new Schema([ - new Field("s", new Struct([new Field("y", new Int32(), true)]), true), - ]); - expect(table3.schema).toEqual(expectedSchema); - }); - - it("will work even if columns are sparsely provided", async function () { - const sparseRecords = [{ a: 1n }, { b: 2n }, { c: 3n }, { d: 4n }]; - const table = makeArrowTable(sparseRecords); - expect(table.numCols).toBe(4); - expect(table.numRows).toBe(4); - - const schema = new Schema([ - new Field("a", new Int64(), true), - new Field("b", new Int32(), true), - new Field("c", new Int64(), true), - new Field("d", new Int16(), true), - ]); - const table2 = makeArrowTable(sparseRecords, { schema }); - expect(table2.numCols).toBe(4); - expect(table2.numRows).toBe(4); - expect(table2.schema).toEqual(schema); - }); - - it("will handle missing columns in schema alignment when using embeddings", async function () { - const schema = new Schema( - [ - new Field("domain", new Utf8(), true), - new Field("name", new Utf8(), true), - new Field("description", new Utf8(), true), - ], - new Map([["embedding_functions", JSON.stringify([])]]), - ); - - const data = [ - { domain: "google.com", name: "Google" }, - { domain: "facebook.com", name: "Facebook" }, - ]; - - const table = await convertToTable(data, undefined, { schema }); - - expect(table.numCols).toBe(3); - expect(table.numRows).toBe(2); - - const descriptionColumn = table.getChild("description"); - expect(descriptionColumn).toBeDefined(); - expect(descriptionColumn?.nullCount).toBe(2); - expect(descriptionColumn?.toArray()).toEqual([null, null]); - - expect(table.getChild("domain")?.toArray()).toEqual([ - "google.com", - "facebook.com", - ]); - expect(table.getChild("name")?.toArray()).toEqual([ - "Google", - "Facebook", - ]); - }); - - it("will handle completely missing nested struct columns", async function () { - const schema = new Schema( - [ - new Field("id", new Utf8(), true), - new Field("name", new Utf8(), true), - new Field( - "metadata", - new Struct([ - new Field("version", new Int32(), true), - new Field("author", new Utf8(), true), - new Field( - "tags", - new List(new Field("item", new Utf8(), true)), - true, - ), - ]), - true, - ), - ], - new Map([["embedding_functions", JSON.stringify([])]]), - ); - - const data = [ - { id: "doc1", name: "Document 1" }, - { id: "doc2", name: "Document 2" }, - ]; - - const table = await convertToTable(data, undefined, { schema }); - - expect(table.numCols).toBe(3); - expect(table.numRows).toBe(2); - - const buf = await fromTableToBuffer(table); - const retrievedTable = tableFromIPC(buf); - - const rows = []; - for (let i = 0; i < retrievedTable.numRows; i++) { - rows.push(retrievedTable.get(i)); - } - - expect(rows[0].metadata.version).toBe(null); - expect(rows[0].metadata.author).toBe(null); - expect(rows[0].metadata.tags).toBe(null); - expect(rows[0].id).toBe("doc1"); - expect(rows[0].name).toBe("Document 1"); - }); - - it("will handle partially missing nested struct fields", async function () { - const schema = new Schema( - [ - new Field("id", new Utf8(), true), - new Field( - "metadata", - new Struct([ - new Field("version", new Int32(), true), - new Field("author", new Utf8(), true), - new Field("created_at", new Utf8(), true), - ]), - true, - ), - ], - new Map([["embedding_functions", JSON.stringify([])]]), - ); - - const data = [ - { id: "doc1", metadata: { version: 1, author: "Alice" } }, - { id: "doc2", metadata: { version: 2 } }, - ]; - - const table = await convertToTable(data, undefined, { schema }); - - expect(table.numCols).toBe(2); - expect(table.numRows).toBe(2); - - const metadataColumn = table.getChild("metadata"); - expect(metadataColumn).toBeDefined(); - expect(metadataColumn?.type.toString()).toBe( - "Struct<{version:Int32, author:Utf8, created_at:Utf8}>", - ); - }); - - it("will handle multiple levels of nested structures", async function () { - const schema = new Schema( - [ - new Field("id", new Utf8(), true), - new Field( - "config", - new Struct([ - new Field("database", new Utf8(), true), - new Field( - "connection", - new Struct([ - new Field("host", new Utf8(), true), - new Field("port", new Int32(), true), - new Field( - "ssl", - new Struct([ - new Field("enabled", new Bool(), true), - new Field("cert_path", new Utf8(), true), - ]), - true, - ), - ]), - true, - ), - ]), - true, - ), - ], - new Map([["embedding_functions", JSON.stringify([])]]), - ); - - const data = [ - { - id: "config1", - config: { - database: "postgres", - connection: { host: "localhost" }, - }, - }, - { - id: "config2", - config: { database: "mysql" }, - }, - { - id: "config3", - }, - ]; - - const table = await convertToTable(data, undefined, { schema }); - - expect(table.numCols).toBe(2); - expect(table.numRows).toBe(3); - - const configColumn = table.getChild("config"); - expect(configColumn).toBeDefined(); - expect(configColumn?.type.toString()).toBe( - "Struct<{database:Utf8, connection:Struct<{host:Utf8, port:Int32, ssl:Struct<{enabled:Bool, cert_path:Utf8}>}>}>", - ); - }); - - it("will handle missing columns in Arrow table input when using embeddings", async function () { - const incompleteTable = makeArrowTable([ - { domain: "google.com", name: "Google" }, - { domain: "facebook.com", name: "Facebook" }, - ]); - - const schema = new Schema( - [ - new Field("domain", new Utf8(), true), - new Field("name", new Utf8(), true), - new Field("description", new Utf8(), true), - ], - new Map([["embedding_functions", JSON.stringify([])]]), - ); - - const buf = await fromDataToBuffer(incompleteTable, undefined, schema); - - expect(buf.byteLength).toBeGreaterThan(0); - - const retrievedTable = tableFromIPC(buf); - expect(retrievedTable.numCols).toBe(3); - expect(retrievedTable.numRows).toBe(2); - - const descriptionColumn = retrievedTable.getChild("description"); - expect(descriptionColumn).toBeDefined(); - expect(descriptionColumn?.nullCount).toBe(2); - expect(descriptionColumn?.toArray()).toEqual([null, null]); - - expect(retrievedTable.getChild("domain")?.toArray()).toEqual([ - "google.com", - "facebook.com", - ]); - expect(retrievedTable.getChild("name")?.toArray()).toEqual([ - "Google", - "Facebook", - ]); - }); - - it("should correctly retain values in nested struct fields", async function () { - const testData = [ - { - id: "doc1", - vector: [1, 2, 3], - metadata: { - filePath: "/path/to/file1.ts", - startLine: 10, - endLine: 20, - text: "function test() { return true; }", - }, - }, - { - id: "doc2", - vector: [4, 5, 6], - metadata: { - filePath: "/path/to/file2.ts", - startLine: 30, - endLine: 40, - text: "function test2() { return false; }", - }, - }, - ]; - - const table = makeArrowTable(testData); - - const metadataField = table.schema.fields.find( - (f) => f.name === "metadata", - ); - expect(metadataField).toBeDefined(); - // biome-ignore lint/suspicious/noExplicitAny: accessing fields in different Arrow versions - const childNames = metadataField?.type.children.map((c: any) => c.name); - expect(childNames).toEqual([ - "filePath", - "startLine", - "endLine", - "text", - ]); - - const buf = await fromTableToBuffer(table); - const retrievedTable = tableFromIPC(buf); - - const rows = []; - for (let i = 0; i < retrievedTable.numRows; i++) { - rows.push(retrievedTable.get(i)); - } - - const firstRow = rows[0]; - expect(firstRow.id).toBe("doc1"); - expect(firstRow.vector.toJSON()).toEqual([1, 2, 3]); - expect(firstRow.metadata.filePath).toBe("/path/to/file1.ts"); - expect(firstRow.metadata.startLine).toBe(10); - expect(firstRow.metadata.endLine).toBe(20); - expect(firstRow.metadata.text).toBe("function test() { return true; }"); - }); - }); - - class DummyEmbedding extends EmbeddingFunction { - toJSON(): Partial { - return {}; - } - - async computeSourceEmbeddings(data: string[]): Promise { - return data.map(() => [0.0, 0.0]); - } - - ndims(): number { - return 2; - } - - embeddingDataType() { - return new Float16(); - } - } - - class DummyEmbeddingWithNoDimension extends EmbeddingFunction { - toJSON(): Partial { - return {}; - } - - embeddingDataType() { - return new Float16(); - } - - async computeSourceEmbeddings(data: string[]): Promise { - return data.map(() => [0.0, 0.0]); - } - } - const dummyEmbeddingConfig: EmbeddingFunctionConfig = { - sourceColumn: "string", - function: new DummyEmbedding(), - }; - - const dummyEmbeddingConfigWithNoDimension: EmbeddingFunctionConfig = { - sourceColumn: "string", - function: new DummyEmbeddingWithNoDimension(), - }; - - describe("convertToTable", function () { - it("will infer data types correctly", async function () { - await checkTableCreation( - // biome-ignore lint/suspicious/noExplicitAny: - async (records) => await (convertToTable)(records), - true, - ); - }); - - it("will allow a schema to be provided", async function () { - await checkTableCreation( - async (records, _, schema) => - // biome-ignore lint/suspicious/noExplicitAny: - await (convertToTable)(records, undefined, { schema }), - false, - ); - }); - - it("will use the field order of any provided schema", async function () { - await checkTableCreation( - async (_, recordsReversed, schema) => - // biome-ignore lint/suspicious/noExplicitAny: - await (convertToTable)(recordsReversed, undefined, { schema }), - false, - ); - }); - - it("will make an empty table", async function () { - await checkTableCreation( - async (_, __, schema) => - // biome-ignore lint/suspicious/noExplicitAny: - await (convertToTable)([], undefined, { schema }), - false, - ); - }); - - it("will apply embeddings", async function () { - const records = sampleRecords(); - const table = await convertToTable(records, dummyEmbeddingConfig); - expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( - true, - ); - expect( - table.getChild("vector")?.type.children[0].type.toString(), - ).toEqual(new Float16().toString()); - }); - - it("will fail if missing the embedding source column", async function () { - await expect( - convertToTable([{ id: 1 }], dummyEmbeddingConfig), - ).rejects.toThrow("'string' was not present"); - }); - - it("use embeddingDimension if embedding missing from table", async function () { - const schema = new Schema([new Field("string", new Utf8(), false)]); - // Simulate getting an empty Arrow table (minus embedding) from some other source - // In other words, we aren't starting with records - const table = makeEmptyTable(schema); - - // If the embedding specifies the dimension we are fine - await fromTableToBuffer(table, dummyEmbeddingConfig); - - // We can also supply a schema and should be ok - const schemaWithEmbedding = new Schema([ - new Field("string", new Utf8(), false), - new Field( - "vector", - new FixedSizeList(2, new Field("item", new Float16(), false)), - false, - ), - ]); - await fromTableToBuffer( - table, - dummyEmbeddingConfigWithNoDimension, - schemaWithEmbedding, - ); - - // Otherwise we will get an error - await expect( - fromTableToBuffer(table, dummyEmbeddingConfigWithNoDimension), - ).rejects.toThrow("does not specify `embeddingDimension`"); - }); - - it("will apply embeddings to an empty table", async function () { - const schema = new Schema([ - new Field("string", new Utf8(), false), - new Field( - "vector", - new FixedSizeList(2, new Field("item", new Float16(), false)), - false, - ), - ]); - const table = await convertToTable([], dummyEmbeddingConfig, { - schema, - }); - expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( - true, - ); - expect( - table.getChild("vector")?.type.children[0].type.toString(), - ).toEqual(new Float16().toString()); - }); - - it("will complain if embeddings present but schema missing embedding column", async function () { - const schema = new Schema([new Field("string", new Utf8(), false)]); - await expect( - convertToTable([], dummyEmbeddingConfig, { schema }), - ).rejects.toThrow("column vector was missing"); - }); - - it("will skip embedding application if already applied", async function () { - const records = sampleRecords(); - const table = await convertToTable(records, dummyEmbeddingConfig); - - // fromTableToBuffer will try and apply the embeddings again - // but should skip since the column already has non-null values - const result = await fromTableToBuffer(table, dummyEmbeddingConfig); - expect(result.byteLength).toBeGreaterThan(0); - }); - }); - - describe("makeEmptyTable", function () { - it("will make an empty table", async function () { - await checkTableCreation( - // biome-ignore lint/suspicious/noExplicitAny: - async (_, __, schema) => (makeEmptyTable)(schema), - false, - ); - }); - - it("will make an empty table with a Map field", async function () { - const schema = new Schema([ - new Field( - "attributes", - new Map_( - new Field( - "entries", - new Struct([ - new Field("key", new Utf8(), false), - new Field("value", new Int32(), true), - ]), - false, - ), - true, - ), - ), - ]); - - const table = makeEmptyTable(schema); - - expectValidMapField(table.schema.fields[0]); - - const buffer = await fromTableToBuffer(table); - const roundTripped = tableFromIPC(buffer); - - expectValidMapField(roundTripped.schema.fields[0]); - }); - - it("preserves string schema metadata", function () { - const metadata = new Map([["source", "fixture"]]); - const schema = new Schema( - [new Field("value", new Int32(), true)], - metadata, - ); - - expect(makeEmptyTable(schema).schema.metadata.get("source")).toBe( - "fixture", - ); - }); - - it.each([ - ["non-string keys", new Map([[42, "fixture"]])], - ["non-string values", new Map([["source", 42]])], - [ - "non-string keys and values", - new Map([[42, false]]), ], - ])("rejects schema metadata with %s", function (_, metadataLike) { - const metadata = metadataLike as unknown as Map; - const schema = new Schema( - [new Field("value", new Int32(), true)], - metadata, - ); + new Map([["embedding_functions", JSON.stringify([])]]), + ); - expect(() => makeEmptyTable(schema)).toThrow( - "Expected metadata, if present, to be a Map but it had non-string keys or values", - ); - }); + const data = [ + { id: "doc1", name: "Document 1" }, + { id: "doc2", name: "Document 2" }, + ]; + + const table = await convertToTable(data, undefined, { schema }); + + expect(table.numCols).toBe(3); + expect(table.numRows).toBe(2); + + const buf = await fromTableToBuffer(table); + const retrievedTable = tableFromIPC(buf); + + const rows = []; + for (let i = 0; i < retrievedTable.numRows; i++) { + rows.push(retrievedTable.get(i)); + } + + expect(rows[0].metadata.version).toBe(null); + expect(rows[0].metadata.author).toBe(null); + expect(rows[0].metadata.tags).toBe(null); + expect(rows[0].id).toBe("doc1"); + expect(rows[0].name).toBe("Document 1"); }); - describe("when using two versions of arrow", function () { - it("preserves a dictionary shared by multiple fields", async function () { - const values = ["alpha", "beta", "alpha"]; - const dictionaryVector = vectorFromArray(values); - const batch = new ArrowRecordBatch({ - first: dictionaryVector.data[0], - second: dictionaryVector.data[0], - }); - const table = new ArrowTable([batch]); - - const sanitized = sanitizeTable(table); - expect([...sanitized.getChild("first")!]).toEqual(values); - expect([...sanitized.getChild("second")!]).toEqual(values); - const firstType = sanitized.schema.fields[0].type as { - dictionary: unknown; - }; - const secondType = sanitized.schema.fields[1].type as { - dictionary: unknown; - }; - expect(secondType.dictionary).toBe(firstType.dictionary); - expect(sanitized.batches[0].data.children[1].dictionary).toBe( - sanitized.batches[0].data.children[0].dictionary, - ); - - const buf = await fromDataToBuffer(table); - const actual = currentTableFromIPC(buf); - expect([...actual.getChild("first")!]).toEqual(values); - expect([...actual.getChild("second")!]).toEqual(values); - }); - - it("preserves shared dictionary data from another Arrow version", async function () { - const values = ["alpha", "beta", "alpha"]; - const dictionaryVector = vectorFromArray(values); - const firstBatch = new ArrowRecordBatch({ - label: dictionaryVector.slice(0, 2).data[0], - }); - const secondBatch = new ArrowRecordBatch({ - label: dictionaryVector.slice(2).data[0], - }); - const table = new ArrowTable([firstBatch, secondBatch]); - - const sanitized = sanitizeTable(table); - expect([...sanitized.getChild("label")!]).toEqual(values); - - const dictionaries = sanitized.batches.map( - (batch) => batch.data.children[0].dictionary, - ); - expect(dictionaries[0]).toBeInstanceOf(CurrentVector); - expect(dictionaries[1]).toBe(dictionaries[0]); - - const buf = await fromDataToBuffer(table); - const actual = currentTableFromIPC(buf); - expect([...actual.getChild("label")!]).toEqual(values); - }); - - it("preserves shared chunks in growing dictionaries", async function () { - const type = new Dictionary(new Utf8(), new Int32(), 42, false); - const firstDictionary = vectorFromArray(["alpha", "beta"], new Utf8()); - const secondDictionary = firstDictionary.concat( - vectorFromArray(["gamma"], new Utf8()), - ); - const firstData = arrowMakeData({ - type, - data: Int32Array.from([0, 1]), - dictionary: firstDictionary, - }); - const secondData = arrowMakeData({ - type, - data: Int32Array.from([2]), - dictionary: secondDictionary, - }); - const table = new ArrowTable([ - new ArrowRecordBatch({ label: firstData }), - new ArrowRecordBatch({ label: secondData }), - ]); - - const sanitized = sanitizeTable(table); - const expected = ["alpha", "beta", "gamma"]; - expect([...sanitized.getChild("label")!]).toEqual(expected); - const firstLocalDictionary = - sanitized.batches[0].data.children[0].dictionary!; - const secondLocalDictionary = - sanitized.batches[1].data.children[0].dictionary!; - expect(secondLocalDictionary.data[0]).toBe( - firstLocalDictionary.data[0], - ); - - const buf = await fromTableToBuffer(sanitized); - const actual = currentTableFromIPC(buf); - expect([...actual.getChild("label")!]).toEqual(expected); - }); - - it("can serialize list data from another Arrow version", async function () { - const values = [["anime", "action"], [], null]; - const vector = vectorFromArray( - values, - new List(new Field("item", new Utf8(), true)), - ); - const table = new ArrowTable({ tags: vector }); - - const buf = await fromDataToBuffer(table); - const actual = currentTableFromIPC(buf); - const actualTags = actual.getChild("tags"); - - expect(actualTags?.get(0)?.toJSON()).toEqual(values[0]); - expect(actualTags?.get(1)?.toJSON()).toEqual(values[1]); - expect(actualTags?.get(2)).toBeNull(); - }); - - it("can still import data", async function () { - const schema = new arrow15.Schema([ - new arrow15.Field("id", new arrow15.Int32()), - new arrow15.Field( - "vector", - new arrow15.FixedSizeList( - 1024, - new arrow15.Field("item", new arrow15.Float32(), true), - ), + it("will handle partially missing nested struct fields", async function () { + const schema = new Schema( + [ + new Field("id", new Utf8(), true), + new Field( + "metadata", + new Struct([ + new Field("version", new Int32(), true), + new Field("author", new Utf8(), true), + new Field("created_at", new Utf8(), true), + ]), + true, ), - new arrow15.Field( - "struct", - new arrow15.Struct([ - new arrow15.Field( - "nested", - new arrow15.Dictionary( - new arrow15.Utf8(), - new arrow15.Int32(), - 1, - true, - ), - ), - new arrow15.Field( - "ts_with_tz", - new arrow15.TimestampNanosecond("some_tz"), - ), - new arrow15.Field( - "ts_no_tz", - new arrow15.TimestampNanosecond(null), + ], + new Map([["embedding_functions", JSON.stringify([])]]), + ); + + const data = [ + { id: "doc1", metadata: { version: 1, author: "Alice" } }, + { id: "doc2", metadata: { version: 2 } }, + ]; + + const table = await convertToTable(data, undefined, { schema }); + + expect(table.numCols).toBe(2); + expect(table.numRows).toBe(2); + + const metadataColumn = table.getChild("metadata"); + expect(metadataColumn).toBeDefined(); + expect(metadataColumn?.type.toString()).toBe( + "Struct<{version:Int32, author:Utf8, created_at:Utf8}>", + ); + }); + + it("will handle multiple levels of nested structures", async function () { + const schema = new Schema( + [ + new Field("id", new Utf8(), true), + new Field( + "config", + new Struct([ + new Field("database", new Utf8(), true), + new Field( + "connection", + new Struct([ + new Field("host", new Utf8(), true), + new Field("port", new Int32(), true), + new Field( + "ssl", + new Struct([ + new Field("enabled", new Bool(), true), + new Field("cert_path", new Utf8(), true), + ]), + true, + ), + ]), + true, ), ]), + true, ), - // biome-ignore lint/suspicious/noExplicitAny: skip - ]) as any; - schema.metadataVersion = arrow15.MetadataVersion.V5; - const table = makeArrowTable([], { schema }); + ], + new Map([["embedding_functions", JSON.stringify([])]]), + ); - const buf = await fromTableToBuffer(table); - expect(buf.byteLength).toBeGreaterThan(0); - const actual = tableFromIPC(buf); - const actualSchema = actual.schema; - expect(actualSchema.fields.length).toBe(3); + const data = [ + { + id: "config1", + config: { + database: "postgres", + connection: { host: "localhost" }, + }, + }, + { + id: "config2", + config: { database: "mysql" }, + }, + { + id: "config3", + }, + ]; - // Deep equality gets hung up on some very minor unimportant differences - // between arrow version 13 and 15 which isn't really what we're testing for - // and so we do our own comparison that just checks name/type/nullability - function compareFields(lhs: arrow15.Field, rhs: arrow15.Field) { - expect(lhs.name).toEqual(rhs.name); - expect(lhs.nullable).toEqual(rhs.nullable); - expect(lhs.typeId).toEqual(rhs.typeId); - if ("children" in lhs.type && lhs.type.children !== null) { - const lhsChildren = lhs.type.children as arrow15.Field[]; - lhsChildren.forEach((child: arrow15.Field, idx) => { - compareFields(child, rhs.type.children[idx]); - }); - } - } + const table = await convertToTable(data, undefined, { schema }); + + expect(table.numCols).toBe(2); + expect(table.numRows).toBe(3); + + const configColumn = table.getChild("config"); + expect(configColumn).toBeDefined(); + expect(configColumn?.type.toString()).toBe( + "Struct<{database:Utf8, connection:Struct<{host:Utf8, port:Int32, ssl:Struct<{enabled:Bool, cert_path:Utf8}>}>}>", + ); + }); + + it("will handle missing columns in Arrow table input when using embeddings", async function () { + const incompleteTable = makeArrowTable([ + { domain: "google.com", name: "Google" }, + { domain: "facebook.com", name: "Facebook" }, + ]); + + const schema = new Schema( + [ + new Field("domain", new Utf8(), true), + new Field("name", new Utf8(), true), + new Field("description", new Utf8(), true), + ], + new Map([["embedding_functions", JSON.stringify([])]]), + ); + + const buf = await fromDataToBuffer(incompleteTable, undefined, schema); + + expect(buf.byteLength).toBeGreaterThan(0); + + const retrievedTable = tableFromIPC(buf); + expect(retrievedTable.numCols).toBe(3); + expect(retrievedTable.numRows).toBe(2); + + const descriptionColumn = retrievedTable.getChild("description"); + expect(descriptionColumn).toBeDefined(); + expect(descriptionColumn?.nullCount).toBe(2); + expect(descriptionColumn?.toArray()).toEqual([null, null]); + + expect(retrievedTable.getChild("domain")?.toArray()).toEqual([ + "google.com", + "facebook.com", + ]); + expect(retrievedTable.getChild("name")?.toArray()).toEqual([ + "Google", + "Facebook", + ]); + }); + + it("should correctly retain values in nested struct fields", async function () { + const testData = [ + { + id: "doc1", + vector: [1, 2, 3], + metadata: { + filePath: "/path/to/file1.ts", + startLine: 10, + endLine: 20, + text: "function test() { return true; }", + }, + }, + { + id: "doc2", + vector: [4, 5, 6], + metadata: { + filePath: "/path/to/file2.ts", + startLine: 30, + endLine: 40, + text: "function test2() { return false; }", + }, + }, + ]; + + const table = makeArrowTable(testData); + + const metadataField = table.schema.fields.find( + (f) => f.name === "metadata", + ); + expect(metadataField).toBeDefined(); + // biome-ignore lint/suspicious/noExplicitAny: accessing fields in different Arrow versions + const childNames = metadataField?.type.children.map((c: any) => c.name); + expect(childNames).toEqual(["filePath", "startLine", "endLine", "text"]); + + const buf = await fromTableToBuffer(table); + const retrievedTable = tableFromIPC(buf); + + const rows = []; + for (let i = 0; i < retrievedTable.numRows; i++) { + rows.push(retrievedTable.get(i)); + } + + const firstRow = rows[0]; + expect(firstRow.id).toBe("doc1"); + expect(firstRow.vector.toJSON()).toEqual([1, 2, 3]); + expect(firstRow.metadata.filePath).toBe("/path/to/file1.ts"); + expect(firstRow.metadata.startLine).toBe(10); + expect(firstRow.metadata.endLine).toBe(20); + expect(firstRow.metadata.text).toBe("function test() { return true; }"); + }); + }); + + class DummyEmbedding extends EmbeddingFunction { + toJSON(): Partial { + return {}; + } + + async computeSourceEmbeddings(data: string[]): Promise { + return data.map(() => [0.0, 0.0]); + } + + ndims(): number { + return 2; + } + + embeddingDataType() { + return new Float16(); + } + } + + class DummyEmbeddingWithNoDimension extends EmbeddingFunction { + toJSON(): Partial { + return {}; + } + + embeddingDataType() { + return new Float16(); + } + + async computeSourceEmbeddings(data: string[]): Promise { + return data.map(() => [0.0, 0.0]); + } + } + const dummyEmbeddingConfig: EmbeddingFunctionConfig = { + sourceColumn: "string", + function: new DummyEmbedding(), + }; + + const dummyEmbeddingConfigWithNoDimension: EmbeddingFunctionConfig = { + sourceColumn: "string", + function: new DummyEmbeddingWithNoDimension(), + }; + + describe("convertToTable", function () { + it("will infer data types correctly", async function () { + await checkTableCreation( // biome-ignore lint/suspicious/noExplicitAny: - actualSchema.fields.forEach((field: any, idx: string | number) => { - compareFields(field, actualSchema.fields[idx]); - }); - }); + async (records) => await (convertToTable)(records), + true, + ); }); - describe("converting record batches to buffers", function () { - it("can convert to buffered record batch and back again", async function () { - const records = [ - { text: "dog", vector: [0.1, 0.2] }, - { text: "cat", vector: [0.3, 0.4] }, - ]; - const table = await convertToTable(records); - const batch = table.batches[0]; - - const buffer = await fromRecordBatchToBuffer(batch); - const result = await fromBufferToRecordBatch(buffer); - - expect(JSON.stringify(batch.toArray())).toEqual( - JSON.stringify(result?.toArray()), - ); - }); - - it("converting from buffer returns null if buffer has no record batches", async function () { - const result = await fromBufferToRecordBatch(Buffer.from([0x01, 0x02])); // bad data - expect(result).toEqual(null); - }); + it("will allow a schema to be provided", async function () { + await checkTableCreation( + async (records, _, schema) => + // biome-ignore lint/suspicious/noExplicitAny: + await (convertToTable)(records, undefined, { schema }), + false, + ); }); - describe("boolean null handling", function () { - it("should handle null values in nullable boolean fields", () => { - const { makeArrowTable } = require("../lancedb/arrow"); - const schema = new Schema([new Field("test", new arrow.Bool(), true)]); - - // Test with all null values - const data = [{ test: null }]; - const table = makeArrowTable(data, { schema }); - - expect(table.numRows).toBe(1); - expect(table.schema.names).toEqual(["test"]); - expect(table.getChild("test")!.get(0)).toBeNull(); - }); - - it("should handle mixed null and non-null boolean values", () => { - const { makeArrowTable } = require("../lancedb/arrow"); - const schema = new Schema([new Field("test", new Bool(), true)]); - - // Test with mixed values - const data = [{ test: true }, { test: null }, { test: false }]; - const table = makeArrowTable(data, { schema }); - - expect(table.numRows).toBe(3); - expect(table.getChild("test")!.get(0)).toBe(true); - expect(table.getChild("test")!.get(1)).toBeNull(); - expect(table.getChild("test")!.get(2)).toBe(false); - }); + it("will use the field order of any provided schema", async function () { + await checkTableCreation( + async (_, recordsReversed, schema) => + // biome-ignore lint/suspicious/noExplicitAny: + await (convertToTable)(recordsReversed, undefined, { schema }), + false, + ); }); - // Test for the undefined values bug fix - describe("undefined values handling", () => { - it("should handle mixed undefined and actual values", () => { - const schema = new Schema([ - new Field("text", new Utf8(), true), // nullable - new Field("number", new Int32(), true), // nullable - new Field("bool", new Bool(), true), // nullable - ]); + it("will make an empty table", async function () { + await checkTableCreation( + async (_, __, schema) => + // biome-ignore lint/suspicious/noExplicitAny: + await (convertToTable)([], undefined, { schema }), + false, + ); + }); - const data = [ - { text: undefined, number: 42, bool: true }, - { text: "hello", number: undefined, bool: false }, - { text: "world", number: 123, bool: undefined }, - ]; - const table = makeArrowTable(data, { schema }); + it("will apply embeddings", async function () { + const records = sampleRecords(); + const table = await convertToTable(records, dummyEmbeddingConfig); + expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( + true, + ); + expect( + table.getChild("vector")?.type.children[0].type.toString(), + ).toEqual(new Float16().toString()); + }); - const result = table.toArray(); - expect(result).toHaveLength(3); - expect(result[0].text).toBe(null); - expect(result[0].number).toBe(42); - expect(result[0].bool).toBe(true); - expect(result[1].text).toBe("hello"); - expect(result[1].number).toBe(null); - expect(result[1].bool).toBe(false); - expect(result[2].text).toBe("world"); - expect(result[2].number).toBe(123); - expect(result[2].bool).toBe(null); + it("will fail if missing the embedding source column", async function () { + await expect( + convertToTable([{ id: 1 }], dummyEmbeddingConfig), + ).rejects.toThrow("'string' was not present"); + }); + + it("use embeddingDimension if embedding missing from table", async function () { + const schema = new Schema([new Field("string", new Utf8(), false)]); + // Simulate getting an empty Arrow table (minus embedding) from some other source + // In other words, we aren't starting with records + const table = makeEmptyTable(schema); + + // If the embedding specifies the dimension we are fine + await fromTableToBuffer(table, dummyEmbeddingConfig); + + // We can also supply a schema and should be ok + const schemaWithEmbedding = new Schema([ + new Field("string", new Utf8(), false), + new Field( + "vector", + new FixedSizeList(2, new Field("item", new Float16(), false)), + false, + ), + ]); + await fromTableToBuffer( + table, + dummyEmbeddingConfigWithNoDimension, + schemaWithEmbedding, + ); + + // Otherwise we will get an error + await expect( + fromTableToBuffer(table, dummyEmbeddingConfigWithNoDimension), + ).rejects.toThrow("does not specify `embeddingDimension`"); + }); + + it("will apply embeddings to an empty table", async function () { + const schema = new Schema([ + new Field("string", new Utf8(), false), + new Field( + "vector", + new FixedSizeList(2, new Field("item", new Float16(), false)), + false, + ), + ]); + const table = await convertToTable([], dummyEmbeddingConfig, { + schema, + }); + expect(DataType.isFixedSizeList(table.getChild("vector")?.type)).toBe( + true, + ); + expect( + table.getChild("vector")?.type.children[0].type.toString(), + ).toEqual(new Float16().toString()); + }); + + it("will complain if embeddings present but schema missing embedding column", async function () { + const schema = new Schema([new Field("string", new Utf8(), false)]); + await expect( + convertToTable([], dummyEmbeddingConfig, { schema }), + ).rejects.toThrow("column vector was missing"); + }); + + it("will skip embedding application if already applied", async function () { + const records = sampleRecords(); + const table = await convertToTable(records, dummyEmbeddingConfig); + + // fromTableToBuffer will try and apply the embeddings again + // but should skip since the column already has non-null values + const result = await fromTableToBuffer(table, dummyEmbeddingConfig); + expect(result.byteLength).toBeGreaterThan(0); + }); + }); + + describe("makeEmptyTable", function () { + it("will make an empty table", async function () { + await checkTableCreation( + // biome-ignore lint/suspicious/noExplicitAny: + async (_, __, schema) => (makeEmptyTable)(schema), + false, + ); + }); + + it("will make an empty table with a Map field", async function () { + const schema = new Schema([ + new Field( + "attributes", + new Map_( + new Field( + "entries", + new Struct([ + new Field("key", new Utf8(), false), + new Field("value", new Int32(), true), + ]), + false, + ), + true, + ), + ), + ]); + + const table = makeEmptyTable(schema); + + expectValidMapField(table.schema.fields[0]); + + const buffer = await fromTableToBuffer(table); + const roundTripped = tableFromIPC(buffer); + + expectValidMapField(roundTripped.schema.fields[0]); + }); + + it("preserves string schema metadata", function () { + const metadata = new Map([["source", "fixture"]]); + const schema = new Schema( + [new Field("value", new Int32(), true)], + metadata, + ); + + expect(makeEmptyTable(schema).schema.metadata.get("source")).toBe( + "fixture", + ); + }); + + it.each([ + ["non-string keys", new Map([[42, "fixture"]])], + ["non-string values", new Map([["source", 42]])], + ["non-string keys and values", new Map([[42, false]])], + ])("rejects schema metadata with %s", function (_, metadataLike) { + const metadata = metadataLike as unknown as Map; + const schema = new Schema( + [new Field("value", new Int32(), true)], + metadata, + ); + + expect(() => makeEmptyTable(schema)).toThrow( + "Expected metadata, if present, to be a Map but it had non-string keys or values", + ); + }); + }); + + describe("when using two versions of arrow", function () { + it("preserves a dictionary shared by multiple fields", async function () { + const values = ["alpha", "beta", "alpha"]; + const dictionaryVector = vectorFromArray(values); + const batch = new ArrowRecordBatch({ + first: dictionaryVector.data[0], + second: dictionaryVector.data[0], + }); + const table = new ArrowTable([batch]); + + const sanitized = sanitizeTable(table); + expect([...sanitized.getChild("first")!]).toEqual(values); + expect([...sanitized.getChild("second")!]).toEqual(values); + const firstType = sanitized.schema.fields[0].type as { + dictionary: unknown; + }; + const secondType = sanitized.schema.fields[1].type as { + dictionary: unknown; + }; + expect(secondType.dictionary).toBe(firstType.dictionary); + expect(sanitized.batches[0].data.children[1].dictionary).toBe( + sanitized.batches[0].data.children[0].dictionary, + ); + + const buf = await fromDataToBuffer(table); + const actual = currentTableFromIPC(buf); + expect([...actual.getChild("first")!]).toEqual(values); + expect([...actual.getChild("second")!]).toEqual(values); + }); + + it("preserves shared dictionary data from another Arrow version", async function () { + const values = ["alpha", "beta", "alpha"]; + const dictionaryVector = vectorFromArray(values); + const firstBatch = new ArrowRecordBatch({ + label: dictionaryVector.slice(0, 2).data[0], + }); + const secondBatch = new ArrowRecordBatch({ + label: dictionaryVector.slice(2).data[0], + }); + const table = new ArrowTable([firstBatch, secondBatch]); + + const sanitized = sanitizeTable(table); + expect([...sanitized.getChild("label")!]).toEqual(values); + + const dictionaries = sanitized.batches.map( + (batch) => batch.data.children[0].dictionary, + ); + expect(dictionaries[0]).toBeInstanceOf(CurrentVector); + expect(dictionaries[1]).toBe(dictionaries[0]); + + const buf = await fromDataToBuffer(table); + const actual = currentTableFromIPC(buf); + expect([...actual.getChild("label")!]).toEqual(values); + }); + + it("preserves shared chunks in growing dictionaries", async function () { + const type = new Dictionary(new Utf8(), new Int32(), 42, false); + const firstDictionary = vectorFromArray(["alpha", "beta"], new Utf8()); + const secondDictionary = firstDictionary.concat( + vectorFromArray(["gamma"], new Utf8()), + ); + const firstData = arrowMakeData({ + type, + data: Int32Array.from([0, 1]), + dictionary: firstDictionary, + }); + const secondData = arrowMakeData({ + type, + data: Int32Array.from([2]), + dictionary: secondDictionary, + }); + const table = new ArrowTable([ + new ArrowRecordBatch({ label: firstData }), + new ArrowRecordBatch({ label: secondData }), + ]); + + const sanitized = sanitizeTable(table); + const expected = ["alpha", "beta", "gamma"]; + expect([...sanitized.getChild("label")!]).toEqual(expected); + const firstLocalDictionary = + sanitized.batches[0].data.children[0].dictionary!; + const secondLocalDictionary = + sanitized.batches[1].data.children[0].dictionary!; + expect(secondLocalDictionary.data[0]).toBe(firstLocalDictionary.data[0]); + + const buf = await fromTableToBuffer(sanitized); + const actual = currentTableFromIPC(buf); + expect([...actual.getChild("label")!]).toEqual(expected); + }); + + it("can serialize list data from another Arrow version", async function () { + const values = [["anime", "action"], [], null]; + const vector = vectorFromArray( + values, + new List(new Field("item", new Utf8(), true)), + ); + const table = new ArrowTable({ tags: vector }); + + const buf = await fromDataToBuffer(table); + const actual = currentTableFromIPC(buf); + const actualTags = actual.getChild("tags"); + + expect(actualTags?.get(0)?.toJSON()).toEqual(values[0]); + expect(actualTags?.get(1)?.toJSON()).toEqual(values[1]); + expect(actualTags?.get(2)).toBeNull(); + }); + + it("can still import data", async function () { + const schema = new arrow15.Schema([ + new arrow15.Field("id", new arrow15.Int32()), + new arrow15.Field( + "vector", + new arrow15.FixedSizeList( + 1024, + new arrow15.Field("item", new arrow15.Float32(), true), + ), + ), + new arrow15.Field( + "struct", + new arrow15.Struct([ + new arrow15.Field( + "nested", + new arrow15.Dictionary( + new arrow15.Utf8(), + new arrow15.Int32(), + 1, + true, + ), + ), + new arrow15.Field( + "ts_with_tz", + new arrow15.TimestampNanosecond("some_tz"), + ), + new arrow15.Field( + "ts_no_tz", + new arrow15.TimestampNanosecond(null), + ), + ]), + ), + // biome-ignore lint/suspicious/noExplicitAny: skip + ]) as any; + schema.metadataVersion = arrow15.MetadataVersion.V5; + const table = makeArrowTable([], { schema }); + + const buf = await fromTableToBuffer(table); + expect(buf.byteLength).toBeGreaterThan(0); + const actual = tableFromIPC(buf); + const actualSchema = actual.schema; + expect(actualSchema.fields.length).toBe(3); + + // Deep equality gets hung up on some very minor unimportant differences + // between arrow version 13 and 15 which isn't really what we're testing for + // and so we do our own comparison that just checks name/type/nullability + function compareFields(lhs: arrow15.Field, rhs: arrow15.Field) { + expect(lhs.name).toEqual(rhs.name); + expect(lhs.nullable).toEqual(rhs.nullable); + expect(lhs.typeId).toEqual(rhs.typeId); + if ("children" in lhs.type && lhs.type.children !== null) { + const lhsChildren = lhs.type.children as arrow15.Field[]; + lhsChildren.forEach((child: arrow15.Field, idx) => { + compareFields(child, rhs.type.children[idx]); + }); + } + } + // biome-ignore lint/suspicious/noExplicitAny: + actualSchema.fields.forEach((field: any, idx: string | number) => { + compareFields(field, actualSchema.fields[idx]); }); }); - }, -); + }); + + describe("converting record batches to buffers", function () { + it("can convert to buffered record batch and back again", async function () { + const records = [ + { text: "dog", vector: [0.1, 0.2] }, + { text: "cat", vector: [0.3, 0.4] }, + ]; + const table = await convertToTable(records); + const batch = table.batches[0]; + + const buffer = await fromRecordBatchToBuffer(batch); + const result = await fromBufferToRecordBatch(buffer); + + expect(JSON.stringify(batch.toArray())).toEqual( + JSON.stringify(result?.toArray()), + ); + }); + + it("converting from buffer returns null if buffer has no record batches", async function () { + const result = await fromBufferToRecordBatch(Buffer.from([0x01, 0x02])); // bad data + expect(result).toEqual(null); + }); + }); + + describe("boolean null handling", function () { + it("should handle null values in nullable boolean fields", () => { + const { makeArrowTable } = require("../lancedb/arrow"); + const schema = new Schema([new Field("test", new arrow.Bool(), true)]); + + // Test with all null values + const data = [{ test: null }]; + const table = makeArrowTable(data, { schema }); + + expect(table.numRows).toBe(1); + expect(table.schema.names).toEqual(["test"]); + expect(table.getChild("test")!.get(0)).toBeNull(); + }); + + it("should handle mixed null and non-null boolean values", () => { + const { makeArrowTable } = require("../lancedb/arrow"); + const schema = new Schema([new Field("test", new Bool(), true)]); + + // Test with mixed values + const data = [{ test: true }, { test: null }, { test: false }]; + const table = makeArrowTable(data, { schema }); + + expect(table.numRows).toBe(3); + expect(table.getChild("test")!.get(0)).toBe(true); + expect(table.getChild("test")!.get(1)).toBeNull(); + expect(table.getChild("test")!.get(2)).toBe(false); + }); + }); + + // Test for the undefined values bug fix + describe("undefined values handling", () => { + it("should handle mixed undefined and actual values", () => { + const schema = new Schema([ + new Field("text", new Utf8(), true), // nullable + new Field("number", new Int32(), true), // nullable + new Field("bool", new Bool(), true), // nullable + ]); + + const data = [ + { text: undefined, number: 42, bool: true }, + { text: "hello", number: undefined, bool: false }, + { text: "world", number: 123, bool: undefined }, + ]; + const table = makeArrowTable(data, { schema }); + + const result = table.toArray(); + expect(result).toHaveLength(3); + expect(result[0].text).toBe(null); + expect(result[0].number).toBe(42); + expect(result[0].bool).toBe(true); + expect(result[1].text).toBe("hello"); + expect(result[1].number).toBe(null); + expect(result[1].bool).toBe(false); + expect(result[2].text).toBe("world"); + expect(result[2].number).toBe(123); + expect(result[2].bool).toBe(null); + }); + }); +}); diff --git a/nodejs/__test__/arrow_versions.ts b/nodejs/__test__/arrow_versions.ts new file mode 100644 index 000000000..bfbbfa5ef --- /dev/null +++ b/nodejs/__test__/arrow_versions.ts @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +// Every Apache Arrow release the package accepts as a peer dependency, as +// separately installed copies so the tests can feed objects built by one +// Arrow version into a LanceDB running on another. When a new Arrow major +// is released, add the alias to `package.json` and extend this list; the +// test files pick the change up from here. +import * as arrow15 from "apache-arrow-15"; +import * as arrow16 from "apache-arrow-16"; +import * as arrow17 from "apache-arrow-17"; +import * as arrow18 from "apache-arrow-18"; +import * as arrow19 from "apache-arrow-19"; +import * as arrow20 from "apache-arrow-20"; +import * as arrow21 from "apache-arrow-21"; + +export { arrow15, arrow16, arrow17, arrow18, arrow19, arrow20, arrow21 }; + +export type ApacheArrow = + | typeof arrow15 + | typeof arrow16 + | typeof arrow17 + | typeof arrow18 + | typeof arrow19 + | typeof arrow20 + | typeof arrow21; + +export const arrowVersions: ApacheArrow[] = [ + arrow15, + arrow16, + arrow17, + arrow18, + arrow19, + arrow20, + arrow21, +]; + +/** The newest supported Arrow release, for tests that only need one. */ +export const latestArrow = arrow21; diff --git a/nodejs/__test__/blob.test.ts b/nodejs/__test__/blob.test.ts index 5c01897b2..b41d5716f 100644 --- a/nodejs/__test__/blob.test.ts +++ b/nodejs/__test__/blob.test.ts @@ -1,9 +1,25 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -import { Field, Int64, List, Schema, Struct, Utf8 } from "apache-arrow"; -import { makeArrowTable } from "../lancedb/arrow"; +import { + Field, + Float, + Float32, + Int32, + Int64, + List, + Schema, + Struct, + Utf8, +} from "apache-arrow"; +import { + convertToTable, + makeArrowTable, + resolveBlobInputs, +} from "../lancedb/arrow"; import { BlobFile, blob, coerceBlobValue, isBlobField } from "../lancedb/blob"; +import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding"; +import { getRegistry } from "../lancedb/embedding/registry"; describe("blob()", () => { it("marks the field as lance.blob.v2", () => { @@ -46,6 +62,13 @@ describe("blob()", () => { }); }); +class BlobStruct { + constructor( + public data?: unknown, + public uri?: unknown, + ) {} +} + describe("coerceBlobValue", () => { it.each([ ["Buffer", Buffer.from("x"), { data: Buffer.from("x"), uri: null }], @@ -65,6 +88,41 @@ describe("coerceBlobValue", () => { { uri: "s3://bucket/key" }, { data: null, uri: "s3://bucket/key" }, ], + [ + "ArrayBuffer", + new Uint8Array([120]).buffer, + { data: new Uint8Array([120]), uri: null }, + ], + [ + "URL", + new URL("https://example.com/blob"), + { data: null, uri: "https://example.com/blob" }, + ], + [ + "data struct with ArrayBuffer", + { data: new Uint8Array([121]).buffer }, + { data: new Uint8Array([121]), uri: null }, + ], + [ + "uri struct with URL", + { uri: new URL("file:///tmp/a.png") }, + { data: null, uri: "file:///tmp/a.png" }, + ], + [ + "class instance with Uint8Array data", + new BlobStruct(new Uint8Array([122])), + { data: new Uint8Array([122]), uri: null }, + ], + [ + "class instance with ArrayBuffer data", + new BlobStruct(new Uint8Array([122]).buffer), + { data: new Uint8Array([122]), uri: null }, + ], + [ + "class instance with URL uri", + new BlobStruct(undefined, new URL("s3://bucket/key")), + { data: null, uri: "s3://bucket/key" }, + ], ["null", null, null], ])("accepts %s", (_name, input, expected) => { expect(coerceBlobValue(input)).toEqual(expected); @@ -76,7 +134,7 @@ describe("coerceBlobValue", () => { [ "Int16Array", new Int16Array([1]), - /Blob data must be Buffer or Uint8Array/, + /Blob data must be Buffer, Uint8Array, ArrayBuffer, or Blob/, ], [ "both data and uri", @@ -88,11 +146,160 @@ describe("coerceBlobValue", () => { { data: null, uri: null }, /exactly one of 'data' or 'uri'/, ], + ["Blob", new Blob(["x"]), /must be read asynchronously/], + ["data struct with Blob", { data: new Blob(["x"]) }, /asynchronously/], + [ + "class instance with Blob data", + new BlobStruct(new Blob(["x"])), + /asynchronously/, + ], + [ + "unknown struct field", + { data: Buffer.from("y"), mime: "image/png" }, + /Blob struct values only support 'data' and 'uri', got 'mime'/, + ], + [ + "misspelled struct field", + { uri: "s3://bucket/key", dat: Buffer.from("y") }, + /got 'dat'/, + ], + [ + "SharedArrayBuffer", + new SharedArrayBuffer(2), + /Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob/, + ], + ["Date", new Date(0), /Blob column values must be/], ])("rejects %s", (_name, input, message) => { expect(() => coerceBlobValue(input)).toThrow(message); }); }); +describe("resolveBlobInputs", () => { + it("reads Blob and File values in blob columns only", async () => { + const schema = new Schema([ + new Field("name", new Utf8(), true), + blob("image"), + new Field("images", new List(blob("image")), true), + ]); + const shared = new Blob(["same"]); + const notABlobColumn = new Blob(["other"]); + const data = [ + { + name: notABlobColumn, + image: new File(["file"], "a.txt"), + images: [shared, { data: shared }, null], + }, + { name: "x", image: Buffer.from("buf"), images: null }, + ]; + const resolved = await resolveBlobInputs(data, schema); + + expect(resolved[0].name).toBe(notABlobColumn); + expect(Buffer.from(resolved[0].image as Uint8Array).toString()).toBe( + "file", + ); + const images = resolved[0].images as [Uint8Array, { data: Uint8Array }]; + expect(Buffer.from(images[0]).toString()).toBe("same"); + expect(Buffer.from(images[1].data).toString()).toBe("same"); + // Rows without Blob values and the caller's records are left alone. + expect(resolved[1]).toBe(data[1]); + expect(data[0].images?.[0]).toBe(shared); + }); + + it("reads each distinct Blob once, bare or inside data structs", async () => { + class CountingBlob extends Blob { + calls = 0; + async arrayBuffer(): Promise { + this.calls++; + return super.arrayBuffer(); + } + } + const shared = new CountingBlob(["same"]); + const schema = new Schema([blob("image")]); + const resolved = await resolveBlobInputs( + [ + { image: { data: shared } }, + { image: { data: shared } }, + { image: shared }, + ], + schema, + ); + expect(shared.calls).toBe(1); + const first = (resolved[0].image as { data: Uint8Array }).data; + expect(Buffer.from(first).toString()).toBe("same"); + expect((resolved[1].image as { data: Uint8Array }).data).toBe(first); + expect(resolved[2].image).toBe(first); + }); + + it("reads a Blob held by a class instance", async () => { + const schema = new Schema([blob("image")]); + const [row] = await resolveBlobInputs( + [{ image: new BlobStruct(new Blob(["hi"])) }], + schema, + ); + expect( + Buffer.from((row.image as { data: Uint8Array }).data).toString(), + ).toBe("hi"); + }); + + it("returns the input when there is nothing to read", async () => { + const data = [{ image: Buffer.from("x") }]; + const schema = new Schema([blob("image")]); + await expect(resolveBlobInputs(data, schema)).resolves.toBe(data); + await expect(resolveBlobInputs(data)).resolves.toBe(data); + }); +}); + +describe("convertToTable blob columns with embedding functions", () => { + // Embedding metadata routes rows through ensureNestedFieldsExist, which must + // leave blob values for the blob coercion instead of treating them as structs. + class LengthEmbedding extends EmbeddingFunction { + ndims() { + return 2; + } + embeddingDataType(): Float { + return new Float32(); + } + async computeSourceEmbeddings(data: string[]) { + return data.map((text) => [text.length, 1]); + } + } + getRegistry().register("blob-test-length")(LengthEmbedding); + const func = getRegistry().get("blob-test-length")!.create(); + const embeddingSchema = LanceSchema({ + text: func.sourceField(new Utf8()), + vector: func.vectorField(), + }); + const schema = new Schema( + [...embeddingSchema.fields, blob("image")], + embeddingSchema.metadata, + ); + + it.each([ + ["Uint8Array", new Uint8Array([104, 105]), "hi", null], + ["ArrayBuffer", new Uint8Array([104, 105]).buffer, "hi", null], + ["Blob", new Blob(["hi"]), "hi", null], + ["URI string", "s3://b/k", null, "s3://b/k"], + ["URL", new URL("s3://b/k"), null, "s3://b/k"], + ["data struct", { data: new Uint8Array([104, 105]) }, "hi", null], + ])("accepts %s input", async (_name, image, data, uri) => { + const table = await convertToTable([{ text: "a", image }], undefined, { + schema, + }); + const column = table.getChild("image")!; + expect(column.nullCount).toBe(0); + const bytes = column.getChild("data")!.get(0); + expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data); + expect(column.getChild("uri")!.get(0)).toBe(uri); + }); + + it("writes a missing blob as null", async () => { + const table = await convertToTable([{ text: "a" }], undefined, { schema }); + const column = table.getChild("image")!; + expect(column.nullCount).toBe(1); + expect(column.get(0)).toBeNull(); + }); +}); + describe("BlobFile", () => { it("rejects constructing BlobFile without a native handle", () => { expect(() => new (BlobFile as unknown as { new (): BlobFile })()).toThrow( @@ -103,26 +310,122 @@ describe("BlobFile", () => { describe("makeArrowTable blob columns", () => { it.each([ - ["ArrayBuffer", new Uint8Array([104]).buffer], - ["Blob", new Blob(["hello"])], - ["File", new File(["hello"], "hello.txt")], - ["URL", new URL("https://example.com/blob")], ["ReadableStream", new ReadableStream()], - ["data struct with ArrayBuffer", { data: new Uint8Array([104]).buffer }], - ["uri struct with URL", { uri: new URL("https://example.com/blob") }], + ["Int16Array", new Int16Array([1])], ])("rejects %s input rather than writing a null blob", (_name, image) => { const schema = new Schema([blob("image")]); expect(() => makeArrowTable([{ image }], { schema })).toThrow(/image/); }); - it("rejects Blob instances with enumerable fields as one invalid value", () => { + it.each([ + ["Blob", new Blob(["hello"])], + ["File", new File(["hello"], "hello.txt")], + ["data struct with Blob", { data: new Blob(["hello"]) }], + [ + "Blob with enumerable fields", + Object.assign(new Blob(["hello"]), { position: 0 }), + ], + ])("points %s input at the async APIs", (_name, image) => { const schema = new Schema([blob("image")]); - const image = Object.assign(new Blob(["hello"]), { position: 0 }); - expect(() => makeArrowTable([{ image }], { schema })).toThrow( - "Unsupported object value for field image at row 0.", + expect(() => makeArrowTable([{ id: 1, image }], { schema })).toThrow( + /blob field image at row 0: Blob and File values must be read asynchronously/, ); }); + it.each([ + ["ArrayBuffer", new Uint8Array([104, 105]).buffer, "hi", null], + [ + "data struct with ArrayBuffer", + { data: new Uint8Array([104, 105]).buffer }, + "hi", + null, + ], + ["URL", new URL("https://example.com/a"), null, "https://example.com/a"], + [ + "uri struct with URL", + { uri: new URL("file:///tmp/a.png") }, + null, + "file:///tmp/a.png", + ], + ])("accepts %s input", (_name, image, data, uri) => { + const schema = new Schema([blob("image")]); + const input = { image }; + const table = makeArrowTable([input], { schema }); + const column = table.getChild("image")!; + expect(column.nullCount).toBe(0); + const bytes = column.getChild("data")!.get(0); + expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data); + expect(column.getChild("uri")!.get(0)).toBe(uri); + // The caller's record is left untouched. + expect(input.image).toBe(image); + }); + + it.each([ + ["Uint8Array", new Uint8Array([104, 105]), "hi", null], + ["URI string", "s3://b/k", null, "s3://b/k"], + ["data struct", { data: new Uint8Array([104, 105]) }, "hi", null], + ])("accepts %s input after a null row", (_name, image, data, uri) => { + const schema = new Schema([blob("image")]); + const table = makeArrowTable([{ image: null }, { image }], { schema }); + const column = table.getChild("image")!; + expect(column.nullCount).toBe(1); + expect(column.get(0)).toBeNull(); + const bytes = column.getChild("data")!.get(1); + expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data); + expect(column.getChild("uri")!.get(1)).toBe(uri); + }); + + it("accepts a null struct before one holding a blob", () => { + const schema = new Schema([ + new Field("info", new Struct([blob("image")]), true), + ]); + const table = makeArrowTable( + [{ info: null }, { info: { image: new Uint8Array([97]) } }], + { schema }, + ); + const info = table.getChild("info")!; + expect(info.get(0)).toBeNull(); + const image = (info.get(1) as { image: { data: Uint8Array } }).image; + expect(Buffer.from(image.data).toString()).toBe("a"); + }); + + it("skips a missing blob column named like an Object.prototype member", () => { + const schema = new Schema([ + new Field("id", new Int32(), true), + blob("constructor"), + ]); + const table = makeArrowTable([{ id: 1 }], { schema }); + expect(table.schema.fields.map((f) => f.name)).toEqual(["id"]); + }); + + it("accepts ArrayBuffer and URL inside lists and structs", () => { + const schema = new Schema([ + new Field("images", new List(blob("image")), true), + new Field("info", new Struct([blob("image")]), true), + ]); + const table = makeArrowTable( + [ + { + images: [new Uint8Array([97]).buffer, new URL("s3://b/k")], + info: { image: new Uint8Array([98]).buffer }, + }, + ], + { schema }, + ); + const images = Array.from( + table.getChild("images")!.get(0) as Iterable<{ + data: Uint8Array | null; + uri: string | null; + }>, + ); + expect(Buffer.from(images[0].data!).toString()).toBe("a"); + expect(images[1].uri).toBe("s3://b/k"); + const info = table.getChild("info")!.get(0) as { + image: { data: Uint8Array }; + }; + expect(Buffer.from(info.image.data).toString()).toBe("b"); + }); + it("coerces Buffer input onto a blob field", () => { const schema = new Schema([ new Field("id", new Int64(), true), diff --git a/nodejs/__test__/catalog.test.ts b/nodejs/__test__/catalog.test.ts index f9158c94f..945956b5b 100644 --- a/nodejs/__test__/catalog.test.ts +++ b/nodejs/__test__/catalog.test.ts @@ -80,9 +80,12 @@ describe("remote catalog", () => { const second = await catalog.connectDatabase("other"); expect(await second.tableNames()).toEqual([]); expect(await first.tableNames()).toEqual([]); - expect( - await catalog.listDatabases({ limit: 1, pageToken: "a/b" }), - ).toEqual({ databases: ["team/search"], pageToken: "next" }); + const names = catalog.listDatabases({ pageLimit: 1, pageToken: "a/b" }); + expect(await names.next()).toEqual({ + done: false, + value: "team/search", + }); + expect(names.pageToken()).toBe("next"); await catalog.dropDatabase("team/search", { ignoreMissing: true }); expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create"); expect(requests[0].body).toEqual({ mode: "ExistOk" }); @@ -129,11 +132,139 @@ describe("remote catalog", () => { ); }); + it("lazily traverses pages, including empty pages, and exposes cached state", async () => { + await withCatalog( + [ + // biome-ignore lint/style/useNamingConvention: server wire format + [200, { namespaces: ["a", "b"], page_token: "empty" }], + // biome-ignore lint/style/useNamingConvention: server wire format + [200, { namespaces: [], page_token: "last" }], + // biome-ignore lint/style/useNamingConvention: server wire format + [200, { namespaces: ["c", "d"], page_token: "" }], + ], + async (catalog, requests) => { + const names = catalog.listDatabases({ pageLimit: 2 }); + expect(names[Symbol.asyncIterator]()).toBe(names); + expect(requests).toHaveLength(0); + expect(names.numPageResults()).toBe(0); + expect(names.pageToken()).toBeUndefined(); + expect(await names.next()).toEqual({ done: false, value: "a" }); + expect(names.numPageResults()).toBe(1); + expect(names.pageToken()).toBe("empty"); + expect(await names.next()).toEqual({ done: false, value: "b" }); + expect(requests).toHaveLength(1); + expect(await names.next()).toEqual({ done: false, value: "c" }); + expect(requests).toHaveLength(3); + expect(names.pageToken()).toBeUndefined(); + expect(names.numPageResults()).toBe(1); + const remaining = []; + for await (const name of names) remaining.push(name); + expect(remaining).toEqual(["d"]); + expect(names.numPageResults()).toBe(0); + expect(await names.next()).toEqual({ done: true, value: undefined }); + expect(requests.map((request) => request.url)).toEqual([ + "/v1/namespace/$/list?limit=2", + "/v1/namespace/$/list?limit=2&page_token=empty", + "/v1/namespace/$/list?limit=2&page_token=last", + ]); + }, + ); + }); + + it("resumes from a saved token after draining the cache", async () => { + await withCatalog( + [ + // biome-ignore lint/style/useNamingConvention: server wire format + [200, { namespaces: ["a", "b"], page_token: "next/token" }], + [200, { namespaces: ["c"] }], + ], + async (catalog, requests) => { + const names = catalog.listDatabases({ + pageToken: "start/token", + pageLimit: 2, + }); + expect(names.pageToken()).toBe("start/token"); + expect(await names.next()).toEqual({ done: false, value: "a" }); + while (names.numPageResults() > 0) await names.next(); + const resumed = catalog.listDatabases({ + pageToken: names.pageToken(), + pageLimit: 2, + }); + const results = []; + for await (const name of resumed) results.push(name); + expect(results).toEqual(["c"]); + expect(requests.map((request) => request.url)).toEqual([ + "/v1/namespace/$/list?limit=2&page_token=start%2Ftoken", + "/v1/namespace/$/list?limit=2&page_token=next%2Ftoken", + ]); + }, + ); + }); + + it("terminates on errors and retains the failed request token", async () => { + await withCatalog( + [ + // biome-ignore lint/style/useNamingConvention: server wire format + [200, { namespaces: ["a"], page_token: "retry" }], + [400, {}], + [200, { namespaces: ["b"] }], + ], + async (catalog, requests) => { + const names = catalog.listDatabases(); + await names.next(); + await expect(names.next()).rejects.toThrow(); + expect(names.pageToken()).toBe("retry"); + expect(names.numPageResults()).toBe(0); + expect(await names.next()).toEqual({ done: true, value: undefined }); + expect(requests).toHaveLength(2); + const resumed = catalog.listDatabases({ pageToken: names.pageToken() }); + expect(await resumed.next()).toEqual({ done: false, value: "b" }); + expect(await resumed.next()).toEqual({ done: true, value: undefined }); + }, + ); + }); + + it("handles an empty listing and concurrent advances", async () => { + await withCatalog( + [ + [200, { namespaces: [] }], + [200, { namespaces: ["a", "b"] }], + ], + async (catalog, requests) => { + const empty = catalog.listDatabases(); + expect(await empty.next()).toEqual({ done: true, value: undefined }); + expect(await empty.next()).toEqual({ done: true, value: undefined }); + const names = catalog.listDatabases(); + const pending = [names.next(), names.next(), names.next()]; + expect(names.numPageResults()).toBe(0); + expect(names.pageToken()).toBeUndefined(); + const results = await Promise.all(pending); + expect( + results + .filter((result) => !result.done) + .map((result) => result.value) + .sort(), + ).toEqual(["a", "b"]); + expect(results.filter((result) => result.done)).toHaveLength(1); + expect(requests).toHaveLength(2); + }, + ); + }); + it("validates endpoints and pagination", async () => { await expect(connectCatalog("/tmp/catalog")).rejects.toThrow(); const catalog = await connectCatalog("http://127.0.0.1:1"); - for (const limit of [0, -1, 1.5, 2147483648]) { - await expect(catalog.listDatabases({ limit })).rejects.toThrow("limit"); + for (const limit of [ + 0, + -1, + 1.5, + 2147483648, + Number.NaN, + Number.POSITIVE_INFINITY, + ]) { + expect(() => catalog.listDatabases({ pageLimit: limit })).toThrow( + "limit", + ); } await expect(catalog.connectDatabase("a$b")).rejects.toThrow( "Invalid database name", diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index 6739320aa..f0c97e13a 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -16,15 +16,11 @@ describe("materialized views", () => { beforeEach(async () => { tmpDir = tmp.dirSync({ unsafeCleanup: true }); db = await connect(tmpDir.name); - await db.createTable( - "people", - [ - { name: "ada", age: 36 }, - { name: "kid", age: 7 }, - { name: "grace", age: 85 }, - ], - { storageOptions: { newTableEnableStableRowIds: "true" } }, - ); + await db.createTable("people", [ + { name: "ada", age: 36 }, + { name: "kid", age: 7 }, + { name: "grace", age: 85 }, + ]); }); afterEach(() => tmpDir.removeCallback()); @@ -52,19 +48,9 @@ describe("materialized views", () => { "SELECT * FROM people", ); - for (const [format, fn] of [ - [3, "vector_duplicate_pairs"], - [4, "vector_dedup"], - ]) { - const native = `SELECT * FROM ${fn}('images', 3, 'phash', 4)`; - expect(read(JSON.stringify({ format, query: native })).query).toBe( - native, - ); - } - // A newer writer's layout is reported, never guessed at. for (const newer of [ - `{"format":5,"query":${JSON.stringify(query)}}`, + `{"format":3,"query":${JSON.stringify(query)}}`, '{"kind":"select_v3","source_table":"people"}', ]) { expect(() => read(newer)).toThrow(/cannot refresh/); @@ -106,7 +92,7 @@ describe("materialized views", () => { ); }); - it("refreshes incrementally after an append", async () => { + it("fully refreshes after an append", async () => { const view = await db.createMaterializedView("copy", "people", { withNoData: true, }); @@ -115,11 +101,11 @@ describe("materialized views", () => { const people = await db.openTable("people"); await people.add([{ name: "alan", age: 41 }]); const result = await view.refresh(); - expect(result.mode).toBe("incremental"); - expect(Number(result.rowsWritten)).toBe(1); + expect(result.mode).toBe("rebuild"); + expect(Number(result.rowsWritten)).toBe(4); expect(await view.table().countRows()).toBe(4); - expect((await view.refresh()).mode).toBe("no_op"); + expect((await view.refresh()).mode).toBe("rebuild"); }); it("lists views and rejects non-views", async () => { @@ -172,9 +158,7 @@ describe("materialized views", () => { }); it("quotes bare select names", async () => { - await db.createTable("odd_names", [{ "order item": "widget" }], { - storageOptions: { newTableEnableStableRowIds: "true" }, - }); + await db.createTable("odd_names", [{ "order item": "widget" }]); const view = await db.createMaterializedView("quoted", "odd_names", { select: ["order item"], withNoData: true, @@ -183,10 +167,9 @@ describe("materialized views", () => { expect(Number(result.rowsWritten)).toBe(1); }); - it("requires stable row ids on the source", async () => { + it("does not require stable row ids on the source", async () => { await db.createTable("plain", [{ x: 1 }]); - await expect(db.createMaterializedView("v", "plain")).rejects.toThrow( - "stable row ids", - ); + const view = await db.createMaterializedView("v", "plain"); + expect(await view.table().countRows()).toBe(1); }); }); diff --git a/nodejs/__test__/query.test.ts b/nodejs/__test__/query.test.ts index 5f3e68b16..2cfecf823 100644 --- a/nodejs/__test__/query.test.ts +++ b/nodejs/__test__/query.test.ts @@ -110,6 +110,129 @@ describe("Query outputSchema", () => { }); }); +describe("Query empty results", () => { + let tmpDir: tmp.DirResult; + let table: Table; + + beforeEach(async () => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + const schema = new Schema( + [ + new Field("id", new Int64(), false), + new Field( + "text", + new Utf8(), + true, + new Map([["description", "searchable text"]]), + ), + new Field( + "vector", + new FixedSizeList(2, new Field("item", new Float32())), + true, + ), + ], + new Map([["source", "empty-query-regression"]]), + ); + table = await db.createTable( + "test", + makeArrowTable([{ id: 1n, text: "hello", vector: [1, 2] }], { schema }), + ); + }); + + afterEach(() => { + tmpDir.removeCallback(); + }); + + it.each([ + { + name: "filtered scan", + query: (table: Table) => table.query().where("id < 0"), + fields: ["id", "text", "vector"], + }, + { + name: "zero limit", + query: (table: Table) => table.query().limit(0), + fields: ["id", "text", "vector"], + }, + { + name: "column projection", + query: (table: Table) => + table.query().where("id < 0").select(["text", "id"]), + fields: ["text", "id"], + }, + { + name: "dynamic projection", + query: (table: Table) => + table.query().where("id < 0").select({ doubled: "id * 2" }), + fields: ["doubled"], + }, + { + name: "filtered vector search", + query: (table: Table) => + table.search([0.5, 0.5]).where("id < 0").limit(3), + fields: ["id", "text", "vector", "_distance"], + }, + { + name: "vector projection with row id", + query: (table: Table) => + table + .vectorSearch([0.5, 0.5]) + .where("id < 0") + .select(["id"]) + .withRowId(), + fields: ["id", "_distance", "_rowid"], + }, + { + name: "fast search without an index", + query: (table: Table) => table.search([0.5, 0.5]).fastSearch(), + fields: ["id", "text", "vector", "_rowid", "_distance"], + }, + { + name: "take with no offsets", + query: (table: Table) => table.takeOffsets([]).select(["text"]), + fields: ["text"], + }, + ])("preserves the output schema for $name", async ({ query, fields }) => { + const result = await query(table).toArrow({ maxBatchLength: 1 }); + + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual(fields); + for (const name of fields) { + expect(result.getChild(name)).not.toBeNull(); + expect(result.getChild(name)?.length).toBe(0); + } + }); + + it("preserves types, nullability, and metadata", async () => { + const result = await table.query().where("id < 0").toArrow(); + const schema = await table.schema(); + + expect(result.schema).toEqual(schema); + expect(result.getChild("id")?.type.toString()).toBe("Int64"); + expect(result.schema.fields[0].nullable).toBe(false); + expect(result.schema.fields[1].metadata.get("description")).toBe( + "searchable text", + ); + expect(result.schema.metadata.get("source")).toBe("empty-query-regression"); + }); + + it("preserves the score column for full-text search with no matches", async () => { + await table.createIndex("text", { config: Index.fts() }); + const result = await table + .search("missing", "fts") + .select(["id"]) + .toArrow(); + + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual([ + "id", + "_score", + ]); + expect(result.getChild("_score")?.type.toString()).toBe("Float32"); + }); +}); + describe("Search pagination", () => { let tmpDir: tmp.DirResult; let table: Table; diff --git a/nodejs/__test__/registry.test.ts b/nodejs/__test__/registry.test.ts index 973ad8a25..5c92941d6 100644 --- a/nodejs/__test__/registry.test.ts +++ b/nodejs/__test__/registry.test.ts @@ -3,13 +3,10 @@ import * as apiArrow from "apache-arrow"; -import * as arrow15 from "apache-arrow-15"; -import * as arrow16 from "apache-arrow-16"; -import * as arrow17 from "apache-arrow-17"; -import * as arrow18 from "apache-arrow-18"; - import * as tmp from "tmp"; +import { arrowVersions } from "./arrow_versions"; + import { connect } from "../lancedb"; import { EmbeddingFunction, @@ -18,7 +15,7 @@ import { } from "../lancedb/embedding"; import { getRegistry, register } from "../lancedb/embedding/registry"; -describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => { +describe.each(arrowVersions)("LanceSchema", (arrow) => { test("should preserve input order", async () => { const schema = LanceSchema({ id: new arrow.Int32(), @@ -29,7 +26,7 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => { }); }); -describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => { +describe.each(arrowVersions)("Registry", (arrow) => { let tmpDir: tmp.DirResult; beforeEach(() => { tmpDir = tmp.dirSync({ unsafeCleanup: true }); @@ -240,7 +237,7 @@ describe("Registry.setVar", () => { return data.map(() => [1, 2, 3]); } embeddingDataType() { - return new arrow18.Float32() as apiArrow.Float; + return new apiArrow.Float32(); } protected getSensitiveKeys() { return ["secretKey"]; diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 898ba533e..ec259b79b 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -14,6 +14,14 @@ import { TlsConfig, connect, } from "../lancedb"; +import { + Table as ArrowTable, + Field, + Int64, + RecordBatch, + Schema, + tableToIPC, +} from "../lancedb/arrow"; import { HeaderProvider, OAuthHeaderProvider, @@ -85,6 +93,98 @@ async function withMockDatabase( } describe("remote connection", () => { + it("rejects nprobes(0) before sending a query", async () => { + const requests: string[] = []; + await withMockDatabase( + (req, res) => { + requests.push(req.url ?? ""); + if (req.url === "/v1/table/test/describe/") { + res + .writeHead(200, { "Content-Type": "application/json" }) + .end(JSON.stringify({ version: 1, schema: { fields: [] } })); + } else { + res.writeHead(404).end(); + } + }, + async (db) => { + const table = await db.openTable("test"); + expect(() => table.vectorSearch([0, 0]).nprobes(0)).toThrow( + "Invalid input, nprobes must be greater than 0", + ); + expect(() => + table.query().nearestTo([0, 0]).fullTextSearch("dog").nprobes(0), + ).toThrow("Invalid input, nprobes must be greater than 0"); + await expect( + table + .vectorSearch(Promise.resolve([0, 0])) + .nprobes(0) + .toArrow(), + ).rejects.toThrow("Invalid input, nprobes must be greater than 0"); + }, + ); + expect(requests).toEqual(["/v1/table/test/describe/"]); + }); + + it.each([false, true])( + "preserves an empty query's schema with an empty batch: %s", + async (withEmptyBatch) => { + const schema = new Schema( + [new Field("doubled", new Int64(), false)], + new Map([["source", "query-output"]]), + ); + const result = new ArrowTable( + schema, + withEmptyBatch ? [new RecordBatch(schema, undefined)] : [], + ); + const response = Buffer.from(tableToIPC(result, "stream")); + let queryRequests = 0; + + await withMockDatabase( + (req, res) => { + if (req.url?.endsWith("/describe/")) { + res.writeHead(200, { "Content-Type": "application/json" }).end( + JSON.stringify({ + name: "items", + version: 1, + schema: { + fields: [ + { name: "id", type: { type: "int64" }, nullable: false }, + ], + }, + }), + ); + } else if (req.url?.endsWith("/query/")) { + queryRequests++; + req.resume(); + req.on("end", () => { + res + .writeHead(200, { + "Content-Type": "application/vnd.apache.arrow.stream", + }) + .end(response); + }); + } else { + res.writeHead(404).end(); + } + }, + async (db) => { + const table = await db.openTable("items"); + const result = await table + .query() + .where("id < 0") + .select({ doubled: "id * 2" }) + .toArrow(); + + expect(result.numRows).toBe(0); + expect(result.schema).toEqual(schema); + expect(result.getChild("doubled")?.length).toBe(0); + }, + ); + + expect(queryRequests).toBe(1); + }, + ); + it("rejects the external blob opt-in without blocking regular adds", async () => { let insertRequests = 0; const describeRequests: string[] = []; diff --git a/nodejs/__test__/sanitize.test.ts b/nodejs/__test__/sanitize.test.ts index 587805d53..97856a55b 100644 --- a/nodejs/__test__/sanitize.test.ts +++ b/nodejs/__test__/sanitize.test.ts @@ -2,7 +2,13 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors import * as arrow from "../lancedb/arrow"; -import { sanitizeField, sanitizeMap, sanitizeType } from "../lancedb/sanitize"; +import { + sanitizeField, + sanitizeMap, + sanitizeSchema, + sanitizeType, +} from "../lancedb/sanitize"; +import { arrow21 } from "./arrow_versions"; describe("sanitize", function () { describe("sanitizeType function", function () { @@ -59,6 +65,8 @@ describe("sanitize", function () { ["durationmicrosecond", new arrow.DurationMicrosecond()], ["durationmillisecond", new arrow.DurationMillisecond()], ["durationsecond", new arrow.DurationSecond()], + ["largebinary", new arrow.LargeBinary()], + ["largeutf8", new arrow.LargeUtf8()], ] as const; it.each(allTypeNameTestCases)( @@ -89,6 +97,110 @@ describe("sanitize", function () { "Unrecognized type name in schema: invalid_type", ); }); + + // These types only exist in Arrow 20 (IntervalMonthDayNano) and 21 (the + // rest), so the objects come from a second copy of the newest release. The + // library itself may be running on an older Arrow (see the CI matrix), in + // which case sanitizing them has to fail with a clear message instead. + const hostArrow = arrow as unknown as Record; + const hostSupportsNewerTypes = typeof hostArrow.Utf8View === "function"; + const whenSupported = hostSupportsNewerTypes ? describe : describe.skip; + const whenUnsupported = hostSupportsNewerTypes ? describe.skip : describe; + + const newerTypeTestCases = [ + ["Utf8View", () => new arrow21.Utf8View()], + ["BinaryView", () => new arrow21.BinaryView()], + ["IntervalMonthDayNano", () => new arrow21.IntervalMonthDayNano()], + [ + "LargeList", + () => + new arrow21.LargeList(new arrow21.Field("item", new arrow21.Utf8())), + ], + ] as const; + + whenSupported("types added in Arrow 20 and 21", function () { + it.each(newerTypeTestCases)( + "should sanitize a %s type from another copy of Arrow", + function (name, makeForeign) { + const foreign = makeForeign(); + const result = sanitizeType(foreign); + // Interval subclasses share `Type.Interval`, so the sanitized type + // is the generic `Interval` carrying the same unit, as for the + // other units. + const expected = name === "IntervalMonthDayNano" ? "Interval" : name; + expect(result).toBeInstanceOf(hostArrow[expected] as () => unknown); + expect(result.typeId).toBe(foreign.typeId); + expect(result.toString()).toBe(foreign.toString()); + }, + ); + + it.each([ + ["binaryview", "BinaryView"], + ["utf8view", "Utf8View"], + ["intervalmonthdaynano", "Interval"], + ])('should map type name "%s" to %s', function (name, expected) { + expect(sanitizeType(name)).toBeInstanceOf( + hostArrow[expected] as () => unknown, + ); + }); + + it("should sanitize the child field of a LargeList", function () { + const foreign = new arrow21.LargeList( + new arrow21.Field("item", new arrow21.Int32(), false), + ); + const result = sanitizeType(foreign) as arrow.List; + expect(result.children[0]).toBeInstanceOf(arrow.Field); + expect(result.children[0].name).toBe("item"); + expect(result.children[0].nullable).toBe(false); + const child = result.children[0].type; + expect(arrow.DataType.isInt(child)).toBe(true); + expect((child as arrow.Int).bitWidth).toBe(32); + expect((child as arrow.Int).isSigned).toBe(true); + }); + }); + + whenUnsupported("types newer than the installed Arrow", function () { + it.each( + newerTypeTestCases.filter(([name]) => name !== "IntervalMonthDayNano"), + )( + "should explain that %s needs a newer apache-arrow", + function (name, makeForeign) { + expect(() => sanitizeType(makeForeign())).toThrow( + `The Arrow type ${name} requires apache-arrow`, + ); + }, + ); + + it("should sanitize IntervalMonthDayNano to a generic Interval", () => { + // The subclass shares `Type.Interval`, which every supported Arrow + // release has; only the unit is new. + const result = sanitizeType(new arrow21.IntervalMonthDayNano()); + expect(result).toBeInstanceOf(arrow.Interval); + expect((result as arrow.Interval).unit).toBe( + arrow21.IntervalUnit.MONTH_DAY_NANO, + ); + }); + }); + + it("should rebuild a schema from a second copy of the same Arrow version", function () { + // Arrow 21 makes `instanceof Schema` succeed across copies of the library + // but not `instanceof `; the schema must still be rebuilt. + const foreign = new arrow21.Schema([ + new arrow21.Field("x", new arrow21.Float32(), true), + ]); + // Only Arrow 21+ overrides `Symbol.hasInstance` this way. + const crossCopyInstanceof = Object.getOwnPropertySymbols( + arrow.Schema, + ).includes(Symbol.hasInstance); + expect(foreign instanceof arrow.Schema).toBe(crossCopyInstanceof); + const result = sanitizeSchema(foreign); + expect(result).not.toBe(foreign); + const type = result.fields[0].type; + expect(type).toBeInstanceOf(arrow.DataType); + expect(arrow.DataType.isFloat(type)).toBe(true); + expect((type as arrow.Float).precision).toBe(arrow.Precision.SINGLE); + expect(() => arrow.makeBuilder({ type })).not.toThrow(); + }); }); describe("sanitizeField function", function () { diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 904c282d4..24cf6e0fd 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -7,13 +7,9 @@ import { pathToFileURL } from "node:url"; import * as path from "path"; import * as tmp from "tmp"; -import * as arrow15 from "apache-arrow-15"; -import * as arrow16 from "apache-arrow-16"; -import * as arrow17 from "apache-arrow-17"; -import * as arrow18 from "apache-arrow-18"; - import { AutoQuery, + type BlobInput, Connection, MatchQuery, PhraseQuery, @@ -56,19 +52,21 @@ import { instanceOfFullTextQuery, } from "../lancedb/query"; import { LocalTable } from "../lancedb/table"; +import { + type ApacheArrow, + arrow18, + arrowVersions, + latestArrow, +} from "./arrow_versions"; -describe.each([arrow15, arrow16, arrow17, arrow18])( +describe.each(arrowVersions)( "Given a table", // biome-ignore lint/suspicious/noExplicitAny: (arrow: any) => { let tmpDir: tmp.DirResult; let table: Table; - const schema: - | import("apache-arrow-15").Schema - | import("apache-arrow-16").Schema - | import("apache-arrow-17").Schema - | import("apache-arrow-18").Schema = new arrow.Schema([ + const schema: InstanceType = new arrow.Schema([ new arrow.Field("id", new arrow.Float64(), true), ]); @@ -742,6 +740,114 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( ); // https://github.com/lancedb/lancedb/issues/1963 +// These types were added in Arrow 21, so they are built with a second copy of +// that release rather than the per-version matrix above, and reading them back +// needs the library's own Arrow to know them too. +const hostArrowSupportsViewTypes = + typeof (arrow as unknown as Record).Utf8View === "function"; +(hostArrowSupportsViewTypes ? it : it.skip)( + "creates and adds tables with LargeList, Utf8View and BinaryView columns", + async () => { + const tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + const makeData = () => + new latestArrow.Table({ + id: latestArrow.vectorFromArray([1, 2], new latestArrow.Int32()), + tags: latestArrow.vectorFromArray( + [["a", "b"], ["c"]], + new latestArrow.LargeList( + new latestArrow.Field("item", new latestArrow.Utf8(), true), + ), + ), + text: latestArrow.vectorFromArray( + ["short", "a string that is longer than the twelve inline bytes"], + new latestArrow.Utf8View(), + ), + bytes: latestArrow.vectorFromArray( + [new Uint8Array([1, 2, 3]), new Uint8Array(20).fill(7)], + new latestArrow.BinaryView(), + ), + }); + const table = await db.createTable("views", makeData()); + await table.add(makeData()); + + const schema = await table.schema(); + // Take the enum from the Arrow 21 copy: the host Arrow may predate + // `Type.LargeList`, and this file still has to type-check there. + expect(schema.fields.map((f) => f.type.typeId)).toEqual([ + latestArrow.Type.Int, + latestArrow.Type.LargeList, + // Lance stores the view types as their offset-based equivalents. + latestArrow.Type.Utf8, + latestArrow.Type.Binary, + ]); + const rows = await table.query().toArrow(); + expect(rows.numRows).toBe(4); + const row = rows.get(1)!.toJSON(); + expect(row.id).toBe(2); + expect(row.tags.toJSON()).toEqual(["c"]); + expect(row.text).toBe( + "a string that is longer than the twelve inline bytes", + ); + expect(Array.from(row.bytes)).toEqual(Array(20).fill(7)); + + // Arrow JS 21 serializes a *sliced* view column with a truncated views + // buffer; the library rebuilds view columns as Utf8/Binary so this works. + const sliced = makeData().slice(1, 2); + expect(sliced.getChild("text")!.data[0].offset).toBe(1); + const slicedTable = await db.createTable("views_sliced", sliced); + await slicedTable.add(sliced); + const slicedRows = await slicedTable.query().toArrow(); + expect(slicedRows.numRows).toBe(2); + for (const i of [0, 1]) { + const r = slicedRows.get(i)!.toJSON(); + expect(r.id).toBe(2); + expect(r.text).toBe( + "a string that is longer than the twelve inline bytes", + ); + expect(Array.from(r.bytes)).toEqual(Array(20).fill(7)); + } + + // Regression: a field with an integer-like name. JavaScript enumerates + // such keys first, so a key-ordered rebuild would pair "1" with the text + // column's values. The batches are built explicitly to pin field order. + const numericSchema = new latestArrow.Schema([ + new latestArrow.Field("text", new latestArrow.Utf8View(), true), + new latestArrow.Field("1", new latestArrow.Utf8View(), true), + new latestArrow.Field("n", new latestArrow.Int32(), true), + ]); + const column = ( + values: unknown[], + type: InstanceType, + ) => latestArrow.vectorFromArray(values, type).data[0]; + const numericBatch = new latestArrow.RecordBatch( + numericSchema, + latestArrow.makeData({ + type: new latestArrow.Struct(numericSchema.fields), + length: 3, + nullCount: 0, + children: [ + column(["t0", "t1", "t2"], new latestArrow.Utf8View()), + column(["one0", "one1", "one2"], new latestArrow.Utf8View()), + column([0, 1, 2], new latestArrow.Int32()), + ], + }), + ); + const numericData = new latestArrow.Table(numericSchema, [numericBatch]); + const numericTable = await db.createTable("views_numeric", numericData); + await numericTable.add(numericData.slice(1)); + const numericRows = (await numericTable.query().toArrow()).toArray(); + expect(numericRows.map((r) => r.toJSON())).toEqual([ + { text: "t0", "1": "one0", n: 0 }, + { text: "t1", "1": "one1", n: 1 }, + { text: "t2", "1": "one2", n: 2 }, + { text: "t1", "1": "one1", n: 1 }, + { text: "t2", "1": "one2", n: 2 }, + ]); + tmpDir.removeCallback(); + }, +); + it("should query documents with LangChain PDF metadata", async () => { const tmpDir = tmp.dirSync({ unsafeCleanup: true }); try { @@ -1106,6 +1212,12 @@ describe("When creating an index", () => { .toArrow(); expect(rst.numRows).toBe(2); + expect(() => tbl.vectorSearch(queryVec).nprobes(0)).toThrow( + "Invalid input, nprobes must be greater than 0", + ); + expect(() => + tbl.query().nearestTo(queryVec).fullTextSearch("dog").nprobes(0), + ).toThrow("Invalid input, nprobes must be greater than 0"); expect(() => tbl.query().nearestTo(queryVec).minimumNprobes(0)).toThrow( "Invalid input, minimum_nprobes must be greater than 0", ); @@ -2472,7 +2584,7 @@ describe("when dealing with blob columns", () => { new Field("id", new Int64(), true), blob("image"), ]); - const row = { id: 1n, image: new Uint8Array([104]).buffer }; + const row = { id: 1n, image: new ReadableStream() }; await expect(db.createTable("invalid", [row], { schema })).rejects.toThrow( /field image at row 0/, @@ -2502,6 +2614,8 @@ describe("when dealing with blob columns", () => { it.each([ ["URI string", (uri: string) => uri], ["URI struct", (uri: string) => ({ uri })], + ["URL", (uri: string) => new URL(uri)], + ["URL struct", (uri: string) => ({ uri: new URL(uri) })], ])( "round-trips an external blob from a %s after reopening", async (_label, blobValue) => { @@ -2595,6 +2709,63 @@ describe("when dealing with blob columns", () => { expect(bytes[0]).toEqual(payload); }); + it("accepts ArrayBuffer, Blob, and File in createTable and add", async () => { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + blob("image"), + ]); + const rows: { id: bigint; image: BlobInput }[] = [ + { id: 1n, image: new TextEncoder().encode("array-buffer").buffer }, + { id: 2n, image: new Blob(["blob"]) }, + { id: 3n, image: { data: new File(["file"], "f.txt") } }, + ]; + const table = await db.createTable("widened", rows.slice(0, 2), { + schema, + }); + await table.add(rows.slice(2)); + + const results = await table.query().select(["id"]).withRowId().toArray(); + results.sort((a, b) => Number(a.id - b.id)); + const bytes = await table.fetchBlobs( + "image", + results.map((row) => row._rowid as bigint), + ); + expect(bytes.map((b) => b?.toString())).toEqual([ + "array-buffer", + "blob", + "file", + ]); + }); + + it("accepts Blob values in mergeInsert", async () => { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + blob("image"), + ]); + const table = await db.createTable( + "merge_blobs", + [{ id: 1n, image: Buffer.from("old") }], + { schema }, + ); + await table + .mergeInsert("id") + .whenMatchedUpdateAll() + .whenNotMatchedInsertAll() + .execute([ + { id: 1n, image: new Blob(["new"]) }, + { id: 2n, image: new Uint8Array([104, 105]).buffer }, + ]); + const results = await table.query().select(["id"]).withRowId().toArray(); + results.sort((a, b) => Number(a.id - b.id)); + const bytes = await table.fetchBlobs( + "image", + results.map((row) => row._rowid as bigint), + ); + expect(bytes.map((b) => b?.toString())).toEqual(["new", "hi"]); + }); + it("discovers blob columns", async () => { const { table } = await openBlobTable(); expect(await table.blobColumns()).toEqual(["image"]); @@ -3203,7 +3374,7 @@ it("passes cleanupOlderThan to the native binding as an absolute timestamp", asy expect(optimize).toHaveBeenCalledWith(cutoff.getTime(), true); }); -describe.each([arrow15, arrow16, arrow17, arrow18])( +describe.each(arrowVersions)( "when optimizing a dataset", // biome-ignore lint/suspicious/noExplicitAny: (arrow: any) => { diff --git a/nodejs/lancedb/arrow.ts b/nodejs/lancedb/arrow.ts index df875a808..3130da8eb 100644 --- a/nodejs/lancedb/arrow.ts +++ b/nodejs/lancedb/arrow.ts @@ -40,7 +40,12 @@ import { } from "apache-arrow"; import { Buffers } from "apache-arrow/data"; import { typedArrayToArrowType } from "./arrow_type"; -import { coerceBlobValue, isBlobField } from "./blob"; +import { + blobToRead, + coerceBlobValue, + isBlobField, + withBlobBytes, +} from "./blob"; import { type EmbeddingFunction } from "./embedding/embedding_function"; import { EmbeddingFunctionConfig, @@ -66,11 +71,20 @@ export type SchemaLike = export type FieldLike = | Field | { - type: string; + type: string | DataTypeLike; name: string; nullable: boolean; metadata?: Map; }; +/** + * A `DataType` from any copy or version of apache-arrow. + * + * Arrow 21 brands its classes with `unique symbol` properties, so a type + * object from a second copy of the library no longer satisfies the `DataType` + * type of this one even though it is structurally identical. Inputs that only + * need to be sanitized accept this looser shape instead. + */ +export type DataTypeLike = DataType | { readonly typeId: number }; /** * Create an Arrow field backed by LanceDB's JSON extension type. @@ -440,6 +454,24 @@ export function makeArrowTable( ); } + if (schema !== undefined) { + // Validate blob values up front and give every one the full + // `{ data, uri }` shape, so inference sees the same struct whether a row + // passed bytes, a URI, or a partial struct. + data = mapBlobInputs(data, schema, (value, field, row) => { + if (value === undefined) { + return value; + } + try { + return coerceBlobValue(value); + } catch (e) { + throw new Error( + `Invalid value for blob field ${field} at row ${row}: ${(e as Error).message}`, + ); + } + }); + } + let schemaMetadata = schema?.metadata || new Map(); if (metadata !== undefined) { schemaMetadata = new Map([...schemaMetadata, ...metadata]); @@ -501,6 +533,133 @@ function containsBlobField(field: Field): boolean { ); } +type BlobInputVisitor = (value: unknown, field: string, row: number) => unknown; + +/** + * Calls `visit` on every value that lands in a blob column of `schema`, + * including blobs nested in structs and lists, and replaces it with the + * result. Records and arrays are copied only where a value changed, so the + * caller's data is never mutated. + */ +function mapBlobInputs( + data: Array>, + schema: Schema, + visit: BlobInputVisitor, +): Array> { + if (!schema.fields.some(containsBlobField)) { + return data; + } + return data.map((record, row) => + isObject(record) + ? mapBlobFields(record, schema.fields, "", row, visit) + : record, + ); +} + +function mapBlobFields( + record: Record, + fields: Field[], + prefix: string, + row: number, + visit: BlobInputVisitor, +): Record { + let out: Record | undefined; + for (const field of fields) { + if (!containsBlobField(field) || !Object.hasOwn(record, field.name)) { + continue; + } + const value = record[field.name]; + const mapped = mapBlobField( + value, + field, + `${prefix}${field.name}`, + row, + visit, + ); + if (mapped !== value) { + out ??= { ...record }; + out[field.name] = mapped; + } + } + return out ?? record; +} + +function mapBlobField( + value: unknown, + field: Field, + label: string, + row: number, + visit: BlobInputVisitor, +): unknown { + if (isBlobField(field)) { + return visit(value, label, row); + } + if (field.type instanceof Struct && isObject(value)) { + return mapBlobFields(value, field.type.children, `${label}.`, row, visit); + } + if (isList(field.type) && Array.isArray(value)) { + const child = field.type.children[0]; + let out: unknown[] | undefined; + for (const [index, element] of value.entries()) { + const mapped = mapBlobField( + element, + child, + `${label}[${index}]`, + row, + visit, + ); + if (mapped !== element) { + out ??= [...value]; + out[index] = mapped; + } + } + return out ?? value; + } + return value; +} + +/** + * Reads `Blob` / `File` values in the blob columns of `schema` into bytes so + * the synchronous conversion in {@link makeArrowTable} can accept them. + * Returns `data` itself when there is nothing to read. + */ +export async function resolveBlobInputs( + data: Array>, + schema?: SchemaLike, +): Promise>> { + if (schema === undefined || schema === null) { + return data; + } + const sanitized = sanitizeSchema(schema); + // Keyed by the Blob itself, so rows sharing one Blob (bare or as + // `{ data }`) read it once and share the bytes. + const reads = new Map>(); + mapBlobInputs(data, sanitized, (value) => { + const source = blobToRead(value); + if (source !== undefined && !reads.has(source)) { + reads.set( + source, + source.arrayBuffer().then((buffer) => new Uint8Array(buffer)), + ); + } + return value; + }); + if (reads.size === 0) { + return data; + } + const bytes = new Map( + await Promise.all( + [...reads].map(async ([source, read]) => [source, await read] as const), + ), + ); + return mapBlobInputs(data, sanitized, (value) => { + const source = blobToRead(value); + return source === undefined + ? value + : withBlobBytes(value, bytes.get(source) as Uint8Array); + }); +} + function isObject(value: unknown): value is Record { return ( typeof value === "object" && @@ -1004,6 +1163,11 @@ export async function convertToTable( ); } + processedData = await resolveBlobInputs( + processedData, + makeTableOptions?.schema, + ); + const table = makeArrowTable(processedData, makeTableOptions); return await applyEmbeddings(table, embeddings, makeTableOptions?.schema); } @@ -1059,6 +1223,71 @@ export async function fromRecordsToStreamBuffer( return Buffer.from(await writer.toUint8Array()); } +// `Type.Utf8View` / `Type.BinaryView` as numbers: the enum members only exist +// in Arrow 21+, and this module compiles against every supported release. +const UTF8_VIEW_TYPE_ID = 24; +const BINARY_VIEW_TYPE_ID = 23; + +/** + * Copy a Utf8View / BinaryView `Data` into a single `Data` of `type`. + */ +function materializeViewData(child: ArrowData, type: DataType): ArrowData { + const builder = makeBuilder({ type, nullValues: [null] }); + for (const value of new Vector([child])) { + builder.append(value); + } + return builder.finish().flush(); +} + +/** + * Rebuild any top-level Utf8View / BinaryView column as Utf8 / Binary. + * + * Lance stores the view types as their offset-based equivalents anyway, so + * nothing is lost. Doing it here also sidesteps an Arrow JS 21 bug: its IPC + * writer emits a truncated views buffer for a *sliced* view array, which the + * Rust reader rejects with "Need at least N bytes in buffers[0]". + * + * The record batches are rebuilt positionally rather than through a + * `Record`: JavaScript enumerates integer-like keys first, so + * a field named e.g. `"1"` would otherwise be paired with the wrong column. + * + * Tables without view columns are returned as-is. + */ +function materializeViewColumns(table: ArrowTable): ArrowTable { + const replacements = new Map(); + table.schema.fields.forEach((field, i) => { + if (field.type.typeId === UTF8_VIEW_TYPE_ID) { + replacements.set(i, new Utf8()); + } else if (field.type.typeId === BINARY_VIEW_TYPE_ID) { + replacements.set(i, new Binary()); + } + }); + if (replacements.size === 0) { + return table; + } + const fields = table.schema.fields.map((field, i) => { + const type = replacements.get(i); + return type === undefined + ? field + : new Field(field.name, type, field.nullable, field.metadata); + }); + const schema = new Schema(fields, table.schema.metadata); + const batches = table.batches.map((batch) => { + const children = batch.data.children.map((child, i) => { + const type = replacements.get(i); + return type === undefined ? child : materializeViewData(child, type); + }); + const data = makeData({ + type: new Struct(fields), + length: batch.numRows, + nullCount: 0, + children, + }); + return new RecordBatch(schema, data); + }); + return new ArrowTable(schema, batches); +} + /** * Serialize an Arrow Table into a buffer using the Arrow IPC File serialization * @@ -1075,7 +1304,9 @@ export async function fromTableToBuffer( if (schema !== undefined && schema !== null) { schema = sanitizeSchema(schema); } - const tableWithEmbeddings = await applyEmbeddings(table, embeddings, schema); + const tableWithEmbeddings = materializeViewColumns( + await applyEmbeddings(table, embeddings, schema), + ); validateBlobSchema(tableWithEmbeddings.schema); const writer = RecordBatchFileWriter.writeAll(tableWithEmbeddings); return Buffer.from(await writer.toUint8Array()); @@ -1146,7 +1377,8 @@ export async function fromRecordBatchToBuffer( export async function fromRecordBatchToStreamBuffer( batch: RecordBatch, ): Promise { - const writer = RecordBatchStreamWriter.writeAll([batch]); + const table = materializeViewColumns(new ArrowTable([batch])); + const writer = RecordBatchStreamWriter.writeAll(table); return Buffer.from(await writer.toUint8Array()); } @@ -1163,7 +1395,9 @@ export async function fromTableToStreamBuffer( embeddings?: EmbeddingFunctionConfig, schema?: SchemaLike, ): Promise { - const tableWithEmbeddings = await applyEmbeddings(table, embeddings, schema); + const tableWithEmbeddings = materializeViewColumns( + await applyEmbeddings(table, embeddings, schema), + ); const writer = RecordBatchStreamWriter.writeAll(tableWithEmbeddings); return Buffer.from(await writer.toUint8Array()); } @@ -1290,7 +1524,7 @@ export function ensureNestedFieldsExist( for (const field of schema.fields) { if (field.name in row) { if ( - field.type.constructor.name === "Struct" && + isPlainStructField(field) && row[field.name] !== null && row[field.name] !== undefined ) { @@ -1307,10 +1541,9 @@ export function ensureNestedFieldsExist( } else { // Keep a missing struct valid while filling each of its children with // null. This is distinct from an explicitly null struct value. - completeRow[field.name] = - field.type.constructor.name === "Struct" - ? ensureStructFieldsExist({}, field.type as Struct) - : null; + completeRow[field.name] = isPlainStructField(field) + ? ensureStructFieldsExist({}, field.type as Struct) + : null; } } @@ -1318,6 +1551,15 @@ export function ensureNestedFieldsExist( }); } +/** + * Blob fields are Arrow structs, but their values are bytes, URIs, or + * `{ data } | { uri }` inputs that the blob coercion in `makeArrowTable` + * handles, so they must not be filled in like ordinary structs. + */ +function isPlainStructField(field: Field): boolean { + return field.type.constructor.name === "Struct" && !isBlobField(field); +} + /** * Recursively ensures that all fields in a struct type exist in the data, * filling missing fields with null values. @@ -1331,7 +1573,7 @@ function ensureStructFieldsExist( for (const childField of structType.children) { if (childField.name in data) { if ( - childField.type.constructor.name === "Struct" && + isPlainStructField(childField) && data[childField.name] !== null && data[childField.name] !== undefined ) { @@ -1347,10 +1589,9 @@ function ensureStructFieldsExist( } else { // Keep a missing struct valid while filling each of its children with // null. This is distinct from an explicitly null struct value. - completeStruct[childField.name] = - childField.type.constructor.name === "Struct" - ? ensureStructFieldsExist({}, childField.type as Struct) - : null; + completeStruct[childField.name] = isPlainStructField(childField) + ? ensureStructFieldsExist({}, childField.type as Struct) + : null; } } diff --git a/nodejs/lancedb/blob.ts b/nodejs/lancedb/blob.ts index e6c780ca9..f3d648f48 100644 --- a/nodejs/lancedb/blob.ts +++ b/nodejs/lancedb/blob.ts @@ -12,7 +12,46 @@ const DEDICATED_SIZE_THRESHOLD_KEY = const PACK_FILE_SIZE_THRESHOLD_KEY = "lance-encoding:blob-pack-file-size-threshold"; -export type BlobInput = { +/** + * Bytes accepted for a blob value. `ArrayBuffer` is wrapped without copying. + * `Blob` and `File` are read with `arrayBuffer()`, which only the async write + * paths ({@link Connection.createTable}, {@link Table.add}, + * {@link Table.mergeInsert}) can do. + */ +export type BlobData = Buffer | Uint8Array | ArrayBuffer | Blob; + +/** A URI accepted for a blob value. A `URL` is stored as its `href`. */ +export type BlobUri = string | URL; + +/** + * A value accepted for a `lance.blob.v2` column (see {@link blob}). + * + * Either inline bytes, a URI pointing at external bytes, or a struct that sets + * exactly one of `data` or `uri`. `null` and `undefined` write a null blob. + * + * @example + * ```ts + * import { pathToFileURL } from "node:url"; + * import type { BlobInput } from "@lancedb/lancedb"; + * + * const rows: { id: bigint; image: BlobInput }[] = [ + * { id: 1n, image: await (await fetch(url)).arrayBuffer() }, + * { id: 2n, image: pathToFileURL("/data/cat.png") }, + * { id: 3n, image: { data: new Blob(["hello"]) } }, + * ]; + * await table.add(rows); + * ``` + */ +export type BlobInput = + | BlobData + | BlobUri + | { data: BlobData; uri?: null } + | { data?: null; uri: BlobUri } + | null + | undefined; + +/** A blob value normalized to the columns of the blob struct. */ +export type BlobValue = { data: Buffer | Uint8Array | null; uri: string | null; }; @@ -248,7 +287,77 @@ export type BlobRange = { end: bigint; }; -export function coerceBlobValue(value: unknown): BlobInput | null { +/** + * Rewrites the synchronous-only widenings of {@link BlobInput} (`ArrayBuffer`, + * `URL`) into `Uint8Array` and URI strings, keeping the value's shape. Other + * values are returned unchanged and validated later by {@link coerceBlobValue}. + * + * Throws on `Blob` / `File`, which must be read into bytes first. + */ +export function normalizeBlobInput(value: unknown): unknown { + if (blobToRead(value) !== undefined) { + throw new Error( + "Blob and File values must be read asynchronously. Pass them to " + + "createTable, add, or mergeInsert, or convert them with " + + "`new Uint8Array(await value.arrayBuffer())`", + ); + } + if (value instanceof ArrayBuffer) { + return new Uint8Array(value); + } + if (value instanceof URL) { + return value.href; + } + if (isBlobStruct(value)) { + // Read `data` and `uri` explicitly: spreading a class instance would drop + // fields defined as getters on its prototype. + const out: Record = { + ...value, + data: value.data, + uri: value.uri, + }; + let changed = false; + if (value.data instanceof ArrayBuffer) { + out.data = new Uint8Array(value.data); + changed = true; + } + if (value.uri instanceof URL) { + out.uri = value.uri.href; + changed = true; + } + return changed ? out : value; + } + return value; +} + +/** + * Returns the `Blob` / `File` a blob value needs read, either the value itself + * or its `data` field, or `undefined` when there is nothing to read. + */ +export function blobToRead(value: unknown): Blob | undefined { + if (isBlobLike(value)) { + return value; + } + if (isBlobStruct(value) && isBlobLike(value.data)) { + return value.data; + } + return undefined; +} + +/** + * Replaces the `Blob` that {@link blobToRead} found in `value` with `bytes`, + * keeping the value's shape. + */ +export function withBlobBytes(value: unknown, bytes: Uint8Array): unknown { + if (isBlobLike(value)) { + return bytes; + } + const record = value as Record; + return { ...record, uri: record.uri, data: bytes }; +} + +export function coerceBlobValue(input: unknown): BlobValue | null { + const value = normalizeBlobInput(input); if (value == null) { return null; } @@ -256,7 +365,9 @@ export function coerceBlobValue(value: unknown): BlobInput | null { return { data: value, uri: null }; } if (ArrayBuffer.isView(value)) { - throw new Error("Blob data must be Buffer or Uint8Array"); + throw new Error( + "Blob data must be Buffer, Uint8Array, ArrayBuffer, or Blob", + ); } if (typeof value === "string") { if (value === "") { @@ -264,23 +375,37 @@ export function coerceBlobValue(value: unknown): BlobInput | null { } return { data: null, uri: value }; } - if (typeof value === "object") { - const record = value as Record; + if (isBlobStruct(value)) { + const record = value; if (!("data" in record) && !("uri" in record)) { throw new Error( "Blob struct values must include a 'data' or 'uri' field", ); } + const unknownKeys = Object.keys(record).filter( + (key) => key !== "data" && key !== "uri", + ); + if (unknownKeys.length > 0) { + throw new Error( + `Blob struct values only support 'data' and 'uri', got ${unknownKeys + .map((key) => `'${key}'`) + .join(", ")}`, + ); + } const uri = record.uri; if (uri === "") { throw new Error("Blob uri cannot be empty"); } if (uri != null && typeof uri !== "string") { - throw new Error(`Blob uri must be a string or null, got ${typeof uri}`); + throw new Error( + `Blob uri must be a string, URL, or null, got ${typeof uri}`, + ); } const data = record.data; if (data != null && !isBlobBytes(data)) { - throw new Error("Blob data must be Buffer, Uint8Array, or null"); + throw new Error( + "Blob data must be Buffer, Uint8Array, ArrayBuffer, Blob, or null", + ); } const bytes = (data as Buffer | Uint8Array | null | undefined) ?? null; const uriValue = uri ?? null; @@ -292,7 +417,39 @@ export function coerceBlobValue(value: unknown): BlobInput | null { return { data: bytes, uri: uriValue }; } throw new Error( - "Blob column values must be Buffer, Uint8Array, a URI string, null, or { data?, uri? }", + "Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob, a URI string or URL, null, or { data?, uri? }", + ); +} + +function isBlobLike(value: unknown): value is Blob { + return typeof Blob !== "undefined" && value instanceof Blob; +} + +/** + * Whether `value` is a `{ data?, uri? }` struct: any plain object, or any + * other object with a `data` or `uri` field (such as a class instance), apart + * from the binary, `Blob`, and `URL` values that are blob inputs themselves. + */ +function isBlobStruct(value: unknown): value is Record { + if ( + typeof value !== "object" || + value === null || + Array.isArray(value) || + ArrayBuffer.isView(value) || + value instanceof ArrayBuffer || + (typeof SharedArrayBuffer !== "undefined" && + value instanceof SharedArrayBuffer) || + isBlobLike(value) || + value instanceof URL + ) { + return false; + } + const prototype = Object.getPrototypeOf(value); + return ( + prototype === Object.prototype || + prototype === null || + "data" in value || + "uri" in value ); } diff --git a/nodejs/lancedb/catalog.ts b/nodejs/lancedb/catalog.ts index 5cfe2c5ea..e344d8187 100644 --- a/nodejs/lancedb/catalog.ts +++ b/nodejs/lancedb/catalog.ts @@ -5,9 +5,9 @@ import { Connection, LocalConnection } from "./connection"; import { HeaderProvider } from "./header"; import { JsHeaderProvider, - ListDatabasesResponse, Catalog as NativeCatalog, CatalogOptions as NativeCatalogOptions, + DatabaseNames as NativeDatabaseNames, } from "./native.js"; import { OAuthConfig } from "./oauth"; @@ -21,7 +21,41 @@ export interface CatalogOptions | (() => Record | Promise>); } -export type { ListDatabasesResponse } from "./native.js"; +/** + * A lazy async iterator of database names with pagination state. + * Returned by {@link Catalog.listDatabases}. + */ +export class DatabaseNames implements AsyncIterableIterator { + /** @hidden */ + constructor(private readonly inner: NativeDatabaseNames) {} + + [Symbol.asyncIterator](): AsyncIterableIterator { + return this; + } + + /** Fetch the next name, requesting another page only when needed. */ + async next(): Promise> { + const value = await this.inner.next(); + return value == null + ? { done: true, value: undefined } + : { done: false, value }; + } + + /** Number of names available without another REST request. */ + numPageResults(): number { + return this.inner.numPageResults(); + } + + /** + * Token for the next REST request. Initially this is the supplied starting token. + * Returns undefined after the final page is fetched, even if names remain cached. + * Drain the cache before saving a token to avoid skipping names when resuming. + * A failed request terminates iteration and retains its token for a new iterator. + */ + pageToken(): string | undefined { + return this.inner.pageToken() ?? undefined; + } +} /** A remote catalog manages databases through the server's root namespace. */ export class Catalog { @@ -56,21 +90,35 @@ export class Catalog { await this.inner.dropDatabase(name, options.ignoreMissing); } - /** List one page of databases; pass pageToken from a response for the next page. */ - async listDatabases( - options: { limit?: number; pageToken?: string } = {}, - ): Promise { + /** + * Iterate lazily over all database names using `for await...of`. + * pageToken resumes from a saved token; omitted starts at the beginning. + * pageLimit limits each REST response, not the total; omitted uses the server default. + * Request errors are raised during iteration and terminate the iterator. + * + * @example + * ```ts + * for await (const name of catalog.listDatabases({ pageLimit: 20 })) { + * console.log(name); + * } + * ``` + */ + listDatabases( + options: { pageLimit?: number; pageToken?: string } = {}, + ): DatabaseNames { if ( - options.limit !== undefined && - (!Number.isInteger(options.limit) || - options.limit <= 0 || - options.limit > 2147483647) + options.pageLimit !== undefined && + (!Number.isInteger(options.pageLimit) || + options.pageLimit <= 0 || + options.pageLimit > 2147483647) ) { throw new Error( "Database list limit must be an integer between 1 and 2147483647", ); } - return this.inner.listDatabases(options.limit, options.pageToken); + return new DatabaseNames( + this.inner.listDatabases(options.pageToken, options.pageLimit), + ); } } @@ -82,7 +130,9 @@ export class Catalog { * ```ts * const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" }); * const db = await catalog.createDatabase("analytics", { existOk: true }); - * const page = await catalog.listDatabases({ limit: 20 }); + * for await (const name of catalog.listDatabases({ pageLimit: 20 })) { + * console.log(name); + * } * ``` */ export async function connectCatalog( diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 5e9641202..708f10fb4 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -8,6 +8,7 @@ import { fromTableToStreamBuffer, isArrowTable, makeArrowTable, + resolveBlobInputs, } from "./arrow"; import { Table as ArrowTable, @@ -324,10 +325,6 @@ export abstract class Connection { * The view is populated before creation returns. Set `withNoData` to create * only its definition and empty backing table. The view is a normal table: * it can be queried, indexed and searched, and it appears in `tableNames`. - * The source table must have stable row ids (create it with - * the `newTableEnableStableRowIds` storage option); they keep the view's - * provenance valid across source compactions and cannot be enabled after - * a table exists. */ abstract createMaterializedView( name: string, @@ -1133,7 +1130,11 @@ async function parseTableData( if (isArrowTable(data)) { table = sanitizeTable(data); } else { - table = makeArrowTable(data as Record[], options); + const records = await resolveBlobInputs( + data as Record[], + options?.schema, + ); + table = makeArrowTable(records, options); } if (streaming) { const buf = await fromTableToStreamBuffer( diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index a1776e3f4..694db9fa7 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -84,7 +84,14 @@ export { } from "./arrow"; export { blob, isBlobField, BlobFile } from "./blob"; -export type { BlobOptions, BlobRange, BlobReadOptions } from "./blob"; +export type { + BlobData, + BlobInput, + BlobOptions, + BlobRange, + BlobReadOptions, + BlobUri, +} from "./blob"; export { Connection, @@ -195,6 +202,7 @@ export * as rerankers from "./rerankers"; export { SchemaLike, TableLike, + DataTypeLike, FieldLike, RecordBatchLike, DataLike, @@ -638,6 +646,6 @@ export async function connectNamespace( export { Catalog, CatalogOptions, - ListDatabasesResponse, + DatabaseNames, connectCatalog, } from "./catalog"; diff --git a/nodejs/lancedb/materialized_view.ts b/nodejs/lancedb/materialized_view.ts index a59a3a5f6..7cf1690e0 100644 --- a/nodejs/lancedb/materialized_view.ts +++ b/nodejs/lancedb/materialized_view.ts @@ -9,10 +9,9 @@ export const DEFINITION_META_KEY = "mv.definition"; /** * The newest stored layout this version reads: `{"format": N, "query": ""}`, - * format 2 being a query with `GROUP BY`, format 3 native pairs and format 4 - * native dedup results. + * format 2 being a query with `GROUP BY`. */ -export const DEFINITION_FORMAT = 4; +export const DEFINITION_FORMAT = 2; /** * The query that defines a materialized view, as stored: @@ -188,23 +187,14 @@ export class MaterializedView { /** * Recompute the view from its source. * - * The refresh is incremental when the source's changes can be reconciled - * into the view -- rows added, changed or removed since the last one -- - * and otherwise rebuilds. `full` forces a rebuild; `sourceVersion` - * refreshes to that source version instead of the latest. - * - * Concurrent refreshes of one view do not duplicate its rows. Two that - * plan the same source rows conflict on commit, and the loser throws - * rather than writing them a second time. + * Every refresh rebuilds the complete view. For local views, + * `sourceVersion` refreshes to that source version instead of the latest; + * remote SQL refreshes do not support source-version pinning. */ async refresh(options?: { - full?: boolean; sourceVersion?: number; }): Promise { validateNonNegativeInteger(options?.sourceVersion, "sourceVersion"); - return await this.inner.refreshMaterializedView( - options?.full, - options?.sourceVersion, - ); + return await this.inner.refreshMaterializedView(options?.sourceVersion); } } diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 75a787fcd..f370e1b31 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -41,26 +41,6 @@ export async function* RecordBatchIterator( } } -class RecordBatchIterable< - NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery, -> implements AsyncIterable -{ - private inner: NativeQueryType; - private options?: QueryExecutionOptions; - - constructor(inner: NativeQueryType, options?: QueryExecutionOptions) { - this.inner = inner; - this.options = options; - } - - // biome-ignore lint/suspicious/noExplicitAny: skip - [Symbol.asyncIterator](): AsyncIterator, any, undefined> { - return RecordBatchIterator( - this.inner.execute(this.options?.maxBatchLength, this.options?.timeoutMs), - ); - } -} - /** * Options that control the behavior of a particular query execution */ @@ -271,13 +251,20 @@ export class QueryBase< return RecordBatchIterator(this.nativeExecute()); } - /** Collect the results as an Arrow @see {@link ArrowTable}. */ + /** + * Collect the results as an Arrow @see {@link ArrowTable}. + * + * Empty results retain the query's output schema. + */ async toArrow(options?: Partial): Promise { const batches = []; - const inner = await this.getInner(); - for await (const batch of new RecordBatchIterable(inner, options)) { + const iterator = this.nativeExecute(options); + for await (const batch of RecordBatchIterator(iterator)) { batches.push(batch); } + if (batches.length === 0) { + return tableFromIPC((await iterator).schema()); + } return new ArrowTable(batches); } @@ -532,6 +519,8 @@ export class VectorQuery extends StandardQueryBase { /** * Set the number of partitions to search (probe) * + * The number of probes must be greater than 0. + * * This argument is only used when the vector column has an IVF PQ index. * If there is no index then this value is ignored. * diff --git a/nodejs/lancedb/sanitize.ts b/nodejs/lancedb/sanitize.ts index 98a144451..6ba4cc881 100644 --- a/nodejs/lancedb/sanitize.ts +++ b/nodejs/lancedb/sanitize.ts @@ -9,6 +9,7 @@ // comes from the exact same library instance. This is not always the case // and so we must sanitize the input to ensure that it is compatible. +import * as arrowModule from "apache-arrow"; import { BufferType, Data, Vector } from "apache-arrow"; import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type"; import { @@ -76,6 +77,66 @@ import { Utf8, } from "./arrow"; +/** + * Whether `value` was created by *this* copy of apache-arrow's `ctor`. + * + * Arrow 21 overrides `Symbol.hasInstance` on `Schema`, `Field`, `Table`, + * `RecordBatch`, `Vector` and `Data` so that `instanceof` also matches objects + * built by a different copy of the library. Those are exactly the objects + * this module has to rebuild: the `DataType` instances nested inside them are + * deliberately *not* given that treatment upstream, so Arrow's own builders and + * writers reject them. Walk the real prototype chain instead of trusting + * `instanceof`. + */ +function isOwnInstance( + value: unknown, + // biome-ignore lint/suspicious/noExplicitAny: matches any class constructor + ctor: abstract new (...args: any[]) => T, +): value is T { + return ( + typeof value === "object" && + value !== null && + Object.prototype.isPrototypeOf.call(ctor.prototype, value) + ); +} + +/** + * Types added to Arrow JS after release 18, keyed by the `Type` enum value + * they use. The library compiles and runs against every Arrow release in the + * peer range, so these are looked up by name at runtime instead of imported: + * an older Arrow simply does not have the class. + */ +type NewerTypeName = + | "IntervalMonthDayNano" + | "LargeList" + | "BinaryView" + | "Utf8View"; +const NEWER_TYPES: ReadonlyArray<{ + name: NewerTypeName; + typeId: number; + since: number; +}> = [ + { name: "IntervalMonthDayNano", typeId: -31, since: 20 }, + { name: "LargeList", typeId: 21, since: 21 }, + { name: "BinaryView", typeId: 23, since: 21 }, + { name: "Utf8View", typeId: 24, since: 21 }, +]; + +// biome-ignore lint/suspicious/noExplicitAny: constructor arity varies by type +type DataTypeConstructor = new (...args: any[]) => DataType; + +function newerTypeConstructor(name: NewerTypeName): DataTypeConstructor { + const ctor = (arrowModule as unknown as Record)[name]; + if (typeof ctor !== "function") { + const since = NEWER_TYPES.find((t) => t.name === name)?.since; + throw new Error( + `The Arrow type ${name} requires apache-arrow ${since} or newer, ` + + "but the installed apache-arrow does not provide it", + ); + } + return ctor as DataTypeConstructor; +} + type SanitizationContext = { types: WeakMap; vectors: WeakMap; @@ -212,19 +273,37 @@ export function sanitizeList(typeLike: object) { return sanitizeListWithContext(typeLike, createSanitizationContext()); } +function sanitizeListChildWithContext( + typeLike: object, + context: SanitizationContext, + typeName: string, +): Field { + if (!("children" in typeLike) || !Array.isArray(typeLike.children)) { + throw Error( + `Expected a ${typeName} type to have an array-like \`children\` property`, + ); + } + if (typeLike.children.length !== 1) { + throw Error(`Expected a ${typeName} type to have exactly one child`); + } + return sanitizeFieldWithContext(typeLike.children[0], context); +} + function sanitizeListWithContext( typeLike: object, context: SanitizationContext, ) { - if (!("children" in typeLike) || !Array.isArray(typeLike.children)) { - throw Error( - "Expected a List type to have an array-like `children` property", - ); - } - if (typeLike.children.length !== 1) { - throw Error("Expected a List type to have exactly one child"); - } - return new List(sanitizeFieldWithContext(typeLike.children[0], context)); + return new List(sanitizeListChildWithContext(typeLike, context, "List")); +} + +function sanitizeLargeListWithContext( + typeLike: object, + context: SanitizationContext, +) { + const LargeList = newerTypeConstructor("LargeList"); + return new LargeList( + sanitizeListChildWithContext(typeLike, context, "LargeList"), + ); } export function sanitizeStruct(typeLike: object) { @@ -451,11 +530,30 @@ function sanitizeTypeWithContext( return type; } +function sanitizeNewerTypeById( + typeLike: object, + typeId: number, + context: SanitizationContext, +): DataType | undefined { + const entry = NEWER_TYPES.find((t) => t.typeId === typeId); + if (entry === undefined) { + return undefined; + } + if (entry.name === "LargeList") { + return sanitizeLargeListWithContext(typeLike, context); + } + return new (newerTypeConstructor(entry.name))(); +} + function sanitizeTypeById( typeLike: object, typeId: Type, context: SanitizationContext, ): DataType { + const newer = sanitizeNewerTypeById(typeLike, typeId, context); + if (newer !== undefined) { + return newer; + } switch (typeId) { case Type.NONE: throw Error("Received a Type with a typeId of NONE"); @@ -572,7 +670,7 @@ function sanitizeFieldWithContext( fieldLike: unknown, context: SanitizationContext, ): Field { - if (fieldLike instanceof Field) { + if (isOwnInstance(fieldLike, Field)) { return fieldLike; } if (typeof fieldLike !== "object" || fieldLike === null) { @@ -626,7 +724,7 @@ function sanitizeSchemaWithContext( schemaLike: SchemaLike, context: SanitizationContext, ): Schema { - if (schemaLike instanceof Schema) { + if (isOwnInstance(schemaLike, Schema)) { return schemaLike; } if (typeof schemaLike !== "object" || schemaLike === null) { @@ -653,7 +751,7 @@ function sanitizeSchemaWithContext( } export function sanitizeTable(tableLike: TableLike): Table { - if (tableLike instanceof Table) { + if (isOwnInstance(tableLike, Table)) { return tableLike; } if (typeof tableLike !== "object" || tableLike === null) { @@ -681,7 +779,7 @@ function sanitizeRecordBatch( batchLike: RecordBatchLike, context: SanitizationContext, ): RecordBatch { - if (batchLike instanceof RecordBatch) { + if (isOwnInstance(batchLike, RecordBatch)) { return batchLike; } if (typeof batchLike !== "object" || batchLike === null) { @@ -710,11 +808,17 @@ type DictionaryDataLike = DataLike & { dictionary?: DictionaryVectorLike; }; +// Arrow 21 added `variadicBuffers` to `Data` for the Utf8View / BinaryView +// types. Older Arrow versions do not have the property, so it is optional. +type VariadicDataLike = DataLike & { + variadicBuffers?: ReadonlyArray; +}; + function sanitizeData( dataLike: DataLike, context: SanitizationContext, ): Data { - if (dataLike instanceof Data) { + if (isOwnInstance(dataLike, Data)) { return dataLike; } const cachedData = context.data.get(dataLike); @@ -732,7 +836,12 @@ function sanitizeData( context.vectors.set(dictionaryLike, dictionary); } } - const data = new Data( + // Arrow 21 added an eighth constructor parameter (`variadicBuffers`); older + // releases ignore the extra argument but their typings do not declare it. + const DataCtor = Data as unknown as new ( + ...args: unknown[] + ) => Data; + const data = new DataCtor( sanitizeTypeWithContext(dataLike.type, context), dataLike.offset, dataLike.length, @@ -745,6 +854,7 @@ function sanitizeData( }, dataLike.children.map((child) => sanitizeData(child, context)), dictionary, + (dataLike as VariadicDataLike).variadicBuffers, ); context.data.set(dataLike, data); return data; @@ -753,7 +863,11 @@ function sanitizeData( const constructorsByTypeName = { null: () => new Null(), binary: () => new Binary(), + largebinary: () => new LargeBinary(), + binaryview: () => new (newerTypeConstructor("BinaryView"))(), utf8: () => new Utf8(), + largeutf8: () => new LargeUtf8(), + utf8view: () => new (newerTypeConstructor("Utf8View"))(), bool: () => new Bool(), int8: () => new Int8(), int16: () => new Int16(), @@ -774,6 +888,8 @@ const constructorsByTypeName = { timesecond: () => new TimeSecond(), intervaldaytime: () => new IntervalDayTime(), intervalyearmonth: () => new IntervalYearMonth(), + intervalmonthdaynano: () => + new (newerTypeConstructor("IntervalMonthDayNano"))(), durationnanosecond: () => new DurationNanosecond(), durationmicrosecond: () => new DurationMicrosecond(), durationmillisecond: () => new DurationMillisecond(), diff --git a/nodejs/lancedb/schema.ts b/nodejs/lancedb/schema.ts index 724a81fc9..3e225464f 100644 --- a/nodejs/lancedb/schema.ts +++ b/nodejs/lancedb/schema.ts @@ -69,12 +69,18 @@ class SchemaInferrer { this.addField(path, value, row); } else if (this.providedSchema === undefined) { this.updateInferredField(path, value, row, current); + } else if (current instanceof DeferredTypeEvidence) { + if (value != null) { + this.addSchemaField(this.providedSchema, path, value, row); + } + } else if (current instanceof FieldTree || DataType.isStruct(current)) { + assertStructValue(path, value, row); } } private addField(path: string[], value: unknown, row: number): void { if (this.providedSchema !== undefined) { - this.addSchemaField(this.providedSchema, path, row); + this.addSchemaField(this.providedSchema, path, value, row); return; } @@ -95,15 +101,31 @@ class SchemaInferrer { } } - private addSchemaField(schema: Schema, path: string[], row: number): void { + private addSchemaField( + schema: Schema, + path: string[], + value: unknown, + row: number, + ): void { const field = fieldAtPath(schema, path); if (field === undefined) { throw new Error( `Found field not in schema: ${path.join(".")} at row ${row}`, ); } + if (DataType.isStruct(field.type)) { + assertStructValue(path, value, row); + } - const conflict = this.fields.set(path, field.type); + // Nulls are kept as deferred evidence so that a struct field seen first as + // null can be replaced by a record in a later row, just as a record seen + // first already ignores later nulls. + const conflict = this.fields.set( + path, + DeferredTypeEvidence.from(value, row) ?? field.type, + (existing) => + existing instanceof DeferredTypeEvidence && existing.isOnlyNulls(), + ); if (conflict !== undefined) { throw branchConflictError(conflict, row, "Struct"); } @@ -483,6 +505,26 @@ function assertSupportedValue( } } +/** + * Rejects primitives, arrays, and typed arrays given for a struct field, which + * would otherwise be written as a struct with every child null. + */ +function assertStructValue(path: string[], value: unknown, row: number): void { + if ( + value == null || + (typeof value === "object" && + !Array.isArray(value) && + !ArrayBuffer.isView(value)) + ) { + return; + } + const kind = + typeof value === "object" ? value.constructor.name : typeof value; + throw new Error( + `Expected a struct value for field ${path.join(".")} at row ${row}, got ${kind}.`, + ); +} + function fieldAtPath(schema: Schema, path: string[]): Field | undefined { let fields = schema.fields; let field: Field | undefined; diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index d4c5db509..6bdc69418 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -685,7 +685,6 @@ export abstract class Table { * @ignore */ abstract refreshMaterializedView( - full?: boolean, sourceVersion?: number, ): Promise; @@ -1420,10 +1419,9 @@ export class LocalTable extends Table { } async refreshMaterializedView( - full?: boolean, sourceVersion?: number, ): Promise { - return await this.inner.refreshMaterializedView(full, sourceVersion); + return await this.inner.refreshMaterializedView(sourceVersion); } async materializedViewDefinition(): Promise { diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 051bdb98f..cb3a592d6 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "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 755035ee2..8146e09db 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "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 d31737229..b097a423c 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "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 618df080d..7c7903607 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "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 87628c0c3..5359f6490 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "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 65c7b500c..e16f9d7e5 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index ad6a81acc..3fd95907f 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 29bbe4314..9ac9ede63 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.40.0-beta.12", + "version": "0.40.0-beta.13", "main": "dist/index.js", "exports": { ".": "./dist/index.js", @@ -49,10 +49,14 @@ "@types/jest": "^29.1.2", "@types/node": "22.7.4", "@types/tmp": "^0.2.6", + "apache-arrow": "21.2.0", "apache-arrow-15": "npm:apache-arrow@15.0.0", "apache-arrow-16": "npm:apache-arrow@16.0.0", "apache-arrow-17": "npm:apache-arrow@17.0.0", "apache-arrow-18": "npm:apache-arrow@18.0.0", + "apache-arrow-19": "npm:apache-arrow@19.0.1", + "apache-arrow-20": "npm:apache-arrow@20.0.0", + "apache-arrow-21": "npm:apache-arrow@21.2.0", "eslint": "^8.57.0", "jest": "^29.7.0", "shx": "^0.3.4", @@ -102,7 +106,7 @@ }, "peerDependencies": { "@types/node": ">=22", - "apache-arrow": ">=15.0.0 <=18.1.0" + "apache-arrow": ">=15.0.0 <22.0.0" }, "peerDependenciesMeta": { "@types/node": { diff --git a/nodejs/pnpm-lock.yaml b/nodejs/pnpm-lock.yaml index a03838999..4c19305c8 100644 --- a/nodejs/pnpm-lock.yaml +++ b/nodejs/pnpm-lock.yaml @@ -14,9 +14,6 @@ importers: '@opentelemetry/api': specifier: ^1.9.0 version: 1.9.1 - apache-arrow: - specifier: '>=15.0.0 <=18.1.0' - version: 18.1.0 reflect-metadata: specifier: ^0.2.2 version: 0.2.2 @@ -54,6 +51,9 @@ importers: '@types/tmp': specifier: ^0.2.6 version: 0.2.6 + apache-arrow: + specifier: 21.2.0 + version: 21.2.0 apache-arrow-15: specifier: npm:apache-arrow@15.0.0 version: apache-arrow@15.0.0 @@ -66,6 +66,15 @@ importers: apache-arrow-18: specifier: npm:apache-arrow@18.0.0 version: apache-arrow@18.0.0 + apache-arrow-19: + specifier: npm:apache-arrow@19.0.1 + version: apache-arrow@19.0.1 + apache-arrow-20: + specifier: npm:apache-arrow@20.0.0 + version: apache-arrow@20.0.0 + apache-arrow-21: + specifier: npm:apache-arrow@21.2.0 + version: apache-arrow@21.2.0 eslint: specifier: ^8.57.0 version: 8.57.1 @@ -1703,6 +1712,9 @@ packages: '@types/node@22.7.4': resolution: {integrity: sha512-y+NPi1rFzDs1NdQHHToqeiX2TIS79SWEAw9GYhkkx8bD0ChpfqC+n2j5OXOCpzfojBEBt6DnEnnG9MY0zk1XLg==} + '@types/node@25.9.8': + resolution: {integrity: sha512-VfMrScDmMhUJQmd5hArdQnFvK0OIeD36uN2Va1FcpYaLB/BgkgM8Ulc50XtYISPMz7APJ30+N0T5EM0jlcdfRw==} + '@types/stack-utils@2.0.3': resolution: {integrity: sha512-9aEbYZ3TbYMznPdcdr3SmIrLXwC/AKZXQeCf9Pgao5CKb8CyHuEX5jzWPTkvregvhRJHcpRO6BFoGW9ycaOkYw==} @@ -1840,8 +1852,16 @@ packages: resolution: {integrity: sha512-gFlPaqN9osetbB83zC29AbbZqGiCuFH1vyyPseJ+B7SIbfBtESV62mMT/CkiIt77W6ykC/nTWFzTXFs0Uldg4g==} hasBin: true - apache-arrow@18.1.0: - resolution: {integrity: sha512-v/ShMp57iBnBp4lDgV8Jx3d3Q5/Hac25FWmQ98eMahUiHPXcvwIMKJD0hBIgclm/FCG+LwPkAKtkRO1O/W0YGg==} + apache-arrow@19.0.1: + resolution: {integrity: sha512-APmMLzS4qbTivLrPdQXexGM4JRr+0g62QDaobzEvip/FdQIrv2qLy0mD5Qdmw4buydtVJgbFeKR8f59I6PPGDg==} + hasBin: true + + apache-arrow@20.0.0: + resolution: {integrity: sha512-JUeK0jFRUd7rbmrhhzR3O2KXjLaZ4YYYFOptyUfxOsMIoZCPi6bZR58gVi/xi3HTBMPseXm9PXyQ2V916930pA==} + hasBin: true + + apache-arrow@21.2.0: + resolution: {integrity: sha512-Hxe6Agq26gQOM954qpzYSllJBPJl+e16U5CkfuMUhLrNba+5nKkttIVlflaovN6oaTratqMGAO8H5u/aNhmHWQ==} hasBin: true argparse@1.0.10: @@ -2033,6 +2053,15 @@ packages: resolution: {integrity: sha512-H4UfQhZyakIjC74I9d34fGYDwk3XpSr17QhEd0Q3I9Xq1CETHo4Hcuo87WyWHpAF1aSLjLRf5lD9ZGX2qStUvg==} engines: {node: '>=4.0.0'} + command-line-args@6.0.2: + resolution: {integrity: sha512-AIjYVxrV9X752LmPDLbVYv8aMCuHPSLZJXEo2qo/xJfv+NYhaZ4sMSF01rM+gHPaMgvPM0l5D/F+Qx+i2WfSmQ==} + engines: {node: '>=12.20'} + peerDependencies: + '@75lb/nature': latest + peerDependenciesMeta: + '@75lb/nature': + optional: true + command-line-usage@7.0.4: resolution: {integrity: sha512-85UdvzTNx/+s5CkSgBm/0hzP80RFHAa7PsfeADE5ezZF3uHz3/Tqj9gIKGT9PTtpycc3Ua64T0oVulGfKxzfqg==} engines: {node: '>=12.20.0'} @@ -2281,6 +2310,15 @@ packages: resolution: {integrity: sha512-6Tb2myMioCAgv5kfvP5/PkZZ/ntTpVK39fHY7WkWBgvbeE+VHd/tZuZ4mrC+bxh4cfOZeYKVPaJIZtZXV7GNCQ==} engines: {node: '>=4.0.0'} + find-replace@5.0.2: + resolution: {integrity: sha512-Y45BAiE3mz2QsrN2fb5QEtO4qb44NcS7en/0y9PEVsg351HsLeVclP8QPMH79Le9sH3rs5RSwJu99W0WPZO43Q==} + engines: {node: '>=14'} + peerDependencies: + '@75lb/nature': latest + peerDependenciesMeta: + '@75lb/nature': + optional: true + find-up@4.1.0: resolution: {integrity: sha512-PpOwAdQ/YlXQ2vj8a3h8IipDuYRi3wceVQQGYWxNINccq40Anw7BlsEXCMbt1Zt+OLA6Fq9suIpIWD0OsnISlw==} engines: {node: '>=8'} @@ -2302,6 +2340,9 @@ packages: flatbuffers@24.12.23: resolution: {integrity: sha512-dLVCAISd5mhls514keQzmEG6QHmUUsNuWsb4tFafIUwvvgDjXhtfAYSKOzt5SWOy+qByV5pbsDZ+Vb7HUOBEdA==} + flatbuffers@25.9.23: + resolution: {integrity: sha512-MI1qs7Lo4Syw0EOzUl0xjs2lsoeqFku44KpngfIduHBYvzm8h2+7K8YMQh1JtVVVrUvhLpNwqVi4DERegUJhPQ==} + flatted@3.4.2: resolution: {integrity: sha512-PjDse7RzhcPkIJwy5t7KPWQSZ9cAbzQXcafsetQoD7sOJRQlGikNbx7yZp2OotDnJyrDcbyRq3Ttb18iYOqkxA==} @@ -3345,6 +3386,9 @@ packages: undici-types@6.21.0: resolution: {integrity: sha512-iwDZqg0QAGrg9Rav5H4n0M64c3mkR59cJ6wQp+7C4nI0gsmExaedaYLNO44eT4AtBBwjbTiGPMlt2Md0T9H9JQ==} + undici-types@7.24.6: + resolution: {integrity: sha512-WRNW+sJgj5OBN4/0JpHFqtqzhpbnV0GuB+OozA9gCL7a993SmU+1JBZCzLNxYsbMfIeDL+lTsphD5jN5N+n0zg==} + unist-util-is@6.0.1: resolution: {integrity: sha512-LsiILbtBETkDz8I9p1dQ0uyRUWuaQzd/cuEeS1hoRSyW5E5XGmTzlwY1OrNzzakGowI9Dr/I8HVaw4hTtnxy8g==} @@ -5505,6 +5549,10 @@ snapshots: dependencies: undici-types: 6.19.8 + '@types/node@25.9.8': + dependencies: + undici-types: 7.24.6 + '@types/stack-utils@2.0.3': {} '@types/tmp@0.2.6': {} @@ -5694,17 +5742,40 @@ snapshots: json-bignum: 0.0.3 tslib: 2.8.1 - apache-arrow@18.1.0: + apache-arrow@19.0.1: dependencies: '@swc/helpers': 0.5.21 '@types/command-line-args': 5.2.3 '@types/command-line-usage': 5.0.4 '@types/node': 20.19.41 - command-line-args: 5.2.1 + command-line-args: 6.0.2 command-line-usage: 7.0.4 flatbuffers: 24.12.23 json-bignum: 0.0.3 tslib: 2.8.1 + transitivePeerDependencies: + - '@75lb/nature' + + apache-arrow@20.0.0: + dependencies: + '@swc/helpers': 0.5.21 + '@types/command-line-args': 5.2.3 + '@types/command-line-usage': 5.0.4 + '@types/node': 20.19.41 + command-line-args: 6.0.2 + command-line-usage: 7.0.4 + flatbuffers: 25.9.23 + json-bignum: 0.0.3 + tslib: 2.8.1 + transitivePeerDependencies: + - '@75lb/nature' + + apache-arrow@21.2.0: + dependencies: + '@types/node': 25.9.8 + flatbuffers: 25.9.23 + json-with-bigint: 3.5.8 + tslib: 2.8.1 argparse@1.0.10: dependencies: @@ -5906,6 +5977,13 @@ snapshots: lodash.camelcase: 4.3.0 typical: 4.0.0 + command-line-args@6.0.2: + dependencies: + array-back: 6.2.3 + find-replace: 5.0.2 + lodash.camelcase: 4.3.0 + typical: 7.3.0 + command-line-usage@7.0.4: dependencies: array-back: 6.2.3 @@ -6178,6 +6256,8 @@ snapshots: dependencies: array-back: 3.1.0 + find-replace@5.0.2: {} + find-up@4.1.0: dependencies: locate-path: 5.0.0 @@ -6201,6 +6281,8 @@ snapshots: flatbuffers@24.12.23: {} + flatbuffers@25.9.23: {} + flatted@3.4.2: {} follow-redirects@1.16.0: {} @@ -7405,6 +7487,8 @@ snapshots: undici-types@6.21.0: {} + undici-types@7.24.6: {} + unist-util-is@6.0.1: dependencies: '@types/unist': 3.0.3 diff --git a/nodejs/src/catalog.rs b/nodejs/src/catalog.rs index c2288bc4a..82edda5da 100644 --- a/nodejs/src/catalog.rs +++ b/nodejs/src/catalog.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Duration; -use lancedb::catalog::{ - CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest, -}; +use futures::{StreamExt, future::poll_fn}; + +use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest}; use napi::bindgen_prelude::*; use napi_derive::napi; @@ -25,10 +25,37 @@ pub struct CatalogOptions { pub oauth_config: Option, } -#[napi(object)] -pub struct ListDatabasesResponse { - pub databases: Vec, - pub page_token: Option, +/// A lazy iterator over database names with inspectable pagination state. +#[napi] +pub struct DatabaseNames { + inner: Mutex, + next_lock: futures::lock::Mutex<()>, +} + +#[napi] +impl DatabaseNames { + /// Number of names cached without another REST request. + #[napi] + pub fn num_page_results(&self) -> u32 { + self.inner.lock().unwrap().num_page_results() as u32 + } + + /// Continuation token for the next REST request. + #[napi] + pub fn page_token(&self) -> Option { + self.inner.lock().unwrap().page_token().map(str::to_owned) + } + + /// Fetch the next name, returning None when exhausted. + #[napi] + pub async fn next(&self) -> Result> { + // Serialize advances but allow synchronous inspection between polls. + let _guard = self.next_lock.lock().await; + poll_fn(|cx| self.inner.lock().unwrap().poll_next_unpin(cx)) + .await + .transpose() + .default_error() + } } #[napi] @@ -106,18 +133,14 @@ impl Catalog { .default_error() } #[napi] - pub async fn list_databases( + pub fn list_databases( &self, - limit: Option, page_token: Option, - ) -> Result { - let mut request = ListDatabasesRequest::default(); - request.limit = limit; - request.page_token = page_token; - let response = self.inner.list_databases(request).await.default_error()?; - Ok(ListDatabasesResponse { - databases: response.databases, - page_token: response.page_token, - }) + page_limit: Option, + ) -> DatabaseNames { + DatabaseNames { + inner: Mutex::new(self.inner.list_databases(page_token, page_limit)), + next_lock: futures::lock::Mutex::new(()), + } } } diff --git a/nodejs/src/iterator.rs b/nodejs/src/iterator.rs index 8d65dfc3f..96d08c765 100644 --- a/nodejs/src/iterator.rs +++ b/nodejs/src/iterator.rs @@ -7,6 +7,8 @@ use lancedb::ipc::batches_to_ipc_file; use napi::bindgen_prelude::*; use napi_derive::napi; +use crate::util::schema_to_buffer; + /** Typescript-style Async Iterator over RecordBatches */ #[napi] pub struct RecordBatchIterator { @@ -19,6 +21,12 @@ impl RecordBatchIterator { Self { inner } } + /// Return the executed stream's output schema as an Arrow IPC file. + #[napi] + pub fn schema(&self) -> napi::Result { + schema_to_buffer(&self.inner.schema()) + } + #[napi(catch_unwind)] pub async unsafe fn next(&mut self) -> napi::Result> { if let Some(rst) = self.inner.next().await { diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index f542fe856..712569b2f 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -289,8 +289,13 @@ impl VectorQuery { } #[napi] - pub fn nprobes(&mut self, nprobe: u32) { - self.inner = self.inner.clone().nprobes(nprobe as usize); + pub fn nprobes(&mut self, nprobe: u32) -> napi::Result<()> { + self.inner = self + .inner + .clone() + .nprobes(nprobe as usize) + .default_error()?; + Ok(()) } #[napi] diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index f5c3b2987..79b335c37 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -469,13 +469,12 @@ impl Table { #[napi(catch_unwind)] pub async fn refresh_materialized_view( &self, - full: Option, source_version: Option, ) -> napi::Result { let view = lancedb::MaterializedView::from_table(self.inner_ref()?.clone()) .await .default_error()?; - let mut builder = view.refresh().full(full.unwrap_or(false)); + let mut builder = view.refresh(); if let Some(version) = source_version { let version = u64::try_from(version).map_err(|_| { napi::Error::from_reason("sourceVersion must be a non-negative integer") @@ -492,11 +491,9 @@ impl Table { let view = lancedb::MaterializedView::from_table(inner) .await .default_error()?; - view.definition().to_json().map_err(|err| { - napi::Error::from_reason(format!( - "failed to serialize materialized-view definition: {err}" - )) - }) + Ok(lancedb::materialized_view::definition_metadata_from_sql( + view.definition_sql(), + )) } #[napi(catch_unwind)] @@ -1608,7 +1605,7 @@ impl From for FunctionErrors { #[napi(object)] pub struct RefreshMaterializedViewResult { - /// How the view was brought up to date: "rebuild", "incremental" or "no_op". + /// How the view was brought up to date: "rebuild". pub mode: String, pub rows_written: i64, pub source_version: i64, @@ -1619,8 +1616,6 @@ impl From for RefreshMaterializedViewRes fn from(value: lancedb::RefreshMaterializedViewResult) -> Self { let mode = match value.mode { lancedb::RefreshMode::Rebuild => "rebuild", - lancedb::RefreshMode::Incremental => "incremental", - lancedb::RefreshMode::NoOp => "no_op", }; Self { mode: mode.to_string(), diff --git a/nodejs/tsconfig.json b/nodejs/tsconfig.json index 9b83da5ec..094fe5003 100644 --- a/nodejs/tsconfig.json +++ b/nodejs/tsconfig.json @@ -10,7 +10,10 @@ "resolveJsonModule": true, "emitDecoratorMetadata": true, "experimentalDecorators": true, - "moduleResolution": "Node" + "moduleResolution": "Node", + // flatbuffers 25 (via apache-arrow 21) ships declarations that need + // TypeScript 5.7+ generic typed arrays; skip checking third-party .d.ts. + "skipLibCheck": true }, "exclude": ["./dist/*", "./examples/*"], "typedocOptions": { diff --git a/python/Cargo.toml b/python/Cargo.toml index 799a5df62..625422d06 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.40.0-beta.12" +version = "0.40.0-beta.13" publish = false edition.workspace = true description = "Python bindings for LanceDB" @@ -8,7 +8,7 @@ license.workspace = true repository.workspace = true keywords.workspace = true categories.workspace = true -rust-version = "1.91.0" +rust-version = "1.94.0" [lib] name = "_lancedb" diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index d3a678359..213794ef3 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -46,7 +46,6 @@ from .materialized_view import ( AsyncMaterializedView, MaterializedView, MaterializedViewDefinition, - vector_dedup as vector_dedup, ) from .view import ViewDescription as ViewDescription from .table import AsyncTable, Table @@ -62,7 +61,6 @@ from .namespace import ( from .catalog import ( AsyncCatalog, Catalog, - ListDatabasesResponse, connect_catalog, connect_catalog_async, ) @@ -578,7 +576,6 @@ async def connect_async( __all__ = [ "Catalog", "AsyncCatalog", - "ListDatabasesResponse", "connect_catalog", "connect_catalog_async", "AsyncMaterializedView", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 176e57ac3..8ff53bcbe 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -255,8 +255,6 @@ class Connection(object): filter: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, - *, - vector_source_json: Optional[str] = None, ) -> Table: ... async def create_materialized_view_async( self, @@ -266,8 +264,6 @@ class Connection(object): filter: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, - *, - vector_source_json: Optional[str] = None, ) -> Job: ... async def list_materialized_views(self) -> List[str]: ... async def drop_materialized_view( @@ -500,10 +496,10 @@ class Table: limit: Optional[int] = None, ) -> FunctionErrors: ... async def refresh_materialized_view( - self, full: bool = False, source_version: Optional[int] = None + self, source_version: Optional[int] = None ) -> RefreshMaterializedViewResult: ... async def refresh_materialized_view_async( - self, full: bool = False, source_version: Optional[int] = None + self, source_version: Optional[int] = None ) -> Job: ... async def materialized_view_definition(self) -> str: ... async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ... @@ -946,9 +942,15 @@ class Catalog: async def drop_database( self, name: str, *, ignore_missing: bool = False ) -> None: ... - async def list_databases( - self, *, limit: Optional[int] = None, page_token: Optional[str] = None - ) -> tuple[list[str], Optional[str]]: ... + def list_databases( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> DatabaseNames: ... + +class DatabaseNames: + def num_page_results(self) -> int: ... + def page_token(self) -> Optional[str]: ... + def __aiter__(self) -> DatabaseNames: ... + async def __anext__(self) -> str: ... async def connect_catalog( endpoint: str, diff --git a/python/python/lancedb/catalog.py b/python/python/lancedb/catalog.py index 4fd941ddd..70aab1bbb 100644 --- a/python/python/lancedb/catalog.py +++ b/python/python/lancedb/catalog.py @@ -3,9 +3,8 @@ """Remote catalogs manage databases through a server's root namespace.""" -from dataclasses import dataclass from datetime import timedelta -from typing import Any, Optional, Union +from typing import Any, AsyncIterator, Iterator, Optional, Union from . import _lancedb from .background_loop import LOOP @@ -14,12 +13,66 @@ from .remote import ClientConfig, OAuthConfig from .remote.db import RemoteDBConnection -@dataclass -class ListDatabasesResponse: - """A page of database names and an optional continuation token.""" +class AsyncDatabaseNames(AsyncIterator[str]): + """An async iterator of database names with pagination state. - databases: list[str] - page_token: Optional[str] = None + Returned by [list_databases][lancedb.catalog.AsyncCatalog.list_databases]. + """ + + def __init__(self, inner: _lancedb.DatabaseNames): + self._inner = inner + + def __aiter__(self) -> "AsyncDatabaseNames": + return self + + async def __anext__(self) -> str: + return await self._inner.__anext__() + + def num_page_results(self) -> int: + """Return the number of names available without another REST request.""" + return self._inner.num_page_results() + + def page_token(self) -> Optional[str]: + """Return the token for the next REST request. + + Before the first request, this is the supplied starting token (``None`` + starts at the beginning). After the final page is fetched, it is ``None``, + even if names remain cached. Drain the cache before saving a token to + avoid skipping those names when resuming. A failed request terminates + iteration and leaves its token available for resuming a new iterator. + """ + return self._inner.page_token() + + +class DatabaseNames(Iterator[str]): + """A synchronous iterator of database names with pagination state. + + Returned by [Catalog.list_databases][lancedb.catalog.Catalog.list_databases]. + """ + + def __init__(self, inner: AsyncDatabaseNames): + self._inner = inner + + def __iter__(self) -> "DatabaseNames": + return self + + def __next__(self) -> str: + try: + return LOOP.run(self._inner.__anext__()) + except StopAsyncIteration: + raise StopIteration from None + + def num_page_results(self) -> int: + """Return the number of names available without another REST request.""" + return self._inner.num_page_results() + + def page_token(self) -> Optional[str]: + """Return the token for the next REST request. + + See [page_token][lancedb.catalog.AsyncDatabaseNames.page_token] + for initial, cached, final-page, and error behavior. + """ + return self._inner.page_token() class AsyncCatalog: @@ -50,14 +103,25 @@ class AsyncCatalog: """Connect to an existing database by its logical name.""" return AsyncConnection(await self._inner.connect_database(name)) - async def list_databases( - self, *, limit: Optional[int] = None, page_token: Optional[str] = None - ) -> ListDatabasesResponse: - """List a page of databases. Pass the returned token for the next page.""" - names, token = await self._inner.list_databases( - limit=limit, page_token=page_token + def list_databases( + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> AsyncDatabaseNames: + """Iterate lazily over all database names using ``async for``. + + ``page_token`` resumes from a saved token; ``None`` starts at the beginning. + ``page_limit`` sets the maximum names per REST response (not the total); + ``None`` uses the server default. Request errors are raised during iteration. + + Examples + -------- + ```python + async for name in catalog.list_databases(): + print(name) + ``` + """ + return AsyncDatabaseNames( + self._inner.list_databases(page_token=page_token, page_limit=page_limit) ) - return ListDatabasesResponse(names, token) async def drop_database(self, name: str, *, ignore_missing: bool = False) -> None: """Drop an empty database. A nonempty database is an error.""" @@ -73,7 +137,8 @@ class Catalog: ```python catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret") db = catalog.create_database("analytics", exist_ok=True) - page = catalog.list_databases(limit=20) + for name in catalog.list_databases(): + print(name) ``` """ @@ -118,10 +183,17 @@ class Catalog: ) def list_databases( - self, *, limit: Optional[int] = None, page_token: Optional[str] = None - ) -> ListDatabasesResponse: - """List a page of databases. Pass the returned token for the next page.""" - return LOOP.run(self._inner.list_databases(limit=limit, page_token=page_token)) + self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None + ) -> DatabaseNames: + """Iterate lazily over database names, fetching more results as needed. + + ``page_token`` resumes from a saved token; ``None`` starts at the beginning. + ``page_limit`` sets the maximum names per REST response (not the total); + ``None`` uses the server default. Request errors are raised during iteration. + """ + return DatabaseNames( + self._inner.list_databases(page_token=page_token, page_limit=page_limit) + ) def drop_database(self, name: str, *, ignore_missing: bool = False) -> None: """Drop an empty database. A nonempty database is an error.""" diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 60e1780f0..3a9815cb4 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -57,7 +57,6 @@ from .materialized_view import ( AsyncMaterializedView, MaterializedView, MaterializedViewSource, - VectorDedupSource, SelectArg, normalize_select, ) @@ -122,7 +121,48 @@ def _view_description( class DBConnection(EnforceOverrides): - """An active LanceDB connection interface.""" + """An active LanceDB connection interface. + + Use [close][lancedb.db.DBConnection.close] to release the connection's + underlying resources, or use the connection as a context manager to close it + automatically when leaving the block, including when an exception is raised. + + Examples + -------- + >>> import lancedb + >>> with lancedb.connect("memory://") as db: + ... assert db.is_open() + >>> db.is_open() + False + """ + + def __enter__(self) -> DBConnection: + return self + + def __exit__(self, *_) -> None: + self.close() + + @abstractmethod + def is_open(self) -> bool: + """Return True if the connection is open.""" + pass + + @abstractmethod + def close(self) -> None: + """Close the connection, releasing any underlying resources. + + It is safe to call this method multiple times. Database operations on a + closed connection raise ``RuntimeError: Connection is closed``. + + Examples + -------- + >>> import lancedb + >>> db = lancedb.connect("memory://") + >>> db.close() + >>> db.is_open() + False + """ + pass def list_namespaces( self, @@ -368,9 +408,10 @@ class DBConnection(EnforceOverrides): - [LanceModel][lancedb.pydantic.LanceModel] mode: str; default "create" The mode to use when creating the table. - Can be either "create" or "overwrite". + Can be "create", "overwrite", or "exist_ok". By default, if the table already exists, an exception is raised. If you want to overwrite the table, use mode="overwrite". + To open an existing table without adding data, use mode="exist_ok". exist_ok: bool, default False If a table by the same name already exists, then raise an exception if exist_ok=False. If exist_ok=True, then open the existing table; @@ -571,21 +612,12 @@ class DBConnection(EnforceOverrides): table. The view is a normal table: it can be queried, indexed and searched, and it appears in ``table_names``. - Ordinary source tables must have stable row ids (create them with the - ``new_table_enable_stable_row_ids`` storage option): they keep the - view's provenance valid across source compactions, and cannot be - enabled after a table exists. Native dedup sources also support physical - row IDs interpreted against their pinned snapshot. - Parameters ---------- name: str The name of the view. - source: str or VectorDedupSource - The name of the source table, or an indexed dedup source declared - with [vector_dedup][lancedb.vector_dedup]. Dedup sources capture a - fixed snapshot and preserve the source table. They accept no - additional select, where or limit options. + source: str + The name of the source table. select: list or dict, optional The view's columns: column names, ``(alias, SQL expression)`` pairs, or a dict of the same. Omitting it selects every source @@ -1275,6 +1307,14 @@ class LanceDBConnection(DBConnection): def __repr__(self) -> str: return f"{self.__class__.__name__}(uri={self._conn.uri!r})" + @override + def is_open(self) -> bool: + return self._conn.is_open() + + @override + def close(self) -> None: + self._conn.close() + @override def serialize(self) -> str: import json @@ -1497,8 +1537,6 @@ class LanceDBConnection(DBConnection): """ if namespace_path is None: namespace_path = [] - if mode.lower() not in ["create", "overwrite"]: - raise ValueError("mode must be either 'create' or 'overwrite'") validate_table_name(name) tbl = LanceTable.create( @@ -1606,10 +1644,7 @@ class LanceDBConnection(DBConnection): Examples -------- >>> import lancedb - >>> db = lancedb.connect( - ... "./.lancedb", - ... storage_options={"new_table_enable_stable_row_ids": "true"}, - ... ) + >>> db = lancedb.connect("./.lancedb") >>> data = [{"name": "ada", "age": 36}, {"name": "kid", "age": 7}] >>> table = db.create_table("people", data) >>> view = db.create_materialized_view( @@ -2301,11 +2336,12 @@ class AsyncConnection(object): - pyarrow.Schema - [LanceModel][lancedb.pydantic.LanceModel] - mode: Literal["create", "overwrite"]; default "create" + mode: Literal["create", "overwrite", "exist_ok"]; default "create" The mode to use when creating the table. - Can be either "create" or "overwrite". + Can be "create", "overwrite", or "exist_ok". By default, if the table already exists, an exception is raised. If you want to overwrite the table, use mode="overwrite". + To open an existing table without adding data, use mode="exist_ok". exist_ok: bool, default False If a table by the same name already exists, then raise an exception if exist_ok=False. If exist_ok=True, then open the existing table; @@ -2576,19 +2612,7 @@ class AsyncConnection(object): where: Optional[str], limit: Optional[int], ) -> Tuple[str, Dict[str, str]]: - if not isinstance(source, VectorDedupSource): - return source, {} - if select is not None or where is not None or limit is not None: - raise ValueError( - "vector_dedup cannot be combined with select, where or limit" - ) - version = source.dataset_version - if version is None: - table = await self.open_table(source.source) - version = await table.version() - return source.source, { - "vector_source_json": source._native_source_json(version) - } + return source, {} async def create_materialized_view( self, diff --git a/python/python/lancedb/materialized_view.py b/python/python/lancedb/materialized_view.py index b4d455870..6342641ff 100644 --- a/python/python/lancedb/materialized_view.py +++ b/python/python/lancedb/materialized_view.py @@ -6,11 +6,11 @@ maintained by refresh. See ``DBConnection.create_materialized_view``.""" from __future__ import annotations -import json -import math -from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union +import json +from dataclasses import dataclass + from .background_loop import LOOP from .job import AsyncJob, Job, _typed_job @@ -21,7 +21,6 @@ if TYPE_CHECKING: from .table import AsyncTable, LanceTable DEFINITION_META_KEY = b"mv.definition" -_MAX_FLOAT32 = float.fromhex("0x1.fffffep+127") SelectArg = Union[ str, @@ -31,104 +30,12 @@ SelectArg = Union[ ] -@dataclass(frozen=True) -class VectorDedupSource: - """A declarative source of retained original rows from an indexed snapshot. - - Construct with [vector_dedup][lancedb.vector_dedup] and pass it as the - ``source`` of ``create_materialized_view`` or - ``create_materialized_view_async``. An omitted ``dataset_version`` is - captured once when creation is submitted, not when this object is made. - Refreshes keep that captured version. The source table is never modified. - """ - - source: str - column: str - distance_threshold: float - dataset_version: Optional[int] = None - - def __post_init__(self): - for name in ("source", "column"): - value = getattr(self, name) - if not isinstance(value, str) or not value: - raise ValueError(f"{name} must be a non-empty string") - if self.dataset_version is not None and ( - type(self.dataset_version) is not int - or not 0 < self.dataset_version < 2**64 - ): - raise ValueError("dataset_version must be a positive uint64 integer") - value = self.distance_threshold - if isinstance(value, bool) or not isinstance(value, (int, float)): - raise TypeError("distance_threshold must be a number") - try: - value = float(value) - except OverflowError as exc: - raise ValueError("distance_threshold must fit a finite float32") from exc - if not math.isfinite(value) or abs(value) > _MAX_FLOAT32: - raise ValueError("distance_threshold must fit a finite float32") - object.__setattr__(self, "distance_threshold", value) - - def _native_source_json(self, version: int) -> str: - return json.dumps( - { - "kind": "Dedup", - "dataset_version": version, - "column": self.column, - "distance_threshold": repr(self.distance_threshold), - }, - allow_nan=False, - ) +MaterializedViewSource = str -def vector_dedup( - source: str, - *, - column: str, - distance_threshold: float, - dataset_version: Optional[int] = None, -) -> VectorDedupSource: - """Declare an indexed dedup source for a materialized view. - - The default policy keeps the smallest unassigned snapshot row ID and removes - only its direct qualifying neighbors: an A-B-C chain without an A-C edge - retains A and C. Rows with no pair, including null vectors, survive. - Only same-segment/same-IVF-partition pairs are considered; quantized distances - use reconstructed index vectors. The index must cover the selected snapshot. - - ``source`` is an exact table name in the connection's root namespace, just - as for ``create_materialized_view``. ``dataset_version=None`` captures the - source's version when creation is submitted. Pass an explicit positive - version to reproduce a previous snapshot. The result contains original - source columns and rows. Do not combine this source with ``select``, - ``where`` or ``limit``; query the resulting table for further transforms. - - Creation follows the usual MV contract: by default it waits for the result. - Use ``with_no_data=True`` to declare only, or - ``create_materialized_view_async`` to get a job handle. On Cloud/Enterprise, - execution uses server-side MV registry jobs; local connections run in process. - - Examples - -------- - Given a connected ``db`` and an indexed ``images.phash`` column:: - - from lancedb import vector_dedup - - view = db.create_materialized_view( - "images_clean", - vector_dedup("images", column="phash", distance_threshold=4), - ) - cleaned = view.table - """ - return VectorDedupSource(source, column, distance_threshold, dataset_version) - - -MaterializedViewSource = Union[str, VectorDedupSource] - - -DEFINITION_FORMAT = 4 +DEFINITION_FORMAT = 2 """The newest stored layout this version reads: ``{"format": N, "query": ""}``, -format 2 being a query with ``GROUP BY``, format 3 native pairs and format 4 -native dedup results. A ``kind`` key beside it is for +format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for readers older than the format number.""" @@ -272,38 +179,31 @@ class AsyncMaterializedView: return _definition_from_json(raw) async def refresh( - self, *, full: bool = False, source_version: Optional[int] = None + self, *, source_version: Optional[int] = None ) -> "RefreshMaterializedViewResult": """Recompute the view from its source. - The refresh is incremental when the source's changes can be - reconciled into the view -- rows added, changed or removed since the - last one -- and otherwise rebuilds. ``full=True`` forces a rebuild; + Every refresh rebuilds the complete view. For local views, ``source_version`` refreshes to that source version instead of the - latest. - - Concurrent refreshes of one view do not duplicate its rows. Two that - plan the same source rows conflict on commit, and the loser raises - rather than writing them a second time. + latest; remote SQL refreshes do not support source-version pinning. """ return await self._table._inner.refresh_materialized_view( - full=full, source_version=source_version + source_version=source_version ) async def refresh_async( - self, *, full: bool = False, source_version: Optional[int] = None + self, *, source_version: Optional[int] = None ) -> "AsyncJob[RefreshMaterializedViewResult]": """Submit a refresh and return its job without waiting. - The job may already be complete for a local view. On LanceDB Cloud - and Enterprise, its ``id`` is the server job identifier returned by - the refresh endpoint. + The job may already be complete for a local view. Remote refreshes are + submitted through the SQL service. """ from ._lancedb import RefreshMaterializedViewResult return _typed_job( await self._table._inner.refresh_materialized_view_async( - full=full, source_version=source_version + source_version=source_version ), RefreshMaterializedViewResult.from_json, ) @@ -335,22 +235,18 @@ class MaterializedView: return LOOP.run(self._async.definition()) def refresh( - self, *, full: bool = False, source_version: Optional[int] = None + self, *, source_version: Optional[int] = None ) -> "RefreshMaterializedViewResult": """Recompute the view from its source. See [AsyncMaterializedView.refresh][lancedb.materialized_view.AsyncMaterializedView.refresh].""" - return LOOP.run(self._async.refresh(full=full, source_version=source_version)) + return LOOP.run(self._async.refresh(source_version=source_version)) def refresh_async( - self, *, full: bool = False, source_version: Optional[int] = None + self, *, source_version: Optional[int] = None ) -> "Job[RefreshMaterializedViewResult]": """Submit a refresh and return its job without waiting. See [AsyncMaterializedView.refresh_async][lancedb.materialized_view.AsyncMaterializedView.refresh_async]. """ - return Job( - LOOP.run( - self._async.refresh_async(full=full, source_version=source_version) - ) - ) + return Job(LOOP.run(self._async.refresh_async(source_version=source_version))) diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index cf1910c68..16be63e2d 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -489,6 +489,14 @@ class LanceNamespaceDBConnection(DBConnection): ) self._uri = self._inner.uri + @override + def is_open(self) -> bool: + return self._inner.is_open() + + @override + def close(self) -> None: + self._inner.close() + @override def serialize(self) -> str: import json diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index db213ac71..b210c6497 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -40,7 +40,7 @@ from .rerankers.base import Reranker from .rerankers.rrf import RRFReranker from .rerankers.util import check_reranker_result from .schema import is_blob_like_field, schema_has_blob_field -from .util import flatten_columns +from .util import _validate_query_vector, flatten_columns from . import _wal_hybrid # WAL-PK-FUSION: delete. from ._blob import ( BLOB_MODE_TO_HANDLING, @@ -907,6 +907,9 @@ class LanceQueryBuilder(ABC): fast_search: bool Skip flat search of unindexed data. """ + if query_type != "fts": + _validate_query_vector(query) + if ordering_field_name is not None: import warnings @@ -1619,6 +1622,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): str_query: Optional[str] = None, fast_search: bool = None, ): + _validate_query_vector(query) super().__init__(table) self._query = query self._distance_type = None @@ -1693,13 +1697,15 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): Parameters ---------- nprobes: int - The number of probes to use. + The number of probes to use. Must be greater than 0. Returns ------- LanceVectorQueryBuilder The LanceQueryBuilder object. """ + if nprobes <= 0: + raise ValueError("Invalid input, nprobes must be greater than 0") self._minimum_nprobes = nprobes self._maximum_nprobes = nprobes return self @@ -2495,13 +2501,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): Parameters ---------- nprobes: int - The number of probes to use. + The number of probes to use. Must be greater than 0. Returns ------- LanceHybridQueryBuilder The LanceHybridQueryBuilder object. """ + if nprobes <= 0: + raise ValueError("Invalid input, nprobes must be greater than 0") self._minimum_nprobes = nprobes self._maximum_nprobes = nprobes return self @@ -3654,6 +3662,8 @@ class AsyncVectorQueryBase: """ Set the number of partitions to search (probe) + The number of probes must be greater than 0. + This argument is only used when the vector column has an IVF-based index. If there is no index then this value is ignored. diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 3773ac547..52323daef 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -1054,6 +1054,10 @@ class RemoteDBConnection(DBConnection): """ return LOOP.run(self._conn.namespace_client()) - async def close(self): - """Close the connection to the database.""" + @override + def is_open(self) -> bool: + return self._conn.is_open() + + @override + def close(self) -> None: self._conn.close() diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 2bdd57fe9..11ca7de84 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2617,6 +2617,10 @@ class Table(ABC): to nullable. Currently, you cannot change a nullable column to non-nullable. + Each alteration must specify at least one of "rename", "data_type", + or "nullable". The legacy key "name" is also accepted for renaming. + Unknown keys raise ValueError before any alterations are applied. + Returns ------- AlterColumnsResult @@ -4332,9 +4336,9 @@ class LanceTable(Table): ) if storage_options is None: storage_options = {} - storage_options["new_table_enable_v2_manifest_paths"] = ( + storage_options["new_table_enable_v2_manifest_paths"] = str( enable_v2_manifest_paths - ) + ).lower() self._table = LOOP.run( self._conn._conn.create_table( @@ -4766,7 +4770,7 @@ class LanceTable(Table): [LanceTable.uses_v2_manifest_paths][lancedb.table.LanceTable.uses_v2_manifest_paths] to check if the table is already using the new path style. """ - LOOP.run(self._table.migrate_v2_manifest_paths()) + LOOP.run(self._table.migrate_manifest_paths_v2()) @deprecation.deprecated( deprecated_in="0.33.1", @@ -6107,7 +6111,17 @@ class AsyncTable: @overload async def search( self, - query: Optional[str] = None, + query: None = None, + vector_column_name: Optional[str] = None, + query_type: QueryType = "auto", + ordering_field_name: Optional[str] = None, + fts_columns: Optional[Union[str, List[str]]] = None, + ) -> AsyncQuery: ... + + @overload + async def search( + self, + query: str, vector_column_name: Optional[str] = None, query_type: Literal["auto"] = ..., ordering_field_name: Optional[str] = None, @@ -6117,7 +6131,7 @@ class AsyncTable: @overload async def search( self, - query: Optional[str] = None, + query: str, vector_column_name: Optional[str] = None, query_type: Literal["hybrid"] = ..., ordering_field_name: Optional[str] = None, @@ -6127,7 +6141,7 @@ class AsyncTable: @overload async def search( self, - query: Optional[Union[VEC, "PIL.Image.Image", Tuple]] = None, + query: Union[VEC, "PIL.Image.Image", Tuple], vector_column_name: Optional[str] = None, query_type: Literal["auto"] = ..., ordering_field_name: Optional[str] = None, @@ -6137,7 +6151,7 @@ class AsyncTable: @overload async def search( self, - query: Optional[str] = None, + query: str, vector_column_name: Optional[str] = None, query_type: Literal["fts"] = ..., ordering_field_name: Optional[str] = None, @@ -6147,9 +6161,7 @@ class AsyncTable: @overload async def search( self, - query: Optional[ - Union[VEC, str, "PIL.Image.Image", Tuple, FullTextQuery] - ] = None, + query: Union[VEC, str, "PIL.Image.Image", Tuple, FullTextQuery], vector_column_name: Optional[str] = None, query_type: Literal["vector"] = ..., ordering_field_name: Optional[str] = None, @@ -6165,7 +6177,7 @@ class AsyncTable: query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, - ) -> Union[AsyncHybridQuery, AsyncFTSQuery, AsyncVectorQuery]: + ) -> Union[AsyncQuery, AsyncHybridQuery, AsyncFTSQuery, AsyncVectorQuery]: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) and [full-text search](https://lancedb.com/docs/search/full-text-search/). @@ -6180,8 +6192,9 @@ class AsyncTable: - *default None*. Acceptable types are: list, np.ndarray, PIL.Image.Image - - If None then the select/where/limit clauses are applied to filter - the table + - If None then a plain [AsyncQuery][lancedb.query.AsyncQuery] is + returned, equivalent to calling [query][lancedb.table.AsyncTable.query]. + The select/where/limit clauses are applied to filter the table. vector_column_name: str, optional The name of the vector column to search. @@ -6209,10 +6222,13 @@ class AsyncTable: Returns ------- - LanceQueryBuilder + AsyncQuery, AsyncHybridQuery, AsyncFTSQuery, or AsyncVectorQuery A query builder object representing the query. """ + if query is None: + return self.query() + def is_embedding(query): return isinstance(query, (list, np.ndarray, pa.Array, pa.ChunkedArray)) @@ -6862,6 +6878,10 @@ class AsyncTable: to nullable. Currently, you cannot change a nullable column to non-nullable. + Each alteration must specify at least one of "rename", "data_type", + or "nullable". The legacy key "name" is also accepted for renaming. + Unknown keys raise ValueError before any alterations are applied. + Returns ------- AlterColumnsResult @@ -6905,9 +6925,11 @@ class AsyncTable: """ versions = await self._inner.list_versions() for v in versions: - ts_nanos = v["timestamp"] - v["timestamp"] = datetime.fromtimestamp(ts_nanos // 1e9) + timedelta( - microseconds=(ts_nanos % 1e9) // 1e3 + # Use integer math: float division on ~1e18 nanosecond + # values loses sub-millisecond precision. + seconds, nanos = divmod(v["timestamp"], 1_000_000_000) + v["timestamp"] = datetime.fromtimestamp(seconds) + timedelta( + microseconds=nanos // 1000 ) return versions diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index 3cf5460a5..3a6a8a102 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -280,6 +280,14 @@ def infer_vector_column_dim(data_type: pa.DataType) -> Optional[int]: return None +def _validate_query_vector(query: Any) -> None: + """Reject empty vector inputs before vector-column inference or execution.""" + if (isinstance(query, list) and not query) or ( + isinstance(query, np.ndarray) and query.size == 0 + ): + raise ValueError("Query vector must not be empty") + + def _query_vector_dim(query: Optional[Any]) -> Optional[int]: if query is None: return None @@ -301,6 +309,9 @@ def infer_vector_column_name( query: Optional[Any], # inferred later in query builder vector_column_name: Optional[str], ): + if query_type != "fts": + _validate_query_vector(query) + if vector_column_name is not None: return vector_column_name diff --git a/python/python/tests/test_catalog.py b/python/python/tests/test_catalog.py index ca5d68fdf..e7b7b1ec2 100644 --- a/python/python/tests/test_catalog.py +++ b/python/python/tests/test_catalog.py @@ -55,7 +55,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server): (200, {}), (200, {"tables": []}), (200, {"tables": []}), - (200, {"namespaces": ["team/search"], "page_token": "next"}), + (200, {"namespaces": ["team/search"]}), (204, None), ] ) @@ -85,13 +85,12 @@ def test_catalog_sync_scope_and_serialization(catalog_server): with pytest.raises(ValueError, match="sql_host_override must use"): connection.execute_query_async("SELECT 1") assert restored.table_names() == [] - page = catalog.list_databases(limit=1, page_token="a/b") - assert page == lancedb.ListDatabasesResponse(["team/search"], "next") + assert list(catalog.list_databases()) == ["team/search"] catalog.drop_database("team/search", ignore_missing=True) assert requests[0][0] == "/v1/namespace/team%2Fsearch/create" assert requests[0][2] == {"mode": "ExistOk"} assert requests[1][0] == "/v1/namespace/team%2Fsearch/describe" - assert requests[4][0] == "/v1/namespace/$/list?limit=1&page_token=a%2Fb" + assert requests[4][0] == "/v1/namespace/$/list" assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"} for i, (_, headers, _) in enumerate(requests): headers = {key.lower(): value for key, value in headers.items()} @@ -138,3 +137,129 @@ async def test_catalog_async_and_errors(catalog_server): def test_catalog_requires_remote_endpoint(endpoint): with pytest.raises(ValueError, match="endpoint"): lancedb.connect_catalog(endpoint) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +@pytest.mark.parametrize("page_token,page_limit", [(None, None), ("start/token", 2)]) +async def test_database_names_iteration( + catalog_server, asynchronous, page_token, page_limit +): + endpoint, requests, responses = catalog_server + responses.extend( + [ + (200, {"namespaces": ["first", "second"], "page_token": "a/b"}), + (200, {"namespaces": [], "page_token": "next"}), + (200, {"namespaces": ["last", "final"], "page_token": ""}), + ] + ) + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + names = catalog.list_databases(page_token=page_token, page_limit=page_limit) + assert requests == [] + assert names.num_page_results() == 0 + assert names.page_token() == page_token + for cached, expected in [(1, "first"), (0, "second")]: + assert (await names.__anext__() if asynchronous else next(names)) == expected + assert names.num_page_results() == cached + assert names.page_token() == "a/b" + assert len(requests) == 1 + assert (await names.__anext__() if asynchronous else next(names)) == "last" + assert names.num_page_results() == 1 + assert names.page_token() is None + assert (await names.__anext__() if asynchronous else next(names)) == "final" + assert names.num_page_results() == 0 + assert names.page_token() is None + for _ in range(2): + with pytest.raises(StopAsyncIteration if asynchronous else StopIteration): + if asynchronous: + await names.__anext__() + else: + next(names) + from urllib.parse import parse_qs, urlsplit + + for request, token in zip(requests, [page_token, "a/b", "next"]): + url = urlsplit(request[0]) + assert url.path == "/v1/namespace/$/list" + expected = {} + if token is not None: + expected["page_token"] = [token] + if page_limit is not None: + expected["limit"] = [str(page_limit)] + assert parse_qs(url.query) == expected + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +@pytest.mark.asyncio +async def test_database_names_empty_or_error(catalog_server, asynchronous, fail): + endpoint, requests, responses = catalog_server + responses.append( + (401, {"error": "unauthorized"}) if fail else (200, {"namespaces": []}) + ) + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + names = catalog.list_databases() + if fail: + with pytest.raises(HttpError): + if asynchronous: + await names.__anext__() + else: + next(names) + assert ([name async for name in names] if asynchronous else list(names)) == [] + assert len(requests) == 1 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_database_names_resume(catalog_server, asynchronous): + endpoint, requests, responses = catalog_server + responses.extend( + [ + (200, {"namespaces": ["first"], "page_token": "resume/token"}), + (200, {"namespaces": ["second"]}), + ] + ) + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + names = catalog.list_databases(page_limit=1) + assert (await names.__anext__() if asynchronous else next(names)) == "first" + assert names.num_page_results() == 0 + resumed = catalog.list_databases(page_token=names.page_token(), page_limit=1) + assert resumed.page_token() == "resume/token" + assert len(requests) == 1 + assert ([name async for name in resumed] if asynchronous else list(resumed)) == [ + "second" + ] + assert resumed.page_token() is None + assert resumed.num_page_results() == 0 + assert requests[1][0] == "/v1/namespace/$/list?limit=1&page_token=resume%2Ftoken" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("page_limit", [0, 2**32 - 1]) +@pytest.mark.asyncio +async def test_database_names_invalid_limit(catalog_server, asynchronous, page_limit): + endpoint, requests, _ = catalog_server + catalog = ( + await lancedb.connect_catalog_async(endpoint) + if asynchronous + else lancedb.connect_catalog(endpoint) + ) + names = catalog.list_databases(page_limit=page_limit) + with pytest.raises(ValueError, match="limit"): + if asynchronous: + await names.__anext__() + else: + next(names) + assert names.num_page_results() == 0 + assert requests == [] diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index aeb9feeb8..ebd86bdc3 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import contextlib import inspect import re import sys @@ -350,6 +351,40 @@ async def test_create_table_from_iterator_async(mem_db_async: lancedb.AsyncConne assert await table.count_rows() == 10 +def test_create_exist_ok_mode(tmp_db: lancedb.DBConnection): + db = tmp_db + data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}] + table = db.create_table("t1", data, mode="exist_ok") + assert table.count_rows() == 2 + original = table.to_arrow() + + reopened = db.create_table("t1", [{"vector": [9.0, 8.0]}], mode="exist_ok") + assert reopened.to_arrow().equals(original) + + +def test_create_empty_exist_ok_mode(tmp_db: lancedb.DBConnection): + schema = pa.schema([pa.field("id", pa.int64())]) + table = tmp_db.create_table("empty", schema=schema, mode="exist_ok") + assert table.schema == schema + assert table.count_rows() == 0 + table.add([{"id": 7}]) + + reopened = tmp_db.create_table("empty", schema=schema, mode="exist_ok") + assert reopened.to_arrow().to_pylist() == [{"id": 7}] + + bad_schema = pa.schema([pa.field("other", pa.int64())]) + with pytest.raises(ValueError): + tmp_db.create_table("empty", schema=bad_schema, mode="exist_ok") + + +def test_create_table_invalid_mode(tmp_db: lancedb.DBConnection): + db = tmp_db + data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}] + with pytest.raises(ValueError, match="Invalid mode bogus"): + db.create_table("tb", data, mode="bogus") + assert "tb" not in db + + def test_create_exist_ok(tmp_db: lancedb.DBConnection): from conftest import pandas_string_type @@ -405,6 +440,49 @@ async def test_connect(tmp_path): assert str(db) == f"ListingDatabase(uri={tmp_path}, read_consistency_interval=5s)" +@pytest.fixture(params=["memory", "local", "namespace", "from_inner"]) +def sync_db(request, tmp_path): + if request.param == "memory": + return lancedb.connect("memory://") + if request.param == "namespace": + return lancedb.connect_namespace("dir", {"root": str(tmp_path)}) + db = lancedb.connect(tmp_path) + if request.param == "from_inner": + return lancedb.db.LanceDBConnection.from_inner(db._inner, None) + return db + + +def test_sync_close(sync_db): + assert sync_db.is_open() + assert sync_db.close() is None + assert not sync_db.is_open() + assert sync_db.close() is None + + with pytest.warns(DeprecationWarning, match="table_names"): + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.table_names() + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.list_tables() + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.open_table("test") + + +@pytest.mark.parametrize("raise_error", [False, True]) +def test_sync_context_manager(sync_db, raise_error): + with contextlib.ExitStack() as stack: + if raise_error: + stack.enter_context(pytest.raises(ValueError, match="test error")) + with sync_db as db: + assert db is sync_db + assert db.is_open() + if raise_error: + raise ValueError("test error") + + assert not sync_db.is_open() + with pytest.raises(RuntimeError, match="Connection is closed"): + sync_db.list_tables() + + @pytest.mark.asyncio async def test_close(mem_db_async: lancedb.AsyncConnection): assert mem_db_async.is_open() @@ -495,6 +573,27 @@ async def test_create_exist_ok_async(tmp_db_async: lancedb.AsyncConnection): # await db.create_table("test", schema=bad_schema, exist_ok=True) +@pytest.mark.parametrize("enable_v2_manifest_paths", [False, True]) +@pytest.mark.parametrize("empty", [False, True]) +def test_create_table_deprecated_v2_manifest_paths( + tmp_path, enable_v2_manifest_paths, empty +): + db = lancedb.connect(tmp_path) + data = None if empty else [{"id": 1}] + + with pytest.warns(DeprecationWarning, match="enable_v2_manifest_paths"): + table = db.create_table( + "test", + data=data, + schema=pa.schema([("id", pa.int64())]), + enable_v2_manifest_paths=enable_v2_manifest_paths, + ) + + assert table.uses_v2_manifest_paths() == enable_v2_manifest_paths + assert table.to_arrow().to_pylist() == ([] if empty else data) + assert db.open_table("test").uses_v2_manifest_paths() == enable_v2_manifest_paths + + @pytest.mark.asyncio async def test_create_table_v2_manifest_paths_async(tmp_path): db_with_v2_paths = await lancedb.connect_async( diff --git a/python/python/tests/test_lsm_write_spec.py b/python/python/tests/test_lsm_write_spec.py index d43cb7532..ac767ea31 100644 --- a/python/python/tests/test_lsm_write_spec.py +++ b/python/python/tests/test_lsm_write_spec.py @@ -52,9 +52,9 @@ def test_set_lsm_write_spec_validates(tmp_path): with pytest.raises(Exception, match="num_buckets"): table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1025)) - # Happy path then mutation rejected. + # Happy path, then a set over the installed spec is refused. table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 4)) - with pytest.raises(Exception, match="mutation"): + with pytest.raises(Exception, match="already set"): table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 8)) @@ -176,18 +176,18 @@ def test_get_lsm_write_spec(tmp_path): table.unset_lsm_write_spec() assert table.get_lsm_write_spec() is None - # Identity round-trips (column recovered from the schema). Leaving the - # maintained set to be inferred picks up the index on the table, so the - # spec reads back naming it rather than as "infer". + # Identity round-trips (column recovered from the schema). An unnamed + # maintained set reads back as None: what is stored is the intent to + # maintain every index, not the set it resolves to now. table.set_lsm_write_spec(LsmWriteSpec.identity("id")) spec = table.get_lsm_write_spec() assert spec.spec_type == "identity" assert spec.column == "id" - assert spec.maintained_indexes == [idx_name] + assert spec.maintained_indexes is None table.unset_lsm_write_spec() # Unsharded round-trips (no routing column). Opting out is distinct from - # the inferred default. + # the unnamed default. table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([])) spec = table.get_lsm_write_spec() assert spec.spec_type == "unsharded" diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index 1ee21fbac..bef173323 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -97,24 +97,6 @@ def test_remote_list_uses_namespace_route(): assert requests == ["/v1/namespace/$/materialized_view/list"] -def test_remote_create_async_returns_server_job(): - with mock_remote_materialized_view_create() as (host, requests): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - client_config={"retry_config": {"retries": 0}}, - ) - job = db.create_materialized_view_async("adults", "people", where="age >= 18") - assert job.id == "mv-create-123" - assert requests == [ - ( - "/v1/materialized_view/adults/create", - {"query": 'SELECT * FROM "people" WHERE age >= 18', "with_no_data": False}, - ) - ] - - def test_remote_drop_async_returns_server_job(): with mock_remote_materialized_view_create() as (host, requests): db = lancedb.connect( @@ -221,29 +203,29 @@ def test_definition_round_trips(tmp_path): ) -def test_incremental_refresh_after_append(tmp_path): +def test_full_refresh_after_append(tmp_path): db = make_db(tmp_path) view = db.create_materialized_view("copy", "people", with_no_data=True) view.refresh() db.open_table("people").add([{"name": "alan", "age": 41}]) result = view.refresh() - assert result.mode == "incremental" - assert result.rows_written == 1 + assert result.mode == "rebuild" + assert result.rows_written == 4 assert view.table.count_rows() == 4 - assert view.refresh().mode == "no_op" + assert view.refresh().mode == "rebuild" -def test_incremental_refresh_after_update(tmp_path): +def test_full_refresh_after_update(tmp_path): db = make_db(tmp_path) view = db.create_materialized_view("copy", "people", with_no_data=True) view.refresh() db.open_table("people").update(where="name = 'kid'", values={"age": 8}) result = view.refresh() - assert result.mode == "incremental" - assert result.rows_written == 1 + assert result.mode == "rebuild" + assert result.rows_written == 3 rows = view.table.search().to_list() assert sorted(row["age"] for row in rows) == [8, 36, 85] @@ -331,8 +313,8 @@ async def test_async_create_and_refresh_jobs(tmp_path): @pytest.mark.asyncio -async def test_async_incremental(tmp_path): - db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS) +async def test_async_refresh_rebuilds(tmp_path): + db = await lancedb.connect_async(tmp_path) await db.create_table("people", [{"name": "ada", "age": 36}]) view = await db.create_materialized_view("copy", "people", with_no_data=True) await view.refresh() @@ -340,15 +322,15 @@ async def test_async_incremental(tmp_path): table = await db.open_table("people") await table.add([{"name": "alan", "age": 41}]) result = await view.refresh() - assert result.mode == "incremental" - assert result.rows_written == 1 + assert result.mode == "rebuild" + assert result.rows_written == 2 -def test_source_requires_stable_row_ids(tmp_path): +def test_source_does_not_require_stable_row_ids(tmp_path): db = lancedb.connect(tmp_path) db.create_table("plain", [{"x": 1}]) - with pytest.raises(Exception, match="stable row ids"): - db.create_materialized_view("v", "plain") + view = db.create_materialized_view("v", "plain") + assert view.table.count_rows() == 1 def test_bare_select_names_are_quoted(tmp_path): @@ -373,7 +355,7 @@ def test_scalar_select_is_one_column(tmp_path): view = db.create_materialized_view("just_name", "people", select="name") view.refresh() rows = view.table.search().to_list() - assert set(rows[0]) - {"__source_row_id"} == {"name"} + assert set(rows[0]) == {"name"} assert sorted(row["name"] for row in rows) == ["ada", "grace", "kid"] @@ -384,7 +366,7 @@ async def test_async_scalar_select_is_one_column(tmp_path): view = await db.create_materialized_view("just_name", "people", select="name") await view.refresh() rows = await view.table.query().to_list() - assert set(rows[0]) - {"__source_row_id"} == {"name"} + assert set(rows[0]) == {"name"} def test_limit_above_i64_max_is_refused(tmp_path): @@ -520,166 +502,10 @@ def test_stored_queries_and_legacy_layouts_are_read(): "SELECT * FROM people" ) - for fmt, function in [(3, "vector_duplicate_pairs"), (4, "vector_dedup")]: - native = f"SELECT * FROM {function}('images', 3, 'phash', 4)" - assert read({"format": fmt, "query": native}).query == native - # A newer writer's layout is reported, never guessed at. for newer in ( - {"format": 5, "query": query}, + {"format": 3, "query": query}, {"kind": "select_v3", "source_table": "people"}, ): with pytest.raises(NotImplementedError, match="cannot refresh"): read(newer) - - -def indexed_dedup_source(db, name="images"): - import pyarrow as pa - from lancedb.index import IvfFlat - - vectors = pa.FixedSizeListArray.from_arrays( - pa.array( - [v for i in range(256) for v in (float(i if i < 3 else i * 10), 0.0)], - type=pa.float32(), - ), - 2, - ) - table = db.create_table(name, pa.table({"id": range(256), "vector": vectors})) - table.create_index("vector", config=IvfFlat(num_partitions=1)) - return table - - -@pytest.mark.parametrize("stable", [False, True]) -def test_vector_dedup_source_captures_snapshot_and_preserves_source(tmp_path, stable): - db = lancedb.connect( - tmp_path, - storage_options={"new_table_enable_stable_row_ids": str(stable).lower()}, - ) - source = indexed_dedup_source(db, "images.with.dot") - version = source.version - expression = lancedb.vector_dedup( - source.name, column="vector", distance_threshold=1 - ) - view = db.create_materialized_view("clean", expression) - assert "vector_dedup(" in view.definition.query - assert f", {version}, 'vector', 1.0)" in view.definition.query - assert view.table.count_rows() == 255 - assert [r["id"] for r in view.table.search().where("id < 3").to_list()] == [0, 2] - assert source.version == version - assert source.count_rows() == 256 - - source.delete("id = 0") - view.refresh_async().wait() - assert view.table.search().where("id = 0").to_list() == [ - {"id": 0, "vector": [0.0, 0.0]} - ] - assert expression.dataset_version is None # reusable immutable declaration - - -def test_vector_dedup_create_job_and_explicit_snapshot(tmp_path): - db = lancedb.connect(tmp_path) - source = indexed_dedup_source(db) - version = source.version - source.delete("id = 0") - job = db.create_materialized_view_async( - "clean", - lancedb.vector_dedup( - "images", column="vector", distance_threshold=1, dataset_version=version - ), - ) - assert job.wait() is None - view = db.open_materialized_view("clean") - assert view.table.count_rows() == 255 - assert len(view.table.search().where("id = 0").to_list()) == 1 - - -@pytest.mark.asyncio -async def test_async_vector_dedup_declaration_and_creation(tmp_path): - sync_db = lancedb.connect(tmp_path) - source = indexed_dedup_source(sync_db) - db = await lancedb.connect_async(tmp_path) - expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1) - view = await db.create_materialized_view("clean", expression, with_no_data=True) - assert await view.table.count_rows() == 0 - # The snapshot is captured at declaration even if refresh comes later. - source.delete("id = 0") - job = await view.refresh_async() - await job.wait() - assert await view.table.count_rows() == 255 - another = await db.create_materialized_view_async( - "later", expression, with_no_data=True - ) - await another.wait() - later = await db.open_materialized_view("later") - assert (await later.definition()).query != (await view.definition()).query - assert await later.table.count_rows() == 0 - - -@pytest.mark.parametrize( - "options", [{"select": ["id"]}, {"select": []}, {"where": "id > 0"}, {"limit": 1}] -) -def test_vector_dedup_rejects_extra_clauses_before_opening_source(tmp_path, options): - db = lancedb.connect(tmp_path) - expression = lancedb.vector_dedup("missing", column="vector", distance_threshold=1) - with pytest.raises(ValueError, match="cannot be combined"): - db.create_materialized_view("bad", expression, **options) - assert db.list_tables().tables == [] - - -@pytest.mark.parametrize("version", [0, -1, True, 1.5, "2", 2**64]) -def test_vector_dedup_rejects_invalid_snapshot(version): - with pytest.raises(ValueError, match="dataset_version"): - lancedb.vector_dedup( - "images", column="vector", distance_threshold=1, dataset_version=version - ) - - -@pytest.mark.parametrize( - "threshold", [float("nan"), float("inf"), -float("inf"), 1e100, 10**1000, True, "4"] -) -def test_vector_dedup_rejects_invalid_threshold(threshold): - with pytest.raises((ValueError, TypeError), match="distance_threshold"): - lancedb.vector_dedup("images", column="vector", distance_threshold=threshold) - - -def test_vector_dedup_descriptor_is_frozen(): - from dataclasses import FrozenInstanceError - - expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1) - with pytest.raises(FrozenInstanceError): - expression.column = "other" - - -@pytest.mark.parametrize("with_no_data", [False, True]) -def test_remote_vector_dedup_uses_mv_endpoint_without_sql_client(with_no_data): - with mock_remote_materialized_view_create() as (host, requests): - db = lancedb.connect( - "db://dev", - api_key="fake", - host_override=host, - sql_host_override="grpc://127.0.0.1:1", - client_config={"retry_config": {"retries": 0}}, - ) - job = db.create_materialized_view_async( - "clean", - lancedb.vector_dedup( - "images.with'dot", - column="vector'field", - distance_threshold=4, - dataset_version=7, - ), - with_no_data=with_no_data, - ) - assert job.id == "mv-create-123" - assert requests == [ - ( - "/v1/materialized_view/clean/create", - { - "query": ( - "SELECT * FROM vector_dedup('`images.with''dot`', " - "7, 'vector''field', 4.0)" - ), - "with_no_data": with_no_data, - }, - ) - ] diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index e3d55917d..d3fedbcc8 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -24,6 +24,7 @@ from lancedb.pydantic import LanceModel, Vector from lancedb.query import ( AsyncFTSQuery, AsyncHybridQuery, + AsyncQuery, AsyncQueryBase, AsyncVectorQuery, ColumnOrdering, @@ -974,6 +975,19 @@ def test_query_builder_with_filter(table): assert all(np.array(rs[0]["vector"]) == [3, 4]) +@pytest.mark.parametrize("query_type", ["vector", "hybrid"]) +@pytest.mark.parametrize("nprobes", [0, -1]) +def test_nprobes_nonpositive_sync(table, query_type, nprobes): + if query_type == "hybrid": + query = table.search(query_type="hybrid").vector([0, 0]).text("a") + else: + query = table.search([0, 0]) + with pytest.raises( + ValueError, match="^Invalid input, nprobes must be greater than 0$" + ): + query.nprobes(nprobes) + + def test_invalid_nprobes_sync(table): with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"): LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(0).to_list() @@ -1005,6 +1019,18 @@ def test_multiple_nprobes_calls_works_sync(table): ).minimum_nprobes(20).to_list() +@pytest.mark.asyncio +@pytest.mark.parametrize("hybrid", [False, True]) +async def test_nprobes_zero_async(table_async: AsyncTable, hybrid): + query = table_async.query().nearest_to([0, 0]) + if hybrid: + query = query.nearest_to_text("dog") + with pytest.raises( + ValueError, match="^Invalid input, nprobes must be greater than 0$" + ): + query.nprobes(0) + + @pytest.mark.asyncio async def test_invalid_nprobes_async(table_async: AsyncTable): with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"): @@ -1279,6 +1305,27 @@ async def test_query_to_polars_async(table_async: AsyncTable): assert df.shape == (0, num_columns) +@pytest.mark.asyncio +@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}]) +@pytest.mark.parametrize("with_vector", [False, True]) +async def test_async_search_without_query(mem_db_async, search_kwargs, with_vector): + data = [{"id": i} for i in range(10)] + if with_vector: + for row in data: + row["vector"] = [float(row["id"]), 1.0] + table = await mem_db_async.create_table("test", data) + + query = await table.search(**search_kwargs) + assert isinstance(query, AsyncQuery) + assert await query.limit(3).to_arrow() == await table.query().limit(3).to_arrow() + + query = await table.search(**search_kwargs) + assert await query.where("id >= 8").select(["id"]).limit(3).to_list() == [ + {"id": 8}, + {"id": 9}, + ] + + @pytest.mark.asyncio async def test_none_query(table_async: AsyncTable): with pytest.raises(ValueError): @@ -2091,6 +2138,12 @@ def test_ensure_vector_query_nested_empty_list(): ensure_vector_query([[]]) +@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)]) +def test_vector_query_builder_empty_vector(table, query): + with pytest.raises(ValueError, match="^Query vector must not be empty$"): + LanceVectorQueryBuilder(table, query, "vector") + + def test_fast_search(tmp_path): db = lancedb.connect(tmp_path) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index cd4b78a1b..a9e411a3b 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -16,8 +16,9 @@ import uuid from packaging.version import Version import lancedb +import numpy as np from lancedb.conftest import MockTextEmbeddingFunction -from lancedb.query import ColumnOrdering +from lancedb.query import AsyncQuery, ColumnOrdering from lancedb.remote import ClientConfig from lancedb.remote.errors import HttpError, RetryError import pytest @@ -119,6 +120,77 @@ async def test_async_remote_db(): assert table_names == [] +@pytest.mark.parametrize( + "alteration, match", + [ + ({"path": "id"}, "One of rename, nullable or data_type"), + ( + {"path": "id", "nulable": False}, # spellchecker:disable-line + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ( + { + "path": "id", + "rename": "new_id", + "nulable": False, # spellchecker:disable-line + }, + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ], +) +def test_remote_alter_columns_rejects_invalid_before_request(alteration, match): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + with mock_lancedb_connection(handler) as db: + table = db.open_table("test") + requests.clear() + + with pytest.raises(ValueError, match=match): + table.alter_columns(alteration) + + assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "alteration, match", + [ + ({"path": "id"}, "One of rename, nullable or data_type"), + ( + {"path": "id", "nulable": False}, # spellchecker:disable-line + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ], +) +async def test_async_remote_alter_columns_rejects_invalid_before_request( + alteration, match +): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + async with mock_lancedb_connection_async(handler) as db: + table = await db.open_table("test") + requests.clear() + + with pytest.raises(ValueError, match=match): + await table.alter_columns(alteration) + + assert requests == [] + + @pytest.mark.asyncio async def test_async_checkout(): def handler(request): @@ -1521,6 +1593,60 @@ def test_query_sync_empty_query(): assert data == expected +@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)]) +@pytest.mark.parametrize("vector_column_name", [None, "vector"]) +@pytest.mark.parametrize("query_type", ["auto", "vector"]) +def test_query_sync_empty_vector(query, vector_column_name, query_type): + def handler(body): + pytest.fail("An empty query vector must be rejected before sending a query") + + with query_test_table(handler) as table: + with pytest.raises(ValueError, match="^Query vector must not be empty$"): + table.search( + query, vector_column_name=vector_column_name, query_type=query_type + ).limit(3).to_arrow() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("search_kwargs", [{}, {"query": None}]) +async def test_async_search_without_query(search_kwargs): + expected = pa.table({"id": [7, 8, 9]}) + + def handler(request): + if request.path == "/v1/table/test/describe/": + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + elif request.path == "/v1/table/test/query/": + body = json.loads( + request.rfile.read(int(request.headers["Content-Length"])) + ) + assert body == { + "k": 3, + "filter": "id >= 7", + "vector": [], + "columns": ["id"], + "prefilter": True, + "version": None, + } + request.send_response(200) + request.send_header("Content-Type", "application/vnd.apache.arrow.file") + request.end_headers() + with pa.ipc.new_file(request.wfile, schema=expected.schema) as writer: + writer.write_table(expected) + else: + request.send_response(404) + request.end_headers() + + async with mock_lancedb_connection_async(handler) as db: + table = await db.open_table("test") + query = await table.search(**search_kwargs) + assert isinstance(query, AsyncQuery) + result = await query.where("id >= 7").select(["id"]).limit(3).to_arrow() + assert result == expected + + def test_query_sync_maximal(): def handler(body): assert body == { @@ -1581,6 +1707,53 @@ def test_query_sync_maximal(): ) +@pytest.mark.parametrize("hybrid", [False, True]) +def test_query_sync_nprobes_zero(hybrid): + query_requests = [] + + def handler(body): + query_requests.append(body) + return pa.table({"id": []}) + + with query_test_table(handler) as table: + if hybrid: + query = table.search(query_type="hybrid").vector([1, 2, 3]).text("dog") + else: + query = table.search([1, 2, 3]) + with pytest.raises( + ValueError, match="^Invalid input, nprobes must be greater than 0$" + ): + query.nprobes(0) + + assert query_requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hybrid", [False, True]) +async def test_query_async_nprobes_zero(hybrid): + requests = [] + + def handler(request): + requests.append(request.path) + if request.path == "/v1/table/test/describe/": + send_json(request, {"version": 1, "schema": {"fields": []}}) + else: + request.send_response(404) + request.end_headers() + + async with mock_lancedb_connection_async(handler) as db: + table = await db.open_table("test") + query = table.query().nearest_to([1, 2, 3]) + if hybrid: + query = query.nearest_to_text("dog") + with pytest.raises( + ValueError, match="^Invalid input, nprobes must be greater than 0$" + ): + query.nprobes(0) + + assert requests == ["/v1/table/test/describe/"] + + def test_query_sync_nprobes(): def handler(body): assert body == { @@ -2145,15 +2318,46 @@ async def test_header_provider_overrides_static_headers(): def test_close(): - """Test that close() works without AttributeError.""" - import asyncio - def handler(req): req.send_response(200) req.end_headers() with mock_lancedb_connection(handler) as db: - asyncio.run(db.close()) + assert db.close() is None + assert not db.is_open() + assert db.close() is None + + with pytest.warns(DeprecationWarning, match="table_names"): + with pytest.raises(RuntimeError, match="Connection is closed"): + db.table_names() + with pytest.raises(RuntimeError, match="Connection is closed"): + db.list_tables() + with pytest.raises(RuntimeError, match="Connection is closed"): + db.open_table("test") + + +@pytest.mark.parametrize("raise_error", [False, True]) +def test_sync_context_manager(raise_error): + def handler(req): + req.send_response(200) + req.send_header("Content-Type", "application/json") + req.end_headers() + req.wfile.write(b'{"tables": []}') + + with mock_lancedb_connection(handler) as db: + with contextlib.ExitStack() as stack: + if raise_error: + stack.enter_context(pytest.raises(ValueError, match="test error")) + with db as conn: + assert conn is db + assert conn.is_open() + assert conn.list_tables().tables == [] + if raise_error: + raise ValueError("test error") + + assert not db.is_open() + with pytest.raises(RuntimeError, match="Connection is closed"): + db.list_tables() @pytest.mark.parametrize("exception", [KeyboardInterrupt, SystemExit, GeneratorExit]) @@ -2391,6 +2595,27 @@ def blob_query_response_table(): ) +def blob_descriptor_take_table(): + """What `fetch_blob_files` reads for rows 10, 20, 30, and 40.""" + queried = blob_query_response_table() + image_field = queried.schema.field("image") + images = pa.StructArray.from_arrays( + [ + pa.array([0, 0, 0, 0], type=pa.uint8()), + pa.array([0, 0, 0, 0], type=pa.uint64()), + pa.array([5, 0, 5, 0], type=pa.uint64()), + pa.array([0, 0, 0, 0], type=pa.uint32()), + pa.array(["", "", "", ""], type=pa.string()), + ], + fields=image_field.type, + mask=pa.array([False, True, False, False]), + ) + return pa.Table.from_arrays( + [images, pa.array([10, 20, 30, 40], type=pa.uint64())], + schema=pa.schema([image_field, pa.field("_rowid", pa.uint64())]), + ) + + @contextlib.contextmanager def blob_remote_table(*, server_version=Version("0.5.0")): def handler(request): @@ -2429,9 +2654,14 @@ def blob_remote_table(*, server_version=Version("0.5.0")): elif request.path == "/v1/table/test/query/": content_len = int(request.headers.get("Content-Length", 0)) body = json.loads(request.rfile.read(content_len)) - assert body["columns"] == ["id", "image"] assert body["with_row_id"] is True - response_table = blob_query_response_table() + if body["columns"] == ["image"]: + # fetch_blob_files sizes its handles from a descriptor take. + assert body["filter"].startswith("_rowid IN") + response_table = blob_descriptor_take_table() + else: + assert body["columns"] == ["id", "image"] + response_table = blob_query_response_table() request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.file") request.end_headers() diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index fa1a494b0..64f4f4ce9 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -113,6 +113,32 @@ def test_basic(mem_db: DBConnection): assert table.to_arrow() == expected_data +@pytest.mark.parametrize("enable_v2", [False, True]) +def test_migrate_v2_manifest_paths(tmp_path, enable_v2): + db = lancedb.connect( + tmp_path, + storage_options={"new_table_enable_v2_manifest_paths": str(enable_v2).lower()}, + ) + table = db.create_table("calls", [{"id": 1, "vector": [1.0, 1.0]}]) + table.add([{"id": 2, "vector": [2.0, 2.0]}]) + expected_data = table.to_arrow() + expected_versions = table.list_versions() + assert table.uses_v2_manifest_paths() == enable_v2 + + # Migration is also safe to repeat on a table already using v2 paths. + for _ in range(2): + table.migrate_v2_manifest_paths() + assert table.uses_v2_manifest_paths() + reopened = db.open_table("calls") + assert reopened.uses_v2_manifest_paths() + assert reopened.to_arrow() == expected_data + assert reopened.list_versions() == expected_versions + + manifests = list((tmp_path / "calls.lance" / "_versions").glob("*.manifest")) + assert len(manifests) == len(expected_versions) + assert all(len(path.stem) == 20 and path.stem.isdigit() for path in manifests) + + def test_search_preserves_nulls_from_sliced_arrow_table(mem_db: DBConnection): data = pa.table( { @@ -967,6 +993,63 @@ def test_polars(mem_db: DBConnection): assert len(filtered_result) == 2 +@pytest.mark.asyncio +async def test_list_versions_timestamp_precision(): + # 2026-10-03T01:03:49.274Z in nanoseconds. Float math turns .274000 into + # .273999, so make sure the conversion is exact. + ts_nanos = 1790989429274000000 + + class FakeInner: + async def list_versions(self): + return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}] + + table = table_module.AsyncTable(FakeInner()) + versions = await table.list_versions() + + expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000) + assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000) + assert versions[0]["timestamp"].microsecond == 274000 + + +def test_list_versions_timestamp_precision_sync(mem_db: DBConnection): + # The sync LanceTable.list_versions delegates to AsyncTable.list_versions, + # which is the path the parity report hit. + ts_nanos = 1790989429274000000 + + class FakeInner: + async def list_versions(self): + return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}] + + table = mem_db.create_table("ts_precision", data=[{"id": 1}]) + table._table = table_module.AsyncTable(FakeInner()) + versions = table.list_versions() + + expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000) + assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000) + assert versions[0]["timestamp"].microsecond == 274000 + + +def test_list_versions_timestamp_precision_remote(): + # RemoteTable.list_versions also delegates to AsyncTable.list_versions. + from lancedb.remote.table import RemoteTable + + ts_nanos = 1790989429274000000 + + class FakeInner: + def name(self): + return "ts_precision" + + async def list_versions(self): + return [{"version": 1, "timestamp": ts_nanos, "metadata": {}}] + + table = RemoteTable(table_module.AsyncTable(FakeInner()), "dev") + versions = table.list_versions() + + expected = datetime.fromtimestamp(ts_nanos // 1_000_000_000) + assert versions[0]["timestamp"] == expected + timedelta(microseconds=274000) + assert versions[0]["timestamp"].microsecond == 274000 + + def test_versioning(mem_db: DBConnection): table = mem_db.create_table( "test", @@ -3751,6 +3834,24 @@ def test_empty_query(mem_db: DBConnection): assert df.num_rows == 42 +@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)]) +@pytest.mark.parametrize("vector_column_name", [None, "vector"]) +@pytest.mark.parametrize("query_type", ["auto", "vector"]) +@pytest.mark.parametrize("multiple_vector_columns", [False, True]) +def test_search_empty_vector( + mem_db, query, vector_column_name, query_type, multiple_vector_columns +): + fields = [pa.field("vector", pa.list_(pa.float32(), 8))] + if multiple_vector_columns: + fields.append(pa.field("vec2", pa.list_(pa.float32(), 4))) + table = mem_db.create_table("empty_vector_query", schema=pa.schema(fields)) + + with pytest.raises(ValueError, match="^Query vector must not be empty$"): + table.search( + query, vector_column_name=vector_column_name, query_type=query_type + ).limit(3).to_arrow() + + def test_search_with_schema_inf_single_vector(mem_db: DBConnection): class MyTable(LanceModel): text: str @@ -4144,14 +4245,74 @@ async def test_add_columns_with_schema(mem_db_async: AsyncConnection): ) -def test_alter_columns(mem_db: DBConnection): +@pytest.mark.parametrize("rename_key", ["rename", "name"]) +def test_alter_columns(mem_db: DBConnection, rename_key): data = pa.table({"id": [0, 1]}) table = mem_db.create_table("my_table", data=data) - alter_columns_res = table.alter_columns({"path": "id", "rename": "new_id"}) + alter_columns_res = table.alter_columns({"path": "id", rename_key: "new_id"}) assert alter_columns_res.version == 2 assert table.to_arrow().column_names == ["new_id"] +INVALID_COLUMN_ALTERATIONS = [ + (({"path": "id"},), "One of rename, nullable or data_type"), + ( + ({"path": "id", "nulable": False},), # spellchecker:disable-line + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ( + ( + { + "path": "id", + "rename": "new_id", + "nulable": False, # spellchecker:disable-line + }, + ), + "Unknown column alteration key 'nulable'", # spellchecker:disable-line + ), + ( + ({"path": "id", "rename": "new_id"}, {"path": "id"}), + "One of rename, nullable or data_type", + ), +] + + +def test_alter_columns_nullable_false(mem_db: DBConnection): + table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]})) + result = table.alter_columns({"path": "id", "nullable": False}) + assert result.version == 2 + assert not table.schema.field("id").nullable + + +@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS) +def test_alter_columns_rejects_invalid(mem_db: DBConnection, alterations, match): + table = mem_db.create_table("my_table", data=pa.table({"id": [0, 1]})) + initial_version = table.version + initial_schema = table.schema + + with pytest.raises(ValueError, match=match): + table.alter_columns(*alterations) + + assert table.version == initial_version + assert table.schema == initial_schema + + +@pytest.mark.asyncio +@pytest.mark.parametrize("alterations, match", INVALID_COLUMN_ALTERATIONS) +async def test_alter_columns_rejects_invalid_async( + mem_db_async: AsyncConnection, alterations, match +): + table = await mem_db_async.create_table("my_table", data=pa.table({"id": [0, 1]})) + initial_version = await table.version() + initial_schema = await table.schema() + + with pytest.raises(ValueError, match=match): + await table.alter_columns(*alterations) + + assert await table.version() == initial_version + assert await table.schema() == initial_schema + + def test_update_field_metadata(mem_db: DBConnection): data = pa.table({"id": [0, 1], "category": ["a", "b"]}) table = mem_db.create_table("my_table", data=data) @@ -4177,10 +4338,11 @@ def test_update_field_metadata(mem_db: DBConnection): @pytest.mark.asyncio -async def test_alter_columns_async(mem_db_async: AsyncConnection): +@pytest.mark.parametrize("rename_key", ["rename", "name"]) +async def test_alter_columns_async(mem_db_async: AsyncConnection, rename_key): data = pa.table({"id": [0, 1]}) table = await mem_db_async.create_table("my_table", data=data) - alter_columns_res = await table.alter_columns({"path": "id", "rename": "new_id"}) + alter_columns_res = await table.alter_columns({"path": "id", rename_key: "new_id"}) assert alter_columns_res.version == 2 assert (await table.to_arrow()).column_names == ["new_id"] alter_columns_res = await table.alter_columns( diff --git a/python/src/catalog.rs b/python/src/catalog.rs index f61b409ff..e7cc56c49 100644 --- a/python/src/catalog.rs +++ b/python/src/catalog.rs @@ -1,12 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use std::sync::{Arc, Mutex}; use std::time::Duration; -use lancedb::catalog::{ - CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest, -}; -use pyo3::exceptions::PyValueError; +use futures::{StreamExt, future::poll_fn}; + +use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest}; +use pyo3::exceptions::{PyStopAsyncIteration, PyValueError}; use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods}; use crate::connection::{Connection, PyClientConfig}; @@ -67,19 +68,48 @@ impl Catalog { }) } - #[pyo3(signature = (*, limit=None, page_token=None))] - fn list_databases<'py>( - self_: PyRef<'py, Self>, - limit: Option, - page_token: Option, - ) -> PyResult> { + #[pyo3(signature = (*, page_token=None, page_limit=None))] + fn list_databases(&self, page_token: Option, page_limit: Option) -> DatabaseNames { + DatabaseNames { + inner: Arc::new(Mutex::new( + self.inner.list_databases(page_token, page_limit), + )), + next_lock: Arc::new(tokio::sync::Mutex::new(())), + } + } +} + +#[pyclass] +pub struct DatabaseNames { + inner: Arc>, + next_lock: Arc>, +} + +#[pymethods] +impl DatabaseNames { + fn num_page_results(&self, py: Python<'_>) -> usize { + py.detach(|| self.inner.lock().unwrap().num_page_results()) + } + + fn page_token(&self, py: Python<'_>) -> Option { + py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned)) + } + + fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> { + self_ + } + + fn __anext__(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); + let next_lock = self_.next_lock.clone(); future_into_py(self_.py(), async move { - let mut request = ListDatabasesRequest::default(); - request.limit = limit; - request.page_token = page_token; - let response = inner.list_databases(request).await.infer_error()?; - Ok((response.databases, response.page_token)) + // Serialize advances while allowing synchronous state inspection + // between polls, including while a REST request is pending. + let _guard = next_lock.lock().await; + poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx)) + .await + .ok_or_else(|| PyStopAsyncIteration::new_err(""))? + .infer_error() }) } } diff --git a/python/src/connection.rs b/python/src/connection.rs index 744ab34b3..737ef05e2 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -407,8 +407,7 @@ impl Connection { }) } - #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))] + #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))] pub fn create_materialized_view( self_: PyRef<'_, Self>, name: String, @@ -417,17 +416,10 @@ impl Connection { filter: Option, limit: Option, with_no_data: bool, - vector_source_json: Option, ) -> PyResult> { let inner = self_.get_inner()?.clone(); future_into_py(self_.py(), async move { let mut builder = inner.create_materialized_view(name, source); - if let Some(source) = vector_source_json { - builder = builder.vector_source( - serde_json::from_str(&source) - .map_err(|e| PyValueError::new_err(e.to_string()))?, - ); - } if let Some(projections) = projections { builder = builder.select(projections); } @@ -443,8 +435,7 @@ impl Connection { }) } - #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))] + #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))] pub fn create_materialized_view_async( self_: PyRef<'_, Self>, name: String, @@ -453,17 +444,10 @@ impl Connection { filter: Option, limit: Option, with_no_data: bool, - vector_source_json: Option, ) -> PyResult> { let inner = self_.get_inner()?.clone(); future_into_py(self_.py(), async move { let mut builder = inner.create_materialized_view(name, source); - if let Some(source) = vector_source_json { - builder = builder.vector_source( - serde_json::from_str(&source) - .map_err(|e| PyValueError::new_err(e.to_string()))?, - ); - } if let Some(projections) = projections { builder = builder.select(projections); } diff --git a/python/src/lib.rs b/python/src/lib.rs index ab04663c0..57d30d183 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -64,6 +64,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { env_logger::init_from_env(env); 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 3a4737ee0..1079fb314 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -995,8 +995,9 @@ impl VectorQuery { self.inner = self.inner.clone().refine_factor(refine_factor); } - pub fn nprobes(&mut self, nprobe: u32) { - self.inner = self.inner.clone().nprobes(nprobe as usize); + pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> { + self.inner = self.inner.clone().nprobes(nprobe as usize).infer_error()?; + Ok(()) } pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> { @@ -1188,8 +1189,8 @@ impl HybridQuery { self.inner_vec.refine_factor(refine_factor); } - pub fn nprobes(&mut self, nprobe: u32) { - self.inner_vec.nprobes(nprobe); + pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> { + self.inner_vec.nprobes(nprobe) } pub fn ef(&mut self, ef: u32) { diff --git a/python/src/table.rs b/python/src/table.rs index e7b6f1c6f..16f62bf8d 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -592,8 +592,6 @@ impl From for RefreshMaterializedViewRes fn from(result: lancedb::RefreshMaterializedViewResult) -> Self { let mode = match result.mode { lancedb::RefreshMode::Rebuild => "rebuild", - lancedb::RefreshMode::Incremental => "incremental", - lancedb::RefreshMode::NoOp => "no_op", }; Self { mode: mode.to_string(), @@ -1817,10 +1815,9 @@ impl Table { }) } - #[pyo3(signature = (full=false, source_version=None))] + #[pyo3(signature = (source_version=None))] pub fn refresh_materialized_view( self_: PyRef<'_, Self>, - full: bool, source_version: Option, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); @@ -1828,7 +1825,7 @@ impl Table { let view = lancedb::MaterializedView::from_table(inner) .await .infer_error()?; - let mut builder = view.refresh().full(full); + let mut builder = view.refresh(); if let Some(version) = source_version { builder = builder.source_version(version); } @@ -1837,10 +1834,9 @@ impl Table { }) } - #[pyo3(signature = (full=false, source_version=None))] + #[pyo3(signature = (source_version=None))] pub fn refresh_materialized_view_async( self_: PyRef<'_, Self>, - full: bool, source_version: Option, ) -> PyResult> { let inner = self_.inner_ref()?.clone(); @@ -1848,7 +1844,7 @@ impl Table { let view = lancedb::MaterializedView::from_table(inner) .await .infer_error()?; - let mut builder = view.refresh().full(full); + let mut builder = view.refresh(); if let Some(version) = source_version { builder = builder.source_version(version); } @@ -1863,11 +1859,9 @@ impl Table { let view = lancedb::MaterializedView::from_table(inner) .await .infer_error()?; - view.definition().to_json().map_err(|err| { - PyRuntimeError::new_err(format!( - "failed to serialize materialized-view definition: {err}" - )) - }) + Ok(lancedb::materialized_view::definition_metadata_from_sql( + view.definition_sql(), + )) }) } @@ -1897,6 +1891,17 @@ impl Table { let alterations = alterations .iter() .map(|alteration| { + for key in alteration.keys().iter() { + let key = key.extract::()?; + if !matches!( + key.as_str(), + "path" | "rename" | "name" | "nullable" | "data_type" + ) { + return Err(PyValueError::new_err(format!( + "Unknown column alteration key '{key}'" + ))); + } + } let path = alteration .get_item("path")? .ok_or_else(|| PyValueError::new_err("Missing path"))? diff --git a/python/vendor/lance-arrow/Cargo.toml b/python/vendor/lance-arrow/Cargo.toml index 83f9f0878..c028668fd 100644 --- a/python/vendor/lance-arrow/Cargo.toml +++ b/python/vendor/lance-arrow/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lance-arrow" -version = "13.1.0-beta.1" +version = "14.0.0-beta.1" edition = "2024" authors = ["Lance Devs "] license = "Apache-2.0" diff --git a/python/vendor/lance-arrow/README.lancedb.md b/python/vendor/lance-arrow/README.lancedb.md index d544b0a6d..7d5c7c7e7 100644 --- a/python/vendor/lance-arrow/README.lancedb.md +++ b/python/vendor/lance-arrow/README.lancedb.md @@ -1,5 +1,5 @@ -This directory contains `rust/lance-arrow` from Lance tag `v13.1.0-beta.1` -(commit `e5a3553b1699a7062ff7afe4ee268d596bbf2e0d`), licensed under Apache-2.0. +This directory contains `rust/lance-arrow` from Lance tag `v14.0.0-beta.1` +(commit `ae2903729f02ec8e519a4e53797a052393f25ee3`), licensed under Apache-2.0. LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps until a Lance release includes the fix. The patch is selected in the root diff --git a/python/vendor/lance-arrow/src/lib.rs b/python/vendor/lance-arrow/src/lib.rs index 3c111a19e..3401a11f4 100644 --- a/python/vendor/lance-arrow/src/lib.rs +++ b/python/vendor/lance-arrow/src/lib.rs @@ -442,13 +442,9 @@ pub trait FixedSizeBinaryArrayExt { /// let int_values = UInt8Array::from_iter(0..10); /// let fixed_size_list_arr = FixedSizeBinaryArray::try_new_from_values(&int_values, 2).unwrap(); /// assert_eq!(fixed_size_list_arr, - /// FixedSizeBinaryArray::from(vec![ - /// Some(vec![0, 1].as_slice()), - /// Some(vec![2, 3].as_slice()), - /// Some(vec![4, 5].as_slice()), - /// Some(vec![6, 7].as_slice()), - /// Some(vec![8, 9].as_slice()) - /// ])) + /// FixedSizeBinaryArray::try_from_iter( + /// [[0u8, 1], [2, 3], [4, 5], [6, 7], [8, 9]].into_iter() + /// ).unwrap()) /// ``` fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result; } @@ -3035,7 +3031,8 @@ mod tests { vec![Some(b"1".as_slice()), Some(b"23".as_slice())] ); - let fixed_size = FixedSizeBinaryArray::from(vec![b"abcd", b"efgh"]); + let fixed_size = + FixedSizeBinaryArray::try_from_iter([b"abcd", b"efgh"].into_iter()).unwrap(); assert_eq!( iter_binary_array(&fixed_size).unwrap().collect::>(), vec![Some(b"abcd".as_slice()), Some(b"efgh".as_slice())] diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index cdb0609d3..b002bdc95 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.40.0-beta.12" +version = "0.40.0-beta.13" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true diff --git a/rust/lancedb/examples/ivf_pq.rs b/rust/lancedb/examples/ivf_pq.rs index 79f674615..b457558ba 100644 --- a/rust/lancedb/examples/ivf_pq.rs +++ b/rust/lancedb/examples/ivf_pq.rs @@ -142,7 +142,7 @@ async fn search_index(table: &Table) -> Result<()> { // Override the default of 10 to get more rows .limit(15) // Override the default of 20 to search more partitions - .nprobes(30) + .nprobes(30)? // Override the default of None to apply a refine step .refine_factor(1) .execute() diff --git a/rust/lancedb/src/catalog.rs b/rust/lancedb/src/catalog.rs index 5eaa1d49e..576eb8e88 100644 --- a/rust/lancedb/src/catalog.rs +++ b/rust/lancedb/src/catalog.rs @@ -15,7 +15,12 @@ //! ``` use std::fmt; +use std::pin::Pin; use std::sync::Arc; +use std::task::{Context, Poll}; + +use futures::future::BoxFuture; +use futures::{FutureExt, Stream}; use crate::Result; use crate::connection::Connection; @@ -118,6 +123,73 @@ pub struct ListDatabasesResponse { pub page_token: Option, } +/// A lazy stream of database names with inspectable pagination state. +/// +/// The continuation token advances when a page is fetched, so resuming from it +/// skips any names still cached in this stream. Drain the cache before saving a +/// token if those names must be included. Errors terminate the stream; the token +/// for the failed request remains available for resuming with a new stream. +pub struct DatabaseNames { + catalog: Arc, + names: std::vec::IntoIter, + page_token: Option, + page_limit: Option, + done: bool, + pending: Option>>, +} + +impl DatabaseNames { + /// Number of names available without another REST request. + pub fn num_page_results(&self) -> usize { + self.names.len() + } + + /// Token for the next REST request, or `None` after the last page is fetched. + /// + /// Before the first request, returns the supplied starting token (`None` + /// means start from the beginning). A final page can still have cached names. + pub fn page_token(&self) -> Option<&str> { + self.page_token.as_deref() + } +} + +impl Stream for DatabaseNames { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + loop { + if let Some(name) = this.names.next() { + return Poll::Ready(Some(Ok(name))); + } + if this.done { + return Poll::Ready(None); + } + let pending = this.pending.get_or_insert_with(|| { + let catalog = this.catalog.clone(); + let request = ListDatabasesRequest { + page_token: this.page_token.clone(), + limit: this.page_limit, + }; + async move { catalog.list_databases(request).await }.boxed() + }); + let response = futures::ready!(pending.as_mut().poll(cx)); + this.pending = None; + match response { + Ok(response) => { + this.page_token = response.page_token.filter(|token| !token.is_empty()); + this.done = this.page_token.is_none(); + this.names = response.databases.into_iter(); + } + Err(error) => { + this.done = true; + return Poll::Ready(Some(Err(error))); + } + } + } + } +} + /// A backend that manages databases. Implementations own database lifecycle semantics. #[async_trait::async_trait] pub trait Catalog: Send + Sync + std::fmt::Debug + 'static { @@ -188,12 +260,35 @@ impl CatalogConnection { self.catalog.drop_database(request.into()).await } - /// List a page of databases. Pass [`ListDatabasesRequest::default`] for the first page. - pub async fn list_databases( + /// Iterate lazily over all database names, fetching more results as needed. + /// + /// `page_token` starts at a saved continuation token; `None` starts at the + /// beginning. `page_limit` limits each REST response (not the whole stream); + /// `None` uses the server default. Errors terminate the listing. + /// + /// ``` + /// # async fn example(catalog: &lancedb::catalog::CatalogConnection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let mut names = catalog.list_databases(None, None); + /// while let Some(name) = names.try_next().await? { + /// println!("{name}"); + /// } + /// # Ok(()) + /// # } + /// ``` + pub fn list_databases( &self, - request: ListDatabasesRequest, - ) -> Result { - self.catalog.list_databases(request).await + page_token: Option, + page_limit: Option, + ) -> DatabaseNames { + DatabaseNames { + catalog: self.catalog.clone(), + names: Vec::new().into_iter(), + page_token, + page_limit, + done: false, + pending: None, + } } /// Connect to a database without creating it if it is missing. diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 985195c07..1d0a4a0ab 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -568,6 +568,7 @@ impl Connection { /// Creates a new table by cloning from an existing source table. /// By default, this performs a shallow clone where the new table shares /// the underlying data files with the source table. + /// The target table name must be unused. An existing table is left unchanged. /// /// # Parameters /// - `target_table_name`: The name of the new table to create @@ -2193,6 +2194,63 @@ mod tests { assert_eq!(tables.len(), 0); } + #[rstest::rstest] + #[tokio::test] + async fn test_clone_table_target_already_exists( + #[values("source", "target")] target: &str, + #[values(false, true)] enable_v2_manifest_paths: bool, + ) { + use crate::query::ExecutableQuery; + use futures::TryStreamExt; + + let tmp_dir = tempdir().unwrap(); + let options = ListingDatabaseOptions::builder() + .enable_v2_manifest_paths(enable_v2_manifest_paths) + .build(); + let db = connect(tmp_dir.path().to_str().unwrap()) + .database_options(&options) + .execute() + .await + .unwrap(); + let source_data = arrow_array::record_batch!(("id", Int32, [0, 1, 2])).unwrap(); + let target_data = arrow_array::record_batch!(("id", Int32, [10, 11])).unwrap(); + db.create_table("source", source_data.clone()) + .execute() + .await + .unwrap(); + db.create_table("target", target_data.clone()) + .execute() + .await + .unwrap(); + + let source_uri = tmp_dir.path().join("source.lance"); + let err = db + .clone_table(target, source_uri.to_str().unwrap()) + .execute() + .await + .unwrap_err(); + assert!( + matches!(&err, Error::TableAlreadyExists { name } if name == target), + "unexpected clone error: {err:?}" + ); + + // Reopen both tables to verify that the failed clone left their data and versions intact. + for (name, expected_data) in [("source", source_data), ("target", target_data)] { + let table = db.open_table(name).execute().await.unwrap(); + assert_eq!(table.version().await.unwrap(), 1); + assert_eq!(table.list_versions().await.unwrap().len(), 1); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(batches, vec![expected_data]); + } + } + #[tokio::test] async fn test_clone_table() { let tmp_dir = tempdir().unwrap(); diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index dc1627116..8255fef4b 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -1107,6 +1107,21 @@ impl Database for ListingDatabase { ..Default::default() }; + let target_uri = self.table_uri(&request.target_table_name)?; + match DatasetBuilder::from_uri(&target_uri) + .with_read_params(read_params.clone()) + .load() + .await + { + Ok(_) => { + return Err(Error::TableAlreadyExists { + name: request.target_table_name, + }); + } + Err(lance::Error::DatasetNotFound { .. }) => {} + Err(err) => return Err(err.into()), + } + let mut source_dataset = DatasetBuilder::from_uri(&request.source_uri) .with_read_params(read_params.clone()) .load() @@ -1122,7 +1137,6 @@ impl Database for ListingDatabase { }), }?; - let target_uri = self.table_uri(&request.target_table_name)?; source_dataset .shallow_clone(&target_uri, version_ref, Some(storage_params)) .await diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 54b46fc11..711a73033 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -3,6 +3,8 @@ //! Handles to operations a server may run asynchronously. +use std::future::Future; +use std::pin::Pin; use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; use arrow_array::RecordBatch; @@ -250,6 +252,18 @@ where pub(crate) fn spawned(task: JoinHandle>) -> Self { Self::new_typed(Box::new(SpawnedJob::new(task))) } + + pub(crate) fn spawned_with_cancellation(task: JoinHandle>, cancel: F) -> Self + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let cancellation: Cancellation = Arc::new(move || Box::pin(cancel())); + Self::new_typed(Box::new(SpawnedJob::new_with_cancellation( + task, + cancellation, + ))) + } } impl Job @@ -501,10 +515,28 @@ impl Outcome { struct SpawnedJob { outcome: watch::Receiver>, abort: AbortHandle, + cancellation: Option, } +type CancellationFuture = Pin> + Send>>; +type Cancellation = Arc CancellationFuture + Send + Sync>; + impl SpawnedJob { fn new(task: JoinHandle>) -> Self + where + T: Serialize + Send + 'static, + { + Self::new_inner(task, None) + } + + fn new_with_cancellation(task: JoinHandle>, cancellation: Cancellation) -> Self + where + T: Serialize + Send + 'static, + { + Self::new_inner(task, Some(cancellation)) + } + + fn new_inner(task: JoinHandle>, cancellation: Option) -> Self where T: Serialize + Send + 'static, { @@ -526,7 +558,11 @@ impl SpawnedJob { }; let _ = tx.send(Some(outcome)); }); - Self { outcome, abort } + Self { + outcome, + abort, + cancellation, + } } } @@ -556,6 +592,19 @@ impl JobHandle for SpawnedJob { } async fn cancel(&self) -> Result<()> { + if let Some(cancellation) = &self.cancellation { + let mut outcome = self.outcome.clone(); + tokio::select! { + biased; + settled = outcome.wait_for(Option::is_some) => { + settled.map_err(|_| Error::Runtime { + message: "job outcome was dropped before cancellation completed".to_string(), + })?; + return Ok(()); + } + result = cancellation() => result?, + } + } self.abort.abort(); Ok(()) } diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index 7dec300a2..70b39506b 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -9,96 +9,20 @@ //! query this version cannot maintain reads back as unrefreshable, not as a //! plain table. Queries, indexes and search work on the view unchanged. -mod duplicate_pairs; mod grouped; -mod grouped_units; -mod partitioned; -pub(crate) mod vector_dedup; -pub use duplicate_pairs::{VectorSource, VectorSourceKind}; pub use grouped::IVF_PARTITION; -/// Refreshing a view grouped by [`IVF_PARTITION`] in units, one per index -/// partition, so the work can be spread over several processes. The caller -/// plans once, computes every unit of that plan wherever it likes, and -/// commits the whole set; a unit is bound to the view and plan it was -/// computed for, and the commit publishes all of them or none. -/// -/// ```no_run -/// # #![recursion_limit = "256"] -/// # use lancedb::materialized_view::{ -/// # commit_grouped_refresh, plan_grouped_refresh, write_grouped_unit, WrittenUnit, -/// # }; -/// # use lancedb::materialized_view::MaterializedView; -/// # async fn refresh_in_units(view: &MaterializedView) -> Result<(), Box> { -/// // The view this refresh was requested for. Every path below carries it, -/// // including the fallbacks: a view dropped and recreated meanwhile is a -/// // different view, and refreshing it in place of this one is the identity -/// // crossing the units path refuses. -/// let incarnation = view.incarnation().map(str::to_string); -/// -/// async fn in_one_pass(view: &MaterializedView, incarnation: Option<&str>) -> lancedb::Result<()> { -/// let mut refresh = view.refresh(); -/// if let Some(token) = incarnation { -/// refresh = refresh.expect_incarnation(token); -/// } -/// refresh.execute().await?; -/// Ok(()) -/// } -/// -/// // No plan means this view cannot be refreshed in units: it is not -/// // grouped by the partition of an indexed column, its index no longer -/// // says which fragments it covers, or it carries no incarnation to bind -/// // the units to. -/// let Some(plan) = plan_grouped_refresh(view.table(), None).await? else { -/// in_one_pass(view, incarnation.as_deref()).await?; -/// return Ok(()); -/// }; -/// -/// // Each unit is independent; this loop stands in for dispatching them. -/// // A unit whose view was dropped and recreated, or whose plan is stale, -/// // fails here rather than producing rows for the wrong view. -/// let mut units: Vec = Vec::new(); -/// for unit in 0..plan.units { -/// units.push(write_grouped_unit(view.table(), unit, &plan).await?); -/// } -/// -/// match commit_grouped_refresh(view.table(), &plan, units, incarnation.as_deref()).await { -/// Ok(result) => println!("{} rows at version {}", result.rows_written, result.version), -/// Err(err) => { -/// // The refresh is unrecorded, but it is not necessarily inert: a -/// // commit that raced a concurrent write can have landed without a -/// // watermark. Refresh from scratch rather than assuming either. -/// eprintln!("refresh in units unrecorded ({err}); refreshing in one pass"); -/// in_one_pass(view, incarnation.as_deref()).await?; -/// } -/// } -/// # Ok(()) -/// # } -/// ``` -pub use grouped_units::{ - GroupedRefreshPlan, WrittenUnit, commit_grouped_refresh, plan_grouped_refresh, - write_grouped_unit, -}; -pub use partitioned::{ - PartitionedRefreshPlan, WrittenPartition, commit_partitioned_refresh, plan_partitioned_refresh, - write_refresh_partition, write_refresh_partition_with_inputs, -}; mod query; pub mod refresh; -#[cfg(test)] -mod differential; - use std::collections::HashMap; use std::sync::Arc; use arrow_schema::{DataType, Field as ArrowField, FieldRef, Schema as ArrowSchema, SchemaRef}; use datafusion_common::ScalarValue; -use lance_core::ROW_ID; use lance_datafusion::planner::Planner; use serde::{Deserialize, Serialize}; use crate::connection::Connection; -use crate::database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS; use crate::database::{CreateTableRequest, Database, OpenTableRequest}; use crate::embeddings::EmbeddingDefinition; use crate::function::FunctionBinding; @@ -116,26 +40,6 @@ pub use refresh::{RefreshMaterializedViewResult, RefreshMode}; /// Schema metadata key holding the view definition; see [`DEFINITION_FORMAT`]. pub const DEFINITION_META_KEY: &str = "mv.definition"; -/// Schema metadata key holding the view's incarnation: a token minted at each -/// physical creation of a view table, so a view dropped and recreated under -/// the same name and definition is still told apart from the one a caller -/// captured. A view whose metadata was replaced wholesale, or one declared -/// before tokens existed, carries none until its next refresh mints one. -pub const INCARNATION_META_KEY: &str = "mv.incarnation"; - -/// Schema metadata key holding the source table version the view was last -/// refreshed to. Absent until the first refresh. -pub const SOURCE_VERSION_META_KEY: &str = "mv.source_version"; - -/// Schema metadata key holding the wall-clock time of the last refresh, -/// in milliseconds since the epoch. -pub const REFRESHED_AT_MS_META_KEY: &str = "mv.refreshed_at_ms"; - -/// Column recording which source row produced each view row: the source's -/// stable `_rowid` at refresh time, which is why sources must keep stable -/// row ids. -pub const SOURCE_ROW_ID_COLUMN: &str = "__source_row_id"; - /// Field metadata namespace for declarations about schema structure, such as /// an unenforced primary key. const SCHEMA_DECLARATION_META_PREFIX: &str = "lance-schema:"; @@ -154,26 +58,16 @@ const COLUMN_DEFINITIONS_META_KEY: &str = "lancedb::column_definitions"; /// The newest layout this version reads under [`DEFINITION_META_KEY`]: /// `{"format": N, "query": ""}`, the query as -/// [`MaterializedViewDefinition::to_sql`] renders it. Format 3 identifies a -/// native duplicate-pair source, format 4 a dedup result source, format 2 a query with `GROUP BY`, and -/// format 1 an ordinary query. A reader refuses a newer format rather than +/// [`MaterializedViewDefinition::to_sql`] renders it. Format 2 identifies a +/// query with `GROUP BY`, and format 1 an ordinary query. A reader refuses a newer format rather than /// guess at it. The layout also carries `"kind": "query"`, which readers /// older than the format number report as an unrefreshable view instead of /// failing to read the metadata. -pub const DEFINITION_FORMAT: u64 = 4; +pub const DEFINITION_FORMAT: u64 = 2; /// The format `definition` is written in; see [`DEFINITION_FORMAT`]. fn format_of(definition: &MaterializedViewDefinition) -> u64 { - if let Some(source) = &definition.vector_source { - match source.kind { - VectorSourceKind::Pairs => 3, - VectorSourceKind::Dedup => 4, - } - } else if definition.is_grouped() { - 2 - } else { - 1 - } + if definition.is_grouped() { 2 } else { 1 } } /// Legacy `kind` tag of the structured layout written before @@ -236,9 +130,8 @@ pub enum LateralSource { /// The list column. column: String, }, - /// `name(args)`: a Function in `FROM` position, returning rows. The - /// server stages its output in a hidden table, recorded under - /// [`STAGING_META_KEY`]; a local database cannot refresh this form. + /// `name(args)`: a Function in `FROM` position, returning rows. Local + /// databases do not execute table functions during refresh. Function { /// The Function's name. name: String, @@ -254,37 +147,17 @@ pub(crate) struct ViewUnnest { pub alias: String, } -/// Schema metadata key holding a [`StagingBinding`], present only on a view -/// whose query calls a Function in `FROM` position. -pub const STAGING_META_KEY: &str = "mv.staging"; - -/// Where a Function in `FROM` position has its output staged: a hidden -/// table carrying every source column plus `column`, the Function's list -/// output. Refresh scans this table in place of the query's source. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct StagingBinding { - /// The staging table's name. - pub table: String, - /// Its namespace path; empty is the root namespace. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub namespace: Vec, - /// The list column holding the Function's output. - pub column: String, -} - -/// The list column refresh unnests for `definition`, given the staging its -/// Function output lives in. `None` when the query has no lateral item. +/// The list column refresh unnests for `definition`. `None` when the query +/// has no lateral item. pub(crate) fn physical_unnest( definition: &MaterializedViewDefinition, - staging: Option<&StagingBinding>, ) -> Result> { let Some(lateral) = &definition.lateral else { return Ok(None); }; - let column = match (&lateral.source, staging) { - (LateralSource::Unnest { column }, _) => column.clone(), - (LateralSource::Function { .. }, Some(staging)) => staging.column.clone(), - (LateralSource::Function { name, .. }, None) => { + let column = match &lateral.source { + LateralSource::Unnest { column } => column.clone(), + LateralSource::Function { name, .. } => { return Err(Error::NotSupported { message: format!( "'{name}' in FROM position is a Function; views over Function rows are \ @@ -299,25 +172,11 @@ pub(crate) fn physical_unnest( })) } -/// Read the staging binding off a view's schema metadata, if it has one. -pub fn read_staging(metadata: &HashMap) -> Result> { - metadata - .get(STAGING_META_KEY) - .map(|raw| { - serde_json::from_str(raw).map_err(|e| Error::Runtime { - message: format!("unreadable materialized view staging binding: {e}"), - }) - }) - .transpose() -} - /// The query that defines a materialized view, in the relational shape /// refresh maintains. Stored as SQL; see [`MaterializedViewDefinition::from_sql`] /// for the shape. #[derive(Debug, Clone, PartialEq, Eq)] pub struct MaterializedViewDefinition { - /// A table-bound native vector source. It is never evaluated per source row. - pub vector_source: Option>, /// Name of the source table, in the same database as the view. pub source_table: String, /// Namespace path holding the source table; empty is the root namespace. @@ -408,10 +267,10 @@ pub enum StoredDefinition { #[doc(hidden)] #[derive(Debug, Clone, PartialEq, Eq)] pub struct MaterializedViewInfo { - /// The parsed view definition. - pub definition: MaterializedViewDefinition, - /// The current physical incarnation, when one has been minted. - pub incarnation: Option, + /// The defining SQL as stored. + pub definition_sql: String, + /// The definition parsed into the local refresh grammar, when supported. + pub parsed_definition: Option, } /// The backend-independent request used to create a remote materialized view. @@ -456,12 +315,78 @@ struct LegacyDefinition { /// Serialize `definition` into the layout stored under /// [`DEFINITION_META_KEY`]. pub(crate) fn definition_to_metadata(definition: &MaterializedViewDefinition) -> Result { - Ok(serde_json::json!({ + Ok(definition_metadata_from_sql_with_format( + &definition.to_sql(), + format_of(definition), + )) +} + +/// Serialize an arbitrary SQL definition into the materialized-view metadata +/// envelope. Sophon uses this for SQL queries that are broader than the local +/// refresh engine's supported shape. +pub fn definition_metadata_from_sql(sql: &str) -> String { + definition_metadata_from_sql_with_format(sql, 1) +} + +/// Serialize an arbitrary SQL definition together with the resolution +/// defaults needed to plan unqualified identifiers on a later refresh. +pub fn definition_metadata_from_sql_with_defaults( + sql: &str, + default_database: &str, + default_namespace_path: &[String], +) -> String { + serde_json::json!({ "kind": "query", - "format": format_of(definition), - "query": definition.to_sql(), + "format": 1, + "query": sql, + "default_database": default_database, + "default_namespace_path": default_namespace_path, }) - .to_string()) + .to_string() +} + +fn definition_metadata_from_sql_with_format(sql: &str, format: u64) -> String { + serde_json::json!({ + "kind": "query", + "format": format, + "query": sql, + }) + .to_string() +} + +/// Return the defining SQL stored in [`DEFINITION_META_KEY`], without +/// requiring the local refresh parser to understand it. +pub fn read_definition_sql(metadata: &HashMap) -> Result> { + Ok(read_definition_sql_with_defaults(metadata)?.map(|definition| definition.query)) +} + +/// A stored SQL definition and the optional name-resolution defaults recorded +/// by a remote SQL engine. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct MaterializedViewSqlDefinition { + /// The SQL query that defines the materialized view. + pub query: String, + /// The database used to resolve unqualified references when the view was created. + #[serde(default)] + pub default_database: Option, + /// The namespace used to resolve unqualified references when the view was created. + #[serde(default)] + pub default_namespace_path: Option>, +} + +/// Read the SQL definition and its optional resolution defaults from +/// [`DEFINITION_META_KEY`]. +pub fn read_definition_sql_with_defaults( + metadata: &HashMap, +) -> Result> { + let Some(raw) = metadata.get(DEFINITION_META_KEY) else { + return Ok(None); + }; + serde_json::from_str(raw) + .map(Some) + .map_err(|e| Error::Runtime { + message: format!("unreadable materialized view definition: {e}"), + }) } /// Read a view definition off a schema metadata map, if it carries one. @@ -516,7 +441,6 @@ pub fn read_definition(metadata: &HashMap) -> Result, ) -> Result { - let incarnation = metadata.get(INCARNATION_META_KEY).cloned(); match read_definition(metadata)? { Some(StoredDefinition::Query(definition)) => Ok(MaterializedViewInfo { - definition, - incarnation, + definition_sql: definition.to_sql(), + parsed_definition: Some(definition), }), Some(StoredDefinition::Newer { format }) => Err(Error::NotSupported { message: format!( @@ -563,19 +486,14 @@ pub(crate) struct Planned { pub fields: Vec, pub lineage: Lineage, /// Source columns the query reads; a read through the unnest alias is - /// recorded as the list column, which is what the source has and what - /// incremental refresh watches. + /// recorded as the list column present in the source schema. pub inputs: Vec, } pub(crate) fn plan( source_schema: SchemaRef, definition: &MaterializedViewDefinition, - staging: Option<&StagingBinding>, ) -> Result { - if definition.vector_source.is_some() { - return duplicate_pairs::plan(source_schema, definition); - } let filter = definition .filter .as_deref() @@ -593,7 +511,7 @@ pub(crate) fn plan( // Projections are typed against the source, or for an unnested view // against the source with the list column replaced by its element // under the alias, where `c.chunk` is an ordinary nested path. - let unnest = physical_unnest(definition, staging)?; + let unnest = physical_unnest(definition)?; let source_schema = match &unnest { None => source_schema, Some(unnest) => { @@ -607,12 +525,9 @@ pub(crate) fn plan( } }; let projections: Vec<(String, String)> = if definition.selects_star() { - // `SELECT *`. A source that is itself a view carries its own - // provenance column; the new view records its own, not a copy. source_schema .fields() .iter() - .filter(|f| f.name() != SOURCE_ROW_ID_COLUMN) .map(|f| (f.name().clone(), query::ident_sql(f.name()))) .collect() } else { @@ -653,25 +568,12 @@ pub(crate) fn plan( name: output.clone(), }); } - if output == SOURCE_ROW_ID_COLUMN || output == ROW_ID { - return Err(Error::InvalidInput { - message: format!("view column name '{output}' is reserved"), - }); - } - let parsed = planner .parse_expr(expression) .map_err(|e| Error::InvalidExpression { column: output.clone(), message: e.to_string(), })?; - // Before optimization: the simplifier folds a stable-but-not-immutable - // call like now() into a literal, hiding it from the check while the - // stored definition keeps the call. - ensure_immutable(&parsed, |message| Error::InvalidExpression { - column: output.clone(), - message, - })?; let expr = planner .optimize_expr(parsed) .map_err(|e| Error::InvalidExpression { @@ -734,7 +636,6 @@ pub(crate) fn plan( } let definition = MaterializedViewDefinition { - vector_source: None, source_table: definition.source_table.clone(), source_namespace: definition.source_namespace.clone(), lateral: definition.lateral.clone(), @@ -760,17 +661,14 @@ pub(crate) fn plan( }) } -/// Check that `filter` is an immutable boolean predicate over `source_schema`, -/// returning the columns it reads. +/// Check that `filter` is a boolean predicate over `source_schema`, returning +/// the columns it reads. pub(crate) fn plan_filter(source_schema: &SchemaRef, filter: &str) -> Result> { let expr = Planner::new(source_schema.clone()) .parse_filter(filter) .map_err(|e| Error::InvalidInput { message: format!("invalid view filter: {e}"), })?; - ensure_immutable(&expr, |message| Error::InvalidInput { - message: format!("invalid view filter: {message}"), - })?; let filter_inputs = resolve_inputs(source_schema, &expr, |message| Error::InvalidInput { message: format!("invalid view filter: {message}"), })?; @@ -839,41 +737,6 @@ pub(crate) fn flattened_schema( Ok(Arc::new(ArrowSchema::new(fields))) } -/// Reject any function that is not immutable: a view definition has to -/// evaluate identically across refreshes, or incremental maintenance would -/// mix rows from different evaluations of the same definition. -fn ensure_immutable(expr: &datafusion_expr::Expr, error: impl Fn(String) -> Error) -> Result<()> { - use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; - use datafusion_expr::Volatility; - - // Labeled immutable but not determined by row values alone: version() - // depends on the build, the arrow_* introspectors on schema state. - const NOT_VALUE_DETERMINED: &[&str] = - &["version", "arrow_typeof", "arrow_field", "arrow_metadata"]; - - let mut offending: Option = None; - expr.apply(|node| { - if let datafusion_expr::Expr::ScalarFunction(function) = node { - let name = function.func.name(); - if function.func.signature().volatility != Volatility::Immutable - || NOT_VALUE_DETERMINED.contains(&name) - { - offending = Some(name.to_string()); - return Ok(TreeNodeRecursion::Stop); - } - } - Ok(TreeNodeRecursion::Continue) - }) - .map_err(|e| error(e.to_string()))?; - match offending { - Some(name) => Err(error(format!( - "function '{name}' is not immutable and would evaluate differently \ - across refreshes" - ))), - None => Ok(()), - } -} - /// The root of a possibly-dotted column path: `metadata.age` -> `metadata`. fn root(path: &str) -> &str { path.split('.').next().unwrap_or(path) @@ -1125,7 +988,7 @@ fn project_schema(schema: &ArrowSchema, columns: &[String]) -> SchemaRef { } /// A validated view declaration, ready to become a table: the projected -/// fields plus [`SOURCE_ROW_ID_COLUMN`], definition stamped in metadata. +/// fields with the definition stamped in metadata. /// Produced only by [`prepare_declaration`]. #[derive(Clone)] pub struct PreparedDeclaration { @@ -1158,8 +1021,7 @@ impl PreparedDeclaration { } /// The schema the view will have: the declared columns in order, any - /// internal projections added by [`PreparedDeclaration::input_column`], - /// then [`SOURCE_ROW_ID_COLUMN`]. + /// internal projections added by [`PreparedDeclaration::input_column`]. pub fn schema(&self) -> &SchemaRef { &self.schema } @@ -1168,11 +1030,6 @@ impl PreparedDeclaration { /// read: the column the view projects it to, if any, otherwise an /// internal projection added here, named by [`input_column_name`]. pub fn input_column(&mut self, source_column: &str) -> Result { - if self.definition.vector_source.is_some() { - return Err(Error::InvalidInput { - message: "native vector views have no per-source-row inputs".into(), - }); - } // A grouped view's rows are groups; no source row carries a value into one. if self.definition.is_grouped() { return Err(Error::InvalidInput { @@ -1194,14 +1051,13 @@ impl PreparedDeclaration { if self.schema.field_with_name(&name).is_ok() { return Err(Error::ColumnAlreadyExists { name }); } - let row_id = self.row_id_index()?; let mut fields: Vec = self .schema .fields() .iter() .map(|f| f.as_ref().clone()) .collect(); - fields.insert(row_id, without_declarations(&field.with_name(name.clone()))); + fields.push(without_declarations(&field.with_name(name.clone()))); self.definition.projections.push(ViewProjection { output: name.clone(), expression: source_column @@ -1273,11 +1129,6 @@ impl PreparedDeclaration { columns: Vec<(usize, ArrowField)>, bindings: &[FunctionBinding], ) -> Result { - if self.definition.vector_source.is_some() { - return Err(Error::InvalidInput { - message: "computed columns on native vector views are not supported".into(), - }); - } let invalid = |message: String| Error::InvalidInput { message }; if columns.is_empty() { return Err(invalid("at least one computed column is needed".into())); @@ -1303,10 +1154,7 @@ impl PreparedDeclaration { columns.sort_by_key(|(position, _)| *position); for (inserted, (position, field)) in columns.iter().enumerate() { let name = field.name().as_str(); - if name == SOURCE_ROW_ID_COLUMN - || name == ROW_ID - || name.starts_with(INPUT_COLUMN_PREFIX) - { + if name.starts_with(INPUT_COLUMN_PREFIX) { return Err(invalid(format!("view column name '{name}' is reserved"))); } if fields.iter().any(|f| f.name() == name) { @@ -1348,16 +1196,8 @@ impl PreparedDeclaration { Ok(self) } - fn row_id_index(&self) -> Result { - self.schema - .index_of(SOURCE_ROW_ID_COLUMN) - .map_err(|e| Error::Runtime { - message: e.to_string(), - }) - } - /// Columns the declaration lists: everything before the internal - /// projections and the provenance column. + /// projections. fn visible_count(&self) -> usize { self.definition.projections.len() - self.internal_inputs + computed_columns(&self.schema).len() @@ -1366,9 +1206,7 @@ impl PreparedDeclaration { /// Create the view table and verify it, consuming the declaration. /// /// The view goes at the root of the source's own database, where refresh - /// resolves the recorded source coordinate. Stable row ids are requested - /// at both levels and verified rather than trusted; nothing is rolled - /// back on failure. + /// resolves the recorded source coordinate. pub async fn create(self, name: &str) -> Result { self.create_in(&[], name).await } @@ -1382,53 +1220,16 @@ impl PreparedDeclaration { ) -> Result { let empty: Vec> = vec![]; - // Minted here, not at preparation: a declaration can be cloned and - // create more than one physical table, and each needs its own token. - let incarnation = uuid::Uuid::new_v4().to_string(); - let mut metadata = self.schema.metadata().clone(); - metadata.insert(INCARNATION_META_KEY.to_string(), incarnation.clone()); - let schema = Arc::new(ArrowSchema::new_with_metadata( - self.schema.fields().clone(), - metadata, - )); let reader: Box = - Box::new(arrow_array::RecordBatchIterator::new(empty, schema)); + Box::new(arrow_array::RecordBatchIterator::new(empty, self.schema)); let mut request = CreateTableRequest::new(name.to_string(), Box::new(reader)); request.namespace_path = namespace_path.to_vec(); - let write_params = request - .write_options - .lance_write_params - .get_or_insert_with(Default::default); - write_params.enable_stable_row_ids = true; - let store_params = write_params - .store_params - .get_or_insert_with(Default::default); - crate::connection::merge_storage_options( - store_params, - [( - OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(), - "true".to_string(), - )], - ); let table = self.database.clone().create_table(request).await?; let table = Table::new(table, self.database); - let stable = match table.as_native() { - Some(native) => native.dataset.get().await?.manifest.uses_stable_row_ids(), - None => false, - }; - if !stable { - return Err(Error::Runtime { - message: format!( - "view '{name}' was created without stable row ids: the database \ - ignored the creation option; the table remains and is not \ - usable as a materialized view" - ), - }); - } Ok(MaterializedView { table, - definition: self.definition, - incarnation: Some(incarnation), + definition_sql: self.definition.to_sql(), + parsed_definition: Some(self.definition), }) } } @@ -1484,7 +1285,6 @@ pub async fn prepare_declaration( limit: Option, ) -> Result { let definition = MaterializedViewDefinition { - vector_source: None, source_table: source.name().to_string(), source_namespace: source.namespace().to_vec(), lateral: None, @@ -1508,7 +1308,7 @@ pub async fn prepare_declaration( /// Validate `definition` over `source`, the table it names. The declaration /// is planned against the source as refresh will reach it, and the result /// creates the view with [`PreparedDeclaration::create`]. A query calling a -/// Function in `FROM` position needs [`prepare_staged_definition`]. +/// Function in `FROM` position is not supported by local refresh. /// /// ``` /// # #![recursion_limit = "256"] @@ -1546,27 +1346,6 @@ pub async fn prepare_declaration( /// # Ok(()) /// # } /// ``` -/// -/// An indexed, pinned source can materialize retained original rows using the -/// default greedy direct-representative policy. A row is removed only when it -/// directly matches a retained row; an A-B-C chain keeps A and C. The source is -/// unchanged. Hosts distribute this through [`plan_partitioned_refresh`] and -/// [`write_refresh_partition_with_inputs`], with a barrier between selection -/// tasks and source-fragment materialization tasks. -/// -/// ``` -/// # #![recursion_limit = "256"] -/// use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition}; -/// # async fn dedup(images: &lancedb::Table) -> Result<(), Box> { -/// let version = images.version().await?; -/// let definition = MaterializedViewDefinition::from_sql(&format!( -/// "SELECT * FROM vector_dedup('images', {version}, 'phash', 4)" -/// ))?; -/// let clean = prepare_definition(images, definition).await?.create("images_clean").await?; -/// clean.refresh().execute().await?; -/// # Ok(()) -/// # } -/// ``` pub async fn prepare_definition( source: &Table, definition: MaterializedViewDefinition, @@ -1583,59 +1362,12 @@ pub async fn prepare_definition( ), }); } - prepare_with(source, definition, None).await -} - -/// Validate a `definition` whose query calls a Function in `FROM` position, -/// planned over `staging`: a table carrying every column of the query's -/// source plus `column`, the Function's list output for that row. The view -/// records the query as written and the staging under [`STAGING_META_KEY`]; -/// refresh scans the staging table and unnests `column`. -/// -/// ``` -/// # #![recursion_limit = "256"] -/// use lancedb::materialized_view::{MaterializedViewDefinition, prepare_staged_definition}; -/// -/// // `staging` holds every column of `docs` plus `chunks`, the list -/// // `chunk(body)` returned for each row. -/// # async fn declare(staging: &lancedb::Table) -> Result<(), Box> { -/// let definition = MaterializedViewDefinition::from_sql( -/// "SELECT id, c.text, c.ordinal FROM docs, chunk(body) AS c", -/// )?; -/// let view = prepare_staged_definition(staging, definition, "chunks") -/// .await? -/// .create("chunks") -/// .await?; -/// view.refresh().execute().await?; -/// # Ok(()) -/// # } -/// ``` -pub async fn prepare_staged_definition( - staging: &Table, - definition: MaterializedViewDefinition, - column: impl Into, -) -> Result { - if !matches!( - definition.lateral.as_ref().map(|l| &l.source), - Some(LateralSource::Function { .. }) - ) { - return Err(Error::InvalidInput { - message: "only a query calling a Function in FROM position takes a staging table" - .into(), - }); - } - let binding = StagingBinding { - table: staging.name().to_string(), - namespace: staging.namespace().to_vec(), - column: column.into(), - }; - prepare_with(staging, definition, Some(binding)).await + prepare_with(source, definition).await } async fn prepare_with( source: &Table, definition: MaterializedViewDefinition, - staging: Option, ) -> Result { let Some(caller_native) = source.as_native() else { return Err(Error::NotSupported { @@ -1682,17 +1414,6 @@ async fn prepare_with( ), }); } - if definition.vector_source.is_none() - && !native.dataset.get().await?.manifest.uses_stable_row_ids() - { - return Err(Error::InvalidInput { - message: format!( - "materialized views require stable row ids on the source table; \ - create '{}' with storage option new_table_enable_stable_row_ids=true", - source.name() - ), - }); - } refresh::ensure_no_mem_wal( native.dataset.get().await?.as_ref(), "source table", @@ -1710,45 +1431,23 @@ async fn prepare_with( message: format!("view column name '{}' is reserved", reserved.output), }); } - let source_schema = if let Some(pairs) = &definition.vector_source { - let ds = native - .dataset - .get() - .await? - .checkout_version(pairs.dataset_version) - .await?; - crate::table::datafusion::udtf::duplicate_pairs::plan_duplicate_pairs( - Arc::new(ds.clone()), - &pairs.config()?, - ) - .await?; - Arc::new(ArrowSchema::from(ds.schema())) - } else { - resolved.schema().await? - }; + let source_schema = resolved.schema().await?; let source_metadata = source_schema.metadata().clone(); let Planned { definition, - mut fields, + fields, lineage, .. - } = plan(source_schema.clone(), &definition, staging.as_ref())?; + } = plan(source_schema.clone(), &definition)?; if definition.is_grouped() { grouped::check(native.dataset.get().await?.as_ref(), &definition).await?; } // What later projections (`input_column`) are planned against: for an // unnested view the flattened schema, where the alias is a column. - let planning_schema = match physical_unnest(&definition, staging.as_ref())? { + let planning_schema = match physical_unnest(&definition)? { None => source_schema.clone(), Some(unnest) => flattened_schema(&source_schema, &unnest)?, }; - if definition.vector_source.is_none() { - fields.push(ArrowField::new( - SOURCE_ROW_ID_COLUMN, - DataType::UInt64, - false, - )); - } // Only column-describing metadata comes along: structural declarations // describe how a table is written, and a view is written by refresh alone. let mut metadata: HashMap = HashMap::new(); @@ -1766,14 +1465,6 @@ async fn prepare_with( DEFINITION_META_KEY.to_string(), definition_to_metadata(&definition)?, ); - if let Some(staging) = &staging { - metadata.insert( - STAGING_META_KEY.to_string(), - serde_json::to_string(staging).map_err(|e| Error::Runtime { - message: format!("failed to serialize the staging binding: {e}"), - })?, - ); - } Ok(PreparedDeclaration { schema: Arc::new(ArrowSchema::new_with_metadata(fields, metadata)), definition, @@ -1792,13 +1483,31 @@ pub struct CreateMaterializedViewBuilder { namespace: Vec, source: String, source_namespace: Vec, - vector_source: Option, projections: Vec<(String, String)>, filter: Option, limit: Option, with_no_data: bool, } +fn source_sql_identifier( + namespace_path: &[String], + source: &str, + namespace_in_schema_position: bool, +) -> String { + fn quote(name: &str) -> String { + format!("\"{}\"", name.replace('"', "\"\"")) + } + + let namespace = if namespace_in_schema_position && !namespace_path.is_empty() { + vec![namespace_path.join("$")] + } else { + namespace_path.to_vec() + }; + let mut parts = namespace.iter().map(|part| quote(part)).collect::>(); + parts.push(quote(source)); + parts.join(".") +} + impl CreateMaterializedViewBuilder { pub(crate) fn new(connection: Connection, name: String, source: String) -> Self { Self { @@ -1807,7 +1516,6 @@ impl CreateMaterializedViewBuilder { namespace: Vec::new(), source, source_namespace: Vec::new(), - vector_source: None, projections: Vec::new(), filter: None, limit: None, @@ -1828,38 +1536,6 @@ impl CreateMaterializedViewBuilder { self } - /// Read native pairs or retained original rows from a fixed indexed snapshot. - /// This source selects its full output schema and cannot be combined with - /// projections, filters or limits. Remote creation uses the ordinary MV - /// endpoint and its registry job; it does not execute a SQL query client-side. - pub fn vector_source(mut self, source: VectorSource) -> Self { - self.vector_source = Some(source); - self - } - - fn native_definition(&self) -> Result> { - let Some(source) = &self.vector_source else { - return Ok(None); - }; - source.config()?; - if !self.projections.is_empty() || self.filter.is_some() || self.limit.is_some() { - return Err(Error::InvalidInput { - message: "native vector sources cannot be combined with select, where or limit" - .into(), - }); - } - Ok(Some(MaterializedViewDefinition { - vector_source: Some(Box::new(source.clone())), - source_table: self.source.clone(), - source_namespace: self.source_namespace.clone(), - projections: vec![ViewProjection::star()], - filter: None, - lateral: None, - group_by: Vec::new(), - limit: None, - })) - } - /// The view's columns, as `(name, SQL expression)` pairs. Not calling /// this selects every source column, expanded at creation time. pub fn select( @@ -1893,9 +1569,6 @@ impl CreateMaterializedViewBuilder { } fn query(&self) -> Result { - if let Some(definition) = self.native_definition()? { - return Ok(definition.to_sql()); - } fn quote(name: &str) -> String { format!("\"{}\"", name.replace('"', "\"\"")) } @@ -1909,13 +1582,11 @@ impl CreateMaterializedViewBuilder { .collect::>() .join(", ") }; - let source = self - .source_namespace - .iter() - .chain(std::iter::once(&self.source)) - .map(|part| quote(part)) - .collect::>() - .join("."); + let source = source_sql_identifier( + &self.source_namespace, + &self.source, + self.connection.uri().starts_with("db://"), + ); let mut query = format!("SELECT {projection} FROM {source}"); if let Some(filter) = &self.filter { query.push_str(" WHERE "); @@ -1928,9 +1599,7 @@ impl CreateMaterializedViewBuilder { } /// Submit creation and initial population, returning a [`Job`] that - /// settles when the view is ready. Ordinary sources must keep stable row - /// ids for provenance across compaction. Native vector sources also support - /// physical row IDs because their definitions pin the dataset version. + /// settles when the view is ready. pub async fn execute_async(self) -> Result { if self.connection.uri().starts_with("db://") { return self @@ -1967,24 +1636,19 @@ impl CreateMaterializedViewBuilder { } async fn execute_native(self) -> Result { - let native_definition = self.native_definition()?; let source = self .connection .open_table(&self.source) .namespace(self.source_namespace.clone()) .execute() .await?; - let prepared = if let Some(definition) = native_definition { - prepare_definition(&source, definition).await? - } else { - prepare_declaration( - &source, - (!self.projections.is_empty()).then_some(self.projections.as_slice()), - self.filter.as_deref(), - self.limit, - ) - .await? - }; + let prepared = prepare_declaration( + &source, + (!self.projections.is_empty()).then_some(self.projections.as_slice()), + self.filter.as_deref(), + self.limit, + ) + .await?; let view = prepared.create_in(&self.namespace, &self.name).await?; if !self.with_no_data { view.refresh().execute().await?; @@ -1993,24 +1657,24 @@ impl CreateMaterializedViewBuilder { } } -/// A handle on a materialized view: the view table plus its parsed definition. +/// A handle on a materialized view: the view table plus its defining SQL. #[derive(Debug, Clone)] pub struct MaterializedView { table: Table, - definition: MaterializedViewDefinition, - incarnation: Option, + definition_sql: String, + parsed_definition: Option, } impl MaterializedView { /// Interpret `table` as a materialized view: [`Error::NotAMaterializedView`] - /// for a plain table, [`Error::NotSupported`] for a query this version - /// cannot refresh. + /// for a plain table. A remote definition broader than the local refresh + /// grammar remains available through [`Self::definition_sql`]. pub async fn from_table(table: Table) -> Result { let info = table.base_table().materialized_view_info().await?; Ok(Self { table, - definition: info.definition, - incarnation: info.incarnation, + definition_sql: info.definition_sql, + parsed_definition: info.parsed_definition, }) } @@ -2024,23 +1688,28 @@ impl MaterializedView { self.table.name() } - /// The query that defines the view. - pub fn definition(&self) -> &MaterializedViewDefinition { - &self.definition + /// The query parsed into the local refresh grammar. + /// + /// Sophon may store SQL broader than the local grammar; use + /// [`Self::definition_sql`] when the original SQL is sufficient. + pub fn definition(&self) -> Result<&MaterializedViewDefinition> { + self.parsed_definition + .as_ref() + .ok_or_else(|| Error::NotSupported { + message: + "this materialized-view definition is supported only by the remote SQL engine" + .to_string(), + }) } - /// The view's incarnation token as of when this handle was opened; see - /// [`RefreshMaterializedViewBuilder::expect_incarnation`]. `None` for a - /// view that has none yet (see [`INCARNATION_META_KEY`]). - pub fn incarnation(&self) -> Option<&str> { - self.incarnation.as_deref() + /// The complete SQL query that defines the view. + pub fn definition_sql(&self) -> &str { + &self.definition_sql } /// Recompute the view from its source. /// - /// By default the refresh is incremental when the source's changes can be - /// reconciled into the view, and otherwise rebuilds; see - /// [`RefreshMaterializedViewBuilder`]. + /// Every refresh rebuilds the view from the selected source version. /// /// ```no_run /// # #![recursion_limit = "256"] @@ -2054,9 +1723,7 @@ impl MaterializedView { pub fn refresh(&self) -> RefreshMaterializedViewBuilder { RefreshMaterializedViewBuilder { view: self.clone(), - full: false, source_version: None, - expected_incarnation: None, } } } @@ -2064,38 +1731,17 @@ impl MaterializedView { /// Builds a refresh. Created by [`MaterializedView::refresh`]. pub struct RefreshMaterializedViewBuilder { view: MaterializedView, - full: bool, source_version: Option, - expected_incarnation: Option, } impl RefreshMaterializedViewBuilder { - /// Rebuild the view even where an incremental refresh would do. - pub fn full(mut self, full: bool) -> Self { - self.full = full; - self - } - - /// Refresh to this source table version instead of the latest. + /// Refresh a local view to this source table version instead of the latest. + /// Remote SQL refreshes do not support source-version pinning. pub fn source_version(mut self, version: u64) -> Self { self.source_version = Some(version); self } - /// Refresh only if the view is still the incarnation that minted `token` - /// (see [`MaterializedView::incarnation`]): a refresh requested against - /// one declaration must not land in a view dropped and recreated since, - /// even under the same name and definition. - /// - /// Best effort. The token is read from the latest stored manifest before - /// planning and again immediately before every commit, but it is not part - /// of the commit's own condition, so a recreation that lands between that - /// final read and the commit is not caught. - pub fn expect_incarnation(mut self, token: impl Into) -> Self { - self.expected_incarnation = Some(token.into()); - self - } - /// Submit the refresh and return a job that settles with its result. pub async fn execute_async(self) -> Result> { if self.view.table.as_native().is_none() { @@ -2103,34 +1749,18 @@ impl RefreshMaterializedViewBuilder { .view .table .base_table() - .refresh_materialized_view_async( - self.full, - self.source_version, - self.expected_incarnation.as_deref(), - ) + .refresh_materialized_view_async(self.source_version) .await; } Ok(Job::spawned(tokio::spawn(async move { - refresh::execute_refresh( - &self.view.table, - self.full, - self.source_version, - self.expected_incarnation.as_deref(), - ) - .await + refresh::execute_refresh(&self.view.table, self.source_version).await }))) } /// Refresh the view, waiting for the job to finish. pub async fn execute(self) -> Result { if self.view.table.as_native().is_some() { - return refresh::execute_refresh( - &self.view.table, - self.full, - self.source_version, - self.expected_incarnation.as_deref(), - ) - .await; + return refresh::execute_refresh(&self.view.table, self.source_version).await; } self.execute_async().await?.wait().await } @@ -2249,2033 +1879,84 @@ mod tests { use super::*; use crate::connect; - use crate::table::WriteOptions; - async fn people_db() -> Connection { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!( - ("name", Utf8, ["ada", "grace", "alan"]), - ("age", Int32, [36, 85, 41]) - ) - .unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - conn - } - - /// Sources must keep stable row ids; see the create-time gate. - pub(super) fn stable_row_ids() -> WriteOptions { - WriteOptions { - lance_write_params: Some(lance::dataset::WriteParams { - enable_stable_row_ids: true, - ..Default::default() - }), - } - } - - /// The error a doomed declaration against `people` produces. - async fn declare_err( - cfg: impl FnOnce(CreateMaterializedViewBuilder) -> CreateMaterializedViewBuilder, - ) -> Error { - let conn = people_db().await; - cfg(conn.create_materialized_view("bad", "people")) - .execute() - .await - .unwrap_err() - } - - #[tokio::test] - async fn test_create_records_the_definition() { - let conn = people_db().await; - let view = conn - .create_materialized_view("adults", "people") - .select([("name", "name"), ("shout", "upper(name)")]) - .only_if("age >= 18") - .limit(10) - .execute() - .await - .unwrap(); - - assert_eq!(view.name(), "adults"); + #[test] + fn remote_source_namespace_uses_one_sql_schema_position() { + let namespace = vec!["analytics".to_string(), "daily".to_string()]; assert_eq!( - view.definition(), - &MaterializedViewDefinition { - vector_source: None, - source_table: "people".into(), - source_namespace: Vec::new(), - projections: vec![ - ViewProjection { - output: "name".into(), - expression: "name".into() - }, - ViewProjection { - output: "shout".into(), - expression: "upper(name)".into() - }, - ], - filter: Some("age >= 18".into()), - group_by: Vec::new(), - limit: Some(10), - lateral: None, - } - ); - - // The definition round-trips off the stored schema, not the handle. - let reopened = conn.open_materialized_view("adults").await.unwrap(); - assert_eq!(reopened.definition(), view.definition()); - } - - #[tokio::test] - async fn test_view_schema_is_derived_from_the_query() { - let conn = people_db().await; - let view = conn - .create_materialized_view("shapes", "people") - .select([("shout", "upper(name)"), ("next_age", "age + 1")]) - .execute() - .await - .unwrap(); - - let schema = view.table().schema().await.unwrap(); - assert_eq!( - schema.field_with_name("shout").unwrap().data_type(), - &DataType::Utf8 + source_sql_identifier(&namespace, "source", true), + "\"analytics$daily\".\"source\"" ); assert_eq!( - schema.field_with_name("next_age").unwrap().data_type(), - &DataType::Int32 - ); - assert_eq!( - schema - .field_with_name(SOURCE_ROW_ID_COLUMN) - .unwrap() - .data_type(), - &DataType::UInt64 - ); - assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - } - - #[tokio::test] - async fn test_with_no_data_skips_initial_refresh() { - let conn = people_db().await; - let view = conn - .create_materialized_view("empty", "people") - .with_no_data(true) - .execute() - .await - .unwrap(); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - } - - #[tokio::test] - async fn test_create_and_refresh_async_jobs() { - let conn = people_db().await; - let create_job = conn - .create_materialized_view("async_view", "people") - .with_no_data(true) - .execute_async() - .await - .unwrap(); - assert!(create_job.id().is_none()); - create_job.wait().await.unwrap(); - - let view = conn.open_materialized_view("async_view").await.unwrap(); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - - let refresh_job = view.refresh().execute_async().await.unwrap(); - assert!(refresh_job.id().is_none()); - let result = refresh_job.wait().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 3); - assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - - let drop_job = conn - .drop_materialized_view_async("async_view", &[]) - .await - .unwrap(); - assert!(drop_job.id().is_none()); - drop_job.wait().await.unwrap(); - assert!(conn.open_table("async_view").execute().await.is_err()); - } - - #[tokio::test] - async fn test_drop_materialized_view_rejects_plain_tables() { - let conn = people_db().await; - let error = conn - .drop_materialized_view("people", &[]) - .await - .unwrap_err(); - assert!(matches!(error, Error::NotAMaterializedView { .. })); - - conn.create_materialized_view("drop_me", "people") - .with_no_data(true) - .execute() - .await - .unwrap(); - conn.drop_materialized_view("drop_me", &[]).await.unwrap(); - assert!(conn.open_table("drop_me").execute().await.is_err()); - } - - #[tokio::test] - async fn test_remote_query_quotes_resource_identifiers() { - let conn = people_db().await; - let query = conn - .create_materialized_view("unused", "odd\"source") - .source_namespace(vec!["raw data".into()]) - .select([("double\"age", "age * 2")]) - .only_if("age >= 18") - .limit(10) - .query() - .unwrap(); - assert_eq!( - query, - "SELECT age * 2 AS \"double\"\"age\" FROM \"raw data\".\"odd\"\"source\" WHERE age >= 18 LIMIT 10" + source_sql_identifier(&namespace, "source", false), + "\"analytics\".\"daily\".\"source\"" ); } - #[tokio::test] - async fn native_vector_builder_preserves_source_and_rejects_extra_clauses() { - let conn = people_db().await; - let source = VectorSource { - kind: VectorSourceKind::Dedup, - dataset_version: 7, - column: "vector'field".into(), - distance_threshold: "4.0".into(), - }; - let builder = || { - conn.create_materialized_view("clean", "odd.source'name") - .source_namespace(vec!["raw space".into()]) - .vector_source(source.clone()) - }; - let query = builder().query().unwrap(); - let definition = MaterializedViewDefinition::from_sql(&query).unwrap(); - assert_eq!(definition.source_table, "odd.source'name"); - assert_eq!(definition.source_namespace, ["raw space"]); - assert_eq!(definition.vector_source.as_deref(), Some(&source)); - assert!(definition.selects_star()); - assert!(builder().select([("id", "id")]).query().is_err()); - assert!(builder().only_if("true").query().is_err()); - assert!(builder().limit(1).query().is_err()); - let invalid_source = VectorSource { - dataset_version: 0, - ..source - }; - assert!( - conn.create_materialized_view("bad", "people") - .vector_source(invalid_source) - .query() - .is_err() - ); - } - - /// No projection selects every source column, expanded now: the schema - /// captured at creation is the definition. - #[tokio::test] - async fn test_default_projection_captures_the_source_schema() { - let conn = people_db().await; - let view = conn - .create_materialized_view("copy", "people") - .execute() - .await - .unwrap(); - assert_eq!( - view.definition() - .projections - .iter() - .map(|p| p.output.as_str()) - .collect::>(), - vec!["name", "age"] - ); - } - - #[tokio::test] - async fn test_unknown_column_fails_at_create_time() { - let conn = people_db().await; - let err = conn - .create_materialized_view("bad", "people") - .select([("x", "missing + 1")]) - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "x")); - let names = conn.table_names().execute().await.unwrap(); - assert!(!names.contains(&"bad".to_string())); - } - - #[tokio::test] - async fn test_unknown_filter_column_fails_at_create_time() { - let err = declare_err(|b| b.only_if("missing > 1")).await; - assert!(matches!(err, Error::InvalidInput { message } if message.contains("missing"))); - } - - #[tokio::test] - async fn test_duplicate_output_is_rejected() { - let err = declare_err(|b| b.select([("dup", "age"), ("dup", "age + 1")])).await; - assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "dup")); - } - - #[tokio::test] - async fn test_reserved_output_name_is_rejected() { - let err = declare_err(|b| b.select([(SOURCE_ROW_ID_COLUMN, "age")])).await; - assert!(matches!(err, Error::InvalidInput { message } if message.contains("reserved"))); - } - - #[tokio::test] - async fn test_missing_source_fails() { - let conn = connect("memory://").execute().await.unwrap(); - let err = conn - .create_materialized_view("v", "nope") - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::TableNotFound { .. })); - } - - /// Provenance has to survive source compactions and updates, and stable - /// row ids cannot be enabled after a table exists -- so the requirement - /// is checked at the last moment the caller can still act on it. - #[tokio::test] - async fn test_source_without_stable_row_ids_is_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("x", Int32, [1, 2])).unwrap(); - conn.create_table("plain", batch).execute().await.unwrap(); - - let err = conn - .create_materialized_view("v", "plain") - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, Error::InvalidInput { message } if message.contains("stable row ids")) - ); - assert!( - !conn - .table_names() - .execute() - .await - .unwrap() - .contains(&"v".to_string()) - ); - } - - #[tokio::test] - async fn test_name_collision_fails() { - let conn = people_db().await; - let err = conn - .create_materialized_view("people", "people") - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::TableAlreadyExists { .. })); - } - - #[tokio::test] - async fn test_a_plain_table_is_not_a_view() { - let conn = people_db().await; - let table = conn.open_table("people").execute().await.unwrap(); - let err = MaterializedView::from_table(table).await.unwrap_err(); - assert!(matches!(err, Error::NotAMaterializedView { name } if name == "people")); - - let err = conn.open_materialized_view("people").await.unwrap_err(); - assert!(matches!(err, Error::NotAMaterializedView { .. })); - } - - /// The reason the kind is tagged: a definition written by a newer version - /// reads back as a view this one cannot refresh, not as a plain table. - #[tokio::test] - async fn test_unrecognized_kind_is_refused_by_name() { - let conn = people_db().await; - conn.create_materialized_view("v", "people") - .execute() - .await - .unwrap(); - let table = conn.open_table("v").execute().await.unwrap(); - table - .as_native() - .unwrap() - .replace_schema_metadata(HashMap::from([( - DEFINITION_META_KEY.to_string(), - r#"{"kind": "join"}"#.to_string(), - )])) - .await - .unwrap(); - - let err = conn.open_materialized_view("v").await.unwrap_err(); - assert!(matches!(err, Error::NotSupported { message } if message.contains("join"))); - } - - #[tokio::test] - async fn test_list_reports_views_and_only_views() { - let conn = people_db().await; - conn.create_materialized_view("adults", "people") - .only_if("age >= 18") - .execute() - .await - .unwrap(); - - let views = conn.list_materialized_views().await.unwrap(); - assert_eq!(views, vec!["adults"]); - let view = conn.open_materialized_view("adults").await.unwrap(); - assert_eq!(view.definition().filter.as_deref(), Some("age >= 18")); - } - - /// The creation option outranks a connection configured to create - /// unstable tables: the view still gets stable row ids, on the same - /// store (no fork -- the table must be reachable through the - /// connection afterwards). - #[tokio::test] - async fn test_view_is_stable_despite_connection_override() { - let conn = connect("memory://") - .storage_options([("new_table_enable_stable_row_ids", "false")]) - .execute() - .await - .unwrap(); - let batch = record_batch!(("x", Int32, [1])).unwrap(); - conn.create_table("src", batch) - .storage_option("new_table_enable_stable_row_ids", "true") - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("v", "src") - .execute() - .await - .unwrap(); - let stable = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .manifest - .uses_stable_row_ids(); - assert!(stable); - conn.open_materialized_view("v").await.unwrap(); - } - - /// A committed filter has to be usable as a predicate. - #[tokio::test] - async fn test_non_boolean_filter_is_rejected() { - let err = declare_err(|b| b.only_if("age + 1")).await; - assert!(matches!(err, Error::InvalidInput { message } if message.contains("boolean"))); - } - - /// Nested references stay dotted paths; resolution is by root field. - #[tokio::test] - async fn test_struct_columns_can_be_declared() { - use arrow_array::{ArrayRef, Int32Array, StructArray}; - - let conn = connect("memory://").execute().await.unwrap(); - let ages = StructArray::from(vec![( - Arc::new(ArrowField::new("age", DataType::Int32, false)), - Arc::new(Int32Array::from(vec![36, 17])) as ArrayRef, + #[test] + fn sql_definition_defaults_share_one_metadata_value() { + let metadata = HashMap::from([( + DEFINITION_META_KEY.to_string(), + definition_metadata_from_sql_with_defaults( + "SELECT * FROM source", + "db", + &["analytics".to_string()], + ), )]); - let batch = - arrow_array::RecordBatch::try_from_iter(vec![("metadata", Arc::new(ages) as ArrayRef)]) - .unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("ages", "people") - .select([("age", "metadata.age")]) - .only_if("metadata.age >= 18") - .execute() - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - assert_eq!( - schema.field_with_name("age").unwrap().data_type(), - &DataType::Int32 - ); - } - - /// A newer-kind view must not disappear from the listing. - #[tokio::test] - async fn test_unrecognized_kind_is_listed_by_name() { - let conn = people_db().await; - conn.create_materialized_view("v", "people") - .execute() - .await - .unwrap(); - let table = conn.open_table("v").execute().await.unwrap(); - table - .as_native() + let definition = read_definition_sql_with_defaults(&metadata) .unwrap() - .replace_schema_metadata(HashMap::from([( - DEFINITION_META_KEY.to_string(), - r#"{"kind": "join"}"#.to_string(), - )])) - .await .unwrap(); - - let views = conn.list_materialized_views().await.unwrap(); - assert_eq!(views, vec!["v"]); - } - - /// A definition must evaluate identically across refreshes; anything - /// less makes incremental maintenance a mix of evaluations. - #[tokio::test] - async fn test_volatile_and_unstable_expressions_are_rejected() { - let conn = people_db().await; - for expression in [ - "random()", - "now()", - "version()", - "arrow_typeof(age)", - "arrow_metadata(age, 'k')", - ] { - let err = conn - .create_materialized_view("bad", "people") - .select([("x", expression)]) - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, Error::InvalidExpression { message, .. } - if message.contains("not immutable")), - "{expression} was not rejected" - ); - } - for filter in ["age > random() * 100", "age >= 0 and now() is not null"] { - let err = conn - .create_materialized_view("bad", "people") - .only_if(filter) - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, Error::InvalidInput { message } if message.contains("not immutable")), - "{filter} was not rejected" - ); - } - } - - /// A column projected as itself stays the column it was: blob discovery - /// and the blob APIs key off field metadata, which a bare rebuild of the - /// field would drop. - #[tokio::test] - async fn test_identity_projection_keeps_field_metadata() { - let conn = connect("memory://").execute().await.unwrap(); - let schema = Arc::new(ArrowSchema::new_with_metadata( - vec![ - ArrowField::new("id", DataType::Int32, true), - crate::blob("payload", true), - ], - HashMap::new(), - )); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("v", "src") - .execute() - .await - .unwrap(); - let view_schema = view.table().schema().await.unwrap(); - - let payload = view_schema.field_with_name("payload").unwrap(); - assert!( - crate::blob::is_blob(payload), - "default projection dropped the blob marker: {:?}", - payload.metadata() - ); + assert_eq!(definition.query, "SELECT * FROM source"); + assert_eq!(definition.default_database.as_deref(), Some("db")); assert_eq!( - view.table().blob_columns().await.unwrap(), - vec!["payload".to_string()], - "blob discovery no longer finds the projected column" - ); - assert!(view_schema.metadata().contains_key(DEFINITION_META_KEY)); - // Structural declarations describe how a table is written; a view is - // written by refresh, and its fields are always nullable. - assert!(!view_schema.metadata().contains_key("lance:primary_key")); - - // A computed column is a new value and carries no source metadata. - let computed = conn - .create_materialized_view("c", "src") - .select([("payload", "payload"), ("n", "id + 1")]) - .execute() - .await - .unwrap(); - let computed_schema = computed.table().schema().await.unwrap(); - assert!(crate::blob::is_blob( - computed_schema.field_with_name("payload").unwrap() - )); - assert!( - computed_schema - .field_with_name("n") - .unwrap() - .metadata() - .is_empty() + definition.default_namespace_path, + Some(vec!["analytics".to_string()]) ); } - /// A nested column projected straight through is still that column, and a - /// declaration buried in a struct child binds as hard as one on top. #[tokio::test] - async fn test_nested_projection_metadata_and_declarations() { - // The schema below carries the legacy v1 blob marker, which Lance only - // allows writing at file version <= 2.1. - let conn = connect("memory://") - .storage_options([( - crate::database::listing::OPT_NEW_TABLE_STORAGE_VERSION, - "2.1", - )]) - .execute() - .await - .unwrap(); - let payload = crate::blob("payload", true).with_metadata(HashMap::from([ - ("lance-encoding:blob".to_string(), "true".to_string()), - ( - "lance-schema:unenforced-primary-key".to_string(), - "0".to_string(), - ), - ])); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, true), - ArrowField::new("meta", DataType::Struct(vec![payload].into()), true), - ])); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - // A nested path is a direct projection: the leaf's metadata comes with - // it, so the blob stays a blob rather than a plain struct. - let lifted = conn - .create_materialized_view("lifted", "src") - .select([("payload", "meta.payload")]) - .execute() - .await - .unwrap(); - let field = lifted.table().schema().await.unwrap(); - let field = field.field_with_name("payload").unwrap().clone(); - assert_eq!( - field.metadata().get("lance-encoding:blob"), - Some(&"true".to_string()), - "nested projection lost the leaf's metadata" - ); - assert!( - !field - .metadata() - .contains_key("lance-schema:unenforced-primary-key"), - "a structural declaration rode along" - ); - - // Projecting the struct whole must not carry the child's declaration - // out to a view whose fields are nullable. - let whole = conn - .create_materialized_view("whole", "src") - .select([("meta", "meta")]) - .execute() - .await - .unwrap(); - let schema = whole.table().schema().await.unwrap(); - let DataType::Struct(children) = schema.field_with_name("meta").unwrap().data_type() else { - panic!("meta is not a struct"); - }; - let child = children.iter().find(|c| c.name() == "payload").unwrap(); - assert!( - !child - .metadata() - .contains_key("lance-schema:unenforced-primary-key"), - "a nested declaration survived: {:?}", - child.metadata() - ); - assert_eq!( - child.metadata().get("lance-encoding:blob"), - Some(&"true".to_string()) - ); - } - - /// A map's entries are fields like any other, and a declaration on one - /// binds the view's writes just as hard as one on top. - #[tokio::test] - async fn test_map_declarations_are_stripped() { + async fn refresh_always_rebuilds_without_source_row_ids() { let conn = connect("memory://").execute().await.unwrap(); - let value = - ArrowField::new("value", DataType::Utf8, false).with_metadata(HashMap::from([( - "lance-schema:unenforced-clustering-key:position".to_string(), - "1".to_string(), - )])); - let entries = ArrowField::new( - "entries", - DataType::Struct(vec![ArrowField::new("key", DataType::Utf8, false), value].into()), - false, - ); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "props", - DataType::Map(Arc::new(entries), false), - true, - )])); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("view", "src") - .execute() - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - let DataType::Map(entries, _) = schema.field_with_name("props").unwrap().data_type() else { - panic!("props is not a map"); - }; - let DataType::Struct(children) = entries.data_type() else { - panic!("map entries are not a struct"); - }; - let value = children.iter().find(|c| c.name() == "value").unwrap(); - assert!( - !value - .metadata() - .contains_key("lance-schema:unenforced-clustering-key:position"), - "a declaration survived inside a map: {:?}", - value.metadata() - ); - } - - /// A computed column is declared by field metadata. Projecting one -- - /// as itself or under an alias -- must carry its description without its - /// declaration, which the target table would reject as foreign. - #[tokio::test] - async fn test_view_over_a_computed_column() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("id", Int32, [1, 2])).unwrap(); let source = conn - .create_table("src", batch) - .write_options(stable_row_ids()) + .create_table("src", record_batch!(("x", Int32, [1, 2])).unwrap()) .execute() .await .unwrap(); - source - .add_columns() - .computed("doubled", "id * 2") - .execute() - .await - .unwrap(); - - // Default projection reaches the computed column too. - let whole = conn - .create_materialized_view("whole", "src") - .execute() - .await - .unwrap(); - let schema = whole.table().schema().await.unwrap(); - let field = schema.field_with_name("doubled").unwrap(); - assert!( - !field - .metadata() - .keys() - .any(|k| k.starts_with("computed_column")), - "a computed-column declaration rode along: {:?}", - field.metadata() - ); - - // And under an alias. - conn.create_materialized_view("aliased", "src") - .select([("twice", "doubled")]) - .execute() - .await - .unwrap(); - } - - /// Embedding configuration names columns. It comes along only for the - /// columns a view actually projects, under the names the view gives them. - #[tokio::test] - async fn test_embedding_config_follows_the_projection() { - let config = r#"[{"name":"f","model":{},"source_column":"text","vector_column":"vec"}]"#; - let conn = connect("memory://").execute().await.unwrap(); - let schema = Arc::new(ArrowSchema::new_with_metadata( - vec![ - ArrowField::new("text", DataType::Utf8, true), - ArrowField::new("vec", DataType::Float32, true), - ], - HashMap::from([("embedding_functions".to_string(), config.to_string())]), - )); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let carried = |view: &MaterializedView| { - let view = view.table().clone(); - async move { - view.schema() - .await - .unwrap() - .metadata() - .get("embedding_functions") - .cloned() - } - }; - - // Both columns projected as themselves: kept as it stands. - let whole = conn - .create_materialized_view("whole", "src") - .execute() - .await - .unwrap(); - let kept = carried(&whole).await.expect("config dropped"); - assert!(kept.contains(r#""source_column":"text""#), "{kept}"); - assert!(kept.contains(r#""vector_column":"vec""#), "{kept}"); - - // Only the source column: the configuration names a vector column the - // view does not have, so it describes nothing and goes. - let partial = conn - .create_materialized_view("partial", "src") - .select([("text", "text")]) - .execute() - .await - .unwrap(); - assert_eq!(carried(&partial).await, None); - - // Renamed: the configuration follows the names the view uses. - let renamed = conn - .create_materialized_view("renamed", "src") - .select([("body", "text"), ("embedding", "vec")]) - .execute() - .await - .unwrap(); - let remapped = carried(&renamed).await.expect("config dropped"); - assert!(remapped.contains(r#""source_column":"body""#), "{remapped}"); - assert!( - remapped.contains(r#""vector_column":"embedding""#), - "{remapped}" - ); - - // The Node bindings spell the same configuration in camelCase, and - // the Rust definition names the destination `dest_column`. - for (config, source_key, dest_key) in [ - ( - r#"[{"name":"f","model":{},"sourceColumn":"text","vectorColumn":"vec"}]"#, - "sourceColumn", - "vectorColumn", - ), - ( - r#"[{"name":"f","model":{},"source_column":"text","dest_column":"vec"}]"#, - "source_column", - "dest_column", - ), - ] { - let schema = Arc::new(ArrowSchema::new_with_metadata( - vec![ - ArrowField::new("text", DataType::Utf8, true), - ArrowField::new("vec", DataType::Float32, true), - ], - HashMap::from([("embedding_functions".to_string(), config.to_string())]), - )); - let name = format!("src_{source_key}"); - conn.create_empty_table(&name, schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view(format!("v_{source_key}"), &name) - .select([("body", "text"), ("embedding", "vec")]) - .execute() - .await - .unwrap(); - let carried = carried(&view).await.expect("config dropped"); - assert!( - carried.contains(&format!(r#""{source_key}":"body""#)), - "{carried}" - ); - assert!( - carried.contains(&format!(r#""{dest_key}":"embedding""#)), - "{carried}" - ); - } - - // A computed column is not the source column under another name. - let computed = conn - .create_materialized_view("computed", "src") - .select([("body", "upper(text)"), ("embedding", "vec")]) - .execute() - .await - .unwrap(); - assert_eq!(carried(&computed).await, None); - - // One column projected twice is two columns in the view, and the - // configuration has to describe both rather than whichever came last. - let twice = conn - .create_materialized_view("twice", "src") - .select([("body", "text"), ("a", "vec"), ("b", "vec")]) - .execute() - .await - .unwrap(); - let carried = carried(&twice).await.expect("config dropped"); - let entries: Vec = serde_json::from_str(&carried).unwrap(); - let mut vectors: Vec<&str> = entries - .iter() - .filter_map(|e| e["vector_column"].as_str()) - .collect(); - vectors.sort_unstable(); - assert_eq!(vectors, ["a", "b"], "{carried}"); - } - - /// The native Rust producer records embeddings as column definitions - /// rather than as `embedding_functions`, and a query embeds its own text - /// through them. They are positional, so the view's list covers every one - /// of its fields. - #[tokio::test] - async fn test_native_column_definitions_follow_the_projection() { - let conn = connect("memory://").execute().await.unwrap(); - let rich = crate::table::TableDefinition::new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("text", DataType::Utf8, true), - ArrowField::new("vector", DataType::Float32, true), - ])), - vec![ - ColumnDefinition { - kind: ColumnKind::Physical, - }, - ColumnDefinition { - kind: ColumnKind::Embedding(EmbeddingDefinition::new( - "text", - "model", - Some("vector"), - )), - }, - ], - ) - .into_rich_schema(); - conn.create_empty_table("src", rich) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("view", "src") - .select([("body", "text"), ("embedding", "vector")]) + .create_materialized_view("copy", "src") .execute() .await .unwrap(); - let schema = view.table().schema().await.unwrap(); - let raw = schema - .metadata() - .get(COLUMN_DEFINITIONS_META_KEY) - .expect("the view dropped the native column definitions"); - let definitions: Vec = serde_json::from_str(raw).unwrap(); - assert_eq!( - definitions.len(), - schema.fields().len(), - "column definitions are positional" - ); - let ColumnKind::Embedding(embedding) = &definitions[1].kind else { - panic!("the embedding column came back physical: {raw}"); - }; - assert_eq!(embedding.source_column, "body"); - assert_eq!(embedding.dest_column.as_deref(), Some("embedding")); - assert_eq!(embedding.embedding_name, "model"); - assert!(matches!(definitions[0].kind, ColumnKind::Physical)); - assert!(matches!(definitions[2].kind, ColumnKind::Physical)); - // Without the column the function reads, the view cannot recompute - // the embedding, so it carries no definition for it. - let partial = conn - .create_materialized_view("partial", "src") - .select([("embedding", "vector")]) - .execute() - .await - .unwrap(); - assert_eq!( - partial - .table() + assert_eq!(view.table().count_rows(None).await.unwrap(), 2); + assert!( + view.table() .schema() .await .unwrap() - .metadata() - .get(COLUMN_DEFINITIONS_META_KEY), - None - ); - } - - /// A scan takes the cap as i64, so a larger one is refused where it is - /// declared rather than at the refresh that cannot run it. What a cap of - /// zero means is a refresh question, tested there. - #[tokio::test] - async fn test_limit_above_i64_max_is_refused_at_creation() { - let conn = people_db().await; - let err = conn - .create_materialized_view("too_big", "people") - .limit(i64::MAX as u64 + 1) - .execute() - .await - .unwrap_err(); - assert!( - matches!(&err, Error::InvalidInput { message } if message.contains("exceeds the maximum")), - "got {err:?}" + .field_with_name("__source_row_id") + .is_err() ); - // The boundary itself is accepted. - conn.create_materialized_view("at_max", "people") - .limit(i64::MAX as u64) + source + .add(record_batch!(("x", Int32, [3])).unwrap()) .execute() .await .unwrap(); - } + let refreshed = view.refresh().execute().await.unwrap(); + assert_eq!(refreshed.mode, RefreshMode::Rebuild); + assert_eq!(refreshed.rows_written, 3); + assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - #[tokio::test] - async fn test_drop_is_drop_table() { - let conn = people_db().await; - conn.create_materialized_view("v", "people") - .execute() - .await - .unwrap(); - conn.drop_table("v", &[]).await.unwrap(); - assert!(conn.list_materialized_views().await.unwrap().is_empty()); - } + let unchanged = view.refresh().execute().await.unwrap(); + assert_eq!(unchanged.mode, RefreshMode::Rebuild); + assert_eq!(unchanged.rows_written, 3); - /// The public declaration contract: prepare validates the source and - /// create consumes the declaration into a verified view table; a - /// source that cannot anchor refresh and a target outside the - /// source's database are both refused. - #[tokio::test] - async fn prepare_and_create_bind_the_declaration_lifecycle() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("id", Int32, [1, 2]), ("value", Int32, [3, 4])).unwrap(); - let source = conn - .create_table("src", batch.clone()) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let projections = [ - ("id".to_string(), "id".to_string()), - ("double".to_string(), "value * 2".to_string()), - ]; - let prepared = prepare_declaration(&source, Some(&projections), Some("value > 0"), None) - .await - .unwrap(); - assert_eq!(prepared.definition().source_table, "src"); - - let view = prepared.create("v").await.unwrap(); - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["id", "double", SOURCE_ROW_ID_COLUMN]); - assert!(schema.metadata().contains_key(DEFINITION_META_KEY)); - - // The same call rejects a source without stable row ids, so an - // external creation path cannot skip the check. - conn.create_table("plain", batch).execute().await.unwrap(); - let plain = conn.open_table("plain").execute().await.unwrap(); - let err = prepare_declaration(&plain, None, None, None) - .await - .unwrap_err(); - assert!(err.to_string().contains("stable row ids"), "{err}"); - - // A handle whose location does not resolve back through its name is - // refused: the definition would record a name reaching other data. - let plain_uri = plain - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - let masquerade = conn - .open_table("src") - .location(plain_uri) - .execute() - .await - .unwrap(); - let err = prepare_declaration(&masquerade, None, None, None) - .await - .unwrap_err(); - assert!( - err.to_string().contains("does not resolve to itself"), - "{err}" - ); - - // A table created at a custom location is refused outright: its - // recorded name reaches nothing at the database root, so the - // canonical reopen fails before any URI comparison. - let custom = conn - .create_table( - "custom_loc", - record_batch!(("id", Int32, [1, 2]), ("value", Int32, [3, 4])).unwrap(), - ) - .location("memory://elsewhere/custom_loc") - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let err = prepare_declaration(&custom, None, None, None) - .await - .unwrap_err(); - assert!(err.to_string().contains("custom_loc"), "{err}"); - } - - /// A view declared over a namespaced source records that namespace, and - /// refresh resolves the source through it -- the coordinate round-trips. - #[tokio::test] - async fn a_namespaced_source_round_trips_through_refresh() { - use lance_namespace::models::CreateNamespaceRequest; - - let tmp = tempfile::tempdir().unwrap(); - let mut properties = std::collections::HashMap::new(); - properties.insert("root".to_string(), tmp.path().to_str().unwrap().to_string()); - let conn = crate::connect_namespace("dir", properties) - .execute() - .await - .unwrap(); - conn.create_namespace(CreateNamespaceRequest { - id: Some(vec!["ns".into()]), - ..Default::default() - }) - .await - .unwrap(); - - let batch = record_batch!( - ("name", Utf8, ["ada", "grace", "alan"]), - ("age", Int32, [36, 85, 41]) - ) - .unwrap(); - conn.create_table("people", batch) - .namespace(vec!["ns".to_string()]) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - // A decoy of the same name at the root: resolving the source at the - // wrong namespace materializes one row here instead of three. - let decoy = record_batch!(("name", Utf8, ["mallory"]), ("age", Int32, [42])).unwrap(); - conn.create_table("people", decoy) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("adults", "people") - .with_no_data(true) - .namespace(vec!["ns".to_string()]) - .source_namespace(vec!["ns".to_string()]) - .select([("name", "name")]) - .only_if("age >= 18") - .execute() - .await - .unwrap(); - - assert_eq!(view.definition().source_table, "people"); - assert_eq!(view.definition().source_namespace, vec!["ns".to_string()]); - assert_eq!(view.table().namespace(), &["ns"]); - - // Refresh resolves the source at the recorded namespace, not at root. - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 3); - } - - fn definition(source_namespace: Vec) -> MaterializedViewDefinition { - MaterializedViewDefinition { - vector_source: None, - source_table: "people".to_string(), - source_namespace, - lateral: None, - projections: vec![ViewProjection { - output: "name".to_string(), - expression: "name".to_string(), - }], - filter: None, - group_by: Vec::new(), - limit: None, - } - } - - fn read(stored: impl Into) -> Result> { - read_definition(&HashMap::from([( - DEFINITION_META_KEY.to_string(), - stored.into(), - )])) - } - - /// What is stored is the query, under a format number; the same query - /// reads back whatever namespace the source sits in. - #[test] - fn the_stored_layout_is_the_canonical_query() { - for (namespace, query) in [ - (Vec::new(), "SELECT name FROM people"), - (vec!["ns".to_string()], "SELECT name FROM ns.people"), - ] { - let stored = definition_to_metadata(&definition(namespace.clone())).unwrap(); - let value: serde_json::Value = serde_json::from_str(&stored).unwrap(); - assert_eq!( - value, - serde_json::json!({"kind": "query", "format": 1, "query": query}) - ); - assert_eq!( - read(stored).unwrap(), - Some(StoredDefinition::Query(definition(namespace))) - ); - } - } - - /// The structured layout written before the format number still reads, - /// under both of its kind tags, and a tag that disagrees with its - /// namespace is an error: under `select` old readers resolved it at root. - #[test] - fn legacy_layouts_read_back() { - let legacy = |kind: &str, namespace: Vec<&str>| { - serde_json::json!({ - "kind": kind, - "source_table": "people", - "source_namespace": namespace, - "projections": [{"output": "name", "expression": "name"}], - "inputs": ["name"], - }) - .to_string() - }; - assert_eq!( - read(legacy(SELECT_KIND, vec![])).unwrap(), - Some(StoredDefinition::Query(definition(Vec::new()))) - ); - assert_eq!( - read(legacy(NAMESPACED_SELECT_KIND, vec!["ns"])).unwrap(), - Some(StoredDefinition::Query(definition(vec!["ns".into()]))) - ); - assert!( - read(r#"{"kind":"select","source_table":"people","projections":[]}"#) - .unwrap() - .is_some(), - "a pre-namespace definition carries no namespace key" - ); - for (kind, namespace) in [(SELECT_KIND, vec!["ns"]), (NAMESPACED_SELECT_KIND, vec![])] { - let err = read(legacy(kind, namespace)).unwrap_err(); - assert!( - err.to_string() - .contains("does not match its source namespace"), - "kind '{kind}': {err}" - ); - } - } - - /// A newer writer's definition is reported as such, never guessed at, - /// and a definition that is not a definition at all is an error rather - /// than a plain table. - #[test] - fn a_newer_format_is_reported_not_guessed() { - assert_eq!( - read(r#"{"format":5,"query":"SELECT name FROM people"}"#).unwrap(), - Some(StoredDefinition::Newer { format: "5".into() }) - ); - assert_eq!( - read(r#"{"kind":"join"}"#).unwrap(), - Some(StoredDefinition::Newer { - format: "kind 'join'".into() - }) - ); - for stored in [ - "{}", - r#"{"format":"one"}"#, - r#"{"format":1}"#, - r#"{"format":1,"query":"SELECT name FROM people GROUP BY name"}"#, - r#"{"format":2,"query":"SELECT name FROM people ORDER BY name"}"#, - ] { - assert!(read(stored).is_err(), "{stored}"); - } - } - - /// Planning records what the query reads of the source: a nested path - /// as itself, a read through an unnest alias as the list column. - #[test] - fn planning_records_the_source_columns_read() { - let element = DataType::Struct( - vec![ - ArrowField::new("chunk", DataType::Utf8, true), - ArrowField::new("ordinal", DataType::Int32, true), - ] - .into(), - ); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int64, false), - ArrowField::new( - "meta", - DataType::Struct(vec![ArrowField::new("title", DataType::Utf8, true)].into()), - true, - ), - ArrowField::new( - "chunks", - DataType::List(Arc::new(ArrowField::new("item", element, true))), - true, - ), - ])); - let planned = plan( - schema.clone(), - &MaterializedViewDefinition::from_sql( - "SELECT id, meta.title, c.chunk FROM docs, UNNEST(chunks) AS c WHERE c.ordinal < 3", - ) - .unwrap(), - None, - ) - .unwrap(); - assert_eq!(planned.inputs, ["chunks", "id", "meta.title"]); - assert_eq!( - planned - .fields - .iter() - .map(|f| f.name().as_str()) - .collect::>(), - ["id", "title", "chunk"] - ); - assert_eq!( - planned.fields[2].data_type(), - &DataType::Utf8, - "the element's field is read through the alias" - ); - - let star = plan( - schema, - &MaterializedViewDefinition::from_sql("SELECT * FROM docs, UNNEST(chunks) AS c") - .unwrap(), - None, - ) - .unwrap(); - assert_eq!( - star.definition.to_sql(), - "SELECT id, meta, c FROM docs, UNNEST(chunks) AS c" - ); - let err = plan( - Arc::new(ArrowSchema::new(vec![ArrowField::new( - "id", - DataType::Int64, - false, - )])), - &MaterializedViewDefinition::from_sql("SELECT id FROM docs, UNNEST(id) AS c").unwrap(), - None, - ) - .unwrap_err(); - assert!(err.to_string().contains("not a list"), "{err}"); - } - - /// A binding as the server records it: one Utf8 input over `input` - /// bound to a nullable parameter, one Int32 output named `output`, with - /// the exact schemas the durable contract requires. - pub fn test_binding(binding_id: &str, input: &str, output: &str) -> FunctionBinding { - let input_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new("text", DataType::Utf8, true), - ])) - .unwrap(); - let output_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new(output, DataType::Int32, true), - ])) - .unwrap(); - let input_type = input_schema.fields[0].r#type.r#type.clone(); - let output_type = output_schema.fields[0].r#type.r#type.clone(); - FunctionBinding::from_json( - &serde_json::json!({ - "binding_id": binding_id, - "function": {"name": "embed", "version": "1", "object_id": "fixture", "location": "memory:///fixture", - "manifest_digest": "sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs": [{ - "parameter": "text", "field_id": -1, "field_path": input, - "arrow_type": input_type, "nullable": true, - }], - "outputs": [{ - "result_field": "$value", "output_name": output, "output_field_id": -1, - "output_ordinal": 0, "arrow_type": output_type, "nullable": false, - }], - "input_schema": serde_json::to_value(input_schema).unwrap(), - "output_schema": serde_json::to_value(output_schema).unwrap(), - }) - .to_string(), - ) - .unwrap() - } - - /// A computed column as the server declares it on a table: bound to a - /// registered Function. - pub fn computed_field(name: &str, binding_id: &str, input: &str) -> ArrowField { - ArrowField::new(name, DataType::Int32, true).with_metadata( - crate::table::computed_columns::function_computed_column_metadata( - binding_id, - 0, - &[input.to_string()], - ), - ) - } - - pub async fn people(conn: &Connection) -> Table { - let batch = - record_batch!(("id", Int32, [1, 2, 3]), ("name", Utf8, ["a", "b", "c"])).unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap() - } - - /// `people` with both columns non-nullable, for nullability cases. - pub async fn strict_people(conn: &Connection) -> Table { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, false), - ArrowField::new("name", DataType::Utf8, false), - ])); - let batch = arrow_array::RecordBatch::try_new( - schema, - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3])), - Arc::new(arrow_array::StringArray::from(vec!["a", "b", "c"])), - ], - ) - .unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap() - } - - async fn prepared_people(conn: &Connection) -> PreparedDeclaration { - let source = people(conn).await; - prepare_declaration( - &source, - Some(&[ - ("id".to_string(), "id".to_string()), - ("name".to_string(), "name".to_string()), - ]), - None, - None, - ) - .await - .unwrap() - } - - #[tokio::test] - async fn a_computed_column_is_declared_null_with_its_binding() { - let conn = connect("memory://").execute().await.unwrap(); - let view = prepared_people(&conn) - .await - .with_computed_columns( - vec![(2, computed_field("emb", "fb_1", "name"))], - &[test_binding("fb_1", "name", "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["id", "name", "emb", SOURCE_ROW_ID_COLUMN]); - let declared: Vec = computed_columns(&schema) - .into_iter() - .map(|c| c.name) - .collect(); - assert_eq!(declared, ["emb"]); - let bindings = crate::table::computed_columns::function_bindings(&schema).unwrap(); - assert_eq!(bindings.len(), 1); - assert_eq!(bindings[0].binding_id(), "fb_1"); - // The stored definition is the query alone; bindings live beside it. - let stored: serde_json::Value = - serde_json::from_str(&schema.metadata()[DEFINITION_META_KEY]).unwrap(); - assert_eq!(stored["format"], 1); - assert_eq!(view.definition().projections.len(), 2); + source.delete("true").await.unwrap(); + let emptied = view.refresh().execute().await.unwrap(); + assert_eq!(emptied.mode, RefreshMode::Rebuild); + assert_eq!(emptied.rows_written, 0); assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - assert_eq!(conn.open_materialized_view("v").await.unwrap().name(), "v"); - } - - #[tokio::test] - async fn computed_column_declarations_are_validated() { - let conn = connect("memory://").execute().await.unwrap(); - let prepared = prepared_people(&conn).await; - let binding = test_binding("fb_1", "name", "emb"); - let fails = |prepared: PreparedDeclaration, - columns: Vec<(usize, ArrowField)>, - bindings: &[FunctionBinding]| { - prepared - .with_computed_columns(columns, bindings) - .err() - .map(|e| e.to_string()) - .expect("the declaration should be refused") - }; - let emb = |binding_id: &str| computed_field("emb", binding_id, "name"); - - let err = fails( - prepared.clone(), - vec![(2, emb("fb_1").with_nullable(false))], - std::slice::from_ref(&binding), - ); - assert!(err.contains("must be nullable"), "{err}"); - - let plain = ArrowField::new("emb", DataType::Int32, true); - let err = fails( - prepared.clone(), - vec![(2, plain)], - std::slice::from_ref(&binding), - ); - assert!( - err.contains("does not carry a computed-column declaration"), - "{err}" - ); - - // The rest is the computed-column contract: a binding the field does - // not name, an output the binding does not map to this field, an - // input the view does not hold. - let err = fails( - prepared.clone(), - vec![(2, emb("fb_other"))], - std::slice::from_ref(&binding), - ); - assert!(err.contains("does not match binding 'fb_1'"), "{err}"); - let err = fails( - prepared.clone(), - vec![(2, emb("fb_1"))], - &[test_binding("fb_1", "name", "different_output")], - ); - assert!(err.contains("different_output"), "{err}"); - let err = fails( - prepared.clone(), - vec![(2, emb("fb_1"))], - &[test_binding("fb_1", "bio", "emb")], - ); - assert!(err.contains("'bio'"), "{err}"); - - let err = fails( - prepared.clone(), - vec![(2, computed_field("name", "fb_1", "name"))], - &[test_binding("fb_1", "name", "name")], - ); - assert!(err.contains("already exists"), "{err}"); - let err = fails( - prepared.clone(), - vec![(2, computed_field(SOURCE_ROW_ID_COLUMN, "fb_1", "name"))], - &[test_binding("fb_1", "name", SOURCE_ROW_ID_COLUMN)], - ); - assert!(err.contains("reserved"), "{err}"); - let err = fails( - prepared.clone(), - vec![(7, emb("fb_1"))], - std::slice::from_ref(&binding), - ); - assert!( - err.contains("placed at 7, past the view's 2 columns"), - "{err}" - ); - let err = fails(prepared, Vec::new(), std::slice::from_ref(&binding)); - assert!(err.contains("at least one computed column"), "{err}"); - } - - /// A source column a computed column reads without the view projecting - /// it becomes an internal projection before the provenance column, with - /// the source's nullability; a projected column is read from its - /// projection. - /// `Some(&[])` is a declaration that projects nothing yet: a view of - /// computed columns alone, whose inputs `input_column` places. `None` - /// is `SELECT *`. The two must not collapse into each other. - #[tokio::test] - async fn an_empty_projection_list_is_not_a_star() { - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - - let mut prepared = prepare_declaration(&source, Some(&[]), None, None) - .await - .unwrap(); - assert!(prepared.definition().projections.is_empty()); - assert_eq!(prepared.input_column("name").unwrap(), "__input_name"); - let view = prepared - .with_computed_columns( - vec![(0, computed_field("emb", "fb_1", "__input_name"))], - &[test_binding("fb_1", "__input_name", "emb")], - ) - .unwrap() - .create("only") - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["emb", "__input_name", SOURCE_ROW_ID_COLUMN]); - assert_eq!( - view.definition().to_sql(), - "SELECT name AS __input_name FROM people" - ); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 3); - let reopened = conn.open_materialized_view("only").await.unwrap(); - assert_eq!(reopened.definition(), view.definition()); - - let all = prepare_declaration(&source, None, None, None) - .await - .unwrap(); - assert!( - !all.definition().selects_star(), - "planning expands the star" - ); - let outputs: Vec<&str> = all - .definition() - .projections - .iter() - .map(|p| p.output.as_str()) - .collect(); - assert_eq!(outputs, ["id", "name"]); - } - - #[tokio::test] - async fn an_unprojected_input_becomes_an_internal_projection() { - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - let mut prepared = prepare_declaration( - &source, - Some(&[("key".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap(); - assert_eq!(prepared.input_column("id").unwrap(), "key"); - assert_eq!(prepared.input_column("name").unwrap(), "__input_name"); - assert_eq!(prepared.input_column("name").unwrap(), "__input_name"); - let err = prepared.input_column("missing").unwrap_err().to_string(); - assert!(err.contains("no column 'missing'"), "{err}"); - - let view = prepared - .with_computed_columns( - vec![(1, computed_field("emb", "fb_1", "__input_name"))], - &[test_binding("fb_1", "__input_name", "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["key", "emb", "__input_name", SOURCE_ROW_ID_COLUMN]); - let input = schema.field_with_name("__input_name").unwrap(); - assert_eq!(input.data_type(), &DataType::Utf8); - assert!( - !input.is_nullable(), - "the copy keeps the source's nullability" - ); - assert!(!schema.field_with_name("key").unwrap().is_nullable()); - let projections: Vec<(&str, &str)> = view - .definition() - .projections - .iter() - .map(|p| (p.output.as_str(), p.expression.as_str())) - .collect(); - assert_eq!(projections, [("key", "id"), ("__input_name", "name")]); - } - - /// Two outputs of one binding land at consecutive positions: each - /// insertion widens the range the next may take. - #[tokio::test] - async fn sibling_computed_columns_take_consecutive_positions() { - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let prepared = prepare_declaration( - &source, - Some(&[("id".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap(); - let metadata = |ordinal: u32| { - crate::table::computed_columns::function_computed_column_metadata( - "fb_pair", - ordinal, - &["id".to_string()], - ) - }; - let left = ArrowField::new("left", DataType::Int32, true).with_metadata(metadata(0)); - let right = ArrowField::new("right", DataType::Int32, true).with_metadata(metadata(1)); - let input_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new("value", DataType::Int32, true), - ])) - .unwrap(); - let output_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new("left", DataType::Int32, true), - ArrowField::new("right", DataType::Int32, true), - ])) - .unwrap(); - let int = input_schema.fields[0].r#type.r#type.clone(); - let binding = FunctionBinding::from_json( - &serde_json::json!({ - "binding_id": "fb_pair", - "function": {"name": "pair", "version": "1", "object_id": "fixture", "location": "memory:///fixture", - "manifest_digest": "sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs": [{"parameter": "value", "field_id": -1, "field_path": "id", - "arrow_type": int, "nullable": true}], - "outputs": [ - {"result_field": "left", "output_name": "left", "output_field_id": -1, - "output_ordinal": 0, "arrow_type": int, "nullable": false}, - {"result_field": "right", "output_name": "right", "output_field_id": -1, - "output_ordinal": 1, "arrow_type": int, "nullable": false}, - ], - "input_schema": serde_json::to_value(input_schema).unwrap(), - "output_schema": serde_json::to_value(output_schema).unwrap(), - }) - .to_string(), - ) - .unwrap(); - let view = prepared - .with_computed_columns(vec![(1, left), (2, right)], &[binding]) - .unwrap() - .create("v") - .await - .unwrap(); - let names: Vec = view - .table() - .schema() - .await - .unwrap() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(names, ["id", "left", "right", SOURCE_ROW_ID_COLUMN]); - } - - /// A SQL declaration as `add_columns().computed()` records it. - pub fn sql_field( - name: &str, - data_type: DataType, - expression: &str, - inputs: &str, - ) -> ArrowField { - use crate::table::computed_columns::{ - COMPUTED_COLUMN_META_KEY, EXPRESSION_META_KEY, INPUTS_META_KEY, KIND_META_KEY, SQL_KIND, - }; - ArrowField::new(name, data_type, true).with_metadata(HashMap::from([ - (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), - (KIND_META_KEY.to_string(), SQL_KIND.to_string()), - (EXPRESSION_META_KEY.to_string(), expression.to_string()), - (INPUTS_META_KEY.to_string(), inputs.to_string()), - ])) - } - - /// A SQL declaration is re-planned at admission: it must parse against - /// the view, yield the declared type, and read the inputs it declares. - #[tokio::test] - async fn a_sql_declaration_is_planned_at_admission() { - let conn = connect("memory://").execute().await.unwrap(); - let fails = |prepared: PreparedDeclaration, field: ArrowField| { - prepared - .with_computed_columns(vec![(1, field)], &[]) - .err() - .map(|e| e.to_string()) - .expect("the declaration should be refused") - }; - let prepared = prepared_people(&conn).await; - let err = fails( - prepared.clone(), - sql_field("bad", DataType::Int32, "missing + 1", r#"["missing"]"#), - ); - assert!(err.contains("missing"), "{err}"); - let err = fails( - prepared.clone(), - sql_field("wide", DataType::Int64, "id + 1", r#"["id"]"#), - ); - assert!( - err.contains("declared as Int64 but its expression yields Int32"), - "{err}" - ); - let err = fails( - prepared.clone(), - sql_field("lying", DataType::Int32, "id + 1", r#"["name"]"#), - ); - assert!(err.contains("declares inputs"), "{err}"); - - let view = prepared - .with_computed_columns( - vec![(1, sql_field("next", DataType::Int32, "id + 1", r#"["id"]"#))], - &[], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let names: Vec = view - .table() - .schema() - .await - .unwrap() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(names, ["id", "next", "name", SOURCE_ROW_ID_COLUMN]); - } - - /// Creation persists a declaration only when it re-plans and the data - /// carries no values for it, whichever door created the table. - #[tokio::test] - async fn a_created_table_cannot_carry_computed_values() { - let conn = connect("memory://").execute().await.unwrap(); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("x", DataType::Int32, false), - sql_field("forged", DataType::Int32, "x + 1", r#"["x"]"#), - ])); - let filled = arrow_array::RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("forged", filled) - .execute() - .await - .unwrap_err() - .to_string(); - assert!(err.contains("cannot be written directly"), "{err}"); - assert!( - !conn - .table_names() - .execute() - .await - .unwrap() - .contains(&"forged".to_string()) - ); - - let unfilled = arrow_array::RecordBatch::try_new( - schema, - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::new_null(1)), - ], - ) - .unwrap(); - let table = conn - .create_table("declared", unfilled) - .execute() - .await - .unwrap(); - assert_eq!(table.refresh_column("forged").await.unwrap().rows_filled, 1); - - // A declaration with only its marker is broken, not absent. - let half = ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([( - crate::table::computed_columns::COMPUTED_COLUMN_META_KEY.to_string(), - "true".to_string(), - )])); - let batch = arrow_array::RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("x", DataType::Int32, false), - half, - ])), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("half", batch) - .execute() - .await - .unwrap_err() - .to_string(); - assert!( - err.contains("incomplete computed-column declaration"), - "{err}" - ); - - let bogus = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("x", DataType::Int32, false), - sql_field("bad", DataType::Int32, "missing + 1", r#"["missing"]"#), - ])); - let batch = arrow_array::RecordBatch::new_empty(bogus); - let err = conn - .create_table("bogus", batch) - .execute() - .await - .unwrap_err() - .to_string(); - assert!(err.contains("missing"), "{err}"); - } - - /// The internal-input prefix is reserved for the declaration, like the - /// provenance column, so an alias cannot masquerade as an internal input. - #[tokio::test] - async fn the_internal_input_prefix_is_reserved() { - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let err = prepare_declaration( - &source, - Some(&[("__input_x".to_string(), "name".to_string())]), - None, - None, - ) - .await - .unwrap_err() - .to_string(); - assert!(err.contains("'__input_x' is reserved"), "{err}"); - } - - /// A computed column may not read another, through any path: a - /// Function bound to a child of a computed struct is refused like a SQL - /// declaration over it. - #[tokio::test] - async fn a_computed_column_cannot_read_a_computed_root() { - let conn = connect("memory://").execute().await.unwrap(); - let prepared = prepared_people(&conn).await; - let payload = sql_field( - "payload", - DataType::Struct(vec![ArrowField::new("value", DataType::Utf8, true)].into()), - "named_struct('value', name)", - r#"["name"]"#, - ); - let err = prepared - .with_computed_columns( - vec![ - (2, payload), - (3, computed_field("emb", "fb_dependent", "payload.value")), - ], - &[test_binding("fb_dependent", "payload.value", "emb")], - ) - .err() - .map(|e| e.to_string()) - .expect("a computed root as a Function input should be refused"); - assert!(err.contains("reads computed column 'payload'"), "{err}"); - } - - /// The root check uses the canonical path parser: a quoted top-level - /// name containing a dot is one root, not two segments. - #[tokio::test] - async fn a_quoted_computed_root_is_still_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let prepared = prepared_people(&conn).await; - let payload = sql_field( - "payload.dot", - DataType::Struct(vec![ArrowField::new("value", DataType::Utf8, true)].into()), - "named_struct('value', name)", - r#"["name"]"#, - ); - let input = "`payload.dot`.value"; - let err = prepared - .with_computed_columns( - vec![ - (2, payload), - (3, computed_field("emb", "fb_dependent", input)), - ], - &[test_binding("fb_dependent", input, "emb")], - ) - .err() - .map(|e| e.to_string()) - .expect("a quoted computed root should be refused"); - assert!(err.contains("reads computed column 'payload.dot'"), "{err}"); - } - - /// Namespace-backed creation admits declarations by the same rule, and - /// refuses before the namespace records the table. - #[tokio::test] - async fn a_namespace_created_table_cannot_carry_computed_values() { - let tmp = tempfile::tempdir().unwrap(); - let mut properties = std::collections::HashMap::new(); - properties.insert("root".to_string(), tmp.path().to_str().unwrap().to_string()); - let conn = crate::connect_namespace("dir", properties) - .execute() - .await - .unwrap(); - let half = ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([( - crate::table::computed_columns::COMPUTED_COLUMN_META_KEY.to_string(), - "true".to_string(), - )])); - let batch = arrow_array::RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, false), - half, - ])), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("malformed", batch) - .execute() - .await - .unwrap_err() - .to_string(); - assert!( - err.contains("incomplete computed-column declaration"), - "{err}" - ); - assert!(conn.table_names().execute().await.unwrap().is_empty()); - - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, false), - sql_field("next", DataType::Int32, "id + 1", r#"["id"]"#), - ])); - let filled = arrow_array::RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("forged", filled) - .execute() - .await - .unwrap_err() - .to_string(); - assert!(err.contains("cannot be written directly"), "{err}"); - - let unfilled = arrow_array::RecordBatch::try_new( - schema, - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::new_null(1)), - ], - ) - .unwrap(); - let table = conn - .create_table("declared", unfilled) - .execute() - .await - .unwrap(); - assert_eq!(table.refresh_column("next").await.unwrap().rows_filled, 1); - } - - /// A projected column keeps its nullability; a computed value is - /// nullable. - #[tokio::test] - async fn an_identity_projection_keeps_source_nullability() { - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - let prepared = prepare_declaration( - &source, - Some(&[ - ("id".to_string(), "id".to_string()), - ("n".to_string(), "name".to_string()), - ("next".to_string(), "id + 1".to_string()), - ]), - None, - None, - ) - .await - .unwrap(); - let nullable: Vec = prepared - .schema() - .fields() - .iter() - .map(|f| f.is_nullable()) - .collect(); - assert_eq!(nullable, [false, false, true, false]); } } diff --git a/rust/lancedb/src/materialized_view/differential.rs b/rust/lancedb/src/materialized_view/differential.rs deleted file mode 100644 index 396b4e65b..000000000 --- a/rust/lancedb/src/materialized_view/differential.rs +++ /dev/null @@ -1,730 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The LanceDB Authors - -//! Differential refresh testing. -//! -//! The refresh contract is a property: after any sequence of source -//! mutations, a view maintained by default (incremental-where-possible) -//! refreshes equals the definition evaluated against the source directly, -//! and so does a forced rebuild. The oracle is an independent read of the -//! source -- plain column scan, filter applied in Rust -- so it shares -//! nothing with the refresh path it checks. -//! -//! The oracle runs after every step, not just at the end: a later mutation -//! that forces a rebuild would silently heal an incremental error, and those -//! transient errors are exactly the bugs this exists to catch. - -use arrow_array::{Float32Array, Int32Array, RecordBatch}; -use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; -use futures::{StreamExt, TryStreamExt}; -use lance::dataset::NewColumnTransform; -use std::sync::Arc; - -use super::MaterializedView; -use super::refresh::RefreshMode; -use crate::connect; -use crate::connection::Connection; -use crate::query::{ExecutableQuery, QueryBase, Select}; -use crate::table::{CompactionOptions, OptimizeAction, Table}; - -/// One source mutation, one per correctness-relevant class. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum SrcOp { - /// Fresh non-colliding ids of both parities, so every other op has - /// view-resident rows to act on: the only op that should refresh - /// incrementally. - AppendNew, - /// Deletion in surviving fragments must break the pure-append check. - DeleteEven, - /// An in-place update; on the filtered shape it crosses the predicate, - /// so rows must leave the view. - UpdateOddScore, - /// Fragment rewrite/renumber must break the pure-append check. - Compact, - /// A column the view does not read must NOT force a rebuild. - AddColumn, - /// merge_insert commits an Update whose by-source arm deletes rows, so a - /// classifier that reads Update as "changed only" loses those deletions. - MergeDropLargest, - /// merge_insert that both changes existing rows and inserts new ones in - /// one transaction. - MergeUpsert, -} - -const ALL_OPS: [SrcOp; 7] = [ - SrcOp::AppendNew, - SrcOp::DeleteEven, - SrcOp::UpdateOddScore, - SrcOp::Compact, - SrcOp::AddColumn, - SrcOp::MergeDropLargest, - SrcOp::MergeUpsert, -]; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum Shape { - /// SELECT id, score. - Identity, - /// SELECT id, score WHERE score > 50: additionally sensitive to rows - /// crossing the predicate. - Filtered, - /// SELECT id, score LIMIT 4. Which rows are held depends on the order - /// they were first materialized, so the oracle checks containment and - /// the cap rather than equality. - Limited, -} - -impl Shape { - fn filter(&self) -> Option<&'static str> { - match self { - Self::Identity | Self::Limited => None, - Self::Filtered => Some("score > 50"), - } - } - - fn matches(&self, score: f32) -> bool { - match self { - Self::Identity | Self::Limited => true, - Self::Filtered => score > 50.0, - } - } - - fn limit(&self) -> Option { - match self { - Self::Limited => Some(4), - _ => None, - } - } -} - -struct Case { - conn: Connection, - source: Table, - view: MaterializedView, - shape: Shape, - next_id: i32, - added_columns: u32, -} - -fn rows_batch(ids: &[i32]) -> RecordBatch { - let scores: Vec = ids.iter().map(|id| (*id * 10) as f32).collect(); - RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, true), - ArrowField::new("score", DataType::Float32, true), - ])), - vec![ - Arc::new(Int32Array::from(ids.to_vec())), - Arc::new(Float32Array::from(scores)), - ], - ) - .unwrap() -} - -fn merge_batch(ids: &[i32]) -> RecordBatch { - let scores: Vec = ids.iter().map(|id| (*id * 10 + 5) as f32).collect(); - RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, true), - ArrowField::new("score", DataType::Float32, true), - ])), - vec![ - Arc::new(Int32Array::from(ids.to_vec())), - Arc::new(Float32Array::from(scores)), - ], - ) - .unwrap() -} - -impl Case { - async fn new(shape: Shape) -> Self { - let conn = connect("memory://").execute().await.unwrap(); - let source = conn - .create_table("src", rows_batch(&[1, 2, 3, 4])) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let mut builder = conn - .create_materialized_view("view", "src") - .select([("id", "id"), ("score", "score")]); - if let Some(filter) = shape.filter() { - builder = builder.only_if(filter); - } - if let Some(limit) = shape.limit() { - builder = builder.limit(limit as u64); - } - let view = builder.execute().await.unwrap(); - Self { - conn, - source, - view, - shape, - next_id: 100, - added_columns: 0, - } - } - - async fn apply(&mut self, op: SrcOp) { - match op { - SrcOp::AppendNew => { - // Mixed parity: the middle id is odd, so UpdateOddScore always - // has a filter-matching appended row to evict. - let ids = vec![self.next_id, self.next_id + 101, self.next_id + 202]; - self.next_id += 303; - self.source.add(rows_batch(&ids)).execute().await.unwrap(); - } - SrcOp::DeleteEven => { - self.source.delete("id % 2 = 0").await.unwrap(); - } - SrcOp::UpdateOddScore => { - self.source - .update() - .column("score", "-1.0") - .only_if("id % 2 = 1") - .execute() - .await - .unwrap(); - } - SrcOp::Compact => { - self.source - .optimize(OptimizeAction::Compact { - options: CompactionOptions::default(), - remap_options: None, - }) - .await - .unwrap(); - } - SrcOp::MergeDropLargest => { - let mut ids = self.source_ids().await; - ids.sort_unstable(); - ids.pop(); - if ids.is_empty() { - return; - } - let batch = rows_batch(&ids); - let reader = - arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); - let mut merge = self.source.merge_insert(&["id"]); - merge.when_not_matched_by_source_delete(None); - merge.execute(Box::new(reader)).await.unwrap(); - } - SrcOp::MergeUpsert => { - let mut ids = self.source_ids().await; - ids.sort_unstable(); - // One row that exists (updated in place) and one that does not. - let existing = ids.first().copied().unwrap_or(self.next_id); - let fresh = self.next_id; - self.next_id += 1; - let batch = merge_batch(&[existing, fresh]); - let reader = - arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); - let mut merge = self.source.merge_insert(&["id"]); - merge - .when_matched_update_all(None) - .when_not_matched_insert_all(); - merge.execute(Box::new(reader)).await.unwrap(); - } - SrcOp::AddColumn => { - self.added_columns += 1; - let field = ArrowField::new( - format!("extra_{}", self.added_columns), - DataType::Int32, - true, - ); - self.source - .add_columns() - .transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( - vec![field], - )))) - .execute() - .await - .unwrap(); - } - } - } - - async fn source_ids(&self) -> Vec { - read_rows( - self.source - .query() - .select(Select::columns(&["id", "score"])), - ) - .await - .into_iter() - .map(|(id, _)| id) - .collect() - } - - /// The definition's result, read independently of the refresh path: - /// plain column scan, filter applied here, sorted. - async fn oracle(&self) -> Vec<(i32, i32)> { - let mut rows = read_rows( - self.source - .query() - .select(Select::columns(&["id", "score"])), - ) - .await - .into_iter() - .filter(|(_, score)| self.shape.matches(*score as f32)) - .collect::>(); - rows.sort_unstable(); - rows - } - - async fn view_rows(&self) -> Vec<(i32, i32)> { - let mut rows = read_rows( - self.view - .table() - .query() - .select(Select::columns(&["id", "score"])), - ) - .await; - rows.sort_unstable(); - rows - } - - async fn check(&self, label: &str) -> Result<(), String> { - let expected = self.oracle().await; - let actual = self.view_rows().await; - let Some(cap) = self.shape.limit() else { - if expected != actual { - return Err(format!( - "{label}: view diverged from oracle\n expected: {expected:?}\n actual: {actual:?}" - )); - } - return Ok(()); - }; - // A capped view holds some subset of the definition's result, never - // more than the cap, and never the same row twice. - if actual.len() > cap { - return Err(format!( - "{label}: view holds {} rows, over its cap of {cap}: {actual:?}", - actual.len() - )); - } - let mut unique = actual.clone(); - unique.dedup(); - if unique.len() != actual.len() { - return Err(format!("{label}: view holds a row twice: {actual:?}")); - } - if let Some(stray) = actual.iter().find(|row| !expected.contains(row)) { - return Err(format!( - "{label}: view holds {stray:?}, which the definition does not select: {expected:?}" - )); - } - // Below the cap the view must be complete, or a row was lost. - if actual.len() < cap.min(expected.len()) { - return Err(format!( - "{label}: view holds {} of {} selectable rows under a cap of {cap}: {actual:?}", - actual.len(), - expected.len() - )); - } - Ok(()) - } -} - -async fn read_rows(query: impl ExecutableQuery) -> Vec<(i32, i32)> { - let batches = query - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - batches - .iter() - .flat_map(|batch| { - let ids = batch["id"].as_any().downcast_ref::().unwrap(); - let scores = batch["score"] - .as_any() - .downcast_ref::() - .unwrap(); - // Scores are integer-valued by construction; compare exactly. - (0..batch.num_rows()) - .map(|i| (ids.value(i), scores.value(i) as i32)) - .collect::>() - }) - .collect() -} - -/// Drive one mutation sequence: refresh + oracle-check after every step, -/// then a forced rebuild checked against the same oracle. -async fn run_sequence(ops: &[SrcOp], shape: Shape) -> Result<(), String> { - let label = format!("{shape:?} {ops:?}"); - let mut case = Case::new(shape).await; - case.view - .refresh() - .execute() - .await - .map_err(|e| format!("{label}: initial refresh failed: {e}"))?; - case.check(&format!("{label} (initial)")).await?; - - for (step, op) in ops.iter().enumerate() { - case.apply(*op).await; - case.view - .refresh() - .execute() - .await - .map_err(|e| format!("{label}: refresh at step {step} failed: {e}"))?; - case.check(&format!("{label} (step {step}, {op:?})")) - .await?; - } - - case.view - .refresh() - .full(true) - .execute() - .await - .map_err(|e| format!("{label}: final full refresh failed: {e}"))?; - case.check(&format!("{label} (final rebuild)")).await?; - // Silence the unused-connection lint without dropping it mid-case. - let _ = &case.conn; - Ok(()) -} - -/// Every op sequence up to `max_len`. -fn all_sequences(max_len: u32) -> Vec> { - let mut sequences = Vec::new(); - for len in 1..=max_len { - for mut index in 0..ALL_OPS.len().pow(len) { - let mut ops = Vec::with_capacity(len as usize); - for _ in 0..len { - ops.push(ALL_OPS[index % ALL_OPS.len()]); - index /= ALL_OPS.len(); - } - sequences.push(ops); - } - } - sequences -} - -async fn run_exhaustive(max_len: u32) { - let mut cases = Vec::new(); - for shape in [Shape::Identity, Shape::Filtered, Shape::Limited] { - for ops in all_sequences(max_len) { - cases.push((ops, shape)); - } - } - let failures: Vec = futures::stream::iter(cases) - .map(|(ops, shape)| async move { run_sequence(&ops, shape).await.err() }) - .buffer_unordered(8) - .filter_map(|failure| async move { failure }) - .collect() - .await; - assert!( - failures.is_empty(), - "{} sequences diverged; first: {}", - failures.len(), - failures[0] - ); -} - -#[tokio::test(flavor = "multi_thread")] -async fn differential_exhaustive() { - run_exhaustive(3).await; -} - -#[tokio::test(flavor = "multi_thread")] -#[ignore = "longer sweep; run manually"] -async fn differential_exhaustive_deep() { - run_exhaustive(4).await; -} - -/// Named interleavings that double as repro handles. The mode assertions pin -/// the classifier, which value comparison alone cannot: a wrongly rebuilt -/// view still matches the oracle. -#[tokio::test(flavor = "multi_thread")] -async fn differential_named_regressions() { - // An append is the one op that must stay incremental. - let mut case = Case::new(Shape::Identity).await; - case.view.refresh().execute().await.unwrap(); - case.apply(SrcOp::AppendNew).await; - let result = case.view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - case.check("append stays incremental").await.unwrap(); - - // A column the view does not read must not force a rebuild. - let mut case = Case::new(Shape::Identity).await; - case.view.refresh().execute().await.unwrap(); - case.apply(SrcOp::AddColumn).await; - let result = case.view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - - // Compaction rearranges rows without changing them: the watermark - // advances and nothing rebuilds. - let mut case = Case::new(Shape::Identity).await; - case.view.refresh().execute().await.unwrap(); - case.apply(SrcOp::AppendNew).await; - case.view.refresh().execute().await.unwrap(); - case.apply(SrcOp::Compact).await; - let result = case.view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - case.check("compaction alone").await.unwrap(); - - // Fragment bookkeeping stays coherent across the compaction: the next - // append is separable and computed alone. - case.apply(SrcOp::AppendNew).await; - let result = case.view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 3); - case.check("compact then append").await.unwrap(); - - // A row updated to no longer match the filter must leave the view -- - // and the fixture must prove the eviction happened, not merely that the - // end state matches: an update that never touched a view-resident row - // would also "match". - let mut case = Case::new(Shape::Filtered).await; - case.apply(SrcOp::AppendNew).await; - case.view.refresh().execute().await.unwrap(); - let before = case.view_rows().await.len(); - case.apply(SrcOp::UpdateOddScore).await; - case.view.refresh().execute().await.unwrap(); - let after = case.view_rows().await.len(); - assert!( - after < before, - "no view-resident row was evicted ({before} -> {after}); the fixture \ - no longer exercises the filtered-update transition" - ); - case.check("update crosses the filter").await.unwrap(); -} - -// --------------------------------------------------------------------------- -// Concurrency -// --------------------------------------------------------------------------- -// -// The sequential cases above cannot observe a cross-process race: the -// per-view refresh lock is process-local, so a second refresh in this -// process queues behind the first. What is missing is not more op -// sequences but a second process. These cases add one, and assert the same -// property the harness always asserts -- the view holds each row once. - -/// Rows the definition selects from the source: every id but the first, -/// read straight from the source, sharing nothing with the refresh path. -async fn concurrency_oracle(conn: &Connection) -> Vec { - let batches: Vec = conn - .open_table("src") - .execute() - .await - .unwrap() - .query() - .select(Select::columns(&["id"])) - .execute() - .await - .unwrap() - .try_collect() - .await - .unwrap(); - let mut ids = Vec::new(); - for batch in &batches { - let column = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - if column.value(i) > 1 { - ids.push(column.value(i)); - } - } - } - ids.sort_unstable(); - ids -} - -/// The view's ids, sorted. -async fn concurrency_view_ids(conn: &Connection) -> Vec { - let batches: Vec = conn - .open_table("mv") - .execute() - .await - .unwrap() - .query() - .select(Select::columns(&["id"])) - .execute() - .await - .unwrap() - .try_collect() - .await - .unwrap(); - let mut ids = Vec::new(); - for batch in &batches { - let column = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - ids.push(column.value(i)); - } - } - ids.sort_unstable(); - ids -} - -/// One refresh of the view at `MV_RACE_DIR`, in its own process. -/// -/// Setup happens before the start barrier so warm-up does not stagger the -/// two processes. What makes the race certain rather than likely is the -/// second barrier inside `refresh()` itself, which holds every participant -/// between staging and commit. -#[tokio::test] -#[ignore = "spawned as a child process by the concurrency cases"] -async fn cross_process_refresh_child() { - let Ok(dir) = std::env::var("MV_RACE_DIR") else { - return; - }; - let dir = std::path::PathBuf::from(dir); - let tag = std::env::var("MV_RACE_TAG").unwrap(); - - let conn = connect(dir.to_str().unwrap()).execute().await.unwrap(); - let table = conn.open_table("mv").execute().await.unwrap(); - let _ = table.schema().await.unwrap(); - let _ = table.count_rows(None).await.unwrap(); - let source = conn.open_table("src").execute().await.unwrap(); - let _ = source.count_rows(None).await.unwrap(); - let view = MaterializedView::from_table(table).await.unwrap(); - - std::fs::write(dir.join(format!("ready-{tag}")), b"1").unwrap(); - while !dir.join("START").exists() { - std::thread::sleep(std::time::Duration::from_millis(2)); - } - - let outcome = match view.refresh().execute().await { - Ok(result) => format!("committed rows={}", result.rows_written), - Err(err) if is_commit_conflict(&err) => "conflicted".to_string(), - Err(err) => format!("failed {err}"), - }; - std::fs::write(dir.join(format!("outcome-{tag}")), outcome).unwrap(); -} - -/// Whether a refresh lost its commit to a concurrent one, as opposed to -/// failing for any other reason. -fn is_commit_conflict(err: &crate::Error) -> bool { - let text = err.to_string(); - text.contains("Retryable commit conflict") || text.contains("preempted by concurrent") -} - -/// Two processes refreshing one view concurrently must leave the view -/// equal to the oracle: each selected row present exactly once. -/// -/// Both plan the same incremental delta from one watermark. A refresh is -/// meant to land on the generation it planned or leave nothing behind, so -/// at most one of them may write. -#[tokio::test] -async fn concurrent_refreshes_hold_each_row_once() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().to_str().unwrap().to_string(); - let conn = connect(&path).execute().await.unwrap(); - conn.create_table("src", rows_batch(&[1, 2, 3, 4])) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("mv", "src") - .select([("id", "id"), ("score", "score")]) - .only_if("id > 1") - .execute() - .await - .unwrap(); - // Seed the watermark so the racing refreshes are both incremental. - view.refresh().execute().await.unwrap(); - - // Large enough that a refresh is real work rather than a formality. - let ids: Vec = (100..200_100).collect(); - conn.open_table("src") - .execute() - .await - .unwrap() - .add(rows_batch(&ids)) - .execute() - .await - .unwrap(); - - let tags = ["a", "b"]; - let exe = std::env::current_exe().unwrap(); - let children: Vec = tags - .iter() - .map(|tag| { - std::process::Command::new(&exe) - .args([ - "--exact", - "materialized_view::differential::cross_process_refresh_child", - "--ignored", - "--nocapture", - ]) - .env("MV_RACE_DIR", dir.path()) - .env("MV_RACE_SYNC", dir.path()) - .env("MV_RACE_PEERS", "2") - .env("MV_RACE_TAG", tag) - .stdout(std::process::Stdio::null()) - .stderr(std::process::Stdio::null()) - .spawn() - .unwrap() - }) - .collect(); - - let deadline = std::time::Instant::now() + std::time::Duration::from_secs(180); - while tags - .iter() - .any(|tag| !dir.path().join(format!("ready-{tag}")).exists()) - { - assert!( - std::time::Instant::now() < deadline, - "children never became ready" - ); - std::thread::sleep(std::time::Duration::from_millis(10)); - } - std::fs::write(dir.path().join("START"), b"1").unwrap(); - for (tag, mut child) in tags.iter().zip(children) { - let status = loop { - match child.try_wait().unwrap() { - Some(status) => break status, - None if std::time::Instant::now() >= deadline => { - child.kill().unwrap(); - panic!("child {tag} never finished"); - } - None => std::thread::sleep(std::time::Duration::from_millis(10)), - } - }; - assert!(status.success(), "child {tag} exited {status}"); - } - - // Both refreshes reached the commit boundary before either committed -- - // the in-refresh barrier guarantees it -- so exactly one may win. - let outcomes: Vec = tags - .iter() - .map(|tag| { - std::fs::read_to_string(dir.path().join(format!("outcome-{tag}"))) - .unwrap_or_else(|_| panic!("child {tag} recorded no outcome")) - }) - .collect(); - for tag in tags { - assert!( - dir.path().join(format!("planned-{tag}")).exists(), - "child {tag} never reached the commit boundary, so nothing was synchronized" - ); - } - let committed = outcomes.iter().filter(|o| o.contains("committed")).count(); - let conflicted = outcomes.iter().filter(|o| o.contains("conflicted")).count(); - assert_eq!( - (committed, conflicted), - (1, 1), - "exactly one refresh may win the generation both planned: {outcomes:?}" - ); - - let expected = concurrency_oracle(&conn).await; - let actual = concurrency_view_ids(&conn).await; - assert_eq!( - actual.len(), - expected.len(), - "the view holds {} rows, the oracle {}: a losing refresh left rows behind", - actual.len(), - expected.len() - ); - assert_eq!(actual, expected, "the view does not match the oracle"); -} diff --git a/rust/lancedb/src/materialized_view/duplicate_pairs.rs b/rust/lancedb/src/materialized_view/duplicate_pairs.rs deleted file mode 100644 index d0dc219ad..000000000 --- a/rust/lancedb/src/materialized_view/duplicate_pairs.rs +++ /dev/null @@ -1,478 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The LanceDB Authors - -//! Native pair generation as partitioned materialized-view refresh work. - -use std::sync::Arc; - -use arrow_schema::{Schema, SchemaRef}; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use datafusion_physical_plan::ExecutionPlan; -use datafusion_sql::sqlparser::ast::{ - Expr, FunctionArg, FunctionArgExpr, TableFunctionArgs, Value, -}; -use futures::TryStreamExt; -use lance::Dataset; -use lance::dataset::transaction::Operation; -use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams}; -use lance_table::format::Fragment; -use serde::{Deserialize, Serialize}; - -use super::grouped_units::GroupedRefreshPlan; -use super::refresh::{ - RefreshMaterializedViewResult, ensure_incarnation, ensure_no_mem_wal, open_source, -}; -use super::{MaterializedViewDefinition, Planned, ViewProjection}; -use crate::table::Table; -use crate::table::datafusion::udtf::duplicate_pairs::{ - DUPLICATE_PAIRS_OPERATOR_VERSION, DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec, - duplicate_pair_schema, plan_duplicate_pairs, -}; -use crate::{Error, Result}; - -pub(super) const FUNCTION_NAME: &str = "vector_duplicate_pairs"; -pub(super) const DEDUP_FUNCTION_NAME: &str = "vector_dedup"; - -/// The result of a pinned native vector source. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -pub enum VectorSourceKind { - /// The qualifying pair relation. - Pairs, - /// Original source rows retained by greedy direct representatives. - Dedup, -} - -/// A native table-function source, pinned for the lifetime of the view. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct VectorSource { - pub kind: VectorSourceKind, - pub dataset_version: u64, - pub column: String, - /// Validated numeric SQL spelling; execution and task identity use f32 bits. - pub distance_threshold: String, -} - -impl VectorSource { - pub(super) fn function_name(&self) -> &'static str { - match self.kind { - VectorSourceKind::Pairs => FUNCTION_NAME, - VectorSourceKind::Dedup => DEDUP_FUNCTION_NAME, - } - } - pub(crate) fn config(&self) -> Result { - let threshold: f32 = self - .distance_threshold - .parse() - .map_err(|_| invalid("invalid pair threshold"))?; - if self.dataset_version == 0 || self.column.is_empty() || !threshold.is_finite() { - return Err(invalid( - "duplicate pairs require a positive snapshot, a column and a finite threshold", - )); - } - Ok(DuplicatePairsConfig { - dataset_version: self.dataset_version, - column: self.column.clone(), - distance_threshold: threshold, - }) - } -} - -fn invalid(message: impl Into) -> Error { - Error::InvalidInput { - message: message.into(), - } -} - -pub(super) fn source( - args: &TableFunctionArgs, - kind: VectorSourceKind, -) -> Result<(Vec, String, VectorSource)> { - if args.settings.is_some() { - return Err(invalid("duplicate pairs do not accept function settings")); - } - let exprs = args - .args - .iter() - .map(|arg| match arg { - FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Ok(expr), - _ => Err(invalid( - "duplicate pairs require four positional literal arguments", - )), - }) - .collect::>>()?; - let [table, version, column, threshold] = exprs.as_slice() else { - return Err(invalid( - "native vector sources require (table, dataset_version, column, distance_threshold)", - )); - }; - let string = |expr: &Expr| match expr { - Expr::Value(value) => match &value.value { - Value::SingleQuotedString(value) => Ok(value.clone()), - _ => Err(invalid("table and column must be string literals")), - }, - _ => Err(invalid("table and column must be string literals")), - }; - let table = string(table)?; - let plain = MaterializedViewDefinition::from_sql(&format!("SELECT * FROM {table}"))?; - if plain.vector_source.is_some() - || plain.lateral.is_some() - || plain.filter.is_some() - || plain.is_grouped() - || plain.limit.is_some() - { - return Err(invalid("duplicate pairs must name one source table")); - } - let dataset_version = match version { - Expr::Value(value) => match &value.value { - Value::Number(number, _) => number - .parse() - .map_err(|_| invalid("invalid source version"))?, - _ => return Err(invalid("source version must be an integer literal")), - }, - _ => return Err(invalid("source version must be an integer literal")), - }; - let config = VectorSource { - kind, - dataset_version, - column: string(column)?, - distance_threshold: threshold.to_string(), - }; - config.config()?; - Ok((plain.source_namespace, plain.source_table, config)) -} - -pub(super) fn check_shape(definition: &MaterializedViewDefinition) -> Result<()> { - if definition - .vector_source - .as_ref() - .is_some_and(|s| s.kind == VectorSourceKind::Dedup) - { - if !definition.selects_star() - || definition.filter.is_some() - || definition.limit.is_some() - || definition.lateral.is_some() - || definition.is_grouped() - { - return Err(invalid( - "a vector_dedup view must select * without other clauses", - )); - } - return Ok(()); - } - let schema = duplicate_pair_schema(); - let projections = schema - .fields() - .iter() - .map(|field| ViewProjection { - output: field.name().clone(), - expression: field.name().clone(), - }) - .collect::>(); - if (!definition.selects_star() && definition.projections != projections) - || definition.filter.is_some() - || definition.limit.is_some() - || definition.lateral.is_some() - || definition.is_grouped() - { - return Err(invalid( - "a native pair view selects *, or row_id_a, row_id_b, distance, without other clauses", - )); - } - Ok(()) -} - -pub(super) fn plan( - source_schema: SchemaRef, - definition: &MaterializedViewDefinition, -) -> Result { - check_shape(definition)?; - let config = definition.vector_source.as_ref().expect("native source"); - config.config()?; - source_schema.field_with_name(&config.column)?; - if config.kind == VectorSourceKind::Dedup { - return Ok(Planned { - definition: definition.clone(), - fields: source_schema - .fields() - .iter() - .map(|f| super::without_declarations(f)) - .collect(), - lineage: source_schema - .fields() - .iter() - .map(|f| (f.name().clone(), vec![f.name().clone()])) - .collect(), - inputs: source_schema - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(), - }); - } - let schema = duplicate_pair_schema(); - let projections = schema - .fields() - .iter() - .map(|field| ViewProjection { - output: field.name().clone(), - expression: field.name().clone(), - }) - .collect(); - let mut definition = definition.clone(); - let inputs = vec![config.column.clone()]; - definition.projections = projections; - Ok(Planned { - definition, - fields: schema.fields().iter().map(|f| f.as_ref().clone()).collect(), - lineage: Default::default(), - inputs, - }) -} - -/// A saved native refresh plan, including every physical partition identity. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct DuplicatePairsRefreshPlan { - pub snapshot: GroupedRefreshPlan, - pub source_uri: String, - pub definition_sql: String, - pub operator_version: u32, - pub threshold_bits: u32, - pub tasks: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WrittenPairsPartition { - pub(super) unit: u32, - pub(super) plan_id: String, - pub(super) task: DuplicatePairTask, - pub(super) fragments: Vec, - pub(super) rows: u64, -} - -impl DuplicatePairsRefreshPlan { - pub(super) fn identity(&self) -> Result { - use sha2::{Digest, Sha256}; - let encoded = serde_json::to_vec(&( - &self.snapshot, - &self.source_uri, - &self.definition_sql, - self.operator_version, - self.threshold_bits, - )) - .map_err(|e| invalid(e.to_string()))?; - Ok(format!("{:x}", Sha256::digest(encoded))) - } -} - -async fn open(view: &Table) -> Result<(Dataset, MaterializedViewDefinition, Arc)> { - let native = view - .as_native() - .ok_or_else(|| invalid("native pair refresh requires a local view"))?; - native.dataset.ensure_mutable()?; - native.dataset.reload().await?; - let view_ds = native.dataset.get().await?.as_ref().clone(); - let Some(super::StoredDefinition::Query(definition)) = - super::read_definition(&view_ds.schema().metadata)? - else { - return Err(invalid("view does not carry a supported definition")); - }; - let config = definition - .vector_source - .as_ref() - .ok_or_else(|| invalid("not a native pair view"))? - .config()?; - let source = Arc::new( - open_source(view, &definition, None) - .await? - .checkout_version(config.dataset_version) - .await?, - ); - ensure_no_mem_wal(&source, "source table", &definition.source_table).await?; - let planned = plan(Arc::new(Schema::from(source.schema())), &definition)?; - let expected = Schema::new(planned.fields); - let physical = Schema::from(view_ds.schema()); - if physical.fields().len() != expected.fields().len() - || physical - .fields() - .iter() - .zip(expected.fields()) - .any(|(a, b)| { - a.name() != b.name() - || a.data_type() != b.data_type() - || a.is_nullable() != b.is_nullable() - }) - { - return Err(invalid("native pair view schema changed since creation")); - } - Ok((view_ds, planned.definition, source)) -} - -pub(super) async fn plan_refresh( - view: &Table, - pinned: Option, -) -> Result { - let (view_ds, definition, source) = open(view).await?; - let config = definition - .vector_source - .as_ref() - .expect("native") - .config()?; - if pinned.is_some_and(|v| v != config.dataset_version) { - return Err(invalid( - "refresh source version differs from the native pair view's fixed snapshot", - )); - } - let tasks = plan_duplicate_pairs(source.clone(), &config).await?; - let units = u32::try_from(tasks.len()).map_err(|_| invalid("too many index partitions"))?; - let incarnation = view_ds - .schema() - .metadata - .get(super::INCARNATION_META_KEY) - .cloned() - .ok_or_else(|| invalid("native pair view has no incarnation"))?; - Ok(DuplicatePairsRefreshPlan { - snapshot: GroupedRefreshPlan { - units, - source_version: config.dataset_version, - view_version: view_ds.version().version, - incarnation, - }, - source_uri: source.uri().to_string(), - definition_sql: definition.to_sql(), - operator_version: DUPLICATE_PAIRS_OPERATOR_VERSION, - threshold_bits: config.distance_threshold.to_bits(), - tasks, - }) -} - -pub(super) async fn open_planned( - view: &Table, - plan: &DuplicatePairsRefreshPlan, -) -> Result<(Dataset, MaterializedViewDefinition, Arc)> { - let (view_ds, definition, source) = open(view).await?; - ensure_incarnation(&view_ds, Some(&plan.snapshot.incarnation), view.name()).await?; - let config = definition - .vector_source - .as_ref() - .expect("native") - .config()?; - if source.uri() != plan.source_uri - || source.version().version != plan.snapshot.source_version - || definition.to_sql() != plan.definition_sql - || config.distance_threshold.to_bits() != plan.threshold_bits - || plan.operator_version != DUPLICATE_PAIRS_OPERATOR_VERSION - || plan.tasks.len() != plan.snapshot.units as usize - { - return Err(invalid( - "native pair task belongs to a different snapshot, configuration or operator", - )); - } - Ok((view_ds, definition, source)) -} - -pub(super) async fn write_unit( - view: &Table, - unit: u32, - plan: &DuplicatePairsRefreshPlan, -) -> Result { - let (view_ds, definition, source) = open_planned(view, plan).await?; - let task = plan - .tasks - .get(unit as usize) - .ok_or_else(|| invalid("pair partition is outside the plan"))? - .clone(); - let config = definition - .vector_source - .as_ref() - .expect("native") - .config()?; - let exec = DuplicatePairsExec::try_new(source, config, vec![task.clone()])?; - let stream = exec.execute(0, Arc::new(TaskContext::default()))?; - let rows = Arc::new(std::sync::atomic::AtomicU64::new(0)); - let count = rows.clone(); - let stream = stream.inspect_ok(move |batch| { - count.fetch_add( - batch.num_rows() as u64, - std::sync::atomic::Ordering::Relaxed, - ); - }); - let stream: SendableRecordBatchStream = Box::pin( - datafusion_physical_plan::stream::RecordBatchStreamAdapter::new( - duplicate_pair_schema(), - stream, - ), - ); - let txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { fragments } = txn.operation else { - return Err(invalid( - "native pair partition did not stage append fragments", - )); - }; - Ok(WrittenPairsPartition { - unit, - plan_id: plan.identity()?, - task, - fragments, - rows: rows.load(std::sync::atomic::Ordering::Relaxed), - }) -} - -pub(super) async fn commit( - view: &Table, - plan: &DuplicatePairsRefreshPlan, - units: Vec, - expected: Option<&str>, -) -> Result { - let (_, definition, source) = open_planned(view, plan).await?; - let config = definition - .vector_source - .as_ref() - .expect("native") - .config()?; - if plan_duplicate_pairs(source.clone(), &config).await? != plan.tasks { - return Err(invalid( - "pair plan does not cover the pinned index partitions exactly", - )); - } - let identity = plan.identity()?; - let mut ordered = vec![None; plan.snapshot.units as usize]; - for written in units { - if written.plan_id != identity - || plan.tasks.get(written.unit as usize) != Some(&written.task) - { - return Err(invalid("pair partition receipt belongs to another plan")); - } - let slot = ordered - .get_mut(written.unit as usize) - .ok_or_else(|| invalid("unknown pair partition receipt"))?; - if slot.is_some() { - return Err(invalid("duplicate pair partition receipt")); - } - *slot = Some(written); - } - let units = ordered - .into_iter() - .collect::>>() - .ok_or_else(|| invalid("missing pair partition receipt"))?; - let rows = units.iter().map(|u| u.rows).sum(); - let fragments = units.into_iter().flat_map(|u| u.fragments).collect(); - super::grouped_units::commit_fragments( - view, - &plan.snapshot, - &definition, - source.manifest.timestamp_nanos, - fragments, - rows, - expected, - ) - .await -} diff --git a/rust/lancedb/src/materialized_view/grouped.rs b/rust/lancedb/src/materialized_view/grouped.rs index 095f20f11..3120e2b43 100644 --- a/rust/lancedb/src/materialized_view/grouped.rs +++ b/rust/lancedb/src/materialized_view/grouped.rs @@ -12,7 +12,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use arrow_array::cast::AsArray; use arrow_array::{Array, FixedSizeListArray, UInt32Array}; use arrow_buffer::NullBuffer; -use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; +use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef}; use datafusion::catalog::default_table_source::provider_as_source; use datafusion::datasource::empty::EmptyTable; use datafusion::execution::FunctionRegistry; @@ -36,20 +36,16 @@ use datafusion_sql::sqlparser::parser::Parser; use futures::StreamExt; use lance::Dataset; use lance::index::{DatasetIndexExt, DatasetIndexInternalExt}; -use lance_core::ROW_ID; use lance_datafusion::exec::SessionContextExt; use lance_index::metrics::NoOpMetricsCollector; -use lance_index::vector::VectorIndex; use lance_index::vector::ivf::{IvfTransformer, new_ivf_transformer}; use lance_linalg::distance::DistanceType; use lance_linalg::kernels::normalize_fsl; -use roaring::RoaringBitmap; use uuid::Uuid; use super::refresh::to_view_batch; use super::{ - MaterializedViewDefinition, Planned, SOURCE_ROW_ID_COLUMN, ViewProjection, ensure_immutable, - plan_filter, query, without_declarations, + MaterializedViewDefinition, Planned, ViewProjection, plan_filter, query, without_declarations, }; use crate::{Error, Result}; @@ -68,50 +64,6 @@ pub(super) struct IvfBinding { pub(super) transformer: Arc, /// A cosine index assigns L2 over unit vectors, as lance's index path does. pub(super) normalize: bool, - /// Partitions of the model; the ids `ivf_partition` returns are below it. - pub(super) partitions: u32, -} - -/// The segments of the one IVF index on a column, opened, and the fragments -/// they cover between them; every other fragment's rows are unindexed. -/// `covered` is `None` when a segment does not say which fragments it holds -/// (`fragment_bitmap: None`, which lance defines as unknown rather than -/// empty): its rows would be read once by the partition reader and again as -/// unindexed, so no caller may split this index. -pub(super) struct IvfSegments { - pub(super) segments: Vec, - pub(super) covered: Option, -} - -/// One opened segment and the fragments it holds rows for. Its postings -/// still name rows of fragments it has since lost -- a column rewrite -/// attaches a new file and takes the fragment out of this bitmap -- so a -/// reader of its partitions keeps only the rows it still owns. -pub(super) struct IvfSegment { - pub(super) index: Arc, - pub(super) fragments: Option, -} - -impl IvfSegments { - pub(super) fn is_unindexed(&self, fragment_id: u64) -> Option { - Some(!self.covered.as_ref()?.contains(fragment_id as u32)) - } -} - -/// Coverage after folding in one more segment. Lance defines a segment -/// without a fragment bitmap as unknown coverage, not empty, and one such -/// segment is enough to lose the coverage of the whole index. -fn coverage( - covered: Option, - bitmap: Option<&RoaringBitmap>, -) -> Option { - match (covered, bitmap) { - (Some(mut covered), Some(bitmap)) => { - covered |= bitmap; - Some(covered) - } - _ => None, - } } /// `bindings` maps a source column to its index; planning runs unbound. @@ -278,16 +230,6 @@ pub(super) async fn bind( } async fn bind_column(source: &Dataset, column: &str) -> Result { - bind_column_with_segments(source, column) - .await - .map(|(binding, _)| binding) -} - -/// The binding plus the opened segments it was read from. -pub(super) async fn bind_column_with_segments( - source: &Dataset, - column: &str, -) -> Result<(IvfBinding, IvfSegments)> { let field = source .schema() .field(column) @@ -298,8 +240,6 @@ pub(super) async fn bind_column_with_segments( // on its own fragments carries its own centroids, and grouping every row // by one segment's model would bucket the other segments' rows wrongly. let mut found: Option<(String, Uuid, DistanceType, FixedSizeListArray)> = None; - let mut segments = Vec::new(); - let mut covered = Some(RoaringBitmap::new()); for index in source.load_indices().await?.iter() { if index.fields != [field.id] { continue; @@ -314,12 +254,6 @@ pub(super) async fn bind_column_with_segments( continue; }; let metric = vector.metric_type(); - covered = coverage(covered, index.fragment_bitmap.as_ref()); - let fragments = index.fragment_bitmap.clone(); - segments.push(IvfSegment { - index: vector.clone(), - fragments, - }); match &found { None => found = Some((index.name.clone(), index.uuid, metric, centroids)), Some((name, _, seen_metric, seen)) if *name == index.name => { @@ -369,32 +303,20 @@ pub(super) async fn bind_column_with_segments( // Lance's index path assigns cosine by L2 over unit vectors. let normalize = metric == DistanceType::Cosine; let distance = if normalize { DistanceType::L2 } else { metric }; - let partitions = u32::try_from(centroids.len()).map_err(|_| Error::InvalidInput { - message: format!("{IVF_PARTITION}({column}): too many partitions"), - })?; - Ok(( - IvfBinding { - index: uuid, - transformer: Arc::new(new_ivf_transformer(centroids, distance, vec![])), - normalize, - partitions, - }, - IvfSegments { segments, covered }, - )) + Ok(IvfBinding { + index: uuid, + transformer: Arc::new(new_ivf_transformer(centroids, distance, vec![])), + normalize, + }) } pub(super) fn empty_source(source_schema: &SchemaRef) -> Arc { - let mut fields = source_schema.fields().to_vec(); - fields.push(Arc::new(ArrowField::new(ROW_ID, DataType::UInt64, false))); - provider_as_source(Arc::new(EmptyTable::new(Arc::new(ArrowSchema::new( - fields, - ))))) + provider_as_source(Arc::new(EmptyTable::new(source_schema.clone()))) } -/// The query DataFusion runs: the view's projections plus the group's -/// smallest source row id, which stands as the row's provenance. `filtered` -/// puts the view's predicate in the query, for a caller whose rows did not -/// come from a lance scan that already applied it. +/// The query DataFusion runs. `filtered` puts the view's predicate in the +/// query for a caller whose rows did not come from a Lance scan that already +/// applied it. fn sql(definition: &MaterializedViewDefinition, filtered: bool) -> String { let items: Vec = definition .projections @@ -406,7 +328,7 @@ fn sql(definition: &MaterializedViewDefinition, filtered: bool) -> String { _ => String::new(), }; format!( - "SELECT {}, min({ROW_ID}) AS {ROW_ID} FROM {SOURCE}{predicate} GROUP BY {}", + "SELECT {} FROM {SOURCE}{predicate} GROUP BY {}", items.join(", "), definition.group_by.join(", ") ) @@ -499,11 +421,6 @@ pub(super) fn plan( }; let mut projections: Vec = Vec::with_capacity(definition.projections.len()); for p in &definition.projections { - if p.output == SOURCE_ROW_ID_COLUMN || p.output == ROW_ID { - return Err(Error::InvalidInput { - message: format!("view column name '{}' is reserved", p.output), - }); - } if projections.iter().any(|seen| seen.output == p.output) { return Err(Error::ColumnAlreadyExists { name: p.output.clone(), @@ -549,14 +466,11 @@ pub(super) fn plan( message: format!("invalid grouped view: {e}"), })?; for expr in &exprs { - ensure_immutable(expr, |message| Error::InvalidInput { - message: format!("invalid grouped view: {message}"), - })?; inputs.extend( expr.column_refs() .into_iter() .map(|c| c.name.clone()) - .filter(|name| name != ROW_ID && source_schema.field_with_name(name).is_ok()), + .filter(|name| source_schema.field_with_name(name).is_ok()), ); } inputs.sort(); @@ -566,7 +480,6 @@ pub(super) fn plan( .schema() .fields() .iter() - .filter(|f| f.name() != ROW_ID) .map(|f| without_declarations(f)) .collect(); Ok(Planned { @@ -586,7 +499,6 @@ pub(super) async fn stream( rows_written: Arc, ) -> Result { let mut scanner = source.scan(); - scanner.with_row_id(); if let Some(filter) = &definition.filter { scanner.filter(filter)?; } @@ -611,24 +523,3 @@ pub(super) async fn stream( }); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, mapped))) } - -#[cfg(test)] -mod tests { - use super::*; - - /// A segment that does not name the fragments it holds leaves the whole - /// index's coverage unknown: reading "everything it does not cover" would - /// read its rows a second time. - #[test] - fn one_segment_without_a_bitmap_loses_the_coverage() { - let first = RoaringBitmap::from_iter([0u32, 1]); - let second = RoaringBitmap::from_iter([2u32]); - let both = coverage( - coverage(Some(RoaringBitmap::new()), Some(&first)), - Some(&second), - ); - assert_eq!(both, Some(RoaringBitmap::from_iter([0u32, 1, 2]))); - assert_eq!(coverage(both, None), None); - assert_eq!(coverage(None, Some(&first)), None); - } -} diff --git a/rust/lancedb/src/materialized_view/grouped_units.rs b/rust/lancedb/src/materialized_view/grouped_units.rs deleted file mode 100644 index 7f408be4d..000000000 --- a/rust/lancedb/src/materialized_view/grouped_units.rs +++ /dev/null @@ -1,1368 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The LanceDB Authors - -//! A grouped view keyed by one `ivf_partition(col)` refreshed in units, one -//! per index partition: the index already holds each partition's row ids, so -//! a unit reads only its rows, computes only its groups and writes its own -//! fragments, and a single commit replaces the view. Nothing is shuffled and -//! no worker holds more than one partition's groups. - -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; - -use arrow_array::Array; -use arrow_array::RecordBatch; -use arrow_array::cast::AsArray; -use arrow_array::types::UInt64Type; -use arrow_schema::{Schema as ArrowSchema, SchemaRef}; -use datafusion::catalog::default_table_source::provider_as_source; -use datafusion::physical_plan::SendableRecordBatchStream; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::prelude::SessionContext; -use futures::{StreamExt, TryStreamExt}; -use lance::Dataset; -use lance::dataset::transaction::{Operation, Transaction}; -use lance::dataset::{ - CommitBuilder, InsertBuilder, ProjectionRequest, WriteDestination, WriteMode, WriteParams, -}; -use lance_core::ROW_ID; -use lance_core::utils::address::RowAddress; -use lance_datafusion::exec::SessionContextExt; -use lance_index::metrics::NoOpMetricsCollector; -use lance_table::format::Fragment; -use serde::{Deserialize, Serialize}; - -use super::grouped::{self, IVF_PARTITION}; -use super::refresh::{ - RefreshMaterializedViewResult, RefreshMode, ensure_incarnation, ensure_no_mem_wal, open_source, - refresh_lock, stamp_watermark, to_view_batch, -}; -use super::{MaterializedViewDefinition, StoredDefinition}; -use crate::table::Table; -use crate::{Error, Result}; - -/// Row ids per `take`: bounds one unit's memory while it reads a partition. -const TAKE_CHUNK: usize = 8192; - -/// Rows per partition, and in total, the plan reads to see whether the index -/// assigns by its own centroids. A sample: [`write_grouped_unit`] proves the -/// rows it takes, and this only keeps a doomed refresh off the fleet. -const PROBE_PER_UNIT: usize = 8; -const PROBE_ROWS: usize = 256; - -/// How a grouped view can be refreshed in units. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct GroupedRefreshPlan { - /// Units `0..units`: unit `p` before the last holds index partition - /// `p`'s groups; the last holds the rows the index cannot place (null or - /// unplaceable vectors), grouped under a NULL key. One per index - /// partition, plus that one. - pub units: u32, - /// The source version every unit reads and the commit records. - pub source_version: u64, - /// The view version the plan was made against; the commit lands on it - /// or not at all. - pub view_version: u64, - /// The view this plan belongs to: its incarnation, which a drop and - /// recreate replaces. Two views of the same shape reach the same - /// counters, so without it a result written for one could be published - /// into the other. - pub incarnation: String, -} - -/// The column a grouped view's single `ivf_partition(column)` key names, if -/// that is its whole `GROUP BY`. -fn ivf_key_column(definition: &MaterializedViewDefinition) -> Option { - let [key] = definition.group_by.as_slice() else { - return None; - }; - let inner = key - .strip_prefix(&format!("{IVF_PARTITION}("))? - .strip_suffix(')')?; - let column = inner.trim().trim_matches('`'); - (!column.is_empty() && !column.contains(',')).then(|| column.to_string()) -} - -/// What a grouped refresh of `view` reads: its definition, source and binding. -struct Grouped { - view_ds: Dataset, - definition: MaterializedViewDefinition, - source: Dataset, - column: String, - inputs: Vec, - binding: grouped::IvfBinding, - segments: grouped::IvfSegments, -} - -async fn open_grouped(view: &Table, pinned: Option) -> Result> { - let view_native = view.as_native().ok_or_else(|| Error::NotSupported { - message: "materialized views are supported only on local tables".into(), - })?; - view_native.dataset.reload().await?; - let view_ds = view_native.dataset.get().await?.as_ref().clone(); - let Some(StoredDefinition::Query(definition)) = - super::read_definition(&view_ds.schema().metadata)? - else { - return Ok(None); - }; - let Some(column) = ivf_key_column(&definition) else { - return Ok(None); - }; - let staging = super::read_staging(&view_ds.schema().metadata)?; - let source = open_source(view, &definition, staging.as_ref()).await?; - let source = match pinned { - Some(version) => source.checkout_version(version).await?, - None => source, - }; - ensure_no_mem_wal(&source, "source table", &definition.source_table).await?; - let source_schema = Arc::new(ArrowSchema::from(source.schema())); - let planned = super::plan(source_schema, &definition, staging.as_ref())?; - let (binding, segments) = grouped::bind_column_with_segments(&source, &column).await?; - Ok(Some(Grouped { - view_ds, - definition: planned.definition, - source, - column, - inputs: planned.inputs, - binding, - segments, - })) -} - -/// The units `view` can be refreshed in against `source_version` (the -/// latest when `None`), or `None` when it is not a grouped view keyed by one -/// `ivf_partition(column)`. -pub async fn plan_grouped_refresh( - view: &Table, - source_version: Option, -) -> Result> { - let Some(grouped) = open_grouped(view, source_version).await? else { - return Ok(None); - }; - // The split reads indexed rows through the index and the rest by - // scanning what it does not cover, so an index that does not say what it - // covers cannot be split at all. The caller refreshes single-pass. - if grouped.segments.covered.is_none() { - return Ok(None); - } - // Without one there is nothing to bind a unit's result to, so the view - // is refreshed in one pass, where the rows never leave the process. - let Some(incarnation) = grouped - .view_ds - .schema() - .metadata - .get(super::INCARNATION_META_KEY) - .cloned() - else { - return Ok(None); - }; - // The split is the index's posting lists and the key is recomputed, so a - // view whose index assigns by something else has no units at all. - if !assigns_by_its_centroids(&grouped).await? { - return Ok(None); - } - Ok(Some(GroupedRefreshPlan { - units: grouped.binding.partitions + 1, - source_version: grouped.source.version().version, - view_version: grouped.view_ds.version().version, - incarnation, - })) -} - -/// Fragments one unit wrote, uncommitted, for [`commit_grouped_refresh`]. -/// Only [`write_grouped_unit`] makes one: it names the unit and the plan, -/// and the commit publishes a set of them as the view's whole contents, so -/// a result belonging to another plan, another view, or to a unit already -/// accounted for has to be refusable. Travels to whoever commits, so it -/// serializes, but its fields are not a caller's to write. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WrittenUnit { - unit: u32, - plan: GroupedRefreshPlan, - fragments: Vec, - rows: u64, -} - -impl WrittenUnit { - /// The unit of the plan this result holds. - pub fn unit(&self) -> u32 { - self.unit - } - - /// The plan it was computed for. - pub fn plan(&self) -> &GroupedRefreshPlan { - &self.plan - } - - /// View rows this unit wrote, its groups. - pub fn rows(&self) -> u64 { - self.rows - } - - /// The fragments it staged, uncommitted until the whole set is published. - pub fn fragments(&self) -> &[Fragment] { - &self.fragments - } -} - -/// The grouped view `plan` was made for, opened at the source version it -/// pins. Every holder of a plan asks the same two questions: is this a -/// grouped view, and is it the view the plan names. -/// -/// The pinned version is also what holds every unit to one index: indices -/// are part of the source manifest, so a rebuild lands in a later version -/// that a unit reading this one never sees, and units written on either -/// side of it still group by the same centroids. -async fn open_planned(view: &Table, plan: &GroupedRefreshPlan) -> Result { - let grouped = open_grouped(view, Some(plan.source_version)) - .await? - .ok_or_else(|| Error::InvalidInput { - message: format!( - "'{}' is not a grouped view keyed by {IVF_PARTITION}", - view.name() - ), - })?; - ensure_incarnation(&grouped.view_ds, Some(&plan.incarnation), view.name()).await?; - Ok(grouped) -} - -/// Compute unit `unit` of `plan` against source version `source_version` -/// and write its groups as uncommitted fragments of the view. -pub async fn write_grouped_unit( - view: &Table, - unit: u32, - plan: &GroupedRefreshPlan, -) -> Result { - let grouped = open_planned(view, plan).await?; - if unit >= plan.units { - return Err(Error::InvalidInput { - message: format!( - "unit {unit} is not in the plan of '{}' ({} units)", - view.name(), - plan.units - ), - }); - } - let schema: SchemaRef = Arc::new(ArrowSchema::from(grouped.view_ds.schema())); - let rows_written = Arc::new(AtomicU64::new(0)); - let rows = unit_rows(&grouped, unit).await?; - let stream = group_rows(&grouped, rows, schema.clone(), rows_written.clone()).await?; - let write_txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(grouped.view_ds))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { fragments } = write_txn.operation else { - return Err(Error::Runtime { - message: "expected an append when writing a grouped unit".into(), - }); - }; - Ok(WrittenUnit { - unit, - plan: plan.clone(), - fragments, - rows: rows_written.load(Ordering::Relaxed), - }) -} - -/// The row ids the index places in partition `unit`, at most `limit` of them. -/// A segment's postings outlive its ownership -- a column rewrite attaches a -/// new file and takes the fragment out of its bitmap while leaving the rows -/// listed -- so each segment's rows are kept only while it still holds their -/// fragment. The fragment [`unit_rows`] scans them from instead is the one -/// the rewrite left them in. -async fn indexed_row_ids(grouped: &Grouped, unit: u32, limit: usize) -> Result> { - let addresses = lance::dataset::rowids::get_row_id_index(&grouped.source).await?; - let fragment_of = |row_id: u64| -> Result> { - Ok(match &addresses { - Some(index) => index.get(row_id)?.map(|address| address.fragment_id()), - None => Some(RowAddress::new_from_u64(row_id).fragment_id()), - }) - }; - let mut row_ids = Vec::new(); - for segment in &grouped.segments.segments { - let owned = segment.fragments.as_ref().ok_or_else(|| Error::Runtime { - message: "the index stopped naming the fragments it covers mid-refresh".into(), - })?; - let mut stream = segment - .index - .partition_reader(unit as usize, false, &NoOpMetricsCollector) - .await?; - while let Some(batch) = stream.try_next().await? { - for row_id in batch[ROW_ID].as_primitive::().values() { - if row_ids.len() >= limit { - return Ok(row_ids); - } - if fragment_of(*row_id)?.is_some_and(|fragment| owned.contains(fragment)) { - row_ids.push(*row_id); - } - } - } - } - Ok(row_ids) -} - -/// The first row of `batch` that `unit` holds but `ivf_partition` does not -/// put there, if there is one. `batch` carries the key column. -fn assigned_elsewhere(grouped: &Grouped, batch: &RecordBatch, unit: u32) -> Result> { - let vectors = batch[grouped.column.as_str()].as_fixed_size_list(); - let assigned = grouped.binding.assign(vectors)?; - Ok((0..batch.num_rows()).find(|&row| !(assigned.is_valid(row) && assigned.value(row) == unit))) -} - -/// Whether the index's posting lists are `ivf_partition`'s buckets, over a -/// sample of each partition. Lance also builds an index from precomputed -/// assignments, and records nowhere that it did, so the only way to ask is -/// to recompute some of them. -async fn assigns_by_its_centroids(grouped: &Grouped) -> Result { - let mut expected = HashMap::new(); - for unit in 0..grouped.binding.partitions { - if expected.len() >= PROBE_ROWS { - break; - } - for row_id in indexed_row_ids(grouped, unit, PROBE_PER_UNIT).await? { - expected.insert(row_id, unit); - } - } - if expected.is_empty() { - return Ok(true); - } - let sample: Vec = expected.keys().copied().collect(); - let columns = [ROW_ID, grouped.column.as_str()]; - let projection = ProjectionRequest::from_columns(columns, grouped.source.schema()); - let batch = grouped.source.take_rows(&sample, projection).await?; - let row_ids = batch[ROW_ID].as_primitive::().clone(); - let assigned = grouped - .binding - .assign(batch[grouped.column.as_str()].as_fixed_size_list())?; - Ok((0..batch.num_rows()).all(|row| { - assigned.is_valid(row) && Some(&assigned.value(row)) == expected.get(&row_ids.value(row)) - })) -} - -/// The source rows of one unit: the index's rows for partition `unit` plus -/// the unindexed rows that assign to it, or for the last unit every row the -/// assigner cannot place. Each batch carries the inputs and `_rowid`. -async fn unit_rows(grouped: &Grouped, unit: u32) -> Result> { - let mut columns: Vec<&str> = vec![ROW_ID]; - columns.extend(grouped.inputs.iter().map(String::as_str)); - if !grouped.inputs.contains(&grouped.column) { - columns.push(&grouped.column); - } - let source = Arc::new(grouped.source.clone()); - let mut batches = Vec::new(); - let unassigned = unit == grouped.binding.partitions; - - if !unassigned { - // Rows the index placed in this partition; a deleted row is dropped - // by the take. - let row_ids = indexed_row_ids(grouped, unit, usize::MAX).await?; - let projection = ProjectionRequest::from_columns(&columns, source.schema()); - for chunk in row_ids.chunks(TAKE_CHUNK) { - let batch = source.take_rows(chunk, projection.clone()).await?; - if batch.num_rows() > 0 { - // What the plan sampled, this unit proves: a row grouped - // here that belongs to another key would be aggregated in - // both units and published as two rows of one group. - if let Some(row) = assigned_elsewhere(grouped, &batch, unit)? { - let row_id = batch[ROW_ID].as_primitive::().value(row); - return Err(Error::NotSupported { - message: format!( - "{IVF_PARTITION}({}): the index lists row {row_id} under \ - partition {unit}, which it does not assign to; this view \ - is refreshed in one pass", - grouped.column - ), - }); - } - batches.push(in_column_order(&batch, &columns)?); - } - } - } - - // Rows the index has not seen assign here; the last unit also takes the - // rows no assignment places, from every fragment. - let mut fragments = Vec::new(); - for fragment in source.get_fragments() { - let unindexed = grouped - .segments - .is_unindexed(fragment.id() as u64) - .ok_or_else(|| Error::Runtime { - message: "the index stopped naming the fragments it covers mid-refresh".into(), - })?; - if unassigned || unindexed { - fragments.push(fragment.metadata().clone()); - } - } - if !fragments.is_empty() { - let mut scanner = source.scan(); - scanner.with_fragments(fragments); - scanner.with_row_id(); - scanner.project(&columns)?; - let mut stream = scanner.try_into_stream().await?; - while let Some(batch) = stream.try_next().await? { - let vectors = batch[grouped.column.as_str()].as_fixed_size_list(); - let assigned = grouped.binding.assign(vectors)?; - let keep: arrow_array::BooleanArray = (0..batch.num_rows()) - .map(|row| { - Some(if unassigned { - assigned.is_null(row) - } else { - assigned.is_valid(row) && assigned.value(row) == unit - }) - }) - .collect(); - let kept = arrow_select::filter::filter_record_batch(&batch, &keep)?; - if kept.num_rows() > 0 { - batches.push(in_column_order(&kept, &columns)?); - } - } - } - Ok(batches) -} - -/// `batch` with its columns in `columns` order: the take and the scan place -/// `_rowid` differently, and the query stream coalesces by position. -fn in_column_order(batch: &RecordBatch, columns: &[&str]) -> Result { - let indices = columns - .iter() - .map(|column| batch.schema().index_of(column)) - .collect::, _>>()?; - Ok(batch.project(&indices)?) -} - -/// The grouped query over `rows`, in the view's schema. -async fn group_rows( - grouped: &Grouped, - rows: Vec, - schema: SchemaRef, - rows_written: Arc, -) -> Result { - let input_schema = match rows.first() { - Some(batch) => batch.schema(), - None => { - return Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::empty(), - ))); - } - }; - let rows: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - input_schema, - futures::stream::iter(rows.into_iter().map(Ok)), - )); - let state = grouped::session(grouped::bind(&grouped.source, &grouped.definition).await?); - let ctx = SessionContext::new_with_state(state.clone()); - let source = ctx.read_one_shot(rows)?.into_view(); - let planned = grouped::logical_plan( - &state, - provider_as_source(source), - &grouped.definition, - // Neither the indexed take nor the fragment scan filters: the rows - // reach this query raw, so the predicate belongs in it. - true, - )?; - let groups = ctx - .execute_logical_plan(planned) - .await? - .execute_stream() - .await?; - let out_schema = schema.clone(); - let mapped = groups.map(move |batch| { - let batch = to_view_batch(&batch?, &out_schema)?; - rows_written.fetch_add(batch.num_rows() as u64, Ordering::Relaxed); - Ok(batch) - }); - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, mapped))) -} - -/// Replace the view's rows with the fragments every unit wrote, in one -/// commit, and record `plan.source_version` as what the view reflects. -/// -/// The commit is refused if the view moved since the plan: a refresh that -/// landed meanwhile replaced rows the units never saw. Concurrent refreshes -/// of one view are serialized by the caller, as for a single-pass rebuild; -/// this check catches the ones that slipped through. -pub async fn commit_grouped_refresh( - view: &Table, - plan: &GroupedRefreshPlan, - units: Vec, - expected_incarnation: Option<&str>, -) -> Result { - let grouped = open_planned(view, plan).await?; - let units = complete_set(view.name(), plan, units)?; - let rows_written = units.iter().map(|unit| unit.rows).sum(); - let new_fragments = units.into_iter().flat_map(|unit| unit.fragments).collect(); - commit_fragments( - view, - plan, - &grouped.definition, - grouped.source.manifest.timestamp_nanos, - new_fragments, - rows_written, - expected_incarnation, - ) - .await -} - -/// The common all-or-nothing unit sink, shared by grouped and native views. -#[allow(clippy::too_many_arguments)] -pub(super) async fn commit_fragments( - view: &Table, - plan: &GroupedRefreshPlan, - definition: &MaterializedViewDefinition, - source_ts: u128, - new_fragments: Vec, - rows_written: u64, - expected_incarnation: Option<&str>, -) -> Result { - let view_native = view.as_native().ok_or_else(|| Error::NotSupported { - message: "materialized views are supported only on local tables".into(), - })?; - view_native.dataset.ensure_mutable()?; - let lock = refresh_lock(view_native.dataset.get().await?.uri()); - let _guard = lock.lock().await; - if let Some(expected) = expected_incarnation - && expected != plan.incarnation - { - return Err(Error::Runtime { - message: format!( - "the refresh of '{}' was requested for another incarnation than the one \ - its units were planned against", - view.name() - ), - }); - } - view_native.dataset.reload().await?; - let ds = view_native.dataset.get().await?.as_ref().clone(); - ensure_incarnation(&ds, Some(&plan.incarnation), view.name()).await?; - if super::read_definition(&ds.schema().metadata)? - != Some(super::StoredDefinition::Query(definition.clone())) - { - return Err(Error::InvalidInput { - message: "view definition changed while its units were written".into(), - }); - } - let ds = Arc::new(ds); - if ds.version().version != plan.view_version { - return Err(Error::Runtime { - message: format!( - "'{}' moved from version {} to {} while its units were written; the \ - refresh is unrecorded and must be replanned", - view.name(), - plan.view_version, - ds.version().version - ), - }); - } - let removed_fragment_ids: Vec = ds.get_fragments().iter().map(|f| f.id() as u64).collect(); - // Fragment ids are assigned at commit; the units' fragments all carry 0. - let transaction = Transaction::new( - plan.view_version, - Operation::Update { - removed_fragment_ids, - updated_fragments: Vec::new(), - new_fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: None, - updated_fragment_offsets: None, - }, - None, - ); - #[cfg(test)] - super::refresh::tests::hold_before_publish(ds.uri()).await; - let committed = CommitBuilder::new(WriteDestination::Dataset(ds)) - .execute(transaction) - .await?; - // The version check above is not the fence: lance rebases this `Update` - // over a concurrent append rather than refusing it, which would leave - // the appended rows in a view certified as the units' whole output. - if committed.version().version != plan.view_version + 1 { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced the refresh of '{}' (view version {}); the \ - refresh is unrecorded and must be replanned", - view.name(), - committed.version().version - ), - }); - } - let version = stamp_watermark( - view_native, - committed, - plan.source_version, - source_ts, - Some(definition), - expected_incarnation, - ) - .await?; - Ok(RefreshMaterializedViewResult { - mode: RefreshMode::Rebuild, - rows_written, - source_version: plan.source_version, - version, - }) -} - -/// The units in plan order, exactly one per unit of `plan`. The commit -/// publishes them as the view's whole contents, so a missing, repeated or -/// foreign result would certify a view nobody computed. -fn complete_set( - view: &str, - plan: &GroupedRefreshPlan, - units: Vec, -) -> Result> { - let mut ordered: Vec> = vec![None; plan.units as usize]; - for written in units { - if &written.plan != plan { - return Err(Error::InvalidInput { - message: format!( - "a unit of another plan of '{view}' was offered to this refresh; the \ - refresh is unrecorded and must be replanned" - ), - }); - } - let slot = ordered - .get_mut(written.unit as usize) - .ok_or_else(|| Error::InvalidInput { - message: format!( - "unit {} is not in the plan of '{view}' ({} units)", - written.unit, plan.units - ), - })?; - if slot.is_some() { - return Err(Error::InvalidInput { - message: format!("unit {} of '{view}' was written twice", written.unit), - }); - } - // A unit holds one key, so it writes that key's one group. - if written.rows > 1 { - return Err(Error::InvalidInput { - message: format!( - "unit {} of '{view}' wrote {} groups; the split did not match the \ - grouping and the refresh must be replanned", - written.unit, written.rows - ), - }); - } - *slot = Some(written); - } - ordered - .into_iter() - .enumerate() - .map(|(unit, written)| { - written.ok_or_else(|| Error::InvalidInput { - message: format!( - "unit {unit} of '{view}' is missing; the refresh holds {} units and \ - publishes all of them or none", - plan.units - ), - }) - }) - .collect() -} - -#[cfg(test)] -mod tests { - use super::super::refresh::tests::{clustered_source, declare, rows}; - use super::*; - use crate::connect; - use crate::index::Index; - use crate::index::vector::IvfFlatIndexBuilder; - use arrow_array::types::Float32Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array, RecordBatchIterator}; - use lance::index::DatasetIndexExt; - use lance::index::vector::VectorIndexParams; - use lance_index::IndexType; - use lance_index::vector::ivf::IvfBuildParams; - - const BUCKETS: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, min(id) AS lo, \ - max(id) AS hi FROM images GROUP BY ivf_partition(vec)"; - - /// The clustered source under a two-partition cosine index, then edited - /// behind the index: a row deleted from each cluster, a row appended to - /// each, and one zero vector cosine cannot place. - async fn edited_source(conn: &crate::Connection) -> Table { - let source = clustered_source(conn).await; - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .execute() - .await - .unwrap(); - source.delete("id IN (3, 13)").await.unwrap(); - let appended = FixedSizeListArray::from_iter_primitive::( - [ - Some(vec![Some(1.5), Some(0.0)]), - Some(vec![Some(101.5), Some(100.0)]), - Some(vec![Some(0.0), Some(0.0)]), - ], - 2, - ); - source - .add( - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from(vec![20, 21, 22])) as ArrayRef, - ), - ("vec", Arc::new(appended) as ArrayRef), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - source - } - - async fn refresh_in_units( - view: &Table, - expected_incarnation: Option<&str>, - ) -> Result<( - GroupedRefreshPlan, - Vec, - RefreshMaterializedViewResult, - )> { - let plan = plan_grouped_refresh(view, None) - .await? - .expect("an ivf-keyed view"); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view, unit, &plan).await?); - } - let result = - commit_grouped_refresh(view, &plan, units.clone(), expected_incarnation).await?; - Ok((plan, units, result)) - } - - /// Every partition's unit reads its indexed rows plus the appended rows - /// that assign to it, the deleted rows are gone, and the last unit holds - /// the unplaceable row: together they are the single-pass result. - #[tokio::test] - async fn units_compose_to_the_single_pass_result() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let single = declare(&source, "single", BUCKETS).await.unwrap(); - single.refresh().execute().await.unwrap(); - let expected = rows(single.table(), &["bucket", "n", "lo", "hi"]).await; - assert_eq!(expected.len(), 3, "{expected:?}"); - assert!(expected.iter().any(|row| row == " 1 22 22"), "{expected:?}"); - - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let (plan, units, result) = refresh_in_units(view.table(), None).await.unwrap(); - assert_eq!(plan.units, 3, "two partitions plus the unplaceable rows"); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected - ); - assert_eq!(units.iter().map(|u| u.rows).sum::(), 3); - assert_eq!(units.iter().map(|u| u.rows).collect::>(), [1, 1, 1]); - assert_eq!(result.rows_written, 3); - assert_eq!(result.source_version, plan.source_version); - assert_eq!(view.table().version().await.unwrap(), result.version); - - // The commit is a refresh: the view reflects the source version it - // read and a second pass finds nothing to do. - let again = view.refresh().execute().await.unwrap(); - assert_eq!(again.rows_written, 0, "{again:?}"); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected - ); - } - - /// A second refresh in units replaces the first one's rows rather than - /// appending to them. - #[tokio::test] - async fn a_refresh_in_units_replaces_the_previous_rows() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - source.delete("id = 22").await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - let mut groups = rows(view.table(), &["n", "lo", "hi"]).await; - groups.sort(); - assert_eq!(groups, ["10 0 20", "10 10 21"]); - } - - #[tokio::test] - async fn a_view_not_keyed_by_ivf_partition_has_no_units() { - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - let view = declare( - &source, - "by_id", - "SELECT id % 2 AS k, count(*) AS n FROM images GROUP BY id % 2", - ) - .await - .unwrap(); - assert_eq!( - plan_grouped_refresh(view.table(), None).await.unwrap(), - None - ); - } - - #[tokio::test] - async fn a_unit_outside_the_plan_is_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let err = write_grouped_unit(view.table(), plan.units, &plan) - .await - .unwrap_err(); - assert!(err.to_string().contains("not in the plan"), "{err}"); - } - - /// A commit that lands on a view someone else refreshed meanwhile is - /// refused rather than recorded over their rows. - #[tokio::test] - async fn a_raced_commit_is_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let units = vec![write_grouped_unit(view.table(), 0, &plan).await.unwrap()]; - view.refresh().execute().await.unwrap(); - let before = rows(view.table(), &["bucket", "n", "lo", "hi"]).await; - let raced = commit_grouped_refresh(view.table(), &plan, units, None).await; - assert!(raced.is_err(), "{raced:?}"); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - before - ); - } - - /// A view with a predicate holds the rows it selects however it is - /// refreshed: the units read their rows through the index and by - /// scanning fragments, neither of which filters, so the predicate has - /// to be in the grouped query itself. - #[tokio::test] - async fn a_filtered_view_refreshes_in_units_like_a_single_pass() { - const FILTERED: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images WHERE id < 5 \ - GROUP BY ivf_partition(vec)"; - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let single = declare(&source, "single", FILTERED).await.unwrap(); - single.refresh().execute().await.unwrap(); - let expected = rows(single.table(), &["bucket", "n", "lo", "hi"]).await; - - let view = declare(&source, "buckets", FILTERED).await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected - ); - } - - /// The commit publishes its units as the view's whole contents, so it - /// takes the plan's units exactly once each and nothing else. - #[tokio::test] - async fn only_a_complete_set_of_this_plan_s_units_is_published() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - - let missing = commit_grouped_refresh(view.table(), &plan, units[..2].to_vec(), None).await; - assert!( - missing.unwrap_err().to_string().contains("is missing"), - "a partial set was published" - ); - let mut twice = units.clone(); - twice[1] = units[0].clone(); - let repeated = commit_grouped_refresh(view.table(), &plan, twice, None).await; - assert!( - repeated.unwrap_err().to_string().contains("written twice"), - "a repeated unit was published" - ); - let mut foreign = units.clone(); - foreign[0].plan.source_version += 1; - let stale = commit_grouped_refresh(view.table(), &plan, foreign, None).await; - assert!( - stale.unwrap_err().to_string().contains("another plan"), - "a unit of another plan was published" - ); - assert_eq!(rows(view.table(), &["bucket"]).await, Vec::::new()); - - // The complete set still lands. - commit_grouped_refresh(view.table(), &plan, units, None) - .await - .unwrap(); - assert_eq!(rows(view.table(), &["bucket"]).await.len(), 3); - } - - /// A column rewrite attaches a new file to the row's fragment and takes - /// that fragment out of the index segment's bitmap, but the segment's - /// postings still name its rows. Taking them anyway, while the fragment - /// is also scanned as unindexed, counts those rows twice and buckets the - /// stale copy by the value the rewrite replaced. - #[tokio::test] - async fn a_rewritten_fragment_is_read_once() { - use lance::dataset::{ - MergeInsertBuilder, MergeInsertWriteMode, WhenMatched, WhenNotMatched, - }; - - let conn = connect("memory://").execute().await.unwrap(); - // A column the rewrite leaves alone: RewriteColumns is only for a - // source that covers some of the row, not all of it. - let vectors = FixedSizeListArray::from_iter_primitive::( - (0..20).map(|i| { - let base = if i < 10 { 0.0 } else { 100.0 }; - Some(vec![Some(base + 1.0 + i as f32 * 0.1), Some(base)]) - }), - 2, - ); - let source = conn - .create_table( - "images", - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(0..20)) as ArrayRef, - ), - ("vec", Arc::new(vectors) as ArrayRef), - ( - "tag", - Arc::new(Int32Array::from_iter_values((0..20).map(|i| i * 10))) as ArrayRef, - ), - ]) - .unwrap(), - ) - .write_options(super::super::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - source - .create_index( - &["vec"], - Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(2)), - ) - .execute() - .await - .unwrap(); - - // Row 0 keeps its id and its fragment; its vector is replaced in a - // file the index segment does not cover. - let batch = RecordBatch::try_from_iter(vec![ - ("id", Arc::new(Int32Array::from(vec![0])) as ArrayRef), - ( - "vec", - Arc::new( - FixedSizeListArray::from_iter_primitive::( - [Some(vec![Some(101.0), Some(100.0)])], - 2, - ), - ) as ArrayRef, - ), - ]) - .unwrap(); - let schema = batch.schema(); - let dataset = source.as_native().unwrap().dataset.get().await.unwrap(); - let mut builder = - MergeInsertBuilder::try_new(Arc::new(dataset.as_ref().clone()), vec!["id".to_string()]) - .unwrap(); - builder - .when_matched(WhenMatched::UpdateAll) - .when_not_matched(WhenNotMatched::DoNothing) - .write_mode(MergeInsertWriteMode::RewriteColumns); - builder - .try_build() - .unwrap() - .execute_reader(arrow_array::RecordBatchIterator::new([Ok(batch)], schema)) - .await - .unwrap(); - - let single = declare(&source, "single", BUCKETS).await.unwrap(); - single.refresh().execute().await.unwrap(); - let expected = rows(single.table(), &["bucket", "n", "lo", "hi"]).await; - - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected, - "the rewritten row was counted twice or bucketed by its old vector" - ); - } - - /// Units written on either side of an index rebuild still group by one - /// model: the plan pins a source version, indices live in the source - /// manifest, and a rebuild lands in a later version the units never - /// read. Without that, two units of one plan could assign their rows by - /// different centroids. - #[tokio::test] - async fn a_rebuilt_index_does_not_reach_the_units_of_a_plan() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let first = write_grouped_unit(view.table(), 0, &plan).await.unwrap(); - - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .replace(true) - .execute() - .await - .unwrap(); - let rebuilt = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - assert_eq!(rebuilt.units, plan.units, "same shape"); - assert!( - rebuilt.source_version > plan.source_version, - "the rebuild is a later source version" - ); - - // The remaining units read the version the plan pinned, so they are - // grouped by the index that version holds, not the new one. - let mut units = vec![first]; - for unit in 1..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - commit_grouped_refresh(view.table(), &plan, units, None) - .await - .unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await.len(), - 3 - ); - } - - /// Two views of the same shape over the same source reach the same plan - /// counters, so the results of one must not satisfy the other: they are - /// fragments written into a different dataset. - #[tokio::test] - async fn units_written_for_another_view_are_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let mine = declare(&source, "buckets", BUCKETS).await.unwrap(); - let theirs = declare(&source, "twin", BUCKETS).await.unwrap(); - let my_plan = plan_grouped_refresh(mine.table(), None) - .await - .unwrap() - .unwrap(); - let their_plan = plan_grouped_refresh(theirs.table(), None) - .await - .unwrap() - .unwrap(); - assert_eq!( - (my_plan.units, my_plan.view_version, my_plan.source_version), - ( - their_plan.units, - their_plan.view_version, - their_plan.source_version - ), - "the counters alone do not tell these views apart" - ); - - let mut mine_units = Vec::new(); - for unit in 0..my_plan.units { - mine_units.push( - write_grouped_unit(mine.table(), unit, &my_plan) - .await - .unwrap(), - ); - } - let misrouted = - commit_grouped_refresh(theirs.table(), &their_plan, mine_units.clone(), None).await; - assert!( - misrouted.unwrap_err().to_string().contains("another plan"), - "another view's units were published" - ); - assert_eq!( - rows(theirs.table(), &["bucket"]).await, - Vec::::new() - ); - - // And a unit cannot be written for a view the plan does not name. - let foreign = write_grouped_unit(theirs.table(), 0, &my_plan).await; - assert!( - foreign - .unwrap_err() - .to_string() - .contains("not the incarnation"), - "a unit was written under another view's plan" - ); - } - - /// A write that lands on the view between the plan check and the commit - /// is not a conflict lance refuses -- the update rebases over it -- so - /// the publication checks the generation it actually landed on. - #[tokio::test(flavor = "multi_thread")] - async fn a_write_racing_the_publication_is_refused() { - use super::super::refresh::tests::{ - DRIFT_LOCK, DRIFT_PLANNED, DRIFT_RELEASED, DRIFT_TARGET, - }; - let _serial = DRIFT_LOCK.lock().await; - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - // Its own name: memory:// uris repeat across tests, so a view called - // like its siblings would arm the hook for one of their commits. - let view = declare(&source, "racing_buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - - let publishing = { - let table = view.table().clone(); - let plan = plan.clone(); - tokio::spawn(async move { commit_grouped_refresh(&table, &plan, units, None).await }) - }; - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the commit never reached the publication boundary"); - view.table() - .add( - RecordBatch::try_from_iter(vec![ - ("bucket", Arc::new(Int32Array::from(vec![7])) as ArrayRef), - ("n", Arc::new(Int32Array::from(vec![1])) as ArrayRef), - ("lo", Arc::new(Int32Array::from(vec![0])) as ArrayRef), - ("hi", Arc::new(Int32Array::from(vec![0])) as ArrayRef), - ( - super::super::SOURCE_ROW_ID_COLUMN, - Arc::new(arrow_array::UInt64Array::from(vec![0u64])) as ArrayRef, - ), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - DRIFT_RELEASED.notify_one(); - - let err = publishing.await.unwrap().unwrap_err(); - assert!(err.to_string().contains("raced the refresh"), "{err}"); - } - - #[test] - fn the_key_column_is_read_off_a_single_ivf_key() { - let mut definition = MaterializedViewDefinition::from_sql( - "SELECT ivf_partition(vec) AS b, count(*) AS n FROM t GROUP BY ivf_partition(vec)", - ) - .unwrap(); - assert_eq!(ivf_key_column(&definition).as_deref(), Some("vec")); - definition.group_by = vec!["k".to_string()]; - assert_eq!(ivf_key_column(&definition), None); - definition.group_by = vec!["ivf_partition(vec)".to_string(), "k".to_string()]; - assert_eq!(ivf_key_column(&definition), None); - } - - /// A source whose index does not assign by its own centroids: lance - /// takes precomputed partitions, so row 1 is stored under partition 1 - /// and `ivf_partition` puts it in 0. - async fn precomputed_source(conn: &crate::Connection, dir: &std::path::Path) -> Table { - let vectors = FixedSizeListArray::from_iter_primitive::( - [ - Some(vec![Some(1.0), Some(0.0)]), - Some(vec![Some(2.0), Some(0.0)]), - Some(vec![Some(100.0), Some(100.0)]), - ], - 2, - ); - let source = conn - .create_table( - "images", - RecordBatch::try_from_iter(vec![ - ("id", Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef), - ("vec", Arc::new(vectors) as ArrayRef), - ]) - .unwrap(), - ) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - - let assignments = arrow_array::record_batch!( - ("row_id", UInt64, [0, 1, 2]), - ("partition", UInt32, [0, 1, 1]) - ) - .unwrap(); - let file = dir.join("partitions.lance"); - let file = file.to_str().unwrap(); - lance::Dataset::write( - RecordBatchIterator::new(vec![Ok(assignments.clone())], assignments.schema()), - file, - None, - ) - .await - .unwrap(); - - let centroids = FixedSizeListArray::from_iter_primitive::( - [ - Some(vec![Some(0.0), Some(0.0)]), - Some(vec![Some(100.0), Some(100.0)]), - ], - 2, - ); - let mut ivf = IvfBuildParams::try_with_centroids(2, Arc::new(centroids)).unwrap(); - ivf.precomputed_partitions_file = Some(file.to_string()); - let mut dataset = source - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .as_ref() - .clone(); - dataset - .create_index( - &["vec"], - IndexType::Vector, - None, - &VectorIndexParams::with_ivf_flat_params( - lance_linalg::distance::DistanceType::L2, - ivf, - ), - true, - ) - .await - .unwrap(); - conn.open_table("images").execute().await.unwrap() - } - - async fn incarnation(view: &Table) -> String { - view.as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .schema() - .metadata - .get(crate::materialized_view::INCARNATION_META_KEY) - .cloned() - .unwrap() - } - - /// The split is the index's posting lists and the key is - /// `ivf_partition`. Where they disagree one group is aggregated in two - /// units and published twice, so such a view has no units at all. - #[tokio::test] - async fn an_index_that_does_not_assign_by_its_centroids_has_no_units() { - let dir = tempfile::tempdir().unwrap(); - let conn = connect("memory://").execute().await.unwrap(); - let source = precomputed_source(&conn, dir.path()).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - assert_eq!( - plan_grouped_refresh(view.table(), None).await.unwrap(), - None - ); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - ["0 2 0 1", "1 1 2 2"] - ); - } - - /// The publication is the last place the split can be caught: a unit - /// that holds one key writes one group, and a set that says otherwise is - /// not published however it was produced. - #[tokio::test] - async fn a_unit_that_wrote_more_than_one_group_is_not_published() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - units[1].rows = 2; - let err = commit_grouped_refresh(view.table(), &plan, units, None) - .await - .unwrap_err(); - assert!(err.to_string().contains("wrote 2 groups"), "{err}"); - assert_eq!(rows(view.table(), &["bucket"]).await, Vec::::new()); - } - - /// The plan reads a sample, so the unit proves the rest: it holds only - /// rows that assign to it, or it refuses rather than write a group - /// another unit holds too. - #[tokio::test] - async fn a_unit_holding_a_row_that_assigns_elsewhere_is_refused() { - let dir = tempfile::tempdir().unwrap(); - let conn = connect("memory://").execute().await.unwrap(); - let source = precomputed_source(&conn, dir.path()).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = GroupedRefreshPlan { - units: 3, - source_version: source.version().await.unwrap(), - view_version: view.table().version().await.unwrap(), - incarnation: incarnation(view.table()).await, - }; - assert_eq!( - write_grouped_unit(view.table(), 0, &plan) - .await - .unwrap() - .rows, - 1 - ); - let err = write_grouped_unit(view.table(), 1, &plan) - .await - .unwrap_err(); - assert!(err.to_string().contains("does not assign to"), "{err}"); - } -} diff --git a/rust/lancedb/src/materialized_view/partitioned.rs b/rust/lancedb/src/materialized_view/partitioned.rs deleted file mode 100644 index 58192089b..000000000 --- a/rust/lancedb/src/materialized_view/partitioned.rs +++ /dev/null @@ -1,192 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The LanceDB Authors - -//! Shared plan/write/commit surface for materialized views with independent units. - -use serde::{Deserialize, Serialize}; - -use super::duplicate_pairs::{ - self, DuplicatePairsRefreshPlan, VectorSourceKind, WrittenPairsPartition, -}; -use super::vector_dedup::{self, VectorDedupPlan, WrittenDedupUnit}; -use super::{GroupedRefreshPlan, RefreshMaterializedViewResult, StoredDefinition, WrittenUnit}; -use crate::table::Table; -use crate::{Error, Result}; - -/// Serializable work description consumed by a host's existing job workers. -/// The grouped variant preserves the previous saved-plan representation. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(untagged)] -pub enum PartitionedRefreshPlan { - VectorDedup(VectorDedupPlan), - DuplicatePairs(DuplicatePairsRefreshPlan), - Grouped(GroupedRefreshPlan), -} - -impl PartitionedRefreshPlan { - /// The shared source snapshot and destination generation fence. - pub fn snapshot(&self) -> &GroupedRefreshPlan { - match self { - Self::VectorDedup(plan) => &plan.snapshot, - Self::DuplicatePairs(plan) => &plan.snapshot, - Self::Grouped(plan) => plan, - } - } - /// Selection tasks that must complete before remaining tasks can start. - pub fn dependency_units(&self) -> u32 { - match self { - Self::VectorDedup(plan) => plan.dependencies(), - _ => 0, - } - } - /// Number of independent units in this immutable plan. - pub fn units(&self) -> u32 { - self.snapshot().units - } -} - -/// One completed unit's uncommitted fragments. Failed attempts produce no receipt. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(untagged)] -pub enum WrittenPartition { - VectorDedup(WrittenDedupUnit), - DuplicatePairs(WrittenPairsPartition), - Grouped(WrittenUnit), -} - -impl WrittenPartition { - /// Number of pair or grouped rows staged by this unit. - pub fn rows(&self) -> u64 { - match self { - Self::VectorDedup(unit) => unit.rows, - Self::DuplicatePairs(unit) => unit.rows, - Self::Grouped(unit) => unit.rows(), - } - } - /// Uncommitted output files, retained only after task completion. - pub fn fragments(&self) -> &[lance_table::format::Fragment] { - match self { - Self::VectorDedup(unit) => &unit.fragments, - Self::DuplicatePairs(unit) => &unit.fragments, - Self::Grouped(unit) => unit.fragments(), - } - } -} - -/// Plan native pair or indexed grouped work; return `None` for ordinary views. -/// `pinned`, when supplied, must match a native view's declared source version. -pub async fn plan_partitioned_refresh( - view: &Table, - pinned: Option, -) -> Result> { - if let Some(StoredDefinition::Query(definition)) = - super::read_definition(view.schema().await?.metadata())? - && definition.vector_source.is_some() - { - if definition - .vector_source - .as_ref() - .is_some_and(|s| s.kind == VectorSourceKind::Dedup) - { - return vector_dedup::plan_refresh(view, pinned) - .await - .map(|p| Some(PartitionedRefreshPlan::VectorDedup(p))); - } - return duplicate_pairs::plan_refresh(view, pinned) - .await - .map(|p| Some(PartitionedRefreshPlan::DuplicatePairs(p))); - } - super::plan_grouped_refresh(view, pinned) - .await - .map(|plan| plan.map(PartitionedRefreshPlan::Grouped)) -} - -/// Stream one unit into uncommitted fragments. Persist the returned receipt only -/// on success. A failed or interrupted call must be retried from the unit's start. -pub async fn write_refresh_partition( - view: &Table, - unit: u32, - plan: &PartitionedRefreshPlan, -) -> Result { - write_refresh_partition_with_inputs(view, unit, plan, &[]).await -} - -/// Write a unit after its durable selection dependencies have completed. -/// Supply exactly the first `plan.dependency_units()` receipts for a dependent -/// unit. A missing, duplicate or foreign receipt fails before output is staged. -pub async fn write_refresh_partition_with_inputs( - view: &Table, - unit: u32, - plan: &PartitionedRefreshPlan, - dependencies: &[WrittenPartition], -) -> Result { - match plan { - PartitionedRefreshPlan::VectorDedup(plan) => { - let dependencies = dependencies - .iter() - .map(|unit| match unit { - WrittenPartition::VectorDedup(unit) => Ok(unit.clone()), - _ => Err(Error::InvalidInput { - message: "foreign dedup dependency".into(), - }), - }) - .collect::>>()?; - vector_dedup::write_unit(view, unit, plan, &dependencies) - .await - .map(WrittenPartition::VectorDedup) - } - PartitionedRefreshPlan::DuplicatePairs(plan) => { - duplicate_pairs::write_unit(view, unit, plan) - .await - .map(WrittenPartition::DuplicatePairs) - } - PartitionedRefreshPlan::Grouped(plan) => super::write_grouped_unit(view, unit, plan) - .await - .map(WrittenPartition::Grouped), - } -} - -/// Publish a complete set of unit receipts through the common MV commit sink. -/// `expected` fences the destination incarnation. No partial set is publishable. -pub async fn commit_partitioned_refresh( - view: &Table, - plan: &PartitionedRefreshPlan, - units: Vec, - expected: Option<&str>, -) -> Result { - let foreign = || Error::InvalidInput { - message: "refresh receipt does not match the plan kind".into(), - }; - match plan { - PartitionedRefreshPlan::VectorDedup(plan) => { - let units = units - .into_iter() - .map(|unit| match unit { - WrittenPartition::VectorDedup(unit) => Ok(unit), - _ => Err(foreign()), - }) - .collect::>>()?; - vector_dedup::commit(view, plan, units, expected).await - } - PartitionedRefreshPlan::DuplicatePairs(plan) => { - let units = units - .into_iter() - .map(|unit| match unit { - WrittenPartition::DuplicatePairs(unit) => Ok(unit), - _ => Err(foreign()), - }) - .collect::>>()?; - duplicate_pairs::commit(view, plan, units, expected).await - } - PartitionedRefreshPlan::Grouped(plan) => { - let units = units - .into_iter() - .map(|unit| match unit { - WrittenPartition::Grouped(unit) => Ok(unit), - _ => Err(foreign()), - }) - .collect::>>()?; - super::commit_grouped_refresh(view, plan, units, expected).await - } - } -} diff --git a/rust/lancedb/src/materialized_view/query.rs b/rust/lancedb/src/materialized_view/query.rs index 4bfd5542e..c1a987b2c 100644 --- a/rust/lancedb/src/materialized_view/query.rs +++ b/rust/lancedb/src/materialized_view/query.rs @@ -264,36 +264,7 @@ fn extract(query: &Query) -> Result { }; let mut from = select.from.iter(); - let mut vector_source = None; let (source_namespace, source_table) = match from.next().map(|f| &f.relation) { - Some(TableFactor::Table { - name, - args: Some(args), - alias: None, - with_hints, - version: None, - with_ordinality: false, - partitions, - json_path: None, - sample: None, - index_hints, - .. - }) if matches!( - name.to_string().as_str(), - super::duplicate_pairs::FUNCTION_NAME | super::duplicate_pairs::DEDUP_FUNCTION_NAME - ) && with_hints.is_empty() - && partitions.is_empty() - && index_hints.is_empty() => - { - let kind = if name.to_string() == super::duplicate_pairs::FUNCTION_NAME { - super::VectorSourceKind::Pairs - } else { - super::VectorSourceKind::Dedup - }; - let (namespace, table, config) = super::duplicate_pairs::source(args, kind)?; - vector_source = Some(Box::new(config)); - (namespace, table) - } Some(TableFactor::Table { args: Some(_), .. }) => { return Err(invalid( "a view reads a table; a Function in FROM position follows it: \ @@ -445,7 +416,6 @@ fn extract(query: &Query) -> Result { }; let definition = MaterializedViewDefinition { - vector_source, source_table, source_namespace, lateral, @@ -454,9 +424,6 @@ fn extract(query: &Query) -> Result { group_by, limit, }; - if definition.vector_source.is_some() { - super::duplicate_pairs::check_shape(&definition)?; - } check_grouping(&definition)?; Ok(definition) } @@ -516,18 +483,7 @@ pub fn render(definition: &MaterializedViewDefinition) -> String { .chain(std::iter::once(&definition.source_table)) .map(|part| ident_sql(part)) .collect(); - if let Some(native) = &definition.vector_source { - sql.push_str(&format!( - "{}('{}', {}, '{}', {})", - native.function_name(), - table.join(".").replace('\'', "''"), - native.dataset_version, - native.column.replace('\'', "''"), - native.distance_threshold - )); - } else { - sql.push_str(&table.join(".")); - } + sql.push_str(&table.join(".")); if let Some(lateral) = &definition.lateral { match &lateral.source { LateralSource::Unnest { column } => { @@ -555,45 +511,6 @@ pub fn render(definition: &MaterializedViewDefinition) -> String { mod tests { use super::*; - #[test] - fn native_pairs_are_a_table_source_with_literal_snapshot_arguments() { - for threshold in ["4", "4.0", "-1", "0.05"] { - let sql = - format!("SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', {threshold})"); - let definition = parse(&sql).unwrap(); - assert_eq!(definition.source_table, "images"); - assert!(definition.lateral.is_none()); - assert_eq!( - definition.vector_source.as_ref().unwrap().dataset_version, - 3 - ); - assert_eq!(parse(&render(&definition)).unwrap(), definition); - let metadata = std::collections::HashMap::from([( - super::super::DEFINITION_META_KEY.to_string(), - definition.to_json().unwrap(), - )]); - assert_eq!( - super::super::read_definition(&metadata).unwrap(), - Some(super::super::StoredDefinition::Query(definition)) - ); - } - for sql in [ - "SELECT row_id_a FROM vector_duplicate_pairs('images', 3, 'phash', 4)", - "SELECT * FROM vector_duplicate_pairs('images', 0, 'phash', 4)", - "SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 1e100)", - "SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', cutoff)", - "SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) LIMIT 1", - "SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) WHERE distance < 2", - "SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) AS p", - "SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4), f(phash) AS p", - ] { - assert!( - parse(sql).is_err(), - "accepted unsupported native view: {sql}" - ); - } - } - #[test] fn a_query_round_trips_through_its_canonical_form() { for (sql, canonical) in [ diff --git a/rust/lancedb/src/materialized_view/refresh.rs b/rust/lancedb/src/materialized_view/refresh.rs index 9f0ea2646..f9634837a 100644 --- a/rust/lancedb/src/materialized_view/refresh.rs +++ b/rust/lancedb/src/materialized_view/refresh.rs @@ -1,63 +1,37 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! Refreshing materialized views. +//! Full refresh for materialized views. //! -//! A refresh pins one source version and brings the view to exactly the -//! definition's result at that version: added rows are computed and appended, -//! removed or changed rows are evicted by provenance id and recomputed in the -//! same pass. Compaction outputs cost nothing, which is only sound while -//! [`SOURCE_ROW_ID_COLUMN`] stays valid across the rewrite. Anything the -//! classifier cannot prove intact rebuilds; an indexed rebuild swaps all -//! fragments in one commit that retains index definitions. -//! -//! The watermark ([`SOURCE_VERSION_META_KEY`]) lands in a follow-up commit; a -//! crash or race between the two leaves the view visibly unstamped and the -//! next refresh rebuilds. In-process refreshes serialize on a per-view lock; -//! across processes the commit's inserted-rows filter carries a shared token, -//! so two refreshes of one view conflict and only one lands. +//! A refresh evaluates the definition at one source version, stages the +//! complete result, and atomically replaces every fragment of the view. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; -use std::time::{SystemTime, UNIX_EPOCH}; use arrow_array::cast::AsArray; -use arrow_array::types::UInt64Type; -use arrow_array::{RecordBatch, UInt64Array, new_null_array}; -use arrow_schema::{DataType, Field as ArrowField, FieldRef, Schema as ArrowSchema, SchemaRef}; -use datafusion::common::ScalarValue; +use arrow_array::{RecordBatch, new_null_array}; +use arrow_schema::{Schema as ArrowSchema, SchemaRef}; use datafusion::error::DataFusionError; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::prelude::{col, lit}; -use futures::{StreamExt, TryStreamExt}; +use futures::StreamExt; use lance::Dataset; use lance::dataset::mem_wal::DatasetMemWalExt; -use lance::dataset::transaction::{Operation, Transaction, UpdateMode}; -use lance::dataset::write::delete::DeleteBuilder; +use lance::dataset::transaction::{Operation, Transaction}; use lance::dataset::write::merge_insert::inserted_rows::{ KeyExistenceFilter, KeyExistenceFilterBuilder, KeyValue, }; use lance::dataset::{CommitBuilder, InsertBuilder, WriteDestination, WriteMode, WriteParams}; -use lance_core::{ROW_CREATED_AT_VERSION, ROW_ID, ROW_LAST_UPDATED_AT_VERSION}; - use lance_datafusion::planner::Planner; -use lance_file::version::ConcreteFileVersion; -use lance_table::format::Fragment; use serde::{Deserialize, Serialize}; -use super::{ - DEFINITION_META_KEY, INCARNATION_META_KEY, MaterializedViewDefinition, - REFRESHED_AT_MS_META_KEY, SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY, - definition_to_metadata, -}; +use super::{MaterializedViewDefinition, StoredDefinition}; use crate::database::OpenTableRequest; -use crate::table::computed_columns::{ - computed_column_from_field, computed_columns, ensure_declarations_are_planned, -}; -use crate::table::{NativeTable, NativeTableExt, Table}; +use crate::table::computed_columns::{computed_column_from_field, ensure_declarations_are_planned}; +use crate::table::{NativeTableExt, Table}; use crate::{Error, Result}; /// How a refresh brought the view up to date. @@ -66,10 +40,6 @@ use crate::{Error, Result}; pub enum RefreshMode { /// The view was recomputed from scratch. Rebuild, - /// Rows from source fragments added since the last refresh were appended. - Incremental, - /// The view was already at the requested source version. - NoOp, } /// The result of refreshing a materialized view. @@ -77,28 +47,15 @@ pub enum RefreshMode { pub struct RefreshMaterializedViewResult { /// How the view was brought up to date. pub mode: RefreshMode, - /// Rows written to the view: everything on a rebuild, and on an - /// incremental refresh both the rows added and the rows recomputed in - /// place of ones the source changed. + /// Rows written to the replacement view. pub rows_written: u64, - /// The source table version the view now reflects. + /// The source table version the view reflects. pub source_version: u64, /// The view table version after the refresh. pub version: u64, } -/// Schema metadata key holding the view table version a successful refresh -/// left behind. Any other commit on the view is drift, and refresh rebuilds. -pub const VIEW_VERSION_META_KEY: &str = "mv.view_version"; - -/// Schema metadata key holding the commit timestamp of the watermark's source -/// manifest. A dropped and recreated source reuses version numbers but never -/// their timestamps, so a mismatch means the watermark describes a different -/// incarnation and refresh rebuilds. -pub const SOURCE_VERSION_TS_META_KEY: &str = "mv.source_version_ts"; - -/// One refresh per view at a time within this process. -pub(super) fn refresh_lock(uri: &str) -> Arc> { +fn refresh_lock(uri: &str) -> Arc> { static LOCKS: OnceLock>>>> = OnceLock::new(); LOCKS @@ -110,57 +67,27 @@ pub(super) fn refresh_lock(uri: &str) -> Arc> { .clone() } -/// Internal implementation of the refresh logic. +/// Recompute a local materialized view from one source version. pub(crate) async fn execute_refresh( view: &Table, - full: bool, pinned: Option, - expected_incarnation: Option<&str>, ) -> Result { - // Native table functions use the same partition writer and commit sink in - // local execution; Sophon drives these units through its job registry. - if let Some(super::StoredDefinition::Query(definition)) = - super::read_definition(view.schema().await?.metadata())? - && definition.vector_source.is_some() - { - let plan = super::plan_partitioned_refresh(view, pinned) - .await? - .ok_or_else(|| Error::Runtime { - message: "native pair view has no partition plan".into(), - })?; - let mut units = Vec::with_capacity(plan.units() as usize); - for unit in 0..plan.units() { - let dependencies = &units[..units.len().min(plan.dependency_units() as usize)]; - units.push( - super::write_refresh_partition_with_inputs(view, unit, &plan, dependencies).await?, - ); - } - return super::commit_partitioned_refresh(view, &plan, units, expected_incarnation).await; - } - let view_native = view.as_native().ok_or_else(|| Error::NotSupported { + let native = view.as_native().ok_or_else(|| Error::NotSupported { message: "materialized views are supported only on local tables".into(), })?; - view_native.dataset.ensure_mutable()?; - let lock = refresh_lock(view_native.dataset.get().await?.uri()); + native.dataset.ensure_mutable()?; + let lock = refresh_lock(native.dataset.get().await?.uri()); let _guard = lock.lock().await; - // Force-load the latest view state under the lock: each handle caches - // lazily, and a second handle would otherwise plan from a snapshot taken - // before another handle's commit -- appending the same rows again or - // reporting NoOp over a mutated view. - view_native.dataset.reload().await?; - let view_ds = view_native.dataset.get().await?.as_ref().clone(); + native.dataset.reload().await?; + let view_ds = native.dataset.get().await?.as_ref().clone(); + ensure_no_mem_wal(&view_ds, "materialized view", view.name()).await?; - ensure_incarnation(&view_ds, expected_incarnation, view.name()).await?; - - // The definition a handle cached at open may since have been replaced; - // what refresh executes and what it stamps must be one generation. let definition = match super::read_definition(&view_ds.schema().metadata)? { - Some(super::StoredDefinition::Query(definition)) => definition, - Some(super::StoredDefinition::Newer { format }) => { + Some(StoredDefinition::Query(definition)) => definition, + Some(StoredDefinition::Newer { format }) => { return Err(Error::NotSupported { message: format!( - "materialized view '{}' is stored in format {format}, which this \ - version of lancedb cannot refresh", + "materialized view '{}' is stored in format {format}, which this version of lancedb cannot refresh", view.name() ), }); @@ -171,572 +98,216 @@ pub(crate) async fn execute_refresh( }); } }; - let definition = &definition; - let staging = super::read_staging(&view_ds.schema().metadata)?; - ensure_no_mem_wal(&view_ds, "materialized view", view.name()).await?; - - let source_ds = open_source(view, definition, staging.as_ref()).await?; - let source_ds = match pinned { - Some(version) => source_ds.checkout_version(version).await?, - None => source_ds, + let source = open_source(view, &definition).await?; + let source = match pinned { + Some(version) => source.checkout_version(version).await?, + None => source, }; - ensure_no_mem_wal(&source_ds, "source table", &definition.source_table).await?; - let source_version = source_ds.version().version; - let source_ts = source_ds.manifest.timestamp_nanos; + ensure_no_mem_wal(&source, "source table", &definition.source_table).await?; + let source_version = source.version().version; - // Re-plan the persisted definition against the current source schema and - // require its planned output to be exactly the view's physical schema: a - // definition the stored table cannot represent must not be certified. - let source_schema = Arc::new(ArrowSchema::from(source_ds.schema())); + let source_schema = Arc::new(ArrowSchema::from(source.schema())); let super::Planned { - definition: replanned, + definition, fields: planned_fields, inputs, .. - } = super::plan(source_schema.clone(), definition, staging.as_ref()).map_err(|e| match e { - // The stored query planned when the view was declared; what changed - // since is the source. + } = super::plan(source_schema, &definition).map_err(|error| match error { Error::InvalidExpression { column, message } => Error::Schema { message: format!( - "view column '{column}' no longer plans against '{}' (a source column \ - was dropped or renamed): {message}", + "view column '{column}' no longer plans against '{}' (a source column was dropped or renamed): {message}", definition.source_table ), }, Error::InvalidInput { message } => Error::Schema { message: format!( - "the stored query no longer plans against '{}' (a source column was \ - dropped or renamed): {message}", + "the stored query no longer plans against '{}' (a source column was dropped or renamed): {message}", definition.source_table ), }, - e => e, + error => error, })?; - let mut planned_fields = planned_fields; - planned_fields.push(arrow_schema::Field::new( - SOURCE_ROW_ID_COLUMN, - arrow_schema::DataType::UInt64, - false, - )); - // A computed column is not planned from the source: refresh writes it - // NULL and its declaration's owner fills it. Its declaration must still - // be complete, and it must be able to hold NULL. - let physical = ArrowSchema::from(view_ds.schema()); - let mut computed = computed_columns(&physical).into_iter().map(|c| c.name); - if let Some(name) = computed.by_ref().find(|name| { - physical - .field_with_name(name) - .is_ok_and(|f| !f.is_nullable()) - }) { + + let physical = Arc::new(ArrowSchema::from(view_ds.schema())); + ensure_declarations_are_planned(&physical)?; + if let Some(field) = physical + .fields() + .iter() + .find(|field| computed_column_from_field(field).is_some() && !field.is_nullable()) + { return Err(Error::Schema { message: format!( - "computed column '{name}' of view '{}' cannot hold NULL; recreate the view", + "computed column '{}' of view '{}' cannot hold NULL; recreate the view", + field.name(), view.name() ), }); } - ensure_declarations_are_planned(&physical)?; - let physical_planned: Vec<&FieldRef> = physical + let stored_fields: Vec<_> = physical .fields() .iter() - .filter(|f| computed_column_from_field(f).is_none()) + .filter(|field| computed_column_from_field(field).is_none()) .collect(); - // A projected column that became nullable at the source still fits the - // view's nullable field; the reverse would not. - let matches = planned_fields.len() == physical_planned.len() - && planned_fields.iter().zip(&physical_planned).all(|(e, p)| { - e.name() == p.name() - && e.data_type() == p.data_type() - && (p.is_nullable() || !e.is_nullable()) - }); + let matches = planned_fields.len() == stored_fields.len() + && planned_fields + .iter() + .zip(stored_fields) + .all(|(planned, stored)| { + planned.name() == stored.name() + && planned.data_type() == stored.data_type() + && (stored.is_nullable() || !planned.is_nullable()) + }); if !matches { return Err(Error::Schema { message: format!( - "the stored definition of view '{}' does not produce this \ - view's schema; recreate the view", + "the stored definition of view '{}' does not produce this view's schema; recreate the view", view.name() ), }); } - // The stored query is rewritten whenever its stored form differs from - // the current one: a legacy layout, or a spelling the canonicalizer no - // longer produces. Whether the rows change is a separate question: a - // legacy raw filter like `"Party" = 'D'` read the double quotes as a - // string literal, so its watermark certifies different rows than the - // canonical predicate, and only a rebuild can replace them. - let current = definition_to_metadata(&replanned)?; - let persist = view_ds.schema().metadata.get(DEFINITION_META_KEY) != Some(¤t); - let unnest = super::physical_unnest(&replanned, staging.as_ref())?; - let definition_changed = !same_meaning(&source_schema, definition, &replanned, unnest.as_ref()); - let definition = &replanned; - let persist = persist.then_some(definition); - if definition_changed { - return rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - unnest.as_ref(), - true, - expected_incarnation, - ) - .await; - } - - let metadata = &view_ds.schema().metadata; - let watermark: Option = metadata - .get(SOURCE_VERSION_META_KEY) - .and_then(|raw| raw.parse().ok()); - let recorded_ts: Option = metadata - .get(SOURCE_VERSION_TS_META_KEY) - .and_then(|raw| raw.parse().ok()); - // The watermark speaks only for the view state its refresh left behind; - // any other commit on the view since then is drift, except a fill of its - // computed columns, which rewrites nothing refresh certifies. - let recorded_view_version = metadata - .get(VIEW_VERSION_META_KEY) - .and_then(|raw| raw.parse::().ok()); - let view_intact = match recorded_view_version { - Some(recorded) if recorded == view_ds.version().version => true, - Some(recorded) if recorded < view_ds.version().version => { - only_computed_rewrites_since(&view_ds, recorded).await? - } - _ => false, - }; - - if !full && watermark == Some(source_version) && view_intact && recorded_ts == Some(source_ts) { - return Ok(RefreshMaterializedViewResult { - mode: RefreshMode::NoOp, - rows_written: 0, - source_version, - version: view_ds.version().version, - }); - } - - // A group spans fragments: there is no per-fragment increment. - if definition.is_grouped() { - return rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - None, - persist.is_some(), - expected_incarnation, - ) - .await; - } - let watermark = watermark.filter(|_| view_intact); - match plan_increment( - &source_ds, - source_version, - watermark, - recorded_ts, - full, - definition, + let rows_written = Arc::new(AtomicU64::new(0)); + let unnest = super::physical_unnest(&definition)?; + let stream = compute_stream( + &source, + &definition, &inputs, + unnest.as_ref(), + physical, + rows_written.clone(), ) - .await - { - Some(increment) => { - let reconciled = incremental( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - increment, - definition, - &inputs, - unnest.as_ref(), - persist, - watermark, - expected_incarnation, - ) - .await?; - match reconciled { - Some(result) => Ok(result), - // The delta was too large to reconcile in bounded memory. - None => { - rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - unnest.as_ref(), - persist.is_some(), - expected_incarnation, - ) - .await - } - } - } - None => { - rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - unnest.as_ref(), - persist.is_some(), - expected_incarnation, - ) - .await - } - } -} - -/// The source fragments whose rows are new since the watermark, or `None` -/// where the view has to rebuild. Two tiers: the transaction walk is exact -/// where it applies; the fragment-signature check is the fallback for deltas -/// the walk cannot read, and under it any fragment churn rebuilds. -async fn plan_increment( - source_ds: &Dataset, - source_version: u64, - watermark: Option, - recorded_ts: Option, - full: bool, - definition: &MaterializedViewDefinition, - inputs: &[String], -) -> Option { - if full { - return None; - } - let watermark = watermark?; - if watermark > source_version { - return None; - } - let old = source_ds.checkout_version(watermark).await.ok()?; - // A recreated source reuses version numbers, never their timestamps: a - // mismatch means the watermark describes a different incarnation. - if recorded_ts != Some(old.manifest.timestamp_nanos) { - return None; - } - let old_ids: HashSet = old.get_fragments().iter().map(|f| f.id() as u64).collect(); - let live: Vec = source_ds - .get_fragments() - .iter() - .map(|f| f.metadata().clone()) - .collect(); - - if let Some(delta) = appends_and_rewrites(source_ds, watermark, source_version).await { - // A rewrite that consumed a fragment neither present at the watermark - // nor produced by an earlier rewrite swallowed a mid-delta append; - // its rows cannot be told apart from already-materialized ones. - let folded = delta - .rewritten - .iter() - .any(|id| !old_ids.contains(id) && !delta.produced.contains(id)); - if folded { - return None; - } - // An update rewrites a whole fragment. If it touched one the watermark - // never saw, that fragment holds rows appended since -- new rows the - // update did not change, which the recompute does not cover and which - // this fragment's exclusion from the append set would drop. - if delta - .updated_in_place - .iter() - .any(|id| !old_ids.contains(id)) - { - return None; - } - // Rows past the cap left every delta when the watermark advanced, so - // a capped view cannot reconcile a removal incrementally. The rebuild - // is cheap for the same reason it is capped: the scan stops there. - if definition.limit.is_some() && (delta.deleted_rows || delta.updated_rows) { - return None; - } - // Legacy storage cannot serve the row-version columns update - // discovery scans; deletes and appends need none of them. - if delta.updated_rows - && source_ds.manifest.data_storage_format.lance_file_format() == ConcreteFileVersion::V1 - { - return None; - } - // Every other fragment new at head is an append: appends and rewrites - // are the only operations in the delta that add fragments, and the - // rewrite outputs are already-materialized rows rearranged. - return Some(Increment { - appended: live - .into_iter() - .filter(|f| !old_ids.contains(&f.id) && !delta.produced.contains(&f.id)) - .collect(), - evict_deleted: delta.deleted_rows, - replace_updated: delta.updated_rows, - }); - } - - is_pure_append(&old, source_ds, &relevant_field_ids(source_ds, inputs)).then(|| Increment { - appended: live - .into_iter() - .filter(|f| !old_ids.contains(&f.id)) - .collect(), - evict_deleted: false, - replace_updated: false, + .await?; + let committed = replace_fragments(view_ds, stream).await?; + let version = committed.version().version; + native.dataset.update(committed); + Ok(RefreshMaterializedViewResult { + mode: RefreshMode::Rebuild, + rows_written: rows_written.load(Ordering::Relaxed), + source_version, + version, }) } -/// Version gap beyond which per-version transaction reads stop being cheaper -/// than one fragment scan. -const MAX_TRANSACTION_WALK: u64 = 512; - -/// Fragment ids moved by the `Rewrite` operations of the delta. Only rewrite -/// ids are real in transaction files (an Append's are placeholders assigned -/// at commit), so appends are derived as new-at-head minus rewrite outputs. -/// What an incremental refresh must do to bring the view up to date. -struct Increment { - /// Source fragments whose rows are not in the view yet. - appended: Vec, - /// Source rows left the range, so the view holds rows to evict. - evict_deleted: bool, - /// Source rows changed in the range, so the view holds rows to replace. - replace_updated: bool, -} - -struct TxnDelta { - /// Fragments consumed by `Rewrite` operations. - rewritten: HashSet, - /// Fragments produced by `Rewrite` operations. - produced: HashSet, - /// The delta removed source rows, so the view holds rows to evict. - deleted_rows: bool, - /// The delta changed source rows in place, so the view holds rows to - /// recompute. - updated_rows: bool, - /// Fragments an update modified in place, as opposed to produced. - updated_in_place: HashSet, -} - -/// Read the delta from the transaction log, `None` where it holds anything -/// but appends and rewrites or cannot be read; `None` only sends the caller -/// to a slower check. `ReserveFragments` moves no rows and rides along. -async fn appends_and_rewrites(cur: &Dataset, from: u64, to: u64) -> Option { - if to <= from || to - from > MAX_TRANSACTION_WALK { - return None; +/// Atomically replace all fragments of a materialized view with a complete +/// query result. This is shared with Sophon's SQL refresh job. +pub async fn replace_materialized_view_fragments( + view: &Table, + stream: SendableRecordBatchStream, +) -> Result { + let native = view.as_native().ok_or_else(|| Error::NotSupported { + message: "materialized views are supported only on local tables".into(), + })?; + native.dataset.ensure_mutable()?; + let lock = refresh_lock(native.dataset.get().await?.uri()); + let _guard = lock.lock().await; + native.dataset.reload().await?; + let dataset = native.dataset.get().await?.as_ref().clone(); + if super::read_definition_sql(&dataset.schema().metadata)?.is_none() { + return Err(Error::NotAMaterializedView { + name: view.name().to_string(), + }); } - let mut delta = TxnDelta { - rewritten: HashSet::new(), - produced: HashSet::new(), - deleted_rows: false, - updated_rows: false, - updated_in_place: HashSet::new(), + let committed = replace_fragments(dataset, stream).await?; + let version = committed.version().version; + native.dataset.update(committed); + Ok(version) +} + +async fn replace_fragments(dataset: Dataset, stream: SendableRecordBatchStream) -> Result { + let dataset = Arc::new(dataset); + let read_version = dataset.version().version; + let removed_fragment_ids = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.id() as u64) + .collect(); + let write = InsertBuilder::new(WriteDestination::Dataset(dataset.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted_stream(stream) + .await?; + let Operation::Append { + fragments: new_fragments, + } = write.operation + else { + return Err(Error::Runtime { + message: "expected an append while staging replacement rows".into(), + }); }; - for version in (from + 1)..=to { - let Ok(Some(txn)) = cur.read_transaction_by_version(version).await else { - return None; - }; - match txn.operation { - Operation::Append { .. } | Operation::ReserveFragments { .. } => {} - // A delete removes source rows without changing the ones that - // remain, so the view's other rows stay valid: the refresh - // evicts exactly the ids that left. - Operation::Delete { .. } => delta.deleted_rows = true, - // Update outputs carry no new rows, so they are excluded from - // the append set like rewrite outputs; the changed rows are - // replaced individually below. + let inserted_rows_filter = refresh_conflict_filter(&dataset)?; + let committed = CommitBuilder::new(WriteDestination::Dataset(dataset)) + .execute(Transaction::new( + read_version, Operation::Update { removed_fragment_ids, + updated_fragments: Vec::new(), new_fragments, - updated_fragments, - .. - } => { - delta.updated_rows = true; - // merge_insert reaches here too, and its by-source arm deletes - // rows rather than changing them. - delta.deleted_rows = true; - delta.rewritten.extend(removed_fragment_ids.iter().copied()); - // Only pre-existing fragment ids are real here; created ones - // are placeholders. Rewritten rows are excluded by creation - // version below, not by fragment identity. - delta - .produced - .extend(updated_fragments.iter().map(|f| f.id)); - delta - .updated_in_place - .extend(updated_fragments.iter().map(|f| f.id)); - let _ = new_fragments; - } - Operation::Rewrite { groups, .. } => { - for group in groups { - delta - .rewritten - .extend(group.old_fragments.iter().map(|f| f.id)); - delta - .produced - .extend(group.new_fragments.iter().map(|f| f.id)); - } - } - _ => return None, - } + fields_modified: Vec::new(), + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: Vec::new(), + update_mode: None, + inserted_rows_filter: Some(inserted_rows_filter), + updated_fragment_offsets: None, + }, + None, + )) + .await?; + if committed.version().version != read_version + 1 { + return Err(Error::Runtime { + message: format!( + "a concurrent commit raced this refresh (view version {}, expected {})", + committed.version().version, + read_version + 1 + ), + }); } - Some(delta) + Ok(committed) } -/// Fallback pure-append check: every old fragment still present with an -/// identical signature over the columns the view reads. Compaction, deletes -/// and updates each break it and force a rebuild; a change to a column the -/// view does not read leaves it alone, which is what lets this tier pass -/// deltas the transaction walk cannot. -fn is_pure_append(old: &Dataset, cur: &Dataset, relevant: &HashSet) -> bool { - let signature = |fragment: &lance::dataset::fragment::FileFragment| { - fragment_signature(fragment.metadata(), relevant) - }; - let current: HashSet<(u64, String)> = cur.get_fragments().iter().map(signature).collect(); - old.get_fragments() - .iter() - .all(|fragment| current.contains(&signature(fragment))) +const REFRESH_TOKEN_ID: u64 = u64::MAX; + +fn refresh_conflict_filter(dataset: &Dataset) -> Result { + let field_id = dataset + .schema() + .fields + .first() + .map(|field| field.id) + .ok_or_else(|| Error::Runtime { + message: "a materialized view must have at least one field".into(), + })?; + let mut filter = KeyExistenceFilterBuilder::new(vec![field_id]); + filter + .insert(KeyValue::UInt64(REFRESH_TOKEN_ID)) + .map_err(|error| Error::Runtime { + message: format!( + "failed to build the materialized-view refresh conflict filter: {error}" + ), + })?; + Ok(filter.build()) } -/// A fragment's identity as the view observes it: data files and overlays -/// touching the columns it reads, plus the deletion file. Overlays change no -/// file path, so they must be part of the signature. -fn fragment_signature(metadata: &Fragment, relevant: &HashSet) -> (u64, String) { - let touches_relevant = - |fields: &[i32]| relevant.is_empty() || fields.iter().any(|id| relevant.contains(id)); - let mut files: Vec<&str> = metadata - .files - .iter() - .filter(|file| touches_relevant(&file.fields)) - .map(|file| file.path.as_str()) - .collect(); - files.sort_unstable(); - let mut overlays: Vec = metadata - .overlays - .iter() - .filter(|overlay| touches_relevant(&overlay.data_file.fields)) - .map(|overlay| format!("{}@{}", overlay.data_file.path, overlay.committed_version)) - .collect(); - overlays.sort_unstable(); - ( - metadata.id, - format!( - "{}|{}|{:?}", - files.join(","), - overlays.join(","), - metadata.deletion_file - ), - ) -} - -/// Field ids (with struct descendants) of the source columns the view reads. -fn relevant_field_ids(source: &Dataset, inputs: &[String]) -> HashSet { - fn collect(field: &lance_core::datatypes::Field, ids: &mut HashSet) { - ids.insert(field.id); - for child in &field.children { - collect(child, ids); - } - } - let mut ids = HashSet::new(); - for input in inputs { - if let Some(field) = source.schema().field(input) { - collect(field, &mut ids); - } - } - ids -} - -/// Whether two plannings of a view compute the same rows and columns: the -/// same source, unnest and limit, and expressions the planner reads as the -/// same logical expression, whatever their spelling. A definition that does -/// not plan compares as different. -fn same_meaning( - source_schema: &SchemaRef, - stored: &MaterializedViewDefinition, - replanned: &MaterializedViewDefinition, - unnest: Option<&super::ViewUnnest>, -) -> bool { - if stored.source_table != replanned.source_table - || stored.source_namespace != replanned.source_namespace - || stored.lateral != replanned.lateral - || stored.limit != replanned.limit - || stored.projections.len() != replanned.projections.len() - || stored.filter.is_some() != replanned.filter.is_some() - || stored.group_by.len() != replanned.group_by.len() - { - return false; - } - // Lance's planner does not read aggregates; a grouped definition is - // compared in its canonical spelling, which planning produced. - if stored.is_grouped() { - return stored.projections == replanned.projections - && stored.group_by == replanned.group_by - && stored.filter == replanned.filter; - } - let schema = match unnest { - None => source_schema.clone(), - Some(unnest) => match super::flattened_schema(source_schema, unnest) { - Ok(schema) => schema, - Err(_) => return false, - }, - }; - let planner = Planner::new(schema); - let same_expr = |a: &str, b: &str| match (planner.parse_expr(a), planner.parse_expr(b)) { - (Ok(a), Ok(b)) => a == b, - _ => false, - }; - stored - .projections - .iter() - .zip(&replanned.projections) - .all(|(a, b)| a.output == b.output && same_expr(&a.expression, &b.expression)) - && match (&stored.filter, &replanned.filter) { - (Some(a), Some(b)) => same_expr(a, b), - _ => true, - } -} - -/// Reject MemWAL/LSM state on a refresh participant: un-compacted tiers are -/// invisible to the fragment-planned refresh scan. An active write spec and -/// retained rows both disqualify; shard directories on storage are the -/// durable evidence of the latter. +/// Reject MemWAL/LSM state because an ordinary dataset scan cannot see its +/// un-compacted tiers. pub(crate) async fn ensure_no_mem_wal(dataset: &Dataset, role: &str, name: &str) -> Result<()> { let retained = !dataset.list_mem_wal_latest_shard_ids().await?.is_empty(); if retained || dataset.mem_wal_index_details().await?.is_some() { return Err(Error::NotSupported { - message: format!( - "{role} '{name}' has an LSM write spec or retained un-compacted \ - rows: rows in un-compacted tiers are invisible to refresh" - ), + message: format!("{role} '{name}' has an LSM write spec or retained un-compacted rows"), }); } Ok(()) } -/// The table refresh scans: the staging table when the query calls a -/// Function in FROM position, otherwise the query's source. -pub(super) async fn open_source( - view: &Table, - definition: &MaterializedViewDefinition, - staging: Option<&super::StagingBinding>, -) -> Result { +async fn open_source(view: &Table, definition: &MaterializedViewDefinition) -> Result { let database = view.database_opt().ok_or_else(|| Error::InvalidInput { message: "the view was not opened through a database connection".into(), })?; - let (name, namespace_path) = match staging { - Some(staging) => (staging.table.clone(), staging.namespace.clone()), - None => ( - definition.source_table.clone(), - definition.source_namespace.clone(), - ), - }; let source = database .open_table(OpenTableRequest { - name, - namespace_path, + name: definition.source_table.clone(), + namespace_path: definition.source_namespace.clone(), index_cache_size: None, lance_read_params: None, location: None, @@ -747,584 +318,7 @@ pub(super) async fn open_source( let native = source.as_native().ok_or_else(|| Error::NotSupported { message: "materialized views are supported only on local tables".into(), })?; - let dataset = native.dataset.get().await?.as_ref().clone(); - if definition.vector_source.is_none() && !dataset.manifest.uses_stable_row_ids() { - return Err(Error::InvalidInput { - message: format!( - "source table '{}' does not have stable row ids; it is not the \ - table this view was declared over", - definition.source_table - ), - }); - } - Ok(dataset) -} - -#[allow(clippy::too_many_arguments)] -async fn incremental( - view_native: &NativeTable, - view_ds: &Dataset, - source_ds: &Dataset, - source_version: u64, - source_ts: u128, - increment: Increment, - definition: &MaterializedViewDefinition, - inputs: &[String], - unnest: Option<&super::ViewUnnest>, - persist: Option<&MaterializedViewDefinition>, - watermark: Option, - expected_incarnation: Option<&str>, -) -> Result> { - let new_fragments = increment.appended; - let watermark_version = watermark.unwrap_or(0); - // Provenance ids this refresh removes: dropped rows, plus changed rows - // recomputed in the same commit. One view row per source row makes the - // eviction exact. Staged, not committed: removals ride with the rows - // that replace them, so a reader never sees the view without either. - let mut eviction = Eviction::new(view_ds, EVICTION_CHUNK); - let mut updated_rows = false; - if (increment.evict_deleted || increment.replace_updated) - && let Some(watermark) = watermark - { - let delta = source_ds - .delta() - .with_begin_version(watermark) - .with_end_version(source_version) - .build()?; - // Reconciling holds the delta's provenance ids in staged deletion - // vectors; past the cap, the streamed rebuild is the bounded path. - let cap = eviction_rebuild_cap(); - let mut evicted = 0usize; - if increment.evict_deleted { - let mut stream = delta.get_deleted_row_ids().await?; - while let Some(batch) = stream.try_next().await? { - let ids = row_ids_of(&batch)?; - evicted += ids.len(); - if evicted > cap { - return Ok(None); - } - eviction.push(ids).await?; - } - } - if increment.replace_updated { - // Ids only: `get_updated_rows` carries every column of every - // updated row, and discovery needs none of them. - let mut scanner = source_ds.scan(); - scanner.with_row_id().project(&[ROW_CREATED_AT_VERSION])?; - // A fixed bound: the configured default could make one discovery - // batch arbitrarily large before the fallback cap is consulted. - scanner.batch_size(8192); - scanner.filter(&format!( - "{ROW_CREATED_AT_VERSION} <= {watermark} - AND {ROW_LAST_UPDATED_AT_VERSION} > {watermark} - AND {ROW_LAST_UPDATED_AT_VERSION} <= {source_version}" - ))?; - let mut stream = scanner.try_into_stream().await?; - while let Some(batch) = stream.try_next().await? { - let ids = row_ids_of(&batch)?; - evicted += ids.len(); - if evicted > cap { - return Ok(None); - } - updated_rows |= !ids.is_empty(); - eviction.push(ids).await?; - } - } - } - let eviction = eviction.finish().await?; - - // The cap counts rows already materialized, in first-materialized order. - let remaining = match definition.limit { - Some(limit) => { - let held = view_ds.count_rows(None).await? as u64; - Some(limit.saturating_sub(held)) - } - None => None, - }; - - let mut result = RefreshMaterializedViewResult { - mode: RefreshMode::Incremental, - rows_written: 0, - source_version, - version: view_ds.version().version, - }; - let nothing_to_add = (new_fragments.is_empty() && !updated_rows) || remaining == Some(0); - if nothing_to_add && eviction.is_none() { - result.version = stamp_watermark( - view_native, - view_ds.clone(), - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - return Ok(Some(result)); - } - // Rows left but none arrive: the removals still have to be published. - if nothing_to_add { - let filter = refresh_filter(&empty_keys(view_ds)?)?; - let published = publish( - view_ds, - eviction, - Vec::new(), - Some(filter), - expected_incarnation, - ) - .await?; - result.version = stamp_watermark( - view_native, - published, - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - return Ok(Some(result)); - } - - // Appends carry the view's schema as it stands; the watermark moves in a - // follow-up commit (see the module docs for the crash window). - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let rows_written = Arc::new(AtomicU64::new(0)); - // compute_stream counts what it produces; the truncation below can drop - // some of that, so the written count comes from the tee instead. - let computed = Arc::new(AtomicU64::new(0)); - let mut stream = compute_stream( - source_ds, - definition, - inputs, - unnest, - RowScope { - fragments: Some(new_fragments), - // An update rewrites whole fragments, so a fragment new at head - // can hold rows the view already has. Their creation version - // does not change, so it -- not fragment identity -- says which - // rows are new. - created_after: increment.replace_updated.then_some(watermark_version), - limit: remaining, - ..Default::default() - }, - schema.clone(), - computed.clone(), - ) - .await?; - - // The updated rows' current values, computed the same way and appended - // in the same commit as the new fragments' rows. - if updated_rows { - let recomputed = compute_stream( - source_ds, - definition, - inputs, - unnest, - RowScope { - updated_between: Some((watermark_version, source_version)), - ..Default::default() - }, - schema.clone(), - computed.clone(), - ) - .await?; - stream = Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), - recomputed.chain(stream), - )); - } - - // Nothing survived the filter: the watermark still has to advance or the - // same fragments would be rescanned forever, but any removals still do. - let Some(first) = stream.try_next().await? else { - let published = if eviction.is_some() { - publish( - view_ds, - eviction, - Vec::new(), - Some(refresh_filter(&empty_keys(view_ds)?)?), - expected_incarnation, - ) - .await? - } else { - view_ds.clone() - }; - result.version = stamp_watermark( - view_native, - published, - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - return Ok(Some(result)); - }; - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::iter([Ok(first)]).chain(stream), - )); - - // Any two refreshes of one view must not both commit: the filter's - // shared token makes lance reject the loser on key overlap however - // their planned rows relate. - let keys = Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new(vec![ - source_row_id_field_id(view_ds)?, - ]))); - let stream = collect_source_row_ids(stream, keys.clone(), rows_written.clone()); - - let ds = Arc::new(view_ds.clone()); - let write_txn = InsertBuilder::new(WriteDestination::Dataset(ds.clone())) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { - fragments: new_fragments, - } = write_txn.operation - else { - return Err(Error::Runtime { - message: "expected an append when staging the view's new rows".into(), - }); - }; - let filter = refresh_filter(&keys)?; - let appended = publish( - view_ds, - eviction, - new_fragments, - Some(filter), - expected_incarnation, - ) - .await?; - result.rows_written = rows_written.load(Ordering::Relaxed); - result.version = stamp_watermark( - view_native, - appended, - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - Ok(Some(result)) -} - -#[allow(clippy::too_many_arguments)] -async fn rebuild( - view_native: &NativeTable, - view_ds: &Dataset, - source_ds: &Dataset, - source_version: u64, - source_ts: u128, - definition: &MaterializedViewDefinition, - inputs: &[String], - unnest: Option<&super::ViewUnnest>, - persist_definition: bool, - expected_incarnation: Option<&str>, -) -> Result { - let rows_written = Arc::new(AtomicU64::new(0)); - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let stream = compute_stream( - source_ds, - definition, - inputs, - unnest, - RowScope { - limit: definition.limit, - ..Default::default() - }, - schema, - rows_written.clone(), - ) - .await?; - let keys = Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new(vec![ - source_row_id_field_id(view_ds)?, - ]))); - let stream = collect_source_row_ids(stream, keys.clone(), Arc::new(AtomicU64::new(0))); - // Every rebuild is one fragment swap, indexed or not: an Update commit - // carries no schema metadata, so it cannot erase a definition update - // that raced in the way an overwrite (which adopts its stream's schema) - // durably would -- and it must land on the planned generation or abort. - let replaced = - replace_retaining_indices(view_ds.clone(), stream, keys, expected_incarnation).await?; - let version = stamp_watermark( - view_native, - replaced, - source_version, - source_ts, - persist_definition.then_some(definition), - expected_incarnation, - ) - .await?; - Ok(RefreshMaterializedViewResult { - mode: RefreshMode::Rebuild, - rows_written: rows_written.load(Ordering::Relaxed), - source_version, - version, - }) -} - -/// Replace all of the view's data in one commit that retains its index -/// definitions: new fragments staged uncommitted, one `Update` removing every -/// old fragment. `Update` prunes index bitmaps only for modified fields and -/// none are modified here, so readers never see the view unindexed or empty. -async fn replace_retaining_indices( - view_ds: Dataset, - stream: SendableRecordBatchStream, - keys: Arc>, - expected_incarnation: Option<&str>, -) -> Result { - let ds = Arc::new(view_ds); - let read_version = ds.version().version; - #[cfg(test)] - tests::hold_before_publish(ds.uri()).await; - ensure_incarnation(&ds, expected_incarnation, ds.uri()).await?; - let removed_fragment_ids: Vec = ds.get_fragments().iter().map(|f| f.id() as u64).collect(); - - let write_txn = InsertBuilder::new(WriteDestination::Dataset(ds.clone())) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { - fragments: new_fragments, - } = write_txn.operation - else { - return Err(Error::Runtime { - message: "expected an append when staging the view's replacement rows".into(), - }); - }; - - // Built only now: the tee fills as the staging drains the stream. - let filter = refresh_filter(&keys)?; - let transaction = Transaction::new( - read_version, - Operation::Update { - removed_fragment_ids, - updated_fragments: Vec::new(), - new_fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - // Two refreshes that materialized the same source rows must not - // both land -- a raced first rebuild would double the view. - inserted_rows_filter: Some(filter), - updated_fragment_offsets: None, - }, - None, - ); - let committed = CommitBuilder::new(WriteDestination::Dataset(ds)) - .execute(transaction) - .await?; - if committed.version().version != read_version + 1 { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced this refresh (view version {}); the - refresh is unrecorded and the next one will rebuild", - committed.version().version - ), - }); - } - Ok(committed) -} - -/// Record that the view now reflects `source_version`, including the view -/// Refuse to act on a view that is not `expected`'s incarnation, judged from -/// the latest stored manifest. Not a commit condition; see -/// `RefreshMaterializedViewBuilder::expect_incarnation`. -pub(super) async fn ensure_incarnation( - view_ds: &Dataset, - expected: Option<&str>, - what: &str, -) -> Result<()> { - let Some(expected) = expected else { - return Ok(()); - }; - let mut latest = view_ds.clone(); - latest.checkout_latest().await?; - match latest.schema().metadata.get(INCARNATION_META_KEY) { - Some(actual) if actual == expected => Ok(()), - Some(_) => Err(Error::Runtime { - message: format!( - "materialized view '{what}' is not the incarnation this refresh was \ - requested for: it was dropped and recreated" - ), - }), - None => Err(Error::Runtime { - message: format!( - "materialized view '{what}' carries no incarnation token: its schema \ - metadata was replaced since the token was captured" - ), - }), - } -} - -/// version this very commit produces. The version is predicted and then -/// verified; on a mismatch another commit raced in between, and the stamp -/// ABORTS rather than certify that commit as the refresh's own generation. -/// The view is left visibly unstamped, so the next refresh rebuilds. -pub(super) async fn stamp_watermark( - view_native: &NativeTable, - mut dataset: Dataset, - source_version: u64, - source_ts: u128, - definition: Option<&MaterializedViewDefinition>, - expected_incarnation: Option<&str>, -) -> Result { - ensure_incarnation(&dataset, expected_incarnation, dataset.uri()).await?; - let predicted = dataset.version().version + 1; - // A view with no token (declared before tokens existed, or its metadata - // replaced wholesale) starts a new incarnation here. - let incarnation = dataset - .schema() - .metadata - .get(INCARNATION_META_KEY) - .cloned() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let mut metadata = vec![(INCARNATION_META_KEY.to_string(), Some(incarnation))]; - if let Some(definition) = definition { - metadata.push(( - DEFINITION_META_KEY.to_string(), - Some(definition_to_metadata(definition)?), - )); - } - metadata.extend([ - ( - SOURCE_VERSION_META_KEY.to_string(), - Some(source_version.to_string()), - ), - ( - SOURCE_VERSION_TS_META_KEY.to_string(), - Some(source_ts.to_string()), - ), - ( - REFRESHED_AT_MS_META_KEY.to_string(), - Some(now_ms().to_string()), - ), - ( - VIEW_VERSION_META_KEY.to_string(), - Some(predicted.to_string()), - ), - ]); - dataset.update_schema_metadata(metadata).await?; - let actual = dataset.version().version; - if actual != predicted { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced this refresh (view version {actual}, \ - expected {predicted}); the refresh is unrecorded and the next \ - one will rebuild" - ), - }); - } - view_native.dataset.update(dataset); - Ok(predicted) -} - -fn now_ms() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() -} - -/// Evaluate the definition over `source`, restricted to `fragments` when -/// given, as batches in the view's schema. The filter is pushed into the -/// scan; projections are `(name, expression)` pairs, never spliced into SQL; -/// [`SOURCE_ROW_ID_COLUMN`] is filled by the scan's row id. -/// Which source rows a compute pass reads. -#[derive(Default)] -struct RowScope { - /// Read only these fragments. - fragments: Option>, - /// Read only rows created after this version. - created_after: Option, - /// Read only rows changed in place after the first version and no later - /// than the second. - updated_between: Option<(u64, u64)>, - /// Stop after this many rows. - limit: Option, -} - -/// Whether every commit on the view after `recorded` is a fill of its -/// computed columns: a column rewrite or data replacement touching only -/// those fields and neither adding nor removing rows, or the freshness -/// stamp a fill leaves on them. A version whose transaction cannot be read -/// is not proven, so it counts as drift. -async fn only_computed_rewrites_since(view_ds: &Dataset, recorded: u64) -> Result { - // A fill may write any field under a computed column, so the whole - // subtree counts, not only the root. - let physical = ArrowSchema::from(view_ds.schema()); - fn subtree(field: &lance_core::datatypes::Field, ids: &mut Vec) { - ids.push(field.id as u32); - for child in &field.children { - subtree(child, ids); - } - } - let mut computed_fields = Vec::new(); - for column in computed_columns(&physical) { - if let Some(field) = view_ds.schema().field(&column.name) { - subtree(field, &mut computed_fields); - } - } - if computed_fields.is_empty() { - return Ok(false); - } - for version in recorded + 1..=view_ds.version().version { - let Some(transaction) = view_ds.read_transaction_by_version(version).await? else { - return Ok(false); - }; - let fill = match &transaction.operation { - Operation::Update { - removed_fragment_ids, - new_fragments, - fields_modified, - update_mode: Some(UpdateMode::RewriteColumns), - .. - } => { - removed_fragment_ids.is_empty() - && new_fragments.is_empty() - && !fields_modified.is_empty() - && fields_modified - .iter() - .all(|field| computed_fields.contains(field)) - } - // What `refresh_column` commits for a SQL declaration. - Operation::DataReplacement { replacements } => { - !replacements.is_empty() - && replacements.iter().all(|group| { - !group.1.fields.is_empty() - && group - .1 - .fields - .iter() - .all(|field| computed_fields.contains(&(*field as u32))) - }) - } - // The stamp `refresh_column` writes after its fill (see - // `table::freshness`): field metadata on computed columns, no data. - Operation::UpdateConfig { - config_updates: None, - table_metadata_updates: None, - schema_metadata_updates: None, - field_metadata_updates, - } => { - !field_metadata_updates.is_empty() - && field_metadata_updates - .keys() - .all(|field| computed_fields.contains(&(*field as u32))) - } - _ => false, - }; - if !fill { - return Ok(false); - } - } - Ok(true) + Ok(native.dataset.get().await?.as_ref().clone()) } async fn compute_stream( @@ -1332,58 +326,23 @@ async fn compute_stream( definition: &MaterializedViewDefinition, inputs: &[String], unnest: Option<&super::ViewUnnest>, - scope: RowScope, schema: SchemaRef, rows_written: Arc, ) -> Result { if definition.is_grouped() { return super::grouped::stream(source, definition, inputs, schema, rows_written).await; } - let RowScope { - fragments, - created_after, - updated_between, - limit, - } = scope; let mut scanner = source.scan(); - if let Some(fragments) = fragments { - scanner.with_fragments(fragments); + if let Some(filter) = definition.filter.as_deref().filter(|_| unnest.is_none()) { + scanner.filter(filter)?; } - scanner.with_row_id(); - // Narrowing keeps the definition's filter, so a row updated out of the - // view simply does not come back. Changed rows are named by the predicate - // `DatasetDelta::get_updated_rows` uses, not its streamed ids: an id list - // grows with the delta, this does not. - let updated_filter = updated_between.map(|(from, to)| { - format!( - "{ROW_CREATED_AT_VERSION} <= {from} \ - AND {ROW_LAST_UPDATED_AT_VERSION} > {from} \ - AND {ROW_LAST_UPDATED_AT_VERSION} <= {to}" - ) - }); - let created_filter = - created_after.map(|version| format!("{ROW_CREATED_AT_VERSION} > {version}")); - let clauses: Vec = definition - .filter - .clone() - .filter(|_| unnest.is_none()) - .map(|f| format!("({f})")) - .into_iter() - .chain(updated_filter) - .chain(created_filter) - .collect(); - if !clauses.is_empty() { - scanner.filter(&clauses.join(" AND "))?; - } - // An expanded view cannot project or filter in the scan: both read the - // unnested element, which exists only after the per-batch expansion. let expanded = match unnest { Some(unnest) => Some(UnnestPlan::new(source, definition, inputs, unnest)?), None => { let transforms: Vec<(&str, &str)> = definition .projections .iter() - .map(|p| (p.output.as_str(), p.expression.as_str())) + .map(|projection| (projection.output.as_str(), projection.expression.as_str())) .collect(); scanner.project_with_transform(&transforms)?; None @@ -1392,54 +351,43 @@ async fn compute_stream( if let Some(expanded) = &expanded { scanner.project(&expanded.raw_inputs)?; } - // A scan reads a limit of zero as no limit at all, so a view capped at - // nothing is answered without one. - if limit == Some(0) { + if definition.limit == Some(0) { return Ok(Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), + schema, futures::stream::empty(), ))); } - if let Some(limit) = limit { + if let Some(limit) = definition.limit { let limit = i64::try_from(limit).map_err(|_| Error::InvalidInput { message: format!("view limit {limit} exceeds the maximum of {}", i64::MAX), })?; scanner.limit(Some(limit), None)?; } - - let out_schema = schema.clone(); + let output_schema = schema.clone(); let mapped = scanner.try_into_stream().await?.map(move |batch| { - let batch = batch.map_err(|e| DataFusionError::External(Box::new(e)))?; + let batch = batch.map_err(|error| DataFusionError::External(Box::new(error)))?; let batch = match &expanded { - None => batch, Some(expanded) => expanded.apply(&batch)?, + None => batch, }; - let batch = to_view_batch(&batch, &out_schema)?; + let batch = to_view_batch(&batch, &output_schema)?; rows_written.fetch_add(batch.num_rows() as u64, Ordering::Relaxed); Ok(batch) }); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, mapped))) } -/// `batch` in the view's physical schema: the row id becomes the -/// provenance column and computed columns are written NULL for their owner -/// to fill. pub(super) fn to_view_batch( batch: &RecordBatch, - out_schema: &SchemaRef, + output_schema: &SchemaRef, ) -> datafusion::common::Result { - let mut columns = Vec::with_capacity(out_schema.fields().len()); - for field in out_schema.fields() { + let mut columns = Vec::with_capacity(output_schema.fields().len()); + for field in output_schema.fields() { if computed_column_from_field(field).is_some() { columns.push(new_null_array(field.data_type(), batch.num_rows())); continue; } - let name = if field.name() == SOURCE_ROW_ID_COLUMN { - ROW_ID - } else { - field.name() - }; - let column = batch.column_by_name(name).ok_or_else(|| { + let column = batch.column_by_name(field.name()).ok_or_else(|| { DataFusionError::Internal(format!( "view column '{}' is not produced by the view's definition", field.name() @@ -1447,13 +395,9 @@ pub(super) fn to_view_batch( })?; columns.push(column.clone()); } - Ok(RecordBatch::try_new(out_schema.clone(), columns)?) + Ok(RecordBatch::try_new(output_schema.clone(), columns)?) } -/// The post-scan half of an unnested view's refresh: the scan reads -/// `raw_inputs` plus the row id, and each batch is unnested on the list -/// column, filtered, then projected by expressions typed against -/// `read_schema`, where the element sits under the alias. struct UnnestPlan { column: String, raw_inputs: Vec, @@ -1469,7 +413,6 @@ impl UnnestPlan { inputs: &[String], unnest: &super::ViewUnnest, ) -> Result { - // Whole root columns: a nested input is projected by the expression. let mut raw_inputs: Vec = inputs .iter() .map(|input| super::root(input).to_string()) @@ -1478,9 +421,7 @@ impl UnnestPlan { raw_inputs.sort(); raw_inputs.dedup(); let flattened = super::flattened_schema(&ArrowSchema::from(source.schema()), unnest)?; - // Physical expressions index columns by position, so the schema is - // exactly the scan's output: `raw_inputs` in order, then the row id. - let mut read_fields = Vec::with_capacity(raw_inputs.len() + 1); + let mut read_fields = Vec::with_capacity(raw_inputs.len()); for name in &raw_inputs { let name = if *name == unnest.column { &unnest.alias @@ -1494,25 +435,27 @@ impl UnnestPlan { })?; read_fields.push(field.clone()); } - read_fields.push(ArrowField::new(ROW_ID, DataType::UInt64, false)); let read_schema = Arc::new(ArrowSchema::new(read_fields)); let planner = Planner::new(read_schema.clone()); let physical = |what: &str, sql: &str| -> Result> { - let err = |e: lance::Error| Error::Runtime { - message: format!("{what}: {e}"), + let error = |error: lance::Error| Error::Runtime { + message: format!("{what}: {error}"), }; - let parsed = planner.parse_expr(sql).map_err(err)?; - let optimized = planner.optimize_expr(parsed).map_err(err)?; - planner.create_physical_expr(&optimized).map_err(err) + let parsed = planner.parse_expr(sql).map_err(error)?; + let optimized = planner.optimize_expr(parsed).map_err(error)?; + planner.create_physical_expr(&optimized).map_err(error) }; - let mut projections = Vec::with_capacity(definition.projections.len()); - for projection in &definition.projections { - let expr = physical( - &format!("view column '{}'", projection.output), - &projection.expression, - )?; - projections.push((projection.output.clone(), expr)); - } + let projections = definition + .projections + .iter() + .map(|projection| { + physical( + &format!("view column '{}'", projection.output), + &projection.expression, + ) + .map(|expression| (projection.output.clone(), expression)) + }) + .collect::>()?; let filter = definition .filter .as_deref() @@ -1529,8 +472,7 @@ impl UnnestPlan { fn apply(&self, batch: &RecordBatch) -> datafusion::common::Result { let unnested = unnest_batch(batch, &self.column) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - // Same columns, renamed: the list column is now the element under the alias. + .map_err(|error| DataFusionError::External(Box::new(error)))?; let unnested = RecordBatch::try_new(self.read_schema.clone(), unnested.columns().to_vec())?; let unnested = match &self.filter { None => unnested, @@ -1544,25 +486,20 @@ impl UnnestPlan { arrow_select::filter::filter_record_batch(&unnested, keep)? } }; - let mut columns = Vec::with_capacity(self.projections.len() + 1); - for (output, expr) in &self.projections { - let value = expr.evaluate(&unnested)?.into_array(unnested.num_rows())?; - columns.push((output.clone(), value)); - } - let row_id = unnested - .column_by_name(ROW_ID) - .expect("scan carries the row id") - .clone(); - columns.push((ROW_ID.to_string(), row_id)); + let columns = self + .projections + .iter() + .map(|(output, expression)| { + expression + .evaluate(&unnested)? + .into_array(unnested.num_rows()) + .map(|value| (output.clone(), value)) + }) + .collect::>>()?; Ok(RecordBatch::try_from_iter(columns)?) } } -/// Expand `list_column` one row per element, repeating every other column -/// for each element; an empty or null list contributes no rows. The list -/// column is replaced by its element type, so a projection reads the -/// element's fields as `alias.field` after this. This is the row-cardinality -/// step of an `expanded_select` view, applied per batch on the scan stream. fn unnest_batch(batch: &RecordBatch, list_column: &str) -> Result { use arrow_array::{Array, ListArray, UInt32Array}; use arrow_select::take::take; @@ -1583,30 +520,24 @@ fn unnest_batch(batch: &RecordBatch, list_column: &str) -> Result { batch.column(list_index).data_type() ), })?; - - // One take index per element, naming the source row it came from. A null - // list has no elements; its offsets are equal, so it repeats nothing. let offsets = list.value_offsets(); let mut repeat = Vec::with_capacity(list.values().len()); for row in 0..list.len() { if list.is_valid(row) { - let count = (offsets[row + 1] - offsets[row]) as usize; - repeat.extend(std::iter::repeat_n(row as u32, count)); + repeat.extend(std::iter::repeat_n( + row as u32, + (offsets[row + 1] - offsets[row]) as usize, + )); } } let repeat = UInt32Array::from(repeat); - // The flattened elements, in the same order as `repeat`: only the ranges - // valid rows cover, so a null row's stale range (if any) is skipped. - let elements = { - let mut ranges = Vec::new(); - for row in 0..list.len() { - if list.is_valid(row) { - ranges.extend((offsets[row] as u32)..(offsets[row + 1] as u32)); - } + let mut ranges = Vec::new(); + for row in 0..list.len() { + if list.is_valid(row) { + ranges.extend((offsets[row] as u32)..(offsets[row + 1] as u32)); } - take(list.values().as_ref(), &UInt32Array::from(ranges), None)? - }; - + } + let elements = take(list.values().as_ref(), &UInt32Array::from(ranges), None)?; let mut fields = Vec::with_capacity(batch.num_columns()); let mut columns = Vec::with_capacity(batch.num_columns()); for (index, field) in batch.schema().fields().iter().enumerate() { @@ -1636,3168 +567,38 @@ fn unnest_batch(batch: &RecordBatch, list_column: &str) -> Result { Ok(RecordBatch::try_new(schema, columns)?) } -/// Commit the view's removals and additions as one change, on the exact -/// generation the refresh planned from. Lance rejects an overlapping -/// provenance key, but an unrelated write to the view is not a key conflict, -/// so the generation is checked here too. -async fn publish( - view_ds: &Dataset, - eviction: Option<(Vec, Vec)>, - new_fragments: Vec, - keys: Option, - expected_incarnation: Option<&str>, -) -> Result { - let planned = view_ds.version().version; - #[cfg(test)] - tests::hold_before_publish(view_ds.uri()).await; - #[cfg(test)] - tests::hold_until_peers_planned(); - ensure_incarnation(view_ds, expected_incarnation, view_ds.uri()).await?; - let (updated_fragments, removed_fragment_ids) = eviction.unwrap_or_default(); - let committed = CommitBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .execute(Transaction::new( - planned, - Operation::Update { - removed_fragment_ids, - updated_fragments, - new_fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: keys, - updated_fragment_offsets: None, - }, - None, - )) - .await?; - if committed.version().version != planned + 1 { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced this refresh (view version {}); - the refresh is unrecorded and the next one will rebuild", - committed.version().version - ), - }); - } - Ok(committed) -} - -/// A delta batch's row id column, borrowed: a delta batch is as large as one -/// fragment's deletions, so copying it out would be the unbounded step the -/// chunking exists to avoid. -fn row_ids_of(batch: &RecordBatch) -> Result<&UInt64Array> { - let column = batch.column_by_name(ROW_ID).ok_or_else(|| Error::Runtime { - message: format!("'{ROW_ID}' is missing from a delta batch"), - })?; - column - .as_primitive_opt::() - .ok_or_else(|| Error::Runtime { - message: "row ids are not UInt64".into(), - }) -} - -/// A provenance id no source row can hold, carried in every refresh's -/// inserted-rows filter: any two refreshes of one view overlap on it, so -/// the loser of a race conflicts at commit whatever rows each planned. -const REFRESH_TOKEN_ID: u64 = u64::MAX; - -/// A builder with no collected keys, for publishes that only remove rows. -fn empty_keys(view_ds: &Dataset) -> Result>> { - Ok(Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new( - vec![source_row_id_field_id(view_ds)?], - )))) -} - -/// An inserted-rows filter holding the refresh token plus whatever the -/// tee collected. -fn refresh_filter(keys: &Arc>) -> Result { - let mut keys = keys.lock().map_err(|_| Error::Runtime { - message: "the provenance key filter was poisoned mid-refresh".into(), - })?; - keys.insert(KeyValue::UInt64(REFRESH_TOKEN_ID)) - .map_err(|e| Error::Runtime { - message: format!("failed to mark the refresh's filter: {e}"), - })?; - Ok(keys.build()) -} - -/// Reconciled ids past this fall back to the streamed rebuild, whose memory -/// does not grow with the delta. -fn eviction_rebuild_cap() -> usize { - #[cfg(test)] - if let Some(cap) = tests::eviction_cap_override() { - return cap; - } - 4 * 1024 * 1024 -} - -/// Provenance ids per staged eviction: the delete predicate carries one -/// literal per id, so a whole delta at once is unbounded. Each chunk costs a -/// pass over the view's provenance column, which is why it is not smaller. -const EVICTION_CHUNK: usize = 64 * 1024; - -/// Accumulates the view's removals from bounded chunks of provenance ids into -/// the one set of fragment changes the refresh publishes. -struct Eviction { - /// The view as the chunks staged so far leave it. A chunk's delete has to - /// see the deletion vectors the earlier ones wrote, or it stages a - /// fragment that drops them. - snapshot: Dataset, - chunk: usize, - updated: HashMap, - removed: Vec, - pending: Vec, - staged: bool, - /// Largest the buffer ever got, which is the bound under test. - #[cfg(test)] - peak: usize, -} - -impl Eviction { - fn new(view_ds: &Dataset, chunk: usize) -> Self { - Self { - snapshot: view_ds.clone(), - chunk, - updated: HashMap::new(), - removed: Vec::new(), - pending: Vec::with_capacity(chunk), - staged: false, - #[cfg(test)] - peak: 0, - } - } - - /// Fill a chunk at a time. Taking the batch whole and splitting it would - /// hold a fragment's worth of ids and recopy the tail per chunk. - async fn push(&mut self, ids: &UInt64Array) -> Result<()> { - for id in ids.values() { - self.pending.push(*id); - #[cfg(test)] - { - self.peak = self.peak.max(self.pending.len()); - } - if self.pending.len() == self.chunk { - self.flush().await?; - } - } - Ok(()) - } - - /// Stage what is buffered, keeping the buffer's allocation. - async fn flush(&mut self) -> Result<()> { - let mut chunk = std::mem::take(&mut self.pending); - let staged = self.stage(&chunk).await; - chunk.clear(); - self.pending = chunk; - staged - } - - /// The staged fragment changes, or `None` where nothing was evicted. - async fn finish(mut self) -> Result, Vec)>> { - if !self.pending.is_empty() { - self.flush().await?; - } - if !self.staged { - return Ok(None); - } - let mut updated: Vec = self.updated.into_values().collect(); - updated.sort_unstable_by_key(|f| f.id); - Ok(Some((updated, self.removed))) - } - - async fn stage(&mut self, ids: &[u64]) -> Result<()> { - let (updated, removed) = stage_eviction(&self.snapshot, ids).await?; - self.snapshot = advance(&self.snapshot, &updated); - for fragment in updated { - self.updated.insert(fragment.id, fragment); - } - self.removed.extend(removed); - self.staged = true; - Ok(()) - } -} - -/// The view as staged removals leave it, without committing them. Fragments -/// are replaced in place so the fragment bitmap stays in step; an emptied one -/// is left alone, since the delta names each provenance id only once. -fn advance(view_ds: &Dataset, updated: &[Fragment]) -> Dataset { - if updated.is_empty() { - return view_ds.clone(); - } - let by_id: HashMap = updated.iter().map(|f| (f.id, f)).collect(); - let fragments = view_ds - .manifest - .fragments - .iter() - .map(|f| by_id.get(&f.id).map_or_else(|| f.clone(), |u| (*u).clone())) - .collect(); - let mut manifest = view_ds.manifest.as_ref().clone(); - manifest.fragments = Arc::new(fragments); - let mut snapshot = view_ds.clone(); - snapshot.manifest = Arc::new(manifest); - snapshot -} - -/// The fragment changes that remove the view's rows for `ids`, staged rather -/// than committed so they can ride in the refresh's single data commit. -async fn stage_eviction(view_ds: &Dataset, ids: &[u64]) -> Result<(Vec, Vec)> { - // An expression rather than SQL text: the id list is a value here, not a - // predicate string that grows with the delta and has to be parsed. - let predicate = col(SOURCE_ROW_ID_COLUMN).in_list( - ids.iter() - .map(|id| lit(ScalarValue::UInt64(Some(*id)))) - .collect(), - false, - ); - let staged = DeleteBuilder::from_expr(Arc::new(view_ds.clone()), predicate) - .execute_uncommitted() - .await?; - let Operation::Delete { - updated_fragments, - deleted_fragment_ids, - .. - } = staged.transaction.operation - else { - return Err(Error::Runtime { - message: "expected a delete when staging the view's evictions".into(), - }); - }; - Ok((updated_fragments, deleted_fragment_ids)) -} - -fn source_row_id_field_id(view_ds: &Dataset) -> Result { - view_ds - .schema() - .field(SOURCE_ROW_ID_COLUMN) - .map(|f| f.id) - .ok_or_else(|| Error::Runtime { - message: format!("the view has no '{SOURCE_ROW_ID_COLUMN}' column"), - }) -} - -/// Tee the provenance ids of everything written into `keys`. -fn collect_source_row_ids( - stream: SendableRecordBatchStream, - keys: Arc>, - written: Arc, -) -> SendableRecordBatchStream { - let schema = stream.schema(); - let mapped = stream.map(move |batch| { - let batch = batch?; - let column = batch - .column_by_name(SOURCE_ROW_ID_COLUMN) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "'{SOURCE_ROW_ID_COLUMN}' is missing from the rows being written" - )) - })? - .as_primitive_opt::() - .ok_or_else(|| { - DataFusionError::Internal(format!("'{SOURCE_ROW_ID_COLUMN}' is not a uint64")) - })?; - let mut keys = keys - .lock() - .map_err(|_| DataFusionError::Internal("provenance key filter poisoned".into()))?; - for id in column.values() { - keys.insert(KeyValue::UInt64(*id)) - .map_err(|e| DataFusionError::Internal(e.to_string()))?; - } - written.fetch_add(batch.num_rows() as u64, Ordering::Relaxed); - Ok(batch) - }); - Box::pin(RecordBatchStreamAdapter::new(schema, mapped)) -} - #[cfg(test)] -pub(crate) mod tests { - - /// Park a refresh between planning and publication so a test can move - /// the view underneath it. Inert unless [`DRIFT_TARGET`] names this view. - pub(in crate::materialized_view) async fn hold_before_publish(uri: &str) { - { - let mut target = DRIFT_TARGET.lock().unwrap(); - if target.as_deref() != Some(uri) { - return; - } - // Take it: memory:// uris are relative and repeat across tests, so - // leaving it armed would park an unrelated refresh forever. - *target = None; - } - DRIFT_PLANNED.notify_one(); - DRIFT_RELEASED.notified().await; - } - - /// The rendezvous below is one global pair, so the cases that use it run - /// one at a time rather than trading each other's signals. - pub(super) static EVICTION_CAP: StdMutex> = StdMutex::new(None); - pub(super) fn eviction_cap_override() -> Option { - *EVICTION_CAP.lock().unwrap() - } - - pub(in crate::materialized_view) static DRIFT_LOCK: tokio::sync::Mutex<()> = - tokio::sync::Mutex::const_new(()); - pub(in crate::materialized_view) static DRIFT_TARGET: StdMutex> = - StdMutex::new(None); - pub(in crate::materialized_view) static DRIFT_PLANNED: tokio::sync::Notify = - tokio::sync::Notify::const_new(); - pub(in crate::materialized_view) static DRIFT_RELEASED: tokio::sync::Notify = - tokio::sync::Notify::const_new(); - - /// Block until every participant in a cross-process race has planned and - /// staged its write, so the commits they then attempt genuinely contend - /// rather than depending on the scheduler to overlap them. Inert unless - /// `MV_RACE_SYNC` names a directory shared by the participants. - pub(super) fn hold_until_peers_planned() { - let (Ok(dir), Ok(tag), Ok(peers)) = ( - std::env::var("MV_RACE_SYNC"), - std::env::var("MV_RACE_TAG"), - std::env::var("MV_RACE_PEERS"), - ) else { - return; - }; - let dir = std::path::PathBuf::from(dir); - let peers: usize = peers.parse().unwrap(); - std::fs::write(dir.join(format!("planned-{tag}")), b"1").unwrap(); - let deadline = std::time::Instant::now() + std::time::Duration::from_secs(120); - while planned_count(&dir) < peers { - assert!( - std::time::Instant::now() < deadline, - "peers never reached the commit boundary" - ); - std::thread::sleep(std::time::Duration::from_millis(2)); - } - } - - fn planned_count(dir: &std::path::Path) -> usize { - std::fs::read_dir(dir) - .map(|entries| { - entries - .filter_map(|e| e.ok()) - .filter(|e| e.file_name().to_string_lossy().starts_with("planned-")) - .count() - }) - .unwrap_or(0) - } - use arrow_array::{Int32Array, record_batch}; - use futures::TryStreamExt; - use lance::dataset::NewColumnTransform; - use lance_file::version::LanceFileVersion; +mod tests { + use arrow_array::record_batch; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use super::*; use crate::connect; - use crate::connection::Connection; - use crate::index::Index; - use crate::index::scalar::BTreeIndexBuilder; - use crate::materialized_view::MaterializedView; - use crate::query::{ExecutableQuery, QueryBase, Select}; - use crate::table::{CompactionOptions, OptimizeAction}; - - async fn db_with_source(values: Vec) -> (Connection, Table) { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("x", Int32, values)).unwrap(); - let table = conn - .create_table("src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - (conn, table) - } - - async fn doubled_view(conn: &Connection) -> MaterializedView { - conn.create_materialized_view("doubled", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap() - } - - /// A refreshed doubled view over a source holding `values`. - async fn refreshed_doubled(values: Vec) -> (Connection, Table, MaterializedView) { - let (conn, source) = db_with_source(values).await; - let view = doubled_view(&conn).await; - view.refresh().execute().await.unwrap(); - (conn, source, view) - } - - /// The per-batch expansion behind an `expanded_select` view: every - /// element of the list column becomes a row, the other columns repeat - /// for each, and an empty or null list contributes no rows at all -- - /// which is exactly "zero rows out" for a table-valued function. - /// Four documents with a `list>` column named - /// `c`: doc 1 has two chunks, doc 2 none, doc 3 a null list, doc 4 one. - fn chunked_batch() -> RecordBatch { - use arrow_array::builder::{Int32Builder, ListBuilder, StringBuilder, StructBuilder}; - use arrow_array::{ArrayRef, Int64Array}; - use arrow_schema::{DataType, Field, Fields}; - - let element_fields = Fields::from(vec![ - Field::new("chunk", DataType::Utf8, true), - Field::new("ordinal", DataType::Int32, true), - ]); - let mut list = ListBuilder::new(StructBuilder::new( - element_fields, - vec![ - Box::new(StringBuilder::new()), - Box::new(Int32Builder::new()), - ], - )); - for chunks in [Some(vec!["a", "b"]), Some(vec![]), None, Some(vec!["c"])] { - match chunks { - Some(chunks) => { - for (i, c) in chunks.iter().enumerate() { - let s = list.values(); - s.field_builder::(0).unwrap().append_value(c); - s.field_builder::(1) - .unwrap() - .append_value(i as i32); - s.append(true); - } - list.append(true); - } - None => list.append(false), - } - } - let meta = arrow_array::StructArray::from(vec![( - Arc::new(Field::new("title", DataType::Utf8, true)), - Arc::new(arrow_array::StringArray::from(vec!["t1", "t2", "t3", "t4"])) as ArrayRef, - )]); - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef, - ), - ("meta", Arc::new(meta) as ArrayRef), - ("c", Arc::new(list.finish()) as ArrayRef), - ]) - .unwrap() - } - - #[test] - fn unnest_repeats_siblings_per_element_and_drops_empty_lists() { - use arrow_array::{StringArray, StructArray}; - use arrow_schema::{DataType, Field, Fields}; - - let batch = chunked_batch(); - let element_fields = Fields::from(vec![ - Field::new("chunk", DataType::Utf8, true), - Field::new("ordinal", DataType::Int32, true), - ]); - let out = unnest_batch(&batch, "c").unwrap(); - assert_eq!(out.num_rows(), 3, "{out:?}"); - let ids: Vec = out["id"] - .as_primitive::() - .values() - .to_vec(); - assert_eq!(ids, [1, 1, 4]); - let element = out["c"].as_any().downcast_ref::().unwrap(); - let chunks: Vec<&str> = element - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect(); - assert_eq!(chunks, ["a", "b", "c"]); - let ordinals: Vec = element - .column(1) - .as_primitive::() - .values() - .to_vec(); - assert_eq!(ordinals, [0, 1, 0]); - // the element column is now the struct itself, not a list of it - assert_eq!( - out.schema().field_with_name("c").unwrap().data_type(), - &DataType::Struct(element_fields) - ); - } - - /// A Function in FROM position is refreshed from its staging table: the - /// query on the view stays the one the user wrote, the staging binding - /// says which table and list column refresh reads, and rows come out one - /// per element as with UNNEST. Without a staging, a local database - /// refuses the query rather than guess. - #[tokio::test] - async fn a_function_in_from_position_refreshes_from_its_staging() { - use arrow_array::StringArray; - - let conn = connect("memory://").execute().await.unwrap(); - let staging = conn - .create_table("docs__chunk", chunked_batch()) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let query = - "SELECT id AS doc, e.chunk AS text, e.ordinal FROM docs, chunk(meta.title, 2) AS e"; - let definition = MaterializedViewDefinition::from_sql(query).unwrap(); - - let err = crate::materialized_view::prepare_definition(&staging, definition.clone()) - .await - .unwrap_err(); - assert!( - matches!(&err, Error::InvalidInput { message } if message.contains("reads 'docs'")), - "{err:?}" - ); - - let view = crate::materialized_view::prepare_staged_definition(&staging, definition, "c") - .await - .unwrap() - .create("chunks") - .await - .unwrap(); - assert_eq!(view.definition().to_sql(), query); - let metadata = view.table().schema().await.unwrap().metadata().clone(); - assert_eq!( - crate::materialized_view::read_staging(&metadata).unwrap(), - Some(crate::materialized_view::StagingBinding { - table: "docs__chunk".into(), - namespace: Vec::new(), - column: "c".into(), - }) - ); - - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "ordinal").await, [0, 0, 1]); - let batches = view - .table() - .query() - .select(Select::columns(&["text"])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let out = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap(); - let texts: Vec<&str> = out["text"] - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect(); - assert_eq!(texts, ["a", "b", "c"]); - - // The reopened view reads the same logical query and refreshes - // incrementally from the staging. - staging.add(chunked_batch()).execute().await.unwrap(); - let reopened = conn.open_materialized_view("chunks").await.unwrap(); - assert_eq!(reopened.definition().to_sql(), query); - let result = reopened.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(reopened.table(), "ordinal").await, [0, 0, 0, 0, 1, 1]); - } - - /// An expanded view materializes one row per list element, with the - /// projections reading the element through the alias and the other - /// source columns repeated alongside; sources with no elements yield - /// no rows. The lineage is the list column, so a change to it is what - /// drives incremental refresh. - #[tokio::test] - async fn an_expanded_view_materializes_one_row_per_element() { - use arrow_array::{Int64Array, StringArray}; - - let conn = connect("memory://").execute().await.unwrap(); - let source = conn - .create_table("docs", chunked_batch()) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let mut view = crate::materialized_view::prepare_definition( - &source, - MaterializedViewDefinition::from_sql( - "SELECT id AS doc, meta.title AS title, e.ordinal + 1 AS nth \ - FROM docs, UNNEST(c) AS e WHERE e.ordinal < 5", - ) - .unwrap(), - ) - .await - .unwrap(); - // A computed column's input read through the alias is an element - // field; the recorded source input is the list column. - let text = view.input_column("e.chunk").unwrap(); - assert!(view.definition.lateral.is_some()); - let view = view.create("chunks").await.unwrap(); - view.refresh().execute().await.unwrap(); - - let batches = view - .table() - .query() - .select(Select::columns(&["doc", &text, "title"])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let out = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap(); - let docs: Vec = out["doc"] - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec(); - let strings = |column: &str| -> Vec { - out[column] - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .map(str::to_string) - .collect() - }; - assert_eq!(docs, [1, 1, 4]); - assert_eq!(strings(&text), ["a", "b", "c"]); - assert_eq!(strings("title"), ["t1", "t1", "t4"]); - assert_eq!(read(view.table(), "nth").await, [1, 1, 2]); - - // Appended documents expand incrementally; the existing rows stay. - source.add(chunked_batch()).execute().await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "nth").await, [1, 1, 1, 1, 2, 2]); - } - - async fn read(table: &Table, column: &str) -> Vec { - let batches = table - .query() - .select(Select::columns(&[column])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let mut values: Vec = batches - .iter() - .flat_map(|batch| { - batch[column] - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect::>() - }) - .collect(); - values.sort(); - values - } - - async fn append(table: &Table, values: Vec) { - let batch = record_batch!(("x", Int32, values)).unwrap(); - table.add(batch).execute().await.unwrap(); - } #[tokio::test] - async fn test_first_refresh_materializes_the_view() { - let (conn, _) = db_with_source(vec![1, 2, 3]).await; - let view = doubled_view(&conn).await; - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 3); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - - // The watermark survives on the stored schema, not just the handle. - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - let again = reopened.refresh().execute().await.unwrap(); - assert_eq!(again.mode, RefreshMode::NoOp); - assert_eq!(again.rows_written, 0); - } - - #[tokio::test] - async fn test_filter_selects_the_source_rows() { - let (conn, _) = db_with_source(vec![1, 20, 3, 40]).await; - let view = conn - .create_materialized_view("big", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 10") + async fn raced_first_full_refresh_lands_one_result() { + let directory = tempfile::tempdir().unwrap(); + let connection = connect(directory.path().to_str().unwrap()) .execute() .await .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "x").await, vec![20, 40]); - } - - #[tokio::test] - async fn test_mixed_case_filter_is_canonicalized_for_lineage_and_refresh() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!( - ("id", Int32, [1, 2, 3]), - ("PartyAbbrev", Utf8, ["D", "R", "D"]) - ) - .unwrap(); - conn.create_table("src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) + let source_batch = record_batch!(("x", Int32, [1, 2, 3])).unwrap(); + connection + .create_table("src", source_batch.clone()) .execute() .await .unwrap(); - conn.create_materialized_view("democrats", "src") - .with_no_data(true) - .select([("id", "id")]) - .only_if(r#""PartyAbbrev" = 'D'"#) - .execute() - .await - .unwrap(); - - // Reopen from schema metadata so these assertions cover the stored - // predicate and lineage, not only the declaration-time handle. - let view = conn.open_materialized_view("democrats").await.unwrap(); - assert_eq!( - view.definition().filter.as_deref(), - Some("`PartyAbbrev` = 'D'") - ); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "id").await, vec![1, 3]); - } - - /// A legacy layout whose query means what the canonical one means is - /// rewritten in the current layout on the next refresh, without a - /// rebuild: the rows it certified are the rows the query produces. - #[tokio::test] - async fn a_legacy_layout_with_the_same_meaning_is_rewritten_without_a_rebuild() { - let (conn, source, view) = refreshed_doubled(vec![1]).await; - let legacy = serde_json::json!({ - "kind": "select", - "source_table": "src", - "projections": [ - {"output": "x", "expression": "`x`"}, - {"output": "twice", "expression": "x*2"}, - ], - "inputs": ["x"], - }) - .to_string(); - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - let predicted = dataset.version().version + 1; - dataset - .update_schema_metadata([ - (DEFINITION_META_KEY.to_string(), Some(legacy)), - ( - VIEW_VERSION_META_KEY.to_string(), - Some(predicted.to_string()), - ), - ]) - .await - .unwrap(); - native.dataset.update(dataset); - - append(&source, vec![2]).await; - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - let result = reopened.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(reopened.table(), "twice").await, vec![2, 4]); - let stored: serde_json::Value = serde_json::from_str( - &reopened.table().schema().await.unwrap().metadata()[DEFINITION_META_KEY], - ) - .unwrap(); - assert_eq!( - stored, - serde_json::json!({ - "kind": "query", - "format": 1, - "query": "SELECT x, x * 2 AS twice FROM src", - }) - ); - } - - #[tokio::test] - async fn test_legacy_raw_filter_rebuilds_and_persists_canonical_definition() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!( - ("id", Int32, [1, 2, 3]), - ("PartyAbbrev", Utf8, ["D", "R", "D"]) - ) - .unwrap(); - conn.create_table("legacy_src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("legacy_view", "legacy_src") - .with_no_data(true) - .select([("id", "id")]) - .only_if(r#""PartyAbbrev" = 'X'"#) - .execute() - .await - .unwrap(); - assert_eq!(view.refresh().execute().await.unwrap().rows_written, 0); - - // Model a definition and up-to-date watermark written before filter - // canonicalization was applied to materialized views. - let legacy = serde_json::json!({ - "kind": "select", - "source_table": "legacy_src", - "projections": [{"output": "id", "expression": "id"}], - "filter": r#""PartyAbbrev" = 'D'"#, - "inputs": ["id"], - }) - .to_string(); - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - let predicted = dataset.version().version + 1; - dataset - .update_schema_metadata([ - (DEFINITION_META_KEY.to_string(), Some(legacy)), - ( - VIEW_VERSION_META_KEY.to_string(), - Some(predicted.to_string()), - ), - ]) - .await - .unwrap(); - native.dataset.update(dataset); - - let reopened = conn.open_materialized_view("legacy_view").await.unwrap(); - let result = reopened.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 2); - assert_eq!(read(reopened.table(), "id").await, vec![1, 3]); - - // A fresh handle proves the migration was stored alongside the new - // watermark and therefore happens only once. - let migrated = conn.open_materialized_view("legacy_view").await.unwrap(); - assert_eq!( - migrated.definition().filter.as_deref(), - Some("`PartyAbbrev` = 'D'") - ); - assert_eq!( - migrated.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - assert_eq!(read(migrated.table(), "id").await, vec![1, 3]); - } - - #[tokio::test] - async fn test_append_refreshes_incrementally() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - - append(&source, vec![5]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 10]); - } - - #[tokio::test] - async fn test_incremental_applies_the_filter() { - let (conn, source) = db_with_source(vec![1, 20]).await; - let view = conn - .create_materialized_view("big", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 10") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - append(&source, vec![3, 30]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "x").await, vec![20, 30]); - } - - /// The watermark has to advance even when no appended row matches, or the - /// same fragments would be rescanned by every later refresh. - #[tokio::test] - async fn test_incremental_with_nothing_matching_advances_the_watermark() { - let (conn, source) = db_with_source(vec![20]).await; - let view = conn - .create_materialized_view("big", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 10") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - append(&source, vec![1, 2]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - - let again = view.refresh().execute().await.unwrap(); - assert_eq!(again.mode, RefreshMode::NoOp); - } - - /// Unlike a computed column, a view reflects source mutation: an update - /// rebuilds rather than going stale. - #[tokio::test] - async fn test_update_replaces_the_rows_it_changed() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - // An update changes rows in place, so the view replaces exactly - // those rows and leaves the rest of what it holds alone. - source - .update() - .column("x", "20") - .only_if("x = 2") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1, "only the changed row is recomputed"); - assert_eq!(read(view.table(), "twice").await, vec![2, 6, 40]); - } - - /// Legacy storage cannot serve the row-version columns update discovery - /// scans; appends stay incremental, updates rebuild rather than fail. - #[tokio::test] - async fn test_a_legacy_storage_source_is_reconciled() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("x", Int32, [1, 2, 3])).unwrap(); - let source = conn - .create_table("legacy_src", batch) - .write_options(crate::table::WriteOptions { - lance_write_params: Some(lance::dataset::WriteParams { - enable_stable_row_ids: true, - data_storage_version: Some(LanceFileVersion::Legacy), - ..Default::default() - }), - }) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("legacy_doubled", "legacy_src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - source - .add(record_batch!(("x", Int32, [4])).unwrap()) - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6, 8]); - - source - .update() - .column("x", "20") - .only_if("x = 2") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 6, 8, 40]); - } - - #[tokio::test] - async fn test_delete_evicts_the_view_rows_it_removed() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - // A delete removes source rows without changing the ones that - // remain, so the view evicts exactly those rows and keeps the rest - // rather than recomputing every row it already held. - source.delete("x = 2").await.unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 6]); - - // A delete and an append in one span: both are applied. - source.delete("x = 1").await.unwrap(); - source - .add(record_batch!(("x", Int32, vec![4])).unwrap()) - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![6, 8]); - } - - /// merge_insert commits an `Update`, and its by-source arm removes rows - /// rather than changing them, so the classifier must treat that - /// transaction form as a source of deletions. - #[tokio::test] - async fn test_merge_insert_by_source_delete_evicts_the_view_rows() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - let batch = record_batch!(("x", Int32, vec![1, 3])).unwrap(); - let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); - let mut merge = source.merge_insert(&["x"]); - merge.when_not_matched_by_source_delete(None); - merge.execute(Box::new(reader)).await.unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 6]); - } - - /// A refresh may only certify the generation it planned from. A write - /// that lands between planning and publication is drift the refresh did - /// not account for, so it aborts rather than stamp it as materialized. - #[tokio::test(flavor = "multi_thread")] - async fn test_refresh_aborts_rather_than_certify_view_drift() { - let _serial = DRIFT_LOCK.lock().await; - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("drifting_view", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - let certified = source_version_of(view.table()).await; - - // The source loses a row, so the next refresh plans an eviction. - source.delete("x = 2").await.unwrap(); - - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let refreshing = tokio::spawn(async move { view.refresh().execute().await }); - - // Move the view once the refresh has planned against it. - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the refresh never reached the publication boundary"); - let drifted = conn.open_table("drifting_view").execute().await.unwrap(); - drifted.delete("twice = 6").await.unwrap(); - DRIFT_RELEASED.notify_one(); - - // Publishing removals and additions as one change makes the drift a - // conflict lance itself rejects; a pure append, which touches no - // existing fragment, still relies on the generation check. - let err = refreshing.await.unwrap().unwrap_err(); - let message = err.to_string(); - assert!( - message.contains("raced this refresh") || message.contains("preempted by concurrent"), - "got {err:?}" - ); - // The watermark still names the generation that was actually proven. - assert_eq!(source_version_of(&drifted).await, certified); - } - - async fn source_version_of(table: &Table) -> Option { - table - .schema() - .await - .unwrap() - .metadata() - .get(SOURCE_VERSION_META_KEY) - .cloned() - } - - /// A transaction file carries placeholder ids for the fragments it - /// creates. Into an empty source those collide with the ids the commit - /// assigns, so treating them as already-materialized drops the very - /// first rows while the watermark still advances past them. - #[tokio::test] - async fn test_merge_into_an_empty_source_is_materialized() { - let (_conn, source, view) = refreshed_doubled(vec![]).await; - assert_eq!(read(view.table(), "twice").await, Vec::::new()); - - let batch = record_batch!(("x", Int32, vec![1, 2])).unwrap(); - let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); - let mut merge = source.merge_insert(&["x"]); - merge - .when_matched_update_all(None) - .when_not_matched_insert_all(); - merge.execute(Box::new(reader)).await.unwrap(); - - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - // A second refresh must not double them either. - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// A refresh publishes what it removes and what it adds as one change, - /// so an update never exposes the view without the rows it is replacing. - #[tokio::test(flavor = "multi_thread")] - async fn test_an_update_is_never_visible_as_a_gap() { - let _serial = DRIFT_LOCK.lock().await; - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("atomic_view", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - source - .update() - .column("x", "x + 10") - .execute() - .await - .unwrap(); - - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let refreshing = tokio::spawn(async move { view.refresh().execute().await }); - - // Read the view while the refresh is staged but not yet published. - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the refresh never reached the publication boundary"); - let midway = conn.open_table("atomic_view").execute().await.unwrap(); - assert_eq!( - read(&midway, "twice").await, - vec![2, 4, 6], - "the pre-refresh rows must still be there in full" - ); - DRIFT_RELEASED.notify_one(); - - refreshing.await.unwrap().unwrap(); - let after = conn.open_table("atomic_view").execute().await.unwrap(); - assert_eq!(read(&after, "twice").await, vec![22, 24, 26]); - } - - async fn provenance_by_x(table: &Table) -> HashMap { - let batches = table - .query() - .select(Select::columns(&["x", SOURCE_ROW_ID_COLUMN])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - batches - .iter() - .flat_map(|batch| { - let xs = batch["x"].as_any().downcast_ref::().unwrap(); - let ids = batch[SOURCE_ROW_ID_COLUMN].as_primitive::(); - (0..batch.num_rows()) - .map(|i| (xs.value(i), ids.value(i))) - .collect::>() - }) - .collect() - } - - /// One delta batch is as large as a fragment's deletions, so it arrives - /// well past a chunk: it has to be drained a chunk at a time without ever - /// holding the batch, and the passes' deletion vectors have to accumulate - /// rather than replace each other. - #[tokio::test] - async fn test_a_chunked_eviction_removes_every_row_it_names() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2, 3, 4, 5, 6]).await; - - let native = view.table().as_native().unwrap(); - let view_ds = native.dataset.get().await.unwrap().as_ref().clone(); - let provenance = provenance_by_x(view.table()).await; - - let mut eviction = Eviction::new(&view_ds, 2); - let batch = UInt64Array::from_iter_values([1, 2, 3, 4].map(|x| provenance[&x])); - eviction.push(&batch).await.unwrap(); - assert_eq!( - eviction.peak, 2, - "a batch past the chunk was buffered whole" - ); - let staged = eviction.finish().await.unwrap(); - assert!(staged.is_some(), "four ids over a chunk of two stage twice"); - publish(&view_ds, staged, Vec::new(), None, None) - .await - .unwrap(); - native.dataset.reload().await.unwrap(); - - assert_eq!(read(view.table(), "x").await, vec![5, 6]); - } - - /// Two first rebuilds materialize the same source rows; the loser must - /// key-conflict at commit and land nothing, or the view doubles. - #[tokio::test] - async fn test_a_raced_first_rebuild_lands_nothing() { - let _guard = DRIFT_LOCK.lock().await; - let (conn, _) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("raced_rebuild", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - let native = view.table().as_native().unwrap(); - let view_ds = native.dataset.get().await.unwrap().as_ref().clone(); - let uri = view_ds.uri().to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let racing = tokio::spawn(async move { view.refresh().execute().await }); - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the rebuild never reached the publication boundary"); - - // The other process's first rebuild, reduced to its commit. Its - // source rows are DISJOINT from ours -- the sentinel alone must make - // the two rebuilds conflict. - let batch = record_batch!( - ("x", Int32, [8, 9]), - ("twice", Int32, [16, 18]), - ("__source_row_id", UInt64, [7u64, 8]) - ) - .unwrap(); - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let batch = RecordBatch::try_new( - schema.clone(), - schema - .fields() - .iter() - .map(|f| batch.column_by_name(f.name()).unwrap().clone()) - .collect(), - ) - .unwrap(); - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), - futures::stream::iter([Ok(batch)]), - )); - let keys = Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new(vec![ - source_row_id_field_id(&view_ds).unwrap(), - ]))); - let stream = collect_source_row_ids(stream, keys.clone(), Arc::new(AtomicU64::new(0))); - let write_txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await - .unwrap(); - let Operation::Append { fragments } = write_txn.operation else { - panic!("expected an append"); - }; - let filter = { - let mut keys = keys.lock().unwrap(); - keys.insert(KeyValue::UInt64(super::REFRESH_TOKEN_ID)) - .unwrap(); - keys.build() - }; - CommitBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .execute(Transaction::new( - view_ds.version().version, - Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments: Vec::new(), - new_fragments: fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: Some(filter), - updated_fragment_offsets: None, - }, - None, - )) - .await - .unwrap(); - DRIFT_RELEASED.notify_one(); - - let err = racing.await.unwrap().unwrap_err(); - // The loser lands nothing: only the winner's rows remain. - let raced = conn.open_table("raced_rebuild").execute().await.unwrap(); - assert_eq!( - read(&raced, "twice").await, - vec![16, 18], - "the losing rebuild must not union with the winner ({err})" - ); - } - - /// An incremental refresh racing any other refresh must land nothing, - /// even when their planned rows are disjoint: the shared token makes the - /// commits conflict. - #[tokio::test] - async fn test_a_raced_incremental_lands_nothing() { - let _guard = DRIFT_LOCK.lock().await; - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("raced_incremental", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - append(&source, vec![4]).await; - - let native = view.table().as_native().unwrap(); - native.dataset.reload().await.unwrap(); - let view_ds = native.dataset.get().await.unwrap().as_ref().clone(); - *DRIFT_TARGET.lock().unwrap() = Some(view_ds.uri().to_string()); - let racing = tokio::spawn(async move { view.refresh().execute().await }); - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the refresh never reached the publication boundary"); - - // The concurrent refresh, reduced to its commit: disjoint rows, the - // shared token alone must collide. - let batch = record_batch!( - ("x", Int32, [9]), - ("twice", Int32, [18]), - ("__source_row_id", UInt64, [8u64]) - ) - .unwrap(); - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let batch = RecordBatch::try_new( - schema.clone(), - schema - .fields() - .iter() - .map(|f| batch.column_by_name(f.name()).unwrap().clone()) - .collect(), - ) - .unwrap(); - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), - futures::stream::iter([Ok(batch)]), - )); - let keys = empty_keys(&view_ds).unwrap(); - let stream = collect_source_row_ids(stream, keys.clone(), Arc::new(AtomicU64::new(0))); - let write_txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await - .unwrap(); - let Operation::Append { fragments } = write_txn.operation else { - panic!("expected an append"); - }; - let filter = refresh_filter(&keys).unwrap(); - CommitBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .execute(Transaction::new( - view_ds.version().version, - Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments: Vec::new(), - new_fragments: fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: Some(filter), - updated_fragment_offsets: None, - }, - None, - )) - .await - .unwrap(); - DRIFT_RELEASED.notify_one(); - - let err = racing.await.unwrap().unwrap_err(); - let raced = conn - .open_table("raced_incremental") - .execute() - .await - .unwrap(); - assert_eq!( - read(&raced, "twice").await, - vec![2, 4, 6, 18], - "the losing increment must not land its rows ({err})" - ); - } - - /// Past the eviction cap the refresh falls back to the streamed rebuild, - /// and the result is identical either way. - #[tokio::test] - async fn test_oversized_delta_falls_back_to_rebuild() { - let (conn, source) = db_with_source((1..=20).collect()).await; - let view = doubled_view(&conn).await; - view.refresh().execute().await.unwrap(); - - source.delete("x <= 10").await.unwrap(); - *tests::EVICTION_CAP.lock().unwrap() = Some(4); - let result = view.refresh().execute().await; - *tests::EVICTION_CAP.lock().unwrap() = None; - let result = result.unwrap(); - assert_eq!( - result.mode, - RefreshMode::Rebuild, - "ten evictions, cap of four" - ); - assert_eq!(read(view.table(), "x").await, (11..=20).collect::>()); - - // Under the cap the same shape stays incremental. - source.delete("x = 11").await.unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "x").await, (12..=20).collect::>()); - } - - /// A cap of zero is a view that holds nothing, not a view without a cap. - #[tokio::test] - async fn test_zero_limit_holds_no_rows() { - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("empty", "src") - .with_no_data(true) - .select([("x", "x")]) - .limit(0) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 0); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - - // Still nothing after the source grows, and the watermark advances. - append(&source, vec![4]).await; - view.refresh().execute().await.unwrap(); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// An update rewrites a whole fragment. When it lands on one appended - /// since the watermark, that fragment also holds rows the update never - /// touched -- rows the recompute does not cover and the append set no - /// longer reaches. - #[tokio::test] - async fn test_update_touching_a_new_fragment_keeps_its_untouched_rows() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - - // One fragment, appended after the watermark, holding both a row the - // update will change and a row it will not. - append(&source, vec![3, 40]).await; - source - .update() - .column("x", "99") - .only_if("x = 3") - .execute() - .await - .unwrap(); - - view.refresh().execute().await.unwrap(); - assert_eq!( - read(view.table(), "twice").await, - vec![2, 4, 80, 198], - "a row appended into the updated fragment went missing" - ); - } - - async fn compact(source: &Table) { - source - .optimize(OptimizeAction::Compact { - options: CompactionOptions::default(), - remap_options: None, - }) - .await - .unwrap(); - } - - /// A compaction rearranges rows without changing them, so it costs the - /// view nothing: the watermark advances and no row is recomputed. - #[tokio::test] - async fn test_compaction_alone_refreshes_incrementally() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - append(&source, vec![3]).await; - view.refresh().execute().await.unwrap(); - - compact(&source).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// Rows appended after a compaction are separable through the transaction - /// log: only the appended fragments are computed. - #[tokio::test] - async fn test_append_after_compaction_stays_incremental() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - append(&source, vec![2]).await; - view.refresh().execute().await.unwrap(); - - compact(&source).await; - append(&source, vec![3]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - } - - /// An append swallowed by a later compaction cannot be told apart from - /// the rows the view already holds, so the refresh rebuilds -- once -- - /// rather than duplicate or drop. - #[tokio::test] - async fn test_append_folded_into_compaction_rebuilds() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - - append(&source, vec![2]).await; - compact(&source).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// The create-time gate holds across a drop-and-recreate of the source - /// under the same name. - #[tokio::test] - async fn test_refresh_refuses_a_recreated_source_without_stable_row_ids() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - - conn.drop_table("src", &[]).await.unwrap(); - let batch = record_batch!(("x", Int32, [9])).unwrap(); - conn.create_table("src", batch).execute().await.unwrap(); - - let err = view.refresh().execute().await.unwrap_err(); - assert!( - matches!(err, Error::InvalidInput { message } if message.contains("stable row ids")) - ); - } - - /// A change to a column the view does not read is not a reason to - /// rebuild: the exact signature check is scoped to the view's inputs. - #[tokio::test] - async fn test_unrelated_column_change_does_not_rebuild() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - - source - .add_columns() - .transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( - vec![arrow_schema::Field::new( - "unrelated", - arrow_schema::DataType::Int32, - true, - )], - )))) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - #[tokio::test] - async fn test_full_forces_a_rebuild() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - - append(&source, vec![2]).await; - let result = view.refresh().full(true).execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// A cap and incremental reconciliation do not compose: rows skipped at - /// the cap fall behind the watermark, so room freed later cannot be - /// refilled from any delta. A capped view rebuilds instead, which its - /// own cap keeps cheap. - #[tokio::test] - async fn test_limited_view_rebuilds_rather_than_reconcile() { - let (conn, source) = db_with_source(vec![1, 2]).await; - let view = conn - .create_materialized_view("capped", "src") - .with_no_data(true) - .select([("x", "x")]) - .limit(2) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "x").await, vec![1, 2]); - - append(&source, vec![3, 4]).await; - source - .update() - .column("x", "11") - .only_if("x = 1") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - let held = read(view.table(), "x").await; - assert_eq!(held.len(), 2, "the cap holds: {held:?}"); - let selectable = read(&source, "x").await; - assert!( - held.iter().all(|x| selectable.contains(x)), - "{held:?} is not a subset of {selectable:?}" - ); - } - - #[tokio::test] - async fn test_limit_caps_the_view() { - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("capped", "src") - .with_no_data(true) - .select([("x", "x")]) - .limit(4) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 3); - - // The cap counts already-held rows, so only one appended row lands. - append(&source, vec![4, 5, 6]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(view.table().count_rows(None).await.unwrap(), 4); - - // At the cap, later appends only move the watermark. - append(&source, vec![7]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 0); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// The whole point of the fragment-swap commit: a rebuild must never - /// leave the view without its index definitions. - #[tokio::test] - async fn test_rebuild_retains_indexes() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - view.table() - .create_index(&["twice"], Index::BTree(BTreeIndexBuilder::default())) - .execute() - .await - .unwrap(); - assert_eq!(view.table().list_indices().await.unwrap().len(), 1); - - source - .update() - .column("x", "x + 10") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(view.table().list_indices().await.unwrap().len(), 1); - assert_eq!(read(view.table(), "twice").await, vec![22, 24, 26]); - - // The swapped-in rows are reachable through an indexed query. - let batches = view - .table() - .query() - .only_if("twice = 24") - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); - } - - #[tokio::test] - async fn test_rebuild_of_an_empty_result_is_an_empty_view() { - let (conn, _) = db_with_source(vec![1, 2]).await; - let view = conn - .create_materialized_view("none", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 100") - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 0); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// Provenance: every view row records the source row that produced it. - #[tokio::test] - async fn test_source_row_ids_are_recorded() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2, 3]).await; - - let batches = view - .table() - .query() - .select(Select::columns(&[SOURCE_ROW_ID_COLUMN])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total, 3); - for batch in &batches { - assert_eq!(batch[SOURCE_ROW_ID_COLUMN].null_count(), 0); - } - } - - #[tokio::test] - async fn test_dropping_a_source_input_fails_the_refresh() { - let (conn, source) = db_with_source(vec![1]).await; - let view = conn - .create_materialized_view("v", "src") - .with_no_data(true) - .select([("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - append(&source, vec![2]).await; - source - .add_columns() - .transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( - vec![arrow_schema::Field::new( - "y", - arrow_schema::DataType::Int32, - true, - )], - )))) - .execute() - .await - .unwrap(); - source.drop_columns(&["x"]).await.unwrap(); - - let err = view.refresh().execute().await.unwrap_err(); - assert!( - matches!(&err, Error::Schema { message } - if message.contains("dropped or renamed") && message.contains("No field named x")), - "{err:?}" - ); - } - - /// A pinned refresh materializes the source as of `version`; catching up - /// to the appends beyond it stays incremental. - #[tokio::test] - async fn test_pinned_refresh_and_catch_up() { - let (conn, source) = db_with_source(vec![1]).await; - let view = doubled_view(&conn).await; - let pinned = source.version().await.unwrap(); - - append(&source, vec![2]).await; - let result = view - .refresh() - .source_version(pinned) - .execute() - .await - .unwrap(); - assert_eq!(result.source_version, pinned); - assert_eq!(read(view.table(), "twice").await, vec![2]); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// Views chain: a view's stable row ids and provenance column make it a - /// source like any other, and the default projection takes its declared - /// columns without copying its provenance. - #[tokio::test] - async fn test_a_view_can_source_another_view() { - let (conn, source) = db_with_source(vec![1, 2, 30]).await; - let first = doubled_view(&conn).await; - first.refresh().execute().await.unwrap(); - - let second = conn - .create_materialized_view("second", "doubled") - .with_no_data(true) - .only_if("twice > 10") - .execute() - .await - .unwrap(); - assert!( - second - .definition() - .projections - .iter() - .all(|p| p.output != SOURCE_ROW_ID_COLUMN) - ); - let result = second.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 1); - assert_eq!(read(second.table(), "twice").await, vec![60]); - - append(&source, vec![50]).await; - first.refresh().execute().await.unwrap(); - let result = second.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(second.table(), "twice").await, vec![60, 100]); - } - - /// The watermark speaks only for the state a refresh left behind: a - /// direct write to the view is drift, and the next refresh rebuilds - /// rather than preserving it as current. - #[tokio::test] - async fn test_direct_view_mutation_forces_a_rebuild() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - view.table().delete("x = 1").await.unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// A dropped and recreated source reuses version numbers but never their - /// timestamps; the watermark must not vouch for the replacement's rows. - #[tokio::test] - async fn test_source_recreation_forces_a_rebuild() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - - conn.drop_table("src", &[]).await.unwrap(); - let batch = record_batch!(("x", Int32, [7])).unwrap(); - conn.create_table("src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![14]); - } - - /// A refresh bound to an incarnation refuses a view dropped and recreated - /// since, even under the same name and definition; the recreated view's - /// own token is accepted, and the token survives a refresh's stamp. - #[tokio::test] - async fn test_refresh_refuses_a_recreated_view_incarnation() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - let token = view.incarnation().unwrap().to_string(); - view.refresh() - .expect_incarnation(&token) - .execute() - .await - .unwrap(); - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - assert_eq!(reopened.incarnation(), Some(token.as_str())); - - conn.drop_table("doubled", &[]).await.unwrap(); - let recreated = doubled_view(&conn).await; - assert_ne!(recreated.incarnation(), Some(token.as_str())); - - let err = recreated - .refresh() - .expect_incarnation(&token) - .execute() - .await - .unwrap_err(); - assert!(err.to_string().contains("dropped and recreated"), "{err}"); - assert_eq!(read(recreated.table(), "twice").await, Vec::::new()); - - recreated - .refresh() - .expect_incarnation(recreated.incarnation().unwrap()) - .execute() - .await - .unwrap(); - assert_eq!(read(recreated.table(), "twice").await, vec![2]); - } - - /// A cloned declaration creates two physical tables; each gets its own - /// token. - #[tokio::test] - async fn test_cloned_declaration_mints_a_fresh_incarnation_per_create() { - let (conn, source) = db_with_source(vec![1]).await; - let prepared = crate::materialized_view::prepare_declaration( - &source, - Some(&[("x".into(), "x".into()), ("twice".into(), "x * 2".into())]), - None, - None, - ) - .await - .unwrap(); - let replacement = prepared.clone(); - let first = prepared.create("cloned").await.unwrap(); - let first_token = first.incarnation().unwrap().to_string(); - - conn.drop_table("cloned", &[]).await.unwrap(); - let second = replacement.create("cloned").await.unwrap(); - assert_ne!(second.incarnation(), Some(first_token.as_str())); - } - - /// A recreation that lands after planning but before publication is - /// caught by the pre-commit read: the stale refresh fails and the - /// replacement stays empty under its own token. - #[tokio::test(flavor = "multi_thread")] - async fn test_bound_refresh_cannot_publish_into_a_raced_recreation() { - let _serial = DRIFT_LOCK.lock().await; - let (conn, _) = db_with_source(vec![1]).await; - let view = doubled_view(&conn).await; - let token = view.incarnation().unwrap().to_string(); - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let refreshing = - tokio::spawn(async move { view.refresh().expect_incarnation(token).execute().await }); - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("refresh never reached publication"); - - conn.drop_table("doubled", &[]).await.unwrap(); - let replacement = doubled_view(&conn).await; - let replacement_token = replacement.incarnation().unwrap().to_string(); - DRIFT_RELEASED.notify_one(); - - let result = refreshing.await.unwrap(); - assert!(result.is_err(), "the stale refresh unexpectedly succeeded"); - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - assert_eq!(reopened.incarnation(), Some(replacement_token.as_str())); - assert_eq!(read(reopened.table(), "twice").await, Vec::::new()); - } - - /// Replacing the schema metadata wholesale drops the token. A refresh - /// bound to the old token is refused for that reason, not as a - /// recreation; an unbound refresh mints the view a fresh one. - #[tokio::test] - async fn test_a_view_whose_metadata_was_replaced_starts_a_new_incarnation() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - let token = view.incarnation().unwrap().to_string(); - let mut metadata = HashMap::new(); - metadata.insert( - crate::materialized_view::DEFINITION_META_KEY.to_string(), - crate::materialized_view::definition_to_metadata(view.definition()).unwrap(), - ); - view.table() - .as_native() - .unwrap() - .replace_schema_metadata(metadata) - .await - .unwrap(); - assert_eq!( - conn.open_materialized_view("doubled") - .await - .unwrap() - .incarnation(), - None - ); - - let err = view - .refresh() - .expect_incarnation(&token) - .execute() - .await - .unwrap_err(); - assert!(err.to_string().contains("no incarnation token"), "{err}"); - - view.refresh().execute().await.unwrap(); - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - assert!(reopened.incarnation().is_some()); - assert_ne!(reopened.incarnation(), Some(token.as_str())); - } - - /// In-process refreshes of one view serialize: the loser of the race - /// observes the winner's watermark instead of appending the same rows. - #[tokio::test(flavor = "multi_thread")] - async fn test_concurrent_refreshes_do_not_duplicate() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - - append(&source, vec![2, 3]).await; - let (a, b) = tokio::join!(view.refresh().execute(), view.refresh().execute()); - let (a, b) = (a.unwrap(), b.unwrap()); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - let modes = [a.mode, b.mode]; - assert!(modes.contains(&RefreshMode::Incremental)); - assert!(modes.contains(&RefreshMode::NoOp)); - } - - /// A second handle's lazy cache must not defeat the lock: after another - /// handle's refresh commits, the stale handle plans from the reloaded - /// state and no-ops instead of appending the same fragments again. - #[tokio::test] - async fn test_a_second_handle_does_not_double_append() { - let (conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - let stale = conn.open_materialized_view("doubled").await.unwrap(); - - append(&source, vec![4]).await; - view.refresh().execute().await.unwrap(); - - let result = stale.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::NoOp); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6, 8]); - } - - /// A commit racing between a refresh's data commit and its stamp must - /// not be certified as the refresh's generation: the stamp aborts, and - /// the next refresh rebuilds from the drifted state. - #[tokio::test] - async fn test_stamp_aborts_on_a_racing_commit() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - let view_native = view.table().as_native().unwrap(); - let stale = view_native.dataset.get().await.unwrap().as_ref().clone(); - view.table().delete("x = 1").await.unwrap(); - - let err = stamp_watermark(view_native, stale, 99, 99, None, None).await; - assert!(err.is_err()); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// What refresh executes and what it stamps must be one generation: a - /// replaced definition wins over whatever a stale handle cached. - #[tokio::test] - async fn test_refresh_uses_the_latest_persisted_definition() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - let replacement = crate::materialized_view::MaterializedViewDefinition { - vector_source: None, - source_table: "src".into(), - source_namespace: Vec::new(), - projections: vec![ - crate::materialized_view::ViewProjection { - output: "x".into(), - expression: "x".into(), - }, - crate::materialized_view::ViewProjection { - output: "twice".into(), - expression: "x * 3".into(), - }, - ], - filter: None, - group_by: Vec::new(), - limit: None, - lateral: None, - }; - let mut metadata = HashMap::new(); - metadata.insert( - crate::materialized_view::DEFINITION_META_KEY.to_string(), - crate::materialized_view::definition_to_metadata(&replacement).unwrap(), - ); - view.table() - .as_native() - .unwrap() - .replace_schema_metadata(metadata) - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![3, 6]); - } - - /// A persisted definition that does not produce this view's schema must - /// not refresh at all, let alone be certified. - #[tokio::test] - async fn test_definition_view_schema_mismatch_is_refused() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - let narrower = crate::materialized_view::MaterializedViewDefinition { - vector_source: None, - source_table: "src".into(), - source_namespace: Vec::new(), - projections: vec![crate::materialized_view::ViewProjection { - output: "x".into(), - expression: "x".into(), - }], - filter: None, - group_by: Vec::new(), - limit: None, - lateral: None, - }; - let mut metadata = HashMap::new(); - metadata.insert( - crate::materialized_view::DEFINITION_META_KEY.to_string(), - crate::materialized_view::definition_to_metadata(&narrower).unwrap(), - ); - view.table() - .as_native() - .unwrap() - .replace_schema_metadata(metadata) - .await - .unwrap(); - - let err = view.refresh().execute().await.unwrap_err(); - assert!(matches!(err, Error::Schema { message } if message.contains("does not produce")),); - } - - /// An overlay replaces cell values without changing any file path; the - /// signature must see it, scoped to the columns the view reads like - /// data files are. - #[test] - fn test_fragment_signature_sees_overlays() { - use lance_file::version::ConcreteFileVersion; - use lance_table::format::DataFile; - use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; - - let base = Fragment::new(7); - let mut file = DataFile::new_unstarted("f0.lance", ConcreteFileVersion::V2_1); - file.fields = vec![0, 1].into(); - let mut with_file = base.clone(); - with_file.files.push(file.clone()); - - let overlay = |field: i32| { - let mut data_file = DataFile::new_unstarted("o0.lance", ConcreteFileVersion::V2_1); - data_file.fields = vec![field].into(); - DataOverlayFile { - data_file, - coverage: OverlayCoverage::PerField(Vec::new()), - committed_version: 9, - } - }; - let relevant: HashSet = [0].into_iter().collect(); - - let mut overlaid_relevant = with_file.clone(); - overlaid_relevant.overlays.push(overlay(0)); - assert_ne!( - fragment_signature(&with_file, &relevant), - fragment_signature(&overlaid_relevant, &relevant), - ); - - let mut overlaid_unrelated = with_file.clone(); - overlaid_unrelated.overlays.push(overlay(5)); - assert_eq!( - fragment_signature(&with_file, &relevant), - fragment_signature(&overlaid_unrelated, &relevant), - ); - } - - /// An output whose name needs quoting flows through as a projection - /// alias, never spliced into SQL text. - #[tokio::test] - async fn test_output_names_needing_quotes() { - let (conn, _) = db_with_source(vec![1, 2]).await; - let view = conn - .create_materialized_view("v", "src") - .with_no_data(true) - .select([("double value", "x * 2")]) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "double value").await, vec![2, 4]); - } - - /// MemWAL tiers are visible to reads but not to the refresh scan, so LSM - /// state disqualifies every participant: the source at create, either - /// side at refresh, and the view can never accept a spec. Retained - /// un-compacted rows (the catch-up flag outlives unset) count as state. - #[tokio::test] - async fn lsm_state_disqualifies_source_and_view() { - use crate::table::LsmWriteSpec; - use arrow_array::RecordBatchIterator; - - let tmp_dir = tempfile::tempdir().unwrap(); - let conn = connect(tmp_dir.path().to_str().unwrap()) - .execute() - .await - .unwrap(); - // Hand-rolled: the LSM primary key must be non-nullable, which - // record_batch! cannot express. - let schema = Arc::new(ArrowSchema::new(vec![ - arrow_schema::Field::new("id", arrow_schema::DataType::Int64, false), - arrow_schema::Field::new("x", arrow_schema::DataType::Int32, false), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(arrow_array::Int64Array::from(vec![1, 2])) as _, - Arc::new(Int32Array::from(vec![1, 2])) as _, - ], - ) - .unwrap(); - let table = conn - .create_table("src", batch.clone()) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - table.set_unenforced_primary_key(["id"]).await.unwrap(); - table - .set_lsm_write_spec(LsmWriteSpec::unsharded()) - .await - .unwrap(); - - // An active-LSM source is refused at create. - let err = conn - .create_materialized_view("v", "src") - .with_no_data(true) - .execute() - .await - .unwrap_err(); - assert!(err.to_string().contains("un-compacted"), "{err}"); - - // Unset with nothing written clears the state: the view creates, - // and a spec can never be installed over it. - table.unset_lsm_write_spec().await.unwrap(); - let view = conn - .create_materialized_view("v", "src") + let view = connection + .create_materialized_view("copy", "src") .with_no_data(true) .execute() .await .unwrap(); - let err = view - .table() - .set_lsm_write_spec(LsmWriteSpec::unsharded()) - .await - .unwrap_err(); - assert!(err.to_string().contains("materialized view"), "{err}"); - - // A source that acquires a spec after create fails refresh. - table - .set_lsm_write_spec(LsmWriteSpec::unsharded()) - .await - .unwrap(); - let err = view.refresh().execute().await.unwrap_err(); - assert!(err.to_string().contains("source table 'src'"), "{err}"); - - // Retained rows outlive unset: write through the WAL, unset, and - // refresh still refuses. - let mut merge = table.merge_insert(&["id"]); - merge - .when_matched_update_all(None) - .when_not_matched_insert_all() - .use_lsm(true); - merge - .execute(Box::new(RecordBatchIterator::new( - vec![Ok(batch.clone())], - batch.schema(), - ))) - .await - .unwrap(); - table.unset_lsm_write_spec().await.unwrap(); - let err = view.refresh().execute().await.unwrap_err(); - assert!(err.to_string().contains("source table 'src'"), "{err}"); - } - - /// A view with a computed column, declared over `people` and refreshed. - async fn refreshed_computed_view(conn: &Connection) -> MaterializedView { - use crate::materialized_view::tests::{computed_field, people, test_binding}; - let source = people(conn).await; - let view = crate::materialized_view::prepare_declaration( - &source, - Some(&[ - ("id".to_string(), "id".to_string()), - ("name".to_string(), "name".to_string()), - ]), - None, - None, - ) - .await - .unwrap() - .with_computed_columns( - vec![(2, computed_field("emb", "fb_1", "name"))], - &[test_binding("fb_1", "name", "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - view - } - - async fn unfilled(view: &MaterializedView) -> usize { - view.table() - .count_rows(Some("emb IS NULL".to_string())) - .await - .unwrap() - } - - async fn append_people(conn: &Connection, ids: Vec, names: Vec<&str>) { - let batch = record_batch!(("id", Int32, ids), ("name", Utf8, names)).unwrap(); - conn.open_table("people") - .execute() - .await - .unwrap() - .add(batch) - .execute() - .await - .unwrap(); - } - - /// Commit the fill job's shape on the view: a column rewrite of - /// `fields`, touching no rows. The data is left as it is; what matters - /// here is how the next refresh classifies the commit. - async fn commit_column_rewrite(view: &MaterializedView, fields: &[&str]) { - let native = view.table().as_native().unwrap(); - native.dataset.reload().await.unwrap(); - let dataset = native.dataset.get().await.unwrap().as_ref().clone(); - let fields_modified = fields - .iter() - .map(|name| dataset.schema().field(name).unwrap().id as u32) - .collect(); - let updated_fragments = dataset - .get_fragments() - .iter() - .map(|fragment| fragment.metadata().clone()) - .collect(); - let operation = Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments, - new_fragments: Vec::new(), - fields_modified, - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: Some(UpdateMode::RewriteColumns), - inserted_rows_filter: None, - updated_fragment_offsets: None, - }; - let read_version = dataset.version().version; - CommitBuilder::new(WriteDestination::Dataset(Arc::new(dataset))) - .execute(Transaction::new(read_version, operation, None)) - .await - .unwrap(); - } - - /// Refresh never computes a computed column: every row it writes, on a - /// rebuild, an append and a rewrite, carries NULL there, and the - /// declaration survives all three. - #[tokio::test] - async fn test_computed_columns_are_written_null_and_kept() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - assert_eq!(unfilled(&view).await, 3); - - append_people(&conn, vec![4, 5], vec!["d", "e"]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(unfilled(&view).await, 5); - - conn.open_table("people") - .execute() - .await - .unwrap() - .update() - .column("name", "'z'") - .only_if("id = 1") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!(unfilled(&view).await, 5); - assert_eq!(read(view.table(), "id").await, vec![1, 2, 3, 4, 5]); - let schema = view.table().schema().await.unwrap(); - assert!( - crate::table::computed_columns::function_bindings(&schema) - .unwrap() - .iter() - .any(|b| b.binding_id() == "fb_1"), - "the binding envelope was lost" - ); - assert!( - computed_column_from_field(schema.field_with_name("emb").unwrap()).is_some(), - "the declaration was lost" - ); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// Field metadata on `field` only, the commit shape of the freshness - /// stamp `refresh_column` leaves after its fill. - async fn commit_field_metadata(view: &MaterializedView, field: &str, key: &str) { - let native = view.table().as_native().unwrap(); - native.dataset.reload().await.unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - dataset - .update_field_metadata() - .update(field, [(key.to_string(), "{}".to_string())]) - .unwrap() - .await - .unwrap(); - } - - /// The stamp is metadata on the computed column and rewrites nothing - /// refresh certifies, so it is not drift; the same commit shape on a - /// projected column is, like any other write to it. - #[tokio::test] - async fn test_a_freshness_stamp_is_not_drift() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - - commit_field_metadata( - &view, - "emb", - crate::table::computed_columns::SOURCE_SIGNATURE_META_KEY, - ) - .await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - - commit_field_metadata(&view, "id", "probe").await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - } - - /// The fill job's commit rewrites only computed columns. It is the one - /// commit on a view that is not drift: the next refresh carries on from - /// its watermark instead of rebuilding, which would null what the fill - /// just wrote. - #[tokio::test] - async fn test_a_computed_column_fill_is_not_drift() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - - commit_column_rewrite(&view, &["emb"]).await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - - commit_column_rewrite(&view, &["emb"]).await; - append_people(&conn, vec![4], vec!["d"]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "id").await, vec![1, 2, 3, 4]); - } - - /// A column rewrite that reaches a projected column is drift like any - /// other write: refresh certifies those columns and must recompute them. - #[tokio::test] - async fn test_a_rewrite_of_a_projected_column_is_drift() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - - commit_column_rewrite(&view, &["emb", "name"]).await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - } - - /// The declaration contract is checked before any refresh mutation: a - /// missing binding envelope and a column that lost its declaration both - /// fail closed. - #[tokio::test] - async fn test_a_broken_declaration_is_refused_before_refresh() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - dataset - .update_schema_metadata(vec![( - crate::table::computed_columns::FUNCTION_BINDINGS_META_KEY.to_string(), - None, - )]) - .await - .unwrap(); - let err = view.refresh().execute().await.unwrap_err().to_string(); - assert!(err.contains("references missing binding 'fb_1'"), "{err}"); - - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - dataset - .replace_field_metadata(vec![( - dataset.schema().field("emb").unwrap().id as u32, - HashMap::new(), - )]) - .await - .unwrap(); - let err = view.refresh().execute().await.unwrap_err().to_string(); - assert!(err.contains("does not match binding 'fb_1'"), "{err}"); - } - - /// An input the view does not project is materialized on every refresh - /// path, before the provenance column, with the source's values. - #[tokio::test] - async fn test_internal_inputs_are_materialized_and_refreshed() { - use crate::materialized_view::tests::{computed_field, strict_people, test_binding}; - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - let mut prepared = crate::materialized_view::prepare_declaration( - &source, - Some(&[("id".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap(); - let input = prepared.input_column("name").unwrap(); - let view = prepared - .with_computed_columns( - vec![(1, computed_field("emb", "fb_1", &input))], - &[test_binding("fb_1", &input, "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let names: Vec = view + let result_batch = + RecordBatch::try_new(schema.clone(), source_batch.columns().to_vec()).unwrap(); + let planned = view .table() - .schema() - .await - .unwrap() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(names, ["id", "emb", "__input_name", SOURCE_ROW_ID_COLUMN]); - - let unfilled_inputs = || async { - view.table() - .count_rows(Some("__input_name IS NULL".to_string())) - .await - .unwrap() - }; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - assert_eq!(unfilled_inputs().await, 0); - - let more = arrow_array::RecordBatch::try_new( - source.schema().await.unwrap(), - vec![ - Arc::new(Int32Array::from(vec![4])), - Arc::new(arrow_array::StringArray::from(vec!["d"])), - ], - ) - .unwrap(); - source.add(more).execute().await.unwrap(); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Incremental - ); - assert_eq!(unfilled_inputs().await, 0); - assert_eq!( - view.table() - .count_rows(Some("__input_name = 'd'".to_string())) - .await - .unwrap(), - 1 - ); - - source - .update() - .column("name", "'z'") - .only_if("id = 1") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - view.table() - .count_rows(Some("__input_name = 'z'".to_string())) - .await - .unwrap(), - 1 - ); - assert_eq!( - unfilled(&view).await, - 4, - "rewritten and new rows are unfilled" - ); - } - - /// A SQL declaration is filled by `refresh_column` on the view, which - /// commits a data replacement and then its freshness stamp; the next - /// refresh continues from its watermark and keeps what the fill wrote, - /// and only rows the view added since come back unfilled. - #[tokio::test] - async fn test_a_sql_fill_is_not_drift() { - use crate::materialized_view::tests::{people, sql_field}; - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let view = crate::materialized_view::prepare_declaration( - &source, - Some(&[("id".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap() - .with_computed_columns( - vec![( - 1, - sql_field("next", arrow_schema::DataType::Int32, "id + 1", r#"["id"]"#), - )], - &[], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let filled = || async { - view.table() - .count_rows(Some("next = id + 1".to_string())) - .await - .unwrap() - }; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - assert_eq!( - view.table() - .refresh_column("next") - .await - .unwrap() - .rows_filled, - 3 - ); - assert_eq!(filled().await, 3); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - assert_eq!(filled().await, 3); - - append_people(&conn, vec![4], vec!["d"]).await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Incremental - ); - assert_eq!(filled().await, 3); - assert_eq!( - view.table() - .refresh_column("next") - .await - .unwrap() - .rows_filled, - 1 - ); - assert_eq!(filled().await, 4); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// A fill of a nested computed column writes its child fields; that is - /// still a fill, not drift. - #[tokio::test] - async fn test_a_nested_sql_fill_is_not_drift() { - use crate::materialized_view::tests::{people, sql_field}; - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let payload = sql_field( - "payload", - arrow_schema::DataType::Struct( - vec![arrow_schema::Field::new( - "value", - arrow_schema::DataType::Utf8, - true, - )] - .into(), - ), - "named_struct('value', name)", - r#"["name"]"#, - ); - let view = crate::materialized_view::prepare_declaration( - &source, - Some(&[("name".to_string(), "name".to_string())]), - None, - None, - ) - .await - .unwrap() - .with_computed_columns(vec![(1, payload)], &[]) - .unwrap() - .create("v") - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - view.table() - .refresh_column("payload") - .await - .unwrap() - .rows_filled, - 3 - ); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - assert_eq!( - view.table() - .count_rows(Some("payload.value = name".to_string())) - .await - .unwrap(), - 3 - ); - } - - async fn grouped_source(conn: &Connection) -> Table { - conn.create_table( - "src", - record_batch!(("k", Int32, [1, 1, 2, 3]), ("x", Int32, [10, 11, 20, 30])).unwrap(), - ) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap() - } - - pub(in crate::materialized_view) async fn declare( - source: &Table, - name: &str, - sql: &str, - ) -> Result { - crate::materialized_view::prepare_definition( - source, - MaterializedViewDefinition::from_sql(sql)?, - ) - .await? - .create(name) - .await - } - - /// Rows of `table` as `columns` rendered to text, sorted. - pub(in crate::materialized_view) async fn rows(table: &Table, columns: &[&str]) -> Vec { - let batches = table - .query() - .select(Select::columns(columns)) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let mut rows = Vec::new(); - for batch in &batches { - let formatters: Vec<_> = batch - .columns() - .iter() - .map(|c| { - arrow_cast::display::ArrayFormatter::try_new(c.as_ref(), &Default::default()) - .unwrap() - }) - .collect(); - for row in 0..batch.num_rows() { - let cells: Vec = formatters - .iter() - .map(|f| f.value(row).to_string()) - .collect(); - rows.push(cells.join(" ")); - } - } - rows.sort(); - rows - } - - #[tokio::test] - async fn a_grouped_view_holds_one_row_per_group() { - let conn = connect("memory://").execute().await.unwrap(); - let source = grouped_source(&conn).await; - let view = declare( - &source, - "groups", - "SELECT k, array_agg(x) AS xs, sum(x) AS total FROM src WHERE x < 30 GROUP BY k", - ) - .await - .unwrap(); - assert_eq!( - view.definition().to_sql(), - "SELECT k, array_agg(x) AS xs, sum(x) AS total FROM src WHERE x < 30 GROUP BY k" - ); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 2); - assert_eq!(rows(view.table(), &["k", "total"]).await, ["1 21", "2 20"]); - let unchanged = view.refresh().execute().await.unwrap(); - assert_eq!(unchanged.mode, RefreshMode::NoOp); - - // Any source change recomputes every group. - source - .add(record_batch!(("k", Int32, [2]), ("x", Int32, [5])).unwrap()) - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(rows(view.table(), &["k", "total"]).await, ["1 21", "2 25"]); - - // A view over the groups expands them again. - let expanded = declare( - view.table(), - "elements", - "SELECT k, e AS x FROM groups, UNNEST(xs) AS e", - ) - .await - .unwrap(); - expanded.refresh().execute().await.unwrap(); - assert_eq!( - rows(expanded.table(), &["k", "x"]).await, - ["1 10", "1 11", "2 20", "2 5"] - ); - } - - /// Provenance is one source row per group, so the view's row-id keys - /// stay unique. - #[tokio::test] - async fn a_grouped_view_records_one_source_row_per_group() { - let conn = connect("memory://").execute().await.unwrap(); - let source = grouped_source(&conn).await; - let view = declare( - &source, - "groups", - "SELECT k, count(*) AS n FROM src GROUP BY k", - ) - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - let ids = rows(view.table(), &[SOURCE_ROW_ID_COLUMN]).await; - assert_eq!(ids.len(), 3); - assert_eq!(ids.iter().collect::>().len(), 3, "{ids:?}"); - } - - #[tokio::test] - async fn a_grouped_view_is_planned_at_declaration() { - let conn = connect("memory://").execute().await.unwrap(); - let source = grouped_source(&conn).await; - for sql in [ - // x is neither a key nor aggregated - "SELECT k, x FROM src GROUP BY k", - "SELECT k, random() AS r FROM src GROUP BY k", - "SELECT k, count(*) AS n FROM src GROUP BY missing", - "SELECT k, count(*) AS __source_row_id FROM src GROUP BY k", - "SELECT k, count(*) AS k FROM src GROUP BY k", - ] { - assert!(declare(&source, "v", sql).await.is_err(), "{sql}"); - } - let mut grouped = crate::materialized_view::prepare_definition( - &source, - MaterializedViewDefinition::from_sql("SELECT k, count(*) AS n FROM src GROUP BY k") - .unwrap(), - ) - .await - .unwrap(); - assert!(grouped.input_column("x").is_err()); - } - - /// Twenty nonzero vectors in two clusters far apart, ids 0-9 and 10-19. - /// Nonzero so every vector has a cosine direction. - pub(in crate::materialized_view) async fn clustered_source(conn: &Connection) -> Table { - use arrow_array::types::Float32Type; - use arrow_array::{FixedSizeListArray, Int32Array}; - - let vectors = FixedSizeListArray::from_iter_primitive::( - (0..20).map(|i| { - let base = if i < 10 { 0.0 } else { 100.0 }; - Some(vec![Some(base + 1.0 + i as f32 * 0.1), Some(base)]) - }), - 2, - ); - let batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(0..20)) as arrow_array::ArrayRef, - ), - ("vec", Arc::new(vectors) as arrow_array::ArrayRef), - ]) - .unwrap(); - conn.create_table("images", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap() - } - - #[tokio::test] - async fn ivf_partition_groups_by_the_index_on_the_column() { - use crate::index::vector::IvfFlatIndexBuilder; - - const BUCKETS: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images GROUP BY ivf_partition(vec)"; - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - let unindexed = declare(&source, "buckets", BUCKETS).await.unwrap_err(); - assert!( - unindexed.to_string().contains("IVF vector index"), - "{unindexed}" - ); - - source - .create_index( - &["vec"], - Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(2)), - ) - .execute() - .await - .unwrap(); - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["n", "lo", "hi"]).await, - ["10 0 9", "10 10 19"] - ); - let buckets = rows(view.table(), &["bucket"]).await; - assert_eq!(buckets, ["0", "1"]); - } - - #[tokio::test] - async fn ivf_partition_takes_a_column() { - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - for sql in [ - "SELECT ivf_partition(id) AS b, count(*) AS n FROM images GROUP BY ivf_partition(id)", - "SELECT ivf_partition(vec, vec) AS b, count(*) AS n FROM images \ - GROUP BY ivf_partition(vec, vec)", - ] { - assert!(declare(&source, "v", sql).await.is_err(), "{sql}"); - } - } - - /// A cosine index assigns by L2 over unit vectors, as lance's own index - /// path does; forwarding cosine to the assigner panics. - #[tokio::test] - async fn ivf_partition_supports_cosine_indices() { - use crate::index::vector::IvfFlatIndexBuilder; - - const BUCKETS: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images GROUP BY ivf_partition(vec)"; - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .execute() - .await - .unwrap(); - let view = declare(&source, "cosine_buckets", BUCKETS).await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["n", "lo", "hi"]).await, - ["10 0 9", "10 10 19"] - ); - } - - /// Segments of one index trained separately carry different centroids; - /// grouping by one of them would bucket the other segments' rows wrongly. - #[tokio::test] - async fn ivf_partition_refuses_segments_with_different_models() { - use arrow_array::types::Float32Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array}; - use lance::index::DatasetIndexExt; - use lance::index::vector::VectorIndexParams; - use lance_index::IndexType; - - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - // A second fragment far from the first, so each segment trains its own centroids. - let far = FixedSizeListArray::from_iter_primitive::( - (0..20).map(|i| Some(vec![Some(-500.0 - i as f32), Some(-500.0)])), - 2, - ); - source - .add( - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(20..40)) as ArrayRef, - ), - ("vec", Arc::new(far) as ArrayRef), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - let dataset = source.as_native().unwrap().dataset.get().await.unwrap(); - let mut dataset = dataset.as_ref().clone(); - let params = VectorIndexParams::ivf_flat(2, lance_linalg::distance::DistanceType::L2); - let mut segments = Vec::new(); - for fragment in dataset.get_fragments() { - let segment = dataset - .create_index_builder(&["vec"], IndexType::Vector, ¶ms) - .name("shared".to_string()) - .fragments(vec![fragment.id() as u32]) - .execute_uncommitted() - .await - .unwrap(); - segments.push(segment); - } - dataset - .commit_existing_index_segments("shared", "vec", segments) - .await - .unwrap(); - assert_eq!(dataset.load_indices().await.unwrap().len(), 2); - - let source = conn.open_table("images").execute().await.unwrap(); - let error = declare( - &source, - "buckets", - "SELECT ivf_partition(vec) AS b, count(*) AS n FROM images GROUP BY ivf_partition(vec)", - ) - .await - .unwrap_err(); - assert!(error.to_string().contains("trained separately"), "{error}"); - } - - /// A vector the assigner cannot place has no bucket: it is grouped under - /// NULL, never under partition 0. - #[tokio::test] - async fn ivf_partition_leaves_an_unassignable_vector_unbucketed() { - use crate::index::vector::IvfFlatIndexBuilder; - use arrow_array::types::Float32Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array}; - - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .execute() - .await - .unwrap(); - // The zero vector has no direction, so cosine cannot place it. - let zero = FixedSizeListArray::from_iter_primitive::( - [Some(vec![Some(0.0), Some(0.0)])], - 2, - ); - source - .add( - RecordBatch::try_from_iter(vec![ - ("id", Arc::new(Int32Array::from(vec![20])) as ArrayRef), - ("vec", Arc::new(zero) as ArrayRef), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - let view = declare( - &source, - "buckets", - "SELECT ivf_partition(vec) AS bucket, count(*) AS n, min(id) AS lo, max(id) AS hi \ - FROM images GROUP BY ivf_partition(vec)", - ) - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - // ArrayFormatter renders NULL as the empty string; partition numbers - // are the index's own and not asserted. - let groups = rows(view.table(), &["bucket", "n", "lo", "hi"]).await; - let (unbucketed, bucketed): (Vec<_>, Vec<_>) = - groups.iter().partition(|row| row.starts_with(' ')); - assert_eq!(unbucketed, [" 1 20 20"], "{groups:?}"); - let mut clusters: Vec<&str> = bucketed - .iter() - .map(|row| row.split_once(' ').unwrap().1) - .collect(); - clusters.sort(); - assert_eq!(clusters, ["10 0 9", "10 10 19"], "{groups:?}"); - } - - /// Byte vectors group by a Hamming index, the phash case. The centroids - /// are given, as lance's own tests do: two-means over a handful of - /// hashes can converge with one partition empty. - #[tokio::test] - async fn ivf_partition_groups_byte_vectors_by_hamming() { - use arrow_array::types::UInt8Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array}; - use lance::index::DatasetIndexExt; - use lance::index::vector::VectorIndexParams; - use lance_index::IndexType; - use lance_index::vector::ivf::IvfBuildParams; - - let hashes = FixedSizeListArray::from_iter_primitive::( - (0..20u8).map(|i| { - let base = if i < 10 { 0x00 } else { 0xff }; - Some(vec![ - Some(base ^ (i % 4)), - Some(base), - Some(base), - Some(base), - ]) - }), - 4, - ); - let batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(0..20)) as ArrayRef, - ), - ("phash", Arc::new(hashes) as ArrayRef), - ]) - .unwrap(); - let conn = connect("memory://").execute().await.unwrap(); - let source = conn - .create_table("images", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let centroids = FixedSizeListArray::from_iter_primitive::( - [Some(vec![Some(0); 4]), Some(vec![Some(0xff); 4])], - 4, - ); - let params = VectorIndexParams::with_ivf_flat_params( - lance_linalg::distance::DistanceType::Hamming, - IvfBuildParams::try_with_centroids(2, Arc::new(centroids)).unwrap(), - ); - let mut dataset = source .as_native() .unwrap() .dataset @@ -4806,19 +607,63 @@ pub(crate) mod tests { .unwrap() .as_ref() .clone(); - dataset - .create_index(&["phash"], IndexType::Vector, None, ¶ms, true) + assert!(planned.get_fragments().is_empty()); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let loser_batch = result_batch.clone(); + let blocked = futures::stream::once(async move { + entered_tx.send(()).unwrap(); + release_rx.await.unwrap(); + Ok::<_, DataFusionError>(loser_batch) + }); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema.clone(), blocked)); + let loser_table = view.table().clone(); + let loser = + tokio::spawn( + async move { replace_materialized_view_fragments(&loser_table, stream).await }, + ); + entered_rx.await.unwrap(); + + let winner_stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter([Ok(result_batch)]), + )); + let planned = Arc::new(planned); + let staged = InsertBuilder::new(WriteDestination::Dataset(planned.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted_stream(winner_stream) .await .unwrap(); - let source = conn.open_table("images").execute().await.unwrap(); + let Operation::Append { fragments } = staged.operation else { + panic!("expected append staging"); + }; + let filter = refresh_conflict_filter(&planned).unwrap(); + CommitBuilder::new(WriteDestination::Dataset(planned.clone())) + .execute(Transaction::new( + planned.version().version, + Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments: Vec::new(), + new_fragments: fragments, + fields_modified: Vec::new(), + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: Vec::new(), + update_mode: None, + inserted_rows_filter: Some(filter), + updated_fragment_offsets: None, + }, + None, + )) + .await + .unwrap(); + release_tx.send(()).unwrap(); - const BUCKETS: &str = "SELECT ivf_partition(phash) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images GROUP BY ivf_partition(phash)"; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - ["0 10 0 9", "1 10 10 19"] - ); + loser.await.unwrap().unwrap_err(); + let current = connection.open_table("copy").execute().await.unwrap(); + assert_eq!(current.count_rows(None).await.unwrap(), 3); } } diff --git a/rust/lancedb/src/materialized_view/vector_dedup.rs b/rust/lancedb/src/materialized_view/vector_dedup.rs deleted file mode 100644 index dc2759433..000000000 --- a/rust/lancedb/src/materialized_view/vector_dedup.rs +++ /dev/null @@ -1,636 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The LanceDB Authors - -//! Two dependent waves in the ordinary refresh fleet: index partition selection, -//! then source fragment materialization. Only the latter publishes data files. - -use std::collections::BTreeMap; -use std::io::Cursor; -use std::sync::Arc; - -use arrow_array::{Array, RecordBatch, UInt64Array}; -use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use datafusion::prelude::{SessionConfig, SessionContext, col}; -use datafusion_common::DataFusionError; -use datafusion_execution::{ - SendableRecordBatchStream, TaskContext, runtime_env::RuntimeEnvBuilder, -}; -use datafusion_physical_plan::{ExecutionPlan, stream::RecordBatchStreamAdapter}; -use futures::TryStreamExt; -use lance::Dataset; -use lance::dataset::rowids::{get_row_id_index, load_row_id_sequence}; -use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap}; -use lance::dataset::transaction::Operation; -use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams}; -use lance_core::datatypes::BlobHandling; -use lance_core::utils::address::RowAddress; -use lance_datafusion::exec::{OneShotExec, SessionContextExt}; -use lance_table::format::Fragment; -use roaring::{RoaringBitmap, RoaringTreemap}; -use serde::{Deserialize, Serialize}; -use sha2::{Digest, Sha256}; -use uuid::Uuid; - -use super::duplicate_pairs::{self, DuplicatePairsRefreshPlan, VectorSourceKind}; -use super::{GroupedRefreshPlan, RefreshMaterializedViewResult}; -use crate::table::Table; -use crate::table::datafusion::udtf::duplicate_pairs::{DuplicatePairsExec, plan_duplicate_pairs}; -use crate::{Error, Result}; - -const POLICY_VERSION: u32 = 1; -const SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024; -const MATERIALIZE_BATCH_ROWS: usize = 1024; -const MASK_FILE: &str = "dedup.bitmap"; -const MASK_MAGIC: &[u8] = b"LANCEDB_DEDUP_1"; - -fn invalid(message: impl Into) -> Error { - Error::InvalidInput { - message: message.into(), - } -} -fn io_error(error: std::io::Error) -> Error { - invalid(format!("invalid dedup bitmap: {error}")) -} - -/// Immutable snapshot, partition tasks and source fragments for both waves. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct VectorDedupPlan { - pub snapshot: GroupedRefreshPlan, - native: DuplicatePairsRefreshPlan, - source_fragments: Vec, - policy_version: u32, -} -impl VectorDedupPlan { - pub(super) fn dependencies(&self) -> u32 { - self.native.snapshot.units - } - fn identity(&self) -> Result { - Ok(format!( - "{:x}", - Sha256::digest(serde_json::to_vec(self).map_err(|e| invalid(e.to_string()))?) - )) - } -} - -/// Small receipt for a durable selection mask. Row IDs and pairs are not stored -/// in the host's task registry. UUID directories use ordinary index-file GC. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -struct MaskFile { - source_fragment: u32, - object_id: Uuid, - bytes: usize, - digest: String, -} - -/// A selection task produces masks; a materialization task produces data files. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WrittenDedupUnit { - unit: u32, - plan_id: String, - masks: Vec, - pub(super) fragments: Vec, - pub(super) rows: u64, -} - -pub(super) async fn plan_refresh(view: &Table, pinned: Option) -> Result { - let native = duplicate_pairs::plan_refresh(view, pinned).await?; - let (_, _, source) = duplicate_pairs::open_planned(view, &native).await?; - let source_fragments = source - .get_fragments() - .iter() - .map(|f| f.id() as u32) - .collect::>(); - let mut snapshot = native.snapshot.clone(); - snapshot.units = snapshot - .units - .checked_add( - u32::try_from(source_fragments.len()) - .map_err(|_| invalid("too many source fragments"))?, - ) - .ok_or_else(|| invalid("too many dedup refresh units"))?; - Ok(VectorDedupPlan { - snapshot, - native, - source_fragments, - policy_version: POLICY_VERSION, - }) -} - -async fn open( - view: &Table, - plan: &VectorDedupPlan, -) -> Result<(Dataset, super::MaterializedViewDefinition, Arc)> { - let (view_ds, definition, source) = duplicate_pairs::open_planned(view, &plan.native).await?; - let mut snapshot = plan.native.snapshot.clone(); - snapshot.units = snapshot - .units - .checked_add( - u32::try_from(plan.source_fragments.len()) - .map_err(|_| invalid("too many dedup fragments"))?, - ) - .ok_or_else(|| invalid("too many dedup units"))?; - if plan.policy_version != POLICY_VERSION - || snapshot != plan.snapshot - || definition.vector_source.as_ref().map(|s| s.kind) != Some(VectorSourceKind::Dedup) - || source - .get_fragments() - .iter() - .map(|f| f.id() as u32) - .collect::>() - != plan.source_fragments - { - return Err(invalid("dedup task belongs to another snapshot or policy")); - } - Ok((view_ds, definition, source)) -} - -fn edge_schema() -> SchemaRef { - Arc::new(Schema::new(vec![ - Field::new("a", DataType::UInt64, false), - Field::new("b", DataType::UInt64, false), - ])) -} - -/// Sorting canonical edges makes the first unremoved endpoint the smallest -/// unassigned representative. The sort spills; only vertex masks stay in RAM. -pub async fn select_representatives( - stream: SendableRecordBatchStream, - memory: usize, -) -> Result<(RoaringTreemap, RoaringTreemap)> { - let schema = edge_schema(); - let batch_schema = schema.clone(); - let canonical = stream - .map_ok(move |batch| { - let a = batch - .column(0) - .as_any() - .downcast_ref::() - .expect("native row id"); - let b = batch - .column(1) - .as_any() - .downcast_ref::() - .expect("native row id"); - RecordBatch::try_new( - batch_schema.clone(), - vec![ - Arc::new(UInt64Array::from_iter_values( - a.values().iter().zip(b.values()).map(|(&a, &b)| a.min(b)), - )), - Arc::new(UInt64Array::from_iter_values( - a.values().iter().zip(b.values()).map(|(&a, &b)| a.max(b)), - )), - ], - ) - .map_err(DataFusionError::from) - }) - .and_then(futures::future::ready); - let canonical = Box::pin(RecordBatchStreamAdapter::new(schema, canonical)); - let mut config = SessionConfig::default().with_batch_size(MATERIALIZE_BATCH_ROWS); - // Leave enough room to merge spilled runs while the final in-memory run - // is still resident. Small batches also bound each merge input buffer. - config.options_mut().execution.sort_spill_reservation_bytes = memory / 2; - let ctx = SessionContext::new_with_config_rt( - config, - RuntimeEnvBuilder::new() - .with_memory_limit(memory, 1.0) - .build_arc()?, - ); - let mut sorted = ctx - .read_one_shot(canonical)? - .sort_by(vec![col("a"), col("b")])? - .execute_stream() - .await?; - let mut removed = RoaringTreemap::new(); - let mut representatives = RoaringTreemap::new(); - while let Some(batch) = sorted.try_next().await? { - let a = batch - .column(0) - .as_any() - .downcast_ref::() - .expect("sorted id"); - let b = batch - .column(1) - .as_any() - .downcast_ref::() - .expect("sorted id"); - for (&a, &b) in a.values().iter().zip(b.values()) { - if a != b && !removed.contains(a) && !removed.contains(b) { - representatives.insert(a); - removed.insert(b); - } - } - } - Ok((removed, representatives)) -} - -#[derive(Default)] -struct Masks { - removed: RoaringBitmap, - representatives: RoaringBitmap, -} - -async fn persist_masks( - view: &Dataset, - source: &Dataset, - removed: RoaringTreemap, - representatives: RoaringTreemap, -) -> Result> { - let row_index = get_row_id_index(source).await?; - let mut masks: BTreeMap = BTreeMap::new(); - for (ids, is_removed) in [(removed, true), (representatives, false)] { - for id in ids { - let address = match &row_index { - Some(index) => index - .get(id)? - .ok_or_else(|| invalid("pair row id missing from pinned source"))?, - None => RowAddress::new_from_u64(id), - }; - let mask = masks.entry(address.fragment_id()).or_default(); - if is_removed { - mask.removed.insert(address.row_offset()); - } else { - mask.representatives.insert(address.row_offset()); - } - } - } - let store = view.object_store(None).await?; - let mut files = Vec::new(); - for (source_fragment, mask) in masks { - if source.get_fragment(source_fragment as usize).is_none() { - return Err(invalid("pair row address is outside the pinned snapshot")); - } - let mut bytes = MASK_MAGIC.to_vec(); - mask.removed.serialize_into(&mut bytes).map_err(io_error)?; - mask.representatives - .serialize_into(&mut bytes) - .map_err(io_error)?; - let object_id = Uuid::new_v4(); - store - .put( - &view - .indices_dir() - .join(object_id.to_string()) - .join(MASK_FILE), - &bytes, - ) - .await?; - files.push(MaskFile { - source_fragment, - object_id, - bytes: bytes.len(), - digest: format!("{:x}", Sha256::digest(&bytes)), - }); - } - Ok(files) -} - -fn validate_receipts<'a>( - plan: &VectorDedupPlan, - receipts: &'a [WrittenDedupUnit], - count: u32, -) -> Result> { - let identity = plan.identity()?; - let mut ordered = vec![None; count as usize]; - for receipt in receipts { - if receipt.plan_id != identity { - return Err(invalid("dedup receipt belongs to another plan")); - } - let slot = ordered - .get_mut(receipt.unit as usize) - .ok_or_else(|| invalid("unexpected dedup receipt"))?; - if slot.is_some() { - return Err(invalid("duplicate dedup receipt")); - } - if receipt.unit < plan.dependencies() { - if !receipt.fragments.is_empty() - || receipt.rows != 0 - || receipt - .masks - .iter() - .any(|m| !plan.source_fragments.contains(&m.source_fragment)) - { - return Err(invalid("invalid selection receipt")); - } - } else if !receipt.masks.is_empty() { - return Err(invalid("invalid materialization receipt")); - } - *slot = Some(receipt); - } - ordered - .into_iter() - .collect::>>() - .ok_or_else(|| invalid("missing dedup dependency receipt")) -} - -async fn load_masks( - view: &Dataset, - source_fragment: u32, - dependencies: &[&WrittenDedupUnit], -) -> Result { - let store = view.object_store(None).await?; - let mut masks = Masks::default(); - for file in dependencies - .iter() - .flat_map(|u| &u.masks) - .filter(|m| m.source_fragment == source_fragment) - { - let reader = store - .open( - &view - .indices_dir() - .join(file.object_id.to_string()) - .join(MASK_FILE), - ) - .await?; - if reader.size().await? != file.bytes { - return Err(invalid("dedup mask length mismatch")); - } - let bytes = reader.get_range(0..file.bytes).await?; - if !bytes.starts_with(MASK_MAGIC) || format!("{:x}", Sha256::digest(&bytes)) != file.digest - { - return Err(invalid("dedup mask checksum mismatch")); - } - let mut cursor = Cursor::new(&bytes[MASK_MAGIC.len()..]); - masks.removed |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?; - masks.representatives |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?; - if cursor.position() as usize != bytes.len() - MASK_MAGIC.len() { - return Err(invalid("dedup mask has trailing data")); - } - } - // Native index scopes normally have disjoint rows. Fail closed if an - // overlapping scope would delete a representative used by another scope. - if !masks.removed.is_disjoint(&masks.representatives) { - return Err(invalid( - "overlapping index scopes disagree on retained representatives", - )); - } - Ok(masks.removed) -} - -pub(super) async fn write_unit( - view: &Table, - unit: u32, - plan: &VectorDedupPlan, - dependencies: &[WrittenDedupUnit], -) -> Result { - let (view_ds, definition, source) = open(view, plan).await?; - let mut written = WrittenDedupUnit { - unit, - plan_id: plan.identity()?, - masks: Vec::new(), - fragments: Vec::new(), - rows: 0, - }; - if unit < plan.dependencies() { - let config = definition - .vector_source - .as_ref() - .expect("dedup source") - .config()?; - let exec = DuplicatePairsExec::try_new( - source.clone(), - config, - vec![plan.native.tasks[unit as usize].clone()], - )?; - let stream = exec.execute(0, Arc::new(TaskContext::default()))?; - let (removed, representatives) = select_representatives(stream, SORT_MEMORY_BYTES).await?; - written.masks = persist_masks(&view_ds, &source, removed, representatives).await?; - return Ok(written); - } - let source_fragment = *plan - .source_fragments - .get((unit - plan.dependencies()) as usize) - .ok_or_else(|| invalid("dedup unit is outside the plan"))?; - let dependencies = validate_receipts(plan, dependencies, plan.dependencies())?; - let removed = load_masks(&view_ds, source_fragment, &dependencies).await?; - let fragment = source - .get_fragment(source_fragment as usize) - .ok_or_else(|| invalid("source fragment missing"))?; - // The scanner prefilter uses snapshot row IDs. Our persisted masks use - // physical offsets so each writer only loads its own fragment's artifacts. - let blocked = if source.manifest.uses_stable_row_ids() { - let sequence = load_row_id_sequence(&source, fragment.metadata()).await?; - if removed - .max() - .is_some_and(|offset| offset as u64 >= sequence.len()) - { - return Err(invalid("dedup mask offset is outside the source fragment")); - } - RowAddrTreeMap::from_iter(sequence.select(removed.iter().map(|offset| offset as usize))) - } else { - RowAddrTreeMap::from_iter( - removed - .iter() - .map(|offset| u64::from(RowAddress::new_from_parts(source_fragment, offset))), - ) - }; - let mut scanner = source.scan(); - scanner - .with_fragments(vec![fragment.metadata().clone()]) - .with_row_addr_prefilter(RowAddrMask::from_block(blocked)) - .blob_handling(BlobHandling::AllBinary) - .batch_size(MATERIALIZE_BATCH_ROWS) - .strict_batch_size(true); - let stream = scanner.try_into_stream().await?; - let schema = lance_io::stream::RecordBatchStream::schema(&stream); - let rows = Arc::new(std::sync::atomic::AtomicU64::new(0)); - let count = rows.clone(); - let stream = stream - .map_err(|e| DataFusionError::External(Box::new(e))) - .inspect_ok(move |batch| { - count.fetch_add( - batch.num_rows() as u64, - std::sync::atomic::Ordering::Relaxed, - ); - }); - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); - // Materialized binary payloads must become the destination's logical blob - // structs before the Lance writer encodes them. Reuse the normal add path's - // coercions, including nulls and nested fields. - let coerced = crate::table::datafusion::cast::cast_to_table_schema( - Arc::new(OneShotExec::new(stream)), - &Schema::from(view_ds.schema()), - )?; - let stream = coerced.execute(0, Arc::new(TaskContext::default()))?; - let txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { fragments } = txn.operation else { - return Err(invalid( - "dedup materialization did not stage append fragments", - )); - }; - written.fragments = fragments; - written.rows = rows.load(std::sync::atomic::Ordering::Relaxed); - Ok(written) -} - -pub(super) async fn commit( - view: &Table, - plan: &VectorDedupPlan, - units: Vec, - expected: Option<&str>, -) -> Result { - let (_, definition, source) = open(view, plan).await?; - let config = definition - .vector_source - .as_ref() - .expect("dedup source") - .config()?; - if plan_duplicate_pairs(source.clone(), &config).await? != plan.native.tasks { - return Err(invalid("dedup plan does not cover the pinned index")); - } - validate_receipts(plan, &units, plan.snapshot.units)?; - let rows = units.iter().map(|u| u.rows).sum(); - let fragments = units.into_iter().flat_map(|u| u.fragments).collect(); - super::grouped_units::commit_fragments( - view, - &plan.snapshot, - &definition, - source.manifest.timestamp_nanos, - fragments, - rows, - expected, - ) - .await -} - -#[cfg(test)] -mod tests { - use super::*; - - fn edges(values: &[(u64, u64)], batch_size: usize) -> SendableRecordBatchStream { - let schema = edge_schema(); - let batches = values - .chunks(batch_size) - .map(|edges| { - Ok(RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.0))), - Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.1))), - ], - ) - .unwrap()) - }) - .collect::>(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::iter(batches), - )) - } - - #[tokio::test] - async fn direct_representatives_handle_chains_orientation_and_uint64() { - let high = u64::MAX - 2; - let pairs = [ - (2, 1), - (1, 0), - (high + 2, high + 1), - (high, high + 1), - (1, 0), - ]; - let (removed, representatives) = - select_representatives(edges(&pairs, 1), SORT_MEMORY_BYTES) - .await - .unwrap(); - assert_eq!(removed.iter().collect::>(), [1, high + 1]); - assert_eq!(representatives.iter().collect::>(), [0, high]); - } - - #[tokio::test] - async fn masks_reject_overlapping_representatives_and_corrupt_artifacts() { - let dir = tempfile::tempdir().unwrap(); - let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::UInt64, false)])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt64Array::from(vec![0, 1, 2]))], - ) - .unwrap(); - let ds = Dataset::write( - arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema), - dir.path().join("source.lance").to_str().unwrap(), - None, - ) - .await - .unwrap(); - let first = persist_masks( - &ds, - &ds, - [1].into_iter().collect(), - [0].into_iter().collect(), - ) - .await - .unwrap(); - let second = persist_masks( - &ds, - &ds, - [2].into_iter().collect(), - [1].into_iter().collect(), - ) - .await - .unwrap(); - let receipt = |masks| WrittenDedupUnit { - unit: 0, - plan_id: "test".into(), - masks, - fragments: vec![], - rows: 0, - }; - let first = receipt(first); - let second = receipt(second); - assert_eq!( - load_masks(&ds, 0, &[&first]).await.unwrap(), - [1].into_iter().collect() - ); - assert!( - load_masks(&ds, 0, &[&first, &second]) - .await - .unwrap_err() - .to_string() - .contains("overlapping") - ); - let file = &first.masks[0]; - ds.object_store(None) - .await - .unwrap() - .put( - &ds.indices_dir() - .join(file.object_id.to_string()) - .join(MASK_FILE), - &vec![0; file.bytes], - ) - .await - .unwrap(); - assert!( - load_masks(&ds, 0, &[&first]) - .await - .unwrap_err() - .to_string() - .contains("checksum") - ); - } - - #[tokio::test] - async fn external_sort_handles_more_edges_than_its_budget() { - let pairs = (0..262144) - .rev() - .map(|i| ((i + 1) as u64, i as u64)) - .collect::>(); - let (removed, reps) = select_representatives(edges(&pairs, 4096), 2 * 1024 * 1024) - .await - .unwrap(); - assert_eq!(removed.len(), 131072); - assert_eq!(reps.len(), 131072); - assert!(removed.iter().all(|id| id % 2 == 1)); - assert!(reps.iter().all(|id| id % 2 == 0)); - } -} diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index c8b19f013..fffb703fb 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1252,10 +1252,17 @@ impl VectorQuery { /// This method sets both the minimum and maximum number of partitions to search. /// For more fine-grained control see [`VectorQuery::minimum_nprobes`] and /// [`VectorQuery::maximum_nprobes`]. - pub fn nprobes(mut self, nprobes: usize) -> Self { + /// + /// Returns an error if `nprobes` is not greater than 0. + pub fn nprobes(mut self, nprobes: usize) -> Result { + if nprobes == 0 { + return Err(Error::InvalidInput { + message: "nprobes must be greater than 0".to_string(), + }); + } self.request.minimum_nprobes = nprobes; self.request.maximum_nprobes = Some(nprobes); - self + Ok(self) } /// Set the minimum number of partitions to search @@ -2361,6 +2368,7 @@ mod tests { .nearest_to(&[9.8, 8.7]) .unwrap() .nprobes(1000) + .unwrap() .postfilter() .distance_type(DistanceType::Cosine) .approx_mode(ApproxMode::Accurate) @@ -2386,6 +2394,31 @@ mod tests { assert_eq!(query.request.refine_factor, Some(999)); } + #[tokio::test] + async fn test_nprobes_validation() { + let conn = connect("memory://").execute().await.unwrap(); + let table = conn + .create_table("my_table", make_test_batches()) + .execute() + .await + .unwrap(); + let query = table.query().nearest_to(&[0.1, 0.2]).unwrap(); + + assert!(matches!( + query.clone().nprobes(0), + 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. + 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)); + } + } + #[test] fn test_approx_mode_serde_parse_default_and_display() { assert_eq!(ApproxMode::default(), ApproxMode::Normal); diff --git a/rust/lancedb/src/remote/catalog.rs b/rust/lancedb/src/remote/catalog.rs index 1c0129889..72b77a87e 100644 --- a/rust/lancedb/src/remote/catalog.rs +++ b/rust/lancedb/src/remote/catalog.rs @@ -416,6 +416,7 @@ mod tests { RemoteCatalog::try_new(&endpoint, options).unwrap(), )); let page = catalog + .catalog() .list_databases(ListDatabasesRequest::default().limit(1).page_token("a/b")) .await .unwrap(); @@ -443,6 +444,7 @@ mod tests { .unwrap(); assert!( catalog + .catalog() .list_databases(ListDatabasesRequest::default()) .await .unwrap() @@ -538,6 +540,128 @@ mod tests { ); } + #[tokio::test] + async fn database_names_follow_pages_lazily() { + use futures::TryStreamExt; + + let (endpoint, task) = server(vec![ + ( + 200, + json!({"namespaces": ["first", "second"], "page_token": "a/b"}), + ), + (200, json!({"namespaces": [], "page_token": "next"})), + (200, json!({"namespaces": ["last"], "page_token": ""})), + ]) + .await; + let catalog = CatalogConnection::new(Arc::new( + RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(), + )); + let mut names = catalog.list_databases(None, None); + drop(catalog); + assert_eq!(names.num_page_results(), 0); + assert_eq!(names.page_token(), None); + assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first")); + assert_eq!(names.num_page_results(), 1); + assert_eq!(names.page_token(), Some("a/b")); + assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second")); + assert_eq!(names.num_page_results(), 0); + assert_eq!(names.page_token(), Some("a/b")); + assert!(!task.is_finished()); + assert_eq!(names.try_next().await.unwrap().as_deref(), Some("last")); + assert_eq!(names.try_next().await.unwrap(), None); + assert_eq!(names.try_next().await.unwrap(), None); + let requests = task.await.unwrap(); + assert_eq!(requests[0].line, "GET /v1/namespace/$/list HTTP/1.1"); + assert_eq!( + requests[1].line, + "GET /v1/namespace/$/list?page_token=a%2Fb HTTP/1.1" + ); + assert_eq!( + requests[2].line, + "GET /v1/namespace/$/list?page_token=next HTTP/1.1" + ); + } + + #[tokio::test] + async fn database_names_stop_after_error() { + use futures::StreamExt; + + let (endpoint, task) = server(vec![ + (200, json!({"namespaces": ["first"], "page_token": "next"})), + (401, json!({"error": "unauthorized"})), + ]) + .await; + let catalog = CatalogConnection::new(Arc::new( + RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(), + )); + let mut names = catalog.list_databases(None, None); + assert_eq!(names.next().await.unwrap().unwrap(), "first"); + assert!(names.next().await.unwrap().is_err()); + assert_eq!(names.num_page_results(), 0); + assert_eq!(names.page_token(), Some("next")); + assert!(names.next().await.is_none()); + task.await.unwrap(); + } + + #[tokio::test] + async fn database_names_can_resume_with_page_limit() { + use futures::TryStreamExt; + let (endpoint, task) = server(vec![ + ( + 200, + json!({"namespaces": ["first", "second"], "page_token": "next/token"}), + ), + (200, json!({"namespaces": ["third", "fourth"]})), + ]) + .await; + let catalog = CatalogConnection::new(Arc::new( + RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(), + )); + let mut names = catalog.list_databases(Some("start/token".into()), Some(2)); + assert_eq!(names.page_token(), Some("start/token")); + assert_eq!(names.num_page_results(), 0); + assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first")); + assert_eq!(names.page_token(), Some("next/token")); + assert_eq!(names.num_page_results(), 1); + assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second")); + assert_eq!(names.num_page_results(), 0); + let token = names.page_token().map(str::to_owned); + drop(names); + let mut resumed = catalog.list_databases(token, Some(2)); + assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("third")); + assert_eq!(resumed.page_token(), None); + assert_eq!(resumed.num_page_results(), 1); + assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("fourth")); + assert_eq!(resumed.num_page_results(), 0); + assert_eq!(resumed.try_next().await.unwrap(), None); + let requests = task.await.unwrap(); + assert_eq!( + requests[0].line, + "GET /v1/namespace/$/list?limit=2&page_token=start%2Ftoken HTTP/1.1" + ); + assert_eq!( + requests[1].line, + "GET /v1/namespace/$/list?limit=2&page_token=next%2Ftoken HTTP/1.1" + ); + } + + #[tokio::test] + async fn database_names_reject_invalid_page_limit() { + use futures::TryStreamExt; + let catalog = CatalogConnection::new(Arc::new( + RemoteCatalog::try_new("http://127.0.0.1:1", RemoteCatalogOptions::default()).unwrap(), + )); + for limit in [0, u32::MAX] { + let mut names = catalog.list_databases(None, Some(limit)); + assert!(matches!( + names.try_next().await, + Err(Error::InvalidInput { .. }) + )); + assert_eq!(names.num_page_results(), 0); + assert_eq!(names.try_next().await.unwrap(), None); + } + } + #[tokio::test] async fn catalog_validates_before_sending_requests() { for endpoint in [ diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 27606cdfe..4f1c3f56c 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -47,6 +47,10 @@ use super::table::RemoteTable; use super::util::parse_server_version; use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id}; +fn quote_sql_identifier(identifier: &str) -> String { + format!("\"{}\"", identifier.replace('"', "\"\"")) +} + // Request structure for the remote clone table API #[derive(serde::Serialize)] struct RemoteCloneTableRequest { @@ -860,42 +864,22 @@ impl Database for RemoteDatabase { &self, request: CreateMaterializedViewRequest, ) -> Result { - let identifier = build_table_identifier(&request.name, &request.namespace_path)?; - let req = self - .client - .post(&format!("/v1/materialized_view/{identifier}/create")) - .json(&serde_json::json!({ - "query": request.query, - "with_no_data": request.with_no_data, - })); - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - let job_id = extract_job_id(&body); - if request.with_no_data { - return Ok(match job_id { - Some(job_id) => Job::new(Box::new(RemoteJob::new(self.client.clone(), job_id))), - None => Job::new_done(), + return Err(Error::NotSupported { + message: "remote materialized views are always populated by their SQL job" + .to_string(), }); } - if status != StatusCode::ACCEPTED { - return Err(Error::Http { - source: "materialized-view creation with data must return 202 Accepted".into(), - request_id, - status_code: Some(status), - }); - } - let job_id = job_id.ok_or_else(|| Error::Http { - source: "materialized-view creation response did not contain a valid job_id".into(), - request_id, - status_code: Some(status), + let client = self.sql_client.clone().ok_or_else(|| Error::NotSupported { + message: "SQL is unavailable for this remote database client".to_string(), })?; - Ok(Job::new(Box::new(RemoteJob::new( - self.client.clone(), - job_id, - )))) + let namespace = request.namespace_path; + let statement = format!( + "CREATE MATERIALIZED VIEW {} AS {}", + quote_sql_identifier(&request.name), + request.query + ); + Ok(client.submit_as_job(statement, namespace, || async { Ok(()) })) } async fn drop_materialized_view_async( @@ -1547,12 +1531,13 @@ impl Database for RemoteDatabase { let version = parse_server_version(&request_id, &rsp)?; let table_identifier = build_table_identifier(&request.name, &request.namespace_path)?; let cache_key = build_cache_key(&request.name, &request.namespace_path); - let table = Arc::new(RemoteTable::new( + let table = Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), request.name.clone(), request.namespace_path.clone(), table_identifier, version.clone(), + self.sql_client.clone(), )); self.table_cache.insert(cache_key, version).await; @@ -1589,12 +1574,13 @@ impl Database for RemoteDatabase { let version = parse_server_version(&request_id, &rsp)?; let cache_key = build_cache_key(&request.target_table_name, &request.target_namespace_path); - let table = Arc::new(RemoteTable::new( + let table = Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), request.target_table_name.clone(), request.target_namespace_path.clone(), table_identifier, version.clone(), + self.sql_client.clone(), )); self.table_cache.insert(cache_key, version).await; @@ -1607,12 +1593,13 @@ impl Database for RemoteDatabase { // Every open gets its own checkout, schema cache, and freshness state. if let Some(version) = self.table_cache.get(&cache_key).await { - Ok(Arc::new(RemoteTable::new( + Ok(Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), request.name, request.namespace_path, identifier, version, + self.sql_client.clone(), ))) } else { // Describe the table to confirm it exists before moving on. @@ -1625,12 +1612,13 @@ impl Database for RemoteDatabase { let rsp = self.client.check_response(&request_id, rsp).await?; let version = parse_server_version(&request_id, &rsp)?; let describe_body = rsp.text().await.ok(); - let table = Arc::new(RemoteTable::new( + let table = Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), request.name.clone(), request.namespace_path.clone(), identifier, version.clone(), + self.sql_client.clone(), )); // This describe already carries the schema, so hand it to the table // instead of making the first schema read fetch it again. A version or @@ -1957,29 +1945,11 @@ mod tests { } #[tokio::test] - async fn test_create_materialized_view_uses_item_route_and_job() { - let db = super::RemoteDatabase::new_mock(|request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/materialized_view/analytics$adults/create" - ); - let body = request - .body() - .and_then(reqwest::Body::as_bytes) - .and_then(|bytes| serde_json::from_slice::(bytes).ok()) - .unwrap(); - assert_eq!( - body["query"], - "SELECT age AS \"age\" FROM \"raw\".\"people\" WHERE age >= 18 LIMIT 10" - ); - assert_eq!(body["with_no_data"], false); - http::Response::builder() - .status(202) - .body(serde_json::json!({"job_id": "j1-mv-create"}).to_string()) - .unwrap() + async fn test_create_materialized_view_requires_sql_client() { + let db = super::RemoteDatabase::new_mock(|request| -> http::Response { + panic!("unexpected REST request: {}", request.url().path()) }); - let job = db + let error = db .create_materialized_view_async(CreateMaterializedViewRequest { name: "adults".into(), namespace_path: vec!["analytics".into()], @@ -1988,8 +1958,8 @@ mod tests { with_no_data: false, }) .await - .unwrap(); - assert_eq!(job.id(), Some("j1-mv-create")); + .unwrap_err(); + assert!(error.to_string().contains("SQL is unavailable")); } #[tokio::test] diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 489158922..e8ddcf6ff 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fs; +use std::future::{Future, pending}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::{Duration, Instant}; @@ -20,12 +21,14 @@ use arrow_schema::{Schema, SchemaRef}; use futures::TryStreamExt; use http::header::{HeaderMap, HeaderName, HeaderValue}; use prost::Message; -use tokio::sync::{Mutex, Notify, OnceCell, mpsc}; +use serde::{Serialize, de::DeserializeOwned}; +use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch}; use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}; use uuid::Uuid; use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; use crate::error::{Error, Result}; +use crate::job::Job; use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; @@ -209,6 +212,42 @@ impl SqlClient { .await } + pub(super) fn submit_as_job( + &self, + statement: String, + namespace: Vec, + finish: F, + ) -> Job + where + T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { + let client = self.clone(); + let (query_tx, query_rx) = watch::channel(None::>); + let task = tokio::spawn(async move { + let query = Arc::new(client.submit(&statement, &namespace).await?); + query_tx.send_replace(Some(query.clone())); + let mut reader = query.reader().await?; + while reader.try_next().await?.is_some() {} + finish().await + }); + Job::spawned_with_cancellation(task, move || { + let mut query_rx = query_rx.clone(); + async move { + loop { + let query = query_rx.borrow().clone(); + if let Some(query) = query { + return query.cancel().await; + } + if query_rx.changed().await.is_err() { + pending::<()>().await; + } + } + } + }) + } + pub(super) async fn describe(&self, query_id: Uuid) -> Result { let query = self .queries diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index d20ba495f..0d7efc9c4 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -21,7 +21,10 @@ use super::*; use crate::database::Database; use crate::remote::RemoteCatalogOptions; use crate::remote::client::HeaderProvider; +use crate::remote::client::test_utils::client_with_handler; use crate::remote::db::RemoteDatabase; +use crate::remote::table::RemoteTable; +use crate::table::BaseTable; #[derive(Debug, Default)] struct DelayedHeaderProvider { @@ -73,6 +76,7 @@ struct TestSqlService { first_continuation_count: Arc, transient_poll_failures: Arc, headers: Arc>>, + gate_refresh_release: Arc, result: RecordBatch, large_result: RecordBatch, dictionary_result: RecordBatch, @@ -121,6 +125,7 @@ impl Default for TestSqlService { first_continuation_count: Arc::new(AtomicUsize::new(0)), transient_poll_failures: Arc::new(AtomicUsize::new(0)), headers: Arc::new(std::sync::Mutex::new(Vec::new())), + gate_refresh_release: Arc::new(Notify::new()), result, large_result, dictionary_result, @@ -279,6 +284,7 @@ impl FlightService for TestSqlService { ) -> std::result::Result::DoGetStream>, Status> { self.do_get_count.fetch_add(1, Ordering::SeqCst); let ticket = request.get_ref().ticket.as_ref(); + let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; let large = ticket == b"SELECT large message"; @@ -290,7 +296,11 @@ impl FlightService for TestSqlService { self.result.clone() }; let schema = result.schema(); + let gate_refresh_release = self.gate_refresh_release.clone(); let input = futures::stream::once(async move { + if gate_refresh { + gate_refresh_release.notified().await; + } if slow { tokio::time::sleep(Duration::from_millis(250)).await; } @@ -446,6 +456,147 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { server.await.unwrap().unwrap(); } +#[tokio::test] +async fn remote_refresh_twice_keeps_latest_selector() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let incoming = futures::stream::try_unfold(listener, |listener| async { + let (socket, _) = listener.accept().await?; + Ok::<_, std::io::Error>(Some((socket, listener))) + }); + let service = TestSqlService::default(); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn( + tonic::transport::Server::builder() + .add_service(FlightServiceServer::new(service)) + .serve_with_incoming_shutdown(incoming, async { + let _ = shutdown_rx.await; + }), + ); + let sql_client = SqlClient::new( + "analytics".into(), + None, + "test-key".into(), + None, + Some(format!("grpc://{address}")), + ClientConfig::default(), + ); + let selectors = Arc::new(std::sync::Mutex::new(Vec::new())); + let captured = selectors.clone(); + let client = client_with_handler(move |request| { + let body = request + .body() + .and_then(reqwest::Body::as_bytes) + .and_then(|bytes| serde_json::from_slice::(bytes).ok()) + .unwrap(); + captured.lock().unwrap().push(body["version"].clone()); + let text = match request.url().path() { + "/v1/table/ns$view/describe/" => r#"{"version": 42, "schema": {"fields": []}}"#, + "/v1/table/ns$view/count_rows/" => "2", + other => panic!("unexpected REST request: {other}"), + }; + http::Response::builder() + .status(200) + .body(text.to_string()) + .unwrap() + }); + let table = RemoteTable::new_with_sql_client( + client, + "view".into(), + vec!["ns".into()], + "ns$view".into(), + Default::default(), + Some(sql_client), + ); + + let first = table + .refresh_materialized_view_async(None) + .await + .unwrap() + .wait() + .await + .unwrap(); + let second = table + .refresh_materialized_view_async(None) + .await + .unwrap() + .wait() + .await + .unwrap(); + + assert_eq!(first.version, 42); + assert_eq!(second.version, 42); + assert!( + selectors + .lock() + .unwrap() + .iter() + .all(serde_json::Value::is_null) + ); + shutdown_tx.send(()).unwrap(); + server.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn cancelling_remote_refresh_cancels_sql_query() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let incoming = futures::stream::try_unfold(listener, |listener| async { + let (socket, _) = listener.accept().await?; + Ok::<_, std::io::Error>(Some((socket, listener))) + }); + let service = TestSqlService::default(); + let cancel_count = service.cancel_count.clone(); + let do_get_count = service.do_get_count.clone(); + let gate_refresh_release = service.gate_refresh_release.clone(); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn( + tonic::transport::Server::builder() + .add_service(FlightServiceServer::new(service)) + .serve_with_incoming_shutdown(incoming, async { + let _ = shutdown_rx.await; + }), + ); + let sql_client = SqlClient::new( + "analytics".into(), + None, + "test-key".into(), + None, + Some(format!("grpc://{address}")), + ClientConfig::default(), + ); + let client = client_with_handler(|request| -> http::Response { + panic!( + "cancelled refresh must not read result metadata: {}", + request.url().path() + ); + }); + let table = RemoteTable::new_with_sql_client( + client, + "cancel_view".into(), + vec![], + "cancel_view".into(), + Default::default(), + Some(sql_client), + ); + let job = table.refresh_materialized_view_async(None).await.unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + while do_get_count.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + job.cancel().await.unwrap(); + assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. }))); + assert_eq!(cancel_count.load(Ordering::SeqCst), 1); + + gate_refresh_release.notify_one(); + shutdown_tx.send(()).unwrap(); + server.await.unwrap().unwrap(); +} + #[tokio::test] async fn submits_polls_fetches_cancels_and_reuses_client() { let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 16c546147..cd9e9ab64 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -8,6 +8,7 @@ use self::insert::{RemoteWriteExec, WriteOp}; use super::client::RequestResultExt; use super::client::{HttpSend, RestfulLanceDbClient, Sender}; use super::db::ServerVersion; +use super::sql::SqlClient; use super::{ARROW_FILE_CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE, extract_job_id}; use crate::blob::BlobFile; use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions}; @@ -18,7 +19,7 @@ use crate::index::scalar::FtsQuery; use crate::index::waiter::wait_for_index; use crate::job::Job; use crate::materialized_view::{ - MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult, ViewProjection, + MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult, }; use crate::query::wal_fusion::PkFusionMemory; // WAL-PK-FUSION: delete. use crate::query::{QueryFilter, QueryRequest, Select, VectorQueryRequest}; @@ -264,7 +265,6 @@ struct FreshnessJob { enum TrackedJobResult { None, RefreshColumn, - MaterializedView, } #[async_trait] @@ -305,11 +305,6 @@ impl crate::job::JobHandle for FreshnessJob { }) }) }), - TrackedJobResult::MaterializedView => result.value().and_then(|value| { - serde_json::from_value::(value.clone()) - .ok() - .map(|result| result.version) - }), } .filter(|version| *version != 0); if let Some(version) = result_version { @@ -347,6 +342,10 @@ fn compute_min_timestamp( } } +fn quote_sql_identifier(identifier: &str) -> String { + format!("\"{}\"", identifier.replace('"', "\"\"")) +} + fn freshness_headers_snapshot( freshness: &Mutex, interval: Option, @@ -496,6 +495,7 @@ pub struct RemoteTable { namespace: Vec, identifier: String, server_version: ServerVersion, + sql_client: Option, version: Arc>>, location: RwLock>, @@ -654,12 +654,13 @@ impl RemoteTable { Ok(job_id) } - pub fn new( + pub(super) fn new_with_sql_client( client: RestfulLanceDbClient, name: String, namespace: Vec, identifier: String, server_version: ServerVersion, + sql_client: Option, ) -> Self { Self { client, @@ -667,6 +668,7 @@ impl RemoteTable { namespace, identifier, server_version, + sql_client, version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), @@ -701,6 +703,7 @@ impl RemoteTable { namespace: self.namespace.clone(), identifier: self.identifier.clone(), server_version: self.server_version.clone(), + sql_client: self.sql_client.clone(), version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), @@ -1619,6 +1622,7 @@ mod test_utils { namespace: vec![], identifier: name, server_version: version.map(ServerVersion).unwrap_or_default(), + sql_client: None, version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), @@ -1644,6 +1648,7 @@ mod test_utils { namespace: vec![], identifier: name, server_version: ServerVersion::default(), + sql_client: None, version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), @@ -1678,6 +1683,7 @@ mod test_utils { namespace: vec![], identifier: name, server_version: version.map(ServerVersion).unwrap_or_default(), + sql_client: None, version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), @@ -2106,29 +2112,9 @@ impl BaseTable for RemoteTable { &self.identifier } async fn materialized_view_info(&self) -> Result { - #[derive(Deserialize)] - struct Projection { - output_column: String, - expression: String, - } - #[derive(Deserialize)] struct DescribeMaterializedViewResponse { - source_table: String, - #[serde(default)] - source_namespace: Vec, - #[serde(default)] - projections: Vec, - #[serde(default)] - filter: Option, - #[serde(default)] - limit: Option, - /// The defining query, once the server describes a view by it; - /// takes precedence over the structured fields. - #[serde(default)] - query: Option, - #[serde(default)] - incarnation: Option, + query: String, } let request = self.client.post(&format!( @@ -2139,79 +2125,64 @@ impl BaseTable for RemoteTable { let response = self.check_table_response(&request_id, response).await?; let response: DescribeMaterializedViewResponse = response.json().await.err_to_http(request_id)?; - let definition = match response.query { - Some(query) => MaterializedViewDefinition::from_sql(&query)?, - None => MaterializedViewDefinition { - vector_source: None, - source_table: response.source_table, - source_namespace: response.source_namespace, - lateral: None, - projections: response - .projections - .into_iter() - .map(|projection| ViewProjection { - output: projection.output_column, - expression: projection.expression, - }) - .collect(), - filter: response.filter, - group_by: Vec::new(), - limit: response.limit, - }, - }; + let parsed_definition = MaterializedViewDefinition::from_sql(&response.query).ok(); Ok(MaterializedViewInfo { - definition, - incarnation: response.incarnation, + definition_sql: response.query, + parsed_definition, }) } async fn refresh_materialized_view_async( &self, - full: bool, source_version: Option, - expected_incarnation: Option<&str>, ) -> Result> { self.check_mutable().await?; - let mut body = serde_json::json!({ "full": full }); - if let Some(source_version) = source_version { - body["source_version"] = source_version.into(); - } - if let Some(expected_incarnation) = expected_incarnation { - body["expected_incarnation"] = expected_incarnation.into(); - } - let request = self - .client - .post(&format!( - "/v1/materialized_view/{}/refresh", - self.identifier - )) - .json(&body); - let freshness_request = self.snapshot_freshness_headers(); - let (request_id, response) = self - .send_with_freshness(request, true, freshness_request) - .await?; - let response = self.check_table_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - if status != StatusCode::ACCEPTED { - return Err(Error::Http { - source: "materialized-view refresh must return 202 Accepted".into(), - request_id, - status_code: Some(status), + if source_version.is_some() { + return Err(Error::NotSupported { + message: "source-version pinning is available only for local materialized views" + .to_string(), }); } - let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { - source: "materialized-view refresh response did not contain a valid job_id".into(), - request_id, - status_code: Some(status), + let sql_client = self.sql_client.clone().ok_or_else(|| Error::NotSupported { + message: "SQL is unavailable for this remote table client".to_string(), })?; - Ok(Job::new_typed(Box::new(FreshnessJob { - inner: RemoteJob::new(self.client.clone(), job_id), - freshness: self.freshness.clone(), - version: self.version.clone(), - tracked_result: TrackedJobResult::MaterializedView, - freshness_request, - }))) + let statement = format!( + "REFRESH MATERIALIZED VIEW {}", + quote_sql_identifier(&self.name) + ); + let namespace = self.namespace.clone(); + let client = self.client.clone(); + let name = self.name.clone(); + let identifier = self.identifier.clone(); + let server_version = self.server_version.clone(); + let freshness = self.freshness.clone(); + let freshness_request = self.snapshot_freshness_headers(); + let job_sql_client = sql_client.clone(); + Ok( + sql_client.submit_as_job(statement, namespace.clone(), move || async move { + let table = Self::new_with_sql_client( + client, + name, + namespace, + identifier, + server_version, + Some(job_sql_client), + ); + table.checkout_latest().await?; + let version = table.version().await?; + let rows_written = + u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime { + message: "materialized-view row count exceeds u64".to_string(), + })?; + freshness_request.observe_version(&freshness, version); + Ok(RefreshMaterializedViewResult { + mode: crate::materialized_view::RefreshMode::Rebuild, + rows_written, + source_version: 0, + version, + }) + }), + ) } async fn query_snapshot(&self) -> Result> { let description = self.describe().await?; @@ -3101,7 +3072,13 @@ impl BaseTable for RemoteTable { // loop can re-execute the plan (and re-stream the body) on each retry. // This mirrors the old `send_streaming(with_retry=true)` path, which // likewise buffered the reader to support retries. - let batches = new_data.collect::, _>>()?; + let schema = RecordBatchReader::schema(new_data.as_ref()); + let mut batches = new_data.collect::, _>>()?; + // An empty reader still carries a schema. Keep it in an empty batch so + // the buffered source remains scannable and can be replayed on retries. + if batches.is_empty() { + batches.push(RecordBatch::new_empty(schema)); + } let source: Box = Box::new(batches); let rescannable = source.rescannable(); let input: Arc = @@ -3517,7 +3494,12 @@ impl BaseTable for RemoteTable { let (request_id, response) = self .send_with_freshness(request, true, freshness_request) .await?; - let response = self.check_table_response(&request_id, response).await?; + // A Function declaration can return 404 for the Function rather than the table. + let response = self + .client + .check_response(&request_id, response) + .await + .inspect_err(|error| self.handle_error_invalidation(error))?; let body = response.text().await.err_to_http(request_id.clone())?; if body.trim().is_empty() { @@ -3557,7 +3539,12 @@ impl BaseTable for RemoteTable { .post(&format!("/v1/table/{}/backfill_column", self.identifier)) .json(&body); let (request_id, response) = self.send(request, true).await?; - let response = self.check_table_response(&request_id, response).await?; + // Preserve dependency errors: a deleted bound Function also returns 404. + let response = self + .client + .check_response(&request_id, response) + .await + .inspect_err(|error| self.handle_error_invalidation(error))?; let body = response.text().await.err_to_http(request_id.clone())?; #[derive(serde::Deserialize)] @@ -4051,7 +4038,7 @@ mod tests { )) }; - // All endpoints should translate 404 to TableNotFound. + // These table operations should translate 404 to TableNotFound. let results: Vec>> = vec![ Box::pin(table.version().map_ok(|_| ())), Box::pin(table.schema().map_ok(|_| ())), @@ -4688,6 +4675,25 @@ mod tests { assert_eq!(result.rows_updated, if old_server { 0 } else { 5 }); } + #[tokio::test] + async fn test_alter_columns_rejects_missing_changes_before_request() { + let table = Table::new_with_handler::("my_table", |request| { + panic!("Unexpected request: {}", request.url().path()) + }); + + for alterations in [ + vec![ColumnAlteration::new("id".into())], + vec![ + ColumnAlteration::new("id".into()).rename("new_id".into()), + ColumnAlteration::new("id".into()), + ], + ] { + let err = table.alter_columns(&alterations).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("path 'id'")); + } + } + #[rstest] #[case(true)] #[case(false)] @@ -4801,6 +4807,104 @@ mod tests { } } + #[rstest] + #[case::no_batches_insert(false, false)] + #[case::empty_batch_insert(true, false)] + #[case::no_batches_delete(false, true)] + #[case::empty_batch_delete(true, true)] + #[tokio::test] + async fn test_merge_insert_empty_source( + #[case] has_batch: bool, + #[case] delete_unmatched: bool, + ) { + let mut fields = vec![Field::new("id", DataType::Int64, false)]; + if !delete_unmatched { + fields.extend([ + Field::new("k", DataType::Int64, false), + Field::new( + "vector", + DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Float32, true)), + 2, + ), + true, + ), + Field::new("s", DataType::Utf8, true), + ]); + } + let schema = Arc::new(Schema::new_with_metadata( + fields, + HashMap::from([("source".to_string(), "empty".to_string())]), + )); + let batches = if has_batch { + vec![Ok(RecordBatch::new_empty(schema.clone()))] + } else { + vec![] + }; + let data: Box = + Box::new(RecordBatchIterator::new(batches, schema.clone())); + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_ref = attempts.clone(); + let num_deleted_rows = if delete_unmatched { 3 } else { 0 }; + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); + assert_eq!(request.headers()[CONTENT_TYPE], ARROW_STREAM_CONTENT_TYPE); + let params = request.url().query_pairs().collect::>(); + assert_eq!(params["on"], "id"); + assert_eq!( + params["when_not_matched_insert_all"], + (!delete_unmatched).to_string() + ); + assert_eq!( + params["when_not_matched_by_source_delete"], + delete_unmatched.to_string() + ); + + let body = request.body().unwrap().as_bytes().unwrap(); + let reader = StreamReader::try_new(Cursor::new(body), None).unwrap(); + assert_eq!(reader.schema(), schema); + for batch in reader { + assert_eq!(batch.unwrap().num_rows(), 0); + } + + // The empty source must retain its schema when replayed after a conflict. + if attempts_ref.fetch_add(1, Ordering::SeqCst) == 0 { + http::Response::builder() + .status(409) + .body(String::new()) + .unwrap() + } else { + http::Response::builder() + .status(200) + .body( + json!({ + "version": 43, + "num_deleted_rows": num_deleted_rows, + "num_inserted_rows": 0, + "num_updated_rows": 0, + }) + .to_string(), + ) + .unwrap() + } + }); + + let mut merge = table.merge_insert(&["id"]); + if delete_unmatched { + merge.when_not_matched_by_source_delete(None); + } else { + merge.when_not_matched_insert_all(); + } + let result = merge.execute(data).await.unwrap(); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert_eq!(result.version, 43); + assert_eq!(result.num_deleted_rows, num_deleted_rows); + assert_eq!(result.num_inserted_rows, 0); + assert_eq!(result.num_updated_rows, 0); + } + #[tokio::test] async fn test_merge_insert_composite_key() { let batch = RecordBatch::try_new( @@ -5937,6 +6041,22 @@ mod tests { assert_eq!(data[0].as_ref().unwrap(), &expected_data); } + #[tokio::test] + async fn test_query_vector_nprobes_zero() { + let table = Table::new_with_handler::<&str>("my_table", |_| { + panic!("invalid nprobes must be rejected before sending a request") + }); + let result = table + .query() + .nearest_to(vec![0.1, 0.2, 0.3]) + .unwrap() + .nprobes(0); + assert!(matches!( + result, + Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0" + )); + } + #[tokio::test] async fn test_query_vector_all_params() { let table = Table::new_with_handler("my_table", |request| { @@ -6010,6 +6130,7 @@ mod tests { .postfilter() .distance_type(crate::DistanceType::Cosine) .nprobes(12) + .unwrap() .refine_factor(2) .bypass_vector_index() .execute() @@ -8653,6 +8774,73 @@ mod tests { ); } + #[rstest] + #[case(false, "Function is unavailable")] + #[case(true, "Function name was not found")] + #[tokio::test] + async fn test_function_column_operations_preserve_dependency_not_found( + #[case] declare_column: bool, + #[case] message: &'static str, + ) { + let schema_requests = Arc::new(AtomicUsize::new(0)); + let requests = schema_requests.clone(); + let table = Table::new_with_handler("my_table", move |request| { + match request.url().path() { + "/v1/table/my_table/describe/" => { + requests.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body( + r#"{"version":1,"schema":{"fields":[{"name":"description","nullable":true,"type":{"type":"string"}}]}}"#.to_string(), + ) + .unwrap() + } + "/v1/table/my_table/backfill_column" | "/v1/table/my_table/add_columns/" => { + http::Response::builder() + .status(404) + .body(json!({"code": 4, "error": message}).to_string()) + .unwrap() + } + path => panic!("unexpected request: {path}"), + } + }); + table.schema().await.unwrap(); + assert_eq!(schema_requests.load(Ordering::SeqCst), 1); + + let error = if declare_column { + let fixture: serde_json::Value = serde_json::from_str(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_fixed_size_declaration_request.json" + )) + .unwrap(); + let application = crate::function::FunctionApplication::from_json( + &fixture["function"]["application"].to_string(), + ) + .unwrap(); + table + .add_columns() + .function_as("embedding", application) + .execute() + .await + .unwrap_err() + } else { + table.refresh_column_async("embedding").await.unwrap_err() + }; + let Error::Http { + source, + status_code, + request_id, + } = error + else { + panic!("dependency 404 was misclassified: {error:?}"); + }; + assert_eq!(status_code, Some(StatusCode::NOT_FOUND)); + assert!(source.to_string().contains(message)); + assert!(!request_id.is_empty()); + + table.schema().await.unwrap(); + assert_eq!(schema_requests.load(Ordering::SeqCst), 2); + } + /// The error listing is table-addressed with optional job and column /// filters, mirroring the server's SQL surface, and the two non-record /// signals come back as their own fields rather than as rows. @@ -12759,54 +12947,16 @@ mod tests { } #[tokio::test] - async fn test_materialized_view_describe_and_refresh() { + async fn test_materialized_view_describe_and_refresh_requires_sql() { + const QUERY: &str = + "SELECT s.x, d.label FROM analytics.source s JOIN analytics.dim d ON s.id = d.id"; let table = Table::new_with_handler("my_table", |request| match request.url().path() { "/v1/materialized_view/my_table/describe" => http::Response::builder() .status(200) .body( json!({ "name": "my_table", - "source_table": "source", - "source_namespace": ["analytics"], - "projections": [{ - "output_column": "double_x", - "expression": "x * 2" - }], - "filter": "x > 0", - "limit": 10, - "inputs": ["x"], - "incarnation": "inc-1" - }) - .to_string(), - ) - .unwrap(), - "/v1/materialized_view/my_table/refresh" => { - assert_eq!(request.method(), "POST"); - assert_eq!( - request_body_json(&request), - json!({ - "full": true, - "source_version": 7, - "expected_incarnation": "inc-1" - }) - ); - http::Response::builder() - .status(202) - .body(json!({"job_id": "j1-mv-refresh"}).to_string()) - .unwrap() - } - "/v1/jobs/describe" => http::Response::builder() - .status(200) - .body( - json!({ - "job_id": "j1-mv-refresh", - "job_state": "DONE", - "result": { - "mode": "rebuild", - "rows_written": 2, - "source_version": 7, - "version": 9 - } + "query": QUERY }) .to_string(), ) @@ -12814,21 +12964,10 @@ mod tests { path => panic!("unexpected request: {path}"), }); let view = crate::MaterializedView::from_table(table).await.unwrap(); - assert_eq!(view.definition().source_table, "source"); - assert_eq!(view.definition().source_namespace, ["analytics"]); - assert_eq!(view.incarnation(), Some("inc-1")); - - let result = view - .refresh() - .full(true) - .source_version(7) - .expect_incarnation("inc-1") - .execute() - .await - .unwrap(); - assert_eq!(result.mode, crate::RefreshMode::Rebuild); - assert_eq!(result.rows_written, 2); - assert_eq!(result.version, 9); + assert_eq!(view.definition_sql(), QUERY); + assert!(view.definition().is_err()); + let error = view.refresh().execute().await.unwrap_err(); + assert!(error.to_string().contains("SQL is unavailable")); } #[tokio::test] diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index 07aeeae63..ac37bf7ae 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -3,23 +3,28 @@ //! Cloud blob column listing, whole-byte fetch, and seekable HTTP byte-range handles. +use std::collections::HashMap; use std::ops::Range; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; -use arrow_array::{Array, LargeBinaryArray}; +use arrow_array::{Array, LargeBinaryArray, StructArray, UInt8Array, UInt64Array}; use arrow_schema::DataType; use bytes::{Bytes, BytesMut}; +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 crate::Error; use crate::blob::BlobFile; use crate::error::Result; +use crate::query::{QueryFilter, QueryRequest, Select}; use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; -use crate::table::BaseTable; +use crate::table::{AnyQuery, BaseTable}; use super::{ FreshnessHeaders, FreshnessState, ReadSnapshot, RemoteTable, VERSION_HEADER, @@ -520,7 +525,11 @@ impl RemoteTable { /// Open seekable handles for `row_ids`. /// - /// Each non-null row is probed once to determine its size. + /// Descriptors for the requested rows come from one `_rowid` take query per + /// [`MAX_FETCH_BLOBS_ROW_IDS`] unique row ids, all read at one exact table + /// version. Each handle is sized from its descriptor and reads that + /// version. Only external blobs whose descriptor records no size are + /// probed, because their length is the whole object's. pub(super) async fn fetch_blob_files_impl( &self, column: &str, @@ -530,46 +539,251 @@ impl RemoteTable { if row_ids.is_empty() { return Ok(Vec::new()); } - if !self.server_version.support_blobs() { - return Err(Error::NotSupported { - message: "fetch_blob_files requires LanceDB Cloud server 0.5.0 or newer".into(), - }); - } + self.ensure_blob_files_supported()?; - let read_snapshot = self.snapshot_read_state().await; - self.fetch_blob_files_with_snapshot(column, row_ids, read_snapshot) - .await + let mut read_snapshot = self.snapshot_read_state().await; + let sizes = self + .take_blob_descriptor_sizes(column, row_ids, &mut read_snapshot) + .await?; + let version = read_snapshot.version; + + let mut files = Vec::with_capacity(row_ids.len()); + let mut probe_indices = Vec::new(); + let mut probe_requesters = Vec::new(); + for (index, row_id) in row_ids.iter().enumerate() { + match sizes[row_id] { + DescriptorSize::Null => files.push(None), + DescriptorSize::Known(size) => { + let requester = self.blob_range_requester(column, *row_id, &read_snapshot); + files.push(Some(RemoteBlobFile::new(requester, size, version).into())); + } + DescriptorSize::Unresolved => { + files.push(None); + probe_indices.push(index); + probe_requesters.push(self.blob_range_requester( + column, + *row_id, + &read_snapshot, + )); + } + } + } + let probed = probe_blob_files(probe_requesters).await?; + for (index, file) in probe_indices.into_iter().zip(probed) { + files[index] = file; + } + Ok(files) } + /// Probe each row once for its size. + /// + /// Serves callers that must not issue a query, such as a single blob that + /// the whole-byte route could not return. async fn fetch_blob_files_with_snapshot( &self, column: &str, row_ids: &[u64], read_snapshot: ReadSnapshot, ) -> Result>> { - let encoded_column = urlencoding::encode(column); let requesters = row_ids .iter() - .map(|row_id| { - let path = format!( - "/v1/table/{}/blob/{encoded_column}/{row_id}/bytes", - self.identifier - ); - let requester: Arc = Arc::new(TableBlobRangeRequester { - client: self.client.clone(), - path, - version: read_snapshot.version, - branch: self.branch.clone(), - freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)), - parent_freshness: self.freshness.clone(), - parent_freshness_request: read_snapshot.freshness, - read_consistency_interval: self.client.read_consistency_interval, - }); - requester - }) + .map(|row_id| self.blob_range_requester(column, *row_id, &read_snapshot)) .collect(); probe_blob_files(requesters).await } + + /// Read the blob descriptor of every unique row in `row_ids`. + /// + /// Pins `read_snapshot` to the version the first take read, so later chunks + /// and the returned handles all see one table version. + async fn take_blob_descriptor_sizes( + &self, + column: &str, + row_ids: &[u64], + read_snapshot: &mut ReadSnapshot, + ) -> Result> { + let mut unique_row_ids = row_ids.to_vec(); + unique_row_ids.sort_unstable(); + unique_row_ids.dedup(); + + let mut sizes = HashMap::with_capacity(unique_row_ids.len()); + for chunk in unique_row_ids.chunks(MAX_FETCH_BLOBS_ROW_IDS) { + let (mut chunk_sizes, response_version) = self + .take_blob_descriptor_chunk(column, chunk, *read_snapshot) + .await?; + if read_snapshot.version.is_none() { + match response_version { + Some(version) => read_snapshot.version = Some(version), + None => { + // A server that does not report its read version may have read + // an older version than a describe now returns, so the take is + // repeated at the described version. + read_snapshot.version = + Some(self.describe_read_snapshot(*read_snapshot).await?.version); + (chunk_sizes, _) = self + .take_blob_descriptor_chunk(column, chunk, *read_snapshot) + .await?; + } + } + } + sizes.extend(chunk_sizes); + } + let found = unique_row_ids + .iter() + .filter(|row_id| sizes.contains_key(row_id)) + .count(); + if found < unique_row_ids.len() { + return Err(Error::InvalidInput { + message: format!( + "blob read for column '{column}' requested {} row ids but only {found} exist \ + in the table; pass row ids collected from this table", + unique_row_ids.len() + ), + }); + } + Ok(sizes) + } + + /// Take one chunk of descriptors, also returning the version the server + /// reported reading. + async fn take_blob_descriptor_chunk( + &self, + column: &str, + row_ids: &[u64], + read_snapshot: ReadSnapshot, + ) -> Result<(HashMap, Option)> { + let query = AnyQuery::Query(QueryRequest { + filter: Some(QueryFilter::Datafusion( + col(ROW_ID).in_list(row_ids.iter().map(|id| lit(*id)).collect(), false), + )), + select: Select::columns(&[column]), + with_row_id: true, + // Row ids name base-table rows, and the MemWAL scanner has no stable `_rowid`. + use_lsm: Some(false), + ..Default::default() + }); + let body = self + .prepare_query_bodies(&query, read_snapshot.version)? + .pop() + .expect("a plain query has one body"); + let request = self + .client + .post(&format!("/v1/table/{}/query/", self.identifier)) + .json(&body); + let (request_id, response) = self + .send_with_freshness(request, true, read_snapshot.freshness) + .await?; + let response_version = response + .headers() + .get(&VERSION_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()); + let mut stream = self.read_arrow_response(&request_id, response).await?; + + let mut sizes = HashMap::with_capacity(row_ids.len()); + let invalid_response = |message: String| Error::Http { + source: message.into(), + request_id: request_id.clone(), + status_code: None, + }; + while let Some(batch) = stream.try_next().await? { + let taken_row_ids = batch + .column_by_name(ROW_ID) + .and_then(|array| array.as_any().downcast_ref::()) + .ok_or_else(|| { + invalid_response(format!("blob descriptor take is missing UInt64 '{ROW_ID}'")) + })?; + let values = batch.column_by_name(column).ok_or_else(|| { + invalid_response(format!("blob descriptor take is missing '{column}'")) + })?; + // Only blob v2 columns read back as descriptors. + let descriptors = values + .as_any() + .downcast_ref::() + .and_then(BlobDescriptorSizes::try_new) + .ok_or_else(|| Error::InvalidInput { + message: format!("column '{column}' is not a blob column"), + })?; + for (index, row_id) in taken_row_ids.values().iter().enumerate() { + sizes.insert(*row_id, descriptors.size(index)?); + } + } + Ok((sizes, response_version)) + } + + fn ensure_blob_files_supported(&self) -> Result<()> { + if self.server_version.support_blobs() { + Ok(()) + } else { + Err(Error::NotSupported { + message: "fetch_blob_files requires LanceDB Cloud server 0.5.0 or newer".into(), + }) + } + } + + fn blob_range_requester( + &self, + column: &str, + row_id: u64, + read_snapshot: &ReadSnapshot, + ) -> Arc { + let path = format!( + "/v1/table/{}/blob/{}/{row_id}/bytes", + self.identifier, + urlencoding::encode(column) + ); + Arc::new(TableBlobRangeRequester { + client: self.client.clone(), + path, + version: read_snapshot.version, + branch: self.branch.clone(), + freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)), + parent_freshness: self.freshness.clone(), + parent_freshness_request: read_snapshot.freshness, + read_consistency_interval: self.client.read_consistency_interval, + }) + } +} + +/// What a blob v2 descriptor says about a row's size. +#[derive(Debug, Clone, Copy)] +enum DescriptorSize { + Null, + Known(u64), + /// An external blob that recorded no size. Its length is the whole + /// object's, which only the server can resolve. + Unresolved, +} + +/// The `kind` and `size` children of a blob v2 descriptor array. +struct BlobDescriptorSizes<'a> { + descriptors: &'a StructArray, + kinds: &'a UInt8Array, + sizes: &'a UInt64Array, +} + +impl<'a> BlobDescriptorSizes<'a> { + fn try_new(descriptors: &'a StructArray) -> Option { + let child = |name: &str| descriptors.column_by_name(name); + Some(Self { + descriptors, + kinds: child("kind")?.as_any().downcast_ref::()?, + sizes: child("size")?.as_any().downcast_ref::()?, + }) + } + + fn size(&self, index: usize) -> Result { + if self.descriptors.is_null(index) || self.kinds.is_null(index) { + return Ok(DescriptorSize::Null); + } + let kind = BlobKind::try_from(self.kinds.value(index))?; + let size = self.sizes.value(index); + Ok(if kind == BlobKind::External && size == 0 { + DescriptorSize::Unresolved + } else { + DescriptorSize::Known(size) + }) + } } /// Probe blob sizes while preserving row order. @@ -689,6 +903,8 @@ mod tests { use reqwest::Request; use semver::Version; + use arrow_array::RecordBatch; + use super::*; const PAYLOAD: &[u8] = b"0123456789abcdefghijklmnopqrstuvwxyz"; @@ -772,10 +988,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); - let mut files = table - .fetch_blob_files_impl("image", &[10, 20]) - .await - .unwrap(); + let mut files = table.probe_files("image", &[10, 20]).await.unwrap(); assert_eq!(files.len(), 2); assert!(files[1].is_none()); @@ -787,12 +1000,368 @@ mod tests { ); } + /// Descriptors as a remote query returns them, one `(kind, size)` per row. + fn descriptors(rows: &[Option<(BlobKind, u64)>]) -> StructArray { + let kinds = rows + .iter() + .map(|row| row.map_or(0, |(kind, _)| kind as u8)) + .collect::(); + let sizes = rows + .iter() + .map(|row| row.map_or(0, |(_, size)| size)) + .collect::(); + let len = rows.len(); + StructArray::new( + lance_core::datatypes::BLOB_V2_DESC_FIELDS.clone(), + vec![ + Arc::new(kinds), + Arc::new(UInt64Array::from(vec![0; len])), + Arc::new(sizes), + Arc::new(arrow_array::UInt32Array::from(vec![0; len])), + Arc::new(arrow_array::StringArray::from(vec![""; len])), + ], + Some(rows.iter().map(Option::is_some).collect()), + ) + } + + impl RemoteTable { + /// Open handles through the per-row size probe alone. + async fn probe_files( + &self, + column: &str, + row_ids: &[u64], + ) -> Result>> { + let read_snapshot = self.snapshot_read_state().await; + self.fetch_blob_files_with_snapshot(column, row_ids, read_snapshot) + .await + } + } + + /// A request the take mock received. + #[derive(Debug, Clone, PartialEq)] + enum TakeMockRequest { + Query { + row_ids: Vec, + version: Option, + }, + Describe, + Range { + row_id: u64, + range: String, + version: Option, + }, + } + + /// Row ids named by a `_rowid IN (...)` filter. + fn filter_row_ids(filter: &str) -> Vec { + let list = filter.split_once('(').unwrap().1.trim_end_matches(')'); + list.split(',') + .map(|id| id.trim().parse().unwrap()) + .collect() + } + + /// A table whose query route answers `_rowid` takes from `rows` and whose + /// byte routes serve [`PAYLOAD`]. Takes return rows in reverse order, and + /// report `query_version` in `x-lancedb-version` when it is set. + fn mock_take_blob_table( + rows: Vec<(u64, Option<(BlobKind, u64)>)>, + query_version: Option<&'static str>, + ) -> ( + RemoteTable, + Arc>>, + ) { + let requests = Arc::new(StdMutex::new(Vec::new())); + let captured = requests.clone(); + let table = RemoteTable::new_mock( + "my_table".to_string(), + move |request| { + let path = request.url().path().to_string(); + if path == "/v1/table/my_table/query/" { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["columns"], serde_json::json!(["image"])); + assert_eq!(body["with_row_id"], serde_json::json!(true)); + assert_eq!(body["use_lsm"], serde_json::json!(false)); + let requested = filter_row_ids(body["filter"].as_str().unwrap()); + captured.lock().unwrap().push(TakeMockRequest::Query { + row_ids: requested.clone(), + version: body["version"].as_u64(), + }); + let taken = rows + .iter() + .rev() + .filter(|(row_id, _)| requested.contains(row_id)) + .collect::>(); + let batch = RecordBatch::try_from_iter([ + ( + "image", + Arc::new(descriptors( + &taken.iter().map(|(_, row)| *row).collect::>(), + )) as Arc, + ), + ( + ROW_ID, + Arc::new(UInt64Array::from_iter_values( + taken.iter().map(|(row_id, _)| *row_id), + )) as Arc, + ), + ]) + .unwrap(); + let mut body = Vec::new(); + let mut writer = + arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + let mut response = http::Response::builder().status(200); + if let Some(version) = query_version { + response = response.header(VERSION_HEADER, version); + } + return response.body(body).unwrap(); + } + if path == "/v1/table/my_table/describe/" { + captured.lock().unwrap().push(TakeMockRequest::Describe); + return http::Response::builder() + .status(200) + .body(r#"{"version":9,"schema":{"fields":[]}}"#.as_bytes().to_vec()) + .unwrap(); + } + let row_id = path + .strip_prefix("/v1/table/my_table/blob/image/") + .and_then(|rest| rest.strip_suffix("/bytes")) + .unwrap_or_else(|| panic!("unexpected path: {path}")) + .parse() + .unwrap(); + captured.lock().unwrap().push(TakeMockRequest::Range { + row_id, + range: request.headers()[header::RANGE] + .to_str() + .unwrap() + .to_string(), + version: request + .url() + .query_pairs() + .find(|(name, _)| name == "version") + .map(|(_, value)| value.into_owned()), + }); + range_response(&request, PAYLOAD) + }, + Some(Version::new(0, 5, 0)), + ); + (table, requests) + } + + #[tokio::test] + async fn remote_blob_files_are_sized_from_one_descriptor_take() { + let payload_size = PAYLOAD.len() as u64; + let (table, requests) = mock_take_blob_table( + vec![ + (10, Some((BlobKind::Inline, payload_size))), + (20, None), + (30, Some((BlobKind::Inline, 0))), + (40, Some((BlobKind::Packed, payload_size))), + (50, Some((BlobKind::Dedicated, payload_size))), + (60, Some((BlobKind::External, 0))), + ], + Some("7"), + ); + + // Reordered and duplicated row ids come back in request order. + let files = table + .fetch_blob_files_impl("image", &[60, 10, 20, 30, 40, 50, 10]) + .await + .unwrap(); + + let sizes = files + .iter() + .map(|file| file.as_ref().map(BlobFile::size)) + .collect::>(); + assert_eq!( + sizes, + [ + Some(payload_size), + Some(payload_size), + None, + Some(0), + Some(payload_size), + Some(payload_size), + Some(payload_size), + ] + ); + // One take, then a probe for the external blob that recorded no size. + assert_eq!( + requests.lock().unwrap().as_slice(), + [ + TakeMockRequest::Query { + row_ids: vec![10, 20, 30, 40, 50, 60], + version: None, + }, + TakeMockRequest::Range { + row_id: 60, + range: "bytes=0-0".into(), + version: Some("7".into()), + }, + ] + ); + + // Sized handles read the version the take read. + assert_eq!( + files[1].as_ref().unwrap().read_range(5..12).await.unwrap(), + &PAYLOAD[5..12] + ); + assert!(files[3].as_ref().unwrap().read().await.unwrap().is_empty()); + assert_eq!( + requests.lock().unwrap().last().unwrap(), + &TakeMockRequest::Range { + row_id: 10, + range: "bytes=5-11".into(), + version: Some("7".into()), + } + ); + } + + #[tokio::test] + async fn remote_blob_files_take_every_chunk_at_the_first_version() { + let row_ids = (0..MAX_FETCH_BLOBS_ROW_IDS as u64 + 1).collect::>(); + let (table, requests) = mock_take_blob_table( + row_ids.iter().map(|row_id| (*row_id, None)).collect(), + Some("7"), + ); + + let files = table + .fetch_blob_files_impl("image", &row_ids) + .await + .unwrap(); + + assert_eq!(files.len(), row_ids.len()); + let versions = requests + .lock() + .unwrap() + .iter() + .map(|request| match request { + TakeMockRequest::Query { row_ids, version } => (row_ids.len(), *version), + other => panic!("unexpected request: {other:?}"), + }) + .collect::>(); + assert_eq!(versions, [(MAX_FETCH_BLOBS_ROW_IDS, None), (1, Some(7))]); + } + + #[tokio::test] + async fn remote_blob_files_retake_at_the_described_version_when_a_take_does_not_report_it() { + let payload_size = PAYLOAD.len() as u64; + let (table, requests) = + mock_take_blob_table(vec![(10, Some((BlobKind::Inline, payload_size)))], None); + + let file = table + .fetch_blob_files_impl("image", &[10]) + .await + .unwrap() + .pop() + .flatten() + .unwrap(); + file.read_range(0..4).await.unwrap(); + + // The unversioned take is repeated at the described version. + assert_eq!( + requests.lock().unwrap().as_slice(), + [ + TakeMockRequest::Query { + row_ids: vec![10], + version: None, + }, + TakeMockRequest::Describe, + TakeMockRequest::Query { + row_ids: vec![10], + version: Some(9), + }, + TakeMockRequest::Range { + row_id: 10, + range: "bytes=0-3".into(), + version: Some("9".into()), + }, + ] + ); + } + + #[tokio::test] + async fn remote_blob_files_take_at_the_checked_out_version() { + let payload_size = PAYLOAD.len() as u64; + let (table, requests) = mock_take_blob_table( + vec![(10, Some((BlobKind::Inline, payload_size)))], + Some("7"), + ); + table.checkout(9).await.unwrap(); + requests.lock().unwrap().clear(); + + table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + + assert_eq!( + requests.lock().unwrap().as_slice(), + [TakeMockRequest::Query { + row_ids: vec![10], + version: Some(9), + }] + ); + } + + #[tokio::test] + async fn remote_blob_files_reject_row_ids_missing_from_the_take() { + let (table, requests) = mock_take_blob_table(vec![(10, None)], Some("7")); + + let error = table + .fetch_blob_files_impl("image", &[10, 99]) + .await + .unwrap_err(); + + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!(error.to_string().contains("only 1 exist"), "{error}"); + assert_eq!(requests.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn remote_blob_files_reject_a_non_blob_column() { + let table = RemoteTable::new_mock( + "my_table".to_string(), + |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/query/"); + let batch = RecordBatch::try_from_iter([ + ( + "id", + Arc::new(arrow_array::Int64Array::from(vec![1])) as Arc, + ), + ( + ROW_ID, + Arc::new(UInt64Array::from(vec![10])) as Arc, + ), + ]) + .unwrap(); + let mut body = Vec::new(); + let mut writer = + arrow_ipc::writer::FileWriter::try_new(&mut body, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + drop(writer); + http::Response::builder().status(200).body(body).unwrap() + }, + Some(Version::new(0, 5, 0)), + ); + + let error = table.fetch_blob_files_impl("id", &[10]).await.unwrap_err(); + + assert!(matches!(error, Error::InvalidInput { .. }), "{error}"); + assert!( + error.to_string().contains("'id' is not a blob column"), + "{error}" + ); + } + #[tokio::test] async fn remote_blob_file_reads_the_requested_range() { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -841,7 +1410,7 @@ mod tests { ); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -890,7 +1459,7 @@ mod tests { table.checkout(5).await.unwrap(); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -910,7 +1479,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -933,10 +1502,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); - let mut files = table - .fetch_blob_files_impl("image", &[10, 20, 30]) - .await - .unwrap(); + let mut files = table.probe_files("image", &[10, 20, 30]).await.unwrap(); assert_eq!(files.len(), 3); assert!(files[1].is_none()); @@ -964,10 +1530,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); - let error = table - .fetch_blob_files_impl("image", &[10]) - .await - .unwrap_err(); + let error = table.probe_files("image", &[10]).await.unwrap_err(); assert!( error .to_string() @@ -991,10 +1554,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); - let error = table - .fetch_blob_files_impl("image", &[10]) - .await - .unwrap_err(); + let error = table.probe_files("image", &[10]).await.unwrap_err(); assert!( error .to_string() @@ -1015,10 +1575,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); - let error = table - .fetch_blob_files_impl("image", &[10]) - .await - .unwrap_err(); + let error = table.probe_files("image", &[10]).await.unwrap_err(); assert!( error .to_string() @@ -1076,7 +1633,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -1116,7 +1673,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -1167,7 +1724,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -1194,7 +1751,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -1273,7 +1830,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); - let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let mut files = table.probe_files("image", &[10]).await.unwrap(); let file = files.remove(0).unwrap(); assert_eq!(file.size(), PAYLOAD.len() as u64); @@ -1288,7 +1845,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); - let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let mut files = table.probe_files("image", &[10]).await.unwrap(); let file = files.remove(0).unwrap(); let probe_requests = requests.lock().unwrap().len(); @@ -1313,7 +1870,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); - let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let mut files = table.probe_files("image", &[10]).await.unwrap(); let file = files.remove(0).unwrap(); let probe_requests = requests.lock().unwrap().len(); @@ -1365,7 +1922,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); - let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let mut files = table.probe_files("image", &[10]).await.unwrap(); let file = files.remove(0).unwrap(); let error = file.read_range(1..3).await.unwrap_err(); @@ -1456,7 +2013,7 @@ mod tests { let requests = Arc::new(StdMutex::new(Vec::new())); let table = mock_remote_blob_table(requests.clone()); let file = table - .fetch_blob_files_impl("image", &[10]) + .probe_files("image", &[10]) .await .unwrap() .pop() @@ -1502,7 +2059,7 @@ mod tests { Some(Version::new(0, 5, 0)), ); - let mut files = table.fetch_blob_files_impl("image", &[10]).await.unwrap(); + let mut files = table.probe_files("image", &[10]).await.unwrap(); let file = files.remove(0).unwrap(); let error = file.read_range(1..3).await.unwrap_err(); assert!(error.to_string().contains("206"), "got: {error}"); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 251baad99..cf7fcc145 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -330,7 +330,8 @@ pub use self::merge::MergeResult; /// date) and [`LsmWriteSpec::with_writer_config_defaults`] (default /// `ShardWriter` configuration recorded in the MemWAL index). /// -/// A fresh spec maintains every index on the table, resolved on install. +/// A fresh spec maintains every index on the table, including ones created +/// after it is installed. /// /// Install a spec with [`Table::set_lsm_write_spec`] and remove it with /// [`Table::unset_lsm_write_spec`]. The actual `merge_insert` dispatch @@ -348,9 +349,8 @@ pub enum LsmWriteSpec { num_buckets: u32, /// Indexes the MemWAL maintains in-memory as rows are appended. /// - /// `None` means every index it can maintain, resolved on install — a - /// snapshot, so indexes created later need the spec unset and re-set. - /// `Some([])` maintains nothing. + /// `None` means every index the table has, including ones created + /// later. `Some([])` maintains nothing. maintained_indexes: Option>, /// Default `ShardWriter` configuration recorded in the MemWAL index. writer_config_defaults: HashMap, @@ -363,9 +363,8 @@ pub enum LsmWriteSpec { column: String, /// Indexes the MemWAL maintains in-memory as rows are appended. /// - /// `None` means every index it can maintain, resolved on install — a - /// snapshot, so indexes created later need the spec unset and re-set. - /// `Some([])` maintains nothing. + /// `None` means every index the table has, including ones created + /// later. `Some([])` maintains nothing. maintained_indexes: Option>, /// Default `ShardWriter` configuration recorded in the MemWAL index. writer_config_defaults: HashMap, @@ -374,9 +373,8 @@ pub enum LsmWriteSpec { Unsharded { /// Indexes the MemWAL maintains in-memory as rows are appended. /// - /// `None` means every index it can maintain, resolved on install — a - /// snapshot, so indexes created later need the spec unset and re-set. - /// `Some([])` maintains nothing. + /// `None` means every index the table has, including ones created + /// later. `Some([])` maintains nothing. maintained_indexes: Option>, /// Default `ShardWriter` configuration recorded in the MemWAL index. writer_config_defaults: HashMap, @@ -422,14 +420,15 @@ impl LsmWriteSpec { /// Set which indexes the MemWAL maintains. /// - /// `None` (the default) resolves to every index on the table at install, - /// failing if one cannot be maintained — name the set to install anyway. A - /// list is verbatim: each name must already exist and be maintainable, and - /// an empty list maintains nothing. + /// `None` (the default) is every index on the table, re-read as the table + /// 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. /// /// ``` /// # use lancedb::table::LsmWriteSpec; - /// // Every index the table has when the spec is installed: + /// // Every index the table has, now and later: /// LsmWriteSpec::unsharded().with_maintained_indexes(None); /// // Exactly these: /// LsmWriteSpec::unsharded().with_maintained_indexes(vec!["id_idx".to_string()]); @@ -564,7 +563,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { fn id(&self) -> &str; /// Get the arrow [Schema] of the table. async fn schema(&self) -> Result; - /// Read this table's materialized-view definition and incarnation. + /// Read this table's materialized-view definition. #[doc(hidden)] async fn materialized_view_info( &self, @@ -579,9 +578,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { #[doc(hidden)] async fn refresh_materialized_view_async( &self, - _full: bool, _source_version: Option, - _expected_incarnation: Option<&str>, ) -> Result> { Err(Error::NotSupported { message: "remote materialized-view refresh is not supported on this table type".into(), @@ -1675,6 +1672,7 @@ impl Table { /// .unwrap() /// .refine_factor(5) /// .nprobes(10) + /// .unwrap() /// .execute() /// .await /// .unwrap(); @@ -1920,11 +1918,28 @@ impl Table { self.inner.function_errors(&request).await } - /// Change a column's name or nullability. + /// Change a column's name, data type, or nullability. + /// + /// Each alteration must specify at least one of `rename`, `data_type`, or + /// `nullable`. Returns [`Error::InvalidInput`] before applying any changes + /// if an alteration does not specify any of these fields. pub async fn alter_columns( &self, alterations: &[ColumnAlteration], ) -> Result { + for alteration in alterations { + if alteration.rename.is_none() + && alteration.nullable.is_none() + && alteration.data_type.is_none() + { + return Err(Error::InvalidInput { + message: format!( + "One of rename, nullable or data_type must be specified for path '{}'", + alteration.path + ), + }); + } + } self.inner.alter_columns(alterations).await } @@ -1988,6 +2003,10 @@ impl Table { /// - [`LsmWriteSpec::identity`] — shard by the raw value of a scalar column. /// - [`LsmWriteSpec::unsharded`] — route every write to a single shard. /// + /// A table carries one spec: this fails while one is installed, since the + /// generations already written were homed under it. + /// [`Table::unset_lsm_write_spec`] removes one. + /// /// # Example /// /// ``` @@ -5664,29 +5683,22 @@ mod tests { assert_eq!(table.get_lsm_write_spec().await.unwrap(), None); // Identity sharding round-trips (column recovered from the schema). - // A spec left at its default maintains every index on the table, so it - // reads back naming the one on the table rather than as "infer". + // A spec left at its default round-trips as the default: what is stored + // is the intent to maintain everything, not the set it resolves to now. let spec = LsmWriteSpec::identity("region"); table.set_lsm_write_spec(spec.clone()).await.unwrap(); - assert_eq!( - table.get_lsm_write_spec().await.unwrap(), - Some(spec.with_maintained_indexes(vec![idx_name.clone()])) - ); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec)); table.unset_lsm_write_spec().await.unwrap(); // Unsharded round-trips (no routing column). let spec = LsmWriteSpec::unsharded(); table.set_lsm_write_spec(spec.clone()).await.unwrap(); - assert_eq!( - table.get_lsm_write_spec().await.unwrap(), - Some(spec.with_maintained_indexes(vec![idx_name])) - ); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec)); } - /// The maintained set defaults to every index on the table, resolved at - /// install. An index the memtable cannot build fails the install rather - /// than being dropped: maintaining it would take the table offline for - /// writes, dropping it would hide that from the caller. + /// The maintained set defaults to every index on the table. A named index + /// the memtable cannot build fails the install, because the caller asked + /// for it; an unnamed one is skipped, because it arrived by existing. #[tokio::test] async fn test_set_lsm_write_spec_infers_maintained_indexes() { let tmp_dir = tempdir().unwrap(); @@ -5739,19 +5751,27 @@ mod tests { ); assert_eq!(table.get_lsm_write_spec().await.unwrap(), None); - // The default covers every index, so the bitmap fails it too. - let err = table + // The default names nothing, so the bitmap is skipped rather than + // refusing the table: it is maintained by existing, not by being asked + // for. The intent is what is stored, not the set it resolves to today. + table .set_lsm_write_spec(LsmWriteSpec::unsharded()) .await - .unwrap_err(); - assert!( - matches!(err, Error::InvalidInput { ref message } - if message.contains("tag_bitmap") && message.contains("maintained_indexes")), - "expected the inferred set to be rejected, got {err:?}" + .unwrap(); + assert_eq!( + table + .get_lsm_write_spec() + .await + .unwrap() + .unwrap() + .maintained_indexes(), + None, + "an unnamed set is stored as the intent to maintain everything" ); - assert_eq!(table.get_lsm_write_spec().await.unwrap(), None); - // Naming the maintainable subset installs. + // Narrowing to a named subset: unset first, since an installed spec + // cannot be set over. + table.unset_lsm_write_spec().await.unwrap(); table .set_lsm_write_spec( LsmWriteSpec::unsharded().with_maintained_indexes(vec!["id_btree".to_string()]), @@ -5783,6 +5803,28 @@ mod tests { .maintained_indexes(), Some([].as_slice()) ); + + // An installed spec is never set over, whatever the new one asks for, + // and the refused call leaves it alone. + let err = table + .set_lsm_write_spec(LsmWriteSpec::bucket("id", 4)) + .await + .unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { ref message } + if message.contains("already set")), + "expected the repeat to be refused, got {err:?}" + ); + assert_eq!( + table + .get_lsm_write_spec() + .await + .unwrap() + .unwrap() + .maintained_indexes(), + Some([].as_slice()), + "a refused call changes nothing" + ); } #[tokio::test] diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index c5d6e072f..916b28ed1 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -160,6 +160,12 @@ impl NativeTable { builder = builder.name(name); } builder.await?; + // A writer already open was built before this index existed, so it has + // to pick it up here or a read needing it is refused until it reopens. + self.dataset + .shard_writer() + .refresh_maintained_indexes(&dataset) + .await; self.dataset.update(dataset); Ok(()) } diff --git a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs index 0d4f23d92..000addbe1 100644 --- a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs +++ b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs @@ -6,11 +6,15 @@ use std::sync::Arc; -use arrow_schema::SchemaRef; +use arrow_array::{RecordBatch, UInt64Array}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; use async_trait::async_trait; +use datafusion::prelude::{SessionConfig, SessionContext, col}; use datafusion_catalog::{Session, TableProvider}; -use datafusion_common::{Result, plan_err}; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_common::{DataFusionError, Result, plan_err}; +use datafusion_execution::{ + SendableRecordBatchStream, TaskContext, runtime_env::RuntimeEnvBuilder, +}; use datafusion_expr::{Expr, TableType}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_plan::{ @@ -25,7 +29,9 @@ use futures::{TryStreamExt, stream}; use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap}; use lance::{Dataset, index::DatasetIndexInternalExt}; use lance_core::datatypes::BlobHandling; +use lance_datafusion::exec::SessionContextExt; use lance_index::metrics::NoOpMetricsCollector; +use roaring::RoaringTreemap; use super::duplicate_pairs::{ DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsOutput, @@ -36,6 +42,84 @@ use crate::table::{NativeTableExt, datafusion::BaseTableAdapter}; const REPRESENTATIVE_SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024; const MATERIALIZE_BATCH_ROWS: usize = 1024; +fn edge_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt64, false), + Field::new("b", DataType::UInt64, false), + ])) +} + +/// Sorting canonical edges makes the first unremoved endpoint the smallest +/// unassigned representative. The sort spills; only vertex masks stay in RAM. +async fn select_representatives( + stream: SendableRecordBatchStream, + memory: usize, +) -> crate::Result<(RoaringTreemap, RoaringTreemap)> { + let schema = edge_schema(); + let batch_schema = schema.clone(); + let canonical = stream + .map_ok(move |batch| { + let a = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("native row id"); + let b = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("native row id"); + RecordBatch::try_new( + batch_schema.clone(), + vec![ + Arc::new(UInt64Array::from_iter_values( + a.values().iter().zip(b.values()).map(|(&a, &b)| a.min(b)), + )), + Arc::new(UInt64Array::from_iter_values( + a.values().iter().zip(b.values()).map(|(&a, &b)| a.max(b)), + )), + ], + ) + .map_err(DataFusionError::from) + }) + .and_then(futures::future::ready); + let canonical = Box::pin(RecordBatchStreamAdapter::new(schema, canonical)); + let mut config = SessionConfig::default().with_batch_size(MATERIALIZE_BATCH_ROWS); + config.options_mut().execution.sort_spill_reservation_bytes = memory / 2; + let ctx = SessionContext::new_with_config_rt( + config, + RuntimeEnvBuilder::new() + .with_memory_limit(memory, 1.0) + .build_arc()?, + ); + let mut sorted = ctx + .read_one_shot(canonical)? + .sort_by(vec![col("a"), col("b")])? + .execute_stream() + .await?; + let mut removed = RoaringTreemap::new(); + let mut representatives = RoaringTreemap::new(); + while let Some(batch) = sorted.try_next().await? { + let a = batch + .column(0) + .as_any() + .downcast_ref::() + .expect("sorted id"); + let b = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("sorted id"); + for (&a, &b) in a.values().iter().zip(b.values()) { + if a != b && !removed.contains(a) && !removed.contains(b) { + representatives.insert(a); + removed.insert(b); + } + } + } + Ok((removed, representatives)) +} + /// Select nonempty physical IVF scopes before dispatch, never by filtering the /// resulting pairs. A count spans all segments; partition IDs are segment-local. #[derive(Debug, Clone, Default)] @@ -322,12 +406,9 @@ impl ExecutionPlan for VectorDedupExec { let dataset = self.dataset.clone(); let pairs = CoalescePartitionsExec::new(self.pairs.clone()).execute(0, context)?; let stream = stream::once(async move { - let (removed, _) = crate::materialized_view::vector_dedup::select_representatives( - pairs, - REPRESENTATIVE_SORT_MEMORY_BYTES, - ) - .await - .map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?; + let (removed, _) = select_representatives(pairs, REPRESENTATIVE_SORT_MEMORY_BYTES) + .await + .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( diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 982818ff6..354eef5a5 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -1505,6 +1505,30 @@ mod lsm_tests { ); } + /// A table with `id`/`text` and a primary key, ready for an LSM write spec. + async fn lsm_text_table(dir: &tempfile::TempDir) -> crate::Table { + let conn = connect(dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let table = conn + .create_table("t", id_text_reader(vec![(1, "alpha")])) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + } + + /// Upsert `rows` through the LSM write path. + async fn upsert_text(table: &crate::Table, rows: Vec<(i64, &str)>) { + let mut builder = table.merge_insert(&[]); + builder + .when_matched_update_all(None) + .when_not_matched_insert_all(); + builder.execute(id_text_reader(rows)).await.unwrap(); + } + /// A reader of `[id: Int64, text: Utf8]` rows. fn id_text_reader(rows: Vec<(i64, &str)>) -> Box { let schema = Arc::new(Schema::new(vec![ @@ -1524,6 +1548,165 @@ mod lsm_tests { Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)) } + /// Under the default set, an index created while a writer is open reaches + /// that writer. + /// + /// The default maintains every index the table has, so creating one changes + /// what the table maintains without anyone touching the spec. A writer left + /// on its old configs would keep building MemTables the new index does not + /// cover. + #[tokio::test] + async fn lsm_default_set_picks_up_an_index_created_while_writing() { + use crate::index::Index; + use lance_index::scalar::FullTextSearchQuery; + + let dir = tempdir().unwrap(); + let table = lsm_text_table(&dir).await; + // The default set, on a table with no FTS index yet. + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + + // Opens the writer, whose MemTable carries no FTS index. + upsert_text(&table, vec![(99, "zebra")]).await; + + // Created with the writer still open and never reopened. + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + + upsert_text(&table, vec![(100, "zebra stripes")]).await; + + let query = FullTextSearchQuery::new("stripes".to_string()) + .with_column("text".to_string()) + .unwrap(); + let batches = table + .query() + .full_text_search(query) + .execute() + .await + .expect("the created index reached the open writer") + .try_collect::>() + .await + .unwrap(); + let found: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(found, 1, "the row written after create_index must be found"); + } + + /// A spec is replaced by unsetting it and setting the new one. + /// + /// `unset_lsm_write_spec` drains the open writer, so the rows it held + /// survive the replacement, and the writer the next write opens builds its + /// MemTables from the new set. + #[tokio::test] + async fn lsm_spec_is_replaced_by_unsetting_it_first() { + use crate::index::Index; + use lance_index::scalar::FullTextSearchQuery; + + let dir = tempdir().unwrap(); + let table = lsm_text_table(&dir).await; + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + let fts_index = table.list_indices().await.unwrap()[0].name.clone(); + + // Installed maintaining nothing, then written to: the writer opens with + // no FTS index. + table + .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes(Vec::new())) + .await + .unwrap(); + upsert_text(&table, vec![(99, "zebra")]).await; + + // Setting over an installed spec is refused; unset, then set. + let err = table + .set_lsm_write_spec( + LsmWriteSpec::unsharded().with_maintained_indexes(vec![fts_index.clone()]), + ) + .await + .expect_err("an installed spec cannot be set over"); + assert!( + err.to_string().contains("already set"), + "unexpected error: {err}" + ); + table.unset_lsm_write_spec().await.unwrap(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes(vec![fts_index])) + .await + .unwrap(); + + // A row written after the replacement is answered by the index the new + // writer maintains. + upsert_text(&table, vec![(100, "zebra stripes")]).await; + + let query = FullTextSearchQuery::new("stripes".to_string()) + .with_column("text".to_string()) + .unwrap(); + let batches = table + .query() + .full_text_search(query) + .execute() + .await + .expect("the replaced spec maintains the index") + .try_collect::>() + .await + .unwrap(); + let found: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + found, 1, + "the row written after the replacement must be found" + ); + } + + /// A row written before the index exists is searchable once it does. + /// + /// The MemTable holding it was built without the index. Creating the index + /// replaces the writer's configs, which seals that MemTable and waits for + /// its flush, so no resident MemTable is left that cannot answer the read. + #[tokio::test] + async fn lsm_a_row_written_before_the_index_is_searchable_after_it() { + use crate::index::Index; + use lance_index::scalar::FullTextSearchQuery; + + let dir = tempdir().unwrap(); + let table = lsm_text_table(&dir).await; + // No index yet, and the default set: maintain whatever the table has. + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + + // Opens the writer, whose MemTable therefore carries no FTS index. + upsert_text(&table, vec![(99, "zebra")]).await; + + // Built afterwards: the spec now resolves to it, the MemTable does not. + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + + let query = FullTextSearchQuery::new("zebra".to_string()) + .with_column("text".to_string()) + .unwrap(); + let batches = table + .query() + .full_text_search(query) + .execute() + .await + .expect("no resident MemTable is left that cannot answer") + .try_collect::>() + .await + .unwrap(); + let found: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(found, 1, "the row written before the index must be found"); + } + #[tokio::test] async fn lsm_read_full_text_search() { use crate::index::Index; diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index fc925c014..0f86fdcac 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -77,21 +77,36 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec) let dataset = table.dataset.get().await?; if dataset.mem_wal_index_details().await?.is_some() { return Err(Error::InvalidInput { - message: "set_lsm_write_spec: an LSM write spec is already set on this table; mutation is not supported".into(), + message: "set_lsm_write_spec: an LSM write spec is already set on this \ + table and cannot be changed" + .into(), }); } } - // Before the builder borrows the dataset clone. `list_indices` merges an - // index's segments into one entry, so the result needs no dedup. - let maintained_indexes = { - let dataset = table.dataset.get().await?; - resolve_maintained_indexes( - &dataset, - &table.list_indices().await?, - spec.maintained_indexes(), - ) - .await? + // A named set is checked against the table here, where the caller can still + // be told which names exist. An unnamed set is not resolved at all: it is an + // intent lance re-reads whenever it builds a MemTable, so an index created + // later is maintained without another call. + let maintained_indexes = match spec.maintained_indexes() { + Some(requested) => { + let indices = table.list_indices().await?; + for name in requested { + if !indices.iter().any(|index| &index.name == name) { + return Err(Error::InvalidInput { + message: format!( + "maintained index '{}' does not exist on this table; it has {}", + name, + index_name_list(&indices), + ), + }); + } + } + let dataset = table.dataset.get().await?; + validate_maintained_indexes(&dataset, requested).await?; + Some(requested.to_vec()) + } + None => None, }; table.checkout_latest().await?; @@ -138,7 +153,9 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec) writer_config_defaults } }; - builder = builder.maintained_indexes(maintained_indexes); + if let Some(maintained_indexes) = maintained_indexes { + builder = builder.maintained_indexes(maintained_indexes); + } for (key, value) in writer_config_defaults { builder = builder.add_writer_config_default(key, value); } @@ -147,48 +164,6 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec) Ok(()) } -/// Resolve a spec's maintained-index selection against `indices`, as reported -/// by [`Table::list_indices`](crate::Table::list_indices). -/// -/// `None` means every index on the table, snapshotted now. Lance validates -/// either selection against its shard-writer rules, so a spec that installs is -/// one the MemWAL can open. -/// -/// An unmaintainable index fails an inferred set rather than being dropped from -/// it — dropping would leave the caller believing it is maintained. -async fn resolve_maintained_indexes( - dataset: &Dataset, - indices: &[IndexConfig], - requested: Option<&[String]>, -) -> Result> { - let Some(requested) = requested else { - let all: Vec = indices.iter().map(|index| index.name.clone()).collect(); - validate_maintained_indexes(dataset, &all) - .await - .map_err(|source| Error::InvalidInput { - message: format!( - "cannot maintain every index on this table: {source}. Set \ - maintained_indexes explicitly to choose from {}", - index_name_list(indices), - ), - })?; - return Ok(all); - }; - for name in requested { - if !indices.iter().any(|index| &index.name == name) { - return Err(Error::InvalidInput { - message: format!( - "maintained index '{}' does not exist on this table; it has {}", - name, - index_name_list(indices), - ), - }); - } - } - validate_maintained_indexes(dataset, requested).await?; - Ok(requested.to_vec()) -} - /// Index names for an error message. fn index_name_list(indices: &[IndexConfig]) -> String { if indices.is_empty() { @@ -298,7 +273,9 @@ fn lsm_write_spec_from_details( }; Ok(base - .with_maintained_indexes(details.maintained_indexes.clone()) + .with_maintained_indexes( + (!details.maintain_all_indexes).then(|| details.maintained_indexes.clone()), + ) .with_writer_config_defaults(details.writer_config_defaults.clone())) } @@ -460,6 +437,55 @@ impl ShardWriterCache { Ok(entry) } + /// Install `dataset`'s maintained index set on the cached writer, if one is + /// open, so a change reaches a table that is already being written to. + /// + /// Reported rather than propagated: the index is committed durably by the + /// time this runs, and the writer this process happens to hold may be fenced + /// or poisoned -- which is not a reason to fail creating an index. Any + /// writer opened afterwards reads the committed set. + /// + /// Holds the entry's read lock across the refresh, so a concurrent close + /// waits rather than retiring the writer mid-seal. Writes take the same + /// lock first, so the order matches theirs. + pub(crate) async fn refresh_maintained_indexes(&self, dataset: &Dataset) { + let entry = { + let guard = self.slot.read().await; + guard.as_ref().map(|(_, entry)| entry.clone()) + }; + let Some(entry) = entry else { + return; + }; + let guard = entry.inner.read().await; + let Some(writer) = guard.as_ref() else { + return; + }; + // Replacing the configs seals the active memtable so the next one carries + // the new indexes. The outgoing one keeps the set it was built with, so + // its flush has to land before a read: while it is still frozen, an + // indexed read sees a resident memtable that cannot answer and is + // refused. + let sealed = match dataset.refresh_mem_wal_index_configs(writer).await { + Ok(sealed) => sealed, + Err(error) => { + log::warn!( + "the maintained index set changed but the open shard writer could \ + not be refreshed, so it keeps the set it opened with: {error}" + ); + return; + } + }; + if let Some(fence) = sealed + && let Err(error) = fence.wait().await + { + log::warn!( + "the shard writer took the new maintained index set, but the memtable \ + it sealed was not flushed, so an indexed read stays refused until \ + it is: {error}" + ); + } + } + /// Snapshot the cached writer's shard for the LSM read path: its shard id, /// authoritative in-memory manifest, and active + frozen memtable refs. /// Returns `None` when no writer is currently cached (e.g. nothing has been diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs index 86c1fe5f2..f3e40fc7b 100644 --- a/rust/lancedb/src/table/query/lsm.rs +++ b/rust/lancedb/src/table/query/lsm.rs @@ -33,7 +33,7 @@ use datafusion_physical_plan::expressions::Column; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use lance::Dataset; -use lance::dataset::mem_wal::scanner::InMemoryMemTables; +use lance::dataset::mem_wal::scanner::{InMemoryMemTableRef, InMemoryMemTables}; use lance::dataset::mem_wal::{ DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriterConfig, }; @@ -468,20 +468,19 @@ async fn fts_plan( .to_string(), })?; - // Without a maintained in-memory FTS index for this column, the active memtable - // arm produces an empty plan (`active_source_can_execute_fts` returns false), so - // the search silently omits un-compacted documents. Reject rather than mislead. - if !index_maintained( - dataset, - column, - &details.maintained_indexes, - "InvertedIndexDetails", - ) - .await? + // Two ways the fresh tier cannot answer: the spec maintains no FTS index for + // this column, or a resident MemTable was built before the index joined the + // set. Either leaves the memtable arm an empty plan + // (`active_source_can_execute_fts` returns false), so the search silently + // 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() + }) { return Err(Error::NotSupported { message: format!( - "the MemWAL LSM scanner full-text search requires the FTS index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes); otherwise un-compacted documents are omitted. set use_lsm(false) to read the base table only" + "the MemWAL LSM scanner full-text search requires the FTS index on '{column}' to be maintained by the write spec (LsmWriteSpec::with_maintained_indexes) and carried by every resident MemTable; otherwise un-compacted documents are omitted. set use_lsm(false) to read the base table only" ), }); } @@ -493,6 +492,27 @@ async fn fts_plan( Ok(scanner.create_plan().await?) } +/// Whether the spec maintains the index named `name`. +/// +/// An unnamed set names nothing and covers everything, so the stored list alone +/// cannot answer this. +fn spec_maintains(details: &MemWalIndexDetails, name: &str) -> bool { + details.maintain_all_indexes || details.maintained_indexes.iter().any(|n| n == name) +} + +/// Whether every resident MemTable carries the index an arm needs. +/// +/// The spec is intent; a MemTable built before the index joined the set +/// contributes nothing to an indexed read, so its rows would go missing. +fn resident_memtables_carry( + in_memory: &HashMap, + carries: impl Fn(&InMemoryMemTableRef) -> bool, +) -> bool { + in_memory + .values() + .all(|tables| carries(&tables.active) && tables.frozen.iter().all(&carries)) +} + /// Whether an index of `type_url_suffix` covering `column` is in the MemWAL spec's /// maintained set. Only a maintained index has its catch-up tracked (so exclusion is /// gated correctly) and its in-memory arm kept current; an unmaintained base index @@ -502,7 +522,7 @@ async fn fts_plan( async fn index_maintained( dataset: &Dataset, column: &str, - maintained: &[String], + details: &MemWalIndexDetails, type_url_suffix: &str, ) -> Result { use lance::index::DatasetIndexExt; @@ -512,7 +532,7 @@ async fn index_maintained( let indices = dataset.load_indices().await?; Ok(indices.iter().any(|idx| { idx.fields.contains(&field.id) - && maintained.iter().any(|m| m == &idx.name) + && spec_maintains(details, &idx.name) && idx .index_details .as_ref() @@ -583,9 +603,7 @@ async fn arm_maintained_index_names( }) .map(|idx| idx.name.clone()) .collect(); - if let Some(name) = - resolve_single_index(segment_names, &details.maintained_indexes, arm, &column)? - { + if let Some(name) = resolve_single_index(segment_names, details, arm, &column)? { names.push(name); } } @@ -604,7 +622,7 @@ async fn arm_maintained_index_names( /// compaction watermark). fn resolve_single_index( mut names: Vec, - maintained: &[String], + details: &MemWalIndexDetails, arm: &str, column: &str, ) -> Result> { @@ -620,7 +638,7 @@ fn resolve_single_index( Ok(names .into_iter() .next() - .filter(|name| maintained.contains(name))) + .filter(|name| spec_maintains(details, name))) } /// Drop the primary-key columns Lance appends internally for dedup when the user's @@ -688,21 +706,19 @@ async fn vector_plan( None => default_vector_column(&arrow_schema, Some(query_vector.len() as i32))?, }; - // The base arm relies on the column's vector index (`fast_search`). Unless it is - // maintained, its catch-up is untracked and exclusion falls back to the - // compaction watermark — 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.maintained_indexes, - "VectorIndexDetails", - ) - .await? + // 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. + // 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() + }) { 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); 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) 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" ), }); } @@ -888,7 +904,10 @@ mod tests { #[test] fn resolve_single_index_dedupes_segments() { - let maintained = vec!["fts_idx".to_string()]; + let maintained = MemWalIndexDetails { + maintained_indexes: vec!["fts_idx".to_string()], + ..Default::default() + }; // Two physical segments of ONE logical index must not count as "multiple". assert_eq!( resolve_single_index( @@ -917,4 +936,33 @@ mod tests { None ); } + + /// A spec that maintains every index names none of them, so the stored list + /// is empty while every index on the table is in fact maintained. Reading + /// the list alone would refuse the read path on the default spec. + #[test] + fn an_unnamed_set_maintains_every_index() { + let all = MemWalIndexDetails { + maintained_indexes: Vec::new(), + maintain_all_indexes: true, + ..Default::default() + }; + assert!(spec_maintains(&all, "fts_idx")); + assert_eq!( + resolve_single_index(vec!["fts_idx".to_string()], &all, "full-text", "text").unwrap(), + Some("fts_idx".to_string()), + ); + + // The same empty list without the flag maintains nothing. + let none = MemWalIndexDetails { + maintained_indexes: Vec::new(), + maintain_all_indexes: false, + ..Default::default() + }; + assert!(!spec_maintains(&none, "fts_idx")); + assert_eq!( + resolve_single_index(vec!["fts_idx".to_string()], &none, "full-text", "text").unwrap(), + None, + ); + } } diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index 4d0e2815d..4990901ae 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -807,6 +807,29 @@ mod tests { // Alter Columns Tests + #[tokio::test] + async fn test_alter_columns_rejects_missing_changes() { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap(); + let table = conn.create_table("test", batch).execute().await.unwrap(); + let initial_version = table.version().await.unwrap(); + let initial_schema = table.schema().await.unwrap(); + + for alterations in [ + vec![ColumnAlteration::new("id".into())], + vec![ + ColumnAlteration::new("id".into()).rename("new_id".into()), + ColumnAlteration::new("id".into()), + ], + ] { + let err = table.alter_columns(&alterations).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("path 'id'")); + assert_eq!(table.version().await.unwrap(), initial_version); + assert_eq!(table.schema().await.unwrap(), initial_schema); + } + } + #[tokio::test] async fn test_alter_column_rename() { let conn = connect("memory://").execute().await.unwrap(); diff --git a/rust/lancedb/tests/duplicate_pairs.rs b/rust/lancedb/tests/duplicate_pairs.rs index f7d7ed3a2..63f1c9f5a 100644 --- a/rust/lancedb/tests/duplicate_pairs.rs +++ b/rust/lancedb/tests/duplicate_pairs.rs @@ -289,127 +289,6 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result ); let retry = ctx.sql(&sql).await?.collect().await?; assert_eq!(pairs(&retry), pairs(&actual)); - // The same table function can be maintained through the MV unit - // lifecycle. Staging attempts do not make any rows visible. - use lancedb::materialized_view::{ - MaterializedViewDefinition, commit_partitioned_refresh, plan_partitioned_refresh, - prepare_definition, write_refresh_partition, - }; - let view_name = format!("pairs_{}", threshold.to_bits()); - let view = prepare_definition(&table, MaterializedViewDefinition::from_sql(&sql)?) - .await? - .create(&view_name) - .await?; - assert_eq!(view.table().schema().await?.fields().len(), 3); - let plan = plan_partitioned_refresh(view.table(), Some(version)) - .await? - .unwrap(); - assert_eq!(plan.units() as usize, tasks.len()); - let plan = serde_json::from_slice(&serde_json::to_vec(&plan)?)?; - let first = write_refresh_partition(view.table(), 0, &plan).await?; - assert_eq!(view.table().count_rows(None).await?, 0); - assert!( - commit_partitioned_refresh( - view.table(), - &plan, - vec![first.clone(), first.clone()], - view.incarnation() - ) - .await - .is_err() - ); - // Lost receipt: its files remain unreferenced; the retry writes - // another attempt, of which only one receipt is ever committed. - drop(write_refresh_partition(view.table(), 0, &plan).await?); - let mut units = vec![first]; - for unit in 1..plan.units() { - units.push(write_refresh_partition(view.table(), unit, &plan).await?); - } - assert_eq!(view.table().count_rows(None).await?, 0); - let outcome = - commit_partitioned_refresh(view.table(), &plan, units.clone(), view.incarnation()) - .await?; - assert_eq!(outcome.rows_written as usize, pairs(&expected).len()); - assert!( - commit_partitioned_refresh(view.table(), &plan, units, view.incarnation()) - .await - .is_err() - ); - use lancedb::query::ExecutableQuery; - let written = view - .table() - .query() - .execute() - .await? - .try_collect::>() - .await?; - assert_eq!(pairs(&written), pairs(&expected)); - // Repeating a refresh replaces the view rather than appending. - view.refresh().execute().await?; - assert_eq!( - view.table().count_rows(None).await? as usize, - pairs(&expected).len() - ); - // The complete declarative pipeline keeps direct representatives - // and materializes original rows, including isolated source rows. - use lancedb::materialized_view::write_refresh_partition_with_inputs; - use lancedb::query::QueryBase; - let definition = MaterializedViewDefinition::from_sql(&format!( - "SELECT * FROM vector_dedup('source', {version}, 'vector', {threshold})" - ))?; - let clean = prepare_definition(&table, definition) - .await? - .create(&format!("clean_{}", threshold.to_bits())) - .await?; - let clean_plan = plan_partitioned_refresh(clean.table(), Some(version)) - .await? - .unwrap(); - assert_eq!(clean_plan.dependency_units() as usize, tasks.len()); - assert_eq!(clean_plan.units() as usize, tasks.len() + 2); - let clean_plan: lancedb::materialized_view::PartitionedRefreshPlan = - serde_json::from_slice(&serde_json::to_vec(&clean_plan)?)?; - assert!( - write_refresh_partition(clean.table(), tasks.len() as u32, &clean_plan) - .await - .is_err() - ); - let mut receipts = Vec::new(); - for unit in 0..clean_plan.units() { - let dependencies = &receipts[..receipts.len().min(tasks.len())]; - if unit == 0 || unit == tasks.len() as u32 { - // Both selection and materialization can lose a receipt. - drop( - write_refresh_partition_with_inputs( - clean.table(), - unit, - &clean_plan, - dependencies, - ) - .await?, - ); - } - let receipt = write_refresh_partition_with_inputs( - clean.table(), - unit, - &clean_plan, - dependencies, - ) - .await?; - receipts.push(serde_json::from_slice(&serde_json::to_vec(&receipt)?)?); - } - assert_eq!(clean.table().count_rows(None).await?, 0); - assert!( - commit_partitioned_refresh( - clean.table(), - &clean_plan, - receipts[..receipts.len() - 1].to_vec(), - clean.incarnation() - ) - .await - .is_err() - ); - commit_partitioned_refresh(clean.table(), &clean_plan, receipts, clean.incarnation()) - .await?; let mut edges = pairs(&expected) .into_iter() .map(|(a, b, _)| (a.min(b), a.max(b))) @@ -449,27 +328,6 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result } } } - let kept = clean - .table() - .query() - .select(lancedb::query::Select::columns(&["id"])) - .execute() - .await? - .try_collect::>() - .await?; - let mut actual_ids = kept - .iter() - .flat_map(|b| { - b.column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .iter() - .copied() - }) - .collect::>(); - actual_ids.sort_unstable(); expected_ids.sort_unstable(); let mut query_ids = clean_query .iter() @@ -486,14 +344,10 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result query_ids.sort_unstable(); assert_eq!(query_ids, expected_ids); - assert_eq!(actual_ids, expected_ids); if threshold == 1.0 { - assert!(actual_ids.contains(&0) && actual_ids.contains(&2)); - assert!(!actual_ids.contains(&1), "A-B-C retains A and C"); + assert!(query_ids.contains(&0) && query_ids.contains(&2)); + assert!(!query_ids.contains(&1), "A-B-C retains A and C"); } - clean.refresh().execute().await?; - assert_eq!(clean.table().count_rows(None).await?, expected_ids.len()); - assert_eq!(table.count_rows(None).await?, n * 2); assert!( DuplicatePairsExec::try_new( ds.clone(), @@ -527,10 +381,8 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result } #[tokio::test] -async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() -> anyhow::Result<()> -{ +async fn dedup_query_materializes_blob_payloads_and_null_vectors() -> anyhow::Result<()> { use arrow_array::{LargeBinaryArray, StringArray, StructArray, types::Float32Type}; - use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition}; for stable in [false, true] { let dir = tempfile::tempdir()?; let conn = lancedb::connect(dir.path().to_str().unwrap()) @@ -607,7 +459,6 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() true, ) .await?; - let version = ds.version().version; use lancedb::table::datafusion::udtf::{ duplicate_pairs::DuplicatePairsOutput, vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable}, @@ -692,59 +543,6 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() assert_eq!(images.value(i), vec![ids.value(i) as u8; 4096]); } } - let definition = MaterializedViewDefinition::from_sql(&format!( - "SELECT * FROM vector_dedup('source', {version}, 'vector', 1)" - ))?; - let clean = prepare_definition(&table, definition) - .await? - .create("clean") - .await?; - // A later source delete does not change the declared result snapshot. - table.delete("id = 0").await?; - let adapter = BaseTableAdapter::try_new(table.base_table().clone()).await?; - let explicit = adapter.vector_query_snapshot(Some(version)).await?; - let latest = adapter.vector_query_snapshot(None).await?; - assert_eq!(explicit.version().version, version); - assert!(latest.version().version > version); - assert_eq!(explicit.count_rows(None).await?, 5); - assert_eq!(latest.count_rows(None).await?, 4); - clean.refresh().execute().await?; - assert_eq!(clean.table().count_rows(None).await?, 4); - let output = Dataset::open(&clean.table().uri().await?).await?; - let batches = output - .scan() - .blob_handling(lance_core::datatypes::BlobHandling::AllBinary) - .try_into_stream() - .await? - .try_collect::>() - .await?; - let mut ids = Vec::new(); - for batch in batches { - let row_ids = batch - .column_by_name("id") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let images = batch - .column_by_name("image") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - for i in 0..batch.num_rows() { - let id = row_ids.value(i); - assert_eq!(images.value(i), vec![id as u8; 4096]); - ids.push(id); - } - } - ids.sort_unstable(); - assert_eq!(ids, [0, 2, 3, 4]); - assert_eq!( - table.count_rows(None).await?, - 4, - "refresh must not mutate the source" - ); } Ok(()) }