diff --git a/.bumpversion.toml b/.bumpversion.toml index 601b14d3f..cab6bb104 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.37.1-beta.0" +current_version = "0.38.0-beta.0" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/.github/workflows/codex-update-lance-dependency.yml b/.github/workflows/codex-update-lance-dependency.yml index 420daf650..79ef84364 100644 --- a/.github/workflows/codex-update-lance-dependency.yml +++ b/.github/workflows/codex-update-lance-dependency.yml @@ -4,14 +4,14 @@ on: workflow_call: inputs: tag: - description: "Tag name from Lance. If omitted, the skill will use the latest Lance release that needs an update." + description: "Tag name from Lance (e.g. `v7.2.0-beta.1`). If omitted, the newest release is resolved automatically — stable releases are preferred over pre-releases — and the run is skipped if it is not newer than the version currently pinned in Cargo.toml." required: false default: "" type: string workflow_dispatch: inputs: tag: - description: "Tag name from Lance. Leave empty to use the latest Lance release that needs an update." + description: "Tag name from Lance (e.g. `v7.2.0-beta.1`). Leave empty to resolve the newest release automatically — stable releases are preferred over pre-releases — and skip the run if it is not newer than the version currently pinned in Cargo.toml." required: false default: "" type: string diff --git a/.github/workflows/docs-link-check.yml b/.github/workflows/docs-link-check.yml index 0e22100eb..1286819bc 100644 --- a/.github/workflows/docs-link-check.yml +++ b/.github/workflows/docs-link-check.yml @@ -36,7 +36,9 @@ jobs: permissions: contents: read outputs: + checker_outcome: ${{ steps.lychee.outcome }} exit_code: ${{ steps.lychee.outputs.exit_code }} + status: ${{ steps.validate.outputs.status }} steps: - name: Checkout uses: actions/checkout@v6 @@ -50,6 +52,7 @@ jobs: - name: Check links id: lychee + continue-on-error: true uses: lycheeverse/lychee-action@e7477775783ea5526144ba13e8db5eec57747ce8 # v2.9.0 with: # Restricted to http(s) on purpose. Much of docs/src is generated @@ -68,38 +71,50 @@ jobs: format: json output: ./lychee/out.json jobSummary: false - # The report, not a red build, is the signal for broken links. The - # validation step below still fails the run if the check itself - # breaks. + # The report issue, not a red workflow run, is the signal for link + # findings and checker failures alike. fail: false - name: Validate report + id: validate # lychee does not reserve exit code 2 for broken links: its CLI # parser also exits 2 on an invalid option, before any link was # checked or any report written. Only a parseable report whose - # counts agree with the exit code counts as a link verdict; anything - # else fails here, and the report job below is skipped entirely, so - # the tracking issue is never touched. Exit 2 covers timeouts as - # well as errors, and a timed-out host is exactly the transient - # unavailability this report exists to surface, so both count as - # findings. Requiring total > 0 also catches a glob that silently - # stopped matching any file. - if: steps.lychee.outputs.exit_code == 0 || steps.lychee.outputs.exit_code == 2 + # counts agree with a completed exit code (0 or 2) counts as a link + # verdict. Everything else becomes a checker-error report instead of + # failing the workflow. Exit 2 covers timeouts as well as errors, and a + # timed-out host is exactly the transient unavailability this report + # exists to surface, so both count as findings. Requiring total > 0 + # also catches a glob that silently stopped matching any file. + if: always() env: + CHECKER_OUTCOME: ${{ steps.lychee.outcome }} EXIT_CODE: ${{ steps.lychee.outputs.exit_code }} run: | - jq -e --argjson code "$EXIT_CODE" ' - (.total > 0) and - (if $code == 0 - then .errors == 0 and .timeouts == 0 - and (.error_map | length == 0) and (.timeout_map | length == 0) - else (.errors + .timeouts) > 0 - and ((.error_map | length) + (.timeout_map | length)) > 0 - end) - ' ./lychee/out.json + status=checker-error + if [[ "$CHECKER_OUTCOME" == success ]] && + [[ "$EXIT_CODE" == 0 || "$EXIT_CODE" == 2 ]] && + jq -e --argjson code "$EXIT_CODE" ' + (.total > 0) and + (if $code == 0 + then .errors == 0 and .timeouts == 0 + and (.error_map | length == 0) and (.timeout_map | length == 0) + else (.errors + .timeouts) > 0 + and ((.error_map | length) + (.timeout_map | length)) > 0 + end) + ' ./lychee/out.json + then + if [[ "$EXIT_CODE" == 0 ]]; then + status=healthy + else + status=findings + fi + fi + echo "status=$status" >> "$GITHUB_OUTPUT" + echo "Validated link check as $status" - name: Upload report - if: steps.lychee.outputs.exit_code == 2 + if: steps.validate.outputs.status == 'findings' uses: actions/upload-artifact@v7 with: name: link-report @@ -115,26 +130,11 @@ jobs: permissions: issues: write env: + CHECKER_OUTCOME: ${{ needs.scan.outputs.checker_outcome }} EXIT_CODE: ${{ needs.scan.outputs.exit_code }} + STATUS: ${{ needs.scan.outputs.status }} GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} steps: - - name: Classify checker result - # lychee exits 0 when every link resolves and 2 when links fail, - # both already cross-checked against the report by the scan job's - # validation step. Anything else (1 runtime, 3 bad config) means the - # check never produced a link verdict, which must surface as a failed - # run rather than be published as "broken documentation links". - run: | - case "$EXIT_CODE" in - 0|2) - echo "lychee exit code $EXIT_CODE" - ;; - *) - echo "::error::lychee exited with '$EXIT_CODE': the link check did not complete. Leaving the report issue untouched." - exit 1 - ;; - esac - - name: Find existing report issue id: report # Matched on title alone, and through search rather than a listing: @@ -144,7 +144,7 @@ jobs: # Closed issues are included because a healthy run closes the report: # an open-only lookup would forget that identity and the next failing # run would open a duplicate. The oldest match stays the canonical - # report and is reopened below when links break again. + # report and is reopened below when a problem recurs. run: | match=$(gh issue list --repo "$GITHUB_REPOSITORY" --state all \ --search "in:title \"$REPORT_TITLE\" author:app/github-actions" \ @@ -154,14 +154,14 @@ jobs: echo "state=$(jq -r '.state // empty' <<<"$match")" >> "$GITHUB_OUTPUT" - name: Download report - if: env.EXIT_CODE == 2 + if: env.STATUS == 'findings' uses: actions/download-artifact@v8 with: name: link-report path: ./lychee - name: Compose report - if: env.EXIT_CODE == 2 + if: env.STATUS == 'findings' run: | run_url="$GITHUB_SERVER_URL/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID" { @@ -185,22 +185,41 @@ jobs: ' ./lychee/out.json } > ./lychee/issue.md + - name: Compose checker error report + if: env.STATUS == 'checker-error' + run: | + mkdir -p ./lychee + run_url="$GITHUB_SERVER_URL/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID" + { + echo "The documentation link check did not complete in [the latest run]($run_url)." + echo + echo "This issue is rewritten by every scheduled run and closed automatically once a trustworthy run finds that all links resolve." + echo + echo "The checker did not produce a trustworthy link verdict. Treat the previous result, if any, as stale until a later run completes." + echo + echo "* Action outcome: \`$CHECKER_OUTCOME\`" + echo "* Exit code: \`${EXIT_CODE:-not reported}\`" + echo "* Verdict validation: \`failed\`" + } > ./lychee/issue.md + - name: Reopen report issue # A healthy run closes the report, and the issue action below only # rewrites the body of whatever number it is given. Without an - # explicit reopen, the 2 -> 0 -> 2 sequence would keep rewriting a - # closed issue while links are broken. A CLOSED state implies the - # lookup found a canonical issue, so no separate emptiness check. - if: env.EXIT_CODE == 2 && steps.report.outputs.state == 'CLOSED' + # explicit reopen, a later finding or checker error would rewrite a + # closed issue. A CLOSED state implies the lookup found a canonical + # issue, so no separate emptiness check. + if: >- + env.STATUS != 'healthy' && + steps.report.outputs.state == 'CLOSED' env: ISSUE_NUMBER: ${{ steps.report.outputs.number }} run: | run_url="$GITHUB_SERVER_URL/$GITHUB_REPOSITORY/actions/runs/$GITHUB_RUN_ID" gh issue reopen "$ISSUE_NUMBER" --repo "$GITHUB_REPOSITORY" \ - --comment "Broken documentation links found again in [the latest run]($run_url)." + --comment "The documentation link checker reported a problem again in [the latest run]($run_url)." - - name: Report broken links - if: env.EXIT_CODE == 2 + - name: Report link-check problem + if: env.STATUS != 'healthy' uses: peter-evans/create-issue-from-file@fca9117c27cdc29c6c4db3b86c48e4115a786710 # v6.0.0 with: # Empty on the first failing run, which creates the issue; afterwards @@ -213,7 +232,9 @@ jobs: - name: Close report issue once links are healthy # An OPEN state implies the lookup found a canonical issue; a report # that is already closed needs nothing. - if: env.EXIT_CODE == 0 && steps.report.outputs.state == 'OPEN' + if: >- + env.STATUS == 'healthy' && + steps.report.outputs.state == 'OPEN' env: ISSUE_NUMBER: ${{ steps.report.outputs.number }} run: | diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml index 74b7d05e6..4f5a927dc 100644 --- a/.github/workflows/pypi-publish.yml +++ b/.github/workflows/pypi-publish.yml @@ -69,6 +69,16 @@ jobs: uses: actions/setup-python@v6 with: python-version: "3.10" + - name: Add swap for Arm fat LTO + if: matrix.config.platform == 'aarch64' + shell: bash + run: | + swap_file="$RUNNER_TEMP/lancedb-swap" + sudo fallocate --length 16G "$swap_file" + sudo chmod 600 "$swap_file" + sudo mkswap "$swap_file" + sudo swapon "$swap_file" + free -h - uses: ./.github/workflows/build_linux_wheel with: python-minor-version: 10 diff --git a/Cargo.lock b/Cargo.lock index f6186672c..b4f11fb65 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3455,8 +3455,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" [[package]] name = "fsst" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "rand 0.9.5", @@ -4815,8 +4815,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a" [[package]] name = "lance" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arc-swap", "arrow", @@ -4832,7 +4832,6 @@ dependencies = [ "async-recursion", "async-trait", "async_cell", - "aws-credential-types", "aws-sdk-dynamodb", "byteorder", "bytes", @@ -4848,7 +4847,6 @@ dependencies = [ "either", "fst", "futures", - "half", "humantime", "itertools 0.14.0", "lance-arrow", @@ -4890,8 +4888,8 @@ dependencies = [ [[package]] name = "lance-arrow" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "arrow-buffer", @@ -4913,7 +4911,7 @@ dependencies = [ [[package]] name = "lance-arrow-scalar" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "arrow-buffer", @@ -4927,7 +4925,7 @@ dependencies = [ [[package]] name = "lance-arrow-stats" version = "58.0.0" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "arrow-schema", @@ -4936,8 +4934,8 @@ dependencies = [ [[package]] name = "lance-bitpacking" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrayref", "crunchy", @@ -4947,8 +4945,8 @@ dependencies = [ [[package]] name = "lance-core" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "arrow-buffer", @@ -4956,12 +4954,10 @@ dependencies = [ "arrow-schema", "async-trait", "blake3", - "byteorder", "bytes", "datafusion-common", "datafusion-sql", "futures", - "itertools 0.14.0", "lance-arrow", "lance-derive", "libc", @@ -4979,7 +4975,6 @@ dependencies = [ "snafu 0.9.0", "tempfile", "tokio", - "tokio-stream", "tokio-util", "tracing", "twox-hash", @@ -4988,8 +4983,8 @@ dependencies = [ [[package]] name = "lance-datafusion" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow", "arrow-array", @@ -5008,7 +5003,6 @@ dependencies = [ "jsonb", "lance-arrow", "lance-core", - "lance-datagen", "log", "pin-project", "prost", @@ -5019,8 +5013,8 @@ dependencies = [ [[package]] name = "lance-datagen" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow", "arrow-array", @@ -5037,8 +5031,8 @@ dependencies = [ [[package]] name = "lance-derive" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "proc-macro2", "quote", @@ -5047,8 +5041,8 @@ dependencies = [ [[package]] name = "lance-encoding" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-arith", "arrow-array", @@ -5073,7 +5067,6 @@ dependencies = [ "num-traits", "prost", "prost-build", - "rand 0.9.5", "tokio", "tracing", "xxhash-rust", @@ -5082,8 +5075,8 @@ dependencies = [ [[package]] name = "lance-file" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-arith", "arrow-array", @@ -5114,8 +5107,8 @@ dependencies = [ [[package]] name = "lance-index" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arc-swap", "arrow", @@ -5130,7 +5123,6 @@ dependencies = [ "async-trait", "bitvec", "bytes", - "chrono", "crossbeam-queue", "datafusion", "datafusion-common", @@ -5148,7 +5140,6 @@ dependencies = [ "lance-bitpacking", "lance-core", "lance-datafusion", - "lance-datagen", "lance-encoding", "lance-file", "lance-index-core", @@ -5177,13 +5168,12 @@ dependencies = [ "tempfile", "tokio", "tracing", - "uuid", ] [[package]] name = "lance-index-core" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "arrow-schema", @@ -5205,8 +5195,8 @@ dependencies = [ [[package]] name = "lance-io" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow", "arrow-array", @@ -5220,7 +5210,6 @@ dependencies = [ "futures", "http 1.5.0", "io-uring", - "lance-arrow", "lance-core", "lance-namespace", "log", @@ -5238,29 +5227,28 @@ dependencies = [ "tokio", "tracing", "url", + "uuid", ] [[package]] name = "lance-linalg" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", - "arrow-buffer", "arrow-schema", "cc", "half", "lance-arrow", "lance-core", "num-traits", - "rand 0.9.5", "rayon", ] [[package]] name = "lance-namespace" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow", "async-trait", @@ -5272,8 +5260,8 @@ dependencies = [ [[package]] name = "lance-namespace-impls" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow", "arrow-ipc", @@ -5326,14 +5314,13 @@ dependencies = [ [[package]] name = "lance-select" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "arrow-buffer", "arrow-schema", "byteorder", - "bytes", "itertools 0.14.0", "lance-core", "roaring", @@ -5342,8 +5329,8 @@ dependencies = [ [[package]] name = "lance-table" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow", "arrow-array", @@ -5383,8 +5370,8 @@ dependencies = [ [[package]] name = "lance-testing" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "arrow-array", "arrow-schema", @@ -5397,8 +5384,8 @@ dependencies = [ [[package]] name = "lance-tokenizer" -version = "11.0.0-beta.3" -source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.3#f7d475539cefbd140cc46a828f3d843e68cd10f1" +version = "11.0.0-beta.11" +source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.11#5ef1e969030b79fe8c6d31fff9e40d18b24ebbd3" dependencies = [ "frostem", "icu_segmenter", @@ -5411,7 +5398,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.37.1-beta.0" +version = "0.38.0-beta.0" dependencies = [ "ahash", "anyhow", @@ -5447,7 +5434,6 @@ dependencies = [ "datafusion-physical-plan", "datafusion-sql", "futures", - "goosefs-sdk", "half", "hf-hub", "http 1.5.0", @@ -5480,6 +5466,7 @@ dependencies = [ "random_word", "regex", "reqwest 0.12.28", + "roaring", "rstest", "semver", "serde", @@ -5499,7 +5486,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.37.1-beta.0" +version = "0.38.0-beta.0" dependencies = [ "arrow-array", "arrow-buffer", @@ -5524,7 +5511,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.37.1-beta.0" +version = "0.38.0-beta.0" dependencies = [ "arrow", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 936660d78..3e332adfc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,20 +13,20 @@ categories = ["database-implementations"] rust-version = "1.91.0" [workspace.dependencies] -lance = { "version" = "=11.0.0-beta.3", default-features = false, "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-core = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-datagen = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-file = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-io = { "version" = "=11.0.0-beta.3", default-features = false, "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-index = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-linalg = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-namespace-impls = { "version" = "=11.0.0-beta.3", default-features = false, "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-table = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-testing = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-datafusion = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-encoding = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=11.0.0-beta.3", "tag" = "v11.0.0-beta.3", "git" = "https://github.com/lance-format/lance.git" } +lance = { "version" = "=11.0.0-beta.11", default-features = false, "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-core = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-datagen = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-file = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-io = { "version" = "=11.0.0-beta.11", default-features = false, "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-index = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-linalg = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-namespace-impls = { "version" = "=11.0.0-beta.11", default-features = false, "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-table = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-testing = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-datafusion = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-encoding = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { "version" = "=11.0.0-beta.11", "tag" = "v11.0.0-beta.11", "git" = "https://github.com/lance-format/lance.git" } ahash = "0.8" # Note that this one does not include pyarrow arrow = { version = "58.0.0", optional = false } diff --git a/deny.toml b/deny.toml index 034b48c25..d94c9d536 100644 --- a/deny.toml +++ b/deny.toml @@ -101,6 +101,13 @@ ignore = [ # https://rustsec.org/advisories/RUSTSEC-2026-0195 { id = "RUSTSEC-2026-0194", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" }, { id = "RUSTSEC-2026-0195", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" }, + # smartstring: unmaintained — the repository was archived by its author on + # 2026-05-03. Not a vulnerability. Reached only transitively through polars + # (polars-core/-io/-ops/-time/-utils); nothing in LanceDB depends on it directly. + # The advisory states no safe upgrade is available: upstream recommends + # compact_str/smol_str, so clearing this requires polars to migrate. + # https://rustsec.org/advisories/RUSTSEC-2026-0249 + { id = "RUSTSEC-2026-0249", reason = "smartstring unmaintained via polars; no fixed upstream release" }, ] # --------------------------------------------------------------------------- diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 11e901ad0..f9a0ea053 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.37.1-beta.0 + 0.38.0-beta.0 ``` diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index fa4e0748a..e4cbc1e96 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -386,6 +386,29 @@ Drop an existing table. *** +### dropTableAsync() + +```ts +abstract dropTableAsync(name, namespacePath?): Promise +``` + +Start dropping a table and return its cleanup job. + +The table may become unavailable before its data files are removed. Wait +on the returned job to know when cleanup has finished. + +#### Parameters + +* **name**: `string` + +* **namespacePath?**: `string`[] + +#### Returns + +`Promise`<[`Job`](Job.md)> + +*** + ### getJob() ```ts diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 11fca32d0..278559cc4 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -69,14 +69,33 @@ abstract addColumns(newColumnTransforms): Promise Add new columns with defined values. +The `{ computed }` form stores the expression rather than evaluating it +now: the column is committed with no values, and rows get them from +[Table#refreshColumn](Table.md#refreshcolumn). Declaring one therefore costs the same on a +large table as on an empty one. + +A refresh does not revisit rows it has already filled, so mutating an +input leaves the value computed at fill time; recomputing means dropping +the column and declaring it again. While a declaration reads a column, +that column cannot be renamed, retyped or dropped. + +Computed columns are local-only: LanceDB Cloud and Enterprise reject a +declaration. + #### Parameters -* **newColumnTransforms**: `Field`<`any`> \| `Field`<`any`>[] \| `Schema`<`any`> \| [`AddColumnsSql`](../interfaces/AddColumnsSql.md)[] +* **newColumnTransforms**: + \| `Field`<`any`> + \| `Field`<`any`>[] + \| `Schema`<`any`> + \| [`AddColumnsSql`](../interfaces/AddColumnsSql.md)[] + \| `object` Either: - An array of objects with column names and SQL expressions to calculate values - A single Arrow Field defining one column with its data type (column will be initialized with null values) - An array of Arrow Fields defining columns with their data types (columns will be initialized with null values) - An Arrow Schema defining columns with their data types (columns will be initialized with null values) + - `{ computed }`, declaring columns defined by a SQL expression whose type and inputs are derived from it #### Returns @@ -85,6 +104,13 @@ Add new columns with defined values. A promise that resolves to an object containing the new version number of the table after adding the columns. +#### Example + +```ts +await table.addColumns({ computed: [{ name: "doubled", valueSql: "x * 2" }] }); +const { rowsFilled } = await table.refreshColumn("doubled"); +``` + *** ### alterColumns() @@ -431,9 +457,10 @@ Read the [LsmWriteSpec](../interfaces/LsmWriteSpec.md) currently installed on th Resolves to `undefined` when the MemWAL LSM write path is not enabled (no spec has been set, or it was removed with [Table#unsetLsmWriteSpec](Table.md#unsetlsmwritespec)). -The returned spec — including its `maintainedIndexes` and -`writerConfigDefaults` — mirrors what was passed to -[Table#setLsmWriteSpec](Table.md#setlsmwritespec). +The returned spec mirrors what was passed to +[Table#setLsmWriteSpec](Table.md#setlsmwritespec), except that `maintainedIndexes` always +reports the concrete list resolved when the spec was set — `undefined` +never round-trips. #### Returns @@ -717,6 +744,32 @@ for await (const batch of table.query()) { *** +### refreshColumn() + +```ts +abstract refreshColumn(column): Promise +``` + +Fill the rows of a computed column that hold no value yet. + +Rows appended since the last refresh are filled by the next one; rows +already filled are left as they are, so the call is idempotent and does +not observe a mutated input. Local tables only. + +#### Parameters + +* **column**: `string` + The name of the computed column to fill. + +#### Returns + +`Promise`<[`RefreshColumnResult`](../interfaces/RefreshColumnResult.md)> + +A promise that resolves to the +number of rows filled and the new version number of the table. + +*** + ### restore() ```ts @@ -806,6 +859,11 @@ All variants require the table to have an unenforced primary key ([Table#setUnenforcedPrimaryKey](Table.md#setunenforcedprimarykey)); bucket sharding additionally requires it to be the single column being bucketed. +Omitting `maintainedIndexes` maintains every index on the table, resolved +here, failing if one cannot be maintained — name them to install anyway. +Naming them pins an exact set, and a still-building index is rejected +rather than quietly omitted. + #### Parameters * **spec**: [`LsmWriteSpec`](../interfaces/LsmWriteSpec.md) diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 7455a81ce..bd2ca54b5 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -105,6 +105,7 @@ - [OptimizeOptions](interfaces/OptimizeOptions.md) - [OptimizeStats](interfaces/OptimizeStats.md) - [QueryExecutionOptions](interfaces/QueryExecutionOptions.md) +- [RefreshColumnResult](interfaces/RefreshColumnResult.md) - [RemovalStats](interfaces/RemovalStats.md) - [RenameTableOptions](interfaces/RenameTableOptions.md) - [RestNamespaceConfig](interfaces/RestNamespaceConfig.md) diff --git a/docs/src/js/interfaces/LsmWriteSpec.md b/docs/src/js/interfaces/LsmWriteSpec.md index 8a588df6a..f2ae91186 100644 --- a/docs/src/js/interfaces/LsmWriteSpec.md +++ b/docs/src/js/interfaces/LsmWriteSpec.md @@ -34,7 +34,9 @@ Bucket and identity variants: the sharding column. optional maintainedIndexes: string[]; ``` -Names of indexes the MemWAL should keep up to date during writes. +Indexes the MemWAL keeps up to date. Omit to maintain every supported +index, resolved on install — a snapshot, so indexes created later are not +maintained. Pass `[]` for none. *** diff --git a/docs/src/js/interfaces/RefreshColumnResult.md b/docs/src/js/interfaces/RefreshColumnResult.md new file mode 100644 index 000000000..d2854fda6 --- /dev/null +++ b/docs/src/js/interfaces/RefreshColumnResult.md @@ -0,0 +1,23 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / RefreshColumnResult + +# Interface: RefreshColumnResult + +## Properties + +### rowsFilled + +```ts +rowsFilled: number; +``` + +*** + +### version + +```ts +version: number; +``` diff --git a/docs/src/js/interfaces/TableStatistics.md b/docs/src/js/interfaces/TableStatistics.md index e19cba119..e2e8ef34d 100644 --- a/docs/src/js/interfaces/TableStatistics.md +++ b/docs/src/js/interfaces/TableStatistics.md @@ -44,4 +44,7 @@ The number of rows in the table totalBytes: number; ``` -The total number of bytes in the table +The total size, in bytes, of the table's data files, index files, and +overlay files + +Read from the manifest, so this excludes deletion files and manifests. diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 3df2ac178..09b088e46 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.37.1-beta.0 + 0.38.0-beta.0 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index e1c39c3e2..9d9fe1f87 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.37.1-beta.0 + 0.38.0-beta.0 pom ${project.artifactId} LanceDB Java SDK Parent POM @@ -28,7 +28,7 @@ UTF-8 15.0.0 - 11.0.0-beta.3 + 11.0.0-beta.11 false 2.30.0 1.7 diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 10e12edc8..2e9373b9b 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.37.1-beta.0" +version = "0.38.0-beta.0" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/__test__/arrow.test.ts b/nodejs/__test__/arrow.test.ts index c05849cb9..29030d4f8 100644 --- a/nodejs/__test__/arrow.test.ts +++ b/nodejs/__test__/arrow.test.ts @@ -6,7 +6,9 @@ import * as arrow17 from "apache-arrow-17"; import * as arrow18 from "apache-arrow-18"; import { + Vector as CurrentVector, convertToTable, + tableFromIPC as currentTableFromIPC, fromBufferToRecordBatch, fromDataToBuffer, fromRecordBatchToBuffer, @@ -19,6 +21,7 @@ import { FunctionOptions, } from "../lancedb/embedding/embedding_function"; import { EmbeddingFunctionConfig } from "../lancedb/embedding/registry"; +import { sanitizeTable } from "../lancedb/sanitize"; // biome-ignore lint/suspicious/noExplicitAny: skip function sampleRecords(): Array> { @@ -64,7 +67,11 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( tableFromIPC, DataType, Dictionary, + RecordBatch: ArrowRecordBatch, + Table: ArrowTable, Uint8: ArrowUint8, + makeData: arrowMakeData, + vectorFromArray, // biome-ignore lint/suspicious/noExplicitAny: } = arrow; type Schema = ApacheArrow["Schema"]; @@ -1054,6 +1061,114 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( }); 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()), diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 68180471a..af471b478 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -89,6 +89,16 @@ describe("given a connection", () => { await db.createTable("test4", [{ id: 1 }, { id: 2 }]); }); + it("should return a completed job when dropping a local table", async () => { + await db.createTable("async-drop", [{ id: 1 }]); + + const job = await db.dropTableAsync("async-drop"); + expect(job.id).toBeNull(); + await expect(job.status()).resolves.toBe("finished"); + await job.wait(); + await expect(db.tableNames()).resolves.toEqual([]); + }); + it("should fail if creating table twice, unless overwrite is true", async () => { let tbl = await db.createTable("test", [{ id: 1 }, { id: 2 }]); await expect(tbl.countRows()).resolves.toBe(2); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 15d6e0804..bc495d24b 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -277,8 +277,16 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( }, numIndices: 0, numRows: 3, - totalBytes: 44, + // Full on-disk size of the two data files, footers and metadata included. + totalBytes: 684, }); + + // Index files count toward totalBytes too (only deletion files and + // manifests are excluded). + await table.createIndex("id", { config: Index.btree() }); + const statsWithIndex = await table.stats(); + expect(statsWithIndex.numIndices).toBe(1); + expect(statsWithIndex.totalBytes).toBeGreaterThan(684); }); it("should overwrite data if asked", async () => { @@ -3332,3 +3340,45 @@ describe("LSM merge insert", () => { await expect(table.query().useLsm(true).toArray()).rejects.toThrow(); }); }); + +describe("computed columns", () => { + let tmpDir: tmp.DirResult; + beforeEach(() => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + }); + afterEach(() => tmpDir.removeCallback()); + + it("declares a column and fills it on refresh", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable("computed", [{ x: 1 }, { x: 2 }]); + + await table.addColumns({ + computed: [{ name: "doubled", valueSql: "x * 2" }], + }); + let rows = await table.query().toArray(); + expect(rows.map((r) => r.doubled)).toEqual([null, null]); + + const result = await table.refreshColumn("doubled"); + expect(result.rowsFilled).toBe(2); + + rows = await table.query().toArray(); + expect(rows.map((r) => r.doubled).sort()).toEqual([2, 4]); + }); + + it("fills rows added since the last refresh", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable("computed_append", [{ x: 1 }]); + + await table.addColumns({ + computed: [{ name: "doubled", valueSql: "x * 2" }], + }); + await table.refreshColumn("doubled"); + await table.add([{ x: 5 }]); + + const result = await table.refreshColumn("doubled"); + expect(result.rowsFilled).toBe(1); + + const rows = await table.query().toArray(); + expect(rows.map((r) => r.doubled).sort()).toEqual([10, 2]); + }); +}); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index e63a7ae65..a81dc0442 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -327,6 +327,14 @@ export abstract class Connection { */ abstract dropTable(name: string, namespacePath?: string[]): Promise; + /** + * Start dropping a table and return its cleanup job. + * + * The table may become unavailable before its data files are removed. Wait + * on the returned job to know when cleanup has finished. + */ + abstract dropTableAsync(name: string, namespacePath?: string[]): Promise; + /** * Drop all tables in the database. * @param {string[]} namespacePath The namespace path to drop tables from (defaults to root namespace). @@ -705,6 +713,10 @@ export class LocalConnection extends Connection { return this.inner.dropTable(name, namespacePath ?? []); } + async dropTableAsync(name: string, namespacePath?: string[]): Promise { + return this.inner.dropTableAsync(name, namespacePath ?? []); + } + async dropAllTables(namespacePath?: string[]): Promise { return this.inner.dropAllTables(namespacePath ?? []); } diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 319222421..9f2e97989 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -50,6 +50,7 @@ export { MergeResult, AddResult, AddColumnsResult, + RefreshColumnResult, AlterColumnsResult, UpdateFieldMetadataResult, DeleteResult, diff --git a/nodejs/lancedb/sanitize.ts b/nodejs/lancedb/sanitize.ts index ae0bc0179..8fb2f1a0a 100644 --- a/nodejs/lancedb/sanitize.ts +++ b/nodejs/lancedb/sanitize.ts @@ -9,7 +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 { BufferType, Data } from "apache-arrow"; +import { BufferType, Data, Vector } from "apache-arrow"; import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type"; import { Binary, @@ -74,6 +74,20 @@ import { Utf8, } from "./arrow"; +type SanitizationContext = { + types: WeakMap; + vectors: WeakMap; + data: WeakMap>; +}; + +function createSanitizationContext(): SanitizationContext { + return { + types: new WeakMap(), + vectors: new WeakMap(), + data: new WeakMap(), + }; +} + export function sanitizeMetadata( metadataLike?: unknown, ): Map | undefined { @@ -186,6 +200,13 @@ export function sanitizeInterval(typeLike: object) { } export function sanitizeList(typeLike: object) { + return sanitizeListWithContext(typeLike, createSanitizationContext()); +} + +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", @@ -194,19 +215,35 @@ export function sanitizeList(typeLike: object) { if (typeLike.children.length !== 1) { throw Error("Expected a List type to have exactly one child"); } - return new List(sanitizeField(typeLike.children[0])); + return new List(sanitizeFieldWithContext(typeLike.children[0], context)); } export function sanitizeStruct(typeLike: object) { + return sanitizeStructWithContext(typeLike, createSanitizationContext()); +} + +function sanitizeStructWithContext( + typeLike: object, + context: SanitizationContext, +) { if (!("children" in typeLike) || !Array.isArray(typeLike.children)) { throw Error( "Expected a Struct type to have an array-like `children` property", ); } - return new Struct(typeLike.children.map((child) => sanitizeField(child))); + return new Struct( + typeLike.children.map((child) => sanitizeFieldWithContext(child, context)), + ); } export function sanitizeUnion(typeLike: object) { + return sanitizeUnionWithContext(typeLike, createSanitizationContext()); +} + +function sanitizeUnionWithContext( + typeLike: object, + context: SanitizationContext, +) { if ( !("typeIds" in typeLike) || !("mode" in typeLike) || @@ -226,7 +263,7 @@ export function sanitizeUnion(typeLike: object) { typeLike.mode, // biome-ignore lint/suspicious/noExplicitAny: skip typeLike.typeIds as any, - typeLike.children.map((child) => sanitizeField(child)), + typeLike.children.map((child) => sanitizeFieldWithContext(child, context)), ); } @@ -234,6 +271,19 @@ export function sanitizeTypedUnion( typeLike: object, // eslint-disable-next-line @typescript-eslint/naming-convention UnionType: typeof DenseUnion | typeof SparseUnion, +) { + return sanitizeTypedUnionWithContext( + typeLike, + UnionType, + createSanitizationContext(), + ); +} + +function sanitizeTypedUnionWithContext( + typeLike: object, + // eslint-disable-next-line @typescript-eslint/naming-convention + UnionType: typeof DenseUnion | typeof SparseUnion, + context: SanitizationContext, ) { if (!("typeIds" in typeLike)) { throw Error( @@ -248,7 +298,7 @@ export function sanitizeTypedUnion( return new UnionType( typeLike.typeIds as Int32Array | number[], - typeLike.children.map((child) => sanitizeField(child)), + typeLike.children.map((child) => sanitizeFieldWithContext(child, context)), ); } @@ -262,6 +312,16 @@ export function sanitizeFixedSizeBinary(typeLike: object) { } export function sanitizeFixedSizeList(typeLike: object) { + return sanitizeFixedSizeListWithContext( + typeLike, + createSanitizationContext(), + ); +} + +function sanitizeFixedSizeListWithContext( + typeLike: object, + context: SanitizationContext, +) { if (!("listSize" in typeLike) || typeof typeLike.listSize !== "number") { throw Error("Expected a FixedSizeList type to have a `listSize` property"); } @@ -275,11 +335,18 @@ export function sanitizeFixedSizeList(typeLike: object) { } return new FixedSizeList( typeLike.listSize, - sanitizeField(typeLike.children[0]), + sanitizeFieldWithContext(typeLike.children[0], context), ); } export function sanitizeMap(typeLike: object) { + return sanitizeMapWithContext(typeLike, createSanitizationContext()); +} + +function sanitizeMapWithContext( + typeLike: object, + context: SanitizationContext, +) { if (!("children" in typeLike) || !Array.isArray(typeLike.children)) { throw Error( "Expected a Map type to have an array-like `children` property", @@ -292,7 +359,10 @@ export function sanitizeMap(typeLike: object) { throw Error("Expected a Map type to have exactly one child"); } - return new Map_(sanitizeField(typeLike.children[0]), typeLike.keysSorted); + return new Map_( + sanitizeFieldWithContext(typeLike.children[0], context), + typeLike.keysSorted, + ); } export function sanitizeDuration(typeLike: object) { @@ -303,6 +373,13 @@ export function sanitizeDuration(typeLike: object) { } export function sanitizeDictionary(typeLike: object) { + return sanitizeDictionaryWithContext(typeLike, createSanitizationContext()); +} + +function sanitizeDictionaryWithContext( + typeLike: object, + context: SanitizationContext, +) { if (!("id" in typeLike) || typeof typeLike.id !== "number") { throw Error("Expected a Dictionary type to have an `id` property"); } @@ -316,8 +393,8 @@ export function sanitizeDictionary(typeLike: object) { throw Error("Expected a Dictionary type to have an `isOrdered` property"); } return new Dictionary( - sanitizeType(typeLike.dictionary), - sanitizeType(typeLike.indices) as TKeys, + sanitizeTypeWithContext(typeLike.dictionary, context), + sanitizeTypeWithContext(typeLike.indices, context) as TKeys, typeLike.id, typeLike.isOrdered, ); @@ -325,12 +402,23 @@ export function sanitizeDictionary(typeLike: object) { // biome-ignore lint/suspicious/noExplicitAny: skip export function sanitizeType(typeLike: unknown): DataType { + return sanitizeTypeWithContext(typeLike, createSanitizationContext()); +} + +function sanitizeTypeWithContext( + typeLike: unknown, + context: SanitizationContext, +): DataType { if (typeof typeLike === "string") { return dataTypeFromName(typeLike); } if (typeof typeLike !== "object" || typeLike === null) { throw Error("Expected a Type but object was null/undefined"); } + const cached = context.types.get(typeLike); + if (cached !== undefined) { + return cached; + } if ( !("typeId" in typeLike) || !( @@ -349,6 +437,16 @@ export function sanitizeType(typeLike: unknown): DataType { throw Error("Type's typeId property was not a function or number"); } + const type = sanitizeTypeById(typeLike, typeId, context); + context.types.set(typeLike, type); + return type; +} + +function sanitizeTypeById( + typeLike: object, + typeId: Type, + context: SanitizationContext, +): DataType { switch (typeId) { case Type.NONE: throw Error("Received a Type with a typeId of NONE"); @@ -375,21 +473,21 @@ export function sanitizeType(typeLike: unknown): DataType { case Type.Interval: return sanitizeInterval(typeLike); case Type.List: - return sanitizeList(typeLike); + return sanitizeListWithContext(typeLike, context); case Type.Struct: - return sanitizeStruct(typeLike); + return sanitizeStructWithContext(typeLike, context); case Type.Union: - return sanitizeUnion(typeLike); + return sanitizeUnionWithContext(typeLike, context); case Type.FixedSizeBinary: return sanitizeFixedSizeBinary(typeLike); case Type.FixedSizeList: - return sanitizeFixedSizeList(typeLike); + return sanitizeFixedSizeListWithContext(typeLike, context); case Type.Map: - return sanitizeMap(typeLike); + return sanitizeMapWithContext(typeLike, context); case Type.Duration: return sanitizeDuration(typeLike); case Type.Dictionary: - return sanitizeDictionary(typeLike); + return sanitizeDictionaryWithContext(typeLike, context); case Type.Int8: return new Int8(); case Type.Int16: @@ -433,9 +531,9 @@ export function sanitizeType(typeLike: unknown): DataType { case Type.TimestampSecond: return sanitizeTypedTimestamp(typeLike, TimestampSecond); case Type.DenseUnion: - return sanitizeTypedUnion(typeLike, DenseUnion); + return sanitizeTypedUnionWithContext(typeLike, DenseUnion, context); case Type.SparseUnion: - return sanitizeTypedUnion(typeLike, SparseUnion); + return sanitizeTypedUnionWithContext(typeLike, SparseUnion, context); case Type.IntervalDayTime: return new IntervalDayTime(); case Type.IntervalYearMonth: @@ -454,6 +552,13 @@ export function sanitizeType(typeLike: unknown): DataType { } export function sanitizeField(fieldLike: unknown): Field { + return sanitizeFieldWithContext(fieldLike, createSanitizationContext()); +} + +function sanitizeFieldWithContext( + fieldLike: unknown, + context: SanitizationContext, +): Field { if (fieldLike instanceof Field) { return fieldLike; } @@ -471,7 +576,7 @@ export function sanitizeField(fieldLike: unknown): Field { } let type: DataType; try { - type = sanitizeType(fieldLike.type); + type = sanitizeTypeWithContext(fieldLike.type, context); } catch (error: unknown) { throw Error( `Unable to sanitize type for field: ${fieldLike.name} due to error: ${error}`, @@ -501,6 +606,13 @@ export function sanitizeField(fieldLike: unknown): Field { * than lancedb is using. */ export function sanitizeSchema(schemaLike: SchemaLike): Schema { + return sanitizeSchemaWithContext(schemaLike, createSanitizationContext()); +} + +function sanitizeSchemaWithContext( + schemaLike: SchemaLike, + context: SanitizationContext, +): Schema { if (schemaLike instanceof Schema) { return schemaLike; } @@ -522,7 +634,7 @@ export function sanitizeSchema(schemaLike: SchemaLike): Schema { ); } const sanitizedFields = schemaLike.fields.map((field) => - sanitizeField(field), + sanitizeFieldWithContext(field, context), ); return new Schema(sanitizedFields, metadata); } @@ -544,13 +656,18 @@ export function sanitizeTable(tableLike: TableLike): Table { "The table passed in does not appear to be a table (no 'columns' property)", ); } - const schema = sanitizeSchema(tableLike.schema); - - const batches = tableLike.batches.map(sanitizeRecordBatch); + const context = createSanitizationContext(); + const schema = sanitizeSchemaWithContext(tableLike.schema, context); + const batches = tableLike.batches.map((batch) => + sanitizeRecordBatch(batch, context), + ); return new Table(schema, batches); } -function sanitizeRecordBatch(batchLike: RecordBatchLike): RecordBatch { +function sanitizeRecordBatch( + batchLike: RecordBatchLike, + context: SanitizationContext, +): RecordBatch { if (batchLike instanceof RecordBatch) { return batchLike; } @@ -567,19 +684,43 @@ function sanitizeRecordBatch(batchLike: RecordBatchLike): RecordBatch { "The record batch passed in does not appear to be a record batch (no 'data' property)", ); } - const schema = sanitizeSchema(batchLike.schema); - const data = sanitizeData(batchLike.data); + const schema = sanitizeSchemaWithContext(batchLike.schema, context); + const data = sanitizeData(batchLike.data, context) as Data; return new RecordBatch(schema, data); } + +type DictionaryVectorLike = { + data: readonly DataLike[]; +}; + +type DictionaryDataLike = DataLike & { + dictionary?: DictionaryVectorLike; +}; + function sanitizeData( dataLike: DataLike, - // biome-ignore lint/suspicious/noExplicitAny: -): import("apache-arrow").Data> { + context: SanitizationContext, +): Data { if (dataLike instanceof Data) { return dataLike; } - return new Data( - dataLike.type, + const cachedData = context.data.get(dataLike); + if (cachedData !== undefined) { + return cachedData; + } + const dictionaryLike = (dataLike as DictionaryDataLike).dictionary; + let dictionary: Vector | undefined; + if (dictionaryLike !== undefined) { + dictionary = context.vectors.get(dictionaryLike); + if (dictionary === undefined) { + dictionary = new Vector( + dictionaryLike.data.map((data) => sanitizeData(data, context)), + ); + context.vectors.set(dictionaryLike, dictionary); + } + } + const data = new Data( + sanitizeTypeWithContext(dataLike.type, context), dataLike.offset, dataLike.length, dataLike.nullCount, @@ -589,7 +730,11 @@ function sanitizeData( [BufferType.VALIDITY]: dataLike.nullBitmap, [BufferType.TYPE]: dataLike.typeIds, }, + dataLike.children.map((child) => sanitizeData(child, context)), + dictionary, ); + context.data.set(dataLike, data); + return data; } const constructorsByTypeName = { diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 3359a2643..5b8d00076 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -33,6 +33,7 @@ import { Job, Branches as NativeBranches, OptimizeStats, + RefreshColumnResult, TableStatistics, Tags, UpdateFieldMetadataResult, @@ -197,7 +198,11 @@ export interface LsmWriteSpec { column?: string; /** Bucket variant: the number of buckets, in `[1, 1024]`. */ numBuckets?: number; - /** Names of indexes the MemWAL should keep up to date during writes. */ + /** + * Indexes the MemWAL keeps up to date. Omit to maintain every supported + * index, resolved on install — a snapshot, so indexes created later are not + * maintained. Pass `[]` for none. + */ maintainedIndexes?: string[]; /** Default `ShardWriter` configuration recorded in the MemWAL index. */ writerConfigDefaults?: Record; @@ -521,18 +526,54 @@ export abstract class Table { abstract vectorSearch(vector: IntoVector | MultiVector): VectorQuery; /** * Add new columns with defined values. + * + * The `{ computed }` form stores the expression rather than evaluating it + * now: the column is committed with no values, and rows get them from + * {@link Table#refreshColumn}. Declaring one therefore costs the same on a + * large table as on an empty one. + * + * A refresh does not revisit rows it has already filled, so mutating an + * input leaves the value computed at fill time; recomputing means dropping + * the column and declaring it again. While a declaration reads a column, + * that column cannot be renamed, retyped or dropped. + * + * Computed columns are local-only: LanceDB Cloud and Enterprise reject a + * declaration. * @param {AddColumnsSql[] | Field | Field[] | Schema} newColumnTransforms Either: * - An array of objects with column names and SQL expressions to calculate values * - A single Arrow Field defining one column with its data type (column will be initialized with null values) * - An array of Arrow Fields defining columns with their data types (columns will be initialized with null values) * - An Arrow Schema defining columns with their data types (columns will be initialized with null values) + * - `{ computed }`, declaring columns defined by a SQL expression whose type and inputs are derived from it * @returns {Promise} A promise that resolves to an object * containing the new version number of the table after adding the columns. + * @example + * ```ts + * await table.addColumns({ computed: [{ name: "doubled", valueSql: "x * 2" }] }); + * const { rowsFilled } = await table.refreshColumn("doubled"); + * ``` */ abstract addColumns( - newColumnTransforms: AddColumnsSql[] | Field | Field[] | Schema, + newColumnTransforms: + | AddColumnsSql[] + | Field + | Field[] + | Schema + | { computed: AddColumnsSql[] }, ): Promise; + /** + * Fill the rows of a computed column that hold no value yet. + * + * Rows appended since the last refresh are filled by the next one; rows + * already filled are left as they are, so the call is idempotent and does + * not observe a mutated input. Local tables only. + * @param {string} column The name of the computed column to fill. + * @returns {Promise} A promise that resolves to the + * number of rows filled and the new version number of the table. + */ + abstract refreshColumn(column: string): Promise; + /** * Alter the name or nullability of columns. * @param {ColumnAlteration[]} columnAlterations One or more alterations to @@ -595,6 +636,11 @@ export abstract class Table { * All variants require the table to have an unenforced primary key * ({@link Table#setUnenforcedPrimaryKey}); bucket sharding additionally * requires it to be the single column being bucketed. + * + * Omitting `maintainedIndexes` maintains every index on the table, resolved + * here, failing if one cannot be maintained — name them to install anyway. + * Naming them pins an exact set, and a still-building index is rejected + * rather than quietly omitted. * @param {LsmWriteSpec} spec The sharding spec to install. * @returns {Promise} * @example @@ -622,9 +668,10 @@ export abstract class Table { * * Resolves to `undefined` when the MemWAL LSM write path is not enabled (no * spec has been set, or it was removed with {@link Table#unsetLsmWriteSpec}). - * The returned spec — including its `maintainedIndexes` and - * `writerConfigDefaults` — mirrors what was passed to - * {@link Table#setLsmWriteSpec}. + * The returned spec mirrors what was passed to + * {@link Table#setLsmWriteSpec}, except that `maintainedIndexes` always + * reports the concrete list resolved when the spec was set — `undefined` + * never round-trips. * @returns {Promise} */ abstract getLsmWriteSpec(): Promise; @@ -1078,8 +1125,22 @@ export class LocalTable extends Table { // TODO: Support BatchUDF async addColumns( - newColumnTransforms: AddColumnsSql[] | Field | Field[] | Schema, + newColumnTransforms: + | AddColumnsSql[] + | Field + | Field[] + | Schema + | { computed: AddColumnsSql[] }, ): Promise { + // Columns defined by an expression are declared, not materialized here. + if ( + typeof newColumnTransforms === "object" && + !Array.isArray(newColumnTransforms) && + "computed" in newColumnTransforms + ) { + return await this.inner.addComputedColumns(newColumnTransforms.computed); + } + // Handle single Field -> convert to array of Fields if (newColumnTransforms instanceof Field) { newColumnTransforms = [newColumnTransforms]; @@ -1114,6 +1175,10 @@ export class LocalTable extends Table { throw new Error("Invalid input type for addColumns"); } + async refreshColumn(column: string): Promise { + return await this.inner.refreshColumn(column); + } + async alterColumns( columnAlterations: ColumnAlteration[], ): Promise { diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 3c93ed470..e0fd7426d 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.37.1-beta.0", + "version": "0.38.0-beta.0", "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 5ade5aaa3..ef281de3d 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.37.1-beta.0", + "version": "0.38.0-beta.0", "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 16bb0edd0..d535820fa 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.37.1-beta.0", + "version": "0.38.0-beta.0", "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 6ee11e4bc..7aa21301e 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.37.1-beta.0", + "version": "0.38.0-beta.0", "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 c2e15bb9f..d220991f7 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.37.1-beta.0", + "version": "0.38.0-beta.0", "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 d2820b1a1..519d7376a 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.37.1-beta.0", + "version": "0.38.0-beta.0", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 601b51380..9f608d6d0 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.37.1-beta.0", + "version": "0.38.0-beta.0", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package-lock.json b/nodejs/package-lock.json index bdbd3cf79..9222bf582 100644 --- a/nodejs/package-lock.json +++ b/nodejs/package-lock.json @@ -1,12 +1,12 @@ { "name": "@lancedb/lancedb", - "version": "0.37.1-beta.0", + "version": "0.38.0-beta.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@lancedb/lancedb", - "version": "0.37.1-beta.0", + "version": "0.38.0-beta.0", "cpu": [ "x64", "arm64" diff --git a/nodejs/package.json b/nodejs/package.json index 671f3f94d..c87af926b 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.37.1-beta.0", + "version": "0.38.0-beta.0", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index c45321aba..c9f5e10ea 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -334,6 +334,22 @@ impl Connection { .default_error() } + /// Start dropping a table and return its cleanup job. + #[napi(catch_unwind)] + pub async fn drop_table_async( + &self, + name: String, + namespace_path: Option>, + ) -> napi::Result { + let ns = namespace_path.unwrap_or_default(); + let job = self + .get_inner()? + .drop_table_async(&name, &ns) + .await + .default_error()?; + Ok(crate::job::Job::new(job)) + } + #[napi(catch_unwind)] pub async fn drop_all_tables(&self, namespace_path: Option>) -> napi::Result<()> { let ns = namespace_path.unwrap_or_default(); diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 2ac2fecb2..40ed7d9f0 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -347,6 +347,30 @@ impl Table { Ok(res.into()) } + #[napi(catch_unwind)] + pub async fn add_computed_columns( + &self, + columns: Vec, + ) -> napi::Result { + let table = self.inner_ref()?; + let mut builder = table.add_columns(); + for column in columns { + builder = builder.computed(column.name, column.value_sql); + } + let res = builder.execute().await.default_error()?; + Ok(res.into()) + } + + #[napi(catch_unwind)] + pub async fn refresh_column(&self, column: String) -> napi::Result { + let res = self + .inner_ref()? + .refresh_column(column) + .await + .default_error()?; + Ok(res.into()) + } + #[napi(catch_unwind)] pub async fn add_columns_with_schema( &self, @@ -772,7 +796,8 @@ pub struct LsmWriteSpec { pub column: Option, /// Bucket variant: the number of buckets, in `[1, 1024]`. pub num_buckets: Option, - /// Names of indexes the MemWAL should keep up to date during writes. + /// Indexes the MemWAL keeps up to date. Omitted resolves every + /// maintainable index on install; an empty array means none. pub maintained_indexes: Option>, /// Default `ShardWriter` configuration recorded in the MemWAL index. pub writer_config_defaults: Option>, @@ -782,7 +807,6 @@ impl TryFrom for lancedb::table::LsmWriteSpec { type Error = napi::Error; fn try_from(value: LsmWriteSpec) -> napi::Result { - let maintained = value.maintained_indexes.unwrap_or_default(); let writer_config_defaults = value.writer_config_defaults.unwrap_or_default(); let spec = match value.spec_type.as_str() { "bucket" => { @@ -809,7 +833,7 @@ impl TryFrom for lancedb::table::LsmWriteSpec { } }; Ok(spec - .with_maintained_indexes(maintained) + .with_maintained_indexes(value.maintained_indexes) .with_writer_config_defaults(writer_config_defaults)) } } @@ -827,7 +851,7 @@ impl From for LsmWriteSpec { spec_type: "bucket".to_string(), column: Some(column), num_buckets: Some(num_buckets), - maintained_indexes: Some(maintained_indexes), + maintained_indexes, writer_config_defaults: Some(writer_config_defaults), }, Native::Identity { @@ -838,7 +862,7 @@ impl From for LsmWriteSpec { spec_type: "identity".to_string(), column: Some(column), num_buckets: None, - maintained_indexes: Some(maintained_indexes), + maintained_indexes, writer_config_defaults: Some(writer_config_defaults), }, Native::Unsharded { @@ -848,7 +872,7 @@ impl From for LsmWriteSpec { spec_type: "unsharded".to_string(), column: None, num_buckets: None, - maintained_indexes: Some(maintained_indexes), + maintained_indexes, writer_config_defaults: Some(writer_config_defaults), }, } @@ -1043,7 +1067,10 @@ impl From for IndexStatistics { #[napi(object)] pub struct TableStatistics { - /// The total number of bytes in the table + /// The total size, in bytes, of the table's data files, index files, and + /// overlay files + /// + /// Read from the manifest, so this excludes deletion files and manifests. pub total_bytes: i64, /// The number of rows in the table @@ -1193,6 +1220,21 @@ pub struct AddColumnsResult { pub version: i64, } +#[napi(object)] +pub struct RefreshColumnResult { + pub rows_filled: i64, + pub version: i64, +} + +impl From for RefreshColumnResult { + fn from(value: lancedb::table::RefreshColumnResult) -> Self { + Self { + rows_filled: value.rows_filled as i64, + version: value.version as i64, + } + } +} + impl From for AddColumnsResult { fn from(value: lancedb::table::AddColumnsResult) -> Self { Self { diff --git a/python/Cargo.toml b/python/Cargo.toml index 5a196e27c..bede2bc37 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.37.1-beta.0" +version = "0.38.0-beta.0" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 69efb2b62..6ea0faab0 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -198,6 +198,9 @@ class Connection(object): async def drop_table( self, name: str, namespace_path: Optional[List[str]] = None ) -> None: ... + async def drop_table_async( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> Job: ... async def drop_all_tables( self, namespace_path: Optional[List[str]] = None ) -> None: ... @@ -335,6 +338,10 @@ class Table: ) -> list[FtsToken]: ... async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ... async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ... + async def add_computed_columns( + self, columns: list[tuple[str, str]] + ) -> AddColumnsResult: ... + async def refresh_column(self, column: str) -> RefreshColumnResult: ... async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ... async def alter_columns( self, columns: list[dict[str, Any]] @@ -654,9 +661,10 @@ class LsmWriteSpec: def identity(column: str) -> "LsmWriteSpec": ... @staticmethod def unsharded() -> "LsmWriteSpec": ... - def with_maintained_indexes(self, indexes: List[str]) -> "LsmWriteSpec": - """Return a copy of this spec asking the MemWAL to keep the named - indexes up to date as rows are appended.""" + def with_maintained_indexes(self, indexes: Optional[List[str]]) -> "LsmWriteSpec": + """Set which indexes the MemWAL keeps up to date. None resolves every + index on the table at install, failing if one cannot be maintained; + a list is verbatim, empty means none.""" ... def with_writer_config_defaults(self, defaults: Dict[str, str]) -> "LsmWriteSpec": """Return a copy of this spec recording the given default @@ -671,13 +679,19 @@ class LsmWriteSpec: @property def num_buckets(self) -> Optional[int]: ... @property - def maintained_indexes(self) -> List[str]: ... + def maintained_indexes(self) -> Optional[List[str]]: + """Indexes the MemWAL keeps up to date, or None for every supported one.""" + ... @property def writer_config_defaults(self) -> Dict[str, str]: ... class AddColumnsResult: version: int +class RefreshColumnResult: + rows_filled: int + version: int + class AlterColumnsResult: version: int diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index eeae8bf50..14b6c0b0d 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -524,6 +524,12 @@ class DBConnection(EnforceOverrides): namespace_path = [] raise NotImplementedError + def drop_table_async( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> Job: + """Start dropping a table and return its cleanup job.""" + raise NotImplementedError + def rename_table( self, cur_name: str, @@ -1186,6 +1192,20 @@ class LanceDBConnection(DBConnection): ) ) + @override + def drop_table_async( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> Job: + """Start dropping a table and return its cleanup job. + + The table may become unavailable before its data files are removed. + Call :meth:`Job.wait` to wait for cleanup to finish. + """ + if namespace_path is None: + namespace_path = [] + job = LOOP.run(self._conn.drop_table_async(name, namespace_path=namespace_path)) + return Job(job if isinstance(job, AsyncJob) else AsyncJob(job)) + @override def drop_all_tables(self, namespace_path: Optional[List[str]] = None): if namespace_path is None: @@ -1963,6 +1983,23 @@ class AsyncConnection(object): if f"Table '{name}' was not found" not in str(e): raise e + async def drop_table_async( + self, + name: str, + *, + namespace_path: Optional[List[str]] = None, + ) -> AsyncJob: + """Start dropping a table and return its cleanup job. + + The table may become unavailable before its data files are removed. + Await :meth:`AsyncJob.wait` to wait for cleanup to finish. + """ + if namespace_path is None: + namespace_path = [] + return AsyncJob( + await self._inner.drop_table_async(name, namespace_path=namespace_path) + ) + async def drop_all_tables(self, namespace_path: Optional[List[str]] = None): """Drop all tables from the database. diff --git a/python/python/lancedb/embeddings/jinaai.py b/python/python/lancedb/embeddings/jinaai.py index 9656f041f..f6ab601b3 100644 --- a/python/python/lancedb/embeddings/jinaai.py +++ b/python/python/lancedb/embeddings/jinaai.py @@ -87,12 +87,13 @@ class JinaEmbeddings(EmbeddingFunction): if isinstance(image, bytes): image_dict = {"image": base64.b64encode(image).decode("utf-8")} elif isinstance(image, (str, Path)): - parsed = urlparse.urlparse(image) - # TODO handle drive letter on windows. + parsed = urlparse(str(image)) PIL_Image = attempt_import_or_raise("PIL.Image", "pillow") if parsed.scheme == "file": pil_image = PIL_Image.open(parsed.path) - elif parsed.scheme == "": + elif parsed.scheme == "" or (os.name == "nt" and len(parsed.scheme) == 1): + # A Windows drive letter parses as a one-character scheme + # ("C:\\img.png" -> scheme="c"), so treat it as a local path. pil_image = PIL_Image.open(image if os.name == "nt" else parsed.path) elif parsed.scheme.startswith("http"): pil_image = PIL_Image.open(io.BytesIO(url_retrieve(image))) diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index b151395cc..0e60bd218 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -49,6 +49,7 @@ from lancedb._lancedb import ( ) from lancedb.background_loop import LOOP from lancedb.db import AsyncConnection, DBConnection +from lancedb.job import AsyncJob, Job from lance_namespace import ( LanceNamespace, connect as namespace_connect, @@ -624,6 +625,18 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path = [] LOOP.run(self._inner.drop_table(name, namespace_path=namespace_path)) + @override + def drop_table_async( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> Job: + """Start dropping a table and return its cleanup job.""" + if namespace_path is None: + namespace_path = [] + job = LOOP.run( + self._inner.drop_table_async(name, namespace_path=namespace_path) + ) + return Job(job if isinstance(job, AsyncJob) else AsyncJob(job)) + @override def rename_table( self, @@ -1134,6 +1147,14 @@ class AsyncLanceNamespaceDBConnection: namespace_path = [] await self._inner.drop_table(name, namespace_path=namespace_path) + async def drop_table_async( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> AsyncJob: + """Start dropping a table and return its cleanup job.""" + if namespace_path is None: + namespace_path = [] + return await self._inner.drop_table_async(name, namespace_path=namespace_path) + async def rename_table( self, cur_name: str, diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 332886590..16ad65dcb 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -23,7 +23,7 @@ import pyarrow as pa from ..common import DATA from ..db import DBConnection, LOOP -from ..job import Job +from ..job import AsyncJob, Job if TYPE_CHECKING: from .._lancedb import JobDescription, JobInfo @@ -663,6 +663,16 @@ class RemoteDBConnection(DBConnection): namespace_path = [] LOOP.run(self._conn.drop_table(name, namespace_path=namespace_path)) + @override + def drop_table_async( + self, name: str, namespace_path: Optional[List[str]] = None + ) -> Job: + """Start dropping a table and return its cleanup job.""" + if namespace_path is None: + namespace_path = [] + job = LOOP.run(self._conn.drop_table_async(name, namespace_path=namespace_path)) + return Job(job if isinstance(job, AsyncJob) else AsyncJob(job)) + @override def rename_table( self, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 71bebc78f..c4dc5ab39 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -968,9 +968,21 @@ class RemoteTable(Table): def count_rows(self, filter: Optional[str] = None) -> int: return LOOP.run(self._table.count_rows(filter)) - def add_columns(self, transforms: Dict[str, str]) -> AddColumnsResult: + def add_columns( + self, + transforms: Dict[str, str] | None = None, + *, + computed: Dict[str, str] | None = None, + ) -> AddColumnsResult: + if computed: + raise NotImplementedError( + "computed columns are supported only on local tables" + ) return LOOP.run(self._table.add_columns(transforms)) + def refresh_column(self, column: str): + raise NotImplementedError("computed columns are supported only on local tables") + def alter_columns( self, *alterations: Iterable[Dict[str, str]] ) -> AlterColumnsResult: diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 525ed3d63..b27e606a4 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -11,6 +11,11 @@ Provides StreamingDataset, a PyTorch IterableDataset that guarantees: - **Resumability**: state_dict / load_state_dict capture per-split consumption counts so training can resume from an exact mid-epoch position even when the distributed topology changes between runs. + +Transform failures on bad rows (e.g. nulls or NaNs from incomplete data) can +be tolerated with ``on_transform_error="skip"``; see the parameter +documentation on StreamingDataset for how this interacts with the guarantees +above. """ import ctypes @@ -22,7 +27,7 @@ import time from collections import deque from concurrent.futures import ThreadPoolExecutor from multiprocessing import RawArray -from typing import Any, Callable, Iterator, Optional +from typing import Any, Callable, Iterator, Optional, Union from torch.utils.data import IterableDataset, get_worker_info @@ -127,6 +132,49 @@ class StreamingDataset(IterableDataset): Maximum number of transforms to run concurrently. Must be greater than zero. When ``None`` (the default), uses ``os.cpu_count()`` or 1 when the CPU count is unavailable. + on_transform_error: + What to do when the transform raises an exception: + + - ``"raise"`` (the default): the exception propagates and iteration + aborts. + - ``"skip"``: the failing rows are dropped and iteration continues. + - ``"warn"``: like ``"skip"``, but a warning is logged for each + failing batch. + - a callable ``handler(exc) -> bool``: called with the exception; + return ``True`` to skip the failing rows or ``False`` to re-raise. + Useful to skip only expected error types (compatible with + ``webdataset.handlers`` style handlers). + + When a batch fails, the transform is re-invoked on each single-row + slice of the batch so that only the rows that actually fail are + dropped. Transforms should therefore be deterministic and accept + batches of any size (including one row). Skipped rows are counted in + ``rows_skipped``. + + Skipping weakens the elastic-determinism guarantee at the end of the + epoch: splits that lose more rows than others run dry earlier, and + each rank's iterator ends at the last cycle where every split *it + owns* still has a row. Because bad rows are not distributed evenly + across splits, this means one rank's iterator can yield noticeably + fewer or more steps than another rank's *in the same run* — there is + no cross-rank coordination that stops every rank at the same global + step. This is generally safe for asynchronous or single-rank use, + but synchronous distributed training (e.g. ranks that call + ``all_reduce`` every step) can hang or deadlock if one rank's + iterator is exhausted while others are still stepping; callers doing + synchronous multi-rank training with ``on_transform_error != "raise"`` + are responsible for their own cross-rank stopping mechanism (e.g. + broadcasting a stop signal on ``StopIteration``). The final few + global steps can also differ across topologies (bounded by the skew + in bad-row counts across splits). The sequence of samples yielded + from each split remains deterministic. Mid-epoch + checkpoints remain exact provided the transform fails + deterministically; in multi-rank training each rank must save its + own ``state_dict`` and the states must be combined with + ``merge_state_dicts`` before resuming on a different topology. + Prefer the ``filter`` parameter when bad rows can be expressed as a + SQL predicate (e.g. ``"col IS NOT NULL"``) — filtering happens before + splits are built, so every guarantee is fully preserved. worker_info_override: If set, used in place of ``torch.utils.data.get_worker_info()`` to determine the DataLoader worker assignment. Intended for unit tests @@ -152,6 +200,7 @@ class StreamingDataset(IterableDataset): filter: Optional[str] = None, transform: Optional[Callable] = None, transform_parallelism: Optional[int] = None, + on_transform_error: Union[str, Callable[[Exception], bool]] = "raise", connection_factory: Optional[Callable[[str], Any]] = None, worker_info_override=None, ): @@ -167,6 +216,13 @@ class StreamingDataset(IterableDataset): ) if transform_parallelism is not None and transform_parallelism <= 0: raise ValueError("transform_parallelism must be greater than 0") + if on_transform_error not in ("raise", "skip", "warn") and not callable( + on_transform_error + ): + raise ValueError( + "on_transform_error must be 'raise', 'skip', 'warn', or a " + f"callable, got {on_transform_error!r}" + ) self._table = table self._num_splits = num_splits @@ -182,6 +238,7 @@ class StreamingDataset(IterableDataset): self._filter = filter self._transform = transform self._transform_parallelism = transform_parallelism + self._on_transform_error = on_transform_error self._connection_factory = connection_factory self._worker_info_override = worker_info_override @@ -199,19 +256,28 @@ class StreamingDataset(IterableDataset): # in the main process. RawArray is picklable via the forkserver # reduction protocol so it survives the dataset pickle round-trip. # Layout: [unscanned_rows, raw_rows, cooked_rows, consumed_rows, - # bytes_loaded, fetch_time_us, transform_time_us] - self._worker_stats: RawArray = RawArray(ctypes.c_int64, 7) + # bytes_loaded, fetch_time_us, transform_time_us, + # rows_skipped] + self._worker_stats: RawArray = RawArray(ctypes.c_int64, 8) # Cumulative bytes of Arrow buffer data fetched across all iterations. self._bytes_loaded: int = 0 # Cumulative seconds spent in LanceDB I/O and in transform functions. self._fetch_time: float = 0.0 self._transform_time: float = 0.0 + # Cumulative rows dropped by on_transform_error across all iterations. + self._rows_skipped: int = 0 # Number of samples each split has already been consumed. At global # step boundaries all splits have consumed this many samples, so a # single scalar captures the topology-independent checkpoint state. self._resume_offset: int = 0 + # Permutation position each split has consumed through, keyed by + # global split index. Equal to _resume_offset for every split unless + # on_transform_error skipped rows, in which case skipped positions + # push the watermark of the affected splits further ahead. Splits + # this instance has never iterated have no entry. + self._resume_positions: dict[int, int] = {} # Build the permutation table once, deterministically. builder = permutation_builder(table) @@ -275,6 +341,7 @@ class StreamingDataset(IterableDataset): # Set identity transform on each Permutation so __getitems__ returns # the raw RecordBatch. Stage 2 applies the real transform. permutations: list[Permutation] = [] + initial_positions: list[int] = [] for split_idx in my_splits: perm = Permutation.from_tables( self._table, self._perm_table, split=split_idx @@ -282,14 +349,20 @@ class StreamingDataset(IterableDataset): if self._columns is not None: perm = perm.select_columns(self._columns) perm = perm.with_transform(lambda batch: batch) - if self._resume_offset > 0: - perm = perm.with_skip(self._resume_offset) + start_pos = self._resume_positions.get(split_idx, self._resume_offset) + if start_pos > 0: + perm = perm.with_skip(start_pos) + initial_positions.append(start_pos) permutations.append(perm) n = len(permutations) split_sizes = [perm.num_rows for perm in permutations] initial_offset = self._resume_offset local_consumed = [0] * n + # Permutation position each split has consumed through (absolute, + # i.e. counted from the start of the unskipped split). Runs ahead of + # initial + local_consumed when rows are skipped. + pos_consumed = list(initial_positions) batch_size = self._read_batch_size max_prefetch = self._prefetch_batches @@ -302,12 +375,14 @@ class StreamingDataset(IterableDataset): self._transform if self._transform is not None else Transforms.arrow2python ) - # Per-split pipeline state. + # Per-split pipeline state. Batches are paired with the absolute + # permutation position of their first row so that skipped rows can be + # accounted for in pos_consumed. fetch_head = [0] * n - io_pending = [deque() for _ in range(n)] # Future[RecordBatch] - raw_batches = [deque() for _ in range(n)] # RecordBatch — fetched, awaiting tx - tx_pending = [deque() for _ in range(n)] # Future[list[Any]] - cooked = [deque() for _ in range(n)] # rows ready to yield + io_pending = [deque() for _ in range(n)] # (abs_start, Future[RecordBatch]) + raw_batches = [deque() for _ in range(n)] # (abs_start, RecordBatch) + tx_pending = [deque() for _ in range(n)] # Future[list[(abs_pos, row)]] + cooked = [deque() for _ in range(n)] # (abs_pos, row) ready to yield # Limit simultaneous transforms to transform_workers across all splits. tx_semaphore = threading.Semaphore(transform_workers) @@ -330,7 +405,8 @@ class StreamingDataset(IterableDataset): fetch_head[i] += fetch perm_i = permutations[i] indices = list(range(start, start + fetch)) - io_pending[i].append(io_pool.submit(_io_call, perm_i, indices)) + abs_start = initial_positions[i] + start + io_pending[i].append((abs_start, io_pool.submit(_io_call, perm_i, indices))) def _fill_io(i: int) -> None: while len(io_pending[i]) < max_prefetch and fetch_head[i] < split_sizes[i]: @@ -338,15 +414,72 @@ class StreamingDataset(IterableDataset): def _drain_io(i: int) -> None: """Move completed I/O futures into raw_batches non-blockingly.""" - while io_pending[i] and io_pending[i][0].done(): - raw_batches[i].append(io_pending[i].popleft().result()) + while io_pending[i] and io_pending[i][0][1].done(): + abs_start, fut = io_pending[i].popleft() + raw_batches[i].append((abs_start, fut.result())) # ── Stage 2 helpers ─────────────────────────────────────────────────── - def _tx_call_guarded(batch): + on_error = self._on_transform_error + + def _should_skip(exc: Exception) -> bool: + if on_error == "raise": + return False + if callable(on_error): + return bool(on_error(exc)) + return True # "skip" or "warn" + + def _check_row_count(rows: list, num_rows: int) -> None: + if len(rows) != num_rows: + raise ValueError( + f"transform returned {len(rows)} rows for a batch of " + f"{num_rows}; transforms must return exactly one output " + "row per input row. To drop bad rows, raise inside the " + "transform and pass on_transform_error='skip'." + ) + + def _transform_isolated(abs_start, batch, batch_exc): + """Re-run the transform on single-row slices, dropping failures.""" + out = [] + skipped = 0 + first_exc = None + for j in range(batch.num_rows): + try: + rows = list(final_transform(batch.slice(j, 1))) + except Exception as exc: + if not _should_skip(exc): + raise + skipped += 1 + if first_exc is None: + first_exc = exc + continue + _check_row_count(rows, 1) + out.append((abs_start + j, rows[0])) + self._rows_skipped += skipped + if skipped and on_error == "warn": + logger.warning( + "Skipped %d of %d rows whose transform failed (first error: %r)", + skipped, + batch.num_rows, + first_exc if first_exc is not None else batch_exc, + ) + return out + + def _transform_batch(abs_start, batch): + """Apply the transform, returning [(abs_pos, row), ...].""" + try: + rows = list(final_transform(batch)) + except Exception as exc: + if not _should_skip(exc): + raise + return _transform_isolated(abs_start, batch, exc) + _check_row_count(rows, batch.num_rows) + return [(abs_start + j, row) for j, row in enumerate(rows)] + + def _tx_call_guarded(abs_start, batch): try: t0 = time.perf_counter() - result = final_transform(batch) + result = _transform_batch(abs_start, batch) self._transform_time += time.perf_counter() - t0 return result finally: @@ -355,8 +488,8 @@ class StreamingDataset(IterableDataset): def _try_submit_tx(i: int) -> None: """Submit transforms for raw_batches[i] up to available capacity.""" while raw_batches[i] and tx_semaphore.acquire(blocking=False): - batch = raw_batches[i].popleft() - tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch)) + abs_start, batch = raw_batches[i].popleft() + tx_pending[i].append(tx_pool.submit(_tx_call_guarded, abs_start, batch)) def _drain_tx(i: int) -> None: """Move completed transform futures into cooked non-blockingly.""" @@ -384,11 +517,14 @@ class StreamingDataset(IterableDataset): # Acquire a transform slot (may block briefly if all # transform_workers are busy with other splits). tx_semaphore.acquire() - batch = raw_batches[i].popleft() - tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch)) + abs_start, batch = raw_batches[i].popleft() + tx_pending[i].append( + tx_pool.submit(_tx_call_guarded, abs_start, batch) + ) elif io_pending[i]: # Block on the oldest in-flight I/O fetch. - raw_batches[i].append(io_pending[i].popleft().result()) + abs_start, fut = io_pending[i].popleft() + raw_batches[i].append((abs_start, fut.result())) _advance(i) else: break # split exhausted @@ -407,15 +543,28 @@ class StreamingDataset(IterableDataset): _fill_io(i) while True: - # Stop when any split is exhausted (all exhaust - # simultaneously: equal split sizes + round-robin). - if any(local_consumed[i] >= split_sizes[i] for i in range(n)): + # A cycle only runs if every split can still produce a + # row. Without skips all splits exhaust simultaneously + # (equal split sizes + round-robin); when + # on_transform_error drops rows a split can run dry + # early, ending the epoch at the last complete cycle. + # This check only sees splits owned by this rank/worker + # (my_splits) — there is no cross-rank coordination, so + # a different rank with fewer skipped rows keeps going; + # see the on_transform_error docstring. + exhausted = False + for i in range(n): + _ensure_cooked(i) + if not cooked[i]: + exhausted = True + break + if exhausted: break for i in range(n): - _ensure_cooked(i) - row = cooked[i].popleft() + pos, row = cooked[i].popleft() local_consumed[i] += 1 + pos_consumed[i] = pos + 1 _advance(i) # After the last split in each cycle: update the @@ -424,21 +573,39 @@ class StreamingDataset(IterableDataset): # even when __iter__ runs in a worker process. if i == n - 1: self._resume_offset = initial_offset + local_consumed[i] + for j, split_idx in enumerate(my_splits): + self._resume_positions[split_idx] = pos_consumed[j] ws = self._worker_stats ws[0] = sum( split_sizes[j] - fetch_head[j] for j in range(n) ) ws[1] = sum( - batch.num_rows for q in raw_batches for batch in q + batch.num_rows + for q in raw_batches + for _, batch in q ) ws[2] = sum(len(q) for q in cooked) ws[3] = sum(local_consumed) ws[4] = self._bytes_loaded ws[5] = int(self._fetch_time * 1_000_000) ws[6] = int(self._transform_time * 1_000_000) + ws[7] = self._rows_skipped yield row finally: + # Final stats flush: the per-cycle write above never runs + # when iteration ends mid-cycle (e.g. a split whose rows + # were all skipped before completing a single cycle), so + # counters like rows_skipped would otherwise be stale. + ws = self._worker_stats + ws[0] = sum(split_sizes[j] - fetch_head[j] for j in range(n)) + ws[1] = 0 # queue-depth properties document 0 when idle + ws[2] = 0 + ws[3] = sum(local_consumed) + ws[4] = self._bytes_loaded + ws[5] = int(self._fetch_time * 1_000_000) + ws[6] = int(self._transform_time * 1_000_000) + ws[7] = self._rows_skipped self._raw_batches_ref = None self._cooked_ref = None self._fetch_head_ref = None @@ -492,7 +659,7 @@ class StreamingDataset(IterableDataset): batches. Returns 0 when not iterating. """ if self._raw_batches_ref is not None: - return sum(batch.num_rows for q in self._raw_batches_ref for batch in q) + return sum(batch.num_rows for q in self._raw_batches_ref for _, batch in q) return int(self._worker_stats[1]) @property @@ -522,6 +689,19 @@ class StreamingDataset(IterableDataset): ) return int(self._worker_stats[0]) + @property + def rows_skipped(self) -> int: + """Number of rows dropped because their transform raised an exception. + + Only ever non-zero when ``on_transform_error`` is set to ``"skip"``, + ``"warn"``, or a callable that returned ``True``. Accumulates across + multiple iterations of the same dataset instance and is never reset + automatically. + """ + if self._raw_batches_ref is not None: + return self._rows_skipped + return int(self._worker_stats[7]) + @property def consumed_rows(self) -> int: """Number of rows already yielded to the caller across all splits. @@ -587,12 +767,27 @@ class StreamingDataset(IterableDataset): every split has been consumed the same number of times (by the round-robin design), so the per-split count is a single uniform value that is identical across all ranks and DataLoader workers. + + ``positions_consumed_per_split`` records how far into each split's + permutation iteration has advanced. It only differs from + ``samples_consumed_per_split`` when ``on_transform_error`` skipped + rows, in which case entries are exact for the splits this instance + iterated and a lower bound (the sample count) for splits owned by + other ranks or workers. Combine the state dicts from all ranks with + [merge_state_dicts][lancedb.streaming.StreamingDataset.merge_state_dicts] + to recover the exact value for every split before resuming on a + different topology. """ + positions = [ + self._resume_positions.get(split, self._resume_offset) + for split in range(self._num_splits) + ] return { "shuffle_seed": self._shuffle_seed, "num_splits": self._num_splits, "epoch": self._epoch, "samples_consumed_per_split": [self._resume_offset] * self._num_splits, + "positions_consumed_per_split": positions, } def load_state_dict(self, state: dict) -> None: @@ -618,3 +813,96 @@ class StreamingDataset(IterableDataset): self._resume_offset = consumed[0] if consumed else 0 else: self._resume_offset = int(consumed) + # Older checkpoints predate positions_consumed_per_split; without + # skipped rows positions equal sample counts, so falling back to + # _resume_offset (the .get default in __iter__) is exact. + positions = state.get("positions_consumed_per_split") + if positions is None: + self._resume_positions = {} + else: + self._resume_positions = { + split: int(pos) for split, pos in enumerate(positions) + } + + @staticmethod + def merge_state_dicts(states: list[dict]) -> dict: + """Merge state dicts saved by different ranks into one exact state. + + Only needed when ``on_transform_error`` skips rows in multi-rank + training: each rank then knows the exact permutation position only for + its own splits, and records a lower bound for the rest. Because + exactly one rank owns each split, the elementwise maximum across all + ranks' ``positions_consumed_per_split`` recovers the exact position of + every split. Without skipped rows every rank's state is already + identical and merging is a no-op. + + Raises ``ValueError`` if the states are empty or were not produced by + the same run (mismatched seed, split count, epoch, or sample counts). + + The merge is always all-to-all and topology-agnostic: collect the + ``state_dict()`` from every rank of the *previous* run into one list, + merge that whole list, and hand the identical merged result to every + rank of the *next* run — regardless of whether the rank count grew, + shrank, or stayed the same. There is no pairwise or subset merging + step, because each split's exact position is only known to whichever + rank owned that split, and the elementwise maximum needs every rank's + contribution to be correct. + + For example, checkpointing 8 ranks and resuming on 4 (the same + pattern applies when growing, e.g. 4 ranks resuming on 8):: + + states = [ds.state_dict() for ds in previous_run_datasets] # 8 + merged = StreamingDataset.merge_state_dicts(states) + for ds in resumed_datasets: # now only 4 ranks + ds.load_state_dict(merged) # same dict on every rank + + The rank count on either side never affects the merge itself, since + ``merge_state_dicts`` only cares about the list of states it is + given. Each split's position is recovered by elementwise maximum; + here rank 0 owned split 0 (and skipped two rows there) while rank 1 + owned split 1 (and skipped one row): + + >>> rank0 = { + ... "shuffle_seed": 0, "num_splits": 2, "epoch": 0, + ... "samples_consumed_per_split": [3, 3], + ... "positions_consumed_per_split": [5, 3], + ... } + >>> rank1 = { + ... "shuffle_seed": 0, "num_splits": 2, "epoch": 0, + ... "samples_consumed_per_split": [3, 3], + ... "positions_consumed_per_split": [3, 4], + ... } + >>> merged = StreamingDataset.merge_state_dicts([rank0, rank1]) + >>> merged["positions_consumed_per_split"] + [5, 4] + """ + if not states: + raise ValueError("merge_state_dicts requires at least one state dict") + first = states[0] + for state in states[1:]: + for key in ("shuffle_seed", "num_splits", "epoch"): + if state[key] != first[key]: + raise ValueError( + f"{key} mismatch across state dicts: " + f"{state[key]} != {first[key]}" + ) + if ( + state["samples_consumed_per_split"] + != first["samples_consumed_per_split"] + ): + raise ValueError( + "samples_consumed_per_split mismatch across state dicts; " + "state_dict() must be called at the same global step " + "boundary on every rank" + ) + merged = dict(first) + all_positions = [ + state.get( + "positions_consumed_per_split", state["samples_consumed_per_split"] + ) + for state in states + ] + merged["positions_consumed_per_split"] = [ + max(per_split) for per_split in zip(*all_positions) + ] + return merged diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 6819ea5f1..106cbd921 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -177,6 +177,7 @@ if TYPE_CHECKING: CompactionStats, Tag, AddColumnsResult, + RefreshColumnResult, AddResult, AlterColumnsResult, UpdateFieldMetadataResult, @@ -1985,7 +1986,14 @@ class Table(ABC): @abstractmethod def add_columns( - self, transforms: Dict[str, str] | pa.Field | List[pa.Field] | pa.Schema + self, + transforms: Dict[str, str] + | pa.Field + | List[pa.Field] + | pa.Schema + | None = None, + *, + computed: Dict[str, str] | None = None, ): """ Add new columns with defined values. @@ -1999,11 +2007,68 @@ class Table(ABC): Alternatively, a pyarrow Field or Schema can be provided to add new columns with the specified data types. The new columns will be initialized with null values. + computed: Dict[str, str], optional + A map of column name to a SQL expression defining the column. The + column's type and inputs are derived from the expression, so no + data type is supplied. + + Unlike ``transforms``, the expression is stored rather than + evaluated now: the column is committed with no values, and rows get + them from [`refresh_column`][lancedb.table.Table.refresh_column]. + Declaring one therefore costs the same on a large table as on an + empty one. + + A refresh does not revisit rows it has already filled, so mutating + an input leaves the value computed at fill time; recomputing means + dropping the column and declaring it again. While a declaration + reads a column, that column cannot be renamed, retyped or dropped. + + Local tables only; LanceDB Cloud and Enterprise raise + ``NotImplementedError``. Cannot be combined with ``transforms``. Returns ------- AddColumnsResult version: the new version number of the table after adding columns. + + Examples + -------- + >>> import lancedb + >>> db = lancedb.connect("./.lancedb") + >>> table = db.create_table("computed_demo", [{"x": 1}, {"x": 2}]) + >>> table.add_columns(computed={"doubled": "x * 2"}) + AddColumnsResult(version=2) + >>> table.refresh_column("doubled") + RefreshColumnResult(rows_filled=2, version=3) + >>> table.to_arrow().sort_by("x").to_pandas() + x doubled + 0 1 2 + 1 2 4 + """ + + @abstractmethod + def refresh_column(self, column: str) -> "RefreshColumnResult": + """ + Fill the rows of a computed column that hold no value yet. + + Declared with ``add_columns(computed=...)``, a column starts empty and + gets its values here. Rows appended since the last refresh are filled + by the next one; rows already filled are left as they are, so the call + is idempotent and does not observe a mutated input. + + Local tables only; LanceDB Cloud and Enterprise raise + ``NotImplementedError``. + + Parameters + ---------- + column: str + The name of the computed column to fill. + + Returns + ------- + RefreshColumnResult + rows_filled: the number of rows given a value. + version: the new version number of the table. """ @abstractmethod @@ -4014,9 +4079,21 @@ class LanceTable(Table): return LOOP.run(self._table.index_stats(index_name)) def add_columns( - self, transforms: Dict[str, str] | pa.field | List[pa.field] | pa.Schema + self, + transforms: Dict[str, str] + | pa.field + | List[pa.field] + | pa.Schema + | None = None, + *, + computed: Dict[str, str] | None = None, ) -> AddColumnsResult: - return LOOP.run(self._table.add_columns(transforms)) + return LOOP.run(self._table.add_columns(transforms, computed=computed)) + + def refresh_column(self, column: str) -> "RefreshColumnResult": + """Fill a computed column's unfilled rows. See + [`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column].""" + return LOOP.run(self._table.refresh_column(column)) def alter_columns( self, *alterations: Iterable[Dict[str, str]] @@ -4751,6 +4828,13 @@ class AsyncTable: via [`set_unenforced_primary_key`]; bucket sharding additionally requires it to be the single column being bucketed. + By default the MemWAL maintains every index on the table, resolved + here — a snapshot, so an index created afterwards needs the spec unset + and set again. This fails if one cannot be maintained; name the set + with ``with_maintained_indexes`` to install anyway. That pins an exact + set (a still-building index is rejected, not omitted); ``[]`` maintains + none. + Parameters ---------- spec : LsmWriteSpec @@ -4777,9 +4861,9 @@ class AsyncTable: Returns ``None`` when the MemWAL LSM write path is not enabled (no spec has been set, or it was removed with `unset_lsm_write_spec`). - The returned spec — including its ``maintained_indexes`` and - ``writer_config_defaults`` — mirrors what was passed to - `set_lsm_write_spec`. + The returned spec mirrors what was passed to `set_lsm_write_spec`, + except that ``maintained_indexes`` always reports the concrete list + resolved when the spec was set — ``None`` never round-trips. """ return await self._inner.get_lsm_write_spec() @@ -5924,7 +6008,14 @@ class AsyncTable: return await self._inner.update(updates_sql, where) async def add_columns( - self, transforms: dict[str, str] | pa.field | List[pa.field] | pa.Schema + self, + transforms: dict[str, str] + | pa.field + | List[pa.field] + | pa.Schema + | None = None, + *, + computed: dict[str, str] | None = None, ) -> AddColumnsResult: """ Add new columns with defined values. @@ -5937,6 +6028,21 @@ class AsyncTable: each row in the table, and can reference existing columns. Alternatively, you can pass a pyarrow field or schema to add new columns with NULLs. + computed: Dict[str, str], optional + A map of column name to a SQL expression defining the column. The + column's type and inputs are derived from the expression. + + Unlike ``transforms``, the expression is stored rather than + evaluated now: the column is committed with no values, and rows get + them from + [`refresh_column`][lancedb.table.AsyncTable.refresh_column]. + + A refresh does not revisit rows it has already filled, so mutating + an input leaves the value computed at fill time. While a + declaration reads a column, that column cannot be renamed, retyped + or dropped. + + Local tables only. Cannot be combined with ``transforms``. Returns ------- @@ -5950,11 +6056,43 @@ class AsyncTable: {isinstance(f, pa.Field) for f in transforms} ): transforms = pa.schema(transforms) + if computed: + if transforms: + raise ValueError( + "add_columns cannot take both transforms and computed columns" + ) + return await self._inner.add_computed_columns(list(computed.items())) + if transforms is None: + raise ValueError("add_columns requires transforms or computed columns") if isinstance(transforms, pa.Schema): return await self._inner.add_columns_with_schema(transforms) else: return await self._inner.add_columns(list(transforms.items())) + async def refresh_column(self, column: str) -> RefreshColumnResult: + """ + Fill the rows of a computed column that hold no value yet. + + Declared with ``add_columns(computed=...)``, a column starts empty and + gets its values here. Rows appended since the last refresh are filled + by the next one; rows already filled are left as they are, so the call + is idempotent and does not observe a mutated input. + + Local tables only; LanceDB Cloud and Enterprise raise + ``NotImplementedError``. + + Parameters + ---------- + column: str + The name of the computed column to fill. + + Returns + ------- + RefreshColumnResult + The number of rows filled and the new version of the table. + """ + return await self._inner.refresh_column(column) + async def alter_columns( self, *alterations: Iterable[dict[str, Any]] ) -> AlterColumnsResult: @@ -6415,7 +6553,9 @@ class TableStatistics: Attributes ---------- total_bytes: int - The total number of bytes in the table. + The total size, in bytes, of the table's data files, index files, and + overlay files. Read from the manifest, so this excludes deletion files + and manifests. num_rows: int The total number of rows in the table. num_indices: int diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index 84e78fd8f..38bbb53fb 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -755,8 +755,7 @@ def test_delete_table(tmp_db: lancedb.DBConnection): assert tmp_db.table_names() == [] -@pytest.mark.asyncio -async def test_delete_table_async(tmp_db: lancedb.DBConnection): +def test_drop_table_async(tmp_db: lancedb.DBConnection): data = pd.DataFrame( { "vector": [[3.1, 4.1], [5.9, 26.5]], @@ -772,7 +771,10 @@ async def test_delete_table_async(tmp_db: lancedb.DBConnection): assert tmp_db.table_names() == ["test"] - tmp_db.drop_table("test") + job = tmp_db.drop_table_async("test") + assert job.id is None + assert job.status() == "finished" + job.wait() assert tmp_db.table_names() == [] tmp_db.create_table("test", data=data) @@ -781,6 +783,17 @@ async def test_delete_table_async(tmp_db: lancedb.DBConnection): tmp_db.drop_table("does_not_exist", ignore_missing=True) +@pytest.mark.asyncio +async def test_drop_table_async_connection(tmp_db_async: lancedb.AsyncConnection): + await tmp_db_async.create_table("test", data=pa.table({"id": [1, 2]})) + + job = await tmp_db_async.drop_table_async("test") + assert job.id is None + assert await job.status() == "finished" + await job.wait() + assert await tmp_db_async.table_names() == [] + + def test_drop_database(tmp_db: lancedb.DBConnection): data = pd.DataFrame( { diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index 22918082b..734f835c6 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -1456,6 +1456,408 @@ def test_shuffle_clump_size_yields_all_rows(lance_table): ) +# --------------------------------------------------------------------------- +# on_transform_error tests +# --------------------------------------------------------------------------- + + +class BadRowError(ValueError): + """Raised by the failing transforms below when a batch contains a bad id.""" + + +def _failing_transform(bad_ids: set): + """A transform that raises BadRowError whenever the batch has a bad id. + + Raises on the full batch and on any single-row slice containing a bad id, + so per-row isolation drops exactly the bad rows. + """ + + def transform(batch: pa.RecordBatch) -> list: + ids = batch.column("id").to_pylist() + bad = sorted(set(ids) & bad_ids) + if bad: + raise BadRowError(f"bad ids in batch: {bad}") + return [{"id": i} for i in ids] + + return transform + + +def _sequential_split_members(table) -> list[list[int]]: + """Return each split's ids in yield order for shuffle=False. + + With a single rank and no workers the round-robin yields one row per split + per cycle, so item k of a clean run belongs to split k % NUM_SPLITS. + """ + ds = StreamingDataset(table, num_splits=NUM_SPLITS, shuffle=False) + members: list[list[int]] = [[] for _ in range(NUM_SPLITS)] + for k, row in enumerate(ds): + members[k % NUM_SPLITS].append(row["id"]) + return members + + +def test_on_transform_error_default_raises(lance_table): + """By default a transform exception propagates and aborts iteration.""" + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + transform=_failing_transform({7}), + ) + with pytest.raises(BadRowError): + list(ds) + + +def test_on_transform_error_invalid_value(lance_table): + with pytest.raises(ValueError, match="on_transform_error"): + StreamingDataset(lance_table, num_splits=NUM_SPLITS, on_transform_error="bogus") + + +def test_on_transform_error_skip_drops_bad_rows(lance_table): + """With one bad row per split, 'skip' yields every good row exactly once + and counts the dropped rows in rows_skipped.""" + members = _sequential_split_members(lance_table) + bad_ids = {members[i][4] for i in range(NUM_SPLITS)} + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="skip", + ) + assert ds.rows_skipped == 0 + + ids = [row["id"] for row in ds] + + assert sorted(ids) == sorted(set(range(NUM_ROWS)) - bad_ids) + assert ds.rows_skipped == NUM_SPLITS + + +def test_on_transform_error_skip_uneven_ends_at_last_complete_cycle(lance_table): + """When one split loses more rows than the others, the epoch ends at the + last cycle where every split still has a row — no crash, no bad rows, and + every step remains one sample per split.""" + members = _sequential_split_members(lance_table) + bad_ids = set(members[0][:3]) # all 3 bad rows in split 0 + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="skip", + ) + items = [row["id"] for row in ds] + + rows_per_split = NUM_ROWS // NUM_SPLITS + expected_cycles = rows_per_split - len(bad_ids) + assert len(items) == expected_cycles * NUM_SPLITS + assert len(set(items)) == len(items), "duplicate samples yielded" + assert not set(items) & bad_ids, "a bad row was yielded" + # Split 0 contributed exactly its surviving rows, in order, one per cycle. + survivors = [i for i in members[0] if i not in bad_ids] + assert items[0::NUM_SPLITS] == survivors[:expected_cycles] + + +def test_on_transform_error_warn_logs(lance_table, caplog): + """'warn' skips like 'skip' but logs a warning for the failing batch.""" + members = _sequential_split_members(lance_table) + bad_ids = {members[i][3] for i in range(NUM_SPLITS)} + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="warn", + ) + with caplog.at_level(logging.WARNING, logger="lancedb.streaming"): + items = list(ds) + + assert len(items) == NUM_ROWS - NUM_SPLITS + assert ds.rows_skipped == NUM_SPLITS + assert "Skipped" in caplog.text + assert "BadRowError" in caplog.text + + +def test_on_transform_error_callable_selective(lance_table): + """A callable handler can skip expected errors and re-raise the rest.""" + members = _sequential_split_members(lance_table) + bad_ids = {members[i][0] for i in range(NUM_SPLITS)} + + handled: list[Exception] = [] + + def handler(exc: Exception) -> bool: + handled.append(exc) + return isinstance(exc, BadRowError) + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error=handler, + ) + items = list(ds) + assert len(items) == NUM_ROWS - NUM_SPLITS + assert handled and all(isinstance(exc, BadRowError) for exc in handled) + + def broken_transform(batch: pa.RecordBatch) -> list: + raise TypeError("boom") + + ds2 = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=broken_transform, + on_transform_error=handler, + ) + with pytest.raises(TypeError, match="boom"): + list(ds2) + + +def test_transform_wrong_row_count_raises(lance_table): + """A transform that returns the wrong number of rows is an error even with + on_transform_error='skip' — silent shrinkage would corrupt accounting.""" + + def drops_rows(batch: pa.RecordBatch) -> list: + return batch.column("id").to_pylist()[:-1] + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + transform=drops_rows, + on_transform_error="skip", + ) + with pytest.raises(ValueError, match="one output row per input row"): + list(ds) + + +def test_skip_deterministic_across_runs(lance_table): + """With a fixed seed, skipping produces the identical sample sequence on + every run — skips are data-dependent, not run-dependent.""" + bad_ids = {5, 17, 46} + + def run() -> tuple[list[int], int]: + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle_seed=SHUFFLE_SEED, + transform=_failing_transform(bad_ids), + on_transform_error="skip", + ) + return [row["id"] for row in ds], ds.rows_skipped + + ids_a, skipped_a = run() + ids_b, skipped_b = run() + assert ids_a == ids_b + assert skipped_a == skipped_b + assert not set(ids_a) & bad_ids + + +def test_skip_elastic_det_across_world_sizes(lance_table): + """With equal bad-row counts per split, skipping preserves the full + elastic-determinism guarantee: identical global batches at every step for + every compatible world_size.""" + members = _sequential_split_members(lance_table) + bad_ids = {members[i][6] for i in range(NUM_SPLITS)} + + def collect(world_size: int) -> list[frozenset[int]]: + micro = GLOBAL_BATCH_SIZE // world_size + iters = [ + iter( + StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + rank=rank, + world_size=world_size, + transform=_failing_transform(bad_ids), + on_transform_error="skip", + ) + ) + for rank in range(world_size) + ] + _STOP = object() + batches: list[frozenset[int]] = [] + while True: + step_samples: set[int] = set() + exhausted = 0 + for it in iters: + for _ in range(micro): + val = next(it, _STOP) + if val is _STOP: + exhausted += 1 + break + step_samples.add(val["id"]) + if exhausted == len(iters): + break + assert exhausted == 0, ( + "Rank iterators exhausted at different steps despite equal " + "bad-row counts per split" + ) + batches.append(frozenset(step_samples)) + return batches + + reference = collect(1) + assert len(reference) == NUM_ROWS // NUM_SPLITS - 1 + for ws in (2, 3, 4): + assert collect(ws) == reference, f"world_size={ws} diverged" + + +def test_resumability_with_skips_same_topology(lance_table): + """Checkpointing mid-epoch with skipped rows resumes exactly: no sample + repeated, no sample lost, skipped rows stay skipped.""" + members = _sequential_split_members(lance_table) + # Uneven skips: positions diverge across splits (2 bad in split 0, 1 in + # split 5), which only a position-based checkpoint can resume exactly. + bad_ids = {members[0][2], members[0][3], members[5][7]} + kwargs = dict( + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="skip", + ) + + reference = [row["id"] for row in StreamingDataset(lance_table, **kwargs)] + rows_per_split = NUM_ROWS // NUM_SPLITS + assert len(reference) == (rows_per_split - 2) * NUM_SPLITS + + steps = 3 + ds = StreamingDataset(lance_table, **kwargs) + it = iter(ds) + consumed = [next(it)["id"] for _ in range(steps * NUM_SPLITS)] + checkpoint = ds.state_dict() + it.close() + + # Split 0 skipped positions 2 and 3 within its first 3 yields; split 5's + # bad row is beyond the checkpoint. Everything else is at 3 = the sample + # count. + positions = checkpoint["positions_consumed_per_split"] + assert positions[0] == 5 + assert positions[1:] == [3] * (NUM_SPLITS - 1) + assert checkpoint["samples_consumed_per_split"] == [3] * NUM_SPLITS + + ds2 = StreamingDataset(lance_table, **kwargs) + ds2.load_state_dict(checkpoint) + resumed = [row["id"] for row in ds2] + + assert consumed == reference[: steps * NUM_SPLITS] + assert resumed == reference[steps * NUM_SPLITS :] + + +def test_resumability_with_skips_elastic_merge(lance_table): + """Elastic resume with skips: each rank's checkpoint knows exact positions + only for its own splits; merge_state_dicts recovers the global state, and + a run on a different world_size continues exactly.""" + members = _sequential_split_members(lance_table) + # Bad rows early in split 0 (rank 0) and split 6 (rank 1 of a ws=2 run) so + # both ranks' position vectors diverge before the checkpoint. + bad_ids = {members[0][0], members[0][2], members[6][1]} + kwargs = dict( + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="skip", + ) + + reference = [row["id"] for row in StreamingDataset(lance_table, **kwargs)] + + steps = 3 + world_size = 2 + micro = GLOBAL_BATCH_SIZE // world_size + datasets = [ + StreamingDataset(lance_table, rank=rank, world_size=world_size, **kwargs) + for rank in range(world_size) + ] + iters = [iter(ds) for ds in datasets] + seen: list[frozenset[int]] = [] + for _ in range(steps): + step_samples = set() + for it in iters: + for _ in range(micro): + step_samples.add(next(it)["id"]) + seen.append(frozenset(step_samples)) + states = [ds.state_dict() for ds in datasets] + for it in iters: + it.close() + + merged = StreamingDataset.merge_state_dicts(states) + expected_positions = [3] * NUM_SPLITS + expected_positions[0] = 5 # skipped positions 0 and 2 + expected_positions[6] = 4 # skipped position 1 + assert merged["positions_consumed_per_split"] == expected_positions + + # The first 3 global batches match the world_size=1 reference. + ref_batches = [ + frozenset(reference[s * NUM_SPLITS : (s + 1) * NUM_SPLITS]) + for s in range(len(reference) // NUM_SPLITS) + ] + assert seen == ref_batches[:steps] + + # Resume on world_size=1 from the merged state. + ds_resume = StreamingDataset(lance_table, **kwargs) + ds_resume.load_state_dict(merged) + resumed = [row["id"] for row in ds_resume] + assert resumed == reference[steps * NUM_SPLITS :] + + +def test_rows_skipped_flushed_when_split_entirely_bad(lance_table): + """A split whose rows all fail never completes a cycle, so the epoch ends + immediately — but rows_skipped must still report the drops after the + iterator exits (the shared-memory counter is flushed on exhaustion).""" + members = _sequential_split_members(lance_table) + bad_ids = set(members[0]) # every row of split 0 is bad + + ds = StreamingDataset( + lance_table, + num_splits=NUM_SPLITS, + shuffle=False, + transform=_failing_transform(bad_ids), + on_transform_error="skip", + ) + assert list(ds) == [] + assert ds.rows_skipped == len(bad_ids) + + +def test_merge_state_dicts_validates_consistency(lance_table): + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + state = ds.state_dict() + other = dict(state, shuffle_seed=SHUFFLE_SEED + 1) + with pytest.raises(ValueError, match="shuffle_seed mismatch"): + StreamingDataset.merge_state_dicts([state, other]) + with pytest.raises(ValueError, match="at least one"): + StreamingDataset.merge_state_dicts([]) + + +def test_load_state_dict_without_positions_key(lance_table): + """Checkpoints from before positions_consumed_per_split existed still + resume exactly (positions equal sample counts when nothing is skipped).""" + reference = [ + row["id"] + for row in StreamingDataset( + lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED + ) + ] + + steps = 4 + ds = StreamingDataset(lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED) + it = iter(ds) + for _ in range(steps * NUM_SPLITS): + next(it) + checkpoint = ds.state_dict() + it.close() + del checkpoint["positions_consumed_per_split"] + + ds2 = StreamingDataset( + lance_table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED + ) + ds2.load_state_dict(checkpoint) + resumed = [row["id"] for row in ds2] + assert resumed == reference[steps * NUM_SPLITS :] + + def test_num_splits_defaults_to_world_size(lance_table): """Omitting num_splits gives world_size splits (one per rank).""" ds = StreamingDataset( diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 678270f19..9850669eb 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -631,3 +631,23 @@ def test_url_retrieve_downloads_image(): image_bytes = url_retrieve(image_url) img = Image.open(io.BytesIO(image_bytes)) assert img.size[0] > 0 and img.size[1] > 0 + + +def test_jina_generate_image_input_dict_local_path(tmp_path): + """ + JinaEmbeddings._generate_image_input_dict must accept a local image path + (str or Path), not just bytes. Previously it crashed with + `AttributeError: 'function' object has no attribute 'urlparse'` on any + str/Path input because it called `urlparse.urlparse(image)` instead of + `urlparse(image)` (urlparse was imported as a function, not a module). + """ + Image = pytest.importorskip("PIL.Image") + from lancedb.embeddings.jinaai import JinaEmbeddings + + image_path = tmp_path / "test.png" + Image.new("RGB", (4, 4), color="red").save(image_path, format="PNG") + + for image in (str(image_path), image_path): + image_dict = JinaEmbeddings._generate_image_input_dict(image) + assert "image" in image_dict + assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0 diff --git a/python/python/tests/test_lsm_write_spec.py b/python/python/tests/test_lsm_write_spec.py index d38918f09..218793b89 100644 --- a/python/python/tests/test_lsm_write_spec.py +++ b/python/python/tests/test_lsm_write_spec.py @@ -83,7 +83,9 @@ def test_lsm_write_spec_repr(): assert s.spec_type == "bucket" assert s.column == "id" assert s.num_buckets == 4 - assert s.maintained_indexes == [] + # A fresh spec defers its maintained set to install time. + assert s.maintained_indexes is None + assert s.with_maintained_indexes([]).maintained_indexes == [] assert "bucket" in repr(s) assert "id" in repr(s) assert "4" in repr(s) @@ -169,18 +171,23 @@ 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). + # 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". 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] table.unset_lsm_write_spec() - # Unsharded round-trips (no routing column). - table.set_lsm_write_spec(LsmWriteSpec.unsharded()) + # Unsharded round-trips (no routing column). Opting out is distinct from + # the inferred default. + table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([])) spec = table.get_lsm_write_spec() assert spec.spec_type == "unsharded" assert spec.column is None + assert spec.maintained_indexes == [] @pytest.mark.asyncio diff --git a/python/python/tests/test_merge_insert_lsm.py b/python/python/tests/test_merge_insert_lsm.py index 5674a05ab..e74c21589 100644 --- a/python/python/tests/test_merge_insert_lsm.py +++ b/python/python/tests/test_merge_insert_lsm.py @@ -544,7 +544,7 @@ def test_lsm_read_fts_unmaintained_index_errors(tmp_path): table.create_index("text", config=FTS()) # No maintained indexes: the active memtable FTS arm cannot serve un-compacted # docs, so the search would silently omit them — reject instead. - table.set_lsm_write_spec(LsmWriteSpec.unsharded()) + table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([])) with pytest.raises(Exception, match="maintained"): table.search("fox", query_type="fts", fts_columns="text").to_arrow() @@ -631,7 +631,7 @@ def test_lsm_read_vector_unmaintained_index_errors(tmp_path): ) # Spec with NO maintained indexes: the base vector index's catch-up is untracked, # so the scanner rejects rather than risk dropping compacted-but-unindexed rows. - table.set_lsm_write_spec(LsmWriteSpec.unsharded()) + table.set_lsm_write_spec(LsmWriteSpec.unsharded().with_maintained_indexes([])) with pytest.raises(Exception, match="maintained"): table.search([1.0] * VECTOR_DIM).to_arrow() diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 06451dbc1..312b1ed49 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -3758,7 +3758,8 @@ def test_stats(mem_db: DBConnection): stats = table.stats() print(f"{stats=}") assert stats == { - "total_bytes": 60, + # Full on-disk size of the data file, footer and metadata included. + "total_bytes": 633, "num_rows": 2, "num_indices": 0, "fragment_stats": { @@ -3776,6 +3777,13 @@ def test_stats(mem_db: DBConnection): }, } + # Index files count toward total_bytes too (only deletion files and + # manifests are excluded). + table.create_index("id", config=BTree()) + stats_with_index = table.stats() + assert stats_with_index["num_indices"] == 1 + assert stats_with_index["total_bytes"] > stats["total_bytes"] + def test_create_table_empty_list_with_schema(mem_db: DBConnection): """Test creating table with empty list data and schema @@ -3891,3 +3899,37 @@ async def test_async_search_runs_embedding_on_dedicated_executor( assert all(name.startswith("lancedb-embedding") for name in captured_threads), ( f"embedding ran off the dedicated executor: {captured_threads}" ) + + +def test_computed_column_declare_and_refresh(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("computed", [{"x": 1}, {"x": 2}]) + + table.add_columns(computed={"doubled": "x * 2"}) + assert table.to_arrow()["doubled"].to_pylist() == [None, None] + + result = table.refresh_column("doubled") + assert result.rows_filled == 2 + assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4] + + table.add([{"x": 5}]) + assert table.refresh_column("doubled").rows_filled == 1 + assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4, 10] + + +def test_computed_column_rejects_transforms_and_computed_together(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("computed_mixed", [{"x": 1}]) + with pytest.raises(ValueError): + table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"}) + + +@pytest.mark.asyncio +async def test_computed_column_async(tmp_path): + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("computed_async", [{"x": 3}]) + + await table.add_columns(computed={"tripled": "x * 3"}) + await table.refresh_column("tripled") + + assert (await table.to_arrow())["tripled"].to_pylist() == [9] diff --git a/python/src/connection.rs b/python/src/connection.rs index b97d48ad8..dbda29ba6 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -346,6 +346,23 @@ impl Connection { }) } + #[pyo3(signature = (name, namespace_path=None))] + pub fn drop_table_async( + self_: PyRef<'_, Self>, + name: String, + namespace_path: Option>, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + let ns_path = namespace_path.unwrap_or_default(); + future_into_py(self_.py(), async move { + inner + .drop_table_async(name, &ns_path) + .await + .infer_error() + .map(crate::job::Job::new) + }) + } + #[pyo3(signature = (namespace_path=None,))] pub fn drop_all_tables( self_: PyRef<'_, Self>, diff --git a/python/src/index.rs b/python/src/index.rs index 8c81dcecf..dd362373e 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -289,7 +289,7 @@ struct IvfHnswFlatParams { target_partition_size: Option, } -#[pyclass(get_all)] +#[pyclass(module = "lancedb._lancedb", get_all)] /// A description of an index currently configured on a column pub struct IndexConfig { /// The type of the index diff --git a/python/src/lib.rs b/python/src/lib.rs index 6b0c0cf97..a19bf172d 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -16,7 +16,8 @@ use query::{FTSQuery, HybridQuery, Query, VectorQuery}; use session::Session; use table::{ AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken, - LsmWriteSpec, MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult, + LsmWriteSpec, MergeResult, PyBlobFile, RefreshColumnResult, Table, UpdateFieldMetadataResult, + UpdateResult, }; pub mod arrow; @@ -57,6 +58,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/session.rs b/python/src/session.rs index 891e61e44..4d58dd269 100644 --- a/python/src/session.rs +++ b/python/src/session.rs @@ -11,7 +11,7 @@ use pyo3::{PyResult, pyclass, pymethods}; /// Sessions allow you to configure cache sizes for index and metadata caches, /// which can significantly impact memory use and performance. They can /// also be re-used across multiple connections to share the same cache state. -#[pyclass(from_py_object)] +#[pyclass(module = "lancedb._lancedb", from_py_object)] #[derive(Clone)] pub struct Session { pub(crate) inner: Arc, diff --git a/python/src/table.rs b/python/src/table.rs index 461ac5822..221161b23 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -341,12 +341,22 @@ impl From for MergeResult { } } +/// Render for `__repr__`, so the default reads as Python's `None` rather than +/// Rust's `Some([..])`. +fn fmt_maintained(maintained: &Option>) -> String { + match maintained { + Some(names) => format!("{:?}", names), + None => "None".to_string(), + } +} + /// Specification selecting Lance's MemWAL LSM-style write path for /// `merge_insert`. /// /// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()` /// classmethods, then optionally chain `with_maintained_indexes(...)` and -/// `with_writer_config_defaults(...)`. +/// `with_writer_config_defaults(...)`. A fresh spec maintains every index the +/// MemWAL supports, resolved on install. #[pyclass(from_py_object)] #[derive(Clone, Debug)] pub struct LsmWriteSpec { @@ -386,11 +396,11 @@ impl LsmWriteSpec { } } - /// Replace the list of indexes the MemWAL should keep up to date as - /// rows are appended. Each name must reference an index that - /// already exists on the table at the time `set_lsm_write_spec` - /// is called. - pub fn with_maintained_indexes(&self, indexes: Vec) -> Self { + /// Set which indexes the MemWAL maintains. `None` (the default) + /// resolves every supported index on install; a list is verbatim, + /// and an empty list maintains nothing. + #[pyo3(signature = (indexes))] + pub fn with_maintained_indexes(&self, indexes: Option>) -> Self { Self { inner: self.inner.clone().with_maintained_indexes(indexes), } @@ -412,23 +422,29 @@ impl LsmWriteSpec { maintained_indexes, writer_config_defaults, } => format!( - "LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={:?}, writer_config_defaults={:?})", - column, num_buckets, maintained_indexes, writer_config_defaults, + "LsmWriteSpec.bucket(column={:?}, num_buckets={}, maintained_indexes={}, writer_config_defaults={:?})", + column, + num_buckets, + fmt_maintained(maintained_indexes), + writer_config_defaults, ), lancedb::table::LsmWriteSpec::Identity { column, maintained_indexes, writer_config_defaults, } => format!( - "LsmWriteSpec.identity(column={:?}, maintained_indexes={:?}, writer_config_defaults={:?})", - column, maintained_indexes, writer_config_defaults, + "LsmWriteSpec.identity(column={:?}, maintained_indexes={}, writer_config_defaults={:?})", + column, + fmt_maintained(maintained_indexes), + writer_config_defaults, ), lancedb::table::LsmWriteSpec::Unsharded { maintained_indexes, writer_config_defaults, } => format!( - "LsmWriteSpec.unsharded(maintained_indexes={:?}, writer_config_defaults={:?})", - maintained_indexes, writer_config_defaults, + "LsmWriteSpec.unsharded(maintained_indexes={}, writer_config_defaults={:?})", + fmt_maintained(maintained_indexes), + writer_config_defaults, ), } } @@ -463,10 +479,10 @@ impl LsmWriteSpec { } } - /// Names of indexes the MemWAL should keep up to date during writes. + /// Indexes the MemWAL keeps up to date, or `None` for every supported one. #[getter] - pub fn maintained_indexes(&self) -> Vec { - self.inner.maintained_indexes().to_vec() + pub fn maintained_indexes(&self) -> Option> { + self.inner.maintained_indexes().map(<[String]>::to_vec) } /// Default `ShardWriter` configuration recorded by this spec. @@ -494,6 +510,32 @@ pub struct AddColumnsResult { pub version: u64, } +#[pyclass(get_all, from_py_object)] +#[derive(Clone, Debug)] +pub struct RefreshColumnResult { + pub rows_filled: u64, + pub version: u64, +} + +#[pymethods] +impl RefreshColumnResult { + pub fn __repr__(&self) -> String { + format!( + "RefreshColumnResult(rows_filled={}, version={})", + self.rows_filled, self.version + ) + } +} + +impl From for RefreshColumnResult { + fn from(result: lancedb::table::RefreshColumnResult) -> Self { + Self { + rows_filled: result.rows_filled, + version: result.version, + } + } +} + #[pymethods] impl AddColumnsResult { pub fn __repr__(&self) -> String { @@ -658,7 +700,7 @@ impl PyBlobFile { } } -#[pyclass(get_all, from_py_object)] +#[pyclass(module = "lancedb._lancedb", get_all, from_py_object)] #[derive(Clone, Debug)] pub struct FtsToken { pub text: String, @@ -1591,6 +1633,29 @@ impl Table { }) } + pub fn add_computed_columns( + self_: PyRef<'_, Self>, + columns: Vec<(String, String)>, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let mut builder = inner.add_columns(); + for (name, expression) in columns { + builder = builder.computed(name, expression); + } + let result = builder.execute().await.infer_error()?; + Ok(AddColumnsResult::from(result)) + }) + } + + pub fn refresh_column(self_: PyRef<'_, Self>, column: String) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let result = inner.refresh_column(column).await.infer_error()?; + Ok(RefreshColumnResult::from(result)) + }) + } + pub fn add_columns_with_schema( self_: PyRef<'_, Self>, schema: PyArrowType, diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index c2137d0b5..23c0dcfd0 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.37.1-beta.0" +version = "0.38.0-beta.0" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true @@ -49,8 +49,6 @@ lance-namespace = { workspace = true } lance-namespace-impls = { workspace = true } metrics = { workspace = true, optional = true } metrics-util = { workspace = true, optional = true } -# Pin the GooseFS SDK to the version required by Lance's OpenDAL dependency. -goosefs-sdk = { version = "=0.1.9", optional = true } moka = { workspace = true } pin-project = { workspace = true } tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] } @@ -100,6 +98,7 @@ anyhow = "1" lance-testing = { workspace = true } tempfile = "3.5.0" random_word = { version = "0.4.3", features = ["en"] } +roaring = "0.11.4" tokio = { version = "1.23", features = ["io-util", "macros", "net", "rt-multi-thread", "sync", "test-util"] } uuid = { version = "1.7.0", features = ["v4"] } walkdir = "2" @@ -135,7 +134,6 @@ azure = [ ] cos = ["lance/tencent", "lance-io/tencent"] goosefs = [ - "dep:goosefs-sdk", "lance/goosefs", "lance-io/goosefs", "lance-namespace-impls/dir-goosefs", @@ -190,6 +188,9 @@ required-features = ["bedrock"] [[example]] name = "bench_streaming_dataloader" +[[example]] +name = "bench_open_missing_table" + [[example]] name = "simple" diff --git a/rust/lancedb/examples/bench_open_missing_table.rs b/rust/lancedb/examples/bench_open_missing_table.rs new file mode 100644 index 000000000..8e6b16e11 --- /dev/null +++ b/rust/lancedb/examples/bench_open_missing_table.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +// Release benchmark for opening a missing table as sibling-table cardinality grows. +// +// The fixture uses real `.lance` directories and marker files. Fixture creation is +// outside the timed section. Defaults intentionally cover 1k, 10k, and 100k siblings +// with 10 warmups and 100 distinct missing-table opens per scale: +// +// ```text +// cargo run --release -p lancedb --example bench_open_missing_table +// ``` +// +// `BENCH_SIBLINGS`, `BENCH_WARMUPS`, and `BENCH_TRIALS` override those defaults. +// Reduced settings are useful only as a smoke test. Performance comparisons require +// the same machine, filesystem, fixture sizes, settings, lockfile, and alternating +// baseline/candidate execution order. + +use std::time::{Duration, Instant}; + +use anyhow::{Context, Result, bail}; +use lancedb::connection::Connection; +use lancedb::{Error, connect}; +use object_store::ObjectStoreExt as _; +use object_store::path::Path; + +const MAX_SIBLINGS: usize = 1_000_000; +const MAX_WARMUPS: usize = 10_000; +const MAX_TRIALS: usize = 100_000; + +fn env_usize(key: &str, default: usize, max: usize) -> Result { + let value = match std::env::var(key) { + Ok(value) => value + .parse() + .with_context(|| format!("invalid {key} value: {value}"))?, + Err(std::env::VarError::NotPresent) => default, + Err(error) => return Err(error).with_context(|| format!("reading {key}")), + }; + if value == 0 || value > max { + bail!("{key} must be between 1 and {max}"); + } + Ok(value) +} + +fn sibling_counts() -> Result> { + let raw = std::env::var("BENCH_SIBLINGS").unwrap_or_else(|_| "1000,10000,100000".into()); + let mut counts = raw + .split(',') + .map(|value| { + value + .trim() + .parse::() + .with_context(|| format!("invalid BENCH_SIBLINGS value: {value}")) + }) + .collect::>>()?; + counts.sort_unstable(); + counts.dedup(); + if counts.is_empty() || counts[0] == 0 || counts[counts.len() - 1] > MAX_SIBLINGS { + bail!("BENCH_SIBLINGS values must be between 1 and {MAX_SIBLINGS}"); + } + Ok(counts) +} + +async fn add_siblings( + store: &object_store::local::LocalFileSystem, + start: usize, + end: usize, +) -> Result<()> { + for index in start..end { + let marker = Path::from(format!("sibling_{index:06}.lance/_marker")); + store + .put(&marker, bytes::Bytes::new().into()) + .await + .with_context(|| format!("creating benchmark marker {marker}"))?; + } + Ok(()) +} + +async fn time_missing_open(db: &Connection, name: &str) -> Result { + let started = Instant::now(); + let result = db.open_table(name).execute().await; + let elapsed = started.elapsed(); + match result { + Err(Error::TableNotFound { .. }) => Ok(elapsed), + Err(error) => bail!("expected TableNotFound for {name}, got {error:?}"), + Ok(_) => bail!("benchmark missing-table name unexpectedly exists: {name}"), + } +} + +fn percentile(sorted: &[Duration], percentile: usize) -> Duration { + let rank = (sorted.len() * percentile).div_ceil(100).saturating_sub(1); + sorted[rank] +} + +#[tokio::main] +async fn main() -> Result<()> { + let counts = sibling_counts()?; + let warmups = env_usize("BENCH_WARMUPS", 10, MAX_WARMUPS)?; + let trials = env_usize("BENCH_TRIALS", 100, MAX_TRIALS)?; + + let fixture = tempfile::tempdir().context("creating benchmark fixture")?; + let database_path = fixture.path(); + let fixture_store = object_store::local::LocalFileSystem::new_with_prefix(database_path) + .context("creating benchmark object store")?; + let db = connect(database_path.to_str().context("non-UTF-8 fixture path")?) + .execute() + .await?; + + println!( + "config: siblings={counts:?} warmups={warmups} trials={trials} profile={} os={} arch={}", + if cfg!(debug_assertions) { + "debug" + } else { + "release" + }, + std::env::consts::OS, + std::env::consts::ARCH, + ); + println!("lower is better; fixture setup and teardown are excluded"); + println!("| siblings | samples | p50 | p95 | max |"); + println!("| ---: | ---: | ---: | ---: | ---: |"); + + let mut created = 0; + for sibling_count in counts { + add_siblings(&fixture_store, created, sibling_count).await?; + created = sibling_count; + + for index in 0..warmups { + let name = format!("__missing_warmup_{sibling_count}_{index}"); + let _ = time_missing_open(&db, &name).await?; + } + + let mut samples = Vec::with_capacity(trials); + for index in 0..trials { + let name = format!("__missing_trial_{sibling_count}_{index}"); + samples.push(time_missing_open(&db, &name).await?); + } + samples.sort_unstable(); + + println!( + "| {sibling_count} | {} | {:?} | {:?} | {:?} |", + samples.len(), + percentile(&samples, 50), + percentile(&samples, 95), + samples[samples.len() - 1], + ); + } + + Ok(()) +} diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index e1c18dd84..d59123ec3 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -17,7 +17,7 @@ use arrow_array::builder::LargeBinaryBuilder; use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams}; use lance_arrow::FieldExt; -use lance_file::version::LanceFileVersion; +use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; use lance_io::object_store::ObjectStore; use object_store::path::Path; @@ -333,7 +333,10 @@ pub(crate) fn ensure_blob_storage_version(schema: &Schema, params: &mut WritePar .data_storage_version .unwrap_or(LanceFileVersion::Stable) .resolve(); - if resolved < LanceFileVersion::V2_2 { + if matches!( + resolved, + ConcreteFileVersion::V1 | ConcreteFileVersion::V2_0 | ConcreteFileVersion::V2_1 + ) { params.data_storage_version = Some(LanceFileVersion::V2_2); } } @@ -499,7 +502,7 @@ mod tests { ensure_blob_storage_version(&blob_schema(), &mut params); assert_eq!( params.data_storage_version.unwrap().resolve(), - LanceFileVersion::V2_2 + ConcreteFileVersion::V2_2 ); } @@ -512,7 +515,7 @@ mod tests { ensure_blob_storage_version(&blob_schema(), &mut params); assert_eq!( params.data_storage_version.unwrap().resolve(), - LanceFileVersion::V2_2 + ConcreteFileVersion::V2_2 ); } diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index dd53a2d2e..12ca306b8 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -409,6 +409,11 @@ impl Connection { /// /// The names will be returned in lexicographical order (ascending) /// + /// Listing databases discover physical `*.lance` entries without opening every + /// dataset. The result is a point-in-time discovery snapshot: an entry may still be + /// under creation, may contain only uncommitted storage, or may be concurrently + /// dropped before it is opened. + /// /// The parameters `page_token` and `limit` can be used to paginate the results pub fn table_names(&self) -> TableNamesBuilder { TableNamesBuilder::new(self.internal.clone()) @@ -456,10 +461,9 @@ impl Connection { /// /// # Returns /// Created [`TableRef`], or [`Error::TableNotFound`] if the table does not exist. - /// If the table's storage is present but holds no readable dataset (for example a - /// `.lance` directory left behind by an interrupted drop and re-create, which - /// [`Self::table_names`] still lists) this returns [`Error::TableCorrupted`] - /// instead. + /// On listing databases, a committed Lance manifest is authoritative for table + /// existence. Uncommitted files or a physical `.lance` directory alone do not + /// make a table openable. pub fn open_table(&self, name: impl Into) -> OpenTableBuilder { OpenTableBuilder::new( self.internal.clone(), @@ -561,6 +565,21 @@ impl Connection { .await } + /// Start dropping a table and return a handle to the cleanup job. + /// + /// The table may become unavailable before its physical data is removed. + /// Call [`crate::job::Job::wait`] to wait for cleanup to finish. Local + /// backends may complete the drop before returning the handle. + pub async fn drop_table_async( + &self, + name: impl AsRef, + namespace_path: &[String], + ) -> Result { + self.internal + .drop_table_async(name.as_ref(), namespace_path) + .await + } + /// Drop the database /// /// This is the same as dropping all of the tables diff --git a/rust/lancedb/src/connection/create_table.rs b/rust/lancedb/src/connection/create_table.rs index b10141beb..39cc82ec0 100644 --- a/rust/lancedb/src/connection/create_table.rs +++ b/rust/lancedb/src/connection/create_table.rs @@ -438,10 +438,9 @@ mod tests { .await .unwrap() .data_storage_format - .lance_file_version() - .unwrap(); + .lance_file_format(); // Compare resolved versions since Stable/Next are aliases that resolve at storage time - assert_eq!(storage_format.resolve(), data_storage_version.resolve()); + assert_eq!(storage_format, data_storage_version.resolve()); } #[tokio::test] diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index f99f6e12a..f52c02439 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -323,6 +323,18 @@ pub trait Database: ) -> Result<()>; /// Drop a table in the database async fn drop_table(&self, name: &str, namespace_path: &[String]) -> Result<()>; + /// Start dropping a table and return a handle to the cleanup job. + /// + /// Backends without asynchronous cleanup complete the drop before + /// returning an already-finished job. + async fn drop_table_async( + &self, + name: &str, + namespace_path: &[String], + ) -> Result { + self.drop_table(name, namespace_path).await?; + Ok(crate::job::Job::new_done()) + } /// Drop all tables in the database async fn drop_all_tables(&self, namespace_path: &[String]) -> Result<()>; fn as_any(&self) -> &dyn std::any::Any; diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 0ab3614e7..f284320c6 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -1291,16 +1291,21 @@ impl Database for ListingDatabase { mod tests { use super::*; use crate::Table; + use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream}; use crate::connection::ConnectRequest; use crate::data::scannable::Scannable; use crate::database::{CreateTableMode, CreateTableRequest}; use crate::query::QueryRequest; use crate::table::{AnyQuery, WriteOptions}; use arrow_array::{Int32Array, RecordBatch, StringArray}; - use arrow_schema::{DataType, Field, Schema}; - use futures::TryStreamExt; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use futures::{TryStreamExt, stream::once}; use std::path::PathBuf; + use std::sync::Arc; + use std::time::Duration; use tempfile::tempdir; + use tokio::sync::Barrier; + use tokio::time::timeout; async fn setup_database() -> (tempfile::TempDir, ListingDatabase) { let tempdir = tempdir().unwrap(); @@ -1324,6 +1329,114 @@ mod tests { (tempdir, db) } + struct BarrierScannable { + batch: RecordBatch, + barrier: Arc, + } + + impl Scannable for BarrierScannable { + fn schema(&self) -> SchemaRef { + self.batch.schema() + } + + fn scan_as_stream(&mut self) -> SendableRecordBatchStream { + let batch = self.batch.clone(); + let schema = batch.schema(); + let barrier = self.barrier.clone(); + Box::pin(SimpleRecordBatchStream { + schema, + stream: once(async move { + barrier.wait().await; + Ok(batch) + }), + }) + } + } + + fn create_request(name: &str, data: Box) -> CreateTableRequest { + CreateTableRequest { + name: name.to_string(), + namespace_path: vec![], + data, + mode: CreateTableMode::Create, + write_options: Default::default(), + location: None, + namespace_client: None, + } + } + + #[tokio::test] + async fn test_create_ignores_uncommitted_storage_without_manifest() { + let (tmp_dir, db) = setup_database().await; + let data_dir = tmp_dir.path().join("test.lance/data"); + std::fs::create_dir_all(&data_dir).unwrap(); + std::fs::write(data_dir.join("orphan.lance"), b"uncommitted").unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); + + let table = db + .create_table(create_request("test", Box::new(batch))) + .await + .unwrap(); + assert_eq!(table.count_rows(None).await.unwrap(), 1); + } + + #[tokio::test] + async fn test_concurrent_create_is_arbitrated_by_manifest_commit() { + let uri = format!("memory:///concurrent-create-{}", uuid::Uuid::new_v4()); + let db = crate::connect(&uri).execute().await.unwrap(); + let store: Arc = + Arc::new(object_store::memory::InMemory::new()); + let table_url = url::Url::parse("memory:///database/test.lance").unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); + let barrier = Arc::new(Barrier::new(2)); + + #[allow(deprecated)] + let request = |batch, barrier| { + let mut request = create_request("test", Box::new(BarrierScannable { batch, barrier })); + request.write_options = WriteOptions { + lance_write_params: Some(lance::dataset::WriteParams { + store_params: Some(ObjectStoreParams { + object_store: Some((store.clone(), table_url.clone())), + ..Default::default() + }), + commit_handler: Some(Arc::new( + lance_table::io::commit::ConditionalPutCommitHandler, + )), + ..Default::default() + }), + }; + request + }; + + let left = db + .database() + .create_table(request(batch.clone(), barrier.clone())); + let right = db.database().create_table(request(batch, barrier)); + let (left, right) = timeout(Duration::from_secs(30), async { tokio::join!(left, right) }) + .await + .expect("concurrent creates deadlocked"); + + let results = [left, right]; + assert_eq!( + results.iter().filter(|result| result.is_ok()).count(), + 1, + "expected one successful create, got {results:?}" + ); + assert_eq!( + results + .iter() + .filter(|result| matches!(result, Err(Error::TableAlreadyExists { .. }))) + .count(), + 1, + "expected one manifest conflict, got {results:?}" + ); + } + #[tokio::test] async fn test_listing_database_root_ops_do_not_create_manifest() { let tempdir = tempdir().unwrap(); diff --git a/rust/lancedb/src/error.rs b/rust/lancedb/src/error.rs index 4a6e6d8d9..6bd1ffa2b 100644 --- a/rust/lancedb/src/error.rs +++ b/rust/lancedb/src/error.rs @@ -71,6 +71,14 @@ pub enum Error { IndexNotFound { name: String }, #[snafu(display("Embedding function '{name}' was not found. : {reason}"))] EmbeddingFunctionNotFound { name: String, reason: String }, + #[snafu(display("Column '{name}' was not found"))] + ColumnNotFound { name: String }, + #[snafu(display("Column '{name}' already exists"))] + ColumnAlreadyExists { name: String }, + #[snafu(display("Column '{name}' is not a computed column"))] + NotAComputedColumn { name: String }, + #[snafu(display("Invalid expression for column '{column}': {message}"))] + InvalidExpression { column: String, message: String }, #[snafu(display("Table '{name}' already exists"))] TableAlreadyExists { name: String }, diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 4b5f8832f..be9d0eef6 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -19,6 +19,15 @@ const ARROW_FILE_CONTENT_TYPE: &str = "application/vnd.apache.arrow.file"; #[cfg(test)] const JSON_CONTENT_TYPE: &str = "application/json"; +fn extract_job_id(body: &str) -> Option { + serde_json::from_str::(body) + .ok()? + .get("job_id")? + .as_str() + .filter(|job_id| !job_id.is_empty()) + .map(str::to_string) +} + pub use client::{ClientConfig, HeaderProvider, RetryConfig, TimeoutConfig, TlsConfig}; pub use db::{RemoteDatabaseOptions, RemoteDatabaseOptionsBuilder}; pub use oauth::{OAuthConfig, OAuthFlow, OAuthHeaderProvider}; diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 839cb3797..45a0bd925 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -9,6 +9,7 @@ use http::StatusCode; use lance_io::object_store::StorageOptions; use lance_namespace_impls::{DynamicContextProvider, OperationInfo}; use moka::future::Cache; +use reqwest::Response; use reqwest::header::CONTENT_TYPE; use lance_namespace::models::{ @@ -23,15 +24,17 @@ use crate::database::{ JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; +use crate::job::Job; +use crate::remote::job::RemoteJob; use crate::remote::util::stream_as_body; use crate::table::BaseTable; -use super::ARROW_STREAM_CONTENT_TYPE; use super::client::{ ClientConfig, HeaderProvider, HttpSend, RequestResultExt, RestfulLanceDbClient, Sender, }; use super::table::RemoteTable; use super::util::parse_server_version; +use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id}; // Request structure for the remote clone table API #[derive(serde::Serialize)] @@ -326,6 +329,22 @@ impl RemoteDatabase { } } +impl RemoteDatabase { + async fn submit_drop_table( + &self, + name: &str, + namespace_path: &[String], + ) -> Result<(String, Response)> { + let identifier = build_table_identifier(name, namespace_path, &self.client.id_delimiter); + let cache_key = build_cache_key(name, namespace_path); + let req = self.client.post(&format!("/v1/table/{}/drop/", identifier)); + let (request_id, resp) = self.client.send(req).await?; + let resp = self.client.check_response(&request_id, resp).await?; + self.table_cache.remove(&cache_key).await; + Ok((request_id, resp)) + } +} + #[cfg(all(test, feature = "remote"))] mod test_utils { use super::*; @@ -894,13 +913,28 @@ impl Database for RemoteDatabase { } async fn drop_table(&self, name: &str, namespace_path: &[String]) -> Result<()> { - let identifier = build_table_identifier(name, namespace_path, &self.client.id_delimiter); - let cache_key = build_cache_key(name, namespace_path); - let req = self.client.post(&format!("/v1/table/{}/drop/", identifier)); - let (request_id, resp) = self.client.send(req).await?; - self.client.check_response(&request_id, resp).await?; - self.table_cache.remove(&cache_key).await; - Ok(()) + self.submit_drop_table(name, namespace_path) + .await + .map(|_| ()) + } + + async fn drop_table_async(&self, name: &str, namespace_path: &[String]) -> Result { + let (request_id, response) = self.submit_drop_table(name, namespace_path).await?; + let status = response.status(); + let body = response.text().await.err_to_http(request_id.clone())?; + let job_id = extract_job_id(&body); + Ok(match job_id { + Some(job_id) => Job::new(Box::new(RemoteJob::new(self.client.clone(), job_id))), + None if status == StatusCode::ACCEPTED => { + return Err(Error::Http { + source: "asynchronous drop-table response did not contain a valid job_id" + .into(), + request_id, + status_code: Some(status), + }); + } + None => Job::new_done(), + }) } async fn drop_all_tables(&self, namespace_path: &[String]) -> Result<()> { @@ -1492,6 +1526,67 @@ mod tests { // NOTE: the API will return 200 even if the table does not exist. So we shouldn't expect 404. } + #[tokio::test] + async fn test_drop_table_does_not_read_response_body() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(200) + .body(vec![0xff]) + .unwrap() + }); + + conn.drop_table("table1", &[]).await.unwrap(); + } + + #[tokio::test] + async fn test_drop_table_async_returns_job() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/table/table1/drop/"); + http::Response::builder() + .status(202) + .body(r#"{"job_id":"drop-job-123"}"#) + .unwrap() + }); + + let job = conn.drop_table_async("table1", &[]).await.unwrap(); + assert_eq!(job.id(), Some("drop-job-123")); + } + + #[tokio::test] + async fn test_drop_table_async_old_server_returns_done_job() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder().status(200).body("").unwrap() + }); + + let job = conn.drop_table_async("table1", &[]).await.unwrap(); + assert_eq!(job.id(), None); + assert_eq!(job.status().await.unwrap(), "finished"); + } + + #[tokio::test] + async fn test_drop_table_async_rejects_accepted_response_without_job_id() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder().status(202).body("{}").unwrap() + }); + + let error = conn.drop_table_async("table1", &[]).await.err().unwrap(); + assert!(error.to_string().contains("valid job_id")); + } + + #[tokio::test] + async fn test_drop_table_async_rejects_empty_job_id() { + let conn = Connection::new_with_handler(|_| { + http::Response::builder() + .status(202) + .body(r#"{"job_id":""}"#) + .unwrap() + }); + + let error = conn.drop_table_async("table1", &[]).await.err().unwrap(); + assert!(error.to_string().contains("valid job_id")); + } + #[tokio::test] async fn test_rename_table() { let conn = Connection::new_with_handler(|request| { diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index f3e872cbe..8ca84a520 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -8,7 +8,7 @@ use self::insert::{RemoteWriteExec, WriteOp}; use super::client::RequestResultExt; use super::client::{HttpSend, RestfulLanceDbClient, Sender}; use super::db::ServerVersion; -use super::{ARROW_FILE_CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE}; +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}; use crate::expr::expr_to_sql_string; @@ -392,13 +392,7 @@ impl RemoteTable { .text() .await .ok() - .and_then(|body| serde_json::from_str::(&body).ok()) - .and_then(|value| { - value - .get("job_id") - .and_then(|id| id.as_str()) - .map(str::to_string) - }); + .and_then(|body| extract_job_id(&body)); if let Some(wait_timeout) = index.wait_timeout { let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); @@ -2520,9 +2514,9 @@ impl BaseTable for RemoteTable { self.check_mutable().await?; // Map the spec onto the server's request DTO. `sharding` is internally - // tagged on `mode` to mirror sophon's `Sharding` enum; `maintained_indexes` - // and `writer_config_defaults` are sent verbatim (an empty list means "no - // maintained indexes", not "default to all"). + // tagged on `mode` to mirror sophon's `Sharding` enum. A null + // `maintained_indexes` asks the server to resolve every maintainable + // index at HEAD; a list is verbatim, an empty one meaning none. let sharding = match &spec { LsmWriteSpec::Bucket { column, @@ -2706,6 +2700,13 @@ impl BaseTable for RemoteTable { Ok(result) } + // A declaration reaches here as AllNulls, which the remote protocol + // has no representation for. + NewColumnTransform::AllNulls(_) => { + return Err(Error::NotSupported { + message: "computed columns are supported only on local tables".into(), + }); + } _ => { return Err(Error::NotSupported { message: "Only SQL expressions are supported for adding columns".into(), @@ -6455,6 +6456,37 @@ mod tests { assert_eq!(result.version, if old_server { 0 } else { 43 }); } + /// Computed columns are local-only. Both halves say so here rather than + /// reaching the wire and failing somewhere less legible. + #[tokio::test] + async fn test_computed_columns_are_refused() { + let table = Table::new_with_handler("my_table", |request| -> http::Response { + panic!("unexpected request: {}", request.url().path()) + }); + + let declared = Arc::new(Schema::new(vec![Field::new( + "doubled", + DataType::Int32, + true, + )])); + let err = table + .add_columns() + .transform(NewColumnTransform::AllNulls(declared)) + .execute() + .await + .unwrap_err(); + assert!( + matches!(&err, Error::NotSupported { message } if message.contains("local tables")), + "{err:?}" + ); + + let err = table.refresh_column("doubled").await.unwrap_err(); + assert!( + matches!(&err, Error::NotSupported { message } if message.contains("local tables")), + "{err:?}" + ); + } + #[tokio::test] async fn test_prewarm_index() { let table = Table::new_with_handler("my_table", |request| { @@ -6599,7 +6631,7 @@ mod tests { .unwrap() }); let spec = crate::table::LsmWriteSpec::unsharded() - .with_maintained_indexes(["id_idx"]) + .with_maintained_indexes(vec!["id_idx".to_string()]) .with_writer_config_defaults([("max_memtable_rows", "1000")]); table.set_lsm_write_spec(spec).await.unwrap(); } @@ -6618,7 +6650,8 @@ mod tests { body["sharding"], serde_json::json!({ "mode": "bucket", "column": "id", "num_buckets": 16 }) ); - assert_eq!(body["maintained_indexes"], serde_json::json!([])); + // An unpinned maintained set sends null: resolve server-side. + assert_eq!(body["maintained_indexes"], serde_json::Value::Null); http::Response::builder().status(200).body("{}").unwrap() }); table @@ -6627,6 +6660,23 @@ mod tests { .unwrap(); } + /// `[]` (none) must stay distinguishable on the wire from null (all). + #[tokio::test] + async fn test_set_lsm_write_spec_no_maintained_indexes() { + let table = Table::new_with_handler("my_table", |request| { + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["maintained_indexes"], serde_json::json!([])); + http::Response::builder().status(200).body("{}").unwrap() + }); + table + .set_lsm_write_spec( + crate::table::LsmWriteSpec::bucket("id", 16).with_maintained_indexes(Vec::new()), + ) + .await + .unwrap(); + } + #[tokio::test] async fn test_set_lsm_write_spec_identity() { let table = Table::new_with_handler("my_table", |request| { @@ -6701,7 +6751,7 @@ mod tests { } => { assert_eq!(column, "id"); assert_eq!(num_buckets, 4); - assert_eq!(maintained_indexes, vec!["id_idx".to_string()]); + assert_eq!(maintained_indexes, Some(vec!["id_idx".to_string()])); assert_eq!( writer_config_defaults .get("durable_write") diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index c409591ef..9d436a0d0 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -50,7 +50,6 @@ use crate::DistanceType; use crate::blob::BlobRangeRequest; use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions}; use crate::database::Database; -use crate::database::listing::LANCE_FILE_EXTENSION; use crate::database::read_freshness::TableFreshness; use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; @@ -69,6 +68,7 @@ pub mod add_columns; mod add_data; pub mod branch_merge; pub mod checkpoint; +pub mod computed_columns; mod create_index; pub mod datafusion; pub(crate) mod dataset; @@ -78,6 +78,7 @@ pub mod merge; pub mod optimize; mod primary_key; pub mod query; +pub mod refresh; pub mod schema_evolution; pub mod update; pub mod write_progress; @@ -91,16 +92,19 @@ pub use branch_merge::{ MergeBranchResult, MergeBranchStatus, MergePreview, RowCountSummary, }; pub use chrono::Duration; +pub use computed_columns::{ + ComputedColumn, ComputedColumnKind, computed_column_from_field, computed_columns, +}; pub use delete::DeleteResult; use futures::future::join_all; pub use lance::dataset::refs::{BranchContents, Ref, TagContents, Tags as LanceTags}; pub use lance::dataset::scanner::DatasetRecordBatchStream; -use lance::dataset::statistics::DatasetStatisticsExt; pub use lance_index::optimize::OptimizeOptions; pub use lsm_stats::{BucketStats, GenerationStats, LsmStats, MemtableStats}; pub use optimize::{ CompactionMode, CompactionOptions, IndexRemapMode, OptimizeAction, OptimizeStats, }; +pub use refresh::RefreshColumnResult; pub use schema_evolution::{ AddColumnsResult, AlterColumnsResult, DropColumnsResult, FieldMetadataUpdate, UpdateFieldMetadataResult, @@ -155,55 +159,6 @@ pub(crate) fn map_namespace_lance_error(err: lance::Error, table_name: &str) -> } } -/// Map a `lance::Error::DatasetNotFound` for the table at `uri` into a `lancedb::Error`. -/// -/// Lance reports "there is nothing at this location" and "there is a table directory -/// here but nothing loadable inside it" with the same error. Only the first is a -/// `TableNotFound`: a `.lance` directory left behind by an interrupted drop and -/// re-create is still reported by `Connection::table_names`, so callers need to be able -/// to tell "never existed" from "exists but is broken". -/// -/// See . -async fn map_dataset_not_found( - uri: &str, - name: &str, - params: ReadParams, - err: lance::Error, -) -> Error { - let name = name.to_string(); - let source = Box::new(err); - if table_dir_exists(uri, params).await.unwrap_or(false) { - Error::TableCorrupted { name, source } - } else { - Error::TableNotFound { name, source } - } -} - -/// Whether a table directory is present at `uri`, even though no dataset could be -/// loaded from it. -/// -/// This looks for a `.lance` entry in the parent directory, which is exactly what -/// `ListingDatabase::table_names` lists, so the two APIs agree on whether a table is -/// present. Probing `uri` itself would not work: object stores have no empty -/// directories to probe, and on a local filesystem the interesting case is precisely an -/// empty directory. -async fn table_dir_exists(uri: &str, params: ReadParams) -> Result { - let (object_store, path, _) = DatasetBuilder::from_uri(uri) - .with_read_params(params) - .build_object_store() - .await?; - // Only `*.lance` entries are ever reported as tables, so nothing else can produce - // the list-then-open mismatch this guards against. - if path.extension() != Some(LANCE_FILE_EXTENSION) { - return Ok(false); - } - let (Some(parent), Some(dir_name)) = (path.parent(), path.filename()) else { - return Ok(false); - }; - let entries = object_store.read_dir(parent).await?; - Ok(entries.iter().any(|entry| entry.as_str() == dir_name)) -} - /// Defines the type of column #[derive(Debug, Clone, Serialize, Deserialize)] pub enum ColumnKind { @@ -373,6 +328,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. +/// /// Install a spec with [`Table::set_lsm_write_spec`] and remove it with /// [`Table::unset_lsm_write_spec`]. The actual `merge_insert` dispatch /// onto the MemWAL writer is a follow-up. @@ -387,9 +344,12 @@ pub enum LsmWriteSpec { Bucket { column: String, num_buckets: u32, - /// Names of indexes (already created on the table) that the - /// MemWAL should maintain in-memory as rows are appended. - maintained_indexes: Vec, + /// 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. + maintained_indexes: Option>, /// Default `ShardWriter` configuration recorded in the MemWAL index. writer_config_defaults: HashMap, }, @@ -399,35 +359,41 @@ pub enum LsmWriteSpec { /// distinct value of `column` becomes its own shard. Identity { column: String, - /// Names of indexes (already created on the table) that the - /// MemWAL should maintain in-memory as rows are appended. - maintained_indexes: Vec, + /// 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. + maintained_indexes: Option>, /// Default `ShardWriter` configuration recorded in the MemWAL index. writer_config_defaults: HashMap, }, /// No sharding — every `merge_insert` call writes to a single MemWAL shard. Unsharded { - /// Names of indexes (already created on the table) that the - /// MemWAL should maintain in-memory as rows are appended. - maintained_indexes: Vec, + /// 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. + maintained_indexes: Option>, /// Default `ShardWriter` configuration recorded in the MemWAL index. writer_config_defaults: HashMap, }, } impl LsmWriteSpec { - /// Construct a hash-bucket sharding spec with no maintained indexes. + /// Construct a hash-bucket sharding spec maintaining every index on the table. pub fn bucket(column: impl Into, num_buckets: u32) -> Self { Self::Bucket { column: column.into(), num_buckets, - maintained_indexes: Vec::new(), + maintained_indexes: None, writer_config_defaults: HashMap::new(), } } /// Construct an identity-sharding spec (shard by the raw value of - /// `column`) with no maintained indexes. + /// `column`) maintaining every index on the table. /// /// `column` must be a deterministic function of the unenforced primary /// key: every row with a given primary key must always produce the same @@ -439,28 +405,37 @@ impl LsmWriteSpec { pub fn identity(column: impl Into) -> Self { Self::Identity { column: column.into(), - maintained_indexes: Vec::new(), + maintained_indexes: None, writer_config_defaults: HashMap::new(), } } - /// Construct an unsharded spec with no maintained indexes. + /// Construct an unsharded spec maintaining every index on the table. pub fn unsharded() -> Self { Self::Unsharded { - maintained_indexes: Vec::new(), + maintained_indexes: None, writer_config_defaults: HashMap::new(), } } - /// Replace the list of indexes the MemWAL should keep up to date as - /// rows are appended. Each name must reference an index that already - /// exists on the table at the time `set_lsm_write_spec` is called. - pub fn with_maintained_indexes(mut self, indexes: I) -> Self - where - I: IntoIterator, - S: Into, - { - let v: Vec = indexes.into_iter().map(Into::into).collect(); + /// 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. + /// + /// ``` + /// # use lancedb::table::LsmWriteSpec; + /// // Every index the table has when the spec is installed: + /// LsmWriteSpec::unsharded().with_maintained_indexes(None); + /// // Exactly these: + /// LsmWriteSpec::unsharded().with_maintained_indexes(vec!["id_idx".to_string()]); + /// // None at all: + /// LsmWriteSpec::unsharded().with_maintained_indexes(Vec::new()); + /// ``` + pub fn with_maintained_indexes(mut self, indexes: impl Into>>) -> Self { + let indexes = indexes.into(); match &mut self { Self::Bucket { maintained_indexes, .. @@ -470,7 +445,7 @@ impl LsmWriteSpec { } | Self::Unsharded { maintained_indexes, .. - } => *maintained_indexes = v, + } => *maintained_indexes = indexes, } self } @@ -506,8 +481,9 @@ impl LsmWriteSpec { self } - /// Borrow the list of index names this spec asks MemWAL to maintain. - pub fn maintained_indexes(&self) -> &[String] { + /// Borrow the list of index names this spec asks MemWAL to maintain, or + /// `None` when it asks for every index on the table. + pub fn maintained_indexes(&self) -> Option<&[String]> { match self { Self::Bucket { maintained_indexes, .. @@ -517,7 +493,7 @@ impl LsmWriteSpec { } | Self::Unsharded { maintained_indexes, .. - } => maintained_indexes, + } => maintained_indexes.as_deref(), } } @@ -669,6 +645,15 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "set_lsm_write_spec is not supported on this table type".into(), }) } + /// Switch this table to required index catch-up, one way. + /// + /// The default implementation returns `NotSupported`. Implementations + /// that support the MemWAL LSM write path must override this. + async fn require_mem_wal_index_catchup(&self) -> Result<()> { + Err(Error::NotSupported { + message: "require_mem_wal_index_catchup is not supported on this table type".into(), + }) + } /// Remove the [`LsmWriteSpec`] from this table. /// /// This is a no-op if no spec is currently set. @@ -764,6 +749,23 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { transforms: NewColumnTransform, read_columns: Option>, ) -> Result; + /// Declare computed columns, each defined by a SQL expression. + async fn add_computed_columns( + &self, + _columns: &[(String, String)], + ) -> Result { + Err(Error::NotSupported { + message: "computed columns are not supported on this table type".into(), + }) + } + /// Fill a computed column's unfilled rows. + /// + /// The default returns `NotSupported`; Lance-backed tables override it. + async fn refresh_column(&self, _column: &str) -> Result { + Err(Error::NotSupported { + message: "computed columns are supported only on local tables".into(), + }) + } /// Alter columns in the table. async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result; /// Drop columns from the table. @@ -1656,6 +1658,29 @@ impl Table { AddColumnsBuilder::new(self.inner.clone()) } + /// Fill the fragments of a computed column that hold no values yet. + /// + /// Declared with + /// [`AddColumnsBuilder::computed`](add_columns::AddColumnsBuilder::computed), + /// a column starts empty and gets its values here. Fragments appended + /// since the last refresh are filled by the next one; fragments already + /// filled are left as they are, so the call is idempotent and does not + /// observe a mutated input. + /// + /// Local tables only. + /// + /// ``` + /// # use lancedb::Table; + /// # async fn refresh(table: &Table) -> Result<(), Box> { + /// let result = table.refresh_column("doubled").await?; + /// println!("filled {} rows at version {}", result.rows_filled, result.version); + /// # Ok(()) + /// # } + /// ``` + pub async fn refresh_column(&self, column: impl AsRef) -> Result { + self.inner.refresh_column(column.as_ref()).await + } + /// Change a column's name or nullability. pub async fn alter_columns( &self, @@ -1715,7 +1740,7 @@ impl Table { /// # async fn example(table: &Table) -> Result<(), Box> { /// table /// .set_lsm_write_spec( - /// LsmWriteSpec::bucket("id", 16).with_maintained_indexes(["id_idx"]), + /// LsmWriteSpec::bucket("id", 16).with_maintained_indexes(vec!["id_idx".to_string()]), /// ) /// .await?; /// # Ok(()) @@ -1725,6 +1750,20 @@ impl Table { self.inner.set_lsm_write_spec(spec).await } + /// Switch this table to required index catch-up, one way. + /// + /// Separate from [`Self::set_lsm_write_spec`] on purpose: a table carrying + /// the bit retains its SSTables until an index records that it holds the + /// compacted rows, so turn it on only once something can repair coverage. + /// A writer that already holds the dataset can call the equivalent on + /// `DatasetMemWalExt` instead; this is the table-level entry point. + /// + /// Errors if no spec is set, or if the table already records SSTable + /// compaction progress from before this protocol. + pub async fn require_mem_wal_index_catchup(&self) -> Result<()> { + self.inner.require_mem_wal_index_catchup().await + } + /// Remove the [`LsmWriteSpec`] from this table, reverting to the standard /// `merge_insert` write path. /// @@ -1737,9 +1776,10 @@ impl Table { /// /// Returns `Ok(None)` when the MemWAL LSM write path is not enabled (no /// spec has been set, or it was removed with [`Table::unset_lsm_write_spec`]). - /// The returned spec — including its [`LsmWriteSpec::maintained_indexes`] and - /// [`LsmWriteSpec::writer_config_defaults`] — mirrors what was passed to - /// [`Table::set_lsm_write_spec`]. + /// The returned spec mirrors what was passed to + /// [`Table::set_lsm_write_spec`], except that + /// [`LsmWriteSpec::maintained_indexes`] always reports the concrete list + /// resolved when the spec was set — `None` never round-trips. /// /// # Example /// @@ -2401,8 +2441,6 @@ impl NativeTable { None => false, }; - // Kept so that a `DatasetNotFound` can be re-checked against storage below. - let recovery_params = params.clone(); let mut builder = DatasetBuilder::from_uri(uri).with_read_params(params); // Set up commit handler when managed_versioning is enabled @@ -2421,7 +2459,12 @@ impl NativeTable { let dataset = match builder.load().await { Ok(dataset) => dataset, Err(e @ lance::Error::DatasetNotFound { .. }) => { - return Err(map_dataset_not_found(uri, name, recovery_params, e).await); + // The manifest load is the existence check. A physical prefix may be + // from a concurrent or abandoned create, so it cannot refine this error. + return Err(Error::TableNotFound { + name: name.to_string(), + source: Box::new(e), + }); } Err(e) => return Err(e.into()), }; @@ -2633,6 +2676,7 @@ impl NativeTable { namespace_client: Option>, pushdown_operations: HashSet, ) -> Result { + computed_columns::ensure_no_foreign_declarations(batches.arrow_schema().fields())?; // Default params uses format v1. let params = params.unwrap_or(WriteParams { ..Default::default() @@ -3081,6 +3125,13 @@ impl BaseTable for NativeTable { let ds = self.dataset.get().await?; let table_schema = Schema::from(&ds.schema().clone()); + computed_columns::ensure_not_written( + &table_schema, + add.data.schema().fields().iter().map(|f| f.name().as_str()), + )?; + if matches!(add.mode, AddDataMode::Overwrite) { + computed_columns::ensure_no_foreign_declarations(add.data.schema().fields())?; + } let num_partitions = if let Some(parallelism) = add.write_parallelism { parallelism @@ -3241,6 +3292,11 @@ impl BaseTable for NativeTable { params: MergeInsertBuilder, new_data: Box, ) -> Result { + let source_schema = arrow_array::RecordBatchReader::schema(&new_data); + computed_columns::ensure_not_written( + &Schema::from(self.dataset.get().await?.schema()), + source_schema.fields().iter().map(|f| f.name().as_str()), + )?; let result = merge::execute_merge_insert(self, params, new_data).await?; self.bump_freshness(); Ok(result) @@ -3254,6 +3310,10 @@ impl BaseTable for NativeTable { merge::lsm::set_lsm_write_spec(self, spec).await } + async fn require_mem_wal_index_catchup(&self) -> Result<()> { + merge::lsm::require_mem_wal_index_catchup(self).await + } + async fn unset_lsm_write_spec(&self) -> Result<()> { merge::lsm::unset_lsm_write_spec(self).await } @@ -3322,6 +3382,18 @@ impl BaseTable for NativeTable { Ok(result) } + async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { + let result = schema_evolution::execute_declare(self, columns).await?; + self.bump_freshness(); + Ok(result) + } + + async fn refresh_column(&self, column: &str) -> Result { + let result = refresh::execute_refresh_column(self, column).await?; + self.bump_freshness(); + Ok(result) + } + async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result { let result = schema_evolution::execute_alter_columns(self, alterations).await?; self.bump_freshness(); @@ -3550,9 +3622,24 @@ impl BaseTable for NativeTable { let num_rows = self.count_rows(None).await?; let num_indices = self.list_indices().await?.len(); let ds = self.dataset.get().await?; - let ds_clone = (*ds).clone(); - let ds_stats = Arc::new(ds_clone).calculate_data_stats().await?; - let total_bytes = ds_stats.fields.iter().map(|f| f.bytes_on_disk).sum::() as usize; + // Sizes come from the manifest. Summing per-field `bytes_on_disk` instead + // would open every data file to read its column metadata, which costs one + // IO per fragment and reports 0 for legacy v1 storage. + // + // The manifest summary covers only the fragments' base data files, so + // overlay files (recorded on each fragment) and index files (recorded in + // the manifest's index section) are added separately. + let mut total_bytes = ds.manifest().summary().total_files_size as usize; + for frag in ds.manifest().fragments.iter() { + for overlay in &frag.overlays { + if let Some(size) = overlay.data_file.file_size_bytes.get() { + total_bytes += size.get() as usize; + } + } + } + for index in ds.load_indices().await?.iter() { + total_bytes += index.total_size_bytes().unwrap_or(0) as usize; + } let frags = ds.get_fragments(); let mut sorted_sizes = join_all( @@ -3624,7 +3711,12 @@ impl BaseTable for NativeTable { #[skip_serializing_none] #[derive(Debug, Deserialize, PartialEq)] pub struct TableStatistics { - /// The total number of bytes in the table + /// The total size, in bytes, of the table's data files, index files, and + /// overlay files + /// + /// Read from the manifest, so this excludes deletion files and manifests, + /// and it excludes any file whose size the manifest does not record + /// (tables and indices written before writers persisted file sizes). pub total_bytes: usize, /// The number of rows in the table @@ -3669,7 +3761,7 @@ pub struct FragmentSummaryStats { #[allow(deprecated)] mod tests { use std::sync::Arc; - use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; use arrow_array::{ @@ -3685,6 +3777,7 @@ mod tests { use super::*; use crate::connect; use crate::connection::ConnectBuilder; + use crate::io::object_store::io_tracking::IoTrackingStore; use crate::query::Select; use crate::query::{ExecutableQuery, QueryBase}; use crate::test_utils::connection::new_test_connection; @@ -3750,73 +3843,50 @@ mod tests { ); } - /// Write a table and then break it, leaving the `.lance` directory in place. - /// - /// `remove_all` reproduces an interrupted drop + re-create (the directory is left - /// empty); otherwise only the manifests are removed, leaving the data files behind. - async fn write_then_corrupt_table(dir: &std::path::Path, remove_all: bool) -> String { - let dataset_path = dir.join("test.lance"); - let uri = dataset_path.to_str().unwrap().to_string(); - - let batch = make_test_batches(); - let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); - Dataset::write(reader, &uri, None).await.unwrap(); - - if remove_all { - for entry in std::fs::read_dir(&dataset_path).unwrap() { - let entry = entry.unwrap(); - if entry.file_type().unwrap().is_dir() { - std::fs::remove_dir_all(entry.path()).unwrap(); - } else { - std::fs::remove_file(entry.path()).unwrap(); - } - } - assert_eq!(std::fs::read_dir(&dataset_path).unwrap().count(), 0); - } else { - let versions = dataset_path.join("_versions"); - assert!(versions.is_dir(), "expected manifests under {versions:?}"); - std::fs::remove_dir_all(&versions).unwrap(); - assert!(std::fs::read_dir(&dataset_path).unwrap().count() > 0); - } - - uri - } - #[tokio::test] - async fn test_open_corrupt_empty_dir() { + async fn test_open_not_found_when_empty_directory_exists() { let tmp_dir = tempdir().unwrap(); - let uri = write_then_corrupt_table(tmp_dir.path(), true).await; + let dataset_path = tmp_dir.path().join("test.lance"); + std::fs::create_dir(&dataset_path).unwrap(); - let err = NativeTable::open(&uri).await.unwrap_err(); + let err = NativeTable::open(dataset_path.to_str().unwrap()) + .await + .unwrap_err(); assert!( - matches!(&err, Error::TableCorrupted { name, .. } if name == "test"), + matches!(&err, Error::TableNotFound { name, .. } if name == "test"), "got {err:?}" ); } #[tokio::test] - async fn test_open_corrupt_missing_manifest() { + async fn test_open_not_found_when_only_uncommitted_storage_exists() { let tmp_dir = tempdir().unwrap(); - let uri = write_then_corrupt_table(tmp_dir.path(), false).await; + let dataset_path = tmp_dir.path().join("test.lance"); + let data_dir = dataset_path.join("data"); + std::fs::create_dir_all(&data_dir).unwrap(); + std::fs::write(data_dir.join("orphan.lance"), b"uncommitted").unwrap(); - let err = NativeTable::open(&uri).await.unwrap_err(); + let err = NativeTable::open(dataset_path.to_str().unwrap()) + .await + .unwrap_err(); assert!( - matches!(&err, Error::TableCorrupted { name, .. } if name == "test"), + matches!(&err, Error::TableNotFound { name, .. } if name == "test"), "got {err:?}" ); } - /// A table listed by `table_names()` must not be reported as missing by - /// `open_table()`. See . + /// Listing databases discover physical `*.lance` entries. That snapshot is not an + /// authoritative table-existence check: only a committed manifest makes a table + /// openable, and the entry could also be concurrently created or dropped. #[tokio::test] - async fn test_open_table_corrupt_is_still_listed() { + async fn test_table_names_may_include_uncommitted_storage() { let tmp_dir = tempdir().unwrap(); let db = connect(tmp_dir.path().to_str().unwrap()) .execute() .await .unwrap(); - write_then_corrupt_table(tmp_dir.path(), true).await; + std::fs::create_dir(tmp_dir.path().join("test.lance")).unwrap(); assert_eq!( db.table_names().execute().await.unwrap(), @@ -3824,12 +3894,177 @@ mod tests { ); let err = db.open_table("test").execute().await.unwrap_err(); assert!( - matches!(&err, Error::TableCorrupted { name, .. } if name == "test"), + matches!(&err, Error::TableNotFound { name, .. } if name == "test"), + "physical storage without a committed manifest is not a table: {err:?}" + ); + } + + #[derive(Debug)] + struct ParentListGuardStore { + inner: Arc, + parent: object_store::path::Path, + parent_list_calls: Arc, + } + + impl std::fmt::Display for ParentListGuardStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("ParentListGuardStore") + } + } + + #[async_trait::async_trait] + #[deny(clippy::missing_trait_methods)] + impl object_store::ObjectStore for ParentListGuardStore { + async fn put_opts( + &self, + location: &object_store::path::Path, + payload: object_store::PutPayload, + opts: object_store::PutOptions, + ) -> object_store::Result { + self.inner.put_opts(location, payload, opts).await + } + + async fn put_multipart_opts( + &self, + location: &object_store::path::Path, + opts: object_store::PutMultipartOptions, + ) -> object_store::Result> { + self.inner.put_multipart_opts(location, opts).await + } + + async fn get_opts( + &self, + location: &object_store::path::Path, + options: object_store::GetOptions, + ) -> object_store::Result { + self.inner.get_opts(location, options).await + } + + async fn get_ranges( + &self, + location: &object_store::path::Path, + ranges: &[std::ops::Range], + ) -> object_store::Result> { + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: futures::stream::BoxStream< + 'static, + object_store::Result, + >, + ) -> futures::stream::BoxStream<'static, object_store::Result> + { + self.inner.delete_stream(locations) + } + + fn list( + &self, + prefix: Option<&object_store::path::Path>, + ) -> futures::stream::BoxStream<'static, object_store::Result> + { + if prefix == Some(&self.parent) { + self.parent_list_calls.fetch_add(1, Ordering::Relaxed); + } + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&object_store::path::Path>, + offset: &object_store::path::Path, + ) -> futures::stream::BoxStream<'static, object_store::Result> + { + if prefix == Some(&self.parent) { + self.parent_list_calls.fetch_add(1, Ordering::Relaxed); + } + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter( + &self, + prefix: Option<&object_store::path::Path>, + ) -> object_store::Result { + if prefix == Some(&self.parent) { + self.parent_list_calls.fetch_add(1, Ordering::Relaxed); + } + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &object_store::path::Path, + to: &object_store::path::Path, + options: object_store::CopyOptions, + ) -> object_store::Result<()> { + self.inner.copy_opts(from, to, options).await + } + + async fn rename_opts( + &self, + from: &object_store::path::Path, + to: &object_store::path::Path, + options: object_store::RenameOptions, + ) -> object_store::Result<()> { + self.inner.rename_opts(from, to, options).await + } + } + + #[derive(Debug)] + struct ParentListGuardWrapper { + parent_list_calls: Arc, + } + + impl WrappingObjectStore for ParentListGuardWrapper { + fn wrap( + &self, + _store_prefix: &str, + inner: Arc, + ) -> Arc { + Arc::new(ParentListGuardStore { + inner, + parent: object_store::path::Path::from("database"), + parent_list_calls: self.parent_list_calls.clone(), + }) + } + } + + #[tokio::test] + async fn test_open_missing_never_lists_database_parent() { + let parent_list_calls = Arc::new(AtomicUsize::new(0)); + let params = ReadParams { + store_options: Some(ObjectStoreParams { + object_store_wrapper: Some(Arc::new(ParentListGuardWrapper { + parent_list_calls: parent_list_calls.clone(), + })), + ..Default::default() + }), + ..Default::default() + }; + + let err = NativeTable::open_with_params( + "memory:///database/missing.lance", + "missing", + Vec::new(), + None, + Some(params), + None, + None, + HashSet::new(), + None, + ) + .await + .unwrap_err(); + + assert!( + matches!(&err, Error::TableNotFound { name, .. } if name == "missing"), "got {err:?}" ); - assert!( - err.to_string().contains("exists but could not be loaded"), - "got {err}" + assert_eq!( + parent_list_calls.load(Ordering::Relaxed), + 0, + "opening one missing table must not enumerate sibling tables" ); } @@ -5067,7 +5302,7 @@ mod tests { // Bucket spec round-trips exactly, including the routing column (recovered // from its field id), maintained indexes, and writer config defaults. let spec = LsmWriteSpec::bucket("id", 4) - .with_maintained_indexes([idx_name]) + .with_maintained_indexes(vec![idx_name.clone()]) .with_writer_config_defaults([("durable_write", "false")]); table.set_lsm_write_spec(spec.clone()).await.unwrap(); assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec)); @@ -5077,15 +5312,125 @@ 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". 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)); + assert_eq!( + table.get_lsm_write_spec().await.unwrap(), + Some(spec.with_maintained_indexes(vec![idx_name.clone()])) + ); 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)); + assert_eq!( + table.get_lsm_write_spec().await.unwrap(), + Some(spec.with_maintained_indexes(vec![idx_name])) + ); + } + + /// 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. + #[tokio::test] + async fn test_set_lsm_write_spec_infers_maintained_indexes() { + let tmp_dir = tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("tag", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + let reader: Box = + Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone())); + let conn = ConnectBuilder::new(uri) + .read_consistency_interval(Duration::from_secs(0)) + .execute() + .await + .unwrap(); + let table = conn.create_table("t", reader).execute().await.unwrap(); + + table + .create_index(&["id"], Index::BTree(Default::default())) + .name("id_btree".to_string()) + .execute() + .await + .unwrap(); + table + .create_index(&["tag"], Index::Bitmap(Default::default())) + .name("tag_bitmap".to_string()) + .execute() + .await + .unwrap(); + + // Explicitly naming the bitmap index fails before anything commits. + let err = table + .set_lsm_write_spec( + LsmWriteSpec::unsharded().with_maintained_indexes(vec!["tag_bitmap".to_string()]), + ) + .await + .unwrap_err(); + assert!( + matches!(err, Error::InvalidInput { ref message } if message.contains("tag_bitmap")), + "expected the bitmap index to be rejected, got {err:?}" + ); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), None); + + // The default covers every index, so the bitmap fails it too. + let err = 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:?}" + ); + assert_eq!(table.get_lsm_write_spec().await.unwrap(), None); + + // Naming the maintainable subset installs. + table + .set_lsm_write_spec( + LsmWriteSpec::unsharded().with_maintained_indexes(vec!["id_btree".to_string()]), + ) + .await + .unwrap(); + assert_eq!( + table + .get_lsm_write_spec() + .await + .unwrap() + .unwrap() + .maintained_indexes(), + Some(["id_btree".to_string()].as_slice()) + ); + + // Opting out entirely is distinct from the default. + table.unset_lsm_write_spec().await.unwrap(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes(Vec::new())) + .await + .unwrap(); + assert_eq!( + table + .get_lsm_write_spec() + .await + .unwrap() + .unwrap() + .maintained_indexes(), + Some([].as_slice()) + ); } #[tokio::test] @@ -5133,12 +5478,16 @@ mod tests { let res = table.stats().await.unwrap(); println!("{:#?}", res); + // `total_bytes` is the full on-disk size of the 11 data files (this table + // has no index or overlay files), so it is well above the 2000 bytes of + // column data these 250 int32 pairs hold: each file carries its own footer + // and metadata. assert_eq!( res, TableStatistics { num_rows: 250, num_indices: 0, - total_bytes: 2300, + total_bytes: 8925, fragment_stats: FragmentStatistics { num_fragments: 11, num_small_fragments: 11, @@ -5178,4 +5527,196 @@ mod tests { } ) } + + /// `total_bytes` counts more than the base data files: index files and + /// overlay files recorded in the manifest are included too. + #[tokio::test] + pub async fn test_stats_includes_index_and_overlay_files() { + use lance::dataset::WriteDestination; + use lance::dataset::transaction::{DataOverlayGroup, Operation}; + use lance_file::version::stable_file_version; + use lance_file::writer::FileWriterOptions; + use lance_io::utils::CachedFileSize; + use lance_table::format::DataFile; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let tmp_dir = tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + let conn = ConnectBuilder::new(uri) + .read_consistency_interval(Duration::from_secs(0)) + .execute() + .await + .unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("foo", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..100)), + Arc::new(Int32Array::from_iter_values(0..100)), + ], + ) + .unwrap(); + let table = conn + .create_table("test_stats_extra_files", batch) + .execute() + .await + .unwrap(); + + let data_only = table.stats().await.unwrap().total_bytes; + assert!(data_only > 0); + + // A scalar index adds index files whose sizes are recorded in the + // manifest's index section. + table + .create_index(&["id"], Index::Auto) + .execute() + .await + .unwrap(); + let with_index = table.stats().await.unwrap().total_bytes; + let dataset = { + let native = table.as_native().unwrap(); + (*native.dataset.get().await.unwrap()).clone() + }; + let index_bytes: usize = dataset + .load_indices() + .await + .unwrap() + .iter() + .map(|idx| idx.total_size_bytes().unwrap_or(0) as usize) + .sum(); + assert!(index_bytes > 0); + assert_eq!(with_index, data_only + index_bytes); + + // Commit an overlay file supplying new `foo` values for the first three + // rows of fragment 0. There is no high-level API that writes overlays + // yet, so write the overlay's data file and commit the `DataOverlay` + // operation by hand. + let read_version = dataset.version().version; + let fragment_id = dataset.get_fragments()[0].id() as u64; + let foo_field_id = dataset.schema().field("foo").unwrap().id; + let overlay_schema = dataset.schema().project_by_ids(&[foo_field_id], true); + let file_version = stable_file_version(); + + let filename = "overlay.lance".to_string(); + let store = dataset.object_store(None).await.unwrap(); + let path = dataset.data_dir().child(filename.clone()); + let obj_writer = store.create(&path).await.unwrap(); + let mut writer = lance_file::versions::create_writer( + file_version, + obj_writer, + overlay_schema, + FileWriterOptions::default(), + ) + .unwrap(); + writer + .write_column(0, Arc::new(Int32Array::from(vec![1000, 1001, 1002])) as _) + .await + .unwrap(); + let summary = writer.finish().await.unwrap(); + let overlay_bytes = summary.size_bytes as usize; + assert!(overlay_bytes > 0); + + let mut data_file = DataFile::new_unstarted(filename, file_version); + data_file.fields = writer + .field_id_to_column_indices() + .iter() + .map(|(field_id, _)| *field_id as i32) + .collect::>() + .into(); + data_file.column_indices = writer + .field_id_to_column_indices() + .iter() + .map(|(_, column_index)| *column_index as i32) + .collect::>() + .into(); + data_file.file_size_bytes = CachedFileSize::new(summary.size_bytes); + + let overlay = DataOverlayFile { + data_file, + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter(0..3)), + committed_version: 0, + }; + Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id, + overlays: vec![overlay], + }], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + + table.checkout_latest().await.unwrap(); + let with_overlay = table.stats().await.unwrap().total_bytes; + assert_eq!(with_overlay, with_index + overlay_bytes); + } + + /// `stats()` must stay manifest-only. Summing per-field `bytes_on_disk` + /// instead opens every data file, so cost would grow with fragment count. + #[tokio::test] + pub async fn test_stats_does_not_read_data_files() { + let tmp_dir = tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + + let conn = ConnectBuilder::new(uri).execute().await.unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10))], + ) + .unwrap(); + + conn.create_table("test_stats_io", batch.clone()) + .execute() + .await + .unwrap(); + let table = conn.open_table("test_stats_io").execute().await.unwrap(); + const NUM_APPENDS: usize = 20; + for _ in 0..NUM_APPENDS { + table.add(batch.clone()).execute().await.unwrap(); + } + + // Reopen through a tracking store so the counters cover `stats()` alone and + // not the writes above. + let (wrapper, io_stats) = IoTrackingStore::new_wrapper(); + let table = conn + .open_table("test_stats_io") + .lance_read_params(ReadParams { + store_options: Some(ObjectStoreParams { + object_store_wrapper: Some(wrapper), + ..Default::default() + }), + ..Default::default() + }) + .execute() + .await + .unwrap(); + io_stats.lock().unwrap().read_iops = 0; + + let stats = table.stats().await.unwrap(); + let read_iops = io_stats.lock().unwrap().read_iops; + + assert_eq!(stats.fragment_stats.num_fragments, NUM_APPENDS + 1); + assert!(stats.total_bytes > 0); + // Reading the fragments' data files would take at least one IOP each. + assert!( + read_iops < stats.fragment_stats.num_fragments as u64, + "stats() issued {} read IOPs across {} fragments", + read_iops, + stats.fragment_stats.num_fragments + ); + } } diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs index 0d410cd04..6aa2ce86a 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -15,6 +15,7 @@ use crate::{Error, Result}; pub struct AddColumnsBuilder { parent: Arc, transform: Option, + computed: Vec<(String, String)>, read_columns: Option>, } @@ -23,6 +24,7 @@ impl std::fmt::Debug for AddColumnsBuilder { f.debug_struct("AddColumnsBuilder") .field("parent", &self.parent) .field("has_transform", &self.transform.is_some()) + .field("computed", &self.computed) .field("read_columns", &self.read_columns) .finish() } @@ -33,19 +35,57 @@ impl AddColumnsBuilder { Self { parent, transform: None, + computed: Vec::new(), read_columns: None, } } - /// Set how the new columns' values are produced. Required. + /// Set how the new columns' values are produced. pub fn transform(mut self, transform: NewColumnTransform) -> Self { self.transform = Some(transform); self } + /// Add a column defined by `expression`, evaluated by a later refresh + /// rather than by this commit. Its type and inputs are derived from the + /// expression. + /// + /// The column is committed with no values, so declaring one costs the same + /// on an empty table as on a large one. Rows get values from + /// [`Table::refresh_column`](super::Table::refresh_column), which fills + /// every fragment that has none -- including fragments appended since the + /// last refresh. + /// + /// Refresh does not revisit a fragment it has filled, so mutating an input + /// leaves the value computed at fill time; recomputing means dropping the + /// column and declaring it again. An input cannot be renamed, retyped or + /// dropped while a declaration reads it, since the expression names it. + /// + /// Local tables only: LanceDB Cloud and Enterprise reject a declaration + /// with `NotSupported`. + /// + /// ``` + /// # use lancedb::Table; + /// # async fn declare(table: &Table) -> Result<(), Box> { + /// table + /// .add_columns() + /// .computed("doubled", "x * 2") + /// .execute() + /// .await?; + /// let filled = table.refresh_column("doubled").await?; + /// println!("filled {} rows", filled.rows_filled); + /// # Ok(()) + /// # } + /// ``` + pub fn computed(mut self, name: impl Into, expression: impl Into) -> Self { + self.computed.push((name.into(), expression.into())); + self + } + /// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper - /// receives. Every other transform determines what it reads, so setting - /// this alongside one is an error rather than a silent no-op. + /// receives. Every other transform, and a computed column, determines what + /// it reads, so setting this alongside one is an error rather than a silent + /// no-op. pub fn read_columns(mut self, columns: impl IntoIterator>) -> Self { self.read_columns = Some(columns.into_iter().map(Into::into).collect()); self @@ -56,24 +96,42 @@ impl AddColumnsBuilder { let Self { parent, transform, + computed, read_columns, } = self; - let Some(transform) = transform else { - return Err(Error::InvalidInput { - message: "add_columns requires a transform".into(), - }); - }; - - if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) { - return Err(Error::InvalidInput { - message: "read_columns applies only to a BatchUDF transform; \ - every other transform determines what it reads" + match (transform, computed.is_empty()) { + (None, true) => Err(Error::InvalidInput { + message: "add_columns requires a transform or a computed column".into(), + }), + // The two commit through different transforms, so one call covering + // both would be two commits and could half-apply. + (Some(_), false) => Err(Error::InvalidInput { + message: "add_columns cannot mix a transform with computed columns; \ + they cannot be added atomically in one call" .into(), - }); + }), + (Some(transform), true) => { + if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) { + return Err(Error::InvalidInput { + message: "read_columns applies only to a BatchUDF transform; \ + every other transform determines what it reads" + .into(), + }); + } + parent.add_columns(transform, read_columns).await + } + (None, false) => { + if read_columns.is_some() { + return Err(Error::InvalidInput { + message: "read_columns applies only to a BatchUDF transform; \ + a computed column's inputs come from its expression" + .into(), + }); + } + parent.add_computed_columns(&computed).await + } } - - parent.add_columns(transform, read_columns).await } } @@ -85,8 +143,8 @@ mod tests { use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{BatchUDF, NewColumnTransform}; - use crate::Table; use crate::connect; + use crate::{Error, Table}; async fn table_with_two_columns(name: &str) -> Table { let conn = connect("memory://").execute().await.unwrap(); @@ -98,10 +156,7 @@ mod tests { async fn test_requires_a_transform() { let table = table_with_two_columns("no_transform").await; let err = table.add_columns().execute().await.unwrap_err(); - assert!( - err.to_string().contains("requires a transform"), - "got: {err}" - ); + assert!(matches!(err, Error::InvalidInput { .. })); } #[tokio::test] @@ -117,7 +172,7 @@ mod tests { .execute() .await .unwrap_err(); - assert!(err.to_string().contains("BatchUDF"), "got: {err}"); + assert!(matches!(err, Error::InvalidInput { .. })); let schema = table.schema().await.unwrap(); assert!( @@ -126,6 +181,47 @@ mod tests { ); } + #[tokio::test] + async fn test_mixing_transform_and_computed_is_rejected() { + let table = table_with_two_columns("mixed_add").await; + let err = table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "eager".into(), + "x * 2".into(), + )])) + .computed("lazy", "x * 3") + .execute() + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + + let schema = table.schema().await.unwrap(); + assert!(schema.field_with_name("eager").is_err()); + assert!(schema.field_with_name("lazy").is_err()); + } + + #[tokio::test] + async fn test_read_columns_with_computed_is_rejected() { + let table = table_with_two_columns("read_cols_computed").await; + let err = table + .add_columns() + .computed("doubled", "x * 2") + .read_columns(["x"]) + .execute() + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + assert!( + table + .schema() + .await + .unwrap() + .field_with_name("doubled") + .is_err() + ); + } + #[tokio::test] async fn test_read_columns_limits_what_a_batch_udf_sees() { let table = table_with_two_columns("read_cols_udf").await; diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs new file mode 100644 index 000000000..9a6a2585d --- /dev/null +++ b/rust/lancedb/src/table/computed_columns.rs @@ -0,0 +1,1348 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Computed columns. +//! +//! A computed column is defined by a rule rather than by values supplied at +//! write time. Declaring one commits the column carrying that rule in field +//! metadata but no data, so the cost does not scale with the table; a later +//! refresh fills the rows. +//! +//! The rule is tagged by kind ([`ComputedColumnKind`]) because kinds differ in +//! where the column's type and inputs come from. A SQL expression is +//! self-describing -- both are derived from the expression, so a caller writes +//! neither -- while a kind resolved through a registry cannot be typed without +//! consulting it. Only SQL exists today; the tag is what lets another kind be +//! added without a second reading of the same key. +//! +//! [`computed_columns`] and [`computed_column_from_field`] read declarations +//! back off a schema. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; +use datafusion_common::tree_node::TreeNode; +use datafusion_physical_plan::PhysicalExpr; +use lance::dataset::NewColumnTransform; +use lance_datafusion::planner::Planner; + +use crate::{Error, Result}; + +/// Field metadata key marking a column as computed. The value is `"true"`. +pub const COMPUTED_COLUMN_META_KEY: &str = "computed_column"; + +/// Field metadata key naming the kind of rule that defines the column. +pub const KIND_META_KEY: &str = "computed_column.kind"; + +/// Field metadata key holding the SQL expression that defines the column. +pub const EXPRESSION_META_KEY: &str = "computed_column.expression"; + +/// Field metadata key holding the column's inputs, as a JSON array of names. +pub const INPUTS_META_KEY: &str = "computed_column.inputs"; + +/// Value of [`KIND_META_KEY`] for a column defined by a SQL expression. +pub const SQL_KIND: &str = "sql"; + +/// The rule that defines a computed column's values. +/// +/// Non-exhaustive: a kind added later is an additive change, and a caller that +/// only handles the kinds it knows keeps compiling. +#[derive(Debug, Clone, PartialEq, Eq)] +#[non_exhaustive] +pub enum ComputedColumnKind { + /// A SQL expression evaluated by DataFusion. It is the whole definition: + /// the column's type and its inputs are both derived from it. + Sql { + /// The expression. + expression: String, + }, + /// A kind this version does not understand, written by a newer one. + /// + /// Reported rather than hidden so a caller can tell a column it cannot + /// refresh apart from one that was never computed. Nothing produces this. + Unrecognized { + /// The kind as it was found in the metadata. + kind: String, + }, +} + +/// A computed column's declaration, as read back from field metadata. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ComputedColumn { + /// Name of the computed column. + pub name: String, + /// The rule that defines it. + pub kind: ComputedColumnKind, + /// Columns the rule reads, recorded at declaration time. + /// + /// Outside the kind because every kind has inputs and the consumers that + /// use them -- refresh planning, dependency ordering -- do not care which + /// kind produced them. Where they come from does differ, and that is + /// settled at declaration: derived from a SQL expression, supplied by the + /// caller for a kind that cannot be parsed. + pub inputs: Vec, +} + +/// Build the field metadata recording a SQL binding. +fn computed_column_metadata(expression: &str, inputs: &[String]) -> HashMap { + 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(), + serde_json::to_string(inputs).unwrap_or_else(|_| "[]".to_string()), + ), + ]) +} + +/// Read a field's computed-column declaration, if it carries one. +/// +/// A field flagged computed but carrying no kind, or a SQL one missing its +/// expression, is not a computed column here: without the rule there is +/// nothing to refresh from, so it is reported as absent rather than as a +/// half-formed declaration. An unrecognized kind is different -- the rule is +/// there and intact, this version just cannot act on it -- and comes back as +/// [`ComputedColumnKind::Unrecognized`]. +pub fn computed_column_from_field(field: &ArrowField) -> Option { + let metadata = field.metadata(); + if metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str) != Some("true") { + return None; + } + let kind = match metadata.get(KIND_META_KEY)?.as_str() { + SQL_KIND => ComputedColumnKind::Sql { + expression: metadata.get(EXPRESSION_META_KEY)?.clone(), + }, + other => ComputedColumnKind::Unrecognized { + kind: other.to_string(), + }, + }; + let inputs = metadata + .get(INPUTS_META_KEY) + .and_then(|raw| serde_json::from_str::>(raw).ok()) + .unwrap_or_default(); + Some(ComputedColumn { + name: field.name().clone(), + kind, + inputs, + }) +} + +/// Read every computed-column declaration carried by `schema`, in field order. +/// +/// Introspection is a pure read of the schema the caller already holds, the +/// way a SQL catalog reports a generation expression as another column of +/// `information_schema.columns`. +pub fn computed_columns(schema: &ArrowSchema) -> Vec { + schema + .fields() + .iter() + .filter_map(|field| computed_column_from_field(field)) + .collect() +} + +/// Reject a schema change to a column some declaration reads. +/// +/// A binding is SQL text naming its inputs, so renaming, retyping or dropping +/// one leaves an expression that no longer resolves. Refusing the change keeps +/// a declaration that survived [`plan`] evaluable for as long as it exists. +/// +/// Paths are compared at their root: a declaration reading `metadata` is +/// invalidated by a change to `metadata.age` just as surely. +pub(crate) fn ensure_not_an_input(schema: &SchemaRef, paths: &[&str]) -> Result<()> { + for declaration in computed_columns(schema) { + // The expression, not stored inputs, is the source of truth; an + // expression that no longer parses proves nothing, so refuse. + let inputs = match &declaration.kind { + ComputedColumnKind::Sql { expression } => Planner::new(schema.clone()) + .parse_expr(expression) + .map(|parsed| Planner::column_names_in_expr(&parsed)) + .map_err(|e| Error::InvalidInput { + message: format!( + "computed column '{}' has an unevaluable expression ({e}); drop it \ + before changing the schema", + declaration.name + ), + })?, + _ => declaration.inputs.clone(), + }; + for path in paths { + // Exact target only: the binding travels with the whole column, + // not with a nested field the expression still shapes. + if declaration.name == *path { + continue; + } + if declaration.name == root(path) { + return Err(Error::InvalidInput { + message: format!( + "'{}' is part of computed column '{}'; drop the column and declare \ + it again", + path, declaration.name + ), + }); + } + if inputs.iter().any(|input| root(input) == root(path)) { + return Err(Error::InvalidInput { + message: format!( + "column '{}' is read by computed column '{}'; drop that column first", + path, declaration.name + ), + }); + } + } + } + Ok(()) +} + +/// Reject a write that supplies values for a computed column directly: +/// only refresh materializes one, and refresh never revisits a filled row. +pub(crate) fn ensure_not_written<'a>( + schema: &ArrowSchema, + written: impl IntoIterator, +) -> Result<()> { + let declared: Vec = computed_columns(schema) + .into_iter() + .map(|declaration| declaration.name) + .collect(); + for name in written { + if declared.iter().any(|declared| declared == root(name)) { + return Err(Error::InvalidInput { + message: format!( + "column '{}' is computed; its values come from refresh and cannot be \ + written directly", + root(name) + ), + }); + } + } + Ok(()) +} + +/// Reject a batch holding values for a computed column. Null slots are the +/// declared state, so planner-padded placeholders pass. +pub(crate) fn ensure_batch_writes_no_computed_values( + declared: &[String], + batch: &arrow_array::RecordBatch, +) -> Result<()> { + for name in declared { + if let Some(column) = batch.column_by_name(name) + && column.null_count() != column.len() + { + return Err(Error::InvalidInput { + message: format!( + "column '{name}' is computed; its values come from refresh and cannot \ + be written directly" + ), + }); + } + } + Ok(()) +} + +/// Reject fields carrying declaration metadata that did not come through +/// [`plan`]. One authority for creation, overwrite and raw transforms. +pub(crate) fn ensure_no_foreign_declarations<'a>( + fields: impl IntoIterator>, +) -> Result<()> { + for field in fields { + if field.metadata().keys().any(|k| is_declaration_key(k)) { + return Err(Error::InvalidInput { + message: format!( + "field '{}' carries computed-column metadata; declare computed columns \ + with add_columns().computed()", + field.name() + ), + }); + } + } + Ok(()) +} + +/// True for field-metadata keys that belong to a computed-column declaration. +/// +/// A declaration is immutable through metadata edits: it is validated as a +/// whole at declare time, and rewriting any piece of it -- the flag, the +/// kind, the expression, the inputs -- would bypass that validation or move +/// a binding out from under a refresh. Drop the column and declare it again. +pub(crate) fn is_declaration_key(key: &str) -> bool { + key == COMPUTED_COLUMN_META_KEY || key.starts_with("computed_column.") +} + +/// Reject retyping a computed column itself. +/// +/// A cast keeps the stored expression while changing the type it must yield +/// -- and lance's cast rewrites the field without its metadata, so the +/// declaration silently stops being one. Dropping and redeclaring is the +/// coherent way to change a computed column's type. +pub(crate) fn ensure_not_retyped(schema: &ArrowSchema, paths: &[&str]) -> Result<()> { + for declaration in computed_columns(schema) { + for path in paths { + if declaration.name == root(path) { + return Err(Error::InvalidInput { + message: format!( + "column '{}' is computed; drop it and declare it again to change \ + its type", + declaration.name + ), + }); + } + } + } + Ok(()) +} + +/// The top-level column a possibly nested input path reads. +pub(crate) fn root(path: &str) -> &str { + path.split('.').next().unwrap_or(path) +} + +/// A declaration's expression bound to a schema, ready to evaluate. +pub(crate) struct BoundExpression { + /// The columns the expression names, as written; nested inputs keep + /// their dotted path. + pub inputs: Vec, + /// The top-level columns evaluation reads, in [`Self::read_schema`] + /// order. A nested input appears through its root. + pub roots: Vec, + /// The projected schema evaluation runs against. + pub read_schema: SchemaRef, + /// The compiled expression. + pub physical: Arc, + /// The type the expression yields. + pub data_type: DataType, +} + +/// Parse, resolve and compile `expression` against `schema`. +/// +/// Inputs come from the expression as written, before optimization: the +/// simplifier can fold a referenced column out entirely (`true OR x > 0`), +/// and the guard protecting the stored SQL has to see every column the text +/// names, not just the ones the simplified form still reads. +pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result { + let invalid = |message: String| Error::InvalidExpression { + column: column.to_string(), + message, + }; + + let planner = Planner::new(schema.clone()); + let parsed = planner + .parse_expr(expression) + .map_err(|e| invalid(e.to_string()))?; + + // A declaration is evaluated more than once -- staging and writing are + // separate passes, and a refresh years later replays the same text -- so + // a function that can answer differently each time has no coherent value + // to declare. + let mut volatile = None; + parsed + .apply(|expr| { + use datafusion_common::tree_node::TreeNodeRecursion; + if let datafusion_expr::Expr::ScalarFunction(function) = expr + && function.func.signature().volatility != datafusion_expr::Volatility::Immutable + { + volatile = Some(function.func.name().to_string()); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .map_err(|e| invalid(e.to_string()))?; + if let Some(function) = volatile { + return Err(invalid(format!( + "'{function}' is not deterministic; a computed column's expression must \ + yield the same value every time it is evaluated" + ))); + } + + let mut inputs = Planner::column_names_in_expr(&parsed); + inputs.sort(); + inputs.dedup(); + + // A nested input is recorded by its path but read through its root + // column; Schema::index_of resolves top-level names only. Resolved here + // rather than left to the planner so an unknown column names itself in + // the error instead of surfacing as a plan failure. + let mut indices = Vec::with_capacity(inputs.len()); + for input in &inputs { + let index = schema + .index_of(root(input)) + .map_err(|_| invalid(format!("unknown column '{input}'")))?; + if !indices.contains(&index) { + indices.push(index); + } + } + indices.sort_unstable(); + + // Physical expressions address columns by position, so the planner that + // compiles the expression has to be built on the projected schema + // evaluation will actually read. + let read_schema = Arc::new( + schema + .project(&indices) + .map_err(|e| invalid(e.to_string()))?, + ); + let roots = read_schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + + let optimized = planner + .optimize_expr(parsed) + .map_err(|e| invalid(e.to_string()))?; + let physical = Planner::new(read_schema.clone()) + .create_physical_expr(&optimized) + .map_err(|e| invalid(e.to_string()))?; + let data_type = physical + .data_type(read_schema.as_ref()) + .map_err(|e| invalid(e.to_string()))?; + + Ok(BoundExpression { + inputs, + roots, + read_schema, + physical, + data_type, + }) +} + +/// Resolve `(name, expression)` pairs against `schema` into fields carrying +/// their bindings. +/// +/// Everything that can be known statically is checked here rather than at +/// refresh time: that the expression parses, that every column it reads +/// exists, and that the target name is free. A declaration that survives this +/// is one a refresh can always act on. +pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result> { + if columns.is_empty() { + return Err(Error::InvalidInput { + message: "at least one computed column is required".into(), + }); + } + + let mut fields = Vec::with_capacity(columns.len()); + let mut declared: Vec<&str> = Vec::with_capacity(columns.len()); + + for (name, expression) in columns { + if schema.field_with_name(name).is_ok() || declared.contains(&name.as_str()) { + return Err(Error::ColumnAlreadyExists { name: name.clone() }); + } + + let bound = bind(schema.clone(), name, expression)?; + + // Declared columns start entirely null, so nullability is a property + // of the declaration rather than of what the expression yields. + fields.push( + ArrowField::new(name, bound.data_type, true) + .with_metadata(computed_column_metadata(expression, &bound.inputs)), + ); + declared.push(name); + } + + Ok(fields) +} + +/// Build the transform that declares `columns` against `schema`. +/// +/// An all-null column is how a binding with no values yet is carried into a +/// commit; that it is spelled `AllNulls` is a detail of the commit, not of the +/// column, which is why this is internal and +/// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) is the +/// public way in. +pub(crate) fn declare( + schema: SchemaRef, + columns: &[(String, String)], +) -> Result { + let fields = plan(schema, columns)?; + Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( + fields, + )))) +} + +/// Commit a declaration of a kind this version does not produce, the way a +/// newer lancedb would leave one behind. Bypasses admission, which exists to +/// stop exactly this through the public API. +#[cfg(test)] +pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &str) { + let field = ArrowField::new(name, DataType::Int32, true).with_metadata(HashMap::from([ + (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), + (KIND_META_KEY.to_string(), kind.to_string()), + (INPUTS_META_KEY.to_string(), r#"["x"]"#.to_string()), + ])); + super::schema_evolution::commit_add_columns( + table.as_native().unwrap(), + NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(vec![field]))), + None, + ) + .await + .unwrap(); +} + +#[cfg(test)] +mod tests { + use arrow_array::record_batch; + use arrow_schema::DataType; + use futures::TryStreamExt; + use lance::dataset::ColumnAlteration; + + use super::*; + use crate::connect; + use crate::query::{ExecutableQuery, QueryBase, Select}; + use crate::{Error, Table}; + + async fn table_with_ints(name: &str) -> Table { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!(("x", Int32, [1, 2, 3])).unwrap(); + conn.create_table(name, batch).execute().await.unwrap() + } + + /// Declare `columns` the way a caller would: plan the expressions, then + /// add them through the ordinary column API. + async fn add_computed(table: &Table, columns: &[(String, String)]) -> Result { + let mut builder = table.add_columns(); + for (name, expression) in columns { + builder = builder.computed(name, expression); + } + Ok(builder.execute().await?.version) + } + + async fn declared(table: &Table) -> Vec { + computed_columns(table.schema().await.unwrap().as_ref()) + } + + #[tokio::test] + async fn test_declare_infers_type_and_inputs() { + let table = table_with_ints("declare_infers").await; + let initial = table.version().await.unwrap(); + + let version = add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + assert!(version > initial); + + let schema = table.schema().await.unwrap(); + let field = schema.field_with_name("doubled").unwrap(); + assert_eq!(field.data_type(), &DataType::Int32); + assert!(field.is_nullable()); + + assert_eq!( + declared(&table).await, + vec![ComputedColumn { + name: "doubled".into(), + kind: ComputedColumnKind::Sql { + expression: "x * 2".into() + }, + inputs: vec!["x".into()], + }] + ); + } + + /// The binding reaches the schema only if `AllNulls` carries per-field + /// metadata through the commit. The whole representation rests on it. + #[tokio::test] + async fn test_all_nulls_preserves_field_metadata() { + let table = table_with_ints("metadata_survives").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let schema = table.schema().await.unwrap(); + let metadata = schema.field_with_name("doubled").unwrap().metadata(); + assert_eq!( + metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str), + Some("true") + ); + assert_eq!(metadata.get(KIND_META_KEY).map(String::as_str), Some("sql")); + assert_eq!( + metadata.get(EXPRESSION_META_KEY).map(String::as_str), + Some("x * 2") + ); + assert_eq!( + metadata.get(INPUTS_META_KEY).map(String::as_str), + Some(r#"["x"]"#) + ); + } + + #[tokio::test] + async fn test_declared_column_is_all_null() { + let table = table_with_ints("declare_is_null").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let batches = table + .query() + .select(Select::columns(&["doubled"])) + .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["doubled"].null_count(), batch.num_rows()); + } + } + + #[tokio::test] + async fn test_unknown_column_fails_at_declare_time() { + let table = table_with_ints("unknown_input").await; + let err = add_computed(&table, &[("bad".into(), "missing + 1".into())]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "bad")); + + let schema = table.schema().await.unwrap(); + assert!(schema.field_with_name("bad").is_err()); + } + + #[tokio::test] + async fn test_unparsable_expression_fails_at_declare_time() { + let table = table_with_ints("bad_syntax").await; + let err = add_computed(&table, &[("bad".into(), "x *".into())]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "bad")); + assert!( + table + .schema() + .await + .unwrap() + .field_with_name("bad") + .is_err() + ); + } + + /// A user-defined function is an expression like any other; only its + /// resolution is missing. When a registry-aware planner exists this + /// becomes a supported declaration rather than a new API. + #[tokio::test] + async fn test_unregistered_function_is_rejected_for_now() { + let table = table_with_ints("udf_not_yet").await; + let err = add_computed(&table, &[("vec".into(), "embed(x)".into())]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "vec")); + assert!( + table + .schema() + .await + .unwrap() + .field_with_name("vec") + .is_err() + ); + } + + #[tokio::test] + async fn test_existing_column_name_is_rejected() { + let table = table_with_ints("name_taken").await; + let err = add_computed(&table, &[("x".into(), "x * 2".into())]) + .await + .unwrap_err(); + assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "x")); + assert!(declared(&table).await.is_empty()); + } + + #[tokio::test] + async fn test_constant_expression_needs_no_inputs() { + let table = table_with_ints("constant").await; + add_computed(&table, &[("answer".into(), "42".into())]) + .await + .unwrap(); + + let declared = declared(&table).await; + assert_eq!(declared.len(), 1); + assert!(declared[0].inputs.is_empty()); + } + + #[tokio::test] + async fn test_multiple_columns_in_one_commit() { + let table = table_with_ints("multi").await; + let initial = table.version().await.unwrap(); + + add_computed( + &table, + &[ + ("plus".into(), "x + 1".into()), + ("squared".into(), "x * x".into()), + ], + ) + .await + .unwrap(); + + assert_eq!(table.version().await.unwrap(), initial + 1); + let declared = declared(&table).await; + assert_eq!(declared.len(), 2); + assert_eq!(declared[0].name, "plus"); + assert_eq!(declared[1].name, "squared"); + } + + #[tokio::test] + async fn test_duplicate_declaration_in_one_call_is_rejected() { + let table = table_with_ints("dupe").await; + let err = add_computed( + &table, + &[ + ("dup".into(), "x + 1".into()), + ("dup".into(), "x + 2".into()), + ], + ) + .await + .unwrap_err(); + assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "dup")); + assert!(declared(&table).await.is_empty()); + } + + /// A column added by an ordinary transform is materialized, not bound, so + /// it carries no declaration to report. + #[tokio::test] + async fn test_ordinary_columns_are_not_reported_as_computed() { + let table = table_with_ints("plain").await; + assert!(declared(&table).await.is_empty()); + + table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "eager".into(), + "x * 2".into(), + )])) + .execute() + .await + .unwrap(); + assert!(declared(&table).await.is_empty()); + } + + /// Built-in functions type the column the same way an operator does. + #[tokio::test] + async fn test_builtin_function_inference() { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!(("name", Utf8, ["ada", "grace"]), ("n", Int32, [-1, 2])).unwrap(); + let table = conn + .create_table("builtins", batch) + .execute() + .await + .unwrap(); + + add_computed( + &table, + &[ + ("shout".into(), "upper(name)".into()), + ("width".into(), "length(name)".into()), + ("magnitude".into(), "abs(n)".into()), + ], + ) + .await + .unwrap(); + + let schema = table.schema().await.unwrap(); + assert_eq!( + schema.field_with_name("shout").unwrap().data_type(), + &DataType::Utf8 + ); + assert_eq!( + schema.field_with_name("magnitude").unwrap().data_type(), + &DataType::Int32 + ); + // length() returns a width-dependent integer type; assert it is one + // rather than pinning which. + assert!( + schema + .field_with_name("width") + .unwrap() + .data_type() + .is_integer() + ); + + let declared = declared(&table).await; + assert_eq!(declared.len(), 3); + assert_eq!(declared[0].inputs, vec!["name".to_string()]); + assert_eq!(declared[2].inputs, vec!["n".to_string()]); + } + + /// The reason the kind is tagged: a declaration written by a newer version + /// has to read back as a computed column this one cannot evaluate, not as + /// an ordinary column. Reported as absent it would be refreshable by + /// nothing and redeclarable over, silently. + #[tokio::test] + async fn test_unrecognized_kind_is_reported_rather_than_hidden() { + let table = table_with_ints("foreign_kind").await; + super::add_foreign_kind(&table, "embedding", "udf").await; + + assert_eq!( + declared(&table).await, + vec![ComputedColumn { + name: "embedding".into(), + kind: ComputedColumnKind::Unrecognized { kind: "udf".into() }, + inputs: vec!["x".into()], + }] + ); + + let err = add_computed(&table, &[("embedding".into(), "x * 2".into())]) + .await + .unwrap_err(); + assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "embedding")); + } + + /// A kind is what makes a declaration readable at all, so the flag alone + /// is half-formed in the same way a missing expression is. + #[test] + fn test_flag_without_a_kind_is_not_a_declaration() { + let field = + ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([( + COMPUTED_COLUMN_META_KEY.to_string(), + "true".to_string(), + )])); + assert_eq!(computed_column_from_field(&field), None); + } + + /// A SQL declaration is its expression; without one there is nothing to + /// refresh from. + #[test] + fn test_sql_kind_without_an_expression_is_not_a_declaration() { + let field = ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([ + (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), + (KIND_META_KEY.to_string(), SQL_KIND.to_string()), + ])); + assert_eq!(computed_column_from_field(&field), None); + } + + #[tokio::test] + async fn test_inputs_are_deduplicated_and_sorted() { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!(("b", Int32, [1, 2]), ("a", Int32, [3, 4])).unwrap(); + let table = conn.create_table("dedupe", batch).execute().await.unwrap(); + + add_computed(&table, &[("total".into(), "b + a + b".into())]) + .await + .unwrap(); + + assert_eq!( + declared(&table).await[0].inputs, + vec!["a".to_string(), "b".to_string()] + ); + } + + #[tokio::test] + async fn test_dropping_an_input_is_refused() { + let table = table_with_ints("drop_input").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let err = table.drop_columns(&["x"]).await.unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("doubled")), + "{err:?}" + ); + } + + #[tokio::test] + async fn test_renaming_an_input_is_refused() { + let table = table_with_ints("rename_input").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let err = table + .alter_columns(&[ColumnAlteration::new("x".into()).rename("y".into())]) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("doubled")), + "{err:?}" + ); + } + + /// Nothing resolves against nullability, so it is not a rebinding. + #[tokio::test] + async fn test_altering_an_input_nullability_is_allowed() { + let table = table_with_ints("nullable_input").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + table + .alter_columns(&[ColumnAlteration::new("x".into()).set_nullable(true)]) + .await + .unwrap(); + } + + /// The gate's reproducer: a volatile function evaluates differently in + /// the counting and writing passes, so the declared value is incoherent. + /// Refused at declare time. + #[tokio::test] + async fn test_a_volatile_expression_is_refused() { + let table = table_with_ints("volatile_expr").await; + let err = add_computed(&table, &[("maybe".into(), "random() < 0.5".into())]) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidExpression { message, .. } + if message.contains("random") && message.contains("deterministic")), + "{err:?}" + ); + } + + /// The gate's reproducer: the simplifier folds `true OR x > 0` to a + /// constant, but the stored SQL still names `x`, so the recorded inputs + /// must too -- otherwise dropping `x` is allowed and refresh breaks. + #[tokio::test] + async fn test_inputs_survive_expression_optimization() { + let table = table_with_ints("optimized_inputs").await; + add_computed(&table, &[("flag".into(), "true OR x > 0".into())]) + .await + .unwrap(); + + assert_eq!(declared(&table).await[0].inputs, vec!["x".to_string()]); + let err = table.drop_columns(&["x"]).await.unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("flag")), + "{err:?}" + ); + } + + /// The gate's reproducer: casting a computed column rewrites the field + /// without its metadata, silently destroying the declaration. + #[tokio::test] + async fn test_retyping_the_computed_column_is_refused() { + use arrow_schema::DataType as ArrowDataType; + + let table = table_with_ints("retype_computed").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let err = table + .alter_columns(&[ColumnAlteration::new("doubled".into()).cast_to(ArrowDataType::Int64)]) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("computed")), + "{err:?}" + ); + + // The declaration survives the refused change. + table.refresh_column("doubled").await.unwrap(); + } + + /// A declaration cannot be edited, fabricated or erased through field + /// metadata: it is validated as a whole at declare time. + #[tokio::test] + async fn test_declaration_metadata_is_immutable() { + use crate::table::FieldMetadataUpdate; + + let table = table_with_ints("metadata_tamper").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + // Moving the binding. + let err = table + .update_field_metadata(&[ + FieldMetadataUpdate::new("doubled").set(EXPRESSION_META_KEY, "x * 3") + ]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); + + // Fabricating a declaration on a plain column. + let err = table + .update_field_metadata(&[FieldMetadataUpdate::new("x") + .set(COMPUTED_COLUMN_META_KEY, "true") + .set(KIND_META_KEY, SQL_KIND) + .set(EXPRESSION_META_KEY, "x")]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); + + // Erasing the declaration wholesale. + let err = table + .update_field_metadata(&[FieldMetadataUpdate::new("doubled") + .set("note", "hi") + .replace()]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}"); + + // Ordinary metadata on a computed column still merges. + table + .update_field_metadata(&[FieldMetadataUpdate::new("doubled").set("note", "hi")]) + .await + .unwrap(); + table.refresh_column("doubled").await.unwrap(); + } + + /// The gate's reproducer: only refresh materializes a declared column; + /// a direct write would store an arbitrary durable value. + #[tokio::test] + async fn test_a_computed_column_cannot_be_written_directly() { + let table = table_with_ints("direct_write").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let batch = record_batch!(("x", Int32, [4]), ("doubled", Int32, [999])).unwrap(); + let err = table.add(batch.clone()).execute().await.unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("refresh")), + "{err:?}" + ); + + let err = table + .update() + .column("doubled", "999") + .execute() + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + + let mut merge = table.merge_insert(&["x"]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all(); + let err = merge + .execute(Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(batch.clone())], + batch.schema(), + ))) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + + // The append that omits the column still works. + let plain = record_batch!(("x", Int32, [4])).unwrap(); + table.add(plain).execute().await.unwrap(); + } + + /// The gate's reproducer: the reciprocal of the declare-under-spec check. + #[tokio::test] + async fn test_installing_an_lsm_spec_over_computed_columns_is_refused() { + use crate::table::LsmWriteSpec; + + let tmp_dir = tempfile::tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "x", + DataType::Int32, + false, + )])); + let batch = arrow_array::RecordBatch::try_new( + schema, + vec![Arc::new(arrow_array::Int32Array::from(vec![1, 2])) as _], + ) + .unwrap(); + let table = conn + .create_table("lsm_after", batch) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["x"]).await.unwrap(); + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let err = table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::NotSupported { message } if message.contains("computed")), + "{err:?}" + ); + assert!(table.get_lsm_write_spec().await.unwrap().is_none()); + } + + /// The gate's reproducer: declaration metadata is admitted only through + /// the validated declare path, never smuggled through a raw transform. + #[tokio::test] + async fn test_forged_declaration_metadata_is_rejected() { + let table = table_with_ints("forged_metadata").await; + let field = + ArrowField::new("doubled", DataType::Int32, 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(), "x * 2".to_string()), + (INPUTS_META_KEY.to_string(), "[]".to_string()), + ])); + let err = table + .add_columns() + .transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( + vec![field], + )))) + .execute() + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("computed()")), + "{err:?}" + ); + assert!(declared(&table).await.is_empty()); + } + + /// The gate's reproducer: SQL INSERT is a write path too. + #[tokio::test] + async fn test_sql_insert_cannot_write_a_computed_column() { + use datafusion::prelude::SessionContext; + + let table = table_with_ints("sql_insert").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let ctx = SessionContext::new(); + let provider = + crate::table::datafusion::BaseTableAdapter::try_new(table.base_table().clone()) + .await + .unwrap(); + ctx.register_table("t", Arc::new(provider)).unwrap(); + + let result = async { + ctx.sql("INSERT INTO t (x, doubled) VALUES (4, 999)") + .await? + .collect() + .await + } + .await; + let err = result.unwrap_err().to_string(); + assert!(err.contains("refresh"), "{err}"); + } + + /// The gate's reproducer: an overwrite must not smuggle in a filled + /// declaration. + #[tokio::test] + async fn test_overwrite_cannot_inject_a_declaration() { + use crate::table::AddDataMode; + + let table = table_with_ints("overwrite_inject").await; + let field = + ArrowField::new("doubled", DataType::Int32, 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(), "x * 2".to_string()), + ])); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("x", DataType::Int32, true), + field, + ])); + let batch = arrow_array::RecordBatch::try_new( + schema, + vec![ + Arc::new(arrow_array::Int32Array::from(vec![1])) as _, + Arc::new(arrow_array::Int32Array::from(vec![999])) as _, + ], + ) + .unwrap(); + + let err = table + .add(batch) + .mode(AddDataMode::Overwrite) + .execute() + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("declare")), + "{err:?}" + ); + } + + #[tokio::test] + async fn test_create_table_cannot_inject_a_declaration() { + let conn = connect("memory://").execute().await.unwrap(); + let field = + ArrowField::new("doubled", DataType::Int32, 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(), "x * 2".to_string()), + ])); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("x", DataType::Int32, true), + field, + ])); + let batch = arrow_array::RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(arrow_array::Int32Array::from(vec![1])) as _, + Arc::new(arrow_array::Int32Array::from(vec![999])) as _, + ], + ) + .unwrap(); + let err = conn + .create_table("forged_create", batch) + .execute() + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("computed()")), + "{err:?}" + ); + } + + #[tokio::test] + async fn test_sql_insert_omitting_computed_is_allowed() { + use datafusion::prelude::SessionContext; + + let table = table_with_ints("sql_insert_omitted").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let ctx = SessionContext::new(); + let provider = + crate::table::datafusion::BaseTableAdapter::try_new(table.base_table().clone()) + .await + .unwrap(); + ctx.register_table("t", Arc::new(provider)).unwrap(); + ctx.sql("INSERT INTO t (x) VALUES (4)") + .await + .unwrap() + .collect() + .await + .unwrap(); + + table.checkout_latest().await.unwrap(); + assert_eq!(table.count_rows(None).await.unwrap(), 4); + } + + #[tokio::test] + async fn test_a_nested_computed_field_cannot_be_renamed() { + let table = table_with_ints("computed_struct_rename").await; + add_computed(&table, &[("payload".into(), "named_struct('a', x)".into())]) + .await + .unwrap(); + + let err = table + .alter_columns(&[ColumnAlteration::new("payload.a".into()).rename("b".into())]) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("payload")), + "{err:?}" + ); + } + + /// Stale handles must not commit the computed/LSM state in either order. + #[tokio::test] + async fn test_stale_handles_cannot_mix_computed_and_lsm() { + use crate::table::LsmWriteSpec; + + let tmp_dir = tempfile::tempdir().unwrap(); + let uri = tmp_dir.path().to_str().unwrap(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "x", + DataType::Int32, + false, + )])); + let batch = arrow_array::RecordBatch::try_new( + schema.clone(), + vec![Arc::new(arrow_array::Int32Array::from(vec![1])) as _], + ) + .unwrap(); + let conn = connect(uri).execute().await.unwrap(); + let table = conn.create_table("mix", batch).execute().await.unwrap(); + table.set_unenforced_primary_key(["x"]).await.unwrap(); + let stale = conn.open_table("mix").execute().await.unwrap(); + + // Declare on one handle; the stale handle must not install a spec. + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + let err = stale + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "install won"); + + // Reverse order on fresh tables. + let batch = arrow_array::RecordBatch::try_new( + schema, + vec![Arc::new(arrow_array::Int32Array::from(vec![1])) as _], + ) + .unwrap(); + let table = conn.create_table("mix2", batch).execute().await.unwrap(); + table.set_unenforced_primary_key(["x"]).await.unwrap(); + let stale = conn.open_table("mix2").execute().await.unwrap(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + let err = add_computed(&stale, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "declare won"); + } + + /// The gate's reproducer: after catch-up activation, an LSM write, and + /// unset, retained SSTable rows survive without a live spec. The catch-up + /// flag is the durable marker; declaration refuses on it. + #[tokio::test] + async fn test_unset_with_retained_lsm_rows_cannot_admit_a_declaration() { + use crate::table::LsmWriteSpec; + use arrow_array::{Int64Array, RecordBatchIterator}; + + let tmp_dir = tempfile::tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int64, false), + ArrowField::new("value", DataType::Int64, false), + ])); + let batch = arrow_array::RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2])) as _, + Arc::new(Int64Array::from(vec![10, 20])) as _, + ], + ) + .unwrap(); + let table = conn + .create_table("t", batch.clone()) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["id"]).await.unwrap(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + table.require_mem_wal_index_catchup().await.unwrap(); + + 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)], schema))) + .await + .unwrap(); + table.unset_lsm_write_spec().await.unwrap(); + + let err = add_computed(&table, &[("doubled".into(), "value * 2".into())]) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::NotSupported { message } if message.contains("LSM")), + "{err:?}" + ); + } + + /// A declaration does not read itself, so it travels with its binding. + #[tokio::test] + async fn test_dropping_the_computed_column_is_allowed() { + let table = table_with_ints("drop_computed").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + table.drop_columns(&["doubled"]).await.unwrap(); + assert!(declared(&table).await.is_empty()); + } +} diff --git a/rust/lancedb/src/table/datafusion/insert.rs b/rust/lancedb/src/table/datafusion/insert.rs index e176c228b..b9bd2396e 100644 --- a/rust/lancedb/src/table/datafusion/insert.rs +++ b/rust/lancedb/src/table/datafusion/insert.rs @@ -17,7 +17,7 @@ use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, }; -use futures::TryStreamExt; +use futures::StreamExt; use lance::Dataset; use lance::dataset::transaction::{Operation, Transaction}; use lance::dataset::{CommitBuilder, InsertBuilder, WriteParams, WriteProgressFn}; @@ -194,12 +194,23 @@ impl ExecutionPlan for InsertExec { let output_bytes = MetricBuilder::new(&self.metrics).output_bytes(partition); let input_schema = input_stream.schema(); + let declared: Vec = crate::table::computed_columns::computed_columns( + &arrow_schema::Schema::from(self.dataset.schema()), + ) + .into_iter() + .map(|declaration| declaration.name) + .collect(); let input_stream: SendableRecordBatchStream = Box::pin(InstrumentedRecordBatchStreamAdapter::new( input_schema, - input_stream.map_ok(move |batch| { + input_stream.map(move |batch| { + let batch = batch?; + crate::table::computed_columns::ensure_batch_writes_no_computed_values( + &declared, &batch, + ) + .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?; output_bytes.add(batch.get_array_memory_size()); - batch + Ok(batch) }), partition, &self.metrics, diff --git a/rust/lancedb/src/table/merge.rs b/rust/lancedb/src/table/merge.rs index 82a1d1473..3a5b6882d 100644 --- a/rust/lancedb/src/table/merge.rs +++ b/rust/lancedb/src/table/merge.rs @@ -1161,7 +1161,7 @@ mod lsm_tests { .unwrap(); let fts_index = table.list_indices().await.unwrap()[0].name.clone(); table - .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes([fts_index])) + .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes(vec![fts_index])) .await .unwrap(); @@ -1254,7 +1254,7 @@ mod lsm_tests { .unwrap(); let vec_index = table.list_indices().await.unwrap()[0].name.clone(); table - .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes([vec_index])) + .set_lsm_write_spec(LsmWriteSpec::unsharded().with_maintained_indexes(vec![vec_index])) .await .unwrap(); diff --git a/rust/lancedb/src/table/merge/lsm.rs b/rust/lancedb/src/table/merge/lsm.rs index 0eb7c0231..5751cd916 100644 --- a/rust/lancedb/src/table/merge/lsm.rs +++ b/rust/lancedb/src/table/merge/lsm.rs @@ -29,6 +29,7 @@ use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef}; use lance::Dataset; use lance::dataset::mem_wal::{ DatasetMemWalExt, ShardWriter, ShardWriterConfig, evaluate_sharding_spec, + validate_maintained_indexes, }; use lance::index::DatasetIndexExt; use lance_core::datatypes::Schema as LanceSchema; @@ -37,8 +38,9 @@ use tokio::sync::RwLock; use uuid::Uuid; use crate::error::{Error, Result}; +use crate::index::IndexConfig; use crate::table::merge::{MergeInsertBuilder, MergeResult}; -use crate::table::{LsmWriteSpec, NativeTable}; +use crate::table::{BaseTable, LsmWriteSpec, NativeTable}; /// Spec id of the sole sharding spec installed by [`set_lsm_write_spec`]. /// Must match Lance's `InitializeMemWalBuilder` (`SHARDING_SPEC_ID`). @@ -80,32 +82,53 @@ pub(crate) async fn set_lsm_write_spec(table: &NativeTable, spec: LsmWriteSpec) } } + // 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? + }; + + table.checkout_latest().await?; let mut dataset = (*table.dataset.get().await?).clone(); + let schema = arrow_schema::Schema::from(dataset.schema()); + if !crate::table::computed_columns::computed_columns(&schema).is_empty() { + return Err(Error::NotSupported { + message: "an LSM write spec cannot be installed on a table with computed \ + columns: rows in un-compacted tiers are invisible to refresh" + .into(), + }); + } let mut builder = dataset.initialize_mem_wal(); - let (maintained_indexes, writer_config_defaults) = match spec { + let writer_config_defaults = match spec { LsmWriteSpec::Bucket { column, num_buckets, - maintained_indexes, writer_config_defaults, + .. } => { builder = builder.bucket_sharding(column, num_buckets); - (maintained_indexes, writer_config_defaults) + writer_config_defaults } LsmWriteSpec::Identity { column, - maintained_indexes, writer_config_defaults, + .. } => { builder = builder.identity_sharding(column); - (maintained_indexes, writer_config_defaults) + writer_config_defaults } LsmWriteSpec::Unsharded { - maintained_indexes, writer_config_defaults, + .. } => { builder = builder.unsharded(); - (maintained_indexes, writer_config_defaults) + writer_config_defaults } }; builder = builder.maintained_indexes(maintained_indexes); @@ -117,6 +140,88 @@ 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() { + return "no indexes".to_string(); + } + let mut names: Vec<&str> = indices.iter().map(|index| index.name.as_str()).collect(); + names.sort_unstable(); + format!("[{}]", names.join(", ")) +} + +// ============================================================================= +// require_mem_wal_index_catchup +// ============================================================================= + +/// Switch this table to required index catch-up, one way. +/// +/// Deliberately **not** part of installing the write spec. Until something can +/// actually repair coverage, a table carrying the bit reports every index as +/// not known to hold the compacted rows, so its SSTables are retained +/// indefinitely -- and the WAL pod trims on the legacy rule meanwhile, leaving +/// readers pointed at files that are gone. Turn this on only once remote +/// maintenance owns the merge and the repair for the table. +/// +/// Lance refuses the activation if the table already records SSTable +/// compaction progress: those numbers predate this protocol and cannot be +/// validated, so such a table must be drained rather than activated. +#[allow(clippy::redundant_pub_crate)] +pub(crate) async fn require_mem_wal_index_catchup(table: &NativeTable) -> Result<()> { + table.dataset.ensure_mutable()?; + let mut dataset = (*table.dataset.get().await?).clone(); + if dataset.mem_wal_index_details().await?.is_none() { + return Err(Error::InvalidInput { + message: "require_mem_wal_index_catchup: no LSM write spec is set on this table".into(), + }); + } + dataset.require_mem_wal_index_catchup().await?; + table.dataset.update(dataset); + Ok(()) +} + // ============================================================================= // unset_lsm_write_spec // ============================================================================= diff --git a/rust/lancedb/src/table/query/lsm.rs b/rust/lancedb/src/table/query/lsm.rs index 074d13476..6155ec095 100644 --- a/rust/lancedb/src/table/query/lsm.rs +++ b/rust/lancedb/src/table/query/lsm.rs @@ -36,6 +36,7 @@ use lance::dataset::mem_wal::{ DatasetMemWalExt, LsmScanner, ShardManifestStore, ShardSnapshot, ShardWriterConfig, }; use lance_index::mem_wal::{MemWalIndexDetails, ShardManifest}; +use lance_table::feature_flags::FLAG_MEM_WAL_INDEX_CATCHUP; use uuid::Uuid; use super::NativeTable; @@ -84,9 +85,8 @@ pub(super) async fn create_lsm_plan( let pk_columns = pk_columns(&ds_ref)?; // The base index an indexed arm relies on may lag compaction; resolve it so the // snapshot retains SSTables the index has not yet caught up to. - let arm_index = arm_maintained_index_name(&ds_ref, &query, &details).await?; - let (snapshots, in_memory) = - build_read_context(table, &ds_ref, &details, arm_index.as_deref()).await?; + let arm_indexes = arm_maintained_index_names(&ds_ref, &query, &details).await?; + let (snapshots, in_memory) = build_read_context(table, &ds_ref, &details, &arm_indexes).await?; let limit = query.base.limit; let offset = query.base.offset; @@ -232,28 +232,44 @@ fn pk_columns(dataset: &Dataset) -> Result> { Ok(pk) } -/// Per-shard SSTable exclusion watermark: the generation at or below which SSTables -/// are safe to drop for this arm. A generation is droppable only once it is -/// compacted into the base table AND covered by `index_name`'s catch-up (for an -/// indexed arm); a plain scan (`index_name == None`) uses the compaction watermark -/// alone. Capping at the index catch-up keeps rows the base index has not yet -/// indexed visible through their SSTable. First occurrence per shard mirrors Lance's -/// `compacted_generation_for_shard`. +/// Per-shard SSTable exclusion watermark: the generation at or below which +/// SSTables are safe to drop for this query. +/// +/// A generation is droppable only once it is compacted into the base table AND +/// covered by the catch-up of every index the query relies on, so the watermark +/// is the minimum across `index_names`. Gating on fewer than all of them would +/// drop SSTables holding rows an uncounted index has not yet indexed, and that +/// arm would silently return fewer rows. +/// +/// See [`arm_maintained_index_names`] for which indexes are collected today: a +/// vector search with a scalar prefilter is not yet among them. +/// +/// An empty `index_names` (a plain scan) uses the compaction watermark alone. +/// First occurrence per shard mirrors Lance's `compacted_generation_for_shard`. fn exclusion_watermarks( details: &MemWalIndexDetails, - index_name: Option<&str>, + index_names: &[String], + catchup_required: bool, ) -> HashMap { let mut exclude: HashMap = HashMap::new(); for entry in &details.compacted_sstables { let mut watermark = entry.generation; - if let Some(name) = index_name - && let Some(caught_up) = details + for name in index_names { + match details .index_catchup .iter() - .find(|icp| icp.index_name == name) + .find(|icp| icp.index_name == *name) .and_then(|icp| icp.caught_up_generation_for_shard(&entry.shard_id)) - { - watermark = watermark.min(caught_up); + { + Some(caught_up) => watermark = watermark.min(caught_up), + // No entry. On a table that requires catch-up this means the + // index is *not* known to hold these rows, and the base arm is + // index-only -- so every generation stays readable from its + // SSTable. Without the bit the field is not maintained at all, + // and absence carries no information. + None if catchup_required => watermark = 0, + None => {} + } } exclude.entry(entry.shard_id).or_insert(watermark); } @@ -267,13 +283,26 @@ fn exclusion_watermarks( /// with a live cached `ShardWriter` (this session's in-flight writes) the /// writer's authoritative in-memory manifest and memtables override the /// on-disk view so a read sees data not yet flushed. +/// Whether this table reads a missing `index_catchup` entry as "not caught up". +/// +/// Both words must be set. A reader honouring the bit while a writer does not +/// would retain SSTables the writer had already trimmed, and the reverse would +/// serve rows from files the writer still expects to be excluded -- so a +/// half-set manifest is treated as legacy, which is the conservative side. +fn requires_index_catchup(dataset: &Dataset) -> bool { + let manifest = dataset.manifest(); + manifest.reader_feature_flags & FLAG_MEM_WAL_INDEX_CATCHUP != 0 + && manifest.writer_feature_flags & FLAG_MEM_WAL_INDEX_CATCHUP != 0 +} + async fn build_read_context( table: &NativeTable, dataset: &Dataset, details: &MemWalIndexDetails, - index_name: Option<&str>, + index_names: &[String], ) -> Result<(Vec, HashMap)> { - let exclude = exclusion_watermarks(details, index_name); + let catchup_required = requires_index_catchup(dataset); + let exclude = exclusion_watermarks(details, index_names, catchup_required); let shard_ids = dataset.list_mem_wal_latest_shard_ids().await?; // Use the dataset's own object store (not `ObjectStore::from_uri`, which @@ -487,19 +516,33 @@ async fn index_maintained( })) } -/// The maintained base index the query's arm relies on (vector index for ANN, FTS -/// index for full-text), used to gate SSTable compaction exclusion by index catch-up. -/// `None` for a plain scan or when no maintained index covers the searched column. -async fn arm_maintained_index_name( +/// Every maintained base index this query relies on, used to gate SSTable +/// exclusion by index catch-up. +/// +/// Returns a list because the watermark must be the lowest across every index a +/// query relies on. Today it never holds more than one: `reject_unsupported` +/// refuses hybrid search, so the vector and full-text arms are mutually +/// exclusive. +/// +/// The case that is genuinely multi-index -- a vector search with a scalar or +/// bitmap prefilter -- is **not collected yet**. Identifying those needs the +/// planner's chosen indexes, not the columns the filter names, and no Lance API +/// exposes them. Until it does, such a query is gated on its vector index alone. +/// +/// Empty for a plain scan, or when no maintained index covers the searched +/// column. +async fn arm_maintained_index_names( dataset: &Dataset, query: &VectorQueryRequest, details: &MemWalIndexDetails, -) -> Result> { +) -> Result> { use lance::index::DatasetIndexExt; - // Resolve the arm's searched column, the index-detail type it relies on, and a + + // Each arm's searched column, the index-detail type it relies on, and a // label for diagnostics — catch-up is taken from the vector/FTS index // specifically, not a BTree on the same column. - let (column, type_url_suffix, arm) = if !query.query_vector.is_empty() { + let mut arms: Vec<(String, &str, &str)> = Vec::new(); + if !query.query_vector.is_empty() { let arrow_schema = ArrowSchema::from(dataset.schema()); let column = match &query.column { Some(column) => column.clone(), @@ -508,31 +551,43 @@ async fn arm_maintained_index_name( default_vector_column(&arrow_schema, dim)? } }; - (column, "VectorIndexDetails", "vector") - } else if let Some(fts) = &query.base.full_text_search { - match fts.columns().into_iter().next() { - Some(column) => (column, "InvertedIndexDetails", "full-text"), - None => return Ok(None), - } - } else { - return Ok(None); - }; - let Some(field) = dataset.schema().field(&column) else { - return Ok(None); - }; + arms.push((column, "VectorIndexDetails", "vector")); + } + if let Some(fts) = &query.base.full_text_search + && let Some(column) = fts.columns().into_iter().next() + { + arms.push((column, "InvertedIndexDetails", "full-text")); + } + if arms.is_empty() { + return Ok(Vec::new()); + } + let indices = dataset.load_indices().await?; - let segment_names: Vec = indices - .iter() - .filter(|idx| { - idx.fields.contains(&field.id) - && idx - .index_details - .as_ref() - .is_some_and(|d| d.type_url.ends_with(type_url_suffix)) - }) - .map(|idx| idx.name.clone()) - .collect(); - resolve_single_index(segment_names, &details.maintained_indexes, arm, &column) + let mut names = Vec::with_capacity(arms.len()); + for (column, type_url_suffix, arm) in arms { + let Some(field) = dataset.schema().field(&column) else { + continue; + }; + let segment_names: Vec = indices + .iter() + .filter(|idx| { + idx.fields.contains(&field.id) + && idx + .index_details + .as_ref() + .is_some_and(|d| d.type_url.ends_with(type_url_suffix)) + }) + .map(|idx| idx.name.clone()) + .collect(); + if let Some(name) = + resolve_single_index(segment_names, &details.maintained_indexes, arm, &column)? + { + names.push(name); + } + } + names.sort(); + names.dedup(); + Ok(names) } /// Resolve the single logical index from the names of its matching physical @@ -734,22 +789,117 @@ mod tests { }; // Plain scan: drop every compacted generation (through 5). - assert_eq!(exclusion_watermarks(&details, None).get(&shard), Some(&5)); + assert_eq!( + exclusion_watermarks(&details, &[], false).get(&shard), + Some(&5) + ); // FTS arm with a lagging index: exclusion is capped at the index catch-up // (2), so SSTable generations 3..=5 are retained until the index covers // them — otherwise those documents would silently vanish from FTS results. assert_eq!( - exclusion_watermarks(&details, Some("fts_idx")).get(&shard), + exclusion_watermarks(&details, &["fts_idx".to_string()], false).get(&shard), Some(&2) ); // A caught-up index — or one untracked in index_catchup — falls back to the // compaction watermark. assert_eq!( - exclusion_watermarks(&details, Some("caught_up_idx")).get(&shard), + exclusion_watermarks(&details, &["caught_up_idx".to_string()], false).get(&shard), Some(&5) ); + + // The same missing entry, once the table requires catch-up: absence now + // means "not known to hold these rows", so nothing may be excluded and + // every generation stays readable from its SSTable. This is the whole + // point of the protocol -- an indexed query against a table whose index + // has not caught up must not silently lose rows. + assert_eq!( + exclusion_watermarks(&details, &["untracked_idx".to_string()], true).get(&shard), + Some(&0) + ); + + // A tracked index is unaffected by the mode: the recorded position is + // information either way, and it still caps the exclusion. + assert_eq!( + exclusion_watermarks(&details, &["fts_idx".to_string()], true).get(&shard), + Some(&2) + ); + + // One missing entry is enough to hold everything back, even alongside an + // index that has caught up. + let mixed = vec!["fts_idx".to_string(), "untracked_idx".to_string()]; + assert_eq!( + exclusion_watermarks(&details, &mixed, true).get(&shard), + Some(&0) + ); + } + + /// A hybrid search reads a vector and a full-text index, and either may lag. + /// Retaining to the lower of the two is what keeps both arms complete; + /// gating on one alone would drop SSTables the other has not indexed. + #[test] + fn exclusion_watermark_takes_the_minimum_across_every_index_used() { + let shard = Uuid::from_u128(1); + let details = MemWalIndexDetails { + compacted_sstables: vec![CompactedSsTable::new(shard, 9)], + index_catchup: vec![ + IndexCatchupProgress::new( + "vec_idx".to_string(), + vec![CompactedSsTable::new(shard, 7)], + ), + IndexCatchupProgress::new( + "fts_idx".to_string(), + vec![CompactedSsTable::new(shard, 4)], + ), + ], + maintained_indexes: vec!["vec_idx".to_string(), "fts_idx".to_string()], + ..Default::default() + }; + + // Each index alone stops at its own catch-up. + assert_eq!( + exclusion_watermarks(&details, &["vec_idx".to_string()], false).get(&shard), + Some(&7) + ); + assert_eq!( + exclusion_watermarks(&details, &["fts_idx".to_string()], false).get(&shard), + Some(&4) + ); + + // Used together, the lower one governs regardless of order. + let both = ["vec_idx".to_string(), "fts_idx".to_string()]; + assert_eq!( + exclusion_watermarks(&details, &both, false).get(&shard), + Some(&4) + ); + let reversed = ["fts_idx".to_string(), "vec_idx".to_string()]; + assert_eq!( + exclusion_watermarks(&details, &reversed, false).get(&shard), + Some(&4) + ); + } + + /// An index with no catch-up entry contributes no cap today, so a lagging + /// sibling must still govern rather than being widened by the untracked one. + #[test] + fn an_untracked_index_does_not_widen_a_lagging_sibling() { + let shard = Uuid::from_u128(1); + let details = MemWalIndexDetails { + compacted_sstables: vec![CompactedSsTable::new(shard, 9)], + index_catchup: vec![IndexCatchupProgress::new( + "fts_idx".to_string(), + vec![CompactedSsTable::new(shard, 4)], + )], + maintained_indexes: vec!["fts_idx".to_string(), "untracked_idx".to_string()], + ..Default::default() + }; + + let both = ["fts_idx".to_string(), "untracked_idx".to_string()]; + assert_eq!( + exclusion_watermarks(&details, &both, false).get(&shard), + Some(&4) + ); } #[test] diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs new file mode 100644 index 000000000..ab4f8e452 --- /dev/null +++ b/rust/lancedb/src/table/refresh.rs @@ -0,0 +1,708 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Filling computed columns. +//! +//! A row without a value gets one; a row that has one keeps it. Refresh is +//! therefore idempotent and does not observe input mutation -- once a row is +//! filled, changing what the expression reads leaves the stored result alone. +//! +//! Two passes per fragment. The first scans only the unfilled live rows and +//! evaluates the expression over them, which yields the exact fill count and +//! decides whether the fragment is staged at all -- a fragment where nothing +//! would change stages nothing, which is what lets an expression yielding +//! null settle instead of restaging forever. The second streams the +//! fragment's physical rows into `write_column` a batch at a time, so peak +//! memory is bounded by a scan batch. The expression is evaluated by this +//! module, never through a projection alias, and only over rows being +//! filled: every other row -- deleted, or already holding a value -- has its +//! inputs masked to null first, so a poison value in a row nobody is filling +//! cannot fail the refresh. + +use std::sync::Arc; + +use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions}; +use arrow_schema::Schema as ArrowSchema; +use datafusion_expr::ColumnarValue; +use futures::{Stream, StreamExt, TryStreamExt}; +use lance::Dataset; +use lance::dataset::WriteDestination; +use lance::dataset::fragment::FileFragment; +use lance::dataset::transaction::Operation; +use lance_core::ROW_ID; +use lance_core::datatypes::Schema as LanceSchema; +use serde::{Deserialize, Serialize}; + +use super::computed_columns::{BoundExpression, ComputedColumnKind, computed_column_from_field}; +use super::{BaseTable, NativeTable}; +use crate::{Error, Result}; + +/// The result of refreshing a computed column. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +pub struct RefreshColumnResult { + /// Rows that had a value computed. + #[serde(default)] + pub rows_filled: u64, + /// The commit version associated with the operation. + #[serde(default)] + pub version: u64, +} + +/// Internal implementation of the refresh logic. +pub(crate) async fn execute_refresh_column( + table: &NativeTable, + column: &str, +) -> Result { + table.dataset.ensure_mutable()?; + ensure_no_lsm_write_spec(table).await?; + let dataset = table.dataset.get().await?; + + let expression = declared_expression(&dataset, column)?; + let schema = Arc::new(ArrowSchema::from(dataset.schema())); + let bound = Arc::new(super::computed_columns::bind(schema, column, &expression)?); + let field = dataset + .schema() + .field(column) + .ok_or_else(|| Error::ColumnNotFound { + name: column.to_string(), + })?; + // The dataset's own field, so the identity write_column checks against the + // manifest holds by construction. + let column_schema = LanceSchema { + fields: vec![field.clone()], + metadata: Default::default(), + }; + + let mut rows_filled = 0u64; + let mut replacements = Vec::new(); + for fragment in dataset.get_fragments() { + let gained = count_fragment_gains(&dataset, &fragment, &bound, column).await?; + if gained == 0 { + continue; + } + rows_filled += gained; + let values = fill_stream(&dataset, &fragment, bound.clone(), column).await?; + replacements.push(fragment.write_column(values, &column_schema).await?); + } + + if replacements.is_empty() { + return Ok(RefreshColumnResult { + rows_filled: 0, + version: dataset.version().version, + }); + } + + let read_version = dataset.version().version; + // The dataset's own session, so registrations and caches survive the + // commit being installed on the handle. + let session = dataset.session(); + let new_dataset = Dataset::commit( + WriteDestination::Dataset(dataset.clone()), + Operation::DataReplacement { replacements }, + Some(read_version), + None, + None, + session, + false, + ) + .await?; + + let version = new_dataset.version().version; + table.dataset.update(new_dataset); + Ok(RefreshColumnResult { + rows_filled, + version, + }) +} + +/// Refuse to refresh under an LSM write spec. +/// +/// Refresh enumerates base fragments, and a write spec keeps visible rows in +/// un-compacted MemWAL tiers it cannot reach -- success would silently omit +/// readable rows. +async fn ensure_no_lsm_write_spec(table: &NativeTable) -> Result<()> { + // The catch-up flag outlives unset and marks retained SSTable rows. + let catchup = table.dataset.get().await?.manifest().reader_feature_flags + & lance_table::feature_flags::FLAG_MEM_WAL_INDEX_CATCHUP + != 0; + if catchup || table.get_lsm_write_spec().await?.is_some() { + return Err(Error::NotSupported { + message: "refresh_column is not supported on a table with an LSM write \ + spec: rows in un-compacted tiers are invisible to refresh" + .into(), + }); + } + Ok(()) +} + +/// The SQL expression `column` is declared with. +fn declared_expression(dataset: &Dataset, column: &str) -> Result { + let schema = ArrowSchema::from(dataset.schema()); + let field = schema + .field_with_name(column) + .map_err(|_| Error::ColumnNotFound { + name: column.to_string(), + })?; + let declaration = + computed_column_from_field(field).ok_or_else(|| Error::NotAComputedColumn { + name: column.to_string(), + })?; + match declaration.kind { + ComputedColumnKind::Sql { expression } => Ok(expression), + ComputedColumnKind::Unrecognized { kind } => Err(Error::NotSupported { + message: format!( + "computed column '{column}' is defined by '{kind}', which this version of \ + lancedb cannot evaluate" + ), + }), + } +} + +/// Quote `name` as a lance SQL identifier. +/// +/// Lance's dialect delimits with backticks, so a double-quoted name would +/// parse as a string literal rather than a column. +fn quote_identifier(name: &str) -> String { + format!("`{}`", name.replace('`', "``")) +} + +/// Assemble the batch evaluation runs against: the bound roots, in read-schema +/// order. Built by name so scan-side column order never matters. +fn evaluation_batch( + batch: &RecordBatch, + bound: &BoundExpression, + mask_out: Option<&BooleanArray>, +) -> lance_core::Result { + let mut columns = Vec::with_capacity(bound.roots.len()); + for name in &bound.roots { + let column = batch.column_by_name(name).ok_or_else(|| { + lance_core::Error::invalid_input(format!( + "refreshing a computed column read no {name} column" + )) + })?; + // Rows outside the mask must not reach the expression: a value in a + // deleted or already-filled row can be one it would choke on. + columns.push(match mask_out { + Some(mask) => arrow::compute::nullif(column, mask)?, + None => column.clone(), + }); + } + Ok(RecordBatch::try_new_with_options( + bound.read_schema.clone(), + columns, + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + )?) +} + +/// Evaluate the expression over `batch`, materializing a constant result to +/// the batch's length. +fn evaluate(bound: &BoundExpression, batch: &RecordBatch) -> lance_core::Result { + let value = bound + .physical + .evaluate(batch) + .map_err(lance_core::Error::from)?; + match value { + ColumnarValue::Array(array) => Ok(array), + scalar => scalar + .into_array(batch.num_rows()) + .map_err(lance_core::Error::from), + } +} + +/// How many rows of one fragment would gain a value. +/// +/// Scans only the unfilled live rows -- deleted rows never reach the +/// expression here, the filter having already excluded them -- and counts the +/// non-null results. Exact, so it is both the staging decision and the +/// fragment's contribution to `rows_filled`. +async fn count_fragment_gains( + dataset: &Dataset, + fragment: &FileFragment, + bound: &BoundExpression, + column: &str, +) -> Result { + let mut scanner = dataset.scan(); + scanner + .with_fragments(vec![fragment.metadata().clone()]) + .with_row_id() + .filter(&format!("{} IS NULL", quote_identifier(column)))? + .project(&bound.roots)?; + + let mut gained = 0u64; + let mut batches = scanner.try_into_stream().await?; + while let Some(batch) = batches.try_next().await? { + let evaluated = evaluate(bound, &evaluation_batch(&batch, bound, None)?)?; + gained += (batch.num_rows() - evaluated.null_count()) as u64; + } + Ok(gained) +} + +/// Stream one fragment's column in physical order, filling the unfilled live +/// rows and keeping every other value. +/// +/// Deleted rows are carried through so the values line up positionally with +/// the fragment's data files; they are never read back, but the column file +/// has to cover them. +async fn fill_stream( + dataset: &Dataset, + fragment: &FileFragment, + bound: Arc, + column: &str, +) -> Result> + Send + use<>> { + let mut projection: Vec = bound.roots.clone(); + projection.push(column.to_string()); + let mut scanner = dataset.scan(); + scanner + .with_fragments(vec![fragment.metadata().clone()]) + .with_row_id() + .include_deleted_rows() + .project(&projection)?; + + let projected = Arc::new(ArrowSchema::new(vec![ + ArrowSchema::from(dataset.schema()) + .field_with_name(column) + .map_err(|_| Error::ColumnNotFound { + name: column.to_string(), + })? + .clone(), + ])); + + let column = column.to_string(); + let batches = scanner.try_into_stream().await?; + Ok(batches.map(move |batch| { + let batch = batch?; + let missing = |name: &str| { + lance_core::Error::invalid_input(format!( + "refreshing a computed column read no {name} column" + )) + }; + let existing = batch + .column_by_name(&column) + .ok_or_else(|| missing(&column))?; + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| missing(ROW_ID))?; + + // Only an unfilled live row gains a value; a deleted row has a null + // row id and keeps its (null) slot. + let unfilled = arrow::compute::is_null(existing.as_ref())?; + let live = arrow::compute::is_not_null(row_ids.as_ref())?; + let fill = arrow::compute::and(&unfilled, &live)?; + let keep = arrow::compute::not(&fill)?; + + let computed = evaluate(&bound, &evaluation_batch(&batch, &bound, Some(&keep))?)?; + let merged = arrow_select::zip::zip(&fill, &computed, existing)?; + Ok(RecordBatch::try_new(projected.clone(), vec![merged])?) + })) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{Int32Array, record_batch}; + use futures::TryStreamExt; + + use crate::connect; + use crate::query::{ExecutableQuery, QueryBase, Select}; + use crate::{Error, Result, Table}; + + async fn table_with(name: &str, values: Vec) -> Table { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!(("x", Int32, values)).unwrap(); + conn.create_table(name, batch).execute().await.unwrap() + } + + async fn declare_doubled(table: &Table) -> Result { + Ok(table + .add_columns() + .computed("doubled", "x * 2") + .execute() + .await? + .version) + } + + 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() + .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_refresh_fills_a_declared_column() { + let table = table_with("refresh_fills", vec![1, 2, 3]).await; + let declared = declare_doubled(&table).await.unwrap(); + assert_eq!(read(&table, "doubled").await, vec![None, None, None]); + + let result = table.refresh_column("doubled").await.unwrap(); + assert!(result.version > declared); + assert_eq!(result.rows_filled, 3); + assert_eq!( + read(&table, "doubled").await, + vec![Some(2), Some(4), Some(6)] + ); + } + + /// Values written after the last refresh must be reachable by another one. + #[tokio::test] + async fn test_refresh_fills_rows_appended_since_the_last_refresh() { + let table = table_with("refresh_appended", vec![1, 2]).await; + declare_doubled(&table).await.unwrap(); + table.refresh_column("doubled").await.unwrap(); + + append(&table, vec![5, 6]).await; + assert_eq!( + read(&table, "doubled").await, + vec![None, None, Some(2), Some(4)] + ); + + let result = table.refresh_column("doubled").await.unwrap(); + assert_eq!(result.rows_filled, 2); + assert_eq!( + read(&table, "doubled").await, + vec![Some(2), Some(4), Some(10), Some(12)] + ); + } + + #[tokio::test] + async fn test_refresh_with_nothing_to_fill() { + let table = table_with("refresh_noop", vec![1, 2, 3]).await; + declare_doubled(&table).await.unwrap(); + table.refresh_column("doubled").await.unwrap(); + + let again = table.refresh_column("doubled").await.unwrap(); + assert_eq!(again.rows_filled, 0); + assert_eq!( + read(&table, "doubled").await, + vec![Some(2), Some(4), Some(6)] + ); + } + + /// A row is filled only by gaining a value, so an expression yielding null + /// settles at once instead of re-selecting the same rows forever. Nothing + /// is staged, so the version does not move either. + #[tokio::test] + async fn test_refresh_converges_on_a_null_result() { + let table = table_with("refresh_null_result", vec![1, 2, 3]).await; + let declared = table + .add_columns() + .computed("maybe", "nullif(x, x)") + .execute() + .await + .unwrap() + .version; + + let first = table.refresh_column("maybe").await.unwrap(); + assert_eq!(first.rows_filled, 0); + assert_eq!(first.version, declared); + assert_eq!(read(&table, "maybe").await, vec![None, None, None]); + + let again = table.refresh_column("maybe").await.unwrap(); + assert_eq!(again.rows_filled, 0); + assert_eq!(again.version, declared); + } + + /// The contract's boundary: a filled fragment is not revisited, so + /// mutating an input leaves the value computed at fill time. + #[tokio::test] + async fn test_refresh_does_not_observe_input_mutation() { + let table = table_with("refresh_mutation", vec![1]).await; + declare_doubled(&table).await.unwrap(); + table.refresh_column("doubled").await.unwrap(); + assert_eq!(read(&table, "doubled").await, vec![Some(2)]); + + table.update().column("x", "3").execute().await.unwrap(); + + let again = table.refresh_column("doubled").await.unwrap(); + assert_eq!(again.rows_filled, 0); + assert_eq!(read(&table, "doubled").await, vec![Some(2)]); + } + + /// A row rewrite before the first refresh materializes the declared + /// column as null behind a covering data file. Those rows are still + /// unfilled and a later refresh has to reach them. + #[tokio::test] + async fn test_update_before_the_first_refresh() { + let table = table_with("refresh_update_first", vec![1]).await; + declare_doubled(&table).await.unwrap(); + + table.update().column("x", "3").execute().await.unwrap(); + + let result = table.refresh_column("doubled").await.unwrap(); + assert_eq!(result.rows_filled, 1); + assert_eq!(read(&table, "doubled").await, vec![Some(6)]); + } + + /// The contract holds row by row, not fragment by fragment: revisiting a + /// fragment to fill one row must not recompute a filled row sitting beside + /// it, even where the input behind it has since changed. + #[tokio::test] + async fn test_refresh_does_not_recompute_a_filled_row_beside_an_unfilled_one() { + let table = table_with("refresh_mixed", vec![1, 2]).await; + declare_doubled(&table).await.unwrap(); + table.refresh_column("doubled").await.unwrap(); + + append(&table, vec![5]).await; + table + .update() + .column("x", "100") + .only_if("x = 1") + .execute() + .await + .unwrap(); + table + .optimize(crate::table::OptimizeAction::Compact { + options: crate::table::CompactionOptions::default(), + remap_options: None, + }) + .await + .unwrap(); + + let result = table.refresh_column("doubled").await.unwrap(); + assert_eq!(result.rows_filled, 1); + // 2 is the mutated row keeping the value it was filled with, not 200. + assert_eq!( + read(&table, "doubled").await, + vec![Some(2), Some(4), Some(10)] + ); + } + + /// Filling a fragment must not disturb the values it already holds, which + /// is what makes a compaction-mixed fragment safe to revisit. + #[tokio::test] + async fn test_refresh_preserves_already_filled_rows() { + let table = table_with("refresh_preserves", vec![1, 2]).await; + declare_doubled(&table).await.unwrap(); + table.refresh_column("doubled").await.unwrap(); + + append(&table, vec![5]).await; + table + .optimize(crate::table::OptimizeAction::Compact { + options: crate::table::CompactionOptions::default(), + remap_options: None, + }) + .await + .unwrap(); + + let result = table.refresh_column("doubled").await.unwrap(); + assert_eq!(result.rows_filled, 1); + assert_eq!( + read(&table, "doubled").await, + vec![Some(2), Some(4), Some(10)] + ); + } + + #[tokio::test] + async fn test_refresh_leaves_deleted_rows_alone() { + let table = table_with("refresh_deleted", vec![1, 2, 3, 4]).await; + declare_doubled(&table).await.unwrap(); + table.delete("x = 2").await.unwrap(); + + let result = table.refresh_column("doubled").await.unwrap(); + assert_eq!(result.rows_filled, 3); + assert_eq!( + read(&table, "doubled").await, + vec![Some(2), Some(6), Some(8)] + ); + } + + #[tokio::test] + async fn test_refresh_a_constant_expression() { + let table = table_with("refresh_constant", vec![1, 2, 3]).await; + table + .add_columns() + .computed("answer", "42") + .execute() + .await + .unwrap(); + + let result = table.refresh_column("answer").await.unwrap(); + assert_eq!(result.rows_filled, 3); + } + + /// A name needing quotes reaches the evaluator intact: it is carried as a + /// projection alias, never spliced into SQL text. + #[tokio::test] + async fn test_refresh_a_column_whose_name_needs_quoting() { + let table = table_with("refresh_quoted", vec![1, 2, 3]).await; + table + .add_columns() + .computed("double value", "x * 2") + .execute() + .await + .unwrap(); + + let result = table.refresh_column("double value").await.unwrap(); + assert_eq!(result.rows_filled, 3); + assert_eq!( + read(&table, "double value").await, + vec![Some(2), Some(4), Some(6)] + ); + } + + /// A fragment spanning several scan batches exercises the streamed fill: + /// the probe buffers only until the first gained value and the rest flows + /// through write_column a batch at a time. + #[tokio::test] + async fn test_refresh_streams_a_multi_batch_fragment() { + let values: Vec = (0..20_000).collect(); + let table = table_with("refresh_multi_batch", values.clone()).await; + declare_doubled(&table).await.unwrap(); + + let result = table.refresh_column("doubled").await.unwrap(); + assert_eq!(result.rows_filled, 20_000); + + let read_back = read(&table, "doubled").await; + assert_eq!(read_back.len(), 20_000); + let mut expected: Vec> = values.iter().map(|v| Some(v * 2)).collect(); + expected.sort(); + assert_eq!(read_back, expected); + } + + /// The gate's reproducer: the commit must reuse the configured session, + /// or registrations and caches vanish from the handle after a refresh. + #[tokio::test] + async fn test_refresh_preserves_the_configured_session() { + let session = Arc::new(lance::session::Session::default()); + let conn = crate::connect("memory://") + .session(session.clone()) + .execute() + .await + .unwrap(); + let batch = record_batch!(("x", Int32, [1, 2])).unwrap(); + let table = conn + .create_table("session_kept", batch) + .execute() + .await + .unwrap(); + declare_doubled(&table).await.unwrap(); + table.refresh_column("doubled").await.unwrap(); + + let dataset = table.as_native().unwrap().dataset.get().await.unwrap(); + assert!(Arc::ptr_eq(&dataset.session(), &session)); + } + + /// Both orders of declare+spec are refused at the source (see the + /// schema_evolution tests); refresh's own check covers a dataset another + /// writer left in that state. + #[tokio::test] + async fn test_refresh_refuses_a_foreign_lsm_state() { + use crate::table::LsmWriteSpec; + + let tmp_dir = tempfile::tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "x", + arrow_schema::DataType::Int32, + false, + )])); + let batch = + arrow_array::RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + let table = conn.create_table("lsm", batch).execute().await.unwrap(); + table.set_unenforced_primary_key(["x"]).await.unwrap(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + super::super::computed_columns::add_foreign_kind(&table, "doubled", "sql").await; + + let err = table.refresh_column("doubled").await.unwrap_err(); + assert!( + matches!(&err, Error::NotSupported { message } if message.contains("LSM")), + "{err:?}" + ); + } + + /// After catch-up activation and unset, no spec remains but the catch-up + /// flag still marks retained SSTable rows; refresh refuses on the flag. + #[tokio::test] + async fn test_refresh_refuses_retained_catchup_state() { + use crate::table::LsmWriteSpec; + + let tmp_dir = tempfile::tempdir().unwrap(); + let conn = connect(tmp_dir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + let schema = Arc::new(arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "x", + arrow_schema::DataType::Int32, + false, + )])); + let batch = arrow_array::RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + let table = conn + .create_table("catchup", batch.clone()) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["x"]).await.unwrap(); + table + .set_lsm_write_spec(LsmWriteSpec::unsharded()) + .await + .unwrap(); + table.require_mem_wal_index_catchup().await.unwrap(); + let mut merge = table.merge_insert(&["x"]); + merge + .when_matched_update_all(None) + .when_not_matched_insert_all() + .use_lsm(true); + merge + .execute(Box::new(arrow_array::RecordBatchIterator::new( + vec![Ok(batch)], + schema, + ))) + .await + .unwrap(); + table.unset_lsm_write_spec().await.unwrap(); + super::super::computed_columns::add_foreign_kind(&table, "doubled", "sql").await; + + let err = table.refresh_column("doubled").await.unwrap_err(); + assert!( + matches!(&err, Error::NotSupported { message } if message.contains("LSM")), + "{err:?}" + ); + } + + /// A declaration of a kind this version cannot evaluate is refused by + /// name, rather than mistaken for a plain column or fed to the SQL path. + #[tokio::test] + async fn test_refresh_rejects_a_kind_it_cannot_evaluate() { + let table = table_with("refresh_foreign", vec![1, 2, 3]).await; + super::super::computed_columns::add_foreign_kind(&table, "embedding", "udf").await; + + let err = table.refresh_column("embedding").await.unwrap_err(); + assert!(matches!(err, Error::NotSupported { message } if message.contains("udf"))); + } +} diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index ce208111a..7503fd790 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -8,12 +8,14 @@ //! - [`alter_columns`](execute_alter_columns): Rename columns, change types, or modify nullability //! - [`drop_columns`](execute_drop_columns): Remove columns from the table +use arrow_schema::Schema as ArrowSchema; use lance::dataset::{ColumnAlteration, NewColumnTransform}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use super::NativeTable; -use crate::Result; +use super::computed_columns; +use super::{BaseTable, NativeTable}; +use crate::{Error, Result}; /// The result of an add columns operation. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] @@ -98,6 +100,48 @@ pub(crate) async fn execute_add_columns( table: &NativeTable, transforms: NewColumnTransform, read_columns: Option>, +) -> Result { + // Declarations are admitted only through [`execute_declare`]. + match &transforms { + NewColumnTransform::AllNulls(schema) => { + computed_columns::ensure_no_foreign_declarations(schema.fields())? + } + NewColumnTransform::BatchUDF(udf) => { + computed_columns::ensure_no_foreign_declarations(udf.output_schema.fields())? + } + _ => {} + } + commit_add_columns(table, transforms, read_columns).await +} + +/// Declare validated computed columns. The only admission path for +/// declaration metadata. +pub(crate) async fn execute_declare( + table: &NativeTable, + columns: &[(String, String)], +) -> Result { + // An LSM write spec keeps visible rows in tiers refresh cannot reach; + // checked against latest committed state, not this handle's snapshot. + // The catch-up flag outlives unset and marks retained SSTable rows. + table.checkout_latest().await?; + let catchup = table.dataset.get().await?.manifest().reader_feature_flags + & lance_table::feature_flags::FLAG_MEM_WAL_INDEX_CATCHUP + != 0; + if catchup || table.get_lsm_write_spec().await?.is_some() { + return Err(Error::NotSupported { + message: "computed columns are not supported on a table with an LSM write \ + spec: rows in un-compacted tiers are invisible to refresh" + .into(), + }); + } + let transform = computed_columns::declare(table.schema().await?, columns)?; + commit_add_columns(table, transform, None).await +} + +pub(crate) async fn commit_add_columns( + table: &NativeTable, + transforms: NewColumnTransform, + read_columns: Option>, ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); @@ -116,6 +160,21 @@ pub(crate) async fn execute_alter_columns( ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); + // Nullability is not part of what an expression resolves against, so only + // a rename or a retype can invalidate a binding. + let schema = std::sync::Arc::new(ArrowSchema::from(dataset.schema())); + let rebinding = alterations + .iter() + .filter(|alteration| alteration.rename.is_some() || alteration.data_type.is_some()) + .map(|alteration| alteration.path.as_str()) + .collect::>(); + computed_columns::ensure_not_an_input(&schema, &rebinding)?; + let retyped = alterations + .iter() + .filter(|alteration| alteration.data_type.is_some()) + .map(|alteration| alteration.path.as_str()) + .collect::>(); + computed_columns::ensure_not_retyped(schema.as_ref(), &retyped)?; dataset.alter_columns(alterations).await?; let version = dataset.version().version; table.dataset.update(dataset); @@ -131,6 +190,10 @@ pub(crate) async fn execute_drop_columns( ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); + computed_columns::ensure_not_an_input( + &std::sync::Arc::new(ArrowSchema::from(dataset.schema())), + columns, + )?; dataset.drop_columns(columns).await?; let version = dataset.version().version; table.dataset.update(dataset); @@ -147,6 +210,44 @@ pub(crate) async fn execute_update_field_metadata( table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); + // A declaration is validated as a whole at declare time; editing its keys + // here would bypass that, fabricate one on a plain column, or move a + // binding out from under a refresh. A replace on a declared column would + // silently erase it. + let schema = ArrowSchema::from(dataset.schema()); + let declared: Vec = computed_columns::computed_columns(&schema) + .into_iter() + .map(|declaration| declaration.name) + .collect(); + for update in updates { + if update + .metadata + .keys() + .any(|key| computed_columns::is_declaration_key(key)) + { + return Err(Error::InvalidInput { + message: format!( + "metadata keys of a computed-column declaration cannot be edited \ + (path '{}'); drop the column and declare it again", + update.path + ), + }); + } + if update.replace + && declared + .iter() + .any(|name| name == computed_columns::root(&update.path)) + { + return Err(Error::InvalidInput { + message: format!( + "replacing all metadata of computed column '{}' would erase its \ + declaration; drop the column and declare it again", + update.path + ), + }); + } + } + let mut builder = dataset.update_field_metadata(); for update in updates { let entries = update.metadata.iter().map(|(k, v)| (k.clone(), v.clone())); diff --git a/rust/lancedb/src/table/update.rs b/rust/lancedb/src/table/update.rs index 61eb93992..fd9fa6828 100644 --- a/rust/lancedb/src/table/update.rs +++ b/rust/lancedb/src/table/update.rs @@ -82,6 +82,10 @@ pub(crate) async fn execute_update( // 1. Snapshot the current dataset let dataset = table.dataset.get().await?; + super::computed_columns::ensure_not_written( + &arrow_schema::Schema::from(dataset.schema()), + update.columns.iter().map(|(name, _)| name.as_str()), + )?; // 2. Initialize the Lance Core builder let mut builder = LanceUpdateBuilder::new(dataset); diff --git a/rust/lancedb/tests/blob_integration.rs b/rust/lancedb/tests/blob_integration.rs index 77d49abd9..b92f961f4 100644 --- a/rust/lancedb/tests/blob_integration.rs +++ b/rust/lancedb/tests/blob_integration.rs @@ -10,7 +10,7 @@ use arrow_array::{ use arrow_schema::{DataType, Field, Fields, Schema}; use futures::TryStreamExt; use lance::Dataset; -use lance_file::version::LanceFileVersion; +use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; use lancedb::{ Connection, Error, Result, Table, blob::{BlobRangeRequest, blob}, @@ -61,7 +61,7 @@ async fn create_inline_blob_table( Ok(table) } -async fn storage_format_version(table: &Table) -> LanceFileVersion { +async fn storage_format_version(table: &Table) -> ConcreteFileVersion { table .as_native() .unwrap() @@ -69,9 +69,14 @@ async fn storage_format_version(table: &Table) -> LanceFileVersion { .await .unwrap() .data_storage_format - .lance_file_version() - .unwrap() - .resolve() + .lance_file_format() +} + +fn supports_blob_v2(version: ConcreteFileVersion) -> bool { + matches!( + version, + ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 + ) } async fn uses_stable_row_ids(table: &Table) -> bool { @@ -112,7 +117,7 @@ async fn declaring_blob_column_bumps_format_and_enables_stable_row_ids() -> Resu .execute() .await?; - assert!(storage_format_version(&table).await >= LanceFileVersion::V2_2); + assert!(supports_blob_v2(storage_format_version(&table).await)); assert!(uses_stable_row_ids(&table).await); Ok(()) } @@ -127,7 +132,7 @@ async fn explicit_stable_row_id_setting_wins_over_blob_default() -> Result<()> { .execute() .await?; - assert!(storage_format_version(&table).await >= LanceFileVersion::V2_2); + assert!(supports_blob_v2(storage_format_version(&table).await)); assert!(!uses_stable_row_ids(&table).await); Ok(()) } @@ -139,7 +144,7 @@ async fn non_blob_table_keeps_default_format_and_row_id_setting() -> Result<()> let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); let table = db.create_empty_table("t", schema).execute().await?; - assert!(storage_format_version(&table).await < LanceFileVersion::V2_2); + assert!(!supports_blob_v2(storage_format_version(&table).await)); assert!(!uses_stable_row_ids(&table).await); Ok(()) } @@ -171,7 +176,7 @@ async fn creating_with_blob_data_bumps_format() -> Result<()> { .unwrap(); let table = db.create_table("t", batch).execute().await?; - assert!(storage_format_version(&table).await >= LanceFileVersion::V2_2); + assert!(supports_blob_v2(storage_format_version(&table).await)); assert!(uses_stable_row_ids(&table).await); assert_eq!(table.count_rows(None).await?, 1); Ok(()) @@ -281,7 +286,7 @@ async fn connection_level_stable_row_id_setting_wins_over_blob_default() -> Resu .execute() .await?; - assert!(storage_format_version(&table).await >= LanceFileVersion::V2_2); + assert!(supports_blob_v2(storage_format_version(&table).await)); assert!(!uses_stable_row_ids(&table).await); Ok(()) } @@ -297,7 +302,7 @@ async fn namespace_create_applies_blob_defaults() -> Result<()> { .execute() .await?; - assert!(storage_format_version(&table).await >= LanceFileVersion::V2_2); + assert!(supports_blob_v2(storage_format_version(&table).await)); assert!(uses_stable_row_ids(&table).await); Ok(()) } @@ -474,7 +479,7 @@ async fn fetch_blobs_round_trips_nested_blob_column() -> Result<()> { let batch = RecordBatch::try_new(schema, vec![Arc::new(info_array) as ArrayRef]).unwrap(); let table = db.create_table("t", batch).execute().await?; - assert!(storage_format_version(&table).await >= LanceFileVersion::V2_2); + assert!(supports_blob_v2(storage_format_version(&table).await)); assert!(uses_stable_row_ids(&table).await); let ids = collect_row_ids(&table).await?; @@ -1305,7 +1310,7 @@ async fn optimize_preserves_blob_v2_null_and_empty_distinction() -> Result<()> { .await?; table.add(null_empty_input_batch()).execute().await?; assert!( - storage_format_version(&table).await >= LanceFileVersion::V2_2, + supports_blob_v2(storage_format_version(&table).await), "blob v2 columns require storage >= 2.2" );