mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-27 08:28:28 +00:00
Compare commits
26 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| cbd5b4e288 | |||
| 897c492531 | |||
| 4f0cf10a9b | |||
| f1626012df | |||
| de5ca64037 | |||
| 9b42f0a8bc | |||
| 0bc081608a | |||
| d6f9f8560e | |||
| 0bd0944062 | |||
| 91f775c093 | |||
| 2ce88f8e02 | |||
| ac99e4dce5 | |||
| 82231bf66d | |||
| 8d2fea9151 | |||
| 2f27aa377b | |||
| 1bf6b3ea7e | |||
| 8450683b2a | |||
| 65cd142c7e | |||
| 5d0a1ef66c | |||
| 82906ecfee | |||
| ab3041e01e | |||
| 7813907eb7 | |||
| f05140f21c | |||
| dfce767f4c | |||
| 7b6ee0d655 | |||
| ca39258342 |
@@ -0,0 +1,20 @@
|
||||
{
|
||||
"name": "lancedb",
|
||||
"interface": {
|
||||
"displayName": "LanceDB"
|
||||
},
|
||||
"plugins": [
|
||||
{
|
||||
"name": "lancedb",
|
||||
"source": {
|
||||
"source": "local",
|
||||
"path": "./plugins/lancedb"
|
||||
},
|
||||
"policy": {
|
||||
"installation": "AVAILABLE",
|
||||
"authentication": "ON_INSTALL"
|
||||
},
|
||||
"category": "Developer Tools"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -5,3 +5,7 @@ This directory contains repo-scoped code agent skills for the LanceDB project.
|
||||
Each skill is a folder that contains a required `SKILL.md` and optional bundled resources.
|
||||
|
||||
Codex discovers skills from `.agents/skills` in the current working directory and parent directories.
|
||||
|
||||
The `lancedb` skill lives in the `plugins/lancedb` plugin (see `plugins/lancedb/skills/lancedb`)
|
||||
so it can be installed via the plugin marketplaces (`.claude-plugin/marketplace.json` and
|
||||
`.agents/plugins/marketplace.json`); the `lancedb` entry here is a symlink into that plugin.
|
||||
|
||||
Symlink
+1
@@ -0,0 +1 @@
|
||||
../../plugins/lancedb/skills/lancedb
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.32.0-beta.2"
|
||||
current_version = "0.33.0-beta.0"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"name": "lancedb",
|
||||
"owner": {
|
||||
"name": "LanceDB"
|
||||
},
|
||||
"description": "LanceDB plugins for Claude Code.",
|
||||
"plugins": [
|
||||
{
|
||||
"name": "lancedb",
|
||||
"source": "./plugins/lancedb",
|
||||
"description": "Write, review, debug, and document LanceDB pipelines in Python and TypeScript that work across local LanceDB OSS tables and remote LanceDB Enterprise/Cloud tables.",
|
||||
"version": "0.1.0",
|
||||
"author": {
|
||||
"name": "LanceDB"
|
||||
},
|
||||
"category": "development"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -18,6 +18,14 @@ inputs:
|
||||
description: "The manylinux version to build for"
|
||||
required: false
|
||||
default: "2_17"
|
||||
package-name:
|
||||
description: "Override [project] name in python/pyproject.toml (e.g. 'lancedb-compat'). Default keeps 'lancedb'."
|
||||
required: false
|
||||
default: "lancedb"
|
||||
rustflags:
|
||||
description: "RUSTFLAGS for the build container, as a single whitespace-free token (e.g. '-Ctarget-cpu=x86-64-v2'). Empty leaves RUSTFLAGS unset, keeping the defaults from .cargo/config.toml."
|
||||
required: false
|
||||
default: ""
|
||||
runs:
|
||||
using: "composite"
|
||||
steps:
|
||||
@@ -27,6 +35,18 @@ runs:
|
||||
ARM_BUILD: ${{ inputs.arm-build }}
|
||||
run: |
|
||||
echo "ARM BUILD: $ARM_BUILD"
|
||||
- name: Patch package name for variant build
|
||||
if: ${{ inputs.package-name != 'lancedb' }}
|
||||
shell: bash
|
||||
env:
|
||||
PACKAGE_NAME: ${{ inputs.package-name }}
|
||||
run: |
|
||||
# Swap the [project] name so this build produces e.g. lancedb-compat
|
||||
# wheels. The package still installs files under the lancedb/
|
||||
# namespace -- import lancedb still works after pip install.
|
||||
sed -i.bak 's/^name = "lancedb"$/name = "'"$PACKAGE_NAME"'"/' python/pyproject.toml
|
||||
rm -f python/pyproject.toml.bak
|
||||
grep '^name = ' python/pyproject.toml
|
||||
- name: Build x86_64 Manylinux wheel
|
||||
if: ${{ inputs.arm-build == 'false' }}
|
||||
uses: PyO3/maturin-action@v1
|
||||
@@ -34,7 +54,7 @@ runs:
|
||||
maturin-version: "1.12.4"
|
||||
command: build
|
||||
working-directory: python
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc ${{ inputs.rustflags != '' && format('-e RUSTFLAGS={0}', inputs.rustflags) || '' }}"
|
||||
target: x86_64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
@@ -51,7 +71,7 @@ runs:
|
||||
maturin-version: "1.12.4"
|
||||
command: build
|
||||
working-directory: python
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc ${{ inputs.rustflags != '' && format('-e RUSTFLAGS={0}', inputs.rustflags) || '' }}"
|
||||
target: aarch64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
|
||||
@@ -87,7 +87,7 @@ jobs:
|
||||
bash ci/update_lockfiles.sh --amend
|
||||
- name: Push new version tag
|
||||
if: ${{ !inputs.dry_run }}
|
||||
uses: ad-m/github-push-action@master
|
||||
uses: ad-m/github-push-action@881a6320fdb16eb5318c5054f31c218aec2b324c # v1.3.0
|
||||
with:
|
||||
# Need to use PAT here too to trigger next workflow. See comment above.
|
||||
github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
|
||||
|
||||
@@ -22,7 +22,7 @@ permissions:
|
||||
|
||||
jobs:
|
||||
linux:
|
||||
name: Python ${{ matrix.config.platform }} manylinux${{ matrix.config.manylinux }}
|
||||
name: Python ${{ matrix.config.package_name }} ${{ matrix.config.platform }} manylinux${{ matrix.config.manylinux }}
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
@@ -31,11 +31,28 @@ jobs:
|
||||
manylinux: "2_28"
|
||||
extra_args: "--features fp16kernels"
|
||||
runner: ubuntu-22.04
|
||||
package_name: "lancedb"
|
||||
rustflags: ""
|
||||
# For successful fat LTO builds, we need a large runner to avoid OOM errors.
|
||||
- platform: aarch64
|
||||
manylinux: "2_28"
|
||||
extra_args: "--features fp16kernels"
|
||||
runner: ubuntu-2404-8x-arm64
|
||||
package_name: "lancedb"
|
||||
rustflags: ""
|
||||
# `lancedb-compat`: pre-Haswell-friendly variant for x86_64 hosts
|
||||
# without AVX2 (Sandy Bridge / Ivy Bridge / Westmere on Intel,
|
||||
# Bulldozer / Piledriver / Steamroller on AMD). Compiled at the
|
||||
# `x86-64-v2` baseline; runtime SIMD dispatch in lance-linalg
|
||||
# picks the appropriate tier (scalar / AVX / AVX+FMA / AVX2+FMA
|
||||
# / AVX-512) at load time. Same import as `lancedb` -- conflicts
|
||||
# at install time, so users pick one.
|
||||
- platform: x86_64
|
||||
manylinux: "2_28"
|
||||
extra_args: ""
|
||||
runner: ubuntu-22.04
|
||||
package_name: "lancedb-compat"
|
||||
rustflags: "-Ctarget-cpu=x86-64-v2"
|
||||
runs-on: ${{ matrix.config.runner }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -52,11 +69,13 @@ jobs:
|
||||
args: "--release --strip ${{ matrix.config.extra_args }}"
|
||||
arm-build: ${{ matrix.config.platform == 'aarch64' }}
|
||||
manylinux: ${{ matrix.config.manylinux }}
|
||||
package-name: ${{ matrix.config.package_name }}
|
||||
rustflags: ${{ matrix.config.rustflags }}
|
||||
- uses: actions/upload-artifact@v7
|
||||
if: startsWith(github.ref, 'refs/tags/python-v')
|
||||
with:
|
||||
name: wheels-linux-${{ matrix.config.platform }}-${{ matrix.config.manylinux }}
|
||||
path: target/wheels/lancedb-*.whl
|
||||
name: wheels-linux-${{ matrix.config.package_name }}-${{ matrix.config.platform }}-${{ matrix.config.manylinux }}
|
||||
path: target/wheels/*.whl
|
||||
if-no-files-found: error
|
||||
mac:
|
||||
timeout-minutes: 90
|
||||
@@ -145,7 +164,7 @@ jobs:
|
||||
FURY_TOKEN: ${{ secrets.FURY_TOKEN }}
|
||||
run: |
|
||||
shopt -s nullglob
|
||||
WHEELS=(target/wheels/lancedb-*.whl)
|
||||
WHEELS=(target/wheels/*.whl)
|
||||
if [[ ${#WHEELS[@]} -eq 0 ]]; then
|
||||
echo "No wheels found in target/wheels/" >&2
|
||||
exit 1
|
||||
|
||||
@@ -98,7 +98,7 @@ jobs:
|
||||
cargo build --profile ci --benches --all-features --tests
|
||||
|
||||
linux:
|
||||
timeout-minutes: 30
|
||||
timeout-minutes: 60
|
||||
# To build all features, we need more disk space than is available
|
||||
# on the free OSS github runner. This is mostly due to the the
|
||||
# sentence-transformers feature.
|
||||
@@ -158,7 +158,7 @@ jobs:
|
||||
run: CARGO_ARGS="--profile ci" make -C ./lancedb remote-tests
|
||||
|
||||
macos:
|
||||
timeout-minutes: 30
|
||||
timeout-minutes: 60
|
||||
strategy:
|
||||
matrix:
|
||||
mac-runner: ["macos-14", "macos-15"]
|
||||
|
||||
Generated
+190
-157
@@ -157,9 +157,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "anyhow"
|
||||
version = "1.0.103"
|
||||
version = "1.0.104"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3"
|
||||
checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
|
||||
|
||||
[[package]]
|
||||
name = "approx"
|
||||
@@ -535,13 +535,13 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.89"
|
||||
version = "0.1.91"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb"
|
||||
checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1750,7 +1750,7 @@ version = "3.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34"
|
||||
dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2288,9 +2288,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "997a31e15872606a49478e670c58302094c97cb96abb0a7d60720f8e92170040"
|
||||
checksum = "754ef4e8f073922a26f5b23133b9db4829342362b09be0bc94309cf261c2f098"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-schema",
|
||||
@@ -2335,9 +2335,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-catalog"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f7dd61161508f8f5fa1107774ea687bd753c22d83a32eebf963549f89de14139"
|
||||
checksum = "06afd1e38dd27bbb1258685a1fc6524df6aff4e07b25b393a47de59635178d99"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -2360,9 +2360,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-catalog-listing"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "897c70f871277f9ce99aa38347be0d679bbe3e617156c4d2a8378cec8a2a0891"
|
||||
checksum = "f0668fb32c12065ec242be0e5b4bc62bd7a06a0be3ecd83791ef877e4be67e02"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -2383,9 +2383,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-common"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "121c9ded5d87d9172319e006f2afdb9928d72dbacd6a90a458d8acb1e3b43a65"
|
||||
checksum = "ca43b263cdff57042cfa8fb817fb3469f4878933380dccff25f5e793580abbf9"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -2407,9 +2407,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-common-runtime"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "981b9dae74f78ee3d9f714fb49b01919eab975461b56149510c3ba9ea11287d1"
|
||||
checksum = "05f0ba2b864792bdca4d76c59a1de0ab6e1b61946596b9936888dbd6360035f2"
|
||||
dependencies = [
|
||||
"futures",
|
||||
"log",
|
||||
@@ -2418,9 +2418,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-datasource"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ffd7d295b2ec7c00d8a56562f41ed41062cf0af75549ed891c12a0a09eddfefe"
|
||||
checksum = "b840a8bce0bcbf5afad02946d438591e7c373f7afccaf3d874c04485772514dd"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -2448,9 +2448,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-datasource-arrow"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "552b0b3f342f7ec41b3fbd70f6339dc82a30cfd0349e7f280e7852528085349f"
|
||||
checksum = "a24cc0b9cf6e367f27f27406eff13abf48a11b72446aaa40b3105c0ded5c17d9"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -2472,9 +2472,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-datasource-csv"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68850aa426b897e879c8b87e512ea8124f1d0a2869a4e51808ddaaddf1bc0ada"
|
||||
checksum = "e1abe56b2a7a2d1d6de5117dd1a203181e267f28529faa5da546947621b697d7"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -2495,9 +2495,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-datasource-json"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "402f93242ae08ef99139ee2c528a49d087efe88d5c7b2c3ff5480855a40ce54f"
|
||||
checksum = "9c3e467f0611ad7bdd5aad17c63c9bb6182d04e5282e5496d897ea2b49c024ba"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -2518,15 +2518,15 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-doc"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cb9e7e5d11130c48c8bd4e80c79a9772dd28ce6dc330baca9246205d245b9e2e"
|
||||
checksum = "d69bb69d8769e34f76839c960dbde24c1ac0c885a79b6c3c2287bdc56ec67891"
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-execution"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "37a8643ab852eb68864e1b72ae789e8066282dce48eea6347ffb0aee33d1ccc0"
|
||||
checksum = "d8eac0a09bc8d263f52025cad9e001da4d8138d633fa288edda4d06b1772eae6"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-buffer",
|
||||
@@ -2546,9 +2546,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-expr"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6932f4d71eed9c8d9341476a2b845aadfabde5495d08dbcd8fc23881f49fa7a0"
|
||||
checksum = "eeb14d374767ee0fc62dc79a5ba8bcf8a63c14e993c7d992d0e63adfa23d77d3"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-schema",
|
||||
@@ -2568,9 +2568,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-expr-common"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0225491839a31b1f7d2cb8092c2d50792e2fe1c1724e4e6d08e011f5feaf4ed2"
|
||||
checksum = "b7b19a8c95522bee8cbb313d74263b85e355d2b52f42e67ef5694bf5de9e9356"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2580,9 +2580,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-functions"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "14872c47bfc3d21e53ec82f57074e6987a15941c1e2f43cde4ac6ae2746634e3"
|
||||
checksum = "5f64c983bbbdcb729d921a2b2ac3375598719b5cc0c30345ad664936f3176fc7"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-buffer",
|
||||
@@ -2612,9 +2612,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-functions-aggregate"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "75a2ca14e1b609be21e657e2d3130b2f446456b08393b377bb721a33952d2e09"
|
||||
checksum = "89bc17041e424a47ed062f43df24d84aab8b57c4c3221e5c1a5eef46d6c5718b"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2633,9 +2633,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-functions-aggregate-common"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1ece74ba09092d2ef9c9b54a38445450aea292a1f8b04faf531936b723a24b3c"
|
||||
checksum = "97dd2a9e865c6108059f5b37b77934f84b50bfb108f837bd0e5c9536e03f0545"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2645,9 +2645,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-functions-nested"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f3e3f9ee8ca59bf70518802107de6f1b88a9509efdc629fadc5de9d6b2d5ef5"
|
||||
checksum = "75f0bdfeef16d96417b9632ef855645376b242e9006a126dfd0bedfc54a93f5f"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ord",
|
||||
@@ -2670,9 +2670,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-functions-table"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "89161dffc22cf2b50f9f4b1bee83b5221d3b4ed7c2e37fd7aa2b22a5297b3a26"
|
||||
checksum = "f4e4941673c917819616877e9993da4503e4f4739812be0bc32c5356184c6383"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -2686,9 +2686,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-functions-window"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7339345b226b3874037708bf5023ba1c2de705128f8457a095aae5ae9cb9c78"
|
||||
checksum = "12dd2e16c12b84b6f6b41b19f55b366dd1c46876bb35b86896c6349067379e8d"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2703,9 +2703,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-functions-window-common"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fa84836dc2392df6f43d6a29d37fb56a8ebdc8b3f4e10ae8dc15861fd20278fb"
|
||||
checksum = "4cdc5e4b6f8b6ef823cc1c761f85088ad4c884fe8df64df3cbcc6b2b84698441"
|
||||
dependencies = [
|
||||
"datafusion-common",
|
||||
"datafusion-physical-expr-common",
|
||||
@@ -2713,9 +2713,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-macros"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "587164e03ad68732aa9e7bfe5686e3f25970d4c64fd4bd80790749840892dae5"
|
||||
checksum = "1a3614234dd93578c92428cb4f408e020874f0d2b7e6c90c928d9d28b5df2ceb"
|
||||
dependencies = [
|
||||
"datafusion-doc",
|
||||
"quote",
|
||||
@@ -2724,9 +2724,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-optimizer"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77f20e8cf9e8654d92f4c16b24c487353ee5bf153ffc12d5772cd399ab8cd281"
|
||||
checksum = "a0635620b050b81bb92764e99250868f654e2cd5ad1bece413283b3f73c83179"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"chrono",
|
||||
@@ -2743,9 +2743,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-physical-expr"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f015a4a82f6f7ff7e1d8d4bf3870a936752fa38b17705dfcc14adef95aa8922c"
|
||||
checksum = "8cabf7a86eb70b816729e33c81bf7767c936ee1226f607a114f5dac2decac8d0"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2764,9 +2764,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-physical-expr-adapter"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "51e6ffff8acdfe54e0ea15ccf38115c4a9184433b0439f42907637928d00a235"
|
||||
checksum = "de222e04f7e6744501555a54ab0abe26bfdfebee380af79a9bdc175704246859"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2779,9 +2779,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-physical-expr-common"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7967a3e171c6a4bf09474b3f7a14f1a3db13ed1714ba12156f33fcce2bba54e8"
|
||||
checksum = "72d0d0057fc5a502d45c870cb6d47c66eb7bdd5edb1bd71ad6f3f724975ac2a8"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"chrono",
|
||||
@@ -2796,9 +2796,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-physical-optimizer"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "59ff803e2a96054cb6d83f35f9e60fd4f42eac515e1932bd1b2dbc91d5fcbf36"
|
||||
checksum = "86046eed10950c5f9aaed9acfd148e9bd2e1dfdfe4f9aef607d1447b271e4183"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2814,9 +2814,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-physical-plan"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "776ee54d47d15bdb126452f9ca17b03761e3b004682914beaedd3f86eb507fbc"
|
||||
checksum = "9bc84da934c903407ba297971ebcc020c4c1a38aafd765d6c144c76eff3fa6a1"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-data",
|
||||
@@ -2847,9 +2847,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-pruning"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d5fb9e5774660aa69c3ba93c610f175f75b65cb8c3776edb3626de8f3a4f4ee3"
|
||||
checksum = "eb63eeac6de19be40f487b65dd84e546195f783c5a9928618e0c4f2a3569b0d7"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"datafusion-common",
|
||||
@@ -2863,9 +2863,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-session"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "15ce715fa2a61f4623cc234bcc14a3ef6a91f189128d5b14b468a6a17cdfc417"
|
||||
checksum = "5f961d209177f91bd014db5cbb2c33b7d28a2597b9003e77f17aeb712964315a"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"datafusion-common",
|
||||
@@ -2877,9 +2877,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion-sql"
|
||||
version = "54.0.0"
|
||||
version = "54.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6094ad36a3ed6d7ac87b20b479b2d0b118250f66cf997603829fdc65b44a7099"
|
||||
checksum = "b1d71cb454da682b2af7488e1fc1ddd72ee1b28f19297b8ccad73f0a21ee9a69"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"bigdecimal",
|
||||
@@ -3006,7 +3006,7 @@ dependencies = [
|
||||
"libc",
|
||||
"option-ext",
|
||||
"redox_users",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3229,7 +3229,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3421,8 +3421,9 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "fsst"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6f1155128aba964cf6925c22a667ed763b52c8c653693c4b46f8a79d509a8c1"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"rand 0.9.5",
|
||||
@@ -4218,7 +4219,7 @@ dependencies = [
|
||||
"libc",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.3",
|
||||
"system-configuration",
|
||||
"tokio",
|
||||
"tower-service",
|
||||
@@ -4517,7 +4518,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46"
|
||||
dependencies = [
|
||||
"hermit-abi",
|
||||
"libc",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4611,7 +4612,7 @@ dependencies = [
|
||||
"portable-atomic-util",
|
||||
"serde_core",
|
||||
"wasm-bindgen",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4777,8 +4778,9 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
|
||||
|
||||
[[package]]
|
||||
name = "lance"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "23d04bed056e254bc6e31264b031c8492507ca57939586f016924081dcf221a9"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -4852,8 +4854,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-arrow"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6e7b87c8183988c40a6bd30a6d8ec588b84e53f702b82f19859dc71ba6c02bc"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4875,7 +4878,8 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "771f68b04b47f3addf781116f65061808de94b05e1e9411c23c18f32d14ebe79"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4889,17 +4893,20 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dd47ec33c90bf29f688fd02118e37d3a5ad5c339caa3163f89e417dc0867001f"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
"half",
|
||||
"lance-arrow-scalar",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-bitpacking"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4a4d84a3f36133c70bf89d306d813a29f8eb8555bba2b84995260867cfafdd5e"
|
||||
dependencies = [
|
||||
"arrayref",
|
||||
"crunchy",
|
||||
@@ -4909,8 +4916,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-core"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "238c8a58308e7718d6bd96b53494eb7953fa299778bc4911cc571c3576e9446d"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4948,8 +4956,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datafusion"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "12f0ac4e1cf2f9b1b2fb5ee11bcd04af617bdd6f6cca13726a41c4212220193e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -4979,8 +4988,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datagen"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dd65e7ea88ab28e5d3e91b58a56c02bfd44c47474caae5f8aed1322df1611476"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -4997,8 +5007,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-derive"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf46656b359786f0b73f13936193583972b7af6e53ccb542da87830be18e94b2"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5007,8 +5018,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-encoding"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4725f14fe6cc1b2a5644786b4afa828d0c243064e208fa17378f839243d049c0"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5043,8 +5055,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-file"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8c161b00eb5813f98f1d6d52e06f1b712f9eebb0a7f535a8dc1e1122ceca7307"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5074,8 +5087,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf7980b0a6287fd46b9308a31e2cf284065d5693bfb43336297f0400172670ad"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5141,8 +5155,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index-core"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a95f46ac3e4cdd710ca6929226cfb0dd343d5d2370ecb4b40e1c452043f7d27a"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5164,8 +5179,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-io"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6deecd351ed6849184cc83000eabb87c8a5bfc519996f92451284498af7b42c"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-arith",
|
||||
@@ -5208,8 +5224,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-linalg"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0e5a9b99bd1f49bc2fe5afb81323141abc506c818f589de95dbab4f6143d9a88"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5225,8 +5242,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8624fd33a894b5f2eb411cb56588d29138137ce100dee8e335843828e249b860"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5238,8 +5256,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-impls"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39f0cb9651a65f8eb411d825f34967fff46ba575a76e6ae9becda6ee31c1a974"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -5293,8 +5312,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-select"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bb5c718c141ea4f067203b11a54ccf57f8effbf963f345bc811f4837a660ad57"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5309,8 +5329,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-table"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ffc125611b4fe46f99f00b5262e71e17bd947a3bc3b81f3a4a604cc9ca290b3f"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5349,8 +5370,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-testing"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2598f24fa7333b49fdfc9945dfe084dafef2fc9ac94b5017ab4fbe4cdcd59318"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5363,8 +5385,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-tokenizer"
|
||||
version = "9.1.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.1.0-beta.2#9139486ed371dd6caeb9c708b39fba7f5649a9b0"
|
||||
version = "9.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73060a844ceda2405759b8f0f16a1c586b5b285dbdf7a3784350d048f991cfd3"
|
||||
dependencies = [
|
||||
"icu_segmenter",
|
||||
"jieba-rs",
|
||||
@@ -5377,7 +5400,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb"
|
||||
version = "0.32.0-beta.2"
|
||||
version = "0.33.0-beta.0"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"anyhow",
|
||||
@@ -5432,7 +5455,6 @@ dependencies = [
|
||||
"lance-namespace-impls",
|
||||
"lance-table",
|
||||
"lance-testing",
|
||||
"lazy_static",
|
||||
"log",
|
||||
"metrics",
|
||||
"metrics-util",
|
||||
@@ -5466,7 +5488,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-nodejs"
|
||||
version = "0.32.0-beta.2"
|
||||
version = "0.33.0-beta.0"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5491,7 +5513,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-python"
|
||||
version = "0.35.0-beta.2"
|
||||
version = "0.36.0-beta.0"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5591,9 +5613,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.186"
|
||||
version = "0.2.189"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
|
||||
checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
|
||||
|
||||
[[package]]
|
||||
name = "libloading"
|
||||
@@ -6064,9 +6086,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi"
|
||||
version = "3.10.5"
|
||||
version = "3.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6826e5ddc15589b2d68c8ad5321c18e85d40488e93e32962f362e572669bccf6"
|
||||
checksum = "de33522036981030a75c231829566bc63414e08101a6f5ff4ac6cef19c8e0941"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"chrono",
|
||||
@@ -6089,9 +6111,9 @@ checksum = "c9c366d2c8c60b86fa632df75f745509b52f9128f91a6bad4c796e44abb505e1"
|
||||
|
||||
[[package]]
|
||||
name = "napi-derive"
|
||||
version = "3.5.10"
|
||||
version = "3.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0fe526e81c105d3640516fcde83909dd1afe757c0d7a15af58830b5bc0fb9a1"
|
||||
checksum = "a49c513341a61a16a10af6efcce46b30d0822ba2d4fb197d24d33dfc199c78d5"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"ctor 1.0.5",
|
||||
@@ -6103,9 +6125,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi-derive-backend"
|
||||
version = "5.1.2"
|
||||
version = "6.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "514281397bcddd9ea9a876c7a21a57bff2374237a000ca9a64ea0211ec1993e2"
|
||||
checksum = "4747005fa3e2c9989ac45a723a514c5db2411238b72981a3cda4c701a9dfea17"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"proc-macro2",
|
||||
@@ -6116,9 +6138,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi-sys"
|
||||
version = "3.2.3"
|
||||
version = "3.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73e43cf2eb0bd1bf95a43c07c076ebd2da5d1e015a71c3d201faeffffcc0ecac"
|
||||
checksum = "85fbf1fa9f1babfe396d74bbbf52b3643770243e8f5b0b46715d4caf7f0dfc9a"
|
||||
dependencies = [
|
||||
"libloading",
|
||||
]
|
||||
@@ -6207,7 +6229,7 @@ version = "0.50.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
|
||||
dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -7586,8 +7608,8 @@ version = "0.14.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7"
|
||||
dependencies = [
|
||||
"heck 0.4.1",
|
||||
"itertools 0.11.0",
|
||||
"heck 0.5.0",
|
||||
"itertools 0.14.0",
|
||||
"log",
|
||||
"multimap",
|
||||
"petgraph",
|
||||
@@ -7606,7 +7628,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"itertools 0.11.0",
|
||||
"itertools 0.14.0",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
@@ -7815,7 +7837,7 @@ dependencies = [
|
||||
"quinn-udp",
|
||||
"rustc-hash",
|
||||
"rustls 0.23.40",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.3",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tracing",
|
||||
@@ -7853,9 +7875,9 @@ dependencies = [
|
||||
"cfg_aliases",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"socket2 0.5.10",
|
||||
"socket2 0.6.3",
|
||||
"tracing",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.60.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8161,9 +8183,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "regex"
|
||||
version = "1.13.0"
|
||||
version = "1.13.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2a0e75113e14dc5acb068cd0786884f214f1312650a3d36d269f5c4f3cdee8a2"
|
||||
checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
|
||||
dependencies = [
|
||||
"aho-corasick",
|
||||
"memchr",
|
||||
@@ -8173,9 +8195,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "regex-automata"
|
||||
version = "0.4.14"
|
||||
version = "0.4.16"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f"
|
||||
checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad"
|
||||
dependencies = [
|
||||
"aho-corasick",
|
||||
"memchr",
|
||||
@@ -8629,7 +8651,7 @@ dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8700,7 +8722,7 @@ dependencies = [
|
||||
"security-framework",
|
||||
"security-framework-sys",
|
||||
"webpki-root-certs",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8905,9 +8927,9 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc"
|
||||
|
||||
[[package]]
|
||||
name = "serde"
|
||||
version = "1.0.228"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
|
||||
checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
"serde_derive",
|
||||
@@ -8915,29 +8937,29 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "serde_core"
|
||||
version = "1.0.228"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
|
||||
checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
|
||||
dependencies = [
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive"
|
||||
version = "1.0.228"
|
||||
version = "1.0.229"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
|
||||
checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.150"
|
||||
version = "1.0.151"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
|
||||
checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"memchr",
|
||||
@@ -9268,7 +9290,7 @@ version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451"
|
||||
dependencies = [
|
||||
"heck 0.4.1",
|
||||
"heck 0.5.0",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
@@ -9280,7 +9302,7 @@ version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40"
|
||||
dependencies = [
|
||||
"heck 0.4.1",
|
||||
"heck 0.5.0",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
@@ -9596,6 +9618,17 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "3.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sync_wrapper"
|
||||
version = "1.0.2"
|
||||
@@ -9713,7 +9746,7 @@ dependencies = [
|
||||
"getrandom 0.4.2",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -9916,9 +9949,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.52.3"
|
||||
version = "1.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe"
|
||||
checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"libc",
|
||||
@@ -10403,9 +10436,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.23.5"
|
||||
version = "1.24.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ea5fab0d6c3c01ae70085a09cb03d4c7a1d6314e2b3e075392783396d724ca0a"
|
||||
checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239"
|
||||
dependencies = [
|
||||
"getrandom 0.4.2",
|
||||
"js-sys",
|
||||
@@ -10690,7 +10723,7 @@ version = "0.1.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||
dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
+14
-15
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=9.0.0", default-features = false }
|
||||
lance-core = "=9.0.0"
|
||||
lance-datagen = "=9.0.0"
|
||||
lance-file = "=9.0.0"
|
||||
lance-io = { "version" = "=9.0.0", default-features = false }
|
||||
lance-index = "=9.0.0"
|
||||
lance-linalg = "=9.0.0"
|
||||
lance-namespace = "=9.0.0"
|
||||
lance-namespace-impls = { "version" = "=9.0.0", default-features = false }
|
||||
lance-table = "=9.0.0"
|
||||
lance-testing = "=9.0.0"
|
||||
lance-datafusion = "=9.0.0"
|
||||
lance-encoding = "=9.0.0"
|
||||
lance-arrow = "=9.0.0"
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
@@ -64,7 +64,6 @@ snafu = "0.8"
|
||||
url = "2"
|
||||
num-traits = "0.2"
|
||||
regex = "1.10"
|
||||
lazy_static = "1"
|
||||
semver = "1.0.25"
|
||||
chrono = "0.4"
|
||||
|
||||
|
||||
+165
-30
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<version>0.33.0-beta.0</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
@@ -249,6 +249,57 @@ try (BufferAllocator allocator = new RootAllocator();
|
||||
}
|
||||
```
|
||||
|
||||
### Creating an Empty Table
|
||||
|
||||
To create an empty table, send an Arrow IPC stream that contains the table schema and no record batches.
|
||||
The schema in the IPC stream becomes the table schema, and rows can be inserted later.
|
||||
|
||||
```java
|
||||
import org.lance.namespace.model.CreateTableRequest;
|
||||
import org.lance.namespace.model.CreateTableResponse;
|
||||
import org.apache.arrow.memory.BufferAllocator;
|
||||
import org.apache.arrow.memory.RootAllocator;
|
||||
import org.apache.arrow.vector.VectorSchemaRoot;
|
||||
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
|
||||
import org.apache.arrow.vector.types.FloatingPointPrecision;
|
||||
import org.apache.arrow.vector.types.pojo.ArrowType;
|
||||
import org.apache.arrow.vector.types.pojo.Field;
|
||||
import org.apache.arrow.vector.types.pojo.FieldType;
|
||||
import org.apache.arrow.vector.types.pojo.Schema;
|
||||
|
||||
import java.io.ByteArrayOutputStream;
|
||||
import java.nio.channels.Channels;
|
||||
import java.util.Arrays;
|
||||
|
||||
Schema schema = new Schema(Arrays.asList(
|
||||
new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null),
|
||||
new Field("name", FieldType.nullable(new ArrowType.Utf8()), null),
|
||||
new Field("embedding",
|
||||
FieldType.nullable(new ArrowType.FixedSizeList(128)),
|
||||
Arrays.asList(new Field("item",
|
||||
FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE)),
|
||||
null)))
|
||||
));
|
||||
|
||||
byte[] emptyTableData;
|
||||
try (BufferAllocator allocator = new RootAllocator();
|
||||
VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
|
||||
root.setRowCount(0);
|
||||
|
||||
ByteArrayOutputStream out = new ByteArrayOutputStream();
|
||||
try (ArrowStreamWriter writer = new ArrowStreamWriter(root, null, Channels.newChannel(out))) {
|
||||
writer.start();
|
||||
writer.end();
|
||||
}
|
||||
emptyTableData = out.toByteArray();
|
||||
}
|
||||
|
||||
CreateTableRequest request = new CreateTableRequest();
|
||||
request.setId(Arrays.asList("my_namespace", "empty_table"));
|
||||
|
||||
CreateTableResponse response = namespaceClient.createTable(request, emptyTableData);
|
||||
```
|
||||
|
||||
### Insert
|
||||
|
||||
```java
|
||||
@@ -431,9 +482,88 @@ query.setVector(vector);
|
||||
byte[] result = namespaceClient.queryTable(query);
|
||||
```
|
||||
|
||||
### Reading Query Results
|
||||
## Indexing
|
||||
|
||||
Query results are returned in Apache Arrow IPC file format. Here's how to read them:
|
||||
The Java SDK exposes the REST namespace index operations through the same `LanceNamespace` client.
|
||||
Index creation runs asynchronously, so use `listTableIndices` or `describeTableIndexStats` to check progress.
|
||||
|
||||
### Creating a Vector Index
|
||||
|
||||
```java
|
||||
import org.lance.namespace.model.CreateTableIndexRequest;
|
||||
import org.lance.namespace.model.CreateTableIndexResponse;
|
||||
|
||||
CreateTableIndexRequest request = new CreateTableIndexRequest();
|
||||
request.setId(Arrays.asList("my_namespace", "my_table"));
|
||||
request.setColumn("embedding");
|
||||
request.setIndexType("IVF_PQ");
|
||||
request.setDistanceType("cosine");
|
||||
request.setName("embedding_idx");
|
||||
|
||||
CreateTableIndexResponse response = namespaceClient.createTableIndex(request);
|
||||
System.out.println("Index transaction: " + response.getTransactionId());
|
||||
```
|
||||
|
||||
### Creating a Scalar Index
|
||||
|
||||
```java
|
||||
import org.lance.namespace.model.CreateTableIndexRequest;
|
||||
import org.lance.namespace.model.CreateTableScalarIndexResponse;
|
||||
|
||||
CreateTableIndexRequest request = new CreateTableIndexRequest();
|
||||
request.setId(Arrays.asList("my_namespace", "my_table"));
|
||||
request.setColumn("category");
|
||||
request.setIndexType("BTREE");
|
||||
request.setName("category_idx");
|
||||
|
||||
CreateTableScalarIndexResponse response = namespaceClient.createTableScalarIndex(request);
|
||||
System.out.println("Index transaction: " + response.getTransactionId());
|
||||
```
|
||||
|
||||
### Creating a Full Text Search Index
|
||||
|
||||
```java
|
||||
import org.lance.namespace.model.CreateTableIndexRequest;
|
||||
import org.lance.namespace.model.CreateTableScalarIndexResponse;
|
||||
|
||||
CreateTableIndexRequest request = new CreateTableIndexRequest();
|
||||
request.setId(Arrays.asList("my_namespace", "my_table"));
|
||||
request.setColumn("text_column");
|
||||
request.setIndexType("FTS");
|
||||
request.setName("text_idx");
|
||||
request.setBaseTokenizer("simple");
|
||||
request.setLowerCase(true);
|
||||
request.setWithPosition(true);
|
||||
|
||||
CreateTableScalarIndexResponse response = namespaceClient.createTableScalarIndex(request);
|
||||
System.out.println("Index transaction: " + response.getTransactionId());
|
||||
```
|
||||
|
||||
### Listing Indexes
|
||||
|
||||
```java
|
||||
import org.lance.namespace.model.IndexContent;
|
||||
import org.lance.namespace.model.ListTableIndicesRequest;
|
||||
import org.lance.namespace.model.ListTableIndicesResponse;
|
||||
|
||||
ListTableIndicesRequest request = new ListTableIndicesRequest();
|
||||
request.setId(Arrays.asList("my_namespace", "my_table"));
|
||||
|
||||
ListTableIndicesResponse response = namespaceClient.listTableIndices(request);
|
||||
for (IndexContent index : response.getIndexes()) {
|
||||
System.out.println(index.getIndexName() + ": " + index.getStatus());
|
||||
}
|
||||
```
|
||||
|
||||
!!! note
|
||||
The current Java namespace API exposes index type, index name, distance type, and full text search tokenizer options.
|
||||
IVF training parameters such as `num_partitions` are not exposed by `CreateTableIndexRequest` yet.
|
||||
To make those configurable from Java, the namespace API must add those fields first.
|
||||
|
||||
## Reading Query Results
|
||||
|
||||
Query results are returned as bytes in Apache Arrow IPC file format. Put the byte-channel
|
||||
adapter behind a small helper so query code can work with `ArrowFileReader` directly:
|
||||
|
||||
```java
|
||||
import org.apache.arrow.vector.ipc.ArrowFileReader;
|
||||
@@ -441,45 +571,50 @@ import org.apache.arrow.vector.VectorSchemaRoot;
|
||||
import org.apache.arrow.memory.BufferAllocator;
|
||||
import org.apache.arrow.memory.RootAllocator;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.channels.SeekableByteChannel;
|
||||
|
||||
// Helper class to read Arrow data from byte array
|
||||
class ByteArraySeekableByteChannel implements SeekableByteChannel {
|
||||
private final byte[] data;
|
||||
private long position = 0;
|
||||
private boolean isOpen = true;
|
||||
|
||||
public ByteArraySeekableByteChannel(byte[] data) {
|
||||
this.data = data;
|
||||
final class ArrowIpc {
|
||||
static ArrowFileReader openFileReader(byte[] data, BufferAllocator allocator) throws IOException {
|
||||
return new ArrowFileReader(new ByteArraySeekableByteChannel(data), allocator);
|
||||
}
|
||||
|
||||
@Override
|
||||
public int read(ByteBuffer dst) {
|
||||
int remaining = dst.remaining();
|
||||
int available = (int) (data.length - position);
|
||||
if (available <= 0) return -1;
|
||||
int toRead = Math.min(remaining, available);
|
||||
dst.put(data, (int) position, toRead);
|
||||
position += toRead;
|
||||
return toRead;
|
||||
}
|
||||
private static final class ByteArraySeekableByteChannel implements SeekableByteChannel {
|
||||
private final byte[] data;
|
||||
private long position = 0;
|
||||
private boolean isOpen = true;
|
||||
|
||||
@Override public long position() { return position; }
|
||||
@Override public SeekableByteChannel position(long newPosition) { position = newPosition; return this; }
|
||||
@Override public long size() { return data.length; }
|
||||
@Override public boolean isOpen() { return isOpen; }
|
||||
@Override public void close() { isOpen = false; }
|
||||
@Override public int write(ByteBuffer src) { throw new UnsupportedOperationException(); }
|
||||
@Override public SeekableByteChannel truncate(long size) { throw new UnsupportedOperationException(); }
|
||||
private ByteArraySeekableByteChannel(byte[] data) {
|
||||
this.data = data;
|
||||
}
|
||||
|
||||
@Override
|
||||
public int read(ByteBuffer dst) {
|
||||
int remaining = dst.remaining();
|
||||
int available = (int) (data.length - position);
|
||||
if (available <= 0) return -1;
|
||||
int toRead = Math.min(remaining, available);
|
||||
dst.put(data, (int) position, toRead);
|
||||
position += toRead;
|
||||
return toRead;
|
||||
}
|
||||
|
||||
@Override public long position() { return position; }
|
||||
@Override public SeekableByteChannel position(long newPosition) { position = newPosition; return this; }
|
||||
@Override public long size() { return data.length; }
|
||||
@Override public boolean isOpen() { return isOpen; }
|
||||
@Override public void close() { isOpen = false; }
|
||||
@Override public int write(ByteBuffer src) { throw new UnsupportedOperationException(); }
|
||||
@Override public SeekableByteChannel truncate(long size) { throw new UnsupportedOperationException(); }
|
||||
}
|
||||
}
|
||||
|
||||
// Read query results
|
||||
byte[] queryResult = namespaceClient.queryTable(query);
|
||||
|
||||
try (BufferAllocator allocator = new RootAllocator();
|
||||
ArrowFileReader reader = new ArrowFileReader(
|
||||
new ByteArraySeekableByteChannel(queryResult), allocator)) {
|
||||
ArrowFileReader reader = ArrowIpc.openFileReader(queryResult, allocator)) {
|
||||
|
||||
for (int i = 0; i < reader.getRecordBlocks().size(); i++) {
|
||||
reader.loadRecordBatch(reader.getRecordBlocks().get(i));
|
||||
|
||||
@@ -83,6 +83,24 @@ Delete a branch.
|
||||
|
||||
***
|
||||
|
||||
### diff()
|
||||
|
||||
```ts
|
||||
diff(fromBranch): Promise<BranchDiff>
|
||||
```
|
||||
|
||||
Compare a branch against main without modifying either branch.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **fromBranch**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`BranchDiff`](../interfaces/BranchDiff.md)>
|
||||
|
||||
***
|
||||
|
||||
### list()
|
||||
|
||||
```ts
|
||||
@@ -94,3 +112,28 @@ List all branches, mapping name to branch metadata.
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Record`<`string`, [`BranchContents`](BranchContents.md)>>
|
||||
|
||||
***
|
||||
|
||||
### merge()
|
||||
|
||||
```ts
|
||||
merge(fromBranch, dryRun): Promise<MergeBranchResult>
|
||||
```
|
||||
|
||||
Merge a branch into main.
|
||||
|
||||
Set `dryRun` to `true` to preview the merge. A rejected merge resolves
|
||||
with `status: "rejected"` instead of throwing.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **fromBranch**: `string`
|
||||
Branch to merge from.
|
||||
|
||||
* **dryRun**: `boolean` = `false`
|
||||
When true, only preview the merge. Defaults to false.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`MergeBranchResult`](../interfaces/MergeBranchResult.md)>
|
||||
|
||||
@@ -52,6 +52,11 @@
|
||||
- [AddDataOptions](interfaces/AddDataOptions.md)
|
||||
- [AddResult](interfaces/AddResult.md)
|
||||
- [AlterColumnsResult](interfaces/AlterColumnsResult.md)
|
||||
- [BranchColumnChange](interfaces/BranchColumnChange.md)
|
||||
- [BranchColumnSummary](interfaces/BranchColumnSummary.md)
|
||||
- [BranchDiff](interfaces/BranchDiff.md)
|
||||
- [BranchIndexSummary](interfaces/BranchIndexSummary.md)
|
||||
- [BranchRowCountSummary](interfaces/BranchRowCountSummary.md)
|
||||
- [ClientConfig](interfaces/ClientConfig.md)
|
||||
- [ColumnAlteration](interfaces/ColumnAlteration.md)
|
||||
- [ColumnOrdering](interfaces/ColumnOrdering.md)
|
||||
@@ -86,6 +91,9 @@
|
||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
||||
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
|
||||
- [MergeBlocker](interfaces/MergeBlocker.md)
|
||||
- [MergeBranchResult](interfaces/MergeBranchResult.md)
|
||||
- [MergePreview](interfaces/MergePreview.md)
|
||||
- [MergeResult](interfaces/MergeResult.md)
|
||||
- [NativeOAuthConfig](interfaces/NativeOAuthConfig.md)
|
||||
- [OAuthConfig](interfaces/OAuthConfig.md)
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / BranchColumnChange
|
||||
|
||||
# Interface: BranchColumnChange
|
||||
|
||||
A column whose definition differs between main and the branch.
|
||||
|
||||
## Properties
|
||||
|
||||
### branch
|
||||
|
||||
```ts
|
||||
branch: BranchColumnSummary;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### main
|
||||
|
||||
```ts
|
||||
main: BranchColumnSummary;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### name
|
||||
|
||||
```ts
|
||||
name: string;
|
||||
```
|
||||
@@ -0,0 +1,33 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / BranchColumnSummary
|
||||
|
||||
# Interface: BranchColumnSummary
|
||||
|
||||
Summary of a column in a branch diff.
|
||||
|
||||
## Properties
|
||||
|
||||
### dataType
|
||||
|
||||
```ts
|
||||
dataType: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### name
|
||||
|
||||
```ts
|
||||
name: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### nullable
|
||||
|
||||
```ts
|
||||
nullable: boolean;
|
||||
```
|
||||
@@ -0,0 +1,129 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / BranchDiff
|
||||
|
||||
# Interface: BranchDiff
|
||||
|
||||
Read-only comparison of a branch against main.
|
||||
|
||||
## Properties
|
||||
|
||||
### addedColumns
|
||||
|
||||
```ts
|
||||
addedColumns: BranchColumnSummary[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### addedIndexes
|
||||
|
||||
```ts
|
||||
addedIndexes: BranchIndexSummary[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### baseMoved
|
||||
|
||||
```ts
|
||||
baseMoved: boolean;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### branchVersion
|
||||
|
||||
```ts
|
||||
branchVersion: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### changedColumns
|
||||
|
||||
```ts
|
||||
changedColumns: BranchColumnChange[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### fromBranch
|
||||
|
||||
```ts
|
||||
fromBranch: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### mainVersion
|
||||
|
||||
```ts
|
||||
mainVersion: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### mergeBlockers
|
||||
|
||||
```ts
|
||||
mergeBlockers: MergeBlocker[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### mergeable
|
||||
|
||||
```ts
|
||||
mergeable: boolean;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### parentVersion
|
||||
|
||||
```ts
|
||||
parentVersion: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### removedColumns
|
||||
|
||||
```ts
|
||||
removedColumns: BranchColumnSummary[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### removedIndexes
|
||||
|
||||
```ts
|
||||
removedIndexes: BranchIndexSummary[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### rowCountBranch
|
||||
|
||||
```ts
|
||||
rowCountBranch: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### rowCountMain
|
||||
|
||||
```ts
|
||||
rowCountMain: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### rowSummary
|
||||
|
||||
```ts
|
||||
rowSummary: BranchRowCountSummary;
|
||||
```
|
||||
@@ -0,0 +1,41 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / BranchIndexSummary
|
||||
|
||||
# Interface: BranchIndexSummary
|
||||
|
||||
Summary of an index in a branch diff.
|
||||
|
||||
## Properties
|
||||
|
||||
### columns
|
||||
|
||||
```ts
|
||||
columns: string[];
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### indexName
|
||||
|
||||
```ts
|
||||
indexName: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### indexType?
|
||||
|
||||
```ts
|
||||
optional indexType: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### status
|
||||
|
||||
```ts
|
||||
status: string;
|
||||
```
|
||||
@@ -0,0 +1,57 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / BranchRowCountSummary
|
||||
|
||||
# Interface: BranchRowCountSummary
|
||||
|
||||
Row-level comparison between main and the branch.
|
||||
|
||||
## Properties
|
||||
|
||||
### deltaAvailable
|
||||
|
||||
```ts
|
||||
deltaAvailable: boolean;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### inputsChanged
|
||||
|
||||
```ts
|
||||
inputsChanged: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### newOnBase
|
||||
|
||||
```ts
|
||||
newOnBase: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### newOnBranch
|
||||
|
||||
```ts
|
||||
newOnBranch: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### staleRecompute
|
||||
|
||||
```ts
|
||||
staleRecompute: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### unchanged
|
||||
|
||||
```ts
|
||||
unchanged: number;
|
||||
```
|
||||
@@ -0,0 +1,25 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / MergeBlocker
|
||||
|
||||
# Interface: MergeBlocker
|
||||
|
||||
A reason why a branch cannot currently be merged.
|
||||
|
||||
## Properties
|
||||
|
||||
### code
|
||||
|
||||
```ts
|
||||
code: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### message
|
||||
|
||||
```ts
|
||||
message: string;
|
||||
```
|
||||
@@ -0,0 +1,46 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / MergeBranchResult
|
||||
|
||||
# Interface: MergeBranchResult
|
||||
|
||||
Result of previewing or attempting a branch merge.
|
||||
|
||||
## Properties
|
||||
|
||||
### diff
|
||||
|
||||
```ts
|
||||
diff: BranchDiff;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### mainVersionAfter?
|
||||
|
||||
```ts
|
||||
optional mainVersionAfter: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### preview
|
||||
|
||||
```ts
|
||||
preview: MergePreview;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### status
|
||||
|
||||
```ts
|
||||
status:
|
||||
| "unknown"
|
||||
| "rejected"
|
||||
| "ready"
|
||||
| "notImplemented"
|
||||
| "merged";
|
||||
```
|
||||
@@ -0,0 +1,17 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / MergePreview
|
||||
|
||||
# Interface: MergePreview
|
||||
|
||||
Changes that would be, or were, promoted by a branch merge.
|
||||
|
||||
## Properties
|
||||
|
||||
### promotedColumns
|
||||
|
||||
```ts
|
||||
promotedColumns: string[];
|
||||
```
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<version>0.33.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.32.0-beta.2</version>
|
||||
<version>0.33.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>9.1.0-beta.2</lance-core.version>
|
||||
<lance-core.version>9.0.0</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.32.0-beta.2"
|
||||
version = "0.33.0-beta.0"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
|
||||
@@ -52,6 +52,7 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
Float64,
|
||||
Struct,
|
||||
List,
|
||||
Map_,
|
||||
Int16,
|
||||
Int32,
|
||||
Int64,
|
||||
@@ -69,6 +70,30 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
type Schema = ApacheArrow["Schema"];
|
||||
type Table = ApacheArrow["Table"];
|
||||
|
||||
function expectValidMapField(
|
||||
// biome-ignore lint/suspicious/noExplicitAny: Arrow Field types vary across supported versions
|
||||
field: any,
|
||||
): void {
|
||||
expect(DataType.isMap(field.type)).toBe(true);
|
||||
expect(field.type.keysSorted).toBe(true);
|
||||
expect(field.type.children).toHaveLength(1);
|
||||
|
||||
const entries = field.type.children[0];
|
||||
expect(entries.name).toBe("entries");
|
||||
expect(entries.nullable).toBe(false);
|
||||
expect(DataType.isStruct(entries.type)).toBe(true);
|
||||
expect(entries.type.children).toHaveLength(2);
|
||||
|
||||
const [key, value] = entries.type.children;
|
||||
expect([key.name, value.name]).toEqual(["key", "value"]);
|
||||
expect(key.nullable).toBe(false);
|
||||
expect(DataType.isUtf8(key.type)).toBe(true);
|
||||
expect(value.nullable).toBe(true);
|
||||
expect(DataType.isInt(value.type)).toBe(true);
|
||||
expect(value.type.bitWidth).toBe(32);
|
||||
expect(value.type.isSigned).toBe(true);
|
||||
}
|
||||
|
||||
// Helper method to verify various ways to create a table
|
||||
async function checkTableCreation(
|
||||
tableCreationMethod: (
|
||||
@@ -938,6 +963,34 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
it("will make an empty table with a Map field", async function () {
|
||||
const schema = new Schema([
|
||||
new Field(
|
||||
"attributes",
|
||||
new Map_(
|
||||
new Field(
|
||||
"entries",
|
||||
new Struct([
|
||||
new Field("key", new Utf8(), false),
|
||||
new Field("value", new Int32(), true),
|
||||
]),
|
||||
false,
|
||||
),
|
||||
true,
|
||||
),
|
||||
),
|
||||
]);
|
||||
|
||||
const table = makeEmptyTable(schema);
|
||||
|
||||
expectValidMapField(table.schema.fields[0]);
|
||||
|
||||
const buffer = await fromTableToBuffer(table);
|
||||
const roundTripped = tableFromIPC(buffer);
|
||||
|
||||
expectValidMapField(roundTripped.schema.fields[0]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("when using two versions of arrow", function () {
|
||||
|
||||
@@ -225,6 +225,113 @@ describe("remote connection", () => {
|
||||
);
|
||||
});
|
||||
|
||||
it("diffs and merges remote branches", async () => {
|
||||
const sampleDiff = {
|
||||
fromBranch: "exp",
|
||||
parentVersion: 1,
|
||||
mainVersion: 2,
|
||||
branchVersion: 3,
|
||||
baseMoved: false,
|
||||
rowCountMain: 3,
|
||||
rowCountBranch: 3,
|
||||
rowSummary: {
|
||||
unchanged: 3,
|
||||
newOnBase: 0,
|
||||
newOnBranch: 0,
|
||||
staleRecompute: 0,
|
||||
inputsChanged: 0,
|
||||
deltaAvailable: false,
|
||||
},
|
||||
addedColumns: [{ name: "tag", dataType: "utf8", nullable: true }],
|
||||
removedColumns: [],
|
||||
changedColumns: [],
|
||||
addedIndexes: [],
|
||||
removedIndexes: [],
|
||||
mergeable: true,
|
||||
mergeBlockers: [],
|
||||
};
|
||||
const mergeBodies: Record<string, unknown>[] = [];
|
||||
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
const path = req.url ?? "";
|
||||
if (path.endsWith("/describe/")) {
|
||||
res.writeHead(200, { "Content-Type": "application/json" }).end(
|
||||
JSON.stringify({
|
||||
name: "t",
|
||||
version: 2,
|
||||
schema: { fields: [] },
|
||||
}),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let raw = "";
|
||||
req.on("data", (chunk) => {
|
||||
raw += chunk;
|
||||
});
|
||||
req.on("end", () => {
|
||||
const body = raw ? JSON.parse(raw) : {};
|
||||
if (path.endsWith("/branches/diff/")) {
|
||||
// biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
|
||||
expect(body).toEqual({ from_branch: "exp" });
|
||||
res
|
||||
.writeHead(200, { "Content-Type": "application/json" })
|
||||
.end(JSON.stringify(sampleDiff));
|
||||
return;
|
||||
}
|
||||
if (path.endsWith("/branches/merge/")) {
|
||||
mergeBodies.push(body);
|
||||
const dryRun = body["dry_run"] === true;
|
||||
const response = {
|
||||
status: dryRun ? "ready" : "rejected",
|
||||
diff: dryRun
|
||||
? sampleDiff
|
||||
: {
|
||||
...sampleDiff,
|
||||
mergeable: false,
|
||||
mergeBlockers: [
|
||||
{ code: "baseMoved", message: "main has advanced" },
|
||||
],
|
||||
},
|
||||
preview: { promotedColumns: dryRun ? ["tag"] : [] },
|
||||
};
|
||||
res
|
||||
.writeHead(dryRun ? 200 : 409, {
|
||||
"Content-Type": "application/json",
|
||||
})
|
||||
.end(JSON.stringify(response));
|
||||
return;
|
||||
}
|
||||
res.writeHead(404).end();
|
||||
});
|
||||
},
|
||||
async (db) => {
|
||||
const table = await db.openTable("t");
|
||||
const branches = await table.branches();
|
||||
|
||||
await expect(branches.diff("exp")).resolves.toEqual(sampleDiff);
|
||||
|
||||
const rejected = await branches.merge("exp");
|
||||
expect(rejected.status).toBe("rejected");
|
||||
expect(rejected.diff.mergeBlockers).toEqual([
|
||||
{ code: "baseMoved", message: "main has advanced" },
|
||||
]);
|
||||
|
||||
const preview = await branches.merge("exp", true);
|
||||
expect(preview.status).toBe("ready");
|
||||
expect(preview.preview.promotedColumns).toEqual(["tag"]);
|
||||
},
|
||||
);
|
||||
|
||||
expect(mergeBodies).toEqual([
|
||||
// biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
|
||||
{ from_branch: "exp", dry_run: false },
|
||||
// biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format
|
||||
{ from_branch: "exp", dry_run: true },
|
||||
]);
|
||||
});
|
||||
|
||||
describe("TlsConfig", () => {
|
||||
it("should create TlsConfig with all fields", () => {
|
||||
const tlsConfig: TlsConfig = {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as arrow from "../lancedb/arrow";
|
||||
import { sanitizeField, sanitizeType } from "../lancedb/sanitize";
|
||||
import { sanitizeField, sanitizeMap, sanitizeType } from "../lancedb/sanitize";
|
||||
|
||||
describe("sanitize", function () {
|
||||
describe("sanitizeType function", function () {
|
||||
@@ -181,4 +181,15 @@ describe("sanitize", function () {
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("sanitizeMap function", function () {
|
||||
it.each([
|
||||
["no children", []],
|
||||
["two children", [{}, {}]],
|
||||
])("should reject a Map type with %s", function (_, children) {
|
||||
expect(() => sanitizeMap({ children, keysSorted: false })).toThrow(
|
||||
"Expected a Map type to have exactly one child",
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -124,6 +124,14 @@ export {
|
||||
export {
|
||||
Table,
|
||||
Branches,
|
||||
BranchColumnSummary,
|
||||
BranchColumnChange,
|
||||
BranchIndexSummary,
|
||||
BranchRowCountSummary,
|
||||
MergeBlocker,
|
||||
BranchDiff,
|
||||
MergePreview,
|
||||
MergeBranchResult,
|
||||
AddDataOptions,
|
||||
UpdateOptions,
|
||||
OptimizeOptions,
|
||||
|
||||
@@ -288,12 +288,11 @@ export function sanitizeMap(typeLike: object) {
|
||||
if (!("keysSorted" in typeLike) || typeof typeLike.keysSorted !== "boolean") {
|
||||
throw Error("Expected a Map type to have a `keysSorted` property");
|
||||
}
|
||||
if (typeLike.children.length !== 1) {
|
||||
throw Error("Expected a Map type to have exactly one child");
|
||||
}
|
||||
|
||||
return new Map_(
|
||||
// biome-ignore lint/suspicious/noExplicitAny: skip
|
||||
typeLike.children.map((field) => sanitizeField(field)) as any,
|
||||
typeLike.keysSorted,
|
||||
);
|
||||
return new Map_(sanitizeField(typeLike.children[0]), typeLike.keysSorted);
|
||||
}
|
||||
|
||||
export function sanitizeDuration(typeLike: object) {
|
||||
|
||||
@@ -1329,6 +1329,76 @@ export interface FieldMetadataUpdate {
|
||||
replace?: boolean;
|
||||
}
|
||||
|
||||
/** Summary of a column in a branch diff. */
|
||||
export interface BranchColumnSummary {
|
||||
name: string;
|
||||
dataType: string;
|
||||
nullable: boolean;
|
||||
}
|
||||
|
||||
/** A column whose definition differs between main and the branch. */
|
||||
export interface BranchColumnChange {
|
||||
name: string;
|
||||
main: BranchColumnSummary;
|
||||
branch: BranchColumnSummary;
|
||||
}
|
||||
|
||||
/** Summary of an index in a branch diff. */
|
||||
export interface BranchIndexSummary {
|
||||
indexName: string;
|
||||
columns: string[];
|
||||
indexType?: string;
|
||||
status: string;
|
||||
}
|
||||
|
||||
/** Row-level comparison between main and the branch. */
|
||||
export interface BranchRowCountSummary {
|
||||
unchanged: number;
|
||||
newOnBase: number;
|
||||
newOnBranch: number;
|
||||
staleRecompute: number;
|
||||
inputsChanged: number;
|
||||
deltaAvailable: boolean;
|
||||
}
|
||||
|
||||
/** A reason why a branch cannot currently be merged. */
|
||||
export interface MergeBlocker {
|
||||
code: string;
|
||||
message: string;
|
||||
}
|
||||
|
||||
/** Read-only comparison of a branch against main. */
|
||||
export interface BranchDiff {
|
||||
fromBranch: string;
|
||||
parentVersion: number;
|
||||
mainVersion: number;
|
||||
branchVersion: number;
|
||||
baseMoved: boolean;
|
||||
rowCountMain: number;
|
||||
rowCountBranch: number;
|
||||
rowSummary: BranchRowCountSummary;
|
||||
addedColumns: BranchColumnSummary[];
|
||||
removedColumns: BranchColumnSummary[];
|
||||
changedColumns: BranchColumnChange[];
|
||||
addedIndexes: BranchIndexSummary[];
|
||||
removedIndexes: BranchIndexSummary[];
|
||||
mergeable: boolean;
|
||||
mergeBlockers: MergeBlocker[];
|
||||
}
|
||||
|
||||
/** Changes that would be, or were, promoted by a branch merge. */
|
||||
export interface MergePreview {
|
||||
promotedColumns: string[];
|
||||
}
|
||||
|
||||
/** Result of previewing or attempting a branch merge. */
|
||||
export interface MergeBranchResult {
|
||||
status: "ready" | "rejected" | "notImplemented" | "merged" | "unknown";
|
||||
diff: BranchDiff;
|
||||
preview: MergePreview;
|
||||
mainVersionAfter?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Branch manager for a {@link Table}.
|
||||
*
|
||||
@@ -1381,4 +1451,28 @@ export class Branches {
|
||||
async delete(name: string): Promise<void> {
|
||||
return await this.#inner.delete(name);
|
||||
}
|
||||
|
||||
/** Compare a branch against main without modifying either branch. */
|
||||
async diff(fromBranch: string): Promise<BranchDiff> {
|
||||
return (await this.#inner.diff(fromBranch)) as unknown as BranchDiff;
|
||||
}
|
||||
|
||||
/**
|
||||
* Merge a branch into main.
|
||||
*
|
||||
* Set `dryRun` to `true` to preview the merge. A rejected merge resolves
|
||||
* with `status: "rejected"` instead of throwing.
|
||||
*
|
||||
* @param fromBranch Branch to merge from.
|
||||
* @param dryRun When true, only preview the merge. Defaults to false.
|
||||
*/
|
||||
async merge(
|
||||
fromBranch: string,
|
||||
dryRun: boolean = false,
|
||||
): Promise<MergeBranchResult> {
|
||||
return (await this.#inner.merge(
|
||||
fromBranch,
|
||||
dryRun,
|
||||
)) as unknown as MergeBranchResult;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.33.0-beta.0",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.33.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.32.0-beta.2",
|
||||
"version": "0.33.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.32.0-beta.2",
|
||||
"version": "0.33.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.32.0-beta.2",
|
||||
"version": "0.33.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.32.0-beta.2",
|
||||
"version": "0.33.0-beta.0",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.33.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.32.0-beta.2",
|
||||
"version": "0.33.0-beta.0",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.33.0-beta.0",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
|
||||
+1
-1
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.33.0-beta.0",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
|
||||
@@ -232,6 +232,10 @@ impl From<ClientConfig> for lancedb::remote::ClientConfig {
|
||||
tls_config: config.tls_config.map(Into::into),
|
||||
header_provider: None, // the header provider is set separately later
|
||||
user_id: config.user_id,
|
||||
// Resolved from LANCE_CLIENT_MAX_BYTES_PER_REQUEST or the default.
|
||||
max_bytes_per_request: None,
|
||||
// Resolved from LANCE_CLIENT_MAX_REQUEST_DURATION or the read timeout.
|
||||
max_request_duration: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+25
-1
@@ -165,7 +165,7 @@ impl Table {
|
||||
if let Some(train) = train {
|
||||
builder = builder.train(train);
|
||||
}
|
||||
builder.execute().await.default_error().map(|_| ())
|
||||
builder.execute().await.default_error()
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
@@ -1355,4 +1355,28 @@ impl Branches {
|
||||
pub async fn delete(&self, name: String) -> napi::Result<()> {
|
||||
self.inner.delete_branch(&name).await.default_error()
|
||||
}
|
||||
|
||||
#[napi(ts_return_type = "Promise<Record<string, unknown>>")]
|
||||
pub async fn diff(&self, from_branch: String) -> napi::Result<serde_json::Value> {
|
||||
let diff = self.inner.diff_branch(&from_branch).await.default_error()?;
|
||||
serde_json::to_value(diff).map_err(|err| {
|
||||
napi::Error::from_reason(format!("failed to serialize branch diff: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[napi(ts_return_type = "Promise<Record<string, unknown>>")]
|
||||
pub async fn merge(
|
||||
&self,
|
||||
from_branch: String,
|
||||
dry_run: Option<bool>,
|
||||
) -> napi::Result<serde_json::Value> {
|
||||
let result = self
|
||||
.inner
|
||||
.merge_branch(&from_branch, dry_run.unwrap_or(false))
|
||||
.await
|
||||
.default_error()?;
|
||||
serde_json::to_value(result).map_err(|err| {
|
||||
napi::Error::from_reason(format!("failed to serialize branch merge result: {err}"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"name": "lancedb",
|
||||
"description": "Write, review, debug, and document LanceDB pipelines in Python and TypeScript that work across local LanceDB OSS tables and remote LanceDB Enterprise/Cloud tables, with idiomatic query/search patterns and performance defaults for ingestion, indexing, filtering, and diagnostics.",
|
||||
"version": "0.1.0",
|
||||
"author": {
|
||||
"name": "LanceDB"
|
||||
},
|
||||
"homepage": "https://www.lancedb.com",
|
||||
"keywords": [
|
||||
"lancedb",
|
||||
"vector-search",
|
||||
"full-text-search",
|
||||
"hybrid-search",
|
||||
"python",
|
||||
"typescript",
|
||||
"pipelines",
|
||||
"ingestion",
|
||||
"indexing",
|
||||
"performance"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
{
|
||||
"name": "lancedb",
|
||||
"version": "0.1.0",
|
||||
"description": "Codex plugin for building LanceDB pipelines in Python and TypeScript.",
|
||||
"author": {
|
||||
"name": "LanceDB"
|
||||
},
|
||||
"keywords": [
|
||||
"lancedb",
|
||||
"vector-search",
|
||||
"full-text-search",
|
||||
"hybrid-search",
|
||||
"python",
|
||||
"typescript",
|
||||
"pipelines"
|
||||
],
|
||||
"skills": "./skills/",
|
||||
"interface": {
|
||||
"displayName": "LanceDB",
|
||||
"shortDescription": "Build LanceDB pipelines in Python and TypeScript.",
|
||||
"longDescription": "Write, review, debug, and document LanceDB pipelines in Python and TypeScript that work across local LanceDB OSS tables and remote LanceDB Enterprise/Cloud tables, with idiomatic query/search patterns and performance defaults for ingestion, indexing, filtering, and diagnostics.",
|
||||
"developerName": "LanceDB",
|
||||
"websiteURL": "https://www.lancedb.com",
|
||||
"category": "Developer Tools",
|
||||
"capabilities": [
|
||||
"Developer Tools"
|
||||
],
|
||||
"defaultPrompt": "Create a LanceDB table, embed sample text, and run a vector search.",
|
||||
"composerIcon": "./assets/logo.png",
|
||||
"logo": "./assets/logo.png",
|
||||
"logoDark": "./assets/logo-dark.png"
|
||||
}
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 10 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 35 KiB |
@@ -1,6 +1,6 @@
|
||||
---
|
||||
name: lancedb
|
||||
description: Use when writing, reviewing, debugging, or documenting LanceDB pipelines in Python or TypeScript, especially code that should work across local LanceDB OSS tables and remote LanceDB Enterprise/Cloud tables. Helps avoid non-portable full-table materialization, choose idiomatic query/search patterns, and apply LanceDB performance defaults for ingestion, indexing, filtering, and diagnostics.
|
||||
description: Use when writing, reviewing, debugging, or documenting LanceDB pipelines in Python or TypeScript, especially code that should work across local LanceDB OSS tables and remote LanceDB Enterprise/Cloud tables. Helps avoid non-portable full-table materialization, choose idiomatic query/search patterns, apply LanceDB performance defaults for ingestion, indexing, filtering, and diagnostics, and resolve connections to the remote server for Enterprise-only operations such as jobs.
|
||||
---
|
||||
|
||||
# Building LanceDB Pipelines
|
||||
@@ -19,7 +19,7 @@ Do NOT assume local-only table helpers exist on remote tables. If the user asks
|
||||
## Workflow
|
||||
|
||||
1. Identify the SDK: Python, TypeScript, or both.
|
||||
2. Identify the table mode: local/embedded OSS, remote Enterprise/Cloud, or portable across both. If the user says "LanceDB Enterprise", choose the remote table path.
|
||||
2. Identify the table mode: local/embedded OSS, remote Enterprise/Cloud, or portable across both. If the user says "LanceDB Enterprise", choose the remote table path. If the task involves jobs in any way (listing, inspecting, creating, or canceling jobs), it is always the remote path and requires a remote server connection — see "Connecting to the LanceDB remote server" below before doing anything else.
|
||||
3. Read the matching language branch before writing or changing code:
|
||||
- Python patterns: `references/python/patterns.md`
|
||||
- Python API quick reference: `references/python/api_reference.md`
|
||||
@@ -29,7 +29,9 @@ Do NOT assume local-only table helpers exist on remote tables. If the user asks
|
||||
- TypeScript performance guidance: `references/typescript/performance.md`
|
||||
- Column metadata authoring (both SDKs): `references/column_metadata.md`
|
||||
- Branch operations (both SDKs): `references/branch_ops.md`
|
||||
4. Start with `patterns.md` for the selected SDK. Read `api_reference.md` when choosing method names or return collectors. Read `performance.md` when the task involves ingestion, indexing, filtering, query tuning, diagnostics, or large datasets. Read `column_metadata.md` when the task is documenting, tagging, classifying, or grouping table columns (field descriptions, `lancedb:tag:*` tags, logical column families). Read `branch_ops.md` when the task involves branch lifecycle (list/create/delete), writing to a non-main branch, or verifying a change stayed off main.
|
||||
- Remote server connection resolution (jobs, raw REST): `references/remote_connect.md`
|
||||
- Job operations REST API (list/describe/cancel/query_events): `references/remote_jobs.md`
|
||||
4. Start with `patterns.md` for the selected SDK. Read `api_reference.md` when choosing method names or return collectors. Read `performance.md` when the task involves ingestion, indexing, filtering, query tuning, diagnostics, or large datasets. Read `column_metadata.md` when the task is documenting, tagging, classifying, or grouping table columns (field descriptions, `lancedb:tag:*` tags, logical column families). Read `branch_ops.md` when the task involves branch lifecycle (list/create/delete), writing to a non-main branch, or verifying a change stayed off main. Read `remote_connect.md` when the task involves jobs or direct REST access to an Enterprise deployment, and `remote_jobs.md` for the job REST methods themselves (list, describe, cancel, query_events).
|
||||
5. For Python schemas, favor Pydantic models and validate records before writing. Use PyArrow schemas when Arrow-native, streaming, or highly dynamic data makes them materially better suited.
|
||||
6. Prefer `search()` or `query()` builders with explicit `select()` and `limit()` for reads.
|
||||
7. Avoid table-level full materialization in remote or portable code. This is the main local-vs-remote read pitfall.
|
||||
@@ -70,6 +72,10 @@ Rules for portable Enterprise ingestion:
|
||||
|
||||
This is Enterprise/Cloud-specific. Local/OSS tables have no separate data plane, so `mode="overwrite"` and immediate same-name reuse are fine there.
|
||||
|
||||
## Connecting to the LanceDB remote server
|
||||
|
||||
LanceDB Enterprise/Cloud deployments are served by a server implementing the lance-namespace OpenAPI spec (<https://github.com/lance-format/lance-namespace/blob/main/docs/src/spec.yaml>). Every remote (`db://...`) connection talks to such a server, and some operations exist only there. In particular, **all operations around jobs (listing, inspecting, creating, or canceling jobs) run server-side** — there is no local/OSS equivalent. Before any job work, or any direct REST call to an Enterprise deployment, read `references/remote_connect.md` to resolve the base URL, credentials, and database header and to validate the connection. Then use the four job REST methods documented in `references/remote_jobs.md` (list, describe, cancel, query_events).
|
||||
|
||||
## Script
|
||||
|
||||
Run the scanner when reviewing or modifying an existing codebase:
|
||||
@@ -0,0 +1,6 @@
|
||||
interface:
|
||||
display_name: "LanceDB"
|
||||
short_description: "Build LanceDB pipelines in Python and TypeScript"
|
||||
default_prompt: "Use $lancedb to create a table, embed sample text, and run a vector search."
|
||||
icon_small: "./assets/icon.png"
|
||||
icon_large: "./assets/icon.png"
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 23 KiB |
+66
-1
@@ -2,7 +2,7 @@
|
||||
|
||||
Manage branches on a LanceDB table: list what exists, create new ones, delete stale ones, and direct read/write operations at a specific branch without touching main. Use for branch lifecycle tasks, experimental/isolated table versions, targeting an operation at a non-main branch, or confirming a mutation did not affect main.
|
||||
|
||||
Works on local/OSS and remote Enterprise/Cloud tables.
|
||||
Works on local/OSS and remote Enterprise/Cloud tables, except merging a branch into main, which is Enterprise-only.
|
||||
|
||||
## The branch model (important)
|
||||
|
||||
@@ -101,6 +101,69 @@ assert b"lancedb:description" not in (table.schema.field("category").metadata or
|
||||
|
||||
Two handles on the same branch see each other's writes (e.g. `table.branches.create("exp")` and `db.open_table(name, branch="exp")`); main stays isolated.
|
||||
|
||||
## Merging a branch into main (Enterprise only)
|
||||
|
||||
Merge is available through the SDKs (`table.branches.merge(...)`) on **Enterprise tables only** — it is not supported on Cloud or local/OSS tables, which raise `NotSupported`.
|
||||
|
||||
`merge` takes the branch to merge **from** and a `dry_run` flag. Both the SDK method and the underlying REST endpoint **actually merge by default** (`dry_run=False`); pass `dry_run=True` to only preview. A rejected merge is **not an exception** — it returns a result with `status="rejected"` rather than raising, so inspect the return value. Use `branches.diff(from_branch)` to inspect a branch's pending diff without attempting a merge.
|
||||
|
||||
```python
|
||||
exp = "experiment-reindex"
|
||||
|
||||
# preview only — returns status="ready" if it would merge cleanly
|
||||
preview = table.branches.merge(exp, dry_run=True)
|
||||
|
||||
# actually merge (default)
|
||||
result = table.branches.merge(exp)
|
||||
if result["status"] == "merged":
|
||||
print("landed at", result["mainVersionAfter"])
|
||||
elif result["status"] == "rejected":
|
||||
print(result["diff"]["mergeBlockers"]) # why it was refused
|
||||
|
||||
# inspect a branch's pending diff without merging
|
||||
diff = table.branches.diff(exp)
|
||||
```
|
||||
|
||||
Async: `await table.branches.merge(exp)`, `await table.branches.diff(exp)`.
|
||||
|
||||
```typescript
|
||||
const branches = await table.branches();
|
||||
const exp = "experiment-reindex";
|
||||
|
||||
// preview only (second arg is dryRun)
|
||||
const preview = await branches.merge(exp, true);
|
||||
|
||||
// actually merge (default)
|
||||
const result = await branches.merge(exp);
|
||||
if (result.status === "merged") {
|
||||
console.log("landed at", result.mainVersionAfter);
|
||||
} else if (result.status === "rejected") {
|
||||
console.log(result.diff.mergeBlockers);
|
||||
}
|
||||
|
||||
const diff = await branches.diff(exp);
|
||||
```
|
||||
|
||||
The result is the wire JSON, containing `status` (`ready` on a passing dry run, `merged` on success, `rejected` when refused — also `notImplemented`/`unknown`), the branch `diff` (including `mergeBlockers` explaining any rejection), a `preview` of the columns that would be promoted, and — after a real merge — `mainVersionAfter`.
|
||||
|
||||
### Merge preconditions
|
||||
|
||||
Merge only **promotes newly added columns** onto main; it does not replay arbitrary commits. Practically, a branch is mergeable only if it has **exactly one commit since it was created, and that commit added a column**. The merge is rejected (`status: "rejected"`, with `mergeBlockers` set) if:
|
||||
|
||||
- the branch was forked from another branch rather than directly from main
|
||||
- main has advanced since the branch was forked
|
||||
- the branch's rows changed since the fork (row counts must match main exactly)
|
||||
- the branch removed columns or changed a column's type/nullability
|
||||
- the branch added no columns (index-only changes are not merged)
|
||||
|
||||
### Adding a column in a single commit
|
||||
|
||||
Because the branch must contain just one column-adding commit, add the column with its values in one operation rather than add-then-backfill:
|
||||
|
||||
1. **SQL transformation** — `add_columns` with a SQL expression computed from existing columns, so the column lands populated in one commit.
|
||||
2. **Precompute the values** — compute the column's values externally, then add the fully-populated column in a single operation (e.g. via `merge_insert`/`add_columns` with the data ready).
|
||||
3. **Lance-format-level data evolution (pylance)** — use Lance's data evolution with backfill, documented at <https://lance.org/guide/data_evolution/#with-data-backfill>.
|
||||
|
||||
## Quick reference
|
||||
|
||||
| Goal | Python | TypeScript |
|
||||
@@ -113,5 +176,7 @@ Two handles on the same branch see each other's writes (e.g. `table.branches.cre
|
||||
| Delete branch | `table.branches.delete(name)` | `await branches.delete(name)` |
|
||||
| Which branch is this handle on? | `table.current_branch()` (`None` = main) | `table.currentBranch()` (`null` = main) |
|
||||
| Target main | use the original (non-branch) handle | use the original (non-branch) handle |
|
||||
| Merge branch into main (Enterprise only) | `table.branches.merge(from_branch, dry_run=False)` | `await branches.merge(fromBranch, dryRun)` |
|
||||
| Preview a branch's pending diff (Enterprise only) | `table.branches.diff(from_branch)` | `await branches.diff(fromBranch)` |
|
||||
|
||||
Branch names must be non-empty; empty names raise a validation error.
|
||||
+9
-2
@@ -4,12 +4,19 @@ Quick method reference for Python LanceDB code. Cross-check source for non-trivi
|
||||
|
||||
## Connect
|
||||
|
||||
If you're connecting to a remote database, use this:
|
||||
```python
|
||||
import lancedb
|
||||
|
||||
db = lancedb.connect("./camelot-db") # local/OSS
|
||||
db = lancedb.connect("db://my-db", api_key=api_key, region=region) # remote
|
||||
db = lancedb.connect("db://my-db", api_key=api_key, host_override=host_override) # remote
|
||||
```
|
||||
(values may be found in LANCEDB_API_KEY and LANCEDB_HOST_OVERRIDE, either in env vars or a .env file)
|
||||
|
||||
If you're connecting to a local table using OSS LanceDB, use this:
|
||||
```python
|
||||
db = lancedb.connect("./camelot-db") # local/OSS
|
||||
```
|
||||
If you're not sure which, or if you can't find the api_key or host_override params, ask the user.
|
||||
|
||||
**Place the local database directory next to the script/entrypoint that opens it** (i.e. resolve the path relative to the script, `Path(__file__).parent / "camelot-db"`), not buried under a shared `data/` folder. The Lance dataset is the database, not a data file — keeping it beside its code makes ownership obvious and paths stable regardless of the working directory the script is launched from.
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
# Connecting to a LanceDB remote server
|
||||
|
||||
LanceDB Enterprise/Cloud deployments are served by a server implementing the
|
||||
lance-namespace OpenAPI spec
|
||||
(<https://github.com/lance-format/lance-namespace/blob/main/docs/src/spec.yaml>).
|
||||
Every remote (`db://...`) connection talks to such a server, and some operations
|
||||
exist only there. In particular, all operations around jobs (listing, inspecting,
|
||||
creating, or canceling jobs) run server-side — there is no local/OSS equivalent, so
|
||||
resolve a server connection before attempting any job work. The job REST methods
|
||||
themselves are documented in `references/remote_jobs.md`.
|
||||
|
||||
Every request needs two things:
|
||||
|
||||
1. **Base URL** — the server endpoint
|
||||
2. **Credentials** — an API key (`x-api-key` header over REST), and usually a database name (`x-lancedb-database` header)
|
||||
|
||||
## Resolution steps
|
||||
|
||||
1. If the user already gave a URL and API key (or said which environment they're working against), use that.
|
||||
2. Otherwise, look for credentials already available in the environment:
|
||||
- Env vars like `LANCEDB_URI` / `LANCEDB_HOST` / `LANCEDB_API_KEY`
|
||||
- A server endpoint already running or port-forwarded locally (the REST default port is 2333, i.e. `http://localhost:2333`)
|
||||
3. If you didn't find both pieces, ask the user directly: **"What's your LanceDB endpoint's URL, and what's your API key?"** Also ask which database to use if it isn't obvious. Don't guess or probe further — the user knows their deployment.
|
||||
|
||||
## Validating the connection
|
||||
|
||||
Make a cheap authenticated request and check the status before starting real work:
|
||||
|
||||
```bash
|
||||
curl -s -w "\n%{http_code}" "{base_url}/v1/table/?limit=1" \
|
||||
-H "x-api-key: <key>" \
|
||||
-H "x-lancedb-database: <database>"
|
||||
```
|
||||
|
||||
- `200` — connection, key, and database header all good
|
||||
- `401` — API key missing or wrong
|
||||
- `400` mentioning a database header — this deployment expects `x-lancedb-database`
|
||||
|
||||
## Non-REST equivalents
|
||||
|
||||
The same credentials work through the SDKs and CLI:
|
||||
|
||||
- Python SDK: `lancedb.connect("db://<database>", api_key="<key>", host_override="<base_url>")`
|
||||
- TypeScript SDK: `await lancedb.connect("db://<database>", { apiKey: "<key>", hostOverride: "<base_url>" })`
|
||||
- `lancedb` CLI: a `[profiles.<name>]` entry in `~/.lancedb/config.toml` with `http_server_url`, `api_key`, `database`
|
||||
@@ -0,0 +1,151 @@
|
||||
# Job operations over the LanceDB remote server REST API
|
||||
|
||||
Jobs are server-side background operations on LanceDB Enterprise/Cloud — index builds,
|
||||
column backfills, materialized view refreshes, and similar async work. Endpoints that
|
||||
trigger async work (e.g. the column backfill or materialized view refresh endpoints)
|
||||
return a `job_id`; these four methods are how you track and manage those jobs.
|
||||
|
||||
Resolve the connection first — see `references/remote_connect.md`. All four methods
|
||||
are **POST** requests under `{base_url}/v1/jobs/` with JSON bodies, and take the usual
|
||||
`x-api-key` / `x-lancedb-database` headers. If every job call returns `501`, job APIs
|
||||
are disabled on that deployment (the server has no job registry configured) — report
|
||||
that rather than retrying.
|
||||
|
||||
## 1. List jobs — `POST /v1/jobs/list`
|
||||
|
||||
The body is optional; an empty body lists everything. All fields are filters:
|
||||
|
||||
```json
|
||||
{
|
||||
"limit": 100,
|
||||
"table_name": "my_table",
|
||||
"job_type": "...",
|
||||
"job_subtype": "...",
|
||||
"state": "...",
|
||||
"page_token": "..."
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
curl -s -X POST "{base_url}/v1/jobs/list" \
|
||||
-H "x-api-key: <key>" -H "x-lancedb-database: <database>" \
|
||||
-H "content-type: application/json" \
|
||||
-d '{"table_name": "my_table"}'
|
||||
```
|
||||
|
||||
Response:
|
||||
|
||||
```json
|
||||
{
|
||||
"jobs": [
|
||||
{
|
||||
"job_id": "...",
|
||||
"table": "my_table",
|
||||
"job_type": "...",
|
||||
"job_subtype": "...",
|
||||
"state": "done",
|
||||
"created_at_millis": 1720000000000
|
||||
}
|
||||
],
|
||||
"page_token": "..."
|
||||
}
|
||||
```
|
||||
|
||||
A `page_token` in the response means there are more results — pass it back in the next
|
||||
request to continue. Note list rows use a lowercase `state` string, while describe uses
|
||||
an uppercase `job_state`.
|
||||
|
||||
## 2. Describe a job — `POST /v1/jobs/describe`
|
||||
|
||||
Body: `{"job_id": "<id>"}`. Returns full detail for one job:
|
||||
|
||||
```json
|
||||
{
|
||||
"job_id": "...",
|
||||
"job_type": "...",
|
||||
"job_subtype": "...",
|
||||
"job_state": "IN_PROGRESS",
|
||||
"creation_ms": 1720000000000,
|
||||
"spec": {},
|
||||
"status": {}
|
||||
}
|
||||
```
|
||||
|
||||
`job_state` is one of `IN_PROGRESS`, `CANCELLED`, `FAILED`, `DONE`. `spec` and `status`
|
||||
are job-type-specific JSON objects (the job's input specification and its current
|
||||
progress/status). Returns `404` for an unknown job id.
|
||||
|
||||
## 3. Cancel a job — `POST /v1/jobs/cancel`
|
||||
|
||||
Body: `{"job_id": "<id>"}`; response echoes `{"job_id": "<id>"}`. Cancellation is a
|
||||
service-level operation requiring the same administrative authorization as the
|
||||
`/admin` routes — a database-scoped API key that can list and describe jobs may still
|
||||
get a permission error here. Other errors: `404` unknown job, `409` state conflict
|
||||
(e.g. already in a terminal state), `429` too much write contention (safe to retry).
|
||||
|
||||
## 4. Query job event history — `POST /v1/jobs/query_events`
|
||||
|
||||
Returns the event history (state transitions, progress updates) for one or more jobs.
|
||||
Body: `{"job_id": "<id>"}` for one job, or `{"job_ids": ["<id>", ...]}` for a batch.
|
||||
Optional fields: `limit` (max event rows), `limit_per_job` (per job in a batch query),
|
||||
and `filter` — a SQL-like expression over the columns `state`, `updated_by`,
|
||||
`owner_component`, and `claim_entity`. (`full_text_search` is reserved and currently
|
||||
rejected as not implemented.)
|
||||
|
||||
The response is **not JSON** — it is an Arrow IPC stream
|
||||
(`content-type: application/vnd.apache.arrow.stream`). Decode it, e.g. in Python:
|
||||
|
||||
```python
|
||||
import pyarrow.ipc
|
||||
import requests
|
||||
|
||||
resp = requests.post(
|
||||
f"{base_url}/v1/jobs/query_events",
|
||||
headers={"x-api-key": key, "x-lancedb-database": database},
|
||||
json={"job_id": job_id},
|
||||
)
|
||||
resp.raise_for_status()
|
||||
events = pyarrow.ipc.open_stream(resp.content).read_all()
|
||||
```
|
||||
|
||||
## Feature engineering (Geneva) jobs
|
||||
|
||||
Feature engineering jobs — UDF column backfills and materialized view refreshes run
|
||||
through Geneva — are tracked **separately** from the `/v1/jobs` registry above. Their
|
||||
records live in a `geneva_jobs` table inside the database itself (in the `__system`
|
||||
namespace), and you access them through a Python `geneva` connection rather than the
|
||||
REST endpoints above:
|
||||
|
||||
```python
|
||||
import geneva
|
||||
from geneva.jobs import JobStateManager
|
||||
|
||||
# Same credentials as lancedb.connect / the REST API
|
||||
conn = geneva.connect("db://<database>", api_key="<key>", host_override="<base_url>")
|
||||
jsm = JobStateManager(conn)
|
||||
|
||||
# List jobs. NOTE: status defaults to "RUNNING"; pass status=None for all jobs.
|
||||
# Statuses: PENDING | RUNNING | DONE | FAILED | CANCELLED
|
||||
jobs = jsm.list_jobs(table_name="my_table", status=None)
|
||||
|
||||
# Fetch one job by id (returns a list of JobRecord)
|
||||
records = jsm.get("<job_id>")
|
||||
```
|
||||
|
||||
Each `JobRecord` has `table_name`, `column_name`, `job_id`, `job_type`, `status`,
|
||||
`launched_at`, `completed_at`, `config`, `launched_by`, `manifest_id`, `cluster_name`,
|
||||
`metrics` (progress counters), `events` (human-readable history), and `updated_at`.
|
||||
For filters `list_jobs` doesn't support (e.g. time ranges), query the underlying table
|
||||
directly: `jsm.get_table(True).search().where("launched_at >= TIMESTAMP '...'")` —
|
||||
pass `True` to check out the latest version, since other processes update job state.
|
||||
|
||||
Stale-status caveat: nothing reaps dead Geneva jobs, so a job can sit in
|
||||
`RUNNING`/`PENDING` forever if its worker died. Treat a job as effectively `FAILED`
|
||||
when it has been running longer than ~36 hours, or its `updated_at` is more than ~2
|
||||
hours old (this matches the heuristic the Geneva console UI applies on read).
|
||||
|
||||
## Workflow tips
|
||||
|
||||
- To wait for async work (a backfill, an index build), poll `describe` until
|
||||
`job_state` leaves `IN_PROGRESS`; on `FAILED`, pull `status` and `query_events` for
|
||||
the failure detail.
|
||||
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.35.0-beta.2"
|
||||
current_version = "0.36.0"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.35.0-beta.2"
|
||||
version = "0.36.0"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
|
||||
@@ -8,6 +8,27 @@ A Python library for [LanceDB](https://github.com/lancedb/lancedb).
|
||||
pip install lancedb
|
||||
```
|
||||
|
||||
### Pre-Haswell x86_64 hosts: `lancedb-compat`
|
||||
|
||||
The default `lancedb` wheel targets `x86-64-haswell` (AVX2 + FMA + F16C) for full performance on modern hardware. Pre-Haswell hosts — Intel Sandy Bridge / Ivy Bridge / Westmere; AMD Bulldozer / Piledriver / Steamroller — don't have AVX2 and crash with `Illegal instruction` at `import lancedb`.
|
||||
|
||||
For those hosts, install the `lancedb-compat` package instead:
|
||||
|
||||
```bash
|
||||
pip install lancedb-compat
|
||||
```
|
||||
|
||||
Same Python API (`import lancedb` works as usual). The compat wheel is compiled at the `x86-64-v2` baseline (Nehalem-class) and uses runtime SIMD dispatch in the embedded lance crate to pick the right kernel tier (scalar / AVX / AVX+FMA / AVX2+FMA / AVX-512) at load time, so it still goes fast on modern hardware while running cleanly on the pre-Haswell silicon. Use `lance.simd_info()` from Python to verify which tier was selected.
|
||||
|
||||
`lancedb` and `lancedb-compat` install to the same `lancedb/` namespace and conflict at install time. Pick one. To switch, `pip uninstall lancedb` first, then `pip install lancedb-compat` (or vice-versa).
|
||||
|
||||
If you need a custom baseline (or `lancedb-compat` isn't yet published for your platform), build from source with the override:
|
||||
|
||||
```bash
|
||||
RUSTFLAGS="-C target-cpu=x86-64-v2" maturin build --release
|
||||
pip install ./target/wheels/lancedb-*.whl
|
||||
```
|
||||
|
||||
### Preview Releases
|
||||
|
||||
Stable releases are created about every 2 weeks. For the latest features and bug fixes, you can install the preview release. These releases receive the same level of testing as stable releases, but are not guaranteed to be available for more than 6 months after they are released. Once your application is stable, we recommend switching to stable releases.
|
||||
|
||||
@@ -63,7 +63,7 @@ tests = [
|
||||
"polars>=0.19, <=1.3.0",
|
||||
"pyarrow<25",
|
||||
"pyarrow-stubs>=16.0",
|
||||
"pylance==9.0.0rc1",
|
||||
"pylance==9.0.0",
|
||||
"requests>=2.31.0",
|
||||
"datafusion>=54,<55",
|
||||
"opentelemetry-sdk>=1.30.0",
|
||||
|
||||
@@ -19,17 +19,6 @@ from .db import AsyncConnection, DBConnection, LanceDBConnection
|
||||
from .remote import ClientConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
from .expr import Expr, col, lit, func
|
||||
from .udf import (
|
||||
udf,
|
||||
table_udf,
|
||||
Udf,
|
||||
Job,
|
||||
JobFailedError,
|
||||
MaterializedView,
|
||||
AsyncJob,
|
||||
AsyncMaterializedView,
|
||||
)
|
||||
from .lineage import Lineage, Node, Edge, FunctionRef
|
||||
from .schema import blob, vector, BlobType
|
||||
from .table import AsyncTable, Table
|
||||
from .types import BaseTokenizerType
|
||||
@@ -502,18 +491,6 @@ async def connect_async(
|
||||
|
||||
|
||||
__all__ = [
|
||||
"udf",
|
||||
"table_udf",
|
||||
"Udf",
|
||||
"Job",
|
||||
"JobFailedError",
|
||||
"MaterializedView",
|
||||
"AsyncJob",
|
||||
"AsyncMaterializedView",
|
||||
"Lineage",
|
||||
"Node",
|
||||
"Edge",
|
||||
"FunctionRef",
|
||||
"connect",
|
||||
"connect_async",
|
||||
"tokenize",
|
||||
|
||||
@@ -219,6 +219,7 @@ class Table:
|
||||
data: pa.RecordBatchReader,
|
||||
mode: Literal["append", "overwrite"],
|
||||
progress: Optional[Any] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
) -> AddResult: ...
|
||||
async def update(
|
||||
self, updates: Dict[str, str], where: Optional[str]
|
||||
@@ -318,6 +319,10 @@ class Branches:
|
||||
) -> Table: ...
|
||||
async def checkout(self, name: str, version: Optional[int] = None) -> Table: ...
|
||||
async def delete(self, name: str) -> None: ...
|
||||
async def diff(self, from_branch: str) -> Dict[str, Any]: ...
|
||||
async def merge(
|
||||
self, from_branch: str, dry_run: bool = False
|
||||
) -> Dict[str, Any]: ...
|
||||
|
||||
class IndexConfig:
|
||||
name: str
|
||||
|
||||
+41
-548
@@ -41,6 +41,7 @@ from lance_namespace import (
|
||||
ListTablesResponse,
|
||||
connect as namespace_connect,
|
||||
)
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
|
||||
from . import __version__
|
||||
from ._lancedb import connect as lancedb_connect # type: ignore
|
||||
@@ -65,7 +66,6 @@ if TYPE_CHECKING:
|
||||
from .common import DATA, URI
|
||||
from .embeddings import EmbeddingFunctionConfig
|
||||
from ._lancedb import Session
|
||||
from .udf import MaterializedView, AsyncMaterializedView
|
||||
|
||||
from .namespace_utils import (
|
||||
_normalize_create_namespace_mode,
|
||||
@@ -563,277 +563,6 @@ class DBConnection(EnforceOverrides):
|
||||
"""
|
||||
raise NotImplementedError("serialize is not supported for this connection type")
|
||||
|
||||
# -- Derived compute: functions, materialized views, jobs -------------
|
||||
# Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
# connections raise NotImplementedError for now.
|
||||
|
||||
def create_function(
|
||||
self,
|
||||
name,
|
||||
language: str = "python",
|
||||
return_type: Optional[str] = None,
|
||||
body: Optional[str] = None,
|
||||
options: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
):
|
||||
"""Register a UDF (CREATE FUNCTION).
|
||||
|
||||
Pass a ``@udf`` / ``@table_udf``-decorated function (preferred):
|
||||
|
||||
db.create_function(embed)
|
||||
|
||||
or the explicit fields:
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name: str or Udf
|
||||
A decorated UDF object, or the function name.
|
||||
language: str
|
||||
Implementation language (currently "python").
|
||||
return_type: str
|
||||
SQL return type, e.g. "FLOAT", "FLOAT[1536]",
|
||||
"STRUCT(a FLOAT, b VARCHAR)", "TABLE(chunk VARCHAR, idx INT)".
|
||||
body: str
|
||||
Function body: source text, or base64 cloudpickle bytes when
|
||||
options["body_format"] == "cloudpickle".
|
||||
options: dict, optional
|
||||
input_columns, pip, num_gpus, batch_size, timeout,
|
||||
error_policy, docker_image, body_format, ...
|
||||
replace: bool
|
||||
Drop an existing function of the same name first.
|
||||
"""
|
||||
from .udf import Udf
|
||||
|
||||
if isinstance(name, Udf):
|
||||
req = name.create_request()
|
||||
name, language, return_type, body, options = (
|
||||
req["name"],
|
||||
req["language"],
|
||||
req["return_type"],
|
||||
req["body"],
|
||||
req["options"],
|
||||
)
|
||||
if replace:
|
||||
try:
|
||||
self.drop_function(name)
|
||||
except Exception:
|
||||
pass
|
||||
LOOP.run(self._conn.create_function(name, language, return_type, body, options))
|
||||
|
||||
def list_functions(self):
|
||||
"""List registered functions (SHOW FUNCTIONS)."""
|
||||
return LOOP.run(self._conn.list_functions())
|
||||
|
||||
def drop_function(self, name: str):
|
||||
"""Drop a registered function (DROP FUNCTION)."""
|
||||
LOOP.run(self._conn.drop_function(name))
|
||||
|
||||
def create_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
source=None,
|
||||
select=None,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
where: Optional[str] = None,
|
||||
auto_refresh: bool = False,
|
||||
with_no_data: bool = False,
|
||||
replace: bool = False,
|
||||
partition_by: Optional[str] = None,
|
||||
) -> "MaterializedView":
|
||||
"""Create a materialized view (CREATE MATERIALIZED VIEW); returns a
|
||||
`MaterializedView` handle (``.wait()`` blocks until it is populated).
|
||||
|
||||
Two ways to specify the view body:
|
||||
|
||||
- ergonomic: pass ``source`` (a table name or table) and ``select``
|
||||
items -- column names, expression strings ("embed(body)"),
|
||||
(alias, expression) tuples, or ``@udf`` / ``@table_udf`` objects.
|
||||
The SELECT is assembled and parsed server-side (one parser, shared
|
||||
with SQL).
|
||||
- raw: pass ``query=`` with a full SELECT, e.g.
|
||||
"SELECT id, embed(body) AS vec FROM articles WHERE id > 1".
|
||||
|
||||
`partition_by` partitions the view's (single) table function on a source
|
||||
column. If that column has an IVF vector index the server partitions by
|
||||
its index clusters (image-dedup style); otherwise it groups by distinct
|
||||
value. (Geneva's `partition_by` and `partition_by_indexed_column` unify
|
||||
here -- the engine picks the strategy from the column.)
|
||||
"""
|
||||
from .udf import build_view_query, MaterializedView
|
||||
|
||||
if query is None:
|
||||
if source is None or select is None:
|
||||
raise ValueError(
|
||||
"create_materialized_view needs either query= or both "
|
||||
"source and select"
|
||||
)
|
||||
query = build_view_query(source, select)
|
||||
if where:
|
||||
query += f" WHERE {where}"
|
||||
if replace:
|
||||
self._drop_view_if_exists(name)
|
||||
job_id = LOOP.run(
|
||||
self._conn.create_materialized_view(
|
||||
name,
|
||||
query=query,
|
||||
auto_refresh=auto_refresh,
|
||||
with_no_data=with_no_data,
|
||||
partition_by=partition_by,
|
||||
)
|
||||
)
|
||||
return MaterializedView(self, name, job_id=job_id)
|
||||
|
||||
def _drop_view_if_exists(self, name: str) -> None:
|
||||
# `replace=True` is "drop if present"; only a not-found error is
|
||||
# benign here. Anything else (perms, server fault) must surface rather
|
||||
# than be masked by a later create failure.
|
||||
try:
|
||||
self.drop_materialized_view(name)
|
||||
except Exception as e:
|
||||
msg = str(e).lower()
|
||||
if "not found" not in msg and "does not exist" not in msg:
|
||||
raise
|
||||
|
||||
def job(self, job_id: str):
|
||||
"""A `Job` for reconnecting to an inflight job by id -- e.g. an
|
||||
id you stored, or one returned from the SQL / REST surface. Submit
|
||||
methods (`refresh_column`, `MaterializedView.refresh`) already return a
|
||||
handle directly, so you do not need this to wait on a fresh submission."""
|
||||
from .udf import Job
|
||||
|
||||
return Job(self, job_id)
|
||||
|
||||
def lineage(
|
||||
self,
|
||||
table: str,
|
||||
column: Optional[str] = None,
|
||||
*,
|
||||
direction: Optional[str] = None,
|
||||
depth: Optional[int] = None,
|
||||
):
|
||||
"""Derived-compute lineage of a table/view, or one of its columns:
|
||||
upstream sources, downstream dependents, and the function version +
|
||||
location that produced each derived column (with a drift flag). Returns
|
||||
a `Lineage`. `direction` is "upstream" | "downstream" | "both" (server
|
||||
default both); `depth` limits column-hops (transitive when omitted)."""
|
||||
# `self._conn` is the AsyncConnection; drive its async `lineage`
|
||||
# (which parses the JSON) on the loop, mirroring create_materialized_view.
|
||||
return LOOP.run(
|
||||
self._conn.lineage(table, column, direction=direction, depth=depth)
|
||||
)
|
||||
|
||||
def _refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Internal: submit a materialized-view refresh, return the job id.
|
||||
The public surface is ``MaterializedView.refresh()`` (which returns a
|
||||
`Job`); this stays private so refresh is only reached through the
|
||||
handle.
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._conn._refresh_materialized_view(
|
||||
name,
|
||||
full=full,
|
||||
src_version=src_version,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
)
|
||||
)
|
||||
|
||||
def explain_refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH). Returns a
|
||||
plan with .has_work / .source_version / .last_refreshed_version /
|
||||
.full_refresh / .rebuild / .units_total. `full=True` plans a full
|
||||
rebuild (incremental planning needs stable row IDs on the source)."""
|
||||
return LOOP.run(
|
||||
self._conn.explain_refresh_materialized_view(
|
||||
name, full=full, src_version=src_version
|
||||
)
|
||||
)
|
||||
|
||||
def alter_materialized_view(self, name: str, *, auto_refresh: bool):
|
||||
"""Update a materialized view's options (ALTER MATERIALIZED VIEW)."""
|
||||
LOOP.run(self._conn.alter_materialized_view(name, auto_refresh=auto_refresh))
|
||||
|
||||
def drop_materialized_view(self, name: str):
|
||||
"""Drop a materialized view definition (DROP MATERIALIZED VIEW)."""
|
||||
LOOP.run(self._conn.drop_materialized_view(name))
|
||||
|
||||
def list_materialized_views(self):
|
||||
"""List registered materialized view definitions."""
|
||||
return LOOP.run(self._conn.list_materialized_views())
|
||||
|
||||
def list_jobs(self):
|
||||
"""List inflight server-side jobs across the database's tables."""
|
||||
return LOOP.run(self._conn.list_jobs())
|
||||
|
||||
def get_job(self, job_id: str, table: "str | None" = None):
|
||||
"""Look up one server-side job by id (the wait()/status poll path).
|
||||
|
||||
Passing ``table`` (the job's table) lets the server answer with an O(1)
|
||||
single-node read instead of scanning the database's active jobs.
|
||||
Returns the job's status, or None if it's unknown or no longer active.
|
||||
"""
|
||||
return LOOP.run(self._conn.get_job(job_id, table))
|
||||
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""Cancel an inflight server-side job by id (CANCEL JOB).
|
||||
|
||||
Returns True if a matching inflight job was found and flagged for
|
||||
cancellation, False if none was inflight (already finished or
|
||||
unknown id) -- cancellation is best-effort.
|
||||
"""
|
||||
return LOOP.run(self._conn.cancel_job(job_id))
|
||||
|
||||
def describe_platform_job(self, platform_job_id: str):
|
||||
"""Describe a platform job (POST /v1/jobs/describe): registry-backed
|
||||
lifecycle state plus the owner-written status payload. None when the
|
||||
registry has no such job."""
|
||||
return LOOP.run(self._conn.describe_platform_job(platform_job_id))
|
||||
|
||||
def resolve_platform_job_id(
|
||||
self, manifest_job_id: str, table: "str | None" = None
|
||||
):
|
||||
"""Resolve a submission (manifest) job id to its platform job id.
|
||||
None until the job has registered (dispatch is async)."""
|
||||
return LOOP.run(self._conn.resolve_platform_job_id(manifest_job_id, table))
|
||||
|
||||
def cancel_platform_job(self, platform_job_id: str) -> None:
|
||||
"""Cancel a platform job (POST /v1/jobs/cancel). Idempotent on
|
||||
already-terminal jobs."""
|
||||
return LOOP.run(self._conn.cancel_platform_job(platform_job_id))
|
||||
|
||||
def job_history(self, job_id: "str | None" = None):
|
||||
"""Durable history of completed server-side jobs (SHOW JOB HISTORY).
|
||||
|
||||
Pass ``job_id`` to narrow to a single job. Unlike :meth:`list_jobs`
|
||||
(live, inflight) these are the terminal records.
|
||||
"""
|
||||
return LOOP.run(self._conn.job_history(job_id))
|
||||
|
||||
def errors(self, job_id: "str | None" = None, table: "str | None" = None):
|
||||
"""Per-row UDF errors recorded by ``error_policy=skip`` (SHOW ERRORS),
|
||||
optionally filtered by ``job_id`` and/or ``table``.
|
||||
"""
|
||||
return LOOP.run(self._conn.errors(job_id, table))
|
||||
|
||||
|
||||
class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
@@ -1018,10 +747,12 @@ class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return self._namespace_conn().list_namespaces(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
limit=limit,
|
||||
return LOOP.run(
|
||||
self._conn.list_namespaces(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
limit=limit,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -1031,10 +762,12 @@ class LanceDBConnection(DBConnection):
|
||||
mode: Optional[str] = None,
|
||||
properties: Optional[Dict[str, str]] = None,
|
||||
) -> CreateNamespaceResponse:
|
||||
return self._namespace_conn().create_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
properties=properties,
|
||||
return LOOP.run(
|
||||
self._conn.create_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
properties=properties,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -1044,19 +777,24 @@ class LanceDBConnection(DBConnection):
|
||||
mode: Optional[str] = None,
|
||||
behavior: Optional[str] = None,
|
||||
) -> DropNamespaceResponse:
|
||||
return self._namespace_conn().drop_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
behavior=behavior,
|
||||
)
|
||||
try:
|
||||
return LOOP.run(
|
||||
self._conn.drop_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
behavior=behavior,
|
||||
)
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Namespace not empty" in str(e):
|
||||
raise NamespaceNotEmptyError(str(e)) from e
|
||||
raise
|
||||
|
||||
@override
|
||||
def describe_namespace(
|
||||
self, namespace_path: List[str]
|
||||
) -> DescribeNamespaceResponse:
|
||||
return self._namespace_conn().describe_namespace(
|
||||
namespace_path=namespace_path,
|
||||
)
|
||||
return LOOP.run(self._conn.describe_namespace(namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_tables(
|
||||
@@ -1085,12 +823,6 @@ class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
if namespace_path:
|
||||
return self._namespace_conn().list_tables(
|
||||
namespace_path=namespace_path,
|
||||
page_token=page_token,
|
||||
limit=limit,
|
||||
)
|
||||
return LOOP.run(
|
||||
self._conn.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
@@ -1188,22 +920,6 @@ class LanceDBConnection(DBConnection):
|
||||
raise ValueError("mode must be either 'create' or 'overwrite'")
|
||||
validate_table_name(name)
|
||||
|
||||
if namespace_path:
|
||||
return self._namespace_conn().create_table(
|
||||
name,
|
||||
data=data,
|
||||
schema=schema,
|
||||
mode=mode,
|
||||
exist_ok=exist_ok,
|
||||
on_bad_vectors=on_bad_vectors,
|
||||
fill_value=fill_value,
|
||||
embedding_functions=embedding_functions,
|
||||
namespace_path=namespace_path,
|
||||
storage_options=storage_options,
|
||||
data_storage_version=data_storage_version,
|
||||
enable_v2_manifest_paths=enable_v2_manifest_paths,
|
||||
)
|
||||
|
||||
tbl = LanceTable.create(
|
||||
self,
|
||||
name,
|
||||
@@ -1216,22 +932,11 @@ class LanceDBConnection(DBConnection):
|
||||
embedding_functions=embedding_functions,
|
||||
namespace_path=namespace_path,
|
||||
storage_options=storage_options,
|
||||
data_storage_version=data_storage_version,
|
||||
enable_v2_manifest_paths=enable_v2_manifest_paths,
|
||||
)
|
||||
return tbl
|
||||
|
||||
def _namespace_conn(self) -> DBConnection:
|
||||
"""Return a LanceNamespaceDBConnection backed by this connection's
|
||||
directory namespace. Used to delegate child-namespace operations."""
|
||||
from lancedb.namespace import LanceNamespaceDBConnection
|
||||
|
||||
return LanceNamespaceDBConnection(
|
||||
self.namespace_client(),
|
||||
read_consistency_interval=self.read_consistency_interval,
|
||||
storage_options=self.storage_options,
|
||||
namespace_client_impl=None,
|
||||
namespace_client_properties=None,
|
||||
)
|
||||
|
||||
@override
|
||||
def open_table(
|
||||
self,
|
||||
@@ -1278,14 +983,7 @@ class LanceDBConnection(DBConnection):
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
if namespace_path:
|
||||
tbl = self._namespace_conn().open_table(
|
||||
name,
|
||||
namespace_path=namespace_path,
|
||||
storage_options=storage_options,
|
||||
index_cache_size=index_cache_size,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
tbl = LanceTable.open(
|
||||
self,
|
||||
name,
|
||||
@@ -1293,6 +991,15 @@ class LanceDBConnection(DBConnection):
|
||||
storage_options=storage_options,
|
||||
index_cache_size=index_cache_size,
|
||||
)
|
||||
except (RuntimeError, ValueError) as e:
|
||||
if namespace_path and (
|
||||
"Table not found" in str(e) or "was not found" in str(e)
|
||||
):
|
||||
table_id = namespace_path + [name]
|
||||
raise TableNotFoundError(
|
||||
f"Table not found: {'$'.join(table_id)}"
|
||||
) from e
|
||||
raise
|
||||
|
||||
if branch is not None:
|
||||
tbl = tbl.branches.checkout(branch, version)
|
||||
@@ -1376,9 +1083,6 @@ class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
if namespace_path:
|
||||
self._namespace_conn().drop_table(name, namespace_path=namespace_path)
|
||||
return
|
||||
LOOP.run(
|
||||
self._conn.drop_table(
|
||||
name, namespace_path=namespace_path, ignore_missing=ignore_missing
|
||||
@@ -1927,7 +1631,7 @@ class AsyncConnection(object):
|
||||
namespace_client=namespace_client,
|
||||
)
|
||||
|
||||
return AsyncTable(new_table, conn=self)
|
||||
return AsyncTable(new_table)
|
||||
|
||||
async def open_table(
|
||||
self,
|
||||
@@ -2000,7 +1704,7 @@ class AsyncConnection(object):
|
||||
namespace_client=namespace_client,
|
||||
managed_versioning=managed_versioning,
|
||||
)
|
||||
tbl = AsyncTable(table, conn=self)
|
||||
tbl = AsyncTable(table)
|
||||
# "main" is the default branch, so treat it as no branch: remote rejects
|
||||
# every branch checkout (even "main"), and the version still applies.
|
||||
if branch is not None and branch != "main":
|
||||
@@ -2057,218 +1761,7 @@ class AsyncConnection(object):
|
||||
source_tag=source_tag,
|
||||
is_shallow=is_shallow,
|
||||
)
|
||||
return AsyncTable(table, conn=self)
|
||||
|
||||
# -- Derived compute: functions, materialized views, jobs -------------
|
||||
# Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
# connections raise NotImplementedError for now.
|
||||
|
||||
async def create_function(
|
||||
self,
|
||||
name,
|
||||
language: str = "python",
|
||||
return_type: Optional[str] = None,
|
||||
body: Optional[str] = None,
|
||||
options: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
):
|
||||
"""Register a UDF (CREATE FUNCTION). Accepts a ``@udf``/``@table_udf``
|
||||
object (preferred) or the explicit (name, language, return_type, body,
|
||||
options)."""
|
||||
from .udf import Udf
|
||||
|
||||
if isinstance(name, Udf):
|
||||
req = name.create_request()
|
||||
name, language, return_type, body, options = (
|
||||
req["name"],
|
||||
req["language"],
|
||||
req["return_type"],
|
||||
req["body"],
|
||||
req["options"],
|
||||
)
|
||||
if replace:
|
||||
try:
|
||||
await self.drop_function(name)
|
||||
except Exception:
|
||||
pass
|
||||
await self._inner.create_function(name, language, return_type, body, options)
|
||||
|
||||
async def list_functions(self):
|
||||
"""List registered functions (SHOW FUNCTIONS)."""
|
||||
return await self._inner.list_functions()
|
||||
|
||||
async def drop_function(self, name: str):
|
||||
"""Drop a registered function (DROP FUNCTION)."""
|
||||
await self._inner.drop_function(name)
|
||||
|
||||
async def create_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
source=None,
|
||||
select=None,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
where: Optional[str] = None,
|
||||
auto_refresh: bool = False,
|
||||
with_no_data: bool = False,
|
||||
replace: bool = False,
|
||||
partition_by: Optional[str] = None,
|
||||
) -> "AsyncMaterializedView":
|
||||
"""Create a materialized view; returns an `AsyncMaterializedView`
|
||||
handle (``.wait()`` blocks until populated). Pass either ``query=`` (a
|
||||
full SELECT) or ``source`` + ``select`` items; `partition_by`
|
||||
partitions the view's table function on a source column (index-cluster
|
||||
if the column is IVF-indexed, else distinct-value). See the sync
|
||||
method for the select grammar."""
|
||||
from .udf import build_view_query, AsyncMaterializedView
|
||||
|
||||
if query is None:
|
||||
if source is None or select is None:
|
||||
raise ValueError(
|
||||
"create_materialized_view needs either query= or both "
|
||||
"source and select"
|
||||
)
|
||||
query = build_view_query(source, select)
|
||||
if where:
|
||||
query += f" WHERE {where}"
|
||||
if replace:
|
||||
try:
|
||||
await self.drop_materialized_view(name)
|
||||
except Exception as e:
|
||||
msg = str(e).lower()
|
||||
if "not found" not in msg and "does not exist" not in msg:
|
||||
raise
|
||||
job_id = await self._inner.create_materialized_view(
|
||||
name,
|
||||
query,
|
||||
auto_refresh=auto_refresh,
|
||||
with_no_data=with_no_data,
|
||||
partition_by=partition_by,
|
||||
)
|
||||
return AsyncMaterializedView(self, name, job_id=job_id)
|
||||
|
||||
def job(self, job_id: str):
|
||||
"""An `AsyncJob` for reconnecting to an inflight job by id (a
|
||||
stored id, or one from the SQL / REST surface). Submit methods already
|
||||
return a handle, so this is only needed to re-attach to an existing
|
||||
job."""
|
||||
from .udf import AsyncJob
|
||||
|
||||
return AsyncJob(self, job_id)
|
||||
|
||||
async def lineage(
|
||||
self,
|
||||
table: str,
|
||||
column: Optional[str] = None,
|
||||
*,
|
||||
direction: Optional[str] = None,
|
||||
depth: Optional[int] = None,
|
||||
):
|
||||
"""Derived-compute lineage of a table/view (or column). See the sync
|
||||
`Connection.lineage`. Returns a `Lineage`."""
|
||||
from .lineage import Lineage
|
||||
|
||||
raw = await self._inner.table_lineage(table, column, direction, depth)
|
||||
return Lineage.from_json(raw)
|
||||
|
||||
async def _refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Internal: submit a refresh, return the job id. The public surface is
|
||||
``AsyncMaterializedView.refresh()`` (returns an `AsyncJob`).
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh.
|
||||
"""
|
||||
return await self._inner.refresh_materialized_view(
|
||||
name,
|
||||
full=full,
|
||||
src_version=src_version,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
)
|
||||
|
||||
async def explain_refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH)."""
|
||||
return await self._inner.explain_refresh_materialized_view(
|
||||
name, full=full, src_version=src_version
|
||||
)
|
||||
|
||||
async def alter_materialized_view(self, name: str, *, auto_refresh: bool):
|
||||
"""Update a materialized view's options."""
|
||||
await self._inner.alter_materialized_view(name, auto_refresh)
|
||||
|
||||
async def drop_materialized_view(self, name: str):
|
||||
"""Drop a materialized view definition."""
|
||||
await self._inner.drop_materialized_view(name)
|
||||
|
||||
async def list_materialized_views(self):
|
||||
"""List registered materialized view definitions."""
|
||||
return await self._inner.list_materialized_views()
|
||||
|
||||
async def list_jobs(self):
|
||||
"""List inflight server-side jobs across the database's tables."""
|
||||
return await self._inner.list_jobs()
|
||||
|
||||
async def get_job(self, job_id: str, table: "str | None" = None):
|
||||
"""Look up one server-side job by id (the wait()/status poll path).
|
||||
``table`` (the job's table) enables an O(1) server-side lookup.
|
||||
Returns the job's status, or None if unknown / no longer active."""
|
||||
return await self._inner.get_job(job_id, table)
|
||||
|
||||
async def cancel_job(self, job_id: str) -> bool:
|
||||
"""Cancel an inflight server-side job by id (CANCEL JOB).
|
||||
|
||||
Returns True if a matching inflight job was found and flagged for
|
||||
cancellation, False otherwise (best-effort).
|
||||
"""
|
||||
return await self._inner.cancel_job(job_id)
|
||||
|
||||
async def describe_platform_job(self, platform_job_id: str):
|
||||
"""Describe a platform job: registry-backed lifecycle state plus the
|
||||
owner-written status payload. None when the registry has no such
|
||||
job."""
|
||||
return await self._inner.describe_platform_job(platform_job_id)
|
||||
|
||||
async def resolve_platform_job_id(
|
||||
self, manifest_job_id: str, table: "str | None" = None
|
||||
):
|
||||
"""Resolve a submission (manifest) job id to its platform job id.
|
||||
None until the job has registered (dispatch is async)."""
|
||||
return await self._inner.resolve_platform_job_id(manifest_job_id, table)
|
||||
|
||||
async def cancel_platform_job(self, platform_job_id: str) -> None:
|
||||
"""Cancel a platform job. Idempotent on already-terminal jobs."""
|
||||
return await self._inner.cancel_platform_job(platform_job_id)
|
||||
|
||||
async def job_history(self, job_id: "str | None" = None):
|
||||
"""Durable history of completed server-side jobs (SHOW JOB HISTORY).
|
||||
|
||||
Reads each table's durable job-history store. Pass ``job_id`` to narrow
|
||||
to a single job. Unlike :meth:`list_jobs` (live, inflight) these are the
|
||||
terminal records, with created/updated/completed timestamps.
|
||||
"""
|
||||
return await self._inner.job_history(job_id)
|
||||
|
||||
async def errors(self, job_id: "str | None" = None, table: "str | None" = None):
|
||||
"""Per-row UDF errors recorded by ``error_policy=skip`` (SHOW ERRORS).
|
||||
|
||||
Optionally filtered by ``job_id`` and/or ``table``.
|
||||
"""
|
||||
return await self._inner.errors(job_id, table)
|
||||
return AsyncTable(table)
|
||||
|
||||
async def rename_table(
|
||||
self,
|
||||
|
||||
@@ -14,29 +14,76 @@ import numpy as np
|
||||
|
||||
DEFAULT_WATSONX_URL = "https://us-south.ml.cloud.ibm.com"
|
||||
|
||||
MODELS_DIMS = {
|
||||
# Models currently available on the watsonx.ai SaaS platform.
|
||||
# These are the IDs advertised to new users via model_names() and shown in
|
||||
# validation error messages. Regional availability and withdrawal dates are
|
||||
# documented at:
|
||||
# https://www.ibm.com/docs/en/watsonx/saas?topic=models-supported-encoder
|
||||
CURRENT_MODELS: dict[str, int] = {
|
||||
"ibm/granite-embedding-278m-multilingual": 768,
|
||||
"ibm/slate-125m-english-rtrvr-v2": 768,
|
||||
"ibm/slate-30m-english-rtrvr-v2": 384,
|
||||
"intfloat/multilingual-e5-large": 1024,
|
||||
}
|
||||
|
||||
# Full dimension map including legacy model IDs from earlier releases.
|
||||
# Kept so that existing tables whose stored metadata uses these names can still
|
||||
# resolve dimensions on load without raising an error. These IDs are NOT
|
||||
# advertised to new users.
|
||||
MODELS_DIMS: dict[str, int] = {
|
||||
**CURRENT_MODELS,
|
||||
# Deprecated — withdrawal announced but still functional until the dates above.
|
||||
"sentence-transformers/all-minilm-l6-v2": 384,
|
||||
# Pre-v2 legacy names retained for metadata compatibility only.
|
||||
"ibm/slate-125m-english-rtrvr": 768,
|
||||
"ibm/slate-30m-english-rtrvr": 384,
|
||||
"sentence-transformers/all-minilm-l12-v2": 384,
|
||||
"intfloat/multilingual-e5-large": 1024,
|
||||
}
|
||||
|
||||
|
||||
@register("watsonx")
|
||||
class WatsonxEmbeddings(TextEmbeddingFunction):
|
||||
"""
|
||||
An embedding function that uses the IBM watsonx.ai Embeddings API.
|
||||
|
||||
API Docs:
|
||||
---------
|
||||
https://cloud.ibm.com/apidocs/watsonx-ai#text-embeddings
|
||||
https://cloud.ibm.com/apidocs/watsonx-ai#text-embeddings
|
||||
|
||||
Supported embedding models:
|
||||
---------------------------
|
||||
https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx
|
||||
https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str, default "ibm/slate-125m-english-rtrvr"
|
||||
The ID of the embedding model to use. For new tables,
|
||||
``"ibm/granite-embedding-278m-multilingual"`` is recommended.
|
||||
api_key : str, optional
|
||||
IBM Cloud API key. Falls back to the ``WATSONX_API_KEY`` environment
|
||||
variable when not provided.
|
||||
project_id : str, optional
|
||||
watsonx.ai project ID. Explicit value takes precedence over the
|
||||
``WATSONX_PROJECT_ID`` environment variable. Mutually exclusive with
|
||||
``space_id`` — exactly one must be supplied.
|
||||
space_id : str, optional
|
||||
watsonx.ai deployment space ID. Explicit value takes precedence over
|
||||
the ``WATSONX_SPACE_ID`` environment variable. Mutually exclusive with
|
||||
``project_id`` — exactly one must be supplied.
|
||||
url : str, optional
|
||||
watsonx.ai service URL. Defaults to
|
||||
``"https://us-south.ml.cloud.ibm.com"``.
|
||||
params : dict, optional
|
||||
Extra parameters forwarded verbatim to ``Embeddings`` (e.g.
|
||||
``{"truncate_input_tokens": 512}``).
|
||||
"""
|
||||
|
||||
# Intentionally kept at the original pre-PR default so that existing tables
|
||||
# whose stored metadata contains model:{} reload with the same model they
|
||||
# were created with. New users should pass name= explicitly, e.g.
|
||||
# name="ibm/granite-embedding-278m-multilingual".
|
||||
name: str = "ibm/slate-125m-english-rtrvr"
|
||||
api_key: Optional[str] = None
|
||||
project_id: Optional[str] = None
|
||||
space_id: Optional[str] = None
|
||||
url: Optional[str] = None
|
||||
params: Optional[Dict] = None
|
||||
|
||||
@@ -46,12 +93,13 @@ class WatsonxEmbeddings(TextEmbeddingFunction):
|
||||
|
||||
@staticmethod
|
||||
def model_names():
|
||||
return [
|
||||
"ibm/slate-125m-english-rtrvr",
|
||||
"ibm/slate-30m-english-rtrvr",
|
||||
"sentence-transformers/all-minilm-l12-v2",
|
||||
"intfloat/multilingual-e5-large",
|
||||
]
|
||||
"""Return the IDs of models currently available for new tables.
|
||||
|
||||
Legacy / deprecated IDs are intentionally excluded. They remain
|
||||
resolvable for dimension lookups on existing tables via ``MODELS_DIMS``,
|
||||
but should not be used when creating new tables.
|
||||
"""
|
||||
return list(CURRENT_MODELS.keys())
|
||||
|
||||
def ndims(self):
|
||||
return self._ndims
|
||||
@@ -59,7 +107,10 @@ class WatsonxEmbeddings(TextEmbeddingFunction):
|
||||
@cached_property
|
||||
def _ndims(self):
|
||||
if self.name not in MODELS_DIMS:
|
||||
raise ValueError(f"Unknown model name {self.name}")
|
||||
raise ValueError(
|
||||
f"Unknown model '{self.name}'. "
|
||||
f"Available models: {list(CURRENT_MODELS.keys())}"
|
||||
)
|
||||
return MODELS_DIMS[self.name]
|
||||
|
||||
def generate_embeddings(
|
||||
@@ -81,27 +132,45 @@ class WatsonxEmbeddings(TextEmbeddingFunction):
|
||||
"ibm_watsonx_ai.foundation_models"
|
||||
)
|
||||
|
||||
kwargs = {"model_id": self.name}
|
||||
# --- credentials ---
|
||||
# Explicit field takes priority; env var is the fallback.
|
||||
api_key = self.api_key or os.environ.get("WATSONX_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"WATSONX_API_KEY not set. Either set it in your environment or "
|
||||
"pass it as `api_key` argument to WatsonxEmbeddings."
|
||||
)
|
||||
credentials = ibm_watsonx_ai.Credentials(
|
||||
api_key=api_key,
|
||||
url=self.url or DEFAULT_WATSONX_URL,
|
||||
)
|
||||
|
||||
# --- project_id / space_id (exactly one required) ---
|
||||
# Explicit field always wins; env var is consulted only when the
|
||||
# corresponding field was not set, so passing project_id= never
|
||||
# conflicts with a stray WATSONX_SPACE_ID env var and vice-versa.
|
||||
space_id, project_id = self.space_id, self.project_id
|
||||
|
||||
if project_id is None and space_id is None:
|
||||
# Neither was passed explicitly — fall back to env vars.
|
||||
project_id = os.environ.get("WATSONX_PROJECT_ID")
|
||||
space_id = os.environ.get("WATSONX_SPACE_ID")
|
||||
|
||||
if project_id and space_id:
|
||||
raise ValueError("Provide either `project_id` or `space_id`, not both.")
|
||||
if not project_id and not space_id:
|
||||
raise ValueError(
|
||||
"Either WATSONX_PROJECT_ID or WATSONX_SPACE_ID must be set. "
|
||||
"Pass one as an argument to WatsonxEmbeddings or set the "
|
||||
"corresponding environment variable."
|
||||
)
|
||||
|
||||
client_kwargs: Dict = dict(model_id=self.name, credentials=credentials)
|
||||
if self.params:
|
||||
kwargs["params"] = self.params
|
||||
if self.project_id:
|
||||
kwargs["project_id"] = self.project_id
|
||||
elif "WATSONX_PROJECT_ID" in os.environ:
|
||||
kwargs["project_id"] = os.environ["WATSONX_PROJECT_ID"]
|
||||
client_kwargs["params"] = self.params
|
||||
if project_id:
|
||||
client_kwargs["project_id"] = project_id
|
||||
else:
|
||||
raise ValueError("WATSONX_PROJECT_ID must be set or passed")
|
||||
client_kwargs["space_id"] = space_id
|
||||
|
||||
creds_kwargs = {}
|
||||
if self.api_key:
|
||||
creds_kwargs["api_key"] = self.api_key
|
||||
elif "WATSONX_API_KEY" in os.environ:
|
||||
creds_kwargs["api_key"] = os.environ["WATSONX_API_KEY"]
|
||||
else:
|
||||
raise ValueError("WATSONX_API_KEY must be set or passed")
|
||||
if self.url:
|
||||
creds_kwargs["url"] = self.url
|
||||
else:
|
||||
creds_kwargs["url"] = DEFAULT_WATSONX_URL
|
||||
kwargs["credentials"] = ibm_watsonx_ai.Credentials(**creds_kwargs)
|
||||
|
||||
return ibm_watsonx_ai_foundation_models.Embeddings(**kwargs)
|
||||
return ibm_watsonx_ai_foundation_models.Embeddings(**client_kwargs)
|
||||
|
||||
@@ -1,177 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""Client-side model of derived-compute lineage.
|
||||
|
||||
`Connection.lineage()` / `Table.lineage()` / `MaterializedView.lineage()` return
|
||||
a `Lineage`: the graph of what a column or materialized view derives from
|
||||
(upstream), what derives from it (downstream), and -- for each derived column --
|
||||
the function that produced it, the version it was produced with, and whether
|
||||
that is stale relative to the function the registry now holds.
|
||||
|
||||
The server returns this as JSON (the wire contract); these classes deserialize
|
||||
it. Nothing here talks to the server.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Union
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionRef:
|
||||
"""The function that produced a derived column, with version + location."""
|
||||
|
||||
name: str
|
||||
#: Version that produced the data (stamped at compute time), if known.
|
||||
as_computed_version: Optional[str] = None
|
||||
#: Version the registry currently holds for this function name.
|
||||
current_version: Optional[str] = None
|
||||
#: True when the column was produced by an older function than the registry
|
||||
#: now holds -- i.e. silently stale; re-refresh to catch up.
|
||||
stale_vs_current: bool = False
|
||||
language: Optional[str] = None
|
||||
docker_image: Optional[str] = None
|
||||
env_digest: Optional[str] = None
|
||||
code_uri: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "FunctionRef":
|
||||
return cls(
|
||||
name=d["name"],
|
||||
as_computed_version=d.get("as_computed_version"),
|
||||
current_version=d.get("current_version"),
|
||||
stale_vs_current=d.get("stale_vs_current", False),
|
||||
language=d.get("language"),
|
||||
docker_image=d.get("docker_image"),
|
||||
env_digest=d.get("env_digest"),
|
||||
code_uri=d.get("code_uri"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
"""A lineage node: a table, view, column, or function."""
|
||||
|
||||
kind: str # "table" | "view" | "column" | "function"
|
||||
id: str # "table", "table.column", or "fn:name@version"
|
||||
table: Optional[str] = None
|
||||
function: Optional[FunctionRef] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "Node":
|
||||
fn = d.get("function")
|
||||
return cls(
|
||||
kind=d["kind"],
|
||||
id=d["id"],
|
||||
table=d.get("table"),
|
||||
function=FunctionRef._from(fn) if fn else None,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Edge:
|
||||
"""`downstream` depends on `upstream`, produced by `via` (a function name,
|
||||
or None for a passthrough)."""
|
||||
|
||||
downstream: str
|
||||
upstream: str
|
||||
via: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "Edge":
|
||||
return cls(downstream=d["downstream"], upstream=d["upstream"], via=d.get("via"))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Lineage:
|
||||
"""A derived-compute lineage graph (nodes + labeled edges)."""
|
||||
|
||||
target: str
|
||||
nodes: List[Node] = field(default_factory=list)
|
||||
edges: List[Edge] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, raw: Union[str, bytes, dict]) -> "Lineage":
|
||||
d = json.loads(raw) if isinstance(raw, (str, bytes)) else raw
|
||||
return cls(
|
||||
target=d.get("target", ""),
|
||||
nodes=[Node._from(n) for n in d.get("nodes", [])],
|
||||
edges=[Edge._from(e) for e in d.get("edges", [])],
|
||||
)
|
||||
|
||||
def functions(self) -> List[FunctionRef]:
|
||||
"""The function nodes in the graph."""
|
||||
return [n.function for n in self.nodes if n.function is not None]
|
||||
|
||||
def stale(self) -> List[FunctionRef]:
|
||||
"""Functions whose as-computed version is behind the current registry
|
||||
version -- the columns they produced are silently out of date."""
|
||||
return [f for f in self.functions() if f.stale_vs_current]
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
def prune(d: dict) -> dict:
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
return {
|
||||
"target": self.target,
|
||||
"nodes": [
|
||||
prune(
|
||||
{
|
||||
"kind": n.kind,
|
||||
"id": n.id,
|
||||
"table": n.table,
|
||||
"function": prune(vars(n.function)) if n.function else None,
|
||||
}
|
||||
)
|
||||
for n in self.nodes
|
||||
],
|
||||
"edges": [prune(vars(e)) for e in self.edges],
|
||||
}
|
||||
|
||||
def to_graphviz(self) -> str:
|
||||
"""Graphviz DOT for the lineage DAG: columns/tables as nodes, function
|
||||
names on edges, drift edges dashed + red."""
|
||||
stale_names = {f.name for f in self.stale()}
|
||||
out = [
|
||||
"digraph lineage {",
|
||||
" rankdir=LR;",
|
||||
' node [fontname="monospace"];',
|
||||
]
|
||||
for n in self.nodes:
|
||||
if n.kind == "function":
|
||||
continue
|
||||
shape = "ellipse" if n.kind in ("table", "view") else "box"
|
||||
out.append(f' "{n.id}" [shape={shape}];')
|
||||
for e in self.edges:
|
||||
attrs = ""
|
||||
if e.via:
|
||||
if e.via in stale_names:
|
||||
attrs = f' [label="{e.via}" color=red style=dashed]'
|
||||
else:
|
||||
attrs = f' [label="{e.via}"]'
|
||||
out.append(f' "{e.upstream}" -> "{e.downstream}"{attrs};')
|
||||
out.append("}")
|
||||
return "\n".join(out)
|
||||
|
||||
def _repr_html_(self) -> str:
|
||||
warn = ""
|
||||
drift = self.stale()
|
||||
if drift:
|
||||
names = ", ".join(sorted({f.name for f in drift}))
|
||||
warn = (
|
||||
f'<p style="color:#b00000"><b>stale vs current:</b> {names} '
|
||||
"(re-refresh to catch up)</p>"
|
||||
)
|
||||
rows = "".join(
|
||||
f"<tr><td><code>{e.downstream}</code></td>"
|
||||
f"<td>← {e.via or ''}</td>"
|
||||
f"<td><code>{e.upstream}</code></td></tr>"
|
||||
for e in self.edges
|
||||
)
|
||||
return (
|
||||
f"<b>lineage: <code>{self.target}</code></b>{warn}"
|
||||
"<table><tr><th>derived</th><th>via</th><th>from</th></tr>"
|
||||
f"{rows}</table>"
|
||||
)
|
||||
@@ -885,7 +885,7 @@ class Permutation:
|
||||
This method refines the current selection, potentially removing columns. It
|
||||
will not add back columns that were previously removed.
|
||||
|
||||
If any of the columns do not exist then an error will be raised
|
||||
If any of the columns do not exist then an error will be raised.
|
||||
|
||||
This does not introduce a post-processing step. It simply reduces the amount
|
||||
of data we read.
|
||||
@@ -898,9 +898,14 @@ class Permutation:
|
||||
for name in columns:
|
||||
value = self.selection.get(name, None)
|
||||
if value is None:
|
||||
raise ValueError(
|
||||
f"Cannot select column `{name}` because it does not exist"
|
||||
)
|
||||
if name == "_rowid":
|
||||
# _rowid is a system column not in the default schema
|
||||
# but can be explicitly selected
|
||||
value = "_rowid"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Cannot select column `{name}` because it does not exist"
|
||||
)
|
||||
new_selection[name] = value
|
||||
return self._with_selection(new_selection)
|
||||
|
||||
|
||||
@@ -3875,16 +3875,14 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
>>> asyncio.run(doctest_example()) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||||
RRFReranker(K=60)
|
||||
ProjectionExec: expr=[vector@0 as vector, text@3 as text, _distance@2 as _distance]
|
||||
Take: columns="vector, _rowid, _distance, (text)"
|
||||
CoalesceBatchesExec: target_batch_size=1024
|
||||
LanceRead: uri=..., projection=[text], source=stream(_rowid)
|
||||
GlobalLimitExec: skip=0, fetch=10
|
||||
FilterExec: _distance@2 IS NOT NULL
|
||||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||||
KNNVectorDistance: metric=l2
|
||||
LanceRead: uri=..., projection=[vector], ...
|
||||
ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score]
|
||||
Take: columns="_rowid, _score, (vector), (text)"
|
||||
CoalesceBatchesExec: target_batch_size=1024
|
||||
LanceRead: uri=..., projection=[vector, text], source=stream(_rowid)
|
||||
GlobalLimitExec: skip=0, fetch=10
|
||||
MatchQuery: column=text, query=[hello]
|
||||
|
||||
|
||||
@@ -13,14 +13,10 @@ from typing import (
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
TYPE_CHECKING,
|
||||
Union,
|
||||
Literal,
|
||||
overload,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..udf import Job
|
||||
import warnings
|
||||
|
||||
from lancedb import __version__
|
||||
@@ -578,6 +574,7 @@ class RemoteTable(Table):
|
||||
on_bad_vectors: str = "error",
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table](Table). It has the same API signature as
|
||||
the OSS version.
|
||||
@@ -603,6 +600,12 @@ class RemoteTable(Table):
|
||||
progress: bool, callable, or tqdm-like, optional
|
||||
A callback or tqdm-compatible progress bar. See
|
||||
:meth:`Table.add` for details.
|
||||
write_parallelism: int, optional
|
||||
Number of partitions to write in parallel. Higher values increase
|
||||
throughput but also peak memory use, since each partition buffers
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -618,6 +621,7 @@ class RemoteTable(Table):
|
||||
on_bad_vectors=on_bad_vectors,
|
||||
fill_value=fill_value,
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
@@ -918,142 +922,8 @@ 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: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
computed: Optional[Dict[str, tuple]] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
result = None
|
||||
if transforms is not None:
|
||||
result = LOOP.run(self._table.add_columns(transforms))
|
||||
if computed:
|
||||
LOOP.run(self._table.add_columns(computed=computed))
|
||||
return result
|
||||
|
||||
def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> "Job":
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
|
||||
The expression is resolved server-side from each column's stored
|
||||
binding; columns bound to the same struct-returning function
|
||||
refresh together. Returns a `Job` to wait on, poll, or cancel
|
||||
(``tbl.refresh_column("c").wait()``). Server-backed feature
|
||||
(LanceDB Enterprise / Cloud).
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh) and override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill).
|
||||
"""
|
||||
from ..udf import Job
|
||||
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
job_id = LOOP.run(
|
||||
self._table.refresh_column(
|
||||
list(columns),
|
||||
where=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
return Job(self._job_conn(), job_id)
|
||||
|
||||
def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Derived-compute lineage of this table, or one of its columns:
|
||||
upstream sources, downstream dependents, and the function version +
|
||||
location that produced each derived column (with a drift flag). Returns
|
||||
a `Lineage`. See `Connection.lineage`."""
|
||||
return self._job_conn().lineage(
|
||||
self._name, column, direction=direction, depth=depth
|
||||
)
|
||||
|
||||
def _job_conn(self):
|
||||
"""A client connection for polling jobs this table spawns. Built lazily
|
||||
from the table's serialized connection state and cached (not pickled --
|
||||
a forked/unpickled table rebuilds it on next use)."""
|
||||
from lancedb import deserialize_conn
|
||||
|
||||
conn = getattr(self, "_job_conn_cache", None)
|
||||
if conn is None:
|
||||
conn = deserialize_conn(self._serialized_connection_state())
|
||||
self._job_conn_cache = conn
|
||||
return conn
|
||||
|
||||
def load_columns(
|
||||
self,
|
||||
source: Union[str, Iterable[str]],
|
||||
pk: str,
|
||||
columns: Union[Iterable[str], Dict[str, str]],
|
||||
*,
|
||||
source_format: str = "parquet",
|
||||
source_pk: Optional[str] = None,
|
||||
on_missing: str = "carry",
|
||||
source_storage_options: Optional[Dict[str, str]] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
commit_granularity: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Fill existing columns from an external source by primary-key join.
|
||||
|
||||
The distributed-job equivalent of Geneva's ``Table.load_columns()``:
|
||||
imports precomputed values (e.g. embeddings) from Parquet/Lance/IPC into
|
||||
this table, matching on a primary key. Returns the load job id.
|
||||
Server-backed feature (LanceDB Enterprise / Cloud).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source: str | list[str]
|
||||
One source URI or a list of URIs.
|
||||
pk: str
|
||||
Destination primary-key column. Also the source key unless
|
||||
``source_pk`` is given.
|
||||
columns: list[str] | dict[str, str]
|
||||
Value columns to load. A list loads same-named columns; a dict maps
|
||||
``{target: source}``.
|
||||
source_format: str
|
||||
``"parquet"`` (default), ``"lance"``, or ``"ipc"``.
|
||||
source_pk: str, optional
|
||||
Source primary-key column when it differs from ``pk``.
|
||||
on_missing: str
|
||||
Behavior for destination rows with no source match:
|
||||
``"carry"`` (default, keep existing), ``"null"``, or ``"error"``.
|
||||
"""
|
||||
if isinstance(source, str):
|
||||
source = [source]
|
||||
if isinstance(columns, dict):
|
||||
mappings = [(target, src) for target, src in columns.items()]
|
||||
else:
|
||||
mappings = [(c, None) for c in columns]
|
||||
return LOOP.run(
|
||||
self._table.load_columns(
|
||||
list(source),
|
||||
source_format,
|
||||
pk,
|
||||
mappings,
|
||||
source_key=source_pk,
|
||||
source_storage_options=source_storage_options,
|
||||
on_missing=on_missing,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
commit_granularity=commit_granularity,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
def add_columns(self, transforms: Dict[str, str]) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms))
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
|
||||
@@ -40,12 +40,12 @@ class WatsonxReranker(Reranker):
|
||||
IBM Cloud API key. Falls back to the ``WATSONX_API_KEY`` environment
|
||||
variable when not provided.
|
||||
project_id : str, optional
|
||||
watsonx.ai project ID. Falls back to the ``WATSONX_PROJECT_ID``
|
||||
environment variable when not provided. Mutually exclusive with
|
||||
watsonx.ai project ID. Explicit value takes precedence over the
|
||||
``WATSONX_PROJECT_ID`` environment variable. Mutually exclusive with
|
||||
``space_id`` — exactly one must be supplied.
|
||||
space_id : str, optional
|
||||
watsonx.ai deployment space ID. Falls back to the ``WATSONX_SPACE_ID``
|
||||
environment variable when not provided. Mutually exclusive with
|
||||
watsonx.ai deployment space ID. Explicit value takes precedence over
|
||||
the ``WATSONX_SPACE_ID`` environment variable. Mutually exclusive with
|
||||
``project_id`` — exactly one must be supplied.
|
||||
url : str, optional
|
||||
watsonx.ai service URL. Defaults to
|
||||
@@ -100,8 +100,16 @@ class WatsonxReranker(Reranker):
|
||||
)
|
||||
|
||||
# --- project_id / space_id (exactly one required) ---
|
||||
project_id = self.project_id or os.environ.get("WATSONX_PROJECT_ID")
|
||||
space_id = self.space_id or os.environ.get("WATSONX_SPACE_ID")
|
||||
# Explicit field always wins; env vars are consulted only when neither
|
||||
# was passed explicitly, so a stray WATSONX_SPACE_ID never overrides an
|
||||
# explicit project_id and vice-versa.
|
||||
project_id = self.project_id
|
||||
space_id = self.space_id
|
||||
|
||||
if project_id is None and space_id is None:
|
||||
# Neither was passed explicitly — fall back to env vars.
|
||||
project_id = os.environ.get("WATSONX_PROJECT_ID")
|
||||
space_id = os.environ.get("WATSONX_SPACE_ID")
|
||||
|
||||
if project_id and space_id:
|
||||
raise ValueError("Provide either `project_id` or `space_id`, not both.")
|
||||
|
||||
@@ -7,10 +7,158 @@ import sys
|
||||
from typing import Callable, Iterator, Optional
|
||||
from lancedb.arrow import to_arrow
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
import pyarrow.dataset as ds
|
||||
|
||||
from .pydantic import LanceModel
|
||||
|
||||
# pyarrow's default scanner settings are tuned for narrow rows. For wide rows
|
||||
# (e.g. embedding columns) they buffer a huge read-ahead window in host memory
|
||||
# and can OOM the client during bulk ingestion. We size the scanner so the
|
||||
# estimated in-flight memory stays within a budget, while leaving narrow
|
||||
# datasets on pyarrow's defaults (no throughput regression).
|
||||
_SCAN_MEMORY_BUDGET_BYTES = 1024 * 1024 * 1024 # ~1 GiB in-flight target
|
||||
_TARGET_BATCH_BYTES = 16 * 1024 * 1024 # ~16 MiB per batch
|
||||
_MIN_BATCH_ROWS = 512
|
||||
# pyarrow defaults (see arrow/dataset ScanOptions); we never exceed these.
|
||||
_PA_DEFAULT_BATCH_ROWS = 131_072
|
||||
_PA_DEFAULT_BATCH_READAHEAD = 16
|
||||
_PA_DEFAULT_FRAGMENT_READAHEAD = 4
|
||||
# Read-ahead used for wide rows. pyarrow reads a whole parquet row group at a
|
||||
# time and keeps `batch_readahead` of them resident, so read-ahead depth (not
|
||||
# batch size) dominates peak memory for wide data; keep both small but leave a
|
||||
# little prefetch for throughput. Tuned empirically against embedding datasets.
|
||||
_WIDE_BATCH_READAHEAD = 2
|
||||
_WIDE_FRAGMENT_READAHEAD = 1
|
||||
# Estimate for variable-width columns (string/binary/list) whose true width is
|
||||
# unknown from the schema alone. Only needs to be large enough to flag "wide".
|
||||
_VARIABLE_WIDTH_ESTIMATE = 128
|
||||
# Rows peeked from a rescannable source to refine the list-length guess for
|
||||
# variable-length list columns (e.g. embeddings stored as `list<float32>`
|
||||
# instead of `list<float32, N>`), whose per-row width the schema can't tell us.
|
||||
_SAMPLE_ROWS = 10
|
||||
|
||||
|
||||
def _observed_list_length(sample: pa.ChunkedArray) -> Optional[int]:
|
||||
"""Average element count per row observed in a list/large_list sample."""
|
||||
if len(sample) == 0:
|
||||
return None
|
||||
mean = pc.mean(pc.list_value_length(sample)).as_py()
|
||||
return None if mean is None else max(1, round(mean))
|
||||
|
||||
|
||||
def _estimate_field_width(
|
||||
dtype: pa.DataType, sample: Optional[pa.ChunkedArray] = None
|
||||
) -> int:
|
||||
if pa.types.is_fixed_size_list(dtype):
|
||||
return dtype.list_size * _estimate_field_width(dtype.value_type)
|
||||
if pa.types.is_struct(dtype):
|
||||
return sum(
|
||||
_estimate_field_width(
|
||||
dtype.field(i).type,
|
||||
pc.struct_field(sample, [i]) if sample is not None else None,
|
||||
)
|
||||
for i in range(dtype.num_fields)
|
||||
)
|
||||
if pa.types.is_dictionary(dtype):
|
||||
return _estimate_field_width(dtype.value_type)
|
||||
if pa.types.is_fixed_size_binary(dtype):
|
||||
return dtype.byte_width
|
||||
if pa.types.is_boolean(dtype):
|
||||
return 1
|
||||
if (pa.types.is_list(dtype) or pa.types.is_large_list(dtype)) and (
|
||||
sample is not None
|
||||
):
|
||||
observed_length = _observed_list_length(sample)
|
||||
if observed_length is not None:
|
||||
return observed_length * _estimate_field_width(dtype.value_type)
|
||||
# Fixed-width scalars (ints, floats, temporal, decimal) expose bit_width;
|
||||
# variable-width types (string, binary, list, map, ...) raise ValueError.
|
||||
try:
|
||||
return max(1, dtype.bit_width // 8)
|
||||
except (ValueError, AttributeError):
|
||||
return _VARIABLE_WIDTH_ESTIMATE
|
||||
|
||||
|
||||
def _estimate_bytes_per_row(
|
||||
schema: pa.Schema, sample: Optional[pa.Table] = None
|
||||
) -> int:
|
||||
return max(
|
||||
1,
|
||||
sum(
|
||||
_estimate_field_width(
|
||||
field.type, sample.column(field.name) if sample is not None else None
|
||||
)
|
||||
for field in schema
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _sample_head(head: Callable[..., pa.Table]) -> Optional[pa.Table]:
|
||||
"""Best-effort peek at a few rows to refine the bytes-per-row estimate.
|
||||
|
||||
Uses a tight batch size and no read-ahead so the peek itself can't trigger
|
||||
the wide-row memory blowup this module exists to avoid. Returns None (fall
|
||||
back to the schema-only estimate) if sampling isn't possible for any
|
||||
reason, e.g. an empty dataset.
|
||||
"""
|
||||
try:
|
||||
sample = head(
|
||||
_SAMPLE_ROWS,
|
||||
batch_size=_SAMPLE_ROWS,
|
||||
batch_readahead=1,
|
||||
fragment_readahead=1,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
return sample if sample.num_rows > 0 else None
|
||||
|
||||
|
||||
def _bounded_scanner_kwargs(
|
||||
schema: pa.Schema, sample: Optional[pa.Table] = None
|
||||
) -> dict:
|
||||
"""Scanner kwargs that cap in-flight memory for wide rows.
|
||||
|
||||
Narrow datasets keep pyarrow's defaults unchanged (no throughput
|
||||
regression). For wide rows (e.g. embedding columns) pyarrow's default
|
||||
read-ahead buffers many large batches/row-groups at once, which can OOM the
|
||||
client during bulk ingestion, so we shrink the batch size and read-ahead to
|
||||
keep the estimated in-flight memory near the budget.
|
||||
|
||||
Read-ahead (not just batch size) has to drop: pyarrow reads a whole parquet
|
||||
row group at a time and keeps `batch_readahead`/`fragment_readahead` of them
|
||||
resident, so a small batch size alone still pins large row-group buffers.
|
||||
|
||||
`sample`, if given, is a small (see `_SAMPLE_ROWS`) table of rows from the
|
||||
source used to refine the estimate for variable-length list columns (e.g.
|
||||
embeddings stored without a fixed size), whose width the schema alone
|
||||
can't tell us.
|
||||
"""
|
||||
bytes_per_row = _estimate_bytes_per_row(schema, sample)
|
||||
|
||||
# If pyarrow's defaults already stay within budget, leave them alone so
|
||||
# narrow datasets keep their throughput. A "unit" of in-flight memory is one
|
||||
# default-sized batch, held `batch_readahead + fragment_readahead` deep.
|
||||
default_in_flight = (
|
||||
_PA_DEFAULT_BATCH_ROWS
|
||||
* bytes_per_row
|
||||
* (_PA_DEFAULT_BATCH_READAHEAD + _PA_DEFAULT_FRAGMENT_READAHEAD)
|
||||
)
|
||||
if default_in_flight <= _SCAN_MEMORY_BUDGET_BYTES:
|
||||
return {}
|
||||
|
||||
# Wide rows: cap batch bytes and pull read-ahead down so only a couple of
|
||||
# large row-group buffers are resident at once.
|
||||
batch_size = min(
|
||||
_PA_DEFAULT_BATCH_ROWS,
|
||||
max(_MIN_BATCH_ROWS, _TARGET_BATCH_BYTES // bytes_per_row),
|
||||
)
|
||||
return {
|
||||
"batch_size": batch_size,
|
||||
"batch_readahead": _WIDE_BATCH_READAHEAD,
|
||||
"fragment_readahead": _WIDE_FRAGMENT_READAHEAD,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Scannable:
|
||||
@@ -56,10 +204,12 @@ def _from_table(data: pa.Table) -> Scannable:
|
||||
|
||||
@to_scannable.register(ds.Dataset)
|
||||
def _from_dataset(data: ds.Dataset) -> Scannable:
|
||||
sample = _sample_head(data.head)
|
||||
scanner_kwargs = _bounded_scanner_kwargs(data.schema, sample)
|
||||
return Scannable(
|
||||
schema=data.schema,
|
||||
num_rows=data.count_rows(),
|
||||
reader=lambda: data.scanner().to_reader(),
|
||||
reader=lambda: data.scanner(**scanner_kwargs).to_reader(),
|
||||
)
|
||||
|
||||
|
||||
@@ -206,10 +356,12 @@ def _register_optional_converters():
|
||||
|
||||
@to_scannable.register(lance.LanceDataset)
|
||||
def _from_lance(data: lance.LanceDataset) -> Scannable:
|
||||
sample = _sample_head(data.head)
|
||||
scanner_kwargs = _bounded_scanner_kwargs(data.schema, sample)
|
||||
return Scannable(
|
||||
schema=data.schema,
|
||||
num_rows=data.count_rows(),
|
||||
reader=lambda: data.scanner().to_reader(),
|
||||
reader=lambda: data.scanner(**scanner_kwargs).to_reader(),
|
||||
)
|
||||
|
||||
|
||||
|
||||
+87
-310
@@ -162,7 +162,6 @@ def _maybe_add_fts_error_note(
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .db import LanceDBConnection
|
||||
from .udf import Job
|
||||
from ._lancedb import (
|
||||
Table as LanceDBTable,
|
||||
OptimizeStats,
|
||||
@@ -703,24 +702,6 @@ def _normalize_progress(progress):
|
||||
return progress, False
|
||||
|
||||
|
||||
def _computed_groups(computed):
|
||||
"""Group computed columns by expression, preserving declaration order
|
||||
(struct-returning functions need their columns adjacent so schema order
|
||||
matches field order). Accepts the ergonomic forms -- `fn("col")` values
|
||||
and tuple keys for struct fan-out -- via `_normalize_computed`."""
|
||||
from .udf import _normalize_computed
|
||||
|
||||
groups = []
|
||||
for name, (sql_type, expression) in _normalize_computed(computed).items():
|
||||
for expr, cols in groups:
|
||||
if expr == expression:
|
||||
cols.append((name, sql_type))
|
||||
break
|
||||
else:
|
||||
groups.append((expression, [(name, sql_type)]))
|
||||
return groups
|
||||
|
||||
|
||||
class Table(ABC):
|
||||
"""
|
||||
A Table is a collection of Records in a LanceDB Database.
|
||||
@@ -826,59 +807,6 @@ class Table(ABC):
|
||||
"""The number of rows in this Table"""
|
||||
return self.count_rows(None)
|
||||
|
||||
def add_computed_column(
|
||||
self,
|
||||
columns,
|
||||
fn,
|
||||
args: Optional[List[str]] = None,
|
||||
types=None,
|
||||
) -> None:
|
||||
"""Declare computed column(s) bound to a UDF -- no compute happens
|
||||
here (the agent fills them lazily, or refresh_column() triggers a run).
|
||||
|
||||
.. deprecated::
|
||||
A computed column is an expression over a registered function, so
|
||||
bind it as one: ``add_columns(computed={"vec": embed("data")})``.
|
||||
``embed("data")`` applies the function to the `data` column and
|
||||
infers the type from the function's return signature -- the
|
||||
function never couples to a particular column. Prefer that form.
|
||||
"""
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"add_computed_column is deprecated; use add_columns(computed="
|
||||
'{"vec": embed("data")}).',
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
from .udf import Udf, struct_field_types
|
||||
|
||||
multi = isinstance(columns, (tuple, list))
|
||||
if isinstance(fn, Udf):
|
||||
expr = fn.expression(*(args or []))
|
||||
if types is None:
|
||||
if multi:
|
||||
if not fn.returns.upper().startswith("STRUCT"):
|
||||
raise ValueError(
|
||||
"several columns need a STRUCT-returning function"
|
||||
)
|
||||
types = struct_field_types(fn.returns)
|
||||
else:
|
||||
types = fn.returns
|
||||
else:
|
||||
if types is None:
|
||||
raise ValueError("pass types= when fn is a name string")
|
||||
expr = f"{fn}({', '.join(args or [])})"
|
||||
if multi:
|
||||
if len(types) != len(columns):
|
||||
raise ValueError(
|
||||
f"{len(columns)} columns but {len(types)} output types"
|
||||
)
|
||||
computed = {c: (t, expr) for c, t in zip(columns, types)}
|
||||
else:
|
||||
computed = {columns: (types, expr)}
|
||||
self.add_columns(computed=computed)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def embedding_functions(self) -> Dict[str, EmbeddingFunctionConfig]:
|
||||
@@ -955,7 +883,7 @@ class Table(ABC):
|
||||
wait_timeout: Optional[timedelta] = ...,
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
) -> "Job": ...
|
||||
) -> None: ...
|
||||
|
||||
# Legacy API overload (deprecated)
|
||||
@overload
|
||||
@@ -979,7 +907,7 @@ class Table(ABC):
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
target_partition_size: Optional[int] = ...,
|
||||
) -> "Job": ...
|
||||
) -> None: ...
|
||||
|
||||
def create_index(
|
||||
self,
|
||||
@@ -1030,14 +958,6 @@ class Table(ABC):
|
||||
train : bool, default True
|
||||
Whether to train the index with existing data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Job
|
||||
A handle on the index build. When the server defers the build to a
|
||||
background job, ``job.wait()`` blocks until it completes; when the
|
||||
build finished within this call, the job is already ``finished``.
|
||||
Prefer ``job.wait()`` over the deprecated ``wait_timeout``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
New API (recommended):
|
||||
@@ -1279,6 +1199,7 @@ class Table(ABC):
|
||||
on_bad_vectors: OnBadVectorsType = "error",
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table](Table).
|
||||
|
||||
@@ -1324,6 +1245,13 @@ class Table(ABC):
|
||||
with tqdm() as pbar:
|
||||
table.add(data, progress=pbar)
|
||||
|
||||
write_parallelism: int, optional
|
||||
Number of partitions to write in parallel. Higher values increase
|
||||
throughput but also peak memory use, since each partition buffers
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
|
||||
Returns
|
||||
-------
|
||||
AddResult
|
||||
@@ -2205,8 +2133,8 @@ class LanceTable(Table):
|
||||
def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import LanceDBConnection
|
||||
|
||||
async_tbl = AsyncTable(tbl)
|
||||
conn = LanceDBConnection.from_inner(tbl.database())
|
||||
async_tbl = AsyncTable(tbl, conn=conn._conn)
|
||||
return cls(
|
||||
conn,
|
||||
async_tbl.name,
|
||||
@@ -2275,7 +2203,7 @@ class LanceTable(Table):
|
||||
namespace_client = self._namespace_client
|
||||
if namespace_client is None:
|
||||
conn_uri = getattr(self._conn, "uri", "")
|
||||
if get_uri_scheme(conn_uri) == "namespace":
|
||||
if get_uri_scheme(conn_uri) == "namespace" or self._namespace_path:
|
||||
namespace_client = self._conn.namespace_client()
|
||||
self._namespace_client = namespace_client
|
||||
|
||||
@@ -2608,7 +2536,7 @@ class LanceTable(Table):
|
||||
wait_timeout: Optional[timedelta] = ...,
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
) -> "Job": ...
|
||||
) -> None: ...
|
||||
|
||||
# Legacy API overload (deprecated)
|
||||
@overload
|
||||
@@ -2634,7 +2562,7 @@ class LanceTable(Table):
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
target_partition_size: Optional[int] = ...,
|
||||
) -> "Job": ...
|
||||
) -> None: ...
|
||||
|
||||
def create_index(
|
||||
self,
|
||||
@@ -2693,14 +2621,6 @@ class LanceTable(Table):
|
||||
train : bool, default True
|
||||
Whether to train the index with existing data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Job
|
||||
A handle on the index build. When the server defers the build to a
|
||||
background job, ``job.wait()`` blocks until it completes; when the
|
||||
build finished within this call, the job is already ``finished``.
|
||||
Prefer ``job.wait()`` over the deprecated ``wait_timeout``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
New API (recommended):
|
||||
@@ -2776,7 +2696,7 @@ class LanceTable(Table):
|
||||
target_partition_size=target_partition_size,
|
||||
)
|
||||
self.checkout_latest()
|
||||
return self._sync_job(None)
|
||||
return
|
||||
else:
|
||||
# New API: metric is the column name
|
||||
column = metric
|
||||
@@ -2813,30 +2733,19 @@ class LanceTable(Table):
|
||||
),
|
||||
)
|
||||
self.checkout_latest()
|
||||
return self._sync_job(None)
|
||||
return
|
||||
|
||||
return self._sync_job(
|
||||
LOOP.run(
|
||||
self._table.create_index(
|
||||
column,
|
||||
replace=replace,
|
||||
config=config,
|
||||
wait_timeout=wait_timeout,
|
||||
name=name,
|
||||
train=train,
|
||||
)
|
||||
return LOOP.run(
|
||||
self._table.create_index(
|
||||
column,
|
||||
replace=replace,
|
||||
config=config,
|
||||
wait_timeout=wait_timeout,
|
||||
name=name,
|
||||
train=train,
|
||||
)
|
||||
)
|
||||
|
||||
def _sync_job(self, ajob) -> "Job":
|
||||
"""Convert an AsyncJob (or None for work done in-process) into a sync
|
||||
Job bound to this table's connection."""
|
||||
from .udf import Job
|
||||
|
||||
if ajob is not None and ajob.id:
|
||||
return Job(self._conn, ajob.id, table=self.name)
|
||||
return Job._completed(self._conn, table=self.name)
|
||||
|
||||
def _is_legacy_create_index_call(
|
||||
self,
|
||||
first_arg: str,
|
||||
@@ -3087,12 +2996,8 @@ class LanceTable(Table):
|
||||
config = LabelList()
|
||||
else:
|
||||
raise ValueError(f"Unknown index type {index_type}")
|
||||
return self._sync_job(
|
||||
LOOP.run(
|
||||
self._table.create_index(
|
||||
column, replace=replace, config=config, name=name
|
||||
)
|
||||
)
|
||||
return LOOP.run(
|
||||
self._table.create_index(column, replace=replace, config=config, name=name)
|
||||
)
|
||||
|
||||
@deprecation.deprecated(
|
||||
@@ -3175,7 +3080,7 @@ class LanceTable(Table):
|
||||
)
|
||||
|
||||
try:
|
||||
ajob = LOOP.run(
|
||||
LOOP.run(
|
||||
self._table.create_index(
|
||||
field_names,
|
||||
replace=replace,
|
||||
@@ -3190,7 +3095,6 @@ class LanceTable(Table):
|
||||
language=config.language,
|
||||
)
|
||||
raise e
|
||||
return self._sync_job(ajob)
|
||||
|
||||
@staticmethod
|
||||
def infer_tokenizer_configs(tokenizer_name: str) -> dict:
|
||||
@@ -3262,6 +3166,7 @@ class LanceTable(Table):
|
||||
on_bad_vectors: OnBadVectorsType = "error",
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add data to the table.
|
||||
If vector columns are missing and the table
|
||||
@@ -3283,6 +3188,12 @@ class LanceTable(Table):
|
||||
progress: bool, callable, or tqdm-like, optional
|
||||
A callback or tqdm-compatible progress bar. See
|
||||
:meth:`Table.add` for details.
|
||||
write_parallelism: int, optional
|
||||
Number of partitions to write in parallel. Higher values increase
|
||||
throughput but also peak memory use, since each partition buffers
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -3298,6 +3209,7 @@ class LanceTable(Table):
|
||||
on_bad_vectors=on_bad_vectors,
|
||||
fill_value=fill_value,
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
@@ -3919,68 +3831,9 @@ 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
|
||||
| None = None,
|
||||
*,
|
||||
computed: Optional[Dict] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
result = None
|
||||
if transforms is not None:
|
||||
result = LOOP.run(self._table.add_columns(transforms))
|
||||
if computed:
|
||||
# computed binds an expression over a registered function to a
|
||||
# column: {col: fn("input_col")} -- fn("input_col") yields the
|
||||
# expression and carries the inferred type; a tuple key fans a
|
||||
# STRUCT return out to several columns. Declares the binding only;
|
||||
# the server fills the values (server-backed). The legacy
|
||||
# {col: (sql_type, expression)} tuple form is still accepted.
|
||||
result_unused = LOOP.run(self._table.add_columns(computed=computed))
|
||||
del result_unused
|
||||
return result
|
||||
|
||||
def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> "Job":
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
|
||||
The expression is resolved server-side from each column's stored
|
||||
binding; columns bound to the same struct-returning function
|
||||
refresh together. Returns a `Job` to wait on, poll, or cancel
|
||||
(``tbl.refresh_column("col").wait()``) -- mirrors
|
||||
`MaterializedView.refresh()`. Server-backed feature (LanceDB
|
||||
Enterprise / Cloud).
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh) and override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill).
|
||||
"""
|
||||
from .udf import Job
|
||||
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
job_id = LOOP.run(
|
||||
self._table.refresh_column(
|
||||
list(columns),
|
||||
where=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
return Job(self._conn, job_id, table=self.name)
|
||||
self, transforms: Dict[str, str] | pa.field | List[pa.field] | pa.Schema
|
||||
) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms))
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
@@ -4596,7 +4449,6 @@ class AsyncTable:
|
||||
self,
|
||||
table: LanceDBTable,
|
||||
*,
|
||||
conn: Optional[Any] = None,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
namespace_client: Optional[Any] = None,
|
||||
pushdown_operations: Optional[set] = None,
|
||||
@@ -4610,9 +4462,6 @@ class AsyncTable:
|
||||
[AsyncConnection.open_table][lancedb.AsyncConnection.open_table] to obtain
|
||||
Table objects."""
|
||||
self._inner = table
|
||||
#: The owning AsyncConnection, when known -- lets index/refresh calls
|
||||
#: hand back AsyncJob handles that can reach the platform jobs API.
|
||||
self._conn = conn
|
||||
self._namespace_path = namespace_path or []
|
||||
self._namespace_client = namespace_client
|
||||
self._pushdown_operations = pushdown_operations or set()
|
||||
@@ -4920,14 +4769,6 @@ class AsyncTable:
|
||||
train: bool, default True
|
||||
Whether to train the index with existing data. Vector indices always train
|
||||
with existing data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
AsyncJob
|
||||
A handle on the index build. When the server defers the build to a
|
||||
background job, ``await job.wait()`` blocks until it completes;
|
||||
when the build finished within this call, the job is already
|
||||
``finished``. Prefer ``await job.wait()`` over ``wait_timeout``.
|
||||
"""
|
||||
if config is not None:
|
||||
if not isinstance(
|
||||
@@ -4953,7 +4794,7 @@ class AsyncTable:
|
||||
+ str(type(config))
|
||||
)
|
||||
try:
|
||||
job_id = await self._inner.create_index(
|
||||
await self._inner.create_index(
|
||||
column,
|
||||
index=config,
|
||||
replace=replace,
|
||||
@@ -4970,12 +4811,6 @@ class AsyncTable:
|
||||
)
|
||||
raise e
|
||||
|
||||
from .udf import AsyncJob
|
||||
|
||||
if job_id:
|
||||
return AsyncJob(self._conn, job_id, table=self.name)
|
||||
return AsyncJob._completed(self._conn, table=self.name)
|
||||
|
||||
async def drop_index(self, name: str) -> None:
|
||||
"""
|
||||
Drop an index from the table.
|
||||
@@ -5117,6 +4952,7 @@ class AsyncTable:
|
||||
on_bad_vectors: Optional[OnBadVectorsType] = None,
|
||||
fill_value: Optional[float] = None,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table](Table).
|
||||
|
||||
@@ -5141,6 +4977,12 @@ class AsyncTable:
|
||||
progress: callable or tqdm-like, optional
|
||||
A callback or tqdm-compatible progress bar. See
|
||||
:meth:`Table.add` for details.
|
||||
write_parallelism: int, optional
|
||||
Number of partitions to write in parallel. Higher values increase
|
||||
throughput but also peak memory use, since each partition buffers
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
|
||||
"""
|
||||
schema = await self.schema()
|
||||
@@ -5172,7 +5014,12 @@ class AsyncTable:
|
||||
data = to_scannable(data)
|
||||
progress, owns = _normalize_progress(progress)
|
||||
try:
|
||||
return await self._inner.add(data, mode or "append", progress=progress)
|
||||
return await self._inner.add(
|
||||
data,
|
||||
mode or "append",
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Cast error" in str(e):
|
||||
raise ValueError(e)
|
||||
@@ -5751,44 +5598,9 @@ class AsyncTable:
|
||||
|
||||
return await self._inner.update(updates_sql, where)
|
||||
|
||||
async def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
Returns the refresh job id. Server-backed feature.
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh); they override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill)."""
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
return await self._inner.refresh_column(
|
||||
list(columns),
|
||||
where_clause=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
|
||||
async def add_columns(
|
||||
self,
|
||||
transforms: dict[str, str]
|
||||
| pa.field
|
||||
| List[pa.field]
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: Optional[Dict] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
self, transforms: dict[str, str] | pa.field | List[pa.field] | pa.Schema
|
||||
) -> AddColumnsResult:
|
||||
"""
|
||||
Add new columns with defined values.
|
||||
|
||||
@@ -5807,7 +5619,6 @@ class AsyncTable:
|
||||
version: the new version number of the table after adding columns.
|
||||
|
||||
"""
|
||||
result = None
|
||||
if isinstance(transforms, pa.Field):
|
||||
transforms = [transforms]
|
||||
if isinstance(transforms, list) and all(
|
||||
@@ -5815,69 +5626,9 @@ class AsyncTable:
|
||||
):
|
||||
transforms = pa.schema(transforms)
|
||||
if isinstance(transforms, pa.Schema):
|
||||
result = await self._inner.add_columns_with_schema(transforms)
|
||||
elif transforms is not None:
|
||||
result = await self._inner.add_columns(list(transforms.items()))
|
||||
if computed:
|
||||
# computed binds an expression over a registered function to a
|
||||
# column: {col: fn("input_col")} -- fn("input_col") yields the
|
||||
# expression and carries the inferred type; a tuple key fans a
|
||||
# STRUCT return out to several columns. Declares the binding only;
|
||||
# the server fills the values (server-backed). The legacy
|
||||
# {col: (sql_type, expression)} tuple form is still accepted.
|
||||
for expression, cols in _computed_groups(computed):
|
||||
await self._inner.add_computed_columns(cols, expression)
|
||||
return result
|
||||
|
||||
async def add_computed_column(
|
||||
self,
|
||||
columns,
|
||||
fn,
|
||||
args: Optional[List[str]] = None,
|
||||
types=None,
|
||||
) -> None:
|
||||
"""Declare computed column(s) bound to a UDF (async).
|
||||
|
||||
.. deprecated::
|
||||
Use ``add_columns(computed={"col": fn("input_col")})`` -- a computed
|
||||
column is an expression over a registered function, so bind it that
|
||||
way instead of coupling the UDF to the column here.
|
||||
"""
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"add_computed_column is deprecated; use add_columns(computed="
|
||||
'{"col": fn("input_col")}).',
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
from .udf import Udf, struct_field_types
|
||||
|
||||
multi = isinstance(columns, (tuple, list))
|
||||
if isinstance(fn, Udf):
|
||||
expr = fn.expression(*(args or []))
|
||||
if types is None:
|
||||
if multi:
|
||||
if not fn.returns.upper().startswith("STRUCT"):
|
||||
raise ValueError(
|
||||
"several columns need a STRUCT-returning function"
|
||||
)
|
||||
types = struct_field_types(fn.returns)
|
||||
else:
|
||||
types = fn.returns
|
||||
return await self._inner.add_columns_with_schema(transforms)
|
||||
else:
|
||||
if types is None:
|
||||
raise ValueError("pass types= when fn is a name string")
|
||||
expr = f"{fn}({', '.join(args or [])})"
|
||||
if multi:
|
||||
if len(types) != len(columns):
|
||||
raise ValueError(
|
||||
f"{len(columns)} columns but {len(types)} output types"
|
||||
)
|
||||
computed = {c: (t, expr) for c, t in zip(columns, types)}
|
||||
else:
|
||||
computed = {columns: (types, expr)}
|
||||
await self.add_columns(computed=computed)
|
||||
return await self._inner.add_columns(list(transforms.items()))
|
||||
|
||||
async def alter_columns(
|
||||
self, *alterations: Iterable[dict[str, Any]]
|
||||
@@ -6516,6 +6267,24 @@ class Branches:
|
||||
"""Delete a branch."""
|
||||
LOOP.run(self._table.branches.delete(name))
|
||||
|
||||
def diff(self, from_branch: str) -> Dict[str, Any]:
|
||||
"""Diff a branch against main."""
|
||||
return LOOP.run(self._table.branches.diff(from_branch))
|
||||
|
||||
def merge(self, from_branch: str, dry_run: bool = False) -> Dict[str, Any]:
|
||||
"""Merge a branch into main, or dry-run.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
from_branch: str
|
||||
Branch to merge from.
|
||||
dry_run: bool, default False
|
||||
When True, only preview. When False, attempt the merge.
|
||||
|
||||
A rejected merge returns ``status="rejected"`` instead of raising.
|
||||
"""
|
||||
return LOOP.run(self._table.branches.merge(from_branch, dry_run))
|
||||
|
||||
def _wrap(
|
||||
self, async_table: "AsyncTable", version: Optional[int] = None
|
||||
) -> "Table":
|
||||
@@ -6626,7 +6395,7 @@ class AsyncBranches:
|
||||
if from_ref == "main":
|
||||
from_ref = None
|
||||
inner = await self._table.branches.create(name, from_ref, from_version)
|
||||
return AsyncTable(inner, conn=self._table._conn)
|
||||
return AsyncTable(inner)
|
||||
|
||||
async def checkout(self, name: str, version: Optional[int] = None) -> "AsyncTable":
|
||||
"""Check out an existing branch and return a handle scoped to it.
|
||||
@@ -6640,11 +6409,19 @@ class AsyncBranches:
|
||||
handle is a read-only view of that version; when omitted it tracks
|
||||
the branch's latest and stays writable.
|
||||
"""
|
||||
return AsyncTable(
|
||||
await self._table.branches.checkout(name, version),
|
||||
conn=self._table._conn,
|
||||
)
|
||||
return AsyncTable(await self._table.branches.checkout(name, version))
|
||||
|
||||
async def delete(self, name: str) -> None:
|
||||
"""Delete a branch."""
|
||||
await self._table.branches.delete(name)
|
||||
|
||||
async def diff(self, from_branch: str) -> Dict[str, Any]:
|
||||
"""Diff a branch against main."""
|
||||
return await self._table.branches.diff(from_branch)
|
||||
|
||||
async def merge(self, from_branch: str, dry_run: bool = False) -> Dict[str, Any]:
|
||||
"""Merge a branch into main, or dry-run.
|
||||
|
||||
A rejected merge returns ``status="rejected"`` instead of raising.
|
||||
"""
|
||||
return await self._table.branches.merge(from_branch, dry_run)
|
||||
|
||||
@@ -1,847 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""UDF authoring for LanceDB derived compute (server-backed).
|
||||
|
||||
`@udf` / `@table_udf` turn a plain Python function into a registrable
|
||||
server-side UDF: a cloudpickled (or source) body, a SQL signature inferred
|
||||
from type hints, and the runtime options (pip deps, GPUs, batching, ...).
|
||||
Register and use them through the existing connection/table API:
|
||||
|
||||
import lancedb
|
||||
from lancedb import udf, table_udf
|
||||
|
||||
db = lancedb.connect("db://my_db", api_key="...", host_override="...")
|
||||
|
||||
@udf(pip=["torch>=2.0"], num_gpus=1)
|
||||
def embed(text: str) -> list[float]:
|
||||
return model.encode(text).tolist()
|
||||
|
||||
db.create_function(embed) # CREATE FUNCTION (once)
|
||||
tbl = db.open_table("docs")
|
||||
tbl.add_columns(computed={"vec": embed("text")}) # bind embed(text) -> vec
|
||||
tbl.refresh_column("vec").wait() # materialize (returns a Job)
|
||||
view = db.create_materialized_view("chunks", tbl, ["id", chunk_fn])
|
||||
|
||||
`embed("text")` applies the registered function to the `text` column and yields
|
||||
the expression `embed(text)`; the function itself stays decoupled from any
|
||||
column, so the same `embed` works on any column or table.
|
||||
|
||||
These operations are server-backed (LanceDB Enterprise / Cloud); the
|
||||
decorator itself works locally (define + call), only registration needs a
|
||||
remote connection.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import dataclasses
|
||||
import functools
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
import textwrap
|
||||
import json
|
||||
import time
|
||||
import typing
|
||||
|
||||
# -- type hints -> SQL type strings -------------------------------------
|
||||
|
||||
_SCALARS = {
|
||||
int: "BIGINT",
|
||||
# Pragmatic default for ML workloads: python float maps to FLOAT
|
||||
# (Float32). Use an explicit `returns=` for DOUBLE.
|
||||
float: "FLOAT",
|
||||
str: "VARCHAR",
|
||||
bool: "BOOLEAN",
|
||||
bytes: "BLOB",
|
||||
}
|
||||
|
||||
|
||||
class TypeInferenceError(TypeError):
|
||||
pass
|
||||
|
||||
|
||||
def sql_type(hint) -> str:
|
||||
"""SQL type string for a python type hint."""
|
||||
if hint in _SCALARS:
|
||||
return _SCALARS[hint]
|
||||
origin = typing.get_origin(hint)
|
||||
if origin in (list, typing.List):
|
||||
(item,) = typing.get_args(hint) or (None,)
|
||||
if item in _SCALARS:
|
||||
return f"{_SCALARS[item]}[]"
|
||||
raise TypeInferenceError(
|
||||
f"unsupported list item type {item!r}; use an explicit returns="
|
||||
)
|
||||
fields = _struct_fields(hint)
|
||||
if fields is not None:
|
||||
inner = ", ".join(f"{name} {sql_type(h)}" for name, h in fields)
|
||||
return f"STRUCT({inner})"
|
||||
raise TypeInferenceError(
|
||||
f"cannot infer a SQL type for {hint!r}; pass an explicit type string"
|
||||
)
|
||||
|
||||
|
||||
def _struct_fields(hint):
|
||||
"""(name, hint) pairs for a TypedDict or dataclass, else None."""
|
||||
if dataclasses.is_dataclass(hint):
|
||||
return [(f.name, f.type) for f in dataclasses.fields(hint)]
|
||||
# TypedDict detection: a dict subclass with __annotations__.
|
||||
if (
|
||||
isinstance(hint, type)
|
||||
and issubclass(hint, dict)
|
||||
and typing.get_type_hints(hint)
|
||||
):
|
||||
return list(typing.get_type_hints(hint).items())
|
||||
return None
|
||||
|
||||
|
||||
def return_type(fn, override: "str | None", table: bool) -> str:
|
||||
"""SQL return type for a function: explicit override wins, else the
|
||||
return annotation. Table functions render as TABLE(...) and accept
|
||||
struct-shaped hints (TypedDict/dataclass, optionally list-wrapped)."""
|
||||
if override is not None:
|
||||
s = override.strip()
|
||||
if table and not s.upper().startswith("TABLE"):
|
||||
if s.upper().startswith("STRUCT"):
|
||||
return "TABLE" + s[len("STRUCT") :]
|
||||
raise TypeInferenceError(
|
||||
"a table function's returns= must be TABLE(...) or STRUCT(...)"
|
||||
)
|
||||
return s
|
||||
|
||||
hints = typing.get_type_hints(fn)
|
||||
ret = hints.get("return")
|
||||
if ret is None:
|
||||
raise TypeInferenceError(
|
||||
f"function {fn.__name__!r} needs a return annotation or returns="
|
||||
)
|
||||
if table:
|
||||
# Accept list[Row] / Row where Row is a TypedDict or dataclass.
|
||||
if typing.get_origin(ret) in (list, typing.List):
|
||||
(ret,) = typing.get_args(ret)
|
||||
fields = _struct_fields(ret)
|
||||
if fields is None:
|
||||
raise TypeInferenceError(
|
||||
"a table function must return rows shaped as a TypedDict or "
|
||||
"dataclass (optionally list-wrapped); or pass returns=..."
|
||||
)
|
||||
inner = ", ".join(f"{name} {sql_type(h)}" for name, h in fields)
|
||||
return f"TABLE({inner})"
|
||||
return sql_type(ret)
|
||||
|
||||
|
||||
def param_types(fn) -> "list[tuple[str, str]]":
|
||||
"""(name, sql type) per parameter, from annotations. Each UDF
|
||||
parameter binds to a source column of the same name by default."""
|
||||
hints = typing.get_type_hints(fn)
|
||||
out = []
|
||||
for name, p in inspect.signature(fn).parameters.items():
|
||||
if p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD):
|
||||
raise TypeInferenceError("*args/**kwargs are not supported in UDFs")
|
||||
hint = hints.get(name)
|
||||
if hint is None:
|
||||
raise TypeInferenceError(
|
||||
f"parameter {name!r} of {fn.__name__!r} needs a type annotation"
|
||||
)
|
||||
out.append((name, sql_type(hint)))
|
||||
return out
|
||||
|
||||
|
||||
# -- column expressions -------------------------------------------------
|
||||
|
||||
|
||||
class ColumnExpr(str):
|
||||
"""A computed-column expression produced by applying a registered
|
||||
function to column names, e.g. ``embed("data") -> "embed(data)"``.
|
||||
|
||||
It IS the expression string everywhere a string is expected (views, SQL,
|
||||
logging), and additionally carries the function's declared return type so
|
||||
``add_columns(computed=...)`` can declare the column without a hand-written
|
||||
type. ``field_types`` holds the per-field SQL types of a STRUCT return, for
|
||||
fanning one expression out to several columns.
|
||||
"""
|
||||
|
||||
data_type: "str | None"
|
||||
field_types: "list[str] | None"
|
||||
|
||||
def __new__(cls, expr: str, data_type=None, field_types=None):
|
||||
obj = super().__new__(cls, expr)
|
||||
obj.data_type = data_type
|
||||
obj.field_types = field_types
|
||||
return obj
|
||||
|
||||
|
||||
def _normalize_computed(computed: dict) -> dict:
|
||||
"""Normalize the user-facing ``computed=`` mapping to the canonical
|
||||
``{name: (sql_type, expression)}`` form.
|
||||
|
||||
Accepts, per entry:
|
||||
- value is a `ColumnExpr` (from ``fn("col")``): the column's SQL type
|
||||
comes from the function's return type -- no hand-written type needed. A
|
||||
tuple key (``("chunk", "idx")``) fans a STRUCT return out to one
|
||||
(type, expression) entry per field, in declared order.
|
||||
- value is a legacy ``(sql_type, expression)`` tuple: passed through (the
|
||||
escape hatch, e.g. bare-name function strings).
|
||||
"""
|
||||
out: dict = {}
|
||||
for key, val in computed.items():
|
||||
if isinstance(val, ColumnExpr):
|
||||
expr = str(val)
|
||||
if isinstance(key, (tuple, list)):
|
||||
if not val.field_types:
|
||||
raise ValueError(
|
||||
f"columns {tuple(key)} need a STRUCT-returning function; "
|
||||
f"{expr} returns a single value"
|
||||
)
|
||||
if len(val.field_types) != len(key):
|
||||
raise ValueError(
|
||||
f"{len(key)} columns but {len(val.field_types)} struct fields "
|
||||
f"in {expr}"
|
||||
)
|
||||
for name, t in zip(key, val.field_types):
|
||||
out[name] = (t, expr)
|
||||
else:
|
||||
if val.data_type is None:
|
||||
raise ValueError(f"cannot infer a type for {expr}; pass types=")
|
||||
out[key] = (val.data_type, expr)
|
||||
else:
|
||||
out[key] = val
|
||||
return out
|
||||
|
||||
|
||||
# -- the @udf / @table_udf decorators -----------------------------------
|
||||
|
||||
|
||||
class Udf:
|
||||
def __init__(
|
||||
self,
|
||||
fn,
|
||||
*,
|
||||
returns: "str | None" = None,
|
||||
table: bool = False,
|
||||
name: "str | None" = None,
|
||||
pip: "list[str] | None" = None,
|
||||
pip_index_url: "str | None" = None,
|
||||
pip_extra_index_urls: "list[str] | None" = None,
|
||||
find_links: "list[str] | None" = None,
|
||||
requirements: "str | list[str] | None" = None,
|
||||
conda: "list[str] | None" = None,
|
||||
conda_channels: "list[str] | None" = None,
|
||||
env: "dict[str, str] | list[str] | None" = None,
|
||||
num_cpus: "int | None" = None,
|
||||
num_gpus: "int | None" = None,
|
||||
batch_size: "int | None" = None,
|
||||
timeout: "float | None" = None,
|
||||
error_policy: "str | None" = None,
|
||||
max_skip_ratio: "float | None" = None,
|
||||
retries: "int | None" = None,
|
||||
docker_image: "str | None" = None,
|
||||
description: "str | None" = None,
|
||||
prefer_source: bool = False,
|
||||
):
|
||||
functools.update_wrapper(self, fn)
|
||||
self.fn = fn
|
||||
self.name = name or fn.__name__
|
||||
self.table = table
|
||||
self.params = param_types(fn)
|
||||
self.returns = return_type(fn, returns, table)
|
||||
self.prefer_source = prefer_source
|
||||
self.options: "dict[str, str]" = {}
|
||||
if conda and (pip or requirements):
|
||||
raise ValueError("pass conda or pip/requirements, not both")
|
||||
if conda_channels and not conda:
|
||||
raise ValueError("conda_channels requires conda")
|
||||
if pip:
|
||||
self.options["pip"] = ",".join(pip)
|
||||
if pip_extra_index_urls:
|
||||
self.options["pip_extra_index_urls"] = ",".join(pip_extra_index_urls)
|
||||
if find_links:
|
||||
self.options["find_links"] = ",".join(find_links)
|
||||
if requirements:
|
||||
self.options["requirements"] = _format_requirements(requirements)
|
||||
if conda:
|
||||
self.options["conda"] = ",".join(conda)
|
||||
if conda_channels:
|
||||
self.options["conda_channels"] = ",".join(conda_channels)
|
||||
if env:
|
||||
self.options["env"] = _format_env(env)
|
||||
for key, val in [
|
||||
("pip_index_url", pip_index_url),
|
||||
("num_cpus", num_cpus),
|
||||
("num_gpus", num_gpus),
|
||||
("batch_size", batch_size),
|
||||
("timeout", timeout),
|
||||
("error_policy", error_policy),
|
||||
("max_skip_ratio", max_skip_ratio),
|
||||
("retries", retries),
|
||||
("docker_image", docker_image),
|
||||
]:
|
||||
if val is not None:
|
||||
self.options[key] = str(val)
|
||||
# Keep the source in the description (when available) so the
|
||||
# catalog stays inspectable even for pickled bodies.
|
||||
if description is not None:
|
||||
self.options["description"] = description
|
||||
else:
|
||||
try:
|
||||
self.options["description"] = textwrap.dedent(inspect.getsource(fn))
|
||||
except (OSError, TypeError):
|
||||
pass
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
"""Call with real values to run locally; call with column-name
|
||||
strings to build an expression for backfills and views, e.g.
|
||||
``embed("data")`` -> the expression ``embed(data)`` (a `ColumnExpr`
|
||||
carrying the function's return type for `add_columns(computed=...)`)."""
|
||||
if args and all(isinstance(a, str) for a in args) and not kwargs:
|
||||
return self.expression(*args)
|
||||
return self.fn(*args, **kwargs)
|
||||
|
||||
def expression(self, *columns: str) -> ColumnExpr:
|
||||
"""The expression applying this function to `columns` (default: the
|
||||
function's own parameter names). Returns a `ColumnExpr` -- a string
|
||||
that also carries the declared return type (and struct field types)."""
|
||||
cols = columns or [p for p, _ in self.params]
|
||||
expr = f"{self.name}({', '.join(cols)})"
|
||||
field_types = None
|
||||
if self.returns.upper().startswith("STRUCT"):
|
||||
field_types = struct_field_types(self.returns)
|
||||
return ColumnExpr(expr, data_type=self.returns, field_types=field_types)
|
||||
|
||||
def _body(self) -> "tuple[str, str]":
|
||||
"""(body literal, body_format). Source when requested and
|
||||
retrievable; cloudpickle otherwise (handles closures)."""
|
||||
if self.prefer_source:
|
||||
try:
|
||||
src = textwrap.dedent(inspect.getsource(self.fn))
|
||||
# Strip the decorator line(s) so the stored body is a
|
||||
# plain function definition.
|
||||
lines = src.splitlines(keepends=True)
|
||||
while lines and lines[0].lstrip().startswith("@"):
|
||||
lines.pop(0)
|
||||
return "".join(lines), "source"
|
||||
except (OSError, TypeError):
|
||||
pass
|
||||
import cloudpickle
|
||||
|
||||
raw = cloudpickle.dumps(self.fn)
|
||||
return base64.b64encode(raw).decode("ascii"), "cloudpickle"
|
||||
|
||||
def _body_and_options(self) -> "tuple[str, dict[str, str]]":
|
||||
"""The body literal plus the finalized options (body_format /
|
||||
python_version / cloudpickle-pip bookkeeping for a non-source
|
||||
body)."""
|
||||
body, body_format = self._body()
|
||||
options = dict(self.options)
|
||||
if body_format != "source":
|
||||
options["body_format"] = body_format
|
||||
# Pickled code objects only load under the same interpreter
|
||||
# minor version; record ours so the worker can fail with a
|
||||
# clear message instead of a bytecode error.
|
||||
options["python_version"] = self.pickle_environment()
|
||||
# The worker deserializes the body with cloudpickle; make sure
|
||||
# the job's pip environment provides it. Conda bakes inject
|
||||
# cloudpickle server-side, so do not create an invalid pip+conda
|
||||
# declaration here.
|
||||
if "conda" not in options:
|
||||
pip = [d for d in options.get("pip", "").split(",") if d]
|
||||
if not any(d.startswith("cloudpickle") for d in pip):
|
||||
pip.append("cloudpickle")
|
||||
options["pip"] = ",".join(pip)
|
||||
return body, options
|
||||
|
||||
def create_request(self) -> dict:
|
||||
"""Keyword arguments for `connection.create_function`."""
|
||||
body, options = self._body_and_options()
|
||||
return {
|
||||
"name": self.name,
|
||||
"language": "python",
|
||||
"return_type": self.returns,
|
||||
"body": body,
|
||||
"options": options,
|
||||
}
|
||||
|
||||
def create_statement(self) -> str:
|
||||
"""The equivalent `CREATE FUNCTION` SQL (for SQL-surface callers)."""
|
||||
params = ", ".join(f"{n} {t}" for n, t in self.params)
|
||||
body, options = self._body_and_options()
|
||||
with_clause = ""
|
||||
if options:
|
||||
rendered = ", ".join(
|
||||
f"{k} = '{_escape(v)}'" for k, v in sorted(options.items())
|
||||
)
|
||||
with_clause = f" WITH ({rendered})"
|
||||
return (
|
||||
f"CREATE FUNCTION {self.name}({params}) RETURNS {self.returns} "
|
||||
f"LANGUAGE python AS '{_escape_body(body)}'{with_clause}"
|
||||
)
|
||||
|
||||
def pickle_environment(self) -> str:
|
||||
"""Python version the body pickles under -- workers should match
|
||||
the minor version for cloudpickle compatibility."""
|
||||
return f"{sys.version_info.major}.{sys.version_info.minor}"
|
||||
|
||||
|
||||
def _escape(s: str) -> str:
|
||||
return str(s).replace("'", "''")
|
||||
|
||||
|
||||
def _format_requirements(requirements: "str | list[str]") -> str:
|
||||
if isinstance(requirements, str):
|
||||
return requirements
|
||||
return "\n".join(str(req) for req in requirements)
|
||||
|
||||
|
||||
def _format_env(env: "dict[str, str] | list[str]") -> str:
|
||||
if isinstance(env, dict):
|
||||
return "; ".join(f"{key}={value}" for key, value in env.items())
|
||||
return "; ".join(str(entry) for entry in env)
|
||||
|
||||
|
||||
def _escape_body(body: str) -> str:
|
||||
# The server unescapes \n / \t in single-quoted bodies; encode real
|
||||
# newlines accordingly and escape quotes.
|
||||
return (
|
||||
body.replace("\\", "\\\\")
|
||||
.replace("'", "''")
|
||||
.replace("\n", "\\n")
|
||||
.replace("\t", "\\t")
|
||||
)
|
||||
|
||||
|
||||
def udf(fn=None, **kwargs):
|
||||
"""Decorate a function as a scalar (or struct-returning) UDF.
|
||||
|
||||
@udf
|
||||
def doubled(val: int) -> float: ...
|
||||
|
||||
@udf(pip=["torch>=2"], num_gpus=1)
|
||||
def embed(body: str) -> list[float]: ...
|
||||
"""
|
||||
if fn is not None:
|
||||
return Udf(fn, **kwargs)
|
||||
return lambda f: Udf(f, **kwargs)
|
||||
|
||||
|
||||
def table_udf(fn=None, **kwargs):
|
||||
"""Decorate a table function (UDTF): each input row may emit zero or
|
||||
more output rows. Only usable in materialized views.
|
||||
|
||||
class Chunk(TypedDict):
|
||||
chunk: str
|
||||
chunk_idx: int
|
||||
|
||||
@table_udf
|
||||
def chunker(body: str) -> list[Chunk]: ...
|
||||
"""
|
||||
kwargs["table"] = True
|
||||
if fn is not None:
|
||||
return Udf(fn, **kwargs)
|
||||
return lambda f: Udf(f, **kwargs)
|
||||
|
||||
|
||||
# -- view / job handles (thin references over a connection) -------------
|
||||
|
||||
|
||||
def struct_field_types(returns: str) -> "list[str]":
|
||||
"""Field type strings of a STRUCT(...) SQL type, in declared order."""
|
||||
inner = returns.strip()[len("STRUCT(") : -1]
|
||||
fields, depth, start = [], 0, 0
|
||||
for i, c in enumerate(inner):
|
||||
if c in "([":
|
||||
depth += 1
|
||||
elif c in ")]":
|
||||
depth -= 1
|
||||
elif c == "," and depth == 0:
|
||||
fields.append(inner[start:i].strip())
|
||||
start = i + 1
|
||||
fields.append(inner[start:].strip())
|
||||
# Each field is "name TYPE"; drop the name.
|
||||
return [f.split(None, 1)[1] for f in fields]
|
||||
|
||||
|
||||
def build_view_query(source, select) -> str:
|
||||
"""Assemble a view SELECT from a source (name or table) and select
|
||||
items: a column name, an expression string, a (alias, expression)
|
||||
tuple, or a @udf/@table_udf object."""
|
||||
src = source.name if hasattr(source, "name") else source
|
||||
items = []
|
||||
for item in select:
|
||||
if isinstance(item, Udf):
|
||||
items.append(item.expression())
|
||||
elif isinstance(item, tuple):
|
||||
alias, expr = item
|
||||
expr = expr.expression() if isinstance(expr, Udf) else expr
|
||||
items.append(f"{expr} AS {alias}")
|
||||
else:
|
||||
items.append(item)
|
||||
return f"SELECT {', '.join(items)} FROM {src}"
|
||||
|
||||
|
||||
def _job_id_matches(handle_id: str, listed_id: str) -> bool:
|
||||
# The refresh/backfill endpoints return the submission id (a uuid), but
|
||||
# the agent names the manifest job "<table>-<type>-<first 8 of the
|
||||
# submission id>" -- which is what list_jobs and cancel report. Match the
|
||||
# canonical id directly, or by that submission prefix.
|
||||
if listed_id == handle_id:
|
||||
return True
|
||||
prefix = handle_id[:8]
|
||||
return len(prefix) >= 4 and prefix in listed_id
|
||||
|
||||
|
||||
class MaterializedView:
|
||||
"""A reference to a materialized view (name + connection). Operations are
|
||||
server-backed connection calls bound to the name.
|
||||
|
||||
``create_materialized_view`` returns one of these; ``job_id`` is the
|
||||
initial-population job (None when the view was created with no data), so
|
||||
``db.create_materialized_view(...).wait()`` blocks until it is populated.
|
||||
"""
|
||||
|
||||
def __init__(self, conn, name: str, job_id: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.name = name
|
||||
#: initial-population job id from create, or None (with_no_data).
|
||||
self.job_id = job_id
|
||||
|
||||
def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
"""Block until the initial-population job (from create) finishes.
|
||||
A no-op when the view was created with no data."""
|
||||
if self.job_id is None:
|
||||
return "finished"
|
||||
return Job(self.conn, self.job_id, table=self.name).wait(
|
||||
timeout=timeout, poll=poll
|
||||
)
|
||||
|
||||
def refresh(self, full: bool = False) -> "Job":
|
||||
"""Refresh the materialized view; returns a `Job` to wait on,
|
||||
poll, or cancel (``view.refresh().wait()``).
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh. A full rebuild preserves
|
||||
the view's indexes -- they are reindexed by the distributed indexer.
|
||||
"""
|
||||
job_id = self.conn._refresh_materialized_view(self.name, full=full)
|
||||
return Job(self.conn, job_id, table=self.name)
|
||||
|
||||
def explain_refresh(self, full: bool = False):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH)."""
|
||||
return self.conn.explain_refresh_materialized_view(self.name, full=full)
|
||||
|
||||
def alter(self, auto_refresh: bool) -> None:
|
||||
self.conn.alter_materialized_view(self.name, auto_refresh=auto_refresh)
|
||||
|
||||
def drop(self) -> None:
|
||||
self.conn.drop_materialized_view(self.name)
|
||||
|
||||
# A materialized view is a first-class table: it can be indexed and
|
||||
# searched like any other. These open the materialized dataset by name and
|
||||
# delegate. Indexes declared this way are recorded against the view, so the
|
||||
# engine re-applies them after a full refresh rebuilds the dataset (a full
|
||||
# refresh overwrites the dataset, which would otherwise drop its indices).
|
||||
def _table(self):
|
||||
return self.conn.open_table(self.name)
|
||||
|
||||
def create_index(self, *args, **kwargs):
|
||||
"""Build an index on the materialized view (see Table.create_index)."""
|
||||
return self._table().create_index(*args, **kwargs)
|
||||
|
||||
def create_scalar_index(self, *args, **kwargs):
|
||||
"""Build a scalar index on the materialized view."""
|
||||
return self._table().create_scalar_index(*args, **kwargs)
|
||||
|
||||
def create_fts_index(self, *args, **kwargs):
|
||||
"""Build a full-text-search index on the materialized view."""
|
||||
return self._table().create_fts_index(*args, **kwargs)
|
||||
|
||||
def search(self, *args, **kwargs):
|
||||
"""Search the materialized view (vector / FTS / hybrid)."""
|
||||
return self._table().search(*args, **kwargs)
|
||||
|
||||
def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Lineage of the materialized view (or one of its columns). Delegates
|
||||
to the backing table; the server already includes the view's sources
|
||||
and downstream dependents. Returns a `Lineage`."""
|
||||
return self._table().lineage(column, direction=direction, depth=depth)
|
||||
|
||||
|
||||
_PROGRESS = re.compile(r"(\d+)/(\d+)")
|
||||
|
||||
|
||||
class JobFailedError(RuntimeError):
|
||||
"""Raised by ``Job.wait()`` when the server reports the job ``failed``.
|
||||
|
||||
Carries the server-side error so a doomed backfill (e.g. a multi-column
|
||||
``REFRESH COLUMN`` of a scalar UDF) surfaces its real cause promptly,
|
||||
instead of the caller blocking until ``wait()``'s timeout.
|
||||
"""
|
||||
|
||||
def __init__(self, job_id: str, error: "str | None"):
|
||||
self.job_id = job_id
|
||||
self.error = error
|
||||
super().__init__(f"job {job_id} failed: {error or 'unknown error'}")
|
||||
|
||||
|
||||
class Job:
|
||||
"""A reference to a server-side job, backed by the platform jobs API.
|
||||
|
||||
Holds the submission (manifest) id and resolves the platform job id
|
||||
lazily; ``status``/``progress``/``wait`` read the registry-backed
|
||||
describe endpoint, so terminal states and errors are first-class.
|
||||
"""
|
||||
|
||||
#: How long an unresolved job is treated as still materializing
|
||||
#: (submission -> dispatch -> registry record is async).
|
||||
GRACE_SECONDS = 20.0
|
||||
|
||||
#: Platform lifecycle state -> the user-facing vocabulary.
|
||||
_STATES = {
|
||||
"IN_PROGRESS": "running",
|
||||
"DONE": "finished",
|
||||
"FAILED": "failed",
|
||||
"CANCELLED": "cancelled",
|
||||
}
|
||||
|
||||
def __init__(self, conn, job_id: str, table: "str | None" = None):
|
||||
self.conn = conn
|
||||
#: The submission (manifest) id the launching call handed out.
|
||||
self.id = job_id
|
||||
#: The job's table, when known -- narrows platform-id resolution.
|
||||
self.table = table
|
||||
self._platform_id: "str | None" = None
|
||||
self._created = time.monotonic()
|
||||
self._finished = False
|
||||
|
||||
@classmethod
|
||||
def _completed(cls, conn=None, table: "str | None" = None) -> "Job":
|
||||
"""A job for work that completed synchronously within the call that
|
||||
returned it (native tables, scalar/FTS builds). ``status``/``wait``
|
||||
report ``finished`` immediately and ``cancel`` is a no-op."""
|
||||
job = cls(conn, "", table)
|
||||
job._finished = True
|
||||
return job
|
||||
|
||||
def _resolve(self) -> "str | None":
|
||||
if self._platform_id is None:
|
||||
self._platform_id = self.conn.resolve_platform_job_id(self.id, self.table)
|
||||
return self._platform_id
|
||||
|
||||
def _describe(self):
|
||||
platform_id = self._resolve()
|
||||
if platform_id is None:
|
||||
return None
|
||||
return self.conn.describe_platform_job(platform_id)
|
||||
|
||||
@staticmethod
|
||||
def _payload(described) -> dict:
|
||||
# Older records carry the status-store URI string instead of a
|
||||
# payload; anything non-dict means "no structured status".
|
||||
try:
|
||||
payload = json.loads(described.status_json)
|
||||
except (TypeError, ValueError):
|
||||
return {}
|
||||
return payload if isinstance(payload, dict) else {}
|
||||
|
||||
def status(self) -> str:
|
||||
"""pending / running / finished / failed / cancelled (or unknown
|
||||
when the job never appeared in the registry)."""
|
||||
if self._finished:
|
||||
return "finished"
|
||||
described = self._describe()
|
||||
if described is not None:
|
||||
return self._STATES.get(described.job_state, described.job_state)
|
||||
if time.monotonic() - self._created < self.GRACE_SECONDS:
|
||||
return "pending"
|
||||
return "unknown"
|
||||
|
||||
def progress(self) -> "tuple[int, int] | None":
|
||||
"""(units_done, units_total) once workers have published progress."""
|
||||
if self._finished:
|
||||
return None
|
||||
described = self._describe()
|
||||
if described is None:
|
||||
return None
|
||||
payload = self._payload(described)
|
||||
if payload.get("units_total") is not None:
|
||||
return payload.get("units_done") or 0, payload["units_total"]
|
||||
return None
|
||||
|
||||
def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
if self._finished:
|
||||
return "finished"
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
described = self._describe()
|
||||
if described is None:
|
||||
if time.monotonic() - self._created > self.GRACE_SECONDS:
|
||||
raise JobFailedError(
|
||||
self.id,
|
||||
"job did not appear in the job registry within the "
|
||||
"grace period",
|
||||
)
|
||||
time.sleep(min(poll, 0.5))
|
||||
continue
|
||||
state = self._STATES.get(described.job_state, described.job_state)
|
||||
if state == "finished":
|
||||
return state
|
||||
if state == "cancelled":
|
||||
return state
|
||||
if state == "failed":
|
||||
raise JobFailedError(self.id, self._payload(described).get("error"))
|
||||
time.sleep(poll)
|
||||
raise TimeoutError(f"job {self.id} still {self.status()} after {timeout}s")
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Request cancellation. Workers drain cooperatively; poll ``status``
|
||||
for the terminal ``cancelled``."""
|
||||
if self._finished:
|
||||
return
|
||||
deadline = time.monotonic() + 5.0
|
||||
while (platform_id := self._resolve()) is None:
|
||||
if time.monotonic() > deadline:
|
||||
raise RuntimeError(
|
||||
f"job {self.id} has not registered yet; retry cancel shortly"
|
||||
)
|
||||
time.sleep(0.5)
|
||||
self.conn.cancel_platform_job(platform_id)
|
||||
|
||||
|
||||
class AsyncMaterializedView:
|
||||
"""Async reference to a materialized view (name + async connection)."""
|
||||
|
||||
def __init__(self, conn, name: str, job_id: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.name = name
|
||||
#: initial-population job id from create, or None (with_no_data).
|
||||
self.job_id = job_id
|
||||
|
||||
async def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
"""Block until the initial-population job (from create) finishes.
|
||||
A no-op when the view was created with no data."""
|
||||
if self.job_id is None:
|
||||
return "finished"
|
||||
return await AsyncJob(self.conn, self.job_id, table=self.name).wait(
|
||||
timeout=timeout, poll=poll
|
||||
)
|
||||
|
||||
async def refresh(self, full: bool = False) -> "AsyncJob":
|
||||
"""Refresh the materialized view; returns an `AsyncJob` to wait
|
||||
on, poll, or cancel.
|
||||
|
||||
``full=True`` forces a full rebuild instead of an incremental refresh
|
||||
(indexes are preserved and reindexed by the distributed indexer).
|
||||
"""
|
||||
job_id = await self.conn._refresh_materialized_view(self.name, full=full)
|
||||
return AsyncJob(self.conn, job_id, table=self.name)
|
||||
|
||||
async def explain_refresh(self, full: bool = False):
|
||||
return await self.conn.explain_refresh_materialized_view(self.name, full=full)
|
||||
|
||||
async def alter(self, auto_refresh: bool) -> None:
|
||||
await self.conn.alter_materialized_view(self.name, auto_refresh=auto_refresh)
|
||||
|
||||
async def drop(self) -> None:
|
||||
await self.conn.drop_materialized_view(self.name)
|
||||
|
||||
async def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Lineage of the materialized view (or column). Returns a `Lineage`."""
|
||||
return await self.conn.lineage(
|
||||
self.name, column, direction=direction, depth=depth
|
||||
)
|
||||
|
||||
|
||||
class AsyncJob:
|
||||
"""Async reference to a server-side job, backed by the platform jobs API.
|
||||
|
||||
Same contract as `Job` with awaitable methods.
|
||||
"""
|
||||
|
||||
GRACE_SECONDS = 20.0
|
||||
_STATES = Job._STATES
|
||||
|
||||
def __init__(self, conn, job_id: str, table: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.id = job_id
|
||||
self.table = table
|
||||
self._platform_id: "str | None" = None
|
||||
self._created = time.monotonic()
|
||||
self._finished = False
|
||||
|
||||
@classmethod
|
||||
def _completed(cls, conn=None, table: "str | None" = None) -> "AsyncJob":
|
||||
"""See ``Job._completed``."""
|
||||
job = cls(conn, "", table)
|
||||
job._finished = True
|
||||
return job
|
||||
|
||||
async def _resolve(self) -> "str | None":
|
||||
if self._platform_id is None:
|
||||
self._platform_id = await self.conn.resolve_platform_job_id(
|
||||
self.id, self.table
|
||||
)
|
||||
return self._platform_id
|
||||
|
||||
async def _describe(self):
|
||||
platform_id = await self._resolve()
|
||||
if platform_id is None:
|
||||
return None
|
||||
return await self.conn.describe_platform_job(platform_id)
|
||||
|
||||
async def status(self) -> str:
|
||||
if self._finished:
|
||||
return "finished"
|
||||
described = await self._describe()
|
||||
if described is not None:
|
||||
return self._STATES.get(described.job_state, described.job_state)
|
||||
if time.monotonic() - self._created < self.GRACE_SECONDS:
|
||||
return "pending"
|
||||
return "unknown"
|
||||
|
||||
async def progress(self) -> "tuple[int, int] | None":
|
||||
if self._finished:
|
||||
return None
|
||||
described = await self._describe()
|
||||
if described is None:
|
||||
return None
|
||||
payload = Job._payload(described)
|
||||
if payload.get("units_total") is not None:
|
||||
return payload.get("units_done") or 0, payload["units_total"]
|
||||
return None
|
||||
|
||||
async def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
if self._finished:
|
||||
return "finished"
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
described = await self._describe()
|
||||
if described is None:
|
||||
if time.monotonic() - self._created > self.GRACE_SECONDS:
|
||||
raise JobFailedError(
|
||||
self.id,
|
||||
"job did not appear in the job registry within the "
|
||||
"grace period",
|
||||
)
|
||||
await asyncio.sleep(min(poll, 0.5))
|
||||
continue
|
||||
state = self._STATES.get(described.job_state, described.job_state)
|
||||
if state in ("finished", "cancelled"):
|
||||
return state
|
||||
if state == "failed":
|
||||
raise JobFailedError(self.id, Job._payload(described).get("error"))
|
||||
await asyncio.sleep(poll)
|
||||
raise TimeoutError(f"job {self.id} still {await self.status()} after {timeout}s")
|
||||
|
||||
async def cancel(self) -> None:
|
||||
if self._finished:
|
||||
return
|
||||
deadline = time.monotonic() + 5.0
|
||||
while (platform_id := await self._resolve()) is None:
|
||||
if time.monotonic() > deadline:
|
||||
raise RuntimeError(
|
||||
f"job {self.id} has not registered yet; retry cancel shortly"
|
||||
)
|
||||
await asyncio.sleep(0.5)
|
||||
await self.conn.cancel_platform_job(platform_id)
|
||||
|
||||
@@ -307,13 +307,19 @@ def infer_vector_column_name(
|
||||
# FTS queries do not require a vector column
|
||||
return None
|
||||
|
||||
if query is not None or query_type == "hybrid":
|
||||
try:
|
||||
vector_column_name = inf_vector_column_query(
|
||||
schema, dim=_query_vector_dim(query)
|
||||
)
|
||||
except Exception as e:
|
||||
raise e
|
||||
if query is None and query_type != "hybrid":
|
||||
# No vector search was requested (e.g. a plain scan), so there's
|
||||
# nothing to infer.
|
||||
return None
|
||||
|
||||
vector_column_name = inf_vector_column_query(schema, dim=_query_vector_dim(query))
|
||||
|
||||
if vector_column_name is None:
|
||||
raise ValueError(
|
||||
"No vector column found in the schema. Please specify the "
|
||||
"vector column name explicitly via the `vector_column_name` "
|
||||
"parameter."
|
||||
)
|
||||
|
||||
return vector_column_name
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
|
||||
@@ -955,6 +956,47 @@ def test_local_namespace_operations(tmp_path):
|
||||
assert db.list_namespaces().namespaces == []
|
||||
|
||||
|
||||
def test_local_sync_namespace_uses_rust_without_python_client(tmp_path, monkeypatch):
|
||||
"""Sync local namespace operations should avoid the Python namespace client."""
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
def fail_namespace_client():
|
||||
raise AssertionError("Python namespace client should not be constructed")
|
||||
|
||||
monkeypatch.setattr(db, "namespace_client", fail_namespace_client)
|
||||
|
||||
db.create_namespace(["child"])
|
||||
assert "child" in db.list_namespaces().namespaces
|
||||
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = db.create_table("tbl", schema=schema, namespace_path=["child"])
|
||||
assert table.namespace == ["child"]
|
||||
assert "tbl" in db.table_names(namespace_path=["child"])
|
||||
assert db.list_tables(namespace_path=["child"]).tables == ["tbl"]
|
||||
|
||||
opened = db.open_table("tbl", namespace_path=["child"])
|
||||
assert opened.namespace == ["child"]
|
||||
|
||||
db.drop_table("tbl", namespace_path=["child"])
|
||||
assert db.list_tables(namespace_path=["child"]).tables == []
|
||||
db.drop_namespace(["child"])
|
||||
assert db.list_namespaces().namespaces == []
|
||||
|
||||
|
||||
def test_local_sync_namespace_preserves_public_errors(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
db.create_namespace(["child"])
|
||||
db.create_table(
|
||||
"tbl", schema=pa.schema([pa.field("id", pa.int64())]), namespace_path=["child"]
|
||||
)
|
||||
|
||||
with pytest.raises(TableNotFoundError, match="child\\$missing"):
|
||||
db.open_table("missing", namespace_path=["child"])
|
||||
|
||||
with pytest.raises(NamespaceNotEmptyError):
|
||||
db.drop_namespace(["child"])
|
||||
|
||||
|
||||
def test_create_namespace_invalid_mode_raises(tmp_path):
|
||||
"""Unrecognized create namespace modes raise a clear error."""
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
@@ -1,209 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""Job / AsyncJob against the platform jobs API.
|
||||
|
||||
The reference resolves its submission (manifest) id to a platform job id,
|
||||
then polls describe for registry-backed state: terminal states are
|
||||
first-class (DONE / FAILED / CANCELLED), progress comes from the
|
||||
owner-written status payload, and a failed job raises ``JobFailedError``
|
||||
promptly with the server error.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from lancedb.udf import Job, AsyncJob, JobFailedError
|
||||
|
||||
|
||||
class FakeDescription:
|
||||
"""Mirror of the pyo3 PlatformJobDescription fields the Job reads."""
|
||||
|
||||
def __init__(self, job_state, status=None):
|
||||
self.job_id = "plat-1"
|
||||
self.job_type = "indexer"
|
||||
self.job_subtype = "udf"
|
||||
self.job_state = job_state
|
||||
self.creation_ms = 0
|
||||
self.status_json = json.dumps(status if status is not None else {})
|
||||
|
||||
|
||||
class FakeConn:
|
||||
"""Scripted timeline: resolve returns None until `resolve_after` calls,
|
||||
then the platform id; describe walks a list of descriptions (holding the
|
||||
last once exhausted)."""
|
||||
|
||||
def __init__(self, descriptions, resolve_after=0):
|
||||
self._descs = list(descriptions)
|
||||
self._resolve_after = resolve_after
|
||||
self.resolve_calls = 0
|
||||
self.describe_calls = 0
|
||||
self.cancelled = []
|
||||
|
||||
def resolve_platform_job_id(self, manifest_job_id, table=None):
|
||||
self.resolve_calls += 1
|
||||
if self.resolve_calls <= self._resolve_after:
|
||||
return None
|
||||
return "plat-1"
|
||||
|
||||
def describe_platform_job(self, platform_job_id):
|
||||
assert platform_job_id == "plat-1"
|
||||
snap = self._descs[min(self.describe_calls, len(self._descs) - 1)]
|
||||
self.describe_calls += 1
|
||||
return snap
|
||||
|
||||
def cancel_platform_job(self, platform_job_id):
|
||||
self.cancelled.append(platform_job_id)
|
||||
|
||||
|
||||
class AsyncFakeConn(FakeConn):
|
||||
async def resolve_platform_job_id(self, manifest_job_id, table=None):
|
||||
return FakeConn.resolve_platform_job_id(self, manifest_job_id, table)
|
||||
|
||||
async def describe_platform_job(self, platform_job_id):
|
||||
return FakeConn.describe_platform_job(self, platform_job_id)
|
||||
|
||||
async def cancel_platform_job(self, platform_job_id):
|
||||
return FakeConn.cancel_platform_job(self, platform_job_id)
|
||||
|
||||
|
||||
def test_status_maps_platform_states():
|
||||
for wire, want in [
|
||||
("IN_PROGRESS", "running"),
|
||||
("DONE", "finished"),
|
||||
("FAILED", "failed"),
|
||||
("CANCELLED", "cancelled"),
|
||||
]:
|
||||
job = Job(FakeConn([FakeDescription(wire)]), "job-1", table="t")
|
||||
assert job.status() == want
|
||||
|
||||
|
||||
def test_status_pending_before_resolution():
|
||||
job = Job(FakeConn([], resolve_after=10_000), "job-1", table="t")
|
||||
assert job.status() == "pending"
|
||||
|
||||
|
||||
def test_progress_from_status_payload():
|
||||
conn = FakeConn(
|
||||
[
|
||||
FakeDescription(
|
||||
"IN_PROGRESS",
|
||||
status={"units_done": 3, "units_total": 8, "rows_committed": 100},
|
||||
)
|
||||
]
|
||||
)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
assert job.progress() == (3, 8)
|
||||
|
||||
|
||||
def test_progress_none_for_uri_only_status():
|
||||
# Older records carry the status-store URI string, not a payload.
|
||||
desc = FakeDescription("IN_PROGRESS")
|
||||
desc.status_json = json.dumps("s3://bucket/job/job_status")
|
||||
job = Job(FakeConn([desc]), "job-1", table="t")
|
||||
assert job.progress() is None
|
||||
|
||||
|
||||
def test_wait_raises_on_failed_promptly():
|
||||
conn = FakeConn(
|
||||
[
|
||||
FakeDescription("IN_PROGRESS"),
|
||||
FakeDescription(
|
||||
"FAILED", status={"error": "multi-column backfill needs a STRUCT"}
|
||||
),
|
||||
]
|
||||
)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
t0 = time.monotonic()
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
job.wait(timeout=30, poll=0.01)
|
||||
assert time.monotonic() - t0 < 5 # prompt, nowhere near the 30s timeout
|
||||
assert "STRUCT" in str(exc.value)
|
||||
assert exc.value.error == "multi-column backfill needs a STRUCT"
|
||||
assert exc.value.job_id == "job-1"
|
||||
|
||||
|
||||
def test_wait_returns_finished_on_done():
|
||||
conn = FakeConn([FakeDescription("IN_PROGRESS"), FakeDescription("DONE")])
|
||||
job = Job(conn, "job-1", table="t")
|
||||
assert job.wait(timeout=30, poll=0.01) == "finished"
|
||||
|
||||
|
||||
def test_wait_returns_cancelled():
|
||||
conn = FakeConn([FakeDescription("CANCELLED")])
|
||||
job = Job(conn, "job-1", table="t")
|
||||
assert job.wait(timeout=30, poll=0.01) == "cancelled"
|
||||
|
||||
|
||||
def test_wait_raises_when_job_never_registers():
|
||||
# An unresolved job past the grace window is a lost submission, not an
|
||||
# eternal "pending" hang.
|
||||
conn = FakeConn([], resolve_after=10_000)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
job.GRACE_SECONDS = 0.05
|
||||
job._created = time.monotonic() - 1.0
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
job.wait(timeout=5, poll=0.01)
|
||||
assert "registry" in str(exc.value)
|
||||
|
||||
|
||||
def test_cancel_resolves_then_cancels():
|
||||
conn = FakeConn([FakeDescription("IN_PROGRESS")], resolve_after=1)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
job.cancel()
|
||||
assert conn.cancelled == ["plat-1"]
|
||||
|
||||
|
||||
def test_async_wait_raises_on_failed_promptly():
|
||||
conn = AsyncFakeConn(
|
||||
[FakeDescription("FAILED", status={"error": "boom"})],
|
||||
)
|
||||
job = AsyncJob(conn, "job-1", table="t")
|
||||
|
||||
async def run():
|
||||
t0 = time.monotonic()
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
await job.wait(timeout=30, poll=0.01)
|
||||
assert time.monotonic() - t0 < 5
|
||||
assert exc.value.error == "boom"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_async_wait_returns_finished():
|
||||
conn = AsyncFakeConn([FakeDescription("IN_PROGRESS"), FakeDescription("DONE")])
|
||||
job = AsyncJob(conn, "job-1", table="t")
|
||||
|
||||
async def run():
|
||||
assert await job.wait(timeout=30, poll=0.01) == "finished"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_completed_job_is_finished_without_conn():
|
||||
job = Job._completed(table="t")
|
||||
assert job.status() == "finished"
|
||||
assert job.wait(timeout=0.01) == "finished"
|
||||
assert job.progress() is None
|
||||
job.cancel() # no-op, must not touch a connection
|
||||
|
||||
|
||||
def test_completed_job_ignores_registry():
|
||||
conn = FakeConn([FakeDescription("IN_PROGRESS")])
|
||||
job = Job._completed(conn, table="t")
|
||||
assert job.wait(timeout=0.01) == "finished"
|
||||
assert conn.resolve_calls == 0
|
||||
assert conn.describe_calls == 0
|
||||
|
||||
|
||||
def test_completed_async_job_is_finished():
|
||||
async def run():
|
||||
job = AsyncJob._completed(table="t")
|
||||
assert await job.status() == "finished"
|
||||
assert await job.wait(timeout=0.01) == "finished"
|
||||
assert await job.progress() is None
|
||||
await job.cancel()
|
||||
|
||||
asyncio.run(run())
|
||||
@@ -128,16 +128,9 @@ def test_split_hash(mem_db):
|
||||
|
||||
def test_split_hash_with_discard(mem_db):
|
||||
"""Test hash-based splitting with discard weight."""
|
||||
total_rows = 1000
|
||||
tbl = mem_db.create_table(
|
||||
"test_table",
|
||||
pa.table(
|
||||
{
|
||||
"id": range(total_rows),
|
||||
"category": [f"category-{i}" for i in range(total_rows)],
|
||||
"value": range(total_rows),
|
||||
}
|
||||
),
|
||||
pa.table({"id": range(100), "category": ["A", "B"] * 50, "value": range(100)}),
|
||||
)
|
||||
|
||||
permutation_tbl = (
|
||||
@@ -149,12 +142,10 @@ def test_split_hash_with_discard(mem_db):
|
||||
.execute()
|
||||
)
|
||||
|
||||
# Should have fewer rows due to discard, but should not be empty.
|
||||
# Should have fewer than 100 rows due to discard
|
||||
row_count = permutation_tbl.count_rows()
|
||||
assert 0 < row_count < total_rows
|
||||
|
||||
data = permutation_tbl.search(None).to_arrow().to_pydict()
|
||||
assert set(data["split_id"]) == {0, 1}
|
||||
assert row_count < 100
|
||||
assert row_count > 0 # But not empty
|
||||
|
||||
|
||||
def test_split_sequential(mem_db):
|
||||
@@ -1145,3 +1136,61 @@ def test_take_offsets_empty_permutation(some_permutation: Permutation):
|
||||
result = some_permutation.take_offsets([])
|
||||
|
||||
assert result == []
|
||||
|
||||
|
||||
def test_select_rowid(some_permutation: Permutation):
|
||||
"""Test that _rowid can be selected alongside regular columns."""
|
||||
perm_with_rowid = some_permutation.select_columns(["_rowid", "id"])
|
||||
assert "_rowid" in perm_with_rowid.column_names
|
||||
batches = list(perm_with_rowid.iter(100, skip_last_batch=False))
|
||||
for batch in batches:
|
||||
assert "_rowid" in batch[0]
|
||||
|
||||
|
||||
def test_select_rowid_only(some_permutation: Permutation):
|
||||
"""Test that _rowid can be selected as the sole column."""
|
||||
perm_rowid_only = some_permutation.select_columns(["_rowid"])
|
||||
assert perm_rowid_only.column_names == ["_rowid"]
|
||||
batches = list(perm_rowid_only.iter(100, skip_last_batch=False))
|
||||
assert len(batches) > 0
|
||||
for batch in batches:
|
||||
assert list(batch[0].keys()) == ["_rowid"]
|
||||
|
||||
|
||||
def test_select_rowid_not_in_default(some_permutation: Permutation):
|
||||
"""Test that _rowid is NOT in the default column_names or schema."""
|
||||
assert "_rowid" not in some_permutation.column_names
|
||||
assert "_rowid" not in some_permutation.schema.names
|
||||
|
||||
|
||||
def test_select_rowid_identity_permutation(mem_db):
|
||||
"""Test that _rowid works with an identity permutation."""
|
||||
tbl = mem_db.create_table(
|
||||
"test_rowid_identity", pa.table({"id": range(10), "value": range(10)})
|
||||
)
|
||||
perm = Permutation.identity(tbl)
|
||||
perm_with_rowid = perm.select_columns(["_rowid", "id"])
|
||||
batches = list(perm_with_rowid.iter(10, skip_last_batch=False))
|
||||
assert len(batches) == 1
|
||||
assert "_rowid" in batches[0][0]
|
||||
|
||||
|
||||
def test_rename_rowid(some_permutation: Permutation):
|
||||
"""Test that _rowid can be selected and then renamed."""
|
||||
perm_with_rowid = some_permutation.select_columns(["_rowid", "id"])
|
||||
renamed = perm_with_rowid.rename_column("_rowid", "my_row_id")
|
||||
assert "my_row_id" in renamed.column_names
|
||||
assert "_rowid" not in renamed.column_names
|
||||
batches = list(renamed.iter(100, skip_last_batch=False))
|
||||
for batch in batches:
|
||||
assert "my_row_id" in batch[0]
|
||||
assert "_rowid" not in batch[0]
|
||||
|
||||
|
||||
def test_remove_rowid_after_select(some_permutation: Permutation):
|
||||
"""Test that _rowid can be selected and then removed."""
|
||||
perm_with_rowid = some_permutation.select_columns(["_rowid", "id"])
|
||||
assert "_rowid" in perm_with_rowid.column_names
|
||||
perm_without_rowid = perm_with_rowid.remove_columns(["_rowid"])
|
||||
assert "_rowid" not in perm_without_rowid.column_names
|
||||
assert perm_without_rowid.column_names == ["id"]
|
||||
|
||||
@@ -236,6 +236,65 @@ def test_remote_table_branches_sync():
|
||||
table.branches.delete("exp")
|
||||
|
||||
|
||||
def test_remote_table_branch_merge_defaults_to_execute():
|
||||
merge_bodies = []
|
||||
diff = {
|
||||
"fromBranch": "exp",
|
||||
"parentVersion": 1,
|
||||
"mainVersion": 2,
|
||||
"branchVersion": 3,
|
||||
"baseMoved": False,
|
||||
"rowCountMain": 3,
|
||||
"rowCountBranch": 3,
|
||||
"rowSummary": {
|
||||
"unchanged": 3,
|
||||
"newOnBase": 0,
|
||||
"newOnBranch": 0,
|
||||
"staleRecompute": 0,
|
||||
"inputsChanged": 0,
|
||||
"deltaAvailable": False,
|
||||
},
|
||||
"addedColumns": [],
|
||||
"removedColumns": [],
|
||||
"changedColumns": [],
|
||||
"addedIndexes": [],
|
||||
"removedIndexes": [],
|
||||
"mergeable": True,
|
||||
"mergeBlockers": [],
|
||||
}
|
||||
|
||||
def handler(request):
|
||||
if request.path.endswith("/describe/"):
|
||||
status = 200
|
||||
body = {"version": 2, "schema": {"fields": []}}
|
||||
else:
|
||||
content_len = int(request.headers.get("Content-Length"))
|
||||
request_body = json.loads(request.rfile.read(content_len))
|
||||
merge_bodies.append(request_body)
|
||||
dry_run = request_body["dry_run"]
|
||||
status = 200 if dry_run else 409
|
||||
body = {
|
||||
"status": "ready" if dry_run else "rejected",
|
||||
"diff": diff,
|
||||
"preview": {"promotedColumns": []},
|
||||
}
|
||||
|
||||
request.send_response(status)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(json.dumps(body).encode())
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
branches = db.open_table("test").branches
|
||||
assert branches.merge("exp")["status"] == "rejected"
|
||||
assert branches.merge("exp", dry_run=True)["status"] == "ready"
|
||||
|
||||
assert merge_bodies == [
|
||||
{"from_branch": "exp", "dry_run": False},
|
||||
{"from_branch": "exp", "dry_run": True},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_remote_open_table_branch_and_version():
|
||||
async with mock_lancedb_connection_async(_branch_open_handler) as db:
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.dataset as ds
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
from lancedb.scannable import (
|
||||
_PA_DEFAULT_BATCH_ROWS,
|
||||
_SAMPLE_ROWS,
|
||||
_VARIABLE_WIDTH_ESTIMATE,
|
||||
_WIDE_BATCH_READAHEAD,
|
||||
_WIDE_FRAGMENT_READAHEAD,
|
||||
_bounded_scanner_kwargs,
|
||||
_estimate_bytes_per_row,
|
||||
_sample_head,
|
||||
to_scannable,
|
||||
)
|
||||
|
||||
|
||||
def test_estimate_bytes_per_row():
|
||||
# fixed-width scalars
|
||||
assert _estimate_bytes_per_row(pa.schema([("a", pa.int64())])) == 8
|
||||
assert (
|
||||
_estimate_bytes_per_row(pa.schema([("a", pa.int32()), ("b", pa.float64())]))
|
||||
== 12
|
||||
)
|
||||
assert _estimate_bytes_per_row(pa.schema([("a", pa.bool_())])) == 1
|
||||
# fixed-size list (embedding) dominates
|
||||
assert (
|
||||
_estimate_bytes_per_row(pa.schema([("v", pa.list_(pa.float32(), 768))]))
|
||||
== 768 * 4
|
||||
)
|
||||
# struct sums its children
|
||||
struct = pa.struct([("x", pa.int32()), ("y", pa.int32())])
|
||||
assert _estimate_bytes_per_row(pa.schema([("s", struct)])) == 8
|
||||
# variable-width columns get a flat estimate, not zero
|
||||
assert _estimate_bytes_per_row(pa.schema([("s", pa.string())])) > 0
|
||||
|
||||
|
||||
def test_estimate_bytes_per_row_uses_sample_for_variable_length_lists():
|
||||
# A vector column without a fixed size (e.g. `list<float32>` instead of
|
||||
# `list<float32, 768>`) has no width the schema alone can tell us.
|
||||
schema = pa.schema([("v", pa.list_(pa.float32()))])
|
||||
assert _estimate_bytes_per_row(schema) == _VARIABLE_WIDTH_ESTIMATE
|
||||
|
||||
sample = pa.table({"v": pa.array([[0.0] * 768], type=pa.list_(pa.float32()))})
|
||||
assert _estimate_bytes_per_row(schema, sample) == 768 * 4
|
||||
|
||||
|
||||
def test_estimate_bytes_per_row_sample_ignores_missing_or_null_lists():
|
||||
schema = pa.schema([("v", pa.list_(pa.float32()))])
|
||||
# an all-null sample column can't tell us anything either
|
||||
sample = pa.table({"v": pa.array([None], type=pa.list_(pa.float32()))})
|
||||
assert _estimate_bytes_per_row(schema, sample) == _VARIABLE_WIDTH_ESTIMATE
|
||||
|
||||
|
||||
def test_bounded_scanner_kwargs_narrow_uses_defaults():
|
||||
# Narrow rows stay on pyarrow defaults (empty kwargs) so throughput is
|
||||
# unchanged.
|
||||
for schema in [
|
||||
pa.schema([("a", pa.int64()), ("b", pa.int32()), ("c", pa.string())]),
|
||||
pa.schema([("a", pa.int64()), ("t", pa.string()), ("u", pa.string())]),
|
||||
# a 100-dim float32 vector is still under the per-row budget
|
||||
pa.schema([("id", pa.int64()), ("v", pa.list_(pa.float32(), 100))]),
|
||||
]:
|
||||
assert _bounded_scanner_kwargs(schema) == {}, schema
|
||||
|
||||
|
||||
def test_bounded_scanner_kwargs_wide_is_bounded():
|
||||
schema = pa.schema(
|
||||
[
|
||||
("uid", pa.string()),
|
||||
("img", pa.list_(pa.float32(), 768)),
|
||||
("txt", pa.list_(pa.float32(), 768)),
|
||||
]
|
||||
)
|
||||
kwargs = _bounded_scanner_kwargs(schema)
|
||||
assert kwargs, "wide schema should be throttled"
|
||||
assert kwargs["batch_readahead"] == _WIDE_BATCH_READAHEAD
|
||||
assert kwargs["fragment_readahead"] == _WIDE_FRAGMENT_READAHEAD
|
||||
# batch is capped well below the pyarrow default for wide rows
|
||||
assert kwargs["batch_size"] < _PA_DEFAULT_BATCH_ROWS
|
||||
|
||||
|
||||
def test_bounded_scanner_kwargs_variable_length_list_needs_sample():
|
||||
# Without a sample, a variable-length (not fixed-size) vector column looks
|
||||
# narrow because its true width is unknown from the schema alone.
|
||||
schema = pa.schema([("uid", pa.string()), ("vec", pa.list_(pa.float32()))])
|
||||
assert _bounded_scanner_kwargs(schema) == {}
|
||||
|
||||
sample = pa.table(
|
||||
{
|
||||
"uid": pa.array(["a"]),
|
||||
"vec": pa.array([[0.0] * 768], type=pa.list_(pa.float32())),
|
||||
}
|
||||
)
|
||||
kwargs = _bounded_scanner_kwargs(schema, sample)
|
||||
assert kwargs, "sample should reveal the wide vector column"
|
||||
assert kwargs["batch_size"] < _PA_DEFAULT_BATCH_ROWS
|
||||
|
||||
|
||||
def _write_wide_dataset(
|
||||
path, *, files=2, rows_per_file=20_000, dim=768, fixed_size=True
|
||||
):
|
||||
rng = np.random.default_rng(0)
|
||||
for i in range(files):
|
||||
emb = rng.standard_normal((rows_per_file, dim), dtype=np.float32)
|
||||
vec_type = pa.list_(pa.float32(), dim) if fixed_size else pa.list_(pa.float32())
|
||||
vec_array = (
|
||||
pa.FixedSizeListArray.from_arrays(pa.array(emb.reshape(-1)), dim)
|
||||
if fixed_size
|
||||
else pa.array(emb.tolist(), type=vec_type)
|
||||
)
|
||||
table = pa.table(
|
||||
{
|
||||
"uid": pa.array([f"{i}_{j}" for j in range(rows_per_file)]),
|
||||
"vec": vec_array,
|
||||
}
|
||||
)
|
||||
pq.write_table(table, f"{path}/part-{i}.parquet")
|
||||
|
||||
|
||||
def test_dataset_reader_respects_bounded_batch_size(tmp_path):
|
||||
# The Dataset path should stream small batches for wide rows, not pyarrow's
|
||||
# 131072-row default, and still return every row.
|
||||
_write_wide_dataset(str(tmp_path))
|
||||
dataset = ds.dataset(str(tmp_path), format="parquet")
|
||||
|
||||
expected = _bounded_scanner_kwargs(dataset.schema)["batch_size"]
|
||||
|
||||
scannable = to_scannable(dataset)
|
||||
assert scannable.rescannable
|
||||
assert scannable.num_rows == 40_000
|
||||
|
||||
total = 0
|
||||
for batch in scannable.reader():
|
||||
assert batch.num_rows <= expected
|
||||
total += batch.num_rows
|
||||
assert total == 40_000
|
||||
|
||||
# factory can be called again (rescannable)
|
||||
assert sum(b.num_rows for b in scannable.reader()) == 40_000
|
||||
|
||||
|
||||
def test_dataset_reader_samples_variable_length_list_width(tmp_path):
|
||||
# A vector column stored without a fixed size (e.g. produced by tools that
|
||||
# don't tag list columns with their length) is invisible to the
|
||||
# schema-only estimate, so `to_scannable` must peek a sample of rows to
|
||||
# detect that it's wide and bound the scanner accordingly.
|
||||
_write_wide_dataset(str(tmp_path), fixed_size=False)
|
||||
dataset = ds.dataset(str(tmp_path), format="parquet")
|
||||
|
||||
schema_only_kwargs = _bounded_scanner_kwargs(dataset.schema)
|
||||
assert schema_only_kwargs == {}, "schema alone can't see the list width"
|
||||
|
||||
scannable = to_scannable(dataset)
|
||||
assert scannable.rescannable
|
||||
assert scannable.num_rows == 40_000
|
||||
|
||||
total = 0
|
||||
for batch in scannable.reader():
|
||||
assert batch.num_rows < _PA_DEFAULT_BATCH_ROWS
|
||||
total += batch.num_rows
|
||||
assert total == 40_000
|
||||
|
||||
|
||||
def test_sample_head_is_bounded_rows(tmp_path):
|
||||
# The peek itself must not read the whole dataset.
|
||||
_write_wide_dataset(str(tmp_path), files=1, rows_per_file=1000, fixed_size=False)
|
||||
dataset = ds.dataset(str(tmp_path), format="parquet")
|
||||
|
||||
sample = _sample_head(dataset.head)
|
||||
assert sample.num_rows == _SAMPLE_ROWS
|
||||
|
||||
|
||||
def test_sample_head_returns_none_for_empty_dataset(tmp_path):
|
||||
table = pa.table({"v": pa.array([], type=pa.list_(pa.float32()))})
|
||||
pq.write_table(table, f"{tmp_path}/empty.parquet")
|
||||
dataset = ds.dataset(str(tmp_path), format="parquet")
|
||||
|
||||
assert _sample_head(dataset.head) is None
|
||||
@@ -434,6 +434,29 @@ def test_add(mem_db: DBConnection):
|
||||
_add(table, schema)
|
||||
|
||||
|
||||
def test_add_write_parallelism(mem_db: DBConnection):
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = mem_db.create_table("test", schema=schema)
|
||||
|
||||
data = pa.table({"id": list(range(1000))}, schema=schema)
|
||||
table.add(data, write_parallelism=4)
|
||||
assert len(table) == 1000
|
||||
|
||||
# invalid parallelism is rejected
|
||||
with pytest.raises(ValueError, match="write_parallelism"):
|
||||
table.add(data, write_parallelism=0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_write_parallelism_async(mem_db_async: AsyncConnection):
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = await mem_db_async.create_table("test", schema=schema)
|
||||
|
||||
data = pa.table({"id": list(range(1000))}, schema=schema)
|
||||
await table.add(data, write_parallelism=4)
|
||||
assert await table.count_rows() == 1000
|
||||
|
||||
|
||||
def test_add_struct(mem_db: DBConnection):
|
||||
# https://github.com/lancedb/lancedb/issues/2114
|
||||
schema = pa.schema(
|
||||
|
||||
@@ -924,3 +924,23 @@ def test_sanitize_data_stream():
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
next(output)
|
||||
|
||||
|
||||
def test_infer_vector_column_raises_clear_error(tmp_path):
|
||||
"""Regression: querying a table with no inferable vector column should raise
|
||||
a clear ValueError, not a cryptic TypeError (issue #1653).
|
||||
|
||||
Previously, inf_vector_column_query silently returned None which then caused
|
||||
a confusing TypeError deep in schema lookup. The fix adds a ValueError guard
|
||||
so the user gets a direct, actionable error message.
|
||||
"""
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table(
|
||||
"no_vec",
|
||||
data=[{"id": 1, "text": "hello"}, {"id": 2, "text": "world"}],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="vector"):
|
||||
# Plain vector search on a table with no vector column should raise
|
||||
# a clear ValueError, not a cryptic TypeError.
|
||||
table.search([1.0, 2.0]).to_list()
|
||||
|
||||
@@ -0,0 +1,506 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Unit tests for WatsonxEmbeddings — no live API calls required."""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from lancedb.embeddings import get_registry
|
||||
from lancedb.embeddings.watsonx import CURRENT_MODELS, MODELS_DIMS, WatsonxEmbeddings
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_func(monkeypatch, env=None, **create_kwargs):
|
||||
"""
|
||||
Return a WatsonxEmbeddings instance with ibm_watsonx_ai mocked out.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
env : dict, optional
|
||||
Environment variables to inject (merged on top of an empty env so that
|
||||
no real WATSONX_* vars from the host bleed into the test).
|
||||
create_kwargs :
|
||||
Forwarded to ``WatsonxEmbeddings.create()``.
|
||||
"""
|
||||
base_env = {
|
||||
k: "" for k in ("WATSONX_API_KEY", "WATSONX_PROJECT_ID", "WATSONX_SPACE_ID")
|
||||
}
|
||||
base_env.update(env or {})
|
||||
# Only keep keys that have non-empty values so that absent vars are truly absent.
|
||||
clean_env = {k: v for k, v in base_env.items() if v}
|
||||
|
||||
mock_embeddings_instance = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
mock_foundation.Embeddings.return_value = mock_embeddings_instance
|
||||
mock_ibm = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
with patch.dict("os.environ", clean_env, clear=True):
|
||||
with patch(
|
||||
"lancedb.embeddings.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
func = get_registry().get("watsonx").create(**create_kwargs)
|
||||
# Force the cached_property to evaluate inside the patch context.
|
||||
_ = func._watsonx_client
|
||||
return func, mock_foundation
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRegistry:
|
||||
def test_watsonx_registered(self):
|
||||
assert get_registry().get("watsonx") is not None
|
||||
|
||||
def test_model_names_returns_only_current_models(self):
|
||||
names = WatsonxEmbeddings.model_names()
|
||||
assert names == list(CURRENT_MODELS.keys())
|
||||
# Current models must all be present.
|
||||
for name in (
|
||||
"ibm/granite-embedding-278m-multilingual",
|
||||
"ibm/slate-125m-english-rtrvr-v2",
|
||||
"ibm/slate-30m-english-rtrvr-v2",
|
||||
"intfloat/multilingual-e5-large",
|
||||
):
|
||||
assert name in names, f"{name!r} missing from model_names()"
|
||||
# Legacy / deprecated IDs must NOT appear in model_names().
|
||||
for legacy in (
|
||||
"ibm/slate-125m-english-rtrvr",
|
||||
"ibm/slate-30m-english-rtrvr",
|
||||
"sentence-transformers/all-minilm-l12-v2",
|
||||
"sentence-transformers/all-minilm-l6-v2",
|
||||
):
|
||||
assert legacy not in names, (
|
||||
f"Legacy model {legacy!r} should not appear in model_names()"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dimensions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDimensions:
|
||||
@pytest.mark.parametrize(
|
||||
"model_name,expected_dims",
|
||||
[
|
||||
("ibm/granite-embedding-278m-multilingual", 768),
|
||||
("ibm/slate-125m-english-rtrvr-v2", 768),
|
||||
("ibm/slate-30m-english-rtrvr-v2", 384),
|
||||
("intfloat/multilingual-e5-large", 1024),
|
||||
("sentence-transformers/all-minilm-l6-v2", 384),
|
||||
],
|
||||
)
|
||||
def test_current_model_dimensions(self, monkeypatch, model_name, expected_dims):
|
||||
func, _ = _make_func(
|
||||
monkeypatch,
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_PROJECT_ID": "proj"},
|
||||
name=model_name,
|
||||
)
|
||||
assert func.ndims() == expected_dims
|
||||
|
||||
def test_unknown_model_raises(self):
|
||||
func = WatsonxEmbeddings(name="not/a-real-model")
|
||||
with pytest.raises(ValueError, match="Unknown model"):
|
||||
func.ndims()
|
||||
|
||||
# -- Backward-compat: legacy names must still resolve dims on table load --
|
||||
@pytest.mark.parametrize(
|
||||
"legacy_name,expected_dims",
|
||||
[
|
||||
("ibm/slate-125m-english-rtrvr", 768),
|
||||
("ibm/slate-30m-english-rtrvr", 384),
|
||||
("sentence-transformers/all-minilm-l12-v2", 384),
|
||||
],
|
||||
)
|
||||
def test_legacy_model_dimensions_still_resolve(self, legacy_name, expected_dims):
|
||||
"""Tables written with old model names must not raise on reload."""
|
||||
assert MODELS_DIMS[legacy_name] == expected_dims
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Scope resolution (project_id / space_id)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestScopeResolution:
|
||||
def test_explicit_project_id(self, monkeypatch):
|
||||
func, mock_foundation = _make_func(
|
||||
monkeypatch,
|
||||
env={"WATSONX_API_KEY": "key"},
|
||||
project_id="explicit-proj",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Embeddings.call_args
|
||||
assert call_kwargs.get("project_id") == "explicit-proj"
|
||||
assert "space_id" not in call_kwargs
|
||||
|
||||
def test_explicit_space_id(self, monkeypatch):
|
||||
func, mock_foundation = _make_func(
|
||||
monkeypatch,
|
||||
env={"WATSONX_API_KEY": "key"},
|
||||
space_id="explicit-space",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Embeddings.call_args
|
||||
assert call_kwargs.get("space_id") == "explicit-space"
|
||||
assert "project_id" not in call_kwargs
|
||||
|
||||
def test_env_project_id_fallback(self, monkeypatch):
|
||||
func, mock_foundation = _make_func(
|
||||
monkeypatch,
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_PROJECT_ID": "env-proj"},
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Embeddings.call_args
|
||||
assert call_kwargs.get("project_id") == "env-proj"
|
||||
|
||||
def test_env_space_id_fallback(self, monkeypatch):
|
||||
func, mock_foundation = _make_func(
|
||||
monkeypatch,
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_SPACE_ID": "env-space"},
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Embeddings.call_args
|
||||
assert call_kwargs.get("space_id") == "env-space"
|
||||
|
||||
def test_explicit_project_id_wins_over_env_space_id(self, monkeypatch):
|
||||
"""Explicit project_id must not be overridden by WATSONX_SPACE_ID in env."""
|
||||
func, mock_foundation = _make_func(
|
||||
monkeypatch,
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_SPACE_ID": "stray-env-space"},
|
||||
project_id="explicit-proj",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Embeddings.call_args
|
||||
assert call_kwargs.get("project_id") == "explicit-proj"
|
||||
assert "space_id" not in call_kwargs
|
||||
|
||||
def test_explicit_space_id_wins_over_env_project_id(self, monkeypatch):
|
||||
"""Explicit space_id must not be overridden by WATSONX_PROJECT_ID in env."""
|
||||
func, mock_foundation = _make_func(
|
||||
monkeypatch,
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_PROJECT_ID": "stray-env-proj"},
|
||||
space_id="explicit-space",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Embeddings.call_args
|
||||
assert call_kwargs.get("space_id") == "explicit-space"
|
||||
assert "project_id" not in call_kwargs
|
||||
|
||||
def test_both_env_vars_raises(self, monkeypatch):
|
||||
"""When both WATSONX_PROJECT_ID and WATSONX_SPACE_ID env vars are set
|
||||
(and neither is passed explicitly), it must raise 'not both'."""
|
||||
func = WatsonxEmbeddings(name="ibm/granite-embedding-278m-multilingual")
|
||||
mock_ibm = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
with patch.dict(
|
||||
"os.environ",
|
||||
{
|
||||
"WATSONX_API_KEY": "key",
|
||||
"WATSONX_PROJECT_ID": "env-proj",
|
||||
"WATSONX_SPACE_ID": "env-space",
|
||||
},
|
||||
clear=True,
|
||||
):
|
||||
with patch(
|
||||
"lancedb.embeddings.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
with pytest.raises(ValueError, match="not both"):
|
||||
_ = func._watsonx_client
|
||||
|
||||
def test_both_explicit_raises(self):
|
||||
func = WatsonxEmbeddings(
|
||||
name="ibm/granite-embedding-278m-multilingual",
|
||||
project_id="p",
|
||||
space_id="s",
|
||||
)
|
||||
# The error surfaces when _watsonx_client is first accessed.
|
||||
mock_ibm = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
with patch.dict("os.environ", {"WATSONX_API_KEY": "key"}, clear=True):
|
||||
with patch(
|
||||
"lancedb.embeddings.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
with pytest.raises(ValueError, match="not both"):
|
||||
_ = func._watsonx_client
|
||||
|
||||
def test_neither_raises(self):
|
||||
func = WatsonxEmbeddings(name="ibm/granite-embedding-278m-multilingual")
|
||||
mock_ibm = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
with patch.dict("os.environ", {"WATSONX_API_KEY": "key"}, clear=True):
|
||||
with patch(
|
||||
"lancedb.embeddings.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError, match="WATSONX_PROJECT_ID or WATSONX_SPACE_ID"
|
||||
):
|
||||
_ = func._watsonx_client
|
||||
|
||||
def test_missing_api_key_raises(self):
|
||||
func = WatsonxEmbeddings(name="ibm/granite-embedding-278m-multilingual")
|
||||
mock_ibm = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
with patch.dict("os.environ", {"WATSONX_PROJECT_ID": "proj"}, clear=True):
|
||||
with patch(
|
||||
"lancedb.embeddings.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
with pytest.raises(ValueError, match="WATSONX_API_KEY"):
|
||||
_ = func._watsonx_client
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Metadata round-trip (backward compat)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMetadataRoundTrip:
|
||||
def test_reload_with_empty_model_metadata_preserves_model(self):
|
||||
"""
|
||||
Reproduce the exact deserialization path used by the registry:
|
||||
|
||||
create(**{}) → safe_model_dump() == {}
|
||||
→ stored as model: {}
|
||||
→ reloaded via create(**{})
|
||||
|
||||
The model must be identical before and after — no silent switch.
|
||||
This guards against changing the class-level default between releases.
|
||||
"""
|
||||
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
|
||||
registry = EmbeddingFunctionRegistry.get_instance()
|
||||
|
||||
# Simulate original table creation with no explicit args.
|
||||
original = registry.get("watsonx").create()
|
||||
stored = original.safe_model_dump() # what gets written to arrow metadata
|
||||
|
||||
assert stored == {}, (
|
||||
f"Expected empty stored args when create() called with no kwargs; "
|
||||
f"got {stored!r}"
|
||||
)
|
||||
|
||||
# Simulate reload: registry calls create(**stored) == create(**{})
|
||||
reloaded = registry.get("watsonx").create(**stored)
|
||||
|
||||
assert reloaded.name == original.name, (
|
||||
f"Model changed on reload: was {original.name!r}, "
|
||||
f"became {reloaded.name!r}. "
|
||||
"The class-level default must not change without a migration path."
|
||||
)
|
||||
|
||||
def test_reload_from_legacy_metadata_explicit(self):
|
||||
"""
|
||||
Deserialize a representative legacy metadata payload and assert the exact
|
||||
model name — this is the real cross-version regression guard.
|
||||
|
||||
Tables created before the v2 rename stored ``model: {"name": ...}`` with
|
||||
the pre-v2 name. Reloading must produce exactly that model, not silently
|
||||
switch to the current class default.
|
||||
"""
|
||||
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
|
||||
registry = EmbeddingFunctionRegistry.get_instance()
|
||||
|
||||
# This is what is stored in Arrow metadata for a table created with the
|
||||
# pre-v2 default model name (no explicit name= was passed at the time).
|
||||
legacy_stored = {"name": "ibm/slate-125m-english-rtrvr"}
|
||||
|
||||
reloaded = registry.get("watsonx").create(**legacy_stored)
|
||||
|
||||
assert reloaded.name == "ibm/slate-125m-english-rtrvr", (
|
||||
f"Legacy metadata reload returned {reloaded.name!r} instead of "
|
||||
"'ibm/slate-125m-english-rtrvr'. "
|
||||
"MODELS_DIMS must keep legacy entries for backward compat."
|
||||
)
|
||||
|
||||
def test_legacy_model_names_resolve_dims(self):
|
||||
"""Legacy names in MODELS_DIMS so ndims() never raises on old tables."""
|
||||
assert MODELS_DIMS["ibm/slate-125m-english-rtrvr"] == 768
|
||||
assert MODELS_DIMS["ibm/slate-30m-english-rtrvr"] == 384
|
||||
assert MODELS_DIMS["sentence-transformers/all-minilm-l12-v2"] == 384
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# WatsonxReranker — scope resolution (project_id / space_id)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_reranker(env=None, **init_kwargs):
|
||||
"""
|
||||
Return a WatsonxReranker with ibm_watsonx_ai mocked out.
|
||||
|
||||
Scope precedence is tested by inspecting what was passed to Rerank().
|
||||
"""
|
||||
from lancedb.rerankers.watsonx import WatsonxReranker
|
||||
|
||||
base_env = {
|
||||
k: "" for k in ("WATSONX_API_KEY", "WATSONX_PROJECT_ID", "WATSONX_SPACE_ID")
|
||||
}
|
||||
base_env.update(env or {})
|
||||
clean_env = {k: v for k, v in base_env.items() if v}
|
||||
|
||||
mock_rerank_instance = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
mock_foundation.Rerank.return_value = mock_rerank_instance
|
||||
mock_ibm = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
reranker = WatsonxReranker(**init_kwargs)
|
||||
with patch.dict("os.environ", clean_env, clear=True):
|
||||
with patch(
|
||||
"lancedb.rerankers.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
_ = reranker._client
|
||||
return reranker, mock_foundation
|
||||
|
||||
|
||||
class TestRerankerScopeResolution:
|
||||
def test_explicit_project_id(self):
|
||||
_, mock_foundation = _make_reranker(
|
||||
env={"WATSONX_API_KEY": "key"},
|
||||
project_id="explicit-proj",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Rerank.call_args
|
||||
assert call_kwargs.get("project_id") == "explicit-proj"
|
||||
assert "space_id" not in call_kwargs
|
||||
|
||||
def test_explicit_space_id(self):
|
||||
_, mock_foundation = _make_reranker(
|
||||
env={"WATSONX_API_KEY": "key"},
|
||||
space_id="explicit-space",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Rerank.call_args
|
||||
assert call_kwargs.get("space_id") == "explicit-space"
|
||||
assert "project_id" not in call_kwargs
|
||||
|
||||
def test_env_project_id_fallback(self):
|
||||
_, mock_foundation = _make_reranker(
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_PROJECT_ID": "env-proj"},
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Rerank.call_args
|
||||
assert call_kwargs.get("project_id") == "env-proj"
|
||||
|
||||
def test_env_space_id_fallback(self):
|
||||
_, mock_foundation = _make_reranker(
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_SPACE_ID": "env-space"},
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Rerank.call_args
|
||||
assert call_kwargs.get("space_id") == "env-space"
|
||||
|
||||
def test_explicit_project_id_wins_over_env_space_id(self):
|
||||
"""Explicit project_id must not be overridden by WATSONX_SPACE_ID in env."""
|
||||
_, mock_foundation = _make_reranker(
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_SPACE_ID": "stray-env-space"},
|
||||
project_id="explicit-proj",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Rerank.call_args
|
||||
assert call_kwargs.get("project_id") == "explicit-proj"
|
||||
assert "space_id" not in call_kwargs
|
||||
|
||||
def test_explicit_space_id_wins_over_env_project_id(self):
|
||||
"""Explicit space_id must not be overridden by WATSONX_PROJECT_ID in env."""
|
||||
_, mock_foundation = _make_reranker(
|
||||
env={"WATSONX_API_KEY": "key", "WATSONX_PROJECT_ID": "stray-env-proj"},
|
||||
space_id="explicit-space",
|
||||
)
|
||||
_, call_kwargs = mock_foundation.Rerank.call_args
|
||||
assert call_kwargs.get("space_id") == "explicit-space"
|
||||
assert "project_id" not in call_kwargs
|
||||
|
||||
def test_both_explicit_raises(self):
|
||||
from lancedb.rerankers.watsonx import WatsonxReranker
|
||||
|
||||
reranker = WatsonxReranker(project_id="p", space_id="s")
|
||||
mock_ibm = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
with patch.dict("os.environ", {"WATSONX_API_KEY": "key"}, clear=True):
|
||||
with patch(
|
||||
"lancedb.rerankers.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
with pytest.raises(ValueError, match="not both"):
|
||||
_ = reranker._client
|
||||
|
||||
def test_neither_raises(self):
|
||||
from lancedb.rerankers.watsonx import WatsonxReranker
|
||||
|
||||
reranker = WatsonxReranker()
|
||||
mock_ibm = MagicMock()
|
||||
mock_foundation = MagicMock()
|
||||
|
||||
def _fake_import(name):
|
||||
if name == "ibm_watsonx_ai":
|
||||
return mock_ibm
|
||||
if name == "ibm_watsonx_ai.foundation_models":
|
||||
return mock_foundation
|
||||
raise ImportError(name)
|
||||
|
||||
with patch.dict("os.environ", {"WATSONX_API_KEY": "key"}, clear=True):
|
||||
with patch(
|
||||
"lancedb.rerankers.watsonx.attempt_import_or_raise",
|
||||
side_effect=_fake_import,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError, match="WATSONX_PROJECT_ID or WATSONX_SPACE_ID"
|
||||
):
|
||||
_ = reranker._client
|
||||
+5
-456
@@ -18,10 +18,7 @@ use lancedb::{
|
||||
connection::Connection as LanceConnection,
|
||||
connection::NamespaceClientPushdownOperation,
|
||||
database::namespace::LanceNamespaceDatabase,
|
||||
database::{
|
||||
CreateFunctionRequest, CreateMaterializedViewRequest, CreateTableMode, Database,
|
||||
ReadConsistency, RefreshMaterializedViewRequest, TableLineageRequest,
|
||||
},
|
||||
database::{CreateTableMode, Database, ReadConsistency},
|
||||
};
|
||||
use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
@@ -30,107 +27,6 @@ use pyo3::{
|
||||
types::{PyDict, PyDictMethods},
|
||||
};
|
||||
|
||||
/// A registered function, as returned by `list_functions`.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct FunctionInfo {
|
||||
pub name: String,
|
||||
pub language: String,
|
||||
pub return_type: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// A registered materialized view definition.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct MaterializedViewInfo {
|
||||
pub name: String,
|
||||
pub source_table: String,
|
||||
pub projection: Vec<String>,
|
||||
pub udf_columns: Vec<String>,
|
||||
pub filter: Option<String>,
|
||||
pub auto_refresh: bool,
|
||||
}
|
||||
|
||||
/// One inflight server-side job.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub age_seconds: Option<i64>,
|
||||
pub command: Option<String>,
|
||||
pub units_done: Option<i64>,
|
||||
pub units_total: Option<i64>,
|
||||
pub committed: bool,
|
||||
pub rows_skipped: u64,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// A described platform job (POST /v1/jobs/describe).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct PlatformJobDescription {
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub job_subtype: String,
|
||||
/// "IN_PROGRESS" | "CANCELLED" | "FAILED" | "DONE".
|
||||
pub job_state: String,
|
||||
pub creation_ms: i64,
|
||||
/// The owner-written status payload as a JSON string (units_done /
|
||||
/// units_total / rows_committed / error when present).
|
||||
pub status_json: String,
|
||||
}
|
||||
|
||||
/// One durable, completed/terminal server-side job record (SHOW JOB HISTORY).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobHistoryEntry {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub created_ms: i64,
|
||||
pub updated_ms: i64,
|
||||
pub completed_ms: Option<i64>,
|
||||
pub rows_processed: Option<i64>,
|
||||
pub rows_skipped: Option<i64>,
|
||||
pub error: Option<String>,
|
||||
pub events: Option<String>,
|
||||
}
|
||||
|
||||
/// One per-row UDF error recorded by `error_policy=skip` (SHOW ERRORS).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobErrorEntry {
|
||||
pub job_id: String,
|
||||
pub table: String,
|
||||
pub column: String,
|
||||
pub error_type: String,
|
||||
pub error_message: String,
|
||||
pub fragment_id: Option<i64>,
|
||||
pub source_row_id: Option<i64>,
|
||||
pub table_version: Option<i64>,
|
||||
pub age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
/// The plan a REFRESH MATERIALIZED VIEW would execute (EXPLAIN REFRESH).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct MvRefreshPlan {
|
||||
pub table_name: String,
|
||||
pub has_work: bool,
|
||||
pub source_version: u64,
|
||||
pub last_refreshed_version: Option<u64>,
|
||||
pub full_refresh: bool,
|
||||
pub rebuild: bool,
|
||||
pub units_total: u64,
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct Connection {
|
||||
inner: Option<LanceConnection>,
|
||||
@@ -414,357 +310,6 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, language, return_type, body, options=None))]
|
||||
pub fn create_function(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
language: String,
|
||||
return_type: String,
|
||||
body: String,
|
||||
options: Option<HashMap<String, String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_function(CreateFunctionRequest {
|
||||
name,
|
||||
language,
|
||||
return_type,
|
||||
body,
|
||||
options: options.unwrap_or_default(),
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let functions = inner.list_functions().await.infer_error()?;
|
||||
Ok(functions
|
||||
.into_iter()
|
||||
.map(|f| FunctionInfo {
|
||||
name: f.name,
|
||||
language: f.language,
|
||||
return_type: f.return_type,
|
||||
description: f.description,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_function(self_: PyRef<'_, Self>, name: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_function(&name).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, query, auto_refresh=false, with_no_data=false, partition_by=None))]
|
||||
pub fn create_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
query: String,
|
||||
auto_refresh: bool,
|
||||
with_no_data: bool,
|
||||
partition_by: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_materialized_view(CreateMaterializedViewRequest {
|
||||
name,
|
||||
query,
|
||||
auto_refresh,
|
||||
with_no_data,
|
||||
partition_by,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, full=false, src_version=None, num_workers=None, max_workers=None))]
|
||||
pub fn refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.refresh_materialized_view(RefreshMaterializedViewRequest {
|
||||
name,
|
||||
full,
|
||||
src_version,
|
||||
num_workers,
|
||||
max_workers,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Derived-compute lineage of a table/view (or column), returned as the
|
||||
/// server's lineage JSON string (the Python layer parses it).
|
||||
pub fn table_lineage(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
column: Option<String>,
|
||||
direction: Option<String>,
|
||||
depth: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.table_lineage(TableLineageRequest {
|
||||
name,
|
||||
column,
|
||||
direction,
|
||||
depth,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, full=false, src_version=None))]
|
||||
pub fn explain_refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let p = inner
|
||||
.explain_refresh_materialized_view(&name, full, src_version)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(MvRefreshPlan {
|
||||
table_name: p.table_name,
|
||||
has_work: p.has_work,
|
||||
source_version: p.source_version,
|
||||
last_refreshed_version: p.last_refreshed_version,
|
||||
full_refresh: p.full_refresh,
|
||||
rebuild: p.rebuild,
|
||||
units_total: p.units_total,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn alter_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
auto_refresh: bool,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.alter_materialized_view(&name, auto_refresh)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_materialized_view(&name).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let views = inner.list_materialized_views().await.infer_error()?;
|
||||
Ok(views
|
||||
.into_iter()
|
||||
.map(|v| MaterializedViewInfo {
|
||||
name: v.name,
|
||||
source_table: v.source_table,
|
||||
projection: v.projection,
|
||||
udf_columns: v.udf_columns,
|
||||
filter: v.filter,
|
||||
auto_refresh: v.auto_refresh,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let jobs = inner.list_jobs().await.infer_error()?;
|
||||
Ok(jobs
|
||||
.into_iter()
|
||||
.map(|j| JobInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
age_seconds: j.age_seconds,
|
||||
command: j.command,
|
||||
units_done: j.units_done,
|
||||
units_total: j.units_total,
|
||||
committed: j.committed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.cancel_job(&job_id).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn describe_platform_job(
|
||||
self_: PyRef<'_, Self>,
|
||||
platform_job_id: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let described = inner
|
||||
.describe_platform_job(&platform_job_id)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(described.map(|d| PlatformJobDescription {
|
||||
job_id: d.job_id,
|
||||
job_type: d.job_type,
|
||||
job_subtype: d.job_subtype,
|
||||
job_state: d.job_state,
|
||||
creation_ms: d.creation_ms,
|
||||
status_json: d.status.to_string(),
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (manifest_job_id, table=None))]
|
||||
pub fn resolve_platform_job_id(
|
||||
self_: PyRef<'_, Self>,
|
||||
manifest_job_id: String,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.resolve_platform_job_id(&manifest_job_id, table.as_deref())
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cancel_platform_job(
|
||||
self_: PyRef<'_, Self>,
|
||||
platform_job_id: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.cancel_platform_job(&platform_job_id)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id, table=None))]
|
||||
pub fn get_job(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: String,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let job = inner
|
||||
.get_job(&job_id, table.as_deref())
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(job.map(|j| JobInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
age_seconds: j.age_seconds,
|
||||
command: j.command,
|
||||
units_done: j.units_done,
|
||||
units_total: j.units_total,
|
||||
committed: j.committed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None))]
|
||||
pub fn job_history(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let rows = inner.job_history(job_id.as_deref()).await.infer_error()?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|r| JobHistoryEntry {
|
||||
table: r.table,
|
||||
job_id: r.job_id,
|
||||
job_type: r.job_type,
|
||||
state: r.state,
|
||||
column: r.column,
|
||||
created_ms: r.created_ms,
|
||||
updated_ms: r.updated_ms,
|
||||
completed_ms: r.completed_ms,
|
||||
rows_processed: r.rows_processed,
|
||||
rows_skipped: r.rows_skipped,
|
||||
error: r.error,
|
||||
events: r.events,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None, table=None))]
|
||||
pub fn errors(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let rows = inner
|
||||
.errors(job_id.as_deref(), table.as_deref())
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|e| JobErrorEntry {
|
||||
job_id: e.job_id,
|
||||
table: e.table,
|
||||
column: e.column,
|
||||
error_type: e.error_type,
|
||||
error_message: e.error_message,
|
||||
fragment_id: e.fragment_id,
|
||||
source_row_id: e.source_row_id,
|
||||
table_version: e.table_version,
|
||||
age_seconds: e.age_seconds,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (cur_name, new_name, cur_namespace_path=None, new_namespace_path=None))]
|
||||
pub fn rename_table(
|
||||
self_: PyRef<'_, Self>,
|
||||
@@ -1255,6 +800,10 @@ impl From<PyClientConfig> for lancedb::remote::ClientConfig {
|
||||
tls_config: value.tls_config.map(Into::into),
|
||||
header_provider,
|
||||
user_id: value.user_id,
|
||||
// Resolved from LANCE_CLIENT_MAX_BYTES_PER_REQUEST or the default.
|
||||
max_bytes_per_request: None,
|
||||
// Resolved from LANCE_CLIENT_MAX_REQUEST_DURATION or the read timeout.
|
||||
max_request_duration: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,12 +42,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
.write_style("LANCEDB_LOG_STYLE");
|
||||
env_logger::init_from_env(env);
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<connection::FunctionInfo>()?;
|
||||
m.add_class::<connection::MaterializedViewInfo>()?;
|
||||
m.add_class::<connection::JobInfo>()?;
|
||||
m.add_class::<connection::PlatformJobDescription>()?;
|
||||
m.add_class::<connection::JobHistoryEntry>()?;
|
||||
m.add_class::<connection::JobErrorEntry>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<PyBlobFile>()?;
|
||||
|
||||
+44
-82
@@ -21,8 +21,7 @@ use lancedb::blob::BlobFile;
|
||||
use lancedb::index::scalar::FtsIndexBuilder;
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
LoadColumnsRequest, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref,
|
||||
Table as LanceDbTable,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use pyo3::{
|
||||
@@ -626,12 +625,13 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (data, mode, progress=None))]
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
|
||||
pub fn add<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
data: PyScannable,
|
||||
mode: String,
|
||||
progress: Option<Py<PyAny>>,
|
||||
write_parallelism: Option<usize>,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let mut op = self_.inner_ref()?.add(data);
|
||||
if mode == "append" {
|
||||
@@ -641,6 +641,9 @@ impl Table {
|
||||
} else {
|
||||
return Err(PyValueError::new_err(format!("Invalid mode: {}", mode)));
|
||||
}
|
||||
if let Some(write_parallelism) = write_parallelism {
|
||||
op = op.write_parallelism(write_parallelism);
|
||||
}
|
||||
if let Some(progress_obj) = progress {
|
||||
let is_callable = Python::attach(|py| progress_obj.bind(py).is_callable());
|
||||
if is_callable {
|
||||
@@ -794,8 +797,8 @@ impl Table {
|
||||
}
|
||||
|
||||
future_into_py(self_.py(), async move {
|
||||
let job_id = op.execute().await.infer_error()?;
|
||||
Ok(job_id)
|
||||
op.execute().await.infer_error()?;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1295,83 +1298,6 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_computed_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
columns: Vec<(String, String)>,
|
||||
expression: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.add_computed_columns(&columns, &expression)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (columns, where_clause=None, num_workers=None, max_workers=None, batch_size=None, priority=None))]
|
||||
pub fn refresh_column(
|
||||
self_: PyRef<'_, Self>,
|
||||
columns: Vec<String>,
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.refresh_column(
|
||||
&columns,
|
||||
where_clause,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
priority,
|
||||
)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
#[pyo3(signature = (source_uris, source_format, target_key, columns, source_key=None, source_storage_options=None, on_missing=None, num_workers=None, max_workers=None, batch_size=None, commit_granularity=None, priority=None))]
|
||||
pub fn load_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
source_uris: Vec<String>,
|
||||
source_format: String,
|
||||
target_key: String,
|
||||
columns: Vec<(String, Option<String>)>,
|
||||
source_key: Option<String>,
|
||||
source_storage_options: Option<std::collections::HashMap<String, String>>,
|
||||
on_missing: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
commit_granularity: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let request = LoadColumnsRequest {
|
||||
source_uris,
|
||||
source_format,
|
||||
source_storage_options,
|
||||
target_key,
|
||||
source_key,
|
||||
columns,
|
||||
on_missing,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
commit_granularity,
|
||||
priority,
|
||||
};
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.load_columns(request).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
definitions: Vec<(String, String)>,
|
||||
@@ -1667,4 +1593,40 @@ impl Branches {
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn diff(self_: PyRef<'_, Self>, from_branch: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let diff = inner.diff_branch(&from_branch).await.infer_error()?;
|
||||
Python::attach(|py| struct_to_wire_py(py, &diff))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (from_branch, dry_run=false))]
|
||||
pub fn merge(
|
||||
self_: PyRef<'_, Self>,
|
||||
from_branch: String,
|
||||
dry_run: bool,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = inner
|
||||
.merge_branch(&from_branch, dry_run)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Python::attach(|py| struct_to_wire_py(py, &result))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode a serde value as the wire JSON object (camelCase keys).
|
||||
fn struct_to_wire_py(py: Python<'_>, value: &impl serde::Serialize) -> PyResult<Py<PyAny>> {
|
||||
let json = py.import("json")?;
|
||||
Ok(json
|
||||
.call_method1(
|
||||
"loads",
|
||||
(serde_json::to_string(value)
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("failed to serialize json: {e}")))?,),
|
||||
)?
|
||||
.unbind())
|
||||
}
|
||||
|
||||
Generated
+9
-9
@@ -1998,7 +1998,7 @@ requires-dist = [
|
||||
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
|
||||
{ name = "pydantic", specifier = ">=1.10" },
|
||||
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0rc1" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0" },
|
||||
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },
|
||||
{ name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" },
|
||||
{ name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" },
|
||||
@@ -3854,8 +3854,8 @@ crypto = [
|
||||
|
||||
[[package]]
|
||||
name = "pylance"
|
||||
version = "9.0.0rc1"
|
||||
source = { registry = "https://pypi.fury.io/lance-format" }
|
||||
version = "9.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
@@ -3863,12 +3863,12 @@ dependencies = [
|
||||
{ name = "pyarrow" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_vEHBE/pylance-9.0.0rc1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:f0b6b02a1808bb3072ee7fe4e36614cae6f86302513e73ec7f55b2234a963b24" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1Jipm4/pylance-9.0.0rc1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:30f0ebf0d88034301819eb964f9236ce555aaa58e7ab89c5975a3e2250bbb405" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_IvKxo/pylance-9.0.0rc1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:44609ea2615ea6e684b85478d1694af2026458f61cf7895ecc75e238bfd17aa8" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2hidj1/pylance-9.0.0rc1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:182167a8dba9eeabffbffd53bd5b8548613d4d459b7cd7b34a840dd00cbb806f" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1dFx3r/pylance-9.0.0rc1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8a63b11e814b7eab758bcaf0d6f97eb05ea86203d9fb0af718c462c24c7d6c9c" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2a8dSh/pylance-9.0.0rc1-cp310-abi3-win_amd64.whl", hash = "sha256:2ff8b953ae2b0550490c1a7efd210aa91bc223d200ffac28849056cfd7436d97" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e9/be/45733acd64801991852aac8e658601fd8fc12f76ceb81e57fca690896b90/pylance-9.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:8257213501d3298c5b6a344d60938e4bbe4de9f00cd3265371a56d1dc3dd15ca", size = 68377982, upload-time = "2026-07-24T16:53:45.247Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d0/3e/1ef707cb215cc7268c63ad84a91344ad6313d3984b343eeb19a9b708698e/pylance-9.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:804eedfa1fda2e703cca8580c76f0b44a1b849b725a8908c06f8acfca811732f", size = 71844362, upload-time = "2026-07-24T16:56:07.6Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/fb/a499e5c53ddb75c7de44100fd2bb1f7cc735966200d20292eed7af9ef552/pylance-9.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0a0b75595e3766c1d5f4c90abdc52337b4de60f1a52d123a4f8c4e5bcbbbfa8f", size = 75663088, upload-time = "2026-07-24T17:10:31.283Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e9/80/0714e09f64a68dbdf62558955e737a7436df763b9834a6aba506861b5352/pylance-9.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:d2f69c5c390ae3710c35a429905fe15f769951777fafb1b72a359e035ec121f7", size = 71866858, upload-time = "2026-07-24T16:56:35.937Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4b/3c/78d3a6d6ca0d843b7c3ac0c30d9cd2cf4635b0c2cabe6ec66583d5bfc1a1/pylance-9.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:836268a7832d62f3d5ccbe1c3fca239621971d90297bcfff14a70b3cb6842aa8", size = 75642656, upload-time = "2026-07-24T17:13:11.879Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/cf/c1/dc9c9a31e171530ec0add024d922488c046437d32a7087c29e254a7eacc7/pylance-9.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:96441d27a5ed3805300388ccf8f31835cd280c98751f80b6a1a11dcd6808fc43", size = 81668288, upload-time = "2026-07-24T17:05:38.707Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.32.0-beta.2"
|
||||
version = "0.33.0-beta.0"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
@@ -34,7 +34,6 @@ datafusion.workspace = true
|
||||
object_store = { workspace = true }
|
||||
snafu = { workspace = true }
|
||||
half = { workspace = true }
|
||||
lazy_static.workspace = true
|
||||
lance = { workspace = true }
|
||||
lance-core = { workspace = true }
|
||||
lance-datafusion.workspace = true
|
||||
|
||||
@@ -118,12 +118,8 @@ async fn create_empty_table(db: &Connection) -> Result<LanceDbTable> {
|
||||
|
||||
async fn create_index(table: &LanceDbTable) -> Result<()> {
|
||||
// --8<-- [start:create_index]
|
||||
table
|
||||
.create_index(&["vector"], Index::Auto)
|
||||
.execute()
|
||||
.await?;
|
||||
table.create_index(&["vector"], Index::Auto).execute().await
|
||||
// --8<-- [end:create_index]
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn search(table: &LanceDbTable) -> Result<Vec<RecordBatch>> {
|
||||
|
||||
@@ -23,10 +23,8 @@ use crate::connection::create_table::CreateTableBuilder;
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::database::listing::ListingDatabase;
|
||||
use crate::database::{
|
||||
CloneTableRequest, CreateFunctionRequest, CreateMaterializedViewRequest, Database,
|
||||
DatabaseOptions, FunctionInfo, JobErrorInfo, JobHistoryInfo, JobInfo, MaterializedViewInfo,
|
||||
MvRefreshPlan, OpenTableRequest, PlatformJobDescription, ReadConsistency,
|
||||
RefreshMaterializedViewRequest, TableLineageRequest, TableNamesRequest,
|
||||
CloneTableRequest, Database, DatabaseOptions, OpenTableRequest, ReadConsistency,
|
||||
TableNamesRequest,
|
||||
};
|
||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
@@ -490,140 +488,6 @@ impl Connection {
|
||||
)
|
||||
}
|
||||
|
||||
// -- Derived compute: functions, materialized views, jobs -------------
|
||||
// Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
// databases return NotSupported for now.
|
||||
|
||||
/// Register a UDF (CREATE FUNCTION).
|
||||
pub async fn create_function(&self, request: CreateFunctionRequest) -> Result<()> {
|
||||
self.internal.create_function(request).await
|
||||
}
|
||||
|
||||
/// List registered functions (SHOW FUNCTIONS).
|
||||
pub async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
self.internal.list_functions().await
|
||||
}
|
||||
|
||||
/// Drop a registered function (DROP FUNCTION).
|
||||
pub async fn drop_function(&self, name: &str) -> Result<()> {
|
||||
self.internal.drop_function(name).await
|
||||
}
|
||||
|
||||
/// Create a materialized view (CREATE MATERIALIZED VIEW). Returns
|
||||
/// the initial-population job id, absent when `with_no_data`.
|
||||
pub async fn create_materialized_view(
|
||||
&self,
|
||||
request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
self.internal.create_materialized_view(request).await
|
||||
}
|
||||
|
||||
/// Refresh a materialized view; returns the refresh job id.
|
||||
pub async fn refresh_materialized_view(
|
||||
&self,
|
||||
request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
self.internal.refresh_materialized_view(request).await
|
||||
}
|
||||
|
||||
/// Derived-compute lineage of a table/view (or column), as server-defined
|
||||
/// JSON. Read-only.
|
||||
pub async fn table_lineage(&self, request: TableLineageRequest) -> Result<String> {
|
||||
self.internal.table_lineage(request).await
|
||||
}
|
||||
|
||||
/// Plan a materialized-view refresh without submitting work
|
||||
/// (EXPLAIN REFRESH).
|
||||
pub async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
name: &str,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
self.internal
|
||||
.explain_refresh_materialized_view(name, full, src_version)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Update a materialized view's options (ALTER MATERIALIZED VIEW).
|
||||
pub async fn alter_materialized_view(&self, name: &str, auto_refresh: bool) -> Result<()> {
|
||||
self.internal
|
||||
.alter_materialized_view(name, auto_refresh)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Drop a materialized view definition (DROP MATERIALIZED VIEW).
|
||||
pub async fn drop_materialized_view(&self, name: &str) -> Result<()> {
|
||||
self.internal.drop_materialized_view(name).await
|
||||
}
|
||||
|
||||
/// List registered materialized view definitions.
|
||||
pub async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
self.internal.list_materialized_views().await
|
||||
}
|
||||
|
||||
/// List inflight server-side jobs across the database's tables.
|
||||
pub async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
self.internal.list_jobs().await
|
||||
}
|
||||
|
||||
/// Cancel an inflight server-side job by id. Returns true if a
|
||||
/// matching inflight job was flagged for cancellation.
|
||||
pub async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
||||
self.internal.cancel_job(job_id).await
|
||||
}
|
||||
|
||||
/// Describe a platform job (`POST /v1/jobs/describe`): registry-backed
|
||||
/// lifecycle state plus the owner-written status payload. `None` when the
|
||||
/// registry has no such job.
|
||||
pub async fn describe_platform_job(
|
||||
&self,
|
||||
platform_job_id: &str,
|
||||
) -> Result<Option<PlatformJobDescription>> {
|
||||
self.internal.describe_platform_job(platform_job_id).await
|
||||
}
|
||||
|
||||
/// Resolve a submission (manifest) job id to its platform job id. `None`
|
||||
/// until the job has registered (dispatch is async).
|
||||
pub async fn resolve_platform_job_id(
|
||||
&self,
|
||||
manifest_job_id: &str,
|
||||
table_hint: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
self.internal
|
||||
.resolve_platform_job_id(manifest_job_id, table_hint)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Cancel a platform job. Idempotent on already-terminal jobs.
|
||||
pub async fn cancel_platform_job(&self, platform_job_id: &str) -> Result<()> {
|
||||
self.internal.cancel_platform_job(platform_job_id).await
|
||||
}
|
||||
|
||||
/// Look up a single server-side job by id -- the `wait()`/status poll path.
|
||||
/// `table_hint` (the job's table) enables an O(1) server-side lookup; `None`
|
||||
/// scans the database's active jobs. A `None` result means unknown / not
|
||||
/// active.
|
||||
pub async fn get_job(&self, job_id: &str, table_hint: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
self.internal.get_job(job_id, table_hint).await
|
||||
}
|
||||
|
||||
/// Durable job history (SHOW JOB HISTORY) across the database's tables.
|
||||
/// Pass `job_id` to narrow to a single job.
|
||||
pub async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
self.internal.job_history(job_id).await
|
||||
}
|
||||
|
||||
/// Per-row UDF errors (SHOW ERRORS) across the database's tables, optionally
|
||||
/// filtered by `job_id` and/or `table`.
|
||||
pub async fn errors(
|
||||
&self,
|
||||
job_id: Option<&str>,
|
||||
table: Option<&str>,
|
||||
) -> Result<Vec<JobErrorInfo>> {
|
||||
self.internal.errors(job_id, table).await
|
||||
}
|
||||
|
||||
/// Rename a table in the database.
|
||||
///
|
||||
/// This is only supported in LanceDB Cloud.
|
||||
|
||||
@@ -27,7 +27,7 @@ use lance_namespace::models::{
|
||||
};
|
||||
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::error::Result;
|
||||
use crate::table::{BaseTable, WriteOptions};
|
||||
|
||||
pub mod listing;
|
||||
@@ -200,222 +200,6 @@ pub enum ReadConsistency {
|
||||
Strong,
|
||||
}
|
||||
|
||||
/// A request to register a UDF (CREATE FUNCTION).
|
||||
///
|
||||
/// Functions are first-class database objects, decoupled from any
|
||||
/// column; computed columns and materialized views reference them by
|
||||
/// name. Server-backed feature (LanceDB Enterprise / Cloud).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateFunctionRequest {
|
||||
/// Function name.
|
||||
pub name: String,
|
||||
/// Implementation language (currently "python").
|
||||
pub language: String,
|
||||
/// SQL return type, e.g. `FLOAT`, `FLOAT[1536]`,
|
||||
/// `STRUCT(a FLOAT, b VARCHAR)`, `TABLE(chunk VARCHAR, idx INT)`.
|
||||
pub return_type: String,
|
||||
/// Function body: source text, or base64 cloudpickle bytes when
|
||||
/// `options["body_format"] = "cloudpickle"`.
|
||||
pub body: String,
|
||||
/// Options: input_columns, pip, num_gpus, batch_size, timeout,
|
||||
/// error_policy, docker_image, body_format, ...
|
||||
pub options: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// A registered function, as returned by `list_functions`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FunctionInfo {
|
||||
pub name: String,
|
||||
pub language: String,
|
||||
pub return_type: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// A request to create a materialized view (CREATE MATERIALIZED VIEW).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateMaterializedViewRequest {
|
||||
/// View name.
|
||||
pub name: String,
|
||||
/// The view's SELECT statement, e.g.
|
||||
/// `SELECT id, embed(body) AS vec FROM articles WHERE id > 1`.
|
||||
/// Bare columns project through; function-call columns compute via
|
||||
/// registered UDFs (a RETURNS TABLE function makes a row-expanding
|
||||
/// chunker view).
|
||||
pub query: String,
|
||||
/// Refresh automatically when the source table changes.
|
||||
pub auto_refresh: bool,
|
||||
/// Register the definition only; skip the initial population.
|
||||
pub with_no_data: bool,
|
||||
/// Optional source column to partition the view's table function on. If the
|
||||
/// column has an IVF vector index the server partitions by its clusters
|
||||
/// (image-dedup style); otherwise it groups by distinct value.
|
||||
pub partition_by: Option<String>,
|
||||
}
|
||||
|
||||
impl CreateMaterializedViewRequest {
|
||||
pub fn new(name: impl Into<String>, query: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
query: query.into(),
|
||||
auto_refresh: false,
|
||||
with_no_data: false,
|
||||
partition_by: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A request to refresh a materialized view.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RefreshMaterializedViewRequest {
|
||||
/// View name.
|
||||
pub name: String,
|
||||
/// Force a full rebuild (recompute and replace every row) instead of the
|
||||
/// default incremental refresh.
|
||||
pub full: bool,
|
||||
/// Pin the refresh to a source-table version; latest when absent.
|
||||
pub src_version: Option<u64>,
|
||||
/// Initial worker count.
|
||||
pub num_workers: Option<u32>,
|
||||
/// Elastic worker ceiling.
|
||||
pub max_workers: Option<u32>,
|
||||
}
|
||||
|
||||
/// A request for the derived-compute lineage of a table/view (or one of its
|
||||
/// columns). The response is server-defined lineage JSON, returned opaque so
|
||||
/// this client need not model the server's lineage schema.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TableLineageRequest {
|
||||
/// Table or view name.
|
||||
pub name: String,
|
||||
/// Column for column-level lineage; whole table/view when absent.
|
||||
pub column: Option<String>,
|
||||
/// "upstream" | "downstream" | "both" (server default when absent).
|
||||
pub direction: Option<String>,
|
||||
/// Column-hops to walk; transitive when absent.
|
||||
pub depth: Option<u32>,
|
||||
}
|
||||
|
||||
impl RefreshMaterializedViewRequest {
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
full: false,
|
||||
src_version: None,
|
||||
num_workers: None,
|
||||
max_workers: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A registered materialized view definition, as returned by
|
||||
/// `list_materialized_views`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MaterializedViewInfo {
|
||||
pub name: String,
|
||||
pub source_table: String,
|
||||
/// Source columns projected through.
|
||||
pub projection: Vec<String>,
|
||||
/// `alias=expression` per UDF-computed column.
|
||||
pub udf_columns: Vec<String>,
|
||||
pub filter: Option<String>,
|
||||
pub auto_refresh: bool,
|
||||
}
|
||||
|
||||
/// A described platform job (`POST /v1/jobs/describe`): the job registry's
|
||||
/// lifecycle state plus the owner-written status payload.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PlatformJobDescription {
|
||||
/// The platform (registry) job id -- what describe/cancel accept.
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub job_subtype: String,
|
||||
/// "IN_PROGRESS" | "CANCELLED" | "FAILED" | "DONE".
|
||||
pub job_state: String,
|
||||
pub creation_ms: i64,
|
||||
/// The owner-written status payload -- `units_done` / `units_total` /
|
||||
/// `rows_committed` / `error` when present. Records whose owner has not
|
||||
/// written a payload yet carry the raw status-store URI string instead.
|
||||
pub status: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A row from `list_jobs`: one inflight server-side job (index build,
|
||||
/// compaction, column refresh, view refresh, ...).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
/// Lifecycle state: "running", "cancelling", or "stale".
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub age_seconds: Option<i64>,
|
||||
pub command: Option<String>,
|
||||
pub units_done: Option<i64>,
|
||||
pub units_total: Option<i64>,
|
||||
/// Whether the job's final commit has completed (output visible).
|
||||
pub committed: bool,
|
||||
pub rows_skipped: u64,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// A row from `job_history`: one durable, completed/terminal server-side job
|
||||
/// record (SHOW JOB HISTORY), read from a table's `_job_history` store. Unlike
|
||||
/// `JobInfo` (live, inflight jobs) this carries created/updated/completed
|
||||
/// timestamps and the lifecycle event log.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobHistoryInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub created_ms: i64,
|
||||
pub updated_ms: i64,
|
||||
pub completed_ms: Option<i64>,
|
||||
pub rows_processed: Option<i64>,
|
||||
pub rows_skipped: Option<i64>,
|
||||
pub error: Option<String>,
|
||||
/// Newline-joined lifecycle event log, oldest first.
|
||||
pub events: Option<String>,
|
||||
}
|
||||
|
||||
/// A row from `errors`: one per-row UDF failure recorded by `error_policy=skip`
|
||||
/// (SHOW ERRORS).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobErrorInfo {
|
||||
pub job_id: String,
|
||||
pub table: String,
|
||||
pub column: String,
|
||||
pub error_type: String,
|
||||
pub error_message: String,
|
||||
pub fragment_id: Option<i64>,
|
||||
pub source_row_id: Option<i64>,
|
||||
pub table_version: Option<i64>,
|
||||
pub age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
/// The plan a `REFRESH MATERIALIZED VIEW` would execute, as returned by
|
||||
/// `explain_refresh_materialized_view` (EXPLAIN REFRESH). No work is run.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MvRefreshPlan {
|
||||
pub table_name: String,
|
||||
/// Whether a refresh would do anything (rebuild or non-empty units).
|
||||
pub has_work: bool,
|
||||
pub source_version: u64,
|
||||
pub last_refreshed_version: Option<u64>,
|
||||
pub full_refresh: bool,
|
||||
/// Source changed non-append-only since the last refresh -> rebuild.
|
||||
pub rebuild: bool,
|
||||
/// Number of row-range work units the refresh would process.
|
||||
pub units_total: u64,
|
||||
}
|
||||
|
||||
fn not_supported<T>(what: &str) -> Result<T> {
|
||||
Err(Error::NotSupported {
|
||||
message: format!("{} is not supported by this database", what),
|
||||
})
|
||||
}
|
||||
|
||||
/// The `Database` trait defines the interface for database implementations.
|
||||
///
|
||||
/// A database is responsible for managing tables and their metadata.
|
||||
@@ -461,126 +245,6 @@ pub trait Database:
|
||||
///
|
||||
/// See [`CloneTableRequest`] for detailed documentation and examples.
|
||||
async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
|
||||
// -- Derived compute: functions, materialized views, jobs -------------
|
||||
//
|
||||
// Server-backed features (LanceDB Enterprise / Cloud). The defaults
|
||||
// return NotSupported; the remote database overrides them. Local
|
||||
// single-node implementations are planned.
|
||||
|
||||
/// Register a UDF (CREATE FUNCTION).
|
||||
async fn create_function(&self, _request: CreateFunctionRequest) -> Result<()> {
|
||||
not_supported("create_function")
|
||||
}
|
||||
/// List registered functions (SHOW FUNCTIONS).
|
||||
async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
not_supported("list_functions")
|
||||
}
|
||||
/// Drop a registered function (DROP FUNCTION).
|
||||
async fn drop_function(&self, _name: &str) -> Result<()> {
|
||||
not_supported("drop_function")
|
||||
}
|
||||
/// Create a materialized view (CREATE MATERIALIZED VIEW). Returns
|
||||
/// the initial-population job id, absent when `with_no_data`.
|
||||
async fn create_materialized_view(
|
||||
&self,
|
||||
_request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
not_supported("create_materialized_view")
|
||||
}
|
||||
/// Refresh a materialized view; returns the refresh job id.
|
||||
async fn refresh_materialized_view(
|
||||
&self,
|
||||
_request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
not_supported("refresh_materialized_view")
|
||||
}
|
||||
/// Derived-compute lineage of a table/view (or column), as server-defined
|
||||
/// JSON. Read-only.
|
||||
async fn table_lineage(&self, _request: TableLineageRequest) -> Result<String> {
|
||||
not_supported("table_lineage")
|
||||
}
|
||||
/// Plan a materialized-view refresh without submitting work
|
||||
/// (EXPLAIN REFRESH). `full` plans a full rebuild (incremental
|
||||
/// planning requires stable row IDs on the source).
|
||||
async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
_name: &str,
|
||||
_full: bool,
|
||||
_src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
not_supported("explain_refresh_materialized_view")
|
||||
}
|
||||
/// Update a materialized view's options (ALTER MATERIALIZED VIEW).
|
||||
async fn alter_materialized_view(&self, _name: &str, _auto_refresh: bool) -> Result<()> {
|
||||
not_supported("alter_materialized_view")
|
||||
}
|
||||
/// Drop a materialized view definition (DROP MATERIALIZED VIEW).
|
||||
async fn drop_materialized_view(&self, _name: &str) -> Result<()> {
|
||||
not_supported("drop_materialized_view")
|
||||
}
|
||||
/// List registered materialized view definitions.
|
||||
async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
not_supported("list_materialized_views")
|
||||
}
|
||||
/// List inflight server-side jobs across the database's tables.
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
not_supported("list_jobs")
|
||||
}
|
||||
|
||||
/// Describe a platform job (`POST /v1/jobs/describe`): registry-backed
|
||||
/// lifecycle state plus the owner-written status payload. `None` when the
|
||||
/// registry has no such job.
|
||||
async fn describe_platform_job(
|
||||
&self,
|
||||
_platform_job_id: &str,
|
||||
) -> Result<Option<PlatformJobDescription>> {
|
||||
not_supported("describe_platform_job")
|
||||
}
|
||||
|
||||
/// Resolve a submission (manifest) job id to its platform job id via the
|
||||
/// registry's manifest-id filter (`POST /v1/jobs/list`). `None` until the
|
||||
/// job has registered (dispatch is async).
|
||||
async fn resolve_platform_job_id(
|
||||
&self,
|
||||
_manifest_job_id: &str,
|
||||
_table_hint: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
not_supported("resolve_platform_job_id")
|
||||
}
|
||||
|
||||
/// Cancel a platform job (`POST /v1/jobs/cancel`). Idempotent: cancelling
|
||||
/// an already-terminal job is a no-op success.
|
||||
async fn cancel_platform_job(&self, _platform_job_id: &str) -> Result<()> {
|
||||
not_supported("cancel_platform_job")
|
||||
}
|
||||
/// Cancel an inflight server-side job by id. Returns true if a
|
||||
/// matching inflight job was found and flagged for cancellation,
|
||||
/// false if none was inflight (best-effort, like SQL `CANCEL JOB`).
|
||||
async fn cancel_job(&self, _job_id: &str) -> Result<bool> {
|
||||
not_supported("cancel_job")
|
||||
}
|
||||
/// Point-access for a single job by id -- the `wait()`/status poll path.
|
||||
/// `table_hint` (the job's table, which `wait()` callers know) enables an
|
||||
/// O(1) server-side lookup. `None` if the job is unknown or not active.
|
||||
async fn get_job(&self, _job_id: &str, _table_hint: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
not_supported("get_job")
|
||||
}
|
||||
/// Durable job history (SHOW JOB HISTORY) across the database's tables,
|
||||
/// optionally narrowed to a single `job_id`.
|
||||
async fn job_history(&self, _job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
not_supported("job_history")
|
||||
}
|
||||
/// Per-row UDF errors (SHOW ERRORS) recorded by `error_policy=skip` across
|
||||
/// the database's tables, optionally filtered by `job_id` and/or `table`.
|
||||
async fn errors(
|
||||
&self,
|
||||
_job_id: Option<&str>,
|
||||
_table: Option<&str>,
|
||||
) -> Result<Vec<JobErrorInfo>> {
|
||||
not_supported("errors")
|
||||
}
|
||||
|
||||
/// Open a table in the database
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
/// Rename a table in the database
|
||||
|
||||
@@ -761,8 +761,8 @@ mod tests {
|
||||
verify_splitter(splitter, test_data(), 50, &[11, 8, 9], false).await;
|
||||
}
|
||||
|
||||
async fn collect_hash_split() -> RecordBatch {
|
||||
let total_rows = 50;
|
||||
#[tokio::test]
|
||||
async fn test_hash_split() {
|
||||
let data = lance_datagen::gen_batch()
|
||||
.with_seed(Seed::from(42))
|
||||
.col(
|
||||
@@ -783,7 +783,7 @@ mod tests {
|
||||
);
|
||||
|
||||
let split_batches = splitter
|
||||
.apply(data, total_rows)
|
||||
.apply(data, 10)
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
@@ -791,35 +791,20 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
let schema = split_batches[0].schema();
|
||||
concat_batches(&schema, &split_batches).unwrap()
|
||||
}
|
||||
let split_batch = concat_batches(&schema, &split_batches).unwrap();
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hash_split() {
|
||||
let total_rows = 50;
|
||||
let split_batch = collect_hash_split().await;
|
||||
let split_batch_again = collect_hash_split().await;
|
||||
// These assertions are all based on fixed seed in data generation but they match
|
||||
// up roughly to what we expect (25% discarded, 25% in split 0, 50% in split 1)
|
||||
|
||||
assert_eq!(split_batch.num_rows(), split_batch_again.num_rows());
|
||||
assert_eq!(split_batch.num_columns(), split_batch_again.num_columns());
|
||||
for (left, right) in split_batch
|
||||
.columns()
|
||||
.iter()
|
||||
.zip(split_batch_again.columns().iter())
|
||||
{
|
||||
assert_eq!(left, right);
|
||||
}
|
||||
|
||||
assert!(split_batch.num_rows() > 0);
|
||||
assert!(split_batch.num_rows() < total_rows);
|
||||
// 8 rows (16%) are discarded because discard_weight is 1
|
||||
assert_eq!(split_batch.num_rows(), 42);
|
||||
assert_eq!(split_batch.num_columns(), 2);
|
||||
|
||||
let split_ids = split_batch.column(1).as_primitive::<UInt64Type>().values();
|
||||
let num_in_split_0 = split_ids.iter().filter(|v| **v == 0).count();
|
||||
let num_in_split_1 = split_ids.iter().filter(|v| **v == 1).count();
|
||||
|
||||
assert_eq!(num_in_split_0 + num_in_split_1, split_batch.num_rows());
|
||||
assert!(num_in_split_0 > 0);
|
||||
assert!(num_in_split_1 > num_in_split_0);
|
||||
assert_eq!(num_in_split_0, 12); // 24%
|
||||
assert_eq!(num_in_split_1, 30); // 60%
|
||||
}
|
||||
}
|
||||
|
||||
+13
-13
@@ -61,29 +61,29 @@ pub fn is_in(expr: Expr, list: Vec<Expr>) -> Expr {
|
||||
expr.in_list(list, false)
|
||||
}
|
||||
|
||||
lazy_static::lazy_static! {
|
||||
static ref FUNC_REGISTRY: std::sync::RwLock<std::collections::HashMap<String, Arc<ScalarUDF>>> = {
|
||||
static FUNC_REGISTRY: std::sync::LazyLock<std::collections::HashMap<String, Arc<ScalarUDF>>> =
|
||||
std::sync::LazyLock::new(|| {
|
||||
let mut m = std::collections::HashMap::new();
|
||||
m.insert("lower".to_string(), datafusion_functions::string::lower());
|
||||
m.insert("upper".to_string(), datafusion_functions::string::upper());
|
||||
m.insert("contains".to_string(), datafusion_functions::string::contains());
|
||||
m.insert(
|
||||
"contains".to_string(),
|
||||
datafusion_functions::string::contains(),
|
||||
);
|
||||
m.insert("btrim".to_string(), datafusion_functions::string::btrim());
|
||||
m.insert("ltrim".to_string(), datafusion_functions::string::ltrim());
|
||||
m.insert("rtrim".to_string(), datafusion_functions::string::rtrim());
|
||||
m.insert("concat".to_string(), datafusion_functions::string::concat());
|
||||
m.insert("octet_length".to_string(), datafusion_functions::string::octet_length());
|
||||
std::sync::RwLock::new(m)
|
||||
};
|
||||
}
|
||||
m.insert(
|
||||
"octet_length".to_string(),
|
||||
datafusion_functions::string::octet_length(),
|
||||
);
|
||||
m
|
||||
});
|
||||
|
||||
pub fn func(name: impl AsRef<str>, args: Vec<Expr>) -> crate::Result<Expr> {
|
||||
let name = name.as_ref();
|
||||
let registry = FUNC_REGISTRY
|
||||
.read()
|
||||
.map_err(|e| crate::Error::InvalidInput {
|
||||
message: format!("lock poisoned: {}", e),
|
||||
})?;
|
||||
let udf = registry
|
||||
let udf = FUNC_REGISTRY
|
||||
.get(name)
|
||||
.ok_or_else(|| crate::Error::InvalidInput {
|
||||
message: format!("unknown function: {}", name),
|
||||
|
||||
@@ -283,10 +283,7 @@ impl IndexBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
/// Returns the server-minted job id when the index build was deferred to
|
||||
/// a background job (remote tables only); `None` when the build completed
|
||||
/// synchronously within this call.
|
||||
pub async fn execute(self) -> Result<Option<String>> {
|
||||
pub async fn execute(self) -> Result<()> {
|
||||
self.parent.clone().create_index(self).await
|
||||
}
|
||||
}
|
||||
|
||||
@@ -47,6 +47,20 @@ pub trait HeaderProvider: Send + Sync + std::fmt::Debug {
|
||||
async fn get_headers(&self) -> Result<HashMap<String, String>>;
|
||||
}
|
||||
|
||||
/// Default maximum bytes per insert request (8 GiB).
|
||||
///
|
||||
/// Sized so a multipart part can hold at least one full Lance data file (the
|
||||
/// default is 1M rows / 90 GB per file), which keeps fragments from being split
|
||||
/// into undersized files across parts. The time-based cut
|
||||
/// ([`DEFAULT_MAX_REQUEST_DURATION_DIVISOR`]) bounds request duration on slow
|
||||
/// uploads, so a large byte budget does not risk the read timeout.
|
||||
const DEFAULT_MAX_BYTES_PER_REQUEST: u64 = 8 * 1024 * 1024 * 1024;
|
||||
|
||||
/// The default max request duration is the read timeout divided by this, leaving
|
||||
/// headroom for the server to finalize and acknowledge a part before the read
|
||||
/// timeout (which also covers the request-body upload) fires.
|
||||
const DEFAULT_MAX_REQUEST_DURATION_DIVISOR: u32 = 2;
|
||||
|
||||
/// Configuration for the LanceDB Cloud HTTP client.
|
||||
#[derive(Clone)]
|
||||
pub struct ClientConfig {
|
||||
@@ -71,6 +85,33 @@ pub struct ClientConfig {
|
||||
/// Alternatively, set `LANCEDB_USER_ID_ENV_KEY` to specify another environment
|
||||
/// variable that contains the user ID value.
|
||||
pub user_id: Option<String>,
|
||||
/// Maximum number of bytes to send in a single insert HTTP request.
|
||||
///
|
||||
/// During a multipart write, each partition's data is split into one or more
|
||||
/// parts of at most this many (Arrow IPC, compressed) bytes, each uploaded as
|
||||
/// a separate request under the shared upload id. This bounds how long any
|
||||
/// one request stays open, so large bulk ingests do not exceed the client
|
||||
/// read timeout while the server streams the part to object storage.
|
||||
///
|
||||
/// The request body is still streamed (not buffered), so this does not
|
||||
/// increase peak memory. Set to `Some(0)` to disable splitting (one request
|
||||
/// per partition). You can also set the `LANCE_CLIENT_MAX_BYTES_PER_REQUEST`
|
||||
/// environment variable. Defaults to 8 GiB.
|
||||
pub max_bytes_per_request: Option<u64>,
|
||||
/// Maximum wall-clock time to spend uploading a single insert HTTP request.
|
||||
///
|
||||
/// Complements [`Self::max_bytes_per_request`]: during a multipart write a
|
||||
/// part is cut when it reaches either the byte budget or this duration,
|
||||
/// whichever comes first. The client read timeout also covers the
|
||||
/// request-body upload, so a slow or throttled upload of a large part can
|
||||
/// hit that timeout before the byte budget is reached; cutting by time keeps
|
||||
/// each request short enough that it completes (and the server acknowledges
|
||||
/// the part) within the read timeout.
|
||||
///
|
||||
/// Set to `Some(Duration::ZERO)` to disable the time-based cut. You can also
|
||||
/// set the `LANCE_CLIENT_MAX_REQUEST_DURATION` environment variable (integer
|
||||
/// seconds). Defaults to half the resolved read timeout.
|
||||
pub max_request_duration: Option<Duration>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ClientConfig {
|
||||
@@ -87,6 +128,8 @@ impl std::fmt::Debug for ClientConfig {
|
||||
&self.header_provider.as_ref().map(|_| "Some(...)"),
|
||||
)
|
||||
.field("user_id", &self.user_id)
|
||||
.field("max_bytes_per_request", &self.max_bytes_per_request)
|
||||
.field("max_request_duration", &self.max_request_duration)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -102,6 +145,8 @@ impl Default for ClientConfig {
|
||||
tls_config: None,
|
||||
header_provider: None,
|
||||
user_id: None,
|
||||
max_bytes_per_request: None,
|
||||
max_request_duration: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -248,6 +293,16 @@ pub struct RestfulLanceDbClient<S: HttpSend = Sender> {
|
||||
/// Connection-level read consistency interval. Drives the
|
||||
/// `x-lancedb-min-timestamp` freshness header sent on read requests.
|
||||
pub(crate) read_consistency_interval: Option<Duration>,
|
||||
// Note the `Option` here means the opposite of the same-named
|
||||
// `ClientConfig` fields: those are pre-resolution, where `None` means "fall
|
||||
// back to env var / default". These are post-resolution (see
|
||||
// `resolve_max_bytes_per_request` / `resolve_max_request_duration`), where a
|
||||
// default has already been applied and `None` means the feature is disabled.
|
||||
/// Maximum bytes per insert request. `None` disables request splitting.
|
||||
pub(crate) max_bytes_per_request: Option<u64>,
|
||||
/// Maximum wall-clock time per insert request. `None` disables the
|
||||
/// time-based part cut.
|
||||
pub(crate) max_request_duration: Option<Duration>,
|
||||
}
|
||||
|
||||
impl<S: HttpSend> std::fmt::Debug for RestfulLanceDbClient<S> {
|
||||
@@ -429,6 +484,10 @@ impl RestfulLanceDbClient<Sender> {
|
||||
};
|
||||
debug!("Created client for host: {}", host);
|
||||
let retry_config = client_config.retry_config.clone().try_into()?;
|
||||
let max_bytes_per_request =
|
||||
Self::resolve_max_bytes_per_request(client_config.max_bytes_per_request)?;
|
||||
let max_request_duration =
|
||||
Self::resolve_max_request_duration(client_config.max_request_duration, read_timeout)?;
|
||||
Ok(Self {
|
||||
client,
|
||||
host,
|
||||
@@ -440,8 +499,52 @@ impl RestfulLanceDbClient<Sender> {
|
||||
.unwrap_or("$".to_string()),
|
||||
header_provider: client_config.header_provider,
|
||||
read_consistency_interval,
|
||||
max_bytes_per_request,
|
||||
max_request_duration,
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve the max bytes per insert request from config, environment, or the
|
||||
/// default. A value of `0` (from either source) disables request splitting.
|
||||
fn resolve_max_bytes_per_request(passed: Option<u64>) -> Result<Option<u64>> {
|
||||
let value = if let Some(value) = passed {
|
||||
value
|
||||
} else if let Ok(env) = std::env::var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST") {
|
||||
env.parse::<u64>().map_err(|_| Error::InvalidInput {
|
||||
message: format!(
|
||||
"LANCE_CLIENT_MAX_BYTES_PER_REQUEST must be a non-negative integer, got '{}'",
|
||||
env
|
||||
),
|
||||
})?
|
||||
} else {
|
||||
DEFAULT_MAX_BYTES_PER_REQUEST
|
||||
};
|
||||
Ok((value > 0).then_some(value))
|
||||
}
|
||||
|
||||
/// Resolve the max request duration from config, environment, or a default
|
||||
/// derived from the read timeout. A zero duration (from either source)
|
||||
/// disables the time-based cut.
|
||||
fn resolve_max_request_duration(
|
||||
passed: Option<Duration>,
|
||||
read_timeout: Duration,
|
||||
) -> Result<Option<Duration>> {
|
||||
let value = if let Some(value) = passed {
|
||||
value
|
||||
} else if let Ok(env) = std::env::var("LANCE_CLIENT_MAX_REQUEST_DURATION") {
|
||||
let secs = env.parse::<u64>().map_err(|_| Error::InvalidInput {
|
||||
message: format!(
|
||||
"LANCE_CLIENT_MAX_REQUEST_DURATION must be a non-negative integer \
|
||||
number of seconds, got '{}'",
|
||||
env
|
||||
),
|
||||
})?;
|
||||
Duration::from_secs(secs)
|
||||
} else {
|
||||
read_timeout / DEFAULT_MAX_REQUEST_DURATION_DIVISOR
|
||||
};
|
||||
Ok((!value.is_zero()).then_some(value))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: HttpSend> RestfulLanceDbClient<S> {
|
||||
@@ -449,6 +552,18 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
|
||||
&self.host
|
||||
}
|
||||
|
||||
/// Maximum bytes per insert request, or `None` if request splitting is
|
||||
/// disabled.
|
||||
pub(crate) fn max_bytes_per_request(&self) -> Option<u64> {
|
||||
self.max_bytes_per_request
|
||||
}
|
||||
|
||||
/// Maximum wall-clock time per insert request, or `None` if the time-based
|
||||
/// cut is disabled.
|
||||
pub(crate) fn max_request_duration(&self) -> Option<Duration> {
|
||||
self.max_request_duration
|
||||
}
|
||||
|
||||
pub fn default_headers(
|
||||
api_key: &str,
|
||||
region: &str,
|
||||
@@ -875,6 +990,8 @@ pub mod test_utils {
|
||||
id_delimiter: "$".to_string(),
|
||||
header_provider: None,
|
||||
read_consistency_interval,
|
||||
max_bytes_per_request: None,
|
||||
max_request_duration: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -900,6 +1017,12 @@ pub mod test_utils {
|
||||
id_delimiter: config.id_delimiter.unwrap_or_else(|| "$".to_string()),
|
||||
header_provider: config.header_provider,
|
||||
read_consistency_interval: None,
|
||||
max_bytes_per_request: config
|
||||
.max_bytes_per_request
|
||||
.and_then(|v| (v > 0).then_some(v)),
|
||||
max_request_duration: config
|
||||
.max_request_duration
|
||||
.and_then(|v| (!v.is_zero()).then_some(v)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1103,6 +1226,8 @@ mod tests {
|
||||
id_delimiter: "+".to_string(),
|
||||
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
|
||||
read_consistency_interval: None,
|
||||
max_bytes_per_request: None,
|
||||
max_request_duration: None,
|
||||
};
|
||||
|
||||
// Apply dynamic headers
|
||||
@@ -1139,6 +1264,8 @@ mod tests {
|
||||
id_delimiter: "+".to_string(),
|
||||
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
|
||||
read_consistency_interval: None,
|
||||
max_bytes_per_request: None,
|
||||
max_request_duration: None,
|
||||
};
|
||||
|
||||
// Apply dynamic headers
|
||||
@@ -1177,6 +1304,8 @@ mod tests {
|
||||
id_delimiter: "+".to_string(),
|
||||
header_provider: Some(Arc::new(provider) as Arc<dyn HeaderProvider>),
|
||||
read_consistency_interval: None,
|
||||
max_bytes_per_request: None,
|
||||
max_request_duration: None,
|
||||
};
|
||||
|
||||
// Header provider errors should fail the request
|
||||
@@ -1288,4 +1417,193 @@ mod tests {
|
||||
std::env::remove_var("LANCEDB_USER_ID");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_max_bytes_passed_value_wins() {
|
||||
// An explicit config value is used verbatim; env/default are not consulted.
|
||||
let resolved =
|
||||
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(1234)).unwrap();
|
||||
assert_eq!(resolved, Some(1234));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_max_bytes_zero_disables() {
|
||||
let resolved =
|
||||
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(0)).unwrap();
|
||||
assert_eq!(resolved, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_bytes_default_when_unset() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
|
||||
}
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
|
||||
assert_eq!(resolved, Some(DEFAULT_MAX_BYTES_PER_REQUEST));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_bytes_from_env() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096");
|
||||
}
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
|
||||
}
|
||||
assert_eq!(resolved, Some(4096));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_bytes_env_zero_disables() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "0");
|
||||
}
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
|
||||
}
|
||||
assert_eq!(resolved, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_bytes_config_overrides_env() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096");
|
||||
}
|
||||
// A config value takes precedence over the environment variable.
|
||||
let resolved =
|
||||
RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(Some(1234)).unwrap();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
|
||||
}
|
||||
assert_eq!(resolved, Some(1234));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_bytes_invalid_env_errors() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "not-a-number");
|
||||
}
|
||||
let err = RestfulLanceDbClient::<Sender>::resolve_max_bytes_per_request(None).unwrap_err();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST");
|
||||
}
|
||||
assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_max_request_duration_passed_value_wins() {
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
|
||||
Some(Duration::from_secs(42)),
|
||||
Duration::from_secs(300),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(resolved, Some(Duration::from_secs(42)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_max_request_duration_zero_disables() {
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
|
||||
Some(Duration::ZERO),
|
||||
Duration::from_secs(300),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(resolved, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_request_duration_default_is_half_read_timeout() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
|
||||
}
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
|
||||
None,
|
||||
Duration::from_secs(300),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(resolved, Some(Duration::from_secs(150)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_request_duration_from_env_seconds() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "30");
|
||||
}
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
|
||||
None,
|
||||
Duration::from_secs(300),
|
||||
)
|
||||
.unwrap();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
|
||||
}
|
||||
assert_eq!(resolved, Some(Duration::from_secs(30)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_request_duration_env_zero_disables() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "0");
|
||||
}
|
||||
let resolved = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
|
||||
None,
|
||||
Duration::from_secs(300),
|
||||
)
|
||||
.unwrap();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
|
||||
}
|
||||
assert_eq!(resolved, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(request_limits_env)]
|
||||
fn test_resolve_max_request_duration_invalid_env_errors() {
|
||||
let _guard = lock_env();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "12.5");
|
||||
}
|
||||
let err = RestfulLanceDbClient::<Sender>::resolve_max_request_duration(
|
||||
None,
|
||||
Duration::from_secs(300),
|
||||
)
|
||||
.unwrap_err();
|
||||
// SAFETY: This is only called in tests
|
||||
unsafe {
|
||||
std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION");
|
||||
}
|
||||
assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,10 +19,8 @@ use lance_namespace::models::{
|
||||
|
||||
use crate::Error;
|
||||
use crate::database::{
|
||||
CloneTableRequest, CreateFunctionRequest, CreateMaterializedViewRequest, CreateTableMode,
|
||||
CreateTableRequest, Database, DatabaseOptions, FunctionInfo, JobErrorInfo, JobHistoryInfo,
|
||||
JobInfo, MaterializedViewInfo, MvRefreshPlan, OpenTableRequest, PlatformJobDescription,
|
||||
ReadConsistency, RefreshMaterializedViewRequest, TableLineageRequest, TableNamesRequest,
|
||||
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
|
||||
OpenTableRequest, ReadConsistency, TableNamesRequest,
|
||||
};
|
||||
use crate::error::Result;
|
||||
use crate::remote::util::stream_as_body;
|
||||
@@ -35,210 +33,6 @@ use super::client::{
|
||||
use super::table::RemoteTable;
|
||||
use super::util::parse_server_version;
|
||||
|
||||
// Wire types for the derived-compute routes (functions, materialized
|
||||
// views, jobs). Field shapes mirror the server's REST contract.
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCreateFunctionRequest {
|
||||
language: String,
|
||||
return_type: String,
|
||||
body: String,
|
||||
options: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteFunctionEntry {
|
||||
name: String,
|
||||
language: String,
|
||||
return_type: String,
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListFunctionsResponse {
|
||||
functions: Vec<RemoteFunctionEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCreateMaterializedViewRequest {
|
||||
query: String,
|
||||
auto_refresh: bool,
|
||||
with_no_data: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
partition_by: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteCreateMaterializedViewResponse {
|
||||
#[serde(default)]
|
||||
job_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteRefreshMaterializedViewRequest {
|
||||
#[serde(skip_serializing_if = "std::ops::Not::not")]
|
||||
full: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
src_version: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
num_workers: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
max_workers: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteRefreshMaterializedViewResponse {
|
||||
job_id: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteExplainRefreshRequest {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
full: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
src_version: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteExplainRefreshResponse {
|
||||
table_name: String,
|
||||
has_work: bool,
|
||||
source_version: u64,
|
||||
last_refreshed_version: Option<u64>,
|
||||
full_refresh: bool,
|
||||
rebuild: bool,
|
||||
units_total: u64,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteAlterMaterializedViewRequest {
|
||||
auto_refresh: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteMaterializedViewEntry {
|
||||
name: String,
|
||||
source_table: String,
|
||||
#[serde(default)]
|
||||
projection: Vec<String>,
|
||||
#[serde(default)]
|
||||
udf_columns: Vec<String>,
|
||||
#[serde(default)]
|
||||
filter: Option<String>,
|
||||
#[serde(default)]
|
||||
auto_refresh: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListMaterializedViewsResponse {
|
||||
views: Vec<RemoteMaterializedViewEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteDescribePlatformJobResponse {
|
||||
job_id: String,
|
||||
job_type: String,
|
||||
#[serde(default)]
|
||||
job_subtype: String,
|
||||
job_state: String,
|
||||
#[serde(default)]
|
||||
creation_ms: i64,
|
||||
#[serde(default)]
|
||||
status: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListPlatformJobsResponse {
|
||||
#[serde(default)]
|
||||
jobs: Vec<RemotePlatformJobRow>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemotePlatformJobRow {
|
||||
job_id: String,
|
||||
#[serde(default)]
|
||||
table: String,
|
||||
#[serde(default)]
|
||||
job_subtype: String,
|
||||
#[serde(default)]
|
||||
state: String,
|
||||
#[serde(default)]
|
||||
created_at_millis: i64,
|
||||
#[serde(default)]
|
||||
status: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Platform list-row state -> the client's job vocabulary.
|
||||
fn platform_state_to_client(state: &str) -> String {
|
||||
match state {
|
||||
"in_progress" => "running",
|
||||
"done" => "finished",
|
||||
other => other,
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Describe job_state -> the client's job vocabulary.
|
||||
fn describe_state_to_client(state: &str) -> String {
|
||||
match state {
|
||||
"IN_PROGRESS" => "running",
|
||||
"DONE" => "finished",
|
||||
"FAILED" => "failed",
|
||||
"CANCELLED" => "cancelled",
|
||||
other => other,
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn payload_i64(status: &serde_json::Value, key: &str) -> Option<i64> {
|
||||
status.get(key).and_then(serde_json::Value::as_i64)
|
||||
}
|
||||
|
||||
fn payload_error(status: &serde_json::Value) -> Option<String> {
|
||||
status
|
||||
.get("error")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_string)
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteErrorEntry {
|
||||
job_id: String,
|
||||
table: String,
|
||||
column: String,
|
||||
error_type: String,
|
||||
error_message: String,
|
||||
#[serde(default)]
|
||||
fragment_id: Option<i64>,
|
||||
#[serde(default)]
|
||||
source_row_id: Option<i64>,
|
||||
#[serde(default)]
|
||||
table_version: Option<i64>,
|
||||
#[serde(default)]
|
||||
age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteErrorsResponse {
|
||||
errors: Vec<RemoteErrorEntry>,
|
||||
}
|
||||
|
||||
impl From<RemoteErrorEntry> for JobErrorInfo {
|
||||
fn from(e: RemoteErrorEntry) -> Self {
|
||||
JobErrorInfo {
|
||||
job_id: e.job_id,
|
||||
table: e.table,
|
||||
column: e.column,
|
||||
error_type: e.error_type,
|
||||
error_message: e.error_message,
|
||||
fragment_id: e.fragment_id,
|
||||
source_row_id: e.source_row_id,
|
||||
table_version: e.table_version,
|
||||
age_seconds: e.age_seconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Request structure for the remote clone table API
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCloneTableRequest {
|
||||
@@ -847,426 +641,6 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
Ok(table)
|
||||
}
|
||||
|
||||
async fn create_function(&self, request: CreateFunctionRequest) -> Result<()> {
|
||||
let body = RemoteCreateFunctionRequest {
|
||||
language: request.language,
|
||||
return_type: request.return_type,
|
||||
body: request.body,
|
||||
options: request.options,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/function/{}/create", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
let req = self.client.get("/v1/function/list");
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListFunctionsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body
|
||||
.functions
|
||||
.into_iter()
|
||||
.map(|f| FunctionInfo {
|
||||
name: f.name,
|
||||
language: f.language,
|
||||
return_type: f.return_type,
|
||||
description: f.description,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn drop_function(&self, name: &str) -> Result<()> {
|
||||
let req = self.client.post(&format!("/v1/function/{}/drop", name));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_materialized_view(
|
||||
&self,
|
||||
request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
let body = RemoteCreateMaterializedViewRequest {
|
||||
query: request.query,
|
||||
auto_refresh: request.auto_refresh,
|
||||
with_no_data: request.with_no_data,
|
||||
partition_by: request.partition_by,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/create", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteCreateMaterializedViewResponse =
|
||||
rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn refresh_materialized_view(
|
||||
&self,
|
||||
request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
let body = RemoteRefreshMaterializedViewRequest {
|
||||
full: request.full,
|
||||
src_version: request.src_version,
|
||||
num_workers: request.num_workers,
|
||||
max_workers: request.max_workers,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/refresh", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteRefreshMaterializedViewResponse =
|
||||
rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn table_lineage(&self, request: TableLineageRequest) -> Result<String> {
|
||||
let mut req = self
|
||||
.client
|
||||
.get(&format!("/v1/table/{}/lineage", request.name));
|
||||
if let Some(column) = &request.column {
|
||||
req = req.query(&[("column", column)]);
|
||||
}
|
||||
if let Some(direction) = &request.direction {
|
||||
req = req.query(&[("direction", direction)]);
|
||||
}
|
||||
if let Some(depth) = request.depth {
|
||||
req = req.query(&[("depth", depth.to_string())]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
// Server-defined lineage JSON, returned opaque (the client does not
|
||||
// model the lineage schema; the Python layer deserializes it).
|
||||
rsp.text().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
name: &str,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
let body = RemoteExplainRefreshRequest {
|
||||
full: Some(full),
|
||||
src_version,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/explain_refresh", name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteExplainRefreshResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(MvRefreshPlan {
|
||||
table_name: body.table_name,
|
||||
has_work: body.has_work,
|
||||
source_version: body.source_version,
|
||||
last_refreshed_version: body.last_refreshed_version,
|
||||
full_refresh: body.full_refresh,
|
||||
rebuild: body.rebuild,
|
||||
units_total: body.units_total,
|
||||
})
|
||||
}
|
||||
|
||||
async fn alter_materialized_view(&self, name: &str, auto_refresh: bool) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/alter", name))
|
||||
.json(&RemoteAlterMaterializedViewRequest { auto_refresh });
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_materialized_view(&self, name: &str) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/drop", name));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
let req = self.client.get("/v1/materialized_view/list");
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListMaterializedViewsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body
|
||||
.views
|
||||
.into_iter()
|
||||
.map(|v| MaterializedViewInfo {
|
||||
name: v.name,
|
||||
source_table: v.source_table,
|
||||
projection: v.projection,
|
||||
udf_columns: v.udf_columns,
|
||||
filter: v.filter,
|
||||
auto_refresh: v.auto_refresh,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/list")
|
||||
.json(&serde_json::json!({ "include_status": true }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListPlatformJobsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0);
|
||||
Ok(body
|
||||
.jobs
|
||||
.into_iter()
|
||||
.map(|row| JobInfo {
|
||||
table: row.table,
|
||||
job_id: row.job_id,
|
||||
// The platform job_type is always "indexer"; the subtype
|
||||
// (udf / mv_refresh / compaction / ...) is the useful label.
|
||||
job_type: row.job_subtype,
|
||||
state: platform_state_to_client(&row.state),
|
||||
column: None,
|
||||
age_seconds: (row.created_at_millis > 0)
|
||||
.then(|| (now_ms - row.created_at_millis) / 1000),
|
||||
command: None,
|
||||
units_done: payload_i64(&row.status, "units_done"),
|
||||
units_total: payload_i64(&row.status, "units_total"),
|
||||
committed: row.state == "done",
|
||||
rows_skipped: payload_i64(&row.status, "rows_skipped").unwrap_or(0) as u64,
|
||||
error: payload_error(&row.status),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn get_job(&self, job_id: &str, table: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
// A point snapshot from the platform API: resolve the submission id,
|
||||
// then describe. The snapshot keeps the caller's id.
|
||||
let Some(platform_id) = self.resolve_platform_job_id(job_id, table).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(described) = self.describe_platform_job(&platform_id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(JobInfo {
|
||||
table: table.unwrap_or_default().to_string(),
|
||||
job_id: job_id.to_string(),
|
||||
job_type: described.job_subtype,
|
||||
state: describe_state_to_client(&described.job_state),
|
||||
column: None,
|
||||
age_seconds: None,
|
||||
command: None,
|
||||
units_done: payload_i64(&described.status, "units_done"),
|
||||
units_total: payload_i64(&described.status, "units_total"),
|
||||
committed: described.job_state == "DONE",
|
||||
rows_skipped: payload_i64(&described.status, "rows_skipped").unwrap_or(0) as u64,
|
||||
error: payload_error(&described.status),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn describe_platform_job(
|
||||
&self,
|
||||
platform_job_id: &str,
|
||||
) -> Result<Option<PlatformJobDescription>> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/describe")
|
||||
.json(&serde_json::json!({ "job_id": platform_job_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
if rsp.status().as_u16() == 404 {
|
||||
return Ok(None);
|
||||
}
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteDescribePlatformJobResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(Some(PlatformJobDescription {
|
||||
job_id: body.job_id,
|
||||
job_type: body.job_type,
|
||||
job_subtype: body.job_subtype,
|
||||
job_state: body.job_state,
|
||||
creation_ms: body.creation_ms,
|
||||
status: body.status,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn resolve_platform_job_id(
|
||||
&self,
|
||||
manifest_job_id: &str,
|
||||
table_hint: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
let req = self.client.post("/v1/jobs/list").json(&serde_json::json!({
|
||||
"manifest_job_id": manifest_job_id,
|
||||
"table_name": table_hint,
|
||||
"job_type": "indexer",
|
||||
}));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListPlatformJobsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.jobs.into_iter().next().map(|row| row.job_id))
|
||||
}
|
||||
|
||||
async fn cancel_platform_job(&self, platform_job_id: &str) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/cancel")
|
||||
.json(&serde_json::json!({ "job_id": platform_job_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
||||
// Resolve the submission id and cancel through the platform API.
|
||||
// False when no matching job has registered (the legacy best-effort
|
||||
// contract).
|
||||
let Some(platform_id) = self.resolve_platform_job_id(job_id, None).await? else {
|
||||
return Ok(false);
|
||||
};
|
||||
self.cancel_platform_job(&platform_id).await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
// One job: describe (identity) plus query_events (timeline). No id:
|
||||
// a registry listing, timeline-free -- pass an id for the event log.
|
||||
let Some(caller_id) = job_id else {
|
||||
let rows = self.list_jobs().await?;
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0);
|
||||
return Ok(rows
|
||||
.into_iter()
|
||||
.map(|j| JobHistoryInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
created_ms: j.age_seconds.map(|a| now_ms - a * 1000).unwrap_or_default(),
|
||||
updated_ms: 0,
|
||||
completed_ms: None,
|
||||
rows_processed: None,
|
||||
rows_skipped: j.rows_skipped.try_into().ok(),
|
||||
error: j.error,
|
||||
events: None,
|
||||
})
|
||||
.collect());
|
||||
};
|
||||
// Accept either a platform id or a submission id.
|
||||
let platform_id = match self.describe_platform_job(caller_id).await? {
|
||||
Some(_) => caller_id.to_string(),
|
||||
None => match self.resolve_platform_job_id(caller_id, None).await? {
|
||||
Some(id) => id,
|
||||
None => return Ok(Vec::new()),
|
||||
},
|
||||
};
|
||||
let Some(described) = self.describe_platform_job(&platform_id).await? else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/query_events")
|
||||
.json(&serde_json::json!({ "job_id": platform_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body = rsp.bytes().await.err_to_http(request_id.clone())?;
|
||||
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None)
|
||||
.map_err(|e| Error::Http {
|
||||
source: format!("failed to read job-events IPC stream: {e}").into(),
|
||||
request_id: request_id.clone(),
|
||||
status_code: None,
|
||||
})?;
|
||||
|
||||
let mut created_ms = i64::MAX;
|
||||
let mut updated_ms = 0i64;
|
||||
let mut completed_ms = None;
|
||||
let mut last_error = None;
|
||||
let mut events = Vec::new();
|
||||
for batch in reader {
|
||||
let batch = batch.map_err(|e| Error::Http {
|
||||
source: format!("failed to decode job-events batch: {e}").into(),
|
||||
request_id: request_id.clone(),
|
||||
status_code: None,
|
||||
})?;
|
||||
let states = batch
|
||||
.column_by_name("state")
|
||||
.and_then(|c| c.as_any().downcast_ref::<arrow_array::StringArray>());
|
||||
let times = batch
|
||||
.column_by_name("updated_at_millis")
|
||||
.and_then(|c| c.as_any().downcast_ref::<arrow_array::Int64Array>());
|
||||
let payloads = batch
|
||||
.column_by_name("payload")
|
||||
.and_then(|c| c.as_any().downcast_ref::<arrow_array::StringArray>());
|
||||
let (Some(states), Some(times)) = (states, times) else {
|
||||
continue;
|
||||
};
|
||||
for i in 0..batch.num_rows() {
|
||||
let state = states.value(i);
|
||||
let ts = times.value(i);
|
||||
created_ms = created_ms.min(ts);
|
||||
updated_ms = updated_ms.max(ts);
|
||||
if matches!(state, "succeeded" | "failed" | "timed_out" | "canceled") {
|
||||
completed_ms = Some(ts);
|
||||
}
|
||||
if let Some(payloads) = payloads {
|
||||
if !arrow_array::Array::is_null(payloads, i) {
|
||||
if let Ok(payload) =
|
||||
serde_json::from_str::<serde_json::Value>(payloads.value(i))
|
||||
{
|
||||
if let Some(e) = payload_error(&payload) {
|
||||
last_error = Some(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
events.push(format!("{state} {ts}"));
|
||||
}
|
||||
}
|
||||
Ok(vec![JobHistoryInfo {
|
||||
table: String::new(),
|
||||
job_id: caller_id.to_string(),
|
||||
job_type: described.job_subtype,
|
||||
state: describe_state_to_client(&described.job_state),
|
||||
column: None,
|
||||
created_ms: if created_ms == i64::MAX {
|
||||
described.creation_ms
|
||||
} else {
|
||||
created_ms
|
||||
},
|
||||
updated_ms,
|
||||
completed_ms,
|
||||
rows_processed: payload_i64(&described.status, "rows_committed"),
|
||||
rows_skipped: payload_i64(&described.status, "rows_skipped"),
|
||||
error: last_error.or_else(|| payload_error(&described.status)),
|
||||
events: (!events.is_empty()).then(|| events.join("\n")),
|
||||
}])
|
||||
}
|
||||
|
||||
async fn errors(&self, job_id: Option<&str>, table: Option<&str>) -> Result<Vec<JobErrorInfo>> {
|
||||
let mut req = self.client.get("/v1/errors");
|
||||
if let Some(j) = job_id {
|
||||
req = req.query(&[("job", j)]);
|
||||
}
|
||||
if let Some(t) = table {
|
||||
req = req.query(&[("table", t)]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteErrorsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.errors.into_iter().map(JobErrorInfo::from).collect())
|
||||
}
|
||||
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>> {
|
||||
let identifier = build_table_identifier(
|
||||
&request.name,
|
||||
@@ -2206,227 +1580,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_derived_compute_routes() {
|
||||
// create_function
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/function/embed/create");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["language"], "python");
|
||||
assert_eq!(body["return_type"], "FLOAT[4]");
|
||||
assert_eq!(body["body"], "def embed(x): ...");
|
||||
assert_eq!(body["options"]["pip"], "torch");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"embed","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.create_function(crate::database::CreateFunctionRequest {
|
||||
name: "embed".into(),
|
||||
language: "python".into(),
|
||||
return_type: "FLOAT[4]".into(),
|
||||
body: "def embed(x): ...".into(),
|
||||
options: [("pip".to_string(), "torch".to_string())].into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// list_functions
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/function/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"functions":[{"name":"embed","language":"python","return_type":"Float32","description":""}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let functions = conn.list_functions().await.unwrap();
|
||||
assert_eq!(functions.len(), 1);
|
||||
assert_eq!(functions[0].name, "embed");
|
||||
|
||||
// drop_function
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/function/embed/drop");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"embed","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.drop_function("embed").await.unwrap();
|
||||
|
||||
// create_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/create");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["query"], "SELECT id, embed(body) AS vec FROM docs");
|
||||
assert_eq!(body["auto_refresh"], true);
|
||||
assert_eq!(body["with_no_data"], false);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","job_id":"j-1"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
let mut request = crate::database::CreateMaterializedViewRequest::new(
|
||||
"mv1",
|
||||
"SELECT id, embed(body) AS vec FROM docs",
|
||||
);
|
||||
request.auto_refresh = true;
|
||||
let job_id = conn.create_materialized_view(request).await.unwrap();
|
||||
assert_eq!(job_id.as_deref(), Some("j-1"));
|
||||
|
||||
// refresh_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/refresh");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["num_workers"], 2);
|
||||
assert!(body.get("src_version").is_none());
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"j-2"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
let mut request = crate::database::RefreshMaterializedViewRequest::new("mv1");
|
||||
request.num_workers = Some(2);
|
||||
let job_id = conn.refresh_materialized_view(request).await.unwrap();
|
||||
assert_eq!(job_id, "j-2");
|
||||
|
||||
// alter_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/alter");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["auto_refresh"], false);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.alter_materialized_view("mv1", false).await.unwrap();
|
||||
|
||||
// drop_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/drop");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.drop_materialized_view("mv1").await.unwrap();
|
||||
|
||||
// list_materialized_views
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"views":[{"name":"mv1","source_table":"docs","projection":["id"],"udf_columns":["vec=embed(body)"],"filter":null,"auto_refresh":true}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let views = conn.list_materialized_views().await.unwrap();
|
||||
assert_eq!(views.len(), 1);
|
||||
assert_eq!(views[0].source_table, "docs");
|
||||
assert!(views[0].auto_refresh);
|
||||
|
||||
// list_jobs: platform listing with status payloads
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/jobs/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"jobs":[{"job_id":"plat-3","table":"docs","job_type":"indexer","job_subtype":"udf","state":"in_progress","created_at_millis":1000,"status":{"units_done":1,"units_total":2}}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let jobs = conn.list_jobs().await.unwrap();
|
||||
assert_eq!(jobs.len(), 1);
|
||||
assert_eq!(jobs[0].state, "running");
|
||||
assert_eq!(jobs[0].job_type, "udf");
|
||||
assert_eq!(jobs[0].units_total, Some(2));
|
||||
|
||||
// cancel_job: resolve via the manifest-id list filter, then cancel
|
||||
let conn = Connection::new_with_handler(|request| match request.url().path() {
|
||||
"/v1/jobs/list" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"jobs":[{"job_id":"plat-3","state":"in_progress"}]}"#)
|
||||
.unwrap(),
|
||||
"/v1/jobs/cancel" => {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id":"plat-3"}"#)
|
||||
.unwrap()
|
||||
}
|
||||
other => panic!("unexpected path {other}"),
|
||||
});
|
||||
assert!(conn.cancel_job("j-3").await.unwrap());
|
||||
|
||||
// cancel_job: never registered -> false, and no cancel request
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"jobs":[]}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert!(!conn.cancel_job("gone").await.unwrap());
|
||||
|
||||
// job_history(None): a registry listing, timeline-free
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"jobs":[{"job_id":"plat-1","table":"docs","job_type":"indexer","job_subtype":"udf","state":"done","created_at_millis":1000,"status":{"rows_skipped":3}}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let hist = conn.job_history(None).await.unwrap();
|
||||
assert_eq!(hist.len(), 1);
|
||||
assert_eq!(hist[0].state, "finished");
|
||||
assert_eq!(hist[0].rows_skipped, Some(3));
|
||||
|
||||
// job_history(id): unknown everywhere -> empty
|
||||
let conn = Connection::new_with_handler(|request| match request.url().path() {
|
||||
"/v1/jobs/describe" => http::Response::builder().status(404).body("").unwrap(),
|
||||
"/v1/jobs/list" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"jobs":[]}"#)
|
||||
.unwrap(),
|
||||
other => panic!("unexpected path {other}"),
|
||||
});
|
||||
assert!(conn.job_history(Some("j-1")).await.unwrap().is_empty());
|
||||
|
||||
// errors: GET /v1/errors with job + table filters
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/errors");
|
||||
assert_eq!(request.url().query(), Some("job=j-1&table=docs"));
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"errors":[{"job_id":"j-1","table":"docs","column":"vec","error_type":"ValueError","error_message":"boom","fragment_id":0,"source_row_id":42,"table_version":7,"age_seconds":5}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let errs = conn.errors(Some("j-1"), Some("docs")).await.unwrap();
|
||||
assert_eq!(errs.len(), 1);
|
||||
assert_eq!(errs[0].error_type, "ValueError");
|
||||
assert_eq!(errs[0].source_row_id, Some(42));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clone_table() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
|
||||
+317
-260
@@ -18,9 +18,11 @@ use crate::index::waiter::wait_for_index;
|
||||
use crate::query::{QueryFilter, QueryRequest, Select, VectorQueryRequest};
|
||||
use crate::table::AddColumnsResult;
|
||||
use crate::table::AddResult;
|
||||
use crate::table::BranchDiff;
|
||||
use crate::table::DeleteResult;
|
||||
use crate::table::DropColumnsResult;
|
||||
use crate::table::LsmWriteSpec;
|
||||
use crate::table::MergeBranchResult;
|
||||
use crate::table::MergeResult;
|
||||
use crate::table::Tags;
|
||||
use crate::table::UpdateResult;
|
||||
@@ -1361,6 +1363,8 @@ impl<S: HttpSend + 'static> RemoteTable<S> {
|
||||
upload_id.to_string(),
|
||||
output.tracker.clone(),
|
||||
self.branch.clone(),
|
||||
self.client.max_bytes_per_request(),
|
||||
self.client.max_request_duration(),
|
||||
));
|
||||
|
||||
let task_ctx = Arc::new(datafusion_execution::TaskContext::default());
|
||||
@@ -1815,6 +1819,79 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn diff_branch(&self, from_branch: &str) -> Result<BranchDiff> {
|
||||
if from_branch.trim().is_empty() {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "from_branch must be a non-empty string".into(),
|
||||
});
|
||||
}
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/branches/diff/", self.identifier))
|
||||
.json(&serde_json::json!({ "from_branch": from_branch }));
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
if response.status() == StatusCode::NOT_FOUND {
|
||||
return Err(Error::TableNotFound {
|
||||
name: format!("{} (branch: {})", self.name, from_branch),
|
||||
source: format!("branch '{}' does not exist", from_branch).into(),
|
||||
});
|
||||
}
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
serde_json::from_str(&body).map_err(|err| Error::Http {
|
||||
source: format!(
|
||||
"Failed to parse diff_branch response: {}, body: {}",
|
||||
err, body
|
||||
)
|
||||
.into(),
|
||||
request_id,
|
||||
status_code: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn merge_branch(&self, from_branch: &str, dry_run: bool) -> Result<MergeBranchResult> {
|
||||
if from_branch.trim().is_empty() {
|
||||
return Err(Error::InvalidInput {
|
||||
message: "from_branch must be a non-empty string".into(),
|
||||
});
|
||||
}
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/branches/merge/", self.identifier))
|
||||
.json(&serde_json::json!({
|
||||
"from_branch": from_branch,
|
||||
"dry_run": dry_run,
|
||||
}));
|
||||
// No retry. 409 rejected merge is final and carries a body.
|
||||
let (request_id, response) = self.send(request, false).await?;
|
||||
let status = response.status();
|
||||
if status == StatusCode::NOT_FOUND {
|
||||
return Err(Error::TableNotFound {
|
||||
name: format!("{} (branch: {})", self.name, from_branch),
|
||||
source: format!("branch '{}' does not exist", from_branch).into(),
|
||||
});
|
||||
}
|
||||
// 200 and 409 both carry MergeBranchResult.
|
||||
if status != StatusCode::OK && status != StatusCode::CONFLICT {
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(Error::Http {
|
||||
source: format!("unexpected status {status} from merge_branch: {body}").into(),
|
||||
request_id,
|
||||
status_code: Some(status),
|
||||
});
|
||||
}
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
serde_json::from_str(&body).map_err(|err| Error::Http {
|
||||
source: format!(
|
||||
"Failed to parse merge_branch response: {}, body: {}",
|
||||
err, body
|
||||
)
|
||||
.into(),
|
||||
request_id,
|
||||
status_code: Some(status),
|
||||
})
|
||||
}
|
||||
|
||||
fn current_branch(&self) -> Option<String> {
|
||||
self.branch.clone()
|
||||
}
|
||||
@@ -1863,25 +1940,33 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
let table_schema = self.schema().await?;
|
||||
let table_def = TableDefinition::try_from_rich_schema(table_schema.clone())?;
|
||||
|
||||
let num_partitions = if let Some(parallelism) = add.write_parallelism {
|
||||
if parallelism > 1 && self.server_version.support_multipart_write() {
|
||||
parallelism
|
||||
} else {
|
||||
1
|
||||
}
|
||||
} else if self.server_version.support_multipart_write() {
|
||||
// Peek at the first batch to estimate write partitions, same as NativeTable.
|
||||
let num_partitions = if self.server_version.support_multipart_write() {
|
||||
// Peek at the first batch to estimate write partitions (same as
|
||||
// NativeTable) and, regardless of `write_parallelism`, to detect a
|
||||
// fully empty input. A multipart write creates its upload session
|
||||
// before any partition executes; if the input turns out to have no
|
||||
// batches at all, no partition ever stages a part (see
|
||||
// `send_multipart_chunked`), so completing the write has nothing to
|
||||
// commit and e.g. `mode=overwrite` would be silently dropped. Route
|
||||
// empty input through the single-request path instead, which always
|
||||
// sends one schema-only request.
|
||||
let mut peeked = PeekedScannable::new(add.data);
|
||||
let n = if let Some(first_batch) = peeked.peek().await {
|
||||
let max_partitions = lance_core::utils::tokio::get_num_compute_intensive_cpus();
|
||||
estimate_write_partitions(
|
||||
first_batch.get_array_memory_size(),
|
||||
first_batch.num_rows(),
|
||||
peeked.num_rows(),
|
||||
max_partitions,
|
||||
)
|
||||
} else {
|
||||
1
|
||||
let n = match peeked.peek().await {
|
||||
Some(first_batch) => match add.write_parallelism {
|
||||
Some(parallelism) if parallelism > 1 => parallelism,
|
||||
Some(_) => 1,
|
||||
None => {
|
||||
let max_partitions =
|
||||
lance_core::utils::tokio::get_num_compute_intensive_cpus();
|
||||
estimate_write_partitions(
|
||||
first_batch.get_array_memory_size(),
|
||||
first_batch.num_rows(),
|
||||
peeked.num_rows(),
|
||||
max_partitions,
|
||||
)
|
||||
}
|
||||
},
|
||||
None => 1,
|
||||
};
|
||||
add.data = Box::new(peeked);
|
||||
n
|
||||
@@ -2109,7 +2194,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
Ok(delete_response)
|
||||
}
|
||||
|
||||
async fn create_index(&self, mut index: IndexBuilder) -> Result<Option<String>> {
|
||||
async fn create_index(&self, mut index: IndexBuilder) -> Result<()> {
|
||||
self.check_mutable().await?;
|
||||
let request = self
|
||||
.client
|
||||
@@ -2202,28 +2287,14 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
|
||||
// The server returns a job id only when the build was deferred to a
|
||||
// background job (pending vector index). Older servers return an
|
||||
// empty body; treat anything unparseable as "no job".
|
||||
#[derive(serde::Deserialize)]
|
||||
struct CreateIndexResponse {
|
||||
job_id: Option<String>,
|
||||
}
|
||||
let job_id = response
|
||||
.text()
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|body| serde_json::from_str::<CreateIndexResponse>(&body).ok())
|
||||
.and_then(|r| r.job_id);
|
||||
self.check_table_response(&request_id, response).await?;
|
||||
|
||||
if let Some(wait_timeout) = index.wait_timeout {
|
||||
let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column));
|
||||
self.wait_for_index(&[&index_name], wait_timeout).await?;
|
||||
}
|
||||
|
||||
Ok(job_id)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Poll until the columns are fully indexed. Will return Error::Timeout if the columns
|
||||
@@ -2436,126 +2507,6 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
message: "optimize is not supported on LanceDB cloud.".into(),
|
||||
})
|
||||
}
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
columns: &[(String, String)],
|
||||
expression: &str,
|
||||
) -> Result<()> {
|
||||
let new_columns: Vec<serde_json::Value> = columns
|
||||
.iter()
|
||||
.map(|(name, data_type)| {
|
||||
serde_json::json!({
|
||||
"name": name,
|
||||
"computed": { "data_type": data_type, "expression": expression },
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/add_columns/", self.identifier))
|
||||
.json(&serde_json::json!({ "new_columns": new_columns }));
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
self.check_table_response(&request_id, response).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn refresh_column(
|
||||
&self,
|
||||
columns: &[String],
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> Result<String> {
|
||||
let mut body = serde_json::json!({ "columns": columns });
|
||||
if let Some(w) = where_clause {
|
||||
body["where_clause"] = serde_json::Value::String(w);
|
||||
}
|
||||
if let Some(n) = num_workers {
|
||||
body["num_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = max_workers {
|
||||
body["max_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = batch_size {
|
||||
body["batch_size"] = n.into();
|
||||
}
|
||||
if let Some(p) = priority {
|
||||
body["priority"] = serde_json::Value::String(p);
|
||||
}
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/refresh_column", self.identifier))
|
||||
.json(&body);
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RefreshColumnResponse {
|
||||
job_id: String,
|
||||
}
|
||||
let body: RefreshColumnResponse = response.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn load_columns(&self, request: crate::table::LoadColumnsRequest) -> Result<String> {
|
||||
let columns: Vec<serde_json::Value> = request
|
||||
.columns
|
||||
.iter()
|
||||
.map(|(target, source)| {
|
||||
serde_json::json!({
|
||||
"target": target,
|
||||
"source": source.clone().unwrap_or_else(|| target.clone()),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let mut source = serde_json::json!({
|
||||
"uris": request.source_uris,
|
||||
"format": request.source_format,
|
||||
});
|
||||
if let Some(opts) = request.source_storage_options {
|
||||
source["storage_options"] = serde_json::to_value(opts).unwrap_or_default();
|
||||
}
|
||||
let mut body = serde_json::json!({
|
||||
"columns": columns,
|
||||
"source": source,
|
||||
"target_key": request.target_key,
|
||||
});
|
||||
if let Some(k) = request.source_key {
|
||||
body["source_key"] = serde_json::Value::String(k);
|
||||
}
|
||||
if let Some(m) = request.on_missing {
|
||||
body["on_missing"] = serde_json::Value::String(m);
|
||||
}
|
||||
if let Some(n) = request.num_workers {
|
||||
body["num_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.max_workers {
|
||||
body["max_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.batch_size {
|
||||
body["batch_size"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.commit_granularity {
|
||||
body["commit_granularity"] = n.into();
|
||||
}
|
||||
if let Some(p) = request.priority {
|
||||
body["priority"] = serde_json::Value::String(p);
|
||||
}
|
||||
let http_request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/load_columns", self.identifier))
|
||||
.json(&body);
|
||||
let (request_id, response) = self.send(http_request, true).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
#[derive(serde::Deserialize)]
|
||||
struct LoadColumnsResponse {
|
||||
job_id: String,
|
||||
}
|
||||
let body: LoadColumnsResponse = response.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn add_columns(
|
||||
&self,
|
||||
transforms: NewColumnTransform,
|
||||
@@ -3050,75 +3001,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_column() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/refresh_column");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["columns"], serde_json::json!(["vec"]));
|
||||
assert_eq!(body["num_workers"], 2);
|
||||
assert!(body.get("where_clause").is_none());
|
||||
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"j-9"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let job_id = table
|
||||
.refresh_column(&["vec".to_string()], None, Some(2), None, None, None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(job_id, "j-9");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_columns() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/load_columns");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
body["columns"],
|
||||
serde_json::json!([{"target": "embedding", "source": "emb"}])
|
||||
);
|
||||
assert_eq!(body["source"]["format"], "parquet");
|
||||
assert_eq!(
|
||||
body["source"]["uris"],
|
||||
serde_json::json!(["s3://b/x.parquet"])
|
||||
);
|
||||
assert_eq!(body["target_key"], "document_id");
|
||||
assert_eq!(body["source_key"], "doc_id");
|
||||
assert_eq!(body["on_missing"], "null");
|
||||
assert_eq!(body["num_workers"], 4);
|
||||
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"lc-7"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let request = crate::table::LoadColumnsRequest {
|
||||
source_uris: vec!["s3://b/x.parquet".to_string()],
|
||||
source_format: "parquet".to_string(),
|
||||
source_storage_options: None,
|
||||
target_key: "document_id".to_string(),
|
||||
source_key: Some("doc_id".to_string()),
|
||||
columns: vec![("embedding".to_string(), Some("emb".to_string()))],
|
||||
on_missing: Some("null".to_string()),
|
||||
num_workers: Some(4),
|
||||
max_workers: None,
|
||||
batch_size: None,
|
||||
commit_granularity: None,
|
||||
priority: None,
|
||||
};
|
||||
let job_id = table.load_columns(request).await.unwrap();
|
||||
assert_eq!(job_id, "lc-7");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_version() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
@@ -4653,42 +4535,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_returns_deferred_job_id() {
|
||||
let table =
|
||||
Table::new_with_handler("my_table", move |request| match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => {
|
||||
let schema = Schema::new(vec![Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||
128,
|
||||
),
|
||||
false,
|
||||
)]);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&schema))
|
||||
.unwrap()
|
||||
}
|
||||
"/v1/table/my_table/create_index/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id": "0a1b2c3d-4e5f-6789-abcd-ef0123456789"}"#.to_string())
|
||||
.unwrap(),
|
||||
path => panic!("Unexpected path: {}", path),
|
||||
});
|
||||
|
||||
let job_id = table
|
||||
.create_index(&["vector"], Index::IvfPq(Default::default()))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
job_id.as_deref(),
|
||||
Some("0a1b2c3d-4e5f-6789-abcd-ef0123456789")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_nested_field_paths() {
|
||||
let schema = nested_index_schema();
|
||||
@@ -7405,6 +7251,76 @@ mod tests {
|
||||
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_write_empty_overwrite_uses_single_partition() {
|
||||
// A multipart write creates its upload session before any partition
|
||||
// executes. If the input has no batches at all, every partition would
|
||||
// stage nothing (see `send_multipart_chunked`), so completing the write
|
||||
// would have nothing to commit and `mode=overwrite` would be silently
|
||||
// dropped. An explicit `write_parallelism` must not force the multipart
|
||||
// path for empty input; it should fall back to the single-request path,
|
||||
// which always sends one schema-only request and carries `mode=overwrite`.
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let multipart_count = Arc::new(AtomicUsize::new(0));
|
||||
|
||||
let insert_count_c = insert_count.clone();
|
||||
let multipart_count_c = multipart_count.clone();
|
||||
|
||||
let table = Table::new_with_handler_version(
|
||||
"my_table",
|
||||
semver::Version::new(0, 4, 0),
|
||||
move |request| {
|
||||
let path = request.url().path();
|
||||
|
||||
if path == "/v1/table/my_table/describe/" {
|
||||
return simple_describe_response();
|
||||
}
|
||||
|
||||
if path.contains("multipart_write") {
|
||||
multipart_count_c.fetch_add(1, Ordering::SeqCst);
|
||||
panic!("Should not use multipart write endpoints for empty input");
|
||||
}
|
||||
|
||||
if path == "/v1/table/my_table/insert/" {
|
||||
let query = request.url().query().unwrap_or("");
|
||||
assert!(
|
||||
!query.contains("upload_id"),
|
||||
"Should not have upload_id for empty input"
|
||||
);
|
||||
assert!(
|
||||
query.contains("mode=overwrite"),
|
||||
"Should carry mode=overwrite, got query: {}",
|
||||
query
|
||||
);
|
||||
insert_count_c.fetch_add(1, Ordering::SeqCst);
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 2}"#.to_string())
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
panic!("Unexpected request path: {}", path);
|
||||
},
|
||||
);
|
||||
|
||||
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)]));
|
||||
let empty_batches: Vec<std::result::Result<RecordBatch, arrow_schema::ArrowError>> =
|
||||
Vec::new();
|
||||
let data: Box<dyn RecordBatchReader + Send> =
|
||||
Box::new(RecordBatchIterator::new(empty_batches, schema));
|
||||
let result = table
|
||||
.add(data)
|
||||
.mode(AddDataMode::Overwrite)
|
||||
.write_parallelism(4)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.version, 2);
|
||||
assert_eq!(multipart_count.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_write_abort_on_insert_failure() {
|
||||
let create_count = Arc::new(AtomicUsize::new(0));
|
||||
@@ -8370,6 +8286,147 @@ mod tests {
|
||||
assert!(matches!(err, Error::TableNotFound { .. }), "got {err:?}");
|
||||
}
|
||||
|
||||
fn sample_branch_diff_json() -> &'static str {
|
||||
r#"{
|
||||
"fromBranch":"exp",
|
||||
"parentVersion":1,
|
||||
"mainVersion":1,
|
||||
"branchVersion":2,
|
||||
"baseMoved":false,
|
||||
"rowCountMain":3,
|
||||
"rowCountBranch":3,
|
||||
"rowSummary":{
|
||||
"unchanged":3,
|
||||
"newOnBase":0,
|
||||
"newOnBranch":0,
|
||||
"staleRecompute":0,
|
||||
"inputsChanged":0,
|
||||
"deltaAvailable":false
|
||||
},
|
||||
"addedColumns":[{"name":"tag","dataType":"utf8","nullable":true}],
|
||||
"removedColumns":[],
|
||||
"changedColumns":[],
|
||||
"addedIndexes":[],
|
||||
"removedIndexes":[],
|
||||
"mergeable":true,
|
||||
"mergeBlockers":[]
|
||||
}"#
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_diff_branch() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/branches/diff/");
|
||||
let body = request_body_json(&request);
|
||||
assert_eq!(body["from_branch"], "exp");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(sample_branch_diff_json())
|
||||
.unwrap()
|
||||
});
|
||||
let diff = table.diff_branch("exp").await.unwrap();
|
||||
assert_eq!(diff.from_branch, "exp");
|
||||
assert!(diff.mergeable);
|
||||
assert_eq!(diff.added_columns.len(), 1);
|
||||
assert_eq!(diff.added_columns[0].name, "tag");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_merge_branch_dry_run() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/branches/merge/");
|
||||
let body = request_body_json(&request);
|
||||
assert_eq!(body["from_branch"], "exp");
|
||||
assert_eq!(body["dry_run"], true);
|
||||
let resp = format!(
|
||||
r#"{{"status":"ready","diff":{},"preview":{{"promotedColumns":["tag"]}}}}"#,
|
||||
sample_branch_diff_json()
|
||||
);
|
||||
http::Response::builder().status(200).body(resp).unwrap()
|
||||
});
|
||||
let result = table.merge_branch("exp", true).await.unwrap();
|
||||
assert_eq!(result.status, crate::table::MergeBranchStatus::Ready);
|
||||
assert_eq!(result.preview.promoted_columns, vec!["tag".to_string()]);
|
||||
assert!(result.main_version_after.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_merge_branch_rejected_returns_ok_with_body() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/branches/merge/");
|
||||
let body = request_body_json(&request);
|
||||
assert_eq!(body["dry_run"], false);
|
||||
let mut diff: serde_json::Value =
|
||||
serde_json::from_str(sample_branch_diff_json()).unwrap();
|
||||
diff["mergeable"] = serde_json::json!(false);
|
||||
diff["mergeBlockers"] = serde_json::json!([{
|
||||
"code": "baseMoved",
|
||||
"message": "main has advanced"
|
||||
}]);
|
||||
let resp = serde_json::json!({
|
||||
"status": "rejected",
|
||||
"diff": diff,
|
||||
"preview": { "promotedColumns": [] }
|
||||
});
|
||||
http::Response::builder()
|
||||
.status(409)
|
||||
.body(resp.to_string())
|
||||
.unwrap()
|
||||
});
|
||||
let result = table.merge_branch("exp", false).await.unwrap();
|
||||
assert_eq!(result.status, crate::table::MergeBranchStatus::Rejected);
|
||||
assert!(!result.diff.mergeable);
|
||||
assert_eq!(result.diff.merge_blockers.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_merge_branch_unknown_blocker_code_parses() {
|
||||
let table = Table::new_with_handler("my_table", |_| {
|
||||
let mut diff: serde_json::Value =
|
||||
serde_json::from_str(sample_branch_diff_json()).unwrap();
|
||||
diff["mergeable"] = serde_json::json!(false);
|
||||
diff["mergeBlockers"] = serde_json::json!([{
|
||||
"code": "multipleCommits",
|
||||
"message": "branch has more than one data commit"
|
||||
}]);
|
||||
let resp = serde_json::json!({
|
||||
"status": "rejected",
|
||||
"diff": diff,
|
||||
"preview": { "operation": "append", "rowsAdded": 2 }
|
||||
});
|
||||
http::Response::builder()
|
||||
.status(409)
|
||||
.body(resp.to_string())
|
||||
.unwrap()
|
||||
});
|
||||
let result = table.merge_branch("exp", false).await.unwrap();
|
||||
assert_eq!(result.status, crate::table::MergeBranchStatus::Rejected);
|
||||
assert_eq!(
|
||||
result.diff.merge_blockers[0].code,
|
||||
crate::table::MergeBlockerCode::Unknown
|
||||
);
|
||||
assert!(result.preview.promoted_columns.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_merge_branch_unexpected_2xx_is_error() {
|
||||
let table = Table::new_with_handler("my_table", |_| {
|
||||
http::Response::builder()
|
||||
.status(204)
|
||||
.body(String::new())
|
||||
.unwrap()
|
||||
});
|
||||
let err = table.merge_branch("exp", false).await.unwrap_err();
|
||||
match err {
|
||||
Error::Http {
|
||||
status_code: Some(code),
|
||||
..
|
||||
} => assert_eq!(code, reqwest::StatusCode::NO_CONTENT),
|
||||
other => panic!("expected Http error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_checkout_branch_validates_via_list() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
//! DataFusion ExecutionPlan for inserting data into remote LanceDB tables.
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use arrow_array::{ArrayRef, RecordBatch, UInt64Array};
|
||||
use arrow_ipc::CompressionType;
|
||||
@@ -15,7 +16,7 @@ use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
|
||||
use datafusion_physical_plan::{
|
||||
DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
|
||||
};
|
||||
use futures::StreamExt;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use http::header::CONTENT_TYPE;
|
||||
use lance::io::exec::utils::InstrumentedRecordBatchStreamAdapter;
|
||||
|
||||
@@ -49,6 +50,15 @@ pub struct RemoteInsertExec<S: HttpSend = Sender> {
|
||||
tracker: Option<Arc<WriteProgressTracker>>,
|
||||
/// Branch to write to via `?branch=`. `None` targets the main branch.
|
||||
branch: Option<String>,
|
||||
/// For multipart writes, split each partition into parts of at most this
|
||||
/// many bytes, each uploaded as a separate request. `None` sends the whole
|
||||
/// partition as a single request.
|
||||
max_bytes_per_request: Option<u64>,
|
||||
/// For multipart writes, also cut a part once it has been uploading for this
|
||||
/// long, even if it has not reached `max_bytes_per_request`. Bounds request
|
||||
/// duration on slow/throttled uploads so no request exceeds the read
|
||||
/// timeout. `None` disables the time-based cut.
|
||||
max_request_duration: Option<Duration>,
|
||||
}
|
||||
|
||||
impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
@@ -63,7 +73,7 @@ impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
branch: Option<String>,
|
||||
) -> Self {
|
||||
Self::new_inner(
|
||||
table_name, identifier, client, input, overwrite, None, tracker, branch,
|
||||
table_name, identifier, client, input, overwrite, None, tracker, branch, None, None,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -82,6 +92,8 @@ impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
upload_id: String,
|
||||
tracker: Option<Arc<WriteProgressTracker>>,
|
||||
branch: Option<String>,
|
||||
max_bytes_per_request: Option<u64>,
|
||||
max_request_duration: Option<Duration>,
|
||||
) -> Self {
|
||||
Self::new_inner(
|
||||
table_name,
|
||||
@@ -92,6 +104,8 @@ impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
Some(upload_id),
|
||||
tracker,
|
||||
branch,
|
||||
max_bytes_per_request,
|
||||
max_request_duration,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -105,6 +119,8 @@ impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
upload_id: Option<String>,
|
||||
tracker: Option<Arc<WriteProgressTracker>>,
|
||||
branch: Option<String>,
|
||||
max_bytes_per_request: Option<u64>,
|
||||
max_request_duration: Option<Duration>,
|
||||
) -> Self {
|
||||
let num_partitions = if upload_id.is_some() {
|
||||
input.output_partitioning().partition_count()
|
||||
@@ -131,6 +147,8 @@ impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
upload_id,
|
||||
tracker,
|
||||
branch,
|
||||
max_bytes_per_request,
|
||||
max_request_duration,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -214,6 +232,238 @@ impl<S: HttpSend + 'static> RemoteInsertExec<S> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared context for the requests of a single partition's multipart upload.
|
||||
/// These values are identical for every part; only the part id and streamed
|
||||
/// body differ between requests. Bundling them keeps the per-part helpers from
|
||||
/// each threading the same handful of arguments.
|
||||
struct PartRequestCtx<'a, S: HttpSend> {
|
||||
client: &'a RestfulLanceDbClient<S>,
|
||||
identifier: &'a str,
|
||||
table_name: &'a str,
|
||||
upload_id: &'a str,
|
||||
branch: Option<&'a str>,
|
||||
overwrite: bool,
|
||||
}
|
||||
|
||||
impl<S: HttpSend + 'static> PartRequestCtx<'_, S> {
|
||||
/// Upload a partition as one or more multipart parts, cutting a new part
|
||||
/// whenever the current one reaches `max_bytes` (Arrow IPC, compressed) or
|
||||
/// has been uploading for `max_duration`, whichever comes first.
|
||||
///
|
||||
/// Each part is a separate `/insert?upload_id=...&upload_part_id=...` request
|
||||
/// whose body is still streamed through a bounded channel, so peak memory
|
||||
/// stays at a couple of batches regardless of `max_bytes`. The server stages
|
||||
/// every part under the shared `upload_id` and merges them atomically when
|
||||
/// the caller completes the multipart write. An empty partition stages
|
||||
/// nothing: the multipart write always has at least one non-empty partition
|
||||
/// to commit.
|
||||
///
|
||||
/// The byte budget targets a good on-disk fragment size; the duration budget
|
||||
/// bounds request time so a slow or throttled upload does not keep a request
|
||||
/// open past the client read timeout (which also covers the request body).
|
||||
async fn send_multipart_chunked(
|
||||
&self,
|
||||
max_bytes: u64,
|
||||
max_duration: Option<Duration>,
|
||||
mut input: SendableRecordBatchStream,
|
||||
tracker: Option<Arc<WriteProgressTracker>>,
|
||||
) -> DataFusionResult<()> {
|
||||
let schema = input.schema();
|
||||
|
||||
// A part always starts from a batch we already hold: the first batch of
|
||||
// the partition, or the look-ahead batch from the previous part. This
|
||||
// keeps empty partitions from staging a part and stops a size cut that
|
||||
// lands exactly on the end of input from emitting a trailing empty part.
|
||||
let mut first = match input.next().await {
|
||||
Some(batch) => batch?,
|
||||
None => return Ok(()),
|
||||
};
|
||||
|
||||
loop {
|
||||
let input_ended = self
|
||||
.send_one_part(
|
||||
&schema,
|
||||
max_bytes,
|
||||
max_duration,
|
||||
first,
|
||||
&mut input,
|
||||
&tracker,
|
||||
)
|
||||
.await?;
|
||||
|
||||
if input_ended {
|
||||
break;
|
||||
}
|
||||
|
||||
first = match input.next().await {
|
||||
Some(batch) => batch?,
|
||||
None => break,
|
||||
};
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build the `/insert` request for a single multipart part.
|
||||
fn build_part_request(&self, part_id: &str, body: reqwest::Body) -> reqwest::RequestBuilder {
|
||||
let mut request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/insert/", self.identifier))
|
||||
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
|
||||
.query(&[("upload_id", self.upload_id)])
|
||||
.query(&[("upload_part_id", part_id)]);
|
||||
// Every part of an overwrite carries `mode=overwrite`. The server records
|
||||
// it against the shared `upload_id` and applies the overwrite once, when
|
||||
// the multipart write is completed, rather than per part.
|
||||
if self.overwrite {
|
||||
request = request.query(&[("mode", "overwrite")]);
|
||||
}
|
||||
if let Some(b) = self.branch {
|
||||
request = request.query(&[("branch", b)]);
|
||||
}
|
||||
request.body(body)
|
||||
}
|
||||
|
||||
/// Send a single part's request and drain the response, mapping HTTP and
|
||||
/// table-not-found errors into `DataFusionError`.
|
||||
async fn send_part_request(&self, request: reqwest::RequestBuilder) -> DataFusionResult<()> {
|
||||
let (request_id, response) = self
|
||||
.client
|
||||
.send(request)
|
||||
.await
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
let response =
|
||||
RemoteTable::<Sender>::handle_table_not_found(self.table_name, response, &request_id)
|
||||
.await
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
let response = self
|
||||
.client
|
||||
.check_response(&request_id, response)
|
||||
.await
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
response.bytes().await.map_err(|e| {
|
||||
DataFusionError::External(Box::new(Error::Http {
|
||||
source: Box::new(e),
|
||||
request_id: request_id.clone(),
|
||||
status_code: None,
|
||||
}))
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Stream one part, starting from `first` and pulling from `input` until the
|
||||
/// part reaches `max_bytes`, has been uploading for `max_duration`, or the
|
||||
/// input ends. The body is streamed through a bounded channel concurrently
|
||||
/// with the request, so peak memory stays at a couple of batches. Wire bytes
|
||||
/// are recorded on `tracker` as each chunk is produced, so progress advances
|
||||
/// smoothly rather than jumping once per completed part. Returns whether the
|
||||
/// input was exhausted while filling this part.
|
||||
async fn send_one_part(
|
||||
&self,
|
||||
schema: &arrow_schema::SchemaRef,
|
||||
max_bytes: u64,
|
||||
max_duration: Option<Duration>,
|
||||
first: RecordBatch,
|
||||
input: &mut SendableRecordBatchStream,
|
||||
tracker: &Option<Arc<WriteProgressTracker>>,
|
||||
) -> DataFusionResult<bool> {
|
||||
let (mut chunk_tx, chunk_rx) =
|
||||
futures::channel::mpsc::channel::<Result<Vec<u8>, std::io::Error>>(2);
|
||||
let body = reqwest::Body::wrap_stream(chunk_rx);
|
||||
|
||||
let part_id = uuid::Uuid::new_v4().to_string();
|
||||
let request = self.build_part_request(&part_id, body);
|
||||
|
||||
// Measured from just before the request is sent, matching the window the
|
||||
// client read timeout applies to the upload.
|
||||
let started = Instant::now();
|
||||
let tracker = tracker.clone();
|
||||
// Unlike `stream_as_http_body`, this producer also cuts the part at the
|
||||
// byte/time budget and reports back whether the input ended, so it drives
|
||||
// its own bounded mpsc channel joined with the request instead of reusing
|
||||
// that helper.
|
||||
let producer = async move {
|
||||
let options = arrow_ipc::writer::IpcWriteOptions::default()
|
||||
.try_with_compression(Some(CompressionType::LZ4_FRAME))
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
let mut writer =
|
||||
arrow_ipc::writer::StreamWriter::try_new_with_options(Vec::new(), schema, options)
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
|
||||
let mut part_bytes: u64 = 0;
|
||||
let mut input_ended = false;
|
||||
let mut pending = Some(first);
|
||||
loop {
|
||||
let batch = match pending.take() {
|
||||
Some(batch) => batch,
|
||||
None => match input.next().await {
|
||||
Some(Ok(batch)) => batch,
|
||||
Some(Err(e)) => {
|
||||
// Abort the body so the server does not treat the
|
||||
// truncated stream as a successful write; the
|
||||
// original error is surfaced to the caller.
|
||||
let _ = chunk_tx
|
||||
.send(Err(std::io::Error::other("input stream error")))
|
||||
.await;
|
||||
return Err(e);
|
||||
}
|
||||
None => {
|
||||
input_ended = true;
|
||||
break;
|
||||
}
|
||||
},
|
||||
};
|
||||
writer
|
||||
.write(&batch)
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
let chunk = std::mem::take(writer.get_mut());
|
||||
let chunk_len = chunk.len();
|
||||
part_bytes += chunk_len as u64;
|
||||
if chunk_tx.send(Ok(chunk)).await.is_err() {
|
||||
// The request finished or failed; stop producing.
|
||||
break;
|
||||
}
|
||||
if let Some(ref t) = tracker {
|
||||
t.record_bytes(chunk_len);
|
||||
}
|
||||
if part_bytes >= max_bytes
|
||||
|| max_duration.is_some_and(|limit| started.elapsed() >= limit)
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
writer
|
||||
.finish()
|
||||
.map_err(|e| DataFusionError::External(Box::new(e)))?;
|
||||
let tail = std::mem::take(writer.get_mut());
|
||||
if !tail.is_empty() {
|
||||
let tail_len = tail.len();
|
||||
if chunk_tx.send(Ok(tail)).await.is_ok()
|
||||
&& let Some(ref t) = tracker
|
||||
{
|
||||
t.record_bytes(tail_len);
|
||||
}
|
||||
}
|
||||
Ok::<bool, DataFusionError>(input_ended)
|
||||
};
|
||||
|
||||
let send = self.send_part_request(request);
|
||||
|
||||
// `join!` rather than `tokio::spawn`: the producer borrows `input` (and
|
||||
// `schema`), so it cannot satisfy the `'static` bound a spawned task
|
||||
// needs. Running both futures on this task lets them make progress
|
||||
// concurrently without that constraint.
|
||||
let (producer_result, send_result) = futures::join!(producer, send);
|
||||
// Prefer the producer error (e.g. NaN rejection) over any HTTP error it
|
||||
// induced.
|
||||
let input_ended = producer_result?;
|
||||
send_result?;
|
||||
|
||||
Ok(input_ended)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: HttpSend + 'static> DisplayAs for RemoteInsertExec<S> {
|
||||
fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match t {
|
||||
@@ -278,6 +528,8 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteInsertExec<S> {
|
||||
self.upload_id.clone(),
|
||||
self.tracker.clone(),
|
||||
self.branch.clone(),
|
||||
self.max_bytes_per_request,
|
||||
self.max_request_duration,
|
||||
)))
|
||||
}
|
||||
|
||||
@@ -310,8 +562,36 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteInsertExec<S> {
|
||||
let upload_id = self.upload_id.clone();
|
||||
let tracker = self.tracker.clone();
|
||||
let branch = self.branch.clone();
|
||||
let max_bytes_per_request = self.max_bytes_per_request;
|
||||
let max_request_duration = self.max_request_duration;
|
||||
|
||||
let stream = futures::stream::once(async move {
|
||||
// Multipart writes with a byte budget split the partition into
|
||||
// several bounded, still-streamed requests so no single request
|
||||
// stays open long enough to hit the client read timeout.
|
||||
if let (Some(upload_id), Some(max_bytes)) =
|
||||
(upload_id.as_deref(), max_bytes_per_request)
|
||||
{
|
||||
let ctx = PartRequestCtx {
|
||||
client: &client,
|
||||
identifier: &identifier,
|
||||
table_name: &table_name,
|
||||
upload_id,
|
||||
branch: branch.as_deref(),
|
||||
overwrite,
|
||||
};
|
||||
ctx.send_multipart_chunked(max_bytes, max_request_duration, input_stream, tracker)
|
||||
.await?;
|
||||
// Count 0 here as for the non-multipart path below: the parts are
|
||||
// only staged, so the real row count is resolved when the caller
|
||||
// completes the multipart write.
|
||||
let count_array: ArrayRef = Arc::new(UInt64Array::from(vec![0u64]));
|
||||
return Ok::<RecordBatch, DataFusionError>(RecordBatch::try_new(
|
||||
COUNT_SCHEMA.clone(),
|
||||
vec![count_array],
|
||||
)?);
|
||||
}
|
||||
|
||||
let mut request = client
|
||||
.post(&format!("/v1/table/{}/insert/", identifier))
|
||||
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE);
|
||||
@@ -423,9 +703,15 @@ mod tests {
|
||||
use arrow_schema::{DataType, Field, Schema as ArrowSchema};
|
||||
use datafusion::prelude::SessionContext;
|
||||
use datafusion_catalog::MemTable;
|
||||
use std::sync::Arc;
|
||||
use datafusion_common::{DataFusionError, Result as DataFusionResult};
|
||||
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
|
||||
use datafusion_physical_expr::EquivalenceProperties;
|
||||
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
|
||||
use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use super::RemoteInsertExec;
|
||||
use crate::Table;
|
||||
use crate::remote::ARROW_STREAM_CONTENT_TYPE;
|
||||
use crate::table::datafusion::BaseTableAdapter;
|
||||
@@ -591,4 +877,489 @@ mod tests {
|
||||
// Verify: should have made exactly one HTTP request despite multiple input partitions
|
||||
assert_eq!(request_count.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
/// Build a single-partition input plan from the given batches.
|
||||
async fn input_plan_from_batches(
|
||||
schema: Arc<ArrowSchema>,
|
||||
batches: Vec<arrow_array::RecordBatch>,
|
||||
) -> Arc<dyn ExecutionPlan> {
|
||||
use datafusion_catalog::TableProvider;
|
||||
let mem = MemTable::try_new(schema, vec![batches]).unwrap();
|
||||
let ctx = SessionContext::new();
|
||||
mem.scan(&ctx.state(), None, &[], None).await.unwrap()
|
||||
}
|
||||
|
||||
/// Build a single-partition input plan from the batches spread across the
|
||||
/// given partitions.
|
||||
async fn input_plan_from_partitions(
|
||||
schema: Arc<ArrowSchema>,
|
||||
partitions: Vec<Vec<arrow_array::RecordBatch>>,
|
||||
) -> Arc<dyn ExecutionPlan> {
|
||||
use datafusion_catalog::TableProvider;
|
||||
let mem = MemTable::try_new(schema, partitions).unwrap();
|
||||
let ctx = SessionContext::new();
|
||||
mem.scan(&ctx.state(), None, &[], None).await.unwrap()
|
||||
}
|
||||
|
||||
fn counting_insert_client(
|
||||
counter: Arc<AtomicUsize>,
|
||||
) -> crate::remote::client::RestfulLanceDbClient<crate::remote::client::test_utils::MockSender>
|
||||
{
|
||||
crate::remote::client::test_utils::client_with_handler(move |request| {
|
||||
let path = request.url().path();
|
||||
assert_eq!(path, "/v1/table/my_table/insert/");
|
||||
let query = request.url().query().unwrap_or("");
|
||||
assert!(query.contains("upload_id=upload-1"), "query: {query}");
|
||||
assert!(query.contains("upload_part_id="), "query: {query}");
|
||||
counter.fetch_add(1, Ordering::SeqCst);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(String::new())
|
||||
.unwrap()
|
||||
})
|
||||
}
|
||||
|
||||
/// Insert handler that records the `upload_part_id` of every part request so
|
||||
/// a test can assert the ids are distinct.
|
||||
fn recording_insert_client(
|
||||
part_ids: Arc<Mutex<Vec<String>>>,
|
||||
) -> crate::remote::client::RestfulLanceDbClient<crate::remote::client::test_utils::MockSender>
|
||||
{
|
||||
crate::remote::client::test_utils::client_with_handler(move |request| {
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/insert/");
|
||||
let part_id = request
|
||||
.url()
|
||||
.query_pairs()
|
||||
.find(|(k, _)| k == "upload_part_id")
|
||||
.map(|(_, v)| v.into_owned())
|
||||
.expect("upload_part_id query param");
|
||||
part_ids.lock().unwrap().push(part_id);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(String::new())
|
||||
.unwrap()
|
||||
})
|
||||
}
|
||||
|
||||
/// Single-partition input plan that yields one good batch and then an error,
|
||||
/// for exercising the mid-part input-error abort path in `send_one_part`.
|
||||
#[derive(Debug)]
|
||||
struct ErroringExec {
|
||||
schema: Arc<ArrowSchema>,
|
||||
properties: Arc<PlanProperties>,
|
||||
}
|
||||
|
||||
impl ErroringExec {
|
||||
fn new() -> Self {
|
||||
let schema = record_batch!(("id", Int32, [1, 2])).unwrap().schema();
|
||||
let properties = PlanProperties::new(
|
||||
EquivalenceProperties::new(schema.clone()),
|
||||
datafusion_physical_plan::Partitioning::UnknownPartitioning(1),
|
||||
datafusion_physical_plan::execution_plan::EmissionType::Incremental,
|
||||
datafusion_physical_plan::execution_plan::Boundedness::Bounded,
|
||||
);
|
||||
Self {
|
||||
schema,
|
||||
properties: Arc::new(properties),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl DisplayAs for ErroringExec {
|
||||
fn fmt_as(
|
||||
&self,
|
||||
_t: DisplayFormatType,
|
||||
f: &mut std::fmt::Formatter<'_>,
|
||||
) -> std::fmt::Result {
|
||||
write!(f, "ErroringExec")
|
||||
}
|
||||
}
|
||||
|
||||
impl ExecutionPlan for ErroringExec {
|
||||
fn name(&self) -> &str {
|
||||
"ErroringExec"
|
||||
}
|
||||
fn properties(&self) -> &Arc<PlanProperties> {
|
||||
&self.properties
|
||||
}
|
||||
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
|
||||
vec![]
|
||||
}
|
||||
fn with_new_children(
|
||||
self: Arc<Self>,
|
||||
_children: Vec<Arc<dyn ExecutionPlan>>,
|
||||
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
|
||||
Ok(self)
|
||||
}
|
||||
fn execute(
|
||||
&self,
|
||||
_partition: usize,
|
||||
_context: Arc<TaskContext>,
|
||||
) -> DataFusionResult<SendableRecordBatchStream> {
|
||||
let batch = record_batch!(("id", Int32, [1, 2])).unwrap();
|
||||
let stream = futures::stream::iter(vec![
|
||||
Ok(batch),
|
||||
Err(DataFusionError::Execution("boom".to_string())),
|
||||
]);
|
||||
Ok(Box::pin(RecordBatchStreamAdapter::new(
|
||||
self.schema.clone(),
|
||||
stream,
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_chunked_splits_into_parts() {
|
||||
use futures::StreamExt;
|
||||
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let client = counting_insert_client(insert_count.clone());
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
true,
|
||||
)]));
|
||||
let batches = vec![
|
||||
record_batch!(("id", Int32, [1, 2])).unwrap(),
|
||||
record_batch!(("id", Int32, [3, 4])).unwrap(),
|
||||
record_batch!(("id", Int32, [5, 6])).unwrap(),
|
||||
];
|
||||
let input = input_plan_from_batches(schema, batches).await;
|
||||
|
||||
// A 1-byte budget forces every batch into its own part.
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(1),
|
||||
None,
|
||||
);
|
||||
|
||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
while stream.next().await.transpose().unwrap().is_some() {}
|
||||
|
||||
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_single_part_when_under_budget() {
|
||||
use futures::StreamExt;
|
||||
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let client = counting_insert_client(insert_count.clone());
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
true,
|
||||
)]));
|
||||
let batches = vec![
|
||||
record_batch!(("id", Int32, [1, 2])).unwrap(),
|
||||
record_batch!(("id", Int32, [3, 4])).unwrap(),
|
||||
record_batch!(("id", Int32, [5, 6])).unwrap(),
|
||||
];
|
||||
let input = input_plan_from_batches(schema, batches).await;
|
||||
|
||||
// A large byte budget and no time limit keep the whole partition in a
|
||||
// single part.
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(64 * 1024 * 1024),
|
||||
None,
|
||||
);
|
||||
|
||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
while stream.next().await.transpose().unwrap().is_some() {}
|
||||
|
||||
assert_eq!(insert_count.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_chunked_splits_by_duration() {
|
||||
use futures::StreamExt;
|
||||
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let client = counting_insert_client(insert_count.clone());
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
true,
|
||||
)]));
|
||||
let batches = vec![
|
||||
record_batch!(("id", Int32, [1, 2])).unwrap(),
|
||||
record_batch!(("id", Int32, [3, 4])).unwrap(),
|
||||
record_batch!(("id", Int32, [5, 6])).unwrap(),
|
||||
];
|
||||
let input = input_plan_from_batches(schema, batches).await;
|
||||
|
||||
// A large byte budget but a tiny duration budget: writing and sending
|
||||
// one batch already takes longer than the limit, so each batch is cut
|
||||
// into its own part on the time check rather than the byte check.
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(64 * 1024 * 1024),
|
||||
Some(std::time::Duration::from_nanos(1)),
|
||||
);
|
||||
|
||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
while stream.next().await.transpose().unwrap().is_some() {}
|
||||
|
||||
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_empty_partition_stages_nothing() {
|
||||
use futures::StreamExt;
|
||||
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let client = counting_insert_client(insert_count.clone());
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
true,
|
||||
)]));
|
||||
// An empty partition should stage no parts; on the multipart path the
|
||||
// write relies on another partition having data to commit.
|
||||
let input = input_plan_from_batches(schema, vec![]).await;
|
||||
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(64 * 1024 * 1024),
|
||||
None,
|
||||
);
|
||||
|
||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
while stream.next().await.transpose().unwrap().is_some() {}
|
||||
|
||||
assert_eq!(insert_count.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_chunked_uses_distinct_part_ids() {
|
||||
use futures::StreamExt;
|
||||
use std::collections::HashSet;
|
||||
|
||||
let part_ids = Arc::new(Mutex::new(Vec::new()));
|
||||
let client = recording_insert_client(part_ids.clone());
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
true,
|
||||
)]));
|
||||
let batches = vec![
|
||||
record_batch!(("id", Int32, [1, 2])).unwrap(),
|
||||
record_batch!(("id", Int32, [3, 4])).unwrap(),
|
||||
record_batch!(("id", Int32, [5, 6])).unwrap(),
|
||||
];
|
||||
let input = input_plan_from_batches(schema, batches).await;
|
||||
|
||||
// A 1-byte budget forces every batch into its own part.
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(1),
|
||||
None,
|
||||
);
|
||||
|
||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
while stream.next().await.transpose().unwrap().is_some() {}
|
||||
|
||||
let ids = part_ids.lock().unwrap().clone();
|
||||
assert_eq!(ids.len(), 3, "expected one part id per part: {ids:?}");
|
||||
assert!(
|
||||
ids.iter().all(|id| !id.is_empty()),
|
||||
"part ids must be non-empty: {ids:?}"
|
||||
);
|
||||
let unique: HashSet<&String> = ids.iter().collect();
|
||||
assert_eq!(unique.len(), 3, "part ids must be distinct: {ids:?}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_chunks_each_partition_independently() {
|
||||
use futures::StreamExt;
|
||||
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let client = counting_insert_client(insert_count.clone());
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
true,
|
||||
)]));
|
||||
let partitions = vec![
|
||||
// Partition 0: two batches, split into two parts by the 1-byte budget.
|
||||
vec![
|
||||
record_batch!(("id", Int32, [1, 2])).unwrap(),
|
||||
record_batch!(("id", Int32, [3, 4])).unwrap(),
|
||||
],
|
||||
// Partition 1: one batch, one part.
|
||||
vec![record_batch!(("id", Int32, [5, 6])).unwrap()],
|
||||
];
|
||||
let input = input_plan_from_partitions(schema, partitions).await;
|
||||
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(1),
|
||||
None,
|
||||
);
|
||||
|
||||
for partition in 0..2 {
|
||||
let mut stream = exec
|
||||
.execute(partition, Arc::new(TaskContext::default()))
|
||||
.unwrap();
|
||||
while stream.next().await.transpose().unwrap().is_some() {}
|
||||
}
|
||||
|
||||
// 2 parts from partition 0 + 1 part from partition 1.
|
||||
assert_eq!(insert_count.load(Ordering::SeqCst), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_input_error_surfaces_original() {
|
||||
use futures::StreamExt;
|
||||
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let client = counting_insert_client(insert_count.clone());
|
||||
|
||||
// A large byte budget keeps the good batch and the following error in
|
||||
// the same part, exercising the mid-part abort path.
|
||||
let input: Arc<dyn ExecutionPlan> = Arc::new(ErroringExec::new());
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
Some(64 * 1024 * 1024),
|
||||
None,
|
||||
);
|
||||
|
||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
let mut err = None;
|
||||
while let Some(item) = stream.next().await {
|
||||
if let Err(e) = item {
|
||||
err = Some(e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let err = err.expect("expected the input stream error to surface");
|
||||
// The original DataFusion error must win over the HTTP error it induces.
|
||||
assert!(
|
||||
err.to_string().contains("boom"),
|
||||
"expected original input error, got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_multipart_records_progress_within_a_part() {
|
||||
use crate::table::write_progress::{ProgressCallback, WriteProgress, WriteProgressTracker};
|
||||
use futures::StreamExt;
|
||||
|
||||
let insert_count = Arc::new(AtomicUsize::new(0));
|
||||
let client = counting_insert_client(insert_count.clone());
|
||||
|
||||
let schema = Arc::new(ArrowSchema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
true,
|
||||
)]));
|
||||
let batches = vec![
|
||||
record_batch!(("id", Int32, [1, 2])).unwrap(),
|
||||
record_batch!(("id", Int32, [3, 4])).unwrap(),
|
||||
record_batch!(("id", Int32, [5, 6])).unwrap(),
|
||||
];
|
||||
let input = input_plan_from_batches(schema, batches).await;
|
||||
|
||||
let observed = Arc::new(Mutex::new(Vec::<usize>::new()));
|
||||
let observed_cb = observed.clone();
|
||||
let callback: ProgressCallback = Arc::new(Mutex::new(move |p: &WriteProgress| {
|
||||
observed_cb.lock().unwrap().push(p.output_bytes());
|
||||
}));
|
||||
let tracker = Arc::new(WriteProgressTracker::new(callback, None));
|
||||
|
||||
// A large byte budget keeps all three batches in one part; smooth
|
||||
// progress therefore requires bytes to be reported per chunk rather than
|
||||
// once when the part completes.
|
||||
let exec = RemoteInsertExec::new_multipart(
|
||||
"my_table".to_string(),
|
||||
"my_table".to_string(),
|
||||
client,
|
||||
input,
|
||||
false,
|
||||
"upload-1".to_string(),
|
||||
Some(tracker),
|
||||
None,
|
||||
Some(64 * 1024 * 1024),
|
||||
None,
|
||||
);
|
||||
|
||||
let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap();
|
||||
while stream.next().await.transpose().unwrap().is_some() {}
|
||||
|
||||
assert_eq!(
|
||||
insert_count.load(Ordering::SeqCst),
|
||||
1,
|
||||
"batches should all land in a single part"
|
||||
);
|
||||
let observed = observed.lock().unwrap();
|
||||
assert!(
|
||||
observed.len() > 1,
|
||||
"expected multiple incremental progress updates within the part: {observed:?}"
|
||||
);
|
||||
assert!(
|
||||
observed.windows(2).all(|w| w[1] >= w[0]),
|
||||
"progress bytes should be monotonic: {observed:?}"
|
||||
);
|
||||
assert!(
|
||||
*observed.last().unwrap() > 0,
|
||||
"final progress should report bytes: {observed:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user