mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-25 21:22:29 +00:00
Merge main into gatekeeper/fix-2325-1
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -245,12 +245,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 +273,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/<triple>/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 }}
|
||||
|
||||
Generated
+4
-4
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
@@ -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.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -653,94 +653,48 @@ See [Listing](Listing.md) for pagination state and error behavior.
|
||||
### listNamespaces()
|
||||
|
||||
```ts
|
||||
abstract listNamespaces(namespacePath?, options?): Promise<ListNamespacesResponse>
|
||||
abstract listNamespaces(namespacePath?, options?): Listing<string>
|
||||
```
|
||||
|
||||
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<ListTablesResponse>
|
||||
abstract listTables(namespacePath?, options?): Listing<string>
|
||||
```
|
||||
|
||||
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<ListTablesResponse>
|
||||
```
|
||||
|
||||
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);
|
||||
}
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -1298,7 +1298,11 @@ Specify exactly one of `column` or `indexName`.
|
||||
|
||||
Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in
|
||||
the client process from index metadata. For remote tables, this means the
|
||||
same tokenizer model files must also exist locally.
|
||||
same tokenizer model files must also exist locally (under
|
||||
`LANCE_LANGUAGE_MODEL_HOME`) and must match the files the server uses. A
|
||||
mismatch cannot be detected: the returned tokens may differ from what the
|
||||
server's index produces. If the model files are missing, an error naming
|
||||
the tokenizer is raised.
|
||||
|
||||
#### Parameters
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -65,7 +65,7 @@ then no error will be raised.
|
||||
### mode
|
||||
|
||||
```ts
|
||||
mode: "overwrite" | "create";
|
||||
mode: "overwrite" | "create" | "exist_ok";
|
||||
```
|
||||
|
||||
The mode to use when creating the table.
|
||||
@@ -76,6 +76,9 @@ happen. Any provided data will be ignored.
|
||||
|
||||
If this is set to "overwrite" then any existing table will be replaced.
|
||||
|
||||
If this is set to "exist_ok" then an existing table is opened and any
|
||||
provided data is ignored.
|
||||
|
||||
***
|
||||
|
||||
### schema?
|
||||
|
||||
@@ -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.
|
||||
@@ -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;
|
||||
```
|
||||
@@ -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.
|
||||
@@ -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[];
|
||||
```
|
||||
@@ -20,12 +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 top-level blob column holds the blob bytes as
|
||||
`LargeBinary`. The query also reads `_rowid` to fetch the bytes, and
|
||||
- `"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 struct or 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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -26,14 +26,17 @@ is also an [asynchronous API client](#connections-asynchronous).
|
||||
|
||||
::: 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 +221,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
|
||||
@@ -303,6 +315,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
|
||||
@@ -520,12 +540,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.
|
||||
@@ -633,6 +705,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
|
||||
|
||||
@@ -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<string>): Promise<string[]> {
|
||||
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' ]
|
||||
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -350,6 +350,7 @@
|
||||
"pages": [
|
||||
"enterprise/index",
|
||||
"enterprise/architecture",
|
||||
"enterprise/distributed-search",
|
||||
"enterprise/deployment/index",
|
||||
"deployment/storage",
|
||||
"deployment/monitoring",
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
@@ -384,7 +384,7 @@ The inner workings of RaBitQ quantization are rather mathematically dense. It ge
|
||||
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
|
||||
##### Using RaBitQ {#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"`.
|
||||
Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`.
|
||||
|
||||
|
||||
@@ -10,6 +10,9 @@ Generate embeddings via the [ollama](https://github.com/ollama/ollama-python) py
|
||||
- [Ollama docs on embeddings](https://github.com/ollama/ollama/blob/main/docs/api.md#generate-embeddings)
|
||||
- [Ollama blog on embeddings](https://ollama.com/blog/embedding-models)
|
||||
|
||||
For direct calls, `compute_source_embeddings` accepts a list of strings or a
|
||||
one-dimensional NumPy string array. Both inputs are sent to Ollama in one batch.
|
||||
|
||||
| Parameter | Type | Default Value | Description |
|
||||
|------------------------|----------------------------|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| `name` | `str` | `nomic-embed-text` | The name of the model. |
|
||||
|
||||
@@ -45,6 +45,8 @@ Supported parameters (to be passed in `create` method) are:
|
||||
| `output_dimension` | `int` | `None` | Output embedding dimension. Only supported by `voyage-multimodal-3.5`. Valid options: 256, 512, 1024 (default), 2048. |
|
||||
|
||||
|
||||
For text models, `compute_source_embeddings()` accepts a single string, a list of strings, a one-dimensional NumPy array, or a PyArrow `Array` or `ChunkedArray`. Batch inputs retain their order.
|
||||
|
||||
Usage Example:
|
||||
|
||||
<CodeGroup>
|
||||
@@ -55,10 +57,15 @@ Usage Example:
|
||||
|
||||
### Multimodal Example {#multimodal-example}
|
||||
|
||||
The `voyage-multimodal-3.5` model can embed text alongside images. You can use image URLs, file paths, or PIL Image objects:
|
||||
The `voyage-multimodal-3.5` model can embed text alongside images. You can use image URLs, local `pathlib.Path` objects, raw image bytes, or PIL Image objects.
|
||||
|
||||
For source embeddings, PIL images and images loaded from bytes or local paths are
|
||||
converted to RGB and encoded as JPEG. This supports RGBA, grayscale with alpha,
|
||||
and palette images; transparency is discarded. Composite transparent images onto
|
||||
your chosen background before embedding if the background color matters.
|
||||
|
||||
<CodeGroup>
|
||||
<CodeBlock filename="Python" language="Python" icon="python">
|
||||
{PyEmbeddingVoyageaiMultimodal}
|
||||
</CodeBlock>
|
||||
</CodeGroup>
|
||||
</CodeGroup>
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -58,6 +58,33 @@ uv add --prerelease allow --index https://pypi.fury.io/lancedb/ lancedb
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
If you also need the optional `lance` Python API, install the `pylance` extra and
|
||||
include the Lance preview index:
|
||||
|
||||
<CodeGroup>
|
||||
```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]"
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
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.
|
||||
|
||||
<Note>
|
||||
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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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<string>): Promise<string[]> {
|
||||
const names: string[] = [];
|
||||
for await (const name of items) names.push(name);
|
||||
return names;
|
||||
}
|
||||
|
||||
describe("when connecting", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
@@ -36,6 +36,87 @@ describe("when connecting", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe.each([
|
||||
[
|
||||
"URI overload",
|
||||
(uri: string, options: { storageOptions?: Record<string, string> }) =>
|
||||
connect(uri, options),
|
||||
],
|
||||
[
|
||||
"object overload",
|
||||
async (
|
||||
uri: string,
|
||||
options: { storageOptions?: Record<string, string> },
|
||||
) => {
|
||||
const input = { uri, ...options };
|
||||
if (Object.isFrozen(options)) {
|
||||
Object.freeze(input);
|
||||
}
|
||||
const db = await connect(input);
|
||||
expect(input).toEqual({ uri, ...options });
|
||||
expect(input.storageOptions).toBe(options.storageOptions);
|
||||
return db;
|
||||
},
|
||||
],
|
||||
[
|
||||
"namespace",
|
||||
(uri: string, options: { storageOptions?: Record<string, string> }) =>
|
||||
connectNamespace("dir", { root: uri }, options),
|
||||
],
|
||||
])("%s connection options", (_name, open) => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("should accept frozen options without storage options", async () => {
|
||||
const options = Object.freeze({});
|
||||
const db = await open(tmpDir.name, options);
|
||||
try {
|
||||
await db.createTable("frozen", [{ id: 1 }]);
|
||||
await expect(db.tableNames()).resolves.toContain("frozen");
|
||||
expect(options).toEqual({});
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("should accept frozen options with storage options", async () => {
|
||||
const storageOptions = Object.freeze({
|
||||
newTableDataStorageVersion: "stable",
|
||||
});
|
||||
const options = Object.freeze({ storageOptions });
|
||||
const db = await open(tmpDir.name, options);
|
||||
try {
|
||||
const table = await db.createTable("frozen", [{ id: 1 }]);
|
||||
await expect(table.countRows()).resolves.toBe(1);
|
||||
expect(options.storageOptions).toBe(storageOptions);
|
||||
expect(storageOptions).toEqual({ newTableDataStorageVersion: "stable" });
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("should preserve reused options and their storage options", async () => {
|
||||
const storageOptions = { newTableDataStorageVersion: "stable" };
|
||||
const options = { storageOptions };
|
||||
for (let i = 0; i < 2; i++) {
|
||||
const db = await open(tmpDir.name, options);
|
||||
try {
|
||||
const table = await db.createTable(`reused_${i}`, [{ id: i }]);
|
||||
await expect(table.countRows()).resolves.toBe(1);
|
||||
expect(options.storageOptions).toBe(storageOptions);
|
||||
expect(storageOptions).toEqual({
|
||||
newTableDataStorageVersion: "stable",
|
||||
});
|
||||
} finally {
|
||||
await db.close();
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("given a connection", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let db: Connection;
|
||||
@@ -54,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",
|
||||
);
|
||||
@@ -117,6 +198,80 @@ describe("given a connection", () => {
|
||||
await expect(tbl.countRows()).resolves.toBe(1);
|
||||
});
|
||||
|
||||
describe.each([undefined, []])("root namespace %p", (namespacePath) => {
|
||||
it("should preserve createTable options after the namespace", async () => {
|
||||
await db.createTable("test", [{ id: 1 }, { id: 2 }]);
|
||||
|
||||
const table = await db.createTable("test", [{ id: 3 }], namespacePath, {
|
||||
mode: "overwrite",
|
||||
});
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(rows.map((row) => row.id)).toEqual([3]);
|
||||
});
|
||||
|
||||
it("should preserve createEmptyTable overwrite after the namespace", async () => {
|
||||
await db.createTable("test", [{ id: 1 }]);
|
||||
const schema = new Schema([new Field("id", new Float64(), true)]);
|
||||
|
||||
const table = await db.createEmptyTable("test", schema, namespacePath, {
|
||||
mode: "overwrite",
|
||||
});
|
||||
|
||||
await expect(table.countRows()).resolves.toBe(0);
|
||||
});
|
||||
|
||||
it("should preserve createEmptyTable existOk after the namespace", async () => {
|
||||
await db.createTable("test", [{ id: 1 }]);
|
||||
const schema = new Schema([new Field("id", new Float64(), true)]);
|
||||
|
||||
const table = await db.createEmptyTable("test", schema, namespacePath, {
|
||||
existOk: true,
|
||||
});
|
||||
|
||||
await expect(table.countRows()).resolves.toBe(1);
|
||||
});
|
||||
|
||||
it("should preserve tableNames options after the namespace", async () => {
|
||||
for (const name of ["a", "b", "c"]) {
|
||||
await db.createTable(name, [{ id: 1 }]);
|
||||
}
|
||||
|
||||
await expect(
|
||||
db.tableNames(namespacePath, { limit: 1, startAfter: "a" }),
|
||||
).resolves.toEqual(["b"]);
|
||||
});
|
||||
|
||||
it("should preserve listTables options after the namespace", async () => {
|
||||
for (const name of ["a", "b", "c"]) {
|
||||
await db.createTable(name, [{ id: 1 }]);
|
||||
}
|
||||
|
||||
const first = db.listTables(namespacePath, { pageLimit: 1 });
|
||||
expect(await first.next()).toEqual({ done: false, value: "a" });
|
||||
expect(first.pageToken()).toBeDefined();
|
||||
|
||||
const second = db.listTables(namespacePath, {
|
||||
pageLimit: 1,
|
||||
pageToken: first.pageToken(),
|
||||
});
|
||||
expect(await second.next()).toEqual({ done: false, value: "b" });
|
||||
});
|
||||
});
|
||||
it("should keep the existing table when mode is exist_ok", async () => {
|
||||
const created = await db.createTable("test", [{ id: 1 }, { id: 2 }], {
|
||||
mode: "exist_ok",
|
||||
});
|
||||
await expect(created.countRows()).resolves.toBe(2);
|
||||
|
||||
const reopened = await db.createTable("test", [{ id: 9 }], {
|
||||
mode: "exist_ok",
|
||||
});
|
||||
await expect(reopened.countRows()).resolves.toBe(2);
|
||||
const rows = (await reopened.query().toArray()) as { id: number }[];
|
||||
expect(rows.map((row) => row.id).sort()).toEqual([1, 2]);
|
||||
});
|
||||
|
||||
it("should respect limit and page token when listing tables", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
@@ -144,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 () => {
|
||||
@@ -167,18 +321,71 @@ 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(await collectNames(db.listTables([], { pageLimit: 2 }))).toEqual(
|
||||
created,
|
||||
);
|
||||
});
|
||||
|
||||
expect(seen).toEqual(created);
|
||||
it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])(
|
||||
"should not wrap a large table listing limit %s",
|
||||
async (limit) => {
|
||||
await db.createTable("a", [{ id: 1 }]);
|
||||
await db.createTable("b", [{ id: 2 }]);
|
||||
for (const names of [
|
||||
db.tableNames({ limit }),
|
||||
db.tableNames([], { limit }),
|
||||
]) {
|
||||
await expect(names).resolves.toEqual(["a", "b"]);
|
||||
}
|
||||
for (const namespacePath of [undefined, []]) {
|
||||
expect(() =>
|
||||
db.listTables(namespacePath, { pageLimit: limit }),
|
||||
).toThrow(/limit/);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
-1,
|
||||
1.5,
|
||||
Number.NaN,
|
||||
Number.POSITIVE_INFINITY,
|
||||
Number.NEGATIVE_INFINITY,
|
||||
Number.MAX_SAFE_INTEGER + 1,
|
||||
])("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/);
|
||||
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 () => {
|
||||
await db.createTable("a", [{ id: 1 }]);
|
||||
await db.createTable("b", [{ id: 2 }]);
|
||||
for (const [limit, names] of [
|
||||
[0, []],
|
||||
[1, ["a"]],
|
||||
[2 ** 32 - 1, ["a", "b"]],
|
||||
] as const) {
|
||||
await expect(db.tableNames({ limit })).resolves.toEqual(names);
|
||||
}
|
||||
await expect(db.tableNames()).resolves.toEqual(["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 () => {
|
||||
@@ -191,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 () => {
|
||||
@@ -215,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 () => {
|
||||
@@ -293,6 +493,69 @@ describe("given a connection", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe.each(["createTable", "createEmptyTable"] as const)(
|
||||
"%s storage options",
|
||||
(method) => {
|
||||
const schema = new Schema([new Field("id", new Float64(), true)]);
|
||||
|
||||
it("should accept frozen creation options without storage options", async () => {
|
||||
const options = Object.freeze({
|
||||
dataStorageVersion: "stable",
|
||||
enableV2ManifestPaths: true,
|
||||
});
|
||||
const table =
|
||||
method === "createTable"
|
||||
? await db.createTable("frozen_options", [{ id: 1 }], options)
|
||||
: await db.createEmptyTable("frozen_options", schema, options);
|
||||
expect(await (table as LocalTable).usesV2ManifestPaths()).toBe(true);
|
||||
expect(options).not.toHaveProperty("storageOptions");
|
||||
});
|
||||
|
||||
it("should accept frozen storage options and preserve override precedence", async () => {
|
||||
const storageOptions = Object.freeze({
|
||||
newTableDataStorageVersion: "legacy",
|
||||
newTableEnableV2ManifestPaths: "false",
|
||||
});
|
||||
const options = Object.freeze({
|
||||
storageOptions,
|
||||
dataStorageVersion: "stable",
|
||||
enableV2ManifestPaths: true,
|
||||
});
|
||||
const table =
|
||||
method === "createTable"
|
||||
? await db.createTable("frozen_storage", [{ id: 1 }], options)
|
||||
: await db.createEmptyTable("frozen_storage", schema, options);
|
||||
expect(await (table as LocalTable).usesV2ManifestPaths()).toBe(true);
|
||||
expect(options.storageOptions).toBe(storageOptions);
|
||||
expect(storageOptions).toEqual({
|
||||
newTableDataStorageVersion: "legacy",
|
||||
newTableEnableV2ManifestPaths: "false",
|
||||
});
|
||||
});
|
||||
|
||||
it("should not retain a previous override when reusing storage options", async () => {
|
||||
const storageOptions = { newTableEnableV2ManifestPaths: "false" };
|
||||
const create = async (name: string, enableV2ManifestPaths?: boolean) =>
|
||||
method === "createTable"
|
||||
? db.createTable(name, [{ id: 1 }], {
|
||||
storageOptions,
|
||||
enableV2ManifestPaths,
|
||||
})
|
||||
: db.createEmptyTable(name, schema, {
|
||||
storageOptions,
|
||||
enableV2ManifestPaths,
|
||||
});
|
||||
const first = await create("override", true);
|
||||
expect(await (first as LocalTable).usesV2ManifestPaths()).toBe(true);
|
||||
const second = await create("without_override");
|
||||
expect(await (second as LocalTable).usesV2ManifestPaths()).toBe(false);
|
||||
expect(storageOptions).toEqual({
|
||||
newTableEnableV2ManifestPaths: "false",
|
||||
});
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
it("should be able to migrate tables to the V2 manifest paths", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const table = (await db.createEmptyTable(
|
||||
@@ -451,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 () => {
|
||||
@@ -474,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",
|
||||
);
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
const assert = require("node:assert/strict");
|
||||
const {
|
||||
TransformersEmbeddingFunction,
|
||||
} = require("../../dist/embedding/transformers");
|
||||
const { getRegistry } = require("../../dist/embedding/registry");
|
||||
|
||||
async function main() {
|
||||
const transformers = await import("@huggingface/transformers");
|
||||
// Never fetch a real model, including if a loader mock stops intercepting.
|
||||
transformers.env.allowRemoteModels = false;
|
||||
transformers.env.allowLocalModels = false;
|
||||
const modelCalls = [];
|
||||
const tokenizerCalls = [];
|
||||
const testCase = Number(process.argv[2]);
|
||||
const defaults = "Xenova/all-MiniLM-L6-v2";
|
||||
const cases = [
|
||||
[{}, defaults, defaults],
|
||||
[{ model: "test/model" }, "test/model", "test/model"],
|
||||
[
|
||||
{ model: "test/model", tokenizer: "test/tokenizer" },
|
||||
"test/model",
|
||||
"test/tokenizer",
|
||||
],
|
||||
[
|
||||
{ model: "test/model", tokenizer: "$var:tokenizer" },
|
||||
"test/model",
|
||||
"test/tokenizer",
|
||||
],
|
||||
[
|
||||
{ model: "test/model", tokenizer: "test/tokenizer" },
|
||||
"test/model",
|
||||
"test/tokenizer",
|
||||
],
|
||||
];
|
||||
assert.ok(
|
||||
Number.isInteger(testCase) && testCase >= 0 && testCase < cases.length,
|
||||
);
|
||||
transformers.AutoModel.from_pretrained = async (...args) => {
|
||||
modelCalls.push(args);
|
||||
return {};
|
||||
};
|
||||
transformers.AutoTokenizer.from_pretrained = async (...args) => {
|
||||
tokenizerCalls.push(args);
|
||||
if (testCase === 4) throw new Error("missing tokenizer");
|
||||
return {};
|
||||
};
|
||||
getRegistry().setVar("tokenizer", "test/tokenizer");
|
||||
const [options, model, tokenizer] = cases[testCase];
|
||||
const fn = new TransformersEmbeddingFunction(options);
|
||||
if (testCase === 4) {
|
||||
await assert.rejects(fn.init(), /error loading tokenizer test\/tokenizer/);
|
||||
} else {
|
||||
await fn.init();
|
||||
}
|
||||
assert.deepEqual(modelCalls, [[model, { dtype: "fp32" }]]);
|
||||
assert.deepEqual(tokenizerCalls, [[tokenizer]]);
|
||||
assert.deepEqual(fn.toJSON(), options);
|
||||
}
|
||||
|
||||
main().catch((error) => {
|
||||
console.error(error);
|
||||
process.exitCode = 1;
|
||||
});
|
||||
@@ -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([]);
|
||||
});
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
});
|
||||
});
|
||||
+218
-47
@@ -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();
|
||||
@@ -3385,6 +3398,106 @@ describe("when dealing with blob columns", () => {
|
||||
expect(bytes[0]!.equals(payload)).toBe(true);
|
||||
});
|
||||
|
||||
it("returns bytes for blobs nested in structs", async () => {
|
||||
const { table, alpha, empty } = await openNestedBlobTable();
|
||||
const rows = await table.query().toArray({ blobMode: "bytes" });
|
||||
const byId = new Map(rows.map((row) => [Number(row.id), row]));
|
||||
expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha);
|
||||
expect(byId.get(1)!.info.caption).toBe("one");
|
||||
expect(byId.get(2)!.info.image).toBeNull();
|
||||
expect(byId.get(2)!.info.caption).toBe("two");
|
||||
expect(byId.get(3)!.info).toBeNull();
|
||||
expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty);
|
||||
});
|
||||
|
||||
it("lines up nested rows and bytes across single-row batches", async () => {
|
||||
const { table, alpha, empty } = await openNestedBlobTable();
|
||||
const result = await table
|
||||
.query()
|
||||
.toArrow({ blobMode: "bytes", maxBatchLength: 1 });
|
||||
expect(result.batches.length).toBe(4);
|
||||
const byId = new Map(result.toArray().map((row) => [Number(row.id), row]));
|
||||
expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha);
|
||||
expect(byId.get(1)!.info.caption).toBe("one");
|
||||
expect(byId.get(2)!.info.image).toBeNull();
|
||||
expect(byId.get(2)!.info.caption).toBe("two");
|
||||
expect(byId.get(3)!.info).toBeNull();
|
||||
expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty);
|
||||
});
|
||||
|
||||
it("handles struct fields named like Object prototypes", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field(
|
||||
"info",
|
||||
new Struct([
|
||||
new Field("constructor", new Struct([blob("image")]), true),
|
||||
]),
|
||||
true,
|
||||
),
|
||||
]);
|
||||
const payload = Buffer.from("proto");
|
||||
const table = await db.createTable(
|
||||
"prototype_named_blobs",
|
||||
[{ id: 1n, info: { constructor: { image: payload } } }],
|
||||
{ schema },
|
||||
);
|
||||
expect(await table.blobColumns()).toEqual(["info.constructor.image"]);
|
||||
const [row] = await table.query().toArray({ blobMode: "bytes" });
|
||||
// Arrow's StructRow shadows the `constructor` property, so read the
|
||||
// field through the JSON form instead of property access.
|
||||
const info = JSON.parse(JSON.stringify(row.info));
|
||||
expect(Buffer.from(Object.values<number>(info.constructor.image))).toEqual(
|
||||
payload,
|
||||
);
|
||||
});
|
||||
|
||||
it("returns nested blob bytes for a struct selection", async () => {
|
||||
const { table, alpha } = await openNestedBlobTable();
|
||||
const [row] = await table
|
||||
.query()
|
||||
.where("id = 1")
|
||||
.select(["info"])
|
||||
.toArray({ blobMode: "bytes" });
|
||||
expect(Object.keys(row)).toEqual(["info"]);
|
||||
expect(Buffer.from(row.info.image)).toEqual(alpha);
|
||||
expect(row.info.caption).toBe("one");
|
||||
});
|
||||
|
||||
it("returns nested blob bytes through aliases", async () => {
|
||||
const { table, alpha } = await openNestedBlobTable();
|
||||
const [row] = await table
|
||||
.query()
|
||||
.where("id = 1")
|
||||
.select(
|
||||
new Map([
|
||||
["picture", "info.image"],
|
||||
["details", "info"],
|
||||
]),
|
||||
)
|
||||
.toArray({ blobMode: "bytes" });
|
||||
expect(Buffer.from(row.picture)).toEqual(alpha);
|
||||
expect(Buffer.from(row.details.image)).toEqual(alpha);
|
||||
expect(row.details.caption).toBe("one");
|
||||
});
|
||||
|
||||
it("returns a nested bytes schema for an empty result", async () => {
|
||||
const { table } = await openNestedBlobTable();
|
||||
const result = await table
|
||||
.query()
|
||||
.where("id > 100")
|
||||
.toArrow({ blobMode: "bytes" });
|
||||
expect(result.numRows).toBe(0);
|
||||
const field = result.schema.fields.find((f) => f.name === "info")!;
|
||||
expect(arrow.DataType.isStruct(field.type)).toBe(true);
|
||||
const image = (field.type as Struct).children.find(
|
||||
(f) => f.name === "image",
|
||||
)!;
|
||||
expect(arrow.DataType.isLargeBinary(image.type)).toBe(true);
|
||||
expect(isBlobField(image)).toBe(false);
|
||||
});
|
||||
|
||||
it("creates list blobs without advertising unfetchable elements", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema([
|
||||
@@ -3542,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 = (
|
||||
@@ -3563,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 };
|
||||
}
|
||||
@@ -3593,6 +3724,31 @@ describe("when dealing with blob columns", () => {
|
||||
return { table, rowIds, alpha, beta };
|
||||
}
|
||||
|
||||
async function openNestedBlobTable() {
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field(
|
||||
"info",
|
||||
new Struct([blob("image"), new Field("caption", new Utf8(), true)]),
|
||||
true,
|
||||
),
|
||||
]);
|
||||
const alpha = Buffer.from("alpha");
|
||||
const empty = Buffer.alloc(0);
|
||||
const table = await db.createTable(
|
||||
"nested_blobs",
|
||||
[{ id: 1n, info: { image: alpha, caption: "one" } }],
|
||||
{ schema },
|
||||
);
|
||||
await table.add([
|
||||
{ id: 2n, info: { image: null, caption: "two" } },
|
||||
{ id: 3n, info: null },
|
||||
{ id: 4n, info: { image: empty, caption: "four" } },
|
||||
]);
|
||||
return { table, alpha, empty };
|
||||
}
|
||||
|
||||
async function openExternalBlobTable() {
|
||||
const db = await connect(path.join(tmpDir.name, "db"));
|
||||
const payload = Buffer.from("external-payload");
|
||||
@@ -3742,6 +3898,21 @@ describe("when optimizing a dataset", () => {
|
||||
expect(stats.prune.oldVersionsRemoved).toBe(2);
|
||||
});
|
||||
|
||||
it("retains tagged versions older than the cleanup cutoff", async () => {
|
||||
await (await table.tags()).create("keep", 1);
|
||||
await table.add([{ id: 3 }]);
|
||||
|
||||
const cutoff = await nextMillisecond();
|
||||
const stats = await table.optimize({ cleanupOlderThan: cutoff });
|
||||
// Versions 2 and 3 are pruned, tagged version 1 is retained
|
||||
expect(stats.prune.oldVersionsRemoved).toBe(2);
|
||||
|
||||
await table.checkout("keep");
|
||||
expect(await table.countRows()).toBe(1);
|
||||
await table.checkoutLatest();
|
||||
expect(await table.countRows()).toBe(3);
|
||||
});
|
||||
|
||||
it("delete unverified", async () => {
|
||||
const version = await table.version();
|
||||
const versionFile = `${tmpDir.name}/${table.name}.lance/_versions/${String(
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { execFileSync } from "node:child_process";
|
||||
import { resolve } from "node:path";
|
||||
|
||||
describe("Transformers embedding tokenizer selection", () => {
|
||||
// Run init's native ESM import outside Jest's CommonJS sandbox, like a user.
|
||||
// The fixture mocks model loading and forbids local/remote model access.
|
||||
test.each([
|
||||
["default model", 0],
|
||||
["default tokenizer", 1],
|
||||
["custom tokenizer", 2],
|
||||
["tokenizer registry variable", 3],
|
||||
["tokenizer loading error", 4],
|
||||
])("handles %s", (_name, testCase) => {
|
||||
execFileSync(
|
||||
process.execPath,
|
||||
[
|
||||
resolve(__dirname, "fixtures", "transformers_tokenizer.cjs"),
|
||||
String(testCase),
|
||||
],
|
||||
{ stdio: "pipe", timeout: 15_000 },
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -61,6 +61,7 @@ import {
|
||||
import { inferSchema } from "./schema";
|
||||
|
||||
export * from "apache-arrow";
|
||||
export type { Data as ArrowData } from "apache-arrow";
|
||||
export type SchemaLike =
|
||||
| Schema
|
||||
| {
|
||||
|
||||
+58
-118
@@ -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";
|
||||
@@ -53,8 +49,11 @@ export interface CreateTableOptions {
|
||||
* happen. Any provided data will be ignored.
|
||||
*
|
||||
* If this is set to "overwrite" then any existing table will be replaced.
|
||||
*
|
||||
* If this is set to "exist_ok" then an existing table is opened and any
|
||||
* provided data is ignored.
|
||||
*/
|
||||
mode: "create" | "overwrite";
|
||||
mode: "create" | "overwrite" | "exist_ok";
|
||||
/**
|
||||
* If this is true and the table already exists and the mode is "create"
|
||||
* then no error will be raised.
|
||||
@@ -138,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 {
|
||||
@@ -154,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";
|
||||
@@ -275,45 +248,18 @@ export abstract class Connection {
|
||||
options?: Partial<TableNamesOptions>,
|
||||
): Promise<string[]>;
|
||||
|
||||
/**
|
||||
* 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<ListTablesOptions>} options - Pagination options
|
||||
* (`pageToken`, `limit`).
|
||||
* @returns {Promise<ListTablesResponse>} A page of table names and an
|
||||
* optional token for the tables after it.
|
||||
*/
|
||||
abstract listTables(
|
||||
options?: Partial<ListTablesOptions>,
|
||||
): Promise<ListTablesResponse>;
|
||||
/**
|
||||
* 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<ListTablesOptions>} options - Pagination options
|
||||
* (`pageToken`, `limit`).
|
||||
* @returns {Promise<ListTablesResponse>} A page of table names and an
|
||||
* optional token for the tables after it.
|
||||
*/
|
||||
abstract listTables(
|
||||
namespacePath?: string[],
|
||||
options?: Partial<ListTablesOptions>,
|
||||
): Promise<ListTablesResponse>;
|
||||
options?: ListingOptions,
|
||||
): Listing<string>;
|
||||
|
||||
/**
|
||||
* Open a table in the database.
|
||||
@@ -539,23 +485,11 @@ export abstract class Connection {
|
||||
namespacePath: string[],
|
||||
): Promise<DescribeNamespaceResponse>;
|
||||
|
||||
/**
|
||||
* 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<ListNamespacesOptions>} options - Pagination options
|
||||
* (`pageToken`, `limit`).
|
||||
* @returns {Promise<ListNamespacesResponse>} 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<ListNamespacesOptions>,
|
||||
): Promise<ListNamespacesResponse>;
|
||||
options?: ListingOptions,
|
||||
): Listing<string>;
|
||||
|
||||
/**
|
||||
* Create a new namespace at the given path.
|
||||
@@ -717,13 +651,13 @@ export class LocalConnection extends Connection {
|
||||
} else {
|
||||
// First argument is options object (backwards compatibility)
|
||||
namespacePath = undefined;
|
||||
tableNamesOptions = namespacePathOrOptions;
|
||||
tableNamesOptions = namespacePathOrOptions ?? options;
|
||||
}
|
||||
|
||||
return this.inner.tableNames(
|
||||
namespacePath ?? [],
|
||||
tableNamesOptions?.startAfter,
|
||||
tableNamesOptions?.limit,
|
||||
normalizeTableListingLimit(tableNamesOptions?.limit),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -818,22 +752,17 @@ export class LocalConnection extends Connection {
|
||||
);
|
||||
}
|
||||
|
||||
async listTables(
|
||||
namespacePathOrOptions?: string[] | Partial<ListTablesOptions>,
|
||||
options?: Partial<ListTablesOptions>,
|
||||
): Promise<ListTablesResponse> {
|
||||
// Detect if first argument is namespacePath array or options object
|
||||
const namespacePath = Array.isArray(namespacePathOrOptions)
|
||||
? namespacePathOrOptions
|
||||
: undefined;
|
||||
const listTablesOptions = Array.isArray(namespacePathOrOptions)
|
||||
? options
|
||||
: namespacePathOrOptions;
|
||||
|
||||
return this.inner.listTables(
|
||||
namespacePath ?? [],
|
||||
listTablesOptions?.pageToken,
|
||||
listTablesOptions?.limit,
|
||||
listTables(
|
||||
namespacePath?: string[],
|
||||
options: ListingOptions = {},
|
||||
): Listing<string> {
|
||||
validateListingOptions(options);
|
||||
return new Listing(
|
||||
this.inner.listTables(
|
||||
namespacePath ?? [],
|
||||
options.pageToken,
|
||||
options.pageLimit,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -890,23 +819,24 @@ export class LocalConnection extends Connection {
|
||||
private getStorageOptions(
|
||||
options?: Partial<CreateTableOptions>,
|
||||
): Record<string, string> | undefined {
|
||||
let storageOptions = options?.storageOptions;
|
||||
if (options?.dataStorageVersion !== undefined) {
|
||||
if (options.storageOptions === undefined) {
|
||||
options.storageOptions = {};
|
||||
}
|
||||
options.storageOptions["newTableDataStorageVersion"] =
|
||||
options.dataStorageVersion;
|
||||
storageOptions = {
|
||||
...storageOptions,
|
||||
newTableDataStorageVersion: options.dataStorageVersion,
|
||||
};
|
||||
}
|
||||
|
||||
if (options?.enableV2ManifestPaths !== undefined) {
|
||||
if (options.storageOptions === undefined) {
|
||||
options.storageOptions = {};
|
||||
}
|
||||
options.storageOptions["newTableEnableV2ManifestPaths"] =
|
||||
options.enableV2ManifestPaths ? "true" : "false";
|
||||
storageOptions = {
|
||||
...storageOptions,
|
||||
newTableEnableV2ManifestPaths: options.enableV2ManifestPaths
|
||||
? "true"
|
||||
: "false",
|
||||
};
|
||||
}
|
||||
|
||||
return cleanseStorageOptions(options?.storageOptions);
|
||||
return cleanseStorageOptions(storageOptions);
|
||||
}
|
||||
|
||||
async createTable(
|
||||
@@ -939,7 +869,7 @@ export class LocalConnection extends Connection {
|
||||
} else {
|
||||
// Third argument is options object (backwards compatibility)
|
||||
namespacePath = undefined;
|
||||
createOptions = namespacePathOrOptions;
|
||||
createOptions = namespacePathOrOptions ?? options;
|
||||
}
|
||||
|
||||
return this._createTableImpl(name, data, namespacePath, createOptions);
|
||||
@@ -986,7 +916,7 @@ export class LocalConnection extends Connection {
|
||||
} else {
|
||||
// Third argument is options object (backwards compatibility)
|
||||
namespacePath = undefined;
|
||||
createOptions = namespacePathOrOptions;
|
||||
createOptions = namespacePathOrOptions ?? options;
|
||||
}
|
||||
|
||||
let mode: string = createOptions?.mode ?? "create";
|
||||
@@ -1035,12 +965,15 @@ export class LocalConnection extends Connection {
|
||||
|
||||
listNamespaces(
|
||||
namespacePath?: string[],
|
||||
options?: Partial<ListNamespacesOptions>,
|
||||
): Promise<ListNamespacesResponse> {
|
||||
return this.inner.listNamespaces(
|
||||
namespacePath ?? [],
|
||||
options?.pageToken,
|
||||
options?.limit,
|
||||
options: ListingOptions = {},
|
||||
): Listing<string> {
|
||||
validateListingOptions(options);
|
||||
return new Listing(
|
||||
this.inner.listNamespaces(
|
||||
namespacePath ?? [],
|
||||
options.pageToken,
|
||||
options.pageLimit,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1103,6 +1036,13 @@ export class LocalConnection extends Connection {
|
||||
}
|
||||
}
|
||||
|
||||
function normalizeTableListingLimit(limit?: number): number | undefined {
|
||||
validateNonNegativeInteger(limit, "limit");
|
||||
// The native binding takes u32: larger JS numbers otherwise wrap to zero or
|
||||
// a smaller page size. A listing limit is an upper bound, so saturate it.
|
||||
return limit === undefined ? undefined : Math.min(limit, 0xffffffff);
|
||||
}
|
||||
|
||||
/**
|
||||
* Takes storage options and makes all the keys snake case.
|
||||
*/
|
||||
|
||||
@@ -38,6 +38,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
#model?: import("@huggingface/transformers").PreTrainedModel;
|
||||
#tokenizer?: import("@huggingface/transformers").PreTrainedTokenizer;
|
||||
#modelName: XenovaTransformerOptions["model"];
|
||||
#tokenizerName: string;
|
||||
#initialized = false;
|
||||
#tokenizerOptions: XenovaTransformerOptions["tokenizerOptions"];
|
||||
#ndims?: number;
|
||||
@@ -58,6 +59,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
|
||||
this.#ndims = options.ndims;
|
||||
this.#modelName = modelName;
|
||||
this.#tokenizerName = options.tokenizer ?? modelName;
|
||||
}
|
||||
|
||||
async init() {
|
||||
@@ -86,11 +88,11 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
}
|
||||
try {
|
||||
this.#tokenizer = await transformers.AutoTokenizer.from_pretrained(
|
||||
this.#modelName,
|
||||
this.#tokenizerName,
|
||||
);
|
||||
} catch (e) {
|
||||
throw new Error(
|
||||
`error loading tokenizer for ${this.#modelName}. Make sure you are using a wasm compatible model:\nReason: ${e}`,
|
||||
`error loading tokenizer ${this.#tokenizerName}. Make sure you are using a wasm compatible model:\nReason: ${e}`,
|
||||
);
|
||||
}
|
||||
this.#initialized = true;
|
||||
|
||||
+8
-13
@@ -97,13 +97,9 @@ export {
|
||||
Connection,
|
||||
CreateTableOptions,
|
||||
TableNamesOptions,
|
||||
ListTablesOptions,
|
||||
OpenTableOptions,
|
||||
ListNamespacesOptions,
|
||||
CreateNamespaceOptions,
|
||||
DropNamespaceOptions,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
CreateNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
@@ -420,10 +416,10 @@ export async function connect(
|
||||
throw new Error("uri is required");
|
||||
}
|
||||
|
||||
finalOptions = (finalOptions as ConnectionOptions) ?? {};
|
||||
(<ConnectionOptions>finalOptions).storageOptions = cleanseStorageOptions(
|
||||
(<ConnectionOptions>finalOptions).storageOptions,
|
||||
);
|
||||
finalOptions = {
|
||||
...finalOptions,
|
||||
storageOptions: cleanseStorageOptions(finalOptions.storageOptions),
|
||||
};
|
||||
|
||||
// Create native header provider if one was provided
|
||||
let nativeProvider: NativeJsHeaderProvider | undefined;
|
||||
@@ -630,11 +626,10 @@ export async function connectNamespace(
|
||||
properties = configOrProperties as Record<string, string>;
|
||||
}
|
||||
|
||||
const finalOptions: ConnectNamespaceOptions = (options ??
|
||||
{}) as ConnectNamespaceOptions;
|
||||
finalOptions.storageOptions = cleanseStorageOptions(
|
||||
finalOptions.storageOptions,
|
||||
);
|
||||
const finalOptions: ConnectNamespaceOptions = {
|
||||
...options,
|
||||
storageOptions: cleanseStorageOptions(options?.storageOptions),
|
||||
};
|
||||
|
||||
const nativeConn = await LanceDbConnection.newWithNamespace(
|
||||
implName,
|
||||
|
||||
@@ -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;
|
||||
|
||||
+108
-37
@@ -1,8 +1,10 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import type { Data as ArrowData } from "apache-arrow";
|
||||
import {
|
||||
Table as ArrowTable,
|
||||
DataType,
|
||||
Field,
|
||||
type IntoVector,
|
||||
LargeBinary,
|
||||
@@ -12,7 +14,6 @@ import {
|
||||
extractVectorBuffer,
|
||||
fromBufferToRecordBatch,
|
||||
fromRecordBatchToBuffer,
|
||||
makeData,
|
||||
tableFromIPC,
|
||||
vectorFromArray,
|
||||
} from "./arrow";
|
||||
@@ -66,20 +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<string, string> {
|
||||
const columns = new Set(blobColumns);
|
||||
const pairs =
|
||||
projection ??
|
||||
blobColumns.map((column): [string, string] => [column, column]);
|
||||
return new Map(pairs.filter(([, source]) => columns.has(source)));
|
||||
): 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,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
return sources;
|
||||
}
|
||||
|
||||
async function readBlobBytes(
|
||||
table: ArrowTable,
|
||||
sources: Map<string, string>,
|
||||
sources: BlobSource[],
|
||||
snapshot: NativeTable,
|
||||
dropRowId: boolean,
|
||||
): Promise<ArrowTable> {
|
||||
@@ -91,18 +114,25 @@ async function readBlobBytes(
|
||||
rowIdColumn === null
|
||||
? []
|
||||
: Array.from(rowIdColumn.toArray() as BigUint64Array);
|
||||
const replacements = new Map<number, (Buffer | null)[]>();
|
||||
for (const [output, source] of sources) {
|
||||
const index = table.schema.fields.findIndex(
|
||||
(field) => field.name === output,
|
||||
);
|
||||
if (index < 0) {
|
||||
const outputFields = new Map(
|
||||
table.schema.fields.map((field) => [field.name, field]),
|
||||
);
|
||||
const replacements = new Map<Field, (Buffer | null)[]>();
|
||||
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(
|
||||
index,
|
||||
field,
|
||||
values.map((value) => value ?? null),
|
||||
);
|
||||
}
|
||||
@@ -117,37 +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<number, (Buffer | null)[]>,
|
||||
replacements: Map<Field, (Buffer | null)[]>,
|
||||
): ArrowTable {
|
||||
if (replacements.size === 0) {
|
||||
return table;
|
||||
}
|
||||
const fields = table.schema.fields.map((field, index) =>
|
||||
replacements.has(index)
|
||||
? new Field(field.name, new LargeBinary(), true)
|
||||
: 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);
|
||||
}
|
||||
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 values = replacements.get(index);
|
||||
if (values === undefined) {
|
||||
return child;
|
||||
const replaceData = (
|
||||
data: ArrowData,
|
||||
field: Field,
|
||||
updated: Field,
|
||||
): ArrowData => {
|
||||
if (field === updated) {
|
||||
return data;
|
||||
}
|
||||
return vectorFromArray(
|
||||
values.slice(offset, offset + batch.numRows),
|
||||
new LargeBinary(),
|
||||
).data[0];
|
||||
});
|
||||
const values = replacements.get(field);
|
||||
if (values !== undefined) {
|
||||
return vectorFromArray(
|
||||
values.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);
|
||||
@@ -175,12 +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 top-level blob column holds the blob bytes as
|
||||
* `LargeBinary`. The query also reads `_rowid` to fetch the bytes, and
|
||||
* - `"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 struct or 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.
|
||||
@@ -433,7 +504,7 @@ export class QueryBase<
|
||||
await snapshot.blobColumns(),
|
||||
context.projection,
|
||||
);
|
||||
if (sources.size === 0) {
|
||||
if (sources.length === 0) {
|
||||
return this.collect(options, false);
|
||||
}
|
||||
const addRowId = !context.withRowId;
|
||||
|
||||
@@ -989,7 +989,11 @@ export abstract class Table {
|
||||
*
|
||||
* Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in
|
||||
* the client process from index metadata. For remote tables, this means the
|
||||
* same tokenizer model files must also exist locally.
|
||||
* same tokenizer model files must also exist locally (under
|
||||
* `LANCE_LANGUAGE_MODEL_HOME`) and must match the files the server uses. A
|
||||
* mismatch cannot be detected: the returned tokens may differ from what the
|
||||
* server's index produces. If the model files are missing, an error naming
|
||||
* the tokenizer is raised.
|
||||
*/
|
||||
abstract tokenize(
|
||||
query: string,
|
||||
|
||||
+19
-82
@@ -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<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct ListNamespacesResponse {
|
||||
pub namespaces: Vec<String>,
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct ListTablesResponse {
|
||||
pub tables: Vec<String>,
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct CreateNamespaceResponse {
|
||||
pub properties: Option<HashMap<String, String>>,
|
||||
@@ -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<Vec<String>>,
|
||||
page_token: Option<String>,
|
||||
limit: Option<u32>,
|
||||
) -> napi::Result<ListTablesResponse> {
|
||||
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<crate::listing::NameListing> {
|
||||
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<Vec<String>>,
|
||||
page_token: Option<String>,
|
||||
limit: Option<u32>,
|
||||
) -> napi::Result<ListNamespacesResponse> {
|
||||
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<crate::listing::NameListing> {
|
||||
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<u32>) -> napi::Result<Option<i32>> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
+6
-6
@@ -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<String> {
|
||||
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<String> {
|
||||
self.inner.job_type()
|
||||
|
||||
@@ -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"])
|
||||
```
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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`.
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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]: ...
|
||||
@@ -581,6 +581,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:
|
||||
@@ -1011,6 +1012,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: ...
|
||||
|
||||
|
||||
+98
-52
@@ -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(
|
||||
@@ -542,7 +548,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 +605,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 +674,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 +748,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 +786,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,7 +814,9 @@ 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:
|
||||
@@ -1230,18 +1254,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")
|
||||
|
||||
@@ -1300,6 +1325,10 @@ class LanceDBConnection(DBConnection):
|
||||
self.storage_options = storage_options
|
||||
self._manifest_enabled = manifest_enabled
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
# The native connection does not expose the Python Session wrapper.
|
||||
# Keep the object supplied by the caller so this property stays
|
||||
# synchronous for debugger inspection.
|
||||
self._session = session
|
||||
if _inner is not None:
|
||||
self._conn = _inner
|
||||
# Native-derived wrappers resolve this in their async reconstruction
|
||||
@@ -1366,7 +1395,12 @@ class LanceDBConnection(DBConnection):
|
||||
|
||||
@property
|
||||
def session(self) -> Optional[Session]:
|
||||
return self._conn.session
|
||||
"""Return the session supplied when this connection was opened.
|
||||
|
||||
None when the caller did not pass a session, including connections
|
||||
reconstructed with :meth:`from_inner`.
|
||||
"""
|
||||
return self._session
|
||||
|
||||
@property
|
||||
def uri(self) -> str:
|
||||
@@ -1427,7 +1461,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
|
||||
@@ -1443,16 +1477,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
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1503,8 +1535,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
|
||||
----------
|
||||
@@ -1520,12 +1552,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
|
||||
)
|
||||
@@ -2277,12 +2309,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
|
||||
@@ -2297,15 +2329,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,
|
||||
@@ -2384,13 +2419,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
|
||||
----------
|
||||
@@ -2406,15 +2441,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,
|
||||
@@ -2832,7 +2870,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
|
||||
@@ -3372,8 +3412,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:
|
||||
|
||||
@@ -38,9 +38,11 @@ class EmbeddingFunction(BaseModel, ABC):
|
||||
3. ndims() which returns the number of dimensions of the vector column
|
||||
"""
|
||||
|
||||
max_retries: int = (
|
||||
7 # Setting 0 disables retires. Maybe this should not be enabled by default,
|
||||
)
|
||||
# Maximum number of retries for transient errors (rate limits, 5xx responses,
|
||||
# connection errors, ...), with exponential backoff capped at 60 seconds between
|
||||
# attempts. Permanent errors (HTTP 4xx other than 408/409/429) are never retried.
|
||||
# Setting 0 disables retries.
|
||||
max_retries: int = 7
|
||||
_ndims: int = PrivateAttr()
|
||||
_original_args: dict = PrivateAttr()
|
||||
|
||||
@@ -239,7 +241,9 @@ class EmbeddingFunction(BaseModel, ABC):
|
||||
Creates a pydantic Field that can automatically annotate
|
||||
the source column for this embedding function
|
||||
"""
|
||||
return Field(json_schema_extra={"source_column_for": self}, **kwargs)
|
||||
json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {})
|
||||
json_schema_extra["source_column_for"] = self
|
||||
return Field(json_schema_extra=json_schema_extra, **kwargs)
|
||||
|
||||
def VectorField(self, **kwargs):
|
||||
"""
|
||||
@@ -250,7 +254,9 @@ class EmbeddingFunction(BaseModel, ABC):
|
||||
static type checkers. LanceDB will infer the fixed vector dimension from
|
||||
this embedding function.
|
||||
"""
|
||||
return Field(json_schema_extra={"vector_column_for": self}, **kwargs)
|
||||
json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {})
|
||||
json_schema_extra["vector_column_for"] = self
|
||||
return Field(json_schema_extra=json_schema_extra, **kwargs)
|
||||
|
||||
def __eq__(self, __value: object) -> bool:
|
||||
if not hasattr(__value, "__dict__"):
|
||||
|
||||
@@ -5,9 +5,10 @@
|
||||
import os
|
||||
import io
|
||||
import base64
|
||||
from functools import cached_property
|
||||
from urllib.parse import urlparse
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, ClassVar, List, Union, Optional, Any, Dict
|
||||
from typing import TYPE_CHECKING, List, Union, Optional, Any, Dict
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
@@ -51,7 +52,6 @@ class JinaEmbeddings(EmbeddingFunction):
|
||||
|
||||
name: str = "jina-clip-v1"
|
||||
api_key: Optional[str] = None
|
||||
_session: ClassVar = None
|
||||
|
||||
def ndims(self):
|
||||
# TODO: fix hardcoding
|
||||
@@ -206,8 +206,7 @@ class JinaEmbeddings(EmbeddingFunction):
|
||||
texts: list[str] or np.ndarray (of str)
|
||||
The texts to embed
|
||||
"""
|
||||
self._init_client()
|
||||
resp = JinaEmbeddings._session.post( # type: ignore
|
||||
resp = self._session.post(
|
||||
API_URL, json={"input": input, "model": self.name}
|
||||
).json()
|
||||
if "data" not in resp:
|
||||
@@ -220,14 +219,15 @@ class JinaEmbeddings(EmbeddingFunction):
|
||||
|
||||
return [result["embedding"] for result in sorted_embeddings]
|
||||
|
||||
def _init_client(self):
|
||||
@cached_property
|
||||
def _session(self):
|
||||
import requests
|
||||
|
||||
if JinaEmbeddings._session is None:
|
||||
if self.api_key is None and os.environ.get("JINA_API_KEY") is None:
|
||||
api_key_not_found_help("jina")
|
||||
api_key = self.api_key or os.environ.get("JINA_API_KEY")
|
||||
JinaEmbeddings._session = requests.Session()
|
||||
JinaEmbeddings._session.headers.update(
|
||||
{"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"}
|
||||
)
|
||||
if self.api_key is None and os.environ.get("JINA_API_KEY") is None:
|
||||
api_key_not_found_help("jina")
|
||||
api_key = self.api_key or os.environ.get("JINA_API_KEY")
|
||||
session = requests.Session()
|
||||
session.headers.update(
|
||||
{"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"}
|
||||
)
|
||||
return session
|
||||
|
||||
@@ -70,6 +70,8 @@ class OllamaEmbeddings(TextEmbeddingFunction):
|
||||
The texts to embed
|
||||
"""
|
||||
# TODO retry, rate limit, token limit
|
||||
if isinstance(texts, np.ndarray):
|
||||
texts = texts.tolist()
|
||||
embeddings = self._compute_embedding(texts)
|
||||
return list(embeddings)
|
||||
|
||||
|
||||
@@ -69,11 +69,15 @@ class EmbeddingFunctionRegistry:
|
||||
def decorator(cls):
|
||||
if not issubclass(cls, EmbeddingFunction):
|
||||
raise TypeError("Must be a subclass of EmbeddingFunction")
|
||||
if cls.__name__ in self._functions:
|
||||
raise KeyError(f"{cls.__name__} was already registered")
|
||||
key = alias or cls.__name__
|
||||
if key in self._functions and (
|
||||
not alias or self._functions[key] is not cls
|
||||
):
|
||||
raise KeyError(f"{key} was already registered")
|
||||
self._functions[key] = cls
|
||||
cls.__embedding_function_registry_alias__ = alias
|
||||
# Store the name the class is registered under, so the table
|
||||
# metadata written for it can be looked up again when reading.
|
||||
cls.__embedding_function_registry_alias__ = key
|
||||
return cls
|
||||
|
||||
return decorator
|
||||
|
||||
@@ -241,15 +241,30 @@ def retry_with_exponential_backoff(
|
||||
exponential_base: float = 2,
|
||||
jitter: bool = True,
|
||||
max_retries: int = 7,
|
||||
max_delay: float = 60,
|
||||
):
|
||||
"""Retry a function with exponential backoff.
|
||||
|
||||
Only errors that may succeed when repeated are retried. Errors that carry an
|
||||
HTTP status code are retried when the status is 408, 409, 429 or >= 500; any
|
||||
other 4xx status (for example 400, 401, 403, 404 or 422) is a permanent failure
|
||||
that is raised immediately. Errors without a recognisable status code (for
|
||||
example connection errors, timeouts, or exceptions raised by a custom embedding
|
||||
function) are retried.
|
||||
|
||||
The delay before the n-th retry is ``initial_delay * exponential_base ** n``,
|
||||
never more than ``max_delay``. With ``jitter`` enabled, a random amount of up to
|
||||
half of that delay is removed ("equal jitter"), so the actual wait is always in
|
||||
the range ``[delay / 2, delay]``.
|
||||
|
||||
Args:
|
||||
func (function): The function to be retried.
|
||||
initial_delay (float): Initial delay in seconds (default is 1).
|
||||
exponential_base (float): The base for exponential backoff (default is 2).
|
||||
jitter (bool): Whether to add jitter to the delay (default is True).
|
||||
max_retries (int): Maximum number of retries (default is 7).
|
||||
max_delay (float): Maximum delay between two attempts, in seconds
|
||||
(default is 60).
|
||||
|
||||
Returns:
|
||||
function: The decorated function.
|
||||
@@ -265,13 +280,12 @@ def retry_with_exponential_backoff(
|
||||
try:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
# Currently retrying on all exceptions as there is no way to know the
|
||||
# format of the error msgs used by different APIs. We'll log the error
|
||||
# and say that it is assumed that if this portion errors out, it's due
|
||||
# to rate limit but the user should check the error message to be sure.
|
||||
# Errors are retried unless they are known to be permanent (see
|
||||
# `_is_non_retryable_error`). Custom embedding functions can raise
|
||||
# arbitrary exceptions, so unknown errors are assumed to be transient.
|
||||
except Exception as e: # noqa: PERF203
|
||||
# Don't retry on authentication errors (e.g., OpenAI 401)
|
||||
# These are permanent failures that won't be fixed by retrying
|
||||
# Don't retry permanent failures (e.g., OpenAI 401 or 404). They
|
||||
# won't be fixed by retrying and would only delay the real error.
|
||||
if _is_non_retryable_error(e):
|
||||
raise
|
||||
|
||||
@@ -280,24 +294,72 @@ def retry_with_exponential_backoff(
|
||||
if num_retries > max_retries:
|
||||
raise Exception(
|
||||
f"Maximum number of retries ({max_retries}) exceeded.", e
|
||||
)
|
||||
) from e
|
||||
|
||||
delay *= exponential_base * (1 + jitter * random.random())
|
||||
delay = min(delay * exponential_base, max_delay)
|
||||
sleep_for = delay
|
||||
if jitter:
|
||||
sleep_for = delay / 2 + random.random() * delay / 2
|
||||
logging.warning(
|
||||
"Error occurred: %s \n Retrying in %s seconds (retry %s of %s) \n",
|
||||
e,
|
||||
delay,
|
||||
sleep_for,
|
||||
num_retries,
|
||||
max_retries,
|
||||
)
|
||||
time.sleep(delay)
|
||||
time.sleep(sleep_for)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
# HTTP statuses in the 4xx range that are worth retrying: request timeout,
|
||||
# conflict (the OpenAI SDK treats it as transient) and rate limiting.
|
||||
_RETRYABLE_CLIENT_STATUSES = (408, 409, 429)
|
||||
|
||||
# AWS error codes that are reported with a 4xx status but are transient.
|
||||
_RETRYABLE_AWS_ERROR_CODES = (
|
||||
"Throttl",
|
||||
"TooManyRequests",
|
||||
"RequestLimitExceeded",
|
||||
"ProvisionedThroughputExceeded",
|
||||
"ModelNotReady",
|
||||
"ModelTimeout",
|
||||
)
|
||||
|
||||
|
||||
def _get_http_status(error: Exception):
|
||||
"""Best-effort extraction of an HTTP status code from an exception.
|
||||
|
||||
Provider SDKs expose it in different ways: ``status_code`` (openai, cohere),
|
||||
``http_status`` (voyageai), ``response.status_code`` (requests, httpx),
|
||||
``code`` (google-api-core, urllib) or ``response["ResponseMetadata"]``
|
||||
(botocore). Returns None if no status can be found.
|
||||
"""
|
||||
candidates = [
|
||||
getattr(error, "status_code", None),
|
||||
getattr(error, "http_status", None),
|
||||
]
|
||||
response = getattr(error, "response", None)
|
||||
if isinstance(response, dict):
|
||||
candidates.append(response.get("ResponseMetadata", {}).get("HTTPStatusCode"))
|
||||
else:
|
||||
candidates.append(getattr(response, "status_code", None))
|
||||
candidates.append(getattr(error, "code", None))
|
||||
|
||||
for status in candidates:
|
||||
if isinstance(status, int) and not isinstance(status, bool):
|
||||
if 100 <= status <= 599:
|
||||
return status
|
||||
return None
|
||||
|
||||
|
||||
def _is_non_retryable_error(error: Exception) -> bool:
|
||||
"""Check if an error should not be retried.
|
||||
|
||||
An error is non-retryable if it is an authentication error or if it carries an
|
||||
HTTP client-error status (4xx) other than 408, 409 and 429. Errors without a
|
||||
recognisable status are treated as retryable.
|
||||
|
||||
Args:
|
||||
error: The exception to check
|
||||
|
||||
@@ -314,13 +376,25 @@ def _is_non_retryable_error(error: Exception) -> bool:
|
||||
if error_type == "AuthenticationError":
|
||||
return True
|
||||
|
||||
# Check for other common non-retryable HTTP status codes
|
||||
# 401 Unauthorized, 403 Forbidden
|
||||
if hasattr(error, "status_code"):
|
||||
if error.status_code in (401, 403):
|
||||
return True
|
||||
# A missing API key or a missing optional dependency fails the same way on
|
||||
# every attempt
|
||||
if isinstance(error, (MissingApiKeyError, ImportError)):
|
||||
return True
|
||||
|
||||
return False
|
||||
status = _get_http_status(error)
|
||||
if status is None or not 400 <= status < 500:
|
||||
return False
|
||||
if status in _RETRYABLE_CLIENT_STATUSES:
|
||||
return False
|
||||
|
||||
# AWS reports throttling with a 400 status; the error code tells them apart.
|
||||
response = getattr(error, "response", None)
|
||||
if isinstance(response, dict):
|
||||
aws_code = str(response.get("Error", {}).get("Code", ""))
|
||||
if aws_code.startswith(_RETRYABLE_AWS_ERROR_CODES):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def url_retrieve(url: str):
|
||||
@@ -337,9 +411,15 @@ def url_retrieve(url: str):
|
||||
raise ConnectionError("could not download {} due to {}".format(url, err))
|
||||
|
||||
|
||||
class MissingApiKeyError(ValueError):
|
||||
"""Raised when an embedding provider's API key is not set."""
|
||||
|
||||
|
||||
def api_key_not_found_help(provider):
|
||||
logging.error("Could not find API key for %s", provider)
|
||||
raise ValueError(f"Please set the {provider.upper()}_API_KEY environment variable.")
|
||||
raise MissingApiKeyError(
|
||||
f"Please set the {provider.upper()}_API_KEY environment variable."
|
||||
)
|
||||
|
||||
|
||||
def is_flash_attn_2_available():
|
||||
|
||||
@@ -75,7 +75,7 @@ def transform_input(input_data: Union[str, bytes, Path]):
|
||||
content = {"type": "text", "text": input_data}
|
||||
elif isinstance(input_data, PIL_Image.Image):
|
||||
buffered = BytesIO()
|
||||
input_data.save(buffered, format="JPEG")
|
||||
input_data.convert("RGB").save(buffered, format="JPEG")
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||||
content = {
|
||||
"type": "image_base64",
|
||||
@@ -84,7 +84,7 @@ def transform_input(input_data: Union[str, bytes, Path]):
|
||||
elif isinstance(input_data, bytes):
|
||||
img = PIL_Image.open(BytesIO(input_data))
|
||||
buffered = BytesIO()
|
||||
img.save(buffered, format="JPEG")
|
||||
img.convert("RGB").save(buffered, format="JPEG")
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||||
content = {
|
||||
"type": "image_base64",
|
||||
@@ -103,7 +103,7 @@ def transform_input(input_data: Union[str, bytes, Path]):
|
||||
else:
|
||||
img = PIL_Image.open(input_data)
|
||||
buffered = BytesIO()
|
||||
img.save(buffered, format="JPEG")
|
||||
img.convert("RGB").save(buffered, format="JPEG")
|
||||
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
|
||||
content = {
|
||||
"type": "image_base64",
|
||||
@@ -145,6 +145,8 @@ def sanitize_text_input(inputs: TEXT) -> List[str]:
|
||||
"""
|
||||
if isinstance(inputs, str):
|
||||
inputs = [inputs]
|
||||
elif isinstance(inputs, (list, np.ndarray)):
|
||||
inputs = list(inputs)
|
||||
elif isinstance(inputs, pa.Array):
|
||||
inputs = inputs.to_pylist()
|
||||
elif isinstance(inputs, pa.ChunkedArray):
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -52,6 +52,9 @@ class PyarrowScannerAdapter(pa.dataset.Scanner):
|
||||
raise NotImplementedError("memory_pool not supported")
|
||||
|
||||
def count_rows(self):
|
||||
if isinstance(self.filter, pa.compute.Expression):
|
||||
# stream the batches so memory does not grow with the matching rows
|
||||
return sum(batch.num_rows for batch in self.to_reader())
|
||||
return self.table.count_rows(self.filter)
|
||||
|
||||
def from_batches(self, **kwargs):
|
||||
@@ -83,6 +86,8 @@ class PyarrowScannerAdapter(pa.dataset.Scanner):
|
||||
return self.to_reader().read_all()
|
||||
|
||||
def to_reader(self, *, limit: Optional[int] = None):
|
||||
if isinstance(self.filter, pa.compute.Expression):
|
||||
return self._to_reader_with_expression(limit=limit)
|
||||
query = self.table.search()
|
||||
# Disable the builtin limit
|
||||
if limit is None:
|
||||
@@ -98,13 +103,30 @@ class PyarrowScannerAdapter(pa.dataset.Scanner):
|
||||
query = query.where(self.filter, prefilter=True)
|
||||
return query.to_batches(batch_size=self.batch_size)
|
||||
|
||||
def _to_reader_with_expression(self, *, limit: Optional[int] = None):
|
||||
# LanceDB filters are SQL strings, but engines such as DuckDB push down
|
||||
# pyarrow compute expressions. Scan the table and apply the filter and
|
||||
# projection with pyarrow, since the filter may reference columns that
|
||||
# are not part of the projection.
|
||||
if limit is not None and limit <= 0:
|
||||
raise ValueError("limit must be positive")
|
||||
query = self.table.search().limit(self.table.count_rows())
|
||||
scanner = pa.dataset.Scanner.from_batches(
|
||||
query.to_batches(batch_size=self.batch_size),
|
||||
columns=self.columns,
|
||||
filter=self.filter,
|
||||
)
|
||||
if limit is None:
|
||||
return scanner.to_reader()
|
||||
return scanner.head(limit).to_reader()
|
||||
|
||||
|
||||
class PyarrowDatasetAdapter(pa.dataset.Dataset):
|
||||
def __init__(self, table: Table):
|
||||
self.table = table
|
||||
|
||||
def count_rows(self, filter: Optional[Filter] = None):
|
||||
return self.table.count_rows(filter)
|
||||
return self.scanner(filter=filter).count_rows()
|
||||
|
||||
def get_fragments(self, filter: Optional[Filter] = None):
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -11,6 +11,7 @@ through a namespace abstraction.
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import warnings
|
||||
from typing import Any, Dict, Iterable, List, Optional, Union
|
||||
from uuid import UUID
|
||||
|
||||
@@ -51,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 (
|
||||
@@ -60,8 +68,6 @@ from lance_namespace import (
|
||||
CreateNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
NamespaceExistsRequest,
|
||||
TableExistsRequest,
|
||||
)
|
||||
@@ -82,6 +88,31 @@ from ._lancedb import Session
|
||||
_MAX_QUERY_K = 2**31 - 1
|
||||
|
||||
|
||||
def _create_table_storage_options(
|
||||
storage_options: Optional[Dict[str, str]],
|
||||
data_storage_version: Optional[str],
|
||||
enable_v2_manifest_paths: Optional[bool],
|
||||
) -> Optional[Dict[str, str]]:
|
||||
if data_storage_version is None and enable_v2_manifest_paths is None:
|
||||
return storage_options
|
||||
options = dict(storage_options or {})
|
||||
for name, value in (
|
||||
("data_storage_version", data_storage_version),
|
||||
("enable_v2_manifest_paths", enable_v2_manifest_paths),
|
||||
):
|
||||
if value is not None:
|
||||
warnings.warn(
|
||||
f"setting {name} directly on create_table is deprecated. "
|
||||
"Use storage_options instead.",
|
||||
DeprecationWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
options[f"new_table_{name}"] = (
|
||||
str(value).lower() if isinstance(value, bool) else value
|
||||
)
|
||||
return options
|
||||
|
||||
|
||||
def _query_to_namespace_request(
|
||||
table_id: List[str],
|
||||
query: "Query",
|
||||
@@ -587,6 +618,9 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
if mode.lower() not in ["create", "overwrite"]:
|
||||
raise ValueError("mode must be either 'create' or 'overwrite'")
|
||||
validate_table_name(name)
|
||||
storage_options = _create_table_storage_options(
|
||||
storage_options, data_storage_version, enable_v2_manifest_paths
|
||||
)
|
||||
async_table = LOOP.run(
|
||||
self._inner.create_table(
|
||||
name,
|
||||
@@ -825,7 +859,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.
|
||||
|
||||
@@ -842,12 +876,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
|
||||
)
|
||||
@@ -994,7 +1028,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.
|
||||
|
||||
@@ -1011,12 +1045,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
|
||||
)
|
||||
@@ -1230,6 +1264,9 @@ class AsyncLanceNamespaceDBConnection:
|
||||
if mode.lower() not in ["create", "overwrite"]:
|
||||
raise ValueError("mode must be either 'create' or 'overwrite'")
|
||||
validate_table_name(name)
|
||||
storage_options = _create_table_storage_options(
|
||||
storage_options, data_storage_version, enable_v2_manifest_paths
|
||||
)
|
||||
table = await self._inner.create_table(
|
||||
name,
|
||||
data,
|
||||
@@ -1427,12 +1464,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.
|
||||
|
||||
@@ -1449,12 +1486,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
|
||||
)
|
||||
|
||||
@@ -1584,12 +1621,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.
|
||||
|
||||
@@ -1606,12 +1643,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
|
||||
)
|
||||
|
||||
@@ -1646,7 +1683,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)
|
||||
|
||||
|
||||
@@ -909,9 +909,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":
|
||||
|
||||
@@ -127,7 +127,11 @@ def Vector(
|
||||
core_schema.list_schema(
|
||||
min_length=dim,
|
||||
max_length=dim,
|
||||
items_schema=core_schema.float_schema(),
|
||||
items_schema=(
|
||||
core_schema.int_schema()
|
||||
if pa.types.is_integer(value_type)
|
||||
else core_schema.float_schema()
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
@@ -208,7 +212,11 @@ def MultiVector(
|
||||
items_schema=core_schema.list_schema(
|
||||
min_length=dim,
|
||||
max_length=dim,
|
||||
items_schema=core_schema.float_schema(),
|
||||
items_schema=(
|
||||
core_schema.int_schema()
|
||||
if pa.types.is_integer(value_type)
|
||||
else core_schema.float_schema()
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -260,7 +268,9 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType:
|
||||
|
||||
def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]:
|
||||
return [
|
||||
_pydantic_to_field(name, field) for name, field in model.model_fields.items()
|
||||
_pydantic_to_field(name, field)
|
||||
for name, field in model.model_fields.items()
|
||||
if field.exclude is not True
|
||||
]
|
||||
|
||||
|
||||
@@ -493,6 +503,9 @@ class LanceModel(pydantic.BaseModel):
|
||||
|
||||
vec_and_function = []
|
||||
for name, field_info in cls.safe_get_fields().items():
|
||||
if field_info.exclude is True:
|
||||
continue
|
||||
|
||||
func = get_extras(field_info, "vector_column_for")
|
||||
if func is not None:
|
||||
vec_and_function.append([name, func])
|
||||
@@ -500,6 +513,8 @@ class LanceModel(pydantic.BaseModel):
|
||||
configs = []
|
||||
for vec, func in vec_and_function:
|
||||
for source, field_info in cls.safe_get_fields().items():
|
||||
if field_info.exclude is True:
|
||||
continue
|
||||
src_func = get_extras(field_info, "source_column_for")
|
||||
if src_func is func:
|
||||
# note we can't use == here since the function is a pydantic
|
||||
@@ -524,4 +539,20 @@ def get_extras(field_info: FieldInfo, key: str) -> Any:
|
||||
|
||||
def model_to_dict(model: pydantic.BaseModel) -> dict[str, Any]:
|
||||
"""Convert a Pydantic model to a dictionary."""
|
||||
return model.model_dump()
|
||||
return _enum_members_to_values(model.model_dump())
|
||||
|
||||
|
||||
def _enum_members_to_values(value: Any) -> Any:
|
||||
# model_dump() keeps Enum members, and Arrow cannot convert a member of a
|
||||
# plain Enum (one without a str or int mixin). Write its value instead,
|
||||
# which is what the Arrow type of the field is derived from.
|
||||
if isinstance(value, Enum):
|
||||
return value.value
|
||||
if isinstance(value, dict):
|
||||
return {k: _enum_members_to_values(v) for k, v in value.items()}
|
||||
if isinstance(value, (list, tuple)) and value:
|
||||
# Skip lists of plain scalars, such as vectors, without walking them.
|
||||
if isinstance(value[0], (int, float, str, bytes)):
|
||||
return value
|
||||
return [_enum_members_to_values(v) for v in value]
|
||||
return value
|
||||
|
||||
@@ -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=..., ...
|
||||
<BLANKLINE>
|
||||
|
||||
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=...]
|
||||
<BLANKLINE>
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -2575,10 +2577,6 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
with_row_ids: bool,
|
||||
offset: Optional[int] = None,
|
||||
) -> pa.Table:
|
||||
if norm == "rank":
|
||||
vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance")
|
||||
fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score")
|
||||
|
||||
# If both result sets are empty (e.g. after hard filtering),
|
||||
# return early to avoid errors in reranking or score restoration.
|
||||
if vector_results.num_rows == 0 and fts_results.num_rows == 0:
|
||||
@@ -2605,15 +2603,23 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
original_score_row_ids = None
|
||||
# normalize the scores to be between 0 and 1, 0 being most relevant
|
||||
# We check whether the results (vector and FTS) are empty, because when
|
||||
# they are, they often are missing the _rowid column, which causes an error
|
||||
# they are, they often are missing the _rowid column, which causes an error.
|
||||
# The original values are captured before any rank conversion so that
|
||||
# they can be restored after reranking.
|
||||
if vector_results.num_rows > 0:
|
||||
distance_i = vector_results.column_names.index("_distance")
|
||||
original_distances = vector_results.column(distance_i)
|
||||
original_distance_row_ids = vector_results.column("_rowid")
|
||||
if norm == "rank":
|
||||
vector_results = LanceHybridQueryBuilder._rank(
|
||||
vector_results, "_distance"
|
||||
)
|
||||
vector_results = vector_results.set_column(
|
||||
distance_i,
|
||||
vector_results.field(distance_i),
|
||||
LanceHybridQueryBuilder._normalize_scores(original_distances),
|
||||
LanceHybridQueryBuilder._normalize_scores(
|
||||
vector_results.column(distance_i)
|
||||
),
|
||||
)
|
||||
|
||||
# In fts higher scores represent relevance. Not inverting them here as
|
||||
@@ -2622,10 +2628,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
score_i = fts_results.column_names.index("_score")
|
||||
original_scores = fts_results.column(score_i)
|
||||
original_score_row_ids = fts_results.column("_rowid")
|
||||
if norm == "rank":
|
||||
fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score")
|
||||
fts_results = fts_results.set_column(
|
||||
score_i,
|
||||
fts_results.field(score_i),
|
||||
LanceHybridQueryBuilder._normalize_scores(original_scores),
|
||||
LanceHybridQueryBuilder._normalize_scores(fts_results.column(score_i)),
|
||||
)
|
||||
|
||||
results = reranker.rerank_hybrid(fts_query, vector_results, fts_results)
|
||||
@@ -2953,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=..., ...
|
||||
<BLANKLINE>
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -3595,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=..., ...
|
||||
<BLANKLINE>
|
||||
<BLANKLINE>
|
||||
|
||||
Parameters
|
||||
@@ -4515,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
|
||||
@@ -4871,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=..., ...
|
||||
<BLANKLINE>
|
||||
<BLANKLINE>
|
||||
|
||||
Parameters
|
||||
|
||||
@@ -36,7 +36,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
|
||||
@@ -54,8 +54,6 @@ from lance_namespace import (
|
||||
CreateNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
)
|
||||
from ..pydantic import LanceModel
|
||||
from ..table import Table
|
||||
@@ -268,7 +266,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 +282,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 +377,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 +394,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
|
||||
)
|
||||
@@ -843,6 +841,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 +1134,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(
|
||||
|
||||
@@ -7,6 +7,7 @@ import logging
|
||||
from functools import cached_property
|
||||
import os
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
@@ -83,6 +84,10 @@ from ..table import (
|
||||
from ..types import BaseTokenizerType
|
||||
from ..util import infer_vector_column_name
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import lance
|
||||
import polars as pl
|
||||
|
||||
|
||||
class RemoteTable(Table):
|
||||
def __init__(
|
||||
@@ -244,12 +249,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, **kwargs) -> "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))
|
||||
@@ -280,8 +305,12 @@ class RemoteTable(Table):
|
||||
"""Tokenize a query using the tokenizer configured on an FTS index.
|
||||
|
||||
Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are
|
||||
rebuilt in the client process from index metadata, so the same tokenizer
|
||||
model files must exist locally.
|
||||
rebuilt in the client process from index metadata. For remote tables,
|
||||
this means the same tokenizer model files must also exist locally
|
||||
(under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the
|
||||
server uses. A mismatch cannot be detected: the returned tokens may
|
||||
differ from what the server's index produces. If the model files are
|
||||
missing, an error naming the tokenizer is raised.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._table.tokenize(query, column=column, index_name=index_name)
|
||||
@@ -1152,11 +1181,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())
|
||||
@@ -1184,6 +1223,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.
|
||||
|
||||
@@ -62,7 +62,7 @@ class CohereReranker(Reranker):
|
||||
"COHERE_API_KEY not set. Either set it in your environment or \
|
||||
pass it as `api_key` argument to the CohereReranker."
|
||||
)
|
||||
return cohere.Client(os.environ.get("COHERE_API_KEY") or self.api_key)
|
||||
return cohere.Client(self.api_key or os.environ.get("COHERE_API_KEY"))
|
||||
|
||||
def _rerank(self, result_set: pa.Table, query: str):
|
||||
result_set = self._handle_empty_results(result_set)
|
||||
|
||||
@@ -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",
|
||||
@@ -70,15 +69,18 @@ class OpenaiReranker(Reranker):
|
||||
],
|
||||
)
|
||||
results = json.loads(response.choices[0].message.content)["documents"]
|
||||
docs, scores = list(
|
||||
zip(*[(result["content"], result["relevance_score"]) for result in results])
|
||||
) # tuples
|
||||
# replace the self.column column with the docs
|
||||
result_set = result_set.drop(self.column)
|
||||
result_set = result_set.append_column(
|
||||
self.column, pa.array(docs, type=pa.string())
|
||||
)
|
||||
# add the scores
|
||||
# The model returns the documents in its own order, so match each score
|
||||
# back to its row by content.
|
||||
scores_by_content = {
|
||||
result["content"]: result["relevance_score"] for result in results
|
||||
}
|
||||
missing = sum(doc not in scores_by_content for doc in docs)
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"The model returned no relevance_score for {missing} of "
|
||||
f"{len(docs)} documents"
|
||||
)
|
||||
scores = [scores_by_content[doc] for doc in docs]
|
||||
result_set = result_set.append_column(
|
||||
"_relevance_score", pa.array(scores, type=pa.float32())
|
||||
)
|
||||
@@ -129,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"))
|
||||
|
||||
@@ -64,7 +64,7 @@ class VoyageAIReranker(Reranker):
|
||||
pass it as `api_key` argument to the VoyageAIReranker."
|
||||
)
|
||||
return voyageai.Client(
|
||||
api_key=os.environ.get("VOYAGE_API_KEY") or self.api_key,
|
||||
api_key=self.api_key or os.environ.get("VOYAGE_API_KEY"),
|
||||
)
|
||||
|
||||
def _rerank(self, result_set: pa.Table, query: str):
|
||||
|
||||
@@ -276,6 +276,14 @@ def _from_iterable(data: Iterator) -> Scannable:
|
||||
_registered_modules: set[str] = set()
|
||||
|
||||
|
||||
def _hf_batches(data):
|
||||
if data._indices is None:
|
||||
yield from data.data.to_batches()
|
||||
else:
|
||||
for table in data.with_format("arrow").iter(batch_size=1000):
|
||||
yield from table.to_batches()
|
||||
|
||||
|
||||
def _register_optional_converters():
|
||||
"""Register converters for optional dependencies that are already imported."""
|
||||
|
||||
@@ -323,7 +331,15 @@ def _register_optional_converters():
|
||||
def _from_hf_dataset(data: HFDataset) -> Scannable:
|
||||
table = data.data.table # Access underlying Arrow table
|
||||
return Scannable(
|
||||
schema=table.schema, num_rows=len(data), reader=table.to_reader
|
||||
schema=table.schema,
|
||||
num_rows=len(data),
|
||||
reader=(
|
||||
table.to_reader
|
||||
if data._indices is None
|
||||
else lambda: pa.RecordBatchReader.from_batches(
|
||||
table.schema, _hf_batches(data)
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
@to_scannable.register(HFDatasetDict)
|
||||
@@ -335,7 +351,7 @@ def _register_optional_converters():
|
||||
|
||||
def gen():
|
||||
for split_name, dataset in data.items():
|
||||
for batch in dataset.data.to_batches():
|
||||
for batch in _hf_batches(dataset):
|
||||
split_arr = pa.array(
|
||||
[split_name] * len(batch), type=pa.string()
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
+190
-45
@@ -15,9 +15,11 @@ from functools import cached_property
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Dict,
|
||||
Iterable,
|
||||
Iterator,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
@@ -29,11 +31,12 @@ from typing import (
|
||||
)
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from lancedb.scannable import _register_optional_converters, to_scannable
|
||||
from lancedb.scannable import _hf_batches, _register_optional_converters, to_scannable
|
||||
|
||||
from . import __version__
|
||||
from ._blob import (
|
||||
BlobFile,
|
||||
_normalize_blob_descriptor,
|
||||
_normalize_blob_row_ids,
|
||||
_wrap_blob_files,
|
||||
strip_auto_row_ids,
|
||||
@@ -486,7 +489,7 @@ def _into_pyarrow_reader(
|
||||
if _check_for_hugging_face(data):
|
||||
if isinstance(data, datasets.Dataset):
|
||||
schema = data.features.arrow_schema
|
||||
return pa.RecordBatchReader.from_batches(schema, data.data.to_batches())
|
||||
return pa.RecordBatchReader.from_batches(schema, _hf_batches(data))
|
||||
elif isinstance(data, datasets.dataset_dict.DatasetDict):
|
||||
schema = _schema_from_hf(data, None)
|
||||
if "split" not in schema.names:
|
||||
@@ -555,7 +558,7 @@ def _into_pyarrow_reader(
|
||||
):
|
||||
return data.collect().to_arrow().to_reader()
|
||||
elif isinstance(data, Iterable):
|
||||
return _iterator_to_reader(data)
|
||||
return _iterator_to_reader(data, schema)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Unknown data type {type(data)}. "
|
||||
@@ -565,10 +568,19 @@ def _into_pyarrow_reader(
|
||||
)
|
||||
|
||||
|
||||
def _iterator_to_reader(data: Iterable) -> pa.RecordBatchReader:
|
||||
def _iterator_to_reader(
|
||||
data: Iterable, schema: Optional[pa.Schema] = None
|
||||
) -> pa.RecordBatchReader:
|
||||
# Each batch is treated as it's own reader, mainly so we can
|
||||
# re-use the _into_pyarrow_reader logic.
|
||||
first = _into_pyarrow_reader(next(data))
|
||||
data = iter(data)
|
||||
try:
|
||||
first_batch = next(data)
|
||||
except StopIteration:
|
||||
if schema is None:
|
||||
raise ValueError("Cannot create table from empty iterable without a schema")
|
||||
return pa.RecordBatchReader.from_batches(schema, [])
|
||||
first = _into_pyarrow_reader(first_batch)
|
||||
schema = first.schema
|
||||
|
||||
def gen():
|
||||
@@ -1281,7 +1293,7 @@ def _to_batches_with_split(data):
|
||||
with an extra `split` column
|
||||
"""
|
||||
for key, dataset in data.items():
|
||||
for batch in dataset.data.to_batches():
|
||||
for batch in _hf_batches(dataset):
|
||||
table = pa.Table.from_batches([batch])
|
||||
if "split" not in table.column_names:
|
||||
table = table.append_column(
|
||||
@@ -1433,13 +1445,13 @@ 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
|
||||
@abstractmethod
|
||||
def version(self) -> int:
|
||||
"""The version of this Table"""
|
||||
raise NotImplementedError
|
||||
raise NotImplementedError("version is not supported for this table type")
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
@@ -1448,7 +1460,7 @@ class Table(ABC):
|
||||
of this Table
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
raise NotImplementedError("schema is not supported for this table type")
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
@@ -1477,7 +1489,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":
|
||||
@@ -1486,11 +1498,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"""
|
||||
@@ -1513,7 +1525,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.
|
||||
@@ -1548,20 +1560,26 @@ 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":
|
||||
"""Return the table as a Polars LazyFrame.
|
||||
|
||||
For remote tables, use ``table.search().to_polars()`` to convert query
|
||||
results to a Polars DataFrame.
|
||||
|
||||
Note
|
||||
----
|
||||
The Polars streaming engine is not supported because it does not currently
|
||||
@@ -1572,7 +1590,7 @@ class Table(ABC):
|
||||
-------
|
||||
polars.LazyFrame
|
||||
"""
|
||||
raise NotImplementedError
|
||||
raise NotImplementedError("to_polars is not supported for this table type")
|
||||
|
||||
# New unified API overload
|
||||
@overload
|
||||
@@ -1677,7 +1695,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,
|
||||
@@ -1695,7 +1713,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:
|
||||
"""
|
||||
@@ -1715,7 +1735,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)
|
||||
@@ -1732,14 +1752,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(
|
||||
@@ -1804,7 +1824,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,
|
||||
@@ -1930,7 +1952,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(
|
||||
@@ -2004,7 +2028,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:
|
||||
"""
|
||||
@@ -2156,7 +2180,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(
|
||||
@@ -2180,11 +2204,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:
|
||||
@@ -2265,11 +2292,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
|
||||
@@ -2349,6 +2378,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,
|
||||
@@ -2437,7 +2504,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(
|
||||
@@ -2504,7 +2571,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(
|
||||
@@ -2650,7 +2717,11 @@ class Table(ABC):
|
||||
|
||||
Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are
|
||||
rebuilt in the client process from index metadata. For remote tables,
|
||||
this means the same tokenizer model files must also exist locally.
|
||||
this means the same tokenizer model files must also exist locally
|
||||
(under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the
|
||||
server uses. A mismatch cannot be detected: the returned tokens may
|
||||
differ from what the server's index produces. If the model files are
|
||||
missing, an error naming the tokenizer is raised.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
@@ -3261,11 +3332,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())
|
||||
@@ -3291,6 +3372,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.
|
||||
@@ -4912,7 +4998,11 @@ class LanceTable(Table):
|
||||
|
||||
Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are
|
||||
rebuilt in the client process from index metadata. For remote tables,
|
||||
this means the same tokenizer model files must also exist locally.
|
||||
this means the same tokenizer model files must also exist locally
|
||||
(under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the
|
||||
server uses. A mismatch cannot be detected: the returned tokens may
|
||||
differ from what the server's index produces. If the model files are
|
||||
missing, an error naming the tokenizer is raised.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._table.tokenize(query, column=column, index_name=index_name)
|
||||
@@ -5415,6 +5505,11 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray
|
||||
else:
|
||||
values_has_nan = pc.is_nan(values)
|
||||
values_indices = pc.list_parent_indices(arr)
|
||||
if arr.null_count:
|
||||
# Match flatten(), which excludes child spans under null parent lists.
|
||||
values_indices = pc.filter(
|
||||
values_indices, pc.take(pc.is_valid(arr), values_indices)
|
||||
)
|
||||
has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan))
|
||||
indices = pa.array(range(len(arr)), type=pa.uint32())
|
||||
return pc.is_in(indices, has_nan_indices)
|
||||
@@ -5503,6 +5598,7 @@ def _infer_target_schema(
|
||||
field.name, # preserve original field name
|
||||
target_type,
|
||||
nullable=field.nullable,
|
||||
metadata=field.metadata,
|
||||
)
|
||||
|
||||
schema = schema.set(i, new_field)
|
||||
@@ -7401,7 +7497,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.
|
||||
|
||||
@@ -7419,15 +7517,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.
|
||||
|
||||
@@ -7448,13 +7554,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()
|
||||
@@ -7485,6 +7596,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.
|
||||
@@ -7614,7 +7755,11 @@ class AsyncTable:
|
||||
|
||||
Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are
|
||||
rebuilt in the client process from index metadata. For remote tables,
|
||||
this means the same tokenizer model files must also exist locally.
|
||||
this means the same tokenizer model files must also exist locally
|
||||
(under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the
|
||||
server uses. A mismatch cannot be detected: the returned tokens may
|
||||
differ from what the server's index produces. If the model files are
|
||||
missing, an error naming the tokenizer is raised.
|
||||
"""
|
||||
return await self._inner.tokenize(query, column=column, index_name=index_name)
|
||||
|
||||
|
||||
@@ -426,7 +426,8 @@ def _(value: date):
|
||||
|
||||
|
||||
@value_to_sql.register(list)
|
||||
def _(value: list):
|
||||
@value_to_sql.register(tuple)
|
||||
def _(value: Union[list, tuple]):
|
||||
return "[" + ", ".join(map(value_to_sql, value)) + "]"
|
||||
|
||||
|
||||
@@ -451,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())
|
||||
@@ -471,6 +481,15 @@ def _(value: np.floating):
|
||||
return value_to_sql(float(value))
|
||||
|
||||
|
||||
@value_to_sql.register(np.datetime64)
|
||||
def _(value: np.datetime64):
|
||||
if np.isnat(value):
|
||||
return value_to_sql(None)
|
||||
if np.datetime_data(value.dtype)[0] in ("ps", "fs", "as"):
|
||||
raise ValueError("SQL timestamps cannot preserve sub-nanosecond precision")
|
||||
return value_to_sql(str(np.datetime_as_string(value, unit="ns", timezone="UTC")))
|
||||
|
||||
|
||||
def deprecated(func):
|
||||
"""This is a decorator which can be used to mark functions
|
||||
as deprecated. It will result in a warning being emitted
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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)
|
||||
@@ -0,0 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import pytest
|
||||
from lancedb.db import AsyncConnection
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["count_rows", "query", "head"])
|
||||
async def test_closed_table_queries_raise_runtime_error(
|
||||
mem_db_async: AsyncConnection, operation
|
||||
):
|
||||
table = await mem_db_async.create_table("closed_query", data=[{"id": 0}])
|
||||
assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}]
|
||||
table.close()
|
||||
table.close()
|
||||
assert not table.is_open()
|
||||
|
||||
with pytest.raises(RuntimeError, match="Table closed_query is closed"):
|
||||
if operation == "query":
|
||||
table.query()
|
||||
else:
|
||||
await getattr(table, operation)()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_open_table_queries(mem_db_async: AsyncConnection):
|
||||
table = await mem_db_async.create_table("open_query", data=[{"id": 0}])
|
||||
assert await table.count_rows() == 1
|
||||
assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}]
|
||||
assert (await table.head()).to_pylist() == [{"id": 0}]
|
||||
table.close()
|
||||
@@ -86,6 +86,7 @@ def test_sync_debugger_inspection_does_not_use_background_loop(tmp_path, monkeyp
|
||||
|
||||
assert members["uri"] == str(tmp_path)
|
||||
assert members["read_consistency_interval"] is None
|
||||
assert members["session"] is None
|
||||
assert repr(db) == f"LanceDBConnection(uri={str(tmp_path)!r})"
|
||||
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
|
||||
|
||||
@@ -105,6 +106,8 @@ def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkey
|
||||
|
||||
assert db.read_consistency_interval == consistency_interval
|
||||
assert db_from_inner.read_consistency_interval == consistency_interval
|
||||
assert db.session is None
|
||||
assert db_from_inner.session is None
|
||||
|
||||
|
||||
def test_ingest_pd(tmp_path):
|
||||
|
||||
@@ -22,3 +22,19 @@ def test_basic_query(tmp_path):
|
||||
|
||||
results = duck_conn.sql("SELECT SUM(y) FROM adapter").fetchall()
|
||||
assert results[0][0] == 26
|
||||
|
||||
|
||||
def test_query_with_filter(tmp_path):
|
||||
data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]})
|
||||
conn = lancedb.connect(tmp_path)
|
||||
tbl = conn.create_table("test", data)
|
||||
|
||||
adapter = PyarrowDatasetAdapter(tbl) # noqa: F841
|
||||
|
||||
duck_conn = duckdb.connect()
|
||||
|
||||
results = duck_conn.sql("SELECT SUM(y) FROM adapter WHERE x > 2").fetchall()
|
||||
assert results[0][0] == 15
|
||||
|
||||
results = duck_conn.sql("SELECT y FROM adapter WHERE x > 2 ORDER BY y").fetchall()
|
||||
assert results == [(7,), (8,)]
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import copy
|
||||
|
||||
import lancedb
|
||||
import pytest
|
||||
from lancedb.embeddings.base import TextEmbeddingFunction
|
||||
from lancedb.embeddings.registry import get_registry
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
|
||||
class LocalEmbedding(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 2
|
||||
|
||||
def generate_embeddings(self, texts):
|
||||
return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,tag",
|
||||
[("SourceField", "source_column_for"), ("VectorField", "vector_column_for")],
|
||||
)
|
||||
@pytest.mark.parametrize("extra", [None, {}, {"x-purpose": "retrieval"}])
|
||||
def test_embedding_field_preserves_extra_metadata(method, tag, extra):
|
||||
function = LocalEmbedding.create(max_retries=0)
|
||||
original = copy.deepcopy(extra)
|
||||
field = getattr(function, method)(json_schema_extra=extra)
|
||||
assert extra == original
|
||||
assert field.json_schema_extra[tag] is function
|
||||
for key, value in (extra or {}).items():
|
||||
assert field.json_schema_extra[key] == value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,tag",
|
||||
[("SourceField", "source_column_for"), ("VectorField", "vector_column_for")],
|
||||
)
|
||||
def test_extra_metadata_cannot_replace_embedding_identity(method, tag):
|
||||
function = LocalEmbedding.create(max_retries=0)
|
||||
supplied = {tag: "caller value", "x-purpose": "retrieval"}
|
||||
field = getattr(function, method)(json_schema_extra=supplied)
|
||||
assert supplied[tag] == "caller value"
|
||||
assert field.json_schema_extra[tag] is function
|
||||
assert field.json_schema_extra["x-purpose"] == "retrieval"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["SourceField", "VectorField"])
|
||||
def test_embedding_field_preserves_other_field_arguments(method):
|
||||
function = LocalEmbedding.create(max_retries=0)
|
||||
field = getattr(function, method)(title="Custom title", description="Custom field")
|
||||
assert field.title == "Custom title"
|
||||
assert field.description == "Custom field"
|
||||
|
||||
|
||||
def test_extra_metadata_preserves_schema_and_local_embedding(tmp_path, monkeypatch):
|
||||
registry = get_registry()
|
||||
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
|
||||
registry.register("test-embedding-field-extra")(LocalEmbedding)
|
||||
function = LocalEmbedding.create(max_retries=0)
|
||||
|
||||
class Row(LanceModel):
|
||||
text: str = function.SourceField(json_schema_extra={"x-purpose": "source"})
|
||||
vector: Vector(2) = function.VectorField(
|
||||
json_schema_extra={"x-purpose": "embedding"}
|
||||
)
|
||||
|
||||
json_schema = Row.model_json_schema()
|
||||
assert json_schema["properties"]["text"]["x-purpose"] == "source"
|
||||
assert json_schema["properties"]["vector"]["x-purpose"] == "embedding"
|
||||
configs = Row.parse_embedding_functions()
|
||||
assert len(configs) == 1
|
||||
assert configs[0].function is function
|
||||
assert configs[0].source_column == "text"
|
||||
assert configs[0].vector_column == "vector"
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table(
|
||||
"rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}]
|
||||
)
|
||||
assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
|
||||
reopened = lancedb.connect(tmp_path).open_table("rows")
|
||||
assert isinstance(reopened.embedding_functions["vector"].function, LocalEmbedding)
|
||||
assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
|
||||
@@ -0,0 +1,88 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import lancedb
|
||||
import pytest
|
||||
from lancedb.embeddings.base import TextEmbeddingFunction
|
||||
from lancedb.embeddings.registry import EmbeddingFunctionRegistry, get_registry
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
|
||||
class FirstEmbedding(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 2
|
||||
|
||||
def generate_embeddings(self, texts):
|
||||
return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts]
|
||||
|
||||
|
||||
class SecondEmbedding(FirstEmbedding):
|
||||
def generate_embeddings(self, texts):
|
||||
return [[0.0, 1.0] if text == "alpha" else [1.0, 0.0] for text in texts]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"first_alias,second_alias,second_class",
|
||||
[
|
||||
("shared", "shared", SecondEmbedding),
|
||||
(None, "FirstEmbedding", SecondEmbedding),
|
||||
("SecondEmbedding", None, SecondEmbedding),
|
||||
(None, None, FirstEmbedding),
|
||||
("", "", FirstEmbedding),
|
||||
],
|
||||
)
|
||||
def test_duplicate_registration_preserves_original(
|
||||
first_alias, second_alias, second_class
|
||||
):
|
||||
registry = EmbeddingFunctionRegistry()
|
||||
registry.register(first_alias)(FirstEmbedding)
|
||||
key = first_alias or "FirstEmbedding"
|
||||
original_alias = FirstEmbedding.__embedding_function_registry_alias__
|
||||
second_attributes = dict(second_class.__dict__)
|
||||
with pytest.raises(KeyError, match=f"{key} was already registered"):
|
||||
registry.register(second_alias)(second_class)
|
||||
assert registry.get(key) is FirstEmbedding
|
||||
assert FirstEmbedding.__embedding_function_registry_alias__ == original_alias
|
||||
assert dict(second_class.__dict__) == second_attributes
|
||||
|
||||
|
||||
def test_class_name_does_not_collide_with_distinct_alias():
|
||||
registry = EmbeddingFunctionRegistry()
|
||||
registry.register()(FirstEmbedding)
|
||||
same_name = type("FirstEmbedding", (SecondEmbedding,), {})
|
||||
registry.register("separate")(same_name)
|
||||
assert registry.get("FirstEmbedding") is FirstEmbedding
|
||||
assert registry.get("separate") is same_name
|
||||
|
||||
|
||||
def test_same_class_explicit_alias_registration_is_idempotent():
|
||||
registry = EmbeddingFunctionRegistry()
|
||||
assert registry.register("shared")(FirstEmbedding) is FirstEmbedding
|
||||
attributes = dict(FirstEmbedding.__dict__)
|
||||
assert registry.register("shared")(FirstEmbedding) is FirstEmbedding
|
||||
assert registry.get("shared") is FirstEmbedding
|
||||
assert dict(FirstEmbedding.__dict__) == attributes
|
||||
|
||||
|
||||
def test_alias_collision_does_not_change_persisted_table(tmp_path, monkeypatch):
|
||||
registry = get_registry()
|
||||
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
|
||||
alias = "test-alias-collision-persisted-table"
|
||||
registry.register(alias)(FirstEmbedding)
|
||||
function = FirstEmbedding.create(max_retries=0)
|
||||
|
||||
class Row(LanceModel):
|
||||
text: str = function.SourceField()
|
||||
vector: Vector(2) = function.VectorField()
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table(
|
||||
"rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}]
|
||||
)
|
||||
assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
|
||||
assert registry.register(alias)(FirstEmbedding) is FirstEmbedding
|
||||
with pytest.raises(KeyError, match=f"{alias} was already registered"):
|
||||
registry.register(alias)(SecondEmbedding)
|
||||
reopened = lancedb.connect(tmp_path).open_table("rows")
|
||||
assert isinstance(reopened.embedding_functions["vector"].function, FirstEmbedding)
|
||||
assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
|
||||
@@ -23,6 +23,12 @@ from lancedb.embeddings.utils import retry
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_embedding_registry(monkeypatch):
|
||||
registry = get_registry()
|
||||
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
|
||||
|
||||
|
||||
def mock_embed_func(input_data):
|
||||
return [np.random.randn(128).tolist() for _ in range(len(input_data))]
|
||||
|
||||
@@ -64,6 +70,37 @@ def test_embedding_function(tmp_path):
|
||||
assert np.allclose(actual, expected)
|
||||
|
||||
|
||||
def test_embedding_function_registered_without_alias(tmp_path):
|
||||
registry = EmbeddingFunctionRegistry.get_instance()
|
||||
|
||||
@registry.register()
|
||||
class NoAliasEmbeddingFunction(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 2
|
||||
|
||||
def generate_embeddings(self, texts):
|
||||
return [np.array([1.0, 2.0]) for _ in texts]
|
||||
|
||||
assert registry.get("NoAliasEmbeddingFunction") is NoAliasEmbeddingFunction
|
||||
|
||||
func = NoAliasEmbeddingFunction.create()
|
||||
|
||||
class Schema(LanceModel):
|
||||
text: str = func.SourceField()
|
||||
vector: Vector(2) = func.VectorField()
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
db.create_table("test", schema=Schema)
|
||||
|
||||
# The class name is written to the table metadata, so the function can be
|
||||
# resolved again when the table is reopened.
|
||||
table = lancedb.connect(tmp_path).open_table("test")
|
||||
table.add([{"text": "hello"}])
|
||||
conf = table.embedding_functions["vector"]
|
||||
assert isinstance(conf.function, NoAliasEmbeddingFunction)
|
||||
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0]]
|
||||
|
||||
|
||||
def test_instructor_ndims_uses_instruction():
|
||||
instructor = get_registry().get("instructor").create()
|
||||
model = MagicMock()
|
||||
@@ -596,6 +633,318 @@ def test_openai_no_retry_on_401(mock_sleep):
|
||||
assert mock_sleep.call_count == 0
|
||||
|
||||
|
||||
class _HttpStatusError(Exception):
|
||||
"""Stand-in for provider SDK errors (openai, cohere, ...) with a status code."""
|
||||
|
||||
def __init__(self, status_code, message="error"):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
class _ResponseStatusError(Exception):
|
||||
"""Stand-in for requests/httpx errors, which expose `response.status_code`."""
|
||||
|
||||
def __init__(self, status_code):
|
||||
super().__init__(f"HTTP {status_code}")
|
||||
self.response = MagicMock(status_code=status_code)
|
||||
|
||||
|
||||
def _run_with_retry(side_effect, **kwargs):
|
||||
"""Run a function through the retry helper with a mocked `time.sleep`.
|
||||
|
||||
Returns (mock_func, sleeps, error).
|
||||
"""
|
||||
from lancedb.embeddings.utils import retry_with_exponential_backoff
|
||||
|
||||
mock_func = MagicMock(side_effect=side_effect)
|
||||
sleeps = []
|
||||
error = None
|
||||
with patch("time.sleep", side_effect=sleeps.append):
|
||||
try:
|
||||
retry_with_exponential_backoff(mock_func, **kwargs)()
|
||||
except Exception as e:
|
||||
error = e
|
||||
return mock_func, sleeps, error
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 401, 403, 404, 405, 410, 413, 422])
|
||||
@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError])
|
||||
def test_retry_does_not_retry_permanent_http_errors(status, make_error):
|
||||
err = make_error(status)
|
||||
mock_func, sleeps, error = _run_with_retry(err, max_retries=7)
|
||||
|
||||
assert error is err
|
||||
assert mock_func.call_count == 1
|
||||
assert sleeps == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [408, 409, 429, 500, 502, 503, 504])
|
||||
@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError])
|
||||
def test_retry_retries_transient_http_errors(status, make_error):
|
||||
mock_func, sleeps, error = _run_with_retry(
|
||||
[make_error(status), make_error(status), "result"], max_retries=7
|
||||
)
|
||||
|
||||
assert error is None
|
||||
assert mock_func.call_count == 3
|
||||
assert len(sleeps) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"exc",
|
||||
[
|
||||
ConnectionError("connection reset"),
|
||||
TimeoutError("timed out"),
|
||||
OSError("network is unreachable"),
|
||||
ValueError("custom embedding function failure"),
|
||||
RuntimeError("boom"),
|
||||
Exception("plain exception"),
|
||||
],
|
||||
)
|
||||
def test_retry_retries_errors_without_http_status(exc):
|
||||
mock_func, sleeps, error = _run_with_retry([exc, exc, "result"], max_retries=7)
|
||||
|
||||
assert error is None
|
||||
assert mock_func.call_count == 3
|
||||
assert len(sleeps) == 2
|
||||
|
||||
|
||||
def test_retry_error_status_attributes_of_other_types_are_ignored():
|
||||
# `code` is a string for openai errors ("model_not_found"), and arbitrary
|
||||
# exceptions may carry unrelated numeric attributes. Neither is an HTTP status.
|
||||
class WeirdError(Exception):
|
||||
code = "model_not_found"
|
||||
status_code = None
|
||||
|
||||
class ExitCodeError(Exception):
|
||||
code = 2
|
||||
|
||||
for exc in (WeirdError("a"), ExitCodeError("b")):
|
||||
mock_func, sleeps, error = _run_with_retry(exc, max_retries=2)
|
||||
assert mock_func.call_count == 3
|
||||
assert len(sleeps) == 2
|
||||
|
||||
|
||||
def test_retry_botocore_style_errors():
|
||||
# botocore ClientError exposes `response["ResponseMetadata"]["HTTPStatusCode"]`.
|
||||
# Bedrock reports throttling with a 400, which must stay retryable.
|
||||
class ClientError(Exception):
|
||||
def __init__(self, status, code):
|
||||
super().__init__(code)
|
||||
self.response = {
|
||||
"ResponseMetadata": {"HTTPStatusCode": status},
|
||||
"Error": {"Code": code},
|
||||
}
|
||||
|
||||
throttled = ClientError(400, "ThrottlingException")
|
||||
mock_func, sleeps, error = _run_with_retry([throttled, "result"], max_retries=3)
|
||||
assert error is None
|
||||
assert mock_func.call_count == 2
|
||||
|
||||
invalid = ClientError(400, "ValidationException")
|
||||
mock_func, sleeps, error = _run_with_retry(invalid, max_retries=3)
|
||||
assert error is invalid
|
||||
assert mock_func.call_count == 1
|
||||
assert sleeps == []
|
||||
|
||||
|
||||
def test_retry_authentication_error_by_name_not_retried():
|
||||
class AuthenticationError(Exception):
|
||||
pass
|
||||
|
||||
mock_func, sleeps, error = _run_with_retry(AuthenticationError("bad key"))
|
||||
assert isinstance(error, AuthenticationError)
|
||||
assert mock_func.call_count == 1
|
||||
assert sleeps == []
|
||||
|
||||
|
||||
def test_retry_gives_up_after_max_retries_with_original_error_attached():
|
||||
original = _HttpStatusError(503, "overloaded")
|
||||
mock_func, sleeps, error = _run_with_retry(original, max_retries=3)
|
||||
|
||||
assert mock_func.call_count == 4
|
||||
assert len(sleeps) == 3
|
||||
assert "Maximum number of retries (3) exceeded" in str(error)
|
||||
assert error.args[1] is original
|
||||
assert error.__cause__ is original
|
||||
|
||||
|
||||
def test_retry_delay_is_capped_by_max_delay():
|
||||
# Default settings (7 retries, 60s cap): the old implementation could sleep for
|
||||
# up to ~47 minutes in a single retry and hours in total.
|
||||
exc = _HttpStatusError(503)
|
||||
for seed in range(50):
|
||||
import random
|
||||
|
||||
random.seed(seed)
|
||||
_, sleeps, _ = _run_with_retry(exc)
|
||||
assert len(sleeps) == 7
|
||||
assert max(sleeps) <= 60
|
||||
assert sum(sleeps) <= 7 * 60
|
||||
|
||||
# Custom cap
|
||||
_, sleeps, _ = _run_with_retry(exc, max_retries=10, max_delay=5)
|
||||
assert len(sleeps) == 10
|
||||
assert max(sleeps) <= 5
|
||||
|
||||
|
||||
def test_retry_delay_without_jitter_is_exact_exponential_then_capped():
|
||||
_, sleeps, _ = _run_with_retry(
|
||||
_HttpStatusError(500),
|
||||
initial_delay=1,
|
||||
exponential_base=2,
|
||||
jitter=False,
|
||||
max_retries=8,
|
||||
max_delay=30,
|
||||
)
|
||||
assert sleeps == [2, 4, 8, 16, 30, 30, 30, 30]
|
||||
|
||||
|
||||
def test_retry_jitter_is_bounded():
|
||||
# Each wait is within [delay / 2, delay], where delay grows by exponential_base
|
||||
# per retry (no random factor is compounded across retries).
|
||||
exc = _HttpStatusError(500)
|
||||
for seed in range(50):
|
||||
import random
|
||||
|
||||
random.seed(seed)
|
||||
_, sleeps, _ = _run_with_retry(exc, max_retries=8, max_delay=100)
|
||||
for i, slept in enumerate(sleeps):
|
||||
delay = min(2 ** (i + 1), 100)
|
||||
assert delay / 2 <= slept <= delay
|
||||
|
||||
# Extremes of the random source map to the ends of the range
|
||||
with patch("random.random", return_value=0.0):
|
||||
_, sleeps, _ = _run_with_retry(exc, max_retries=3)
|
||||
assert sleeps == [1, 2, 4]
|
||||
with patch("random.random", return_value=1.0):
|
||||
_, sleeps, _ = _run_with_retry(exc, max_retries=3)
|
||||
assert sleeps == [2, 4, 8]
|
||||
|
||||
|
||||
def test_retry_large_max_retries_does_not_overflow():
|
||||
_, sleeps, _ = _run_with_retry(_HttpStatusError(503), max_retries=1500)
|
||||
assert len(sleeps) == 1500
|
||||
assert max(sleeps) <= 60
|
||||
|
||||
|
||||
def _register_failing_embedding(name, error):
|
||||
"""Register an embedding function that always raises `error`.
|
||||
|
||||
Returns the class and a list that records every call. The function is
|
||||
re-created from table metadata, so the call count has to live outside of it.
|
||||
"""
|
||||
calls = []
|
||||
|
||||
@register(name)
|
||||
class FailingEmbeddingFunction(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 4
|
||||
|
||||
def generate_embeddings(self, texts):
|
||||
calls.append(len(texts))
|
||||
raise error
|
||||
|
||||
return FailingEmbeddingFunction, calls
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 404, 422])
|
||||
def test_permanent_embedding_error_is_not_retried(tmp_path, status):
|
||||
"""A permanent provider error (e.g. a misspelled model name) must surface
|
||||
immediately from both `table.add` and `table.search` instead of being retried
|
||||
with minutes of backoff."""
|
||||
error = _HttpStatusError(status, "The model `typo-embed` does not exist")
|
||||
cls, calls = _register_failing_embedding(f"failing-embedding-{status}", error)
|
||||
func = cls.create() # default max_retries=7
|
||||
|
||||
class Schema(LanceModel):
|
||||
text: str = func.SourceField()
|
||||
vector: Vector(func.ndims()) = func.VectorField()
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("t", schema=Schema)
|
||||
|
||||
with patch("time.sleep") as mock_sleep:
|
||||
with pytest.raises(Exception, match="The model `typo-embed` does not exist"):
|
||||
table.add([{"text": "hello"}])
|
||||
assert len(calls) == 1
|
||||
|
||||
with pytest.raises(_HttpStatusError, match="does not exist"):
|
||||
table.search("hello").limit(1).to_list()
|
||||
assert len(calls) == 2
|
||||
|
||||
mock_sleep.assert_not_called()
|
||||
|
||||
|
||||
def test_transient_embedding_error_is_retried_in_table_add(tmp_path):
|
||||
cls, calls = _register_failing_embedding(
|
||||
"failing-embedding-503", _HttpStatusError(503)
|
||||
)
|
||||
func = cls.create(max_retries=2)
|
||||
|
||||
class Schema(LanceModel):
|
||||
text: str = func.SourceField()
|
||||
vector: Vector(func.ndims()) = func.VectorField()
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("t", schema=Schema)
|
||||
|
||||
with patch("time.sleep") as mock_sleep:
|
||||
with pytest.raises(Exception, match="Maximum number of retries"):
|
||||
table.add([{"text": "hello"}])
|
||||
|
||||
assert len(calls) == 3
|
||||
assert mock_sleep.call_count == 2
|
||||
assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list)
|
||||
|
||||
|
||||
@patch("time.sleep")
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[
|
||||
pytest.param("missing_key", id="missing-api-key"),
|
||||
pytest.param(ImportError("No module named 'voyageai'"), id="missing-package"),
|
||||
],
|
||||
)
|
||||
def test_no_retry_on_deterministic_errors(mock_sleep, error):
|
||||
from lancedb.embeddings.utils import (
|
||||
api_key_not_found_help,
|
||||
retry_with_exponential_backoff,
|
||||
)
|
||||
|
||||
def func():
|
||||
if error == "missing_key":
|
||||
api_key_not_found_help("voyageai")
|
||||
raise error
|
||||
|
||||
mock_func = MagicMock(side_effect=func)
|
||||
wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3)
|
||||
|
||||
with pytest.raises((ValueError, ImportError)):
|
||||
wrapped_func()
|
||||
|
||||
assert mock_func.call_count == 1
|
||||
assert mock_sleep.call_count == 0
|
||||
|
||||
|
||||
@patch("time.sleep")
|
||||
def test_retry_on_wrapped_provider_error(mock_sleep):
|
||||
"""Providers such as Bedrock wrap throttling in a plain ValueError."""
|
||||
from lancedb.embeddings.utils import retry_with_exponential_backoff
|
||||
|
||||
mock_func = MagicMock(
|
||||
side_effect=[
|
||||
ValueError("Error raised by boto3 client: ThrottlingException"),
|
||||
[0.1],
|
||||
]
|
||||
)
|
||||
wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3)
|
||||
|
||||
assert wrapped_func() == [0.1]
|
||||
assert mock_func.call_count == 2
|
||||
assert mock_sleep.call_count == 1
|
||||
|
||||
|
||||
def test_openai_bad_request_partial_batch_recovery():
|
||||
# Regression test for issue #1677
|
||||
|
||||
@@ -708,3 +1057,28 @@ def test_jina_generate_image_input_dict_local_path(tmp_path):
|
||||
image_dict = JinaEmbeddings._generate_image_input_dict(image)
|
||||
assert "image" in image_dict
|
||||
assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0
|
||||
|
||||
|
||||
def test_jina_embeddings_use_their_own_api_key(monkeypatch):
|
||||
"""Each JinaEmbeddings instance must send its own API key."""
|
||||
pytest.importorskip("requests")
|
||||
monkeypatch.delenv("JINA_API_KEY", raising=False)
|
||||
registry = get_registry()
|
||||
registry.set_var("jina_key_a", "key-a")
|
||||
registry.set_var("jina_key_b", "key-b")
|
||||
func_a = registry.get("jina").create(api_key="$var:jina_key_a")
|
||||
func_b = registry.get("jina").create(api_key="$var:jina_key_b")
|
||||
|
||||
sent = []
|
||||
|
||||
def fake_post(session, url, json=None, **kwargs):
|
||||
sent.append(session.headers["Authorization"])
|
||||
response = MagicMock()
|
||||
response.json.return_value = {"data": [{"index": 0, "embedding": [0.0]}]}
|
||||
return response
|
||||
|
||||
with patch("requests.Session.post", fake_post):
|
||||
func_a.compute_source_embeddings(["hello"])
|
||||
func_b.compute_source_embeddings(["hello"])
|
||||
|
||||
assert sent == ["Bearer key-a", "Bearer key-b"]
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -80,6 +80,67 @@ class TestExprConstruction:
|
||||
e = lit(d)
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value, decimal_type",
|
||||
[
|
||||
("1E-7", pa.decimal128(12, 7)),
|
||||
("-1.23E-7", pa.decimal128(12, 9)),
|
||||
("1E+3", pa.decimal128(12, 0)),
|
||||
("-1.23E+5", pa.decimal128(12, 0)),
|
||||
("0E-7", pa.decimal128(12, 7)),
|
||||
("-0E+3", pa.decimal128(12, 0)),
|
||||
("1E-38", pa.decimal128(38, 38)),
|
||||
("1E+37", pa.decimal128(38, 0)),
|
||||
("1.234567890123456789012345678901234567E+37", pa.decimal128(38, 0)),
|
||||
("0E+128", pa.decimal128(1, 0)),
|
||||
],
|
||||
)
|
||||
def test_lit_decimal_exponent_query(self, tmp_path, value, decimal_type):
|
||||
expected = Decimal(value)
|
||||
table = lancedb.connect(tmp_path).create_table(
|
||||
"decimals",
|
||||
pa.table(
|
||||
{
|
||||
"id": [1, 2],
|
||||
# Arrow's Python ingestion rejects 0E+128 even though its
|
||||
# numeric value is representable. Store canonical zero to
|
||||
# exercise the expression literal independently of ingestion.
|
||||
"price": pa.array(
|
||||
[Decimal(0) if expected.is_zero() else expected, None],
|
||||
type=decimal_type,
|
||||
),
|
||||
}
|
||||
),
|
||||
)
|
||||
predicate = col("price") == lit(expected)
|
||||
result = table.search().where(predicate).to_arrow()
|
||||
assert result["id"].to_pylist() == [1]
|
||||
assert result["price"].to_pylist() == [expected]
|
||||
# Exercise the SQL serializer and parser as well as the typed query path.
|
||||
sql_result = table.search().where(predicate.to_sql()).to_arrow()
|
||||
assert sql_result["id"].to_pylist() == [1]
|
||||
assert sql_result["price"].to_pylist() == [expected]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value",
|
||||
[
|
||||
"NaN",
|
||||
"sNaN",
|
||||
"Infinity",
|
||||
"-Infinity",
|
||||
"1E-39",
|
||||
"1E+38",
|
||||
"1E+129",
|
||||
"1E+1000000000",
|
||||
"0E+129",
|
||||
"12345678901234567890123456789012345678E+1",
|
||||
"9" * 39,
|
||||
],
|
||||
)
|
||||
def test_lit_decimal_unrepresentable_raises(self, value):
|
||||
with pytest.raises(ValueError):
|
||||
lit(Decimal(value))
|
||||
|
||||
|
||||
class TestExprOperators:
|
||||
def test_eq_operator(self):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -129,3 +129,106 @@ def test_generator(tmp_path: Path):
|
||||
|
||||
assert len(tbl) == 2
|
||||
assert tbl.schema == ds.features.arrow_schema
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("operation", ["create", "add"])
|
||||
@pytest.mark.parametrize(
|
||||
"selection",
|
||||
["normal", "contiguous", "selected", "duplicates", "filtered", "shuffled", "empty"],
|
||||
)
|
||||
async def test_write_selected_hf_dataset(tmp_path, asynchronous, operation, selection):
|
||||
source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")})
|
||||
data = {
|
||||
"normal": source,
|
||||
"contiguous": source.select(range(1, 4)),
|
||||
"selected": source.select([4, 1]),
|
||||
"duplicates": source.select([2, 2, 0]),
|
||||
"filtered": source.filter(lambda row: row["id"] % 2 == 1),
|
||||
"shuffled": source.shuffle(seed=42),
|
||||
"empty": source.select([]),
|
||||
}[selection]
|
||||
expected = data.to_list()
|
||||
if asynchronous:
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
if operation == "create":
|
||||
table = await db.create_table("selected", data)
|
||||
else:
|
||||
table = await db.create_table("selected", schema=data.features.arrow_schema)
|
||||
await table.add(data)
|
||||
actual = await table.to_arrow()
|
||||
else:
|
||||
db = lancedb.connect(tmp_path)
|
||||
if operation == "create":
|
||||
table = db.create_table("selected", data)
|
||||
else:
|
||||
table = db.create_table("selected", schema=data.features.arrow_schema)
|
||||
table.add(data)
|
||||
actual = table.to_arrow()
|
||||
assert actual.to_pylist() == expected
|
||||
assert actual.schema.equals(data.features.arrow_schema, check_metadata=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize("operation", ["create", "add"])
|
||||
async def test_write_selected_hf_splits(tmp_path, asynchronous, operation):
|
||||
source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")})
|
||||
data = datasets.DatasetDict(
|
||||
{
|
||||
"train": source.select([4, 1]),
|
||||
"test": source.select([3, 0, 2]),
|
||||
"empty": source.select([]),
|
||||
}
|
||||
)
|
||||
expected = [
|
||||
dict(row, split=split)
|
||||
for split, dataset in data.items()
|
||||
for row in dataset.to_list()
|
||||
]
|
||||
schema = source.features.arrow_schema.append(pa.field("split", pa.string()))
|
||||
if asynchronous:
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
if operation == "create":
|
||||
table = await db.create_table("splits", data)
|
||||
else:
|
||||
table = await db.create_table("splits", schema=schema)
|
||||
await table.add(data)
|
||||
actual = await table.to_arrow()
|
||||
else:
|
||||
db = lancedb.connect(tmp_path)
|
||||
if operation == "create":
|
||||
table = db.create_table("splits", data)
|
||||
else:
|
||||
table = db.create_table("splits", schema=schema)
|
||||
table.add(data)
|
||||
actual = table.to_arrow()
|
||||
assert actual.to_pylist() == expected
|
||||
assert actual.schema.equals(schema, check_metadata=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("format_type", [None, "numpy", "pandas", "arrow"])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
def test_selected_hf_reader_is_rescannable(format_type, empty):
|
||||
from lancedb.scannable import _register_optional_converters, to_scannable
|
||||
|
||||
source = datasets.Dataset.from_dict(
|
||||
{"id": list(range(2105)), "text": [str(i) for i in range(2105)]}
|
||||
)
|
||||
indices = [] if empty else list(reversed(range(2105))) + [1, 1]
|
||||
data = source.select(indices).with_format(format_type, columns=["id"])
|
||||
original_format = data.format
|
||||
expected = source.select(indices).to_list()
|
||||
_register_optional_converters()
|
||||
scannable = to_scannable(data)
|
||||
assert scannable.rescannable
|
||||
assert scannable.num_rows == len(expected)
|
||||
for _ in range(2):
|
||||
batches = list(scannable.reader())
|
||||
if not empty:
|
||||
assert all(batch.num_rows <= 1000 for batch in batches)
|
||||
actual = pa.Table.from_batches(batches, schema=scannable.schema)
|
||||
assert actual.to_pylist() == expected
|
||||
assert actual.schema.equals(data.features.arrow_schema, check_metadata=True)
|
||||
assert data.format == original_format
|
||||
|
||||
@@ -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])
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
import tempfile
|
||||
import shutil
|
||||
import importlib
|
||||
import warnings
|
||||
import http.server
|
||||
import json
|
||||
import threading
|
||||
@@ -25,6 +26,125 @@ PUSHDOWN_DATA = pa.table(
|
||||
)
|
||||
|
||||
|
||||
class TestNamespaceCreateTableOptions:
|
||||
@pytest.mark.parametrize("version", ["legacy", "stable"])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
def test_deprecated_storage_version(self, tmp_path, version, empty):
|
||||
db = lancedb.connect_namespace("dir", {"root": str(tmp_path / "namespace")})
|
||||
data = None if empty else [{"id": 1}]
|
||||
schema = pa.schema([("id", pa.int64())])
|
||||
other_version = "stable" if version == "legacy" else "legacy"
|
||||
options = {"new_table_data_storage_version": other_version}
|
||||
control = lancedb.connect(tmp_path / "control").create_table(
|
||||
"test",
|
||||
data=data,
|
||||
schema=schema,
|
||||
storage_options={"new_table_data_storage_version": version},
|
||||
)
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always", DeprecationWarning)
|
||||
table = db.create_table(
|
||||
"test",
|
||||
data=data,
|
||||
schema=schema,
|
||||
storage_options=options,
|
||||
data_storage_version=version,
|
||||
)
|
||||
|
||||
assert (
|
||||
table.to_lance().data_storage_version
|
||||
== control.to_lance().data_storage_version
|
||||
)
|
||||
assert table.to_arrow().to_pylist() == ([] if empty else data)
|
||||
assert options == {"new_table_data_storage_version": other_version}
|
||||
assert any("data_storage_version" in str(w.message) for w in caught)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("version", ["legacy", "stable"])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
async def test_deprecated_storage_version_async(self, tmp_path, version, empty):
|
||||
root = str(tmp_path / "namespace")
|
||||
db = lancedb.connect_namespace_async("dir", {"root": root})
|
||||
data = None if empty else [{"id": 1}]
|
||||
schema = pa.schema([("id", pa.int64())])
|
||||
other_version = "stable" if version == "legacy" else "legacy"
|
||||
options = {"new_table_data_storage_version": other_version}
|
||||
control = lancedb.connect(tmp_path / "control").create_table(
|
||||
"test",
|
||||
data=data,
|
||||
schema=schema,
|
||||
storage_options={"new_table_data_storage_version": version},
|
||||
)
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always", DeprecationWarning)
|
||||
table = await db.create_table(
|
||||
"test",
|
||||
data=data,
|
||||
schema=schema,
|
||||
storage_options=options,
|
||||
data_storage_version=version,
|
||||
)
|
||||
|
||||
reopened = lancedb.connect_namespace("dir", {"root": root}).open_table("test")
|
||||
assert (
|
||||
reopened.to_lance().data_storage_version
|
||||
== control.to_lance().data_storage_version
|
||||
)
|
||||
assert (await table.to_arrow()).to_pylist() == ([] if empty else data)
|
||||
assert options == {"new_table_data_storage_version": other_version}
|
||||
assert any("data_storage_version" in str(w.message) for w in caught)
|
||||
|
||||
@pytest.mark.parametrize("enabled", [False, True])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
def test_deprecated_manifest_paths(self, tmp_path, enabled, empty):
|
||||
db = lancedb.connect_namespace("dir", {"root": str(tmp_path)})
|
||||
data = None if empty else [{"id": 1}]
|
||||
options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()}
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always", DeprecationWarning)
|
||||
table = db.create_table(
|
||||
"test",
|
||||
data=data,
|
||||
schema=pa.schema([("id", pa.int64())]),
|
||||
storage_options=options,
|
||||
enable_v2_manifest_paths=enabled,
|
||||
)
|
||||
|
||||
assert table.uses_v2_manifest_paths() == enabled
|
||||
assert table.to_arrow().to_pylist() == ([] if empty else data)
|
||||
assert db.open_table("test").uses_v2_manifest_paths() == enabled
|
||||
assert options == {
|
||||
"new_table_enable_v2_manifest_paths": str(not enabled).lower()
|
||||
}
|
||||
assert any("enable_v2_manifest_paths" in str(w.message) for w in caught)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("enabled", [False, True])
|
||||
@pytest.mark.parametrize("empty", [False, True])
|
||||
async def test_deprecated_manifest_paths_async(self, tmp_path, enabled, empty):
|
||||
db = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
|
||||
data = None if empty else [{"id": 1}]
|
||||
options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()}
|
||||
with warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always", DeprecationWarning)
|
||||
table = await db.create_table(
|
||||
"test",
|
||||
data=data,
|
||||
schema=pa.schema([("id", pa.int64())]),
|
||||
storage_options=options,
|
||||
enable_v2_manifest_paths=enabled,
|
||||
)
|
||||
|
||||
assert await table.uses_v2_manifest_paths() == enabled
|
||||
assert (await table.to_arrow()).to_pylist() == ([] if empty else data)
|
||||
reopened = await db.open_table("test")
|
||||
assert await reopened.uses_v2_manifest_paths() == enabled
|
||||
assert options == {
|
||||
"new_table_enable_v2_manifest_paths": str(not enabled).lower()
|
||||
}
|
||||
assert any("enable_v2_manifest_paths" in str(w.message) for w in caught)
|
||||
|
||||
|
||||
def _ipc_file(table: pa.Table = PUSHDOWN_DATA) -> bytes:
|
||||
sink = pa.BufferOutputStream()
|
||||
with pa.ipc.new_file(sink, table.schema) as writer:
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from lancedb.embeddings import get_registry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ollama_embeddings():
|
||||
httpx = pytest.importorskip("httpx")
|
||||
pytest.importorskip("ollama")
|
||||
requests = []
|
||||
|
||||
def respond(request):
|
||||
body = json.loads(request.content)
|
||||
requests.append(body)
|
||||
if body["model"] == "missing-model":
|
||||
return httpx.Response(404, json={"error": "model not found"})
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"embeddings": [
|
||||
[float(index), float(len(text))]
|
||||
for index, text in enumerate(body["input"])
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
function = (
|
||||
get_registry()
|
||||
.get("ollama")
|
||||
.create(
|
||||
name="nomic-embed-text",
|
||||
options={"num_ctx": 4096},
|
||||
keep_alive="10m",
|
||||
ollama_client_kwargs={
|
||||
"transport": httpx.MockTransport(respond),
|
||||
"trust_env": False,
|
||||
},
|
||||
)
|
||||
)
|
||||
yield function, requests
|
||||
function._ollama_client._client.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"texts",
|
||||
[
|
||||
pytest.param(["hello", "世界"], id="list"),
|
||||
pytest.param(np.array(["hello", "世界"]), id="numpy-unicode"),
|
||||
pytest.param(np.array(["hello", "世界"], dtype=object), id="numpy-object"),
|
||||
pytest.param(
|
||||
np.array([b"hello", "世界".encode("utf-8")]), id="numpy-utf8-bytes"
|
||||
),
|
||||
pytest.param([b"hello", "世界".encode("utf-8")], id="list-utf8-bytes"),
|
||||
pytest.param(np.array(["hello", "unused", "世界"])[::2], id="numpy-slice"),
|
||||
pytest.param(pa.array(["hello", "世界"]), id="arrow"),
|
||||
pytest.param(pa.chunked_array([["hello"], ["世界"]]), id="chunked-arrow"),
|
||||
],
|
||||
)
|
||||
def test_ollama_source_embeddings_inputs(ollama_embeddings, texts):
|
||||
function, requests = ollama_embeddings
|
||||
|
||||
assert function.compute_source_embeddings(texts) == [[0.0, 5.0], [1.0, 2.0]]
|
||||
assert requests == [
|
||||
{
|
||||
"model": "nomic-embed-text",
|
||||
"input": ["hello", "世界"],
|
||||
"options": {"num_ctx": 4096},
|
||||
"keep_alive": "10m",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"texts",
|
||||
[
|
||||
pytest.param(np.array([["hello"], ["world"]]), id="numpy-two-dimensional"),
|
||||
pytest.param([["hello"], ["world"]], id="list-two-dimensional"),
|
||||
pytest.param(np.array([b"\xff"]), id="numpy-invalid-utf8"),
|
||||
pytest.param([b"\xff"], id="list-invalid-utf8"),
|
||||
],
|
||||
)
|
||||
def test_ollama_source_embeddings_rejects_invalid_input_locally(
|
||||
ollama_embeddings, texts
|
||||
):
|
||||
function, requests = ollama_embeddings
|
||||
|
||||
with pytest.raises(ValidationError, match="valid string"):
|
||||
function.compute_source_embeddings(texts)
|
||||
|
||||
assert requests == []
|
||||
|
||||
|
||||
def test_ollama_generate_embeddings_empty_array(ollama_embeddings):
|
||||
function, requests = ollama_embeddings
|
||||
|
||||
assert function.generate_embeddings(np.array([], dtype=str)) == []
|
||||
assert len(requests) == 1
|
||||
assert requests[0]["input"] == []
|
||||
|
||||
|
||||
def test_ollama_query_embeddings(ollama_embeddings):
|
||||
function, requests = ollama_embeddings
|
||||
|
||||
assert function.compute_query_embeddings("hello") == [[0.0, 5.0]]
|
||||
assert len(requests) == 1
|
||||
assert requests[0]["input"] == ["hello"]
|
||||
|
||||
|
||||
def test_ollama_source_embeddings_preserves_api_error(ollama_embeddings):
|
||||
ollama = pytest.importorskip("ollama")
|
||||
function, requests = ollama_embeddings
|
||||
function.name = "missing-model"
|
||||
|
||||
with pytest.raises(ollama.ResponseError, match="model not found") as error:
|
||||
function.compute_source_embeddings(np.array(["hello"]))
|
||||
|
||||
assert error.value.status_code == 404
|
||||
assert len(requests) == 1
|
||||
@@ -0,0 +1,66 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
from lancedb.rerankers import OpenaiReranker
|
||||
|
||||
|
||||
def reranker_replying(documents):
|
||||
"""Return an OpenaiReranker whose chat client replies with these documents."""
|
||||
reply = json.dumps({"documents": documents})
|
||||
|
||||
def create(**kwargs):
|
||||
message = SimpleNamespace(content=reply)
|
||||
return SimpleNamespace(choices=[SimpleNamespace(message=message)])
|
||||
|
||||
reranker = OpenaiReranker(api_key="test-key")
|
||||
reranker.__dict__["_client"] = SimpleNamespace(
|
||||
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
|
||||
)
|
||||
return reranker
|
||||
|
||||
|
||||
def vector_results():
|
||||
return pa.table(
|
||||
{
|
||||
"id": [1, 2, 3],
|
||||
"text": ["dogs bark", "cats purr", "birds sing"],
|
||||
"_distance": [0.1, 0.2, 0.3],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_scores_stay_with_their_rows_when_the_model_reorders():
|
||||
reranker = reranker_replying(
|
||||
[
|
||||
{"content": "cats purr", "relevance_score": 0.9},
|
||||
{"content": "birds sing", "relevance_score": 0.5},
|
||||
{"content": "dogs bark", "relevance_score": 0.1},
|
||||
]
|
||||
)
|
||||
|
||||
result = reranker.rerank_vector("which animal purrs?", vector_results())
|
||||
|
||||
assert result.select(["id", "text"]).to_pylist() == [
|
||||
{"id": 2, "text": "cats purr"},
|
||||
{"id": 3, "text": "birds sing"},
|
||||
{"id": 1, "text": "dogs bark"},
|
||||
]
|
||||
assert result["_relevance_score"].to_pylist() == pytest.approx([0.9, 0.5, 0.1])
|
||||
|
||||
|
||||
def test_a_row_the_model_leaves_out_raises():
|
||||
reranker = reranker_replying(
|
||||
[
|
||||
{"content": "cats purr", "relevance_score": 0.9},
|
||||
{"content": "dogs bark", "relevance_score": 0.1},
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="1 of 3"):
|
||||
reranker.rerank_vector("which animal purrs?", vector_results())
|
||||
@@ -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
|
||||
@@ -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())]
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
|
||||
import lancedb
|
||||
from lancedb.integrations.pyarrow import PyarrowDatasetAdapter
|
||||
@@ -48,3 +49,19 @@ def test_dataset_adapter(tmp_path):
|
||||
assert adapter.head(10).num_rows == 0
|
||||
|
||||
assert adapter.scanner().projected_schema == pa.schema({"x": pa.int64()})
|
||||
|
||||
|
||||
def test_dataset_adapter_expression_filter(tmp_path):
|
||||
data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]})
|
||||
conn = lancedb.connect(tmp_path)
|
||||
tbl = conn.create_table("test", data)
|
||||
|
||||
adapter = PyarrowDatasetAdapter(tbl)
|
||||
filt = pc.field("x") > 2
|
||||
|
||||
assert adapter.count_rows(filt) == 2
|
||||
assert adapter.to_table(filter=filt) == data.slice(2, 2)
|
||||
assert adapter.head(1, filter=filt) == data.slice(2, 1)
|
||||
# The filter can reference a column that is not projected
|
||||
assert adapter.to_table(columns=["y"], filter=filt) == pa.table({"y": [7, 8]})
|
||||
assert adapter.scanner(filter=filt).count_rows() == 2
|
||||
|
||||
@@ -415,6 +415,67 @@ def test_nullable_vector():
|
||||
assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("vector_type", [Vector, MultiVector])
|
||||
@pytest.mark.parametrize(
|
||||
"value_type, values",
|
||||
[
|
||||
(pa.int64(), [-(2**53 + 1), 2**53 + 3]),
|
||||
(pa.uint64(), [2**53 + 1, 2**64 - 1]),
|
||||
(pa.int8(), [-128, 127]),
|
||||
(pa.uint8(), [0, 255]),
|
||||
(pa.float32(), [1.25, 2.5]),
|
||||
(pa.float64(), [1.25, 2.5]),
|
||||
],
|
||||
)
|
||||
def test_typed_vector_round_trip(tmp_path, vector_type, value_type, values):
|
||||
import lancedb
|
||||
|
||||
annotation = vector_type(2, value_type=value_type)
|
||||
expected = [values, values] if vector_type is MultiVector else values
|
||||
|
||||
class Row(LanceModel):
|
||||
id: int
|
||||
payload: annotation
|
||||
|
||||
row = Row(id=1, payload=expected)
|
||||
table = lancedb.connect(tmp_path).create_table("typed_vectors", [row], schema=Row)
|
||||
table.add([Row(id=2, payload=expected)])
|
||||
assert table.schema == Row.to_arrow_schema()
|
||||
assert table.search().to_arrow().to_pylist() == [
|
||||
{"id": 1, "payload": expected},
|
||||
{"id": 2, "payload": expected},
|
||||
]
|
||||
assert json.loads(row.model_dump_json())["payload"] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("vector_type", [Vector, MultiVector])
|
||||
@pytest.mark.parametrize("nullable", [False, True])
|
||||
def test_integer_vector_validation(vector_type, nullable):
|
||||
annotation = vector_type(2, value_type=pa.int64(), nullable=nullable)
|
||||
|
||||
class Row(LanceModel):
|
||||
payload: annotation
|
||||
|
||||
def shape(values):
|
||||
return [values] if vector_type is MultiVector else values
|
||||
|
||||
assert Row.to_arrow_schema().field("payload").nullable is nullable
|
||||
assert Row(payload=shape([1.0, -2.0])).model_dump()["payload"] == shape([1, -2])
|
||||
with pytest.raises(pydantic.ValidationError, match="valid integer"):
|
||||
Row(payload=shape([1.5, 2]))
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
Row(payload=shape([1]))
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
Row(payload=shape([1, 2, 3]))
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
Row(payload=None)
|
||||
|
||||
class OptionalRow(LanceModel):
|
||||
payload: Optional[annotation] = None
|
||||
|
||||
assert OptionalRow(payload=None).payload is None
|
||||
|
||||
|
||||
def test_bare_vector_raises_clear_error():
|
||||
namespace = {
|
||||
"__name__": "test_model_without_pyarrow",
|
||||
@@ -733,3 +794,52 @@ def test_enum_types():
|
||||
assert schema.field("priority").type == pa.int64()
|
||||
assert schema.field("opt_status").type == pa.dictionary(pa.int32(), pa.utf8())
|
||||
assert schema.field("opt_status").nullable
|
||||
|
||||
|
||||
def test_plain_enum_values_are_written(mem_db):
|
||||
"""Members of an Enum without a str/int mixin are written as their values."""
|
||||
|
||||
class Color(Enum):
|
||||
RED = "red"
|
||||
BLUE = "blue"
|
||||
|
||||
class Priority(Enum):
|
||||
LOW = 1
|
||||
HIGH = 2
|
||||
|
||||
class Tag(LanceModel):
|
||||
color: Color
|
||||
|
||||
class TestModel(LanceModel):
|
||||
color: Color
|
||||
priority: Priority
|
||||
opt_color: Optional[Color] = None
|
||||
colors: List[Color]
|
||||
tag: Tag
|
||||
vector: Vector(2)
|
||||
|
||||
rows = [
|
||||
TestModel(
|
||||
color=Color.RED,
|
||||
priority=Priority.HIGH,
|
||||
opt_color=Color.BLUE,
|
||||
colors=[Color.BLUE, Color.RED],
|
||||
tag=Tag(color=Color.BLUE),
|
||||
vector=[1.0, 2.0],
|
||||
),
|
||||
TestModel(
|
||||
color=Color.BLUE,
|
||||
priority=Priority.LOW,
|
||||
colors=[],
|
||||
tag=Tag(color=Color.RED),
|
||||
vector=[3.0, 4.0],
|
||||
),
|
||||
]
|
||||
tbl = mem_db.create_table("test", schema=TestModel)
|
||||
tbl.add(rows)
|
||||
|
||||
assert tbl.to_arrow().select(["color", "priority", "opt_color"]).to_pylist() == [
|
||||
{"color": "red", "priority": 2, "opt_color": "blue"},
|
||||
{"color": "blue", "priority": 1, "opt_color": None},
|
||||
]
|
||||
assert tbl.search().to_pydantic(TestModel) == rows
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user