rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-25 21:22:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: merge main to pick up Tokio advisory fix

Browse Source
This commit is contained in:
Xuanwo
2026-10-23 21:48:31 +02:00
parent 3d4f93e431 e74a7c15b5
commit 0ba5e9637f
135 changed files with 14115 additions and 1011 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.agents/skills/lancedb-update-lance-dependency/SKILL.md
+27
View File
@@ -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.
.github/workflows/ci-scripts.yml
+6
View File
@@ -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
.github/workflows/docs.yml
+1
View File
@@ -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
.github/workflows/rust.yml
+37 -6
View File
@@ -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 }}
Cargo.lock
Generated
+5 -5
View File
@@ -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",
]
@@ -5136,7 +5136,6 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "15.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5834,6 +5833,7 @@ dependencies = [
"datafusion-common",
"env_logger",
"futures",
"lance-arrow",
"lance-core",
"lance-io",
"lance-namespace",
@@ -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",
Cargo.toml
+6 -1
View File
@@ -19,6 +19,7 @@ lance-datagen = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git"
lance-file = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-jsonb = { version = "=15.0.0-beta.1", tag = "v15.0.0-beta.1", git = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
@@ -26,7 +27,7 @@ lance-table = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" =
lance-testing = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { version = "=15.0.0-beta.1", path = "python/vendor/lance-arrow" }
lancedb = { path = "rust/lancedb", default-features = false }
ahash = "0.8"
# Note that this one does not include pyarrow
@@ -79,6 +80,10 @@ tonic = { version = "0.14", features = ["tls-native-roots", "tls-aws-lc"] }
uuid = { version = "1.7.0", features = ["v4", "v7"] }
chrono = { version = "0.4", default-features = false, features = ["clock"] }
[patch."https://github.com/lance-format/lance.git"]
# Carry the offset-safe validity fix until it is available in a Lance release.
lance-arrow = { path = "python/vendor/lance-arrow" }
[patch.crates-io]
# Lance v14.0.0-beta.18 requires multipart write modes from this upstream revision.
object_store = { git = "https://github.com/apache/arrow-rs-object-store.git", rev = "c10758c207c0103355f940f0c736fb8ecb157b04" }
ci/tests/test_python_dependencies.py
+47
View File
@@ -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()
docs/src/js/classes/Connection.md
+17 -63
View File
@@ -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`&lt;[`ListNamespacesOptions`](../interfaces/ListNamespacesOptions.md)&gt;
Pagination options
(`pageToken`, `limit`).
* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md)
#### Returns
`Promise`&lt;[`ListNamespacesResponse`](../interfaces/ListNamespacesResponse.md)&gt;
Child namespace names and
an optional token for fetching the next page.
[`Listing`](Listing.md)&lt;`string`&gt;
***
### 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`&lt;[`ListTablesOptions`](../interfaces/ListTablesOptions.md)&gt;
Pagination options
(`pageToken`, `limit`).
##### Returns
`Promise`&lt;[`ListTablesResponse`](../interfaces/ListTablesResponse.md)&gt;
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`&lt;[`ListTablesOptions`](../interfaces/ListTablesOptions.md)&gt;
Pagination options
(`pageToken`, `limit`).
* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md)
##### Returns
#### Returns
`Promise`&lt;[`ListTablesResponse`](../interfaces/ListTablesResponse.md)&gt;
[`Listing`](Listing.md)&lt;`string`&gt;
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);
}
```
***
docs/src/js/classes/Table.md
+5 -1
View File
@@ -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
docs/src/js/globals.md
-4
View File
@@ -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)
docs/src/js/interfaces/CreateTableOptions.md
+4 -1
View File
@@ -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?
docs/src/js/interfaces/ListNamespacesOptions.md
-27
View File
@@ -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.
docs/src/js/interfaces/ListNamespacesResponse.md
-23
View File
@@ -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;
```
docs/src/js/interfaces/ListTablesOptions.md
-35
View File
@@ -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.
docs/src/js/interfaces/ListTablesResponse.md
-23
View File
@@ -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[];
```
docs/src/js/interfaces/QueryExecutionOptions.md
+5 -3
View File
@@ -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.
docs/src/js/interfaces/TableNamesOptions.md
+1 -1
View File
@@ -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
docs/src/python/python.md
+18 -5
View File
@@ -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)
@@ -631,6 +634,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
docs/web-tests/ts/connection.test.ts
+11 -5
View File
@@ -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' ]
docs/web-tests/ts/quickstart.test.ts
+1 -6
View File
@@ -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;
}
docs/web/docs.json
+1
View File
@@ -350,6 +350,7 @@
"pages": [
"enterprise/index",
"enterprise/architecture",
"enterprise/distributed-search",
"enterprise/deployment/index",
"deployment/storage",
"deployment/monitoring",
docs/web/enterprise/architecture.mdx
+3
View File
@@ -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.
docs/web/enterprise/distributed-search.mdx
+33
View File
@@ -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.
docs/web/indexing/vector-index.mdx
+1 -1
View File
@@ -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`.
docs/web/integrations/embedding/ollama.mdx
+3
View File
@@ -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. |
docs/web/integrations/embedding/voyageai.mdx
+9 -2
View File
@@ -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>
docs/web/quickstart.mdx
+27
View File
@@ -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
docs/web/search/vector-search.mdx
+1 -1
View File
@@ -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.
nodejs/.npmignore
+1
View File
@@ -15,6 +15,7 @@ renovate.json
src
lancedb
examples
scripts
nodejs-artifacts
Cargo.toml
biome.json
nodejs/__test__/connection.test.ts
+332 -56
View File
@@ -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",
);
nodejs/__test__/fixtures/transformers_tokenizer.cjs
+66
View File
@@ -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;
});
nodejs/__test__/native_symbols.test.ts
+47
View File
@@ -0,0 +1,47 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
const {
checkNativeBinary,
findForbiddenUndefinedSymbols,
} = require("../scripts/check-native-symbols.js");
test("detects the unresolved AVX-512 symbol from broken musl binaries", () => {
const output = [
" U napi_create_function",
" U sum_4bit_dist_table_32bytes_batch_avx512",
" U strlen",
].join("\n");
expect(findForbiddenUndefinedSymbols(output)).toEqual([
"sum_4bit_dist_table_32bytes_batch_avx512",
]);
});
test("accepts native binaries without unresolved internal Lance symbols", () => {
const runNm = jest.fn(() => ({
status: 0,
stdout:
" U napi_create_function\n U strlen\n",
stderr: "",
}));
expect(() => checkNativeBinary("lancedb.node", runNm)).not.toThrow();
expect(runNm).toHaveBeenCalledWith(
"nm",
["-D", "--undefined-only", "lancedb.node"],
{ encoding: "utf8" },
);
});
test("rejects native binaries with the unresolved AVX-512 symbol", () => {
const runNm = jest.fn(() => ({
status: 0,
stdout: " U sum_4bit_dist_table_32bytes_batch_avx512\n",
stderr: "",
}));
expect(() => checkNativeBinary("lancedb.node", runNm)).toThrow(
"lancedb.node contains unresolved internal Lance symbols: sum_4bit_dist_table_32bytes_batch_avx512",
);
});
nodejs/__test__/nested_blob.test.ts
+327
View File
@@ -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,
);
});
});
nodejs/__test__/table.test.ts
+218 -47
View File
@@ -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(
nodejs/__test__/transformers.test.ts
+26
View File
@@ -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 },
);
});
});
nodejs/lancedb/arrow.ts
+1
View File
@@ -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
| {
nodejs/lancedb/connection.ts
+58 -118
View File
@@ -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.
*/
nodejs/lancedb/embedding/transformers.ts
+4 -2
View File
@@ -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;
nodejs/lancedb/index.ts
+8 -13
View File
@@ -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,
nodejs/lancedb/query.ts
+108 -37
View File
@@ -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;
nodejs/lancedb/table.ts
+5 -1
View File
@@ -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,
nodejs/package.json
+2 -1
View File
@@ -90,7 +90,8 @@
"postdocs": "node typedoc_post_process.js",
"lint": "biome check . && biome format .",
"lint-fix": "biome check --write . && biome format --write .",
"prepublishOnly": "napi prepublish -t npm",
"check:native-symbols": "node scripts/check-native-symbols.js",
"prepublishOnly": "pnpm check:native-symbols && napi prepublish -t npm",
"test": "jest --verbose",
"integration": "S3_TEST=1 pnpm test",
"universal": "napi universalize",
nodejs/scripts/check-native-symbols.js
+77
View File
@@ -0,0 +1,77 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
const { spawnSync } = require("node:child_process");
const { existsSync } = require("node:fs");
const path = require("node:path");
const FORBIDDEN_UNDEFINED_SYMBOLS = new Set([
"sum_4bit_dist_table_32bytes_batch_avx512",
]);
function findForbiddenUndefinedSymbols(output) {
const found = new Set();
for (const line of output.split(/\r?\n/)) {
const columns = line.trim().split(/\s+/);
const symbol = columns.at(-1)?.split("@")[0];
if (symbol && FORBIDDEN_UNDEFINED_SYMBOLS.has(symbol)) {
found.add(symbol);
}
}
return [...found].sort();
}
function checkNativeBinary(binaryPath, runNm = spawnSync) {
const result = runNm("nm", ["-D", "--undefined-only", binaryPath], {
encoding: "utf8",
});
if (result.error) {
throw new Error(`Unable to inspect ${binaryPath}: ${result.error.message}`);
}
if (result.status !== 0) {
throw new Error(
`Unable to inspect ${binaryPath}: nm exited with status ${result.status}\n${result.stderr}`,
);
}
const forbidden = findForbiddenUndefinedSymbols(result.stdout);
if (forbidden.length > 0) {
throw new Error(
`${binaryPath} contains unresolved internal Lance symbols: ${forbidden.join(
", ",
)}`,
);
}
}
function main() {
const binaryPath = path.resolve(
__dirname,
"..",
"npm",
"linux-x64-musl",
"lancedb.linux-x64-musl.node",
);
if (!existsSync(binaryPath)) {
throw new Error(
`Missing ${binaryPath}; assemble the native artifacts before publishing`,
);
}
checkNativeBinary(binaryPath);
}
if (require.main === module) {
try {
main();
} catch (error) {
console.error(error instanceof Error ? error.message : error);
process.exitCode = 1;
}
}
module.exports = { checkNativeBinary, findForbiddenUndefinedSymbols };
nodejs/src/connection.rs
+19 -52
View File
@@ -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: limit.map(|l| l as i32),
..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)]
python/Cargo.toml
+2
View File
@@ -22,6 +22,8 @@ bytes.workspace = true
lancedb.workspace = true
datafusion-common.workspace = true
lance-core.workspace = true
# The workspace path dependency includes the patched crate in Maturin sources.
lance-arrow.workspace = true
lance-namespace.workspace = true
lance-namespace-impls.workspace = true
lance-io.workspace = true
python/README.md
+27
View File
@@ -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
python/pyproject.toml
+3 -2
View File
@@ -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",
python/python/lancedb/_lancedb.pyi
+5 -6
View File
@@ -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]],
@@ -1010,6 +1008,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: ...
python/python/lancedb/db.py
+61 -33
View File
@@ -5,6 +5,7 @@
from __future__ import annotations
from abc import abstractmethod
import asyncio
from datetime import timedelta
from pathlib import Path
import sys
@@ -31,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
@@ -170,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
@@ -186,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,
@@ -316,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
----------
@@ -333,7 +340,7 @@ class DBConnection(EnforceOverrides):
Returns
-------
ListTablesResponse
TableListing
Response containing table names and optional page_token for pagination.
"""
raise NotImplementedError(
@@ -1299,6 +1306,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
@@ -1365,7 +1376,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:
@@ -1426,7 +1442,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
@@ -1442,16 +1458,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
)
)
@@ -1502,8 +1516,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
----------
@@ -1519,12 +1533,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
)
@@ -2276,12 +2290,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
@@ -2296,15 +2310,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,
@@ -2383,13 +2400,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
----------
@@ -2405,15 +2422,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,
@@ -2652,8 +2672,16 @@ class AsyncConnection(object):
if fill_value is None:
fill_value = 0.0
data, schema = sanitize_create_table(
data, schema, metadata, on_bad_vectors, fill_value
# Input preparation may advance a user-provided iterator. Keep that work
# off the background event loop so an iterator can use the synchronous
# LanceDB API without blocking the loop that API needs to make progress.
data, schema = await asyncio.to_thread(
sanitize_create_table,
data,
schema,
metadata,
on_bad_vectors,
fill_value,
)
validate_schema(schema)
python/python/lancedb/embeddings/base.py
+11 -5
View File
@@ -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__"):
python/python/lancedb/embeddings/jinaai.py
+13 -13
View File
@@ -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
python/python/lancedb/embeddings/ollama.py
+2
View File
@@ -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)
python/python/lancedb/embeddings/openai.py
+18 -7
View File
@@ -101,24 +101,35 @@ class OpenAIEmbeddings(TextEmbeddingFunction):
valid_indices.append(idx)
# TODO retry, rate limit, token limit
try:
def embed_batch(texts_batch, indices_batch):
kwargs = {
"input": valid_texts,
"input": texts_batch,
"model": self.name,
}
if self.name != "text-embedding-ada-002":
kwargs["dimensions"] = self.dim
rs = self._openai_client.embeddings.create(**kwargs)
valid_embeddings = {
idx: v.embedding for v, idx in zip(rs.data, valid_indices)
}
return {idx: v.embedding for v, idx in zip(rs.data, indices_batch)}
def embed_one_or_log(idx, text):
# Retry a single rejected input
try:
return embed_batch([text], [idx])
except openai.BadRequestError:
logging.exception("Bad request for input at index %d: %r", idx, text)
return {}
try:
valid_embeddings = embed_batch(valid_texts, valid_indices)
except openai.AuthenticationError:
logging.error("Authentication failed: Invalid API key provided")
raise
except openai.BadRequestError:
logging.exception("Bad request: %s", texts)
return [None] * len(texts)
# Retry each input individually
valid_embeddings = {}
for idx, text in zip(valid_indices, valid_texts):
valid_embeddings.update(embed_one_or_log(idx, text))
except Exception:
logging.exception("OpenAI embeddings error")
raise
python/python/lancedb/embeddings/registry.py
+7 -3
View File
@@ -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
python/python/lancedb/embeddings/utils.py
+97 -17
View File
@@ -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():
python/python/lancedb/embeddings/voyageai.py
+5 -3
View File
@@ -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):
python/python/lancedb/integrations/pyarrow.py
+23 -1
View File
@@ -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
python/python/lancedb/listing.py
+171
View File
@@ -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
python/python/lancedb/namespace.py
+54 -17
View File
@@ -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
)
python/python/lancedb/permutation.py
+5 -3
View File
@@ -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":
python/python/lancedb/pydantic.py
+35 -4
View File
@@ -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
python/python/lancedb/query.py
+20 -11
View File
@@ -82,6 +82,7 @@ if TYPE_CHECKING:
T = TypeVar("T", bound="LanceModel")
AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
_SYNC_VECTOR_QUERY_TYPES = (list, tuple, np.ndarray, pa.Array, pa.ChunkedArray)
# Number of rows a hybrid query returns when no limit was set on it. This
# mirrors the default the Rust query builder applies to its sub-queries.
@@ -1024,10 +1025,12 @@ class LanceQueryBuilder(ABC):
fast_search=fast_search,
)
if isinstance(query, list):
if isinstance(query, (list, tuple)):
query = np.array(query, dtype=np.float32)
elif isinstance(query, np.ndarray):
query = query.astype(np.float32)
elif isinstance(query, (pa.Array, pa.ChunkedArray)):
query = np.array(query.to_pylist(), dtype=np.float32)
else:
raise TypeError(f"Unsupported query type: {type(query)}")
@@ -1049,7 +1052,7 @@ class LanceQueryBuilder(ABC):
query = cls._query_to_vector(table, query, vector_column_name)
return query, query_type
elif query_type == "auto":
if isinstance(query, (list, np.ndarray)):
if isinstance(query, _SYNC_VECTOR_QUERY_TYPES):
return query, "vector"
else:
conf = table.embedding_functions.get(vector_column_name)
@@ -1065,7 +1068,7 @@ class LanceQueryBuilder(ABC):
@classmethod
def _query_to_vector(cls, table, query, vector_column_name):
if isinstance(query, (list, np.ndarray)):
if isinstance(query, _SYNC_VECTOR_QUERY_TYPES):
return query
conf = table.embedding_functions.get(vector_column_name)
if conf is not None:
@@ -2438,7 +2441,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
if query is not None and (vector is not None or text is not None):
raise ValueError(
"You can either provide a string query in search() method "
"or set `vector()` and `text()` explicitly for hybrid search."
"or set `vector()` and `text()` explicitly for hybrid search. "
"But not both."
)
@@ -2572,10 +2575,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:
@@ -2602,15 +2601,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
@@ -2619,10 +2626,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)
python/python/lancedb/remote/db.py
+7 -9
View File
@@ -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
)
python/python/lancedb/remote/header.py
+1 -1
View File
@@ -141,7 +141,7 @@ class OAuthProvider(HeaderProvider):
# Set expiration if provided
expires_in = token_data.get("expires_in")
if expires_in:
if expires_in is not None:
self._token_expires_at = time.time() + expires_in
else:
# Token doesn't expire or expiration unknown
python/python/lancedb/remote/table.py
+6 -2
View File
@@ -279,8 +279,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)
python/python/lancedb/rerankers/base.py
+4 -1
View File
@@ -241,6 +241,7 @@ class Reranker(ABC):
This op is slower than just keeping relevance score but can be useful
for debugging.
"""
combined_schema = pa.unify_schemas([vector_results.schema, fts_results.schema])
# add nulls to fts results for _distance
if "_distance" not in fts_results.column_names:
fts_results = fts_results.append_column(
@@ -265,7 +266,9 @@ class Reranker(ABC):
else:
vector_results_dict[key] = value
combined = pa.Table.from_pylist(list(vector_results_dict.values()))
combined = pa.Table.from_pylist(
list(vector_results_dict.values()), schema=combined_schema
)
return combined
def _keep_relevance_score(self, combined_results: pa.Table):
python/python/lancedb/rerankers/cohere.py
+1 -1
View File
@@ -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)
python/python/lancedb/rerankers/openai.py
+12 -9
View File
@@ -70,15 +70,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())
)
python/python/lancedb/rerankers/voyageai.py
+2 -2
View File
@@ -32,7 +32,7 @@ class VoyageAIReranker(Reranker):
return_score : str, default "relevance"
options are "relevance" or "all". Only "relevance" is supported for now.
api_key : str, default None
The API key to use. If None, will use the OPENAI_API_KEY environment variable.
The API key to use. If None, will use the VOYAGE_API_KEY environment variable.
truncation : Optional[bool], default None
"""
@@ -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):
python/python/lancedb/scannable.py
+18 -2
View File
@@ -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()
)
python/python/lancedb/table.py
+95 -18
View File
@@ -28,7 +28,7 @@ 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 (
@@ -119,6 +119,33 @@ def _should_push_down_query_table(
return namespace_client is not None and "QueryTable" in pushdown_operations
_CUDA_UNAVAILABLE_MESSAGE = (
"CUDA acceleration is unavailable in this Python environment. Install a "
"CUDA-enabled PyTorch build and verify that the CUDA device is accessible, or "
"omit `accelerator` (or set `accelerator=None`) to create the index on CPU."
)
def _validate_index_accelerator(accelerator: Any) -> None:
if not isinstance(accelerator, str) or accelerator.split(":", 1)[0] != "cuda":
return
try:
import torch
cuda_available = torch.cuda.is_available()
except (ImportError, RuntimeError, AssertionError) as err:
raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err
if not cuda_available:
raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE)
try:
torch.cuda.init()
except (RuntimeError, AssertionError) as err:
raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err
def _polars_predicate_pushdown_barrier(frame: Any) -> Any:
"""Return a Polars frame unchanged while blocking predicate pushdown."""
return frame
@@ -373,7 +400,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:
@@ -442,7 +469,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)}. "
@@ -452,10 +479,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():
@@ -1168,7 +1204,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(
@@ -2531,7 +2567,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
@@ -3600,7 +3640,9 @@ class LanceTable(Table):
# Handle accelerator through pylance
if accelerator is not None:
self.to_lance().create_index(
dataset = self.to_lance()
_validate_index_accelerator(accelerator)
dataset.create_index(
column=column,
index_type=index_type,
metric=metric,
@@ -3625,6 +3667,8 @@ class LanceTable(Table):
acc = getattr(config, "accelerator", None)
if acc is not None:
# Dispatch to pylance for GPU acceleration
dataset = self.to_lance()
_validate_index_accelerator(acc)
index_type_map = {
"IvfFlat": "IVF_FLAT",
"IvfSq": "IVF_SQ",
@@ -3636,7 +3680,7 @@ class LanceTable(Table):
cfg_type = type(config).__name__
lance_index_type = index_type_map.get(cfg_type, "IVF_PQ")
self.to_lance().create_index(
dataset.create_index(
column=column,
index_type=lance_index_type,
metric=getattr(config, "distance_type", "l2"),
@@ -4782,7 +4826,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)
@@ -5285,6 +5333,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)
@@ -5335,6 +5388,12 @@ def _infer_target_schema(
# Use the most common length of the list as the dimensions
dim = _modal_list_size(peeked.column(i))
if dim is None:
# Every list entry is null (e.g. embeddings not computed yet),
# so the dimension is unknowable. Leave the column as a
# variable-length list instead of guessing a fixed size.
continue
# Determine target type based on value type
if pa.types.is_floating(field.type.value_type):
target_type = pa.list_(pa.float32(), dim)
@@ -5367,6 +5426,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)
@@ -5374,9 +5434,14 @@ def _infer_target_schema(
return schema, reader
def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> int:
# Use the most common length of the list as the dimensions
return pc.mode(pc.list_value_length(arr))[0].as_py()["mode"]
def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> Optional[int]:
# Use the most common length of the list as the dimensions. pyarrow's
# `mode` skips nulls, so an all-null column yields an empty result; return
# None in that case so the caller can keep a variable-length list.
modes = pc.mode(pc.list_value_length(arr))
if len(modes) == 0:
return None
return modes[0].as_py()["mode"]
def _infer_vector_dim(arr: Union[pa.Array, pa.ChunkedArray]) -> Optional[int]:
@@ -6213,8 +6278,12 @@ class AsyncTable:
if mode == "overwrite":
# For overwrite, apply the same preprocessing as create_table
# so vector columns are inferred as FixedSizeList.
data, _ = sanitize_create_table(
data, None, on_bad_vectors=on_bad_vectors, fill_value=fill_value
data, _ = await asyncio.to_thread(
sanitize_create_table,
data,
None,
on_bad_vectors=on_bad_vectors,
fill_value=fill_value,
)
if blob_v2_column_paths(schema):
# Plain bytes cannot be inferred as blob v2. Keep matching blob
@@ -6224,7 +6293,8 @@ class AsyncTable:
elif on_bad_vectors != "error" or (
schema.metadata is not None and b"embedding_functions" in schema.metadata
):
data = _sanitize_data(
data = await asyncio.to_thread(
_sanitize_data,
data,
schema,
metadata=schema.metadata,
@@ -6233,7 +6303,10 @@ class AsyncTable:
allow_subschema=True,
)
_register_optional_converters()
data = to_scannable(data)
# Converting an iterator peeks at its first item. A synchronous query in
# that iterator schedules work on LOOP, so peeking on LOOP's own thread
# would deadlock waiting for itself.
data = await asyncio.to_thread(to_scannable, data)
progress, owns = _normalize_progress(progress)
try:
return await self._inner.add(
@@ -7454,7 +7527,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)
python/python/lancedb/util.py
+29 -9
View File
@@ -34,20 +34,20 @@ def safe_import_adlfs():
adlfs = safe_import_adlfs()
def get_uri_scheme(uri: str) -> str:
def get_uri_scheme(uri: Union[str, pathlib.Path]) -> str:
"""
Get the scheme of a URI. If the URI does not have a scheme, assume it is a file URI.
Parameters
----------
uri : str
uri : str or pathlib.Path
The URI to parse.
Returns
-------
str: The scheme of the URI.
"""
parsed = urlparse(uri)
parsed = urlparse(str(uri))
scheme = parsed.scheme
if not scheme:
scheme = "file"
@@ -61,25 +61,26 @@ def get_uri_scheme(uri: str) -> str:
return scheme
def get_uri_location(uri: str) -> str:
def get_uri_location(uri: Union[str, pathlib.Path]) -> str:
"""
Get the location of a URI. If the parameter is not a url, assumes it is just a path
Parameters
----------
uri : str
uri : str or pathlib.Path
The URI to parse.
Returns
-------
str: Location part of the URL, without scheme
"""
parsed = urlparse(uri)
uri_str = str(uri)
parsed = urlparse(uri_str)
if len(parsed.scheme) == 1:
# Windows drive names are parsed as the scheme
# e.g. "c:\path" -> ParseResult(scheme="c", netloc="", path="/path", ...)
# So we add special handling here for schemes that are a single character
return uri
return uri_str
if not parsed.netloc:
return parsed.path
@@ -309,7 +310,16 @@ def _query_vector_dim(query: Optional[Any]) -> Optional[int]:
if query.ndim == 0:
return None
return query.shape[-1]
if isinstance(query, list) and query:
if isinstance(query, (pa.Array, pa.ChunkedArray)):
if len(query) == 0:
return None
if pa.types.is_fixed_size_list(query.type):
return query.type.list_size
if pa.types.is_list(query.type) or pa.types.is_large_list(query.type):
first = query[0].as_py()
return len(first) if first is not None else None
return len(query)
if isinstance(query, (list, tuple)) and query:
first = query[0]
if isinstance(first, (list, tuple, np.ndarray)):
return len(first)
@@ -416,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)) + "]"
@@ -461,6 +472,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
python/python/tests/test_closed_table_query.py
+32
View File
@@ -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()
python/python/tests/test_db.py
+3
View File
@@ -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):
python/python/tests/test_duckdb.py
+16
View File
@@ -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,)]
python/python/tests/test_embedding_field_metadata.py
+84
View File
@@ -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"
python/python/tests/test_embedding_registry.py
+88
View File
@@ -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"
python/python/tests/test_embeddings.py
+419
View File
@@ -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,363 @@ 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
# Fake exception types, avoids requiring the optional openai/httpx deps
class FakeBadRequestError(Exception):
pass
class FakeAuthenticationError(Exception):
pass
fake_openai_module = MagicMock()
fake_openai_module.BadRequestError = FakeBadRequestError
fake_openai_module.AuthenticationError = FakeAuthenticationError
registry = get_registry()
model = registry.get("openai").create(name="text-embedding-ada-002")
def fake_create(input, model=None, **kwargs):
# Full batch always rejected
if len(input) > 1:
raise FakeBadRequestError("bad request")
if input[0] == "bad text":
raise FakeBadRequestError("bad request")
embedding = MagicMock()
embedding.embedding = [0.1, 0.2, 0.3]
result = MagicMock()
result.data = [embedding]
return result
mock_client = MagicMock()
mock_client.embeddings.create.side_effect = fake_create
with (
patch(
"lancedb.embeddings.openai.attempt_import_or_raise",
return_value=fake_openai_module,
),
patch.object(type(model), "_openai_client", new_callable=lambda: mock_client),
):
embeddings = model.generate_embeddings(["good text", "bad text"])
assert embeddings[0] == [0.1, 0.2, 0.3]
assert embeddings[1] is None
def test_ollama_embeddings_pickle():
"""OllamaEmbeddings must pickle even after the cached client is created."""
registry = get_registry()
@@ -663,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"]
python/python/tests/test_expr.py
+61
View File
@@ -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):
python/python/tests/test_header_provider.py
+24
View File
@@ -114,6 +114,30 @@ class TestOAuthProvider:
"x-lancedb-credential-type": "oidc",
}
def test_zero_lifetime_token_is_refreshed(self):
"""A token that expires immediately must not be cached indefinitely."""
call_count = 0
def fetcher():
nonlocal call_count
call_count += 1
return {
"access_token": f"token{call_count}",
"expires_in": 0,
}
provider = OAuthProvider(fetcher, refresh_buffer_seconds=0)
assert provider.get_headers() == {
"Authorization": "Bearer token1",
"x-lancedb-credential-type": "oidc",
}
assert provider.get_headers() == {
"Authorization": "Bearer token2",
"x-lancedb-credential-type": "oidc",
}
assert call_count == 2
def test_missing_access_token(self):
"""Test error handling when access_token is missing."""
python/python/tests/test_huggingface.py
+103
View File
@@ -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
python/python/tests/test_hybrid_query.py
+12
View File
@@ -536,3 +536,15 @@ def test_normalize_scores():
assert pc.equal(result, expected), (
f"Expected {expected} but got {result} for invert={invert}"
)
def test_hybrid_query_rejects_query_with_vector(sync_table: Table):
# Pins the full message: the space between "hybrid search." and "But" is
# load-bearing (see the _validate_query literal concatenation).
with pytest.raises(ValueError) as excinfo:
sync_table.search("dog", query_type="hybrid").vector([0.5, -0.5]).to_arrow()
assert str(excinfo.value) == (
"You can either provide a string query in search() method "
"or set `vector()` and `text()` explicitly for hybrid search. "
"But not both."
)
python/python/tests/test_index.py
+22
View File
@@ -430,6 +430,28 @@ async def test_create_ivfrq_index(some_table: AsyncTable):
assert indices[0].name == "vector_idx"
@pytest.mark.asyncio
async def test_read_ivfrq_index_created_with_lance(db_async, some_table: AsyncTable):
pytest.importorskip("lance")
dataset = await some_table.to_lance()
dataset.create_index(
"vector",
index_type="IVF_RQ",
metric="cosine",
num_partitions=1,
)
# Open a fresh LanceDB handle so the index is discovered from disk.
table = await db_async.open_table("some_table")
indices = await table.list_indices()
assert len(indices) == 1
assert indices[0].index_type == "IvfRq"
assert indices[0].columns == ["vector"]
results = await table.query().nearest_to(list(range(DIM))).limit(8).to_arrow()
assert len(results) == 8
@pytest.mark.asyncio
async def test_create_hnswpq_index(some_table: AsyncTable):
await some_table.create_index("vector", config=HnswPq(num_partitions=10))
python/python/tests/test_listing.py
+137
View File
@@ -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])
python/python/tests/test_namespace.py
+120
View File
@@ -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:
python/python/tests/test_ollama_embeddings.py
+126
View File
@@ -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
python/python/tests/test_openai_reranker.py
+66
View File
@@ -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())
python/python/tests/test_permutation.py
+7
View File
@@ -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())]
python/python/tests/test_pyarrow.py
+17
View File
@@ -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
python/python/tests/test_pydantic.py
+110
View File
@@ -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
python/python/tests/test_pydantic_excluded_fields.py
+128
View File
@@ -0,0 +1,128 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import lancedb
import pytest
from pydantic import BaseModel, Field
from lancedb.pydantic import LanceModel
class Details(BaseModel):
label: str
transient: str = Field(default="nested default", exclude=True)
class Row(LanceModel):
id: int
details: Details
transient: str = Field(default="local default", exclude=True)
generated: list[str] = Field(default_factory=lambda: ["local"], exclude=True)
retained: str = Field(default="persisted", exclude=False)
def rows():
return [
Row(
id=i,
details=Details(label=str(i), transient="private"),
transient="private",
generated=["private"],
)
for i in (1, 2)
]
def assert_schema(schema):
assert schema.names == ["id", "details", "retained"]
assert [field.name for field in schema.field("details").type] == ["label"]
assert Row.field_names() == ["id", "details", "transient", "generated", "retained"]
def assert_restored(restored):
assert sorted(row.id for row in restored) == [1, 2]
for row in restored:
assert row.transient == "local default"
assert row.generated == ["local"]
assert row.details.transient == "nested default"
assert row.details.label == str(row.id)
assert row.retained == "persisted"
@pytest.mark.parametrize("explicit_schema", [False, True])
def test_excluded_fields_round_trip(tmp_path, explicit_schema):
data = rows()
kwargs = {"schema": Row} if explicit_schema else {}
db = lancedb.connect(tmp_path)
table = db.create_table("rows", data=data[:1], **kwargs)
table.add(data[1:])
assert_schema(table.schema)
reopened = lancedb.connect(tmp_path).open_table("rows")
assert_restored(reopened.search().to_pydantic(Row))
assert data[0].transient == "private"
assert data[0].generated == ["private"]
@pytest.mark.asyncio
@pytest.mark.parametrize("explicit_schema", [False, True])
async def test_excluded_fields_round_trip_async(tmp_path, explicit_schema):
data = rows()
kwargs = {"schema": Row} if explicit_schema else {}
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("rows", data=data[:1], **kwargs)
await table.add(data[1:])
assert_schema(await table.schema())
reopened_db = await lancedb.connect_async(tmp_path)
reopened = await reopened_db.open_table("rows")
assert_restored(await reopened.query().to_pydantic(Row))
assert data[0].details.transient == "private"
@pytest.mark.parametrize("excluded", ["transient", "source", "vector"])
def test_excluded_fields_keep_embedding_metadata_consistent(
tmp_path, monkeypatch, excluded
):
from lancedb.embeddings.base import TextEmbeddingFunction
from lancedb.embeddings.registry import get_registry
from lancedb.pydantic import Vector
calls = []
class LocalEmbedding(TextEmbeddingFunction):
def ndims(self):
return 2
def generate_embeddings(self, texts):
calls.extend(texts)
return [[1.0, 0.0] for _ in texts]
registry = get_registry()
monkeypatch.setattr(registry, "_functions", registry._functions.copy())
registry.register("test-excluded-fields")(LocalEmbedding)
function = LocalEmbedding.create(max_retries=0)
class EmbeddedRow(LanceModel):
text: str = function.SourceField(default="alpha", exclude=excluded == "source")
vector: Vector(2) = function.VectorField(
default_factory=lambda: [1.0, 0.0], exclude=excluded == "vector"
)
transient: str = Field(default="local", exclude=True)
configs = EmbeddedRow.parse_embedding_functions()
assert len(configs) == (1 if excluded == "transient" else 0)
db = lancedb.connect(tmp_path)
data = [{"text": "alpha"}] if excluded == "transient" else [EmbeddedRow()]
table = db.create_table("embeddings", schema=EmbeddedRow, data=data)
assert "transient" not in table.schema.names
assert ("text" in table.schema.names) == (excluded != "source")
assert ("vector" in table.schema.names) == (excluded != "vector")
reopened = lancedb.connect(tmp_path).open_table("embeddings")
assert len(reopened.embedding_functions) == len(configs)
restored = reopened.search().to_pydantic(EmbeddedRow)
assert restored[0].text == "alpha"
assert restored[0].vector == [1.0, 0.0]
assert restored[0].transient == "local"
if excluded == "transient":
assert calls == ["alpha"]
assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha"
else:
assert calls == []
python/python/tests/test_query.py
+22
View File
@@ -833,6 +833,28 @@ def test_with_row_id(table: lancedb.table.Table):
assert rs["_rowid"].to_pylist() == [0, 1]
def test_filter_preserves_deeply_nested_struct_fields(tmp_db):
# Regression test for https://github.com/lancedb/lancedb/issues/2217.
data = [
{
"item": "test",
"sub_obj": {
"second_sub_obj": [
{
"third_sub_obj": {"int_val": 1},
"str_val": "test",
}
]
},
}
]
table = tmp_db.create_table("deeply_nested_struct", data=data)
result = table.search().where("item = 'test'").limit(1).to_arrow()
assert result.to_pylist() == data
def test_blob_v2_query_omits_auto_row_id(tmp_db):
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid")
python/python/tests/test_read_consistency_interval.py
+83
View File
@@ -0,0 +1,83 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from datetime import timedelta
import lancedb
import pytest
from lance_namespace import connect as connect_namespace_client
from lancedb import _lancedb
from lancedb.namespace import LanceNamespaceDBConnection
@pytest.mark.parametrize("seconds", [-1.0, -0.000001])
@pytest.mark.parametrize("kind", ["sync", "namespace", "namespace_async", "client"])
def test_negative_read_consistency_interval(tmp_path, seconds, kind):
interval = timedelta(seconds=seconds)
with pytest.raises(ValueError, match="read_consistency_interval"):
if kind == "sync":
lancedb.connect(tmp_path, read_consistency_interval=interval)
elif kind == "client":
client = connect_namespace_client("dir", {"root": str(tmp_path)})
LanceNamespaceDBConnection(client, read_consistency_interval=interval)
else:
connect = (
lancedb.connect_namespace
if kind == "namespace"
else lancedb.connect_namespace_async
)
connect("dir", {"root": str(tmp_path)}, read_consistency_interval=interval)
@pytest.mark.asyncio
@pytest.mark.parametrize("seconds", [-1.0, -0.000001])
async def test_negative_async_read_consistency_interval(tmp_path, seconds):
with pytest.raises(ValueError, match="read_consistency_interval"):
await lancedb.connect_async(
tmp_path, read_consistency_interval=timedelta(seconds=seconds)
)
@pytest.mark.asyncio
@pytest.mark.parametrize("seconds", [float("nan"), float("inf"), -float("inf"), 1e30])
@pytest.mark.parametrize("kind", ["connect", "namespace", "client"])
async def test_binding_rejects_invalid_read_consistency_interval(
tmp_path, seconds, kind
):
with pytest.raises(ValueError, match="read_consistency_interval"):
if kind == "connect":
await _lancedb.connect(str(tmp_path), read_consistency_interval=seconds)
elif kind == "namespace":
_lancedb.connect_namespace(
"dir", {"root": str(tmp_path)}, read_consistency_interval=seconds
)
else:
client = connect_namespace_client("dir", {"root": str(tmp_path)})
_lancedb.connect_namespace_client(client, read_consistency_interval=seconds)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"interval", [None, timedelta(0), timedelta(microseconds=1), timedelta(seconds=5)]
)
async def test_valid_read_consistency_intervals(tmp_path, interval):
db = lancedb.connect(tmp_path / "sync", read_consistency_interval=interval)
assert db.read_consistency_interval == interval
db.close()
db = await lancedb.connect_async(
tmp_path / "async", read_consistency_interval=interval
)
assert await db.get_read_consistency_interval() == interval
db.close()
db = lancedb.connect_namespace(
"dir", {"root": str(tmp_path / "namespace")}, read_consistency_interval=interval
)
assert db.read_consistency_interval == interval
db.close()
client = connect_namespace_client("dir", {"root": str(tmp_path / "client")})
db = LanceNamespaceDBConnection(client, read_consistency_interval=interval)
assert db.read_consistency_interval == interval
db.close()
python/python/tests/test_remote_db.py
+73
View File
@@ -1793,6 +1793,79 @@ def test_table_wait_for_index_timeout():
table.wait_for_index(["id_idx"], timedelta(seconds=1))
def _jieba_fts_handler(request):
"""Serve a table with a ``jieba/default`` FTS index and no other state."""
if request.path == "/v1/table/test/describe/":
body = {
"version": 1,
"schema": {
"fields": [
{"name": "text", "type": {"type": "string"}, "nullable": True}
]
},
}
elif request.path == "/v1/table/test/index/list/":
index_details = {
"base_tokenizer": "jieba/default",
"language": "English",
"with_position": False,
"max_token_length": 40,
"lower_case": True,
"stem": False,
"remove_stop_words": False,
"ascii_folding": False,
}
body = {
"indexes": [
{
"index_name": "text_idx",
"columns": ["text"],
"index_type": "FTS",
"index_details": json.dumps(index_details),
}
]
}
else:
request.send_response(404)
request.end_headers()
return
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(body).encode())
_MODEL_FILES_ERROR = (
r"(?s)Invalid directory path.*"
r"full text search index 'text_idx'.*"
r"tokenize\(\) runs on the client.*'jieba/default'.*"
r"LANCE_LANGUAGE_MODEL_HOME.*match the files the server uses"
)
def test_remote_tokenize_model_backed_without_model_files(tmp_path, monkeypatch):
monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path))
with mock_lancedb_connection(_jieba_fts_handler) as db:
table = db.open_table("test")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
table.tokenize("我们是中国人", index_name="text_idx")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
table.tokenize("我们是中国人", column="text")
@pytest.mark.asyncio
async def test_remote_tokenize_model_backed_without_model_files_async(
tmp_path, monkeypatch
):
monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path))
async with mock_lancedb_connection_async(_jieba_fts_handler) as db:
table = await db.open_table("test")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
await table.tokenize("我们是中国人", index_name="text_idx")
with pytest.raises(ValueError, match=_MODEL_FILES_ERROR):
await table.tokenize("我们是中国人", column="text")
def test_stats():
stats = {
"total_bytes": 38,
python/python/tests/test_reranker_credentials.py
+74
View File
@@ -0,0 +1,74 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import sys
import types
import pytest
from lancedb.rerankers import CohereReranker, VoyageAIReranker
RERANKERS = [
pytest.param("cohere", "COHERE_API_KEY", CohereReranker, {}, id="cohere"),
pytest.param(
"voyageai",
"VOYAGE_API_KEY",
VoyageAIReranker,
{"model_name": "rerank-2.5"},
id="voyageai",
),
]
def fake_sdk(monkeypatch, module_name):
"""Replace the SDK module with one whose Client records the API key."""
keys = []
def client(api_key=None):
keys.append(api_key)
return object()
module = types.ModuleType(module_name)
module.Client = client
monkeypatch.setitem(sys.modules, module_name, module)
return keys
@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS)
@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_reranker_api_key_precedence(
monkeypatch,
module_name,
env_var,
reranker_cls,
kwargs,
explicit_key,
environment_key,
expected_key,
):
keys = fake_sdk(monkeypatch, module_name)
if environment_key is None:
monkeypatch.delenv(env_var, raising=False)
else:
monkeypatch.setenv(env_var, environment_key)
reranker_cls(api_key=explicit_key, **kwargs)._client
assert keys == [expected_key]
@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS)
def test_reranker_requires_api_key(
monkeypatch, module_name, env_var, reranker_cls, kwargs
):
keys = fake_sdk(monkeypatch, module_name)
monkeypatch.delenv(env_var, raising=False)
with pytest.raises(ValueError, match=f"{env_var} not set"):
reranker_cls(**kwargs)._client
assert keys == []
python/python/tests/test_rerankers.py
+148
View File
@@ -785,6 +785,154 @@ def test_empty_hybrid_result_reranker():
assert "_rowid" in result.column_names
@pytest.mark.parametrize("norm", ["score", "rank"])
def test_hybrid_return_all_restores_original_scores(norm):
from lancedb.query import LanceHybridQueryBuilder
vector_results = pa.table(
{
"text": ["a", "b", "c"],
"_rowid": pa.array([0, 1, 2], type=pa.uint64()),
"_distance": pa.array([0.5, 2.5, 10.0], type=pa.float32()),
}
)
fts_results = pa.table(
{
"text": ["b", "a"],
"_rowid": pa.array([1, 0], type=pa.uint64()),
"_score": pa.array([3.25, 1.75], type=pa.float32()),
}
)
result = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
vector_results=vector_results,
norm=norm,
fts_query="query",
reranker=RRFReranker(return_score="all"),
limit=10,
with_row_ids=True,
)
# The raw distances and scores are returned, not the normalized values
# or ranks used for reranking.
rows = {row["_rowid"]: row for row in result.to_pylist()}
assert rows[0]["_distance"] == 0.5
assert rows[1]["_distance"] == 2.5
assert rows[2]["_distance"] == 10.0
assert rows[0]["_score"] == 1.75
assert rows[1]["_score"] == 3.25
assert rows[2]["_score"] is None
@pytest.mark.parametrize("return_score", ["relevance", "all"])
@pytest.mark.parametrize("first_row_id", [1, 2**63])
@pytest.mark.parametrize("score_type", [pa.float32(), pa.float64()])
def test_model_reranker_preserves_arrow_schema(
mocker, return_score, first_row_id, score_type
):
common_schema = pa.schema(
[
pa.field("_rowid", pa.uint64()),
pa.field("text", pa.string()),
pa.field("vector", pa.list_(pa.float32(), 2)),
],
metadata={b"source": b"reranker-test"},
)
vector_results = pa.Table.from_pylist(
[
{
"_rowid": first_row_id,
"text": "vector",
"vector": [1, 0],
"_distance": 0.1,
},
{
"_rowid": first_row_id + 1,
"text": "both",
"vector": [0, 1],
"_distance": 0.2,
},
],
schema=common_schema.append(pa.field("_distance", score_type)),
)
fts_results = pa.Table.from_pylist(
[
{
"_rowid": first_row_id + 1,
"text": "both",
"vector": [0, 1],
"_score": 0.9,
},
{
"_rowid": first_row_id + 2,
"text": "fts",
"vector": [1, 1],
"_score": 0.8,
},
],
schema=common_schema.append(pa.field("_score", score_type)),
)
reranker = CohereReranker(return_score=return_score)
reranker._client = mocker.Mock()
reranker._client.rerank.return_value.results = [
mocker.Mock(index=index, relevance_score=score)
for index, score in [(2, 0.9), (1, 0.8), (0, 0.7)]
]
result = reranker.rerank_hybrid("query", vector_results, fts_results)
assert result.select(common_schema.names).schema.equals(
common_schema, check_metadata=True
)
assert result["_rowid"].to_pylist() == [
first_row_id + 2,
first_row_id + 1,
first_row_id,
]
assert result["text"].to_pylist() == ["fts", "both", "vector"]
if return_score == "all":
assert (
result["_distance"]
.combine_chunks()
.equals(pa.array([None, 0.2, 0.1], type=score_type))
)
assert (
result["_score"]
.combine_chunks()
.equals(pa.array([0.8, 0.9, None], type=score_type))
)
@pytest.mark.parametrize("return_score", ["relevance", "all"])
def test_model_reranker_empty_results_preserve_schema(mocker, return_score):
common_schema = pa.schema(
[
("_rowid", pa.uint64()),
("text", pa.string()),
("vector", pa.list_(pa.float32(), 2)),
],
metadata={b"source": b"reranker-test"},
)
vector_results = pa.Table.from_pylist(
[], schema=common_schema.append(pa.field("_distance", pa.float32()))
)
fts_results = pa.Table.from_pylist(
[], schema=common_schema.append(pa.field("_score", pa.float32()))
)
reranker = CohereReranker(return_score=return_score)
reranker._client = mocker.Mock()
result = reranker.rerank_hybrid("query", vector_results, fts_results)
assert len(result) == 0
assert result.select(common_schema.names).schema.equals(
common_schema, check_metadata=True
)
assert result.schema.field("_relevance_score").type == pa.float32()
reranker._client.rerank.assert_not_called()
def test_cross_encoder_reranker_return_all(tmp_path):
pytest.importorskip("sentence_transformers")
reranker = CrossEncoderReranker(return_score="all")
python/python/tests/test_session.py
+10
View File
@@ -36,3 +36,13 @@ def test_session_cache_configuration(tmp_path):
assert final_cache_size > initial_cache_size # Cache should have grown
assert final_cache_items >= initial_cache_items # Items should not decrease
assert initial_cache_size < index_cache_size + metadata_cache_size
def test_connection_session_returns_supplied_session(tmp_path):
"""Sync connections must expose the session they were opened with."""
db = lancedb.connect(tmp_path)
assert db.session is None
session = lancedb.Session()
db = lancedb.connect(tmp_path, session=session)
assert db.session is session
python/python/tests/test_table.py
+327 -1
View File
@@ -15,7 +15,7 @@ from datetime import date, datetime, timedelta
from decimal import Decimal, localcontext
from time import sleep
from typing import List
from unittest.mock import patch
from unittest.mock import MagicMock, patch
import lancedb
from lancedb import table as table_module
@@ -280,6 +280,30 @@ def test_create_table_infers_large_int_vectors(mem_db: DBConnection):
assert vector_column.to_pylist() == [[0.0, 300.0]]
def test_create_table_all_null_vector_column(mem_db: DBConnection):
# A vector column can legitimately arrive with every value null, e.g. when
# the embeddings have not been computed yet. The dimension is unknowable in
# that case, so schema inference must not crash and should leave the column
# as a variable-length list rather than a fixed-size one.
data = pa.table(
{
"id": pa.array([1, 2]),
"vector": pa.array([None, None], type=pa.list_(pa.float32())),
}
)
table = mem_db.create_table("all_null_vector", data=data)
vector_field = table.schema.field("vector")
assert pa.types.is_list(vector_field.type)
assert vector_field.type.value_type == pa.float32()
result = table.to_arrow()
assert result.num_rows == 2
assert result.column("vector").to_pylist() == [None, None]
assert table.count_rows() == 2
@pytest.mark.asyncio
async def test_create_table_async_infers_large_int_vectors(
mem_db_async: AsyncConnection,
@@ -518,6 +542,30 @@ def test_add(mem_db: DBConnection):
_add(table, schema)
def test_add_from_iterator_that_queries_table(mem_db: DBConnection):
source = mem_db.create_table("source", data=pa.table({"id": range(16)}))
target = mem_db.create_table("target", schema=source.schema)
def batches():
for _ in range(5):
yield source.search().limit(10).to_arrow()
target.add(batches())
assert target.count_rows() == 50
def test_create_table_from_iterator_that_queries_table(mem_db: DBConnection):
source = mem_db.create_table("source", data=pa.table({"id": range(16)}))
def batches():
yield source.search().limit(10).to_arrow()
target = mem_db.create_table("target", data=batches())
assert target.count_rows() == 10
def test_add_releases_arrow_buffers_without_gc(mem_db: DBConnection):
"""Regression test for https://github.com/lancedb/lancedb/issues/2512."""
schema = pa.schema([pa.field("x", pa.int64())])
@@ -1585,6 +1633,25 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection):
job.cancel()
@pytest.mark.parametrize("method", ["create_index", "create_index_async"])
@pytest.mark.parametrize("wait_timeout", ["invalid", 1, []])
def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout):
table = mem_db.create_table("invalid_timeout", [{"id": 1}])
with pytest.raises(TypeError):
getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout)
assert table.list_indices() == []
@pytest.mark.parametrize("method", ["create_index", "create_index_async"])
@pytest.mark.parametrize("wait_timeout", [None, timedelta(seconds=1)])
def test_create_index_valid_wait_timeout(mem_db: DBConnection, method, wait_timeout):
table = mem_db.create_table("valid_timeout", [{"id": 1}])
result = getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout)
if result is not None:
result.wait()
assert len(table.list_indices()) == 1
@patch("lancedb.table.AsyncTable.create_index")
def test_create_index_method(mock_create_index, mem_db: DBConnection):
table = mem_db.create_table(
@@ -1725,6 +1792,51 @@ def test_create_index_method(mock_create_index, mem_db: DBConnection):
)
@pytest.mark.parametrize("api", ["legacy", "unified"])
@pytest.mark.parametrize("cuda_failure", ["unavailable", "init_error"])
def test_create_index_cuda_failure_suggests_cpu(
api, cuda_failure, mem_db: DBConnection
):
table = mem_db.create_table(
f"test_cuda_unavailable_{api}",
data=[{"vector": [3.1, 4.1]}, {"vector": [5.9, 26.5]}],
)
torch = MagicMock()
torch.cuda.is_available.return_value = cuda_failure == "init_error"
if cuda_failure == "init_error":
torch.cuda.init.side_effect = RuntimeError("CUDA driver initialization failed")
with (
patch.dict(sys.modules, {"torch": torch}),
patch.object(table, "to_lance") as to_lance,
):
if api == "legacy":
with (
pytest.warns(DeprecationWarning, match="create_index"),
pytest.raises(RuntimeError, match="accelerator=None.*CPU"),
):
table.create_index(
metric="dot",
num_partitions=1,
num_sub_vectors=1,
vector_column_name="vector",
accelerator="cuda",
)
else:
with pytest.raises(RuntimeError, match="accelerator=None.*CPU"):
table.create_index(
"vector",
config=IvfPq(
distance_type="dot",
num_partitions=1,
num_sub_vectors=1,
accelerator="cuda",
),
)
to_lance.return_value.create_index.assert_not_called()
@patch("lancedb.table.AsyncTable.create_index")
def test_create_index_name_and_train_parameters(
mock_create_index, mem_db: DBConnection
@@ -3673,6 +3785,42 @@ def test_multiple_vector_columns(mem_db: DBConnection):
assert result1["text"][0] != result2["text"][0]
@pytest.mark.parametrize(
"query",
[
(0.0, 0.0),
pa.array([0.0, 0.0]),
pa.chunked_array([[0.0], [0.0]]),
],
)
@pytest.mark.parametrize("query_type", ["auto", "vector"])
def test_sync_search_accepts_supported_vector_containers(
mem_db: DBConnection, query, query_type
):
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("small", pa.list_(pa.float32(), 2)),
pa.field("large", pa.list_(pa.float32(), 3)),
]
)
table = mem_db.create_table(
"supported_vector_containers",
data=pa.table(
{
"id": [1, 2],
"small": [[0.0, 0.0], [1.0, 1.0]],
"large": [[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]],
},
schema=schema,
),
)
result = table.search(query, query_type=query_type).limit(1).to_arrow()
assert result["id"].to_pylist() == [1]
def test_create_scalar_index(mem_db: DBConnection):
vec_array = pa.array(
[[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]], pa.list_(pa.float32(), 2)
@@ -4378,6 +4526,58 @@ async def test_alter_columns_rejects_invalid_async(
assert await table.schema() == initial_schema
@pytest.mark.asyncio
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("list_type", [pa.list_, pa.large_list])
@pytest.mark.parametrize(
"value_type, values, inferred_type",
[
(pa.float64(), [1.25, 2.5], pa.float32()),
(pa.int32(), [0, 255], pa.uint8()),
(pa.int32(), [-1, 256], pa.float32()),
],
)
async def test_inferred_vector_field_metadata(
tmp_db: DBConnection,
tmp_db_async: AsyncConnection,
asynchronous,
list_type,
value_type,
values,
inferred_type,
):
metadata = {
"lancedb:description": "Embedding values",
"lancedb:tag:model": "test-model",
}
schema = pa.schema(
[
pa.field(
"vector", list_type(value_type), nullable=False, metadata=metadata
),
pa.field("text", pa.string(), metadata={"description": "Source text"}),
],
metadata={"description": "Source dataset"},
)
rows = [{"vector": values, "text": "first"}, {"vector": values, "text": "second"}]
data = pa.Table.from_pylist(rows, schema=schema)
if asynchronous:
await tmp_db_async.create_table("vector_metadata", data=data)
table = await tmp_db_async.open_table("vector_metadata")
actual_schema = await table.schema()
actual_data = await table.to_arrow()
else:
tmp_db.create_table("vector_metadata", data=data)
table = tmp_db.open_table("vector_metadata")
actual_schema = table.schema
actual_data = table.to_arrow()
expected_schema = schema.set(
0, schema.field("vector").with_type(pa.list_(inferred_type, 2))
)
assert actual_schema.equals(expected_schema, check_metadata=True)
assert actual_data.to_pylist() == rows
def test_update_field_metadata(mem_db: DBConnection):
data = pa.table({"id": [0, 1], "category": ["a", "b"]})
table = mem_db.create_table("my_table", data=data)
@@ -4592,6 +4792,57 @@ def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection):
assert all("concurrent" in str(w.message) for w in seen)
def test_optimize_retains_tagged_versions(tmp_db: DBConnection):
# A tagged version older than the cleanup cutoff must be retained, not make
# optimize() fail part-way (after compaction, before the index step).
table = tmp_db.create_table(
"test", data=pa.table({"id": list(range(100)), "text": ["a"] * 100})
)
table.create_index("id", config=BTree())
table.tags.create("release-1", table.version)
tagged_version = table.version
for i in range(3):
table.add(pa.table({"id": [1000 + i], "text": ["new"]}))
assert table.stats()["fragment_stats"]["num_fragments"] == 4
assert table.index_stats("id_idx").num_unindexed_rows == 3
sleep(1.1)
table.optimize(cleanup_older_than=timedelta(seconds=1))
# The index step ran after the prune step
assert table.index_stats("id_idx").num_unindexed_rows == 0
# It stays optimizable
table.optimize(cleanup_older_than=timedelta(seconds=1))
assert table.stats()["fragment_stats"]["num_fragments"] == 1
# Untagged old versions were pruned, the tagged one was retained
versions = [v["version"] for v in table.list_versions()]
assert tagged_version in versions
for added in range(tagged_version + 1, tagged_version + 4):
assert added not in versions
assert len(table) == 103
table.checkout("release-1")
assert len(table) == 100
assert table.to_arrow().num_rows == 100
@pytest.mark.asyncio
async def test_optimize_retains_tagged_versions_async(tmp_db_async: AsyncConnection):
table = await tmp_db_async.create_table("test", data=[{"x": [1]}])
await table.tags.create("keep", await table.version())
await table.add(data=[{"x": [2]}])
await table.add(data=[{"x": [3]}])
sleep(1.1)
stats = await table.optimize(cleanup_older_than=timedelta(seconds=1))
assert stats.prune.old_versions_removed >= 1
await table.checkout("keep")
assert await table.count_rows() == 1
assert await table.query().to_arrow() == pa.table({"x": [[1]]})
def test_replace_field_metadata(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("my_table", data=[{"x": 0}])
@@ -4647,6 +4898,81 @@ def test_stats(mem_db: DBConnection):
assert stats_with_index["fragment_stats"]["num_fragments"] == 1
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("empty", [False, True])
@pytest.mark.parametrize("with_schema", [False, True])
def test_create_table_recordbatch_iterable(
mem_db: DBConnection, iterable_type, empty, with_schema
):
schema = pa.schema([pa.field("id", pa.int64())])
batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)]
data = iterable_type(batches)
if empty and not with_schema:
with pytest.raises(ValueError, match="empty iterable without a schema"):
mem_db.create_table("test", data=data)
return
table = mem_db.create_table(
"test", data=data, schema=schema if with_schema else None
)
assert table.schema == schema
assert table.to_arrow().to_pylist() == ([] if empty else [{"id": 1}])
@pytest.mark.asyncio
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("empty", [False, True])
@pytest.mark.parametrize("with_schema", [False, True])
async def test_create_table_recordbatch_iterable_async(
mem_db_async: AsyncConnection, iterable_type, empty, with_schema
):
schema = pa.schema([pa.field("id", pa.int64())])
batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)]
data = iterable_type(batches)
if empty and not with_schema:
with pytest.raises(ValueError, match="empty iterable without a schema"):
await mem_db_async.create_table("test", data=data)
return
table = await mem_db_async.create_table(
"test", data=data, schema=schema if with_schema else None
)
assert await table.schema() == schema
assert (await table.to_arrow()).to_pylist() == ([] if empty else [{"id": 1}])
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
def test_add_empty_recordbatch_iterable(
mem_db: DBConnection, iterable_type, on_bad_vectors
):
table = mem_db.create_table("test", [{"id": 1}])
if on_bad_vectors == "error":
# The default scannable path still requires an input schema. This
# fix does not change its validation policy for schema-less input.
with pytest.raises(ValueError, match="empty iterator"):
table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
else:
# Sanitization has the existing table schema available.
table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
assert table.to_arrow().to_pylist() == [{"id": 1}]
@pytest.mark.asyncio
@pytest.mark.parametrize("iterable_type", [iter, tuple])
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
async def test_add_empty_recordbatch_iterable_async(
mem_db_async: AsyncConnection, iterable_type, on_bad_vectors
):
table = await mem_db_async.create_table("test", [{"id": 1}])
if on_bad_vectors == "error":
with pytest.raises(ValueError, match="empty iterator"):
await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
else:
await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors)
assert (await table.to_arrow()).to_pylist() == [{"id": 1}]
def test_create_table_empty_list_with_schema(mem_db: DBConnection):
"""Test creating table with empty list data and schema
python/python/tests/test_util.py
+32
View File
@@ -21,6 +21,7 @@ from lancedb.table import (
_infer_target_schema,
_merge_metadata,
_sanitize_data,
has_nan_values,
sanitize_create_table,
)
import pyarrow as pa
@@ -30,6 +31,7 @@ import pytest
import lancedb
from lancedb.util import (
flatten_columns,
get_uri_location,
get_uri_scheme,
inf_vector_column_query,
join_uri,
@@ -86,6 +88,22 @@ def test_normalize_uri():
assert parsed_scheme == expected_scheme
def test_get_uri_scheme_and_location_pathlib():
# Test pathlib.Path instances
abs_path = pathlib.Path("/tmp/data")
rel_path = pathlib.Path("relative/data")
assert get_uri_scheme(abs_path) == "file"
assert get_uri_scheme(rel_path) == "file"
assert get_uri_location(abs_path) == str(abs_path)
assert get_uri_location(rel_path) == str(rel_path)
# String URIs still behave identically
assert get_uri_scheme("s3://bucket/data") == "s3"
assert get_uri_location("s3://bucket/data") == "bucket/data"
def test_join_uri_remote():
schemes = ["s3", "az", "gs"]
for scheme in schemes:
@@ -472,6 +490,20 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
assert output["vector"].combine_chunks() == expected
@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()])
@pytest.mark.parametrize("large", [False, True])
@pytest.mark.parametrize("chunked", [False, True])
def test_has_nan_values_skips_hidden_children_of_null_lists(dtype, large, chunked):
array_type = pa.LargeListArray if large else pa.ListArray
values = pa.array([float("nan"), None, float("nan"), 3.0, 4.0, 5.0], type=dtype)
arr = array_type.from_arrays(
[0, 2, 4, 6], values, mask=pa.array([True, False, False])
)
if chunked:
arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2, 1)])
assert has_nan_values(arr).to_pylist() == [False, True, False]
def test_handle_bad_vectors_noop():
# ChunkedArray should be preserved as-is
vector = pa.chunked_array(
python/python/tests/test_util_update_values.py
+215
View File
@@ -0,0 +1,215 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from datetime import date, datetime
import lancedb
import numpy as np
import pyarrow as pa
import pytest
ARRAY_VALUES = [
((1.0, 2.0), [1.0, 2.0]),
([1.0, 2.0], [1.0, 2.0]),
(np.array([1.0, 2.0]), [1.0, 2.0]),
(None, None),
]
@pytest.mark.parametrize("value,expected", ARRAY_VALUES)
def test_update_array_values(tmp_path, value, expected):
dtype = pa.list_(pa.float64())
db = lancedb.connect(tmp_path)
table = db.create_table(
"values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)})
)
table.update(values={"array": value})
assert table.to_arrow()["array"].to_pylist() == [expected]
@pytest.mark.asyncio
@pytest.mark.parametrize("value,expected", ARRAY_VALUES)
async def test_update_array_values_async(tmp_path, value, expected):
dtype = pa.list_(pa.float64())
db = await lancedb.connect_async(tmp_path)
table = await db.create_table(
"values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)})
)
await table.update({"array": value})
assert (await table.to_arrow())["array"].to_pylist() == [expected]
TIME_VALUES = [
np.datetime64("2025", "Y"),
np.datetime64("2025-01", "M"),
np.datetime64("2025-01-02", "D"),
np.datetime64("2025-01-02T03", "h"),
np.datetime64("2025-01-02T03:04", "m"),
np.datetime64("2025-01-02T03:04:05", "s"),
np.datetime64("2025-01-02T03:04:05.123", "ms"),
np.datetime64("2025-01-02T03:04:05.123456", "us"),
np.datetime64("2025-01-02T03:04:05.123456789", "ns"),
np.datetime64("NaT", "ns"),
datetime(2025, 1, 2, 3, 4, 5, 123456),
None,
]
def expected_time(value):
if value is None or isinstance(value, np.datetime64) and np.isnat(value):
return None
return int(np.datetime64(value, "ns").astype(np.int64))
@pytest.mark.parametrize("value", TIME_VALUES)
def test_update_timestamp_values(tmp_path, value):
db = lancedb.connect(tmp_path)
data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))})
table = db.create_table("times", data)
table.update(values={"time": value})
assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [
expected_time(value)
]
@pytest.mark.asyncio
@pytest.mark.parametrize("value", TIME_VALUES)
async def test_update_timestamp_values_async(tmp_path, value):
db = await lancedb.connect_async(tmp_path)
data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))})
table = await db.create_table("times", data)
await table.update({"time": value})
assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [
expected_time(value)
]
@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []])
def test_nested_array_tuple_matches_list(tmp_path, values):
db = lancedb.connect(tmp_path)
data = pa.table(
{"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))}
)
control = db.create_table("list_values", data)
candidate = db.create_table("tuple_values", data)
tuples = tuple(tuple(value) for value in values)
try:
control.update(values={"array": values})
except ValueError as error:
reason = (
"Expected a literal value in array" if values else "concat requires input"
)
assert reason in str(error)
with pytest.raises(ValueError, match=reason):
candidate.update(values={"array": tuples})
assert control.to_arrow().equals(data)
assert candidate.to_arrow().equals(data)
else:
candidate.update(values={"array": tuples})
assert candidate.to_arrow().equals(control.to_arrow())
@pytest.mark.asyncio
@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []])
async def test_nested_array_tuple_matches_list_async(tmp_path, values):
db = await lancedb.connect_async(tmp_path)
data = pa.table(
{"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))}
)
control = await db.create_table("list_values", data)
candidate = await db.create_table("tuple_values", data)
tuples = tuple(tuple(value) for value in values)
try:
await control.update({"array": values})
except ValueError as error:
reason = (
"Expected a literal value in array" if values else "concat requires input"
)
assert reason in str(error)
with pytest.raises(ValueError, match=reason):
await candidate.update({"array": tuples})
assert (await control.to_arrow()).equals(data)
assert (await candidate.to_arrow()).equals(data)
else:
await candidate.update({"array": tuples})
assert (await candidate.to_arrow()).equals(await control.to_arrow())
@pytest.mark.parametrize("unit", ["ps", "fs", "as"])
def test_subnanosecond_update_does_not_round(tmp_path, unit):
db = lancedb.connect(tmp_path)
data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))})
table = db.create_table("precise", data)
with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"):
table.update(values={"time": np.datetime64(1, unit)})
assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [0]
@pytest.mark.asyncio
@pytest.mark.parametrize("unit", ["ps", "fs", "as"])
async def test_subnanosecond_update_does_not_round_async(tmp_path, unit):
db = await lancedb.connect_async(tmp_path)
data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))})
table = await db.create_table("precise", data)
with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"):
await table.update({"time": np.datetime64(1, unit)})
assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [0]
@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"])
def test_update_numpy_date_column(tmp_path, unit):
value = np.datetime64("2025-01-02", unit)
db = lancedb.connect(tmp_path)
table = db.create_table(
"dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())})
)
table.update(values={"day": value})
assert table.to_arrow()["day"].to_pylist() == [value.astype("datetime64[D]").item()]
@pytest.mark.asyncio
@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"])
async def test_update_numpy_date_column_async(tmp_path, unit):
value = np.datetime64("2025-01-02", unit)
db = await lancedb.connect_async(tmp_path)
table = await db.create_table(
"dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())})
)
await table.update({"day": value})
assert (await table.to_arrow())["day"].to_pylist() == [
value.astype("datetime64[D]").item()
]
@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"])
@pytest.mark.parametrize(
"value",
[np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")],
)
def test_update_numpy_timestamp_matches_arrow_ingestion(tmp_path, timezone, value):
dtype = pa.timestamp("ns", tz=timezone)
ingested = pa.array(np.array([value]), type=dtype)
db = lancedb.connect(tmp_path)
table = db.create_table("timestamps", pa.table({"time": ingested}))
expected = table.to_arrow()["time"].cast(pa.int64()).to_pylist()
table.update(values={"time": value})
assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == expected
@pytest.mark.asyncio
@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"])
@pytest.mark.parametrize(
"value",
[np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")],
)
async def test_update_numpy_timestamp_matches_arrow_ingestion_async(
tmp_path, timezone, value
):
dtype = pa.timestamp("ns", tz=timezone)
ingested = pa.array(np.array([value]), type=dtype)
db = await lancedb.connect_async(tmp_path)
table = await db.create_table("timestamps", pa.table({"time": ingested}))
expected = (await table.to_arrow())["time"].cast(pa.int64()).to_pylist()
await table.update({"time": value})
assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == expected

Some files were not shown because too many files have changed in this diff Show More