diff --git a/.agents/skills/lancedb-update-lance-dependency/SKILL.md b/.agents/skills/lancedb-update-lance-dependency/SKILL.md index e97fae7ac..3cba2ae75 100644 --- a/.agents/skills/lancedb-update-lance-dependency/SKILL.md +++ b/.agents/skills/lancedb-update-lance-dependency/SKILL.md @@ -48,6 +48,23 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r This updates the Rust workspace Lance dependencies through `ci/set_lance_version.py`, updates `java/pom.xml`, refreshes Cargo metadata, and prints JSON metadata containing `branch_name`, `commit_message`, and `pr_title`. + Also synchronize the Python requirements in `python/pyproject.toml`: set the + `tests` extra's exact `pylance` pin to the matching published Python release, + and update the `pylance` extra's minimum when needed to exclude older bundled + Lance majors. The minimum must still accept the test pin. Python prerelease + spelling differs from Rust (for example, `15.0.0-beta.1` becomes `15.0.0b1`). + Confirm the matching Python release is published before changing the pin. + + Refresh and check `python/uv.lock` from the repository root: + + ```bash + uv lock --project python --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + uv lock --project python --check --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + ``` + + This strategy considers versions across indexes; use it only with trusted + indexes. Commit the Python requirements and lockfile together with the bump. + 5. Run validation: ```bash @@ -57,6 +74,16 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r Fix real diagnostics and rerun clippy until it succeeds. Do not skip warnings. + Run the Python dependency guards with Python 3.11+ and `packaging` installed: + + ```bash + python3 -m unittest discover -s ci/tests -p test_python_dependencies.py -v + ``` + + These checks reject older Lance majors and ensure the `pylance` and `tests` + extras remain compatible. The full CI-script suite uses + `python3 -m unittest discover -s ci/tests -v`. + 6. Inspect `git status --short` and `git diff` to ensure only the Lance dependency update and required compatibility fixes are present. 7. If the task only asks to prepare local changes, stop here and report the changed files and validation result. diff --git a/.bumpversion.toml b/.bumpversion.toml index 0909579cb..7dcefa69c 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.41.0-beta.13" +current_version = "0.41.0-beta.14" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/.github/workflows/ci-scripts.yml b/.github/workflows/ci-scripts.yml index f7f0dbb07..d45289c2d 100644 --- a/.github/workflows/ci-scripts.yml +++ b/.github/workflows/ci-scripts.yml @@ -5,6 +5,8 @@ on: branches: - main paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -12,6 +14,8 @@ on: - .github/workflows/ci-scripts.yml pull_request: paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -30,5 +34,7 @@ jobs: - uses: actions/setup-python@v6 with: python-version: "3.13" + - name: Install test dependencies + run: python -m pip install packaging - name: Run tests run: python -m unittest discover -s ci/tests -v diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index d0ec583bf..fb3fb1d6a 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -38,6 +38,7 @@ jobs: uses: actions/checkout@v6 - name: Install dependencies needed for ubuntu run: | + sudo apt update sudo apt install -y protobuf-compiler libssl-dev rustup update && rustup default - name: Set up Python diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml index e71564473..1784c4c8a 100644 --- a/.github/workflows/pypi-publish.yml +++ b/.github/workflows/pypi-publish.yml @@ -15,6 +15,9 @@ on: - ci/check_cargo_lock.py - ci/check_sdist.py - python/pyproject.toml + - python/Cargo.toml + - python/vendor/** + - rust/lancedb/Cargo.toml - Cargo.toml # Change in dependency frequently breaks builds - Cargo.lock diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index cbfefc83d..17fe3a7b4 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -29,6 +29,63 @@ env: RUST_BACKTRACE: "1" jobs: + dependency-overrides: + name: Downstream Lance override + runs-on: ubuntu-24.04 + timeout-minutes: 15 + steps: + - uses: actions/checkout@v6 + - uses: actions-rust-lang/setup-rust-toolchain@v1 + - name: Read the pinned Lance tag + id: lance + run: | + python3 - <<'PYTHON' + import os + import tomllib + from pathlib import Path + + manifest = tomllib.loads(Path("Cargo.toml").read_text()) + tag = manifest["workspace"]["dependencies"]["lance"]["tag"] + with open(os.environ["GITHUB_OUTPUT"], "a") as output: + output.write(f"tag={tag}\n") + PYTHON + - uses: actions/checkout@v6 + with: + repository: lance-format/lance + ref: ${{ steps.lance.outputs.tag }} + path: local-lance + - name: Resolve a downstream workspace with a local Lance override + run: | + mkdir -p downstream/src + touch downstream/src/lib.rs + cat > downstream/Cargo.toml <<'TOML' + [package] + name = "lancedb-local-lance-consumer" + version = "0.1.0" + edition = "2024" + + [workspace] + + [dependencies] + lancedb = { path = "../rust/lancedb", default-features = false } + lance-arrow = { path = "../local-lance/rust/lance-arrow" } + + [patch."https://github.com/lance-format/lance.git"] + lance-arrow = { path = "../local-lance/rust/lance-arrow" } + TOML + cargo metadata --manifest-path downstream/Cargo.toml --format-version 1 > downstream/graph.json + python3 - <<'PYTHON' + import json + from pathlib import Path + + graph = json.loads(Path("downstream/graph.json").read_text()) + packages = [p for p in graph["packages"] if p["name"] == "lance-arrow"] + assert len(packages) == 1, [p["id"] for p in packages] + assert Path(packages[0]["manifest_path"]).resolve() == Path( + "local-lance/rust/lance-arrow/Cargo.toml" + ).resolve() + PYTHON + lint: timeout-minutes: 30 runs-on: ubuntu-24.04 @@ -245,12 +302,6 @@ jobs: include: - target: x86_64-pc-windows-msvc runner: windows-2022 - # windows-11-arm is a standard runner, so it is free on public repos. - # Running natively lets the aarch64 tests actually execute -- this - # job used to cross-compile them and then skip the test step, paying - # full codegen and link cost for a compile check. - - target: aarch64-pc-windows-msvc - runner: windows-11-arm runs-on: ${{ matrix.runner }} defaults: run: @@ -279,6 +330,43 @@ jobs: # dependency graph a second time. cargo test --profile ci --features aws,remote --locked --target ${{ matrix.target }} + # Run the aarch64 Windows tests only on main so PRs wait for a single + # Windows platform. windows-11-arm is a standard runner, so it is free on + # public repos. Running natively lets the aarch64 tests actually execute -- + # this job used to cross-compile them and then skip the test step, paying + # full codegen and link cost for a compile check. + windows-arm: + # Keep the historical name so the check shows up as before on main. + name: windows (aarch64-pc-windows-msvc, windows-11-arm) + if: github.ref == 'refs/heads/main' + runs-on: windows-11-arm + defaults: + run: + working-directory: rust/lancedb + steps: + - uses: actions/checkout@v6 + - name: Set target + run: rustup target add aarch64-pc-windows-msvc + - uses: Swatinem/rust-cache@v2 + with: + # Restore everywhere, but only save from main. Per-PR saves are + # unreadable outside their own branch anyway, since GitHub scopes + # caches to the creating ref. + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Install Protoc v21.12 + run: choco install --no-progress protoc + - name: Build + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + cargo build --profile ci --features aws,remote --tests --locked --target aarch64-pc-windows-msvc + - name: Run tests + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + # `--target` has to match the build step above. Without it cargo uses + # target/ci/ rather than target//ci/ and rebuilds the entire + # dependency graph a second time. + cargo test --profile ci --features aws,remote --locked --target aarch64-pc-windows-msvc + msrv: # Check the minimum supported Rust version name: MSRV Check - Rust v${{ matrix.msrv }} diff --git a/Cargo.lock b/Cargo.lock index 34ee05cb1..03f4b0864 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2230,9 +2230,9 @@ dependencies = [ [[package]] name = "ctutils" -version = "0.4.2" +version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +checksum = "12218662ccac37d9714f2191252553d67585b97df8c839cbe49e0bde278dc865" dependencies = [ "cmov", ] @@ -5700,7 +5700,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "ahash", "anyhow", @@ -5798,7 +5798,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "arrow-array", "arrow-buffer", @@ -5823,7 +5823,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" dependencies = [ "arc-swap", "arrow", @@ -10446,9 +10446,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.53.1" +version = "1.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "ce3335fa71841cda333a58d7615b03901380ecf09d59b3296d21f8bbac0dde4e" dependencies = [ "bytes", "libc", diff --git a/ci/set_lance_version.py b/ci/set_lance_version.py index c7bf41f55..bef97f8ad 100644 --- a/ci/set_lance_version.py +++ b/ci/set_lance_version.py @@ -2,6 +2,8 @@ import argparse import re import sys import json +import tomllib +from pathlib import Path LANCE_GIT_URL = "https://github.com/lance-format/lance.git" @@ -148,6 +150,18 @@ def update_cargo_toml(line_updater): with open("Cargo.toml", "w") as f: f.writelines(new_lines) + core_manifest = Path("rust/lancedb/Cargo.toml") + if core_manifest.exists(): + lines = core_manifest.read_text().splitlines(keepends=True) + for index, line in enumerate(lines): + if line.startswith("lance-arrow ="): + updated = tomllib.loads(line_updater(line))["lance-arrow"] + config = {"version": updated} if isinstance(updated, str) else updated + if "path" in config: + config["path"] = "../../" + config["path"] + lines[index] = dict_to_toml_line("lance-arrow", config) + core_manifest.write_text("".join(lines)) + def set_stable_version(version: str): """ diff --git a/ci/tests/test_python_dependencies.py b/ci/tests/test_python_dependencies.py new file mode 100644 index 000000000..97fb672c6 --- /dev/null +++ b/ci/tests/test_python_dependencies.py @@ -0,0 +1,47 @@ +import tomllib +import unittest +from pathlib import Path + +from packaging.requirements import Requirement +from packaging.version import Version + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +class PythonDependenciesTest(unittest.TestCase): + def setUp(self): + with (REPO_ROOT / "python/pyproject.toml").open("rb") as f: + self.project = tomllib.load(f)["project"] + self.pylance = Requirement(self.project["optional-dependencies"]["pylance"][0]) + + def test_pylance_extra_excludes_older_lance_majors(self): + with (REPO_ROOT / "Cargo.toml").open("rb") as f: + lance = tomllib.load(f)["workspace"]["dependencies"]["lance"] + if "version" not in lance: + self.skipTest("Local Lance checkout has no version requirement") + major = Version(lance["version"].lstrip("=")).major + for older_major in range(major): + for suffix in (".0.0", ".0.0b6", ".999.999"): + version = f"{older_major}{suffix}" + with self.subTest(version=version): + self.assertFalse( + self.pylance.specifier.contains(version, prereleases=True), + f"The pylance extra must exclude Lance {version}; " + f"the bundled Lance is major {major}", + ) + + def test_pylance_extra_accepts_tested_version(self): + requirements = [ + Requirement(dep) for dep in self.project["optional-dependencies"]["tests"] + ] + tested = next(dep for dep in requirements if dep.name == "pylance") + pin = next(spec for spec in tested.specifier if spec.operator == "==") + self.assertTrue( + self.pylance.specifier.contains(pin.version, prereleases=True), + f"The tests and pylance extras must be installable together: {tested}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/ci/tests/test_set_lance_version.py b/ci/tests/test_set_lance_version.py index 1493fc59d..1173eab84 100644 --- a/ci/tests/test_set_lance_version.py +++ b/ci/tests/test_set_lance_version.py @@ -4,6 +4,7 @@ import subprocess import sys import tempfile import textwrap +import tomllib import unittest from pathlib import Path @@ -83,6 +84,13 @@ lance_datafusion = {{ "version" = "=8.2.0-beta.4", "features" = ["substrait"], " with self.subTest(version=version), tempfile.TemporaryDirectory() as tmp: workdir = Path(tmp) (workdir / "Cargo.toml").write_text(CARGO_TOML) + core_manifest = workdir / "rust/lancedb/Cargo.toml" + core_manifest.parent.mkdir(parents=True) + core_manifest.write_text( + "[dependencies]\nlance-arrow = { " + 'version = "=1.0.0", git = "https://github.com/lance-format/lance.git", ' + 'tag = "v1.0.0" }\narrow.workspace = true\n' + ) command_log = workdir / "commands.log" fake_bin = workdir / "bin" fake_bin.mkdir() @@ -111,6 +119,26 @@ lance_datafusion = {{ "version" = "=8.2.0-beta.4", "features" = ["substrait"], " + updated_dependencies + UNTOUCHED_DEPENDENCIES, ) + core = tomllib.loads(core_manifest.read_text())["dependencies"] + dependency = core["lance-arrow"] + if isinstance(dependency, str): + dependency = {"version": dependency} + if version == "local": + self.assertEqual( + dependency["path"], "../../../lance/rust/lance-arrow" + ) + else: + expected_version = { + "stable": "9.9.9", + "preview": "10.0.0-beta.3", + }.get(version, version.removeprefix("v")) + self.assertEqual(dependency["version"], f"={expected_version}") + if "beta" in expected_version: + self.assertEqual(dependency["git"], LANCE_GIT_URL) + self.assertEqual(dependency["tag"], f"v{expected_version}") + else: + self.assertNotIn("git", dependency) + self.assertEqual(core["arrow"], {"workspace": True}) commands = command_log.read_text().splitlines() for command in expected_commands: self.assertTrue( diff --git a/deny.toml b/deny.toml index 7d60ae63f..e4982c149 100644 --- a/deny.toml +++ b/deny.toml @@ -110,8 +110,18 @@ ignore = [ # compact_str/smol_str, so clearing this requires polars to migrate. # https://rustsec.org/advisories/RUSTSEC-2026-0249 { id = "RUSTSEC-2026-0249", reason = "smartstring unmaintained via polars; no fixed upstream release" }, + # crypto-bigint < 0.7.6: timing side-channels in constant-time ops, + # incl. subtraction, zero counting, division and Montgomery multiplication. + # Reached only through aws-sdk-s3 sigv4a -> aws-sigv4 (p256/ecdsa); signed + # requests carry the SigV4a credential-derived key. The patched 0.7.6 was + # published 2026-10-23, but no published aws-sigv4 release uses it yet + # (latest 1.6.3 pins crypto-bigint ^0.5), so clearing requires smithy-rs to + # ship an updated aws-sigv4. + # https://rustsec.org/advisories/RUSTSEC-2026-0387 + { id = "RUSTSEC-2026-0387", reason = "crypto-bigint via aws-sigv4 (sigv4a); no fixed upstream aws-sigv4 release yet" }, ] + # --------------------------------------------------------------------------- # Licenses: only allow licenses we've reviewed as compatible with Apache-2.0. # --------------------------------------------------------------------------- @@ -167,7 +177,9 @@ wildcards = "deny" # `[workspace.dependencies]` table: any crate used by more than one member must # go through `workspace = true`, and entries nothing uses are an error. This # keeps versions from drifting between the core crate and the bindings. -workspace-dependencies = { duplicates = "deny", unused = "deny" } +# Python bundles a patched path crate while the Rust SDK uses the overridable +# upstream source. Check shared external declarations, not these path bindings. +workspace-dependencies = { duplicates = "deny", unused = "deny", include-path-dependencies = false } # Internal workspace crates reference each other via `path = "..."`, which # cargo-deny sees as a wildcard version. That's fine for private workspace # members (not published to crates.io), so allow it specifically for paths. diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 0c887f66f..9dbccbc81 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.41.0-beta.13 + 0.41.0-beta.14 ``` diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index d5ab8c3d1..c110eb38c 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -653,94 +653,48 @@ See [Listing](Listing.md) for pagination state and error behavior. ### listNamespaces() ```ts -abstract listNamespaces(namespacePath?, options?): Promise +abstract listNamespaces(namespacePath?, options?): Listing ``` -List the immediate child namespaces under the given parent. - -Results may be paginated. To retrieve subsequent pages, pass the -`pageToken` returned by a previous call. +Lazily list immediate child namespaces, 100 per page by default. #### Parameters * **namespacePath?**: `string`[] - The parent namespace path. Defaults - to the root namespace if omitted. -* **options?**: `Partial`<[`ListNamespacesOptions`](../interfaces/ListNamespacesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) #### Returns -`Promise`<[`ListNamespacesResponse`](../interfaces/ListNamespacesResponse.md)> - -Child namespace names and - an optional token for fetching the next page. +[`Listing`](Listing.md)<`string`> *** ### listTables() -#### listTables(options) - ```ts -abstract listTables(options?): Promise +abstract listTables(namespacePath?, options?): Listing ``` -List a page of the tables in this database. +Lazily list tables in a namespace, requesting 100 per page by default. -To retrieve the tables after the page, pass the `pageToken` the response -carries back in. A page can be shorter than `limit` without being the last -one, so walk until a response carries no page token: - -```ts -const names = []; -let pageToken = undefined; -do { - const page = await conn.listTables({ pageToken, limit: 100 }); - names.push(...page.tables); - pageToken = page.pageToken; -} while (pageToken); -``` - -##### Parameters - -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). - -##### Returns - -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> - -A page of table names and an - optional token for the tables after it. - -#### listTables(namespacePath, options) - -```ts -abstract listTables(namespacePath?, options?): Promise -``` - -List a page of the tables in this database. - -##### Parameters +#### Parameters * **namespacePath?**: `string`[] - The namespace path to list tables from - (defaults to root namespace) -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) -##### Returns +#### Returns -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> +[`Listing`](Listing.md)<`string`> -A page of table names and an - optional token for the tables after it. +#### Example + +```ts +for await (const name of conn.listTables([], { pageLimit: 20 })) { + console.log(name); +} +``` *** diff --git a/docs/src/js/classes/Job.md b/docs/src/js/classes/Job.md index cfbe454ec..74e1c00e2 100644 --- a/docs/src/js/classes/Job.md +++ b/docs/src/js/classes/Job.md @@ -17,7 +17,8 @@ extra round trip on every call: - [Job.refresh](Job.md#refresh) and [Job.status](Job.md#status) fetch the whole record. - [Job.wait](Job.md#wait) records the terminal state it establishes, but not the rest of the record. -- Everything is null until one of those runs. +- Registry details are null until one of those runs. Query handles report + jobType "query" immediately. ## Accessors @@ -57,8 +58,10 @@ get id(): null | string Identifies the operation on the server that is running it. -Operations that run in this process have no server id. The value is -opaque: parsing it or storing it to resume the job later is not supported. +Registry handles expose a registry job ID; query handles expose a +connection-scoped SQL query UUID string. Local handles have no ID. +Query handles report jobType "query". Their IDs are not registry IDs or +durable resume tokens. #### Returns @@ -72,8 +75,8 @@ opaque: parsing it or storing it to resume the job later is not supported. get jobType(): null | string ``` -The job's type, as the server names it. Null for an in-process job, which -has no server-side record. +"query" immediately for SQL-backed handles; otherwise the cached +server-reported job type. Null for in-process handles. #### Returns diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 646341190..e3830e9ad 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -108,10 +108,6 @@ - [JobEventsOptions](interfaces/JobEventsOptions.md) - [JobFailureInfo](interfaces/JobFailureInfo.md) - [JobInfo](interfaces/JobInfo.md) -- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) -- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) -- [ListTablesOptions](interfaces/ListTablesOptions.md) -- [ListTablesResponse](interfaces/ListTablesResponse.md) - [ListingOptions](interfaces/ListingOptions.md) - [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md) - [MemtableStats](interfaces/MemtableStats.md) diff --git a/docs/src/js/interfaces/IvfRqOptions.md b/docs/src/js/interfaces/IvfRqOptions.md index 6b48cf653..298565a62 100644 --- a/docs/src/js/interfaces/IvfRqOptions.md +++ b/docs/src/js/interfaces/IvfRqOptions.md @@ -58,7 +58,8 @@ Number of bits per dimension for residual quantization. This value controls how much each residual component is compressed. The more bits, the more accurate the index will be but the slower search. Typical values -are small integers; the default is 1 bit per dimension. +are small integers. If unset, uses Lance's default (currently 5 bits per +dimension). Set to 1 for maximum compression. *** diff --git a/docs/src/js/interfaces/ListNamespacesOptions.md b/docs/src/js/interfaces/ListNamespacesOptions.md deleted file mode 100644 index 79680417f..000000000 --- a/docs/src/js/interfaces/ListNamespacesOptions.md +++ /dev/null @@ -1,27 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesOptions - -# Interface: ListNamespacesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An optional limit to the number of results to return. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response for pagination. diff --git a/docs/src/js/interfaces/ListNamespacesResponse.md b/docs/src/js/interfaces/ListNamespacesResponse.md deleted file mode 100644 index a99ab480f..000000000 --- a/docs/src/js/interfaces/ListNamespacesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesResponse - -# Interface: ListNamespacesResponse - -## Properties - -### namespaces - -```ts -namespaces: string[]; -``` - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` diff --git a/docs/src/js/interfaces/ListTablesOptions.md b/docs/src/js/interfaces/ListTablesOptions.md deleted file mode 100644 index b8dd8b5ef..000000000 --- a/docs/src/js/interfaces/ListTablesOptions.md +++ /dev/null @@ -1,35 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesOptions - -# Interface: ListTablesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An upper bound on how many tables to return. -Defaults to 100. Zero returns an empty page without a continuation token. - -A page may hold fewer than this and still not be the last one, so keep -going while the response carries a page token rather than while pages are -full. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response, to resume listing where it left off. - -The token is opaque: it carries whatever the database needs to resume, and -callers should not construct or interpret one. diff --git a/docs/src/js/interfaces/ListTablesResponse.md b/docs/src/js/interfaces/ListTablesResponse.md deleted file mode 100644 index 76cac2b23..000000000 --- a/docs/src/js/interfaces/ListTablesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesResponse - -# Interface: ListTablesResponse - -## Properties - -### pageToken? - -```ts -optional pageToken: string; -``` - -*** - -### tables - -```ts -tables: string[]; -``` diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index c0b33dcf0..f18ada7aa 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -20,13 +20,14 @@ How to return blob v2 columns. Applies to `toArrow()` and `toArray()`. - `"descriptions"` (default): each blob is a descriptor. Read the bytes with [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles). -- `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - including blobs nested in structs. The query also reads `_rowid` to - fetch the bytes, and leaves it out of the result unless - [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was called. All the bytes are held in - memory. +- `"bytes"`: each blob column, including blobs nested in structs, holds + the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + leaves it out of the result unless [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was + called. All the bytes are held in memory. -Blobs nested in a list keep their descriptors. +Selecting a parent struct or directly aliasing a blob or struct also returns +bytes. Blobs inside lists keep their descriptors. Arbitrary projection +expressions are not materialized as bytes. The rows and the bytes come from the same table version. If the table changes while the query runs, the query runs again, up to three times. diff --git a/docs/src/js/interfaces/TableNamesOptions.md b/docs/src/js/interfaces/TableNamesOptions.md index 123aa6c9f..e1fd8c382 100644 --- a/docs/src/js/interfaces/TableNamesOptions.md +++ b/docs/src/js/interfaces/TableNamesOptions.md @@ -8,7 +8,7 @@ ## Deprecated -Use [ListTablesOptions](ListTablesOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) +Use [ListingOptions](ListingOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) instead. ## Properties diff --git a/docs/src/python/python.md b/docs/src/python/python.md index f4c697497..356b49ae7 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -20,20 +20,34 @@ pip install lancedb The following methods describe the synchronous API client. There is also an [asynchronous API client](#connections-asynchronous). +Connections to `db://` databases reject local-only options with +`NotImplementedError`. These include `session` and `manifest_enabled=True` on +connect, per-table storage and write options on create, and `index_cache_size`, +`storage_options`, or `location` on open (where the method accepts them). +Omit these options to use server-managed storage and caching. Connection-level +`storage_options` supported by an Enterprise deployment remain available. + +The synchronous remote `create_table` uses the same positional order as the +local API: `name, data, schema, mode, exist_ok, on_bad_vectors, fill_value, +embedding_functions`. Prefer keyword arguments for optional parameters. + ## Connections (Synchronous) ::: lancedb.connect ::: lancedb.db.DBConnection -Resource listings (`list_secrets`, `list_views`, `list_jobs`, `list_functions`, and -`list_materialized_views`) return lazy iterators. Use `list(db.list_views())` to -collect synchronous results, or `[name async for name in db.list_views()]` with -an asynchronous connection. `page_limit` controls each request, and `page_token` -resumes a listing. Drain the iterator's cached results before saving its token. +Resource listings return iterators. Use `list(db.list_tables())` to collect all +table names. `db.list_tables().tables` is a list of the first page; `limit` +controls its page size and `page_token` resumes at a saved token. The same +patterns apply to `list_namespaces()`. Other resource listings use `page_limit`. ::: lancedb.listing.Listing +::: lancedb.listing.TableListing + +::: lancedb.listing.NamespaceListing + ::: lancedb.Session ## Catalogs (Synchronous) @@ -218,10 +232,19 @@ The connected database and `default_namespace_path=["public"]` are used for unqualified tables. Fully qualified references can still query other databases and namespaces available to the same deployment. `execute_query` returns a reader as soon as its initial result stream is available. `execute_query_async` -returns a query handle immediately; use it to inspect progress, open a reader, +returns a query handle once submission is accepted; use it to inspect progress, open a reader, or cancel the query. The SQL client is initialized by the first query and retained for the lifetime of the remote connection. Query ids are random, -connection-scoped references rather than encoded SQL or durable resume tokens: +connection-scoped references rather than encoded SQL or durable resume tokens. + +Both APIs start with one Flight call. The server can stream results on that +call or return a polling continuation that the SDK follows automatically. +This behavior is the same with and without parameters. Open the reader directly; +a directly streamed query needs its reader to make progress. Dropping an unread +handle or its reader stops direct streaming. Deferred execution continues until +explicitly cancelled. + +For example: ```python import lancedb @@ -285,6 +308,25 @@ listing a storage directory. ## Tables (Synchronous) +Local and remote tables use the same arguments for index creation, deletion, +index statistics, and search. For example, `table.delete(where="id = 1")`, +`table.create_index("cosine", 2)` (the deprecated vector API), and +`table.search([1.0, 1.0], fast_search=True)` work with either backend. +`fast_search=True` excludes rows that have not been indexed. + +`create_scalar_index` defaults to `replace=True` on both backends. Both scalar +and full-text index creation accept `wait_timeout`. Remote maintenance methods +accept the local options: cleanup and compaction warn that they are no-ops, +while `optimize` reports an unsupported operation. Removed Tantivy options +produce the same explanatory errors on both backends. + +The old remote keyword names `predicate` for `delete`, `index_name` for +`drop_index`, `index_uuid` for `index_stats`, and `column` for `create_fts_index` +remain available with deprecation warnings. Use `where`, `name`, `index_name`, +and `field_names`, respectively. For positional calls, use the argument order +documented by `Table` below; specify index options by keyword when migrating +code that used the former remote-only argument order. + ::: lancedb.table.Table ::: lancedb.table.FragmentStatistics @@ -301,6 +343,14 @@ listing a storage directory. ## Functions and Jobs +Annotate a blob input with `lancedb.BlobReader` for lazy, seekable access, or +with `bytes` to materialize its contents before the Function runs. Reader +inputs infer the Blob V2 type without an explicit input schema. Readers are +scoped to each invocation and support `read`, `readinto`, `seek`, `tell`, +`size`, and cursor-independent `read_range`. + +::: lancedb.BlobReader + ::: lancedb.functions.FunctionArtifact ::: lancedb.functions.FunctionParameter @@ -518,12 +568,64 @@ instead of being materialized with the rest of the row. `lancedb.BlobType` is `lance.blob.BlobType` when pylance is installed. Without pylance, LanceDB uses a matching `lance.blob.v2` extension type so blob columns -still work. Queries return descriptors. Call +still work. Native table queries return descriptors. Call [`fetch_blob_files`][lancedb.table.Table.fetch_blob_files] for lazy reads or [`fetch_blobs`][lancedb.table.Table.fetch_blobs] for eager bytes. One missing or deleted row ID rejects the entire blob-read batch; the APIs do not return partial results. +Remote SQL returns blob bytes by default. Select `blob_descriptor(video)` to +return an openable reference instead, then use +[`open_blob_file`][lancedb.table.Table.open_blob_file] or the lazy iterable +[`open_blob_files`][lancedb.table.Table.open_blob_files]: + +```python +rows = db.execute_query(""" + SELECT f.timestamp, blob_descriptor(v.video) AS video + FROM videos v JOIN frames f ON v.id = f.video_id +""").read_all() +videos = db.open_table("videos") +for row in rows.to_pylist(): + if row["video"] is not None: + with videos.open_blob_file(row["video"]) as video: + frame = extract_frame(video, row["timestamp"]) +``` + +This example uses a remote SQL connection that produces source-bearing +descriptors. The descriptor contains `kind`, `position`, `size`, `blob_id`, +`blob_uri`, and a `source` struct with `uri`, `version`, `committed_at_micros`, +`branch`, `column`, `row_reference`, and `row_is_address`. The opener verifies +the source table and commit and reads the captured version without changing +the table handle's checkout. Raw native storage descriptors omit `source`; +use `fetch_blob_files(column, row_ids)` for those query results. + +The singular API accepts a dictionary or Arrow struct scalar. The iterable +accepts Arrow arrays/chunked arrays or another iterable, opens handles as +iteration advances, and preserves order, duplicates, and nulls as `None`. +Close each non-null handle after use. Async tables provide +`await table.open_blob_file(...)` and +`async for file in table.open_blob_files(...)`. + +```python +for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): + if video is not None: + with video: + frame = extract_frame(video, row["timestamp"]) +``` + +For an async table, the iterable is consumed with `async for`: + +```python +async for video in async_videos.open_blob_files(rows["video"]): + if video is not None: + with video: + payload = await video.aread() +``` + +Payload bytes load on reads. Opening an external blob with unspecified size +can issue a metadata-only HEAD request. Remote reads use authenticated range +requests that validate the captured URI and commit timestamp. + Blobs inside lists can be written and queried as descriptors. They are omitted from `blob_columns()` because row ids do not identify individual list elements, so the fetch APIs cannot read their bytes yet. @@ -631,6 +733,16 @@ can be used to create, list, or open tables. ::: lancedb.listing.AsyncListing +`async for name in db.list_tables()` loads pages as needed. To inspect the +first page, use `page = await db.list_tables()` and then `page.tables`. Reading +the page list does not consume the iterator. `page.page_token` is the token for +the next request; drain cached results before using it to resume in a new +listing. The same behavior applies to namespaces. + +::: lancedb.listing.AsyncTableListing + +::: lancedb.listing.AsyncNamespaceListing + ## Namespaces (Asynchronous) ::: lancedb.connect_namespace_async diff --git a/docs/web-tests/ts/connection.test.ts b/docs/web-tests/ts/connection.test.ts index 891501d4c..651be7d36 100644 --- a/docs/web-tests/ts/connection.test.ts +++ b/docs/web-tests/ts/connection.test.ts @@ -4,6 +4,12 @@ import { afterEach, describe, expect, jest, test } from "@jest/globals"; import * as lancedb from "@lancedb/lancedb"; import { withTempDirectory } from "./util.ts"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + test("connect to a local database", async () => { await withTempDirectory(async (tempDir) => { const cwd = process.cwd(); @@ -13,12 +19,12 @@ test("connect to a local database", async () => { const db = await lancedb.connect("ex_lancedb"); // --8<-- [end:connect] await db.createTable("probe", [{ id: 1 }]); - expect((await db.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(db.listTables())).toEqual(["probe"]); } finally { process.chdir(cwd); } const reopened = await lancedb.connect(`${tempDir}/ex_lancedb`); - expect((await reopened.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(reopened.listTables())).toEqual(["probe"]); }); }); @@ -109,8 +115,8 @@ async function namespaceTableOpsExample() { ); // Verify - console.log((await db.listNamespaces()).namespaces); // ["prod"] - console.log((await db.listNamespaces(["prod"])).namespaces); // ["recommendations", "search"] + console.log(await collectNames(db.listNamespaces())); // ["prod"] + console.log(await collectNames(db.listNamespaces(["prod"]))); // ["recommendations", "search"] console.log(await db.tableNames(["prod", "search"])); // ["user"] console.log(await db.tableNames(["prod", "recommendations"])); // ["user"] // --8<-- [end:namespace_table_ops] @@ -124,7 +130,7 @@ async function namespaceAdminOpsExample() { await db.createNamespace(["prod"]); await db.createNamespace(["prod", "search"]); - const childNamespaces = (await db.listNamespaces(["prod"])).namespaces; + const childNamespaces = await collectNames(db.listNamespaces(["prod"])); console.log(`Child namespaces under ${JSON.stringify(namespace)}:`, childNamespaces); // Child namespaces under ["prod","search"]: [ 'search' ] diff --git a/docs/web-tests/ts/quickstart.test.ts b/docs/web-tests/ts/quickstart.test.ts index e32a14f86..2c3ee8252 100644 --- a/docs/web-tests/ts/quickstart.test.ts +++ b/docs/web-tests/ts/quickstart.test.ts @@ -20,12 +20,7 @@ const connections = async function tableNames(db: lancedb.Connection) { const names: string[] = []; - let pageToken: string | undefined; - do { - const page = await db.listTables({ pageToken }); - names.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); + for await (const name of db.listTables()) names.push(name); return names; } diff --git a/docs/web/basics/jobs.mdx b/docs/web/basics/jobs.mdx index 0577cd120..a87663a51 100644 --- a/docs/web/basics/jobs.mdx +++ b/docs/web/basics/jobs.mdx @@ -20,19 +20,20 @@ failures. Methods such as Python's `create_index_async` and TypeScript's `createIndexAsync` return job handles. In OSS, the work runs in your application -process; in Enterprise, it runs on server workers. A returned handle can -represent an operation that has already finished, so check its status before +process; in Enterprise, it runs on server workers or through the SQL query +backend. A returned handle can represent an operation that has already finished, +so check its status before waiting or requesting cancellation. See the operation's API reference for when the method returns. ## Submit and wait for a job {#waiting} 1. **Submit the operation.** Keep the returned job handle. For an Enterprise - job, also keep its ID if another process needs to look it up later. + registry job, also keep its ID if another process needs to look it up later. 2. **Check progress.** Read the current status while the operation runs. 3. **Wait for completion.** Check that the job succeeded before using its - result. If it failed or was cancelled, inspect the failure details and any - results it committed before stopping. + result. If it failed or was cancelled, inspect the failure details and any + results it committed before stopping. In Python, `job.status()` reports the current state and `job.wait()` waits for completion, raising if the job fails or is cancelled. If `job.wait()` times out, the client stops waiting. @@ -59,13 +60,37 @@ history. ## OSS and Enterprise execution {#in-process} -| | OSS | Enterprise | +The following table describes local operations and Enterprise registry jobs. +SQL-backed job handles follow the query lifecycle described below. + +| | OSS | Enterprise registry job | | --- | --- | --- | | Execution | In your application process | On server workers | | Job identity | No server ID | A server ID for looking up the job | | Event history | No server event history | Available through the jobs API | | Reattaching after a client restart | No server job to reopen | Reopen by ID while the record is retained | +### SQL-backed job handles + +Remote materialized-view CREATE and REFRESH use the SQL query backend. Python's +`create_materialized_view_async()` and `mv.refresh_async()` still return job +handles with `status()`, `wait()` and `cancel()`. Synchronous creation and refresh +wait on the same SQL execution before returning. + +The existing `job.job_type` identifies query-backed handles as `"query"` from +submission, and `job.id` contains the connection-scoped SQL query UUID string. +Use `describe_query(job.id)` on the submitting connection to inspect the query. +Registry handles retain their server-reported job types (available after +`status()` or `refresh()`) and registry IDs. Local handles have no ID or job type. + +SQL-backed handles do not expose registry records or event history and cannot +be reopened through `open_job()` or reattached after a client restart. Their IDs +are not durable resume tokens. Cancellation follows the SQL server's +capabilities and can be unsupported. + +The server decides whether SQL executes on its serving fleet or starts SQL job +workers. That placement does not change the handle's query-backed identity. + ## Learn more {#next} - [Build and repair indexes](/indexing/reindexing): submit index work. diff --git a/docs/web/basics/sdk.mdx b/docs/web/basics/sdk.mdx index fc386bc7a..b20fc1662 100644 --- a/docs/web/basics/sdk.mdx +++ b/docs/web/basics/sdk.mdx @@ -92,6 +92,13 @@ manages background maintenance. [background maintenance](/maintenance/background-jobs) rather than calling the embedded table's `optimize()` method. +{/* Draft for the remote option validation release. +- **Remove local storage and cache settings.** `db://` connections reject local + sessions, table cache settings, and per-table storage and write options. Omit + these options so Enterprise manages storage and caching. Connection-level + storage options supported by your Enterprise deployment remain available. +*/} + For SQL connection details, see [SQL](/basics/sql). For the operating model and deployment options, see [LanceDB Enterprise](/enterprise). diff --git a/docs/web/basics/sql.mdx b/docs/web/basics/sql.mdx index 1c4ad46cd..c89bd7a20 100644 --- a/docs/web/basics/sql.mdx +++ b/docs/web/basics/sql.mdx @@ -82,6 +82,8 @@ In application code, load credentials from environment variables or a secrets manager. The [connection guide](/basics/connect#enterprise) shows that setup. `execute_query()` returns a blocking Arrow reader when results start arriving. +The server can stream rows on the first call or return result tickets that the +SDK follows automatically, with the same behavior for parameterized queries. The example uses `read_all()` to collect the small result in memory. For larger results, consume the reader in batches. To submit a query and check its status separately, use `execute_query_async()`. See [Queries](/basics/query) for query diff --git a/docs/web/docs.json b/docs/web/docs.json index 9c9468a3b..54203f1a0 100644 --- a/docs/web/docs.json +++ b/docs/web/docs.json @@ -350,6 +350,7 @@ "pages": [ "enterprise/index", "enterprise/architecture", + "enterprise/distributed-search", "enterprise/deployment/index", "deployment/storage", "deployment/monitoring", diff --git a/docs/web/enterprise/architecture.mdx b/docs/web/enterprise/architecture.mdx index c47dcbd2e..8f887c142 100644 --- a/docs/web/enterprise/architecture.mdx +++ b/docs/web/enterprise/architecture.mdx @@ -18,3 +18,6 @@ run on machines that can be replaced without any of them owning the data. [Architecture](/architecture) covers this alongside the embedded form, which is the comparison most readers want first. + +See [Distributed search](/enterprise/distributed-search) for how Enterprise +distributes index execution and table-data reads. diff --git a/docs/web/enterprise/distributed-search.mdx b/docs/web/enterprise/distributed-search.mdx new file mode 100644 index 000000000..ebac960ba --- /dev/null +++ b/docs/web/enterprise/distributed-search.mdx @@ -0,0 +1,33 @@ +--- +title: "Distributed search" +sidebarTitle: "Distributed search" +description: "How LanceDB Enterprise distributes index execution and table-data reads." +icon: "network-wired" +availability: + oss: unavailable + enterprise: available + summary: >- + LanceDB Enterprise distributes index execution and table-data reads across + plan executors. Embedded LanceDB executes searches in the application process. +--- + +Distributed search spreads index execution, table-data reads, and caching across +multiple serving machines in LanceDB Enterprise. Query nodes plan queries and +combine results. Plan executors cache index structures and table data and execute +the remote work. See [Enterprise architecture](/enterprise/architecture). + +## Embedded search {#embedded-search} + +LanceDB OSS runs inside your application process. It searches indexes and reads +table data using that process's resources, including its index cache. Reading a +table from object storage does not by itself distribute query execution across +machines. + +## Search interfaces {#search-interfaces} + +The shared search guides describe the query APIs and index settings: + +- [Vector search](/search/vector-search) and [vector indexes](/indexing/vector-index). +- [Full-text search](/search/full-text-search) and [FTS indexes](/indexing/fts-index). +- [Filtering](/search/filtering) and [scalar indexes](/indexing/scalar-index). +- [Query plans](/search/explain) for inspecting execution and runtime measurements. diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index 2742a992f..037d75f8d 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -58,7 +58,7 @@ Lance with `index_type`. | If your priority is... | Use this index | Why | Approx. compression ratio | Python config class | | :--- | :--- | :--- | :--- | :--- | -| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size | `IvfRq` | +| Maximum compression | `IVF_RQ` | IVF indexing with RaBitQ quantization | Around `1/32` of raw size with `num_bits=1` | `IvfRq` | | Higher accuracy at small dimensions (`dimension <= 256`) | `IVF_PQ` | IVF indexing with product quantization | Usually `1/64` to `1/16` of raw size (depends on `num_sub_vectors`) | `IvfPq` | | Best recall/latency trade-off | `IVF_HNSW_SQ` | IVF-HNSW indexing with scalar | Typically a little larger than `1/4` of raw size | `IvfHnswSq` | | Highest recall / no quantization | `IVF_HNSW_FLAT` | IVF-HNSW indexing with no quantization; embedded library only | Around raw vector size plus HNSW graph overhead | `IvfHnswFlat` | @@ -377,22 +377,27 @@ which has been compressed to $4$ chunks $\times 8$-bit integers $= 32$ bits of q RaBitQ is an advanced quantization technique that outperforms PQ in several ways. It needs no codebook to train, estimates distances very quickly at query-time, -and crucially, quantizes each vector in (with some small overhead) just **one bit per dimension!** +and can quantize each vector in (with some small overhead) as little as **one bit per dimension!** In practice, RaBitQ compresses a $1024$-dimensional `float32` vector into just a few thousand bits, while maintaining good recall. The inner workings of RaBitQ quantization are rather mathematically dense. It generates a quantization codebook by applying a uniformly random, approximately distance-preserving orthogonal transformation of the vertices of the $d$-dimensional hypercube, where $d$ is the dimensionality of the dataset. We defer the details, and an elegant theoretical error bound, to [the original paper.](https://arxiv.org/pdf/2405.12497) -##### Using RaBitQ -Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, call `create_index` with `index_type="IVF_RQ"`. +##### Using RaBitQ {#using-rabitq} +Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, import `IvfRq` from `lancedb.index` and call `table.create_index("vector", config=IvfRq())`. Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`. `num_bits` determines how many bits are used to quantize each dimension. -`1` is the standard RaBitQ setting. Increase to `2`, `4`, or `8` bits to achieve better recall for additional storage and query-time compute. +By default, Python, Rust, and Node leave this unset and inherit Lance's default, currently `5` bits. +Higher values improve recall at the cost of additional storage and query-time compute. +To retain the previous one-bit default and maximum compression, explicitly set `num_bits=1` (for example, `IvfRq(num_bits=1)` in Python). +Five-bit quantized vector codes use roughly five times the storage of one-bit codes; total index size also includes other overhead. -RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout, and cannot be read by some older LanceDB versions. +RaBitQ-quantized indexes computed with `num_bits >= 2` use a newer on-disk layout and cannot be read by some older LanceDB versions. +The new five-bit default therefore affects reader compatibility as well as index size. +Upgrade readers to versions that support multi-bit IVF_RQ before creating indexes with the default settings, or explicitly set `num_bits=1` when sharing indexes with readers that support only one-bit IVF_RQ. See this [blog post](https://lancedb.com/blog/feature-rabitq-quantization/) for further discussion and benchmarking of LanceDB's RaBitQ implementation. @@ -410,7 +415,7 @@ among all integral multiples $\{j \times 10^5: -128 \leq j \leq 127 \}$. | Parameter | Description | | :--- | :--- | -| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (default `1`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | +| `num_bits` | Bits per dimension for quantization. Only applies to `IVF_PQ`/`IVF_HNSW_PQ` (default `8`) and `IVF_RQ` (`IvfRq()` inherits Lance's default, currently `5`, RaBitQ); not used by `IVF_FLAT`/`IVF_SQ`/`IVF_HNSW_FLAT`/`IVF_HNSW_SQ`. Higher values improve accuracy at the cost of additional storage. | `max_iterations` and `sample_rate` also affect quantizer training, but since they apply to every IVF/HNSW index type (not just quantized ones), they're documented as general [Build-time Parameters](/indexing/vector-index#build-time-parameters) instead. All three are passed as keyword arguments to `create_index`, alongside `index_type`: diff --git a/docs/web/ingestion/create-table.mdx b/docs/web/ingestion/create-table.mdx index 7d4363be7..0921de22e 100644 --- a/docs/web/ingestion/create-table.mdx +++ b/docs/web/ingestion/create-table.mdx @@ -148,9 +148,11 @@ non-nullable column, or provide actual nulls for it, ingestion fails; nullable c omitted or written with null values. Without an explicit schema, Python infers list-like vector values as fixed-size `float32` vector fields from the observed dimension. -For Python ingest, malformed vector values fail by default. If you expect occasional wrong-length, -null, or NaN vectors, choose an `on_bad_vectors` policy: `"drop"` removes those rows, `"fill"` writes -`fill_value`, and `"null"` writes nulls. +For Python ingestion, malformed vectors fail by default (`on_bad_vectors="error"`). This includes +vectors with incorrect sizes, NaNs, or null components, including pandas NaNs converted to Arrow +nulls. Use `"drop"` to remove affected rows, or `"null"` to replace affected vectors with nulls. +Use `"fill"` to replace NaN or null components with `fill_value` while preserving valid components. +This also pads short vectors with `fill_value` and truncates long vectors to the expected dimension. ### From an Arrow Table {#from-an-arrow-table} You can also create LanceDB tables directly from Arrow tables. diff --git a/docs/web/integrations/reranking/openai.mdx b/docs/web/integrations/reranking/openai.mdx index 6ceb8a030..71fe63117 100644 --- a/docs/web/integrations/reranking/openai.mdx +++ b/docs/web/integrations/reranking/openai.mdx @@ -26,7 +26,7 @@ credentials from environment variables or a secrets manager. ## Accepted Arguments {#accepted-arguments} | Argument | Type | Default | Description | | --- | --- | --- | --- | -| `model_name` | `str` | `"gpt-4-turbo-preview"` | The name of the reranker model to use.| +| `model_name` | `str` | `"gpt-5.6-sol"` | The name of the reranker model to use.| | `column` | `str` | `"text"` | The name of the column to use as input to the cross encoder model. | | `return_score` | `str` | `"relevance"` | Options are "relevance" or "all". The type of score to return. If "relevance", will return only the `_relevance_score. If "all" is supported, will return relevance score along with the vector and/or fts scores depending on query type. | | `api_key` | `str` | `None` | The API key to use. If None, will use the OPENAI_API_KEY environment variable. diff --git a/docs/web/management/function.mdx b/docs/web/management/function.mdx index 8e5b53e4a..6df99b5fd 100644 --- a/docs/web/management/function.mdx +++ b/docs/web/management/function.mdx @@ -205,6 +205,40 @@ Function output blobs follow the same compaction and cleanup lifecycle as ingested Managed blobs. Retained versions continue to own their payloads; cleanup reclaims an object only after its retained references are gone. +### Draft: incremental blob inputs (upcoming) + +This feature requires an SDK that exports `BlobReader` and a compatible +Enterprise runtime with reader input support enabled. Coordinate the runtime +upgrade before registering or executing reader Functions. + +Annotate a row Function parameter with `BlobReader` to read a Blob v2 input +incrementally. The signature selects the access mode without another +registration option. For example, read the first MP4 box type without loading +the whole video: + +```python +from lancedb import BlobReader, udf + + +@udf +def first_box_type(video: BlobReader) -> bytes: + video.seek(4) + return video.read(4) +``` + +The worker passes a seekable reader backed by the stored blob. `read(size)` +reads from the current cursor; `read_range(offset, length)` reads a byte range +without moving it. `size()` returns the payload size without reading it. +Each row call receives its own reader, which closes when the call returns or +raises. Do not retain it for another invocation. + +Use `video: bytes` when the Function needs the complete payload. Bytes inputs +materialize before invocation and remain subject to eager input limits. +Reader inputs keep the unread payload out of the input batch, but calling +`video.read()` without a size loads all remaining bytes in Python. Reader +inputs are supported in row mode. Return values remain ordinary values such +as `bytes`; this feature does not add streamed outputs. + ## Declare the environment {#environment} The function runs remotely, so its dependencies are part of its definition: diff --git a/docs/web/quickstart.mdx b/docs/web/quickstart.mdx index 788dd6585..3f466041d 100644 --- a/docs/web/quickstart.mdx +++ b/docs/web/quickstart.mdx @@ -58,6 +58,33 @@ uv add --prerelease allow --index https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra and +include the Lance preview index: + + +```bash Python (pip) icon="terminal" +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +```bash Python (uv) icon="terminal" +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + Pre-release builds receive the same level of testing as stable releases, but their availability is not guaranteed for more than 6 months after release. For real-world workloads, we recommend you use the latest stable release diff --git a/docs/web/search/vector-search.mdx b/docs/web/search/vector-search.mdx index 3519df553..cc6ae5421 100644 --- a/docs/web/search/vector-search.mdx +++ b/docs/web/search/vector-search.mdx @@ -235,7 +235,7 @@ The table below summarizes the behavior of `_distance` in search results based o See [vector indexes](/indexing/vector-index) for index configuration. The following sections cover query-time settings. -### Tuning `approx_mode` +### Tuning `approx_mode` {#tuning-approx-mode} Use `approx_mode` when you want to adjust the speed/recall trade-off for approximate vector search at query time. This setting currently applies only to RQ-quantized indexes, such as `IVF_RQ`; other index types ignore it. diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index 87c5d4cb8..7c34a801a 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.13 + 0.41.0-beta.14 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 1f3e0efb0..2dc959887 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.41.0-beta.13 + 0.41.0-beta.14 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 1b9d47e46..e343d68ca 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index bbd0c3c52..73f99803a 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -5,15 +5,15 @@ import { mkdirSync, readdirSync } from "fs"; import { join } from "path"; import { Field, Float64, Schema } from "apache-arrow"; import * as tmp from "tmp"; -import { - Connection, - ListTablesResponse, - Table, - connect, - connectNamespace, -} from "../lancedb"; +import { Connection, Table, connect, connectNamespace } from "../lancedb"; import { LocalTable } from "../lancedb/table"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + describe("when connecting", () => { let tmpDir: tmp.DirResult; beforeEach(() => { @@ -135,7 +135,7 @@ describe("given a connection", () => { await db.close(); expect(db.isOpen()).toBe(false); await expect(db.tableNames()).rejects.toThrow("Connection is closed"); - await expect(db.listTables()).rejects.toThrow("Connection is closed"); + expect(() => db.listTables()).toThrow("Connection is closed"); await expect(db.renameTable("a", "b")).rejects.toThrow( "Connection is closed", ); @@ -247,15 +247,15 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const first = await db.listTables(namespacePath, { limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables(namespacePath, { pageLimit: 1 }); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables(namespacePath, { - limit: 1, - pageToken: first.pageToken, + const second = db.listTables(namespacePath, { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); }); it("should keep the existing table when mode is exist_ok", async () => { @@ -299,19 +299,18 @@ describe("given a connection", () => { await db.createTable("a", [{ id: 1 }]); await db.createTable("c", [{ id: 1 }]); - const all = await db.listTables(); - expect(all.tables).toEqual(["a", "b", "c"]); - expect(all.pageToken).toBeUndefined(); + expect(await collectNames(db.listTables())).toEqual(["a", "b", "c"]); - const first = await db.listTables({ limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables({ - limit: 1, - pageToken: first.pageToken, + const second = db.listTables([], { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); it("should visit every table exactly once when walking pages", async () => { @@ -322,18 +321,9 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const seen: string[] = []; - let pageToken: string | undefined = undefined; - do { - const page: ListTablesResponse = await db.listTables({ - limit: 2, - pageToken, - }); - seen.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); - - expect(seen).toEqual(created); + expect(await collectNames(db.listTables([], { pageLimit: 2 }))).toEqual( + created, + ); }); it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])( @@ -347,11 +337,10 @@ describe("given a connection", () => { ]) { await expect(names).resolves.toEqual(["a", "b"]); } - for (const page of [ - db.listTables({ limit }), - db.listTables([], { limit }), - ]) { - await expect(page).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + expect(() => + db.listTables(namespacePath, { pageLimit: limit }), + ).toThrow(/limit/); } }, ); @@ -366,8 +355,11 @@ describe("given a connection", () => { ])("should reject an invalid table listing limit %s", async (limit) => { await expect(db.tableNames({ limit })).rejects.toThrow(/limit/); await expect(db.tableNames([], { limit })).rejects.toThrow(/limit/); - await expect(db.listTables({ limit })).rejects.toThrow(/limit/); - await expect(db.listTables([], { limit })).rejects.toThrow(/limit/); + for (const namespacePath of [undefined, []]) { + expect(() => db.listTables(namespacePath, { pageLimit: limit })).toThrow( + /limit/, + ); + } }); it("should preserve zero, small and unsigned-32-bit table listing limits", async () => { @@ -379,11 +371,21 @@ describe("given a connection", () => { [2 ** 32 - 1, ["a", "b"]], ] as const) { await expect(db.tableNames({ limit })).resolves.toEqual(names); - const page = await db.listTables({ limit }); - expect(page.tables).toEqual(names); } await expect(db.tableNames()).resolves.toEqual(["a", "b"]); - await expect(db.listTables()).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + for (const pageLimit of [1, 2 ** 31 - 1]) { + expect( + await collectNames(db.listTables(namespacePath, { pageLimit })), + ).toEqual(["a", "b"]); + } + for (const pageLimit of [0, 2 ** 31, 2 ** 32 - 1]) { + expect(() => db.listTables(namespacePath, { pageLimit })).toThrow( + /limit/, + ); + } + } + expect(await collectNames(db.listTables())).toEqual(["a", "b"]); }); it("should list all names and use a default page size of 100", async () => { @@ -396,20 +398,17 @@ describe("given a connection", () => { } await expect(db.tableNames()).resolves.toEqual(names); await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]); - const first = await db.listTables(); - expect(first.tables).toEqual(names.slice(0, 100)); - expect(first.pageToken).toBeDefined(); - const second = await db.listTables({ pageToken: first.pageToken }); - expect(second.tables).toEqual(names.slice(100)); - expect(second.pageToken).toBeUndefined(); - await expect(db.listTables({ limit: 200 })).resolves.toEqual({ - tables: names, - pageToken: undefined, - }); - await expect(db.listTables({ limit: 0 })).resolves.toEqual({ - tables: [], - pageToken: undefined, - }); + const first = db.listTables(); + const firstPage: string[] = []; + for (let i = 0; i < 100; i++) firstPage.push((await first.next()).value!); + expect(firstPage).toEqual(names.slice(0, 100)); + expect(first.pageToken()).toBeDefined(); + expect(first.numPageResults()).toBe(0); + expect(await collectNames(first)).toEqual(names.slice(100)); + expect(await collectNames(db.listTables([], { pageLimit: 200 }))).toEqual( + names, + ); + expect(() => db.listTables([], { pageLimit: 0 })).toThrow(/limit/); }); it("should list tables in a namespace", async () => { @@ -420,12 +419,8 @@ describe("given a connection", () => { await db.createNamespace(["child"]); await db.createTable("nested", [{ id: 1 }], ["child"]); - await expect(db.listTables(["child"])).resolves.toEqual( - expect.objectContaining({ tables: ["nested"] }), - ); - await expect(db.listTables()).resolves.toEqual( - expect.objectContaining({ tables: [] }), - ); + expect(await collectNames(db.listTables(["child"]))).toEqual(["nested"]); + expect(await collectNames(db.listTables())).toEqual([]); }); it("should create tables in v2 mode", async () => { @@ -719,22 +714,35 @@ describe("namespaces", () => { it("should list namespaces created at the root", async () => { await db.createNamespace(["alpha"]); await db.createNamespace(["beta"]); - const list = await db.listNamespaces(); - expect(list.namespaces).toEqual(expect.arrayContaining(["alpha", "beta"])); + const list = await collectNames(db.listNamespaces()); + expect(list).toEqual(expect.arrayContaining(["alpha", "beta"])); }); it("should list child namespaces under a parent", async () => { await db.createNamespace(["parent"]); await db.createNamespace(["parent", "child"]); - const list = await db.listNamespaces(["parent"]); - expect(list.namespaces).toContain("child"); + const list = await collectNames(db.listNamespaces(["parent"])); + expect(list).toContain("child"); + }); + + it("should resume a namespace listing from its page token", async () => { + await db.createNamespace(["alpha"]); + await db.createNamespace(["beta"]); + const first = db.listNamespaces([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect((await first.next()).value).toBe("alpha"); + const second = db.listNamespaces([], { + pageLimit: 1, + pageToken: first.pageToken(), + }); + expect(await collectNames(second)).toEqual(["beta"]); }); it("should drop a namespace", async () => { await db.createNamespace(["ephemeral"]); await db.dropNamespace(["ephemeral"]); - const list = await db.listNamespaces(); - expect(list.namespaces).not.toContain("ephemeral"); + const list = await collectNames(db.listNamespaces()); + expect(list).not.toContain("ephemeral"); }); it("should raise an error on any namespace op after close", async () => { @@ -742,7 +750,7 @@ describe("namespaces", () => { await expect(db.describeNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); - await expect(db.listNamespaces()).rejects.toThrow("Connection is closed"); + expect(() => db.listNamespaces()).toThrow("Connection is closed"); await expect(db.createNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index 3e47a00e9..3852af717 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -129,6 +129,7 @@ describe("materialized views", () => { const job = await db.dropMaterializedViewAsync("adults"); expect(job.id).toBeNull(); + expect(job.jobType).toBeNull(); await job.wait(); expect(await collect(db.listMaterializedViews())).toEqual([]); }); diff --git a/nodejs/__test__/nested_blob.test.ts b/nodejs/__test__/nested_blob.test.ts new file mode 100644 index 000000000..508b57c77 --- /dev/null +++ b/nodejs/__test__/nested_blob.test.ts @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import * as tmp from "tmp"; +import { type Table, blob, connect, isBlobField } from "../lancedb"; +import { + DataType, + Field, + Int64, + List, + Schema, + Struct, + Utf8, +} from "../lancedb/arrow"; + +describe('nested blobs with blobMode "bytes"', () => { + let tmpDir: tmp.DirResult; + let table: Table; + + const caption = new Field( + "caption", + new Utf8(), + false, + new Map([["label", "caption"]]), + ); + const info = new Field( + "info", + new Struct([ + blob("image"), + caption, + new Field("preview", new Struct([blob("image")]), true), + ]), + true, + new Map([["label", "info"]]), + ); + const schema = new Schema( + [new Field("id", new Int64(), false), info, blob("cover")], + new Map([["source", "nested-blobs"]]), + ); + + beforeEach(async () => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + table = await db.createTable( + "nested_blobs", + [ + { + id: 1n, + info: { + image: Buffer.from("hello"), + caption: "first", + preview: { image: Buffer.from("preview") }, + }, + cover: Buffer.from("cover"), + }, + { id: 2n, info: null, cover: null }, + { + id: 3n, + info: { image: null, caption: "third", preview: null }, + cover: null, + }, + { + id: 4n, + info: { + image: Buffer.alloc(0), + caption: "fourth", + preview: { image: null }, + }, + cover: Buffer.alloc(0), + }, + ], + { schema }, + ); + }); + + afterEach(() => tmpDir.removeCallback()); + + it.each([false, true])( + "reads struct leaves across batches (select parent: %s)", + async (selectParent) => { + const query = table + .query() + .orderBy({ columnName: "id", ascending: true }); + if (selectParent) { + query.select(["info"]); + } + const result = await query.toArrow({ + blobMode: "bytes", + maxBatchLength: 1, + }); + expect(result.batches.length).toBeGreaterThan(1); + expect(result.schema.fields.map((field) => field.name)).toEqual( + selectParent ? ["info"] : ["id", "info", "cover"], + ); + expect(result.schema.metadata).toEqual(schema.metadata); + const infoField = result.schema.fields.find( + (field) => field.name === "info", + )!; + expect(infoField.nullable).toBe(true); + expect(infoField.metadata).toEqual(info.metadata); + const fields = (infoField.type as Struct).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[0].metadata.size).toBe(0); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + + const rows = result.toArray(); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(Buffer.from(rows[0].info.preview.image).toString()).toBe( + "preview", + ); + expect(rows[1].info).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(rows[2].info.caption).toBe("third"); + expect(rows[2].info.preview).toBeNull(); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + expect(rows[3].info.preview.image).toBeNull(); + if (!selectParent) { + expect(Buffer.from(rows[0].cover).toString()).toBe("cover"); + expect(rows[1].cover).toBeNull(); + } + }, + ); + + it("reads direct leaf and struct aliases", async () => { + const rows = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select({ + image: "info.image", + details: "info", + preview: "info.preview", + caption: "info.caption", + }) + .toArray({ blobMode: "bytes", maxBatchLength: 1 }); + expect(Object.keys(rows[0])).toEqual([ + "image", + "details", + "preview", + "caption", + ]); + expect(Buffer.from(rows[0].image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.preview.image).toString()).toBe( + "preview", + ); + expect(Buffer.from(rows[0].preview.image).toString()).toBe("preview"); + expect(rows[0].caption).toBe("first"); + expect(rows[1].details).toBeNull(); + expect(rows[1].image).toBeNull(); + expect(rows[2].image).toBeNull(); + expect(Buffer.from(rows[3].image)).toEqual(Buffer.alloc(0)); + }); + + it.each([false, true])( + "preserves a literal scalar alias matching a nested blob path (reverse: %s)", + async (reverse) => { + const projection: [string, string][] = [ + ["info", "info"], + ["info.image", "id"], + ]; + if (reverse) { + projection.reverse(); + } + const [row] = await table + .query() + .where("id = 1") + .select(new Map(projection)) + .toArray({ blobMode: "bytes" }); + expect(row["info.image"]).toBe(1n); + expect(Buffer.from(row.info.image).toString()).toBe("hello"); + }, + ); + + it.each([false, true])( + "materializes every overlapping blob projection (reverse: %s)", + async (reverse) => { + const columns = ["info.image", "info"]; + if (reverse) { + columns.reverse(); + } + const result = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select(columns) + .toArrow({ blobMode: "bytes", maxBatchLength: 1 }); + expect(result.schema.fields.map((field) => field.name)).toEqual(columns); + expect(DataType.isLargeBinary(result.getChild("info.image")!.type)).toBe( + true, + ); + expect( + DataType.isLargeBinary( + result.getChild("info")!.getChild("image")!.type, + ), + ).toBe(true); + const rows = result.toArray(); + expect(Buffer.from(rows[0]["info.image"]).toString()).toBe("hello"); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(rows[1]["info.image"]).toBeNull(); + expect(rows[1].info).toBeNull(); + expect(rows[2]["info.image"]).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(Buffer.from(rows[3]["info.image"])).toEqual(Buffer.alloc(0)); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + }, + ); + + it("materializes dotted aliases independently of nested fields", async () => { + const [row] = await table + .query() + .where("id = 1") + .select({ + details: "info", + "details.image": "info.image", + "details.preview": "info.preview", + "details.preview.image": "info.preview.image", + }) + .toArray({ blobMode: "bytes" }); + expect(Buffer.from(row.details.image).toString()).toBe("hello"); + expect(Buffer.from(row["details.image"]).toString()).toBe("hello"); + expect(Buffer.from(row.details.preview.image).toString()).toBe("preview"); + expect(Buffer.from(row["details.preview"].image).toString()).toBe( + "preview", + ); + expect(Buffer.from(row["details.preview.image"]).toString()).toBe( + "preview", + ); + }); + + it.each([false, true])( + "converts the empty result schema (alias: %s)", + async (alias) => { + const query = table.query().where("id < 0"); + if (alias) { + query.select({ image: "info.image", details: "info" }); + } + const result = await query.toArrow({ blobMode: "bytes" }); + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual( + alias ? ["image", "details"] : ["id", "info", "cover"], + ); + const fields = ( + result.schema.fields.find( + (field) => field.name === (alias ? "details" : "info"), + )!.type as Struct + ).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + if (alias) { + expect(DataType.isLargeBinary(result.schema.fields[0].type)).toBe(true); + } + }, + ); + + it("preserves requested row IDs and leaves subsequent descriptors intact", async () => { + const query = table.query().where("id = 1").select(["info"]); + const [bytes] = await query.toArray({ blobMode: "bytes" }); + expect(Object.keys(bytes)).toEqual(["info"]); + const [description] = await query.toArray(); + expect(Object.keys(description)).toEqual(["info"]); + expect(Number(description.info.image.size)).toBe(5); + const [withRowId] = await query.withRowId().toArray({ blobMode: "bytes" }); + expect(typeof withRowId._rowid).toBe("bigint"); + expect(await table.fetchBlobs("info.image", [withRowId._rowid])).toEqual([ + Buffer.from(withRowId.info.image), + ]); + }); + + it("keeps descriptors in description mode", async () => { + const [description] = await table + .query() + .where("id = 1") + .toArray({ blobMode: "descriptions" }); + expect(Number(description.info.image.size)).toBe(5); + }); + + it("keeps list elements as descriptors alongside struct blob bytes", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable( + "list_blobs", + [ + { + info: { image: Buffer.from("hello"), images: [Buffer.from("list")] }, + }, + ], + { + schema: new Schema([ + new Field( + "info", + new Struct([ + blob("image"), + new Field("images", new List(blob("item")), true), + ]), + true, + ), + ]), + }, + ); + const [bytes] = await table.query().toArray({ blobMode: "bytes" }); + expect(Buffer.from(bytes.info.image).toString()).toBe("hello"); + expect(Number(bytes.info.images.get(0).size)).toBe(4); + const [row] = await table + .query() + .select(["info.images"]) + .toArray({ blobMode: "bytes" }); + expect(Number(row["info.images"].get(0).size)).toBe(4); + expect(Object.keys(row)).toEqual(["info.images"]); + const result = await table.query().toArrow({ blobMode: "bytes" }); + const fields = (result.schema.fields[0].type as Struct).children; + const descriptions = await table.query().toArrow(); + expect(fields[1]).toEqual( + (descriptions.schema.fields[0].type as Struct).children[1], + ); + expect(DataType.isStruct((fields[1].type as List).children[0].type)).toBe( + true, + ); + }); +}); diff --git a/nodejs/__test__/query.test.ts b/nodejs/__test__/query.test.ts index 786f04015..068e55c2f 100644 --- a/nodejs/__test__/query.test.ts +++ b/nodejs/__test__/query.test.ts @@ -538,11 +538,17 @@ describe("IVF search effort", () => { const db = await connect(directory.name); const rows = Array.from({ length: 512 }, (_, id) => ({ id, - vector: Array.from( - { length: 16 }, - (_, j) => - (Math.floor(id / 32) === j ? 1 : 0) + Math.sin(id * 17 + j) * 0.001, - ), + // Keep all eight exact neighbors together even if random IVF training + // splits the first cluster. This tests effort, not training recall. + vector: + id < 8 + ? [...vector] + : Array.from( + { length: 16 }, + (_, j) => + (Math.floor(id / 32) === j ? 1 : 0) + + Math.sin(id * 17 + j) * 0.001, + ), })); table = await db.createTable("effort", rows.slice(0, 256)); await table.add(rows.slice(256)); diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 7e9479f94..d443537b0 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -98,6 +98,55 @@ async function withMockDatabase( } describe("remote connection", () => { + it.each([ + { dataStorageVersion: "2.0" }, + { enableV2ManifestPaths: false }, + { storageOptions: { bogusOptionXyz: "1" } }, + ])("rejects local create options %j before sending data", async (options) => { + const requests: string[] = []; + await withMockDatabase( + (req, res) => { + requests.push(req.url ?? ""); + res.writeHead(200).end(); + }, + async (db) => { + await expect( + db.createTable("test", [{ value: 1 }], { + ...options, + mode: "overwrite", + }), + ).rejects.toThrow(/write_options.*remote/); + await expect( + db.createEmptyTable( + "test", + new Schema([new Field("value", new Int64(), true)]), + options, + ), + ).rejects.toThrow(/write_options.*remote/); + }, + ); + expect(requests).toEqual([]); + }); + + it.each([{ indexCacheSize: 5 }, { storageOptions: { x: "y" } }])( + "rejects local open options %j before sending a request", + async (options) => { + const requests: string[] = []; + await withMockDatabase( + (req, res) => { + requests.push(req.url ?? ""); + res.writeHead(200).end(); + }, + async (db) => { + await expect( + db.openTable("test", undefined, options), + ).rejects.toThrow(/not supported for remote/); + }, + ); + expect(requests).toEqual([]); + }, + ); + it.each([false, true])( "normalizes remote blob writes (range target: %s)", async (wide) => { diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 92e3df456..d848676cd 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3314,51 +3314,64 @@ describe("when dealing with blob columns", () => { ).rejects.toThrow(/blobMode must be/); }); - it("reruns the query when the table changes before the version check", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - await update(2n, "new"); - return liveVersion(); - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(2n); - expect(Buffer.from(row.image).toString()).toBe("new"); - }); + it.each([false, true])( + "reruns the query when the table changes before the version check (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + await update(2n, "new"); + return liveVersion(); + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(2n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "new", + ); + }, + ); - it("reads bytes from the query's version when a write lands before the fetch", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - const version = await liveVersion(); - await update(2n, "new"); - return version; - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(1n); - expect(Buffer.from(row.image).toString()).toBe("old"); - }); + it.each([false, true])( + "reads bytes from the query's version when a write lands before the fetch (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + const version = await liveVersion(); + await update(2n, "new"); + return version; + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(1n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "old", + ); + }, + ); - it("fails when the table changes during every attempt", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - let generation = 1n; - native.version = async () => { - generation += 1n; - await update(generation, `generation ${generation}`); - return liveVersion(); - }; - try { - await expect( - table.query().toArray({ blobMode: "bytes" }), - ).rejects.toThrow(/changed during each of 3 attempts/); - expect(generation).toBe(4n); - } finally { - native.version = liveVersion; - } - }); + it.each([false, true])( + "fails when the table changes during every attempt (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + let generation = 1n; + native.version = async () => { + generation += 1n; + await update(generation, `generation ${generation}`); + return liveVersion(); + }; + try { + await expect( + table.query().toArray({ blobMode: "bytes" }), + ).rejects.toThrow(/changed during each of 3 attempts/); + expect(generation).toBe(4n); + } finally { + native.version = liveVersion; + } + }, + ); it("rejects fetchBlobs on a non-blob column", async () => { const { table, rowIds } = await openBlobTable(); @@ -3642,18 +3655,28 @@ describe("when dealing with blob columns", () => { ).map((value) => (value == null ? null : Number(value.size))); } - async function openVersionedBlobTable() { + async function openVersionedBlobTable(nested = false) { const db = await connect(tmpDir.name, { storageOptions: { newTableEnableStableRowIds: "true" }, }); const schema = new Schema([ new Field("id", new Int64(), false), new Field("generation", new Int64(), true), - blob("image"), + nested + ? new Field("info", new Struct([blob("image")]), true) + : blob("image"), ]); const table = await db.createTable( "versions", - [{ id: 1n, generation: 1n, image: Buffer.from("old") }], + [ + { + id: 1n, + generation: 1n, + ...(nested + ? { info: { image: Buffer.from("old") } } + : { image: Buffer.from("old") }), + }, + ], { schema }, ); const native = ( @@ -3663,7 +3686,15 @@ describe("when dealing with blob columns", () => { await table .mergeInsert("id") .whenMatchedUpdateAll() - .execute([{ id: 1n, generation, image: Buffer.from(image) }]); + .execute([ + { + id: 1n, + generation, + ...(nested + ? { info: { image: Buffer.from(image) } } + : { image: Buffer.from(image) }), + }, + ]); }; return { table, native, update }; } diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index c93414fab..5f067eeef 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -30,16 +30,12 @@ import type { DescribeNamespaceResponse, DropNamespaceResponse, JobInfo, - ListNamespacesResponse, - ListTablesResponse, } from "./native"; import { ViewDescription, viewDescriptionFromNative } from "./view"; export type { CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, }; import { sanitizeTable } from "./sanitize"; import { LocalTable, Table } from "./table"; @@ -141,7 +137,7 @@ export interface OpenTableOptions { } /** - * @deprecated Use {@link ListTablesOptions} with {@link Connection.listTables} + * @deprecated Use {@link ListingOptions} with {@link Connection.listTables} * instead. */ export interface TableNamesOptions { @@ -157,32 +153,6 @@ export interface TableNamesOptions { limit?: number; } -export interface ListTablesOptions { - /** - * Token from a previous response, to resume listing where it left off. - * - * The token is opaque: it carries whatever the database needs to resume, and - * callers should not construct or interpret one. - */ - pageToken?: string; - /** - * An upper bound on how many tables to return. - * Defaults to 100. Zero returns an empty page without a continuation token. - * - * A page may hold fewer than this and still not be the last one, so keep - * going while the response carries a page token rather than while pages are - * full. - */ - limit?: number; -} - -export interface ListNamespacesOptions { - /** Token from a previous response for pagination. */ - pageToken?: string; - /** An optional limit to the number of results to return. */ - limit?: number; -} - export interface CreateNamespaceOptions { /** Creation mode. */ mode?: "create" | "exist_ok" | "overwrite"; @@ -278,45 +248,18 @@ export abstract class Connection { options?: Partial, ): Promise; - /** - * List a page of the tables in this database. - * - * To retrieve the tables after the page, pass the `pageToken` the response - * carries back in. A page can be shorter than `limit` without being the last - * one, so walk until a response carries no page token: - * + /** Lazily list tables in a namespace, requesting 100 per page by default. + * @example * ```ts - * const names = []; - * let pageToken = undefined; - * do { - * const page = await conn.listTables({ pageToken, limit: 100 }); - * names.push(...page.tables); - * pageToken = page.pageToken; - * } while (pageToken); + * for await (const name of conn.listTables([], { pageLimit: 20 })) { + * console.log(name); + * } * ``` - * - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. - */ - abstract listTables( - options?: Partial, - ): Promise; - /** - * List a page of the tables in this database. - * - * @param {string[]} namespacePath - The namespace path to list tables from - * (defaults to root namespace) - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. */ abstract listTables( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Open a table in the database. @@ -542,23 +485,11 @@ export abstract class Connection { namespacePath: string[], ): Promise; - /** - * List the immediate child namespaces under the given parent. - * - * Results may be paginated. To retrieve subsequent pages, pass the - * `pageToken` returned by a previous call. - * - * @param {string[]} namespacePath - The parent namespace path. Defaults - * to the root namespace if omitted. - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} Child namespace names and - * an optional token for fetching the next page. - */ + /** Lazily list immediate child namespaces, 100 per page by default. */ abstract listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Create a new namespace at the given path. @@ -821,22 +752,17 @@ export class LocalConnection extends Connection { ); } - async listTables( - namespacePathOrOptions?: string[] | Partial, - options?: Partial, - ): Promise { - // Detect if first argument is namespacePath array or options object - const namespacePath = Array.isArray(namespacePathOrOptions) - ? namespacePathOrOptions - : undefined; - const listTablesOptions = Array.isArray(namespacePathOrOptions) - ? options - : (namespacePathOrOptions ?? options); - - return this.inner.listTables( - namespacePath ?? [], - listTablesOptions?.pageToken, - normalizeTableListingLimit(listTablesOptions?.limit), + listTables( + namespacePath?: string[], + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listTables( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } @@ -1039,12 +965,15 @@ export class LocalConnection extends Connection { listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise { - return this.inner.listNamespaces( - namespacePath ?? [], - options?.pageToken, - options?.limit, + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listNamespaces( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 4ad8b2a93..0f1b5f333 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -97,13 +97,9 @@ export { Connection, CreateTableOptions, TableNamesOptions, - ListTablesOptions, OpenTableOptions, - ListNamespacesOptions, CreateNamespaceOptions, DropNamespaceOptions, - ListNamespacesResponse, - ListTablesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index 43d1cf1b5..0e383e984 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -131,7 +131,8 @@ export interface IvfRqOptions { * * This value controls how much each residual component is compressed. The more * bits, the more accurate the index will be but the slower search. Typical values - * are small integers; the default is 1 bit per dimension. + * are small integers. If unset, uses Lance's default (currently 5 bits per + * dimension). Set to 1 for maximum compression. */ numBits?: number; diff --git a/nodejs/lancedb/job.ts b/nodejs/lancedb/job.ts index 0baa0f571..1c8d74ad2 100644 --- a/nodejs/lancedb/job.ts +++ b/nodejs/lancedb/job.ts @@ -24,7 +24,8 @@ export interface JobEventsOptions { * - {@link Job.refresh} and {@link Job.status} fetch the whole record. * - {@link Job.wait} records the terminal state it establishes, but not the * rest of the record. - * - Everything is null until one of those runs. + * - Registry details are null until one of those runs. Query handles report + * jobType "query" immediately. * * @hideconstructor */ @@ -38,8 +39,10 @@ export class Job { /** * Identifies the operation on the server that is running it. * - * Operations that run in this process have no server id. The value is - * opaque: parsing it or storing it to resume the job later is not supported. + * Registry handles expose a registry job ID; query handles expose a + * connection-scoped SQL query UUID string. Local handles have no ID. + * Query handles report jobType "query". Their IDs are not registry IDs or + * durable resume tokens. */ get id(): string | null { return this.inner.id ?? null; @@ -51,8 +54,8 @@ export class Job { } /** - * The job's type, as the server names it. Null for an in-process job, which - * has no server-side record. + * "query" immediately for SQL-backed handles; otherwise the cached + * server-reported job type. Null for in-process handles. */ get jobType(): string | null { return this.inner.jobType ?? null; diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 823239ef9..b96e4fda6 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +import type { Data as ArrowData } from "apache-arrow"; import { - type ArrowData, Table as ArrowTable, DataType, Field, @@ -14,7 +14,6 @@ import { extractVectorBuffer, fromBufferToRecordBatch, fromRecordBatchToBuffer, - makeData, tableFromIPC, vectorFromArray, } from "./arrow"; @@ -68,84 +67,42 @@ export interface QueryBlobContext { const BLOB_BYTES_ATTEMPTS = 3; +interface BlobSource { + // Output names are literal: an alias containing dots is still one field. + output: string; + path: string[]; + source: string; +} + function blobSources( blobColumns: string[], projection: [string, string][] | undefined, -): Map { - const holds = (source: string) => - blobColumns.some( - (column) => column === source || column.startsWith(`${source}.`), - ); - const pairs = - projection ?? - Array.from(new Set(blobColumns.map((column) => column.split(".")[0]))).map( - (name): [string, string] => [name, name], - ); - return new Map(pairs.filter(([, source]) => holds(source))); -} - -// Fetched bytes for the blobs under one output column. A blob column is a -// leaf; a struct holds the blobs nested in it, keyed by field name. A Map -// keeps struct field names from reaching shared JavaScript prototypes. -type BlobBytes = (Buffer | null)[] | Map; - -function isLeaf(bytes: BlobBytes): bytes is (Buffer | null)[] { - return Array.isArray(bytes); -} - -function nestedFieldType(type: Struct, bytes: Map): Struct { - return new Struct( - type.children.map((field) => { - const inner = bytes.get(field.name); - if (inner === undefined) { - return field; +): BlobSource[] { + if (projection === undefined) { + return blobColumns.map((source) => { + const [output, ...path] = source.split("."); + return { output, path, source }; + }); + } + const sources: BlobSource[] = []; + for (const [output, source] of projection) { + for (const column of blobColumns) { + if (column === source || column.startsWith(`${source}.`)) { + sources.push({ + output, + path: + column === source ? [] : column.slice(source.length + 1).split("."), + source: column, + }); } - if (isLeaf(inner)) { - return new Field(field.name, new LargeBinary(), true); - } - return new Field( - field.name, - nestedFieldType(field.type as Struct, inner), - field.nullable, - ); - }), - ); -} - -// Rebuilds a struct with the blobs under it as bytes. The struct keeps its own -// validity, so a null struct stays null. -function nestedData( - data: ArrowData, - bytes: Map, - offset: number, -): ArrowData { - const type = data.type as Struct; - const children = data.children.map((child, index) => { - const inner = bytes.get(type.children[index].name); - if (inner === undefined) { - return child; } - if (isLeaf(inner)) { - return vectorFromArray( - inner.slice(offset, offset + data.length), - new LargeBinary(), - ).data[0]; - } - return nestedData(child, inner, offset); - }); - return makeData({ - type: nestedFieldType(type, bytes), - length: data.length, - nullCount: data.nullCount, - nullBitmap: data.nullBitmap, - children, - }); + } + return sources; } async function readBlobBytes( table: ArrowTable, - sources: Map, - blobColumns: string[], + sources: BlobSource[], snapshot: NativeTable, dropRowId: boolean, ): Promise { @@ -157,42 +114,27 @@ async function readBlobBytes( rowIdColumn === null ? [] : Array.from(rowIdColumn.toArray() as BigUint64Array); - const fetch = async (column: string) => - (rowIds.length === 0 ? [] : await snapshot.fetchBlobs(column, rowIds)).map( - (value) => value ?? null, + const outputFields = new Map( + table.schema.fields.map((field) => [field.name, field]), + ); + const replacements = new Map(); + for (const { output, path, source } of sources) { + let field = outputFields.get(output); + for (const name of path) { + field = + field !== undefined && DataType.isStruct(field.type) + ? field.type.children.find((child) => child.name === name) + : undefined; + } + if (field === undefined) { + continue; + } + const values = + rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds); + replacements.set( + field, + values.map((value) => value ?? null), ); - const replacements = new Map(); - for (const [output, source] of sources) { - const index = table.schema.fields.findIndex( - (field) => field.name === output, - ); - if (index < 0) { - continue; - } - if (blobColumns.includes(source)) { - replacements.set(index, await fetch(source)); - continue; - } - if (!DataType.isStruct(table.schema.fields[index].type)) { - continue; - } - const nested = new Map(); - for (const column of blobColumns.filter((c) => - c.startsWith(`${source}.`), - )) { - const path = column.slice(source.length + 1).split("."); - let level = nested; - for (const name of path.slice(0, -1)) { - let next = level.get(name); - if (next === undefined) { - next = new Map(); - level.set(name, next); - } - level = next as Map; - } - level.set(path[path.length - 1], await fetch(column)); - } - replacements.set(index, nested); } const result = withBinaryColumns(table, replacements); if (!dropRowId) { @@ -205,48 +147,76 @@ async function readBlobBytes( ); } -// A fresh field, not a clone: the descriptor field's metadata must not follow -// the bytes. function withBinaryColumns( table: ArrowTable, - replacements: Map, + replacements: Map, ): ArrowTable { if (replacements.size === 0) { return table; } - const fields = table.schema.fields.map((field, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return field; + const replaceField = (field: Field): Field => { + if (replacements.has(field)) { + // The descriptor's extension and encoding metadata must not follow bytes. + return new Field(field.name, new LargeBinary(), true); } - return isLeaf(bytes) - ? new Field(field.name, new LargeBinary(), true) - : new Field( - field.name, - nestedFieldType(field.type as Struct, bytes), - field.nullable, - ); - }); + if (DataType.isStruct(field.type)) { + const children = field.type.children.map(replaceField); + if (children.some((child, i) => child !== field.type.children[i])) { + return field.clone({ type: new Struct(children) }); + } + } + return field; + }; + const fields = table.schema.fields.map(replaceField); const schema = new Schema(fields, table.schema.metadata); let offset = 0; const batches = table.batches.map((batch) => { - const children = batch.data.children.map((child, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return child; + const replaceData = ( + data: ArrowData, + field: Field, + updated: Field, + ): ArrowData => { + if (field === updated) { + return data; } - if (!isLeaf(bytes)) { - return nestedData(child, bytes, offset); + const values = replacements.get(field); + if (values !== undefined) { + return vectorFromArray( + values.slice(offset, offset + batch.numRows), + new LargeBinary(), + ).data[0]; } - return vectorFromArray( - bytes.slice(offset, offset + batch.numRows), - new LargeBinary(), - ).data[0]; - }); + const children = data.children.map((child, index) => + replaceData( + child, + field.type.children[index], + updated.type.children[index], + ), + ); + // Keep the struct's validity bitmap and offset, including null parents. + return data.clone( + updated.type, + undefined, + undefined, + undefined, + undefined, + children, + ); + }; + const children = batch.data.children.map((child, index) => + replaceData(child, table.schema.fields[index], fields[index]), + ); offset += batch.numRows; return new RecordBatch( schema, - makeData({ type: new Struct(fields), length: batch.numRows, children }), + batch.data.clone( + new Struct(fields), + undefined, + undefined, + undefined, + undefined, + children, + ), ); }); return new ArrowTable(schema, batches); @@ -274,13 +244,14 @@ export interface QueryExecutionOptions { * * - `"descriptions"` (default): each blob is a descriptor. Read the bytes * with {@link Table.fetchBlobs} or {@link Table.fetchBlobFiles}. - * - `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - * including blobs nested in structs. The query also reads `_rowid` to - * fetch the bytes, and leaves it out of the result unless - * {@link QueryBase.withRowId} was called. All the bytes are held in - * memory. + * - `"bytes"`: each blob column, including blobs nested in structs, holds + * the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + * leaves it out of the result unless {@link QueryBase.withRowId} was + * called. All the bytes are held in memory. * - * Blobs nested in a list keep their descriptors. + * Selecting a parent struct or directly aliasing a blob or struct also returns + * bytes. Blobs inside lists keep their descriptors. Arbitrary projection + * expressions are not materialized as bytes. * * The rows and the bytes come from the same table version. If the table * changes while the query runs, the query runs again, up to three times. @@ -529,15 +500,17 @@ export class QueryBase< // unchanged version afterwards proves it read the snapshot's version. const snapshot = await context.table.querySnapshot(); const version = await snapshot.version(); - const blobColumns = await snapshot.blobColumns(); - const sources = blobSources(blobColumns, context.projection); - if (sources.size === 0) { + const sources = blobSources( + await snapshot.blobColumns(), + context.projection, + ); + if (sources.length === 0) { return this.collect(options, false); } const addRowId = !context.withRowId; const table = await this.collect(options, addRowId); if ((await context.table.version()) === version) { - return readBlobBytes(table, sources, blobColumns, snapshot, addRowId); + return readBlobBytes(table, sources, snapshot, addRowId); } if (attempt === BLOB_BYTES_ATTEMPTS) { throw new Error( diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index bd4b95463..343483a8e 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index 6e4e57684..4c3b93c81 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index f77c50ee2..7679a422d 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index b671b7ebf..6a62f8e8a 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index b75738ff1..d68341382 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 7158748f9..42a517310 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 83d0fb6f2..b116f2819 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package.json b/nodejs/package.json index 21f903235..f746f6daf 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.41.0-beta.13", + "version": "0.41.0-beta.14", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 67c4706e4..914f9e4c9 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -16,8 +16,7 @@ use crate::table::Table; use lancedb::connection::{ConnectBuilder, Connection as LanceDBConnection, connect_namespace}; use lance_namespace::models::{ - CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest, - ListTablesRequest, + CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, }; use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema}; @@ -31,18 +30,6 @@ pub struct DescribeNamespaceResponse { pub properties: Option>, } -#[napi(object)] -pub struct ListNamespacesResponse { - pub namespaces: Vec, - pub page_token: Option, -} - -#[napi(object)] -pub struct ListTablesResponse { - pub tables: Vec, - pub page_token: Option, -} - #[napi(object)] pub struct CreateNamespaceResponse { pub properties: Option>, @@ -239,31 +226,20 @@ impl Connection { op.execute().await.default_error() } - /// List a page of tables in the database. + /// Lazily list tables in the database. #[napi(catch_unwind)] - pub async fn list_tables( + pub fn list_tables( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let request = ListTablesRequest { - // The root namespace is an empty path, not an absent one: a namespace-backed - // database rejects a request that names no namespace. - id: Some(namespace_path.unwrap_or_default()), - page_token, - limit: limit.map(|limit| i32::try_from(limit).unwrap_or(i32::MAX)), - ..Default::default() - }; - let response = self - .get_inner()? - .list_tables(request) - .await - .default_error()?; - Ok(ListTablesResponse { - tables: response.tables, - page_token: response.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Create table from a Apache Arrow IPC (file) buffer. @@ -663,28 +639,19 @@ impl Connection { } #[napi(catch_unwind)] - /// List child namespaces under the given namespace path - pub async fn list_namespaces( + /// Lazily list child namespaces under the given namespace path + pub fn list_namespaces( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let req = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: checked_limit(limit)?, - ..Default::default() - }; - let resp = self - .get_inner()? - .list_namespaces(req) - .await - .default_error()?; - Ok(ListNamespacesResponse { - namespaces: resp.namespaces, - page_token: resp.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[napi(catch_unwind)] @@ -787,33 +754,3 @@ impl Connection { .default_error() } } - -fn checked_limit(limit: Option) -> napi::Result> { - limit - .map(i32::try_from) - .transpose() - .map_err(|_| napi::Error::from_reason("limit must not exceed 2147483647")) -} - -#[cfg(test)] -mod tests { - use super::checked_limit; - - #[test] - fn checked_limit_rejects_values_above_i32_max() { - // u32 values above i32::MAX wrapped negative through `as i32`, - // sending a negative page size to the server. - assert!(checked_limit(Some(3_000_000_000)).is_err()); - assert!(checked_limit(Some(u32::MAX)).is_err()); - } - - #[test] - fn checked_limit_passes_boundary_and_missing_values() { - assert_eq!( - checked_limit(Some(i32::MAX as u32)).unwrap(), - Some(i32::MAX) - ); - assert_eq!(checked_limit(Some(0)).unwrap(), Some(0)); - assert_eq!(checked_limit(None).unwrap(), None); - } -} diff --git a/nodejs/src/job.rs b/nodejs/src/job.rs index 214e687b4..320812128 100644 --- a/nodejs/src/job.rs +++ b/nodejs/src/job.rs @@ -29,8 +29,8 @@ impl Job { #[napi] impl Job { - /// Identifies the operation on the server that is running it. Operations - /// that run in this process have no server id. The value is opaque. + /// The operation's opaque ID. Interpret registry and query IDs using job_type(). + /// Local handles have no ID. #[napi(getter)] pub fn id(&self) -> Option { self.inner.id().map(str::to_string) @@ -62,8 +62,8 @@ impl Job { /// Ask the backend for this job's current state, and for a server-side job /// its full record, then cache it for the getters below. /// - /// They are all null until this runs, because submitting an operation - /// returns only a job id. {@link Job.status} fetches the whole record too; + /// Registry details are null until this runs; query handles already expose + /// job_type "query". {@link Job.status} fetches the whole registry record too; /// {@link Job.wait} records only the terminal state it establishes. #[napi(catch_unwind)] pub async fn refresh(&self) -> napi::Result<()> { @@ -76,8 +76,8 @@ impl Job { self.inner.state() } - /// The job's type, as the server names it. Null for an in-process job, - /// which has no server-side record. + /// "query" immediately for SQL-backed handles; otherwise the cached + /// server-reported job type. Null for in-process handles. #[napi(getter)] pub fn job_type(&self) -> Option { self.inner.job_type() diff --git a/python/Cargo.toml b/python/Cargo.toml index 076299a7d..265e4d270 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" publish = false edition.workspace = true description = "Python bindings for LanceDB" @@ -22,7 +22,6 @@ bytes.workspace = true lancedb.workspace = true datafusion-common.workspace = true lance-core.workspace = true -# The workspace path dependency includes the patched crate in Maturin sources. lance-arrow.workspace = true lance-namespace.workspace = true lance-namespace-impls.workspace = true diff --git a/python/README.md b/python/README.md index f941c9e20..0b640d27c 100644 --- a/python/README.md +++ b/python/README.md @@ -40,6 +40,33 @@ Stable releases are created about every 2 weeks. For the latest features and bug pip install --pre --extra-index-url https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra +and include the Lance preview index: + +```bash +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +With uv, also select the best matching version across the indexes: + +```bash +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + ### Threading in CPU-limited containers LanceDB uses separate pools for compute work and storage I/O. On a container with @@ -74,3 +101,52 @@ print(results) ### Development See [CONTRIBUTING.md](./CONTRIBUTING.md) for information on how to contribute to LanceDB. + +### Opening SQL blob descriptors + +A SQL `blob_descriptor(video)` result keeps the Blob V2 storage metadata and a +`source` struct containing its table URI, version, commit timestamp, branch, +column and row identity. This example uses a remote SQL connection whose SQL +service produces those source-bearing descriptors. Open the reference on the +source table to read incrementally: + +```python +rows = db.execute_query(""" + SELECT f.timestamp, blob_descriptor(v.video) AS video + FROM videos v JOIN frames f ON v.id = f.video_id +""").read_all() + +videos = db.open_table("videos") +for row in rows.to_pylist(): + if row["video"] is not None: + with videos.open_blob_file(row["video"]) as video: + frame = extract_frame(video, row["timestamp"]) +``` + +`open_blob_file` accepts the dictionary or an Arrow struct scalar directly. It +keeps the query's version and branch without changing the table handle's checkout. +Payload bytes load on reads. An external blob with an unspecified size may need a +metadata-only HEAD request to open. + +`videos.open_blob_files(rows["video"])` returns an iterator over the same handles, +preserving order, duplicates and nulls (`None`). It accepts Arrow arrays/chunked +arrays or another iterable and opens each handle as iteration advances. Close each +non-null handle after use. Async tables offer `await videos.open_blob_file(...)` +and `async for video in videos.open_blob_files(...)`. + +```python +for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + header = video.read(4096) +``` + +To extract a frame while retaining each row's timestamp, iterate over the +handles and rows together: + +```python +for row, video in zip(rows.to_pylist(), videos.open_blob_files(rows["video"])): + if video is not None: + with video: + frame = extract_frame(video, row["timestamp"]) +``` diff --git a/python/pyproject.toml b/python/pyproject.toml index ed54ae536..a0de6509c 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -45,7 +45,8 @@ repository = "https://github.com/lancedb/lancedb" [project.optional-dependencies] pylance = [ - "pylance>=5.0.0b5", + # Keep the minimum in the bundled Lance major and compatible with tests below. + "pylance>=15.0.0b1", ] # A library only needs the OpenTelemetry API; the application supplies and # configures the SDK (the actual exporter/reader). See @@ -63,7 +64,7 @@ tests = [ "polars>=0.19, <=1.32.3", "pyarrow<25", "pyarrow-stubs>=16.0", - "pylance==14.0.0b12", + "pylance==15.0.0b1", "requests>=2.31.0", "datafusion>=54,<55", "opentelemetry-sdk>=1.30.0", diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 6b902c466..5e879f61d 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -24,6 +24,7 @@ from .remote import ClientConfig from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector +from ._blob_reader import BlobReader as BlobReader from .job import AsyncJob, Job from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery @@ -280,6 +281,21 @@ def connect( "uri is required when not connecting through a namespace client" ) if isinstance(uri, str) and uri.startswith("db://"): + # Preserve the deprecated remote timeout options, but validate all other + # unknown keywords just as the local connection does. + unknown_kwargs = { + key: value + for key, value in kwargs.items() + if key not in ("connection_timeout", "read_timeout") + } + if unknown_kwargs: + raise ValueError(f"Unknown keyword arguments: {unknown_kwargs}") + if session is not None: + raise NotImplementedError("session is not supported for remote connections") + if manifest_enabled: + raise NotImplementedError( + "manifest_enabled is not supported for remote connections" + ) if api_key is None: api_key = os.environ.get("LANCEDB_API_KEY") if api_key is None: diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 1e9110dd0..46462553c 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -7,7 +7,7 @@ from __future__ import annotations import io import operator -from collections.abc import Awaitable, Callable, Iterable +from collections.abc import Awaitable, Callable, Iterable, Mapping from typing import TYPE_CHECKING, Optional, Union import pyarrow as pa @@ -31,6 +31,18 @@ FetchBlobsSync = Callable[[str, list[int]], pa.Array | pa.ChunkedArray] FetchBlobsAsync = Callable[[str, list[int]], Awaitable[pa.Array | pa.ChunkedArray]] +def _normalize_blob_descriptor(value): + if isinstance(value, pa.Scalar): + value = value.as_py() + if value is None: + return None + if not isinstance(value, Mapping): + raise TypeError( + "blob descriptor must be a mapping, an Arrow struct scalar, or None" + ) + return dict(value) + + class BlobFile(io.RawIOBase): """Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`. diff --git a/python/python/lancedb/_blob_reader.py b/python/python/lancedb/_blob_reader.py new file mode 100644 index 000000000..2fd0f00ca --- /dev/null +++ b/python/python/lancedb/_blob_reader.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +"""Seekable binary input for Functions that consume Blob V2 values lazily.""" + +import io +import operator +import urllib.request +from collections.abc import Callable +from http import HTTPStatus + +_READ_CHUNK_BYTES = 1024 * 1024 +_READ_TIMEOUT_SECONDS = 300 + + +class BlobReader(io.RawIOBase): + """Read a blob incrementally inside a Function. + + Annotating a Function input with ``BlobReader`` requests a lazy, seekable + reader. Annotating it with ``bytes`` requests materialized bytes instead. + ``read_range`` does not change the sequential cursor. Readers are scoped + to the invocation and must not be retained for a later call. + + For local calls, construct a reader from a range-reading callable and the + blob's size. The callable accepts an offset and length and returns bytes. + """ + + def __init__(self, read_range: Callable[[int, int], bytes], size: int): + super().__init__() + self._read_range = read_range + self._size = operator.index(size) + if self._size < 0: + raise ValueError("blob size must be nonnegative") + self._position = 0 + + @classmethod + def _from_uri(cls, uri: str, size: int): + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + + def read_range(offset, length): + request = urllib.request.Request( + uri, headers={"Range": f"bytes={offset}-{offset + length - 1}"} + ) + with opener.open(request, timeout=_READ_TIMEOUT_SECONDS) as response: + if response.status != HTTPStatus.PARTIAL_CONTENT: + raise OSError( + "blob range service did not return a partial response" + ) + data = response.read(length + 1) + if len(data) != length: + raise OSError( + "blob range service returned an invalid payload length" + ) + return data + + return cls(read_range, size) + + def readable(self) -> bool: + return not self.closed + + def seekable(self) -> bool: + return not self.closed + + def size(self) -> int: + """Return the blob size without reading its payload.""" + self._checkClosed() + return self._size + + def tell(self) -> int: + self._checkClosed() + return self._position + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + self._checkClosed() + offset = operator.index(offset) + if whence == io.SEEK_SET: + position = offset + elif whence == io.SEEK_CUR: + position = self._position + offset + elif whence == io.SEEK_END: + position = self._size + offset + else: + raise ValueError(f"invalid whence: {whence}") + if position < 0: + raise ValueError("negative seek position") + self._position = position + return position + + def read_range(self, offset: int, length: int) -> bytes: + """Read a byte range without changing the cursor; truncate at EOF.""" + self._checkClosed() + offset, length = operator.index(offset), operator.index(length) + if offset < 0 or length < 0: + raise ValueError("offset and length must be nonnegative") + remaining = min(length, max(0, self._size - offset)) + chunks = [] + while remaining: + count = min(remaining, _READ_CHUNK_BYTES) + chunk = self._read_range(offset, count) + if len(chunk) != count: + raise OSError("blob range reader returned an invalid payload length") + chunks.append(chunk) + offset += count + remaining -= count + return b"".join(chunks) + + def read(self, size: int = -1) -> bytes: + self._checkClosed() + size = -1 if size is None else operator.index(size) + if size < 0: + size = max(0, self._size - self._position) + data = self.read_range(self._position, size) + self._position += len(data) + return data + + def readinto(self, buffer) -> int: + view = memoryview(buffer).cast("B") + data = self.read(len(view)) + view[: len(data)] = data + return len(data) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 96eef064c..0b4fc82f1 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -19,11 +19,9 @@ from .index import ( FTS, ) from lance_namespace import ( - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, ) from .remote import ClientConfig from .types import BaseTokenizerType @@ -114,12 +112,12 @@ class Connection(object): uri: str async def is_open(self): ... async def close(self): ... - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: ... + ) -> NameListing: ... async def create_namespace( self, namespace_path: List[str], @@ -136,12 +134,12 @@ class Connection(object): self, namespace_path: List[str], ) -> DescribeNamespaceResponse: ... - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: ... + ) -> NameListing: ... async def table_names( self, namespace_path: Optional[List[str]], @@ -322,6 +320,8 @@ class Job: @property def id(self) -> Optional[str]: ... @property + def job_type(self) -> Optional[str]: ... + @property def _state(self) -> Optional[str]: ... @property def _description(self) -> Optional[JobDescription]: ... @@ -580,6 +580,7 @@ class Table: async def fetch_blob_files( self, column: str, row_ids: list[int] ) -> list[Optional[BlobFile]]: ... + async def open_blob_file(self, descriptor: str) -> BlobFile: ... def vector_search(self) -> VectorQuery: ... class Tags: @@ -1010,6 +1011,7 @@ async def connect_catalog( class NameListing: def num_page_results(self) -> int: ... def page_token(self) -> Optional[str]: ... + async def next_page(self) -> Optional[dict[str, Any]]: ... def __aiter__(self) -> NameListing: ... async def __anext__(self) -> str: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 85d18de75..30cd25a73 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -32,18 +32,24 @@ if sys.version_info >= (3, 12): else: from overrides import EnforceOverrides, override # type: ignore -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, + validate_page_limit, +) from lancedb.embeddings.registry import EmbeddingFunctionRegistry from lancedb.common import data_to_reader, sanitize_uri, validate_schema from lancedb.background_loop import LOOP from lance_namespace import ( LanceNamespace, - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, connect as namespace_connect, ) from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError @@ -171,7 +177,7 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -187,12 +193,12 @@ class DBConnection(EnforceOverrides): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return ListNamespacesResponse(namespaces=[], page_token=None) + return NamespaceListing(AsyncNamespaceListing(None, empty=True)) def create_namespace( self, @@ -317,8 +323,8 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -334,7 +340,7 @@ class DBConnection(EnforceOverrides): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ raise NotImplementedError( @@ -422,10 +428,13 @@ class DBConnection(EnforceOverrides): it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved, and vectors longer than the expected + dimension are truncated. storage_options: dict, optional Additional options for the storage backend. Options already set on the connection will be inherited by the table, but can be overridden here. @@ -542,7 +551,9 @@ class DBConnection(EnforceOverrides): LanceTable(name='table4', ...) """ - raise NotImplementedError + raise NotImplementedError( + "create_table is not supported for this connection type" + ) def __getitem__(self, name: str) -> LanceTable: return self.open_table(name) @@ -597,7 +608,9 @@ class DBConnection(EnforceOverrides): ------- A LanceTable object representing the table. """ - raise NotImplementedError + raise NotImplementedError( + "open_table is not supported for this connection type" + ) def create_materialized_view( self, @@ -664,10 +677,11 @@ class DBConnection(EnforceOverrides): ) -> Job[None]: """Submit materialized-view creation and return its job. - The job may already be complete for a local database. On LanceDB - Cloud and Enterprise, its ``id`` is the server job identifier from - the ``202 Accepted`` create response. Wait for the job before opening - or querying the view. + The job may already be complete for a local database. Remote creation + submits CREATE MATERIALIZED VIEW through the SQL server. Its job has + ``job_type == "query"`` and ``id`` identifies the connection-scoped + query. Registry event history is unavailable. + Wait for the job before opening or querying the view. """ raise NotImplementedError( "materialized views are not supported on this connection type" @@ -737,13 +751,17 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_table is not supported for this connection type" + ) def drop_table_async( self, name: str, namespace_path: Optional[List[str]] = None ) -> Job: """Start dropping a table and return its cleanup job.""" - raise NotImplementedError + raise NotImplementedError( + "drop_table_async is not supported for this connection type" + ) def rename_table( self, @@ -771,19 +789,26 @@ class DBConnection(EnforceOverrides): cur_namespace_path = [] if new_namespace_path is None: new_namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "rename_table is not supported for this connection type" + ) def drop_database(self): """ Drop database This is the same thing as dropping all the tables """ - raise NotImplementedError + raise NotImplementedError( + "drop_database is not supported for this connection type" + ) def drop_all_tables(self, namespace_path: Optional[List[str]] = None): """ Drop all tables from the database + Remote connections do not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + Parameters ---------- namespace_path: List[str], optional @@ -792,12 +817,32 @@ class DBConnection(EnforceOverrides): """ if namespace_path is None: namespace_path = [] - raise NotImplementedError + raise NotImplementedError( + "drop_all_tables is not supported for this connection type" + ) @property def uri(self) -> str: + """The URI of this database.""" return self._uri + def _all_table_names(self) -> Generator[str, None, None]: + page_token = None + while True: + response = self.list_tables(page_token=page_token) + yield from response.tables + page_token = response.page_token + if not page_token: + return + + def __len__(self) -> int: + """Return the number of tables in the root namespace across all pages.""" + return sum(1 for _ in self._all_table_names()) + + def __contains__(self, name: str) -> bool: + """Return whether the root namespace contains the named table.""" + return name in self._all_table_names() + def namespace_client(self) -> LanceNamespace: """Get the equivalent namespace client for this connection. @@ -1230,18 +1275,19 @@ class DBConnection(EnforceOverrides): parameters=[np.array([0.1, 0.2, 0.3], dtype=np.float32), 10], ) - Parameters are supported in queries, not in DDL or DML. A - parameterized query runs on the call that returns its rows: - cancelling it, or dropping its handle before reading, stops it on the - server. Its status follows that stream as well: ``describe()`` - reports ``running`` until the reader has received every row, so open - ``reader()`` directly rather than waiting for ``finished``. + Parameters are supported in queries, not in DDL or DML. With or + without parameters, the server chooses whether to stream rows on the + initial call or return a polling continuation. The reader follows + either path transparently. Dropping an unread handle or its reader + stops a directly streamed query; deferred execution requires explicit + cancellation. Open ``reader()`` directly: a streamed query needs its + reader to make progress and reports ``running`` until the stream ends. Local connections do not support SQL. """ raise NotImplementedError("SQL is not supported for this connection type") - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" raise NotImplementedError("SQL is not supported for this connection type") @@ -1436,7 +1482,7 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -1452,16 +1498,14 @@ class LanceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( - namespace_path=namespace_path, - page_token=page_token, - limit=limit, + namespace_path=namespace_path, page_token=page_token, limit=limit ) ) @@ -1512,8 +1556,8 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -1529,12 +1573,12 @@ class LanceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1583,21 +1627,6 @@ class LanceDBConnection(DBConnection): ) ) - def _all_table_names(self) -> Generator[str, None, None]: - page_token = None - while True: - response = self.list_tables(page_token=page_token) - yield from response.tables - page_token = response.page_token - if not page_token: - return - - def __len__(self) -> int: - return sum(1 for _ in self._all_table_names()) - - def __contains__(self, name: str) -> bool: - return name in self._all_table_names() - @override def create_table( self, @@ -2286,12 +2315,12 @@ class AsyncConnection(object): else: return None - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -2306,15 +2335,18 @@ class AsyncConnection(object): Returns ------- - ListNamespacesResponse + AsyncNamespaceListing Response containing namespace names and optional pagination token """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_namespaces( + validate_page_limit(limit) + inner = self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListNamespacesResponse(**result) + if limit == 0: + return AsyncNamespaceListing(None, empty=True) + return AsyncNamespaceListing(inner, page_token) async def create_namespace( self, @@ -2393,13 +2425,13 @@ class AsyncConnection(object): result = await self._inner.describe_namespace(namespace_path) return DescribeNamespaceResponse(**result) - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> AsyncTableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -2415,15 +2447,18 @@ class AsyncConnection(object): Returns ------- - ListTablesResponse + AsyncTableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_tables( + validate_page_limit(limit) + inner = self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListTablesResponse(**result) + if limit == 0: + return AsyncTableListing(None, empty=True) + return AsyncTableListing(inner, page_token) async def table_names( self, @@ -2521,10 +2556,13 @@ class AsyncConnection(object): it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float The value to use when filling vectors. Only used if on_bad_vectors="fill". + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved, and vectors longer than the expected + dimension are truncated. storage_options: dict, optional Additional options for the storage backend. Options already set on the connection will be inherited by the table, but can be overridden here. @@ -2841,7 +2879,9 @@ class AsyncConnection(object): ) -> AsyncJob[None]: """Submit materialized-view creation and return its job. - Wait for the returned job before opening or querying the view. + Wait for the returned job before opening or querying the view. Remote + creation uses SQL; ``job_type`` is ``"query"`` and ``id`` + identifies that query on this connection. """ source, native_options = await self._materialized_view_source( source, select, where, limit @@ -3381,8 +3421,14 @@ class AsyncConnection(object): ) ) - async def describe_query(self, query_id: UUID) -> QueryDescription: - """Describe a submitted SQL query by its connection-scoped id.""" + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: + """Describe a submitted SQL query by its UUID or string ID. + + For a handle with ``job_type == "query"``, pass its ``id`` on the + connection that submitted it. Registry job IDs are not query IDs. + """ + if isinstance(query_id, str): + query_id = UUID(query_id) return await self._inner.describe_query(query_id) async def namespace_client(self) -> LanceNamespace: diff --git a/python/python/lancedb/embeddings/open_clip.py b/python/python/lancedb/embeddings/open_clip.py index 8b3f0eb06..6199ef271 100644 --- a/python/python/lancedb/embeddings/open_clip.py +++ b/python/python/lancedb/embeddings/open_clip.py @@ -49,6 +49,9 @@ class OpenClipEmbeddings(EmbeddingFunction): self.name, pretrained=self.pretrained ) model.to(self.device) + # open_clip returns the model in training mode, which keeps dropout, + # stochastic depth and batch norm in their training behaviour. + model.eval() self._model, self._preprocess = model, preprocess self._tokenizer = open_clip.get_tokenizer(self.name) self._ndims = None diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 7f5776b32..42c4b27e1 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -33,6 +33,7 @@ from typing import ( Annotated, Any, Callable, + Literal, Optional, Union, get_args, @@ -52,7 +53,9 @@ from pydantic import ( model_validator, ) +from .schema import blob as _blob_field from .schema import is_blob_v2_field as _is_blob_v2_field +from ._blob_reader import BlobReader from .secrets import EnvVarSecret _Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1) @@ -213,6 +216,7 @@ class FunctionParameter(_RemoteValue): name: str arrow_type: str nullable: bool + blob_access: Optional[Literal["bytes", "reader"]] = None class FunctionResultField(_OpenRemoteValue): @@ -1096,6 +1100,43 @@ def _infer_signature( if (input_schema is None) != (output_schema is None): raise ValueError("input_schema and output_schema must be provided together") + try: + annotations = get_type_hints(function, include_extras=True) + except Exception as error: + raise TypeError(f"failed to resolve Function annotations: {error}") from error + + def input_parameter(field): + annotation = annotations.get(field.name) + if get_origin(annotation) in (Union, types.UnionType): + choices = [ + value for value in get_args(annotation) if value is not type(None) + ] + if len(choices) == 1: + annotation = choices[0] + arrow_type = _canonical_arrow_field(field) + if annotation is BlobReader: + if arrow_type != _FUNCTION_BLOB_V2_TYPE: + raise TypeError( + f"BlobReader input '{field.name}' requires a Blob V2 field" + ) + access = "reader" + elif annotation is bytes: + if arrow_type != _FUNCTION_BLOB_V2_TYPE and not ( + pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type) + ): + raise TypeError( + f"bytes input '{field.name}' requires binary or Blob V2" + ) + access = "bytes" + else: + access = "bytes" if arrow_type == _FUNCTION_BLOB_V2_TYPE else None + return FunctionParameter( + name=field.name, + arrow_type=arrow_type, + nullable=field.nullable, + blob_access=access, + ) + if input_schema is not None: if not isinstance(input_schema, pa.Schema): raise TypeError("input_schema must be a PyArrow Schema") @@ -1110,20 +1151,9 @@ def _infer_signature( "input_schema fields must exactly match callable parameters in order: " f"expected {expected!r}, got {actual!r}" ) - inputs = tuple( - FunctionParameter( - name=field.name, - arrow_type=_canonical_arrow_field(field), - nullable=field.nullable, - ) - for field in input_schema - ) + inputs = tuple(input_parameter(field) for field in input_schema) return FunctionSignature(inputs=inputs, output=_function_output(output_schema)) - try: - annotations = get_type_hints(function, include_extras=True) - except Exception as error: - raise TypeError(f"failed to resolve Function annotations: {error}") from error missing = [ parameter.name for parameter in parameters if parameter.name not in annotations ] @@ -1132,15 +1162,21 @@ def _infer_signature( raise TypeError(f"missing Function annotations: {names!r}") inputs = [] for parameter in parameters: - data_type, nullable = _annotation_type(annotations[parameter.name]) - inputs.append( - FunctionParameter( - name=parameter.name, - arrow_type=_canonical_arrow_field( - pa.field(parameter.name, data_type, nullable=nullable) - ), - nullable=nullable, + annotation = annotations[parameter.name] + base = annotation + nullable = False + if get_origin(base) in (Union, types.UnionType): + choices = [value for value in get_args(base) if value is not type(None)] + if len(choices) == 1: + base, nullable = choices[0], True + if base is BlobReader: + inputs.append( + input_parameter(_blob_field(parameter.name, nullable=nullable)) ) + continue + data_type, nullable = _annotation_type(annotation) + inputs.append( + input_parameter(pa.field(parameter.name, data_type, nullable=nullable)) ) output_type, output_nullable = _annotation_type(annotations["return"]) if output_nullable: @@ -1239,7 +1275,7 @@ _PYTHON_CALLABLE_ARTIFACT = "python_callable" _PYTHON_BUNDLE_ARTIFACT = "python_bundle" # Modules the Function image places beside user code. _RESERVED_SOURCE_MODULES = frozenset( - {"user_function", "lance_source", "lance_callable"} + {"user_function", "lance_source", "lance_callable", "lance_blob"} ) # The service accepts bundle paths made of these components only. _SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") @@ -1499,6 +1535,8 @@ class _SourcePackager: """One module-level line that rebinds `name` to `value` in the artifact: an import for modules and importable classes/functions, a literal otherwise.""" + if value is BlobReader: + return f"from lance_blob import BlobReader as {name}" if isinstance(value, types.ModuleType): if value.__name__.split(".")[0] in _NAMESPACE_MODULES: raise ValueError( diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index cf3bf8e4f..80179f4a7 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -882,8 +882,12 @@ class IvfRq: num_partitions: int, default sqrt(num_rows) Number of IVF partitions to create. - num_bits: int, default 1 + num_bits: int, optional Number of bits to encode each dimension in the RabitQ codebook. + If None, uses Lance's default (currently 5), matching Rust and Node. + Set to 1 to retain the previous Python default and maximum compression. + Five-bit vector codes use roughly five times the storage of one-bit codes, + in addition to other index overhead. max_iterations: int, default 50 Max iterations to train kmeans when computing IVF partitions. @@ -897,7 +901,7 @@ class IvfRq: distance_type: Literal["l2", "cosine", "dot"] = "l2" num_partitions: Optional[int] = None - num_bits: int = 1 + num_bits: Optional[int] = None max_iterations: int = 50 sample_rate: int = 256 target_partition_size: Optional[int] = None diff --git a/python/python/lancedb/job.py b/python/python/lancedb/job.py index e57fde0ea..e61a188f2 100644 --- a/python/python/lancedb/job.py +++ b/python/python/lancedb/job.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -"""Handles to operations a server may run asynchronously.""" +"""Handles to asynchronous registry jobs, SQL queries, and local operations.""" import asyncio import json @@ -44,12 +44,12 @@ class AsyncJob(Generic[T]): @property def id(self) -> Optional[str]: - """Identifies the operation on the server that is running it. + """The operation's opaque ID, interpreted according to :attr:`job_type`. - Returned for correlating with server logs or the jobs API. Operations - that run in this process have no server id and return `None`. The value - is opaque: parsing it or storing it to resume the job later is not - supported. + Registry jobs expose a registry job ID; SQL-backed jobs expose a + connection-scoped query ID as a string. Local operations return ``None``. + Query IDs cannot be used with registry job APIs or to resume after a + client restart. """ return self._inner.id if self._inner is not None else None @@ -94,9 +94,11 @@ class AsyncJob(Generic[T]): """Ask the backend for this job's current state, and for a server-side job its full record, then cache it for the properties below. - The properties are all `None` until this runs, because submitting an - operation returns only a job id. `status` fetches the whole record too; - `wait` records only the terminal state it establishes. + Registry details are `None` until this runs, because submitting an + operation does not fetch its full registry record. Query-backed jobs + expose only lifecycle state through this method. For registry jobs, + `status` fetches the whole record too; `wait` records only the terminal + state it establishes. """ if self._inner is None: return @@ -114,11 +116,13 @@ class AsyncJob(Generic[T]): @property def job_type(self) -> Optional[str]: - """The job's type, as the server names it. + """The job's kind of work, without contacting the backend. - `None` for an in-process job, which has no server-side record. + SQL-backed handles report ``"query"`` immediately. Registry handles + report the server's job type after :meth:`refresh` or :meth:`status`. + In-process handles return ``None``. """ - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: @@ -294,7 +298,7 @@ class Job(Generic[T]): @property def job_type(self) -> Optional[str]: """The job's type. See :attr:`AsyncJob.job_type`.""" - return self._field("job_type") + return self._inner.job_type if self._inner is not None else None @property def creation_ms(self) -> Optional[int]: diff --git a/python/python/lancedb/listing.py b/python/python/lancedb/listing.py index 0620aa7d8..8abc9d85b 100644 --- a/python/python/lancedb/listing.py +++ b/python/python/lancedb/listing.py @@ -5,6 +5,7 @@ from __future__ import annotations +import asyncio from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar from .background_loop import LOOP @@ -83,3 +84,173 @@ class Listing(Iterator[T], Generic[T]): for initial, cached, final-page, and error behavior. """ return self._inner.page_token() + + +def validate_page_limit(limit: Optional[int]) -> None: + """Validate the legacy limit argument before passing it to the bindings.""" + if limit is not None and ( + not isinstance(limit, int) or limit < 0 or limit > 2**31 - 1 + ): + raise ValueError("limit must be between 0 and 2147483647") + + +class _AsyncPagedNameListing(AsyncIterator[str]): + """An async listing with an inspectable copy of its current backend page.""" + + def __init__( + self, inner: Any, page_token: Optional[str] = None, *, empty: bool = False + ): + self._inner = inner + self._token = page_token + self._page: Optional[list[str]] = [] if empty else None + self._position = 0 + self._done = empty + self._failure: Optional[BaseException] = None + self._lock = asyncio.Lock() + + @property + def page_token(self) -> Optional[str]: + """The continuation token for the next backend request.""" + return self._token + + def num_page_results(self) -> int: + """Count the remaining names in the current page.""" + return 0 if self._page is None else len(self._page) - self._position + + def _page_copy(self) -> list[str]: + if self._page is None: + raise RuntimeError("Await or iterate the listing before reading its page") + return self._page.copy() + + async def _fetch_page(self) -> None: + try: + page = await self._inner.next_page() + except asyncio.CancelledError: + raise + except BaseException as exc: + self._failure = exc + self._done = True + raise + if page is None: + self._page = [] + self._token = None + self._done = True + else: + self._page = page["items"] + self._token = page["page_token"] + self._done = self._token is None + self._position = 0 + + async def _load_initial(self): + async with self._lock: + if self._failure is not None: + raise self._failure + if self._page is None: + await self._fetch_page() + return self + + def __await__(self): + return self._load_initial().__await__() + + def __aiter__(self): + return self + + async def __anext__(self) -> str: + async with self._lock: + if self._failure is not None: + raise StopAsyncIteration + if self._page is None: + await self._fetch_page() + assert self._page is not None + while self._position >= len(self._page): + if self._done: + raise StopAsyncIteration + await self._fetch_page() + assert self._page is not None + item = self._page[self._position] + self._position += 1 + return item + + +class AsyncTableListing(_AsyncPagedNameListing): + """Awaitable table iterator; awaiting loads one page. + + For example:: + + page = await db.list_tables(limit=20) + print(page.tables, page.page_token) + async for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + """A copy of every table name in the most recently fetched page.""" + return self._page_copy() + + +class AsyncNamespaceListing(_AsyncPagedNameListing): + """Awaitable namespace iterator; awaiting loads one page. + + For example:: + + page = await db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + """A copy of every name in the most recently fetched page.""" + return self._page_copy() + + +class _PagedNameListing(Iterator[str]): + def __init__(self, inner: _AsyncPagedNameListing): + self._inner = inner + LOOP.run(inner._load_initial()) + + def __iter__(self): + return self + + def __next__(self) -> str: + try: + return LOOP.run(self._inner.__anext__()) + except StopAsyncIteration: + raise StopIteration from None + + @property + def page_token(self) -> Optional[str]: + return self._inner.page_token + + def num_page_results(self) -> int: + return self._inner.num_page_results() + + +class TableListing(_PagedNameListing): + """Table iterator retaining the current page as a list. + + For example:: + + page = db.list_tables(limit=20) + print(page.tables, page.page_token) + for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + return self._inner.tables + + +class NamespaceListing(_PagedNameListing): + """Namespace iterator retaining the current page as a list. + + For example:: + + page = db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + return self._inner.namespaces diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index 24253d214..630b17f3c 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -52,7 +52,14 @@ from lancedb.background_loop import LOOP from lancedb.arrow import AsyncRecordBatchReader from lancedb.db import AsyncConnection, DBConnection from lancedb.job import AsyncJob, Job -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, +) from lancedb.sql import AsyncQuery as AsyncSqlQuery from lancedb.sql import QueryDescription, QueryParameters from lance_namespace import ( @@ -61,8 +68,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, NamespaceExistsRequest, TableExistsRequest, ) @@ -875,7 +880,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """ List child namespaces under the given namespace. @@ -892,12 +897,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1044,7 +1049,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """ List all tables in this database with pagination support. @@ -1061,12 +1066,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1483,12 +1488,12 @@ class AsyncLanceNamespaceDBConnection: namespace_path = [] await self._inner.drop_all_tables(namespace_path=namespace_path) - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """ List child namespaces under the given namespace. @@ -1505,12 +1510,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_namespaces( + return self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1640,12 +1645,12 @@ class AsyncLanceNamespaceDBConnection: except TableNotFoundError: return False - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> AsyncTableListing: """ List all tables in this database with pagination support. @@ -1662,12 +1667,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_tables( + return self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1702,7 +1707,7 @@ class AsyncLanceNamespaceDBConnection: parameters=parameters, ) - async def describe_query(self, query_id: UUID) -> QueryDescription: + async def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query when supported.""" return await self._inner.describe_query(query_id) diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index f89ea4d8c..bd8302ad0 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -916,9 +916,11 @@ class Permutation: f"Cannot rename column `{old_name}` to `{new_name}` because a column " "with that name already exists" ) - new_selection = self.selection.copy() - new_selection[new_name] = new_selection[old_name] - del new_selection[old_name] + # Rebuild the mapping so the renamed column keeps its position + new_selection = { + (new_name if name == old_name else name): value + for name, value in self.selection.items() + } return self._with_selection(new_selection) def rename_columns(self, column_map: dict[str, str]) -> "Permutation": diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 9d7610f72..52f38d077 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -1561,9 +1561,10 @@ class LanceQueryBuilder(ABC): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] 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] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -1619,8 +1620,8 @@ class LanceQueryBuilder(ABC): GlobalLimitExec: elapsed=..., skip=0, fetch=10, metrics=[output_rows=..., elapsed_compute=..., output_bytes=...] FilterExec: elapsed=..., _distance@2 IS NOT NULL, metrics=[...] - SortExec: elapsed=..., TopK(fetch=10), expr=[...], - preserve_partitioning=[...], + SortExec: elapsed=..., TopK(fetch=1), expr=[...], + preserve_partitioning=[...], filter=[...], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., row_replacements=...] KNNVectorDistance: elapsed=..., metric=l2, @@ -1633,6 +1634,7 @@ class LanceQueryBuilder(ABC): metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., fragments_scanned=..., ranges_scanned=1, rows_scanned=1, bytes_read=..., iops=..., requests=..., task_wait_time=...] + Parameters ---------- @@ -2959,9 +2961,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] 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] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -3601,9 +3604,10 @@ class AsyncQueryBase(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] 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] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters @@ -4521,9 +4525,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): 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] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score, _rowid@0 as _rowid] LanceRead: uri=..., projection=[vector, text], source=stream(_rowid) GlobalLimitExec: skip=0, fetch=10 @@ -4877,9 +4881,10 @@ class BaseQueryBuilder(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] 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] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 89c7f8a78..291cc52d6 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -7,7 +7,6 @@ from __future__ import annotations from dataclasses import replace from datetime import timedelta import json -import logging from concurrent.futures import ThreadPoolExecutor import sys from typing import ( @@ -36,7 +35,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig import pyarrow as pa from ..common import DATA -from lancedb.listing import Listing +from lancedb.listing import Listing, NamespaceListing, TableListing from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job @@ -47,15 +46,13 @@ from ..secrets import EnvVarSecret, SecretInfo from ..view import ViewDescription if TYPE_CHECKING: - from .._lancedb import JobInfo + from .._lancedb import JobInfo, Session from ..embeddings import EmbeddingFunctionConfig from lance_namespace import ( LanceNamespace, CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, ) from ..pydantic import LanceModel from ..table import Table @@ -242,6 +239,23 @@ class RemoteDBConnection(DBConnection): def __repr__(self) -> str: return f"RemoteConnect(name={self.db_name})" + @property + @override + def uri(self) -> str: + return self._conn.uri + + @property + def session(self) -> Optional["Session"]: + """Remote connections do not expose a local cache session.""" + raise NotImplementedError("session is not supported by remote connections") + + @property + def read_consistency_interval(self) -> Optional[timedelta]: + """Remote connections do not expose a local read consistency interval.""" + raise NotImplementedError( + "read_consistency_interval is not supported by remote connections" + ) + @override def serialize(self) -> str: if getattr(self, "_catalog_oauth", False): @@ -268,7 +282,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -284,12 +298,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -379,7 +393,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """List a page of tables in this database. Parameters @@ -396,12 +410,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -487,18 +501,14 @@ class RemoteDBConnection(DBConnection): if namespace_path is None: namespace_path = [] - if storage_options is not None: - logging.info( - "storage_options is ignored for remote tables" - " (storage is managed; set storage_options on connect() instead)" + table = LOOP.run( + self._conn.open_table( + name, + namespace_path=namespace_path, + storage_options=storage_options, + index_cache_size=index_cache_size, ) - if index_cache_size is not None: - logging.info( - "index_cache_size is ignored for remote tables" - " (there is no local cache to configure)" - ) - - table = LOOP.run(self._conn.open_table(name, namespace_path=namespace_path)) + ) tbl = RemoteTable( table, self.db_name, @@ -571,13 +581,16 @@ class RemoteDBConnection(DBConnection): name: str, data: DATA = None, schema: Optional[Union[pa.Schema, LanceModel]] = None, - on_bad_vectors: str = "error", - fill_value: float = 0.0, mode: Optional[str] = None, exist_ok: bool = False, + on_bad_vectors: str = "error", + fill_value: float = 0.0, embedding_functions: Optional[List[EmbeddingFunctionConfig]] = None, *, namespace_path: Optional[List[str]] = None, + storage_options: Optional[Dict[str, str]] = None, + data_storage_version: Optional[str] = None, + enable_v2_manifest_paths: Optional[bool] = None, ) -> Table: """Create a [Table][lancedb.table.Table] in the database. @@ -617,6 +630,12 @@ class RemoteDBConnection(DBConnection): embedding_functions: list of EmbeddingFunctionConfig, optional The embedding functions to store in the table schema and use to generate vectors when creating the table or adding data. + storage_options: dict, optional + Not supported for remote table creation. Raises NotImplementedError. + data_storage_version: str, optional + Not supported for remote table creation. Raises NotImplementedError. + enable_v2_manifest_paths: bool, optional + Not supported for remote table creation. Raises NotImplementedError. Returns ------- @@ -683,6 +702,14 @@ class RemoteDBConnection(DBConnection): LanceTable(table4) """ + for option, value in ( + ("data_storage_version", data_storage_version), + ("enable_v2_manifest_paths", enable_v2_manifest_paths), + ): + if value is not None: + raise NotImplementedError( + f"{option} is not supported for remote table creation" + ) if exist_ok: if mode == "create": mode = "exist_ok" @@ -704,6 +731,7 @@ class RemoteDBConnection(DBConnection): on_bad_vectors=on_bad_vectors, fill_value=fill_value, embedding_functions=embedding_functions, + storage_options=storage_options, ) ) return RemoteTable( @@ -843,6 +871,20 @@ class RemoteDBConnection(DBConnection): job = LOOP.run(self._conn.drop_table_async(name, namespace_path=namespace_path)) return Job(job if isinstance(job, AsyncJob) else AsyncJob(job)) + @override + def drop_all_tables(self, namespace_path: Optional[List[str]] = None): + """Drop all tables from the database. + + The remote API does not currently support this operation. Use + ``list_tables()`` and ``drop_table()`` to drop tables individually. + """ + LOOP.run(self._conn.drop_all_tables(namespace_path=namespace_path)) + + @override + def drop_database(self): + """Deprecated alias for ``drop_all_tables()``.""" + LOOP.run(self._conn.drop_database()) + @override def rename_table( self, @@ -1122,7 +1164,7 @@ class RemoteDBConnection(DBConnection): ) @override - def describe_query(self, query_id: UUID) -> QueryDescription: + def describe_query(self, query_id: Union[UUID, str]) -> QueryDescription: """Describe a submitted SQL query by its connection-scoped id.""" return LOOP.run( self._conn.describe_query( diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 86e8fb8a0..6b7835470 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -4,9 +4,10 @@ from datetime import timedelta import deprecation import logging -from functools import cached_property +from functools import cached_property, wraps import os from typing import ( + TYPE_CHECKING, Any, Callable, Dict, @@ -41,19 +42,12 @@ from lancedb.index import ( FTS, BTree, Bitmap, - HnswFlat, - HnswPq, - HnswSq, - IvfFlat, - IvfPq, - IvfRq, - IvfSq, LabelList, ) from lancedb.job import Job from lancedb.functions import FunctionApplication, RefreshColumnResult from lancedb.remote.db import LOOP -from lancedb.table import IndexConfigType, KNOWN_METRICS +from lancedb.table import IndexConfigType, _maybe_add_fts_error_note import pyarrow as pa from lancedb.common import DATA, VEC, VECTOR_COLUMN_NAME @@ -68,7 +62,6 @@ from ..query import ( FullTextQuery, LanceQueryBuilder, LanceTakeQueryBuilder, - LanceVectorQueryBuilder, ) from ..table import ( AsyncTable, @@ -79,9 +72,35 @@ from ..table import ( Table, Tags, ) -from ..types import BaseTokenizerType +from ..types import BaseTokenizerType, ScalarIndexType from ..util import infer_vector_column_name +if TYPE_CHECKING: + import lance + import polars as pl + + +def _deprecated_keyword(old_name: str, new_name: str): + """Keep old remote-only keyword names without changing the Table signature.""" + + def decorate(method): + @wraps(method) + def wrapped(self, *args, **kwargs): + if old_name in kwargs: + if new_name in kwargs: + raise TypeError(f"Cannot specify both {old_name} and {new_name}") + warnings.warn( + f"{old_name} is deprecated; use {new_name} instead", + DeprecationWarning, + stacklevel=2, + ) + kwargs[new_name] = kwargs.pop(old_name) + return method(self, *args, **kwargs) + + return wrapped + + return decorate + class RemoteTable(Table): def __init__( @@ -173,6 +192,18 @@ class RemoteTable(Table): """The name of the table""" return self._name + @property + def namespace(self) -> List[str]: + return self._namespace_path + + @deprecation.deprecated( + deprecated_in="0.33.1", + current_version=__version__, + details="Use update_field_metadata() instead.", + ) + def replace_field_metadata(self, field_name: str, new_metadata: Dict[str, str]): + LOOP.run(self._table.replace_field_metadata(field_name, new_metadata)) + def __repr__(self) -> str: return f"RemoteTable({self.db_name}.{self.name})" @@ -243,12 +274,32 @@ class RemoteTable(Table): def to_arrow(self, *, blob_mode: BlobMode = "lazy") -> pa.Table: """to_arrow() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_arrow() is not supported for remote tables.") + raise NotImplementedError( + "to_arrow() is not supported for remote tables. " + "Use table.search().to_arrow() instead." + ) def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs): """to_pandas() is not supported for remote tables.""" validate_blob_mode(blob_mode) - raise NotImplementedError("to_pandas() is not supported for remote tables.") + raise NotImplementedError( + "to_pandas() is not supported for remote tables. " + "Use table.search().to_pandas() instead." + ) + + def to_lance(self, **kwargs) -> "lance.LanceDataset": + """Direct Lance dataset access is not supported for remote tables. + + Use ``table.search().to_arrow()`` to read query results through the server. + """ + return LOOP.run(self._table.to_lance(**kwargs)) + + def to_polars(self, batch_size=None) -> "pl.LazyFrame": + """Use ``table.search().to_polars()`` to read remote query results.""" + raise NotImplementedError( + "to_polars() is not supported for remote tables. " + "Use table.search().to_polars() instead." + ) def checkout(self, version: Union[int, str]): result = LOOP.run(self._table.checkout(version)) @@ -290,9 +341,10 @@ class RemoteTable(Table): self._table.tokenize(query, column=column, index_name=index_name) ) - def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]: + @_deprecated_keyword("index_uuid", "index_name") + def index_stats(self, index_name: str) -> Optional[IndexStatistics]: """List all the stats of a specified index""" - return LOOP.run(self._table.index_stats(index_uuid)) + return LOOP.run(self._table.index_stats(index_name)) @deprecation.deprecated( deprecated_in="0.25.0", @@ -302,43 +354,38 @@ class RemoteTable(Table): def create_scalar_index( self, column: str, - index_type: Literal["BTREE", "BITMAP", "LABEL_LIST", "scalar"] = "scalar", *, - replace: bool = False, + replace: bool = True, + index_type: ScalarIndexType = "BTREE", wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): - """Creates a scalar index. + """Create a scalar index on a column. .. deprecated:: 0.25.0 Use :meth:`create_index` with a BTree, Bitmap, or LabelList config instead. Example: ``table.create_index("column", config=BTree())`` - - Parameters - ---------- - column : str - The column to be indexed. Must be a boolean, integer, float, - or string column. - index_type : str - The index type of the scalar index. Must be "scalar" (BTREE), - "BTREE", "BITMAP", or "LABEL_LIST", - replace : bool - If True, replace the existing index with the new one. """ - if index_type == "scalar" or index_type == "BTREE": + if index_type == "scalar": + warnings.warn( + 'index_type="scalar" is deprecated; use "BTREE" instead', + DeprecationWarning, + stacklevel=2, + ) + index_type = "BTREE" + if index_type == "BTREE": config = BTree() elif index_type == "BITMAP": config = Bitmap() elif index_type == "LABEL_LIST": config = LabelList() else: - raise ValueError(f"Unknown index type: {index_type}") - - LOOP.run( + raise ValueError(f"Unknown index type {index_type}") + return LOOP.run( self._table.create_index( column, - config=config, replace=replace, + config=config, wait_timeout=wait_timeout, name=name, ) @@ -349,12 +396,16 @@ class RemoteTable(Table): current_version=__version__, details="Use create_index() with config=FTS() instead.", ) + @_deprecated_keyword("column", "field_names") def create_fts_index( self, - column: str, + field_names: Union[str, List[str]], *, + ordering_field_names: Optional[Union[str, List[str]]] = None, replace: bool = False, - wait_timeout: Optional[timedelta] = None, + writer_heap_size: Optional[int] = 1024 * 1024 * 1024, + use_tantivy: bool = False, + tokenizer_name: Optional[str] = None, with_position: bool = False, # tokenizer configs: base_tokenizer: BaseTokenizerType = "simple", @@ -370,6 +421,7 @@ class RemoteTable(Table): prefix_only: bool = False, block_size: int = 128, document_granularity: DocumentGranularity = DocumentGranularity.ROW, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a full-text search index on a column. @@ -378,31 +430,70 @@ class RemoteTable(Table): Use :meth:`create_index` with an FTS config instead. Example: ``table.create_index("text_column", config=FTS())`` """ + if use_tantivy: + raise ValueError( + "Tantivy-based FTS has been removed. " + "Remove use_tantivy and recreate the index with native FTS." + ) + if ordering_field_names is not None: + raise ValueError( + "ordering_field_names was only supported by the removed " + "Tantivy-based FTS implementation." + ) + if writer_heap_size != 1024 * 1024 * 1024: + raise ValueError( + "writer_heap_size was only supported by the removed " + "Tantivy-based FTS implementation." + ) + if not isinstance(field_names, str): + raise ValueError( + "Native FTS indexes can only be created on a single field " + "at a time. To search over multiple text fields, create a " + "separate FTS index for each field." + ) + + if tokenizer_name is None: + tokenizer_configs = { + "base_tokenizer": base_tokenizer, + "language": language, + "with_position": with_position, + "max_token_length": max_token_length, + "lower_case": lower_case, + "stem": stem, + "remove_stop_words": remove_stop_words, + "custom_stop_words": custom_stop_words, + "ascii_folding": ascii_folding, + "ngram_min_length": ngram_min_length, + "ngram_max_length": ngram_max_length, + "prefix_only": prefix_only, + } + else: + tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name) + tokenizer_configs["custom_stop_words"] = custom_stop_words + config = FTS( - with_position=with_position, - base_tokenizer=base_tokenizer, - language=language, - max_token_length=max_token_length, - lower_case=lower_case, - stem=stem, - remove_stop_words=remove_stop_words, - custom_stop_words=custom_stop_words, - ascii_folding=ascii_folding, - ngram_min_length=ngram_min_length, - ngram_max_length=ngram_max_length, - prefix_only=prefix_only, block_size=block_size, document_granularity=document_granularity, + **tokenizer_configs, ) - LOOP.run( - self._table.create_index( - column, - config=config, - replace=replace, - wait_timeout=wait_timeout, - name=name, + + try: + LOOP.run( + self._table.create_index( + field_names, + replace=replace, + config=config, + name=name, + wait_timeout=wait_timeout, + ) ) - ) + except (ValueError, RuntimeError) as e: + _maybe_add_fts_error_note( + e, + base_tokenizer=config.base_tokenizer, + language=config.language, + ) + raise e # New unified API overload @overload @@ -412,6 +503,7 @@ class RemoteTable(Table): /, *, config: IndexConfigType, + replace: bool = ..., wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., @@ -422,38 +514,56 @@ class RemoteTable(Table): def create_index( self, metric: Literal["l2", "cosine", "dot", "hamming"] = ..., - vector_column_name: str = ..., - index_cache_size: Optional[int] = ..., num_partitions: Optional[int] = ..., num_sub_vectors: Optional[int] = ..., - replace: Optional[bool] = ..., + vector_column_name: str = ..., + replace: bool = ..., accelerator: Optional[str] = ..., - index_type: Literal[ - "VECTOR", "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" - ] = ..., - wait_timeout: Optional[timedelta] = ..., - *, + index_cache_size: Optional[int] = ..., num_bits: int = ..., + index_type: Literal[ + "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" + ] = ..., + max_iterations: int = ..., + sample_rate: int = ..., + m: int = ..., + ef_construction: int = ..., + *, + wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., + target_partition_size: Optional[int] = ..., ) -> None: ... def create_index( self, metric: str = "l2", - vector_column_name: str = VECTOR_COLUMN_NAME, - index_cache_size: Optional[int] = None, num_partitions: Optional[int] = None, num_sub_vectors: Optional[int] = None, - replace: Optional[bool] = None, + vector_column_name: str = VECTOR_COLUMN_NAME, + replace: bool = True, accelerator: Optional[str] = None, - index_type="vector", - wait_timeout: Optional[timedelta] = None, - *, + index_cache_size: Optional[int] = None, num_bits: int = 8, + index_type: Literal[ + "IVF_FLAT", + "IVF_SQ", + "IVF_PQ", + "IVF_RQ", + "IVF_HNSW_SQ", + "IVF_HNSW_PQ", + "IVF_HNSW_FLAT", + ] = "IVF_PQ", + max_iterations: int = 50, + sample_rate: int = 256, + m: int = 20, + ef_construction: int = 300, + *, config: Optional[IndexConfigType] = None, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, train: bool = True, + target_partition_size: Optional[int] = None, ): """Create an index on a column. @@ -464,6 +574,24 @@ class RemoteTable(Table): argument plus separate ``vector_column_name`` / ``num_partitions`` / etc. parameters, and emits a ``DeprecationWarning``. + Parameters + ---------- + metric : str + For new API: the column name to index. + For legacy API: the distance metric ("l2", "cosine", "dot", "hamming"). + config : IndexConfigType, optional + The index configuration object. If provided, uses the new unified API. + Can be one of: IvfFlat, IvfPq, IvfSq, IvfRq, HnswPq, HnswSq, + BTree, Bitmap, LabelList, Fm, FTS. + replace : bool, default True + Whether to replace an existing index on this column. + wait_timeout : timedelta, optional + Timeout to wait for async indexing to complete. + name : str, optional + Custom name for the index. + train : bool, default True + Whether to train the index with existing data. + Examples -------- New API (recommended): @@ -489,7 +617,6 @@ class RemoteTable(Table): vector_column_name, accelerator, index_cache_size, - replace, ) if is_legacy: @@ -505,6 +632,7 @@ class RemoteTable(Table): stacklevel=2, ) + # Legacy API: first arg is the distance metric column = vector_column_name if accelerator is not None: @@ -513,45 +641,27 @@ class RemoteTable(Table): "and will be ignored." ) - idx_type = index_type.upper() - if idx_type == "VECTOR" or idx_type == "IVF_PQ": - config = IvfPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - ) - elif idx_type == "IVF_RQ": - config = IvfRq( - distance_type=metric, - num_partitions=num_partitions, - num_bits=num_bits, - ) - elif idx_type == "IVF_SQ": - config = IvfSq(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_HNSW_PQ": - config = HnswPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - ) - elif idx_type == "IVF_HNSW_SQ": - config = HnswSq(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_HNSW_FLAT": - config = HnswFlat(distance_type=metric, num_partitions=num_partitions) - elif idx_type == "IVF_FLAT": - config = IvfFlat(distance_type=metric, num_partitions=num_partitions) - else: - raise ValueError( - f"Unknown vector index type: {idx_type}. Valid options are" - " 'IVF_FLAT', 'IVF_PQ', 'IVF_RQ', 'IVF_SQ'," - " 'IVF_HNSW_PQ', 'IVF_HNSW_SQ', 'IVF_HNSW_FLAT'" - ) + # Build config from legacy parameters + config = self._build_vector_config_from_legacy_params( + metric=metric, + index_type="IVF_PQ" + if index_type.upper() == "VECTOR" + else index_type.upper(), + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=None, + ) + else: column = metric - LOOP.run( + return LOOP.run( self._table.create_index( column, replace=replace, @@ -590,37 +700,6 @@ class RemoteTable(Table): ) ) - def _is_legacy_create_index_call( - self, - first_arg: str, - config: Optional[IndexConfigType], - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - vector_column_name: str, - accelerator: Optional[str], - index_cache_size: Optional[int], - replace: Optional[bool], - ) -> bool: - """Detect if this is a legacy create_index call.""" - if config is not None: - return False - if any( - x is not None - for x in ( - num_partitions, - num_sub_vectors, - accelerator, - index_cache_size, - replace, - ) - ): - return True - if vector_column_name != VECTOR_COLUMN_NAME: - return True - if first_arg.lower() in KNOWN_METRICS: - return True - return False - def add( self, data: DATA, @@ -692,9 +771,10 @@ class RemoteTable(Table): query: Union[VEC, str] = None, vector_column_name: Optional[str] = None, query_type="auto", + ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, fast_search: Optional[bool] = None, - ) -> LanceVectorQueryBuilder: + ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -785,6 +865,7 @@ class RemoteTable(Table): query, query_type, vector_column_name=vector_column_name, + ordering_field_name=ordering_field_name, fts_columns=fts_columns or [], fast_search=fast_search, ) @@ -844,7 +925,8 @@ class RemoteTable(Table): self._table._do_merge(merge, new_data, on_bad_vectors, fill_value) ) - def delete(self, predicate: str) -> DeleteResult: + @_deprecated_keyword("predicate", "where") + def delete(self, where: Union[str, Expr]) -> DeleteResult: """Delete rows from the table. This can be used to delete a single row, many rows, all rows, or @@ -852,7 +934,7 @@ class RemoteTable(Table): Parameters ---------- - predicate: str + where: Union[str, Expr] The SQL where clause to use when deleting rows. - For example, 'x = 2' or 'x IN (1, 2, 3)'. @@ -896,7 +978,7 @@ class RemoteTable(Table): x vector _distance # doctest: +SKIP 0 2 [3.0, 4.0] 85.0 # doctest: +SKIP """ - return LOOP.run(self._table.delete(predicate)) + return LOOP.run(self._table.delete(where)) def update( self, @@ -958,7 +1040,12 @@ class RemoteTable(Table): self._table.update(where=where, updates=values, updates_sql=values_sql) ) - def cleanup_old_versions(self, *_): + def cleanup_old_versions( + self, + older_than: Optional[timedelta] = None, + *, + delete_unverified: bool = False, + ): """ cleanup_old_versions() is a no-op for remote tables. """ @@ -968,7 +1055,7 @@ class RemoteTable(Table): ) pass - def compact_files(self, *_): + def compact_files(self, *args, **kwargs): """ compact_files() is a no-op for remote tables. """ @@ -983,12 +1070,14 @@ class RemoteTable(Table): *, cleanup_older_than: Optional[timedelta] = None, delete_unverified: bool = False, + retrain: bool = False, ): """optimize() is not supported for remote tables.""" return LOOP.run( self._table.optimize( cleanup_older_than=cleanup_older_than, delete_unverified=delete_unverified, + retrain=retrain, ) ) @@ -1096,8 +1185,9 @@ class RemoteTable(Table): """No-op for remote tables (no local shard writers).""" return LOOP.run(self._table.close_rte_writers()) - def drop_index(self, index_name: str): - return LOOP.run(self._table.drop_index(index_name)) + @_deprecated_keyword("index_name", "name") + def drop_index(self, name: str) -> None: + return LOOP.run(self._table.drop_index(name)) def prewarm_index(self, name: str) -> None: """Prewarm an index in the table. @@ -1154,11 +1244,21 @@ class RemoteTable(Table): """ return LOOP.run(self._table.uri()) - def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + if with_row_id: + builder.with_row_id() + return builder - def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + if with_row_id: + builder.with_row_id() + return builder def uses_v2_manifest_paths(self) -> bool: return LOOP.run(self._table.uses_v2_manifest_paths()) @@ -1186,6 +1286,11 @@ class RemoteTable(Table): ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + return LOOP.run(self._table.open_blob_file(descriptor)) + def head(self, n=5) -> pa.Table: """ Return the first `n` rows of the table. diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index da0394d1c..75efbbfc7 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -21,7 +21,7 @@ class OpenaiReranker(Reranker): Parameters ---------- - model_name : str, default "gpt-4-turbo-preview" + model_name : str, default "gpt-5.6-sol" The name of the cross encoder model to use. column : str, default "text" The name of the column to use as input to the cross encoder model. @@ -33,7 +33,7 @@ class OpenaiReranker(Reranker): def __init__( self, - model_name: str = "gpt-4-turbo-preview", + model_name: str = "gpt-5.6-sol", column: str = "text", return_score="relevance", api_key: Optional[str] = None, @@ -54,7 +54,6 @@ class OpenaiReranker(Reranker): response = self._client.chat.completions.create( model=self.model_name, response_format={"type": "json_object"}, - temperature=0, messages=[ { "role": "system", @@ -132,4 +131,4 @@ class OpenaiReranker(Reranker): "OPENAI_API_KEY not set. Either set it in your environment or \ pass it as `api_key` argument to the CohereReranker." ) - return openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY") or self.api_key) + return openai.OpenAI(api_key=self.api_key or os.environ.get("OPENAI_API_KEY")) diff --git a/python/python/lancedb/schema.py b/python/python/lancedb/schema.py index 5edda5610..8bec17058 100644 --- a/python/python/lancedb/schema.py +++ b/python/python/lancedb/schema.py @@ -4,6 +4,7 @@ """Schema helpers for Lance blob columns.""" +import ctypes import importlib from typing import TYPE_CHECKING, Optional @@ -16,6 +17,7 @@ if TYPE_CHECKING: _BLOB_EXTENSION_NAME = "lance.blob.v2" _BLOB_V1_KEY = "lance-encoding:blob" _ARROW_EXT_NAME_KEY = "ARROW:extension:name" +_MAX_RUST_USIZE = ctypes.c_size_t(-1).value _BLOB_V2_STORAGE_TYPE = pa.struct( [ pa.field("data", pa.large_binary(), nullable=True), @@ -209,32 +211,44 @@ def blob( Maximum size in bytes of a single packed blob sidecar file before a new one is started. Must be positive. - The thresholds are forwarded to ``lance.blob_field``, which validates - them and stores them as field metadata, so pylance is required whenever - any threshold is set. When all thresholds are omitted the behavior is - unchanged. + Thresholds are validated and stored as Arrow field metadata, with the + same format as ``lance.blob_field``. They do not require pylance. """ - if ( - inline_size_threshold is None - and dedicated_size_threshold is None - and pack_file_size_threshold is None + metadata = {} + for threshold, value, allow_zero, key in ( + ( + "inline_size_threshold", + inline_size_threshold, + True, + "lance-encoding:blob-inline-size-threshold", + ), + ( + "dedicated_size_threshold", + dedicated_size_threshold, + False, + "lance-encoding:blob-dedicated-size-threshold", + ), + ( + "pack_file_size_threshold", + pack_file_size_threshold, + False, + "lance-encoding:blob-pack-file-size-threshold", + ), ): - blob_type = _resolve_blob_type() - return pa.field(name, blob_type(), nullable=nullable) - try: - lance = importlib.import_module("lance") - except ModuleNotFoundError as err: - raise ValueError( - "Setting blob tier thresholds requires pylance to be installed" - ) from err - _resolve_blob_type() - return lance.blob_field( - name, - nullable=nullable, - inline_size_threshold=inline_size_threshold, - dedicated_size_threshold=dedicated_size_threshold, - pack_file_size_threshold=pack_file_size_threshold, - ) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{threshold} must be an int, got {type(value).__name__}") + if value < 0 or (value == 0 and not allow_zero): + bound = "non-negative" if allow_zero else "positive" + raise ValueError(f"{threshold} must be {bound}") + if value > _MAX_RUST_USIZE: + raise OverflowError(f"{threshold} must fit in a Rust usize") + metadata[key] = str(value) + + blob_type = _resolve_blob_type() + field = pa.field(name, blob_type(), nullable=nullable) + return field.with_metadata(metadata) if metadata else field def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType: diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 6f7cf6f17..4c116f857 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -17,9 +17,11 @@ from functools import cached_property from typing import ( TYPE_CHECKING, Any, + AsyncIterator, Callable, Dict, Iterable, + Iterator, List, Literal, Optional, @@ -35,6 +37,7 @@ from lancedb.scannable import _hf_batches, _register_optional_converters, to_sca from . import __version__ from ._blob import ( BlobFile, + _normalize_blob_descriptor, _normalize_blob_row_ids, _wrap_blob_files, strip_auto_row_ids, @@ -264,10 +267,8 @@ if TYPE_CHECKING: OnBadVectorsType, AddMode, CreateMode, - VectorIndexType, ScalarIndexType, BaseTokenizerType, - DistanceType, ) # Type alias for index configuration objects @@ -546,10 +547,11 @@ def _sanitize_data( metadata : Optional[dict], default None The embedding metadata to add to the schema. on_bad_vectors : Literal["error", "drop", "fill", "null"], default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. fill_value : float, default 0.0 The value to use when filling vectors. Only used if on_bad_vectors="fill". - All entries in the vector will be set to this value. + NaN or null elements and missing trailing elements are replaced with this + value. Valid elements are preserved. """ # At this point, the table might not match the schema we are targeting: # 1. There might be embedding columns missing that will be added @@ -1358,13 +1360,23 @@ class Table(ABC): @abstractmethod def name(self) -> str: """The name of this Table""" - raise NotImplementedError + raise NotImplementedError("name is not supported for this table type") + + @property + def namespace(self) -> List[str]: + """Return the namespace path of the table (empty for the root namespace).""" + return [] + + @property + def id(self) -> str: + """Return the full identifier of the table (namespace$name).""" + return "$".join(self.namespace + [self.name]) @property @abstractmethod def version(self) -> int: """The version of this Table""" - raise NotImplementedError + raise NotImplementedError("version is not supported for this table type") @property @abstractmethod @@ -1373,7 +1385,7 @@ class Table(ABC): of this Table """ - raise NotImplementedError + raise NotImplementedError("schema is not supported for this table type") @property @abstractmethod @@ -1402,7 +1414,7 @@ class Table(ABC): tags = table.tags.list() """ - raise NotImplementedError + raise NotImplementedError("tags is not supported for this table type") @property def branches(self) -> "Branches": @@ -1411,11 +1423,11 @@ class Table(ABC): Branches are isolated, writable lines of history forked from another branch (or version). Writes on a branch do not affect ``main``. """ - raise NotImplementedError + raise NotImplementedError("branches is not supported for this table type") def current_branch(self) -> Optional[str]: """The branch this table handle is scoped to, or ``None`` for ``main``.""" - raise NotImplementedError + raise NotImplementedError("current_branch is not supported for this table type") def __len__(self) -> int: """The number of rows in this Table""" @@ -1438,7 +1450,7 @@ class Table(ABC): filter: str, optional A SQL where clause to filter the rows to count. """ - raise NotImplementedError + raise NotImplementedError("count_rows is not supported for this table type") def to_pandas(self, blob_mode: BlobMode = "lazy", **kwargs) -> "pandas.DataFrame": """Return the table as a pandas DataFrame. @@ -1473,20 +1485,32 @@ class Table(ABC): ------- pa.Table """ - raise NotImplementedError + raise NotImplementedError("to_arrow is not supported for this table type") def to_lance(self, **kwargs) -> lance.LanceDataset: - """Return the table as a lance.LanceDataset. + """Return a local table as a lance.LanceDataset. + + Remote tables cannot be opened as Lance datasets by the client. Use + ``table.search().to_arrow()`` to read them through the server. Returns ------- lance.LanceDataset """ - raise NotImplementedError + raise NotImplementedError("to_lance is not supported for this table type") - def to_polars(self, **kwargs) -> "pl.LazyFrame": + def to_polars(self, batch_size=None) -> "pl.LazyFrame": """Return the table as a Polars LazyFrame. + For remote tables, use ``table.search().to_polars()`` to convert query + results to a Polars DataFrame. + + Parameters + ---------- + batch_size: int, optional + Passed to Polars as the maximum row count for scanned PyArrow + record batches. + Note ---- The Polars streaming engine is not supported because it does not currently @@ -1497,7 +1521,194 @@ class Table(ABC): ------- polars.LazyFrame """ - raise NotImplementedError + raise NotImplementedError("to_polars is not supported for this table type") + + def _is_legacy_create_index_call( + self, + first_arg: str, + config: Optional[IndexConfigType], + num_partitions: Optional[int], + num_sub_vectors: Optional[int], + vector_column_name: str, + accelerator: Optional[str], + index_cache_size: Optional[int], + ) -> bool: + """Detect if this is a legacy create_index call.""" + # If config is provided, it's definitely the new API + if config is not None: + return False + + # If old-style parameters were explicitly set, it's legacy + if any( + x is not None + for x in (num_partitions, num_sub_vectors, accelerator, index_cache_size) + ): + return True + + # If vector_column_name differs from default, it's legacy + if vector_column_name != VECTOR_COLUMN_NAME: + return True + + # If first arg is a known metric, assume legacy + if first_arg.lower() in KNOWN_METRICS: + return True + + # Otherwise assume new API + return False + + def _build_vector_config_from_legacy_params( + self, + metric: str, + index_type: str, + num_partitions: Optional[int], + num_sub_vectors: Optional[int], + num_bits: int, + max_iterations: int, + sample_rate: int, + m: int, + ef_construction: int, + target_partition_size: Optional[int], + accelerator: Optional[str], + ) -> IndexConfigType: + """Build an index config object from legacy parameters.""" + if index_type == "IVF_FLAT": + return IvfFlat( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_SQ": + return IvfSq( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_PQ": + return IvfPq( + distance_type=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_RQ": + return IvfRq( + distance_type=metric, + num_partitions=num_partitions, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_PQ": + return HnswPq( + distance_type=metric, + num_partitions=num_partitions, + num_sub_vectors=num_sub_vectors, + num_bits=num_bits, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_SQ": + return HnswSq( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + accelerator=accelerator, + ) + elif index_type == "IVF_HNSW_FLAT": + return HnswFlat( + distance_type=metric, + num_partitions=num_partitions, + max_iterations=max_iterations, + sample_rate=sample_rate, + m=m, + ef_construction=ef_construction, + target_partition_size=target_partition_size, + ) + else: + raise ValueError(f"Unknown index type {index_type}") + + @staticmethod + def infer_tokenizer_configs(tokenizer_name: str) -> dict: + if tokenizer_name == "default": + return { + "base_tokenizer": "simple", + "language": "English", + "max_token_length": 40, + "lower_case": True, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + elif tokenizer_name == "raw": + return { + "base_tokenizer": "raw", + "language": "English", + "max_token_length": None, + "lower_case": False, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + elif tokenizer_name == "whitespace": + return { + "base_tokenizer": "whitespace", + "language": "English", + "max_token_length": None, + "lower_case": False, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } + + # or it's with language stemming with pattern like "en_stem" + if len(tokenizer_name) != 7: + raise ValueError(f"Invalid tokenizer name {tokenizer_name}") + lang = tokenizer_name[:2] + if tokenizer_name[-5:] != "_stem": + raise ValueError(f"Invalid tokenizer name {tokenizer_name}") + if lang not in lang_mapping: + raise ValueError(f"Invalid language code {lang}") + return { + "base_tokenizer": "simple", + "language": lang_mapping[lang], + "max_token_length": 40, + "lower_case": True, + "stem": True, + "remove_stop_words": False, + "ascii_folding": False, + "ngram_min_length": 3, + "ngram_max_length": 3, + "prefix_only": False, + } # New unified API overload @overload @@ -1524,14 +1735,16 @@ class Table(ABC): replace: bool = ..., accelerator: Optional[str] = ..., index_cache_size: Optional[int] = ..., - *, - index_type: VectorIndexType = ..., - wait_timeout: Optional[timedelta] = ..., num_bits: int = ..., + index_type: Literal[ + "IVF_FLAT", "IVF_SQ", "IVF_PQ", "IVF_RQ", "IVF_HNSW_SQ", "IVF_HNSW_PQ" + ] = ..., max_iterations: int = ..., sample_rate: int = ..., m: int = ..., ef_construction: int = ..., + *, + wait_timeout: Optional[timedelta] = ..., name: Optional[str] = ..., train: bool = ..., target_partition_size: Optional[int] = ..., @@ -1539,22 +1752,30 @@ class Table(ABC): def create_index( self, - metric: DistanceType = "l2", + metric: str = "l2", num_partitions: Optional[int] = None, num_sub_vectors: Optional[int] = None, vector_column_name: str = VECTOR_COLUMN_NAME, replace: bool = True, accelerator: Optional[str] = None, index_cache_size: Optional[int] = None, - *, - index_type: VectorIndexType = "IVF_PQ", - wait_timeout: Optional[timedelta] = None, num_bits: int = 8, + index_type: Literal[ + "IVF_FLAT", + "IVF_SQ", + "IVF_PQ", + "IVF_RQ", + "IVF_HNSW_SQ", + "IVF_HNSW_PQ", + "IVF_HNSW_FLAT", + ] = "IVF_PQ", max_iterations: int = 50, sample_rate: int = 256, m: int = 20, ef_construction: int = 300, + *, config: Optional[IndexConfigType] = None, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, train: bool = True, target_partition_size: Optional[int] = None, @@ -1602,7 +1823,7 @@ class Table(ABC): ... "l2", vector_column_name="vector" ... ) """ - raise NotImplementedError + raise NotImplementedError("create_index is not supported for this table type") def create_index_async( self, @@ -1620,7 +1841,9 @@ class Table(ABC): be complete when returned; callers must not assume the index exists until :meth:`Job.wait` returns. """ - raise NotImplementedError + raise NotImplementedError( + "create_index_async is not supported for this table type" + ) def drop_index(self, name: str) -> None: """ @@ -1640,7 +1863,7 @@ class Table(ABC): Use [list_indices][lancedb.table.Table.list_indices] to find the names of the indices. """ - raise NotImplementedError + raise NotImplementedError("drop_index is not supported for this table type") def wait_for_index( self, index_names: Iterable[str], timeout: timedelta = timedelta(seconds=300) @@ -1657,14 +1880,14 @@ class Table(ABC): timeout: timedelta Timeout to wait for asynchronous indexing. The default is 5 minutes. """ - raise NotImplementedError + raise NotImplementedError("wait_for_index is not supported for this table type") @abstractmethod def stats(self) -> TableStatistics: """ Retrieve table and fragment statistics. """ - raise NotImplementedError + raise NotImplementedError("stats is not supported for this table type") @abstractmethod def create_scalar_index( @@ -1729,7 +1952,9 @@ class Table(ABC): ``my_col = 0 OR not_indexed = 1`` will not be able to use any scalar index on ``my_col``. """ - raise NotImplementedError + raise NotImplementedError( + "create_scalar_index is not supported for this table type" + ) def create_fts_index( self, @@ -1855,7 +2080,9 @@ class Table(ABC): ``LANCE_LANGUAGE_MODEL_HOME`` to override the default platform data directory under ``lance/language_models``. """ - raise NotImplementedError + raise NotImplementedError( + "create_fts_index is not supported for this table type" + ) @abstractmethod def add( @@ -1885,8 +2112,8 @@ class Table(ABC): "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. - One of "error", "drop", "fill". + How to handle vectors with incorrect sizes, NaNs, or null elements. + One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". progress: bool, callable, or tqdm-like, optional @@ -1929,7 +2156,7 @@ class Table(ABC): AddResult An object containing the new version number of the table after adding data. """ - raise NotImplementedError + raise NotImplementedError("add is not supported for this table type") def merge_insert(self, on: Union[str, Iterable[str]]) -> LanceMergeInsertBuilder: """ @@ -2005,6 +2232,7 @@ class Table(ABC): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -2081,7 +2309,7 @@ class Table(ABC): - and also the "_distance" column which is the distance between the query vector and the returned vector. """ - raise NotImplementedError + raise NotImplementedError("search is not supported for this table type") @abstractmethod def take_offsets( @@ -2105,11 +2333,14 @@ class Table(ABC): ---------- offsets: list[int] The offsets to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - pa.RecordBatch - A record batch containing the rows at the given offsets. + LanceTakeQueryBuilder + A query builder that can be executed to get the rows at the given + offsets. """ def __getitems__(self, offsets: list[int]) -> pa.RecordBatch: @@ -2190,11 +2421,13 @@ class Table(ABC): ---------- row_ids: list[int] The row ids to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - AsyncTakeQuery - A query object that can be executed to get the rows. + LanceTakeQueryBuilder + A query builder that can be executed to get the rows. """ @abstractmethod @@ -2274,6 +2507,44 @@ class Table(ABC): handles are returned. """ + @abstractmethod + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + """Open a seekable file from a SQL ``blob_descriptor`` result. + + Accepts the descriptor dictionary or an Arrow struct scalar. Its source + pins the table version, branch, column and row identity, so opening does + not change this table's checkout. Null descriptors return ``None``. + Payload bytes load only when the returned handle is read. An external + blob with an unspecified size may require a metadata request on opening. + + ```python + rows = db.execute_query( + "SELECT blob_descriptor(video) AS video FROM videos" + ).read_all() + with videos.open_blob_file(rows["video"][0]) as video: + header = video.read(4096) + ``` + """ + + def open_blob_files(self, descriptors: Iterable) -> Iterator[Optional[BlobFile]]: + """Iterate over lazy files from descriptor dictionaries or Arrow scalars. + + Accepts an Arrow array/chunked array or any iterable. Opens each handle + as iteration advances, preserving input order and nulls as ``None``. + Close each non-null handle after use, normally with a ``with`` block. + + ```python + for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + header = video.read(4096) + ``` + """ + for descriptor in descriptors: + yield self.open_blob_file(descriptor) + @abstractmethod def _execute_query( self, @@ -2362,7 +2633,7 @@ class Table(ABC): x vector 0 3 [5.0, 6.0] """ - raise NotImplementedError + raise NotImplementedError("delete is not supported for this table type") @abstractmethod def update( @@ -2429,7 +2700,7 @@ class Table(ABC): 1 4 [5.0, 6.0] 2 3 [10.0, 10.0] """ - raise NotImplementedError + raise NotImplementedError("update is not supported for this table type") @abstractmethod def cleanup_old_versions( @@ -2847,6 +3118,29 @@ class Table(ABC): version: the new table version after the update. """ + @deprecation.deprecated( + deprecated_in="0.33.1", + current_version=__version__, + details="Use update_field_metadata() instead.", + ) + def replace_field_metadata(self, field_name: str, new_metadata: Dict[str, str]): + """Replace all metadata of a field in the schema. + + .. deprecated:: 0.33.1 + Use [update_field_metadata][lancedb.table.Table.update_field_metadata] + with ``replace=True`` instead. + + Parameters + ---------- + field_name: str + The name of the field whose metadata is replaced. + new_metadata: dict + The new metadata, replacing all existing keys. + """ + self.update_field_metadata( + {"path": field_name, "metadata": new_metadata, "replace": True} + ) + @abstractmethod def drop_columns(self, columns: Iterable[str]) -> DropColumnsResult: """ @@ -3273,13 +3567,6 @@ class LanceTable(Table): """Return the namespace path of the table.""" return self._namespace_path - @property - def id(self) -> str: - """Return the full identifier of the table (namespace$name).""" - if self._namespace_path: - return "$".join(self._namespace_path + [self.name]) - return self.name - @classmethod async def from_inner(cls, tbl: LanceDBTable): from .db import AsyncConnection, LanceDBConnection @@ -3401,11 +3688,21 @@ class LanceTable(Table): """Get the current version of the table""" return LOOP.run(self._table.version()) - def take_offsets(self, offsets: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_offsets(offsets)) + if with_row_id: + builder.with_row_id() + return builder - def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder: - return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> LanceTakeQueryBuilder: + builder = LanceTakeQueryBuilder(self._table.take_row_ids(row_ids)) + if with_row_id: + builder.with_row_id() + return builder def blob_columns(self) -> list[str]: return LOOP.run(self._table.blob_columns()) @@ -3431,6 +3728,11 @@ class LanceTable(Table): ) -> "list[Optional[BlobFile]]": return LOOP.run(self._table.fetch_blob_files(column, row_ids)) + def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + return LOOP.run(self._table.open_blob_file(descriptor)) + @property def tags(self) -> Tags: """Tag management for the table. @@ -3968,6 +4270,11 @@ class LanceTable(Table): cfg_type = type(config).__name__ lance_index_type = index_type_map.get(cfg_type, "IVF_PQ") + quantization_kwargs = {} + num_bits = getattr(config, "num_bits", 8) + if num_bits is not None: + quantization_kwargs["num_bits"] = num_bits + dataset.create_index( column=column, index_type=lance_index_type, @@ -3976,12 +4283,12 @@ class LanceTable(Table): num_sub_vectors=getattr(config, "num_sub_vectors", None), replace=replace, accelerator=acc, - num_bits=getattr(config, "num_bits", 8), m=getattr(config, "m", 20), ef_construction=getattr(config, "ef_construction", 300), target_partition_size=getattr( config, "target_partition_size", None ), + **quantization_kwargs, ) self.checkout_latest() return @@ -4025,130 +4332,6 @@ class LanceTable(Table): ) ) - def _is_legacy_create_index_call( - self, - first_arg: str, - config: Optional[IndexConfigType], - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - vector_column_name: str, - accelerator: Optional[str], - index_cache_size: Optional[int], - ) -> bool: - """Detect if this is a legacy create_index call.""" - # If config is provided, it's definitely the new API - if config is not None: - return False - - # If old-style parameters were explicitly set, it's legacy - if any( - x is not None - for x in (num_partitions, num_sub_vectors, accelerator, index_cache_size) - ): - return True - - # If vector_column_name differs from default, it's legacy - if vector_column_name != VECTOR_COLUMN_NAME: - return True - - # If first arg is a known metric, assume legacy - if first_arg.lower() in KNOWN_METRICS: - return True - - # Otherwise assume new API - return False - - def _build_vector_config_from_legacy_params( - self, - metric: str, - index_type: str, - num_partitions: Optional[int], - num_sub_vectors: Optional[int], - num_bits: int, - max_iterations: int, - sample_rate: int, - m: int, - ef_construction: int, - target_partition_size: Optional[int], - accelerator: Optional[str], - ) -> IndexConfigType: - """Build an index config object from legacy parameters.""" - if index_type == "IVF_FLAT": - return IvfFlat( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_SQ": - return IvfSq( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_PQ": - return IvfPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_RQ": - return IvfRq( - distance_type=metric, - num_partitions=num_partitions, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_PQ": - return HnswPq( - distance_type=metric, - num_partitions=num_partitions, - num_sub_vectors=num_sub_vectors, - num_bits=num_bits, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_SQ": - return HnswSq( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - accelerator=accelerator, - ) - elif index_type == "IVF_HNSW_FLAT": - return HnswFlat( - distance_type=metric, - num_partitions=num_partitions, - max_iterations=max_iterations, - sample_rate=sample_rate, - m=m, - ef_construction=ef_construction, - target_partition_size=target_partition_size, - ) - else: - raise ValueError(f"Unknown index type {index_type}") - def drop_index(self, name: str) -> None: """ Drops an index from the table @@ -4259,6 +4442,7 @@ class LanceTable(Table): *, replace: bool = True, index_type: ScalarIndexType = "BTREE", + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a scalar index on a column. @@ -4267,6 +4451,13 @@ class LanceTable(Table): Use :meth:`create_index` with a BTree, Bitmap, or LabelList config instead. Example: ``table.create_index("column", config=BTree())`` """ + if index_type == "scalar": + warnings.warn( + 'index_type="scalar" is deprecated; use "BTREE" instead', + DeprecationWarning, + stacklevel=2, + ) + index_type = "BTREE" if index_type == "BTREE": config = BTree() elif index_type == "BITMAP": @@ -4276,7 +4467,13 @@ class LanceTable(Table): else: raise ValueError(f"Unknown index type {index_type}") return LOOP.run( - self._table.create_index(column, replace=replace, config=config, name=name) + self._table.create_index( + column, + replace=replace, + config=config, + wait_timeout=wait_timeout, + name=name, + ) ) @deprecation.deprecated( @@ -4308,6 +4505,7 @@ class LanceTable(Table): prefix_only: bool = False, block_size: int = 128, document_granularity: DocumentGranularity = DocumentGranularity.ROW, + wait_timeout: Optional[timedelta] = None, name: Optional[str] = None, ): """Create a full-text search index on a column. @@ -4372,6 +4570,7 @@ class LanceTable(Table): replace=replace, config=config, name=name, + wait_timeout=wait_timeout, ) ) except (ValueError, RuntimeError) as e: @@ -4382,69 +4581,6 @@ class LanceTable(Table): ) raise e - @staticmethod - def infer_tokenizer_configs(tokenizer_name: str) -> dict: - if tokenizer_name == "default": - return { - "base_tokenizer": "simple", - "language": "English", - "max_token_length": 40, - "lower_case": True, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - elif tokenizer_name == "raw": - return { - "base_tokenizer": "raw", - "language": "English", - "max_token_length": None, - "lower_case": False, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - elif tokenizer_name == "whitespace": - return { - "base_tokenizer": "whitespace", - "language": "English", - "max_token_length": None, - "lower_case": False, - "stem": False, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - - # or it's with language stemming with pattern like "en_stem" - if len(tokenizer_name) != 7: - raise ValueError(f"Invalid tokenizer name {tokenizer_name}") - lang = tokenizer_name[:2] - if tokenizer_name[-5:] != "_stem": - raise ValueError(f"Invalid tokenizer name {tokenizer_name}") - if lang not in lang_mapping: - raise ValueError(f"Invalid language code {lang}") - return { - "base_tokenizer": "simple", - "language": lang_mapping[lang], - "max_token_length": 40, - "lower_case": True, - "stem": True, - "remove_stop_words": False, - "ascii_folding": False, - "ngram_min_length": 3, - "ngram_max_length": 3, - "prefix_only": False, - } - def add( self, data: DATA, @@ -4469,7 +4605,7 @@ class LanceTable(Table): "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -4595,6 +4731,7 @@ class LanceTable(Table): query_type: Literal["vector"] = "vector", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceVectorQueryBuilder: ... @overload @@ -4605,6 +4742,7 @@ class LanceTable(Table): query_type: Literal["fts"] = "fts", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceFtsQueryBuilder: ... @overload @@ -4617,6 +4755,7 @@ class LanceTable(Table): query_type: Literal["hybrid"] = "hybrid", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceHybridQueryBuilder: ... @overload @@ -4627,6 +4766,7 @@ class LanceTable(Table): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceEmptyQueryBuilder: ... def search( @@ -4638,6 +4778,7 @@ class LanceTable(Table): query_type: QueryType = "auto", ordering_field_name: Optional[str] = None, fts_columns: Optional[Union[str, List[str]]] = None, + fast_search: Optional[bool] = None, ) -> LanceQueryBuilder: """Create a search query to find the nearest neighbors of the given query vector. We currently support [vector search](https://lancedb.com/docs/search/vector-search/) @@ -4722,6 +4863,7 @@ class LanceTable(Table): vector_column_name=vector_column_name, ordering_field_name=ordering_field_name, fts_columns=fts_columns or [], + fast_search=fast_search, ) @classmethod @@ -4786,7 +4928,7 @@ class LanceTable(Table): otherwise just open the table, it will not add the provided data but will validate against any schema that's specified. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -5479,7 +5621,7 @@ def _handle_bad_vector_column( vector_column_name: str The name of the vector column. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0.0 The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -5498,7 +5640,8 @@ def _handle_bad_vector_column( data = data.set_column(position, vector_column_name, vec_arr) if expected_value_type is not None and ( - pa.types.is_integer(vec_arr.type.value_type) + pa.types.is_null(vec_arr.type.value_type) + or pa.types.is_integer(vec_arr.type.value_type) or pa.types.is_unsigned_integer(vec_arr.type.value_type) ): vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type)) @@ -5542,7 +5685,7 @@ def _handle_bad_vector_column( ) else: raise ValueError( - f"Vector column '{vector_column_name}' has NaNs. " + f"Vector column '{vector_column_name}' has NaNs or null values. " "Set on_bad_vectors='drop' to remove them, " "set on_bad_vectors='fill' and fill_value= to replace them, " "or set on_bad_vectors='null' to replace them with null." @@ -5578,29 +5721,11 @@ def _fill_bad_vector_values( arr = arr.combine_chunks() # A fixed-size slice truncates long vectors and pads short vectors with nulls. - # Slice an array marking the original child nulls in parallel so padding nulls - # can be distinguished from null values that were already present. + # Both padding and existing null elements need to be filled. sliced = pc.list_slice(arr, 0, dim, return_fixed_size_list=True) - child_nulls = pc.is_null(arr.values) - parent_nulls = pc.is_null(arr) - if pa.types.is_list(arr.type): - original_child_nulls = pa.ListArray.from_arrays( - arr.offsets, child_nulls, mask=parent_nulls - ) - elif pa.types.is_large_list(arr.type): - original_child_nulls = pa.LargeListArray.from_arrays( - arr.offsets, child_nulls, mask=parent_nulls - ) - else: - original_child_nulls = pa.FixedSizeListArray.from_arrays( - child_nulls, arr.type.list_size, mask=parent_nulls - ) - sliced_child_nulls = pc.list_slice( - original_child_nulls, 0, dim, return_fixed_size_list=True - ) - needs_fill = pc.is_null(sliced_child_nulls.values) values = sliced.values + needs_fill = pc.is_null(values) if pa.types.is_floating(values.type): values_for_nan_check = ( values.cast(pa.float32()) if pa.types.is_float16(values.type) else values @@ -5614,6 +5739,7 @@ def _fill_bad_vector_values( def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray: + """Find vectors with NaN or null components, including pandas-converted NaNs.""" if isinstance(arr, pa.ChunkedArray): values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks]) else: @@ -5624,6 +5750,9 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray values_has_nan = pc.is_nan(values.cast(pa.float32())) else: values_has_nan = pc.is_nan(values) + # Pandas represents NaN components as Arrow nulls. is_nan returns null + # for these values, so count them as invalid before filtering parent indices. + values_has_nan = pc.fill_null(values_has_nan, True) values_indices = pc.list_parent_indices(arr) if arr.null_count: # Match flatten(), which excludes child spans under null parent lists. @@ -6537,7 +6666,7 @@ class AsyncTable: "append" and "overwrite". Overwrite adopts explicitly declared blob layouts, while inferred bytes retain the existing blob v2 fields. on_bad_vectors: str, default "error" - What to do if any of the vectors are not the same size or contains NaNs. + How to handle vectors with incorrect sizes, NaNs, or null elements. One of "error", "drop", "fill", "null". fill_value: float, default 0. The value to use when filling vectors. Only used if on_bad_vectors="fill". @@ -7617,7 +7746,9 @@ class AsyncTable: """ await self._inner.restore(version) - def take_offsets(self, offsets: list[int]) -> AsyncTakeQuery: + def take_offsets( + self, offsets: list[int], *, with_row_id: bool = False + ) -> AsyncTakeQuery: """ Take a list of offsets from the table. @@ -7635,15 +7766,23 @@ class AsyncTable: ---------- offsets: list[int] The offsets to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- - pa.RecordBatch - A record batch containing the rows at the given offsets. + AsyncTakeQuery + A query object that can be executed to get the rows at the given + offsets. """ - return AsyncTakeQuery(self._inner.take_offsets(offsets), self) + query = AsyncTakeQuery(self._inner.take_offsets(offsets), self) + if with_row_id: + query.with_row_id() + return query - def take_row_ids(self, row_ids: list[int]) -> AsyncTakeQuery: + def take_row_ids( + self, row_ids: list[int], *, with_row_id: bool = False + ) -> AsyncTakeQuery: """ Take a list of row ids from the table. @@ -7664,13 +7803,18 @@ class AsyncTable: ---------- row_ids: list[int] The row ids to take. + with_row_id: bool, default False + If True, include the "_rowid" column in the results. Returns ------- AsyncTakeQuery A query object that can be executed to get the rows. """ - return AsyncTakeQuery(self._inner.take_row_ids(row_ids), self) + query = AsyncTakeQuery(self._inner.take_row_ids(row_ids), self) + if with_row_id: + query.with_row_id() + return query async def blob_columns(self) -> list[str]: return await self._inner.blob_columns() @@ -7701,6 +7845,36 @@ class AsyncTable: ) return _wrap_blob_files(handles) + async def open_blob_file( + self, descriptor: Union[dict, pa.StructScalar, None] + ) -> Optional[BlobFile]: + """Open a snapshot-bound descriptor as a seekable file. + + See [Table.open_blob_file][lancedb.table.Table.open_blob_file]. + """ + descriptor = _normalize_blob_descriptor(descriptor) + if descriptor is None: + return None + return BlobFile(await self._inner.open_blob_file(json.dumps(descriptor))) + + async def open_blob_files( + self, descriptors: Iterable + ) -> AsyncIterator[Optional[BlobFile]]: + """Asynchronously iterate over files in descriptor order, preserving nulls. + + Use ``async for file in table.open_blob_files(descriptors)``. Each + non-null file must be closed after use. + + ```python + async for video in videos.open_blob_files(rows["video"]): + if video is not None: + with video: + payload = await video.aread() + ``` + """ + for descriptor in descriptors: + yield await self.open_blob_file(descriptor) + @property def tags(self) -> AsyncTags: """Tag management for the dataset. diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index a043f0c8d..36536f499 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -452,6 +452,15 @@ def _(value: pa.Scalar): return value_to_sql(value.as_py()) +@value_to_sql.register(pa.MapScalar) +def _(value: pa.MapScalar): + if not value.is_valid: + return value_to_sql(None) + keys = value_to_sql(list(value.values.field(0))) + values = value_to_sql(list(value.values.field(1))) + return f"map({keys}, {values})" + + @value_to_sql.register(np.ndarray) def _(value: np.ndarray): return value_to_sql(value.tolist()) diff --git a/python/python/tests/docs/test_basic.py b/python/python/tests/docs/test_basic.py index 35d7aac10..ef68114aa 100644 --- a/python/python/tests/docs/test_basic.py +++ b/python/python/tests/docs/test_basic.py @@ -91,19 +91,13 @@ def test_quickstart(tmp_path): } ) # --8<-- [end:alter_columns_vector] - # Change it back since we can get a panic with fp16 - tbl.alter_columns( - { - "path": "vector", - "data_type": pa.list_(pa.float32(), list_size=2), - } - ) # --8<-- [start:drop_columns] tbl.drop_columns(["dbl_price"]) # --8<-- [end:drop_columns] # --8<-- [start:create_index] tbl.create_index(num_sub_vectors=1) # --8<-- [end:create_index] + tbl.search([100, 100]).limit(2).to_pandas() # --8<-- [start:delete_rows] tbl.delete("item = 'fizz'") # --8<-- [end:delete_rows] @@ -185,13 +179,6 @@ async def test_quickstart_async(tmp_path): } ) # --8<-- [end:alter_columns_async_vector] - # Change it back since we can get a panic with fp16 - await tbl.alter_columns( - { - "path": "vector", - "data_type": pa.list_(pa.float32(), list_size=2), - } - ) # --8<-- [start:drop_columns_async] await tbl.drop_columns(["dbl_price"]) # --8<-- [end:drop_columns_async] @@ -200,6 +187,7 @@ async def test_quickstart_async(tmp_path): # --8<-- [start:create_index_async] await tbl.create_index("vector") # --8<-- [end:create_index_async] + await tbl.vector_search([100, 100]).limit(2).to_pandas() # --8<-- [start:delete_rows_async] await tbl.delete("item = 'fizz'") # --8<-- [end:delete_rows_async] diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index ec1b8b27e..da362f053 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors +import ctypes import io import subprocess import sys @@ -14,6 +15,7 @@ import pytest from lance.blob import Blob, BlobType as LanceBlobType import lancedb +from lancedb import schema as schema_module from lancedb._blob import ( _set_blob_column, blob_v2_projection_sources, @@ -139,6 +141,41 @@ def test_blob_without_thresholds_has_no_threshold_metadata(): assert not any(b"threshold" in key for key in metadata) +@pytest.mark.parametrize("nullable", [False, True]) +@pytest.mark.parametrize( + "thresholds", + [ + {}, + {"inline_size_threshold": 0}, + {"inline_size_threshold": 8192}, + {"dedicated_size_threshold": 1 << 20}, + {"pack_file_size_threshold": 1 << 26}, + { + "inline_size_threshold": 8192, + "dedicated_size_threshold": 1 << 20, + "pack_file_size_threshold": 1 << 26, + }, + { + "inline_size_threshold": ctypes.c_size_t(-1).value, + "dedicated_size_threshold": ctypes.c_size_t(-1).value, + "pack_file_size_threshold": ctypes.c_size_t(-1).value, + }, + ], +) +def test_blob_thresholds_match_pylance(monkeypatch, nullable, thresholds): + expected = lance.blob_field("payload", nullable=nullable, **thresholds) + actual = lancedb.blob("payload", nullable=nullable, **thresholds) + assert type(actual.type) is LanceBlobType + assert actual.equals(expected, check_metadata=True) + + # Both extension classes must encode the same Arrow field on the wire. + monkeypatch.setattr( + schema_module, "_resolve_blob_type", lambda: schema_module._FallbackBlobType + ) + fallback = lancedb.blob("payload", nullable=nullable, **thresholds) + assert pa.schema([fallback]).serialize() == pa.schema([expected]).serialize() + + def test_blob_thresholds_are_validated(): with pytest.raises(ValueError, match="inline_size_threshold must be non-negative"): lancedb.blob("payload", inline_size_threshold=-1) @@ -2659,3 +2696,135 @@ def test_create_table_with_python_blob_dict_and_object(): b"from-dict", b"from-blob", ] + + +def _openable_descriptors(table, *, by_address): + hits = ( + table.search() + .select(["id", "image", "_rowaddr"]) + .with_row_id(True) + .limit(100) + .to_arrow(blob_mode="descriptions") + ) + version = table.version + committed_at = next( + item["timestamp"] + for item in table.list_versions() + if item["version"] == version + ) + descriptors = {} + for row in hits.to_pylist(): + descriptor = row["image"] + if descriptor is not None and descriptor.get("kind") is not None: + descriptor = dict(descriptor) + descriptor["source"] = { + "uri": table.to_lance().uri, + "version": version, + "committed_at_micros": round(committed_at.timestamp() * 1_000_000), + "branch": None, + "column": "image", + "row_reference": row["_rowaddr" if by_address else "_rowid"], + "row_is_address": by_address, + } + else: + descriptor = None + descriptors[row["id"]] = descriptor + return descriptors + + +@pytest.mark.parametrize("by_address", [False, True]) +def test_open_blob_files_preserve_snapshot_order_nulls_and_empty_values( + tmp_path, by_address +): + db = lancedb.connect(str(tmp_path)) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("videos", schema=schema) + table.add( + [ + {"id": 1, "image": b"alpha"}, + {"id": 2, "image": None}, + {"id": 3, "image": b"beta"}, + {"id": 4, "image": b""}, + ] + ) + descriptors = _openable_descriptors(table, by_address=by_address) + table.add([{"id": 1, "image": b"replacement"}], mode="overwrite") + latest_version = table.version + values = pa.array( + [descriptors[3], descriptors[2], descriptors[1], descriptors[3], descriptors[4]] + ) + files = table.open_blob_files( + pa.chunked_array([values.slice(0, 2), values.slice(2)]) + ) + assert iter(files) is files + actual = [] + for file in files: + if file is None: + actual.append(None) + else: + with file: + actual.append(file.read()) + assert file.closed + assert actual == [b"beta", None, b"alpha", b"beta", b""] + with table.open_blob_file(descriptors[1]) as file: + file.seek(2) + assert file.read_range(0, 2) == b"al" + assert file.tell() == 2 + assert file.read(2) == b"ph" + assert table.version == latest_version + assert table.open_blob_file(pa.scalar(None, type=values.type)) is None + other = db.create_table("other", schema=schema) + with pytest.raises(ValueError, match="different table"): + other.open_blob_file(descriptors[1]) + stale = dict(descriptors[1], source=dict(descriptors[1]["source"])) + stale["source"]["committed_at_micros"] += 1 + with pytest.raises(ValueError, match="snapshot has been replaced"): + table.open_blob_file(stale) + raw = {key: value for key, value in descriptors[1].items() if key != "source"} + with pytest.raises(ValueError, match="source"): + table.open_blob_file(raw) + with pytest.raises(TypeError, match="descriptor"): + table.open_blob_file("not a descriptor") + + +def test_open_blob_files_does_not_open_ahead(monkeypatch): + table = _blob_table("lazy_descriptor_iteration", [{"id": 1, "image": b"x"}]) + opened = [] + + def open_file(descriptor): + opened.append(descriptor) + return descriptor + + monkeypatch.setattr(table, "open_blob_file", open_file) + files = table.open_blob_files(iter([1, None, 2])) + assert opened == [] + assert next(files) == 1 + assert opened == [1] + assert list(files) == [None, 2] + assert opened == [1, None, 2] + + +@pytest.mark.asyncio +async def test_open_blob_files_async_iterates_snapshot_handles(tmp_path): + db = lancedb.connect(str(tmp_path)) + table = db.create_table( + "videos", + schema=pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]), + ) + table.add([{"id": 1, "image": b"alpha"}, {"id": 2, "image": None}]) + descriptors = _openable_descriptors(table, by_address=True) + async_db = await lancedb.connect_async(str(tmp_path)) + videos = await async_db.open_table("videos") + file = await videos.open_blob_file(pa.scalar(descriptors[1])) + with file: + assert file.read(2) == b"al" + actual = [] + async for file in videos.open_blob_files( + [descriptors[2], descriptors[1], descriptors[1]] + ): + if file is None: + actual.append(None) + else: + with file: + actual.append(await file.aread()) + assert actual == [None, b"alpha", b"alpha"] diff --git a/python/python/tests/test_blob_reader.py b/python/python/tests/test_blob_reader.py new file mode 100644 index 000000000..dd12409af --- /dev/null +++ b/python/python/tests/test_blob_reader.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import io + +import pytest + +from lancedb import BlobReader + + +def test_reader_is_lazy_and_seekable(): + requests = [] + size = 64 * 1024 * 1024 + + def read_range(offset, length): + requests.append((offset, length)) + return bytes((index % 251 for index in range(offset, offset + length))) + + with BlobReader(read_range, size) as reader: + assert requests == [] + assert reader.size() == size + reader.seek(-4, io.SEEK_END) + assert reader.read(10) == bytes(index % 251 for index in range(size - 4, size)) + assert reader.read(1) == b"" + assert requests == [(size - 4, 4)] + assert reader.read_range(1, 3) == b"\x01\x02\x03" + assert reader.tell() == size + reader.seek(2) + target = bytearray(3) + assert reader.readinto(target) == 3 + assert target == b"\x02\x03\x04" + reader.seek(-1, io.SEEK_CUR) + assert reader.tell() == 4 + with pytest.raises(ValueError, match="closed"): + reader.read(1) + + +def test_reader_supports_buffered_io_and_bounds_range_requests(): + requests = [] + payload = b"x" * (3 * 1024 * 1024 + 7) + + def read_range(offset, length): + requests.append((offset, length)) + return payload[offset : offset + length] + + with io.BufferedReader(BlobReader(read_range, len(payload))) as reader: + assert reader.read() == payload + assert requests + assert all(length <= 1024 * 1024 for _, length in requests) + + +def test_reader_rejects_invalid_ranges_and_short_reads(): + reader = BlobReader(lambda offset, length: b"", 10) + with pytest.raises(ValueError): + reader.seek(-1) + with pytest.raises(ValueError): + reader.read_range(-1, 1) + with pytest.raises(OSError, match="payload length"): + reader.read(1) diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 207a3d253..3c968a821 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -1039,6 +1039,39 @@ def test_open_clip_opens_percent_encoded_file_uri(tmp_path): assert image.size == (4, 4) +def test_open_clip_model_in_eval_mode(): + """ + open_clip.create_model_and_transforms returns the model in training mode, + which leaves dropout, stochastic depth and batch norm in their training + behaviour. OpenCLIP must switch the model to eval mode before embedding. + """ + from lancedb.embeddings.open_clip import OpenClipEmbeddings + + class FakeModel: + training = True + + def to(self, device): + return self + + def eval(self): + self.training = False + return self + + open_clip = MagicMock() + open_clip.create_model_and_transforms.return_value = (FakeModel(), None, None) + + with patch( + "lancedb.embeddings.open_clip.attempt_import_or_raise", + return_value=open_clip, + ): + func = OpenClipEmbeddings(name="ViTamin-B", pretrained="datacomp1b") + + open_clip.create_model_and_transforms.assert_called_once_with( + "ViTamin-B", pretrained="datacomp1b" + ) + assert func._model.training is False + + def test_jina_generate_image_input_dict_local_path(tmp_path): """ JinaEmbeddings._generate_image_input_dict must accept a local image path diff --git a/python/python/tests/test_errors.py b/python/python/tests/test_errors.py index 338b6a0d7..0e362e5d8 100644 --- a/python/python/tests/test_errors.py +++ b/python/python/tests/test_errors.py @@ -3,7 +3,28 @@ import pickle +import pytest + +from lancedb.db import DBConnection from lancedb.remote.errors import HttpError, LanceDBClientError, RetryError +from lancedb.table import Table + + +@pytest.mark.parametrize( + "base, method", + [ + (Table, "to_lance"), + (Table, "to_polars"), + (DBConnection, "drop_all_tables"), + (DBConnection, "drop_database"), + ], +) +def test_base_unsupported_methods_have_messages(base, method): + with pytest.raises(NotImplementedError) as exc_info: + getattr(base, method)(None) + message = str(exc_info.value) + assert method in message + assert "not supported" in message def test_pickle_lancedb_client_error(): diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 9ee3e388a..61e511345 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -23,6 +23,7 @@ import pyarrow as pa import pytest import lancedb +from lancedb import BlobReader from lancedb.functions import ( PythonRuntimeSpec, SecretBinding, @@ -966,6 +967,60 @@ def test_blob_fields_use_the_scalar_function_semantic_type(): assert signature.output.arrow_type == "blob_v2" +def test_blob_input_access_is_inferred_from_annotations(): + @udf + def read_blob(video: BlobReader, timestamp_s: float) -> bytes: + video.seek(int(timestamp_s)) + return video.read(4) + + @udf + def read_bytes(video: bytes) -> int: + return len(video) + + reader = read_blob.registration_request.signature + assert reader.inputs[0].arrow_type == "blob_v2" + assert reader.inputs[0].blob_access == "reader" + assert reader.inputs[1].blob_access is None + assert reader.output.arrow_type == "binary" + eager = read_bytes.registration_request.signature.inputs[0] + assert eager.arrow_type == "binary" + assert eager.blob_access == "bytes" + encoded = json.loads(read_blob.registration_request.to_canonical_json()) + assert encoded["signature"]["inputs"][0]["blob_access"] == "reader" + + +def test_nullable_blob_reader_and_explicit_blob_bytes(): + @udf + def nullable(video: Optional[BlobReader]) -> int: + return 0 if video is None else video.size() + + @udf( + input_schema=pa.schema([lancedb.blob("video", nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def eager(video: bytes) -> int: + return len(video) + + parameter = nullable.registration_request.signature.inputs[0] + assert parameter.nullable is True + assert parameter.blob_access == "reader" + parameter = eager.registration_request.signature.inputs[0] + assert parameter.arrow_type == "blob_v2" + assert parameter.blob_access == "bytes" + + +def test_blob_reader_rejects_an_incompatible_explicit_schema(): + def read(video: BlobReader) -> int: + return video.size() + + with pytest.raises(TypeError, match="requires a Blob V2 field"): + udf( + read, + input_schema=pa.schema([pa.field("video", pa.int64(), nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + + def test_whole_named_struct_function_can_include_a_blob_result_field(): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=False)]), @@ -1674,7 +1729,11 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip(): ) registration = db.create_function_async(normalize_score) assert registration.id == "job-register" + assert registration.job_type is None + assert registration.status() == "finished" + assert registration.job_type == "create_function" created = registration.wait() + assert registration.job_type == "create_function" reopened = db.get_function("normalize_score", version=created.version) assert created == reopened diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index d7fa2e0ec..1dc2fd9a1 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -391,6 +391,20 @@ def test_custom_stop_words_list(table): ) +def test_list_indices_with_missing_fts_files(table): + # Regression test for https://github.com/lancedb/lancedb/issues/4250 + table.create_fts_index("text") + table.create_scalar_index("id") + + fts = next(i for i in table.list_indices() if i.index_type == "FTS") + shutil.rmtree(Path(table.uri) / "_indices" / fts.index_uuid) + + indices = {i.name: i for i in table.list_indices()} + assert set(indices) == {"text_idx", "id_idx"} + assert indices["text_idx"].index_type == "FTS" + assert indices["text_idx"].columns == ["text"] + + def test_search_fts(table): table.create_fts_index("text") results = table.search("puppy").select(["id", "text"]).limit(5).to_list() diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 0252a09ad..2ec14f6f8 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -421,13 +421,23 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable): @pytest.mark.asyncio -async def test_create_ivfrq_index(some_table: AsyncTable): - await some_table.create_index("vector", config=IvfRq(num_bits=1)) +@pytest.mark.parametrize( + "config, expected_bits", + [ + pytest.param(IvfRq(), 5, id="default"), + pytest.param(IvfRq(num_bits=None), 5, id="none"), + pytest.param(IvfRq(num_bits=1), 1, id="one-bit"), + pytest.param(IvfRq(num_bits=4), 4, id="multi-bit"), + ], +) +async def test_create_ivfrq_index(some_table: AsyncTable, config, expected_bits): + await some_table.create_index("vector", config=config) indices = await some_table.list_indices() assert len(indices) == 1 assert indices[0].index_type == "IvfRq" assert indices[0].columns == ["vector"] assert indices[0].name == "vector_idx" + assert indices[0].index_details["compression"]["num_bits"] == expected_bits @pytest.mark.asyncio diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py index 4a2745eb8..b0b973bbf 100644 --- a/python/python/tests/test_listing.py +++ b/python/python/tests/test_listing.py @@ -88,6 +88,143 @@ async def advance(items, asynchronous): KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"] +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_compatible_name_listing_pages(listing_server, kind, asynchronous): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a", "b"], "next")), + (200, response_page(kind, ["c"], None)), + ] + ) + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)( + namespace_path=["team"], page_token="start", limit=2 + ) + field = kind + assert len(requests) == (0 if asynchronous else 1) + if asynchronous: + with pytest.raises(RuntimeError, match="Await or iterate"): + getattr(listing, field) + assert listing.page_token == "start" + assert await listing is listing + assert len(requests) == 1 + page_names = getattr(listing, field) + assert type(page_names) is list + assert page_names == ["a", "b"] + page_names.append("changed") + assert getattr(listing, field) == ["a", "b"] + assert listing.page_token == "next" + assert listing.num_page_results() == 2 + assert await advance(listing, asynchronous) == "a" + assert getattr(listing, field) == ["a", "b"] + assert await advance(listing, asynchronous) == "b" + assert len(requests) == 1 + assert await advance(listing, asynchronous) == "c" + assert getattr(listing, field) == ["c"] + assert listing.page_token is None + assert len(requests) == 2 + assert ([item async for item in listing] if asynchronous else list(listing)) == [] + if asynchronous: + assert await listing is listing + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_skips_empty_page_only_during_iteration( + listing_server, kind +): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, [], "next")), + (200, response_page(kind, ["a"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + await listing + assert getattr(listing, kind) == [] + assert listing.page_token == "next" + assert len(requests) == 1 + assert [name async for name in listing] == ["a"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_name_listing_zero_limit_uses_no_request( + listing_server, kind, asynchronous +): + endpoint, requests, _ = listing_server + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)(limit=0) + if asynchronous: + assert await listing is listing + assert [name async for name in listing] == [] + else: + assert list(listing) == [] + assert getattr(listing, kind) == [] + assert listing.page_token is None + assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_failure_retains_token(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((401, {})) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(page_token="retry") + with pytest.raises(HttpError): + await listing + assert listing.page_token == "retry" + with pytest.raises(HttpError): + await listing + assert [name async for name in listing] == [] + assert len(requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_direct_async_iteration(listing_server, kind): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a"], "next")), + (200, response_page(kind, ["b"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(limit=1) + assert requests == [] + assert [name async for name in listing] == ["a", "b"] + assert getattr(listing, kind) == ["b"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_serializes_concurrent_advances(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((200, response_page(kind, ["a", "b"], None))) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + results = await asyncio.gather( + listing.__anext__(), + listing.__anext__(), + listing.__anext__(), + return_exceptions=True, + ) + assert results[:2] == ["a", "b"] + assert isinstance(results[2], StopAsyncIteration) + assert len(requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", KINDS) @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index 46972a15c..e2be9c0d4 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -8,6 +8,7 @@ import threading import lancedb import pytest +from lancedb.db import AsyncConnection from lancedb.materialized_view import MaterializedViewDefinition from lancedb.remote.db import RemoteDBConnection @@ -56,7 +57,7 @@ def mock_remote_materialized_views(): @contextlib.contextmanager -def mock_remote_materialized_view_create(): +def mock_remote_materialized_view_drop(): requests = [] class Handler(http.server.BaseHTTPRequestHandler): @@ -67,7 +68,8 @@ def mock_remote_materialized_view_create(): length = int(self.headers.get("Content-Length", 0)) body = json.loads(self.rfile.read(length) or b"{}") requests.append((self.path, body)) - job_id = "mv-drop-123" if self.path.endswith("/drop") else "mv-create-123" + assert self.path.endswith("/drop") + job_id = "mv-drop-123" encoded = json.dumps({"job_id": job_id}).encode() self.send_response(202) self.send_header("Content-Type", "application/json") @@ -98,7 +100,7 @@ def test_remote_list_uses_namespace_route(): def test_remote_drop_async_returns_server_job(): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect( "db://dev", api_key="fake", @@ -192,11 +194,13 @@ def test_create_and_refresh_jobs(tmp_path): "adults", "people", where="age >= 18", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert create_job.wait() is None view = db.open_materialized_view("adults") refresh_job = view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 2 @@ -337,11 +341,13 @@ async def test_async_create_and_refresh_jobs(tmp_path): "adults", "people", with_no_data=True ) assert create_job.id is None + assert create_job.job_type is None assert await create_job.wait() is None view = await db.open_materialized_view("adults") refresh_job = await view.refresh_async() assert refresh_job.id is None + assert refresh_job.job_type is None result = await refresh_job.wait() assert result.mode == "rebuild" assert result.rows_written == 1 @@ -567,6 +573,27 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): assert list(resumed) == [] +class QueryBackedCreateJob: + id = "0198f1b2-c3d4-7e5f-8123-456789abcdef" + + job_type = "query" + + async def wait(self): + return None + + async def status(self): + return "finished" + + +class CapturingCreateConnection: + def __init__(self): + self.calls = [] + + async def create_materialized_view_async(self, name, source, **options): + self.calls.append((name, source, options)) + return QueryBackedCreateJob() + + @pytest.mark.parametrize( "options", [ @@ -576,16 +603,28 @@ def test_local_materialized_view_listing_pages_and_resume(tmp_path): ], ) def test_remote_create_forwards_refresh_options_without_interval_default(options): - with mock_remote_materialized_view_create() as (host, requests): - db = lancedb.connect("db://dev", api_key="fake", host_override=host) - job = db.create_materialized_view_async("adults", "people", **options) - assert job.id == "mv-create-123" - assert requests == [ + native = CapturingCreateConnection() + db = RemoteDBConnection.__new__(RemoteDBConnection) + db._conn = AsyncConnection(native) + job = db.create_materialized_view_async("adults", "people", **options) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert job.wait() is None + assert job.status() == "finished" + assert native.calls == [ ( - "/v1/materialized_view/adults/create", + "adults", + "people", { - "query": 'SELECT * FROM "people"', - **options, + "projections": None, + "filter": None, + "limit": None, + "with_no_data": False, + "auto_refresh_enabled": options.get("auto_refresh_enabled"), + "auto_refresh_interval_seconds": options.get( + "auto_refresh_interval_seconds" + ), }, ) ] @@ -593,18 +632,18 @@ def test_remote_create_forwards_refresh_options_without_interval_default(options @pytest.mark.asyncio async def test_async_remote_create_keeps_interval_unspecified(): - with mock_remote_materialized_view_create() as (host, requests): - db = await lancedb.connect_async("db://dev", api_key="fake", host_override=host) - job = await db.create_materialized_view_async( - "adults", - "people", - auto_refresh_enabled=True, - ) - assert job.id == "mv-create-123" - assert requests[0][1] == { - "query": 'SELECT * FROM "people"', - "auto_refresh_enabled": True, - } + native = CapturingCreateConnection() + db = AsyncConnection(native) + job = await db.create_materialized_view_async( + "adults", "people", auto_refresh_enabled=True + ) + assert job.id == QueryBackedCreateJob.id + assert job.job_type == "query" + assert not hasattr(job, "query_id") + assert await job.wait() is None + assert await job.status() == "finished" + assert native.calls[0][2]["auto_refresh_enabled"] is True + assert native.calls[0][2]["auto_refresh_interval_seconds"] is None @pytest.mark.parametrize( @@ -625,7 +664,7 @@ def test_embedded_create_refuses_automatic_refresh_before_creating_view( @pytest.mark.parametrize("seconds", [0, -1, float("nan"), float("inf")]) def test_remote_create_rejects_invalid_interval_without_request(seconds): - with mock_remote_materialized_view_create() as (host, requests): + with mock_remote_materialized_view_drop() as (host, requests): db = lancedb.connect("db://dev", api_key="fake", host_override=host) with pytest.raises(ValueError, match="positive, finite"): db.create_materialized_view_async( diff --git a/python/python/tests/test_ollama_embeddings.py b/python/python/tests/test_ollama_embeddings.py index 9d8568ca7..de36dc24b 100644 --- a/python/python/tests/test_ollama_embeddings.py +++ b/python/python/tests/test_ollama_embeddings.py @@ -3,7 +3,6 @@ import json -import httpx import numpy as np import pyarrow as pa import pytest @@ -14,6 +13,7 @@ from lancedb.embeddings import get_registry @pytest.fixture def ollama_embeddings(): + httpx = pytest.importorskip("httpx") pytest.importorskip("ollama") requests = [] diff --git a/python/python/tests/test_openai_reranker_credentials.py b/python/python/tests/test_openai_reranker_credentials.py new file mode 100644 index 000000000..aa11239b2 --- /dev/null +++ b/python/python/tests/test_openai_reranker_credentials.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pytest + +from lancedb.rerankers import OpenaiReranker + + +@pytest.mark.parametrize( + "explicit_key, environment_key, expected_key", + [ + ("explicit-test-key", "environment-test-key", "explicit-test-key"), + ("explicit-test-key", None, "explicit-test-key"), + (None, "environment-test-key", "environment-test-key"), + ], +) +def test_openai_reranker_api_key_precedence( + monkeypatch, explicit_key, environment_key, expected_key +): + pytest.importorskip("openai") + if environment_key is None: + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + else: + monkeypatch.setenv("OPENAI_API_KEY", environment_key) + + reranker = OpenaiReranker(api_key=explicit_key) + client = reranker._client + try: + assert client.api_key == expected_key + finally: + client.close() + + +def test_openai_reranker_requires_api_key(monkeypatch): + pytest.importorskip("openai") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with pytest.raises(ValueError, match="OPENAI_API_KEY not set"): + OpenaiReranker()._client diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index ed1894340..50ca53061 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -811,6 +811,13 @@ def test_rename_column(some_permutation: Permutation): some_permutation.rename_column("non_existent", "new_value") +def test_rename_column_keeps_position(some_permutation: Permutation): + renamed = some_permutation.rename_column("id", "new_id") + assert renamed.schema == pa.schema([("new_id", pa.int64()), ("value", pa.int64())]) + batch = renamed.with_format("arrow").__getitems__([0]) + assert batch.schema.names == ["new_id", "value"] + + def test_rename_columns(some_permutation: Permutation): assert some_permutation.rename_columns({"value": "new_value"}).schema == pa.schema( [("id", pa.int64()), ("new_value", pa.int64())] diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 58ea7deb5..cef985aa6 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -2324,6 +2324,33 @@ def test_take_queries(tmp_path): ] +def test_take_queries_with_row_id(tmp_path): + db = lancedb.connect(tmp_path) + table = db.create_table("test", pa.table({"idx": range(100)})) + + by_row_id = table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_row_id.column_names + by_offset = table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_offset.column_names + + # The column is opt-in and must stay out of the default projection. + assert "_rowid" not in table.take_row_ids([5, 2, 17]).to_arrow().column_names + + +@pytest.mark.asyncio +async def test_take_queries_with_row_id_async(tmp_path): + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("test", pa.table({"idx": range(100)})) + + by_row_id = await table.take_row_ids([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_row_id.column_names + by_offset = await table.take_offsets([5, 2, 17], with_row_id=True).to_arrow() + assert "_rowid" in by_offset.column_names + + default = await table.take_row_ids([5, 2, 17]).to_arrow() + assert "_rowid" not in default.column_names + + def test_take_queries_to_batches(tmp_path): # Regression test for the sync take-query path: `to_batches` previously # raised ``AttributeError: 'AsyncTakeQuery' object has no attribute diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index a45e1aaef..32f999d18 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -5,6 +5,7 @@ import contextlib from datetime import timedelta import http.server import io +import inspect import json import multiprocessing as mp import pickle @@ -765,6 +766,15 @@ def test_create_table_embedding_functions(with_schema): ] +def test_create_table_positional_parameters_match_base_contract(): + from lancedb.db import DBConnection + from lancedb.remote.db import RemoteDBConnection + + expected = list(inspect.signature(DBConnection.create_table).parameters)[:9] + actual = list(inspect.signature(RemoteDBConnection.create_table).parameters)[:9] + assert actual == expected + + def test_create_table_exist_ok(): def handler(request): if request.path == "/v1/table/test/create/?mode=exist_ok": @@ -857,10 +867,116 @@ def test_table_unimplemented_functions(): with mock_lancedb_connection(handler) as db: table = db.create_table("test", [{"id": 1}]) - with pytest.raises(NotImplementedError): - table.to_arrow() - with pytest.raises(NotImplementedError): - table.to_pandas() + for method in ["to_arrow", "to_pandas", "to_polars"]: + with pytest.raises(NotImplementedError) as exc_info: + getattr(table, method)() + message = str(exc_info.value) + assert method in message + assert "not" in message and "supported" in message + assert "remote tables" in message + assert "cloud" not in message.lower() + assert f"search().{method}()" in message + + +@pytest.mark.parametrize("method", ["to_arrow", "to_pandas"]) +@pytest.mark.parametrize("blob_mode", ["lazy", "descriptions", "bytes", "invalid"]) +def test_remote_conversion_blob_mode_errors(method, blob_mode): + with blob_remote_table() as table: + if blob_mode == "invalid": + with pytest.raises(ValueError, match="blob_mode must be one of"): + getattr(table, method)(blob_mode=blob_mode) + else: + with pytest.raises(NotImplementedError) as error: + getattr(table, method)(blob_mode=blob_mode) + assert f"search().{method}()" in str(error.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [None, 1]) +@pytest.mark.parametrize("expose_location", [False, True]) +async def test_remote_to_lance_rejection( + tmp_path, version, expose_location, monkeypatch +): + import lance + + location = tmp_path / "test.lance" + lance.write_dataset(pa.table({"id": [1]}), location) + lance.write_dataset(pa.table({"id": [2]}), location, mode="append") + requests = [] + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) + + def handler(request): + requests.append(request.path) + assert request.path == "/v1/table/test/describe/" + content_len = int(request.headers.get("Content-Length", 0)) + body = json.loads(request.rfile.read(content_len)) if content_len else {} + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + response = { + "version": body.get("version") or 2, + "schema": {"fields": []}, + } + if expose_location: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) + + options = {"default_scan_options": {"with_row_id": True}} + with mock_lancedb_connection(handler) as db: + table = db.open_table("test", version=version) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as sync_error: + table.to_lance(**options) + assert requests == opened_requests + + async with mock_lancedb_connection_async(handler) as db: + table = await db.open_table("test", version=version) + opened_requests = list(requests) + with pytest.raises(NotImplementedError) as async_error: + await table.to_lance(**options) + assert requests == opened_requests + + assert "to_lance() is not supported for remote tables" in str(sync_error.value) + assert "query the server" in str(sync_error.value) + assert str(sync_error.value) == str(async_error.value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method, kwargs", + [ + ("drop_all_tables", {}), + ("drop_all_tables", {"namespace_path": []}), + ("drop_all_tables", {"namespace_path": ["team"]}), + ("drop_database", {}), + ], +) +async def test_remote_drop_all_tables_not_supported(method, kwargs): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(500) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError) as sync_error: + getattr(db, method)(**kwargs) + + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError) as async_error: + await getattr(db, method)(**kwargs) + + assert "Dropping all tables is not currently supported in the remote API" in str( + sync_error.value + ) + assert "drop_table" in str(sync_error.value) + assert str(sync_error.value) == str(async_error.value) + assert requests == [] def test_table_to_pandas_not_supported(): @@ -1039,8 +1155,9 @@ def test_table_create_indices(): # Test create_fts_index with custom name (legacy method) with pytest.warns(DeprecationWarning, match="create_fts_index"): table.create_fts_index( - "text", + field_names="text", wait_timeout=timedelta(seconds=2), + tokenizer_name="en_stem", block_size=256, custom_stop_words=["cloud"], name="custom_fts_idx", @@ -1049,6 +1166,8 @@ def test_table_create_indices(): # Test create_index with custom name (legacy form: vector_column_name kwarg) with pytest.warns(DeprecationWarning, match="create_index"): table.create_index( + "cosine", + 2, vector_column_name="vector", wait_timeout=timedelta(seconds=10), name="custom_vector_idx", @@ -1061,7 +1180,8 @@ def test_table_create_indices(): scalar_req = received_requests[0] assert "name" in scalar_req assert scalar_req["name"] == "custom_scalar_idx" - assert scalar_req["replace"] is False + # The wire protocol omits replace=True (the server default). + assert "replace" not in scalar_req # Check FTS index request has custom name fts_req = received_requests[1] @@ -1070,18 +1190,39 @@ def test_table_create_indices(): assert fts_req["replace"] is False assert fts_req["block_size"] == 256 assert fts_req["custom_stop_words"] == ["cloud"] + assert fts_req["language"] == "English" + assert fts_req["stem"] is True # Check vector index request has custom name vector_req = received_requests[2] assert "name" in vector_req assert vector_req["name"] == "custom_vector_idx" assert "replace" not in vector_req + assert vector_req["metric_type"] == "cosine" + assert vector_req["num_partitions"] == 2 + + with pytest.warns(DeprecationWarning, match="create_index"): + table.create_index( + index_type="IVF_HNSW_SQ", + num_partitions=1, + m=10, + ef_construction=50, + max_iterations=12, + sample_rate=64, + ) + tuned_req = received_requests[-1] + assert tuned_req["index_type"] == "IVF_HNSW_SQ" + assert tuned_req["num_partitions"] == 1 + assert tuned_req["m"] == 10 + assert tuned_req["ef_construction"] == 50 + assert tuned_req["max_iterations"] == 12 + assert tuned_req["sample_rate"] == 64 table.wait_for_index(["custom_scalar_idx"], timedelta(seconds=2)) table.wait_for_index( ["custom_fts_idx", "custom_vector_idx"], timedelta(seconds=2) ) - table.drop_index("custom_vector_idx") + table.drop_index(name="custom_vector_idx") table.drop_index("custom_scalar_idx") table.drop_index("custom_fts_idx") @@ -1477,7 +1618,10 @@ def test_remote_json_fts_config(server_version): assert request["block_size"] == 256 -def test_remote_create_index_new_api(caplog): +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +def test_remote_create_index_new_api(caplog, rq_kwargs): received_requests = [] def handler(request): @@ -1550,7 +1694,7 @@ def test_remote_create_index_new_api(caplog): table.create_index("category", config=BTree()) table.create_index("text", config=FTS(block_size=256)) # IvfRq via new API - table.create_index("vector", config=IvfRq(distance_type="l2")) + table.create_index("vector", config=IvfRq(distance_type="l2", **rq_kwargs)) table.create_index( "vector", config=IvfPq(distance_type="l2"), replace=False ) @@ -1575,6 +1719,10 @@ def test_remote_create_index_new_api(caplog): "vector", ] assert received_requests[2]["block_size"] == 256 + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in received_requests[3] + else: + assert received_requests[3]["num_bits"] == rq_kwargs["num_bits"] assert received_requests[4]["replace"] is False assert received_requests[5]["replace"] is False assert "replace is not supported" not in caplog.text @@ -1675,16 +1823,19 @@ def test_remote_unsupported_legacy_hnsw_pq(): ) -def test_remote_python_only_messages(caplog): +def test_remote_python_only_messages(): with mock_lancedb_connection(remote_unsupported_handler) as db: - with caplog.at_level("INFO"): - table = db.open_table("test", storage_options={}, index_cache_size=1) - assert "storage_options is ignored for remote tables" in caplog.text - assert "index_cache_size is ignored for remote tables" in caplog.text - for method in ["to_arrow", "to_pandas"]: + for options in [{"storage_options": {}}, {"index_cache_size": 1}]: + with pytest.raises(NotImplementedError, match=next(iter(options))): + db.open_table("test", **options) + table = db.open_table("test") + for method in ["to_arrow", "to_pandas", "to_polars"]: with pytest.raises(NotImplementedError) as error: getattr(table, method)() - assert str(error.value) == f"{method}() is not supported for remote tables." + assert str(error.value) == ( + f"{method}() is not supported for remote tables. " + f"Use table.search().{method}() instead." + ) for method in ["compact_files", "cleanup_old_versions"]: with pytest.warns(UserWarning) as warning: getattr(table, method)() @@ -2225,6 +2376,23 @@ def test_head(): assert data == pa.table({"id": [1, 2, 3]}) +def test_remote_take_row_ids_with_row_id(): + bodies = [] + + def handler(body): + bodies.append(body) + return pa.table({"id": [1, 2], "_rowid": pa.array([0, 1], pa.uint64())}) + + with query_test_table(handler) as table: + result = table.take_row_ids([0, 1], with_row_id=True).to_arrow() + assert "_rowid" in result.column_names + assert bodies[-1]["with_row_id"] is True + + with query_test_table(handler) as table: + table.take_row_ids([0, 1]).to_arrow() + assert "with_row_id" not in bodies[-1] + + def test_query_sync_minimal(): def handler(body): assert body == { @@ -3473,14 +3641,17 @@ def blob_descriptor_take_table(): @contextlib.contextmanager -def blob_remote_table(*, server_version=Version("0.5.0")): +def blob_remote_table(*, server_version=Version("0.5.0"), location=None): def handler(request): if request.path == "/v1/table/test/describe/": request.send_response(200) request.send_header("Content-Type", "application/json") request.send_header("phalanx-version", str(server_version)) request.end_headers() - request.wfile.write(json.dumps(BLOB_DESCRIBE_RESPONSE).encode()) + response = dict(BLOB_DESCRIBE_RESPONSE) + if location is not None: + response["location"] = str(location) + request.wfile.write(json.dumps(response).encode()) elif request.path.startswith("/v1/table/test/blob/image/"): path = request.path.partition("?")[0] row_id = int(path.split("/")[-2]) @@ -3515,7 +3686,6 @@ def blob_remote_table(*, server_version=Version("0.5.0")): if columns == ["id"]: response_table = blob_query_response_table().select(["id"]) else: - assert body["with_row_id"] is True if columns == ["image"]: # Blob handles are sized from a descriptor take. assert body["filter"].startswith("_rowid IN") @@ -3528,6 +3698,8 @@ def blob_remote_table(*, server_version=Version("0.5.0")): ) else: response_table = blob_query_response_table() + if not body.get("with_row_id", False): + response_table = response_table.drop(["_rowid"]) request.send_response(200) request.send_header("Content-Type", "application/vnd.apache.arrow.file") request.send_header("x-lancedb-version", "1") @@ -3562,6 +3734,32 @@ def test_remote_blob_columns_and_fetch(): assert blobs.to_pylist() == [b"alpha", None, b"gamma"] +@pytest.mark.parametrize("blob_mode", ["descriptions", "bytes", "lazy"]) +def test_sync_remote_blob_pandas_avoids_storage_access( + tmp_path, monkeypatch, blob_mode +): + import lance + + location = tmp_path / "remote.lance" + lance.write_dataset(pa.table({"id": [7], "image": [b"server-owned"]}), location) + + def fail_dataset(*args, **kwargs): + raise AssertionError("remote storage must not be opened") + + monkeypatch.setattr(lance, "dataset", fail_dataset) + with blob_remote_table(location=location) as table: + monkeypatch.setattr(table, "to_lance", fail_dataset) + query = table.search().select(["id", "image"]) + if blob_mode == "descriptions": + df = query.to_pandas(blob_mode=blob_mode) + assert df["id"].tolist() == [1, 2, 3] + assert "_rowid" not in df.columns + assert df["image"][0]["size"] == 5 + else: + with pytest.raises(NotImplementedError, match="stable table snapshot"): + query.to_pandas(blob_mode=blob_mode) + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", [False, True]) async def test_remote_fetch_blob_ranges(use_async): diff --git a/python/python/tests/test_remote_options.py b/python/python/tests/test_remote_options.py new file mode 100644 index 000000000..2022aa29d --- /dev/null +++ b/python/python/tests/test_remote_options.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import inspect +from urllib.parse import parse_qs, urlparse + +import lancedb +import pyarrow as pa +import pytest + +from test_remote_db import mock_lancedb_connection, mock_lancedb_connection_async + + +@pytest.mark.parametrize("uri", ["memory://", "db://dev"]) +def test_unknown_connect_options(uri): + with pytest.raises(ValueError, match="Unknown keyword arguments.*bogus"): + lancedb.connect(uri, api_key="fake", bogus=1) + + +@pytest.mark.parametrize("option", ["session", "manifest_enabled"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_remote_connect_rejects_local_options(option, asynchronous): + value = lancedb.Session() if option == "session" else True + with pytest.raises(NotImplementedError, match=option): + if asynchronous: + await lancedb.connect_async("db://dev", api_key="fake", **{option: value}) + else: + lancedb.connect("db://dev", api_key="fake", **{option: value}) + + +def test_remote_create_signature(): + from lancedb.db import DBConnection + from lancedb.remote.db import RemoteDBConnection + + expected = inspect.signature(DBConnection.create_table).parameters + actual = inspect.signature(RemoteDBConnection.create_table).parameters + assert list(actual) == list(expected) + for name in expected: + assert actual[name].kind == expected[name].kind + if name == "mode": + # Remote connections retain their existing server-default mode. + assert actual[name].default is None + else: + assert actual[name].default == expected[name].default + + +def test_remote_create_positional_mode(): + modes = [] + + def handler(request): + assert urlparse(request.path).path == "/v1/table/ct_pos/create/" + modes.append(parse_qs(urlparse(request.path).query)["mode"]) + request.send_response(200) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + data = [{"value": 1}] + db.create_table("ct_pos", data, None, "overwrite") + db.create_table("ct_pos", data, None, "overwrite") + db.create_table("ct_pos", data, None, "create", True) + with pytest.raises(ValueError, match="Invalid mode"): + db.create_table("ct_pos", data, None, "drop") + assert modes == [["overwrite"], ["overwrite"], ["exist_ok"]] + + +CREATE_OPTIONS = [ + {"new_table_data_storage_version": "2.0"}, + {"new_table_enable_v2_manifest_paths": "false"}, + {"bogus_option_xyz": "1"}, +] + + +@pytest.mark.parametrize("storage_options", CREATE_OPTIONS) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.asyncio +async def test_remote_create_rejects_write_options(storage_options, empty): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + + kwargs = {"schema": pa.schema([("value", pa.int64())])} + if not empty: + kwargs["data"] = [{"value": 1}] + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError, match="write_options.*remote"): + await db.create_table("t", storage_options=storage_options, **kwargs) + assert requests == [] + + +@pytest.mark.parametrize( + "options", + [ + {"data_storage_version": "2.0"}, + {"enable_v2_manifest_paths": False}, + {"storage_options": {"x": "1"}}, + ], +) +def test_sync_remote_create_rejects_write_options(options): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError, match="remote"): + db.create_table("t", [{"value": 1}], **options) + assert requests == [] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "options", [{"index_cache_size": 10}, {"storage_options": {"x": "y"}}] +) +@pytest.mark.asyncio +async def test_remote_open_rejects_local_options(options, asynchronous): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + if asynchronous: + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError, match="remote"): + await db.open_table("t", **options) + else: + with mock_lancedb_connection(handler) as db: + with pytest.raises(NotImplementedError, match="remote"): + db.open_table("t", **options) + assert requests == [] + + +@pytest.mark.asyncio +async def test_remote_open_rejects_location(): + requests = [] + + def handler(request): + requests.append(request.path) + request.send_response(200) + request.end_headers() + request.wfile.write(b'{"version": 1, "schema": {"fields": []}}') + + async with mock_lancedb_connection_async(handler) as db: + with pytest.raises(NotImplementedError, match="location.*remote"): + await db.open_table("t", location="/tmp/zzz") + assert requests == [] diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 743669367..dd6cb4ab5 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import os +from unittest.mock import MagicMock import random import lancedb @@ -547,6 +548,24 @@ def test_openai_reranker(tmp_path): _run_test_reranker(reranker, table, "single player experience", None, schema) +def test_openai_reranker_does_not_send_temperature(): + reranker = OpenaiReranker(api_key="test-key") + create = MagicMock() + create.return_value.choices = [ + MagicMock( + message=MagicMock( + content='{"documents": [{"content": "a", "relevance_score": 0.5}]}' + ) + ) + ] + reranker.__dict__["_client"] = MagicMock() + reranker.__dict__["_client"].chat.completions.create = create + + reranker._rerank(pa.table({"text": ["a"]}), "query") + + assert "temperature" not in create.call_args.kwargs + + @pytest.mark.skipif( os.environ.get("JINA_API_KEY") is None, reason="JINA_API_KEY not set" ) diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index e6f5fd3c7..539b88647 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -137,9 +137,23 @@ async def test_async_namespace_connection_rejects_sql(tmp_path): await connection.describe_query(NIL_QUERY_ID) -def test_describe_query_requires_uuid(): - with pytest.raises(TypeError, match="UUID"): - remote_connection().describe_query(str(NIL_QUERY_ID)) +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", [NIL_QUERY_ID, str(NIL_QUERY_ID)]) +async def test_describe_query_accepts_uuid_and_job_id_string(identity): + calls = [] + + class NativeConnection: + async def describe_query(self, query_id): + calls.append(query_id) + + connection = AsyncConnection(NativeConnection()) + await connection.describe_query(identity) + assert calls == [NIL_QUERY_ID] + + +def test_describe_query_rejects_registry_job_id(): + with pytest.raises(ValueError, match="UUID"): + remote_connection().describe_query("registry-job-id") @pytest.mark.parametrize( diff --git a/python/python/tests/test_sync_api_parity.py b/python/python/tests/test_sync_api_parity.py new file mode 100644 index 000000000..a12efd0c9 --- /dev/null +++ b/python/python/tests/test_sync_api_parity.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import timedelta +import json +from urllib.parse import parse_qs, urlparse + +import lancedb +from lancedb.table import Table +import pytest + +from test_remote_db import mock_lancedb_connection + + +@pytest.fixture +def legacy_table(): + """Implement the pre-existing abstract API without the two new members.""" + + def unsupported(self, *args, **kwargs): + raise NotImplementedError + + def update_field_metadata(self, updates): + self.metadata_updates.append(updates) + return object() # The deprecated wrapper must not expose this result. + + # Supply concrete implementations through normal ABC class creation. Leave + # namespace and replace_field_metadata inherited, as an old subclass would. + implementations = { + name: unsupported + for name in Table.__abstractmethods__ + if name not in {"namespace", "replace_field_metadata"} + } + implementations.update( + name=property(lambda self: "legacy"), + update_field_metadata=update_field_metadata, + ) + legacy_class = type("LegacyTable", (Table,), implementations) + table = legacy_class() + table.metadata_updates = [] + return table + + +def test_legacy_table_namespace_and_id(legacy_table): + assert legacy_table.namespace == [] + assert legacy_table.id == "legacy" + namespace = legacy_table.namespace + namespace.append("changed") + assert legacy_table.namespace == [] + assert legacy_table.id == "legacy" + + +def test_legacy_table_replace_field_metadata(legacy_table): + with pytest.warns(DeprecationWarning, match="update_field_metadata"): + result = legacy_table.replace_field_metadata("nested.field", {"a": "b"}) + assert result is None + assert legacy_table.metadata_updates == [ + {"path": "nested.field", "metadata": {"a": "b"}, "replace": True} + ] + + +def respond(request, body): + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(body).encode()) + + +def test_remote_uri(): + with mock_lancedb_connection(lambda request: respond(request, {})) as db: + assert db.uri == "db://dev" + + +@pytest.mark.parametrize("attribute", ["session", "read_consistency_interval"]) +def test_remote_unsupported_properties(attribute): + with mock_lancedb_connection(lambda request: respond(request, {})) as db: + with pytest.raises(NotImplementedError, match=attribute): + getattr(db, attribute) + + +@pytest.mark.parametrize("tables", [[], ["first", "last"]]) +def test_remote_len_and_contains_follow_pages(tables): + requests = [] + + def handler(request): + requests.append(request.path) + query = parse_qs(urlparse(request.path).query) + if tables and "page_token" not in query: + respond(request, {"tables": [tables[0]], "page_token": "opaque-next"}) + else: + if tables: + assert query["page_token"] == ["opaque-next"] + respond(request, {"tables": tables[1:]}) + + with mock_lancedb_connection(handler) as db: + assert len(db) == len(tables) + assert ("last" in db) == bool(tables) + assert "missing" not in db + requests.clear() + assert ("first" in db) == bool(tables) + assert len(requests) == 1 + + +@pytest.mark.parametrize("namespace", [[], ["parent", "child"]]) +def test_remote_table_identifiers(namespace): + def handler(request): + respond(request, {"version": 1, "schema": {"fields": []}}) + + with mock_lancedb_connection(handler) as db: + table = db.open_table("rfm", namespace_path=namespace) + assert table.namespace == namespace + assert table.id == "$".join(namespace + ["rfm"]) + + +def test_remote_replace_field_metadata(): + updates = [] + + def handler(request): + if request.path.endswith("/describe/"): + respond(request, {"version": 1, "schema": {"fields": []}}) + else: + assert request.path == "/v1/table/rfm/update_field_metadata/" + body = request.rfile.read(int(request.headers["Content-Length"])) + updates.append(json.loads(body)) + respond(request, {"version": 2}) + + with mock_lancedb_connection(handler) as db: + table = db.open_table("rfm") + with pytest.warns(DeprecationWarning, match="update_field_metadata"): + assert table.replace_field_metadata("id", {"a": "b"}) is None + assert updates == [ + {"updates": [{"path": "id", "metadata": {"a": "b"}, "replace": True}]} + ] + + +@pytest.mark.parametrize("custom_session", [False, True]) +def test_local_properties_and_metadata(tmp_path, custom_session): + supplied = lancedb.Session() if custom_session else None + db = lancedb.connect( + tmp_path, session=supplied, read_consistency_interval=timedelta(seconds=5) + ) + session = db.session + assert session is supplied + before = session.size_bytes if session is not None else 0 + table = db.create_table("rfm", [{"id": i} for i in range(3)]) + assert table.to_arrow().num_rows == 3 + if supplied is not None: + assert supplied.size_bytes > before + assert db.read_consistency_interval == timedelta(seconds=5) + assert table.namespace == [] + assert table.id == "rfm" + with pytest.warns(DeprecationWarning): + table.replace_field_metadata("id", {"old": "value"}) + with pytest.warns(DeprecationWarning): + table.replace_field_metadata("id", {"a": "b"}) + assert table.schema.field("id").metadata == {b"a": b"b"} + + +def test_namespace_connection_properties_remain_writable(tmp_path): + session = lancedb.Session() + interval = timedelta(seconds=5) + with lancedb.connect_namespace( + "dir", + {"root": str(tmp_path)}, + session=session, + read_consistency_interval=interval, + ) as db: + assert db.session is session + assert db.read_consistency_interval == interval + db.read_consistency_interval = timedelta(0) + assert db.read_consistency_interval == timedelta(0) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 3776c5ef2..0bd4b5b26 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -60,6 +60,124 @@ def _blob_v2_table(db: DBConnection, name: str): return table +@pytest.mark.parametrize( + "threshold", + [ + "inline_size_threshold", + "dedicated_size_threshold", + "pack_file_size_threshold", + ], +) +@pytest.mark.parametrize( + "value,error", + [ + (-1, ValueError), + (True, TypeError), + (False, TypeError), + (1.5, TypeError), + ("8192", TypeError), + (ctypes.c_size_t(-1).value + 1, OverflowError), + ], +) +def test_blob_thresholds_reject_invalid_values(threshold, value, error): + # Also runs in CI with pylance actually uninstalled. + with pytest.raises(error, match=f"{threshold} must"): + lancedb.blob("blob", **{threshold: value}) + + +@pytest.mark.parametrize( + "threshold", ["dedicated_size_threshold", "pack_file_size_threshold"] +) +def test_blob_thresholds_reject_zero_for_positive_thresholds(threshold): + with pytest.raises(ValueError, match=f"{threshold} must be positive"): + lancedb.blob("blob", **{threshold: 0}) + + +def test_blob_thresholds_accept_integer_boundaries(): + maximum = ctypes.c_size_t(-1).value + field = lancedb.blob( + "blob", + nullable=False, + inline_size_threshold=0, + dedicated_size_threshold=maximum, + pack_file_size_threshold=maximum, + ) + assert not field.nullable + assert field.metadata == { + b"lance-encoding:blob-inline-size-threshold": b"0", + b"lance-encoding:blob-dedicated-size-threshold": str(maximum).encode(), + b"lance-encoding:blob-pack-file-size-threshold": str(maximum).encode(), + } + + +def test_blob_thresholds_control_writes_after_reopen(tmp_path): + # Keep this in test_table.py: CI runs it without pylance or pandas installed. + field = lancedb.blob( + "blob", + inline_size_threshold=4, + dedicated_size_threshold=12, + pack_file_size_threshold=8, + ) + metadata = { + b"lance-encoding:blob-inline-size-threshold": b"4", + b"lance-encoding:blob-dedicated-size-threshold": b"12", + b"lance-encoding:blob-pack-file-size-threshold": b"8", + } + assert field.metadata == metadata + schema = pa.schema([pa.field("id", pa.int64()), field]) + db = lancedb.connect(tmp_path) + table = db.create_table("blobs", schema=schema) + # Every payload fits in the default inline tier; custom thresholds must + # instead send the larger values to packed and dedicated storage. + rows = [ + {"id": 0, "blob": b"a" * 2}, + {"id": 1, "blob": b"b" * 8}, + {"id": 2, "blob": b"c" * 16}, + {"id": 3, "blob": b"d" * 8}, + ] + + def check_writes(table, expected): + assert table.schema.field("blob").metadata.items() >= metadata.items() + hits = table.search().with_row_id(True).limit(10).to_arrow() + actual = {row["id"]: row for row in hits.to_pylist()} + assert actual.keys() == {row["id"] for row in expected} + for row in expected: + size = len(row["blob"]) + descriptor = actual[row["id"]]["blob"] + assert descriptor["size"] == size + assert descriptor["kind"] == (2 if size > 12 else 1 if size > 4 else 0) + row_ids = [actual[row["id"]]["_rowid"] for row in expected] + assert table.fetch_blobs("blob", row_ids).to_pylist() == [ + row["blob"] for row in expected + ] + # Two packed values reach the configured sidecar limit independently. + packed = [ + actual[row["id"]]["blob"] for row in expected if len(row["blob"]) == 8 + ] + assert len({desc["blob_id"] for desc in packed}) == len(packed) + + table.add(rows) + check_writes(table, rows) + table = lancedb.connect(tmp_path).open_table("blobs") + check_writes(table, rows) + updates = [ + {"id": 0, "blob": b"e" * 16}, + {"id": 1, "blob": b"f" * 2}, + {"id": 2, "blob": b"g" * 8}, + {"id": 3, "blob": b"h" * 8}, + {"id": 4, "blob": b"i" * 16}, + ] + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(updates) + ) + assert result.num_updated_rows == 4 + assert result.num_inserted_rows == 1 + check_writes(lancedb.connect(tmp_path).open_table("blobs"), updates) + + async def _blob_v2_table_async(db: AsyncConnection, name: str): schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")]) table = await db.create_table(name, schema=schema) @@ -1634,6 +1752,28 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection): job.cancel() +@pytest.mark.parametrize( + "rq_kwargs", [{}, {"num_bits": None}, {"num_bits": 1}, {"num_bits": 4}] +) +@patch("lancedb.table.LanceTable.to_lance") +def test_create_ivfrq_index_accelerator(mock_to_lance, mem_db: DBConnection, rq_kwargs): + from lancedb.index import IvfRq + + table = mem_db.create_table("test", [{"vector": [1.0] * 8}]) + torch = MagicMock() + torch.cuda.is_available.return_value = True + with patch.dict(sys.modules, {"torch": torch}): + table.create_index("vector", config=IvfRq(accelerator="cuda", **rq_kwargs)) + + kwargs = mock_to_lance.return_value.create_index.call_args.kwargs + assert kwargs["index_type"] == "IVF_RQ" + assert kwargs["accelerator"] == "cuda" + if rq_kwargs.get("num_bits") is None: + assert "num_bits" not in kwargs + else: + assert kwargs["num_bits"] == rq_kwargs["num_bits"] + + @pytest.mark.parametrize("method", ["create_index", "create_index_async"]) @pytest.mark.parametrize("wait_timeout", ["invalid", 1, []]) def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout): @@ -2053,6 +2193,71 @@ def test_add_with_nans(mem_db: DBConnection): assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0])) +@pytest.mark.parametrize("operation", ["create", "append", "overwrite"]) +@pytest.mark.parametrize("input_type", ["pandas", "rows", "arrow"]) +@pytest.mark.parametrize("on_bad_vectors", [None, "error", "drop", "fill", "null"]) +def test_write_vectors_with_null_elements( + mem_db: DBConnection, operation, input_type, on_bad_vectors +): + schema = pa.schema({"vector": pa.list_(pa.float32(), 2), "id": pa.int64()}) + vectors = [[3.0, 4.0], [np.nan, 1.0], [np.nan, np.nan]] + if input_type == "pandas": + pd = pytest.importorskip("pandas") + + data = pd.DataFrame( + { + "vector": [np.array(v, dtype=np.float32) for v in vectors], + "id": [0, 1, 2], + } + ) + else: + # Pandas turns the NaNs above into null child elements. Explicit nulls + # in other input formats must follow the same bad-vector policy. + data = [ + {"vector": [3.0, 4.0], "id": 0}, + {"vector": [None, 1.0], "id": 1}, + {"vector": [None, None], "id": 2}, + ] + if input_type == "arrow": + data = pa.Table.from_pylist(data, schema=schema) + + options = {} if on_bad_vectors is None else {"on_bad_vectors": on_bad_vectors} + if on_bad_vectors == "fill": + options["fill_value"] = 42.0 + + table = None + if operation != "create": + table = mem_db.create_table("test", schema=schema) + + def write(): + if operation == "create": + return mem_db.create_table("test", data=data, schema=schema, **options) + table.add(data, mode=operation, **options) + return table + + if on_bad_vectors in (None, "error"): + # Create/overwrite can wrap errors from lazy Arrow readers in RuntimeError. + expected_error = ( + ValueError if operation == "append" else (ValueError, RuntimeError) + ) + with pytest.raises(expected_error, match="NaN|null"): + write() + if table is not None: + assert table.count_rows() == 0 + return + + table = write() + if on_bad_vectors == "drop": + expected = [{"vector": [3.0, 4.0], "id": 0}] + else: + expected = [ + {"vector": [3.0, 4.0], "id": 0}, + {"vector": [42.0, 1.0] if on_bad_vectors == "fill" else None, "id": 1}, + {"vector": [42.0, 42.0] if on_bad_vectors == "fill" else None, "id": 2}, + ] + assert table.to_arrow().to_pylist() == expected + + def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection): class Schema(LanceModel): text: str @@ -2979,6 +3184,57 @@ def test_update_with_arrow_scalar(mem_db: DBConnection): assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0, 3.0, 4.0]] +@pytest.mark.asyncio +@pytest.mark.parametrize("async_api", [False, True]) +@pytest.mark.parametrize("entries", [[("o'brien", 4), ("nil", None)], [], None]) +async def test_update_map_arrow_scalar(tmp_path, async_api, entries): + map_type = pa.map_(pa.string(), pa.int64()) + data = pa.table( + { + "id": [1, 2], + "attrs": pa.array([[("original", 1)], entries], type=map_type), + } + ) + if async_api: + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("maps", data) + source = await table.query().where("id = 2").to_arrow() + result = await table.update({"attrs": source["attrs"][0]}, where="id = 1") + actual = await table.query().to_arrow() + else: + db = lancedb.connect(tmp_path) + table = db.create_table("maps", data) + source = table.search().where("id = 2").to_arrow() + result = table.update(values={"attrs": source["attrs"][0]}, where="id = 1") + actual = table.to_arrow() + assert result.rows_updated == 1 + assert actual.schema.field("attrs").type == map_type + assert actual.sort_by("id")["attrs"].to_pylist() == [entries, entries] + + +@pytest.mark.parametrize( + "map_type, entries", + [ + (pa.map_(pa.int64(), pa.string()), []), + ( + pa.map_(pa.string(), pa.map_(pa.string(), pa.int64())), + [("outer", [("inner", 3)]), ("null", None)], + ), + ], +) +def test_update_typed_map_arrow_scalar(mem_db: DBConnection, map_type, entries): + table = mem_db.create_table( + "maps", + pa.table({"id": [1, 2], "attrs": pa.array([None, entries], type=map_type)}), + ) + scalar = table.search().where("id = 2").to_arrow()["attrs"][0] + result = table.update(values={"attrs": scalar}, where="id = 1") + actual = table.to_arrow().sort_by("id")["attrs"] + assert result.rows_updated == 1 + assert actual.type == map_type + assert actual.to_pylist() == [entries, entries] + + @pytest.mark.parametrize("use_arrow_scalar", [False, True]) @pytest.mark.parametrize( "decimal_type, value", @@ -3964,26 +4220,42 @@ def test_create_f16_table_from_arrow_data(mem_db: DBConnection): assert "s-2" in expected["text"].to_pylist() -def test_create_f16_table(mem_db: DBConnection): +@pytest.mark.parametrize("accelerator", [None, "cuda"]) +def test_create_f16_table(tmp_path, accelerator): + if accelerator == "cuda": + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + class MyTable(LanceModel): text: str vector: Vector(32, value_type=pa.float16()) rng = np.random.default_rng(42) + vectors = rng.standard_normal((512, 32)).astype(np.float16) df = pa.table( { "text": [f"s-{i}" for i in range(512)], - "vector": [rng.standard_normal(32).astype(np.float16) for _ in range(512)], + "vector": list(vectors), } ) - table = mem_db.create_table( + db = lancedb.connect(tmp_path) + table = db.create_table( "f16_tbl", schema=MyTable, ) table.add(df) - table.create_index(num_partitions=2, num_sub_vectors=2) + table.create_index( + "vector", + config=IvfPq( + num_partitions=2, + num_sub_vectors=2, + accelerator=accelerator, + ), + ) - query = df["vector"][2].as_py() + # Match the issue's float64 query against an explicitly typed float16 column. + query = vectors[2].astype(np.float64) expected = table.search(query).limit(2).to_arrow() assert "s-2" in expected["text"].to_pylist() diff --git a/python/python/tests/test_table_signatures.py b/python/python/tests/test_table_signatures.py new file mode 100644 index 000000000..346935b2b --- /dev/null +++ b/python/python/tests/test_table_signatures.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import timedelta +import inspect +from unittest.mock import AsyncMock, MagicMock + +import lancedb +from lancedb.expr import col +from lancedb.index import BTree, FTS, HnswSq, IvfPq +from lancedb.remote.table import RemoteTable +from lancedb.table import LanceTable, Table +import pyarrow as pa +import pytest + + +@pytest.mark.parametrize( + "method", + [ + "delete", + "drop_index", + "index_stats", + "create_index", + "create_scalar_index", + "create_fts_index", + "search", + "cleanup_old_versions", + "compact_files", + "optimize", + "to_polars", + ], +) +def test_sync_table_signatures(method): + """A Table call must bind the same arguments on either backend.""" + + def parameters(cls): + return [ + (p.name, p.kind, p.default) + for p in inspect.signature(getattr(cls, method)).parameters.values() + ] + + assert parameters(Table) == parameters(LanceTable) == parameters(RemoteTable) + + +@pytest.fixture(params=[LanceTable, RemoteTable], ids=["local", "remote"]) +def sync_table(request, monkeypatch): + """Exercise real sync wrappers with a recording async backend.""" + handle = MagicMock() + handle.name = "test" + handle.create_index = AsyncMock() + handle.delete = AsyncMock() + handle.drop_index = AsyncMock() + handle.index_stats = AsyncMock() + if request.param is RemoteTable: + table = RemoteTable(handle, "test_db") + else: + table = LanceTable.__new__(LanceTable) + table._table = handle + monkeypatch.setattr(table, "_ensure_no_legacy_fts_index", lambda: None) + return table, handle + + +def test_sync_keyword_names(sync_table): + table, handle = sync_table + predicate = col("id") == 1 + table.delete(where=predicate) + handle.delete.assert_awaited_once() + expected = predicate._inner if isinstance(table, LanceTable) else predicate + assert handle.delete.await_args.args[0] is expected + table.drop_index(name="vector_idx") + handle.drop_index.assert_awaited_once_with("vector_idx") + table.index_stats(index_name="vector_idx") + handle.index_stats.assert_awaited_once_with("vector_idx") + + +def test_legacy_vector_positional_arguments(sync_table): + table, handle = sync_table + with pytest.warns(DeprecationWarning): + table.create_index("cosine", 2) + args, kwargs = handle.create_index.await_args + assert args == ("vector",) + assert isinstance(kwargs["config"], IvfPq) + assert kwargs["config"].distance_type == "cosine" + assert kwargs["config"].num_partitions == 2 + assert kwargs["replace"] is True + + +def test_vector_tuning_options(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_index( + index_type="IVF_HNSW_SQ", + num_partitions=1, + m=10, + ef_construction=50, + max_iterations=12, + sample_rate=64, + target_partition_size=1000, + wait_timeout=timeout, + name="custom_idx", + replace=False, + train=False, + ) + kwargs = handle.create_index.await_args.kwargs + config = kwargs.pop("config") + assert isinstance(config, HnswSq) + for name, value in { + "num_partitions": 1, + "m": 10, + "ef_construction": 50, + "max_iterations": 12, + "sample_rate": 64, + "target_partition_size": 1000, + }.items(): + assert getattr(config, name) == value + assert kwargs == dict( + wait_timeout=timeout, name="custom_idx", replace=False, train=False + ) + + +def test_scalar_defaults_and_wait_timeout(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_scalar_index("id", wait_timeout=timeout) + kwargs = handle.create_index.await_args.kwargs + assert isinstance(kwargs["config"], BTree) + assert kwargs["replace"] is True + assert kwargs["wait_timeout"] == timeout + + +def test_fts_tokenizer_and_wait_timeout(sync_table): + table, handle = sync_table + timeout = timedelta(seconds=2) + with pytest.warns(DeprecationWarning): + table.create_fts_index( + field_names="text", + tokenizer_name="en_stem", + replace=True, + wait_timeout=timeout, + custom_stop_words=["example"], + ) + args, kwargs = handle.create_index.await_args + assert args == ("text",) + assert kwargs["replace"] is True + assert kwargs["wait_timeout"] == timeout + config = kwargs["config"] + assert isinstance(config, FTS) + assert config.language == "English" + assert config.stem is True + assert config.custom_stop_words == ["example"] + + +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"use_tantivy": True}, "Tantivy-based FTS has been removed"), + ({"ordering_field_names": "id"}, "ordering_field_names"), + ({"writer_heap_size": 123}, "writer_heap_size"), + ({"tokenizer_name": "invalid"}, "Invalid tokenizer name"), + ], +) +def test_fts_unsupported_options(sync_table, kwargs, message): + table, handle = sync_table + with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match=message): + table.create_fts_index("text", **kwargs) + handle.create_index.assert_not_awaited() + + +@pytest.mark.parametrize( + "method, kwargs", + [ + ( + "cleanup_old_versions", + {"older_than": timedelta(days=1), "delete_unverified": True}, + ), + ("compact_files", {"target_rows_per_fragment": 1000}), + ], +) +def test_remote_maintenance_options_warn(method, kwargs): + table = RemoteTable(MagicMock(name="test"), "test_db") + with pytest.warns(UserWarning, match="no-op"): + getattr(table, method)(**kwargs) + + +def test_remote_optimize_options(): + handle = MagicMock() + handle.optimize = AsyncMock(side_effect=NotImplementedError("remote optimize")) + table = RemoteTable(handle, "test_db") + age = timedelta(days=1) + with pytest.raises(NotImplementedError, match="remote optimize"): + table.optimize(cleanup_older_than=age, delete_unverified=True, retrain=True) + handle.optimize.assert_awaited_once_with( + cleanup_older_than=age, delete_unverified=True, retrain=True + ) + + +@pytest.mark.parametrize( + "method, old, new, value", + [ + ("delete", "predicate", "where", "id = 1"), + ("drop_index", "index_name", "name", "vector_idx"), + ("index_stats", "index_uuid", "index_name", "vector_idx"), + ("create_fts_index", "column", "field_names", "text"), + ], +) +def test_deprecated_remote_keywords(method, old, new, value): + handle = MagicMock() + handle.name = "test" + target = "create_index" if method == "create_fts_index" else method + setattr(handle, target, AsyncMock()) + table = RemoteTable(handle, "test_db") + with pytest.warns(DeprecationWarning, match=f"{old} is deprecated"): + getattr(table, method)(**{old: value}) + assert getattr(handle, target).await_args.args == (value,) + with pytest.raises(TypeError, match="Cannot specify both"): + getattr(table, method)(**{old: value, new: value}) + + +@pytest.mark.parametrize("remote_wrapper", [False, True], ids=["local", "remote"]) +def test_sync_search_and_polars(tmp_path, remote_wrapper): + """Use the native engine through both sync wrappers, without a cloud service.""" + local = lancedb.connect(tmp_path).create_table( + "test", + [{"id": 1, "text": "quick fox", "vector": [1.0, 1.0]}], + ) + table = RemoteTable(local._table, "test_db") if remote_wrapper else local + table.create_index("text", config=FTS()) + assert table.search("quick", query_type="fts").to_arrow()["id"].to_pylist() == [1] + # Both backends must accept the deprecated ordering argument in the same slot. + with pytest.warns(DeprecationWarning, match="ordering_field_name"): + table.search("quick", None, "fts", "id", ["text"]) + fast = table.search([1.0, 1.0], fast_search=True).to_arrow() + assert len(fast) == 0 # No vector index: fast_search excludes unindexed rows. + assert table.search([1.0, 1.0]).to_arrow()["id"].to_pylist() == [1] + if remote_wrapper: + with pytest.raises(NotImplementedError, match=r"table.search\(\).to_polars"): + table.to_polars(batch_size=1) + assert table.search().to_polars()["id"].to_list() == [1] + else: + assert table.to_polars(batch_size=1).collect()["id"].to_list() == [1] + table.delete(where="id = 1") + assert table.count_rows() == 0 + + +def test_index_replacement_defaults(tmp_path): + local = lancedb.connect(tmp_path).create_table("test", pa.table({"id": [1]})) + table = RemoteTable(local._table, "test_db") + for sync in [local, table]: + with pytest.warns(DeprecationWarning): + sync.create_scalar_index("id") + assert len(list(table.list_indices())) == 1 diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 4fb16bb31..e85afb5bc 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -406,7 +406,7 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): ).read_all() if on_bad_vectors == "drop": - expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]]) + expected = pa.array([[1.0, 2.0], [4.0, 5.0]]) elif on_bad_vectors == "fill": expected = pa.array( [ @@ -414,12 +414,12 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): [3.0, 42.0], [4.0, 5.0], [6.0, 7.0], - [None, 9.0], + [42.0, 9.0], [42.0, 42.0], ] ) elif on_bad_vectors == "null": - expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None]) + expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, None, None]) assert output["vector"].combine_chunks() == expected @@ -430,17 +430,17 @@ def test_handle_bad_vectors_jagged(on_bad_vectors): ( pa.list_(pa.float64()), [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]], - [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]], + [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]], ), ( pa.large_list(pa.float64()), [[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]], - [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]], + [[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [42.0, 3.0], [4.0, 5.0]], ), ( pa.list_(pa.float64(), 2), [[1.0, float("nan")], None, [None, 3.0]], - [[1.0, 42.0], [42.0, 42.0], [None, 3.0]], + [[1.0, 42.0], [42.0, 42.0], [42.0, 3.0]], ), ], ) @@ -467,7 +467,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors): ).read_all() output = exception_output(e) assert output == ( - "ValueError: Vector column 'vector' has NaNs. Set " + "ValueError: Vector column 'vector' has NaNs or null values. Set " "on_bad_vectors='drop' to remove them, set on_bad_vectors='fill' " "and fill_value= to replace them, or set on_bad_vectors='null' " "to replace them with null." @@ -490,6 +490,46 @@ def test_handle_bad_vectors_nan(on_bad_vectors): assert output["vector"].combine_chunks() == expected +@pytest.mark.parametrize("value_type", [pa.float16(), pa.float32(), pa.float64()]) +@pytest.mark.parametrize( + "list_type", [pa.list_, pa.large_list, lambda t: pa.list_(t, 2)] +) +@pytest.mark.parametrize("chunked", [False, True]) +def test_has_nan_values_with_null_elements(value_type, list_type, chunked): + arr = pa.array( + [[None, None], [1.0, 2.0], [None, 3.0], None, [4.0, float("nan")], [5.0, 6.0]], + type=list_type(value_type), + ).slice(1, 5) + if chunked: + arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2)]) + # Whole null vectors have no components to inspect. Their handling is + # separate from missing or NaN components in a valid vector. + assert has_nan_values(arr).to_pylist() == [False, True, False, True, False] + + +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +def test_sanitize_vectors_with_only_null_elements(on_bad_vectors): + data = pa.table({"vector": [[None, None]]}) + schema = pa.schema({"vector": pa.list_(pa.float32(), 2)}) + + def sanitize(): + return _sanitize_data( + data, schema, on_bad_vectors=on_bad_vectors, fill_value=42.0 + ).read_all() + + if on_bad_vectors == "error": + with pytest.raises(ValueError, match="null"): + sanitize() + return + + expected = { + "drop": [], + "fill": [{"vector": [42.0, 42.0]}], + "null": [{"vector": None}], + } + assert sanitize().to_pylist() == expected[on_bad_vectors] + + @pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()]) @pytest.mark.parametrize("large", [False, True]) @pytest.mark.parametrize("chunked", [False, True]) diff --git a/python/src/connection.rs b/python/src/connection.rs index baeb29780..2cdf71f4a 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -590,29 +590,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_namespaces( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListNamespacesRequest; - let request = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_namespaces(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("namespaces", response.namespaces)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[pyo3(signature = (namespace_path, mode=None, properties=None))] @@ -725,29 +713,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_tables( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListTablesRequest; - let request = ListTablesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_tables(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("tables", response.tables)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Get the configuration for constructing an equivalent namespace client. diff --git a/python/src/index.rs b/python/src/index.rs index d4e16a67e..c6c3fe17d 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -136,8 +136,10 @@ pub fn extract_index_params(source: &Option>) -> PyResult, - num_bits: u32, + num_bits: Option, max_iterations: u32, sample_rate: u32, target_partition_size: Option, diff --git a/python/src/job.rs b/python/src/job.rs index e22b2f897..1740b3e00 100644 --- a/python/src/job.rs +++ b/python/src/job.rs @@ -81,6 +81,11 @@ impl Job { self.inner.id().map(str::to_string) } + #[getter] + pub fn job_type(&self) -> Option { + self.inner.job_type() + } + pub fn status(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.inner.clone(); future_into_py( diff --git a/python/src/listing.rs b/python/src/listing.rs index 2d1a3916e..79a7effdd 100644 --- a/python/src/listing.rs +++ b/python/src/listing.rs @@ -3,10 +3,11 @@ use crate::error::PythonErrorExt; use crate::runtime::future_into_py; -use futures::{StreamExt, future::poll_fn}; +use futures::{Future, StreamExt, future::poll_fn}; use lancedb::listing::{Listing, ListingOptions}; use pyo3::exceptions::PyStopAsyncIteration; -use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyDict, PyDictMethods}; +use pyo3::{Bound, Py, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; use std::sync::{Arc, Mutex}; pub fn options(page_token: Option, page_limit: Option) -> ListingOptions { @@ -52,6 +53,31 @@ impl NameListing { Ok(value) }) } + + fn next_page(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + let next_lock = self_.next_lock.clone(); + future_into_py(self_.py(), async move { + let _guard = next_lock.lock().await; + let page = futures::future::poll_fn(|cx| { + let mut inner = inner.lock().unwrap(); + let future = inner.next_page(); + futures::pin_mut!(future); + future.poll(cx) + }) + .await + .infer_error()?; + Python::attach(|py| -> PyResult>> { + page.map(|page| { + let dict = PyDict::new(py); + dict.set_item("items", page.items)?; + dict.set_item("page_token", page.page_token)?; + Ok(dict.unbind()) + }) + .transpose() + }) + }) + } } #[pyclass] diff --git a/python/src/table.rs b/python/src/table.rs index e77c58578..e58b8a139 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -24,7 +24,7 @@ use arrow::{ ffi_stream::ArrowArrayStreamReader, pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; -use lancedb::blob::{BlobFile, BlobRangeRequest}; +use lancedb::blob::{BlobDescriptor, BlobFile, BlobRangeRequest}; use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, @@ -1551,6 +1551,23 @@ impl Table { }) } + #[pyo3(signature = (descriptor))] + pub fn open_blob_file<'py>( + self_: PyRef<'py, Self>, + descriptor: &str, + ) -> PyResult> { + let descriptor: BlobDescriptor = serde_json::from_str(descriptor) + .map_err(|error| PyValueError::new_err(format!("invalid blob descriptor: {error}")))?; + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let file = inner.open_blob_file(&descriptor).await.infer_error()?; + Ok(PyBlobFile { + inner: Arc::new(file), + closed: AtomicBool::new(false), + }) + }) + } + /// Optimize the on-disk data by compacting and pruning old data, for better performance. #[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))] pub fn optimize( diff --git a/python/uv.lock b/python/uv.lock index e808f5f9e..99bcb439c 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -2006,8 +2006,8 @@ requires-dist = [ { name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" }, { name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" }, { name = "pydantic", specifier = ">=2.7.4,<3" }, - { name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" }, - { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" }, + { name = "pylance", marker = "extra == 'pylance'", specifier = ">=15.0.0b1" }, + { name = "pylance", marker = "extra == 'tests'", specifier = "==15.0.0b1" }, { 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" }, @@ -3938,7 +3938,7 @@ crypto = [ [[package]] name = "pylance" -version = "14.0.0b12" +version = "15.0.0b1" source = { registry = "https://pypi.fury.io/lance-format/" } dependencies = [ { name = "lance-namespace" }, @@ -3948,12 +3948,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" }, - { url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" }, - { url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" }, - { url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" }, - { url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" }, - { url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" }, + { url = "https://pypi.fury.io/lance-format/-/ver_23L1Nd/pylance-15.0.0b1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:19605d3daca1224cc0db079054be5d4eec403dcf286bf4617430bfc259ab0a52" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1Ho3RK/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be22b0b6f1e1d250ffaa82b449e982c817db33103c483d1005020285310623cf" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1QmyF0/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5b2761d60b1b9db112058213f1f7b2de82b08cbe65888fab08b0a56e702482e" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1in2Ph/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:294a0fedae312ed32d1b66cddbf7e69770ec9b4c82a873af60c619ed88450718" }, + { url = "https://pypi.fury.io/lance-format/-/ver_21xpLH/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ff20a16703d4d3f3bc675b25ce1871ba1d54137929dc46305bffdd07f7b35a6b" }, + { url = "https://pypi.fury.io/lance-format/-/ver_AAGzF/pylance-15.0.0b1-cp310-abi3-win_amd64.whl", hash = "sha256:3b134e24a5becfc940ca134cf78aa3bae8b80ec40f8927a1b856273ab76213f2" }, ] [[package]] diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index 486eb1e8e..c18db11fe 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.41.0-beta.13" +version = "0.41.0-beta.14" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true @@ -46,7 +46,7 @@ lance-index = { workspace = true, features = ["tokenizer-jieba", "tokenizer-lind lance-table = { workspace = true } lance-linalg = { workspace = true } lance-encoding = { workspace = true } -lance-arrow = { workspace = true } +lance-arrow = { version = "=15.0.0-beta.1", tag = "v15.0.0-beta.1", git = "https://github.com/lance-format/lance.git" } lance-namespace = { workspace = true } lance-namespace-impls = { workspace = true } metrics = { workspace = true, optional = true } diff --git a/rust/lancedb/src/blob.rs b/rust/lancedb/src/blob.rs index 05650d7c2..142b2264e 100644 --- a/rust/lancedb/src/blob.rs +++ b/rust/lancedb/src/blob.rs @@ -20,11 +20,120 @@ use arrow_schema::{DataType, Field, Schema}; use lance::dataset::{BlobRangeRequest as LanceBlobRangeRequest, Dataset, WriteParams}; use lance_arrow::FieldExt; use lance_file::version::{ConcreteFileVersion, LanceFileVersion}; -use lance_io::object_store::ObjectStore; +use lance_io::object_store::{ObjectStore, uri_to_url}; use object_store::path::Path; +use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; +/// A stored Blob V2 descriptor with the source needed to reopen its snapshot. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct BlobDescriptor { + /// Canonical Blob V2 storage kind (0 through 5). + pub kind: u8, + /// Stored offset within the backing object. + pub position: u64, + /// Stored length; zero on an external blob can mean unspecified. + pub size: u64, + /// Stored sidecar or base identifier. + pub blob_id: u32, + /// Stored external URI or managed path. + pub blob_uri: String, + /// Source identity captured by the query. + pub source: BlobDescriptorSource, +} + +/// The pinned source of a descriptor returned by SQL. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct BlobDescriptorSource { + /// URI of the dataset on the source branch. + pub uri: String, + /// Exact table version read by the query. + pub version: u64, + /// Commit timestamp in microseconds, distinguishing recreated tables. + pub committed_at_micros: u64, + /// Source branch; `None` identifies main. + pub branch: Option, + /// Blob column path in the source table. + pub column: String, + /// Row ID or row address, as specified by `row_is_address`. + pub row_reference: u64, + /// Whether the row reference is a physical address. + pub row_is_address: bool, +} + +impl BlobDescriptor { + pub(crate) fn validate(&self) -> Result<()> { + if self.kind > 5 + || self.source.version == 0 + || self.source.uri.is_empty() + || self.source.column.is_empty() + { + return Err(Error::InvalidInput { + message: "invalid stored blob descriptor or missing pinned source".into(), + }); + } + Ok(()) + } +} + +pub(crate) async fn open_native_blob_file( + dataset: &Dataset, + descriptor: &BlobDescriptor, +) -> Result { + descriptor.validate()?; + let source = &descriptor.source; + let expected_uri = dataset + .branch_location() + .find_branch(source.branch.as_deref())? + .uri; + if uri_to_url(&expected_uri)? != uri_to_url(&source.uri)? { + return Err(Error::InvalidInput { + message: "blob descriptor belongs to a different table".into(), + }); + } + let dataset = Arc::new( + dataset + .checkout_version(lance::dataset::refs::Ref::Version( + source.branch.clone(), + Some(source.version), + )) + .await?, + ); + if uri_to_url(dataset.uri())? != uri_to_url(&source.uri)? { + return Err(Error::InvalidInput { + message: "blob descriptor belongs to a different table".into(), + }); + } + if dataset.manifest().timestamp_nanos / 1_000 != u128::from(source.committed_at_micros) { + return Err(Error::InvalidInput { + message: "blob descriptor source snapshot has been replaced".into(), + }); + } + ensure_blob_v2_column(dataset.schema(), &source.column)?; + let mut files = if source.row_is_address { + dataset + .take_blobs_by_addresses(&[source.row_reference], &source.column) + .await? + } else { + dataset + .take_blobs(&[source.row_reference], &source.column) + .await? + }; + if files.len() != 1 { + return Err(Error::InvalidInput { + message: "blob descriptor does not resolve to one stored value".into(), + }); + } + files + .pop() + .flatten() + .map(Into::into) + .ok_or_else(|| Error::InvalidInput { + message: "non-null blob descriptor resolved to a null value".into(), + }) +} + /// Seekable handle for one blob value, backed by local storage or a remote /// HTTP byte-range endpoint. #[derive(Debug)] diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 8a0e489c2..bff97bf31 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -13,7 +13,7 @@ use lance::dataset::refs::MAIN_BRANCH; use lance_namespace::models::{ CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest, DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest, - ListNamespacesResponse, ListTablesRequest, ListTablesResponse, + ListTablesRequest, }; #[cfg(feature = "aws")] use object_store::aws::AwsCredential; @@ -28,7 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; -use crate::listing::{Listing, ListingOptions}; +use crate::listing::{Listing, ListingOptions, ListingPage}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -174,6 +174,8 @@ impl OpenTableBuilder { /// Advanced parameters that can be used to customize table reads /// + /// Remote connections reject these parameters with [`Error::NotSupported`]. + /// /// If set, these will take precedence over any overlapping `OpenTableOptions` options pub fn lance_read_params(mut self, params: ReadParams) -> Self { self.request.lance_read_params = Some(params); @@ -372,12 +374,10 @@ impl ExecuteQueryAsyncBuilder { /// Do not follow a named placeholder directly with `$`: `$a$` starts a /// dollar-quoted string. /// - /// A parameterized statement runs on the call that returns its rows rather - /// than detached from it: cancelling the query, or dropping its handle - /// before reading it, stops the statement on the server. Its status follows - /// that stream too -- it reports `Running` until the reader has received - /// every row -- so open the reader directly rather than waiting for - /// `Finished`. + /// With or without parameters, the server may stream rows on the initial + /// call or return a polling continuation. The reader follows either path + /// transparently. Open it directly instead of waiting for `Finished`: + /// directly streamed queries need their reader to make progress. /// /// # Example /// @@ -1126,12 +1126,43 @@ impl Connection { self.internal.drop_all_tables(namespace_path).await } - /// List immediate child namespace names in the given namespace - pub async fn list_namespaces( + /// Lazily list immediate child namespace names. Each request asks for at + /// most 100 names unless `options.page_limit` is set. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_namespaces(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_namespaces( &self, - request: ListNamespacesRequest, - ) -> Result { - self.internal.list_namespaces(request).await + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_namespaces(ListNamespacesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.namespaces, + page_token: response.page_token, + }) + } + }) } /// Create a new namespace @@ -1182,23 +1213,43 @@ impl Connection { self.internal.namespace_client_config().await } - /// List tables with pagination support + /// Lazily list tables. Each request asks for at most 100 names unless + /// `options.page_limit` is set. /// - /// The default limit is 100 tables per page for both local and remote connections, - /// including namespaces. Zero returns an empty page without a continuation token. - /// Follow the response's opaque `page_token` until it is absent to retrieve every table; - /// a page can contain fewer than the limit even when more tables remain. - pub async fn list_tables(&self, mut request: ListTablesRequest) -> Result { - // Apply the SDK's page size here so every backend receives the same request. - let limit = request.limit.get_or_insert(100); - if *limit <= 0 { - return Ok(ListTablesResponse { - context: None, - tables: Vec::new(), - page_token: None, - }); - } - self.internal.list_tables(request).await + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_tables(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_tables( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_tables(ListTablesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.tables, + page_token: response.page_token, + }) + } + }) } /// Get the in-memory embedding registry. @@ -1219,6 +1270,7 @@ pub struct ConnectRequest { /// - `/path/to/database` - local database on file system. /// - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud object store /// - `db://dbname` - LanceDB Cloud + /// - `db://` with a host override - remote tables in the storage root pub uri: String, #[cfg(feature = "remote")] @@ -1331,6 +1383,8 @@ impl ConnectBuilder { /// /// This option is only used when connecting to LanceDB Cloud (db:// URIs) /// and will be ignored for other URIs. + /// Use the URI `db://` together with a host override to connect to remote + /// tables stored directly in the storage root. /// /// # Arguments /// @@ -1478,6 +1532,7 @@ impl ConnectBuilder { /// Enable or disable manifest-backed directory namespace mode for local /// native connections. + /// Remote connections reject enabling this option with [`Error::NotSupported`]. /// /// When enabled, the connection uses the directory namespace database /// directly for all table operations and forces @@ -1513,6 +1568,7 @@ impl ConnectBuilder { } /// Set a custom session for object stores and caching. + /// Remote connections reject this option with [`Error::NotSupported`]. /// /// By default, a new session with default configuration will be created. /// This method allows you to provide a custom session with your own @@ -1544,6 +1600,21 @@ impl ConnectBuilder { fn execute_remote(self) -> Result { use crate::remote::db::RemoteDatabaseOptions; + for (option, provided) in [ + ("session", self.request.session.is_some()), + ("manifest_enabled", self.request.manifest_enabled), + ( + "namespace_client_properties", + !self.request.namespace_client_properties.is_empty(), + ), + ] { + if provided { + return Err(Error::NotSupported { + message: format!("{option} is not supported for remote connections"), + }); + } + } + let mut merged_options = self.request.options.clone(); Self::apply_env_defaults(&ENV_VARS_TO_STORAGE_OPTS, &mut merged_options); let sql_host_override = merged_options.get(OPT_REMOTE_SQL_HOST_OVERRIDE).cloned(); @@ -1956,6 +2027,37 @@ mod tests { } } + #[cfg(feature = "remote")] + #[tokio::test] + async fn test_connect_remote_storage_root() { + let conn = ConnectBuilder::new("db://") + .region("us-east-1") + .api_key("my-api-key") + .host_override("https://example.com") + .execute() + .await + .unwrap(); + + let (impl_name, properties) = conn.namespace_client_config().await.unwrap(); + assert_eq!(impl_name, "rest"); + assert_eq!(properties["uri"], "https://example.com"); + assert!( + !properties.contains_key("header.x-lancedb-database"), + "storage-root connections must not send the x-lancedb-database header" + ); + + let result = ConnectBuilder::new("db://") + .region("us-east-1") + .api_key("my-api-key") + .execute() + .await; + assert!(matches!( + result, + Err(Error::InvalidInput { message }) + if message.contains("A host override is required") + )); + } + #[cfg(feature = "remote")] #[tokio::test] async fn test_connect_rejects_header_provider_with_oauth_config() { @@ -2260,26 +2362,11 @@ mod tests { } names.sort(); - // Walking in pages has to reach every table exactly once, with nothing lost at a - // page boundary. - let mut seen = Vec::with_capacity(names.len()); - let mut page_token = None; - loop { - let page = db - .list_tables(ListTablesRequest { - id: Some(Vec::new()), - limit: Some(2), - page_token, - ..Default::default() - }) - .await - .unwrap(); - seen.extend(page.tables); - page_token = page.page_token.filter(|token| !token.is_empty()); - if page_token.is_none() { - break; - } - } + let seen = futures::TryStreamExt::try_collect::>( + db.list_tables(&[], ListingOptions::default().page_limit(2)), + ) + .await + .unwrap(); assert_eq!(seen, names); } @@ -2305,37 +2392,20 @@ mod tests { .is_empty() ); - let first = db.list_tables(ListTablesRequest::default()).await.unwrap(); - assert_eq!(first.tables, names[..100]); + let mut listing = db.list_tables(&[], ListingOptions::default()); + assert_eq!(listing.num_page_results(), 0); + let first = listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items, names[..100]); assert!(first.page_token.is_some()); - let second = db - .list_tables(ListTablesRequest { - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(second.tables, names[100..]); + let second = listing.next_page().await.unwrap().unwrap(); + assert_eq!(second.items, names[100..]); assert!(second.page_token.is_none()); + assert!(listing.next_page().await.unwrap().is_none()); - let all = db - .list_tables(ListTablesRequest { - limit: Some(200), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(all.tables, names); - assert!(all.page_token.is_none()); - let empty = db - .list_tables(ListTablesRequest { - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(empty.tables.is_empty()); - assert!(empty.page_token.is_none()); + let mut all = db.list_tables(&[], ListingOptions::default().page_limit(200)); + assert_eq!(all.next_page().await.unwrap().unwrap().items, names); + let mut invalid = db.list_tables(&[], ListingOptions::default().page_limit(0)); + assert!(invalid.next_page().await.is_err()); } #[tokio::test] diff --git a/rust/lancedb/src/connection/create_table.rs b/rust/lancedb/src/connection/create_table.rs index 39cc82ec0..9d4408ba7 100644 --- a/rust/lancedb/src/connection/create_table.rs +++ b/rust/lancedb/src/connection/create_table.rs @@ -45,6 +45,9 @@ impl CreateTableBuilder { } /// Apply the given write options when writing the initial data + /// + /// Remote connections reject explicit Lance write parameters with + /// [`Error::NotSupported`] instead of silently using server defaults. pub fn write_options(mut self, write_options: WriteOptions) -> Self { self.request.write_options = write_options; self @@ -70,6 +73,8 @@ impl CreateTableBuilder { /// Set multiple options for the storage layer. /// + /// Table-level storage options are not supported by remote connections. + /// /// Options already set on the connection will be inherited by the table, /// but can be overridden here. /// diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 98d915118..36e6c30c8 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -156,11 +156,20 @@ pub struct FunctionArtifact { } /// One ordered Arrow input parameter. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum BlobAccess { + Bytes, + Reader, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionParameter { pub name: String, pub arrow_type: String, pub nullable: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub blob_access: Option, } /// One field of an ordered named-struct result. @@ -956,6 +965,20 @@ mod conda_environment_tests { mod tests { use super::*; + #[test] + fn blob_reader_access_survives_signature_round_trip() { + let parameter: FunctionParameter = serde_json::from_value(serde_json::json!({ + "name": "video", "arrow_type": "blob_v2", "nullable": false, + "blob_access": "reader" + })) + .unwrap(); + assert_eq!(parameter.blob_access, Some(BlobAccess::Reader)); + assert_eq!( + serde_json::to_value(parameter).unwrap()["blob_access"], + "reader" + ); + } + /// Canonical form is what the FunctionVersion hash is taken over, so key /// order must come from the keys and not from however serde happened to /// emit them. Nesting is included because the sort is recursive. diff --git a/rust/lancedb/src/index/vector.rs b/rust/lancedb/src/index/vector.rs index bce77a610..3e6794b54 100644 --- a/rust/lancedb/src/index/vector.rs +++ b/rust/lancedb/src/index/vector.rs @@ -362,6 +362,9 @@ impl IvfRqIndexBuilder { impl_distance_type_setter!(); impl_ivf_params_setter!(); + /// Set the number of bits used to encode each vector dimension. + /// + /// If unset, uses Lance's default (currently 5). Set to 1 for maximum compression. pub fn num_bits(mut self, num_bits: u32) -> Self { self.num_bits = Some(num_bits); self diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 6865e59ee..5c27587dc 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -13,6 +13,7 @@ use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; use tokio::sync::watch; use tokio::task::{AbortHandle, JoinHandle}; +use uuid::Uuid; use crate::database::JobDescription; use crate::error::{Error, JobFailure, Result}; @@ -55,10 +56,13 @@ fn job_detail_not_supported(what: &str) -> Result { /// Backend-specific tracking for an asynchronous operation. #[async_trait] pub(crate) trait JobHandle: Send + Sync { - /// Server-assigned id, when the backend has one. + /// Backend-specific id, when the operation has one. fn id(&self) -> Option<&str> { None } + fn job_type(&self) -> Option<&str> { + None + } async fn status(&self) -> Result; async fn wait(&self) -> Result; async fn cancel(&self) -> Result<()>; @@ -167,14 +171,15 @@ struct JobCache { /// the endpoint's successful terminal result; unit-result operations use the /// default `Job<()>`. /// -/// The detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the +/// Registry detail accessors ([`Job::state`], [`Job::job_type`], ...) read what the /// handle last observed. Submitting an operation returns only a job id, so /// populating them eagerly would cost an extra round trip on every call: /// /// - [`Job::refresh`] and [`Job::status`] fetch the whole record. /// - [`Job::wait`] records the terminal state it establishes, but not the rest /// of the record; call [`Job::refresh`] for that. -/// - Everything is `None` until one of those runs. +/// - Registry details are `None` until one of those runs. Query-backed handles +/// report [`Job::job_type`] as `"query"` immediately. pub struct Job where T: Clone + Send + Sync + 'static, @@ -273,13 +278,18 @@ where Self::new_typed(Box::new(SpawnedJob::new(task))) } - pub(crate) fn spawned_with_cancellation(task: JoinHandle>, cancel: F) -> Self + pub(crate) fn from_query_task( + query_id: Uuid, + task: JoinHandle>, + cancel: F, + ) -> Self where F: Fn() -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { let cancellation: Cancellation = Arc::new(move || Box::pin(cancel())); Self::new_typed(Box::new(SpawnedJob::new_with_cancellation( + query_id, task, cancellation, ))) @@ -290,12 +300,10 @@ impl Job where T: Clone + Send + Sync + 'static, { - /// Identifies the operation on the server that is running it. - /// - /// Returned for correlating with server logs or the jobs API. Operations - /// that run in this process have no server id and return `None`. The - /// value is opaque: parsing it or storing it to resume the job later is - /// not supported. + /// The operation's opaque identity, interpreted according to [`Job::job_type`]. + /// Registry handles expose a registry job ID; query handles expose the + /// connection-scoped SQL query UUID as a string. Local handles return None. + /// A query ID cannot be used with registry APIs or to resume after restart. pub fn id(&self) -> Option<&str> { match &self.inner { JobInner::Handle { handle, .. } => handle.id(), @@ -369,8 +377,14 @@ where self.cache_read().description.clone() } - /// The job's type, as the server names it. `None` for an in-process job. + /// The job's type: `"query"` for SQL-backed handles, otherwise the cached + /// server-reported kind of work. `None` for an in-process job. pub fn job_type(&self) -> Option { + if let JobInner::Handle { handle, .. } = &self.inner + && let Some(job_type) = handle.job_type() + { + return Some(job_type.to_string()); + } self.with_description(|description| description.job_type.clone()) } @@ -536,6 +550,7 @@ struct SpawnedJob { outcome: watch::Receiver>, abort: AbortHandle, cancellation: Option, + id: Option, } type CancellationFuture = Pin> + Send>>; @@ -549,11 +564,17 @@ impl SpawnedJob { Self::new_inner(task, None) } - fn new_with_cancellation(task: JoinHandle>, cancellation: Cancellation) -> Self + fn new_with_cancellation( + query_id: Uuid, + task: JoinHandle>, + cancellation: Cancellation, + ) -> Self where T: Serialize + Send + 'static, { - Self::new_inner(task, Some(cancellation)) + let mut handle = Self::new_inner(task, Some(cancellation)); + handle.id = Some(query_id.to_string()); + handle } fn new_inner(task: JoinHandle>, cancellation: Option) -> Self @@ -582,12 +603,21 @@ impl SpawnedJob { outcome, abort, cancellation, + id: None, } } } #[async_trait] impl JobHandle for SpawnedJob { + fn id(&self) -> Option<&str> { + self.id.as_deref() + } + + fn job_type(&self) -> Option<&str> { + self.id.as_ref().map(|_| "query") + } + async fn status(&self) -> Result { let label = match &*self.outcome.borrow() { None => "running", @@ -636,15 +666,43 @@ mod tests { use super::*; + #[test] + fn local_typed_result_decode_errors_are_runtime_errors() { + assert!(matches!( + TerminalResult::local(Value::Null).decode::(), + Err(Error::Runtime { .. }) + )); + } + #[tokio::test] async fn mapped_spawned_job_reuses_outcome() { let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1); + assert_eq!(job.job_type(), None); + assert_eq!(job.id(), None); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.wait().await.unwrap(), 42); assert_eq!(job.status().await.unwrap(), "finished"); } + #[tokio::test] + async fn mapped_query_job_preserves_identity_and_type() { + let query_id = Uuid::new_v4(); + let job = Job::from_query_task(query_id, tokio::spawn(async { Ok(41_u64) }), || async { + Ok(()) + }) + .map(|value| value + 1); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.wait().await.unwrap(), 42); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!(job.id(), Some(query_id.to_string().as_str())); + let completed = Job::new_done(); + assert_eq!(completed.job_type(), None); + assert_eq!(completed.id(), None); + } + #[tokio::test] async fn mapped_spawned_job_preserves_cancellation() { let job = Job::spawned(tokio::spawn(async { pending::>().await })) diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index e15017e7f..68521ff2c 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -54,6 +54,7 @@ //! - `/path/to/database` - local database on file system. //! - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud object store //! - `db://dbname` - Lance Cloud +//! - `db://` with a host override - remote tables in the storage root //! //! You can also use [`ConnectBuilder`] to configure the connection to the database. //! diff --git a/rust/lancedb/src/listing.rs b/rust/lancedb/src/listing.rs index 710202d9d..a0f6e5d5e 100644 --- a/rust/lancedb/src/listing.rs +++ b/rust/lancedb/src/listing.rs @@ -51,8 +51,9 @@ type FetchPage = /// A lazy stream of resources with inspectable pagination state. /// -/// No requests are made until the stream is polled. Empty pages with continuation -/// tokens are skipped. Errors terminate the stream and retain the failed request's +/// No requests are made until the stream is polled or [`Self::next_page`] is called. +/// Item iteration skips empty pages with continuation tokens; `next_page` returns +/// them. Errors terminate the listing and retain the failed request's /// token, allowing a new stream to retry it. Repeated tokens are rejected. /// /// A token advances when its page is fetched. Drain [`Self::num_page_results`] @@ -100,6 +101,64 @@ impl Listing { pub fn page_token(&self) -> Option<&str> { self.options.page_token.as_deref() } + + /// Return exactly one page without consuming a following page. + /// + /// If item iteration has started, this returns the remaining items in its + /// cached page. An empty page with a continuation token is returned as is. + pub async fn next_page(&mut self) -> Result>> { + if self.items.len() == 0 { + if self.done { + return Ok(None); + } + futures::future::poll_fn(|cx| self.poll_fetch_page(cx)).await?; + } + let items = std::mem::replace(&mut self.items, Vec::new().into_iter()).collect(); + Ok(Some(ListingPage { + items, + page_token: self.options.page_token.clone(), + })) + } + + fn poll_fetch_page(&mut self, cx: &mut Context<'_>) -> Poll> { + if self + .options + .page_limit + .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) + { + self.done = true; + return Poll::Ready(Err(Error::InvalidInput { + message: "Listing page limit must be between 1 and 2147483647".into(), + })); + } + let pending = self + .pending + .get_or_insert_with(|| (self.fetch)(self.options.clone())); + let response = futures::ready!(pending.as_mut().poll(cx)); + self.pending = None; + match response { + Ok(response) => { + let token = response.page_token.filter(|token| !token.is_empty()); + if token + .as_ref() + .is_some_and(|token| !self.seen_tokens.insert(token.clone())) + { + self.done = true; + return Poll::Ready(Err(Error::Runtime { + message: "Listing response repeated a page_token".into(), + })); + } + self.done = token.is_none(); + self.options.page_token = token; + self.items = response.items.into_iter(); + Poll::Ready(Ok(())) + } + Err(error) => { + self.done = true; + Poll::Ready(Err(error)) + } + } + } } impl Stream for Listing { @@ -114,41 +173,8 @@ impl Stream for Listing { if this.done { return Poll::Ready(None); } - if this - .options - .page_limit - .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) - { - this.done = true; - return Poll::Ready(Some(Err(Error::InvalidInput { - message: "Listing page limit must be between 1 and 2147483647".into(), - }))); - } - let pending = this - .pending - .get_or_insert_with(|| (this.fetch)(this.options.clone())); - let response = futures::ready!(pending.as_mut().poll(cx)); - this.pending = None; - match response { - Ok(response) => { - let token = response.page_token.filter(|token| !token.is_empty()); - if token - .as_ref() - .is_some_and(|token| !this.seen_tokens.insert(token.clone())) - { - this.done = true; - return Poll::Ready(Some(Err(Error::Runtime { - message: "Listing response repeated a page_token".into(), - }))); - } - this.done = token.is_none(); - this.options.page_token = token; - this.items = response.items.into_iter(); - } - Err(error) => { - this.done = true; - return Poll::Ready(Some(Err(error))); - } + if let Err(error) = futures::ready!(this.poll_fetch_page(cx)) { + return Poll::Ready(Some(Err(error))); } } } @@ -224,6 +250,36 @@ mod tests { assert!(listing.next().await.is_none()); } + #[tokio::test] + async fn next_page_preserves_empty_pages_and_remaining_items() { + let requests = Arc::new(Mutex::new(0)); + let count = requests.clone(); + let mut listing = Listing::new(ListingOptions::default(), move |options| { + *count.lock().unwrap() += 1; + async move { + let (items, page_token) = match options.page_token.as_deref() { + None => (vec![], Some("second".to_string())), + Some("second") => (vec!["a", "b"], Some("last".to_string())), + Some("last") => (vec!["c"], None), + _ => unreachable!(), + }; + Ok(ListingPage { items, page_token }) + } + }); + assert_eq!(*requests.lock().unwrap(), 0); + assert!(listing.next_page().await.unwrap().unwrap().items.is_empty()); + assert_eq!(*requests.lock().unwrap(), 1); + assert_eq!(listing.try_next().await.unwrap(), Some("a")); + assert_eq!(*requests.lock().unwrap(), 2); + let remaining = listing.next_page().await.unwrap().unwrap(); + assert_eq!(remaining.items, vec!["b"]); + assert_eq!(remaining.page_token.as_deref(), Some("last")); + assert_eq!(*requests.lock().unwrap(), 2); + assert_eq!(listing.next_page().await.unwrap().unwrap().items, vec!["c"]); + assert!(listing.next_page().await.unwrap().is_none()); + assert_eq!(*requests.lock().unwrap(), 3); + } + #[tokio::test] async fn invalid_limits_do_not_fetch() { for limit in [0, u32::MAX] { diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index ecdd1783b..7ee76a160 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -1716,7 +1716,8 @@ impl CreateMaterializedViewBuilder { } /// Submit creation and initial population, returning a [`Job`] that - /// settles when the view is ready. + /// settles when the view is ready. Remote creation submits SQL; the handle's + /// [`Job::job_type`] is `"query"` and its ID identifies that query. pub async fn execute_async(self) -> Result { if self.connection.uri().starts_with("db://") { return self diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 00a2b88f3..9bac51209 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -2714,6 +2714,159 @@ mod tests { } } + #[tokio::test] + async fn test_limit_usize_max() { + let conn = connect("memory://foo").execute().await.unwrap(); + let table = conn + .create_table("my_table", make_non_empty_batches()) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .offset(1) + .nearest_to(&[0.1; 4]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected - 1); + } + + #[tokio::test] + async fn test_indexed_vector_limit_usize_max() { + use crate::index::vector::IvfFlatIndexBuilder; + + let conn = connect("memory://indexed_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_non_empty_batches()) + .execute() + .await + .unwrap(); + table + .create_index( + &["vector"], + Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(1)), + ) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + for limit in [expected, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, usize::MAX] { + for (bypass, refine_factor) in [(false, None), (false, Some(2)), (true, None)] { + let query = table + .query() + .limit(limit) + .offset(offset) + .nearest_to(&[0.1; 4]) + .unwrap(); + let query = if bypass { + query.bypass_vector_index() + } else { + query + }; + let query = if let Some(factor) = refine_factor { + query.refine_factor(factor) + } else { + query + }; + let batches: Vec = + query.execute().await.unwrap().try_collect().await.unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected.saturating_sub(offset), + "limit={limit}, offset={offset}, bypass={bypass}, refine={refine_factor:?}" + ); + } + } + } + } + + #[tokio::test] + async fn test_empty_vector_limit_usize_max() { + let conn = connect("memory://empty_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_test_batches()) + .execute() + .await + .unwrap(); + let batches: Vec = table + .query() + .limit(usize::MAX) + .nearest_to(&[0.1; 128]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + + #[tokio::test] + async fn test_fts_limit_usize_max_with_offset() { + let conn = connect("memory://fts_limit").execute().await.unwrap(); + let data = RecordBatch::try_from_iter([( + "text", + Arc::new(StringArray::from_iter_values((0..32).map(|_| "hello"))) + as arrow_array::ArrayRef, + )]) + .unwrap(); + let table = conn + .create_table("documents", data) + .execute() + .await + .unwrap(); + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + + for limit in [32, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, i64::MAX as usize, usize::MAX] { + let batches: Vec = table + .query() + .full_text_search(FullTextSearchQuery::new("hello".into())) + .limit(limit) + .offset(offset) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 32usize.saturating_sub(offset), + "limit={limit}, offset={offset}" + ); + } + } + } + #[tokio::test] async fn test_repeated_only_if_combines_with_and() { use crate::expr::{col, lit}; diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index 783d2ef84..f54e2fe3e 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -378,23 +378,30 @@ pub struct ParsedDbUrl { /// Parse a database URL and extract the database name and optional prefix. /// -/// Expected format: `db://db_name` or `db://db_name/prefix` +/// Expected format: `db://db_name`, `db://db_name/prefix`, or `db://` when +/// connecting to the storage root through a host override. pub fn parse_db_url(db_url: &str) -> Result { let parsed_url = url::Url::parse(db_url).map_err(|err| Error::InvalidInput { message: format!("db_url is not a valid URL. '{db_url}'. Error: {err}"), })?; debug_assert_eq!(parsed_url.scheme(), "db"); - if !parsed_url.has_host() { - return Err(Error::InvalidInput { - message: format!("Invalid database URL (missing host) '{}'", db_url), - }); - } - let db_name = urlencoding::decode(parsed_url.host_str().unwrap()) - .map_err(|err| Error::InvalidInput { - message: format!("Invalid encoded database name: {err}"), - })? - .into_owned(); - validate_database_name(&db_name)?; + let db_name = match parsed_url.host_str() { + Some(db_name) => { + let db_name = urlencoding::decode(db_name) + .map_err(|err| Error::InvalidInput { + message: format!("Invalid encoded database name: {err}"), + })? + .into_owned(); + validate_database_name(&db_name)?; + db_name + } + None if matches!(parsed_url.path(), "" | "/") => String::new(), + None => { + return Err(Error::InvalidInput { + message: format!("Invalid database URL (missing host) '{}'", db_url), + }); + } + }; let db_prefix = { let prefix = parsed_url.path().trim_start_matches('/'); if prefix.is_empty() { @@ -802,6 +809,11 @@ impl RestfulLanceDbClient { self.client.get(full_uri) } + pub fn head(&self, uri: &str) -> RequestBuilder { + let full_uri = format!("{}{}", self.host, uri); + self.client.head(full_uri) + } + pub fn post(&self, uri: &str) -> RequestBuilder { let full_uri = format!("{}{}", self.host, uri); self.client.post(full_uri) diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 4a8cbf32e..5c299dc5a 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -63,15 +63,6 @@ struct RemoteCloneTableRequest { is_shallow: Option, } -#[derive(serde::Serialize)] -struct RemoteCreateMaterializedViewRequest { - query: String, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_enabled: Option, - #[serde(skip_serializing_if = "Option::is_none")] - auto_refresh_interval_seconds: Option, -} - // the versions of the server that we support // for any new feature that we need to change the SDK behavior, we should bump the server version, // and add a feature flag as method of `ServerVersion` here. @@ -318,6 +309,13 @@ impl RemoteDatabase { read_consistency_interval: Option, ) -> Result { let parsed = super::client::parse_db_url(uri)?; + if parsed.db_name.is_empty() && host_overrides.rest.is_none() { + return Err(Error::InvalidInput { + message: + "A host override is required when connecting to the storage root with 'db://'" + .to_string(), + }); + } Self::try_new_with_identity( uri, api_key, @@ -587,6 +585,11 @@ mod test_utils { } } + pub(in crate::remote) fn with_sql_client(mut self, client: SqlClient) -> Self { + self.sql_client = Some(client); + self + } + pub fn new_mock_with_config(handler: F, config: ClientConfig) -> Self where F: Fn(reqwest::Request) -> http::Response + Send + Sync + 'static, @@ -878,8 +881,7 @@ impl Database for RemoteDatabase { ) -> Result { if request.with_no_data { return Err(Error::NotSupported { - message: "remote materialized views are always populated by their SQL job" - .to_string(), + message: "remote materialized views must be populated during creation".to_string(), }); } if let Some(seconds) = request.auto_refresh_interval_seconds @@ -889,36 +891,37 @@ impl Database for RemoteDatabase { message: "auto_refresh_interval_seconds must be positive, finite seconds".into(), }); } - let identifier = build_table_identifier(&request.name, &request.namespace_path)?; - let body = RemoteCreateMaterializedViewRequest { - query: request.query, - auto_refresh_enabled: request.auto_refresh_enabled, - auto_refresh_interval_seconds: request.auto_refresh_interval_seconds, - }; - let request = self - .client - .post(&format!("/v1/materialized_view/{identifier}/create")) - .json(&body); - let (request_id, response) = self.client.send(request).await?; - let response = self.client.check_response(&request_id, response).await?; - let status = response.status(); - let body = response.text().await.err_to_http(request_id.clone())?; - if status != StatusCode::ACCEPTED { - return Err(Error::Http { - source: "materialized-view creation must return 202 Accepted".into(), - request_id, - status_code: Some(status), - }); + validate_table_name(&request.name)?; + validate_namespace(&request.namespace_path)?; + let client = self + .sql_client + .as_ref() + .ok_or_else(|| Error::NotSupported { + message: "SQL is unavailable for this remote database client".into(), + })?; + let mut options = Vec::new(); + if let Some(enabled) = request.auto_refresh_enabled { + options.push(format!("auto_refresh_enabled = {enabled}")); } - let job_id = extract_job_id(&body).ok_or_else(|| Error::Http { - source: "materialized-view create response did not contain a valid job_id".into(), - request_id, - status_code: Some(status), - })?; - Ok(Job::new(Box::new(RemoteJob::new( - self.client.clone(), - job_id, - )))) + if let Some(seconds) = request.auto_refresh_interval_seconds { + options.push(format!("auto_refresh_interval_seconds = {seconds}")); + } + let name = request.name.replace('"', "\"\""); + let options = if options.is_empty() { + String::new() + } else { + format!(" WITH ({})", options.join(", ")) + }; + client + .submit_as_job( + format!( + "CREATE MATERIALIZED VIEW \"{name}\"{options} AS {}", + request.query + ), + request.namespace_path, + || async { Ok(()) }, + ) + .await } async fn drop_materialized_view_async( @@ -1478,6 +1481,23 @@ impl Database for RemoteDatabase { } async fn create_table(&self, mut request: CreateTableRequest) -> Result> { + // The remote create endpoint does not transmit Lance write parameters. + // Reject them before consuming data or issuing a request, especially + // for overwrite: silently using the server's format can lose intent. + for (option, provided) in [ + ( + "write_options (including storage_options)", + request.write_options.lance_write_params.is_some(), + ), + ("location", request.location.is_some()), + ("namespace_client", request.namespace_client.is_some()), + ] { + if provided { + return Err(Error::NotSupported { + message: format!("{option} is not supported for remote table creation"), + }); + } + } let data_schema = request.data.schema(); let body = stream_as_body(request.data.scan_as_stream())?; @@ -1602,11 +1622,39 @@ impl Database for RemoteDatabase { } async fn open_table(&self, request: OpenTableRequest) -> Result> { + // Validate before looking in the cache as well as before sending HTTP. + for (option, provided) in [ + ("index_cache_size", request.index_cache_size.is_some()), + ( + "lance_read_params (including storage_options)", + request.lance_read_params.is_some(), + ), + ("location", request.location.is_some()), + ("namespace_client", request.namespace_client.is_some()), + ("managed_versioning", request.managed_versioning.is_some()), + ] { + if provided { + return Err(Error::NotSupported { + message: format!("{option} is not supported for remote table access"), + }); + } + } let identifier = build_table_identifier(&request.name, &request.namespace_path)?; let cache_key = build_cache_key(&request.name, &request.namespace_path); // Every open gets its own checkout, schema cache, and freshness state. if let Some(version) = self.table_cache.get(&cache_key).await { + // A cached version can outlive a drop from another connection. Verify that + // the table still exists before returning a handle for it. + let req = self + .client + .post(&format!("/v1/table/{}/describe/", identifier)) + .json(&serde_json::json!({})); + let (request_id, rsp) = self.client.send_with_retry(req, None, true).await?; + let rsp = + RemoteTable::::handle_table_not_found(&request.name, rsp, &request_id).await?; + self.client.check_response(&request_id, rsp).await?; + Ok(Arc::new(RemoteTable::new_with_sql_client( self.client.clone(), request.name, @@ -1707,7 +1755,9 @@ impl Database for RemoteDatabase { // TODO: Implement namespace-aware drop_all_tables let _namespace_path = namespace_path; // Suppress unused warning for now Err(crate::Error::NotSupported { - message: "Dropping all tables is not currently supported in the remote API".to_string(), + message: "Dropping all tables is not currently supported in the remote API. \ + List tables and call drop_table for each table instead." + .to_string(), }) } @@ -1996,72 +2046,22 @@ mod tests { } #[tokio::test] - async fn test_create_materialized_view_uses_rest_and_preserves_optional_settings() { - for (enabled, interval) in [(None, None), (Some(true), None), (Some(false), Some(0.5))] { - let db = super::RemoteDatabase::new_mock(move |request| { - assert_eq!(request.method(), "POST"); - assert_eq!( - request.url().path(), - "/v1/materialized_view/analytics$adults/create" - ); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - assert_eq!(body["query"], "SELECT * FROM people"); - assert_eq!( - body.get("auto_refresh_enabled"), - enabled.map(serde_json::Value::from).as_ref() - ); - assert_eq!( - body.get("auto_refresh_interval_seconds"), - interval.map(serde_json::Value::from).as_ref() - ); - http::Response::builder() - .status(202) - .body(serde_json::json!({"job_id": "j1-mv-create"}).to_string()) - .unwrap() - }); - let job = db - .create_materialized_view_async(CreateMaterializedViewRequest { - name: "adults".into(), - namespace_path: vec!["analytics".into()], - query: "SELECT * FROM people".into(), - with_no_data: false, - auto_refresh_enabled: enabled, - auto_refresh_interval_seconds: interval, - }) - .await - .unwrap(); - assert_eq!(job.id(), Some("j1-mv-create")); - } - } - - #[tokio::test] - async fn test_create_materialized_view_rejects_invalid_acceptance() { - for (status, body) in [ - (200, r#"{"job_id":"j1"}"#), - (202, "{}"), - (202, r#"{"job_id":""}"#), - ] { - let db = super::RemoteDatabase::new_mock(move |_| { - http::Response::builder().status(status).body(body).unwrap() - }); - let error = db + async fn test_create_materialized_view_validates_before_submission() { + let db = super::RemoteDatabase::new_mock(|_| -> http::Response { + panic!("materialized-view creation must not send REST requests") + }); + for seconds in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let result = db .create_materialized_view_async(CreateMaterializedViewRequest { name: "adults".into(), namespace_path: vec![], query: "SELECT * FROM people".into(), with_no_data: false, auto_refresh_enabled: None, - auto_refresh_interval_seconds: None, + auto_refresh_interval_seconds: Some(seconds), }) - .await - .err() - .unwrap(); - assert!(error.to_string().contains(if status == 200 { - "202 Accepted" - } else { - "job_id" - })); + .await; + assert!(matches!(result, Err(Error::InvalidInput { .. }))); } } @@ -2349,16 +2349,8 @@ mod tests { .unwrap() .is_empty() ); - let page = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(page.tables.is_empty()); - assert!(page.page_token.is_none()); + let mut listing = conn.list_tables(&namespace, ListingOptions::default().page_limit(0)); + assert!(listing.next_page().await.is_err()); } #[rstest::rstest] @@ -2393,25 +2385,17 @@ mod tests { ) .unwrap() }); - let first = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(first.tables.len(), 100); + let mut first_listing = conn.list_tables(&namespace, ListingOptions::default()); + let first = first_listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items.len(), 100); assert_eq!(first.page_token.as_deref(), Some("opaque-token")); - let second = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); + let mut second_listing = conn.list_tables( + &namespace, + ListingOptions::default().page_token(first.page_token.unwrap()), + ); + let second = second_listing.next_page().await.unwrap().unwrap(); let expected: Vec<_> = (100..130).map(|i| format!("t{i:03}")).collect(); - assert_eq!(second.tables, expected); + assert_eq!(second.items, expected); assert!(second.page_token.is_none()); } @@ -2554,14 +2538,45 @@ mod tests { let table = conn.open_table("table1").execute().await.unwrap(); assert_eq!(table.name(), "table1"); - // Storage options should be ignored. - let table = conn + // Per-table storage options must not be silently ignored. + let result = conn .open_table("table1") .storage_option("key", "value") .execute() - .await - .unwrap(); - assert_eq!(table.name(), "table1"); + .await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains("storage_options"))); + } + + #[tokio::test] + async fn test_open_table_checks_that_a_cached_table_still_exists() { + let describe_calls = Arc::new(AtomicUsize::new(0)); + let calls = describe_calls.clone(); + let mut db = super::RemoteDatabase::new_mock(move |request| { + assert_eq!(request.url().path(), "/v1/table/table1/describe/"); + match calls.fetch_add(1, Ordering::SeqCst) { + 0 => http::Response::builder() + .status(200) + .body(r#"{"version": 1}"#.to_string()) + .unwrap(), + 1 => http::Response::builder() + .status(404) + .body("table not found".to_string()) + .unwrap(), + call => panic!("unexpected describe call: {call}"), + } + }); + db.table_cache = moka::future::Cache::new(10); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + + conn.open_table("table1").execute().await.unwrap(); + let result = conn.open_table("table1").execute().await; + + assert!(matches!(result, Err(Error::TableNotFound { .. }))); + assert_eq!(describe_calls.load(Ordering::SeqCst), 2); } #[tokio::test] @@ -2754,6 +2769,126 @@ mod tests { ); } + #[rstest] + #[case("new_table_data_storage_version", "2.0")] + #[case("new_table_enable_v2_manifest_paths", "false")] + #[case("bogus_option_xyz", "1")] + #[tokio::test] + async fn test_create_rejects_local_storage_options(#[case] key: &str, #[case] value: &str) { + let conn = Connection::new_with_handler(|_| -> http::Response<&'static str> { + panic!("must not send a create request") + }); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let data = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + for builder in [ + conn.create_table("table1", data), + conn.create_empty_table("table1", schema), + ] { + let result = builder + .mode(CreateTableMode::Overwrite) + .storage_option(key, value) + .execute() + .await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains("write_options") && message.contains("remote"))); + } + } + + #[tokio::test] + async fn test_create_rejects_local_write_params_and_location() { + let conn = Connection::new_with_handler(|_| -> http::Response<&'static str> { + panic!("must not send a create request") + }); + let schema = Arc::new(Schema::empty()); + for builder in [ + conn.create_empty_table("table1", schema.clone()) + .write_options(crate::table::WriteOptions { + lance_write_params: Some(lance::dataset::WriteParams::default()), + }), + conn.create_empty_table("table1", schema) + .location("/tmp/elsewhere"), + ] { + assert!(matches!( + builder.execute().await, + Err(Error::NotSupported { .. }) + )); + } + } + + #[rstest] + #[tokio::test] + async fn test_open_rejects_local_options_even_when_cached(#[values(false, true)] cached: bool) { + let mut db = super::RemoteDatabase::new_mock(|request| { + assert_eq!(request.url().path(), "/v1/table/table1/create/"); + http::Response::builder().status(200).body("").unwrap() + }); + db.table_cache = moka::future::Cache::new(10); + let table_cache = db.table_cache.clone(); + let conn = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + if cached { + conn.create_empty_table("table1", Arc::new(Schema::empty())) + .execute() + .await + .unwrap(); + } + assert_eq!( + table_cache + .get(&build_cache_key("table1", &[])) + .await + .is_some(), + cached + ); + for (option, builder) in [ + ( + "index_cache_size", + conn.open_table("table1").index_cache_size(10), + ), + ( + "storage_options", + conn.open_table("table1").storage_option("x", "y"), + ), + ("location", conn.open_table("table1").location("/tmp/zzz")), + ( + "managed_versioning", + conn.open_table("table1").managed_versioning(false), + ), + ( + "lance_read_params", + conn.open_table("table1") + .lance_read_params(Default::default()), + ), + ] { + let result = builder.execute().await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains(option) && message.contains("remote"))); + } + } + + #[tokio::test] + async fn test_connect_rejects_local_options() { + for (option, builder) in [ + ( + "manifest_enabled", + ConnectBuilder::new("db://dev").manifest_enabled(true), + ), + ( + "namespace_client_properties", + ConnectBuilder::new("db://dev").namespace_client_properties([("root", "/tmp/db")]), + ), + ] { + let result = builder.api_key("fake").region("us-east-1").execute().await; + assert!(matches!(result, Err(Error::NotSupported { message }) + if message.contains(option) && message.contains("remote"))); + } + } + #[rstest] #[case::matching(Field::new("a", DataType::Int32, false), true)] #[case::different_name(Field::new("x", DataType::Int32, false), false)] @@ -3680,14 +3815,9 @@ mod tests { assert!(table.is_ok(), "Failed to create table: {:?}", table.err()); // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .expect("Failed to list tables"); - assert_eq!(list_response.tables, vec!["test_table"]); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names, vec!["test_table"]); // Get namespace client and verify it can also list tables let namespace_client = conn.namespace_client().await.unwrap(); @@ -3754,17 +3884,12 @@ mod tests { } // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(list_response.tables.len(), 3); - assert!(list_response.tables.contains(&"table1".to_string())); - assert!(list_response.tables.contains(&"table2".to_string())); - assert!(list_response.tables.contains(&"table3".to_string())); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names.len(), 3); + assert!(names.contains(&"table1".to_string())); + assert!(names.contains(&"table2".to_string())); + assert!(names.contains(&"table3".to_string())); } } diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 41d8971ac..7d2e70ac9 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -18,7 +18,9 @@ use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. const INITIAL_POLL_INTERVAL: Duration = Duration::from_millis(200); -const MAX_POLL_INTERVAL: Duration = Duration::from_secs(5); +/// A job's completion is seen at most this long after it happens. A 5 s cap +/// added 2.5 s on average to every short refresh, as much as the job itself. +const MAX_POLL_INTERVAL: Duration = Duration::from_secs(1); #[derive(Debug, Clone, PartialEq, Eq)] enum JobState { @@ -268,13 +270,18 @@ impl JobHandle for RemoteJob { #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + use async_trait::async_trait; use crate::Result; use crate::function::{FunctionVersion, RefreshColumnResult}; use crate::job::{Job, JobHandle, TerminalResult}; + use crate::remote::client::test_utils::client_with_handler; - use super::DescribeJobResponse; + use super::{DescribeJobResponse, RemoteJob}; const FUNCTION_JOB: &str = include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); @@ -333,6 +340,31 @@ mod tests { ); } + /// Five in-progress polls then done: 0.2 + 0.4 + 0.8 + 1 + 1 s of waiting. + /// Under the old 5 s cap the same job would have waited 10 s. + #[tokio::test(start_paused = true)] + async fn wait_polls_at_most_every_second() { + let describes = Arc::new(AtomicUsize::new(0)); + let seen = describes.clone(); + let client = client_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/jobs/describe"); + let n = seen.fetch_add(1, Ordering::SeqCst); + let body = if n < 5 { + r#"{"job_state":"IN_PROGRESS"}"#.to_string() + } else { + UNIT_JOB.to_string() + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let started = tokio::time::Instant::now(); + RemoteJob::new(client, "job-1".to_string()) + .wait() + .await + .expect("unit job completes"); + assert_eq!(describes.load(Ordering::SeqCst), 6); + assert_eq!(started.elapsed(), Duration::from_millis(3400)); + } + #[test] fn remote_wire_unknown_fields_are_forward_decodable() { let response: DescribeJobResponse = diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index a58a9afad..846e852f5 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::fs; -use std::future::{Future, pending}; +use std::future::Future; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::{Duration, Instant}; @@ -12,6 +12,7 @@ use arrow_array::RecordBatch; use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::error::FlightError; use arrow_flight::flight_service_client::FlightServiceClient; +#[cfg(test)] use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; use arrow_flight::{ Action, CancelFlightInfoRequest, CancelFlightInfoResult, CancelStatus, FlightClient, @@ -22,7 +23,7 @@ use futures::TryStreamExt; use http::header::{HeaderMap, HeaderName, HeaderValue}; use prost::Message; use serde::{Serialize, de::DeserializeOwned}; -use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch}; +use tokio::sync::{Mutex, Notify, OnceCell, mpsc}; use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity}; use uuid::Uuid; @@ -34,9 +35,9 @@ use crate::remote::client::{ClientConfig, TlsConfig}; use crate::remote::retry::ResolvedRetryConfig; use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; -mod exchange; +mod stream; -use exchange::{ExchangeQuery, ExchangeQueryHandle}; +use stream::{StreamingQuery, StreamingQueryHandle, Submission}; const DEFAULT_SQL_PORT: u16 = 10025; const DEFAULT_SQL_TLS_PORT: u16 = 10026; @@ -191,22 +192,42 @@ impl SqlClient { } pub(super) async fn execute(&self, request: ExecuteQueryRequest) -> Result { - match request.parameters { - Some(parameters) => { - self.submit_with_parameters( - &request.query, - parameters, - &request.default_namespace_path, - ) - .await + let timeout = self.inner.overall_timeout()?; + with_overall_timeout(timeout, "SQL query submission", async { + validate_namespace_path(&request.default_namespace_path)?; + let query_id = Uuid::now_v7(); + let submission = StreamingQuery::open( + self.inner.clone(), + &request.query, + request.parameters, + &request.default_namespace_path, + ) + .await?; + match submission { + Submission::Stream(response) => { + let query = + Arc::new(StreamingQuery::new(query_id, self.inner.clone(), response)); + self.queries + .insert(query_id, RegisteredQuery::Streamed(query.clone())); + Ok(Query::new(Arc::new(StreamingQueryHandle::new(query)))) + } + Submission::Deferred(poll_info) => { + let query = Arc::new(RemoteQuery::new( + query_id, + self.inner.clone(), + request.default_namespace_path, + poll_info, + )?); + self.queries + .insert(query_id, RegisteredQuery::Polled(query.clone())); + Ok(Query::new(Arc::new(RemoteQueryHandle::new(query)))) + } } - None => { - self.submit(&request.query, &request.default_namespace_path) - .await - } - } + }) + .await } + #[cfg(test)] pub(super) async fn submit( &self, query: &str, @@ -235,69 +256,36 @@ impl SqlClient { .await } - pub(super) fn submit_as_job( + pub(super) async fn submit_as_job( &self, statement: String, namespace: Vec, finish: F, - ) -> Job + ) -> Result> where T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, F: FnOnce() -> Fut + Send + 'static, Fut: Future> + Send + 'static, { - let client = self.clone(); - let (query_tx, query_rx) = watch::channel(None::>); + let query = Arc::new( + self.execute(ExecuteQueryRequest { + query: statement, + default_namespace_path: namespace, + parameters: None, + }) + .await?, + ); + let query_id = query.id(); + let execution = query.clone(); let task = tokio::spawn(async move { - let query = Arc::new(client.submit(&statement, &namespace).await?); - query_tx.send_replace(Some(query.clone())); - let mut reader = query.reader().await?; + let mut reader = execution.reader().await?; while reader.try_next().await?.is_some() {} finish().await }); - Job::spawned_with_cancellation(task, move || { - let mut query_rx = query_rx.clone(); - async move { - loop { - let query = query_rx.borrow().clone(); - if let Some(query) = query { - return query.cancel().await; - } - if query_rx.changed().await.is_err() { - pending::<()>().await; - } - } - } - }) - } - - /// Values cannot ride on a poll, so a statement that has them is sent on - /// its own exchange instead. - pub(super) async fn submit_with_parameters( - &self, - query: &str, - parameters: RecordBatch, - default_namespace_path: &[String], - ) -> Result { - let timeout = self.inner.overall_timeout()?; - with_overall_timeout(timeout, "SQL query submission", async { - validate_namespace_path(default_namespace_path)?; - let query_id = Uuid::now_v7(); - let query = Arc::new( - ExchangeQuery::open( - query_id, - self.inner.clone(), - query, - parameters, - default_namespace_path, - ) - .await?, - ); - self.queries - .insert(query_id, RegisteredQuery::Exchanged(query.clone())); - Ok(Query::new(Arc::new(ExchangeQueryHandle::new(query)))) - }) - .await + Ok(Job::from_query_task(query_id, task, move || { + let query = query.clone(); + async move { query.cancel().await } + })) } pub(super) async fn describe(&self, query_id: Uuid) -> Result { @@ -325,6 +313,7 @@ impl SqlClientInner { ) } + #[cfg(test)] async fn poll( &self, descriptor: FlightDescriptor, @@ -586,14 +575,14 @@ impl SqlClientInner { #[derive(Clone)] enum RegisteredQuery { Polled(Arc), - Exchanged(Arc), + Streamed(Arc), } impl RegisteredQuery { fn touch(&self) { match self { Self::Polled(query) => query.touch(), - Self::Exchanged(query) => query.touch(), + Self::Streamed(query) => query.touch(), } } @@ -602,14 +591,14 @@ impl RegisteredQuery { fn registry_expired(&self) -> bool { match self { Self::Polled(query) => query.registry_expired(Arc::strong_count(query) == 1), - Self::Exchanged(query) => query.registry_expired(Arc::strong_count(query) == 1), + Self::Streamed(query) => query.registry_expired(Arc::strong_count(query) == 1), } } async fn describe(&self) -> Result { match self { Self::Polled(query) => query.describe().await, - Self::Exchanged(query) => query.describe(), + Self::Streamed(query) => query.describe(), } } } diff --git a/rust/lancedb/src/remote/sql/exchange.rs b/rust/lancedb/src/remote/sql/stream.rs similarity index 69% rename from rust/lancedb/src/remote/sql/exchange.rs rename to rust/lancedb/src/remote/sql/stream.rs index 523aaf5d7..ed85c6c92 100644 --- a/rust/lancedb/src/remote/sql/exchange.rs +++ b/rust/lancedb/src/remote/sql/stream.rs @@ -1,26 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! SQL statements sent with parameters. -//! -//! A statement without parameters is polled: `PollFlightInfo` starts it -//! detached from the connection and `DoGet` reads its endpoints. Neither call -//! carries Arrow data, so a statement's values would have to be spelled into -//! its text. A statement with parameters is sent on one `DoExchange` instead: -//! the command and the parameter schema, then one row, and the result rows -//! come back on the same call. The call is the statement's lifetime, so -//! cancelling it -- or dropping the handle before reading -- stops the -//! statement on the server. +//! SQL submitted on one call, with either streamed rows or a polling handoff. use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; use std::time::Instant; use arrow_array::RecordBatch; -use arrow_flight::FlightDescriptor; +use arrow_flight::decode::FlightRecordBatchStream; use arrow_flight::encode::FlightDataEncoderBuilder; -use arrow_flight::sql::{CommandStatementQuery, ProstMessageExt}; +use arrow_flight::sql::{Any, CommandStatementQuery, ProstMessageExt}; +use arrow_flight::{FlightDescriptor, PollInfo, Ticket}; use arrow_schema::SchemaRef; +use futures::TryStreamExt; use prost::Message; use tokio::sync::{Notify, mpsc}; use uuid::Uuid; @@ -51,8 +44,8 @@ fn failure(error: &Error) -> Lifecycle { }) } -/// A statement running on its own `DoExchange` call. -pub(super) struct ExchangeQuery { +/// A statement streaming on its initial `DoGet` or `DoExchange` call. +pub(super) struct StreamingQuery { id: Uuid, client: Arc, /// The call's rows, until a reader takes them or a cancellation drops them. @@ -63,7 +56,7 @@ pub(super) struct ExchangeQuery { last_accessed: StdMutex, } -struct PreparedExchange { +struct PreparedStream { schema: SchemaRef, first: Option, rest: Option, @@ -76,7 +69,7 @@ struct PreparedExchange { /// this one did. A query cancelled meanwhile has its rows dropped instead, /// which ends the call. struct Borrowed<'a> { - query: &'a ExchangeQuery, + query: &'a StreamingQuery, rows: Option, } @@ -104,18 +97,37 @@ impl Drop for Borrowed<'_> { } } -impl ExchangeQuery { - /// Send the statement and its parameters, returning once the server has - /// planned it: a statement it refuses is refused here, as a polled one is - /// on its first poll. +pub(super) const PROTOCOL_HEADER: &str = "lancedb-sql-protocol"; +pub(super) const PROTOCOL_VERSION: &str = "adaptive-v1"; +const RESPONSE_TYPE_URL: &str = "type.googleapis.com/lancedb.sql.QueryResponse"; + +#[derive(Clone, PartialEq, Message)] +pub(super) struct QueryResponse { + #[prost(message, optional, tag = "1")] + pub poll_info: Option, +} + +#[derive(Clone, PartialEq, Message)] +struct CommandTicket { + #[prost(bytes = "vec", tag = "2")] + command: Vec, +} + +pub(super) enum Submission { + Stream(ResultEndpointStream), + Deferred(PollInfo), +} + +impl StreamingQuery { pub(super) async fn open( - id: Uuid, client: Arc, query: &str, - parameters: RecordBatch, + parameters: Option, default_namespace_path: &[String], - ) -> Result { - if parameters.num_rows() != 1 { + ) -> Result { + if let Some(parameters) = ¶meters + && parameters.num_rows() != 1 + { return Err(Error::InvalidInput { message: format!( "SQL query parameters must be exactly one row, got {}", @@ -130,39 +142,102 @@ impl ExchangeQuery { Some(DEFAULT_READ_TIMEOUT), )? .unwrap(); - let descriptor = FlightDescriptor::new_cmd( - CommandStatementQuery { - query: query.to_string(), - transaction_id: None, + let command = CommandStatementQuery { + query: query.to_string(), + transaction_id: None, + } + .as_any() + .encode_to_vec(); + let connection = client.connection(&request_id).await?; + let headers = client.headers(default_namespace_path, &request_id).await?; + let mut metadata = super::client_with_headers(connection.client.clone(), &headers)? + .metadata() + .clone(); + metadata.insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + let mut flight = connection.client.clone(); + let response = if let Some(parameters) = parameters { + let messages = FlightDataEncoderBuilder::new() + .with_flight_descriptor(Some(FlightDescriptor::new_cmd(command))) + .with_schema(parameters.schema()) + .build(futures::stream::iter([Ok(parameters)])) + .try_collect::>() + .await + .map_err(|error| sql_error(&request_id, error))?; + let mut request = tonic::Request::new(futures::stream::iter(messages)); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_exchange(request)).await + } else { + let mut request = + tonic::Request::new(Ticket::new(CommandTicket { command }.encode_to_vec())); + *request.metadata_mut() = metadata; + tokio::time::timeout(read_timeout, flight.do_get(request)).await + } + .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? + .map_err(|error| sql_error(&request_id, error))?; + let negotiated = response + .metadata() + .get(PROTOCOL_HEADER) + .is_some_and(|value| value == PROTOCOL_VERSION); + let mut response_stream = response.into_inner(); + if negotiated { + let first = tokio::time::timeout(read_timeout, response_stream.message()) + .await + .map_err(|_| sql_error(&request_id, "SQL query handoff timed out"))? + .map_err(|error| sql_error(&request_id, error))? + .ok_or_else(|| { + sql_error(&request_id, "SQL response omitted its execution decision") + })?; + if !first.data_header.is_empty() || !first.data_body.is_empty() { + return Err(sql_error( + &request_id, + "SQL response sent rows before its execution decision", + )); } - .as_any() - .encode_to_vec(), + let envelope = + Any::decode(first.app_metadata).map_err(|error| sql_error(&request_id, error))?; + if envelope.type_url != RESPONSE_TYPE_URL { + return Err(sql_error( + &request_id, + "SQL response returned an unknown execution decision", + )); + } + let decision = QueryResponse::decode(envelope.value) + .map_err(|error| sql_error(&request_id, error))?; + if let Some(info) = decision.poll_info { + if info.info.is_none() && info.flight_descriptor.is_none() { + return Err(sql_error( + &request_id, + "SQL handoff omitted result and continuation information", + )); + } + return Ok(Submission::Deferred(info)); + } + } + let stream = FlightRecordBatchStream::new_from_flight_data( + response_stream + .map_err(|status| arrow_flight::error::FlightError::Tonic(Box::new(status))), ); - // The descriptor rides on the schema message, and the row follows it. - let request = FlightDataEncoderBuilder::new() - .with_flight_descriptor(Some(descriptor)) - .with_schema(parameters.schema()) - .build(futures::stream::iter([Ok(parameters)])); - let mut flight = client - .client_with_headers(default_namespace_path, &request_id) - .await?; - let stream = tokio::time::timeout(read_timeout, flight.do_exchange(request)) - .await - .map_err(|_| sql_error(&request_id, "SQL query submission timed out"))? - .map_err(|err| sql_error(&request_id, err))?; - Ok(Self { + Ok(Submission::Stream(ResultEndpointStream { + stream, + request_id, + read_timeout, + })) + } + + pub(super) fn new( + id: Uuid, + client: Arc, + response: ResultEndpointStream, + ) -> Self { + Self { id, client, - response: StdMutex::new(Some(ResultEndpointStream { - stream, - request_id, - read_timeout, - })), + response: StdMutex::new(Some(response)), lifecycle: StdMutex::new(Lifecycle::Running), cancelled: Notify::new(), terminal_at: OnceLock::new(), last_accessed: StdMutex::new(Instant::now()), - }) + } } pub(super) fn touch(&self) { @@ -242,7 +317,7 @@ impl ExchangeQuery { } /// Take the rows and wait for the first batch, which carries the schema. - async fn prepare(&self) -> Result { + async fn prepare(&self) -> Result { let Some(rows) = self.response.lock().unwrap().take() else { // Taken by an earlier attempt that failed, or dropped by a cancel. return Err(match &*self.lifecycle.lock().unwrap() { @@ -275,7 +350,7 @@ impl ExchangeQuery { if first.is_none() { self.settle(Lifecycle::Finished); } - Ok(PreparedExchange { + Ok(PreparedStream { schema, rest: first.is_some().then_some(rows), first, @@ -284,7 +359,7 @@ impl ExchangeQuery { async fn run( self: Arc, - mut prepared: PreparedExchange, + mut prepared: PreparedStream, sender: mpsc::Sender>, ) -> Result<()> { if let Some(batch) = prepared.first.take() @@ -343,13 +418,13 @@ impl ExchangeQuery { } } -pub(super) struct ExchangeQueryHandle { - query: Arc, +pub(super) struct StreamingQueryHandle { + query: Arc, result_started: AtomicBool, } -impl ExchangeQueryHandle { - pub(super) fn new(query: Arc) -> Self { +impl StreamingQueryHandle { + pub(super) fn new(query: Arc) -> Self { Self { query, result_started: AtomicBool::new(false), @@ -357,7 +432,7 @@ impl ExchangeQueryHandle { } } -impl Drop for ExchangeQueryHandle { +impl Drop for StreamingQueryHandle { // Rows nobody took are rows nobody will read. Keeping the call open would // hold the statement running on the server, stalled on flow control, // until the registry forgot it. @@ -369,7 +444,7 @@ impl Drop for ExchangeQueryHandle { } #[async_trait::async_trait] -impl QueryHandle for ExchangeQueryHandle { +impl QueryHandle for StreamingQueryHandle { fn id(&self) -> Uuid { self.query.touch(); self.query.id diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 79ca339db..1368dd753 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -19,8 +19,10 @@ use futures::{StreamExt, TryStreamExt}; use tonic::metadata::MetadataMap; use tonic::{Request, Response, Status, Streaming}; +use super::stream::{PROTOCOL_HEADER, PROTOCOL_VERSION, QueryResponse}; use super::*; use crate::database::Database; +use crate::materialized_view::CreateMaterializedViewRequest; use crate::remote::RemoteCatalogOptions; use crate::remote::client::HeaderProvider; use crate::remote::client::test_utils::client_with_handler; @@ -70,7 +72,11 @@ type ReceivedExchanges = Arc)> #[derive(Clone)] struct TestSqlService { + adaptive: bool, + defer: bool, + invalid_decision: bool, query_count: Arc, + statements: Arc>>, do_get_count: Arc, cancel_count: Arc, cancel_denied_count: Arc, @@ -120,6 +126,51 @@ impl TestSqlService { .unwrap_or_default(), }); } + fn adaptive_response( + &self, + query: &str, + rows: BoxStream<'static, std::result::Result>, + ) -> Response>> { + if !self.adaptive { + return Response::new(rows); + } + let poll_info = + self.defer.then(|| PollInfo { + info: Some(FlightInfo::new().with_endpoint( + FlightEndpoint::new().with_ticket(Ticket::new(query.to_string())), + )), + flight_descriptor: Some(FlightDescriptor::new_cmd(format!("poll:1:{query}"))), + progress: Some(0.25), + expiration_time: None, + }); + let first = FlightData { + app_metadata: if self.invalid_decision { + b"invalid protobuf".to_vec().into() + } else { + Any { + type_url: "type.googleapis.com/lancedb.sql.QueryResponse".to_string(), + value: QueryResponse { poll_info }.encode_to_vec().into(), + } + .encode_to_vec() + .into() + }, + ..Default::default() + }; + let rest = if self.defer { + futures::stream::empty().boxed() + } else { + rows + }; + let mut response = Response::new( + futures::stream::once(async { Ok(first) }) + .chain(rest) + .boxed(), + ); + response + .metadata_mut() + .insert(PROTOCOL_HEADER, PROTOCOL_VERSION.parse().unwrap()); + response + } } impl Default for TestSqlService { @@ -154,7 +205,11 @@ impl Default for TestSqlService { let dictionary_result = RecordBatch::try_new(dictionary_schema, vec![Arc::new(dictionary)]).unwrap(); Self { + adaptive: false, + defer: false, + invalid_decision: false, query_count: Arc::new(AtomicUsize::new(0)), + statements: Arc::new(std::sync::Mutex::new(Vec::new())), do_get_count: Arc::new(AtomicUsize::new(0)), cancel_count: Arc::new(AtomicUsize::new(0)), cancel_denied_count: Arc::new(AtomicUsize::new(0)), @@ -217,6 +272,7 @@ impl FlightService for TestSqlService { .and_then(|any| any.unpack::().ok().flatten()); let (query, stage) = if let Some(command) = command { self.query_count.fetch_add(1, Ordering::SeqCst); + self.statements.lock().unwrap().push(command.query.clone()); (command.query, 0_u8) } else { let continuation = std::str::from_utf8(request.get_ref().cmd.as_ref()) @@ -237,6 +293,9 @@ impl FlightService for TestSqlService { .ok_or_else(|| Status::invalid_argument("invalid continuation"))?; (query.to_string(), stage) }; + if query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } if (query == "SELECT slow" || query == "SELECT cancelling") && stage > 0 { tokio::time::sleep(Duration::from_millis(250)).await; } @@ -308,8 +367,28 @@ impl FlightService for TestSqlService { request: Request, ) -> std::result::Result::DoGetStream>, Status> { self.do_get_count.fetch_add(1, Ordering::SeqCst); - let ticket = request.get_ref().ticket.as_ref(); - let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\""; + let initial = Any::decode(request.get_ref().ticket.clone()) + .ok() + .and_then(|envelope| Any::decode(envelope.value).ok()) + .and_then(|command| command.unpack::().ok().flatten()); + if let Some(command) = &initial { + self.statements.lock().unwrap().push(command.query.clone()); + if command.query == "SELECT rejected" { + return Err(Status::invalid_argument("invalid statement")); + } + self.capture_headers(request.metadata()); + assert_eq!( + request.metadata().get(PROTOCOL_HEADER).unwrap(), + PROTOCOL_VERSION + ); + } + let query = initial + .as_ref() + .map(|command| command.query.as_bytes()) + .unwrap_or(request.get_ref().ticket.as_ref()); + let ticket = query; + let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"" + || ticket == b"SELECT job cancel denied"; let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; let large = ticket == b"SELECT large message"; @@ -337,7 +416,11 @@ impl FlightService for TestSqlService { encoder = encoder.with_max_flight_data_size(8 * 1024 * 1024); } let stream = encoder.build(input).map_err(Status::from); - Ok(Response::new(Box::pin(stream))) + if let Some(initial) = initial { + Ok(self.adaptive_response(&initial.query, stream.boxed())) + } else { + Ok(Response::new(stream.boxed())) + } } async fn do_put( @@ -387,7 +470,9 @@ impl FlightService for TestSqlService { { return Err(Status::not_found("query cancellation completed")); } - let status = if query == "SELECT cancel unspecified" { + let status = if query == "SELECT job cancel denied" { + CancelStatus::NotCancellable + } else if query == "SELECT cancel unspecified" { CancelStatus::Unspecified } else if query == "SELECT cancelling" { if self @@ -486,7 +571,7 @@ impl FlightService for TestSqlService { let _dropped = &dropped; message }); - Ok(Response::new(Box::pin(stream))) + Ok(self.adaptive_response(&query, stream.boxed())) } } @@ -809,6 +894,132 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { server.await.unwrap().unwrap(); } +#[tokio::test] +async fn remote_create_uses_sql_and_preserves_options_and_namespace() { + for defer in [false, true] { + let service = TestSqlService { + adaptive: true, + defer, + ..Default::default() + }; + let statements = service.statements.clone(); + let headers = service.headers.clone(); + let count = service.query_count.clone(); + let (sql_client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + panic!("CREATE must use SQL, not REST: {}", request.url().path()); + }) + .with_sql_client(sql_client.clone()); + for (enabled, interval, expected) in [ + (None, None, ""), + (Some(true), None, " WITH (auto_refresh_enabled = true)"), + ( + Some(false), + Some(0.5), + " WITH (auto_refresh_enabled = false, auto_refresh_interval_seconds = 0.5)", + ), + ] { + let job = db + .create_materialized_view_async(CreateMaterializedViewRequest { + name: "adults".into(), + namespace_path: vec!["team".into(), "reports".into()], + query: "SELECT * FROM people".into(), + with_no_data: false, + auto_refresh_enabled: enabled, + auto_refresh_interval_seconds: interval, + }) + .await + .unwrap(); + assert_eq!(job.job_type().as_deref(), Some("query")); + let query_id = + Uuid::parse_str(job.id().expect("SQL-backed job must expose its ID")).unwrap(); + job.wait().await.unwrap(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + assert!(job.events(Default::default()).await.is_err()); + assert_eq!(job.job_type().as_deref(), Some("query")); + assert_eq!( + sql_client.describe(query_id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + statements.lock().unwrap().last().unwrap(), + &format!("CREATE MATERIALIZED VIEW \"adults\"{expected} AS SELECT * FROM people") + ); + } + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(statements.lock().unwrap().len(), 3); + assert!( + headers + .lock() + .unwrap() + .iter() + .all(|header| header.namespace_path == "team$reports") + ); + shutdown.send(()).unwrap(); + } +} + +#[tokio::test] +async fn query_backed_jobs_propagate_submission_and_execution_failures() { + let (client, shutdown) = start_sql_service( + TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }, + ClientConfig::default(), + ) + .await; + let rejected = client + .submit_as_job("SELECT rejected".into(), vec![], || async { Ok(()) }) + .await; + assert!(rejected.is_err()); + let finished = Arc::new(AtomicBool::new(false)); + let callback_finished = finished.clone(); + let job = client + .submit_as_job("SELECT fail".into(), vec![], move || async move { + callback_finished.store(true, Ordering::SeqCst); + Ok(()) + }) + .await + .unwrap(); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert!(matches!(job.wait().await, Err(Error::JobFailed { .. }))); + assert_eq!(job.status().await.unwrap(), "failed"); + assert!(!finished.load(Ordering::SeqCst)); + shutdown.send(()).unwrap(); +} + +#[tokio::test] +async fn rejected_job_cancellation_does_not_abort_sql_completion() { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let do_get_count = service.do_get_count.clone(); + let release = service.gate_refresh_release.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let job = client + .submit_as_job("SELECT job cancel denied".into(), vec![], || async { + Ok(()) + }) + .await + .unwrap(); + wait_for(&do_get_count, 1, "job did not begin reading its SQL result").await; + assert!(matches!( + job.cancel().await, + Err(Error::NotSupported { .. }) + )); + assert_eq!(job.status().await.unwrap(), "running"); + release.notify_one(); + job.wait().await.unwrap(); + assert_eq!(job.status().await.unwrap(), "finished"); + job.cancel().await.unwrap(); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn remote_refresh_twice_keeps_latest_selector() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -898,7 +1109,11 @@ async fn cancelling_remote_refresh_cancels_sql_query() { let (socket, _) = listener.accept().await?; Ok::<_, std::io::Error>(Some((socket, listener))) }); - let service = TestSqlService::default(); + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; let cancel_count = service.cancel_count.clone(); let do_get_count = service.do_get_count.clone(); let gate_refresh_release = service.gate_refresh_release.clone(); @@ -1609,3 +1824,138 @@ fn validates_metadata_with_header_map() { assert!(insert_header(&mut headers, "bad header", "value").is_err()); assert!(insert_header(&mut headers, "valid-header", "bad\nvalue").is_err()); } + +#[tokio::test] +async fn adaptive_streams_start_on_one_rpc_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = parameters.clone().unwrap_or_else(|| service.result.clone()); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT echo") + }) + .await + .unwrap(); + assert_eq!(query.describe().await.unwrap().status, QueryStatus::Running); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + assert_eq!(polls.load(Ordering::SeqCst), 0); + assert_eq!(gets.load(Ordering::SeqCst), usize::from(!with_parameters)); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_share_polling_and_survive_handle_drop() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let gets = service.do_get_count.clone(); + let exchanges = service.exchanges.clone(); + let expected = service.result.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let with_parameters = parameters.is_some(); + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT adaptive") + }) + .await + .unwrap(); + let id = query.id(); + assert_eq!(collect_result(&query).await.unwrap(), vec![expected]); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Finished + ); + assert!(query.reader().await.is_err()); + drop(query); + assert_eq!( + client.describe(id).await.unwrap().status, + QueryStatus::Finished + ); + assert_eq!( + polls.load(Ordering::SeqCst), + 0, + "the original SQL must not be resubmitted through polling" + ); + assert_eq!( + gets.load(Ordering::SeqCst), + 1 + usize::from(!with_parameters) + ); + assert_eq!( + exchanges.lock().unwrap().len(), + usize::from(with_parameters) + ); + } +} + +#[tokio::test] +async fn adaptive_handoffs_use_server_cancellation_for_both_inputs() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + defer: true, + ..Default::default() + }; + let cancels = service.cancel_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + let query = client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT cancel") + }) + .await + .unwrap(); + query.cancel().await.unwrap(); + assert_eq!(cancels.load(Ordering::SeqCst), 1); + assert_eq!( + query.describe().await.unwrap().status, + QueryStatus::Cancelled + ); + assert!(query.reader().await.is_err()); + } +} + +#[tokio::test] +async fn malformed_adaptive_decision_never_resubmits_the_statement() { + for parameters in [None, Some(parameter_row())] { + let service = TestSqlService { + adaptive: true, + invalid_decision: true, + ..Default::default() + }; + let polls = service.query_count.clone(); + let (client, _shutdown) = start_sql_service(service, ClientConfig::default()).await; + assert!( + client + .execute(ExecuteQueryRequest { + parameters, + ..ExecuteQueryRequest::new("SELECT invalid") + }) + .await + .is_err() + ); + assert_eq!(polls.load(Ordering::SeqCst), 0); + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index b722dc4a1..181ce61f9 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2247,8 +2247,8 @@ impl BaseTable for RemoteTable { let freshness = self.freshness.clone(); let freshness_request = self.snapshot_freshness_headers(); let job_sql_client = sql_client.clone(); - Ok( - sql_client.submit_as_job(statement, namespace.clone(), move || async move { + sql_client + .submit_as_job(statement, namespace.clone(), move || async move { let table = Self::new_with_sql_client( client, name, @@ -2270,8 +2270,8 @@ impl BaseTable for RemoteTable { source_version: 0, version, }) - }), - ) + }) + .await } async fn query_snapshot(&self) -> Result> { let description = self.describe().await?; @@ -2927,6 +2927,10 @@ impl BaseTable for RemoteTable { self.fetch_blob_files_impl(column, row_ids).await } + async fn open_blob_file(&self, descriptor: &crate::blob::BlobDescriptor) -> Result { + self.open_blob_file_impl(descriptor).await + } + async fn explain_plan(&self, query: &AnyQuery, verbose: bool) -> Result { if let AnyQuery::Query(request) = query && let Some(offsets) = &request.take_offsets diff --git a/rust/lancedb/src/remote/table/blobs.rs b/rust/lancedb/src/remote/table/blobs.rs index bcaf47765..e3cdfc6a6 100644 --- a/rust/lancedb/src/remote/table/blobs.rs +++ b/rust/lancedb/src/remote/table/blobs.rs @@ -21,7 +21,7 @@ use reqwest::{Response, StatusCode, header}; use tokio::sync::{Mutex, OwnedSemaphorePermit}; use crate::Error; -use crate::blob::{BlobFile, BlobRangeRequest}; +use crate::blob::{BlobDescriptor, BlobDescriptorSource, BlobFile, BlobRangeRequest}; use crate::error::Result; use crate::query::{QueryFilter, QueryRequest, Select}; use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; @@ -49,6 +49,7 @@ fn is_fetch_blobs_byte_limit_error(error: &Error) -> bool { enum RangeRequestMode { SizeProbe, DataRead, + MetadataRead, } #[async_trait::async_trait] @@ -86,6 +87,9 @@ struct TableBlobRangeRequester { path: String, version: Option, branch: Option, + source_uri: Option, + source_commit_micros: Option, + row_is_address: bool, freshness: Arc>, parent_freshness: Arc>, parent_freshness_request: FreshnessHeaders, @@ -109,15 +113,27 @@ impl BlobRangeRequester for TableBlobRangeRequester { .expect("blob request semaphore is never closed"); let freshness_request = freshness_headers_snapshot(&self.freshness, self.read_consistency_interval); - let mut request = freshness_request - .apply(self.client.get(&self.path)) - .header(header::RANGE, range_header); + let mut request = if matches!(mode, RangeRequestMode::MetadataRead) { + freshness_request.apply(self.client.head(&self.path)) + } else { + freshness_request + .apply(self.client.get(&self.path)) + .header(header::RANGE, range_header) + }; if let Some(version) = version.or(self.version) { request = request.query(&[("version", version)]); } if let Some(branch) = &self.branch { request = request.query(&[("branch", branch)]); } + if let Some(source_uri) = &self.source_uri { + request = request + .query(&[("source_uri", source_uri)]) + .query(&[("row_is_address", self.row_is_address)]); + } + if let Some(timestamp) = self.source_commit_micros { + request = request.query(&[("source_commit_micros", timestamp)]); + } let (request_id, response) = self.client.send_with_retry(request, None, true).await?; // Preserve 416 size-probe responses so the caller can detect empty blobs. if response.status() == StatusCode::RANGE_NOT_SATISFIABLE @@ -131,8 +147,10 @@ impl BlobRangeRequester for TableBlobRangeRequester { } let response = self.client.check_response(&request_id, response).await?; freshness_request.observe_headers(&self.freshness, response.headers()); - self.parent_freshness_request - .observe_headers(&self.parent_freshness, response.headers()); + if self.source_uri.is_none() { + self.parent_freshness_request + .observe_headers(&self.parent_freshness, response.headers()); + } Ok(BlobRangeResponse { request_id, response, @@ -630,6 +648,51 @@ impl RemoteTable { Ok(files) } + pub(super) async fn open_blob_file_impl( + &self, + descriptor: &BlobDescriptor, + ) -> Result { + descriptor.validate()?; + self.ensure_blob_files_supported()?; + let source = &descriptor.source; + let snapshot = ReadSnapshot { + version: Some(source.version), + freshness_state: FreshnessState { + pinned: true, + ..Default::default() + }, + ..Default::default() + }; + let requester = self.blob_range_requester_with_source( + &source.column, + source.row_reference, + &snapshot, + Some(source), + ); + let size = if descriptor.kind == 3 && descriptor.size == 0 { + let metadata = requester + .request_range("", RangeRequestMode::MetadataRead, Some(source.version)) + .await?; + if metadata.response.status() != StatusCode::OK { + return Err(Error::InvalidInput { + message: "non-null blob descriptor did not resolve to a stored value".into(), + }); + } + metadata + .response + .headers() + .get(header::CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .ok_or_else(|| Error::InvalidInput { + message: "blob metadata response has no valid Content-Length".into(), + })? + } else { + descriptor.size + }; + Ok(RemoteBlobFile::new(requester, size, Some(source.version)).into()) + } + /// Probe each row once for its size. /// /// Serves callers that must not issue a query, such as a single blob that @@ -793,6 +856,16 @@ impl RemoteTable { column: &str, row_id: u64, read_snapshot: &ReadSnapshot, + ) -> Arc { + self.blob_range_requester_with_source(column, row_id, read_snapshot, None) + } + + fn blob_range_requester_with_source( + &self, + column: &str, + row_id: u64, + read_snapshot: &ReadSnapshot, + source: Option<&BlobDescriptorSource>, ) -> Arc { let path = format!( "/v1/table/{}/blob/{}/{row_id}/bytes", @@ -803,7 +876,10 @@ impl RemoteTable { client: self.client.clone(), path, version: read_snapshot.version, - branch: self.branch.clone(), + branch: source.map_or_else(|| self.branch.clone(), |source| source.branch.clone()), + source_uri: source.map(|source| source.uri.clone()), + source_commit_micros: source.map(|source| source.committed_at_micros), + row_is_address: source.is_some_and(|source| source.row_is_address), freshness: Arc::new(std::sync::Mutex::new(read_snapshot.freshness_state)), parent_freshness: self.freshness.clone(), parent_freshness_request: read_snapshot.freshness, @@ -1024,6 +1100,99 @@ mod tests { .unwrap() } + #[tokio::test] + async fn open_descriptor_files_are_lazy_and_pin_ranges_for_every_storage_kind() { + for kind in 0..=5 { + let requests = Arc::new(StdMutex::new(Vec::new())); + let captured = requests.clone(); + let table = RemoteTable::new_mock( + "videos".to_string(), + move |request| { + for name in [ + "x-lancedb-min-version", + "x-lancedb-min-read-version", + "x-lancedb-min-timestamp", + ] { + assert!(request.headers().get(name).is_none()); + } + assert_eq!( + request.url().path(), + "/v1/table/videos/blob/video/12884901889/bytes" + ); + let parameters = request + .url() + .query_pairs() + .into_owned() + .collect::>(); + assert_eq!(parameters["version"], "7"); + assert_eq!(parameters["source_commit_micros"], "1234"); + assert_eq!(parameters["branch"], "dev"); + assert_eq!(parameters["row_is_address"], "true"); + assert_eq!( + parameters["source_uri"], + "s3://bucket/videos.lance/tree/dev" + ); + captured.lock().unwrap().push(request.method().clone()); + if request.method() == reqwest::Method::HEAD { + assert!(request.headers().get(header::RANGE).is_none()); + http::Response::builder() + .status(200) + .header(header::CONTENT_LENGTH, PAYLOAD.len().to_string()) + .body(Vec::new()) + .unwrap() + } else { + range_response(&request, PAYLOAD) + } + }, + Some(Version::new(0, 5, 0)), + ); + { + let mut freshness = table.freshness.lock().unwrap(); + freshness.min_version = Some(99); + freshness.min_read_version = Some(2); + } + let descriptor = BlobDescriptor { + kind, + position: 0, + size: if kind == 3 { 0 } else { PAYLOAD.len() as u64 }, + blob_id: 0, + blob_uri: String::new(), + source: BlobDescriptorSource { + uri: "s3://bucket/videos.lance/tree/dev".into(), + version: 7, + committed_at_micros: 1234, + branch: Some("dev".into()), + column: "video".into(), + row_reference: 12884901889, + row_is_address: true, + }, + }; + let file = table.open_blob_file(&descriptor).await.unwrap(); + let expected = if kind == 3 { + vec![reqwest::Method::HEAD] + } else { + vec![] + }; + assert_eq!(*requests.lock().unwrap(), expected); + assert_eq!(file.size(), PAYLOAD.len() as u64); + file.seek(4).await.unwrap(); + assert_eq!(file.read_up_to(5).await.unwrap().as_ref(), &PAYLOAD[4..9]); + assert_eq!(file.tell().await.unwrap(), 9); + assert_eq!(table.freshness.lock().unwrap().min_read_version, Some(2)); + assert_eq!(file.read_range(0..3).await.unwrap().as_ref(), &PAYLOAD[..3]); + assert_eq!(file.tell().await.unwrap(), 9); + assert_eq!( + requests + .lock() + .unwrap() + .iter() + .filter(|method| **method == reqwest::Method::GET) + .count(), + 2 + ); + } + } + fn mock_remote_blob_table( requests: Arc>>, ) -> RemoteTable { diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql.rs index c2ef13ba5..fee2d5d79 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql.rs @@ -18,9 +18,9 @@ pub enum QueryStatus { Running, /// The server has made the complete result available. /// - /// A query submitted with parameters streams its rows on the call that - /// runs it, so it reports `Finished` only once its reader has received - /// the last batch. Read it rather than waiting for this status. + /// Directly streamed queries report `Finished` when their stream reaches + /// its end. Deferred queries report the server's execution state. Open the + /// reader directly instead of waiting for this status before reading. Finished, /// The server accepted cancellation but has not confirmed it yet. Cancelling, @@ -64,9 +64,10 @@ pub(crate) trait QueryHandle: Send + Sync { /// A handle to a submitted SQL query. /// /// The handle can be inspected, opened as an Arrow reader, or cancelled. -/// Dropping it does not cancel the server-side query, unless the query was -/// submitted with parameters: those run on the call that returns their rows, -/// so dropping the handle before reading, or dropping the reader, stops them. +/// The server may stream results on the initial call or defer execution and +/// return result tickets. Dropping an unread handle or its reader stops a +/// directly streamed query; deferred execution continues until explicitly +/// cancelled. This applies equally with and without parameters. /// Identifier lookup is scoped to the connection that submitted the query and /// is not a durable resume mechanism. pub struct Query { @@ -95,10 +96,9 @@ impl Query { /// Get a point-in-time description of the query. /// - /// For a query submitted with parameters the status follows its result - /// stream: it stays [`QueryStatus::Running`] until the reader has received - /// every row, so open [`Self::reader`] directly instead of polling for - /// [`QueryStatus::Finished`] first. + /// Directly streamed queries stay [`QueryStatus::Running`] until the result + /// stream ends. Deferred queries expose the server's execution state. Open + /// [`Self::reader`] directly instead of waiting for [`QueryStatus::Finished`]. pub async fn describe(&self) -> Result { self.handle.describe().await } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 3123eb776..3fcf4bdb7 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3,7 +3,7 @@ //! LanceDB Table APIs -use crate::blob::BlobFile; +use crate::blob::{BlobDescriptor, BlobFile}; use arrow_array::{LargeBinaryArray, RecordBatch, RecordBatchReader}; use arrow_schema::{Schema, SchemaRef}; use async_trait::async_trait; @@ -796,6 +796,12 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "fetch_blob_files is not supported on this table type".into(), }) } + /// Open a descriptor at its pinned source version. See [`Table::open_blob_file`]. + async fn open_blob_file(&self, _descriptor: &BlobDescriptor) -> Result { + Err(Error::NotSupported { + message: "open_blob_file is not supported on this table type".into(), + }) + } /// Gets the table tag manager. async fn tags(&self) -> Result>; /// Optimize the dataset. @@ -1377,6 +1383,26 @@ impl Table { self.inner.fetch_blob_files(column.as_ref(), row_ids).await } + /// Open a seekable file from a SQL `blob_descriptor` result. + /// + /// The descriptor carries its table version, branch, column and row identity. + /// Opening does not materialize the payload or change this table's checkout. + /// Remote reads use the table's authenticated HTTP byte-range endpoint. + /// + /// ``` + /// use lancedb::{Table, blob::BlobDescriptor}; + /// + /// # async fn read_header(videos: &Table, descriptor: &BlobDescriptor) -> lancedb::Result<()> { + /// let video = videos.open_blob_file(descriptor).await?; + /// let header = video.read_range(0..video.size().min(4096)).await?; + /// video.close().await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn open_blob_file(&self, descriptor: &BlobDescriptor) -> Result { + self.inner.open_blob_file(descriptor).await + } + /// Insert new records into this Table /// /// # Arguments @@ -3764,6 +3790,11 @@ impl BaseTable for NativeTable { crate::blob::take_blob_files_aligned(&dataset, column, row_ids).await } + async fn open_blob_file(&self, descriptor: &BlobDescriptor) -> Result { + let dataset = self.dataset.get().await?; + crate::blob::open_native_blob_file(&dataset, descriptor).await + } + /// Delete rows from the table async fn delete(&self, predicate: Predicate<'_>) -> Result { let result = delete::execute_delete(self, predicate).await?; diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index 12404ec73..445f29a18 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -45,10 +45,12 @@ pub struct AddResult { #[derive(Debug, Default, Clone, Copy)] pub enum NaNVectorBehavior { - /// Reject any vectors containing NaN values (the default) + /// Reject any vectors containing NaN values or null elements (the default). + /// Whole null vectors are allowed. #[default] Error, - /// Allow NaN values to be added, but they will not be indexed for search + /// Allow NaN values and null elements to be added, but they will not be indexed + /// for search. Keep, } @@ -104,11 +106,12 @@ impl AddDataBuilder { self } - /// Configure how to handle NaN values in vector columns. + /// Configure how to handle NaN values and null elements in vector columns. /// - /// By default, any vectors containing NaN values will be rejected with an - /// error, since NaNs cannot be indexed for search. Setting this to `Keep` - /// will allow NaN values to be added to the table, but they will not be + /// By default, any vectors containing NaN values or null elements will be + /// rejected with an error, since these vectors cannot be indexed for search. + /// Whole null vectors are allowed. Setting this to `Keep` will allow NaN values + /// and null elements to be added to the table, but these vectors will not be /// indexed and will not be searchable. pub fn on_nan_vectors(mut self, behavior: NaNVectorBehavior) -> Self { self.on_nan_vectors = behavior; @@ -953,8 +956,11 @@ mod tests { ); } + #[rstest::rstest] + #[case::nan(Some(f32::NAN))] + #[case::null(None)] #[tokio::test] - async fn test_add_rejects_nan_vectors() { + async fn test_add_rejects_nan_or_null_vectors(#[case] bad_value: Option) { let schema = Arc::new(Schema::new(vec![Field::new( "embedding", DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4), @@ -974,7 +980,12 @@ mod tests { FixedSizeListArray::try_new( Arc::new(Field::new("item", DataType::Float32, true)), 4, - Arc::new(Float32Array::from(vec![0.1, 0.2, f32::NAN, 0.4])), + Arc::new(Float32Array::from(vec![ + Some(0.1), + Some(0.2), + bad_value, + Some(0.4), + ])), None, ) .unwrap(), @@ -984,8 +995,8 @@ mod tests { let res = table.add(batch.clone()).execute().await; let err = res.unwrap_err(); assert!( - err.to_string().contains("NaN"), - "Expected error mentioning NaN values, but got: {err:?}" + err.to_string().contains("NaN or null"), + "Expected error mentioning NaN or null values, but got: {err:?}" ); table diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 3ac759bdd..1795b778e 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -355,7 +355,10 @@ impl NativeTable { index.sample_rate, index.max_iterations, ); - let rq_params = RQBuildParams::new(index.num_bits.unwrap_or(1) as u8); + let rq_params = index + .num_bits + .map(|num_bits| RQBuildParams::new(num_bits as u8)) + .unwrap_or_default(); let lance_idx_params = VectorIndexParams::with_ivf_rq_params( index.distance_type.into(), ivf_params, diff --git a/rust/lancedb/src/table/datafusion.rs b/rust/lancedb/src/table/datafusion.rs index 39f3cd4c6..e6a76c3f5 100644 --- a/rust/lancedb/src/table/datafusion.rs +++ b/rust/lancedb/src/table/datafusion.rs @@ -25,10 +25,11 @@ use datafusion_physical_plan::{ }; use futures::{TryFutureExt, TryStreamExt}; use lance::dataset::{WriteMode, WriteParams}; +use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD}; use super::{AnyQuery, BaseTable}; use crate::{ - Result, + Error, Result, query::{QueryExecutionOptions, QueryFilter, QueryRequest, Select}, }; use arrow_schema::{DataType, Field}; @@ -186,6 +187,46 @@ impl BaseTableAdapter { }) } + /// Expose the requested Lance system columns to DataFusion. + /// + /// Supports [`ROW_ID`] (`_rowid`) and [`ROW_ADDR`] (`_rowaddr`), both nullable + /// `UInt64` fields. Columns are appended in the requested order, skipping any + /// already exposed by this adapter. Other names return [`Error::InvalidInput`]. + /// + /// System columns are hidden by default so existing `SELECT *` queries keep + /// their data-only schema. After opting in, `SELECT *` includes these columns. + /// + /// ``` + /// use lance_core::{ROW_ADDR, ROW_ID}; + /// use lancedb::table::datafusion::BaseTableAdapter; + /// + /// # async fn adapter(table: &lancedb::Table) -> lancedb::Result { + /// BaseTableAdapter::try_new(table.base_table().clone()) + /// .await? + /// .with_system_columns(&[ROW_ID, ROW_ADDR]) + /// # } + /// ``` + pub fn with_system_columns(mut self, columns: &[impl AsRef]) -> Result { + let mut fields = self.schema.fields().to_vec(); + for column in columns { + let column = column.as_ref(); + let field = match column { + ROW_ID => &*ROW_ID_FIELD, + ROW_ADDR => &*ROW_ADDR_FIELD, + _ => { + return Err(Error::InvalidInput { + message: format!("Unsupported system column: {column}"), + }); + } + }; + if !fields.iter().any(|field| field.name() == column) { + fields.push(Arc::new(field.clone())); + } + } + self.schema = Arc::new(ArrowSchema::new(fields)); + Ok(self) + } + /// Set the RTE read routing of every scan; see [`QueryRequest::use_rte`]. pub fn with_use_rte(mut self, use_rte: Option) -> Self { self.use_rte = use_rte; @@ -242,6 +283,15 @@ impl TableProvider for BaseTableAdapter { .map(|i| self.schema.field(*i).name().clone()) .collect(); query.select = Select::Columns(field_names); + } else { + // Select::All excludes system columns, which may be part of our schema. + query.select = Select::Columns( + self.schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(), + ); } if !filters.is_empty() { let first = filters.first().unwrap().clone(); @@ -319,6 +369,7 @@ pub mod tests { BinaryArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, UInt32Array, }; use arrow_schema::{DataType, Field, Schema}; + use arrow_select::concat::concat_batches; use datafusion::{ datasource::provider_as_source, prelude::{SessionConfig, SessionContext}, @@ -328,11 +379,13 @@ pub mod tests { use datafusion_execution::SendableRecordBatchStream; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col, lit}; use futures::{StreamExt, TryStreamExt}; + use lance_core::{ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD}; use tempfile::tempdir; use crate::{ connect, index::{Index, scalar::BTreeIndexBuilder}, + query::{ExecutableQuery, Query, QueryBase, Select}, table::datafusion::BaseTableAdapter, }; @@ -509,6 +562,173 @@ pub mod tests { } } + #[tokio::test] + async fn test_system_columns_opt_in() { + let fixture = TestFixture::new().await; + let ctx = SessionContext::new(); + ctx.register_table("foo", fixture.adapter.clone()).unwrap(); + + let batches = ctx + .sql("SELECT * FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0].num_columns(), 2); + assert_eq!(batches[0].schema(), fixture.adapter.schema()); + for column in [ROW_ID, ROW_ADDR] { + assert!(ctx.sql(&format!("SELECT {column} FROM foo")).await.is_err()); + } + + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ADDR, ROW_ADDR]) + .unwrap(); + assert_eq!(adapter.schema().fields().len(), 3); + assert_eq!(adapter.schema().field(2), &*ROW_ADDR_FIELD); + assert!(adapter.schema().index_of(ROW_ID).is_err()); + + let adapter = adapter + .with_system_columns(&[ROW_ID.to_string(), ROW_ADDR.to_string()]) + .unwrap(); + assert_eq!(adapter.schema().fields().len(), 4); + assert_eq!(adapter.schema().field(3), &*ROW_ID_FIELD); + assert!(adapter.schema().metadata().is_empty()); + + let err = adapter.with_system_columns(&["_unknown"]).unwrap_err(); + assert!(matches!(err, crate::Error::InvalidInput { .. })); + + // Default adapters must remain compatible with INSERT ... SELECT *. + ctx.sql("INSERT INTO foo SELECT * FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + let batches = ctx + .sql("SELECT COUNT(*) FROM foo") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + batches[0] + .column(0) + .as_primitive::() + .value(0), + 20 + ); + } + + #[tokio::test] + async fn test_system_columns_sql() { + let fixture = TestFixture::new().await; + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ID, ROW_ADDR]) + .unwrap(); + let ctx = SessionContext::new(); + ctx.register_table("foo", Arc::new(adapter)).unwrap(); + + // Compare SQL results with the native query's row IDs and row addresses. + let native = Query::new(fixture.adapter.table.clone()) + .select(Select::columns(&["i", "indexed", ROW_ADDR])) + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let schema = Arc::new( + native[0] + .schema() + .as_ref() + .clone() + .with_metadata(HashMap::new()), + ); + let expected = concat_batches(&schema, &native).unwrap(); + + for columns in [ + vec!["i", ROW_ID], + vec![ROW_ID, "i"], + vec!["i", ROW_ADDR], + vec![ROW_ADDR, "i"], + vec![ROW_ADDR, ROW_ID], + vec![ROW_ID], + vec![ROW_ADDR], + vec!["i", "indexed", ROW_ID, ROW_ADDR], + ] { + let sql = format!("SELECT {} FROM foo ORDER BY i", columns.join(", ")); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let actual = concat_batches(&batches[0].schema(), &batches).unwrap(); + let indices = columns + .iter() + .map(|name| schema.index_of(name).unwrap()) + .collect::>(); + assert_eq!(actual, expected.project(&indices).unwrap(), "{sql}"); + } + + let batches = ctx + .sql("SELECT * FROM foo ORDER BY i") + .await + .unwrap() + .collect() + .await + .unwrap(); + let indices = ["i", "indexed", ROW_ID, ROW_ADDR].map(|name| schema.index_of(name).unwrap()); + assert_eq!(batches[0], expected.project(&indices).unwrap()); + + // Filtering on the ID must work even when the ID is not in the output. + let row_id = expected + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .value(3); + let batches = ctx + .sql(&format!("SELECT i FROM foo WHERE _rowid = {row_id}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches[0], expected.project(&[0]).unwrap().slice(3, 1)); + } + + #[tokio::test] + async fn test_system_columns_scan_without_projection() { + let fixture = TestFixture::new().await; + let adapter = BaseTableAdapter::try_new(fixture.adapter.table.clone()) + .await + .unwrap() + .with_system_columns(&[ROW_ADDR, ROW_ID]) + .unwrap(); + let ctx = SessionContext::new(); + let plan = adapter.scan(&ctx.state(), None, &[], None).await.unwrap(); + assert_eq!(plan.schema(), adapter.schema()); + let batches = datafusion_physical_plan::collect(plan, ctx.task_ctx()) + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 10); + for batch in batches { + assert_eq!(batch.schema(), adapter.schema()); + assert_eq!( + batch + .column(2) + .as_primitive::() + .values(), + batch + .column(3) + .as_primitive::() + .values() + ); + } + } + #[tokio::test] async fn test_batch_size() { let fixture = TestFixture::new().await; diff --git a/rust/lancedb/src/table/datafusion/reject_nan.rs b/rust/lancedb/src/table/datafusion/reject_nan.rs index 084c9722b..462216929 100644 --- a/rust/lancedb/src/table/datafusion/reject_nan.rs +++ b/rust/lancedb/src/table/datafusion/reject_nan.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! A DataFusion projection that rejects vectors containing NaN values. +//! A DataFusion projection that rejects vectors containing NaN or null elements. use std::sync::{Arc, LazyLock}; @@ -31,10 +31,12 @@ fn is_vector_field(field: &Field) -> bool { } } -/// Wraps the input plan with a projection that checks vector columns for NaN values. +/// Wraps the input plan with a projection that checks vector columns for NaN or +/// null elements. /// /// Non-vector columns pass through unchanged. Vector columns are wrapped with a -/// UDF that returns the column as-is if no NaNs are present, or errors otherwise. +/// UDF that returns the column as-is if no NaN or null elements are present, or +/// errors otherwise. Whole null vectors pass through unchanged. pub fn reject_nan_vectors(input: Arc) -> Result> { let schema = input.schema(); let config = Arc::new(ConfigOptions::default()); @@ -70,7 +72,7 @@ pub fn reject_nan_vectors(input: Arc) -> Result datafusion_common::Result<()> { ) })?; - // Only inspect elements that are both in a valid parent row and non-null - // themselves. Values backing null parent rows or null child elements may - // contain garbage (including NaN) per the Arrow spec. - let has_nan = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| { + // Values backing null parent rows may contain garbage per the Arrow spec. + // A null child in a valid row is an invalid vector component, however: + // pandas can represent NaN components as Arrow nulls. + let has_bad_values = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| { let row = fsl.value(i); + if row.null_count() > 0 { + return true; + } match row.data_type() { DataType::Float16 => row .as_any() @@ -150,10 +155,10 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> { } }); - if has_nan { + if has_bad_values { return Err(datafusion_common::DataFusionError::ArrowError( Box::new(arrow_schema::ArrowError::ComputeError( - "Vector column contains NaN values".to_string(), + "Vector column contains NaN or null values".to_string(), )), None, )); @@ -165,7 +170,7 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> { #[cfg(test)] mod tests { use super::*; - use arrow_array::Float32Array; + use arrow_array::{Float16Array, Float32Array, Float64Array}; #[test] fn test_passes_clean_vectors() { @@ -209,23 +214,43 @@ mod tests { } #[test] - fn test_skips_null_elements_within_valid_row() { - // A valid row with null child elements: the underlying buffer may hold - // NaN but the null bitmap says they're absent — should not reject. - let values = Float32Array::from(vec![ - Some(1.0), - None, // null element — buffer may contain NaN - Some(3.0), - None, // null element - ]); + fn test_rejects_null_elements_within_valid_row() { + let arrays: Vec> = vec![ + Arc::new(Float16Array::from(vec![Some(half::f16::ONE), None])), + Arc::new(Float32Array::from(vec![Some(1.0), None])), + Arc::new(Float64Array::from(vec![Some(1.0), None])), + ]; + for values in arrays { + let fsl = FixedSizeListArray::try_new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + 2, + values, + None, + ) + .unwrap(); + let err = check_no_nans(&fsl).unwrap_err(); + assert!(err.to_string().contains("null")); + } + } + + #[test] + fn test_skips_null_elements_in_null_rows_and_outside_slice() { let fsl = FixedSizeListArray::try_new( Arc::new(Field::new("item", DataType::Float32, true)), 2, - Arc::new(values), - None, // both rows are valid + Arc::new(Float32Array::from(vec![ + None, + Some(1.0), + Some(2.0), + Some(3.0), + None, + None, + ])), + Some(vec![true, true, false].into()), ) .unwrap(); - assert!(check_no_nans(&fsl).is_ok()); + assert!(check_no_nans(&fsl.slice(1, 2)).is_ok()); + assert!(check_no_nans(&fsl.slice(0, 2)).is_err()); } #[test] diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index f625a4aff..4ad840350 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -280,7 +280,13 @@ pub async fn create_plan( let (_, element_type) = lance::index::vector::utils::get_vector_type(schema, &column)?; let is_binary = matches!(element_type, DataType::UInt8); - let top_k = query.base.limit.unwrap_or(DEFAULT_TOP_K) + query.base.offset.unwrap_or(0); + let top_k = query + .base + .limit + .unwrap_or(DEFAULT_TOP_K) + .saturating_add(query.base.offset.unwrap_or(0)) + // Lance allocates candidates from k; no search can return more rows than the table. + .min(ds_ref.count_rows(None).await?.max(1)); if is_binary { let query_vector = arrow::compute::cast(&query_vector, &DataType::UInt8)?; @@ -317,9 +323,18 @@ pub async fn create_plan( // For a batch query, `nearest` already applies k to each query vector. // Adding Scanner's global limit would truncate the combined result to k rows. if !is_batch_query { + let offset = query + .base + .offset + .map(|offset| i64::try_from(offset).unwrap_or(i64::MAX)); scanner.limit( - query.base.limit.map(|limit| limit as i64), - query.base.offset.map(|offset| offset as i64), + // Lance's FTS planner adds limit + offset as i64. + query.base.limit.map(|limit| { + i64::try_from(limit) + .unwrap_or(i64::MAX) + .min(i64::MAX - offset.unwrap_or(0)) + }), + offset, )?; }