mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 03:58:26 +00:00
Merge origin/main into gatekeeper/fix-2325-1
This commit is contained in:
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.37.1-beta.0"
|
||||
current_version = "0.38.0-beta.0"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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: |
|
||||
|
||||
@@ -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
|
||||
|
||||
Generated
+47
-60
@@ -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",
|
||||
|
||||
+14
-14
@@ -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 }
|
||||
|
||||
@@ -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" },
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.37.1-beta.0</version>
|
||||
<version>0.38.0-beta.0</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -386,6 +386,29 @@ Drop an existing table.
|
||||
|
||||
***
|
||||
|
||||
### dropTableAsync()
|
||||
|
||||
```ts
|
||||
abstract dropTableAsync(name, namespacePath?): Promise<Job>
|
||||
```
|
||||
|
||||
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
|
||||
|
||||
@@ -69,14 +69,33 @@ abstract addColumns(newColumnTransforms): Promise<AddColumnsResult>
|
||||
|
||||
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<RefreshColumnResult>
|
||||
```
|
||||
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -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;
|
||||
```
|
||||
@@ -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.
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.37.1-beta.0</version>
|
||||
<version>0.38.0-beta.0</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.37.1-beta.0</version>
|
||||
<version>0.38.0-beta.0</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>${project.artifactId}</name>
|
||||
<description>LanceDB Java SDK Parent POM</description>
|
||||
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>11.0.0-beta.3</lance-core.version>
|
||||
<lance-core.version>11.0.0-beta.11</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
@@ -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<Record<string, any>> {
|
||||
@@ -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: <explanation>
|
||||
} = <any>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()),
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -327,6 +327,14 @@ export abstract class Connection {
|
||||
*/
|
||||
abstract dropTable(name: string, namespacePath?: string[]): Promise<void>;
|
||||
|
||||
/**
|
||||
* 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<Job>;
|
||||
|
||||
/**
|
||||
* 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<Job> {
|
||||
return this.inner.dropTableAsync(name, namespacePath ?? []);
|
||||
}
|
||||
|
||||
async dropAllTables(namespacePath?: string[]): Promise<void> {
|
||||
return this.inner.dropAllTables(namespacePath ?? []);
|
||||
}
|
||||
|
||||
@@ -50,6 +50,7 @@ export {
|
||||
MergeResult,
|
||||
AddResult,
|
||||
AddColumnsResult,
|
||||
RefreshColumnResult,
|
||||
AlterColumnsResult,
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
|
||||
+174
-29
@@ -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<object, DataType>;
|
||||
vectors: WeakMap<object, Vector>;
|
||||
data: WeakMap<object, Data<DataType>>;
|
||||
};
|
||||
|
||||
function createSanitizationContext(): SanitizationContext {
|
||||
return {
|
||||
types: new WeakMap(),
|
||||
vectors: new WeakMap(),
|
||||
data: new WeakMap(),
|
||||
};
|
||||
}
|
||||
|
||||
export function sanitizeMetadata(
|
||||
metadataLike?: unknown,
|
||||
): Map<string, string> | 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<any> {
|
||||
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<any> {
|
||||
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<any> {
|
||||
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<any> {
|
||||
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<any> {
|
||||
}
|
||||
|
||||
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<Struct>;
|
||||
return new RecordBatch(schema, data);
|
||||
}
|
||||
|
||||
type DictionaryVectorLike = {
|
||||
data: readonly DataLike[];
|
||||
};
|
||||
|
||||
type DictionaryDataLike = DataLike & {
|
||||
dictionary?: DictionaryVectorLike;
|
||||
};
|
||||
|
||||
function sanitizeData(
|
||||
dataLike: DataLike,
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
): import("apache-arrow").Data<Struct<any>> {
|
||||
context: SanitizationContext,
|
||||
): Data<DataType> {
|
||||
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 = {
|
||||
|
||||
+71
-6
@@ -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<string, string>;
|
||||
@@ -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<AddColumnsResult>} 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<AddColumnsResult>;
|
||||
|
||||
/**
|
||||
* 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<RefreshColumnResult>} A promise that resolves to the
|
||||
* number of rows filled and the new version number of the table.
|
||||
*/
|
||||
abstract refreshColumn(column: string): Promise<RefreshColumnResult>;
|
||||
|
||||
/**
|
||||
* 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<void>}
|
||||
* @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<LsmWriteSpec | undefined>}
|
||||
*/
|
||||
abstract getLsmWriteSpec(): Promise<LsmWriteSpec | undefined>;
|
||||
@@ -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<AddColumnsResult> {
|
||||
// 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<RefreshColumnResult> {
|
||||
return await this.inner.refreshColumn(column);
|
||||
}
|
||||
|
||||
async alterColumns(
|
||||
columnAlterations: ColumnAlteration[],
|
||||
): Promise<AlterColumnsResult> {
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.37.1-beta.0",
|
||||
"version": "0.38.0-beta.0",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -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",
|
||||
|
||||
Generated
+2
-2
@@ -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"
|
||||
|
||||
+1
-1
@@ -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",
|
||||
|
||||
@@ -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<Vec<String>>,
|
||||
) -> napi::Result<crate::job::Job> {
|
||||
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<Vec<String>>) -> napi::Result<()> {
|
||||
let ns = namespace_path.unwrap_or_default();
|
||||
|
||||
+49
-7
@@ -347,6 +347,30 @@ impl Table {
|
||||
Ok(res.into())
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn add_computed_columns(
|
||||
&self,
|
||||
columns: Vec<AddColumnsSql>,
|
||||
) -> napi::Result<AddColumnsResult> {
|
||||
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<RefreshColumnResult> {
|
||||
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<String>,
|
||||
/// Bucket variant: the number of buckets, in `[1, 1024]`.
|
||||
pub num_buckets: Option<u32>,
|
||||
/// 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<Vec<String>>,
|
||||
/// Default `ShardWriter` configuration recorded in the MemWAL index.
|
||||
pub writer_config_defaults: Option<HashMap<String, String>>,
|
||||
@@ -782,7 +807,6 @@ impl TryFrom<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
type Error = napi::Error;
|
||||
|
||||
fn try_from(value: LsmWriteSpec) -> napi::Result<Self> {
|
||||
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<LsmWriteSpec> 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<lancedb::table::LsmWriteSpec> 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<lancedb::table::LsmWriteSpec> 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<lancedb::table::LsmWriteSpec> 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<lancedb::index::IndexStatistics> 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<lancedb::table::RefreshColumnResult> for RefreshColumnResult {
|
||||
fn from(value: lancedb::table::RefreshColumnResult) -> Self {
|
||||
Self {
|
||||
rows_filled: value.rows_filled as i64,
|
||||
version: value.version as i64,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::AddColumnsResult> for AddColumnsResult {
|
||||
fn from(value: lancedb::table::AddColumnsResult) -> Self {
|
||||
Self {
|
||||
|
||||
+1
-1
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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)))
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
{
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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<Vec<String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
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>,
|
||||
|
||||
+1
-1
@@ -289,7 +289,7 @@ struct IvfHnswFlatParams {
|
||||
target_partition_size: Option<u32>,
|
||||
}
|
||||
|
||||
#[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
|
||||
|
||||
+3
-1
@@ -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::<VectorQuery>()?;
|
||||
m.add_class::<RecordBatchStream>()?;
|
||||
m.add_class::<AddColumnsResult>()?;
|
||||
m.add_class::<RefreshColumnResult>()?;
|
||||
m.add_class::<AlterColumnsResult>()?;
|
||||
m.add_class::<UpdateFieldMetadataResult>()?;
|
||||
m.add_class::<AddResult>()?;
|
||||
|
||||
@@ -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<LanceSession>,
|
||||
|
||||
+81
-16
@@ -341,12 +341,22 @@ impl From<lancedb::table::MergeResult> for MergeResult {
|
||||
}
|
||||
}
|
||||
|
||||
/// Render for `__repr__`, so the default reads as Python's `None` rather than
|
||||
/// Rust's `Some([..])`.
|
||||
fn fmt_maintained(maintained: &Option<Vec<String>>) -> 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<String>) -> 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<Vec<String>>) -> 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<String> {
|
||||
self.inner.maintained_indexes().to_vec()
|
||||
pub fn maintained_indexes(&self) -> Option<Vec<String>> {
|
||||
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<lancedb::table::RefreshColumnResult> 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<Bound<'_, PyAny>> {
|
||||
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<Bound<'_, PyAny>> {
|
||||
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<Schema>,
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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<usize> {
|
||||
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<Vec<usize>> {
|
||||
let raw = std::env::var("BENCH_SIBLINGS").unwrap_or_else(|_| "1000,10000,100000".into());
|
||||
let mut counts = raw
|
||||
.split(',')
|
||||
.map(|value| {
|
||||
value
|
||||
.trim()
|
||||
.parse::<usize>()
|
||||
.with_context(|| format!("invalid BENCH_SIBLINGS value: {value}"))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
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<Duration> {
|
||||
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(())
|
||||
}
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
/// `<name>.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 `<name>.lance` directory alone do not
|
||||
/// make a table openable.
|
||||
pub fn open_table(&self, name: impl Into<String>) -> 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<str>,
|
||||
namespace_path: &[String],
|
||||
) -> Result<crate::job::Job> {
|
||||
self.internal
|
||||
.drop_table_async(name.as_ref(), namespace_path)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Drop the database
|
||||
///
|
||||
/// This is the same as dropping all of the tables
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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<crate::job::Job> {
|
||||
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;
|
||||
|
||||
@@ -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<Barrier>,
|
||||
}
|
||||
|
||||
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<dyn Scannable>) -> 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<dyn object_store::ObjectStore> =
|
||||
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();
|
||||
|
||||
@@ -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 },
|
||||
|
||||
@@ -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<String> {
|
||||
serde_json::from_str::<serde_json::Value>(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};
|
||||
|
||||
@@ -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<S: HttpSend> RemoteDatabase<S> {
|
||||
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<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
}
|
||||
|
||||
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<Job> {
|
||||
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| {
|
||||
|
||||
@@ -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<S: HttpSend> RemoteTable<S> {
|
||||
.text()
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|body| serde_json::from_str::<serde_json::Value>(&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<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
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<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
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<String> {
|
||||
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")
|
||||
|
||||
+684
-143
File diff suppressed because it is too large
Load Diff
@@ -15,6 +15,7 @@ use crate::{Error, Result};
|
||||
pub struct AddColumnsBuilder {
|
||||
parent: Arc<dyn BaseTable>,
|
||||
transform: Option<NewColumnTransform>,
|
||||
computed: Vec<(String, String)>,
|
||||
read_columns: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
@@ -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<dyn std::error::Error>> {
|
||||
/// 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<String>, expression: impl Into<String>) -> 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<Item = impl Into<String>>) -> 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;
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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<String> = 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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
@@ -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<Vec<String>> {
|
||||
let Some(requested) = requested else {
|
||||
let all: Vec<String> = 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
|
||||
// =============================================================================
|
||||
|
||||
@@ -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<Vec<String>> {
|
||||
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<Uuid, u64> {
|
||||
let mut exclude: HashMap<Uuid, u64> = 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<ShardSnapshot>, HashMap<Uuid, InMemoryMemTables>)> {
|
||||
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<Option<String>> {
|
||||
) -> Result<Vec<String>> {
|
||||
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<String> = 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<String> = 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]
|
||||
|
||||
@@ -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<RefreshColumnResult> {
|
||||
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<String> {
|
||||
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<RecordBatch> {
|
||||
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<ArrayRef> {
|
||||
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<u64> {
|
||||
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<BoundExpression>,
|
||||
column: &str,
|
||||
) -> Result<impl Stream<Item = lance_core::Result<RecordBatch>> + Send + use<>> {
|
||||
let mut projection: Vec<String> = 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<i32>) -> 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<u64> {
|
||||
Ok(table
|
||||
.add_columns()
|
||||
.computed("doubled", "x * 2")
|
||||
.execute()
|
||||
.await?
|
||||
.version)
|
||||
}
|
||||
|
||||
async fn read(table: &Table, column: &str) -> Vec<Option<i32>> {
|
||||
let batches = table
|
||||
.query()
|
||||
.select(Select::columns(&[column]))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
let mut values: Vec<Option<i32>> = batches
|
||||
.iter()
|
||||
.flat_map(|batch| {
|
||||
batch[column]
|
||||
.as_any()
|
||||
.downcast_ref::<Int32Array>()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.collect();
|
||||
values.sort();
|
||||
values
|
||||
}
|
||||
|
||||
async fn append(table: &Table, values: Vec<i32>) {
|
||||
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<i32> = (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<Option<i32>> = 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")));
|
||||
}
|
||||
}
|
||||
@@ -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<Vec<String>>,
|
||||
) -> Result<AddColumnsResult> {
|
||||
// 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<AddColumnsResult> {
|
||||
// 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<Vec<String>>,
|
||||
) -> Result<AddColumnsResult> {
|
||||
table.dataset.ensure_mutable()?;
|
||||
let mut dataset = (*table.dataset.get().await?).clone();
|
||||
@@ -116,6 +160,21 @@ pub(crate) async fn execute_alter_columns(
|
||||
) -> Result<AlterColumnsResult> {
|
||||
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::<Vec<_>>();
|
||||
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::<Vec<_>>();
|
||||
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<DropColumnsResult> {
|
||||
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<String> = 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()));
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user