Compare commits

..

42 Commits

Author SHA1 Message Date
lancedb-gatefixer[bot] 9d3962686e fix(node): accept Arrow metadata across JavaScript realms (#3904)
## Summary

- accept genuine Arrow metadata maps created in another JavaScript realm
- validate every metadata entry and clone it into a local Map
- cover an Arrow 15 VM-realm table through the public fromDataToBuffer
boundary
- retain structural typing for nested and dictionary Arrow data

## Root cause

The sanitizer used a local-realm instanceof Map check for schema and
field metadata. A genuine Map created in another JavaScript realm has
the required internal Map state but fails that identity check, so
fromDataToBuffer rejected the foreign table before serializing its rows.

## Scope

This fixes the distinct JavaScript-realm sanitizer failure identified
during review. It does not establish the cause of the S3/compaction
panic reported in #1525, so that issue remains open.

## Validation

- pnpm test --runInBand (707 passed, 5 skipped)
- pnpm test --runInBand __test__/arrow.test.ts (189 passed)
- pnpm build
- pnpm lint
- pnpm run docs

Related to #1525

<!-- lance-gatekeeper-fix:v1 agent=b522628ad3bae914eb7266ccd899d508
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 23:35:02 +08:00
lancedb-gatefixer[bot] 25645d82d4 feat(python): accept expressions in update filters (#3876)
## Summary
- allow Python sync, async, and remote table updates to accept type-safe
`Expr` filters
- serialize expression filters before invoking the existing update
implementation
- cover numeric-looking text and apostrophe-containing text in sync and
async regression tests

## Root cause
`Table.update` was the remaining Python write path that required callers
to construct a raw SQL predicate. Dynamic text interpolated without SQL
literal encoding could therefore be parsed as an integer, float, or
unterminated string instead of Utf8. The expression API already encodes
literals safely for query and delete filters.

## Validation
- `cd python && .venv/bin/pytest
python/tests/test_table.py::test_update_async
python/tests/test_table.py::test_update_expr_filter_literals_async
python/tests/test_table.py::test_update
python/tests/test_table.py::test_update_expr_filter_literals -q`
- `cd python && .venv/bin/pytest python/tests/test_expr.py -q`
- `cd python && .venv/bin/ruff format --check .`
- `cd python && .venv/bin/ruff check .`

Fixes #1869

<!-- lance-gatekeeper-fix:v1 agent=01f1e7b69c65e8b6d3b3c1e1a7918179
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 20:28:21 +08:00
lancedb-gatefixer[bot] 0dd9dfdfc7 test(python): cover arithmetic with distance projections (#3862)
## Summary

- add Python regression coverage for integer and double arithmetic
against the generated _distance column
- merge the current main base containing Lance v11.0.0-beta.3 from #3896
- verify both expressions retain the generated scoring field Float32
type and compute the expected values

## Root cause

Lance parsed dynamic projection expressions before vector search added
its generated Float32 _distance field. Without a typed provisional
field, expression discovery rejected mixed numeric arithmetic. Lance
upstream fixed discovery and final-schema replanning in
lance-format/lance#8163, and the current base consumes that fix through
Lance v11.0.0-beta.3.

## Validation

- uv run --extra tests pytest
python/tests/test_query.py::test_select_arithmetic_with_distance -vv
--maxfail=2 — 2 passed
- python/.venv/bin/ruff format --check python/python/tests/test_query.py
- python/.venv/bin/ruff check .

Fixes #2618

<!-- lance-gatekeeper-fix:v1 agent=816a060090517471edfb73652bb5c9fe
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 17:10:51 +08:00
lancedb-gatefixer[bot] d24b2dcacc fix: show nested fields in query schema errors (#3849)
## Summary

- enrich local query field-not-found errors with recursively qualified
Arrow struct leaf paths
- preserve all other Lance and DataFusion errors unchanged
- add a regression test for the Python-visible filter error described in
the issue

## Root cause

DataFusion builds `FieldNotFound` candidates from the top-level Arrow
schema even though Lance supports dotted struct-field filters. As a
result, the error listed only the struct container and hid its valid
nested leaves.

## Validation

- `cargo test --quiet --features remote -p lancedb
table::query::tests::test_missing_filter_field_lists_nested_fields --
--exact`
- `cargo check --quiet --features remote --tests --examples`
- `cargo clippy --quiet --features remote --tests --examples` (passes
with pre-existing unrelated warnings)
- `cargo fmt --all -- --check`

Fixes #951

<!-- lance-gatekeeper-fix:v1 agent=7893f7a181fd8bc1ad00acc62d1a85c2
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 15:47:18 +08:00
lancedb-gatefixer[bot] 2deccf21cf fix(node): read Python embedding metadata (#3836)
## Summary

- normalize Python snake_case and TypeScript camelCase embedding
metadata
- use the normalized metadata for schema validation and embedding lookup
- cover appending through `Table.add()` with a Python-authored schema
fixture

## Root cause

Python writes embedding source and vector column names as
`source_column` and `vector_column`, but the TypeScript SDK only read
`sourceColumn` and `vectorColumn`. The missing source name reached the
add path as `undefined`, preventing JavaScript rows from being embedded
and appended.

## Validation

- `pnpm lint`
- `pnpm test __test__/embedding.test.ts __test__/arrow.test.ts
__test__/registry.test.ts --runInBand` (201 passed, 1 skipped)
- `pnpm build`
- `pnpm run docs`

Fixes #1289

<!-- lance-gatekeeper-fix:v1 agent=b71c18a5e33d26f4d138972e91d34e66
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 13:42:30 +08:00
Lance Release ead4d27bfc Bump version: 0.38.0-beta.10 → 0.38.0-beta.11 2026-08-27 04:31:57 +00:00
lancedb-gatefixer[bot] 5153e5a023 fix(node): preserve JSON field metadata when adding data (#4064)
## Summary

- preserve Arrow field metadata when matching record data to a provided
schema
- retain metadata on partially reconstructed nested struct fields
- add a regression test for lance.json metadata through Arrow IPC
serialization

## Root cause

The TypeScript schema inferrer rebuilt fields selected from a provided
schema without copying their metadata. JSON columns therefore kept their
LargeBinary physical type but lost the lance.json extension marker
before insert, causing the schema mismatch reported in the issue.

## Validation

- pnpm lint
- pnpm build
- pnpm tsc
- pnpm run docs
- pnpm test --runInBand (18 suites and 798 tests passed; 5 tests
skipped)

Fixes #4062

<!-- lance-gatekeeper-fix:v1 agent=3ec52632b71563f53d199b22629f8c4f
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
2026-08-26 16:18:57 -07:00
lancedb-gatefixer[bot] 79f626b09e fix: support double-quoted filter identifiers (#3825)
## Summary

- tokenize predicates with the same GenericDialect lexical rules Lance
delegates to
- rewrite only SQL-standard double-quoted identifier tokens to Lance
backticks
- apply one predicate contract to query, count, update, delete, and both
merge conditions
- cover mixed-case identifiers, ordinary literals, comments, and every
filter-bearing table operation

## Root cause

Lance plans double-quoted tokens as string literals for compatibility.
As a result, `"PartyAbbrev" = 'D'` compared two literals and silently
evaluated to false instead of filtering the mixed-case column.

## Validation

- `cargo fmt --all -- --check`
- `cargo test --locked --quiet --features remote -p lancedb
expr::sql::tests`
- `cargo test --locked --quiet --features remote -p lancedb
test_double_quoted_predicates_across_table_operations`
- `cargo check --quiet --features remote --tests --examples`
- `cargo clippy --quiet --features remote --tests --examples`

Fixes #2057

<!-- lance-gatekeeper-fix:v1 agent=a44b6567cfd8890abb4f7395ff71971a
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 06:17:28 +08:00
lancedb-gatefixer[bot] ae81d73563 fix: share scans across batched vector queries (#3805)
<!-- lance-gatekeeper-fix:v1 agent=d30696bc46eb32f04c9927b0792e35d3
generation=1 -->

## Summary

- use the Lance native batch KNN path so fixed-size batch vector
searches share one flat table scan
- validate consistent query-vector dimensions and retain the per-vector
plan when offsets require its existing semantics
- add Rust and Python regressions and update Rust, Python, and
TypeScript API documentation

## Root cause

LanceDB expanded every vector in a batch into a separate scan plan and
joined the plans with `UnionExec`. For unindexed tables on S3, a batch
of ten vectors therefore ran ten concurrent full scans, amplifying CPU
and retained data enough to produce the reported memory spike.

The native Lance batch KNN path performs bounded-memory selection for
all query vectors over one flat scan. LanceDB now supplies the vectors
as a batch and avoids applying a global scanner limit to the combined
per-query results. Batch queries with a nonzero offset keep the previous
plan because the native batch API does not support per-query offsets.

## Validation

- targeted Rust batch-query plan and execution tests
- `cargo check --quiet --features remote --tests --examples`
- `cargo clippy --quiet --features remote --tests --examples`
- `cargo fmt --all -- --check`
- targeted Python batch-vector regression after rebuilding the extension
- Ruff formatting/checks for the touched Python files
- Node.js build, lint, docs generation, and targeted batch-vector Jest
test
- `git diff --check`

Fixes #2468

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-27 04:23:28 +08:00
Dan Tasse 8b7e13b0c6 docs: add comments about metadata conventions (#4054)
In LanceDB Enterprise, we've adopted these conventions to give some
"canonical" metadata paths. This lets us display them in a certain way
in the UI or let agents standardize on them, to assume they'll find info
in a certain place. This PR (only comments/docs) just documents those
choices.
2026-08-26 14:19:44 -04:00
Xuanwo b78f2a5044 feat: expose list-element FTS document granularity (#4050)
## Summary

LanceDB could not request Lance's list-element FTS document granularity
through Python or Remote APIs, and generic nested-field resolution
exposed Arrow's internal `item` segment instead of the public field
path.

This exposes typed `row | list_element` configuration for Python FTS
index creation and match/phrase queries, preserves `_doc_index`, and
keeps nested FTS paths public (for example, `docs.content`). Remote
list-element requests require server API version 0.6.0 so older servers
cannot silently execute them with row semantics; explicit row requests
remain compatible.

## Compatibility

Omitted index and query parameters retain row granularity. Remote row
index creation omits the new wire field.

## Tracking


[ENT-2342](https://linear.app/lancedb/issue/ENT-2342/expose-list-element-fts-document-granularity-end-to-end)
2026-08-26 23:27:39 +08:00
Wyatt Alt 06872463cf feat: declare conda environments on Functions (#4057)
A Function's remote environment can now be conda instead of pip.
`@udf(conda=[...], conda_channels=[...])` registers one; pip and conda
are exclusive, channels are priority-ordered and require conda. The Rust
and Python `PythonEnvironmentSpec` models gain `channels`, dropped from
the canonical JSON when empty so existing pip registrations keep their
digests.
2026-08-26 06:29:43 -07:00
lancedb-gatefixer[bot] 2fbf6d6211 test(python): cover concurrent S3 table opens (#3833)
## Summary

- add regression coverage for the reported synchronous Python workload
with 32 simultaneous `open_table` calls
- verify every independently opened S3-backed table handle can read
through the connection's shared session and object-store client

## Root cause

In Python v0.13.0, each synchronous table handle lazily constructed its
own Lance dataset. Opening many handles in parallel therefore triggered
independent S3 client construction and bucket-region resolution, which
failed under thread pressure. The current Rust-backed connection path
owns a shared Lance session and retains its object-store handle, so
table opens reuse the existing S3 client; these tests lock in that
behavior through the public Python API and a causal Session-registry
invariant.

## Validation

- `uvx --from 'ruff==0.15.20' ruff format --check
python/tests/test_s3.py`
- `uvx --from 'ruff==0.15.20' ruff check .`
- `cargo fmt --all`
- `cargo test --quiet --features remote -p lancedb
test_concurrent_open_table_reuses_connection_object_store`
- `cargo check --quiet --features remote --tests --examples`
- equivalent 32-thread `open_table(...).count_rows()` workload against a
local database
- targeted S3 test collected successfully locally; execution requires
the CI LocalStack service, which is unavailable in this runner

Fixes #1786

<!-- lance-gatekeeper-fix:v1 agent=d311f3c7151f77ae22b4997702e7b7db
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-26 14:56:57 +08:00
Jack Ye 391cac9034 fix(remote): centralize timeline consistency (#4053)
Centralizes remote table freshness fencing and response-version tracking
in the default transport path.

Covers schema and blob bypass paths, keeps explicit time-travel and
cross-timeline operations unfenced, and advances freshness after refresh
and index job completion.
2026-08-26 12:54:36 +08:00
LanceDB Robot 21530432a0 chore: update lance dependency to v12.0.0-beta.2 (#4056)
Updates the Rust workspace Lance crates and Java lance-core dependency
to
[v12.0.0-beta.2](https://github.com/lance-format/lance/releases/tag/v12.0.0-beta.2).
No compatibility fixes were required; full-workspace Clippy passes with
all features and warnings denied.
2026-08-25 21:56:26 -05:00
lancedb-gatefixer[bot] 9b825c5f29 fix(node): route auto search using table embeddings (#3832)
## Summary

- Resolve automatic string-search routing from the active table schema
whenever the query executes.
- Defer embedding-provider construction while leaving explicit vector
and FTS routes unchanged.
- Cover unrelated global registrations and metadata transitions across
repeated executions of one query builder.

## Root cause

LocalTable.search used the number of globally registered embedding
providers to choose between vector and full-text search. A provider
registered for any other table therefore sent a plain FTS table down the
vector path. A wrapper-lifetime metadata snapshot avoided that
contamination but became stale after time travel or read-consistency
refreshes. The query now records fluent builder operations and creates
the appropriate native vector or FTS query from the active schema on
each execution.

## Validation

- pnpm build
- pnpm tsc
- pnpm lint
- pnpm run docs
- pnpm test --runInBand (681 passed, 5 skipped)

Fixes #1557

<!-- lance-gatekeeper-fix:v1 agent=b6183df8296db4aabdc5d19a2256b029
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-26 10:08:48 +08:00
LanceDB Robot 8083232dd5 chore: update lance dependency to v12.0.0-beta.1 (#4055)
Updates the Lance Rust workspace dependencies and Java lance-core
dependency to
[v12.0.0-beta.1](https://github.com/lance-format/lance/releases/tag/v12.0.0-beta.1).

Includes compatibility updates for the renamed shard-manifest API and
paginated object-store wrappers.
2026-08-25 16:49:17 -07:00
lancedb-gatefixer[bot] 302b21aa94 test(node): cover nested PDF metadata queries (#3827)
## Summary

- add an end-to-end Node regression matching LangChain PDFLoader
metadata
- verify create/query round trips rich nested `loc` and `pdf.info`
fields against the currently configured Apache Arrow peer

## Root cause

LanceDB v0.14 delegated nested object inference to Apache Arrow. Nested
strings were dictionary-encoded with colliding dictionary IDs, so
serializing query results as an IPC file failed with a
dictionary-replacement error. Current `main` recursively infers nested
fields and avoids those invalid dictionaries, but the reported LangChain
path had no end-to-end regression coverage.

## Validation

- `pnpm build`
- `pnpm lint`
- `pnpm run docs`
- `pnpm test --runInBand` (678 passed, 5 skipped)

Fixes #1963

<!-- lance-gatekeeper-fix:v1 agent=bf8d489db7db2e17678b143f9f0a36d2
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-26 06:42:29 +08:00
lancedb-gatefixer[bot] 35b5d015ac fix(node): preserve embedding registration in server bundles (#3806)
## Summary

- lazily initialize built-in OpenAI and Hugging Face providers when
consumers call the public embedding registry API
- choose automatic vector versus FTS search from embedding metadata on a
fresh pinned table revision for every execution
- expose automatic string searches as an `AutoQuery` with only
operations common to both native query families
- keep the registry shared and built-in registration safe across
duplicated module graphs

## Root cause

Nitro treats dependency modules as side-effect-free and removes the bare
OpenAI provider import from its generated route. Registration therefore
never runs, so `getRegistry().get("openai")` remains undefined even when
the registry itself is shared globally. Bundlers may also duplicate the
provider and registry module graphs.

The public embedding entry point now initializes built-in providers only
when `getRegistry()` is explicitly called, keeping initialization on a
live path that Nitro retains. Each terminal automatic-search execution
pins the exact table revision visible at dispatch, reads embedding
metadata and computes an embedding from that snapshot, replays the
builder operations, and constructs and executes the selected native
query against the same snapshot. Pinned native snapshots execute locally
when namespace pushdown cannot carry their revision, while remote
snapshots are seeded directly from one version-and-schema response. The
public `AutoQuery` builder exposes only the operations shared by FTS and
vector search, so runtime class narrowing cannot expose invalid
vector-only methods. Repeated built-in registration replaces stale
constructors from duplicated module graphs while public `register()`
retains its duplicate-alias error.

## Validation

- `cargo fmt --all`
- `cargo check --quiet --features remote --tests --examples`
- `cargo clippy --quiet --features remote --tests --examples`
- `pnpm build`
- `pnpm lint`
- `pnpm run docs`
- `pnpm test --runInBand` (783 passed, 5 skipped)
- serial examples suite with a local OpenAI mock (11 passed), including
`sentence-transformers.test.ts`
- packaged Nitro 2.13.4 server route using the reported imports returned
`{"registered":true}`
- fresh-process FTS fixture initialized both public built-ins and
confirmed automatic string search still returned the indexed row
- schema-consistency regressions cover read-consistency refresh,
checkout, checkoutLatest, restore, runtime class narrowing, concurrent
overwrite during embedding computation, and reused automatic-search
builders
- focused regressions confirm pinned native snapshots bypass unversioned
namespace pushdown and remote snapshots use one describe request

Fixes #2429

<!-- lance-gatekeeper-fix:v1 agent=2adf0f21b8bfb634606ed8897a849e30
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-26 04:17:32 +08:00
lancedb-gatefixer[bot] a57fb68891 docs(python): fix Azure storage options examples (#3899)
## Summary

- document that Azure Blob Storage credentials can be passed directly
through `storage_options`
- provide valid quoted `account_name` and `account_key` examples for
both sync and async Python connections
- execute the option dictionaries during doctests so the original
unquoted-key mistake is caught

## Root cause

The historical Python storage guide used `account_name` and
`account_key` as bare identifiers in dictionary literals. Following that
example either raised `NameError` or, when those names were predefined,
produced incorrect option keys. The runtime already accepts direct Azure
credentials, but the current Python API reference did not contain a
corrected Azure example.

## Validation

- `python/.venv/bin/ruff format --check
python/python/lancedb/__init__.py`
- `python/.venv/bin/ruff check .`
- `cd python && uv run --no-sync pytest --doctest-modules
python/lancedb/__init__.py -q`
- `cd python && uv run --no-sync pytest python/tests/test_import.py -q`

Fixes #2236

<!-- lance-gatekeeper-fix:v1 agent=7a7a9e009eb2ebe52ac9b1adf2e8afb2
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-26 00:31:14 +08:00
Drew a614400755 feat: accept blob URI writes (#3954)
#3528 added blob declarations and binary coercion. String values were
still rejected. They now coerce to the blob `uri` child.

```python
table.add([{"id": 1, "image": "s3://bucket/media/cat.jpg"}])

payload = table.fetch_blobs("image", table.search().to_arrow())
```

A URI under a registered base writes with no extra options. An
unregistered URI fails. `allow_external_blob_outside_bases` is a local
escape hatch that stores an absolute URI. It does not register a base.
Remote `add` rejects that flag before making a request. String input
still coerces and is sent as a `uri` struct.

`add_bases` is a follow-up. `merge_insert` does not coerce string blob
input.

### Testing

- `cargo test -p lancedb --test blob_integration`
- `cargo test -p lancedb blob_coerce`
- `cargo test -p lancedb --features remote --lib
add_rejects_external_blob_flag add_string_blob_becomes_uri_struct`
- `cd python && uv run --extra tests pytest python/tests/test_blob.py -k
uri -q`

Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-25 22:59:53 +08:00
Xuanwo 1d880f11ff fix(python): use dev profile for editable builds (#4049)
## Problem

`uv run ... maturin develop` synchronizes the project as an editable
package before running the command. Maturin's editable build otherwise
uses the release profile, which enables the repository's fat LTO
configuration during local bootstrap.

## Behavior

Editable Python builds now explicitly use Cargo's dev profile. The
minimum maturin version is raised to 1.10, where `editable-profile`
support was introduced.
2026-08-25 19:51:40 +08:00
Lance Release ec4ad54ba2 Bump version: 0.38.0-beta.9 → 0.38.0-beta.10 2026-08-25 10:37:09 +00:00
Xuanwo d0bcc6c6fe fix: remove unsupported Function secrets contract (#4047)
## Problem

The unreleased First-Class Function authoring API exposed
`secrets=[...]` and serialized `required_secrets`, promising runtime
resolution and injection that Sophon does not implement.

## Behavior

Remove the secrets dimension from the public Python decorator, Python
and Rust registration/version models, and shared wire fixtures. Existing
successfully registered Functions retain stable identity: the field
could only be empty and empty values were already omitted from canonical
serialization.

A stable LanceDB release has not shipped this Function API, so this
contracts the surface before it becomes a published compatibility
commitment.

## Validation gap

The Rust shared-golden Function suites and Python formatting/lint checks
pass locally. Python pytest was not run because the local environment
lacks its runtime dependencies and the frozen native editable build did
not complete in practical time.
2026-08-25 18:35:27 +08:00
Lance Release 81c3f108ce Bump version: 0.38.0-beta.8 → 0.38.0-beta.9 2026-08-25 10:32:12 +00:00
Xuanwo c988e4848d refactor: simplify Function binding identity (#4046)
Function applications and bindings currently encode `group_id` and
binding `revision` even though `binding_id` already owns the complete
immutable binding lifecycle and `outputs` already defines the atomic
multi-output set.

Make `binding_id` the sole binding identity, remove the redundant fields
from the Rust and Python client contracts, and describe multi-output
declarations directly. This intentionally replaces the removed wire
fields without a compatibility path.
2026-08-25 18:30:09 +08:00
Lance Release 2fea7cd48d Bump version: 0.38.0-beta.7 → 0.38.0-beta.8 2026-08-25 06:26:35 +00:00
Wyatt Alt 0e65123bd8 fix: package ordinary @udf bodies and emit only the V1 type grammar (#4044)
Registering a real (embedding) Function failed on the client for three
reasons:

- `_package_source` treated `inspect.getclosurevars().unbound` as
  "unresolved globals"; CPython puts attribute names there, so any body
with `np.linalg.norm(...)` or `body.split()` was rejected. Module-scope
  references now come from Python's own scope analysis (`symtable`) over
  the function source, recursively, and each is resolved the way the
  interpreter would: the function's globals first (a module global may
  shadow a builtin), then builtins. Free variables of nested scopes stay
  lexical; postponed annotations are not runtime loads. A genuinely
  missing global still fails.

- `_canonical_arrow_type` emitted spellings the server's frozen grammar
  rejects (`fixed_size_list<T>[n]`, `timestamp[us]`, `struct<...>`,
  zero-sized lists). It now emits exactly the grammar, with the server's
  `fixed_size_list<item, size>` form, and the Rust declaration planner
  parses that form too.

A shared golden
(`tests/fixtures/first_class_functions/v1/arrow_types.json`)
enumerates every grammar type, nested forms and rejected spellings; the
Python emitter and Rust parser are tested against it, and the same file
is under test in sophon. Packaging tests execute the shipped artifact in
a fresh namespace.


Contract changes (hence `breaking-change`):

- `@udf` now rejects namespace acquisition structurally
(`globals()`/`eval`/... by name, plus `import
sys`/`builtins`/`importlib`/`inspect` inside the body), requires the
function's captured `__builtins__` to be the standard mapping itself
(identity, so neither lookups nor implicit hooks such as `__import__`
can differ), rejects module globals that are namespace-bearing modules
(`builtins`, `sys`, ...), and treats the function's own name as
recursion only when the module binds it to the function or to the exact
`UdfDefinition` the decorator produced; it resolves module globals
through the function's real namespace (a module global may shadow a
builtin) and ships importable classes/functions as imports.
- List outputs must declare a non-nullable, metadata-free child named
`item` (`pa.list_(pa.field("item", t, nullable=False))`); that is what
the grammar means, and pyarrow's default nullable child was being
silently collapsed into it.

Contract, stated in the `udf` docstring: the artifact is a snapshot of
the function source plus exactly the module names it references.
Reaching the module namespace by another route is rejected where a
static packager can see it and is otherwise unsupported; there is no
dynamic-access detection beyond that.
2026-08-25 14:15:07 +08:00
Jack Ye 6ed3074d4c feat: pin the base table version for data loader reads (#3982)
A permutation stores `_rowid`s, which are row addresses unless stable
row ids are enabled. Nothing in the data loader pinned a table version,
so a compaction between building a permutation and reading it can
resolve those ids to different rows.

The exposure differs by backend but exists on both:

- Remote never pins. `prepare_query_bodies` stamps `"version":
current_version()` on every request, but `current_version()` is `None`
unless `checkout` was called, so every request means "latest".
- Native pins implicitly by holding an `Arc<Dataset>` under
`ConsistencyMode::Lazy`, but `StreamingDataset.__setstate__` reopens the
table in each DataLoader worker, so each worker pins to whatever is
latest at fork time.

## Changes

`Table::at_version` returns an independent handle pinned to a version
without mutating the receiver. `checkout` cannot serve this: on remote
the version cell is an `Arc<RwLock<Option<u64>>>` shared across clones,
so pinning through it would silently pin the caller's table too.

`PermutationBuilder::build` pins for the whole build and records the
version in the permutation table's schema metadata, alongside the
existing split names. `PermutationReader` pins the base table to that
version before any take.

Because the reader pins on construction, the Python worker fork is
covered without touching the pickle format — `Permutation.__setstate__`
drops the reader and `_ensure_open` rebuilds it, which re-pins.

## Behaviour change

A permutation is now bound to the version it was built against, so rows
appended to the base table afterwards are not visible through an
existing permutation. That is the intended semantics — the permutation
only addresses rows that existed when it was built — but it is a change
worth flagging.

Permutations written before this carry no version key and read exactly
as they did before.
2026-08-24 17:16:42 -07:00
lancedb-gatefixer[bot] c1a8c3f089 fix(node): validate inferred types across records (#3786)
## Summary

- compare inferred Arrow types by their semantic representation across
records
- throw the schema inference error when a later record has an
incompatible type
- cover compatible and incompatible multi-record inference across
supported Arrow versions

## Root cause

Schema inference compared newly allocated Arrow DataType objects by
identity, so equivalent inferred types did not compare equal. The
mismatch path also constructed an Error without throwing it, which
silently accepted incompatible values.

## Validation

- pnpm test __test__/arrow.test.ts --runInBand (176 tests passed)
- pnpm lint
- pnpm build
- pnpm run docs

Fixes #3781

<!-- lance-gatekeeper-fix:v1 agent=00ec4f61a3fd82694fa4fb9bb2b37aa8
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
2026-08-24 16:41:05 -07:00
Will Jones fce45ba9fc feat(nodejs): add listTables, deprecate tableNames (#4041)
`table_names` is being replaced by `list_tables` across the SDKs, but
TypeScript only had `tableNames`. This PR adds `listTables`, which
returns a page of table names together with the token that resumes after
it, and marks `tableNames` and `TableNamesOptions` deprecated in favor
of it.

It binds the `Connection::list_tables` that already exists, so nothing
in the Rust API changes and nothing existing breaks. `pageToken` is
documented as opaque rather than as a table name, since what resumes a
listing is the database's to decide — that keeps callers off a detail
that is going to change.

Stacked on #4040, which fixes a table being dropped at every page
boundary. The page-walking test here needs that fix to pass. Review the
last commit only until #4040 lands.

## 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);
```

A namespace can be listed by passing its path first, mirroring
`tableNames`:

```ts
const page = await conn.listTables(["analytics"], { limit: 100 });
```

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-24 16:25:14 -07:00
lancedb-gatefixer[bot] 5013c176dd fix: pin remote snapshots during permutation construction (#4022)
Fixes #4015

<!-- lance-gatekeeper-fix:v1 agent=1c6407eec60f0af2336ae2dee811e1be
generation=1 -->

## Root cause

PermutationBuilder issued count_rows and the projected row-ID scan
through an unpinned RemoteTable handle. Each request could independently
resolve latest, so a concurrent table update could make the count and
scanned rows come from different snapshots.

## Fix

- add a backend hook for obtaining an independent handle pinned to the
currently selected version
- resolve latest once for remote tables while preserving an explicit
checkout and leaving the caller handle unchanged
- build the count, filtered projection, and scan from that pinned handle
while retaining native-table behavior
- add a remote mock regression that advances latest between count and
scan and covers explicit checkout preservation

## Validation

- cargo test --quiet --features remote -p lancedb --lib
test_remote_permutation_builder_pins_snapshot
- cargo test --quiet --features remote -p lancedb --lib
dataloader::permutation::builder::tests
- cargo check --quiet --features remote --tests --examples
- cargo clippy --quiet --features remote --tests --examples
- cargo fmt --all -- --check

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-25 07:14:33 +08:00
Lance Release 71f85a8d9f Bump version: 0.38.0-beta.6 → 0.38.0-beta.7 2026-08-24 21:09:12 +00:00
Wyatt Alt c72f5b2960 feat: bind a materialized view refresh to the view incarnation (#4043)
A caller that queues a refresh and executes it later can only tell the
view it captured from a drop-and-recreate by comparing the definition
and
version. A recreated view with the same definition and an equal or
higher
version passes that check, and the check runs before the refresh reloads
the view, so it never sees the state it commits against.

This mints an `mv.incarnation` token in the schema metadata at each
physical creation of a view table.
`RefreshMaterializedViewBuilder::expect_incarnation` carries the
captured
token into the refresh, which compares it against the latest stored
manifest before planning and again immediately before each commit
(publish, fragment swap, watermark stamp), refusing to land in a
different
incarnation. A view with no token -- declared before tokens existed, or
its metadata replaced wholesale -- is refused under a bound refresh with
its own wording and is minted one by its next unbound refresh. The token
is exposed through `MaterializedView::incarnation`; refreshes without an
expectation are unchanged.

This is best effort: the token is not part of lance's commit condition,
so a recreation landing between the final pre-commit read and the commit
itself is not caught. Closing that window needs a base-manifest
precondition in lance's commit path.
2026-08-24 14:06:40 -07:00
Will Jones 93f47b8aab fix(remote): stop table_names inventing a page token for a namespace (#4039)
`table_names` paginates using a `start_after` table name. This works for
the `/v1/table` endpoint, which guarantees table-order. But the
`/v1/namespace/{id}/table/list` does not. We change that caller to
instead collect all table names, sort, and apply the pagination locally.

We are deprecating this API, so this is just an interim fix. For good
performance, users should move to the `list_tables` API instead, which
uses opaque tokens that don't rely on lexical sorting.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-24 13:35:31 -07:00
lancedb-gatefixer[bot] 105fd73bc6 fix(python): commit streaming worker checkpoints on consumption (#4023)
## Summary

- add `StreamingDataLoader`, which transports worker snapshots with
prefetched batches and commits them to the parent dataset only when the
trainer receives each batch
- preserve exact non-uniform per-split progress and resume lagging
splits without replaying already-consumed rows
- reject stale parent checkpoints after a standard multi-process
`DataLoader` has started, with guidance to use the consumer-aware loader
- document the new public loader and merge non-uniform state across
ranks

## Root cause

PyTorch runs `StreamingDataset.__iter__` in private worker-process
copies, while callers invoke `state_dict()` on the parent dataset.
Sharing producer counters would still be incorrect because DataLoader
prefetch can advance workers beyond batches returned to the trainer.

## Validation

- `uv run --extra tests pytest python/tests/test_elastic_dataloader.py
-q` (154 passed)
- focused non-uniform merge regression (1 passed)
- `uv run --project python --extra tests --extra dev ruff format .`
- `uv run --project python --extra tests --extra dev ruff check .`
- `cd docs && PYTHONPATH=. ../python/.venv/bin/mkdocs build`

Fixes #3967

<!-- lance-gatekeeper-fix:v1 agent=572be272619660b97e87fd5c85188341
generation=1 -->

---------

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
2026-08-25 04:22:22 +08:00
Will Jones 94d484f539 fix(listing): don't drop a table at a page boundary (#4040)
Listing tables a page at a time against a local database silently
skipped one table at every page boundary. `ListingDatabase::list_tables`
returned the first name of the *next* page as that page's token, but
resuming from a token drops every name at or before it — so the table
the token named was never handed to the caller. Walking `[a, b, c, d,
e]` with a limit of 2 returned `[a, b, d, e]`.

This PR returns the last name of the page as the token instead, which is
what resuming after the token expects.

This is reachable from Python today through
`db.list_tables(page_token=...)` on a local connection; it also affects
`len(db)` and `name in db`, which walk the pages. Remote and
namespace-backed connections page on the server and were never affected.

## Example

```python
db = lancedb.connect(tmp_path)
for name in ["a", "b", "c", "d", "e"]:
    db.create_table(name, [{"id": 1}])

names, token = [], None
while True:
    page = db.list_tables(page_token=token, limit=2)
    names += page.tables
    token = page.page_token
    if not token:
        break

# before: ['a', 'b', 'd', 'e']
# after:  ['a', 'b', 'c', 'd', 'e']
```

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-24 12:00:48 -07:00
Xuanwo b0dae5eb0b feat: return typed refresh job results (#4013)
## Problem

`refresh_column_async` returned a unit-result job even though durable
refresh jobs carry a canonical terminal result. Python callers could not
obtain row counts or source and published versions through the public
`Job` API, and local and remote refresh jobs exposed different result
semantics.

## Behavior

`refresh_column_async` now returns `Job[RefreshColumnResult]` for local
and remote tables. The general typed-job bridge binds each endpoint to
its public result model while preserving unit-result jobs and existing
status, wait, cancel, and timeout behavior. A local no-op refresh
reports no published version.

The Node.js API continues to resolve `wait()` as `void`; its binding
erases the Rust result type internally to preserve the existing public
contract.

## Ownership and integration boundary

LanceDB owns the language-neutral `Job<T>` contract and language-binding
decode. Sophon owns production and durable persistence of terminal
payloads. Sophon #7348 and #7378 now publish the canonical refresh
result for Function-backed and expression-backed refresh jobs,
respectively. The remote client fixture matches the merged server
schema; live deployment and end-to-end demo acceptance remain separate
rollout checks.
2026-08-25 00:01:32 +08:00
Ayush Chaurasia 242ade8017 feat(python): support sequence packing in streaming dataset (#3920)
## How packing works

  Consider four tokenized documents:

  
  [1]
  [2]
  [10, 11, 12, 13, 14, 15, 16, 17]
  [20]

  With:
```
  StreamingDataset(
      table,
      shuffle=False,
      columns=["tokens"],
      num_splits=2,
      pack_sequences=5,
      eos_id=9,
      pad_id=0,
      blocks_per_epoch=6,
  )
```
the documents are assigned to two fixed logical splits. Each split
maintains an independent token buffer, appends eos_id after every
document, and emits blocks of five tokens.

  Because blocks_per_epoch=6, each split emits exactly three blocks:

  Cycl/e 1:
    Split 0: [1, 9, 2, 9, 0] # 9 is eos, 0 is padding
    Split 1: [10, 11, 12, 13, 14]

  Cycle 2:
    Split 0: [0, 0, 0, 0, 0]
    Split 1: [15, 16, 17, 9, 20]

  Cycle 3:
    Split 0: [0, 0, 0, 0, 0]
    Split 1: [9, 0, 0, 0, 0]

If a split runs out of tokens early, it emits padded blocks through the
fixed budget. This prevents one rank from finishing before another.

Logical splits are independent of rank and worker ownership. A
checkpoint records each split’s consumed-document count, emitted-block
count, remaining tokens, and document boundaries. Merging
those per-split states allows the same packed stream to resume after the
topology changes.

doc_ids identifies document segments, including continuations across
block boundaries. It is not a padding mask: padding retains the
preceding document ID, so callers must mask padding using a
  reserved pad_id.

blocks_per_epoch="auto" is also available. It estimates the budget from
a deterministic bounded sample and warns that the result is approximate.

WIP pre-training tests:
```
  ┌────────────────────────────────────┬────────────────────────┬─────────────────────────────────────┐
  │                                    │       GPT-2 124M       │          GPT-2 medium 354M          │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Corpus                             │ 2.4M docs / 12GB table │ 9.67M docs / 45GB table             │
  │ Tokens (Chinchilla)                │ 2.43B                  │ 7.0B                                │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Data prep (ingest→curate→tokenize) │ ~12 min                │ ~51 min                             │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Training wall time                 │ ~50 min                │ 3h 06m                              │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Throughput / MFU                   │ 1.60M tok/s / 35%      │ 684k tok/s / 42.0%, │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Final val loss                     │ 3.230                  │ 2.840                               │
  └────────────────────────────────────┴────────────────────────┴─────────────────────────────────────┘
```

---------

Co-authored-by: OpenAI Codex <codex@openai.com>
2026-08-24 15:56:36 +08:00
Lance Release 40d4d012e7 Bump version: 0.38.0-beta.5 → 0.38.0-beta.6 2026-08-23 17:33:19 +00:00
LanceDB Robot 000e3b506b chore: update lance dependency to v11.0.0-beta.22 (#4036)
Updates the Rust workspace Lance dependencies and Java lance-core
dependency to v11.0.0-beta.22, including the refreshed Cargo lockfile.
No compatibility fixes were required; see the [Lance
tag](https://github.com/lance-format/lance/releases/tag/v11.0.0-beta.22).
2026-08-23 10:31:32 -07:00
Lance Release 1b950188c3 Bump version: 0.38.0-beta.4 → 0.38.0-beta.5 2026-08-23 08:07:48 +00:00
136 changed files with 12145 additions and 1814 deletions
+1 -1
View File
@@ -1,5 +1,5 @@
[tool.bumpversion]
current_version = "0.38.0-beta.5"
current_version = "0.38.0-beta.11"
parse = """(?x)
(?P<major>0|[1-9]\\d*)\\.
(?P<minor>0|[1-9]\\d*)\\.
Generated
+49 -45
View File
@@ -3455,8 +3455,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "fsst"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"rand 0.9.5",
@@ -4815,8 +4815,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
[[package]]
name = "lance"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arc-swap",
"arrow",
@@ -4888,8 +4888,8 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4911,7 +4911,7 @@ dependencies = [
[[package]]
name = "lance-arrow-scalar"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4925,7 +4925,7 @@ dependencies = [
[[package]]
name = "lance-arrow-stats"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -4934,8 +4934,8 @@ dependencies = [
[[package]]
name = "lance-bitpacking"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrayref",
"crunchy",
@@ -4945,8 +4945,8 @@ dependencies = [
[[package]]
name = "lance-core"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4983,8 +4983,8 @@ dependencies = [
[[package]]
name = "lance-datafusion"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5013,8 +5013,8 @@ dependencies = [
[[package]]
name = "lance-datagen"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5031,8 +5031,8 @@ dependencies = [
[[package]]
name = "lance-derive"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"proc-macro2",
"quote",
@@ -5041,8 +5041,8 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5075,8 +5075,8 @@ dependencies = [
[[package]]
name = "lance-file"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5107,8 +5107,8 @@ dependencies = [
[[package]]
name = "lance-index"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arc-swap",
"arrow",
@@ -5172,8 +5172,8 @@ dependencies = [
[[package]]
name = "lance-index-core"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5195,8 +5195,8 @@ dependencies = [
[[package]]
name = "lance-io"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5222,7 +5222,11 @@ dependencies = [
"pin-project",
"prost",
"rand 0.9.5",
"reqsign-core",
"reqsign-file-read-tokio",
"reqsign-google",
"serde",
"serde_json",
"tempfile",
"tokio",
"tracing",
@@ -5232,8 +5236,8 @@ dependencies = [
[[package]]
name = "lance-linalg"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5247,8 +5251,8 @@ dependencies = [
[[package]]
name = "lance-namespace"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"async-trait",
@@ -5260,8 +5264,8 @@ dependencies = [
[[package]]
name = "lance-namespace-impls"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-ipc",
@@ -5314,8 +5318,8 @@ dependencies = [
[[package]]
name = "lance-select"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5329,8 +5333,8 @@ dependencies = [
[[package]]
name = "lance-table"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow",
"arrow-array",
@@ -5370,8 +5374,8 @@ dependencies = [
[[package]]
name = "lance-testing"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5384,8 +5388,8 @@ dependencies = [
[[package]]
name = "lance-tokenizer"
version = "11.0.0-beta.21"
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-beta.21#dd08336cf61b117701a7f5bbf76a7f7080f7e210"
version = "12.0.0-beta.2"
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
dependencies = [
"frostem",
"icu_segmenter",
@@ -5398,7 +5402,7 @@ dependencies = [
[[package]]
name = "lancedb"
version = "0.38.0-beta.4"
version = "0.38.0-beta.11"
dependencies = [
"ahash",
"anyhow",
@@ -5486,7 +5490,7 @@ dependencies = [
[[package]]
name = "lancedb-nodejs"
version = "0.38.0-beta.4"
version = "0.38.0-beta.11"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5511,7 +5515,7 @@ dependencies = [
[[package]]
name = "lancedb-python"
version = "0.38.0-beta.4"
version = "0.38.0-beta.11"
dependencies = [
"arrow",
"async-trait",
+14 -14
View File
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
rust-version = "1.91.0"
[workspace.dependencies]
lance = { "version" = "=11.0.0-beta.21", default-features = false, "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=11.0.0-beta.21", default-features = false, "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=11.0.0-beta.21", default-features = false, "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=11.0.0-beta.21", "tag" = "v11.0.0-beta.21", "git" = "https://github.com/lance-format/lance.git" }
lance = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
lancedb = { path = "rust/lancedb", default-features = false }
ahash = "0.8"
# Note that this one does not include pyarrow
+1 -1
View File
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
<dependency>
<groupId>com.lancedb</groupId>
<artifactId>lancedb-core</artifactId>
<version>0.38.0-beta.5</version>
<version>0.38.0-beta.11</version>
</dependency>
```
+518
View File
@@ -0,0 +1,518 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / AutoQuery
# Class: AutoQuery
A builder for automatic string searches.
Automatic search determines whether to use full-text or vector search from
the table revision selected for each execution. This builder exposes the
common operations supported by both query families.
## Extends
- `StandardQueryBase`&lt;`NativeQuery` \| `NativeVectorQuery`&gt;
## Properties
### inner
```ts
protected inner: Query | VectorQuery | Promise<Query | VectorQuery>;
```
#### Inherited from
`StandardQueryBase.inner`
## Methods
### analyzePlan()
```ts
analyzePlan(distributedMetrics?): Promise<string>
```
Executes the query and returns the physical query plan annotated with runtime metrics.
This is useful for debugging and performance analysis, as it shows how the query was executed
and includes metrics such as elapsed time, rows processed, and I/O statistics.
#### Parameters
* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md)
How distributed worker metrics are displayed for remote query plans.
Defaults to `"aggregate"`.
#### Returns
`Promise`&lt;`string`&gt;
A query execution plan with runtime metrics for each step.
#### Example
```ts
import * as lancedb from "@lancedb/lancedb"
const db = await lancedb.connect("./.lancedb");
const table = await db.createTable("my_table", [
{ vector: [1.1, 0.9], id: "1" },
]);
const plan = await table.query().nearestTo([0.5, 0.2]).analyzePlan();
Example output (with runtime metrics inlined):
AnalyzeExec verbose=true, metrics=[]
ProjectionExec: expr=[id@3 as id, vector@0 as vector, _distance@2 as _distance], metrics=[output_rows=1, elapsed_compute=3.292µs]
Take: columns="vector, _rowid, _distance, (id)", metrics=[output_rows=1, elapsed_compute=66.001µs, batches_processed=1, bytes_read=8, iops=1, requests=1]
CoalesceBatchesExec: target_batch_size=1024, metrics=[output_rows=1, elapsed_compute=3.333µs]
GlobalLimitExec: skip=0, fetch=10, metrics=[output_rows=1, elapsed_compute=167ns]
FilterExec: _distance@2 IS NOT NULL, metrics=[output_rows=1, elapsed_compute=8.542µs]
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST], metrics=[output_rows=1, elapsed_compute=63.25µs, row_replacements=1]
KNNVectorDistance: metric=l2, metrics=[output_rows=1, elapsed_compute=114.333µs, output_batches=1]
LanceScan: uri=/path/to/data, projection=[vector], row_id=true, row_addr=false, ordered=false, metrics=[output_rows=1, elapsed_compute=103.626µs, bytes_read=549, iops=2, requests=2]
```
#### Inherited from
`StandardQueryBase.analyzePlan`
***
### execute()
```ts
protected execute(options?): AsyncGenerator<RecordBatch<any>, void, unknown>
```
Execute the query and return the results as an
#### Parameters
* **options?**: `Partial`&lt;[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)&gt;
#### Returns
`AsyncGenerator`&lt;`RecordBatch`&lt;`any`&gt;, `void`, `unknown`&gt;
#### See
- AsyncIterator
of
- RecordBatch.
By default, LanceDb will use many threads to calculate results and, when
the result set is large, multiple batches will be processed at one time.
This readahead is limited however and backpressure will be applied if this
stream is consumed slowly (this constrains the maximum memory used by a
single query)
#### Inherited from
`StandardQueryBase.execute`
***
### explainPlan()
```ts
explainPlan(verbose): Promise<string>
```
Generates an explanation of the query execution plan.
#### Parameters
* **verbose**: `boolean` = `false`
If true, provides a more detailed explanation. Defaults to false.
#### Returns
`Promise`&lt;`string`&gt;
A Promise that resolves to a string containing the query execution plan explanation.
#### Example
```ts
import * as lancedb from "@lancedb/lancedb"
const db = await lancedb.connect("./.lancedb");
const table = await db.createTable("my_table", [
{ vector: [1.1, 0.9], id: "1" },
]);
const plan = await table.query().nearestTo([0.5, 0.2]).explainPlan();
```
#### Inherited from
`StandardQueryBase.explainPlan`
***
### fastSearch()
```ts
fastSearch(): this
```
Skip searching un-indexed data. This can make search faster, but will miss
any data that is not yet indexed.
Use [Table#optimize](Table.md#optimize) to index all un-indexed data.
#### Returns
`this`
#### Inherited from
`StandardQueryBase.fastSearch`
***
### ~~filter()~~
```ts
filter(predicate): this
```
A filter statement to be applied to this query.
#### Parameters
* **predicate**: `string`
#### Returns
`this`
#### See
where
#### Deprecated
Use `where` instead
#### Inherited from
`StandardQueryBase.filter`
***
### fullTextSearch()
```ts
fullTextSearch(query, options?): this
```
#### Parameters
* **query**: `string` \| [`FullTextQuery`](../interfaces/FullTextQuery.md)
* **options?**: `Partial`&lt;[`FullTextSearchOptions`](../interfaces/FullTextSearchOptions.md)&gt;
#### Returns
`this`
#### Inherited from
`StandardQueryBase.fullTextSearch`
***
### limit()
```ts
limit(limit): this
```
Set the maximum number of results to return.
By default, a plain search has no limit. If this method is not
called then every valid row from the table will be returned.
#### Parameters
* **limit**: `number`
#### Returns
`this`
#### Inherited from
`StandardQueryBase.limit`
***
### offset()
```ts
offset(offset): this
```
Set the number of rows to skip before returning results.
This is useful for pagination.
#### Parameters
* **offset**: `number`
#### Returns
`this`
#### Inherited from
`StandardQueryBase.offset`
***
### orderBy()
```ts
orderBy(ordering): this
```
Sort the results by the specified column(s).
#### Parameters
* **ordering**: [`ColumnOrdering`](../interfaces/ColumnOrdering.md) \| [`ColumnOrdering`](../interfaces/ColumnOrdering.md)[]
#### Returns
`this`
This query builder.
#### Inherited from
`StandardQueryBase.orderBy`
***
### outputSchema()
```ts
outputSchema(): Promise<Schema<any>>
```
Returns the schema of the output that will be returned by this query.
This can be used to inspect the types and names of the columns that will be
returned by the query before executing it.
#### Returns
`Promise`&lt;`Schema`&lt;`any`&gt;&gt;
An Arrow Schema describing the output columns.
#### Inherited from
`StandardQueryBase.outputSchema`
***
### select()
```ts
select(columns): this
```
Return only the specified columns.
By default a query will return all columns from the table. However, this can have
a very significant impact on latency. LanceDb stores data in a columnar fashion. This
means we can finely tune our I/O to select exactly the columns we need.
As a best practice you should always limit queries to the columns that you need. If you
pass in an array of column names then only those columns will be returned.
You can also use this method to create new "dynamic" columns based on your existing columns.
For example, you may not care about "a" or "b" but instead simply want "a + b". This is often
seen in the SELECT clause of an SQL query (e.g. `SELECT a+b FROM my_table`).
To create dynamic columns you can pass in a Map<string, string>. A column will be returned
for each entry in the map. The key provides the name of the column. The value is
an SQL string used to specify how the column is calculated.
For example, an SQL query might state `SELECT a + b AS combined, c`. The equivalent
input to this method would be:
#### Parameters
* **columns**: `string` \| `string`[] \| `Record`&lt;`string`, `string`&gt; \| `Map`&lt;`string`, `string`&gt;
#### Returns
`this`
#### Example
```ts
new Map([["combined", "a + b"], ["c", "c"]])
Columns will always be returned in the order given, even if that order is different than
the order used when adding the data.
Note that you can pass in a `Record<string, string>` (e.g. an object literal). This method
uses `Object.entries` which should preserve the insertion order of the object. However,
object insertion order is easy to get wrong and `Map` is more foolproof.
```
#### Inherited from
`StandardQueryBase.select`
***
### toArray()
```ts
toArray(options?): Promise<any[]>
```
Collect the results as an array of objects.
#### Parameters
* **options?**: `Partial`&lt;[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)&gt;
#### Returns
`Promise`&lt;`any`[]&gt;
#### Inherited from
`StandardQueryBase.toArray`
***
### toArrow()
```ts
toArrow(options?): Promise<Table<any>>
```
Collect the results as an Arrow
#### Parameters
* **options?**: `Partial`&lt;[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)&gt;
#### Returns
`Promise`&lt;`Table`&lt;`any`&gt;&gt;
#### See
ArrowTable.
#### Inherited from
`StandardQueryBase.toArrow`
***
### useLsm()
```ts
useLsm(enable): this
```
Control MemWAL read routing for this query.
By default (unset), when the table carries a MemWAL write spec (see
[Table#setLsmWriteSpec](Table.md#setlsmwritespec)), reads are routed through the LSM scanner so
they also return data written via the `mergeInsert` LSM path that has not yet
been compacted into the base table (the active/frozen in-memory memtables and
the flushed generations), deduplicated by primary key; a table without a spec
reads the base table.
#### Parameters
* **enable**: `boolean`
`true` forces the LSM scanner and errors if the table has no
MemWAL write spec. `false` bypasses the MemWAL and reads the base table only,
even when a spec is present.
Note: the LSM scanner does not support every query shape (e.g. reranking,
hybrid search, `orderBy`). On a MemWAL table those shapes error unless
`useLsm(false)` is set, because a base-only read would silently exclude
un-compacted MemWAL data.
#### Returns
`this`
#### Inherited from
`StandardQueryBase.useLsm`
***
### where()
```ts
where(predicate): this
```
A filter statement to be applied to this query.
The filter should be supplied as an SQL query string. For example:
#### Parameters
* **predicate**: `string`
#### Returns
`this`
#### Example
```ts
x > 10
y > 0 AND y < 100
x > 5 OR y = 'test'
Filtering performance can often be improved by creating a scalar index
on the filter column(s).
Calling this multiple times combines the filters with a logical AND rather
than replacing the previous filter.
```
#### Inherited from
`StandardQueryBase.where`
***
### withRowId()
```ts
withRowId(): this
```
Whether to return the row id in the results.
This column can be used to match results between different queries. For
example, to match results from a full text search and a vector search in
order to perform hybrid search.
#### Returns
`this`
#### Inherited from
`StandardQueryBase.withRowId`
+73 -1
View File
@@ -584,6 +584,70 @@ Child namespace names and
***
### listTables()
#### listTables(options)
```ts
abstract listTables(options?): Promise<ListTablesResponse>
```
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:
```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
* **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`).
##### Returns
`Promise`&lt;[`ListTablesResponse`](../interfaces/ListTablesResponse.md)&gt;
A page of table names and an
optional token for the tables after it.
***
### openMaterializedView()
```ts
@@ -660,7 +724,7 @@ a "not supported" error.
***
### tableNames()
### ~~tableNames()~~
#### tableNames(options)
@@ -682,6 +746,10 @@ Tables will be returned in lexicographical order.
`Promise`&lt;`string`[]&gt;
##### Deprecated
Use [Connection.listTables](Connection.md#listtables) instead.
#### tableNames(namespacePath, options)
```ts
@@ -704,3 +772,7 @@ Tables will be returned in lexicographical order.
##### Returns
`Promise`&lt;`string`[]&gt;
##### Deprecated
Use [Connection.listTables](Connection.md#listtables) instead.
+14 -2
View File
@@ -942,7 +942,7 @@ Get the schema of the table.
abstract search(
query,
queryType?,
ftsColumns?): Query | VectorQuery
ftsColumns?): Query | VectorQuery | AutoQuery
```
Create a search query to find the nearest neighbors
@@ -964,7 +964,7 @@ of the given query
#### Returns
[`Query`](Query.md) \| [`VectorQuery`](VectorQuery.md)
[`Query`](Query.md) \| [`VectorQuery`](VectorQuery.md) \| [`AutoQuery`](AutoQuery.md)
***
@@ -1292,6 +1292,18 @@ abstract updateFieldMetadata(updates): Promise<UpdateFieldMetadataResult>
Update per-field (column) metadata.
The following keys are treated specially, by convention, and should be
used when appropriate:
- `lancedb:description`: for a human-readable description of a field.
- `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
names the tag category; e.g. `lancedb:tag:model: "clip"`.
- `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
`feature_v2` might be in the same logical column.
- `lancedb:status`: for status options (`production`, `candidate`,
`deprecated`, `archived`) to designate the current life cycle state of
this column.
#### Parameters
* **updates**: [`FieldMetadataUpdate`](../interfaces/FieldMetadataUpdate.md)[]
+3
View File
@@ -18,6 +18,7 @@
## Classes
- [AutoQuery](classes/AutoQuery.md)
- [BooleanQuery](classes/BooleanQuery.md)
- [BoostQuery](classes/BoostQuery.md)
- [BranchContents](classes/BranchContents.md)
@@ -100,6 +101,8 @@
- [JobInfo](interfaces/JobInfo.md)
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
- [ListTablesOptions](interfaces/ListTablesOptions.md)
- [ListTablesResponse](interfaces/ListTablesResponse.md)
- [LsmStats](interfaces/LsmStats.md)
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
- [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md)
@@ -17,7 +17,8 @@ metadata: Record<string, null | string>;
```
Metadata key/value pairs. Merged into the field's existing metadata by
default; a value of `null` deletes that key.
default; a value of `null` deletes that key. See
[Table.updateFieldMetadata](../classes/Table.md#updatefieldmetadata) for the conventional `lancedb:*` keys.
***
@@ -0,0 +1,34 @@
[**@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.
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.
@@ -0,0 +1,23 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / ListTablesResponse
# Interface: ListTablesResponse
## Properties
### pageToken?
```ts
optional pageToken: string;
```
***
### tables
```ts
tables: string[];
```
+8 -3
View File
@@ -4,11 +4,16 @@
[@lancedb/lancedb](../globals.md) / TableNamesOptions
# Interface: TableNamesOptions
# Interface: ~~TableNamesOptions~~
## Deprecated
Use [ListTablesOptions](ListTablesOptions.md) with [Connection.listTables](../classes/Connection.md#listtables)
instead.
## Properties
### limit?
### ~~limit?~~
```ts
optional limit: number;
@@ -18,7 +23,7 @@ An optional limit to the number of results to return.
***
### startAfter?
### ~~startAfter?~~
```ts
optional startAfter: string;
@@ -10,16 +10,12 @@
function getRegistry(): EmbeddingFunctionRegistry
```
Utility function to get the global instance of the registry
Get the global embedding function registry.
LanceDB built-in providers are initialized when this public API is first
used, so importing the root package does not change automatic search
selection for tables without embedding metadata.
## Returns
[`EmbeddingFunctionRegistry`](../classes/EmbeddingFunctionRegistry.md)
`EmbeddingFunctionRegistry` The global instance of the registry
## Example
```ts
const registry = getRegistry();
const openai = registry.get("openai").create();
+4
View File
@@ -159,6 +159,8 @@ and combined with [BooleanQuery][lancedb.query.BooleanQuery].
::: lancedb.query.FullTextOperator
::: lancedb.query.DocumentGranularity
::: lancedb.query.Occur
## Embeddings
@@ -261,6 +263,8 @@ instead of being materialized with the rest of the row.
::: lancedb.streaming.StreamingDataset
::: lancedb.streaming.StreamingDataLoader
::: lancedb.permutation.permutation_builder
::: lancedb.permutation.PermutationBuilder
+1 -1
View File
@@ -8,7 +8,7 @@
<parent>
<groupId>com.lancedb</groupId>
<artifactId>lancedb-parent</artifactId>
<version>0.38.0-beta.5</version>
<version>0.38.0-beta.11</version>
<relativePath>../pom.xml</relativePath>
</parent>
+2 -2
View File
@@ -6,7 +6,7 @@
<groupId>com.lancedb</groupId>
<artifactId>lancedb-parent</artifactId>
<version>0.38.0-beta.5</version>
<version>0.38.0-beta.11</version>
<packaging>pom</packaging>
<name>${project.artifactId}</name>
<description>LanceDB Java SDK Parent POM</description>
@@ -28,7 +28,7 @@
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<arrow.version>15.0.0</arrow.version>
<lance-core.version>11.0.0-beta.21</lance-core.version>
<lance-core.version>12.0.0-beta.2</lance-core.version>
<spotless.skip>false</spotless.skip>
<spotless.version>2.30.0</spotless.version>
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
+1 -1
View File
@@ -1,7 +1,7 @@
[package]
name = "lancedb-nodejs"
edition.workspace = true
version = "0.38.0-beta.5"
version = "0.38.0-beta.11"
publish = false
license.workspace = true
description.workspace = true
+189
View File
@@ -1,11 +1,16 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import * as fs from "node:fs";
import * as vm from "node:vm";
import * as arrow15 from "apache-arrow-15";
import * as arrow16 from "apache-arrow-16";
import * as arrow17 from "apache-arrow-17";
import * as arrow18 from "apache-arrow-18";
import {
Field as CurrentField,
LargeBinary as CurrentLargeBinary,
Schema as CurrentSchema,
Vector as CurrentVector,
convertToTable,
tableFromIPC as currentTableFromIPC,
@@ -36,6 +41,59 @@ function sampleRecords(): Array<Record<string, any>> {
},
];
}
it("serializes an Arrow Table created in another JavaScript realm", async () => {
const context = vm.createContext({
TextDecoder,
TextEncoder,
console,
setTimeout,
clearTimeout,
});
vm.runInContext(
fs.readFileSync(
require.resolve("apache-arrow-15/Arrow.es2015.min"),
"utf8",
),
context,
);
const foreignTable: unknown = vm.runInContext(
"Arrow.tableFromArrays({ id: new Int32Array([1, 2, 3]), text: ['foo', 'bar', 'baz'] })",
context,
);
const foreignMetadata = (
foreignTable as { schema: { metadata: Map<string, string> } }
).schema.metadata;
expect(foreignMetadata).not.toBeInstanceOf(Map);
const buf = await fromDataToBuffer(
foreignTable as Parameters<typeof fromDataToBuffer>[0],
);
const actual = currentTableFromIPC(buf);
expect(actual.numRows).toBe(3);
expect(actual.getChild("id")?.toJSON()).toEqual([1, 2, 3]);
expect(actual.getChild("text")?.toJSON()).toEqual(["foo", "bar", "baz"]);
});
it("preserves field metadata from a provided schema", async function () {
const jsonMetadata = new Map([["ARROW:extension:name", "lance.json"]]);
const schema = new CurrentSchema([
new CurrentField("meta", new CurrentLargeBinary(), true, jsonMetadata),
]);
const table = makeArrowTable(
[{ meta: Buffer.from(JSON.stringify({ source: "test" })) }],
{ schema },
);
expect(table.schema.fields[0].metadata).toEqual(jsonMetadata);
const roundTripped = currentTableFromIPC(await fromTableToBuffer(table));
expect(roundTripped.schema.fields[0].metadata).toEqual(jsonMetadata);
});
describe.each([arrow15, arrow16, arrow17, arrow18])(
"Arrow",
(
@@ -515,6 +573,137 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
);
});
it("will allow matching inferred types across records", function () {
expect(() =>
makeArrowTable([{ value: 1 }, { value: 2 }]),
).not.toThrow();
});
it("will reject mismatched inferred types across records", function () {
expect(() => makeArrowTable([{ value: 1 }, { value: "two" }])).toThrow(
"Failed to infer schema for data. Previously inferred type Float64 but found Utf8 for field value at row 1. Consider providing an explicit schema.",
);
});
it("will ignore generated dictionary IDs when comparing inferred types", function () {
const table = makeArrowTable([{ str: "a" }, { str: "b" }], {
dictionaryEncodeStrings: true,
});
expect(table.getChild("str")?.toJSON()).toEqual(["a", "b"]);
});
it("will preserve null values without treating them as type mismatches", function () {
for (const records of [
[{ vector: [1, 2, 3] }, { vector: null }],
[{ vector: null }, { vector: [1, 2, 3] }],
]) {
const table = makeArrowTable(records);
expect(table.numRows).toBe(2);
expect(table.getChild("vector")?.nullCount).toBe(1);
}
});
it("will preserve empty variable-size lists", function () {
for (const records of [
[{ items: [1] }, { items: [] }],
[{ items: [] }, { items: [1] }],
]) {
const table = makeArrowTable(records);
expect(
table
.getChild("items")
?.toJSON()
.map((value) => value.toJSON()),
).toEqual(records.map((record) => record.items));
}
});
it("will propagate deferred evidence through nested lists", function () {
for (const records of [
[{ items: [1] }, { items: [null] }],
[{ items: [null] }, { items: [1] }],
[{ items: [null, 1] }, { items: [2, null] }],
]) {
const table = makeArrowTable(records);
expect(
table
.getChild("items")
?.toJSON()
.map((value) => value.toJSON()),
).toEqual(records.map((record) => record.items));
}
const nestedRecords = [{ items: [[1]] }, { items: [[null]] }];
const nestedTable = makeArrowTable(nestedRecords);
expect(
nestedTable
.getChild("items")
?.toJSON()
.map((value) =>
value
.toJSON()
.map((nestedValue: { toJSON: () => unknown[] }) =>
nestedValue.toJSON(),
),
),
).toEqual(nestedRecords.map((record) => record.items));
});
it("will reject incompatible deferred evidence within a list", function () {
for (const items of [
[[], 1],
[1, []],
[[null], 1],
[1, [null]],
]) {
expect(() => makeArrowTable([{ items }])).toThrow(
"Failed to infer data type for field items at row 0.",
);
}
});
it("will reject empty fixed-size lists", function () {
expect(() =>
makeArrowTable([{ vector: [1, 2, 3] }, { vector: [] }]),
).toThrow(
"Failed to infer schema for data. Previously inferred type FixedSizeList[3]<Float32> but found List[0] for field vector at row 1.",
);
});
it("will reject inferred leaf and branch shape changes", function () {
expect(() =>
makeArrowTable([{ value: 1 }, { value: { nested: 2 } }]),
).toThrow(
"Failed to infer schema for data. Previously inferred type Float64 but found Struct for field value at row 1.",
);
expect(() =>
makeArrowTable([{ value: { nested: 1 } }, { value: 2 }]),
).toThrow(
"Failed to infer schema for data. Previously inferred type Struct but found Float64 for field value at row 1.",
);
});
it("will allow null values around inferred struct values", function () {
for (const { records, nullIndex } of [
{
records: [{ value: null }, { value: { nested: 2 } }],
nullIndex: 0,
},
{
records: [{ value: { nested: 1 } }, { value: null }],
nullIndex: 1,
},
]) {
const table = makeArrowTable(records);
const values = table.getChild("value");
expect(values?.nullCount).toBe(1);
expect(values?.get(nullIndex)).toBeNull();
}
});
it("will allow a schema to be provided", async function () {
await checkTableCreation(
async (records, _, schema) =>
+68 -1
View File
@@ -4,7 +4,13 @@
import { readdirSync } from "fs";
import { Field, Float64, Schema } from "apache-arrow";
import * as tmp from "tmp";
import { Connection, Table, connect, connectNamespace } from "../lancedb";
import {
Connection,
ListTablesResponse,
Table,
connect,
connectNamespace,
} from "../lancedb";
import { LocalTable } from "../lancedb/table";
describe("when connecting", () => {
@@ -47,6 +53,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");
await expect(db.renameTable("a", "b")).rejects.toThrow(
"Connection is closed",
);
@@ -129,6 +136,66 @@ describe("given a connection", () => {
expect(tables).toEqual(["b", "c"]);
});
it("should respect limit and page token when listing tables", async () => {
const db = await connect(tmpDir.name);
await db.createTable("b", [{ id: 1 }]);
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();
const first = await db.listTables({ limit: 1 });
expect(first.tables).toEqual(["a"]);
expect(first.pageToken).toBeDefined();
const second = await db.listTables({
limit: 1,
pageToken: first.pageToken,
});
expect(second.tables).toEqual(["b"]);
});
it("should visit every table exactly once when walking pages", async () => {
const db = await connect(tmpDir.name);
const created = ["a", "b", "c", "d", "e"];
for (const name of created) {
await db.createTable(name, [{ id: 1 }]);
}
const seen: string[] = [];
let pageToken: string | undefined = undefined;
do {
const page: ListTablesResponse = await db.listTables({
limit: 2,
pageToken,
});
seen.push(...page.tables);
pageToken = page.pageToken;
} while (pageToken);
expect(seen).toEqual(created);
});
it("should list tables in a namespace", async () => {
const db = await connect(tmpDir.name, {
// biome-ignore lint/style/useNamingConvention: opaque backend property key, must match Rust
namespaceClientProperties: { manifest_enabled: "true" },
});
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: [] }),
);
});
it("should create tables in v2 mode", async () => {
const db = await connect(tmpDir.name);
const data = [...Array(10000).keys()].map((i) => ({ id: i }));
+52
View File
@@ -187,6 +187,58 @@ describe("embedding functions", () => {
const vector0 = JSON.parse(JSON.stringify(arr[0].vector));
expect(vector0).toEqual([1, 2, 3]);
});
it("should append multiple Python embeddings with the same alias", async () => {
@register("python-mock")
// biome-ignore lint/correctness/noUnusedVariables: the decorator registers this class
class MockEmbeddingFunction extends EmbeddingFunction<string> {
ndims() {
return 3;
}
embeddingDataType(): Float {
return new Float32();
}
async computeQueryEmbeddings(_data: string) {
return [1, 2, 3];
}
async computeSourceEmbeddings(data: string[]) {
return data.map((value) =>
value === "hello world" ? [1, 2, 3] : [4, 5, 6],
);
}
}
const metadata = new Map([
[
"embedding_functions",
'[{"source_column":"text1","vector_column":"vector1","name":"python-mock","model":{}},{"source_column":"text2","vector_column":"vector2","name":"python-mock","model":{}}]',
],
]);
const schema = new Schema(
[
new Field("text1", new Utf8(), true),
new Field("text2", new Utf8(), true),
new Field(
"vector1",
new FixedSizeList(3, new Field("item", new Float32(), true)),
true,
),
new Field(
"vector2",
new FixedSizeList(3, new Field("item", new Float32(), true)),
true,
),
],
metadata,
);
const db = await connect(tmpDir.name);
const table = await db.createEmptyTable("test", schema);
await table.add([{ text1: "hello world", text2: "goodbye world" }]);
const rows = await table.query().toArray();
expect(JSON.parse(JSON.stringify(rows[0].vector1))).toEqual([1, 2, 3]);
expect(JSON.parse(JSON.stringify(rows[0].vector2))).toEqual([4, 5, 6]);
});
it("should append generated vectors to a non-nullable schema", async () => {
@register("non_nullable_schema_test")
@@ -0,0 +1,95 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import { execFileSync } from "node:child_process";
import { resolve } from "node:path";
import type { OpenAIEmbeddingFunction } from "../lancedb/embedding/openai";
import type { EmbeddingFunctionRegistry } from "../lancedb/embedding/registry";
type EmbeddingModule = typeof import("../lancedb/embedding");
type OpenAIModule = typeof import("../lancedb/embedding/openai");
type RegistryModule = typeof import("../lancedb/embedding/registry");
describe("embedding function registry", () => {
const registries: EmbeddingFunctionRegistry[] = [];
afterEach(() => {
for (const registry of registries) {
registry.reset();
}
registries.length = 0;
});
it("defers built-in providers until the public registry API is used", () => {
jest.isolateModules(() => {
const embedding = require("../lancedb/embedding") as EmbeddingModule;
const { getRegistry: getInternalRegistry } =
require("../lancedb/embedding/registry") as RegistryModule;
const registry = getInternalRegistry();
registries.push(registry);
expect(registry.length()).toBe(0);
expect(embedding.getRegistry()).toBe(registry);
expect(registry.get("openai")).toBeDefined();
expect(registry.get("huggingface")).toBeDefined();
});
});
it("preserves automatic FTS search in a fresh process", () => {
execFileSync(
process.execPath,
[resolve(__dirname, "fixtures", "auto_fts_search.cjs")],
{ stdio: "pipe" },
);
});
it("shares registrations across duplicated provider module graphs", () => {
let registeringRegistry: EmbeddingFunctionRegistry | undefined;
let latestOpenAIConstructor: typeof OpenAIEmbeddingFunction | undefined;
jest.isolateModules(() => {
require("../lancedb/embedding/openai");
const { getRegistry } =
require("../lancedb/embedding/registry") as RegistryModule;
registeringRegistry = getRegistry();
registries.push(registeringRegistry);
expect(registeringRegistry.get("openai")).toBeDefined();
});
expect(() => {
jest.isolateModules(() => {
const { OpenAIEmbeddingFunction } =
require("../lancedb/embedding/openai") as OpenAIModule;
latestOpenAIConstructor = OpenAIEmbeddingFunction;
const { getRegistry } =
require("../lancedb/embedding/registry") as RegistryModule;
registries.push(getRegistry());
});
}).not.toThrow();
const previousApiKey = process.env.OPENAI_API_KEY;
process.env.OPENAI_API_KEY = "test";
try {
const latestOpenAI = registeringRegistry!
.get<OpenAIEmbeddingFunction>("openai")!
.create();
expect(latestOpenAI).toBeInstanceOf(latestOpenAIConstructor!);
} finally {
if (previousApiKey === undefined) {
delete process.env.OPENAI_API_KEY;
} else {
process.env.OPENAI_API_KEY = previousApiKey;
}
}
jest.isolateModules(() => {
const { getRegistry } =
require("../lancedb/embedding") as EmbeddingModule;
const publicRegistry = getRegistry();
registries.push(publicRegistry);
expect(publicRegistry).toBe(registeringRegistry);
expect(publicRegistry.get("openai")).toBeDefined();
});
});
});
@@ -0,0 +1,33 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
const assert = require("node:assert/strict");
const tmp = require("tmp");
const { connect, embedding, Index } = require("../../dist");
const { getRegistry } = require("../../dist/embedding/registry");
async function main() {
assert.equal(typeof embedding.getRegistry, "function");
assert.equal(getRegistry().length(), 0);
assert.equal(embedding.getRegistry(), getRegistry());
assert.equal(getRegistry().length(), 2);
const dir = tmp.dirSync({ unsafeCleanup: true });
let db;
try {
db = await connect(dir.name);
const table = await db.createTable("docs", [{ text: "hello world" }]);
await table.createIndex("text", { config: Index.fts() });
const rows = await table.search("hello").toArray();
assert.equal(rows[0].text, "hello world");
} finally {
db?.close();
dir.removeCallback();
}
}
main().catch((error) => {
console.error(error);
process.exitCode = 1;
});
+604 -1
View File
@@ -11,10 +11,13 @@ import * as arrow17 from "apache-arrow-17";
import * as arrow18 from "apache-arrow-18";
import {
AutoQuery,
Connection,
MatchQuery,
PhraseQuery,
Query,
Table,
VectorQuery,
connect,
tokenize,
} from "../lancedb";
@@ -682,6 +685,56 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
},
);
// https://github.com/lancedb/lancedb/issues/1963
it("should query documents with LangChain PDF metadata", async () => {
const tmpDir = tmp.dirSync({ unsafeCleanup: true });
try {
const db = await connect(tmpDir.name);
const documents = [
{
text: "first page",
vector: [1, 0],
source: "first.pdf",
loc: { pageNumber: 1, lines: { from: 1, to: 12 } },
pdf: {
version: "1.10.100",
info: {
format: "PDF 1.7",
producer: "pdf.js",
creator: "Writer",
},
totalPages: 2,
},
},
{
text: "second page",
vector: [0, 1],
source: "second.pdf",
loc: { pageNumber: 2, lines: { from: 13, to: 24 } },
pdf: {
version: "1.10.100",
info: {
format: "PDF 1.7",
producer: "pdf.js",
creator: "Writer",
},
totalPages: 2,
},
},
];
const documentsTable = await db.createTable("documents", documents);
const results = await documentsTable.query().toArray();
expect(results).toHaveLength(2);
expect(results[0].source).toBe("first.pdf");
expect(results[0].pdf.info.producer).toBe("pdf.js");
expect(results[1].loc.pageNumber).toBe(2);
} finally {
tmpDir.removeCallback();
}
});
describe("merge insert", () => {
let tmpDir: tmp.DirResult;
let table: Table;
@@ -1777,6 +1830,194 @@ describe("Read consistency interval", () => {
});
});
describe("automatic search schema consistency", () => {
let tmpDir: tmp.DirResult;
class SchemaRefreshEmbedding extends EmbeddingFunction<string> {
ndims() {
return 2;
}
embeddingDataType() {
return new Float32();
}
async computeSourceEmbeddings(data: string[]) {
return data.map((value) => [value.length, 1]);
}
async computeQueryEmbeddings(value: string) {
return [value.length, 1];
}
}
function embeddingSchema() {
const func = new SchemaRefreshEmbedding();
return LanceSchema({
text: func.sourceField(new Utf8()),
vector: func.vectorField(),
});
}
beforeEach(() => {
getRegistry().reset();
register("schema-refresh")(SchemaRefreshEmbedding);
tmpDir = tmp.dirSync({ unsafeCleanup: true });
});
afterEach(() => {
getRegistry().reset();
tmpDir.removeCallback();
});
it("uses the schema refreshed from another connection", async () => {
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
try {
const stale = await first.createTable("docs", [{ text: "before" }], {
schema: embeddingSchema(),
});
const replacement = await second.createTable(
"docs",
[{ text: "after hello" }],
{ mode: "overwrite" },
);
await replacement.createIndex("text", { config: Index.fts() });
const search = stale.search("hello");
expect(search).toBeInstanceOf(AutoQuery);
expect(search).not.toBeInstanceOf(Query);
expect(search).not.toBeInstanceOf(VectorQuery);
expect("nprobes" in search).toBe(false);
const rows = await search.toArray();
expect(rows[0].text).toBe("after hello");
expect((await stale.schema()).metadata.has("embedding_functions")).toBe(
false,
);
} finally {
first.close();
second.close();
}
});
it("tracks embedding metadata across checkout and restore", async () => {
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
try {
await first.createTable("docs", [{ text: "before" }], {
schema: embeddingSchema(),
});
const table = await second.createTable(
"docs",
[{ text: "after hello" }],
{ mode: "overwrite" },
);
await table.createIndex("text", { config: Index.fts() });
await table.checkout(1);
expect((await table.search("before").toArray())[0].text).toBe("before");
await table.checkoutLatest();
expect((await table.search("hello").toArray())[0].text).toBe(
"after hello",
);
await table.checkout(1);
await table.restore();
expect((await table.search("before").toArray())[0].text).toBe("before");
} finally {
first.close();
second.close();
}
});
it("pins automatic search while computing an embedding", async () => {
let markStarted!: () => void;
let releaseEmbedding!: () => void;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
const released = new Promise<void>((resolve) => {
releaseEmbedding = resolve;
});
class BlockingEmbedding extends SchemaRefreshEmbedding {
async computeQueryEmbeddings(value: string) {
markStarted();
await released;
return [value.length, 1];
}
}
register("schema-refresh-blocking")(BlockingEmbedding);
const func = new BlockingEmbedding();
const schema = LanceSchema({
text: func.sourceField(new Utf8()),
vector: func.vectorField(),
});
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
try {
const table = await first.createTable(
"docs",
[{ text: "hello before" }],
{ schema },
);
const pending = table.search("hello").toArray();
await started;
const replacement = await second.createTable(
"docs",
[{ text: "hello after" }],
{ mode: "overwrite" },
);
await replacement.createIndex("text", { config: Index.fts() });
releaseEmbedding();
expect((await pending)[0].text).toBe("hello before");
} finally {
releaseEmbedding();
first.close();
second.close();
}
});
it("refreshes a reused automatic search for every execution", async () => {
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
try {
const table = await first.createTable("docs", [
{ text: "hello before", marker: "before" },
]);
await table.createIndex("text", { config: Index.fts() });
const search = table.search("hello").select(["text"]);
const before = (await search.toArray())[0];
expect(before.text).toBe("hello before");
expect(before.marker).toBeUndefined();
const replacement = await second.createTable(
"docs",
[{ text: "hello after", marker: "after" }],
{ mode: "overwrite" },
);
await replacement.createIndex("text", { config: Index.fts() });
const after = (await search.toArray())[0];
expect(after.text).toBe("hello after");
expect(after.marker).toBeUndefined();
} finally {
first.close();
second.close();
}
});
});
describe("schema evolution", function () {
let tmpDir: tmp.DirResult;
beforeEach(() => {
@@ -2344,7 +2585,24 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
);
});
test("full text search if no embedding function provided", async () => {
test("full text search if only an unrelated embedding function is registered", async () => {
register("unused")(
class extends EmbeddingFunction<string> {
ndims() {
return 3;
}
embeddingDataType() {
return new Float32();
}
async computeQueryEmbeddings(_data: string) {
return [1, 2, 3];
}
async computeSourceEmbeddings(data: string[]) {
return data.map(() => [1, 2, 3]);
}
},
);
const db = await connect(tmpDir.name);
const data = [
{ text: "hello world", vector: [0.1, 0.2, 0.3] },
@@ -2366,6 +2624,306 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
expect(results2[0].text).toBe(data[1].text);
});
test("auto search stays consistent with the active revision", async () => {
let initCalls = 0;
let queryCalls = 0;
let markStarted!: () => void;
const started = new Promise<void>((resolve) => {
markStarted = resolve;
});
let releaseEmbedding!: () => void;
const embeddingReleased = new Promise<void>((resolve) => {
releaseEmbedding = resolve;
});
@register("refresh-test")
class TestEmbedding extends EmbeddingFunction<string> {
async init() {
initCalls += 1;
}
ndims() {
return 1;
}
embeddingDataType() {
return new arrow.Float32();
}
async computeQueryEmbeddings(value: string) {
queryCalls += 1;
if (value === "blocked") {
markStarted();
await embeddingReleased;
}
return value === "greetings" ? [0.1] : [0.2];
}
async computeSourceEmbeddings(values: string[]) {
return values.map((value) =>
value === "hello world" ? [0.1] : [0.2],
);
}
}
const writer = await connect(tmpDir.name);
await writer.createTable("test", [{ text: "plain", vector: [0.0] }]);
const reader = await connect(tmpDir.name, {
readConsistencyInterval: 0,
});
const tracked = await reader.openTable("test");
type SnapshotCountingNative = {
querySnapshot: () => Promise<unknown>;
};
const native = (tracked as unknown as { inner: SnapshotCountingNative })
.inner;
const querySnapshot = native.querySnapshot.bind(native);
let snapshotCalls = 0;
native.querySnapshot = async () => {
snapshotCalls += 1;
return await querySnapshot();
};
const autoQuery = tracked.search("greetings").select(["text"]).limit(1);
const func = new TestEmbedding();
const schema = LanceSchema({
text: func.sourceField(new arrow.Utf8()),
vector: func.vectorField(),
});
const data = [{ text: "hello world" }, { text: "goodbye world" }];
await writer.createTable("test", data, { mode: "overwrite", schema });
const baselineInitCalls = initCalls;
expect(
(await tracked.schema()).metadata.get("embedding_functions"),
).toBeDefined();
const results = await autoQuery.toArray();
expect(results[0].text).toBe(data[0].text);
expect(initCalls).toBe(baselineInitCalls + 1);
expect(queryCalls).toBe(1);
expect(snapshotCalls).toBe(1);
const repeatedResults = await autoQuery.toArray();
expect(repeatedResults[0].text).toBe(data[0].text);
expect(initCalls).toBe(baselineInitCalls + 1);
expect(queryCalls).toBe(1);
expect(snapshotCalls).toBe(2);
const pending = tracked
.search("blocked")
.select(["text"])
.limit(1)
.toArray();
await started;
const ftsData = [
{ text: "greetings from full text", vector: [0.0] },
{ text: "blocked from full text", vector: [0.0] },
];
const ftsTable = await writer.createTable("test", ftsData, {
mode: "overwrite",
});
await ftsTable.createIndex("text", { config: Index.fts() });
releaseEmbedding();
const pendingResults = await pending;
expect(pendingResults[0].text).toBe(data[1].text);
expect(
(await tracked.schema()).metadata.get("embedding_functions"),
).toBeUndefined();
const ftsResults = await autoQuery.toArray();
expect(ftsResults[0].text).toBe(ftsData[0].text);
});
test("auto search keeps newer preparation during a revision race", async () => {
let aCalls = 0;
let bCalls = 0;
let markAStarted!: () => void;
const aStarted = new Promise<void>((resolve) => {
markAStarted = resolve;
});
let releaseA!: () => void;
const aReleased = new Promise<void>((resolve) => {
releaseA = resolve;
});
let markBStarted!: () => void;
const bStarted = new Promise<void>((resolve) => {
markBStarted = resolve;
});
let releaseB!: () => void;
const bReleased = new Promise<void>((resolve) => {
releaseB = resolve;
});
@register("race-a")
class EmbeddingA extends EmbeddingFunction<string> {
ndims() {
return 1;
}
embeddingDataType() {
return new arrow.Float32();
}
async computeQueryEmbeddings() {
aCalls += 1;
markAStarted();
await aReleased;
return [0.1];
}
async computeSourceEmbeddings(values: string[]) {
return values.map(() => [0.1]);
}
}
@register("race-b")
class EmbeddingB extends EmbeddingFunction<string> {
ndims() {
return 1;
}
embeddingDataType() {
return new arrow.Float32();
}
async computeQueryEmbeddings() {
bCalls += 1;
markBStarted();
await bReleased;
return [0.2];
}
async computeSourceEmbeddings(values: string[]) {
return values.map(() => [0.2]);
}
}
const writer = await connect(tmpDir.name);
const embeddingA = new EmbeddingA();
const schemaA = LanceSchema({
text: embeddingA.sourceField(new arrow.Utf8()),
vector: embeddingA.vectorField(),
});
await writer.createTable("race", [{ text: "revision a" }], {
schema: schemaA,
});
const reader = await connect(tmpDir.name, {
readConsistencyInterval: 0,
});
const tracked = await reader.openTable("race");
const query = tracked.search("query");
const first = query.toArray();
await aStarted;
const embeddingB = new EmbeddingB();
const schemaB = LanceSchema({
text: embeddingB.sourceField(new arrow.Utf8()),
vector: embeddingB.vectorField(),
});
await writer.createTable("race", [{ text: "revision b" }], {
mode: "overwrite",
schema: schemaB,
});
const second = query.toArray();
await bStarted;
releaseA();
releaseB();
await Promise.all([first, second]);
expect(aCalls).toBe(1);
expect(bCalls).toBe(1);
});
test("stale FTS routing keeps newer vector preparation", async () => {
let vectorCalls = 0;
let markVectorStarted!: () => void;
const vectorStarted = new Promise<void>((resolve) => {
markVectorStarted = resolve;
});
let releaseVector!: () => void;
const vectorReleased = new Promise<void>((resolve) => {
releaseVector = resolve;
});
@register("stale-fts-race")
class RaceEmbedding extends EmbeddingFunction<string> {
ndims() {
return 1;
}
embeddingDataType() {
return new arrow.Float32();
}
async computeQueryEmbeddings() {
vectorCalls += 1;
markVectorStarted();
await vectorReleased;
return [0.1];
}
async computeSourceEmbeddings(values: string[]) {
return values.map(() => [0.1]);
}
}
const writer = await connect(tmpDir.name);
const ftsTable = await writer.createTable("stale_fts", [
{ text: "hello", vector: [0.0] },
]);
await ftsTable.createIndex("text", { config: Index.fts() });
const reader = await connect(tmpDir.name, {
readConsistencyInterval: 0,
});
const tracked = await reader.openTable("stale_fts");
type Snapshot = {
schema: () => Promise<Buffer>;
};
type NativeWithSnapshot = {
querySnapshot: () => Promise<Snapshot>;
};
const native = (tracked as unknown as { inner: NativeWithSnapshot })
.inner;
const querySnapshot = native.querySnapshot.bind(native);
let snapshotCalls = 0;
let markStaleSchemaStarted!: () => void;
const staleSchemaStarted = new Promise<void>((resolve) => {
markStaleSchemaStarted = resolve;
});
let releaseStaleSchema!: () => void;
const staleSchemaReleased = new Promise<void>((resolve) => {
releaseStaleSchema = resolve;
});
native.querySnapshot = async () => {
const snapshot = await querySnapshot();
snapshotCalls += 1;
if (snapshotCalls === 1) {
const schema = snapshot.schema.bind(snapshot);
snapshot.schema = async () => {
markStaleSchemaStarted();
await staleSchemaReleased;
return await schema();
};
}
return snapshot;
};
const query = tracked.search("hello");
const staleFtsExecution = query.toArray();
await staleSchemaStarted;
const embedding = new RaceEmbedding();
const vectorSchema = LanceSchema({
text: embedding.sourceField(new arrow.Utf8()),
vector: embedding.vectorField(),
});
await writer.createTable("stale_fts", [{ text: "hello" }], {
mode: "overwrite",
schema: vectorSchema,
});
const vectorExecution = query.toArray();
await vectorStarted;
releaseStaleSchema();
await staleFtsExecution;
releaseVector();
await vectorExecution;
await query.toArray();
expect(vectorCalls).toBe(1);
});
test("tokenizes FTS queries by column or index name", async () => {
const db = await connect(tmpDir.name);
const data = [
@@ -2916,6 +3474,30 @@ describe("column name options", () => {
expect(results[1].query_index).toBe(1);
});
test("observes promised additional vectors while the query is pending", async () => {
const initialVector = new Promise<number[]>(() => undefined);
const query = table.query().nearestTo(initialVector);
const unhandled: unknown[] = [];
const onUnhandled = (reason: unknown) => unhandled.push(reason);
process.on("unhandledRejection", onUnhandled);
try {
query.addQueryVector(Promise.reject(new Error("extra vector failed")));
await new Promise<void>((resolve) => setImmediate(resolve));
expect(unhandled).toEqual([]);
const rejectedQuery = table
.query()
.nearestTo([0.1, 0.2])
.addQueryVector(Promise.reject(new Error("consumed vector failed")));
await expect(rejectedQuery.toArray()).rejects.toThrow(
"consumed vector failed",
);
} finally {
process.off("unhandledRejection", onUnhandled);
}
});
test("index and search multivectors", async () => {
const db = await connect(tmpDir.name);
const data = [];
@@ -2979,6 +3561,27 @@ describe("when creating an empty table", () => {
expect((actualSchema.fields[1].type as Float64).precision).toBe(2);
});
it("can add and query JSON data", async () => {
const schema = new Schema([
new Field("id", new Int32(), true),
new Field(
"meta",
new Utf8(),
true,
new Map([["ARROW:extension:name", "arrow.json"]]),
),
]);
const table = await con.createEmptyTable("json", schema);
const meta = JSON.stringify({ x: 1 });
await table.add([{ id: 1, meta }]);
const rows = await table.query().toArray();
expect(rows).toHaveLength(1);
expect(rows[0].id).toBe(1);
expect(rows[0].meta).toBe(meta);
});
it("can create an empty table from schema that specifies field types by name", async () => {
const schemaLike = {
fields: [
+1 -1
View File
@@ -170,7 +170,7 @@ test("basic table examples", async () => {
// --8<-- [end:create_index]
// --8<-- [start:delete_rows]
await tbl.delete('item = "fizz"');
await tbl.delete("item = 'fizz'");
// --8<-- [end:delete_rows]
// --8<-- [start:drop_table]
+35 -309
View File
@@ -5,7 +5,6 @@ import {
Data as ArrowData,
Table as ArrowTable,
Binary,
Bool,
BufferType,
DataType,
DateUnit,
@@ -18,12 +17,7 @@ import {
FixedSizeList,
Float,
Float32,
Float64,
Int,
Int8,
Int16,
Int32,
Int64,
LargeBinary,
List,
Null,
@@ -36,17 +30,16 @@ import {
Struct,
Timestamp,
Type,
Uint8,
Uint16,
Uint32,
Utf8,
Vector,
makeVector as arrowMakeVector,
util as arrowUtil,
vectorFromArray as badVectorFromArray,
makeBuilder,
makeData,
} from "apache-arrow";
import { Buffers } from "apache-arrow/data";
import { typedArrayToArrowType } from "./arrow_type";
import { type EmbeddingFunction } from "./embedding/embedding_function";
import {
EmbeddingFunctionConfig,
@@ -59,14 +52,7 @@ import {
sanitizeTable,
sanitizeType,
} from "./sanitize";
/**
* Check if a field name indicates a vector column.
*/
function nameSuggestsVectorColumn(fieldName: string): boolean {
const nameLower = fieldName.toLowerCase();
return nameLower.includes("vector") || nameLower.includes("embedding");
}
import { inferSchema } from "./schema";
export * from "apache-arrow";
export type SchemaLike =
@@ -86,8 +72,7 @@ export type FieldLike =
};
export type DataLike =
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
| import("apache-arrow").Data<Struct<any>>
| import("apache-arrow").Data
| {
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
type: any;
@@ -96,6 +81,7 @@ export type DataLike =
stride: number;
nullable: boolean;
children: DataLike[];
dictionary?: { data: readonly DataLike[] };
get nullCount(): number;
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
values: Buffers<any>[BufferType.DATA];
@@ -459,110 +445,6 @@ export function makeArrowTable(
return new ArrowTable(inferredSchema, finalColumns);
}
function inferSchema(
data: Array<Record<string, unknown>>,
schema: Schema | undefined,
opts: MakeArrowTableOptions,
): Schema {
// We will collect all fields we see in the data.
const pathTree = new PathTree<DataType>();
for (const [rowI, row] of data.entries()) {
for (const [path, value] of rowPathsAndValues(row)) {
if (!pathTree.has(path)) {
// First time seeing this field.
if (schema !== undefined) {
const field = getFieldForPath(schema, path);
if (field === undefined) {
throw new Error(
`Found field not in schema: ${path.join(".")} at row ${rowI}`,
);
} else {
pathTree.set(path, field.type);
}
} else {
const inferredType = inferType(value, path, opts);
if (inferredType === undefined) {
throw new Error(`Failed to infer data type for field ${path.join(
".",
)} at row ${rowI}. \
Consider providing an explicit schema.`);
}
pathTree.set(path, inferredType);
}
} else if (schema === undefined) {
const currentType = pathTree.get(path);
const newType = inferType(value, path, opts);
if (currentType !== newType) {
new Error(`Failed to infer schema for data. Previously inferred type \
${currentType} but found ${newType} at row ${rowI}. Consider \
providing an explicit schema.`);
}
}
}
}
if (schema === undefined) {
function fieldsFromPathTree(pathTree: PathTree<DataType>): Field[] {
const fields = [];
for (const [name, value] of pathTree.map.entries()) {
if (value instanceof PathTree) {
const children = fieldsFromPathTree(value);
fields.push(new Field(name, new Struct(children), true));
} else {
fields.push(new Field(name, value, true));
}
}
return fields;
}
const fields = fieldsFromPathTree(pathTree);
return new Schema(fields);
} else {
function takeMatchingFields(
fields: Field[],
pathTree: PathTree<DataType>,
): Field[] {
const outFields = [];
for (const field of fields) {
if (pathTree.map.has(field.name)) {
const value = pathTree.get([field.name]);
if (value instanceof PathTree) {
const struct = field.type as Struct;
const children = takeMatchingFields(struct.children, value);
outFields.push(
new Field(field.name, new Struct(children), field.nullable),
);
} else {
outFields.push(
new Field(field.name, value as DataType, field.nullable),
);
}
}
}
return outFields;
}
const fields = takeMatchingFields(schema.fields, pathTree);
return new Schema(fields);
}
}
function* rowPathsAndValues(
row: Record<string, unknown>,
basePath: string[] = [],
): Generator<[string[], unknown]> {
for (const [key, value] of Object.entries(row)) {
if (isObject(value)) {
yield* rowPathsAndValues(value, [...basePath, key]);
} else {
// Skip undefined values - they should be treated the same as missing fields
// for embedding function purposes
if (value !== undefined) {
yield [[...basePath, key], value];
}
}
}
}
function isObject(value: unknown): value is Record<string, unknown> {
return (
typeof value === "object" &&
@@ -577,146 +459,19 @@ function isObject(value: unknown): value is Record<string, unknown> {
);
}
function getFieldForPath(schema: Schema, path: string[]): Field | undefined {
let current: Field | Schema = schema;
function valueAtPath(datum: Record<string, unknown>, path: string[]): unknown {
let current: unknown = datum;
for (const key of path) {
if (current instanceof Schema) {
const field: Field | undefined = current.fields.find(
(f) => f.name === key,
);
if (field === undefined) {
return undefined;
}
current = field;
} else if (current instanceof Field && DataType.isStruct(current.type)) {
const struct: Struct = current.type;
const field = struct.children.find((f) => f.name === key);
if (field === undefined) {
return undefined;
}
current = field;
if (current == null) {
return null;
}
if (isObject(current) && (Object.hasOwn(current, key) || key in current)) {
current = current[key];
} else {
return undefined;
}
}
if (current instanceof Field) {
return current;
} else {
return undefined;
}
}
/**
* Try to infer which Arrow type to use for a given value.
*
* May return undefined if the type cannot be inferred.
*/
function inferType(
value: unknown,
path: string[],
opts: MakeArrowTableOptions,
): DataType | undefined {
if (typeof value === "bigint") {
return new Int64();
} else if (typeof value === "number") {
// Even if it's an integer, it's safer to assume Float64. Users can
// always provide an explicit schema or use BigInt if they mean integer.
return new Float64();
} else if (typeof value === "string") {
if (opts.dictionaryEncodeStrings) {
return new Dictionary(new Utf8(), new Int32());
} else {
return new Utf8();
}
} else if (typeof value === "boolean") {
return new Bool();
} else if (value instanceof Buffer) {
return new Binary();
} else if (ArrayBuffer.isView(value) && !(value instanceof DataView)) {
const info = typedArrayToArrowType(value);
if (info !== undefined) {
const child = new Field("item", info.elementType, true);
return new FixedSizeList(info.length, child);
}
return undefined;
} else if (Array.isArray(value)) {
if (value.length === 0) {
return undefined; // Without any values we can't infer the type
}
if (path.length === 1 && Object.hasOwn(opts.vectorColumns, path[0])) {
const floatType = sanitizeType(opts.vectorColumns[path[0]].type);
return new FixedSizeList(
value.length,
new Field("item", floatType, true),
);
}
const valueType = inferType(value[0], path, opts);
if (valueType === undefined) {
return undefined;
}
// Try to automatically detect embedding columns.
if (nameSuggestsVectorColumn(path[path.length - 1])) {
// Check if value is a Uint8Array for integer vector type determination
if (value instanceof Uint8Array) {
// For integer vectors, we default to Uint8 (matching Python implementation)
const child = new Field("item", new Uint8(), true);
return new FixedSizeList(value.length, child);
} else {
// For float vectors, we default to Float32
const child = new Field("item", new Float32(), true);
return new FixedSizeList(value.length, child);
}
} else {
const child = new Field("item", valueType, true);
return new List(child);
}
} else {
// TODO: timestamp
return undefined;
}
}
class PathTree<V> {
map: Map<string, V | PathTree<V>>;
constructor(entries?: [string[], V][]) {
this.map = new Map();
if (entries !== undefined) {
for (const [path, value] of entries) {
this.set(path, value);
}
}
}
has(path: string[]): boolean {
let ref: PathTree<V> = this;
for (const part of path) {
if (!(ref instanceof PathTree) || !ref.map.has(part)) {
return false;
}
ref = ref.map.get(part) as PathTree<V>;
}
return true;
}
get(path: string[]): V | undefined {
let ref: PathTree<V> = this;
for (const part of path) {
if (!(ref instanceof PathTree) || !ref.map.has(part)) {
return undefined;
}
ref = ref.map.get(part) as PathTree<V>;
}
return ref as V;
}
set(path: string[], value: V): void {
let ref: PathTree<V> = this;
for (const part of path.slice(0, path.length - 1)) {
if (!ref.map.has(part)) {
ref.map.set(part, new PathTree<V>());
}
ref = ref.map.get(part) as PathTree<V>;
}
ref.map.set(path[path.length - 1], value);
}
return current;
}
function transposeData(
@@ -724,37 +479,26 @@ function transposeData(
field: Field,
path: string[] = [],
): Vector {
const valuesPath = [...path, field.name];
const values = data.map((datum) => valueAtPath(datum, valuesPath));
if (field.type instanceof Struct) {
const childFields = field.type.children;
const fullPath = [...path, field.name];
const childVectors = childFields.map((child) => {
return transposeData(data, child, fullPath);
return transposeData(data, child, valuesPath);
});
const nullCount = values.filter((value) => value === null).length;
const structData = makeData({
type: field.type,
length: values.length,
nullCount,
nullBitmap:
nullCount > 0
? arrowUtil.packBools(values.map((value) => value !== null))
: undefined,
children: childVectors as unknown as ArrowData<DataType>[],
});
return arrowMakeVector(structData);
} else {
const valuesPath = [...path, field.name];
const values = data.map((datum) => {
let current: unknown = datum;
for (const key of valuesPath) {
if (current == null) {
return null;
}
if (
isObject(current) &&
(Object.hasOwn(current, key) || key in current)
) {
current = current[key];
} else {
return null;
}
}
return current;
});
return makeVector(values, field.type, undefined, field.nullable);
}
}
@@ -797,32 +541,6 @@ function makeListVector(lists: unknown[][]): Vector<unknown> {
return listBuilder.finish().toVector();
}
/**
* Map a JS TypedArray instance to the corresponding Arrow element DataType
* and its length. Returns undefined if the value is not a recognized TypedArray.
*/
function typedArrayToArrowType(
value: ArrayBufferView,
): { elementType: DataType; length: number } | undefined {
if (value instanceof Float32Array)
return { elementType: new Float32(), length: value.length };
if (value instanceof Float64Array)
return { elementType: new Float64(), length: value.length };
if (value instanceof Uint8Array)
return { elementType: new Uint8(), length: value.length };
if (value instanceof Uint16Array)
return { elementType: new Uint16(), length: value.length };
if (value instanceof Uint32Array)
return { elementType: new Uint32(), length: value.length };
if (value instanceof Int8Array)
return { elementType: new Int8(), length: value.length };
if (value instanceof Int16Array)
return { elementType: new Int16(), length: value.length };
if (value instanceof Int32Array)
return { elementType: new Int32(), length: value.length };
return undefined;
}
/** Helper function to convert an Array of JS values to an Arrow Vector */
function makeVector(
values: unknown[],
@@ -1462,8 +1180,12 @@ export function ensureNestedFieldsExist(
completeRow[field.name] = row[field.name];
}
} else {
// Field is missing from the data - set to null
completeRow[field.name] = null;
// Keep a missing struct valid while filling each of its children with
// null. This is distinct from an explicitly null struct value.
completeRow[field.name] =
field.type.constructor.name === "Struct"
? ensureStructFieldsExist({}, field.type as Struct)
: null;
}
}
@@ -1498,8 +1220,12 @@ function ensureStructFieldsExist(
completeStruct[childField.name] = data[childField.name];
}
} else {
// Field is missing - set to null
completeStruct[childField.name] = null;
// Keep a missing struct valid while filling each of its children with
// null. This is distinct from an explicitly null struct value.
completeStruct[childField.name] =
childField.type.constructor.name === "Struct"
? ensureStructFieldsExist({}, childField.type as Struct)
: null;
}
}
+40
View File
@@ -0,0 +1,40 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import {
type DataType,
Float32,
Float64,
Int8,
Int16,
Int32,
Uint8,
Uint16,
Uint32,
} from "apache-arrow";
/**
* Map a JS TypedArray instance to the corresponding Arrow element type and
* length. Returns undefined when the view is not a supported TypedArray.
*/
export function typedArrayToArrowType(
value: ArrayBufferView,
): { elementType: DataType; length: number } | undefined {
if (value instanceof Float32Array)
return { elementType: new Float32(), length: value.length };
if (value instanceof Float64Array)
return { elementType: new Float64(), length: value.length };
if (value instanceof Uint8Array)
return { elementType: new Uint8(), length: value.length };
if (value instanceof Uint16Array)
return { elementType: new Uint16(), length: value.length };
if (value instanceof Uint32Array)
return { elementType: new Uint32(), length: value.length };
if (value instanceof Int8Array)
return { elementType: new Int8(), length: value.length };
if (value instanceof Int16Array)
return { elementType: new Int16(), length: value.length };
if (value instanceof Int32Array)
return { elementType: new Int32(), length: value.length };
return undefined;
}
+85
View File
@@ -31,12 +31,14 @@ import type {
JobDescription,
JobInfo,
ListNamespacesResponse,
ListTablesResponse,
} from "./native";
export type {
CreateNamespaceResponse,
DescribeNamespaceResponse,
DropNamespaceResponse,
ListNamespacesResponse,
ListTablesResponse,
};
import { sanitizeTable } from "./sanitize";
import { LocalTable, Table } from "./table";
@@ -134,6 +136,10 @@ export interface OpenTableOptions {
indexCacheSize?: number;
}
/**
* @deprecated Use {@link ListTablesOptions} with {@link Connection.listTables}
* instead.
*/
export interface TableNamesOptions {
/**
* If present, only return names that come lexicographically after the
@@ -147,6 +153,24 @@ 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.
*
* 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;
@@ -231,6 +255,7 @@ export abstract class Connection {
* @param {Partial<TableNamesOptions>} options - options to control the
* paging / start point (backwards compatibility)
*
* @deprecated Use {@link Connection.listTables} instead.
*/
abstract tableNames(options?: Partial<TableNamesOptions>): Promise<string[]>;
/**
@@ -241,12 +266,53 @@ export abstract class Connection {
* @param {Partial<TableNamesOptions>} options - options to control the
* paging / start point
*
* @deprecated Use {@link Connection.listTables} instead.
*/
abstract tableNames(
namespacePath?: string[],
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:
*
* ```ts
* const names = [];
* let pageToken = undefined;
* do {
* const page = await conn.listTables({ pageToken, limit: 100 });
* names.push(...page.tables);
* pageToken = page.pageToken;
* } while (pageToken);
* ```
*
* @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>;
/**
* Open a table in the database.
* @param {string} name - The name of the table
@@ -601,6 +667,25 @@ export class LocalConnection extends Connection {
return await this.inner.listMaterializedViews();
}
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,
);
}
async openTable(
name: string,
namespacePath?: string[],
+42 -2
View File
@@ -4,7 +4,15 @@
import { Field, Schema } from "../arrow";
import { sanitizeType } from "../sanitize";
import { EmbeddingFunction } from "./embedding_function";
import { EmbeddingFunctionConfig, getRegistry } from "./registry";
import {
EmbeddingFunctionConfig,
EmbeddingFunctionRegistry,
getRegistry as getGlobalRegistry,
registerBuiltIn,
} from "./registry";
type OpenAIModule = typeof import("./openai");
type TransformersModule = typeof import("./transformers");
export {
FieldOptions,
@@ -14,7 +22,39 @@ export {
EmbeddingFunctionConstructor,
} from "./embedding_function";
export * from "./registry";
export {
EmbeddingFunctionRegistry,
parseEmbeddingMetadata,
register,
} from "./registry";
export type {
CreateReturnType,
EmbeddingFunctionConfig,
EmbeddingFunctionCreate,
EmbeddingMetadataEntry,
ResolvedEmbeddingFunctionConfig,
} from "./registry";
function initializeBuiltInProviders() {
const { OpenAIEmbeddingFunction } = require("./openai") as OpenAIModule;
const { TransformersEmbeddingFunction } =
require("./transformers") as TransformersModule;
registerBuiltIn("openai", OpenAIEmbeddingFunction);
registerBuiltIn("huggingface", TransformersEmbeddingFunction);
}
/**
* Get the global embedding function registry.
*
* LanceDB built-in providers are initialized when this public API is first
* used, so importing the root package does not change automatic search
* selection for tables without embedding metadata.
*/
export function getRegistry(): EmbeddingFunctionRegistry {
initializeBuiltInProviders();
return getGlobalRegistry();
}
/**
* Create a schema with embedding functions.
+3 -2
View File
@@ -5,14 +5,13 @@ import type OpenAI from "openai";
import type { EmbeddingCreateParams } from "openai/resources/index";
import { Float, Float32 } from "../arrow";
import { EmbeddingFunction } from "./embedding_function";
import { register } from "./registry";
import { registerBuiltIn } from "./registry";
export type OpenAIOptions = {
apiKey: string;
model: EmbeddingCreateParams["model"];
};
@register("openai")
export class OpenAIEmbeddingFunction extends EmbeddingFunction<
string,
Partial<OpenAIOptions>
@@ -100,3 +99,5 @@ export class OpenAIEmbeddingFunction extends EmbeddingFunction<
return response.data[0].embedding;
}
}
registerBuiltIn("openai", OpenAIEmbeddingFunction);
+59 -1
View File
@@ -7,6 +7,10 @@ import {
} from "./embedding_function";
import "reflect-metadata";
const builtInFunctionsKey = Symbol.for(
"@lancedb/lancedb::embedding-built-in-functions::v1",
);
export type CreateReturnType<T> = T extends { init: () => Promise<void> }
? Promise<T>
: T;
@@ -59,6 +63,15 @@ export class EmbeddingFunctionRegistry {
};
}
/** @ignore */
setBuiltIn<
T extends EmbeddingFunctionConstructor = EmbeddingFunctionConstructor,
>(name: string, ctor: T): T {
this.#functions.set(name, ctor);
Reflect.defineMetadata("lancedb::embedding::name", name, ctor);
return ctor;
}
get<T extends EmbeddingFunction<unknown>>(
name: string,
): EmbeddingFunctionCreate<T> | undefined;
@@ -96,6 +109,7 @@ export class EmbeddingFunctionRegistry {
*/
reset(this: EmbeddingFunctionRegistry) {
this.#functions.clear();
getBuiltInFunctions(this).clear();
}
/**
@@ -183,12 +197,56 @@ export class EmbeddingFunctionRegistry {
}
}
const _REGISTRY = new EmbeddingFunctionRegistry();
function getBuiltInFunctions(registry: EmbeddingFunctionRegistry): Set<string> {
const registryWithBuiltIns = registry as EmbeddingFunctionRegistry & {
[key: symbol]: Set<string> | undefined;
};
let builtInFunctions = registryWithBuiltIns[builtInFunctionsKey];
if (builtInFunctions === undefined) {
builtInFunctions = new Set<string>();
registryWithBuiltIns[builtInFunctionsKey] = builtInFunctions;
}
return builtInFunctions;
}
// Server bundlers can load the side-effect embedding entry points and the public
// embedding API from separate module graphs. Keep their registry shared.
const registryKey = Symbol.for(
"@lancedb/lancedb::embedding-function-registry::v1",
);
const registryGlobal = globalThis as typeof globalThis & {
[key: symbol]: EmbeddingFunctionRegistry | undefined;
};
function getGlobalRegistry(): EmbeddingFunctionRegistry {
const existingRegistry = registryGlobal[registryKey];
if (existingRegistry !== undefined) {
return existingRegistry;
}
const registry = new EmbeddingFunctionRegistry();
registryGlobal[registryKey] = registry;
return registry;
}
const _REGISTRY = getGlobalRegistry();
export function register(name?: string) {
return _REGISTRY.register(name);
}
/** @ignore */
export function registerBuiltIn<
T extends EmbeddingFunctionConstructor = EmbeddingFunctionConstructor,
>(name: string, ctor: T): T {
const builtInFunctions = getBuiltInFunctions(_REGISTRY);
if (builtInFunctions.has(name)) {
return _REGISTRY.setBuiltIn(name, ctor);
}
_REGISTRY.register(name)(ctor);
builtInFunctions.add(name);
return ctor;
}
/**
* Utility function to get the global instance of the registry
* @returns `EmbeddingFunctionRegistry` The global instance of the registry
+3 -2
View File
@@ -3,7 +3,7 @@
import { Float, Float32 } from "../arrow";
import { EmbeddingFunction } from "./embedding_function";
import { register } from "./registry";
import { registerBuiltIn } from "./registry";
export type XenovaTransformerOptions = {
/** The wasm compatible model to use */
@@ -31,7 +31,6 @@ export type XenovaTransformerOptions = {
};
};
@register("huggingface")
export class TransformersEmbeddingFunction extends EmbeddingFunction<
string,
Partial<XenovaTransformerOptions>
@@ -158,6 +157,8 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
}
}
registerBuiltIn("huggingface", TransformersEmbeddingFunction);
const tensorDiv = (
src: import("@huggingface/transformers").Tensor,
divBy: number,
+3
View File
@@ -81,11 +81,13 @@ export {
Connection,
CreateTableOptions,
TableNamesOptions,
ListTablesOptions,
OpenTableOptions,
ListNamespacesOptions,
CreateNamespaceOptions,
DropNamespaceOptions,
ListNamespacesResponse,
ListTablesResponse,
CreateNamespaceResponse,
DropNamespaceResponse,
DescribeNamespaceResponse,
@@ -101,6 +103,7 @@ export {
} from "./native.js";
export {
AutoQuery,
ExecutableQuery,
Query,
QueryBase,
+205 -106
View File
@@ -100,6 +100,29 @@ export interface FullTextSearchOptions {
columns?: string | string[];
}
function nearestToNative(
inner: NativeQuery,
vector: Awaited<IntoVector>,
): NativeVectorQuery {
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
if (raw) {
return inner.nearestToRaw(raw.data, raw.dtype);
}
return inner.nearestTo(Float32Array.from(vector as number[]));
}
function addQueryVectorToNative(
inner: NativeVectorQuery,
vector: Awaited<IntoVector>,
) {
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
if (raw) {
inner.addQueryVectorRaw(raw.data, raw.dtype);
} else {
inner.addQueryVector(Float32Array.from(vector as number[]));
}
}
/** Common methods supported by all query types
*
* @see {@link Query}
@@ -111,13 +134,15 @@ export class QueryBase<
NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery,
> implements AsyncIterable<RecordBatch>
{
protected inner!: NativeQueryType | Promise<NativeQueryType>;
/**
* @hidden
*/
protected constructor(
protected inner: NativeQueryType | Promise<NativeQueryType>,
) {
// intentionally empty
protected constructor(inner?: NativeQueryType | Promise<NativeQueryType>) {
if (inner !== undefined) {
this.inner = inner;
}
}
// call a function on the inner (either a promise or the actual object)
@@ -135,6 +160,15 @@ export class QueryBase<
}
}
/**
* Return the native query used by the next terminal operation.
*
* @hidden
*/
protected async getInner(): Promise<NativeQueryType> {
return this.inner;
}
/**
* Return only the specified columns.
*
@@ -207,16 +241,11 @@ export class QueryBase<
/**
* @hidden
*/
protected nativeExecute(
protected async nativeExecute(
options?: Partial<QueryExecutionOptions>,
): Promise<NativeBatchIterator> {
if (this.inner instanceof Promise) {
return this.inner.then((inner) =>
inner.execute(options?.maxBatchLength, options?.timeoutMs),
);
} else {
return this.inner.execute(options?.maxBatchLength, options?.timeoutMs);
}
const inner = await this.getInner();
return inner.execute(options?.maxBatchLength, options?.timeoutMs);
}
/**
@@ -245,12 +274,7 @@ export class QueryBase<
/** Collect the results as an Arrow @see {@link ArrowTable}. */
async toArrow(options?: Partial<QueryExecutionOptions>): Promise<ArrowTable> {
const batches = [];
let inner;
if (this.inner instanceof Promise) {
inner = await this.inner;
} else {
inner = this.inner;
}
const inner = await this.getInner();
for await (const batch of new RecordBatchIterable(inner, options)) {
batches.push(batch);
}
@@ -279,11 +303,8 @@ export class QueryBase<
* @returns A Promise that resolves to a string containing the query execution plan explanation.
*/
async explainPlan(verbose = false): Promise<string> {
if (this.inner instanceof Promise) {
return this.inner.then((inner) => inner.explainPlan(verbose));
} else {
return this.inner.explainPlan(verbose);
}
const inner = await this.getInner();
return inner.explainPlan(verbose);
}
/**
@@ -321,13 +342,8 @@ export class QueryBase<
distributedMetrics?: AnalyzePlanDistributedMetrics,
): Promise<string> {
const distributedMetricsMode = distributedMetrics ?? "aggregate";
if (this.inner instanceof Promise) {
return this.inner.then((inner) =>
inner.analyzePlan(distributedMetricsMode),
);
} else {
return this.inner.analyzePlan(distributedMetricsMode);
}
const inner = await this.getInner();
return inner.analyzePlan(distributedMetricsMode);
}
/**
@@ -339,12 +355,8 @@ export class QueryBase<
* @returns An Arrow Schema describing the output columns.
*/
async outputSchema(): Promise<import("./arrow").Schema> {
let schemaBuffer: Buffer;
if (this.inner instanceof Promise) {
schemaBuffer = await this.inner.then((inner) => inner.outputSchema());
} else {
schemaBuffer = await this.inner.outputSchema();
}
const inner = await this.getInner();
const schemaBuffer = await inner.outputSchema();
const schema = tableFromIPC(schemaBuffer).schema;
return schema;
}
@@ -356,7 +368,7 @@ export class StandardQueryBase<
extends QueryBase<NativeQueryType>
implements ExecutableQuery
{
constructor(inner: NativeQueryType | Promise<NativeQueryType>) {
constructor(inner?: NativeQueryType | Promise<NativeQueryType>) {
super(inner);
}
@@ -510,6 +522,13 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
super(inner);
}
/**
* @hidden
*/
protected doVectorCall(fn: (inner: NativeVectorQuery) => void) {
super.doCall(fn);
}
/**
* Set the number of partitions to search (probe)
*
@@ -537,7 +556,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* the minimum and maximum to the same value.
*/
nprobes(nprobes: number): VectorQuery {
super.doCall((inner) => inner.nprobes(nprobes));
this.doVectorCall((inner) => inner.nprobes(nprobes));
return this;
}
@@ -551,7 +570,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* but will also increase latency.
*/
minimumNprobes(minimumNprobes: number): VectorQuery {
super.doCall((inner) => inner.minimumNprobes(minimumNprobes));
this.doVectorCall((inner) => inner.minimumNprobes(minimumNprobes));
return this;
}
@@ -565,7 +584,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* potential false negatives.
*/
maximumNprobes(maximumNprobes: number): VectorQuery {
super.doCall((inner) => inner.maximumNprobes(maximumNprobes));
this.doVectorCall((inner) => inner.maximumNprobes(maximumNprobes));
return this;
}
@@ -578,7 +597,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* `undefined` means no lower or upper bound.
*/
distanceRange(lowerBound?: number, upperBound?: number): VectorQuery {
super.doCall((inner) => inner.distanceRange(lowerBound, upperBound));
this.doVectorCall((inner) => inner.distanceRange(lowerBound, upperBound));
return this;
}
@@ -592,7 +611,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* also increase the latency of your query. The default value is 1.5*limit.
*/
ef(ef: number): VectorQuery {
super.doCall((inner) => inner.ef(ef));
this.doVectorCall((inner) => inner.ef(ef));
return this;
}
@@ -606,7 +625,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* whose data type is a fixed-size-list of floats.
*/
column(column: string): VectorQuery {
super.doCall((inner) => inner.column(column));
this.doVectorCall((inner) => inner.column(column));
return this;
}
@@ -627,7 +646,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
distanceType(
distanceType: Required<IvfPqOptions>["distanceType"],
): VectorQuery {
super.doCall((inner) => inner.distanceType(distanceType));
this.doVectorCall((inner) => inner.distanceType(distanceType));
return this;
}
@@ -661,7 +680,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* distance between the query vector and the actual uncompressed vector.
*/
refineFactor(refineFactor: number): VectorQuery {
super.doCall((inner) => inner.refineFactor(refineFactor));
this.doVectorCall((inner) => inner.refineFactor(refineFactor));
return this;
}
@@ -686,7 +705,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* factor can often help restore some of the results lost by post filtering.
*/
postfilter(): VectorQuery {
super.doCall((inner) => inner.postfilter());
this.doVectorCall((inner) => inner.postfilter());
return this;
}
@@ -700,7 +719,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* calculate your recall to select an appropriate value for nprobes.
*/
bypassVectorIndex(): VectorQuery {
super.doCall((inner) => inner.bypassVectorIndex());
this.doVectorCall((inner) => inner.bypassVectorIndex());
return this;
}
@@ -708,43 +727,39 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
* Add a query vector to the search
*
* This method can be called multiple times to add multiple query vectors
* to the search. If multiple query vectors are added, then they will be searched
* in parallel, and the results will be concatenated. A column called `query_index`
* will be added to indicate the index of the query vector that produced the result.
*
* Performance wise, this is equivalent to running multiple queries concurrently.
* to the search. A column called `query_index` will be added to indicate the index
* of the query vector that produced the result. Flat searches share one table scan
* across the query vectors, avoiding the scan and memory amplification of running
* multiple queries concurrently. Indexed searches may still perform per-vector
* index work.
*/
addQueryVector(vector: IntoVector): VectorQuery {
if (vector instanceof Promise) {
// Observe the promise as soon as it is accepted. The existing native
// query may still be pending, and delaying observation until it resolves
// can otherwise surface a fast rejection as unhandled.
const settledVector = vector.then(
(value) => ({ status: "fulfilled" as const, value }),
(reason) => ({ status: "rejected" as const, reason }),
);
const res = (async () => {
try {
const v = await vector;
// biome-ignore lint/suspicious/noExplicitAny: we need to get the `inner`, but js has no package scoping
const value: any = this.addQueryVector(v);
const inner = value.inner as
| NativeVectorQuery
| Promise<NativeVectorQuery>;
return inner;
} catch (e) {
return Promise.reject(e);
const inner = await this.getInner();
const outcome = await settledVector;
if (outcome.status === "rejected") {
throw outcome.reason;
}
addQueryVectorToNative(inner, outcome.value);
return inner;
})();
return new VectorQuery(res);
} else {
super.doCall((inner) => {
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
if (raw) {
inner.addQueryVectorRaw(raw.data, raw.dtype);
} else {
inner.addQueryVector(Float32Array.from(vector as number[]));
}
});
this.doVectorCall((inner) => addQueryVectorToNative(inner, vector));
return this;
}
}
rerank(reranker: Reranker): VectorQuery {
super.doCall((inner) =>
this.doVectorCall((inner) =>
inner.rerank(async (args) => {
const vecResults = await fromBufferToRecordBatch(args.vecResults);
const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
@@ -763,6 +778,71 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
}
}
/**
* Create a string query whose vector/FTS routing is resolved against the active
* table schema when the query executes.
*
* @hidden
*/
export function createAutoQuery(
table: NativeTable,
query: string,
columns: string[] | null,
getVector: (metadata: string) => Promise<Awaited<IntoVector>>,
): AutoQuery {
type RouteSnapshot = {
table: NativeTable;
embeddingMetadata: string | undefined;
};
type CachedPreparation = {
metadata: string;
vector: Promise<Awaited<IntoVector>>;
};
let cachedPreparation: CachedPreparation | undefined;
const snapshotRoute = async (): Promise<RouteSnapshot> => {
const snapshot = await table.querySnapshot();
const schema = tableFromIPC(await snapshot.schema()).schema;
return {
table: snapshot,
embeddingMetadata: schema.metadata.get("embedding_functions"),
};
};
const createInner = async (): Promise<NativeQuery | NativeVectorQuery> => {
const route = await snapshotRoute();
if (route.embeddingMetadata === undefined) {
const inner = route.table.query();
inner.fullTextSearch({ query, columns });
return inner;
}
const metadata = route.embeddingMetadata;
if (cachedPreparation?.metadata !== metadata) {
cachedPreparation = {
metadata,
vector: Promise.resolve().then(() => getVector(metadata)),
};
}
const preparation = cachedPreparation;
let vector: Awaited<IntoVector>;
try {
vector = await preparation.vector;
} catch (error) {
if (cachedPreparation === preparation) {
cachedPreparation = undefined;
}
throw error;
}
return nearestToNative(route.table.query(), vector);
};
return new AutoQuery(createInner);
}
/**
* A query that returns a subset of the rows in the table.
*
@@ -788,6 +868,51 @@ export class TakeQuery extends QueryBase<NativeTakeQuery> {
}
}
/**
* A builder for automatic string searches.
*
* Automatic search determines whether to use full-text or vector search from
* the table revision selected for each execution. This builder exposes the
* common operations supported by both query families.
*
* @hideconstructor
*/
export class AutoQuery extends StandardQueryBase<
NativeQuery | NativeVectorQuery
> {
private readonly calls: Array<
(inner: NativeQuery | NativeVectorQuery) => void
> = [];
/** @hidden */
constructor(
private readonly createInner: () => Promise<
NativeQuery | NativeVectorQuery
>,
) {
super();
}
/** @hidden */
protected override doCall(
fn: (inner: NativeQuery | NativeVectorQuery) => void,
) {
this.calls.push(fn);
}
/** @hidden */
protected override async getInner(): Promise<
NativeQuery | NativeVectorQuery
> {
const calls = [...this.calls];
const inner = await this.createInner();
for (const call of calls) {
call(inner);
}
return inner;
}
}
/** A builder for LanceDB queries.
*
* @see {@link Table#query}, {@link Table#search}
@@ -840,45 +965,19 @@ export class Query extends StandardQueryBase<NativeQuery> {
* a default `limit` of 10 will be used. @see {@link Query#limit}
*/
nearestTo(vector: IntoVector): VectorQuery {
const callNearestTo = (
inner: NativeQuery,
resolved: Float32Array | Float64Array | Uint8Array | number[],
): NativeVectorQuery => {
const raw = Array.isArray(resolved)
? null
: extractVectorBuffer(resolved);
if (raw) {
return inner.nearestToRaw(raw.data, raw.dtype);
}
return inner.nearestTo(Float32Array.from(resolved as number[]));
};
if (this.inner instanceof Promise) {
const nativeQuery = this.inner.then(async (inner) => {
const resolved = vector instanceof Promise ? await vector : vector;
return callNearestTo(inner, resolved);
});
const inner = this.inner;
if (inner instanceof Promise) {
const nativeQuery = inner.then(async (resolvedInner) =>
nearestToNative(resolvedInner, await vector),
);
return new VectorQuery(nativeQuery);
}
if (vector instanceof Promise) {
const res = (async () => {
try {
const v = await vector;
// biome-ignore lint/suspicious/noExplicitAny: we need to get the `inner`, but js has no package scoping
const value: any = this.nearestTo(v);
const inner = value.inner as
| NativeVectorQuery
| Promise<NativeVectorQuery>;
return inner;
} catch (e) {
return Promise.reject(e);
}
})();
return new VectorQuery(res);
} else {
const vectorQuery = callNearestTo(this.inner, vector);
return new VectorQuery(vectorQuery);
return new VectorQuery(
vector.then((resolvedVector) => nearestToNative(inner, resolvedVector)),
);
}
return new VectorQuery(nearestToNative(inner, vector));
}
nearestToText(query: string | FullTextQuery, columns?: string[]): Query {
+11 -4
View File
@@ -94,17 +94,24 @@ export function sanitizeMetadata(
if (metadataLike === undefined || metadataLike === null) {
return undefined;
}
if (!(metadataLike instanceof Map)) {
let entries: IterableIterator<[unknown, unknown]>;
try {
entries = Map.prototype.entries.call(metadataLike);
} catch {
throw Error("Expected metadata, if present, to be a Map<string, string>");
}
for (const item of metadataLike) {
if (typeof item[0] !== "string" || typeof item[1] !== "string") {
const metadata = new Map<string, string>();
for (const [key, value] of entries) {
if (typeof key !== "string" || typeof value !== "string") {
throw Error(
"Expected metadata, if present, to be a Map<string, string> but it had non-string keys or values",
);
}
metadata.set(key, value);
}
return metadataLike as Map<string, string>;
return metadata;
}
export function sanitizeInt(typeLike: object) {
+567
View File
@@ -0,0 +1,567 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
import {
Binary,
Bool,
DataType,
Dictionary,
Field,
FixedSizeList,
Float32,
Float64,
Int32,
Int64,
List,
Schema,
Struct,
Utf8,
util as arrowUtil,
} from "apache-arrow";
import { typedArrayToArrowType } from "./arrow_type";
import { sanitizeType } from "./sanitize";
type InferenceOptions = {
dictionaryEncodeStrings: boolean;
vectorColumns: Record<string, { type: unknown }>;
};
/**
* Infer the Arrow schema represented by a set of records.
*
* This is the intentionally small interface to schema inference. The stateful
* details of combining partial type evidence are encapsulated below so callers
* only need to provide records, an optional schema, and inference options.
*/
export function inferSchema(
data: Array<Record<string, unknown>>,
schema: Schema | undefined,
options: InferenceOptions,
): Schema {
return new SchemaInferrer(schema, options).infer(data);
}
class SchemaInferrer {
private readonly fields = new FieldTree();
constructor(
private readonly providedSchema: Schema | undefined,
private readonly options: InferenceOptions,
) {}
infer(data: Array<Record<string, unknown>>): Schema {
for (const [row, record] of data.entries()) {
for (const [path, value] of recordPathsAndValues(record)) {
this.observe(path, value, row);
}
}
return this.providedSchema === undefined
? new Schema(fieldsFromTree(this.fields))
: new Schema(matchingFields(this.providedSchema.fields, this.fields));
}
private observe(path: string[], value: unknown, row: number): void {
const current = this.fields.get(path);
if (current === undefined) {
this.addField(path, value, row);
} else if (this.providedSchema === undefined) {
this.updateInferredField(path, value, row, current);
}
}
private addField(path: string[], value: unknown, row: number): void {
if (this.providedSchema !== undefined) {
this.addSchemaField(this.providedSchema, path, row);
return;
}
const evidence =
this.inferType(value, path) ?? DeferredTypeEvidence.from(value, row);
if (evidence === undefined) {
throw typeInferenceError(path, row);
}
const conflict = this.fields.set(
path,
evidence,
(existing) =>
existing instanceof DeferredTypeEvidence && existing.isOnlyNulls(),
);
if (conflict !== undefined) {
throw branchConflictError(conflict, row, "Struct");
}
}
private addSchemaField(schema: Schema, path: string[], row: number): void {
const field = fieldAtPath(schema, path);
if (field === undefined) {
throw new Error(
`Found field not in schema: ${path.join(".")} at row ${row}`,
);
}
const conflict = this.fields.set(path, field.type);
if (conflict !== undefined) {
throw branchConflictError(conflict, row, "Struct");
}
}
private updateInferredField(
path: string[],
value: unknown,
row: number,
current: FieldNode,
): void {
const newType = this.inferType(value, path);
const deferred = DeferredTypeEvidence.from(value, row);
if (current instanceof FieldTree) {
if (deferred?.isOnlyNulls()) {
return;
}
throw schemaInferenceError(
path,
row,
"Struct",
describeEvidence(newType ?? deferred),
);
}
if (current instanceof DeferredTypeEvidence) {
this.resolveDeferredField(path, row, current, newType, deferred);
return;
}
if (newType !== undefined) {
if (!inferredTypesEqual(current, newType)) {
throw schemaInferenceError(
path,
row,
describeEvidence(current),
describeEvidence(newType),
);
}
return;
}
if (deferred === undefined || !deferred.matches(current)) {
throw schemaInferenceError(
path,
row,
describeEvidence(current),
describeEvidence(deferred),
);
}
}
private resolveDeferredField(
path: string[],
row: number,
current: DeferredTypeEvidence,
newType: DataType | undefined,
deferred: DeferredTypeEvidence | undefined,
): void {
if (newType !== undefined) {
if (!current.matches(newType)) {
throw schemaInferenceError(
path,
row,
current.describe(),
describeEvidence(newType),
);
}
this.fields.set(path, newType);
return;
}
if (deferred !== undefined) {
this.fields.set(path, current.merge(deferred));
return;
}
throw schemaInferenceError(
path,
row,
current.describe(),
describeEvidence(newType),
);
}
private inferType(value: unknown, path: string[]): DataType | undefined {
if (typeof value === "bigint") {
return new Int64();
}
if (typeof value === "number") {
return new Float64();
}
if (typeof value === "string") {
return this.options.dictionaryEncodeStrings
? new Dictionary(new Utf8(), new Int32())
: new Utf8();
}
if (typeof value === "boolean") {
return new Bool();
}
if (value instanceof Buffer) {
return new Binary();
}
if (ArrayBuffer.isView(value) && !(value instanceof DataView)) {
const typedArray = typedArrayToArrowType(value);
return typedArray === undefined
? undefined
: new FixedSizeList(
typedArray.length,
new Field("item", typedArray.elementType, true),
);
}
if (!Array.isArray(value) || value.length === 0) {
return undefined;
}
const configuredVector =
path.length === 1 ? this.options.vectorColumns[path[0]] : undefined;
if (configuredVector !== undefined) {
return new FixedSizeList(
value.length,
new Field("item", sanitizeType(configuredVector.type), true),
);
}
const itemType = this.inferArrayItemType(value, path);
if (itemType === undefined) {
return undefined;
}
return nameSuggestsVectorColumn(path[path.length - 1])
? new FixedSizeList(value.length, new Field("item", new Float32(), true))
: new List(new Field("item", itemType, true));
}
private inferArrayItemType(
values: unknown[],
path: string[],
): DataType | undefined {
let itemType: DataType | undefined;
const deferredItems: unknown[] = [];
for (const value of values) {
const candidate = this.inferType(value, path);
if (candidate === undefined) {
if (!isDeferredValue(value)) {
return undefined;
}
deferredItems.push(value);
} else if (itemType === undefined) {
itemType = candidate;
} else if (!inferredTypesEqual(itemType, candidate)) {
return undefined;
}
}
if (itemType === undefined) {
return undefined;
}
return deferredItems.every((value) =>
deferredValueMatchesType(value, itemType),
)
? itemType
: undefined;
}
}
/** Nulls and empty/all-null lists that do not determine a type by themselves. */
class DeferredTypeEvidence {
private constructor(
private readonly values: Array<{ value: unknown; row: number }>,
) {}
static from(value: unknown, row: number): DeferredTypeEvidence | undefined {
return isDeferredValue(value)
? new DeferredTypeEvidence([{ value, row }])
: undefined;
}
isOnlyNulls(): boolean {
return this.values.every(({ value }) => value == null);
}
matches(type: DataType): boolean {
return this.values.every(({ value }) =>
deferredValueMatchesType(value, type),
);
}
merge(other: DeferredTypeEvidence): DeferredTypeEvidence {
return new DeferredTypeEvidence([...this.values, ...other.values]);
}
describe(): string {
const list = this.values.find(({ value }) => Array.isArray(value));
return list === undefined
? "null"
: `List[${(list.value as unknown[]).length}]`;
}
firstRow(): number {
return this.values[0].row;
}
}
type FieldNode = DataType | DeferredTypeEvidence | FieldTree;
type LeafNode = Exclude<FieldNode, FieldTree>;
type FieldConflict = { path: string[]; value: FieldNode };
/** Nested field state, kept separate from Arrow's eventual Struct types. */
class FieldTree {
private readonly children = new Map<string, FieldNode>();
get(path: string[]): FieldNode | undefined {
let current: FieldNode = this;
for (const part of path) {
if (!(current instanceof FieldTree)) {
return undefined;
}
const child = current.children.get(part);
if (child === undefined) {
return undefined;
}
current = child;
}
return current;
}
set(
path: string[],
value: LeafNode,
canReplaceLeaf: (value: LeafNode) => boolean = () => false,
): FieldConflict | undefined {
let branch: FieldTree = this;
for (const [index, part] of path.slice(0, -1).entries()) {
const child = branch.children.get(part);
if (child === undefined || (isLeaf(child) && canReplaceLeaf(child))) {
const nextBranch = new FieldTree();
branch.children.set(part, nextBranch);
branch = nextBranch;
} else if (child instanceof FieldTree) {
branch = child;
} else {
return { path: path.slice(0, index + 1), value: child };
}
}
const name = path[path.length - 1];
const current = branch.children.get(name);
if (current instanceof FieldTree) {
return { path, value: current };
}
branch.children.set(name, value);
return undefined;
}
entries(): IterableIterator<[string, FieldNode]> {
return this.children.entries();
}
has(name: string): boolean {
return this.children.has(name);
}
}
function isLeaf(value: FieldNode): value is LeafNode {
return !(value instanceof FieldTree);
}
function fieldsFromTree(tree: FieldTree, path: string[] = []): Field[] {
const fields: Field[] = [];
for (const [name, value] of tree.entries()) {
if (value instanceof FieldTree) {
fields.push(
new Field(
name,
new Struct(fieldsFromTree(value, [...path, name])),
true,
),
);
} else if (value instanceof DeferredTypeEvidence) {
throw typeInferenceError([...path, name], value.firstRow());
} else {
fields.push(new Field(name, value, true));
}
}
return fields;
}
function matchingFields(fields: Field[], tree: FieldTree): Field[] {
const matches: Field[] = [];
for (const field of fields) {
if (!tree.has(field.name)) {
continue;
}
const value = tree.get([field.name]);
if (value instanceof FieldTree) {
const struct = field.type as Struct;
matches.push(
new Field(
field.name,
new Struct(matchingFields(struct.children, value)),
field.nullable,
field.metadata,
),
);
} else {
matches.push(field);
}
}
return matches;
}
function* recordPathsAndValues(
record: Record<string, unknown>,
path: string[] = [],
): Generator<[string[], unknown]> {
for (const [name, value] of Object.entries(record)) {
if (isRecord(value)) {
yield* recordPathsAndValues(value, [...path, name]);
} else if (value !== undefined) {
yield [[...path, name], value];
}
}
}
function isRecord(value: unknown): value is Record<string, unknown> {
return (
typeof value === "object" &&
value !== null &&
!Array.isArray(value) &&
!(value instanceof RegExp) &&
!(value instanceof Date) &&
!(value instanceof Set) &&
!(value instanceof Map) &&
!(value instanceof Buffer) &&
!ArrayBuffer.isView(value)
);
}
function fieldAtPath(schema: Schema, path: string[]): Field | undefined {
let fields = schema.fields;
let field: Field | undefined;
for (const [index, name] of path.entries()) {
field = fields.find((candidate) => candidate.name === name);
if (field === undefined || index === path.length - 1) {
return field;
}
if (!DataType.isStruct(field.type)) {
return undefined;
}
fields = field.type.children;
}
return field;
}
function isDeferredValue(value: unknown): boolean {
return (
value == null || (Array.isArray(value) && value.every(isDeferredValue))
);
}
function deferredValueMatchesType(value: unknown, type: DataType): boolean {
if (value == null) {
return true;
}
if (!Array.isArray(value)) {
return false;
}
if (DataType.isList(type)) {
return value.every((item) =>
deferredValueMatchesType(item, type.valueType),
);
}
if (DataType.isFixedSizeList(type)) {
return (
value.length === type.listSize &&
value.every((item) => deferredValueMatchesType(item, type.valueType))
);
}
return false;
}
function inferredTypesEqual(current: DataType, candidate: DataType): boolean {
if (DataType.isDictionary(current)) {
return (
DataType.isDictionary(candidate) &&
current.isOrdered === candidate.isOrdered &&
inferredTypesEqual(current.indices, candidate.indices) &&
inferredTypesEqual(current.dictionary, candidate.dictionary)
);
}
if (DataType.isList(current)) {
return (
DataType.isList(candidate) &&
current.valueField.name === candidate.valueField.name &&
current.valueField.nullable === candidate.valueField.nullable &&
inferredTypesEqual(current.valueType, candidate.valueType)
);
}
if (DataType.isFixedSizeList(current)) {
return (
DataType.isFixedSizeList(candidate) &&
current.listSize === candidate.listSize &&
current.valueField.name === candidate.valueField.name &&
current.valueField.nullable === candidate.valueField.nullable &&
inferredTypesEqual(current.valueType, candidate.valueType)
);
}
return arrowUtil.compareTypes(current, candidate);
}
function describeEvidence(
evidence: DataType | DeferredTypeEvidence | undefined,
): string {
if (evidence === undefined) {
return "an unsupported value";
}
return evidence instanceof DeferredTypeEvidence
? evidence.describe()
: evidence.toString();
}
function branchConflictError(
conflict: FieldConflict,
row: number,
candidate: string,
): Error {
return schemaInferenceError(
conflict.path,
row,
conflict.value instanceof FieldTree
? "Struct"
: describeEvidence(conflict.value),
candidate,
);
}
function schemaInferenceError(
path: string[],
row: number,
currentType: string,
newType: string,
): Error {
return new Error(
`Failed to infer schema for data. Previously inferred type ${currentType} ` +
`but found ${newType} for field ${path.join(".")} at row ${row}. ` +
"Consider providing an explicit schema.",
);
}
function typeInferenceError(path: string[], row: number): Error {
return new Error(
`Failed to infer data type for field ${path.join(".")} at row ${row}. ` +
"Consider providing an explicit schema.",
);
}
function nameSuggestsVectorColumn(name: string): boolean {
const normalized = name.toLowerCase();
return normalized.includes("vector") || normalized.includes("embedding");
}
+45 -15
View File
@@ -43,10 +43,12 @@ import {
Table as _NativeTable,
} from "./native";
import {
AutoQuery,
FullTextQuery,
Query,
TakeQuery,
VectorQuery,
createAutoQuery,
instanceOfFullTextQuery,
} from "./query";
import { sanitizeType } from "./sanitize";
@@ -523,7 +525,7 @@ export abstract class Table {
query: string | IntoVector | MultiVector | FullTextQuery,
queryType?: string,
ftsColumns?: string | string[],
): VectorQuery | Query;
): VectorQuery | Query | AutoQuery;
/**
* Search the table with a given query vector.
*
@@ -628,6 +630,18 @@ export abstract class Table {
/**
* Update per-field (column) metadata.
*
* The following keys are treated specially, by convention, and should be
* used when appropriate:
*
* - `lancedb:description`: for a human-readable description of a field.
* - `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
* names the tag category; e.g. `lancedb:tag:model: "clip"`.
* - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
* `feature_v2` might be in the same logical column.
* - `lancedb:status`: for status options (`production`, `candidate`,
* `deprecated`, `archived`) to designate the current life cycle state of
* this column.
* @param {FieldMetadataUpdate[]} updates One or more per-field updates. Each
* update's metadata is merged into the field's existing metadata by default;
* a value of `null` deletes that key, and `replace: true` swaps the whole map.
@@ -975,10 +989,11 @@ export class LocalTable extends Table {
return this.inner.display();
}
private async getEmbeddingFunctions(): Promise<
Map<string, EmbeddingFunctionConfig>
> {
const schema = await this.schema();
private async getEmbeddingFunctions(
inner: _NativeTable = this.inner,
): Promise<Map<string, EmbeddingFunctionConfig>> {
const schemaBuf = await inner.schema();
const schema = tableFromIPC(schemaBuf).schema;
const registry = getRegistry();
return registry.parseFunctions(schema.metadata);
}
@@ -1160,7 +1175,7 @@ export class LocalTable extends Table {
query: string | IntoVector | MultiVector | FullTextQuery,
queryType: string = "auto",
ftsColumns?: string | string[],
): VectorQuery | Query {
): VectorQuery | Query | AutoQuery {
if (typeof query !== "string" && !instanceOfFullTextQuery(query)) {
if (queryType === "fts") {
throw new Error("Cannot perform full text search on a vector query");
@@ -1175,14 +1190,28 @@ export class LocalTable extends Table {
});
}
// The query type is auto or vector
// fall back to full text search if no embedding functions are defined and the query is a string
if (
queryType === "auto" &&
(getRegistry().length() === 0 || instanceOfFullTextQuery(query))
) {
return this.query().fullTextSearch(query, {
columns: ftsColumns,
if (queryType === "auto") {
if (instanceOfFullTextQuery(query)) {
return this.query().fullTextSearch(query, {
columns: ftsColumns,
});
}
const columns =
typeof ftsColumns === "string" ? [ftsColumns] : (ftsColumns ?? null);
return createAutoQuery(this.inner, query, columns, async (metadata) => {
const functions = await getRegistry().parseFunctions(
new Map([["embedding_functions", metadata]]),
);
// TODO: Support multiple embedding functions
const embeddingFunc: EmbeddingFunctionConfig | undefined = functions
.values()
.next().value;
// The route only calls this callback when embedding metadata exists.
// parseFunctions either yields a provider or reports malformed metadata.
if (!embeddingFunc)
throw new Error("Invalid embedding function metadata");
return await embeddingFunc.function.computeQueryEmbeddings(query);
});
}
@@ -1538,7 +1567,8 @@ export interface FieldMetadataUpdate {
path: string;
/**
* Metadata key/value pairs. Merged into the field's existing metadata by
* default; a value of `null` deletes that key.
* default; a value of `null` deletes that key. See
* {@link Table.updateFieldMetadata} for the conventional `lancedb:*` keys.
*/
metadata: Record<string, string | null>;
/** If true, replace the field's entire metadata map instead of merging. */
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-darwin-arm64",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"os": ["darwin"],
"cpu": ["arm64"],
"main": "lancedb.darwin-arm64.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-arm64-gnu",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"os": ["linux"],
"cpu": ["arm64"],
"main": "lancedb.linux-arm64-gnu.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-arm64-musl",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"os": ["linux"],
"cpu": ["arm64"],
"main": "lancedb.linux-arm64-musl.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-x64-gnu",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"os": ["linux"],
"cpu": ["x64"],
"main": "lancedb.linux-x64-gnu.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-x64-musl",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"os": ["linux"],
"cpu": ["x64"],
"main": "lancedb.linux-x64-musl.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-win32-arm64-msvc",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"os": [
"win32"
],
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-win32-x64-msvc",
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"os": ["win32"],
"cpu": ["x64"],
"main": "lancedb.win32-x64-msvc.node",
+2 -2
View File
@@ -1,12 +1,12 @@
{
"name": "@lancedb/lancedb",
"version": "0.38.0-beta.4",
"version": "0.38.0-beta.11",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "@lancedb/lancedb",
"version": "0.38.0-beta.4",
"version": "0.38.0-beta.11",
"cpu": [
"x64",
"arm64"
+1 -1
View File
@@ -11,7 +11,7 @@
"ann"
],
"private": false,
"version": "0.38.0-beta.5",
"version": "0.38.0-beta.11",
"main": "dist/index.js",
"exports": {
".": "./dist/index.js",
+34
View File
@@ -17,6 +17,7 @@ use lancedb::connection::{ConnectBuilder, Connection as LanceDBConnection, conne
use lance_namespace::models::{
CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest,
ListTablesRequest,
};
use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema};
@@ -36,6 +37,12 @@ pub struct ListNamespacesResponse {
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>>,
@@ -206,6 +213,33 @@ impl Connection {
op.execute().await.default_error()
}
/// List a page of tables in the database.
#[napi(catch_unwind)]
pub async 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,
})
}
/// Create table from a Apache Arrow IPC (file) buffer.
///
/// Parameters:
+5 -2
View File
@@ -14,9 +14,12 @@ pub struct Job {
}
impl Job {
pub(crate) fn new(inner: lancedb::Job) -> Self {
pub(crate) fn new<T>(inner: lancedb::Job<T>) -> Self
where
T: Clone + Send + Sync + 'static,
{
Self {
inner: Arc::new(inner),
inner: Arc::new(inner.map(|_| ())),
}
}
}
+13
View File
@@ -278,6 +278,13 @@ impl Table {
Ok(Query::new(self.inner_ref()?.query()))
}
/// Return a read-only table handle pinned to the current query revision.
#[napi(catch_unwind)]
pub async fn query_snapshot(&self) -> napi::Result<Self> {
let snapshot = self.inner_ref()?.query_snapshot().await.default_error()?;
Ok(Self::new(snapshot))
}
#[napi(catch_unwind)]
pub fn take_offsets(&self, offsets: Vec<i64>) -> napi::Result<TakeQuery> {
Ok(TakeQuery::new(
@@ -554,6 +561,12 @@ impl Table {
.default_error()
}
#[napi(catch_unwind)]
pub async fn checkout_current(&self) -> napi::Result<Self> {
let table = self.inner_ref()?.checkout_current().await.default_error()?;
Ok(Self::new(table))
}
#[napi(catch_unwind)]
pub async fn checkout(&self, version: i64) -> napi::Result<()> {
self.inner_ref()?
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.38.0-beta.5"
version = "0.38.0-beta.11"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
+4 -1
View File
@@ -101,9 +101,12 @@ azure = ["adlfs>=2024.2.0"]
[tool.maturin]
python-source = "python"
module-name = "lancedb._lancedb"
# uv installs the project as an editable package before `uv run`, so keep that
# bootstrap build consistent with `maturin develop`.
editable-profile = "dev"
[build-system]
requires = ["maturin>=1.9.4"]
requires = ["maturin>=1.10"]
build-backend = "maturin"
[tool.ruff.lint]
+22
View File
@@ -29,6 +29,7 @@ from .functions import (
FunctionRegistrationRequest as FunctionRegistrationRequest,
FunctionVersion as FunctionVersion,
PythonRuntimeSpec as PythonRuntimeSpec,
RefreshColumnResult as RefreshColumnResult,
UdfDefinition as UdfDefinition,
udf as udf,
)
@@ -178,6 +179,18 @@ def connect(
... },
... )
For Azure Blob Storage, credentials can be passed directly without setting
environment variables:
>>> azure_storage_options = {
... "account_name": "some-account",
... "account_key": "some-key",
... }
>>> db = lancedb.connect( # doctest: +SKIP
... "az://my-container/my-database",
... storage_options=azure_storage_options,
... )
For tests and temporary data, use an in-memory database:
>>> db = lancedb.connect("memory://")
@@ -464,6 +477,10 @@ async def connect_async(
--------
>>> import lancedb
>>> azure_storage_options = {
... "account_name": "some-account",
... "account_key": "some-key",
... }
>>> async def doctest_example():
... # For a local directory, provide a path to the database
... db = await lancedb.connect_async("~/.lancedb")
@@ -471,6 +488,11 @@ async def connect_async(
... db = await lancedb.connect_async("s3://my-bucket/lancedb",
... storage_options={
... "aws_access_key_id": "***"})
... # Azure credentials can also be passed directly
... db = await lancedb.connect_async(
... "az://my-container/my-database",
... storage_options=azure_storage_options,
... )
... # For tests and temporary data, use an in-memory database
... db = await lancedb.connect_async("memory://")
... # Connect to LanceDB cloud
+4 -2
View File
@@ -270,7 +270,8 @@ def _iter_projection_pairs(
if isinstance(expr, str):
yield name, expr
elif isinstance(expr, Expr):
yield name, expr.to_sql()
source = expr._column_name()
yield name, source if source is not None else expr.to_sql()
return
for column in projection:
if isinstance(column, str):
@@ -280,7 +281,8 @@ def _iter_projection_pairs(
if isinstance(expr, str):
yield name, expr
elif isinstance(expr, Expr):
yield name, expr.to_sql()
source = expr._column_name()
yield name, source if source is not None else expr.to_sql()
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
+5 -9
View File
@@ -87,6 +87,7 @@ class PyExpr:
def contains(self, substr: "PyExpr") -> "PyExpr": ...
def isin(self, values: List["PyExpr"]) -> "PyExpr": ...
def cast(self, data_type: pa.DataType) -> "PyExpr": ...
def column_name(self) -> Optional[str]: ...
def to_sql(self) -> str: ...
def expr_col(name: str) -> PyExpr: ...
@@ -147,7 +148,7 @@ class Connection(object):
limit: Optional[int],
) -> list[str]: ... # Deprecated: Use list_tables instead
def job(self, job_id: str) -> Job: ...
async def create_function_async(self, request_json: str) -> FunctionJob: ...
async def create_function_async(self, request_json: str) -> Job: ...
async def get_function(self, name: str, version: str) -> str: ...
async def list_jobs(self) -> List[JobInfo]: ...
async def get_job(self, job_id: str) -> Optional[JobDescription]: ...
@@ -234,14 +235,7 @@ class Job:
@property
def id(self) -> Optional[str]: ...
async def status(self) -> str: ...
async def wait(self) -> None: ...
async def cancel(self) -> None: ...
class FunctionJob:
@property
def id(self) -> Optional[str]: ...
async def status(self) -> str: ...
async def wait(self) -> str: ...
async def wait(self) -> Optional[str]: ...
async def cancel(self) -> None: ...
class JobInfo:
@@ -290,6 +284,7 @@ class Table:
mode: Literal["append", "overwrite"],
progress: Optional[Any] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult: ...
async def update(
self, updates: Dict[str, str], where: Optional[str]
@@ -614,6 +609,7 @@ class PyQueryRequest:
filter: Optional[Union[str, bytes]]
full_text_search: Optional[FullTextQuery]
select: Optional[Union[str, List[str]]]
select_source_columns: Optional[Dict[str, str]]
fast_search: Optional[bool]
with_row_id: Optional[bool]
use_lsm: Optional[bool]
+2 -2
View File
@@ -46,7 +46,7 @@ from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
from . import __version__
from ._lancedb import connect as lancedb_connect # type: ignore
from .functions import FunctionVersion, UdfDefinition
from .job import AsyncJob, Job, _function_job
from .job import AsyncJob, Job, _typed_job
from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
@@ -2237,7 +2237,7 @@ class AsyncConnection(object):
inner = await self._inner.create_function_async(
definition.registration_request.to_canonical_json()
)
return _function_job(inner)
return _typed_job(inner, FunctionVersion.from_json)
async def get_function(self, name: str, *, version: str) -> FunctionVersion:
"""Open one exact immutable Function version from the remote catalog."""
+5 -1
View File
@@ -249,6 +249,10 @@ class Expr:
# ── utilities ────────────────────────────────────────────────────────────
def _column_name(self) -> str | None:
"""Return the source name when this is a bare column expression."""
return self._inner.column_name()
def to_sql(self) -> str:
"""Render the expression as a SQL string (useful for debugging)."""
return self._inner.to_sql()
@@ -312,7 +316,7 @@ def func(name: str, *args: ExprLike) -> Expr:
--------
>>> from lancedb.expr import col, func
>>> func("lower", col("name"))
Expr(lower(name))
Expr(lower(`name`))
"""
inner_args = [_coerce(a)._inner for a in args]
return Expr(expr_func(name, inner_args))
+231 -104
View File
@@ -1,26 +1,30 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""Canonical values exchanged with LanceDB Enterprise Function services.
"""Canonical Function values exchanged with LanceDB Enterprise services.
These immutable models contain client/wire state only. Catalog persistence,
environment bake, secret resolution, and execution are owned by Sophon.
environment bake, and execution are owned by Sophon.
``RefreshColumnResult`` is also the backend-neutral result of a local
expression-backed refresh job.
"""
from __future__ import annotations
import ast
import builtins
import base64
import functools
import hashlib
import importlib
import inspect
import symtable
import json
import math
import re
import sys
import textwrap
import types
import uuid
from collections.abc import Mapping
from datetime import date, datetime
from typing import (
@@ -218,13 +222,14 @@ class PythonEnvironmentSpec(_RemoteValue):
kind: str
packages: tuple[str, ...] = ()
channels: tuple[str, ...] = ()
path: Optional[str] = None
modules: tuple[str, ...] = ()
image: Optional[str] = None
class PythonRuntimeSpec(_RemoteValue):
"""Remote runtime definition with non-secret environment values.
"""Remote runtime definition with environment values.
V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
their unknown payload fields are intentionally not retained by the client.
@@ -263,7 +268,6 @@ class FunctionVersion(_RemoteValue):
runtime: PythonRuntimeSpec
runtime_digest: str
environment_digest: str
required_secrets: tuple[str, ...] = ()
created_at: str
def __call__(self, **inputs: Any) -> FunctionApplication:
@@ -271,7 +275,7 @@ class FunctionVersion(_RemoteValue):
Every input must be a direct [lancedb.col][lancedb.expr.col]
reference. The returned application is immutable and retains a
named-struct output as one sibling group, so every row's sibling values
named-struct output as one binding, so every row's sibling values
come from one logical Function evaluation. Map result fields to table
columns with
[FunctionApplication.rename][lancedb.functions.FunctionApplication.rename],
@@ -321,22 +325,16 @@ class FunctionVersion(_RemoteValue):
function=FunctionVersionRef(name=self.name, version=self.version),
inputs=tuple(bindings),
output=self.signature.output,
group_id=f"fg_{uuid.uuid4().hex}",
)
class FunctionRegistrationRequest(_RemoteValue):
"""Stable remote registration envelope produced by :func:`udf`.
Only secret names are represented. Secret values are resolved inside the
remote service and have no client request field.
"""
"""Stable remote registration envelope produced by :func:`udf`."""
name: str
artifact: FunctionArtifactRequest
signature: FunctionSignature
runtime: PythonRuntimeSpec
required_secrets: tuple[str, ...] = ()
class FunctionVersionRef(_OpenRemoteValue):
@@ -365,7 +363,7 @@ class ApplicationInput(_OpenRemoteValue):
class FunctionApplication(_OpenRemoteValue):
"""Immutable pre-declaration application of an exact Function version.
A named-struct output remains one grouped application through table
A named-struct output remains one application through table
declaration and execution.
[FunctionApplication.rename][lancedb.functions.FunctionApplication.rename]
records the result-field to table-column mapping without splitting sibling
@@ -375,7 +373,6 @@ class FunctionApplication(_OpenRemoteValue):
function: FunctionVersionRef
inputs: tuple[ApplicationInput, ...]
output: FunctionOutput
group_id: str
columns: Mapping[str, str] = Field(default_factory=dict)
def _known_dict(self) -> dict[str, Any]:
@@ -447,12 +444,10 @@ class OutputMapping(_RemoteValue):
class FunctionBinding(_RemoteValue):
"""Immutable grouped binding persisted by the Enterprise table service."""
"""Immutable Function binding persisted by the Enterprise table service."""
binding_id: str
revision: _UInt64
function: FunctionVersionRef
group_id: str
inputs: tuple[InputBinding, ...]
outputs: tuple[OutputMapping, ...]
input_schema: Optional[Mapping[str, Any]] = None
@@ -460,7 +455,11 @@ class FunctionBinding(_RemoteValue):
class RefreshColumnResult(_RemoteValue):
"""Terminal result of a remote Function-column refresh Job."""
"""Terminal result of an expression-backed or Function-backed refresh Job.
Local jobs produce this value in process. LanceDB Cloud and Enterprise
decode the same value from the durable server-job terminal payload.
"""
rows_assigned: _UInt64
rows_failed: _UInt64
@@ -480,62 +479,60 @@ class RefreshColumnResult(_RemoteValue):
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
_SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
_GRAMMAR_PRIMITIVES = (
(pa.bool_(), "bool"),
(pa.int8(), "int8"),
(pa.int16(), "int16"),
(pa.int32(), "int32"),
(pa.int64(), "int64"),
(pa.uint8(), "uint8"),
(pa.uint16(), "uint16"),
(pa.uint32(), "uint32"),
(pa.uint64(), "uint64"),
(pa.float16(), "float16"),
(pa.float32(), "float32"),
(pa.float64(), "float64"),
(pa.string(), "utf8"),
(pa.binary(), "binary"),
(pa.date32(), "date32"),
(pa.date64(), "date64"),
)
def _canonical_arrow_type(data_type: pa.DataType) -> str:
primitive_types = (
(pa.bool_(), "bool"),
(pa.int8(), "int8"),
(pa.int16(), "int16"),
(pa.int32(), "int32"),
(pa.int64(), "int64"),
(pa.uint8(), "uint8"),
(pa.uint16(), "uint16"),
(pa.uint32(), "uint32"),
(pa.uint64(), "uint64"),
(pa.float16(), "float16"),
(pa.float32(), "float32"),
(pa.float64(), "float64"),
(pa.string(), "utf8"),
(pa.large_utf8(), "large_utf8"),
(pa.binary(), "binary"),
(pa.large_binary(), "large_binary"),
(pa.date32(), "date32"),
(pa.date64(), "date64"),
)
for candidate, name in primitive_types:
"""The server's V1 Function type grammar. Anything outside it is rejected
here rather than at registration."""
for candidate, name in _GRAMMAR_PRIMITIVES:
if data_type == candidate:
return name
if pa.types.is_fixed_size_binary(data_type):
return f"fixed_size_binary[{data_type.byte_width}]"
if pa.types.is_list(data_type):
return f"list<{_canonical_arrow_type(data_type.value_type)}>"
if pa.types.is_large_list(data_type):
return f"large_list<{_canonical_arrow_type(data_type.value_type)}>"
if pa.types.is_fixed_size_list(data_type):
if pa.types.is_list(data_type) or pa.types.is_large_list(data_type):
prefix = "list" if pa.types.is_list(data_type) else "large_list"
return f"{prefix}<{_canonical_list_item(data_type)}>"
if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0:
return (
f"fixed_size_list<{_canonical_arrow_type(data_type.value_type)}>"
f"[{data_type.list_size}]"
f"fixed_size_list<{_canonical_list_item(data_type)}, {data_type.list_size}>"
)
if pa.types.is_struct(data_type):
fields = ",".join(
f"{field.name}:{_canonical_arrow_type(field.type)}" for field in data_type
)
return f"struct<{fields}>"
if pa.types.is_timestamp(data_type):
timezone = f",tz={data_type.tz}" if data_type.tz is not None else ""
return f"timestamp[{data_type.unit}{timezone}]"
if pa.types.is_time32(data_type) or pa.types.is_time64(data_type):
return f"time[{data_type.unit}]"
if pa.types.is_duration(data_type):
return f"duration[{data_type.unit}]"
if pa.types.is_decimal(data_type):
bit_width = data_type.bit_width
return f"decimal{bit_width}({data_type.precision},{data_type.scale})"
raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
def _canonical_list_item(data_type: pa.DataType) -> str:
"""The grammar names only the item type; it always means a non-nullable
child called `item`, so any other child metadata cannot be represented."""
child = data_type.value_field
if child.name != "item" or child.nullable or child.metadata:
raise TypeError(
"unsupported Arrow type for Function signature: list items must be a "
f"non-nullable field named 'item', got {child}"
)
return _canonical_arrow_type(child.type)
def _list_of(item: pa.DataType) -> pa.DataType:
return pa.list_(pa.field("item", item, nullable=False))
def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
nullable = False
origin = get_origin(annotation)
@@ -583,7 +580,7 @@ def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
value_type, value_nullable = _annotation_type(arguments[0])
if value_nullable:
raise TypeError("nullable Function list elements are not supported")
return pa.list_(value_type), nullable
return _list_of(value_type), nullable
raise TypeError(f"unsupported Function annotation: {annotation!r}")
@@ -730,6 +727,104 @@ def _literal_source(value: Any) -> str:
)
_DYNAMIC_NAMESPACE_ACCESS = frozenset(
{"globals", "locals", "vars", "eval", "exec", "compile", "__import__"}
)
# Modules that hand out namespaces (`sys.modules`, `builtins`, importers,
# introspection). The artifact's module namespace holds only the names it was
# packaged with, so reaching around it cannot be represented.
_NAMESPACE_MODULES = frozenset(
{"sys", "builtins", "importlib", "inspect", "gc", "ctypes", "types"}
)
def _namespace_acquisition(
definition: ast.FunctionDef, references: set[str]
) -> list[str]:
found = set(references & _DYNAMIC_NAMESPACE_ACCESS)
for node in ast.walk(definition):
if isinstance(node, ast.Import):
found.update(
alias.name
for alias in node.names
if alias.name.split(".")[0] in _NAMESPACE_MODULES
)
elif isinstance(node, ast.ImportFrom) and node.module:
if node.module.split(".")[0] in _NAMESPACE_MODULES:
found.add(node.module)
return sorted(found)
def _module_references(module_source: str) -> set[str]:
"""Names any scope in `module_source` binds or loads at module scope.
Python's own scope analysis on the exact text that ships: free variables
belong to an enclosing scope inside the function, and postponed
annotations are not runtime loads."""
def visit(table: symtable.SymbolTable, found: set[str]) -> None:
for symbol in table.get_symbols():
if symbol.is_global() and (
symbol.is_referenced() or symbol.is_declared_global()
):
found.add(symbol.get_name())
for child in table.get_children():
visit(child, found)
found: set[str] = set()
for table in symtable.symtable(module_source, "<udf>", "exec").get_children():
visit(table, found)
return found
def _global_source(name: str, value: Any) -> str:
"""One module-level line that rebinds `name` to `value` in the artifact:
an import for modules and importable classes/functions, a literal otherwise."""
if isinstance(value, types.ModuleType):
if value.__name__.split(".")[0] in _NAMESPACE_MODULES:
raise ValueError(
f"@udf cannot package dynamic namespace access: {value.__name__!r}"
)
try:
imported = importlib.import_module(value.__name__)
except ImportError:
imported = None
if imported is not value:
raise TypeError(
f"Function source references module {name!r} that does not import "
f"as {value.__name__!r}"
)
return f"import {value.__name__} as {name}"
module_name = getattr(value, "__module__", None)
qualname = getattr(value, "__qualname__", None)
if (
isinstance(module_name, str)
and isinstance(qualname, str)
and module_name != "__main__"
and "." not in qualname
and "<" not in qualname
):
try:
imported = getattr(importlib.import_module(module_name), qualname)
except (ImportError, AttributeError):
imported = None
if imported is value:
return f"from {module_name} import {qualname} as {name}"
return f"{name} = {_literal_source(value)}"
def _is_recursive_reference(function: Callable[..., Any], name: str) -> bool:
"""`name` inside the body means the function itself unless the module has
since bound it to something else."""
if name != function.__name__:
return False
bound = function.__globals__.get(name, function)
if bound is function:
return True
# The decorator's own result is the one wrapper known to call `function`
# unchanged; any other binding may behave differently from a self-call.
return type(bound) is UdfDefinition and bound._function is function
def _package_source(function: Callable[..., Any]) -> bytes:
if not inspect.isfunction(function) or inspect.iscoroutinefunction(function):
raise TypeError("@udf requires a synchronous Python function")
@@ -754,23 +849,46 @@ def _package_source(function: Callable[..., Any]) -> bytes:
closure = inspect.getclosurevars(function)
if closure.nonlocals:
raise ValueError("@udf cannot package functions that capture closure values")
if closure.unbound:
raise ValueError(
f"@udf source contains unresolved global names: {sorted(closure.unbound)!r}"
)
globals_source = []
for name, value in sorted(closure.globals.items()):
if isinstance(value, types.ModuleType):
globals_source.append(f"import {value.__name__} as {name}")
else:
globals_source.append(f"{name} = {_literal_source(value)}")
function_source = ast.unparse(definition)
parts = ["from __future__ import annotations"]
module_header = "from __future__ import annotations"
references = _module_references(f"{module_header}\n\n{function_source}\n")
dynamic = _namespace_acquisition(definition, references)
if dynamic:
raise ValueError(f"@udf cannot package dynamic namespace access: {dynamic!r}")
# Resolve every module-scope reference the way the interpreter would: the
# function's own globals first (a module global may shadow a builtin, and
# nested scopes are not visible to getclosurevars), then its builtins.
# The artifact runs under the standard builtins; only the exact mapping is
# provably equivalent (a subclass or copy can change lookups and hooks).
if function.__builtins__ is not vars(builtins):
raise ValueError("@udf cannot package a non-standard builtins environment")
globals_source = []
unresolved = []
for name in sorted(references):
if name == function.__name__:
if not _is_recursive_reference(function, name):
raise ValueError(
f"@udf cannot package {name!r}: the module binds that name to "
"another value, which the artifact's own definition would shadow"
)
continue
if name in function.__globals__:
globals_source.append(_global_source(name, function.__globals__[name]))
elif hasattr(builtins, name):
pass
else:
unresolved.append(name)
if unresolved:
raise ValueError(
f"@udf source contains unresolved global names: {unresolved!r}"
)
parts = [module_header]
if globals_source:
parts.extend(["", *globals_source])
parts.extend(["", function_source, ""])
return "\n".join(parts).encode("utf-8")
packaged = "\n".join(parts)
return packaged.encode("utf-8")
class UdfDefinition:
@@ -791,33 +909,32 @@ class UdfDefinition:
output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
pip: tuple[str, ...],
env: Mapping[str, str],
secrets: tuple[str, ...],
python_version: Optional[str],
conda: tuple[str, ...] = (),
conda_channels: tuple[str, ...] = (),
):
function_name = name or function.__name__
if not _FUNCTION_NAME.fullmatch(function_name):
raise ValueError(f"invalid Function name: {function_name!r}")
packages = tuple(sorted(set(pip)))
if pip and conda:
raise ValueError("a Function environment is pip or conda, not both")
if conda_channels and not conda:
raise ValueError("conda_channels requires conda packages")
packages = tuple(sorted(set(conda if conda else pip)))
if any(not package or package != package.strip() for package in packages):
raise ValueError("pip requirements must be non-empty and trimmed")
raise ValueError("package requirements must be non-empty and trimmed")
if conda:
environment_spec = PythonEnvironmentSpec(
kind="conda", packages=packages, channels=tuple(conda_channels)
)
else:
environment_spec = PythonEnvironmentSpec(kind="pip", packages=packages)
environment = dict(env)
if any(
not isinstance(key, str) or not isinstance(value, str)
for key, value in environment.items()
):
raise TypeError("Function env keys and values must be strings")
required_secrets = tuple(sorted(set(secrets)))
invalid_secrets = [
secret for secret in required_secrets if not _SECRET_NAME.fullmatch(secret)
]
if invalid_secrets:
raise ValueError(f"invalid Function secret names: {invalid_secrets!r}")
overlap = set(environment) & set(required_secrets)
if overlap:
raise ValueError(
f"Function env and secret names must be disjoint: {sorted(overlap)!r}"
)
signature = _infer_signature(function, input_schema, output_schema)
source = _package_source(function)
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
@@ -825,7 +942,7 @@ class UdfDefinition:
kind="python",
python_version=python_version
or f"{sys.version_info.major}.{sys.version_info.minor}",
environment=PythonEnvironmentSpec(kind="pip", packages=packages),
environment=environment_spec,
env=environment,
)
self._function = function
@@ -846,7 +963,6 @@ class UdfDefinition:
),
signature=signature,
runtime=runtime,
required_secrets=required_secrets,
)
functools.update_wrapper(self, function)
@@ -872,8 +988,9 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
@@ -885,8 +1002,9 @@ def udf(
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
pip: tuple[str, ...] | list[str] = (),
env: Optional[Mapping[str, str]] = None,
secrets: tuple[str, ...] | list[str] = (),
python_version: Optional[str] = None,
conda: tuple[str, ...] | list[str] = (),
conda_channels: tuple[str, ...] | list[str] = (),
):
"""Prepare a scalar Python callable for remote Function registration.
@@ -909,14 +1027,22 @@ def udf(
provided together with ``input_schema``.
pip : sequence of str, optional
Pip requirements for the remote environment.
conda : sequence of str, optional
Conda packages for the remote environment, instead of ``pip``.
conda_channels : sequence of str, optional
Conda channels in priority order; requires ``conda``.
env : mapping of str to str, optional
Non-secret environment variables. Use ``secrets`` for credentials.
secrets : sequence of str, optional
Names of secrets resolved by the remote service. Secret values are not
accepted by this API or included in the registration request.
Environment variables included in the Function definition.
python_version : str, optional
Remote Python major/minor version. Defaults to the client version.
The packaged artifact is a snapshot: the function source plus exactly
the module-level names it references (modules as imports, importable
classes and functions as imports, literals inline). Code that reaches the
module namespace another way -- ``globals()``/``eval``, ``sys.modules``,
``builtins`` -- is rejected where it can be seen and otherwise
unsupported; closures and a non-standard ``__builtins__`` are rejected.
Returns
-------
UdfDefinition
@@ -928,7 +1054,7 @@ def udf(
Examples
--------
>>> from lancedb import udf
>>> @udf(pip=["numpy==2.2.0"], secrets=["MODEL_TOKEN"])
>>> @udf(pip=["numpy==2.2.0"])
... def score(value: float) -> float:
... return value * 2
>>> score(1.5)
@@ -943,8 +1069,9 @@ def udf(
output_schema=output_schema,
pip=tuple(pip),
env={} if env is None else env,
secrets=tuple(secrets),
python_version=python_version,
conda=tuple(conda),
conda_channels=tuple(conda_channels),
)
if function is None:
+12
View File
@@ -7,6 +7,7 @@ from typing import List, Literal, Optional
from ._lancedb import (
IndexConfig,
)
from .query import DocumentGranularity
from .types import BaseTokenizerType
lang_mapping = {
@@ -121,6 +122,11 @@ class FTS:
>>> config = FTS(block_size=256)
Create an index that treats each deepest-list element as one document:
>>> from lancedb.query import DocumentGranularity
>>> config = FTS(document_granularity=DocumentGranularity.LIST_ELEMENT)
Attributes
----------
with_position : bool, default False
@@ -172,6 +178,11 @@ class FTS:
roughly half of the available CPU cores. The effective value is
limited by the available compute capacity. This build-only setting is
not persisted with the index and does not apply to remote tables.
document_granularity : DocumentGranularity, default ROW
``ROW`` treats the selected text in one table row as one document.
``LIST_ELEMENT`` treats each element of the deepest list on the indexed
field path as one document and returns its physical coordinates in
``_doc_index`` for matching queries.
Notes
-----
@@ -196,6 +207,7 @@ class FTS:
custom_stop_words: Optional[List[str]] = None
memory_limit: Optional[int] = None
num_workers: Optional[int] = None
document_granularity: DocumentGranularity = DocumentGranularity.ROW
@dataclass
+28 -30
View File
@@ -5,12 +5,11 @@
import asyncio
from datetime import timedelta
from typing import Any, Generic, Optional, TypeVar, cast
from typing import Any, Callable, Generic, Optional, TypeVar, cast
from lancedb.background_loop import LOOP
from . import _lancedb
from .functions import FunctionVersion
T = TypeVar("T")
@@ -18,11 +17,18 @@ T = TypeVar("T")
class AsyncJob(Generic[T]):
"""A handle to an operation that may still be running.
The operation may already be complete when the handle is created.
The operation may already be complete when the handle is created. ``T``
is the endpoint's terminal result type; unit-result jobs resolve to
``None``.
"""
def __init__(self, inner: Optional[Any]):
def __init__(
self,
inner: Optional[Any],
result_decoder: Optional[Callable[[Any], T]] = None,
):
self._inner = inner
self._result_decoder = result_decoder
@property
def id(self) -> Optional[str]:
@@ -50,17 +56,21 @@ class AsyncJob(Generic[T]):
async def wait(self, timeout: Optional[timedelta] = None) -> T:
"""Wait until the operation reaches a terminal state.
Returns the endpoint's typed result, or ``None`` for a unit-result
job.
Raises `JobFailedError` if the operation failed, `JobCancelledError`
if it was cancelled, and `TimeoutError` if `timeout` elapses first.
"""
if self._inner is None:
return cast(T, None)
if timeout is None:
return cast(T, await self._inner.wait())
return cast(
T,
await asyncio.wait_for(self._inner.wait(), timeout.total_seconds()),
)
result = await self._inner.wait()
else:
result = await asyncio.wait_for(self._inner.wait(), timeout.total_seconds())
if self._result_decoder is not None:
return self._result_decoder(result)
return cast(T, result)
async def cancel(self):
"""Request cancellation. Cancelling a finished operation is a no-op."""
@@ -70,7 +80,7 @@ class AsyncJob(Generic[T]):
class Job(Generic[T]):
"""Synchronous counterpart of `AsyncJob`."""
"""Synchronous counterpart of `AsyncJob` with the same result type."""
def __init__(self, inner: Optional[AsyncJob[T]]):
self._inner = inner
@@ -96,6 +106,9 @@ class Job(Generic[T]):
def wait(self, timeout: Optional[timedelta] = None) -> T:
"""Block until the operation reaches a terminal state.
Returns the endpoint's typed result, or ``None`` for a unit-result
job.
Raises `JobFailedError` if the operation failed, `JobCancelledError`
if it was cancelled, and `TimeoutError` if `timeout` elapses first.
"""
@@ -110,23 +123,8 @@ class Job(Generic[T]):
LOOP.run(self._inner.cancel())
class _FunctionJobAdapter:
def __init__(self, inner: "_lancedb.FunctionJob"):
self._inner = inner
@property
def id(self) -> Optional[str]:
return self._inner.id
async def status(self) -> str:
return await self._inner.status()
async def wait(self) -> FunctionVersion:
return FunctionVersion.from_json(await self._inner.wait())
async def cancel(self):
await self._inner.cancel()
def _function_job(inner: "_lancedb.FunctionJob") -> AsyncJob[FunctionVersion]:
return AsyncJob(_FunctionJobAdapter(inner))
def _typed_job(
inner: "_lancedb.Job", result_decoder: Callable[[str], T]
) -> AsyncJob[T]:
"""Bind an internal JSON-producing job to its public result model."""
return AsyncJob(inner, result_decoder)
+21 -5
View File
@@ -391,6 +391,15 @@ def _table_to_pickle_state(table: Table) -> dict[str, Any]:
}
def _drop_base_version(permutation_data: pa.Table) -> pa.Table:
"""Strip the recorded base version so the reader leaves the base table unpinned."""
metadata = dict(permutation_data.schema.metadata or {})
if metadata.pop(b"base_version", None) is None:
return permutation_data
metadata.pop(b"base_branch", None)
return permutation_data.replace_schema_metadata(metadata)
def _table_from_pickle_state(state: dict[str, Any]) -> Table:
from . import connect
@@ -679,11 +688,15 @@ class Permutation:
from . import connect
connection_factory = state["connection_factory"]
rebuilt_base = False
if connection_factory is not None:
base_table = connection_factory(state["base_table_name"])
elif "base_table_state" in state:
base_table = _table_from_pickle_state(state["base_table_state"])
base_state = state["base_table_state"]
rebuilt_base = base_state["kind"] == "memory"
base_table = _table_from_pickle_state(base_state)
elif "base_table_data" in state:
rebuilt_base = True
# In-memory base table inlined into the pickle; rebuild the same
# way we rebuild the in-memory permutation table.
mem_db = connect("memory://")
@@ -701,11 +714,14 @@ class Permutation:
)
permutation_table: Optional[Table] = None
if state["permutation_data"] is not None:
permutation_data = state["permutation_data"]
if permutation_data is not None:
if rebuilt_base:
# The base table was materialized from Arrow, so it is a fresh
# single-version dataset and the recorded pin cannot resolve on it.
permutation_data = _drop_base_version(permutation_data)
mem_db = connect("memory://")
permutation_table = mem_db.create_table(
"permutation", state["permutation_data"]
)
permutation_table = mem_db.create_table("permutation", permutation_data)
self.base_table = base_table
self.permutation_table = permutation_table
+41 -9
View File
@@ -167,6 +167,12 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
return {"columns": projection}
def _query_request_projection(req: "PyQueryRequest") -> QueryProjection:
if req.select_source_columns is not None:
return req.select_source_columns
return req.select
def _scanner_kwargs_for_query(
query: Query,
blob_mode: BlobMode,
@@ -375,6 +381,13 @@ class FullTextOperator(str, Enum):
OR = "OR"
class DocumentGranularity(str, Enum):
"""The unit treated as one full-text-search document."""
ROW = "row"
LIST_ELEMENT = "list_element"
class Occur(str, Enum):
SHOULD = "SHOULD"
MUST = "MUST"
@@ -478,6 +491,10 @@ class MatchQuery(FullTextQuery):
prefix_length : int, optional
The number of beginning characters being unchanged for fuzzy matching.
This is useful to achieve prefix matching.
document_granularity : DocumentGranularity, optional
Explicitly select row or deepest-list-element documents. If omitted,
the indexed granularity is inferred. When both granularities are indexed
for the field, this must be specified. With no index, row granularity is used.
"""
query: str
@@ -487,6 +504,9 @@ class MatchQuery(FullTextQuery):
max_expansions: int = pydantic.Field(50, kw_only=True)
operator: FullTextOperator = pydantic.Field(FullTextOperator.OR, kw_only=True)
prefix_length: int = pydantic.Field(0, kw_only=True)
document_granularity: Optional[DocumentGranularity] = pydantic.Field(
None, kw_only=True
)
def query_type(self) -> FullTextQueryType:
return FullTextQueryType.MATCH
@@ -503,11 +523,20 @@ class PhraseQuery(FullTextQuery):
The query string to match against.
column : str
The name of the column to match against.
slop : int, default 0
The maximum number of intervening positions permitted in the phrase.
document_granularity : DocumentGranularity, optional
Explicitly select row or deepest-list-element documents. If omitted,
the indexed granularity is inferred. When both granularities are indexed
for the field, this must be specified. With no index, row granularity is used.
"""
query: str
column: str
slop: int = pydantic.Field(0, kw_only=True)
document_granularity: Optional[DocumentGranularity] = pydantic.Field(
None, kw_only=True
)
def query_type(self) -> FullTextQueryType:
return FullTextQueryType.MATCH_PHRASE
@@ -2776,15 +2805,16 @@ class AsyncQueryBase(object):
req = self._inner.to_query_request()
schema = await self._table.schema()
projection = _query_request_projection(req)
self._blob_auto_row_id = blob_auto_row_id_for_scan(
schema,
req.select,
projection,
with_row_id=self._with_row_id,
)
if not self._blob_auto_row_id:
self._blob_paths = ()
return
self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
self._blob_paths = tuple(blob_v2_projection_sources(schema, projection).keys())
self._inner.with_row_id()
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
@@ -3378,9 +3408,10 @@ class AsyncQuery(AsyncStandardQuery):
pass in multiple vectors. When multiple vectors are passed in, if the vector
column is with multivector type, then the vectors will be treated as a single
query. Or the vectors will be treated as multiple queries, this can be useful
if you want to find the nearest vectors to multiple query vectors.
This is not expected to be faster than making multiple queries concurrently;
it is just a convenience method. If multiple vectors are passed in then
if you want to find the nearest vectors to multiple query vectors. Flat
searches share one table scan across the query vectors, avoiding the scan
and memory amplification of making multiple queries concurrently. If
multiple vectors are passed in then
an additional column `query_index` will be added to the results. This column
will contain the index of the query vector that the result is nearest to.
"""
@@ -3509,8 +3540,8 @@ class AsyncFTSQuery(AsyncStandardQuery):
Typically, a single vector is passed in as the query. However, you can also
pass in multiple vectors. This can be useful if you want to find the nearest
vectors to multiple query vectors. This is not expected to be faster than
making multiple queries concurrently; it is just a convenience method.
vectors to multiple query vectors. Flat searches share one table scan across
the query vectors instead of issuing concurrent full scans.
If multiple vectors are passed in then an additional column `query_index`
will be added to the results. This column will contain the index of the
query vector that the result is nearest to.
@@ -3870,14 +3901,15 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
blob_paths: tuple[str, ...] = ()
if self._table is not None:
schema = await self._table.schema()
projection = _query_request_projection(req)
blob_auto_row_id = blob_auto_row_id_for_scan(
schema,
req.select,
projection,
with_row_id=self._with_row_id,
)
if blob_auto_row_id:
blob_paths = tuple(
blob_v2_projection_sources(schema, req.select).keys()
blob_v2_projection_sources(schema, projection).keys()
)
self._blob_auto_row_id = blob_auto_row_id
self._blob_paths = blob_paths
+16 -6
View File
@@ -36,6 +36,7 @@ from lancedb._lancedb import (
UpdateResult,
)
from lancedb.embeddings.base import EmbeddingFunctionConfig
from lancedb.expr import Expr
from lancedb.index import (
FTS,
BTree,
@@ -49,7 +50,7 @@ from lancedb.index import (
LabelList,
)
from lancedb.job import Job
from lancedb.functions import FunctionApplication
from lancedb.functions import FunctionApplication, RefreshColumnResult
from lancedb.remote.db import LOOP
from lancedb.table import IndexConfigType, KNOWN_METRICS
import pyarrow as pa
@@ -61,6 +62,7 @@ from lancedb.table import _normalize_progress
from ..query import (
AnalyzePlanDistributedMetrics,
DocumentGranularity,
LanceQueryBuilder,
LanceTakeQueryBuilder,
LanceVectorQueryBuilder,
@@ -349,6 +351,7 @@ class RemoteTable(Table):
ngram_max_length: int = 3,
prefix_only: bool = False,
block_size: int = 128,
document_granularity: DocumentGranularity = DocumentGranularity.ROW,
name: Optional[str] = None,
):
"""Create a full-text search index on a column.
@@ -371,6 +374,7 @@ class RemoteTable(Table):
ngram_max_length=ngram_max_length,
prefix_only=prefix_only,
block_size=block_size,
document_granularity=document_granularity,
)
LOOP.run(
self._table.create_index(
@@ -610,6 +614,7 @@ class RemoteTable(Table):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [Table][lancedb.table.Table].
@@ -642,6 +647,8 @@ class RemoteTable(Table):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Not supported on LanceDB Cloud. Setting this raises.
Returns
-------
@@ -658,6 +665,7 @@ class RemoteTable(Table):
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
)
finally:
@@ -856,7 +864,7 @@ class RemoteTable(Table):
def update(
self,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -867,9 +875,11 @@ class RemoteTable(Table):
Parameters
----------
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
@@ -972,7 +982,7 @@ class RemoteTable(Table):
def refresh_column(self, column: str):
return LOOP.run(self._table.refresh_column(column))
def refresh_column_async(self, column: str) -> Job:
def refresh_column_async(self, column: str) -> Job[RefreshColumnResult]:
return Job(LOOP.run(self._table.refresh_column_async(column)))
def alter_columns(
File diff suppressed because it is too large Load Diff
+101 -28
View File
@@ -40,7 +40,7 @@ from ._blob import (
from .types import BlobMode
from lancedb.arrow import peek_reader
from lancedb.background_loop import LOOP, embedding_executor
from lancedb.job import AsyncJob, Job
from lancedb.job import AsyncJob, Job, _typed_job
from .dependencies import (
_check_for_hugging_face,
_check_for_lance,
@@ -72,7 +72,10 @@ from .index import (
FTS,
)
from .expr import Expr
from .functions import FunctionApplication
from .functions import (
FunctionApplication,
RefreshColumnResult as RefreshColumnJobResult,
)
from .merge import LanceMergeInsertBuilder
from .pydantic import LanceModel, model_to_dict
from .query import (
@@ -82,6 +85,7 @@ from .query import (
AsyncQuery,
AsyncTakeQuery,
AsyncVectorQuery,
DocumentGranularity,
FullTextQuery,
LanceEmptyQueryBuilder,
LanceFtsQueryBuilder,
@@ -1165,6 +1169,7 @@ class Table(ABC):
ngram_max_length: int = 3,
prefix_only: bool = False,
block_size: int = 128,
document_granularity: DocumentGranularity = DocumentGranularity.ROW,
wait_timeout: Optional[timedelta] = None,
name: Optional[str] = None,
):
@@ -1243,6 +1248,11 @@ class Table(ABC):
The number of documents per compressed posting block. Must be 128
or 256. A value of 256 uses the experimental FTS V3 format and
may introduce breaking changes.
document_granularity: DocumentGranularity, default ROW
``ROW`` treats the selected text in one table row as one document.
``LIST_ELEMENT`` treats each element of the deepest list on the field
path as one document and returns its physical coordinates in
``_doc_index`` for matching queries.
wait_timeout: timedelta, optional
The timeout to wait if indexing is asynchronous.
name: str, optional
@@ -1266,6 +1276,7 @@ class Table(ABC):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [Table][lancedb.table.Table].
@@ -1317,6 +1328,10 @@ class Table(ABC):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Store blob URIs that sit outside registered blob bases. The row
keeps a reference, so the object has to stay readable. Local
tables only.
Returns
-------
@@ -1729,7 +1744,7 @@ class Table(ABC):
@abstractmethod
def update(
self,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -1744,9 +1759,11 @@ class Table(ABC):
Parameters
----------
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
@@ -1764,6 +1781,7 @@ class Table(ABC):
Examples
--------
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
>>> db = lancedb.connect("./.lancedb")
@@ -1773,7 +1791,7 @@ class Table(ABC):
0 1 [1.0, 2.0]
1 2 [3.0, 4.0]
2 3 [5.0, 6.0]
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
UpdateResult(rows_updated=1, version=2)
>>> table.to_pandas()
x vector
@@ -1969,7 +1987,7 @@ class Table(ABC):
A mapping with one ``FunctionApplication`` value keeps its scalar
or named-struct result in the named table column. A bare
named-struct application expands its ordered result fields as one
atomic sibling group; aliases come from ``rename(columns=...)``.
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
@@ -2039,7 +2057,7 @@ class Table(ABC):
"""
@abstractmethod
def refresh_column_async(self, column: str) -> Job:
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
@@ -2050,6 +2068,12 @@ class Table(ABC):
than failing the job. On local tables the job runs in-process; on
LanceDB Cloud and Enterprise it is the server's backfill job.
Returns
-------
Job[RefreshColumnResult]
A job whose successful ``wait`` returns row counts plus the source
and published table versions.
Examples
--------
>>> import lancedb
@@ -2058,7 +2082,9 @@ class Table(ABC):
>>> table.add_columns(computed={"doubled": "x * 2"})
AddColumnsResult(version=2)
>>> job = table.refresh_column_async("doubled")
>>> job.wait()
>>> result = job.wait()
>>> result.rows_assigned
2
>>> job.status()
'finished'
"""
@@ -2104,12 +2130,25 @@ class Table(ABC):
----------
updates : dict
One or more dicts, each with:
- "path": str dot-path to the field (e.g. "embedding" or "a.b.c").
- "metadata": dict[str, str | None] keys to set; a value of ``None``
deletes that key.
- "replace": bool, optional replace the field's whole metadata map
instead of merging (default False).
The following keys are treated specially, by convention, and should
be used when appropriate:
- "lancedb:description": for a human-readable description of a field.
- ``"lancedb:tag:<name>"`` for a user-defined key-value tag, where the
suffix names the tag category; e.g. "lancedb:tag:model": "clip".
- "lancedb:logical-column" for a column grouping; e.g. "feature_v1"
and "feature_v2" might be in the same logical column.
- "lancedb:status" for status options ("production", "candidate",
"deprecated", "archived") to designate the current life cycle
state of this column.
Returns
-------
UpdateFieldMetadataResult
@@ -3257,6 +3296,7 @@ class LanceTable(Table):
ngram_max_length: int = 3,
prefix_only: bool = False,
block_size: int = 128,
document_granularity: DocumentGranularity = DocumentGranularity.ROW,
name: Optional[str] = None,
):
"""Create a full-text search index on a column.
@@ -3308,7 +3348,11 @@ class LanceTable(Table):
tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name)
tokenizer_configs["custom_stop_words"] = custom_stop_words
config = FTS(block_size=block_size, **tokenizer_configs)
config = FTS(
block_size=block_size,
document_granularity=document_granularity,
**tokenizer_configs,
)
try:
LOOP.run(
@@ -3398,6 +3442,7 @@ class LanceTable(Table):
fill_value: float = 0.0,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add data to the table.
If vector columns are missing and the table
@@ -3425,6 +3470,9 @@ class LanceTable(Table):
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Allow blob URIs outside registered bases. See :meth:`Table.add`.
Local tables only.
Returns
-------
@@ -3441,6 +3489,7 @@ class LanceTable(Table):
fill_value=fill_value,
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
)
finally:
@@ -3795,7 +3844,7 @@ class LanceTable(Table):
def update(
self,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
values: Optional[dict] = None,
*,
values_sql: Optional[Dict[str, str]] = None,
@@ -3806,9 +3855,11 @@ class LanceTable(Table):
Parameters
----------
where: str, optional
The SQL where clause to use when updating rows. For example, 'x = 2'
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
error.
values: dict, optional
The values to update. The keys are the column names and the values
are the values to set.
@@ -3826,6 +3877,7 @@ class LanceTable(Table):
Examples
--------
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
>>> db = lancedb.connect("./.lancedb")
@@ -3835,7 +3887,7 @@ class LanceTable(Table):
0 1 [1.0, 2.0]
1 2 [3.0, 4.0]
2 3 [5.0, 6.0]
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
UpdateResult(rows_updated=1, version=2)
>>> table.to_pandas()
x vector
@@ -4082,7 +4134,7 @@ class LanceTable(Table):
[`AsyncTable.refresh_column`][lancedb.AsyncTable.refresh_column]."""
return LOOP.run(self._table.refresh_column(column))
def refresh_column_async(self, column: str) -> Job:
def refresh_column_async(self, column: str) -> Job[RefreshColumnJobResult]:
"""Fill a computed column's unfilled rows, returning a handle to the
refresh job. See
[`Table.refresh_column_async`][lancedb.table.Table.refresh_column_async].
@@ -5354,6 +5406,7 @@ class AsyncTable:
fill_value: Optional[float] = None,
progress: Optional[Union[bool, Callable, Any]] = None,
write_parallelism: Optional[int] = None,
allow_external_blob_outside_bases: bool = False,
) -> AddResult:
"""Add more data to the [AsyncTable][lancedb.table.AsyncTable].
@@ -5384,6 +5437,9 @@ class AsyncTable:
data in flight. Defaults to an estimate based on the data size,
capped at the number of CPU cores. Lower this if bulk ingestion is
using too much memory.
allow_external_blob_outside_bases: bool, default False
Allow blob URIs outside registered bases. See :meth:`Table.add`.
Local tables only.
"""
schema = await self.schema()
@@ -5420,6 +5476,7 @@ class AsyncTable:
mode or "append",
progress=progress,
write_parallelism=write_parallelism,
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
)
except RuntimeError as e:
if "Cast error" in str(e):
@@ -5944,7 +6001,7 @@ class AsyncTable:
self,
updates: Optional[Dict[str, Any]] = None,
*,
where: Optional[str] = None,
where: Optional[Union[str, Expr]] = None,
updates_sql: Optional[Dict[str, str]] = None,
) -> UpdateResult:
"""
@@ -5959,9 +6016,11 @@ class AsyncTable:
The updates to apply. The keys should be the name of the column to
update. The values should be the new values to assign. This is
required unless updates_sql is supplied.
where: str, optional
An SQL filter that controls which rows are updated. For example, 'x = 2'
or 'x IN (1, 2, 3)'. Only rows that satisfy this filter will be udpated.
where: str or [Expr][lancedb.expr.Expr], optional
The filter condition. Can be a SQL string or a type-safe
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
[lit][lancedb.expr.lit]. Only rows that satisfy this filter will
be updated.
updates_sql: dict, optional
The updates to apply, expressed as SQL expression strings. The keys should
be column names. The values should be SQL expressions. These can be SQL
@@ -5979,13 +6038,14 @@ class AsyncTable:
--------
>>> import asyncio
>>> import lancedb
>>> from lancedb.expr import col
>>> import pandas as pd
>>> async def demo_update():
... data = pd.DataFrame({"x": [1, 2], "vector": [[1, 2], [3, 4]]})
... db = await lancedb.connect_async("./.lancedb")
... table = await db.create_table("my_table", data)
... # x is [1, 2], vector is [[1, 2], [3, 4]]
... await table.update({"vector": [10, 10]}, where="x = 2")
... await table.update({"vector": [10, 10]}, where=col("x") == 2)
... # x is [1, 2], vector is [[1, 2], [10, 10]]
... await table.update(updates_sql={"x": "x + 1"})
... # x is [2, 3], vector is [[1, 2], [10, 10]]
@@ -5999,7 +6059,8 @@ class AsyncTable:
if updates is not None:
updates_sql = {k: value_to_sql(v) for k, v in updates.items()}
return await self._inner.update(updates_sql, where)
predicate = where.to_sql() if isinstance(where, Expr) else where
return await self._inner.update(updates_sql, predicate)
async def add_columns(
self,
@@ -6027,7 +6088,7 @@ class AsyncTable:
A mapping with one ``FunctionApplication`` value keeps its scalar
or named-struct result in the named table column. A bare
named-struct application expands its ordered result fields as one
atomic sibling group; aliases come from ``rename(columns=...)``.
atomic binding; aliases come from ``rename(columns=...)``.
Function columns are supported only on LanceDB Cloud and
Enterprise.
computed: Dict[str, str], optional
@@ -6064,7 +6125,7 @@ class AsyncTable:
isinstance(value, FunctionApplication) for value in transforms.values()
):
raise ValueError(
"one add_columns call declares exactly one Function sibling group"
"one add_columns call declares exactly one Function binding"
)
function_output_name, function_application = next(iter(transforms.items()))
@@ -6122,7 +6183,9 @@ class AsyncTable:
"""
return await self._inner.refresh_column(column)
async def refresh_column_async(self, column: str) -> AsyncJob:
async def refresh_column_async(
self, column: str
) -> AsyncJob[RefreshColumnJobResult]:
"""
Like :meth:`refresh_column`, but returns a handle to the refresh job
instead of blocking until it completes.
@@ -6134,6 +6197,12 @@ class AsyncTable:
in-process; on LanceDB Cloud and Enterprise it is the server's
backfill job.
Returns
-------
AsyncJob[RefreshColumnResult]
A job whose successful ``wait`` returns row counts plus the source
and published table versions.
Examples
--------
>>> import asyncio
@@ -6143,12 +6212,16 @@ class AsyncTable:
... table = await db.create_table("computed_job_async_demo", [{"x": 1}])
... await table.add_columns(computed={"doubled": "x * 2"})
... job = await table.refresh_column_async("doubled")
... await job.wait()
... result = await job.wait()
... assert result.rows_assigned == 1
... return await job.status()
>>> asyncio.run(refresh_in_background())
'finished'
"""
return AsyncJob(await self._inner.refresh_column_async(column))
return _typed_job(
await self._inner.refresh_column_async(column),
RefreshColumnJobResult.from_json,
)
async def alter_columns(
self, *alterations: Iterable[dict[str, Any]]
+2 -2
View File
@@ -105,7 +105,7 @@ def test_quickstart(tmp_path):
tbl.create_index(num_sub_vectors=1)
# --8<-- [end:create_index]
# --8<-- [start:delete_rows]
tbl.delete('item = "fizz"')
tbl.delete("item = 'fizz'")
# --8<-- [end:delete_rows]
# --8<-- [start:drop_table]
db.drop_table("my_table")
@@ -201,7 +201,7 @@ async def test_quickstart_async(tmp_path):
await tbl.create_index("vector")
# --8<-- [end:create_index_async]
# --8<-- [start:delete_rows_async]
await tbl.delete('item = "fizz"')
await tbl.delete("item = 'fizz'")
# --8<-- [end:delete_rows_async]
# --8<-- [start:drop_table_async]
await db.drop_table("my_table_async")
@@ -266,7 +266,7 @@ def test_table():
tbl.add(pydantic_model_items)
# --8<-- [end:add_table_from_pydantic]
# --8<-- [start:delete_row]
tbl.delete('item = "fizz"')
tbl.delete("item = 'fizz'")
# --8<-- [end:delete_row]
# --8<-- [start:delete_specific_row]
data = [
@@ -538,7 +538,7 @@ async def test_table_async():
await async_tbl.add(pydantic_model_items)
# --8<-- [end:add_table_async_from_pydantic]
# --8<-- [start:delete_row_async]
await async_tbl.delete('item = "fizz"')
await async_tbl.delete("item = 'fizz'")
# --8<-- [end:delete_row_async]
# --8<-- [start:delete_specific_row_async]
data = [
+140 -1
View File
@@ -8,7 +8,12 @@ import pyarrow.compute as pc
import pytest
import lancedb
from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids
from lancedb._blob import (
blob_v2_projection_sources,
read_row_ids_from_hits,
stash_auto_row_ids,
)
from lancedb.expr import col
from lancedb.index import FTS
from lancedb.schema import blob_column_paths, blob_v2_column_paths
@@ -70,6 +75,14 @@ def test_blob_v2_column_paths_include_list_children():
]
def test_blob_v2_projection_sources_use_typed_column_name():
schema = pa.schema([lancedb.blob("blob")])
assert blob_v2_projection_sources(schema, {"blob_alias": col("blob")}) == {
"blob_alias": "blob"
}
def _legacy_v1_table(name):
db = lancedb.connect("memory:///")
schema = pa.schema(
@@ -166,6 +179,20 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
@pytest.mark.asyncio
async def test_async_typed_blob_projection_preserves_source_column():
db = await lancedb.connect_async("memory:///typed_blob_projection")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
table = await db.create_table("typed_blob_projection", schema=schema)
await table.add([{"id": 1, "blob": b"alpha"}])
hits = await table.query().select({"blob_alias": col("blob")}).to_arrow()
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
blobs = await table.fetch_blobs("blob", hits)
assert blobs.to_pylist() == [b"alpha"]
def test_fetch_blobs_round_trip():
table = _blob_table(
"round_trip",
@@ -403,6 +430,50 @@ async def test_blob_v2_hybrid_fetch_blobs_async():
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
@pytest.mark.asyncio
async def test_async_hybrid_typed_blob_projection_preserves_source_column():
db = await lancedb.connect_async("memory:///hybrid_typed_blob")
schema = pa.schema(
[
pa.field("id", pa.int64()),
pa.field("text", pa.utf8()),
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
lancedb.blob("blob"),
]
)
table = await db.create_table("hybrid_typed_blob", schema=schema)
await table.add(
[
{
"id": 1,
"text": "hello alpha",
"vector": [1.0, 0.0],
"blob": b"alpha",
},
{
"id": 2,
"text": "hello beta",
"vector": [0.9, 0.1],
"blob": b"beta",
},
]
)
await table.create_index("text", config=FTS(with_position=False))
hits = await (
table.query()
.nearest_to([1.0, 0.0])
.nearest_to_text("hello")
.select({"blob_alias": col("blob")})
.limit(2)
.to_arrow()
)
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
blobs = await table.fetch_blobs("blob", hits)
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
def test_blob_file_seek_read_and_read_range():
payload = _identifiable_payload(1024)
table = _blob_table("seek_read", [{"id": 1, "image": payload}])
@@ -617,3 +688,71 @@ def test_fetch_blobs_nested_path_survives_sort_after_query():
def _identifiable_payload(size: int) -> bytes:
block = 256
return b"".join(bytes([i % 256]) * block for i in range(size // block))
def _external_uri_blob_array(uris):
blob_type = lancedb.blob("image").type
storage_type = blob_type.storage_type
child_names = [field.name for field in storage_type]
assert "uri" in child_names, "blob layout no longer has a uri child"
children = [
pa.array(uris if field.name == "uri" else [None] * len(uris), type=field.type)
for field in storage_type
]
storage = pa.StructArray.from_arrays(children, fields=list(storage_type))
return pa.ExtensionArray.from_storage(blob_type, storage)
def _external_uri_table_and_rows(name, uris):
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table(name, schema=schema)
rows = pa.Table.from_arrays(
[
pa.array(range(len(uris)), type=pa.int64()),
_external_uri_blob_array(uris),
],
schema=schema,
)
return table, rows
def test_add_external_uri_struct_round_trips_with_flag(tmp_path):
payload = b"external-uri-bytes"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
table, rows = _external_uri_table_and_rows("external_struct", [blob_path.as_uri()])
table.add(rows, allow_external_blob_outside_bases=True)
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
def test_add_external_uri_without_flag_raises(tmp_path):
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(b"unreachable")
table, rows = _external_uri_table_and_rows("external_no_flag", [blob_path.as_uri()])
with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
table.add(rows)
assert table.count_rows() == 0
def test_add_external_uri_string_round_trips_with_flag(tmp_path):
payload = b"external-uri-bytes"
blob_path = tmp_path / "payload.bin"
blob_path.write_bytes(payload)
db = lancedb.connect("memory:///")
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
table = db.create_table("external_string", schema=schema)
table.add(
[{"id": 1, "image": blob_path.as_uri()}],
allow_external_blob_outside_bases=True,
)
hits = table.search().to_arrow()
blobs = table.fetch_blobs("image", hits)
assert blobs[0].as_py() == payload
+2 -2
View File
@@ -774,7 +774,7 @@ def test_drop_table_async(tmp_db: lancedb.DBConnection):
job = tmp_db.drop_table_async("test")
assert job.id is None
assert job.status() == "finished"
job.wait()
assert job.wait() is None
assert tmp_db.table_names() == []
tmp_db.create_table("test", data=data)
@@ -790,7 +790,7 @@ async def test_drop_table_async_connection(tmp_db_async: lancedb.AsyncConnection
job = await tmp_db_async.drop_table_async("test")
assert job.id is None
assert await job.status() == "finished"
await job.wait()
assert await job.wait() is None
assert await tmp_db_async.table_names() == []
@@ -32,6 +32,7 @@ Parameters used throughout:
import dataclasses
import logging
import threading
from unittest.mock import patch
import lancedb
@@ -46,6 +47,7 @@ from utils import (
torch = pytest.importorskip("torch")
streaming = pytest.importorskip("lancedb.streaming")
StreamingDataset = streaming.StreamingDataset
StreamingDataLoader = streaming.StreamingDataLoader
# ---------------------------------------------------------------------------
# Dataset parameters
@@ -92,6 +94,27 @@ class FakeWorkerInfo:
num_workers: int
def _collate_with_first_batch_error(samples):
ids = [sample["id"] for sample in samples]
if ids == [0, 1]:
raise ValueError("first batch fails")
return ids
def _collate_with_first_batch_stop(samples):
ids = [sample["id"] for sample in samples]
if ids == [0, 1]:
raise StopIteration("first batch stopped")
return ids
def _collate_with_first_batch_interrupt(samples):
ids = [sample["id"] for sample in samples]
if ids == [0, 1]:
raise KeyboardInterrupt("first batch interrupted")
return ids
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@@ -1008,6 +1031,565 @@ def test_multi_worker_elastic_det_across_worker_counts(lance_table):
# ── Resumability with num_workers ─────────────────────────────────────────────
def test_streaming_dataloader_commits_only_consumed_worker_batches(tmp_path):
"""Prefetched worker state is committed only as the trainer receives it."""
db = lancedb.connect(tmp_path)
table = db.create_table(
"worker_commit", pa.table({"id": [1, 2, 3, 4, 10, 20, 30, 40]})
)
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=4,
)
iterator = iter(loader)
try:
first = next(iterator)["id"].tolist()
assert first == [1, 2]
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2, 0]
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
second = next(iterator)["id"].tolist()
assert second == [10, 20]
checkpoint = dataset.state_dict()
assert checkpoint["samples_consumed_per_split"] == [2, 2]
uninterrupted = [batch["id"].tolist() for batch in iterator]
finally:
iterator._shutdown_workers()
resumed = StreamingDataset(table, num_splits=2, shuffle=False)
resumed.load_state_dict(checkpoint)
resumed_loader = StreamingDataLoader(
resumed,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=4,
)
resumed_iterator = iter(resumed_loader)
try:
remaining = [batch["id"].tolist() for batch in resumed_iterator]
finally:
resumed_iterator._shutdown_workers()
assert remaining == uninterrupted == [[3, 4], [30, 40]]
def test_distributed_checkpoint_uses_rank_local_worker_boundary(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("rank_boundary", pa.table({"id": list(range(8))}))
dataset = StreamingDataset(
table,
num_splits=4,
shuffle=False,
rank=0,
world_size=2,
)
loader = StreamingDataLoader(
dataset,
batch_size=1,
num_workers=2,
multiprocessing_context="spawn",
)
iterator = iter(loader)
try:
assert next(iterator)["id"].tolist() == [0]
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [
1,
0,
0,
0,
]
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
assert next(iterator)["id"].tolist() == [2]
checkpoint = dataset.state_dict()
remaining = [batch["id"].tolist() for batch in iterator]
finally:
iterator._shutdown_workers()
assert checkpoint["samples_consumed_per_split"] == [1, 1, 0, 0]
assert remaining == [[1], [3]]
def test_standard_dataloader_rejects_stale_parent_checkpoint(tmp_path):
"""A standard DataLoader must not expose prefetched producer progress."""
db = lancedb.connect(tmp_path)
table = db.create_table("untracked_workers", pa.table({"id": [1, 2, 10, 20]}))
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
# Merely constructing the checkpoint-aware loader must not authorize a
# later plain DataLoader's worker progress.
StreamingDataLoader(dataset, batch_size=2, num_workers=0)
loader = torch.utils.data.DataLoader(
dataset,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
)
iterator = iter(loader)
try:
assert next(iterator)["id"].tolist() == [1, 2]
with pytest.raises(RuntimeError, match="Use StreamingDataLoader"):
dataset.state_dict()
list(iterator)
finally:
iterator._shutdown_workers()
def test_streaming_dataloader_rejects_persistent_workers(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("persistent_workers", pa.table({"id": [1, 2]}))
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
with pytest.raises(ValueError, match="persistent_workers=True"):
StreamingDataLoader(
dataset,
batch_size=1,
num_workers=2,
persistent_workers=True,
)
def test_collate_failure_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table(
"collate_failure", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]})
)
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=2,
multiprocessing_context="spawn",
collate_fn=_collate_with_first_batch_error,
prefetch_factor=2,
)
iterator = iter(loader)
try:
with pytest.raises(ValueError, match="first batch fails"):
next(iterator)
assert next(iterator) == [100, 101]
assert next(iterator) == [2, 3]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
list(iterator)
finally:
iterator._shutdown_workers()
def test_collate_stop_iteration_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("collate_stop", pa.table({"id": list(range(6))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=0,
collate_fn=_collate_with_first_batch_stop,
)
iterator = iter(loader)
with pytest.raises(RuntimeError, match="collate_fn raised StopIteration"):
next(iterator)
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
assert list(iterator) == [[2, 3], [4, 5]]
def test_batch_base_exception_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("collate_interrupt", pa.table({"id": list(range(6))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=0,
collate_fn=_collate_with_first_batch_interrupt,
)
iterator = iter(loader)
with pytest.raises(KeyboardInterrupt, match="first batch interrupted"):
next(iterator)
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
assert list(iterator) == [[2, 3], [4, 5]]
def test_parent_commit_base_exception_invalidates_consumer_checkpoint(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("commit_interrupt", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
iterator = iter(loader)
real_commit = dataset._commit_worker_state
def interrupt_after_commit(state, *, require_uniform):
real_commit(state, require_uniform=require_uniform)
raise KeyboardInterrupt("after parent commit")
with patch.object(
dataset, "_commit_worker_state", side_effect=interrupt_after_commit
):
with pytest.raises(KeyboardInterrupt, match="after parent commit"):
next(iterator)
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
with pytest.raises(RuntimeError, match="failed before it was returned"):
dataset.state_dict()
def test_direct_iteration_surfaces_prefetch_failure_before_committing_row(
tmp_path, monkeypatch
):
db = lancedb.connect(tmp_path)
table = db.create_table("prefetch_failure", pa.table({"id": list(range(4))}))
release = threading.Event()
failed = threading.Event()
real_getitems = streaming.Permutation.__getitems__
def controlled_getitems(permutation, indices):
if indices and indices[0] >= 2:
assert release.wait(timeout=5)
failed.set()
raise RuntimeError("later prefetched I/O failed")
return real_getitems(permutation, indices)
class SignalDict(dict):
def __setitem__(self, key, value):
super().__setitem__(key, value)
release.set()
assert failed.wait(timeout=5)
monkeypatch.setattr(streaming.Permutation, "__getitems__", controlled_getitems)
dataset = StreamingDataset(
table,
num_splits=1,
shuffle=False,
read_batch_size=2,
io_queue_depth=2,
)
dataset._resume_positions = SignalDict()
iterator = iter(dataset)
assert next(iterator)["id"] == 0
with pytest.raises(RuntimeError, match="later prefetched I/O failed"):
next(iterator)
checkpoint = dataset.state_dict()
assert checkpoint["samples_consumed_per_split"] == [1]
assert checkpoint["positions_consumed_per_split"] == [1]
@pytest.mark.parametrize("workers", [0, 1, 2])
def test_streaming_dataloader_rejects_drop_last(tmp_path, workers):
db = lancedb.connect(tmp_path)
table = db.create_table("drop_last", pa.table({"id": [0, 1, 2]}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
worker_options = {"multiprocessing_context": "spawn"} if workers else {}
with pytest.raises(ValueError, match="drop_last=True"):
StreamingDataLoader(
dataset,
batch_size=2,
num_workers=workers,
drop_last=True,
**worker_options,
)
def test_streaming_dataloader_owns_one_iterator_until_teardown(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("iterator_owner", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(
dataset,
batch_size=2,
num_workers=1,
multiprocessing_context="spawn",
)
first = iter(loader)
try:
assert next(first)["id"].tolist() == [0, 1]
with pytest.raises(RuntimeError, match="concurrent iteration"):
iter(loader)
finally:
first._shutdown_workers()
second = iter(loader)
try:
assert [batch["id"].tolist() for batch in second] == [[2, 3]]
except BaseException:
second._shutdown_workers()
raise
# Natural exhaustion releases ownership too.
third = iter(loader)
try:
assert list(third) == []
finally:
third._shutdown_workers()
def test_zero_worker_shutdown_closes_inner_iterator_before_release(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("zero_worker_shutdown", pa.table({"id": list(range(6))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
first = iter(loader)
assert next(first)["id"].tolist() == [0, 1]
first._shutdown_workers()
assert dataset._consumer_iterator_active is False
assert dataset._raw_batches_ref is None
second = iter(loader)
try:
with pytest.raises(StopIteration):
next(first)
assert next(second)["id"].tolist() == [2, 3]
finally:
second._shutdown_workers()
def test_direct_and_loader_admission_share_one_atomic_lease(tmp_path, monkeypatch):
db = lancedb.connect(tmp_path)
table = db.create_table("direct_loader_lease", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
entered = threading.Event()
release = threading.Event()
direct_result = []
direct_error = []
contender = []
real_resolve = dataset._resolve_my_splits
def controlled_resolve():
if threading.current_thread().name == "direct-start":
entered.set()
assert release.wait(timeout=5)
return real_resolve()
def advance_direct(iterator):
try:
direct_result.append(next(iterator)["id"])
except BaseException as exc:
direct_error.append(exc)
monkeypatch.setattr(dataset, "_resolve_my_splits", controlled_resolve)
direct = iter(dataset)
thread = threading.Thread(
target=advance_direct, args=(direct,), name="direct-start"
)
thread.start()
assert entered.wait(timeout=5)
try:
with pytest.raises(RuntimeError, match="concurrent iteration"):
contender.append(iter(loader))
finally:
release.set()
thread.join(timeout=5)
if contender:
contender[0]._shutdown_workers()
direct.close()
assert not thread.is_alive()
assert direct_error == []
assert direct_result == [0]
def test_loader_acquires_before_snapshot_and_cleans_interrupted_acquire(
tmp_path, monkeypatch
):
db = lancedb.connect(tmp_path)
table = db.create_table("lease_snapshot", pa.table({"id": list(range(4))}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
first = iter(loader)
assert next(first)["id"].tolist() == [0, 1]
entered = threading.Event()
release = threading.Event()
pending = []
pending_errors = []
observed_snapshots = []
real_acquire = dataset._acquire_consumer_iterator
real_snapshot = dataset._checkpoint_snapshot
def controlled_acquire():
if threading.current_thread().name == "stale-start":
entered.set()
assert release.wait(timeout=5)
return real_acquire()
def recording_snapshot():
state = real_snapshot()
if threading.current_thread().name == "stale-start":
observed_snapshots.append(state["samples_consumed_per_split"])
return state
def create_pending_iterator():
try:
pending.append(iter(loader))
except BaseException as exc:
pending_errors.append(exc)
monkeypatch.setattr(dataset, "_acquire_consumer_iterator", controlled_acquire)
monkeypatch.setattr(dataset, "_checkpoint_snapshot", recording_snapshot)
thread = threading.Thread(target=create_pending_iterator, name="stale-start")
thread.start()
assert entered.wait(timeout=5)
assert next(first)["id"].tolist() == [2, 3]
with pytest.raises(StopIteration):
next(first)
release.set()
thread.join(timeout=5)
assert not thread.is_alive()
assert pending_errors == []
assert observed_snapshots == [[4]]
assert len(pending) == 1
assert list(pending[0]) == []
assert dataset.state_dict()["samples_consumed_per_split"] == [4]
def interrupted_acquire():
real_acquire()
raise KeyboardInterrupt("after acquire")
monkeypatch.setattr(dataset, "_acquire_consumer_iterator", interrupted_acquire)
with pytest.raises(KeyboardInterrupt, match="after acquire"):
iter(loader)
assert dataset._consumer_iterator_active is False
def test_consumer_iterator_lease_publication_is_atomic(tmp_path, monkeypatch):
db = lancedb.connect(tmp_path)
table = db.create_table("atomic_lease", pa.table({"id": [0, 1]}))
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
loader = StreamingDataLoader(dataset, batch_size=1, num_workers=0)
real_get_ident = streaming.threading.get_ident
calls = 0
def interrupt_during_publication():
nonlocal calls
calls += 1
if calls == 1:
raise KeyboardInterrupt("during lease mutation")
return real_get_ident()
monkeypatch.setattr(streaming.threading, "get_ident", interrupt_during_publication)
with pytest.raises(KeyboardInterrupt, match="during lease mutation"):
iter(loader)
monkeypatch.setattr(streaming.threading, "get_ident", real_get_ident)
assert dataset._consumer_iterator_active is False
iterator = iter(loader)
try:
assert next(iterator)["id"].tolist() == [0]
finally:
iterator._shutdown_workers()
def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path):
class CustomizedDataset(StreamingDataset):
def __iter__(self):
return iter([1000, 1001])
db = lancedb.connect(tmp_path)
table = db.create_table("custom_iteration", pa.table({"id": [0, 1, 2]}))
dataset = CustomizedDataset(table, num_splits=1, shuffle=False)
assert list(dataset) == [1000, 1001]
with pytest.raises(TypeError, match="override __iter__"):
StreamingDataLoader(
dataset,
batch_size=2,
num_workers=0,
collate_fn=list,
)
def test_interleaved_adapters_do_not_authorize_plain_iteration(tmp_path):
db = lancedb.connect(tmp_path)
table_a = db.create_table("adapter_a", pa.table({"id": [0, 1]}))
table_b = db.create_table("adapter_b", pa.table({"id": [10, 11]}))
dataset_a = StreamingDataset(table_a, num_splits=1, shuffle=False)
dataset_b = StreamingDataset(table_b, num_splits=1, shuffle=False)
initial_state = dataset_a.state_dict()
owner_a = dataset_a._acquire_consumer_iterator()
owner_b = dataset_b._acquire_consumer_iterator()
try:
iterator_a = iter(streaming._StreamingDatasetAdapter(dataset_a))
iterator_b = iter(streaming._StreamingDatasetAdapter(dataset_b))
assert next(iterator_a).data["id"] == 0
assert next(iterator_b).data["id"] == 10
assert [sample.data["id"] for sample in iterator_a] == [1]
assert [sample.data["id"] for sample in iterator_b] == [11]
finally:
dataset_a._release_consumer_iterator(owner_a)
dataset_b._release_consumer_iterator(owner_b)
dataset_a.load_state_dict(initial_state)
with patch(
"lancedb.streaming.get_worker_info",
return_value=FakeWorkerInfo(id=0, num_workers=1),
):
plain_iterator = iter(dataset_a)
assert next(plain_iterator)["id"] == 0
plain_iterator.close()
assert dataset_a._untracked_worker_iteration[0] == 1
with pytest.raises(RuntimeError, match="Use StreamingDataLoader"):
dataset_a.state_dict()
def test_resume_from_partial_split_cycle_preserves_remaining_order(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table("partial_cycle", pa.table({"id": [1, 2, 10, 20]}))
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
iterator = iter(dataset)
assert next(iterator)["id"] == 1
checkpoint = dataset.state_dict()
iterator.close()
assert checkpoint["samples_consumed_per_split"] == [1, 0]
resumed = StreamingDataset(table, num_splits=2, shuffle=False)
resumed.load_state_dict(checkpoint)
assert [row["id"] for row in resumed] == [10, 2, 20]
def test_partial_cycle_resume_preserves_skip_truncation(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table(
"partial_skip", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]})
)
kwargs = dict(
num_splits=2,
shuffle=False,
transform=_failing_transform({1, 2, 3}),
on_transform_error="skip",
)
dataset = StreamingDataset(table, **kwargs)
iterator = iter(dataset)
assert next(iterator)["id"] == 0
checkpoint = dataset.state_dict()
uninterrupted = [row["id"] for row in iterator]
resumed = StreamingDataset(table, **kwargs)
resumed.load_state_dict(checkpoint)
assert [row["id"] for row in resumed] == uninterrupted == [100]
def test_multi_worker_resumability_same_topology(lance_table):
"""Checkpoint with num_workers=2, resume with num_workers=2: exact continuation."""
world_size = 1
@@ -2018,6 +2600,23 @@ def test_merge_state_dicts_validates_consistency(lance_table):
StreamingDataset.merge_state_dicts([])
def test_merge_state_dicts_combines_nonuniform_consumer_progress(lance_table):
dataset = StreamingDataset(
lance_table, num_splits=2, shuffle=False, shuffle_seed=SHUFFLE_SEED
)
rank0 = dataset.state_dict()
rank0["samples_consumed_per_split"] = [2, 0]
rank0["positions_consumed_per_split"] = [2, 0]
rank1 = dataset.state_dict()
rank1["samples_consumed_per_split"] = [0, 2]
rank1["positions_consumed_per_split"] = [0, 2]
merged = StreamingDataset.merge_state_dicts([rank0, rank1])
assert merged["samples_consumed_per_split"] == [2, 2]
assert merged["positions_consumed_per_split"] == [2, 2]
def test_load_state_dict_without_positions_key(lance_table):
"""Checkpoints from before positions_consumed_per_split existed still
resume exactly (positions equal sample counts when nothing is skipped)."""
@@ -2113,6 +2712,273 @@ def test_shuffle_seed_none_generates_stable_seed(lance_table):
assert first == second, "Same resolved seed must produce the same ordering"
# Sequence packing tests
def _create_token_table(tmp_path, documents):
db = lancedb.connect(tmp_path)
tokens = pa.array(documents, type=pa.list_(pa.int64()))
return db.create_table("tokens", pa.table({"tokens": tokens}))
def _packed_dataset(table, pack_sequences, *, blocks_per_epoch, pad_id=0, **kwargs):
return StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
pack_sequences=pack_sequences,
eos_id=9,
pad_id=pad_id,
blocks_per_epoch=blocks_per_epoch,
**kwargs,
)
def test_pack_sequences_emits_blocks_and_pads_final_tail(tmp_path):
table = _create_token_table(tmp_path, [[1, 2], [3, 4], [5]])
dataset = _packed_dataset(table, 6, blocks_per_epoch=2)
blocks = list(dataset)
assert len(blocks) == 2
assert blocks[0]["input_ids"].tolist() == [1, 2, 9, 3, 4, 9]
assert blocks[0]["doc_ids"].tolist() == [0, 0, 0, 1, 1, 1]
assert blocks[1]["input_ids"].tolist() == [5, 9, 0, 0, 0, 0]
assert blocks[1]["doc_ids"].tolist() == [0, 0, 0, 0, 0, 0]
assert blocks[0]["input_ids"].dtype == torch.int64
assert blocks[0]["doc_ids"].dtype == torch.int64
def test_pack_sequences_pads_lagging_splits(tmp_path):
table = _create_token_table(
tmp_path,
[[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]],
)
dataset = _packed_dataset(table, 5, blocks_per_epoch=6, num_splits=2)
input_ids = [block["input_ids"].tolist() for block in dataset]
# Split 0 has four real tokens including EOS markers, while split 1 has
# eleven. Packing must emit three complete two-split cycles.
assert input_ids == [
[1, 9, 2, 9, 0],
[10, 11, 12, 13, 14],
[0, 0, 0, 0, 0],
[15, 16, 17, 9, 20],
[0, 0, 0, 0, 0],
[9, 0, 0, 0, 0],
]
per_rank = []
for rank in range(2):
rank_dataset = _packed_dataset(
table,
5,
blocks_per_epoch=6,
num_splits=2,
world_size=2,
rank=rank,
)
per_rank.append([block["input_ids"].tolist() for block in rank_dataset])
assert [len(blocks) for blocks in per_rank] == [3, 3]
sharded = [block for cycle in zip(*per_rank) for block in cycle]
assert sharded == input_ids
def test_pack_sequences_auto_estimates_filtered_token_column(tmp_path):
db = lancedb.connect(tmp_path)
table = db.create_table(
"tokens",
pa.table(
{
"tokens": pa.array([[1] * 4, [2] * 9], type=pa.list_(pa.int64())),
"keep": [True, False],
}
),
)
table.add(
pa.table(
{
"tokens": pa.array([[3] * 4, [4] * 9], type=pa.list_(pa.int64())),
"keep": [True, False],
}
)
)
with pytest.warns(UserWarning, match="approximate token-count sample"):
dataset = _packed_dataset(
table,
5,
blocks_per_epoch="auto",
num_splits=2,
filter="keep",
)
# Two kept documents contain 8 tokens plus 2 EOS tokens: two blocks.
assert dataset.state_dict()["blocks_per_epoch"] == 2
def test_pack_sequences_checkpoint_resumes_on_new_topology(tmp_path):
table = _create_token_table(
tmp_path,
[[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]],
)
kwargs = dict(pack_sequences=5, blocks_per_epoch=6, num_splits=2)
reference = list(_packed_dataset(table, **kwargs))
datasets = [
_packed_dataset(table, world_size=2, rank=rank, **kwargs) for rank in range(2)
]
iterators = [iter(dataset) for dataset in datasets]
first_cycle = [next(iterator) for iterator in iterators]
checkpoint = StreamingDataset.merge_state_dicts(
[dataset.state_dict() for dataset in datasets]
)
for iterator in iterators:
iterator.close()
resumed = _packed_dataset(table, **kwargs)
resumed.load_state_dict(checkpoint)
actual_remaining = list(resumed)
assert [block["input_ids"].tolist() for block in first_cycle] == [
[1, 9, 2, 9, 0],
[10, 11, 12, 13, 14],
]
assert checkpoint["blocks_emitted_per_split"] == [1, 1]
assert [block["input_ids"].tolist() for block in actual_remaining] == [
block["input_ids"].tolist() for block in reference[2:]
]
assert [block["doc_ids"].tolist() for block in actual_remaining] == [
block["doc_ids"].tolist() for block in reference[2:]
]
def test_packed_checkpoint_requires_complete_split_cycle(tmp_path):
table = _create_token_table(tmp_path, [[1], [2], [10], [20]])
dataset = _packed_dataset(table, pack_sequences=3, blocks_per_epoch=4, num_splits=2)
iterator = iter(dataset)
next(iterator)
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
next(iterator)
assert dataset.state_dict()["blocks_emitted_per_split"] == [1, 1]
iterator.close()
def test_streaming_dataloader_commits_consumed_packed_batches(tmp_path):
table = _create_token_table(
tmp_path,
[[1], [2], [3], [4], [10], [20], [30], [40]],
)
kwargs = dict(pack_sequences=4, blocks_per_epoch=4, num_splits=2)
dataset = _packed_dataset(table, **kwargs)
loader = StreamingDataLoader(
dataset,
batch_size=1,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=2,
)
iterator = iter(loader)
try:
next(iterator)
with pytest.raises(RuntimeError, match="complete logical step boundary"):
dataset.state_dict()
next(iterator)
checkpoint = dataset.state_dict()
uninterrupted = [batch["input_ids"].tolist() for batch in iterator]
finally:
iterator._shutdown_workers()
resumed = _packed_dataset(table, **kwargs)
resumed.load_state_dict(checkpoint)
resumed_loader = StreamingDataLoader(
resumed,
batch_size=1,
num_workers=2,
multiprocessing_context="spawn",
prefetch_factor=2,
)
resumed_iterator = iter(resumed_loader)
try:
remaining = [batch["input_ids"].tolist() for batch in resumed_iterator]
finally:
resumed_iterator._shutdown_workers()
assert checkpoint["blocks_emitted_per_split"] == [1, 1]
assert remaining == uninterrupted
def test_pack_sequences_validates_configuration_and_tokens(tmp_path):
table = _create_token_table(tmp_path, [[1, 2]])
with pytest.raises(ValueError, match="pad_id is required"):
StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
pack_sequences=4,
eos_id=9,
)
with pytest.raises(ValueError, match="blocks_per_epoch is required"):
StreamingDataset(
table,
shuffle=False,
columns=["tokens"],
pack_sequences=4,
eos_id=9,
pad_id=0,
)
with pytest.raises(ValueError, match="must be divisible"):
_packed_dataset(table, 4, blocks_per_epoch=3, num_splits=2)
with pytest.raises(ValueError, match="positive integer or 'auto'"):
_packed_dataset(table, 4, blocks_per_epoch="estimate")
checkpoint = _packed_dataset(table, 4, blocks_per_epoch=1).state_dict()
resumed = _packed_dataset(table, 4, blocks_per_epoch=1, pad_id=8)
with pytest.raises(ValueError, match="pad_id mismatch"):
resumed.load_state_dict(checkpoint)
float_db = lancedb.connect(tmp_path / "float")
float_table = float_db.create_table(
"tokens",
pa.table({"tokens": pa.array([[1.5, 2.5]], type=pa.list_(pa.float64()))}),
)
with pytest.raises(ValueError, match="token column with integer values"):
_packed_dataset(float_table, 4, blocks_per_epoch=1)
null_db = lancedb.connect(tmp_path / "null")
null_table = null_db.create_table(
"tokens",
pa.table({"tokens": pa.array([None], type=pa.list_(pa.int64()))}),
)
with pytest.raises(ValueError, match="does not support null token lists"):
list(_packed_dataset(null_table, 4, blocks_per_epoch=1))
null_value_db = lancedb.connect(tmp_path / "null_value")
null_value_table = null_value_db.create_table(
"tokens",
pa.table(
{"tokens": pa.array([[1], [2, None], [3]], type=pa.list_(pa.int64()))}
),
)
blocks = list(
_packed_dataset(
null_value_table,
2,
blocks_per_epoch=2,
on_transform_error="skip",
)
)
assert [block["input_ids"].tolist() for block in blocks] == [[1, 9], [3, 9]]
# ---------------------------------------------------------------------------
# Doc examples — each test mirrors the code snippet in index.mdx so that
# broken doc examples are caught before they ship.
+21 -21
View File
@@ -52,7 +52,7 @@ class TestExprConstruction:
def test_func(self):
e = func("lower", col("name"))
assert isinstance(e, Expr)
assert e.to_sql() == "lower(name)"
assert e.to_sql() == "lower(`name`)"
def test_func_unknown_raises(self):
with pytest.raises(Exception):
@@ -115,7 +115,7 @@ class TestExprOperators:
def test_and_operator(self):
e = (col("age") > lit(18)) & (col("status") == lit("active"))
assert isinstance(e, Expr)
assert e.to_sql() == "((age > 18) AND (status = 'active'))"
assert e.to_sql() == "((age > 18) AND (`status` = 'active'))"
def test_or_operator(self):
e = (col("a") == lit(1)) | (col("b") == lit(2))
@@ -166,7 +166,7 @@ class TestExprOperators:
def test_coerce_plain_str(self):
e = col("name") == "alice"
assert isinstance(e, Expr)
assert e.to_sql() == "(name = 'alice')"
assert e.to_sql() == "(`name` = 'alice')"
def test_reflexive_comparisons(self):
# 10 < col("age") swaps to col("age") > 10
@@ -198,85 +198,85 @@ class TestExprBytesLiteral:
def test_bytes_equality_expr_sql(self):
e = col("data") == lit(b"\xca\xfe")
assert e.to_sql() == "(data = X'CAFE')"
assert e.to_sql() == "(`data` = X'CAFE')"
def test_bytes_ne_expr_sql(self):
e = col("data") != lit(b"\xff")
assert e.to_sql() == "(data <> X'FF')"
assert e.to_sql() == "(`data` <> X'FF')"
def test_bytes_compound_expr_sql(self):
e = (col("data") == lit(b"\x01")) & (col("id") > lit(5))
assert e.to_sql() == "((data = X'01') AND (id > 5))"
assert e.to_sql() == "((`data` = X'01') AND (id > 5))"
def test_bytes_in_function_call(self):
# Regression test: binary literals inside scalar function calls
# used to fail because DataFusion's unparser does not support Binary
# scalars. Now handled via a placeholder-substitution rewrite.
e = func("contains", col("data"), lit(b"\xff"))
assert e.to_sql() == "contains(data, X'FF')"
assert e.to_sql() == "contains(`data`, X'FF')"
def test_bytes_in_not(self):
e = ~(col("data") == lit(b"\xff"))
assert e.to_sql() == "NOT (data = X'FF')"
assert e.to_sql() == "NOT (`data` = X'FF')"
class TestExprStringMethods:
def test_lower(self):
e = col("name").lower()
assert isinstance(e, Expr)
assert e.to_sql() == "lower(name)"
assert e.to_sql() == "lower(`name`)"
def test_upper(self):
e = col("name").upper()
assert isinstance(e, Expr)
assert e.to_sql() == "upper(name)"
assert e.to_sql() == "upper(`name`)"
def test_contains(self):
e = col("text").contains(lit("hello"))
assert isinstance(e, Expr)
assert e.to_sql() == "contains(text, 'hello')"
assert e.to_sql() == "contains(`text`, 'hello')"
def test_contains_with_str_coerce(self):
e = col("text").contains("hello")
assert isinstance(e, Expr)
assert e.to_sql() == "contains(text, 'hello')"
assert e.to_sql() == "contains(`text`, 'hello')"
def test_chained_lower_eq(self):
e = col("name").lower() == lit("alice")
assert isinstance(e, Expr)
assert e.to_sql() == "(lower(name) = 'alice')"
assert e.to_sql() == "(lower(`name`) = 'alice')"
class TestExprCast:
def test_cast_string(self):
e = col("id").cast("string")
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(id AS VARCHAR)"
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
def test_cast_int32(self):
e = col("score").cast("int32")
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(score AS INTEGER)"
assert e.to_sql() == "arrow_cast(score, 'Int32')"
def test_cast_float64(self):
e = col("val").cast("float64")
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(val AS DOUBLE)"
assert e.to_sql() == "arrow_cast(val, 'Float64')"
def test_cast_pyarrow_type(self):
e = col("score").cast(pa.int32())
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(score AS INTEGER)"
assert e.to_sql() == "arrow_cast(score, 'Int32')"
def test_cast_pyarrow_float64(self):
e = col("val").cast(pa.float64())
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(val AS DOUBLE)"
assert e.to_sql() == "arrow_cast(val, 'Float64')"
def test_cast_pyarrow_string(self):
e = col("id").cast(pa.string())
assert isinstance(e, Expr)
assert e.to_sql() == "CAST(id AS VARCHAR)"
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
def test_cast_pyarrow_and_string_equivalent(self):
# pa.int32() and "int32" should produce equivalent SQL
@@ -597,14 +597,14 @@ class TestExprIsin:
def test_isin_strs(self):
assert (
col("status").isin(["active", "pending"]).to_sql()
== "status IN ('active', 'pending')"
== "`status` IN ('active', 'pending')"
)
def test_isin_coerces_and_mixes(self):
assert col("id").isin([lit(1), 2]).to_sql() == "id IN (1, 2)"
def test_isin_empty(self):
assert col("id").isin([]).to_sql() == "id IN ()"
assert col("id").isin([]).to_sql() == "false"
def test_isin_filter(self, simple_table):
result = simple_table.search().where(col("id").isin([1, 3, 5])).to_arrow()
@@ -37,21 +37,6 @@ def job_result(name: str) -> dict:
return json.loads(fixture(name))["result"]
def assert_no_secret_values(value):
if isinstance(value, dict):
for key, child in value.items():
assert key not in {
"secret_value",
"secret_values",
"resolved_secret",
"resolved_secrets",
}
assert_no_secret_values(child)
elif isinstance(value, list):
for child in value:
assert_no_secret_values(child)
def test_public_function_values_are_in_api_reference():
docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md"
rendered = docs.read_text()
@@ -109,7 +94,6 @@ def test_function_version_identity_is_immutable_and_exact():
version = FunctionVersion.from_json(json.dumps(value))
assert version.name == "embed"
assert version.version == "fv_01K3EXACT"
assert version.required_secrets == ("HF_TOKEN",)
with pytest.raises((TypeError, ValueError)):
version.version = "fv_changed"
@@ -121,7 +105,7 @@ def test_function_version_identity_is_immutable_and_exact():
assert FunctionVersion(**changed) != version
def test_function_version_binds_named_columns_as_one_immutable_group():
def test_function_version_binds_named_columns_as_one_immutable_application():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
@@ -131,13 +115,10 @@ def test_function_version_binds_named_columns_as_one_immutable_group():
assert application.function.name == version.name
assert application.function.version == version.version
assert application.output is version.signature.output
assert application.group_id.startswith("fg_")
assert [
(value.parameter, value.kind, value.value["path"])
for value in application.inputs
] == [("text", "column", "documents.body")]
with pytest.raises((TypeError, ValueError)):
application.group_id = "fg_changed"
def test_function_version_binding_validates_names_and_direct_columns():
@@ -156,7 +137,7 @@ def test_function_version_binding_validates_names_and_direct_columns():
def test_function_version_keeps_named_struct_outputs_in_one_application():
value = job_result("remote_function_job.json")
value["name"] = "text_features"
value["version"] = "fv_grouped"
value["version"] = "fv_multi_output"
value["signature"] = {
"inputs": [
{"name": "title", "arrow_type": "utf8", "nullable": True},
@@ -221,7 +202,6 @@ def test_function_application_uses_rename_columns_only():
assert application.columns["normalized_text"] == "search_text"
assert renamed.columns["normalized_text"] == "body_normalized"
assert renamed.function == application.function
assert renamed.group_id == application.group_id
assert not hasattr(application, "rename_outputs")
with pytest.raises(TypeError, match="immutable"):
renamed.columns["normalized_text"] = "changed"
@@ -242,7 +222,6 @@ def test_function_application_uses_rename_columns_only():
def test_binding_and_refresh_result_keep_stable_remote_fields():
binding = FunctionBinding.from_json(fixture("remote_function_binding.json"))
assert binding.revision == 3
assert binding.function.version == "fv_01K3TEXT"
assert [output.output_ordinal for output in binding.outputs] == [0, 1]
assert binding.input_schema is not None
@@ -297,15 +276,6 @@ def test_refresh_result_rejects_non_u64_values(field):
RefreshColumnResult.from_json(json.dumps(value))
def test_canonical_client_values_contain_secret_names_only():
version = FunctionVersion.from_json(
json.dumps(job_result("remote_function_job.json"))
)
canonical = json.loads(version.to_canonical_json())
assert canonical["required_secrets"] == ["HF_TOKEN"]
assert_no_secret_values(canonical)
class _FunctionDeclarationInner:
def __init__(self):
self.calls = []
@@ -322,7 +292,7 @@ def known_application() -> FunctionApplication:
@pytest.mark.asyncio
async def test_add_columns_routes_struct_as_one_and_grouped_expansion_atomically():
async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomically():
inner = _FunctionDeclarationInner()
table = AsyncTable(inner)
application = known_application()
@@ -343,12 +313,12 @@ async def test_add_columns_routes_struct_as_one_and_grouped_expansion_atomically
@pytest.mark.asyncio
async def test_add_columns_rejects_mixed_groups_and_unknown_newer_application():
async def test_add_columns_rejects_multiple_bindings_and_unknown_newer_application():
inner = _FunctionDeclarationInner()
table = AsyncTable(inner)
application = known_application()
with pytest.raises(ValueError, match="exactly one Function sibling group"):
with pytest.raises(ValueError, match="exactly one Function binding"):
await table.add_columns({"a": application, "b": application})
future = json.loads(fixture("remote_function_application.json"))
@@ -376,7 +346,6 @@ def test_rename_requires_named_struct_and_keeps_partial_mapping_immutable():
"arrow_type": "list<float32>",
"nullable": False,
},
"group_id": "fg_scalar",
}
)
)
@@ -3,7 +3,12 @@
from __future__ import annotations
import base64
import contextlib
import functools
import importlib.util
import types
from datetime import date
import http.server
import json
from pathlib import Path
@@ -16,6 +21,9 @@ import pytest
import lancedb
from lancedb.functions import UdfDefinition, udf
THRESHOLD = 20
_CACHE = None
FIXTURES = (
Path(__file__).parents[3]
@@ -31,28 +39,12 @@ FIXTURES = (
@udf(
pip=["numpy>=2"],
env={"MODE": "test"},
secrets=["API_TOKEN"],
python_version="3.12",
)
def normalize_score(value: float) -> float:
return value / 100.0
def _assert_no_secret_values(value):
if isinstance(value, dict):
for key, child in value.items():
assert key not in {
"secret_value",
"secret_values",
"resolved_secret",
"resolved_secrets",
}
_assert_no_secret_values(child)
elif isinstance(value, list):
for child in value:
_assert_no_secret_values(child)
def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
assert isinstance(normalize_score, UdfDefinition)
assert normalize_score(25.0) == 0.25
@@ -67,13 +59,417 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
"kind": "scalar_to_arrow_batch",
"version": 1,
}
assert request["required_secrets"] == ["API_TOKEN"]
_assert_no_secret_values(request)
def _run_packaged(definition, *args):
"""Execute the shipped artifact in a fresh namespace, as a worker would."""
source = base64.b64decode(definition.registration_request.artifact.content.data)
namespace: dict = {}
exec(compile(source, "<udf>", "exec"), namespace)
return namespace[definition.registration_request.artifact.entrypoint](*args)
def test_udf_conda_environment():
@udf(conda=["scipy", "numpy"], conda_channels=["conda-forge", "defaults"])
def halve(value: float) -> float:
return value / 2
request = json.loads(halve.registration_request.to_canonical_json())
assert request["runtime"]["environment"] == {
"kind": "conda",
"packages": ["numpy", "scipy"],
"channels": ["conda-forge", "defaults"],
}
pip_request = json.loads(normalize_score.registration_request.to_canonical_json())
assert "channels" not in pip_request["runtime"]["environment"]
with pytest.raises(ValueError, match="not both"):
udf(name="both", pip=["numpy"], conda=["numpy"])(lambda value: value)
with pytest.raises(ValueError, match="requires conda"):
udf(name="channels", conda_channels=["conda-forge"])(lambda value: value)
def test_udf_packages_attribute_access_and_body_imports():
@udf
def word_norm(body: str) -> float:
import numpy as np
try:
words = body.split()
except AttributeError as error:
raise ValueError(str(error)) from error
return float(np.linalg.norm([len(w) for w in words]))
assert _run_packaged(word_norm, "aa bb") == pytest.approx(8**0.5)
def test_udf_packages_module_globals_and_global_caches():
@udf
def label(value: int) -> str:
return "big" if value >= THRESHOLD else "small"
assert _run_packaged(label, 21) == "big"
@udf
def cached(value: int) -> int:
global _CACHE
if _CACHE is None:
_CACHE = 40
return _CACHE + value
assert _run_packaged(cached, 2) == 42
def test_udf_annotations_are_not_runtime_names():
@udf
def identity(value: date) -> date:
return value
assert _run_packaged(identity, date(2026, 8, 25)) == date(2026, 8, 25)
def test_udf_nested_scopes_resolve_lexically():
@udf
def score(value: int) -> int:
offset = 2
def add_offset() -> int:
return value + offset
return add_offset() + sum(v for v in [0])
assert _run_packaged(score, 3) == 5
def test_udf_resolves_module_globals_before_builtins(tmp_path):
module_path = tmp_path / "shadowing_udfs.py"
module_path.write_text(
"max = 7\n"
"len = lambda _: 99\n"
"\n"
"def uses_literal_shadow(value: int) -> int:\n"
" def nested() -> int:\n"
" return max\n"
" return nested() + value\n"
"\n"
"def uses_callable_shadow(value: int) -> int:\n"
" def nested() -> int:\n"
" return len([1])\n"
" return nested() + value\n"
)
spec = importlib.util.spec_from_file_location("shadowing_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
# The module's `max = 7` is what the interpreter would use, so it ships.
assert _run_packaged(udf(module.uses_literal_shadow), 1) == 8
# A callable global cannot ship; it must not be silently swapped for the builtin.
with pytest.raises(TypeError, match="unsupported global value of type function"):
udf(module.uses_callable_shadow)
def test_canonical_arrow_type_is_exactly_the_grammar():
from lancedb.functions import _GRAMMAR_PRIMITIVES, _canonical_arrow_type
golden = json.loads(
(
Path(__file__).parents[3]
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
).read_text()
)
primitives = [
case["arrow_type"] for case in golden["valid"] if "<" not in case["arrow_type"]
]
assert [name for _, name in _GRAMMAR_PRIMITIVES] == primitives
for outside in [
pa.timestamp("us"),
pa.decimal128(10, 2),
pa.large_string(),
pa.large_binary(),
pa.binary(4),
pa.duration("s"),
pa.struct([pa.field("a", pa.int32())]),
pa.list_(pa.float32(), 0),
pa.list_(pa.timestamp("us")),
]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(outside)
def test_udf_nested_annotations_are_postponed_in_the_artifact():
@udf
def score(value: int) -> int:
def identity(item: date) -> date:
return item
identity(date(2026, 8, 25))
return value
assert _run_packaged(score, 3) == 3
def test_udf_ships_globals_the_body_deletes():
@udf
def clear(value: int) -> int:
global _CACHE
del _CACHE
return value
assert _run_packaged(clear, 3) == 3
def test_udf_rejects_a_module_global_that_does_not_import_as_itself(tmp_path):
module_path = tmp_path / "fake_module_udfs.py"
module_path.write_text(
"import types\n"
"np = types.ModuleType('numpy')\n"
"np.sqrt = lambda x: 0\n"
"\n"
"def score(value: int) -> int:\n"
" return int(np.sqrt(value))\n"
)
spec = importlib.util.spec_from_file_location("fake_module_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
with pytest.raises(TypeError, match="does not import as 'numpy'"):
udf(module.score)
def test_udf_rejects_a_module_level_namespace_alias(tmp_path):
module_path = tmp_path / "aliasing_udfs.py"
module_path.write_text(
"import builtins as b\n"
"THRESHOLD = 5\n"
"\n"
"def score(value: int) -> int:\n"
" return value + b.vars(b.__import__('aliasing_udfs'))['THRESHOLD']\n"
)
spec = importlib.util.spec_from_file_location("aliasing_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
with pytest.raises(ValueError, match="dynamic namespace access"):
udf(module.score)
@pytest.mark.parametrize(
"access",
[
"globals()['THRESHOLD']",
"eval('THRESHOLD')",
"(lambda g: g()['THRESHOLD'])(globals)",
"__import__('sys').modules[__name__].THRESHOLD",
"sys.modules[__name__].THRESHOLD",
],
)
def test_udf_rejects_dynamic_namespace_access(access):
namespace: dict = {}
exec(
f"def score(value: int) -> int:\n return value + {access}\n",
{"THRESHOLD": 5},
namespace,
)
with pytest.raises(ValueError, match="dynamic namespace access"):
_package_from_text(
"def score(value: int) -> int:\n"
" import sys\n"
f" return value + {access}\n"
)
def _package_from_text(source: str, module_globals: dict | None = None):
"""Load `source` as a real module file so the packager can inspect it."""
import tempfile
directory = tempfile.mkdtemp()
path = Path(directory) / "generated_udf_module.py"
path.write_text(source)
spec = importlib.util.spec_from_file_location(f"generated_udf_{id(source)}", path)
module = importlib.util.module_from_spec(spec)
if module_globals:
module.__dict__.update(module_globals)
spec.loader.exec_module(module)
functions = [
value
for value in vars(module).values()
if callable(value) and getattr(value, "__module__", None) == module.__name__
]
return udf(functions[0])
def test_udf_rejects_a_non_standard_builtins_environment():
def score(value: int) -> int:
return len([1]) + value
score.__globals__ # noqa: B018 -- real function, real globals
import builtins
patched = types.FunctionType(
score.__code__,
{"__builtins__": {**vars(builtins), "len": lambda _: 99}},
"score",
)
patched.__annotations__ = score.__annotations__
assert patched(3) == 102
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(patched)
class ReportingDict(dict): # reports standard entries, resolves differently
def __missing__(self, key):
return vars(builtins)[key]
disguised = types.FunctionType(
score.__code__, {"__builtins__": ReportingDict(len=lambda _: 99)}, "score"
)
disguised.__annotations__ = score.__annotations__
assert disguised(3) == 102
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(disguised)
hooked = types.FunctionType(
score.__code__,
{"__builtins__": {**vars(builtins), "__import__": lambda *a, **k: None}},
"score",
)
hooked.__annotations__ = score.__annotations__
with pytest.raises(ValueError, match="non-standard builtins environment"):
udf(hooked)
def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
module_path = tmp_path / "rebound_udfs.py"
module_path.write_text(
"def fact(value: int) -> int:\n"
" return 1 if value <= 1 else value * fact(value - 1)\n"
"\n"
"def score(value: int) -> int:\n"
" return score + value\n"
)
spec = importlib.util.spec_from_file_location("rebound_udfs", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
assert _run_packaged(udf(module.fact), 5) == 120
raw = module.score
module.score = 10
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw)
# A wrapper that merely exposes __wrapped__ is not the function.
module.score = functools.wraps(raw)(lambda value: 41)
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw)
# The decorator's own result is; a subclass of it is not.
module.fact = udf(module.fact)
assert _run_packaged(module.fact, 4) == 24
class Twisted(UdfDefinition):
def __call__(self, *args, **kwargs):
return 41
raw_fact = module.fact._function
module.fact = Twisted(
raw_fact,
name=None,
input_schema=None,
output_schema=None,
pip=(),
env={},
python_version=None,
)
with pytest.raises(ValueError, match="binds that name to another value"):
udf(raw_fact)
def test_canonical_arrow_type_rejects_unrepresentable_list_children():
from lancedb.functions import _canonical_arrow_type
for outside in [
pa.list_(pa.float32()), # pyarrow default: nullable child
pa.list_(pa.field("custom", pa.float32(), nullable=False)),
pa.list_(pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"})),
pa.list_(pa.field("item", pa.float32(), nullable=False), 0),
]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(outside)
assert (
_canonical_arrow_type(
pa.list_(pa.field("item", pa.float32(), nullable=False), 3)
)
== "fixed_size_list<float32, 3>"
)
def _calls_missing(value: int) -> int:
return missing(value) # noqa: F821
def _shadows_missing_in_a_comprehension(value: int) -> int:
return missing(value) + sum(missing for missing in ()) # noqa: F821
def _shadows_missing_in_a_lambda(value: int) -> int:
return (lambda missing: missing)(value) + missing # noqa: F821
@pytest.mark.parametrize(
"function",
[_calls_missing, _shadows_missing_in_a_comprehension, _shadows_missing_in_a_lambda],
)
def test_udf_rejects_a_truly_unresolved_global(function):
with pytest.raises(ValueError, match=r"unresolved global names: \['missing'\]"):
udf(function)
def _arrow_type_from_golden(spec: dict) -> pa.DataType:
kind = spec["type"]
if kind in ("list", "large_list", "fixed_size_list"):
item = _arrow_type_from_golden(spec["fields"][0]["type"])
field = pa.field("item", item, nullable=False)
if kind == "list":
return pa.list_(field)
if kind == "large_list":
return pa.large_list(field)
return pa.list_(field, spec["length"])
return {
"null": pa.null(),
"bool": pa.bool_(),
"utf8": pa.string(),
"binary": pa.binary(),
"float16": pa.float16(),
"float32": pa.float32(),
"float64": pa.float64(),
"date32": pa.date32(),
"date64": pa.date64(),
}.get(kind) or getattr(pa, kind)()
def test_arrow_type_grammar_matches_the_shared_golden():
golden = json.loads(
(
Path(__file__).parents[3]
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
).read_text()
)
from lancedb.functions import _canonical_arrow_type
emitted = {
case["arrow_type"]: _canonical_arrow_type(_arrow_type_from_golden(case["json"]))
for case in golden["valid"]
}
assert emitted == {
case["arrow_type"]: case["arrow_type"] for case in golden["valid"]
}
assert not set(emitted) & set(golden["invalid"])
for case in golden["server_only"]:
with pytest.raises(TypeError, match="unsupported Arrow type"):
_canonical_arrow_type(_arrow_type_from_golden(case["json"]))
def test_explicit_arrow_schema_is_deterministic():
input_schema = pa.schema([pa.field("value", pa.float32(), nullable=True)])
output_schema = pa.field("embedding", pa.list_(pa.float32(), 3), nullable=False)
output_schema = pa.field(
"embedding",
pa.list_(pa.field("item", pa.float32(), nullable=False), 3),
nullable=False,
)
@udf(input_schema=input_schema, output_schema=output_schema)
def explicit(value):
@@ -82,7 +478,7 @@ def test_explicit_arrow_schema_is_deterministic():
signature = explicit.registration_request.signature
assert signature.inputs[0].arrow_type == "float32"
assert signature.inputs[0].nullable is True
assert signature.output.arrow_type == "fixed_size_list<float32>[3]"
assert signature.output.arrow_type == "fixed_size_list<float32, 3>"
assert signature.output.nullable is False
@@ -130,14 +526,6 @@ def test_annotation_and_explicit_schema_validation_fail_closed():
return value
def test_environment_rejects_secret_value_overlap():
with pytest.raises(ValueError, match="must be disjoint"):
@udf(env={"TOKEN": "plaintext"}, secrets=["TOKEN"])
def overlapping(value: int) -> int:
return value
def test_local_function_catalog_operations_are_not_supported(tmp_path):
db = lancedb.connect(tmp_path)
message = "Function catalog operations are not supported by this database"
@@ -174,7 +562,6 @@ def _mock_remote_function_catalog():
"runtime": body["runtime"],
"runtime_digest": "sha256:runtime",
"environment_digest": "sha256:environment",
"required_secrets": body.get("required_secrets", []),
"created_at": "2026-08-21T00:00:00Z",
}
response = {"job_id": "job-register"}
@@ -187,7 +574,7 @@ def _mock_remote_function_catalog():
"job_state": "DONE",
"result": state["version"],
}
elif self.path == "/v1/functions/get":
elif self.path == "/v1/functions/describe":
assert body == {
"name": "normalize_score",
"version": "fv_exact",
@@ -233,7 +620,6 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
assert create_request == json.loads(
normalize_score.registration_request.to_canonical_json()
)
_assert_no_secret_values(create_request)
def test_blocking_remote_registration_returns_function_version():
+77
View File
@@ -25,6 +25,7 @@ from lancedb.db import DBConnection
from lancedb.index import FTS
from lancedb.query import (
BoostQuery,
DocumentGranularity,
MatchQuery,
MultiMatchQuery,
PhraseQuery,
@@ -245,6 +246,55 @@ def test_create_inverted_index_rejects_invalid_block_size(table):
table.create_index("text", config=FTS(block_size=129))
def test_list_element_document_granularity(tmp_path):
docs_type = pa.list_(pa.struct([pa.field("content", pa.string())]))
docs = pa.array(
[
[
{"content": "alpha beta"},
None,
{"content": ""},
{"content": "the and"},
{"content": "alpha beta"},
]
],
type=docs_type,
)
table = ldb.connect(tmp_path).create_table(
"list_element_docs", pa.table({"id": [0], "docs": docs})
)
row_table = ldb.connect(tmp_path).create_table(
"row_docs", pa.table({"id": [0], "docs": docs})
)
row_table.create_index("docs.content", config=FTS())
row_result = row_table.search(MatchQuery("alpha", "docs.content")).to_arrow()
assert row_result.num_rows == 1
assert "_doc_index" not in row_result.column_names
granularity = DocumentGranularity.LIST_ELEMENT
table.create_index(
"docs.content",
config=FTS(with_position=True, document_granularity=granularity),
)
assert table.list_indices()[0].columns == ["docs.content"]
def coordinates(query):
result = table.search(query).limit(10).to_arrow()
doc_index_type = result.schema.field("_doc_index").type
assert pa.types.is_list(doc_index_type)
assert doc_index_type.value_type == pa.uint32()
return sorted(result["_doc_index"].to_pylist())
assert coordinates(
MatchQuery("alpha", "docs.content", document_granularity=granularity)
) == [[0], [4]]
assert coordinates(
PhraseQuery("alpha beta", "docs.content", document_granularity=granularity)
) == [[0], [4]]
assert coordinates(MatchQuery("alpha", "docs.content")) == [[0], [4]]
assert FTS().document_granularity is DocumentGranularity.ROW
def test_create_inverted_index_respects_build_memory_limit(table):
with pytest.raises(ValueError, match="exceeds worker memory limit"):
table.create_index(
@@ -1089,6 +1139,20 @@ def test_fts_query_to_json():
)
assert json_str == expected
# Test MatchQuery with list-element document granularity
match_query = MatchQuery(
"hello world",
"text",
document_granularity=DocumentGranularity.LIST_ELEMENT,
)
json_str = match_query.to_json()
expected = (
'{"match":{"column":"text","terms":"hello world","boost":1.0,'
'"fuzziness":0,"max_expansions":50,"operator":"Or","prefix_length":0,'
'"document_granularity":"list_element"}}'
)
assert json_str == expected
# Test MatchQuery with options
match_query = MatchQuery("puppy", "text", fuzziness=2, boost=1.5, prefix_length=3)
json_str = match_query.to_json()
@@ -1098,6 +1162,19 @@ def test_fts_query_to_json():
)
assert json_str == expected
# Test PhraseQuery with list-element document granularity
phrase_query = PhraseQuery(
"quick brown fox",
"title",
document_granularity=DocumentGranularity.LIST_ELEMENT,
)
json_str = phrase_query.to_json()
expected = (
'{"phrase":{"column":"title","terms":"quick brown fox","slop":0,'
'"document_granularity":"list_element"}}'
)
assert json_str == expected
# Test PhraseQuery
phrase_query = PhraseQuery("quick brown fox", "title")
json_str = phrase_query.to_json()
+1 -1
View File
@@ -88,7 +88,7 @@ async def binary_table(db_async):
async def test_create_index_async_returns_done_job(some_table: AsyncTable):
job = await some_table.create_index_async("id", config=BTree())
assert job.id is None
await job.wait()
assert await job.wait() is None
assert len(await some_table.list_indices()) == 1
await job.cancel()
+25
View File
@@ -56,6 +56,31 @@ def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
assert permutation_tbl._conn.read_consistency_interval is None
def test_pickled_permutation_reads_pinned_version(tmp_path):
"""An unpickled copy must still read the pinned version, which also covers the
version surviving the ``to_arrow()`` round trip in ``__getstate__``."""
import pickle
db = connect(tmp_path)
tbl = db.create_table("base", pa.table({"idx": range(20)}))
permutation_tbl = permutation_builder(tbl).execute()
perm = Permutation.from_tables(tbl, permutation_tbl)
payload = pickle.dumps(perm)
# Compact so the stored row addresses no longer describe these rows at latest.
tbl.delete("true")
tbl.optimize()
assert tbl.count_rows() == 0
# Unpickle after the mutation: __setstate__ reopens at latest, so this only
# passes if the recorded version is applied on reopen.
restored = pickle.loads(payload)
assert len(restored) == 20
rows = restored.__getitems__(list(range(20)))
assert sorted(row["idx"] for row in rows) == list(range(20))
def test_split_random_counts(mem_db):
"""Test random splitting with absolute counts."""
tbl = mem_db.create_table(
+32
View File
@@ -675,6 +675,21 @@ def test_distance_range(table: lancedb.table.Table):
assert res["_distance"].to_pylist() == [min_dist, max_dist]
@pytest.mark.parametrize("expression", ["1 - _distance", "1.0 - _distance"])
def test_select_arithmetic_with_distance(table, expression):
result = (
table.search([10, 10])
.select({"similarity": expression, "_distance": "_distance"})
.distance_type("cosine")
.to_arrow()
)
assert result.schema.field("similarity").type == pa.float32()
assert result["similarity"].to_pylist() == pytest.approx(
[1 - distance for distance in result["_distance"].to_pylist()]
)
@pytest.mark.asyncio
async def test_distance_range_async(table_async: AsyncTable):
q = [0, 0]
@@ -897,6 +912,23 @@ def test_query_builder_batches(table):
assert rs_list["id"][1] == 2
def test_batch_vector_query_shares_filtered_flat_scan(table):
query = (
table.search([[1.0, 2.0], [3.0, 4.0]])
.where("id > 0", prefilter=True)
.limit(1)
.select(["id"])
)
plan = query.explain_plan(verbose=True)
assert "KNNVectorDistance: queries=2" in plan
assert "UnionExec" not in plan
results = query.to_arrow()
assert len(results) == 2
assert results["query_index"].to_pylist() == [0, 1]
def test_dynamic_projection(table):
rs = (
LanceVectorQueryBuilder(table, [0, 0], "vector")
+118 -1
View File
@@ -875,11 +875,85 @@ def test_remote_create_index_async_returns_job():
table = db.create_table("test", [{"id": 1}])
job = table.create_index_async("id", config=BTree())
assert job.id == "job-1"
job.wait(timeout=timedelta(seconds=30))
assert job.wait(timeout=timedelta(seconds=30)) is None
assert len(describe_calls) == 2
job.cancel()
def test_remote_refresh_async_returns_typed_terminal_result():
terminal_result = {
"rows_assigned": 12,
"rows_failed": 0,
"rows_remaining": 0,
"source_version": 7,
"published_version": 8,
}
def handler(request):
content_len = int(request.headers.get("Content-Length", 0))
body = request.rfile.read(content_len) if content_len > 0 else b""
if request.path == "/v1/table/test/backfill_column":
assert json.loads(body)["column"] == "derived"
request.send_response(202)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"job_id": "refresh-1"}')
elif request.path == "/v1/jobs/describe":
assert json.loads(body)["job_id"] == "refresh-1"
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
{
"job_id": "refresh-1",
"job_type": "function_refresh",
"job_state": "DONE",
"result": terminal_result,
}
).encode()
)
elif request.path == "/v1/table/test/create/?mode=create":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b"{}")
elif request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(
json.dumps(
{
"version": 1,
"schema": {
"fields": [
{
"name": "id",
"type": {"type": "int64"},
"nullable": False,
}
]
},
}
).encode()
)
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
table = db.create_table("test", [{"id": 1}])
job = table.refresh_column_async("derived")
assert job.id == "refresh-1"
result = job.wait(timeout=timedelta(seconds=30))
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.model_dump() == terminal_result
assert result.rows_filled == 12
assert result.version == 8
def test_remote_job_wait_raises_on_failure():
from lancedb.exceptions import JobFailedError
from lancedb.index import BTree
@@ -1544,6 +1618,49 @@ def test_query_sync_fts():
)
def test_query_sync_fts_document_granularity():
from lancedb.query import DocumentGranularity, MatchQuery
def handler(body):
assert body == {
"full_text_query": {
"query": {
"match": {
"column": "docs.content",
"terms": "alpha",
"boost": 1.0,
"fuzziness": 0,
"max_expansions": 50,
"operator": "Or",
"prefix_length": 0,
"document_granularity": "list_element",
}
}
},
"k": 10,
"prefilter": True,
"vector": [],
"version": None,
}
return pa.table(
{
"id": [1, 1],
"_doc_index": pa.array([[0], [4]], type=pa.list_(pa.uint32())),
}
)
with query_test_table(handler, server_version=Version("0.6.0")) as table:
result = table.search(
MatchQuery(
"alpha",
"docs.content",
document_granularity=DocumentGranularity.LIST_ELEMENT,
)
).to_arrow()
assert result["_doc_index"].to_pylist() == [[0], [4]]
def test_query_sync_hybrid():
def handler(body):
if "full_text_query" in body:
+20
View File
@@ -4,6 +4,7 @@
import asyncio
import copy
from concurrent.futures import ThreadPoolExecutor
from datetime import timedelta
import threading
@@ -86,6 +87,25 @@ def test_s3_lifecycle(s3_bucket: str):
asyncio.run(test())
@pytest.mark.s3_test
def test_concurrent_open_table(s3_bucket: str):
uri = f"s3://{s3_bucket}/test_concurrent_open_table"
db = lancedb.connect(uri, storage_options=copy.copy(CONFIG))
db.create_table("test", pa.table({"x": [1, 2, 3]}))
num_workers = 32
barrier = threading.Barrier(num_workers)
def open_and_count(_):
barrier.wait()
return db.open_table("test").count_rows()
with ThreadPoolExecutor(max_workers=num_workers) as pool:
row_counts = list(pool.map(open_and_count, range(num_workers)))
assert row_counts == [3] * num_workers
@pytest.fixture()
def kms_key():
kms = get_boto3_client("kms", endpoint_url=CONFIG["aws_endpoint"])
+176 -3
View File
@@ -11,6 +11,7 @@ import warnings
import weakref
from concurrent.futures import ThreadPoolExecutor
from datetime import date, datetime, timedelta
from decimal import Decimal
from time import sleep
from typing import List
from unittest.mock import patch
@@ -336,6 +337,21 @@ async def test_update_async(mem_db_async: AsyncConnection):
assert await table.count_rows("id == 10") == 1
@pytest.mark.asyncio
async def test_update_expr_filter_literals_async(mem_db_async: AsyncConnection):
values = ["5", "4.66e-84", "it's"]
table = await mem_db_async.create_table(
"update_expr_literals",
data=[{"field": value, "result": "original"} for value in values],
)
for value in values:
update_res = await table.update({"result": value}, where=col("field") == value)
assert update_res.rows_updated == 1
assert (await table.to_arrow())["result"].to_pylist() == values
def test_create_table(mem_db: DBConnection):
schema = pa.schema(
{
@@ -1467,7 +1483,7 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection):
table = mem_db.create_table("job_test", [{"id": i} for i in range(10)])
job = table.create_index_async("id", config=BTree())
assert job.id is None
job.wait()
assert job.wait() is None
assert len(table.list_indices()) == 1
job.cancel()
@@ -2343,6 +2359,148 @@ def test_update(mem_db: DBConnection):
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
def test_update_expr_filter_literals(mem_db: DBConnection):
values = ["5", "4.66e-84", "it's"]
table = mem_db.create_table(
"update_expr_literals",
data=[{"field": value, "result": "original"} for value in values],
)
for value in values:
update_res = table.update(where=col("field") == value, values={"result": value})
assert update_res.rows_updated == 1
assert table.to_arrow()["result"].to_pylist() == values
def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection):
low = Decimal("1.234567890123456789")
high = Decimal("1.234567890123456790")
decimal_schema = pa.schema(
[("val", pa.decimal128(19, 18)), ("result", pa.string())]
)
decimal_table = mem_db.create_table(
"update_expr_decimal",
pa.table(
{"val": [low, high], "result": ["old", "old"]},
schema=decimal_schema,
),
)
predicate = col("val") < lit(high)
assert decimal_table.search().where(predicate).to_arrow().num_rows == 1
result = decimal_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
keyword_table = mem_db.create_table(
"update_expr_keyword", [{"null": 1, "result": "old"}]
)
predicate = col("null") == 1
assert keyword_table.search().where(predicate).to_arrow().num_rows == 1
result = keyword_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
empty_in_table = mem_db.create_table(
"update_expr_empty_in", [{"id": 1, "result": "old"}]
)
predicate = col("id").isin([])
assert empty_in_table.search().where(predicate).to_arrow().num_rows == 0
result = empty_in_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 0
marker = "__lancedb_binary_placeholder_0__"
binary_schema = pa.schema(
[("payload", pa.binary()), ("text", pa.string()), ("result", pa.string())]
)
binary_table = mem_db.create_table(
"update_expr_binary",
pa.table(
{
"payload": [b"\x01", b"\x02"],
"text": ["other", marker],
"result": ["old", "old"],
},
schema=binary_schema,
),
)
predicate = (col("payload") == lit(b"\x01")) | (col("text") == marker)
assert binary_table.search().where(predicate).to_arrow().num_rows == 2
result = binary_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 2
nonfinite_table = mem_db.create_table(
"update_expr_nonfinite",
[{"x": 1.0, "result": "old"}, {"x": 2.0, "result": "old"}],
)
predicate = col("x") < float("inf")
assert nonfinite_table.search().where(predicate).to_arrow().num_rows == 2
result = nonfinite_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 2
float16_table = mem_db.create_table(
"update_expr_float16",
[{"x": 1.0, "result": "old"}, {"x": 3.0, "result": "old"}],
)
predicate = col("x").cast(pa.float16()) < 2.0
assert float16_table.search().where(predicate).to_arrow().num_rows == 1
result = float16_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
string_cast_table = mem_db.create_table(
"update_expr_string_cast",
[{"x": 1, "result": "old"}, {"x": 2, "result": "old"}],
)
predicate = col("x").cast("string") == "1"
assert string_cast_table.search().where(predicate).to_arrow().num_rows == 1
result = string_cast_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
quoted_identifier_schema = pa.schema(
[("payload", pa.binary()), ("odd'name", pa.int64()), ("result", pa.string())]
)
quoted_identifier_table = mem_db.create_table(
"update_expr_quoted_identifier",
pa.table(
{"payload": [b"\x01"], "odd'name": [1], "result": ["old"]},
schema=quoted_identifier_schema,
),
)
predicate = (col("payload") == lit(b"\x01")) & (col("odd'name") == 1)
assert quoted_identifier_table.search().where(predicate).to_arrow().num_rows == 1
result = quoted_identifier_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
decimal256_schema = pa.schema(
[("val", pa.decimal256(40, 2)), ("result", pa.string())]
)
decimal256_table = mem_db.create_table(
"update_expr_decimal256",
pa.table(
{
"val": [Decimal("1.00"), Decimal("3.00")],
"result": ["old", "old"],
},
schema=decimal256_schema,
),
)
predicate = col("val") < lit(Decimal("2.00")).cast(pa.decimal256(40, 2))
assert decimal256_table.search().where(predicate).to_arrow().num_rows == 1
result = decimal256_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 1
binary_empty_table = mem_db.create_table(
"update_expr_binary_empty",
pa.table(
{"payload": [b"\x01", b"\x02"], "result": ["old", "old"]},
schema=pa.schema([("payload", pa.binary()), ("result", pa.string())]),
),
)
predicate = (col("payload") == lit(b"\x01")).isin([])
assert binary_empty_table.search().where(predicate).to_arrow().num_rows == 0
assert predicate.to_sql() == "false"
result = binary_empty_table.update(where=predicate, values={"result": "new"})
assert result.rows_updated == 0
def test_update_with_arrow_scalar(mem_db: DBConnection):
schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})
table = mem_db.create_table("my_table", schema=schema)
@@ -3947,10 +4105,21 @@ def test_refresh_column_async_returns_job(tmp_path):
job = table.refresh_column_async("doubled")
assert job.id is None # in-process jobs have no server id
assert job.wait() is None
result = job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 2
assert result.rows_failed == 0
assert result.rows_remaining == 0
assert result.source_version == 2
assert result.published_version == 3
assert job.status() == "finished"
assert sorted(table.to_arrow()["doubled"].to_pylist()) == [2, 4]
no_op = table.refresh_column_async("doubled").wait()
assert no_op.rows_assigned == 0
assert no_op.source_version == 3
assert no_op.published_version is None
# Bad input raises at the call, not through the job.
with pytest.raises(Exception, match="not a computed column"):
table.refresh_column_async("x")
@@ -3963,6 +4132,10 @@ async def test_refresh_column_async_job_async_table(tmp_path):
await table.add_columns(computed={"tripled": "x * 3"})
job = await table.refresh_column_async("tripled")
assert await job.wait() is None
result = await job.wait()
assert isinstance(result, lancedb.RefreshColumnResult)
assert result.rows_assigned == 1
assert result.source_version == 2
assert result.published_version == 3
assert await job.status() == "finished"
assert (await table.to_arrow())["tripled"].to_pylist() == [9]
+1 -1
View File
@@ -609,7 +609,7 @@ impl Connection {
.create_function_async(request)
.await
.infer_error()
.map(crate::job::FunctionJob::new)
.map(crate::job::Job::new_typed)
})
}
+8
View File
@@ -130,6 +130,14 @@ impl PyExpr {
// ── utilities ────────────────────────────────────────────────────────────
/// Return the referenced column name for a bare column expression.
fn column_name(&self) -> Option<String> {
match &self.0 {
DfExpr::Column(column) if column.relation.is_none() => Some(column.name.clone()),
_ => None,
}
}
/// Render the expression as a SQL string (useful for debugging).
fn to_sql(&self) -> PyResult<String> {
lancedb::expr::expr_to_sql_string(&self.0).map_err(|e| PyValueError::new_err(e.to_string()))
+8 -2
View File
@@ -8,7 +8,7 @@ use lancedb::index::vector::{
};
use lancedb::index::{
Index as LanceDbIndex,
scalar::{BTreeIndexBuilder, FmIndexBuilder, FtsIndexBuilder},
scalar::{BTreeIndexBuilder, DocumentGranularity, FmIndexBuilder, FtsIndexBuilder},
};
use pyo3::IntoPyObject;
use pyo3::types::PyStringMethods;
@@ -60,7 +60,11 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
.ngram_min_length(params.ngram_min_length)
.ngram_max_length(params.ngram_max_length)
.ngram_prefix_only(params.prefix_only)
.custom_stop_words(params.custom_stop_words);
.custom_stop_words(params.custom_stop_words)
.document_granularity(
DocumentGranularity::try_from(params.document_granularity.as_str())
.map_err(|err| PyValueError::new_err(err.to_string()))?,
);
if let Some(memory_limit) = params.memory_limit {
inner_opts = inner_opts.memory_limit_mb(memory_limit);
}
@@ -221,6 +225,7 @@ struct FtsParams {
block_size: usize,
memory_limit: Option<u64>,
num_workers: Option<usize>,
document_granularity: String,
}
#[derive(FromPyObject)]
@@ -481,6 +486,7 @@ mod tests {
block_size = 128
memory_limit = 2048
num_workers = 7
document_granularity = 'row'
config = FTS()",
None,
+18 -55
View File
@@ -5,72 +5,33 @@ use std::sync::Arc;
use crate::runtime::future_into_py;
use pyo3::{Bound, PyAny, PyRef, PyResult, pyclass, pymethods};
use serde::Serialize;
use crate::error::PythonErrorExt;
#[pyclass]
pub struct Job {
inner: Arc<lancedb::Job>,
}
/// Python bridge for a typed remote Function registration job.
///
/// The public Python layer decodes the canonical JSON returned by `wait`
/// into its immutable `FunctionVersion` model.
#[pyclass]
pub struct FunctionJob {
inner: Arc<lancedb::Job<lancedb::function::FunctionVersion>>,
}
impl FunctionJob {
pub(crate) fn new(inner: lancedb::Job<lancedb::function::FunctionVersion>) -> Self {
Self {
inner: Arc::new(inner),
}
}
inner: Arc<lancedb::Job<std::result::Result<Option<String>, String>>>,
}
impl Job {
pub(crate) fn new(inner: lancedb::Job) -> Self {
Self {
inner: Arc::new(inner),
inner: Arc::new(inner.map(|()| Ok(None))),
}
}
}
#[pymethods]
impl FunctionJob {
#[getter]
pub fn id(&self) -> Option<String> {
self.inner.id().map(str::to_string)
}
pub fn status(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(
self_.py(),
async move { inner.status().await.infer_error() },
)
}
pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner
.wait()
.await
.infer_error()?
.to_canonical_json()
.infer_error()
})
}
pub fn cancel(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.cancel().await.infer_error()?;
Ok(())
})
pub(crate) fn new_typed<T>(inner: lancedb::Job<T>) -> Self
where
T: Clone + Serialize + Send + Sync + 'static,
{
Self {
inner: Arc::new(inner.map(|result| {
serde_json::to_string(&result)
.map(Some)
.map_err(|error| format!("failed to serialize typed job result: {error}"))
})),
}
}
}
@@ -92,8 +53,10 @@ impl Job {
pub fn wait(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
future_into_py(self_.py(), async move {
inner.wait().await.infer_error()?;
Ok(None::<()>)
let result = inner.wait().await.infer_error()?;
result
.map_err(|message| lancedb::Error::Runtime { message })
.infer_error()
})
}
-1
View File
@@ -47,7 +47,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<Session>()?;
m.add_class::<Table>()?;
m.add_class::<crate::job::Job>()?;
m.add_class::<crate::job::FunctionJob>()?;
m.add_class::<crate::job::JobInfo>()?;
m.add_class::<crate::job::JobDescription>()?;
m.add_class::<crate::job::JobFailureInfo>()?;
+68 -12
View File
@@ -1,6 +1,7 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
@@ -16,8 +17,8 @@ use arrow::pyarrow::FromPyArrow;
use arrow::pyarrow::IntoPyArrow;
use arrow::pyarrow::ToPyArrow;
use lancedb::index::scalar::{
BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur,
Operator, PhraseQuery,
BooleanQuery, BoostQuery, DocumentGranularity, FtsQuery, FullTextSearchQuery, MatchQuery,
MultiMatchQuery, Occur, Operator, PhraseQuery,
};
use lancedb::query::AnalyzePlanDistributedMetrics;
use lancedb::query::QueryBase;
@@ -76,8 +77,16 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
let max_expansions = ob.getattr("max_expansions")?.extract()?;
let operator = ob.getattr("operator")?.extract::<String>()?;
let prefix_length = ob.getattr("prefix_length")?.extract()?;
let document_granularity = ob
.getattr("document_granularity")?
.extract::<Option<String>>()?
.map(|value| {
DocumentGranularity::try_from(value.as_str())
.map_err(|err| PyValueError::new_err(err.to_string()))
})
.transpose()?;
Ok(Self(
let mut query =
MatchQuery::new(query)
.with_column(Some(column))
.with_boost(boost)
@@ -86,21 +95,32 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
.with_operator(Operator::try_from(operator.as_str()).map_err(|e| {
PyValueError::new_err(format!("Invalid operator: {}", e))
})?)
.with_prefix_length(prefix_length)
.into(),
))
.with_prefix_length(prefix_length);
if let Some(document_granularity) = document_granularity {
query = query.with_document_granularity(document_granularity);
}
Ok(Self(query.into()))
}
"PhraseQuery" => {
let query = ob.getattr("query")?.extract()?;
let column = ob.getattr("column")?.extract()?;
let slop = ob.getattr("slop")?.extract()?;
let document_granularity = ob
.getattr("document_granularity")?
.extract::<Option<String>>()?
.map(|value| {
DocumentGranularity::try_from(value.as_str())
.map_err(|err| PyValueError::new_err(err.to_string()))
})
.transpose()?;
Ok(Self(
PhraseQuery::new(query)
.with_column(Some(column))
.with_slop(slop)
.into(),
))
let mut query = PhraseQuery::new(query)
.with_column(Some(column))
.with_slop(slop);
if let Some(document_granularity) = document_granularity {
query = query.with_document_granularity(document_granularity);
}
Ok(Self(query.into()))
}
"BoostQuery" => {
let positive: Self = ob.getattr("positive")?.extract()?;
@@ -167,6 +187,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
kwargs.set_item("max_expansions", query.max_expansions)?;
kwargs.set_item::<_, &str>("operator", query.operator.into())?;
kwargs.set_item("prefix_length", query.prefix_length)?;
if let Some(document_granularity) = query.document_granularity {
let value = match document_granularity {
DocumentGranularity::Row => "row",
DocumentGranularity::ListElement => "list_element",
};
kwargs.set_item("document_granularity", value)?;
}
namespace
.getattr(intern!(py, "MatchQuery"))?
.call((query.terms, query.column.unwrap()), Some(&kwargs))
@@ -174,6 +201,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
FtsQuery::Phrase(query) => {
let kwargs = PyDict::new(py);
kwargs.set_item("slop", query.slop)?;
if let Some(document_granularity) = query.document_granularity {
let value = match document_granularity {
DocumentGranularity::Row => "row",
DocumentGranularity::ListElement => "list_element",
};
kwargs.set_item("document_granularity", value)?;
}
namespace
.getattr(intern!(py, "PhraseQuery"))?
.call((query.terms, query.column.unwrap()), Some(&kwargs))
@@ -292,6 +326,7 @@ pub struct PyQueryRequest {
pub filter: Option<PyQueryFilter>,
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
pub select: PySelect,
pub select_source_columns: Option<HashMap<String, String>>,
pub fast_search: Option<bool>,
pub with_row_id: Option<bool>,
pub use_lsm: Option<bool>,
@@ -322,6 +357,7 @@ impl From<AnyQuery> for PyQueryRequest {
full_text_search: query_request
.full_text_search
.map(|fts| PyLanceDB(fts.query)),
select_source_columns: PySelect::source_columns(&query_request.select),
select: PySelect(query_request.select),
fast_search: Some(query_request.fast_search),
with_row_id: Some(query_request.with_row_id),
@@ -347,6 +383,7 @@ impl From<AnyQuery> for PyQueryRequest {
offset: vector_query.base.offset,
filter: vector_query.base.filter.map(PyQueryFilter),
full_text_search: None,
select_source_columns: PySelect::source_columns(&vector_query.base.select),
select: PySelect(vector_query.base.select),
fast_search: Some(vector_query.base.fast_search),
with_row_id: Some(vector_query.base.with_row_id),
@@ -379,6 +416,25 @@ impl From<AnyQuery> for PyQueryRequest {
#[derive(Clone)]
pub struct PySelect(Select);
impl PySelect {
fn source_columns(select: &Select) -> Option<HashMap<String, String>> {
match select {
Select::Expr(pairs) => Some(
pairs
.iter()
.filter_map(|(output, expr)| match expr {
lancedb::expr::DfExpr::Column(column) if column.relation.is_none() => {
Some((output.clone(), column.name.clone()))
}
_ => None,
})
.collect(),
),
_ => None,
}
}
}
impl<'py> IntoPyObject<'py> for PySelect {
type Target = PyAny;
type Output = Bound<'py, Self::Target>;
+7 -3
View File
@@ -780,15 +780,19 @@ impl Table {
})
}
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false))]
pub fn add<'a>(
self_: PyRef<'a, Self>,
data: PyScannable,
mode: String,
progress: Option<Py<PyAny>>,
write_parallelism: Option<usize>,
allow_external_blob_outside_bases: bool,
) -> PyResult<Bound<'a, PyAny>> {
let mut op = self_.inner_ref()?.add(data);
let mut op = self_
.inner_ref()?
.add(data)
.allow_external_blob_outside_bases(allow_external_blob_outside_bases);
if mode == "append" {
op = op.mode(AddDataMode::Append);
} else if mode == "overwrite" {
@@ -1619,7 +1623,7 @@ impl Table {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
let job = inner.refresh_column_async(column).await.infer_error()?;
Ok(crate::job::Job::new(job))
Ok(crate::job::Job::new_typed(job))
})
}
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb"
version = "0.38.0-beta.5"
version = "0.38.0-beta.11"
edition.workspace = true
description = "LanceDB: A serverless, low-latency vector database for AI applications"
license.workspace = true
+44
View File
@@ -1679,6 +1679,50 @@ mod tests {
assert_eq!(tables, names[..7]);
}
#[tokio::test]
async fn test_list_tables_walks_page_boundaries() {
let tc = new_test_connection().await.unwrap();
if tc.is_remote {
// What resumes a page is the server's to decide, and asserting it here would be
// asserting the server's contract rather than this one.
return;
}
let db = tc.connection;
let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
let mut names = Vec::with_capacity(5);
for _ in 0..5 {
let name = uuid::Uuid::new_v4().to_string();
names.push(name.clone());
db.create_empty_table(name, schema.clone())
.execute()
.await
.unwrap();
}
names.sort();
// Walking in pages has to reach every table exactly once, with nothing lost at a
// page boundary.
let mut seen = Vec::with_capacity(names.len());
let mut page_token = None;
loop {
let page = db
.list_tables(ListTablesRequest {
id: Some(Vec::new()),
limit: Some(2),
page_token,
..Default::default()
})
.await
.unwrap();
seen.extend(page.tables);
page_token = page.page_token.filter(|token| !token.is_empty());
if page_token.is_none() {
break;
}
}
assert_eq!(seen, names);
}
#[tokio::test]
async fn test_open_table() {
let tc = new_test_connection().await.unwrap();
+61 -10
View File
@@ -974,17 +974,15 @@ impl Database for ListingDatabase {
f.drain(0..index);
}
// Determine if there's a next page
let next_page_token = if let Some(limit) = request.limit {
if f.len() > limit as usize {
let token = f[limit as usize].clone();
// Determine if there's a next page. The token is the last name of this page,
// not the first of the next one: the next page resumes strictly after the
// token, so naming the next page's first entry would skip it.
let next_page_token = match request.limit {
Some(limit) if f.len() > limit as usize => {
f.truncate(limit as usize);
Some(token)
} else {
None
f.last().cloned()
}
} else {
None
_ => None,
};
Ok(ListTablesResponse {
@@ -1478,7 +1476,7 @@ mod tests {
use crate::table::{AnyQuery, WriteOptions};
use arrow_array::{Int32Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use futures::{TryStreamExt, stream::once};
use futures::{TryStreamExt, future::try_join_all, stream::once};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
@@ -1616,6 +1614,59 @@ mod tests {
);
}
#[tokio::test]
async fn test_concurrent_open_table_reuses_connection_object_store() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let session = Arc::new(lance::session::Session::default());
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: Some(session.clone()),
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "test".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let before = session.store_registry().stats();
let opened_tables = try_join_all((0..32).map(|_| {
db.open_table(OpenTableRequest {
name: "test".to_string(),
namespace_path: vec![],
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
})
}))
.await
.unwrap();
let after = session.store_registry().stats();
assert_eq!(opened_tables.len(), 32);
assert_eq!(after.misses, before.misses);
assert_eq!(after.active_stores, before.active_stores);
assert!(after.hits >= before.hits + 32);
}
#[tokio::test]
async fn test_listing_database_root_ops_do_not_create_manifest() {
let tempdir = tempdir().unwrap();
@@ -27,6 +27,12 @@ pub const SRC_ROW_ID_COL: &str = "row_id";
pub const SPLIT_NAMES_CONFIG_KEY: &str = "split_names";
/// Base table version the permutation was built against.
pub const BASE_VERSION_CONFIG_KEY: &str = "base_version";
/// Base table branch the permutation was built against. Absent means main.
pub const BASE_BRANCH_CONFIG_KEY: &str = "base_branch";
pub const DEFAULT_MEMORY_LIMIT: usize = 100 * 1024 * 1024;
/// Where to store the permutation table
@@ -214,21 +220,11 @@ impl PermutationBuilder {
Ok(Box::pin(SimpleRecordBatchStream { schema, stream }))
}
fn add_split_names(
fn add_config_metadata(
data: SendableRecordBatchStream,
split_names: &[String],
metadata: HashMap<String, String>,
) -> Result<SendableRecordBatchStream> {
let schema = data
.schema()
.as_ref()
.clone()
.with_metadata(HashMap::from([(
SPLIT_NAMES_CONFIG_KEY.to_string(),
serde_json::to_string(split_names).map_err(|e| Error::Other {
message: format!("Failed to serialize split names: {}", e),
source: Some(e.into()),
})?,
)]));
let schema = data.schema().as_ref().clone().with_metadata(metadata);
let schema = Arc::new(schema);
let schema_clone = schema.clone();
let stream = data.map_ok(move |batch| batch.with_schema(schema.clone()).unwrap());
@@ -239,7 +235,20 @@ impl PermutationBuilder {
}
/// Builds the permutation table and stores it in the given database.
pub async fn build(self) -> Result<Table> {
pub async fn build(mut self) -> Result<Table> {
// Remote tables resolve latest independently for each request. Use a
// separate pinned handle so count, projection, and scan all refer to one
// snapshot without changing the caller's table checkout state. Native
// tables return `None` here and retain their existing behavior.
if let Some(snapshot) = self
.base_table
.base_table()
.snapshot_at_current_version()
.await?
{
self.base_table = Table::from(snapshot);
}
// Unflushed rows have no row id, so a permutation cannot address them.
match self.base_table.base_table().get_lsm_write_spec().await {
Ok(Some(_)) => {
@@ -256,9 +265,14 @@ impl PermutationBuilder {
Err(err) => return Err(err),
}
// The handle above is already pinned to one version. Record which one, so a
// reader -- in a DataLoader worker, against a table that has since moved --
// resolves these row addresses against the same snapshot.
let base_version = self.base_table.version().await?;
let base_branch = self.base_table.current_branch();
// First pass, apply filter and load row ids. `Shuffler` permutes positions, so
// every rank must scan the rows in the same order to build the same permutation.
// TODO: pin the version resolved here; remote does not implement Lazy.
let mut rows = self.base_table.query().select(Select::columns(&[ROW_ID]));
if let Some(filter) = &self.config.filter {
@@ -318,11 +332,24 @@ impl PermutationBuilder {
// Rename _rowid to row_id
let renamed = rename_column(sorted, ROW_ID, SRC_ROW_ID_COL)?;
let streaming_data = if let Some(split_names) = &self.config.split_names {
Self::add_split_names(renamed, split_names)?
} else {
renamed
};
let mut metadata = HashMap::from([(
BASE_VERSION_CONFIG_KEY.to_string(),
base_version.to_string(),
)]);
// Version numbers are per-branch, so the branch is part of the coordinate.
if let Some(branch) = &base_branch {
metadata.insert(BASE_BRANCH_CONFIG_KEY.to_string(), branch.clone());
}
if let Some(split_names) = &self.config.split_names {
metadata.insert(
SPLIT_NAMES_CONFIG_KEY.to_string(),
serde_json::to_string(split_names).map_err(|e| Error::Other {
message: format!("Failed to serialize split names: {}", e),
source: Some(e.into()),
})?,
);
}
let streaming_data = Self::add_config_metadata(renamed, metadata)?;
let (name, database) = match &self.config.destination {
PermutationDestination::Permanent(database, table_name) => {
@@ -409,6 +436,253 @@ mod tests {
assert!(table.base_table().scan_order_is_deterministic());
}
#[cfg(feature = "remote")]
#[tokio::test]
async fn test_remote_permutation_builder_pins_snapshot() {
use std::sync::{
Mutex,
atomic::{AtomicU64, Ordering},
};
use arrow_array::{RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field, Schema};
let row_ids = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
ROW_ID,
DataType::UInt64,
false,
)])),
vec![Arc::new(UInt64Array::from(vec![100]))],
)
.unwrap();
let mut query_body = Vec::new();
{
let mut writer =
arrow_ipc::writer::FileWriter::try_new(&mut query_body, &row_ids.schema()).unwrap();
writer.write(&row_ids).unwrap();
writer.finish().unwrap();
}
let latest = Arc::new(AtomicU64::new(7));
let expected_snapshot = Arc::new(AtomicU64::new(7));
let planning_versions = Arc::new(Mutex::new(Vec::new()));
let latest_ref = latest.clone();
let expected_snapshot_ref = expected_snapshot.clone();
let planning_versions_ref = planning_versions.clone();
let table = Table::new_with_handler("remote_base", move |request| {
let path = request.url().path();
let body = request
.body()
.and_then(|body| body.as_bytes())
.map(|body| serde_json::from_slice::<serde_json::Value>(body).unwrap());
match path {
"/v1/table/remote_base/describe/" => {
let requested = body.as_ref().and_then(|body| body["version"].as_u64());
let version = requested.unwrap_or_else(|| latest_ref.load(Ordering::SeqCst));
http::Response::builder()
.status(200)
.body(
format!(r#"{{"version":{version},"schema":{{"fields":[]}}}}"#)
.into_bytes(),
)
.unwrap()
}
"/v1/table/remote_base/get_lsm_write_spec/" => http::Response::builder()
.status(200)
.body(br#"{"lsm_write_spec":null}"#.to_vec())
.unwrap(),
"/v1/table/remote_base/count_rows/" => {
let body = body.unwrap();
let version = body["version"].as_u64().unwrap();
assert_eq!(version, expected_snapshot_ref.load(Ordering::SeqCst));
assert_eq!(body["predicate"], "value > 0");
planning_versions_ref.lock().unwrap().push(version);
// Simulate a concurrent append after count_rows. An unpinned
// scan would now resolve version 8 and include different rows.
latest_ref.store(8, Ordering::SeqCst);
http::Response::builder()
.status(200)
.body(b"1".to_vec())
.unwrap()
}
"/v1/table/remote_base/query/" => {
let body = body.unwrap();
let version = body["version"].as_u64().unwrap();
assert_eq!(version, expected_snapshot_ref.load(Ordering::SeqCst));
assert_eq!(body["filter"], "value > 0");
assert_eq!(body["columns"], serde_json::json!([ROW_ID]));
planning_versions_ref.lock().unwrap().push(version);
http::Response::builder()
.status(200)
.header("content-type", "application/vnd.apache.arrow.file")
.body(query_body.clone())
.unwrap()
}
_ => panic!("unexpected request: {path}"),
}
});
let permutation = PermutationBuilder::new(table.clone())
.with_filter("value > 0".to_string())
.build()
.await
.unwrap();
assert_eq!(permutation.count_rows(None).await.unwrap(), 1);
// Building uses a separate handle and must not pin the caller's table.
assert_eq!(table.version().await.unwrap(), 8);
// An explicit checkout is copied as-is and remains checked out afterward.
expected_snapshot.store(6, Ordering::SeqCst);
table.checkout(6).await.unwrap();
let permutation = PermutationBuilder::new(table.clone())
.with_filter("value > 0".to_string())
.build()
.await
.unwrap();
assert_eq!(permutation.count_rows(None).await.unwrap(), 1);
assert_eq!(table.version().await.unwrap(), 6);
assert_eq!(*planning_versions.lock().unwrap(), vec![7, 7, 6, 6]);
}
#[tokio::test]
async fn test_permutation_records_base_version() {
let temp_dir = tempfile::tempdir().unwrap();
let db = connect(temp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap();
let initial_data = lance_datagen::gen_batch()
.col("col_a", lance_datagen::array::step::<Int32Type>())
.into_ldb_stream(RowCount::from(100), BatchCount::from(2));
let data_table = db
.create_table("base_tbl", initial_data)
.execute()
.await
.unwrap();
let build_version = data_table.version().await.unwrap();
let permutation_table = PermutationBuilder::new(data_table.clone())
.build()
.await
.unwrap();
let recorded = permutation_table
.schema()
.await
.unwrap()
.metadata
.get(BASE_VERSION_CONFIG_KEY)
.expect("permutation should record the base version")
.parse::<u64>()
.unwrap();
assert_eq!(recorded, build_version);
// Advancing the base table must not move the recorded version.
let more_data = lance_datagen::gen_batch()
.col("col_a", lance_datagen::array::step::<Int32Type>())
.into_ldb_stream(RowCount::from(50), BatchCount::from(1));
data_table.add(more_data).execute().await.unwrap();
assert!(data_table.version().await.unwrap() > recorded);
assert_eq!(
permutation_table
.schema()
.await
.unwrap()
.metadata
.get(BASE_VERSION_CONFIG_KEY)
.unwrap()
.parse::<u64>()
.unwrap(),
recorded,
);
}
/// Version numbers are per-branch, so a permutation built on a branch must record
/// it -- a worker reopens by name and lands on main at the same number.
#[tokio::test]
async fn test_permutation_records_base_branch() {
let temp_dir = tempfile::tempdir().unwrap();
let db = connect(temp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap();
let initial_data = lance_datagen::gen_batch()
.col("col_a", lance_datagen::array::step::<Int32Type>())
.into_ldb_stream(RowCount::from(10), BatchCount::from(1));
let data_table = db
.create_table("base_tbl", initial_data)
.execute()
.await
.unwrap();
let branch = data_table
.create_branch("exp", lance::dataset::refs::Ref::from(("main", 1)))
.await
.unwrap();
let permutation_table = PermutationBuilder::new(branch.clone())
.build()
.await
.unwrap();
let metadata = permutation_table.schema().await.unwrap().metadata.clone();
assert_eq!(
metadata.get(BASE_BRANCH_CONFIG_KEY).map(String::as_str),
Some("exp")
);
// Main records nothing, so an absent key keeps meaning main.
let main_permutation = PermutationBuilder::new(data_table.clone())
.build()
.await
.unwrap();
assert!(
!main_permutation
.schema()
.await
.unwrap()
.metadata
.contains_key(BASE_BRANCH_CONFIG_KEY)
);
}
#[tokio::test]
async fn test_build_does_not_pin_the_callers_table() {
let temp_dir = tempfile::tempdir().unwrap();
let db = connect(temp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap();
let initial_data = lance_datagen::gen_batch()
.col("col_a", lance_datagen::array::step::<Int32Type>())
.into_ldb_stream(RowCount::from(100), BatchCount::from(1));
let data_table = db
.create_table("base_tbl", initial_data)
.execute()
.await
.unwrap();
PermutationBuilder::new(data_table.clone())
.build()
.await
.unwrap();
// The builder pins its own handle; the caller's must still track latest.
let more_data = lance_datagen::gen_batch()
.col("col_a", lance_datagen::array::step::<Int32Type>())
.into_ldb_stream(RowCount::from(50), BatchCount::from(1));
data_table.add(more_data).execute().await.unwrap();
assert_eq!(data_table.count_rows(None).await.unwrap(), 150);
}
#[tokio::test]
async fn test_permutation_builder() {
let temp_dir = tempfile::tempdir().unwrap();
@@ -8,7 +8,9 @@
//! the rows from a source table that correspond to row IDs stored in a separate table.
use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream};
use crate::dataloader::permutation::builder::SRC_ROW_ID_COL;
use crate::dataloader::permutation::builder::{
BASE_BRANCH_CONFIG_KEY, BASE_VERSION_CONFIG_KEY, SRC_ROW_ID_COL,
};
use crate::dataloader::permutation::split::SPLIT_ID_COLUMN;
use crate::error::Error;
use crate::query::{
@@ -23,6 +25,7 @@ use arrow_array::{RecordBatch, UInt64Array};
use arrow_schema::SchemaRef;
use datafusion_expr::{Expr, col, lit};
use futures::{StreamExt, TryStreamExt};
use lance::dataset::refs::MAIN_BRANCH;
use lance::dataset::scanner::DatasetRecordBatchStream;
use lance::io::RecordBatchStream;
use lance_arrow::RecordBatchExt;
@@ -69,6 +72,10 @@ impl PermutationReader {
permutation_table: Option<Arc<dyn BaseTable>>,
split: u64,
) -> Result<Self> {
let base_table = match &permutation_table {
Some(permutation_table) => Self::pin_base_table(base_table, permutation_table).await?,
None => base_table,
};
let mut slf = Self {
base_table,
permutation_table,
@@ -89,6 +96,34 @@ impl PermutationReader {
Ok(slf)
}
/// Pins the base table to the version the permutation was built against.
/// Permutations written before that was recorded carry no key and stay unpinned.
async fn pin_base_table(
base_table: Arc<dyn BaseTable>,
permutation_table: &Arc<dyn BaseTable>,
) -> Result<Arc<dyn BaseTable>> {
let schema = permutation_table.schema().await?;
let Some(raw) = schema.metadata.get(BASE_VERSION_CONFIG_KEY) else {
return Ok(base_table);
};
let version = raw.parse::<u64>().map_err(|e| Error::InvalidInput {
message: format!(
"Permutation table has an unreadable {} of {:?}: {}",
BASE_VERSION_CONFIG_KEY, raw, e
),
})?;
// The recorded branch, not the handle's: a worker reopens by name and lands
// on main, and version numbers are per-branch.
let branch = schema
.metadata
.get(BASE_BRANCH_CONFIG_KEY)
.map(String::as_str)
.unwrap_or(MAIN_BRANCH);
base_table
.checkout_branch_version(branch, Some(version))
.await
}
pub async fn try_from_tables(
base_table: Arc<dyn BaseTable>,
permutation_table: Arc<dyn BaseTable>,
@@ -511,9 +546,13 @@ mod tests {
use lance_datagen::{BatchCount, RowCount};
use rand::seq::SliceRandom;
// Aliased: `test_utils::datagen` exports a trait of the same name.
use crate::arrow::LanceDbDatagenExt as _;
use crate::{
Table,
arrow::SendableRecordBatchStream,
connect,
dataloader::permutation::builder::PermutationBuilder,
query::{ExecutableQuery, QueryBase},
test_utils::datagen::{LanceDbDatagenExt, virtual_table},
};
@@ -545,6 +584,58 @@ mod tests {
.await
}
/// Compaction moves row addresses, so the reader must read the pinned version.
#[tokio::test]
async fn test_reader_pins_base_version() {
let temp_dir = tempfile::tempdir().unwrap();
let db = connect(temp_dir.path().to_str().unwrap())
.execute()
.await
.unwrap();
let data = lance_datagen::gen_batch()
.col("idx", lance_datagen::array::step::<Int32Type>())
.into_ldb_stream(RowCount::from(20), BatchCount::from(1));
let base_table = db.create_table("base_tbl", data).execute().await.unwrap();
let permutation_table = PermutationBuilder::new(base_table.clone())
.build()
.await
.unwrap();
base_table.delete("true").await.unwrap();
base_table
.optimize(crate::table::OptimizeAction::All)
.await
.unwrap();
assert_eq!(base_table.count_rows(None).await.unwrap(), 0);
let reader = PermutationReader::try_from_tables(
base_table.base_table().clone(),
permutation_table.base_table().clone(),
0,
)
.await
.unwrap();
let values = collect_from_stream::<Int32Type>(
reader
.read(
Select::Columns(vec!["idx".to_string()]),
QueryExecutionOptions::default(),
)
.await
.unwrap(),
"idx",
)
.await;
assert_eq!(
values.len(),
20,
"reader should still see the pinned version"
);
}
#[tokio::test]
async fn test_permutation_reader() {
let base_table = lance_datagen::gen_batch()
+121 -4
View File
@@ -19,6 +19,7 @@
mod sql;
pub(crate) use sql::canonicalize_sql_predicate;
pub use sql::expr_to_sql_string;
use std::sync::Arc;
@@ -156,7 +157,7 @@ mod tests {
use datafusion_common::ScalarValue;
let expr = col("data").eq(lit(ScalarValue::Binary(Some(vec![0xca, 0xfe]))));
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(sql, "(data = X'CAFE')");
assert_eq!(sql, "(`data` = X'CAFE')");
}
#[test]
@@ -166,7 +167,7 @@ mod tests {
let int_expr = col("id").gt(lit(5i64));
let combined = bin_expr.and(int_expr);
let sql = expr_to_sql_string(&combined).unwrap();
assert_eq!(sql, "((data = X'01') AND (id > 5))");
assert_eq!(sql, "((`data` = X'01') AND (id > 5))");
}
#[test]
@@ -184,7 +185,7 @@ mod tests {
// serialized correctly (regression test for placeholder rewrite path).
let expr = contains(col("data"), lit(ScalarValue::Binary(Some(vec![0xff]))));
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(sql, "contains(data, X'FF')");
assert_eq!(sql, "contains(`data`, X'FF')");
}
#[test]
@@ -195,7 +196,7 @@ mod tests {
.eq(lit(ScalarValue::Binary(Some(vec![0xab, 0xcd]))))
.not();
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(sql, "NOT (data = X'ABCD')");
assert_eq!(sql, "NOT (`data` = X'ABCD')");
}
#[test]
@@ -205,6 +206,122 @@ mod tests {
assert!(sql.contains("IN"), "expected IN in: {}", sql);
}
#[test]
fn test_empty_is_in() {
let expr = is_in(col("id"), vec![]);
assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
}
#[test]
fn test_empty_is_in_discards_binary_children() {
use datafusion_common::ScalarValue;
let expr = is_in(
col("payload").eq(lit(ScalarValue::Binary(Some(vec![0x01])))),
vec![],
);
assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
}
#[test]
fn test_keyword_identifier() {
let expr = col("null").eq(lit(1i64));
assert_eq!(expr_to_sql_string(&expr).unwrap(), "(`null` = 1)");
}
#[test]
fn test_decimal_literal_preserves_type() {
use datafusion_common::ScalarValue;
let expr = col("val").lt(lit(ScalarValue::Decimal128(
Some(1_234_567_890_123_456_790),
19,
18,
)));
let sql = expr_to_sql_string(&expr).unwrap();
assert_eq!(
sql,
"(val < arrow_cast('1.234567890123456790', 'Decimal128(19, 18)'))"
);
}
#[test]
fn test_non_finite_float_literal_preserves_type() {
let expr = col("x").lt(lit(f64::INFINITY));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"(x < arrow_cast('inf', 'Float64'))"
);
}
#[test]
fn test_cast_uses_arrow_type_name() {
let string = expr_cast(col("x"), DataType::Utf8);
assert_eq!(
expr_to_sql_string(&string).unwrap(),
"arrow_cast(x, 'Utf8')"
);
let int32 = expr_cast(col("x"), DataType::Int32);
assert_eq!(
expr_to_sql_string(&int32).unwrap(),
"arrow_cast(x, 'Int32')"
);
let expr = expr_cast(col("x"), DataType::Float16).lt(lit(2.0));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"(arrow_cast(x, 'Float16') < 2.0)"
);
let decimal = expr_cast(lit("2.00"), DataType::Decimal256(40, 2));
assert_eq!(
expr_to_sql_string(&decimal).unwrap(),
"arrow_cast('2.00', 'Decimal256(40, 2)')"
);
}
#[test]
fn test_binary_placeholder_does_not_rewrite_user_string() {
use datafusion_common::ScalarValue;
let marker = "__lancedb_binary_placeholder_0__";
let expr = col("payload")
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
.or(col("text").eq(lit(marker)));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"((payload = X'01') OR (`text` = '__lancedb_binary_placeholder_0__'))"
);
}
#[test]
fn test_binary_binding_skips_quoted_identifiers() {
use datafusion_common::ScalarValue;
let expr = col("payload")
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
.and(col("odd'name").eq(lit(1i64)))
.and(col("odd`'name").eq(lit(2i64)));
assert_eq!(
expr_to_sql_string(&expr).unwrap(),
"(((payload = X'01') AND (`odd'name` = 1)) AND (`odd``'name` = 2))"
);
}
#[test]
fn test_binary_placeholder_collision_search_is_linear() {
use datafusion_common::ScalarValue;
let collision_shaped = format!("__lancedb_binary_placeholder_0__{}", "_".repeat(64_000));
let expr = col("payload")
.eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
.and(col("text").eq(lit(collision_shaped.clone())));
let sql = expr_to_sql_string(&expr).unwrap();
assert!(sql.contains("X'01'"));
assert!(sql.contains(&format!("'{collision_shaped}'")));
}
#[test]
fn test_multiple_binary_literals() {
use datafusion_common::ScalarValue;
+330 -43
View File
@@ -1,10 +1,27 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::{
any::TypeId,
collections::{HashMap, HashSet},
};
use arrow_array::types::{
Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType,
};
use arrow_schema::DataType;
use datafusion_common::ScalarValue;
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_expr::Expr;
use datafusion_sql::unparser::{self, dialect::Dialect};
use datafusion_functions::core::expr_fn::{
arrow_cast as datafusion_arrow_cast, arrow_try_cast as datafusion_arrow_try_cast,
};
use datafusion_sql::sqlparser::{
dialect::{Dialect as SqlParserDialect, GenericDialect},
keywords::ALL_KEYWORDS,
tokenizer::{Token, Tokenizer},
};
use datafusion_sql::unparser::{self, dialect::Dialect as UnparserDialect};
/// Unparser dialect that matches the quoting style expected by the Lance SQL
/// parser. Lance uses backtick (`` ` ``) as the only delimited-identifier
@@ -19,17 +36,74 @@ use datafusion_sql::unparser::{self, dialect::Dialect};
/// lower-case by the SQL parser, which would break case-sensitive schemas).
struct LanceSqlDialect;
impl Dialect for LanceSqlDialect {
impl UnparserDialect for LanceSqlDialect {
fn identifier_quote_style(&self, identifier: &str) -> Option<char> {
let needs_quote = identifier.chars().any(|c| c.is_ascii_uppercase())
|| !identifier
.chars()
.enumerate()
.all(|(i, c)| c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit()));
let identifier_upper = identifier.to_ascii_uppercase();
let needs_quote =
(identifier_upper != "ID" && ALL_KEYWORDS.contains(&identifier_upper.as_str()))
|| identifier.chars().any(|c| c.is_ascii_uppercase())
|| !identifier.chars().enumerate().all(|(i, c)| {
c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit())
});
if needs_quote { Some('`') } else { None }
}
}
/// Lance's tokenizer dialect with SQL-standard double-quoted identifiers added.
///
/// Keep this deliberately small: Lance's parser wraps `GenericDialect` and
/// delegates only identifier recognition, leaving every other dialect option at
/// its default. In particular, `/*! ... */` remains an ordinary block comment.
#[derive(Debug, Default)]
struct PredicateDialect(GenericDialect);
impl SqlParserDialect for PredicateDialect {
fn dialect(&self) -> TypeId {
self.0.dialect()
}
fn is_identifier_start(&self, ch: char) -> bool {
self.0.is_identifier_start(ch)
}
fn is_identifier_part(&self, ch: char) -> bool {
self.0.is_identifier_part(ch)
}
fn is_delimited_identifier_start(&self, ch: char) -> bool {
ch == '"' || ch == '`'
}
}
/// Canonicalize a raw SQL predicate for Lance's parser.
///
/// Lance wraps [`GenericDialect`] for identifier recognition while retaining the
/// default dialect behavior for every other lexical option. [`PredicateDialect`]
/// mirrors that contract and additionally recognizes `"` as an identifier
/// delimiter, allowing this function to rewrite only those identifier tokens.
pub fn canonicalize_sql_predicate(predicate: &str) -> crate::Result<String> {
let dialect = PredicateDialect::default();
let tokens = Tokenizer::new(&dialect, predicate)
.with_unescape(false)
.tokenize()
.map_err(|err| crate::Error::InvalidInput {
message: format!("invalid SQL predicate: {err}"),
})?;
Ok(tokens
.into_iter()
.map(|token| match token {
Token::Word(word) if word.quote_style == Some('"') => {
// with_unescape(false) retains doubled double quotes. Decode
// those before escaping any backticks for Lance's delimiter.
let identifier = word.value.replace("\"\"", "\"").replace('`', "``");
format!("`{identifier}`")
}
other => other.to_string(),
})
.collect())
}
/// Prefix for placeholder strings inserted in place of binary literals. Chosen
/// to be extremely unlikely to occur in user data.
const BINARY_PLACEHOLDER_PREFIX: &str = "__lancedb_binary_placeholder_";
@@ -39,24 +113,128 @@ fn bytes_to_hex_sql(bytes: &[u8]) -> String {
format!("X'{hex}'")
}
/// Returns true if *expr* contains a `Binary` or `LargeBinary` scalar literal
/// anywhere in its subtree. DataFusion's SQL unparser cannot serialize those
/// variants, so we route such expressions through a placeholder-substitution
/// path that emits SQL `X'...'` byte-string literals.
fn has_binary_literal(expr: &Expr) -> bool {
let mut found = false;
fn string_literals(expr: &Expr) -> HashSet<String> {
let mut literals = HashSet::new();
let _ = expr.apply(&mut |e: &Expr| {
if matches!(
e,
Expr::Literal(ScalarValue::Binary(_) | ScalarValue::LargeBinary(_), _)
) {
found = true;
Ok(TreeNodeRecursion::Stop)
} else {
Ok(TreeNodeRecursion::Continue)
if let Expr::Literal(
ScalarValue::Utf8(Some(value))
| ScalarValue::LargeUtf8(Some(value))
| ScalarValue::Utf8View(Some(value)),
_,
) = e
{
literals.insert(value.clone());
}
Ok(TreeNodeRecursion::Continue)
});
found
literals
}
fn typed_string_literal(value: String, data_type: DataType) -> Expr {
datafusion_arrow_cast(
Expr::Literal(ScalarValue::Utf8(Some(value)), None),
Expr::Literal(ScalarValue::Utf8(Some(data_type.to_string())), None),
)
}
fn next_binary_placeholder(user_strings: &HashSet<String>, next_id: &mut usize) -> String {
loop {
let placeholder = format!("{BINARY_PLACEHOLDER_PREFIX}{}__", *next_id);
*next_id += 1;
if !user_strings.contains(&placeholder) {
return placeholder;
}
}
}
fn bind_binary_literals(
sql: &str,
mut bindings: HashMap<String, Vec<u8>>,
) -> crate::Result<String> {
let bytes = sql.as_bytes();
let mut output = Vec::with_capacity(bytes.len());
let mut index = 0;
// Walk SQL string tokens once. Placeholders are plain, unescaped string
// literals, so this remains linear even when user strings are large or
// deliberately resemble the placeholder prefix.
while index < bytes.len() {
if bytes[index] == b'`' {
let identifier_start = index;
index += 1;
let mut identifier_end = None;
while index < bytes.len() {
if bytes[index] == b'`' {
if index + 1 < bytes.len() && bytes[index + 1] == b'`' {
index += 2;
} else {
index += 1;
identifier_end = Some(index);
break;
}
} else {
index += 1;
}
}
let Some(identifier_end) = identifier_end else {
return Err(crate::Error::InvalidInput {
message: "unterminated identifier while binding binary literal".to_string(),
});
};
output.extend_from_slice(&bytes[identifier_start..identifier_end]);
continue;
}
if bytes[index] != b'\'' {
output.push(bytes[index]);
index += 1;
continue;
}
let literal_start = index;
index += 1;
let content_start = index;
let mut escaped = false;
let mut content_end = None;
while index < bytes.len() {
if bytes[index] == b'\'' {
if index + 1 < bytes.len() && bytes[index + 1] == b'\'' {
escaped = true;
index += 2;
} else {
content_end = Some(index);
index += 1;
break;
}
} else {
index += 1;
}
}
let Some(content_end) = content_end else {
return Err(crate::Error::InvalidInput {
message: "unterminated string while binding binary literal".to_string(),
});
};
let placeholder = &sql[content_start..content_end];
if !escaped && let Some(value) = bindings.remove(placeholder) {
output.extend_from_slice(bytes_to_hex_sql(&value).as_bytes());
} else {
output.extend_from_slice(&bytes[literal_start..index]);
}
}
if !bindings.is_empty() {
return Err(crate::Error::InvalidInput {
message: "failed to bind binary literal while serializing expression".to_string(),
});
}
String::from_utf8(output).map_err(|e| crate::Error::InvalidInput {
message: format!("failed to bind binary literal: {e}"),
})
}
fn run_unparser(expr: &Expr) -> crate::Result<String> {
@@ -69,25 +247,37 @@ fn run_unparser(expr: &Expr) -> crate::Result<String> {
}
pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
// Fast path: no binary literals — DataFusion's unparser handles everything.
if !has_binary_literal(expr) {
return run_unparser(expr);
}
// Slow path: DataFusion's unparser cannot serialize `Binary`/`LargeBinary`
// scalars, so we rewrite each one to a unique string-literal placeholder,
// let the unparser do the rest of the work, then substitute the SQL
// `X'...'` byte-string literal back in. This keeps the operator/function
// serialization logic centralized in DataFusion and works for every
// expression node type the unparser supports.
let mut bindings: Vec<Vec<u8>> = Vec::new();
// DataFusion's unparser needs a few adaptations before its SQL can be
// reparsed by Lance without changing the typed expression's semantics:
//
// * decimal literals need an explicit cast to preserve precision and scale;
// * casts need exact Arrow type names rather than SQL type aliases;
// * an empty IN list is valid in DataFusion but invalid SQL;
// * binary literals are unsupported by the unparser and need placeholders.
// Eliminate empty membership expressions before visiting their children.
// Otherwise a discarded binary child could leave behind a stale binding.
let rewritten = expr
.clone()
.transform(|e: Expr| match e {
Expr::InList(in_list) if in_list.list.is_empty() => Ok(Transformed::yes(
Expr::Literal(ScalarValue::Boolean(Some(in_list.negated)), None),
)),
other => Ok(Transformed::no(other)),
})
.map_err(|e| crate::Error::InvalidInput {
message: format!("failed to rewrite expression: {e}"),
})?
.data;
let user_strings = string_literals(&rewritten);
let mut next_placeholder_id = 0;
let mut binary_bindings = HashMap::new();
let rewritten = rewritten
.transform(|e: Expr| match e {
Expr::Literal(ScalarValue::Binary(Some(bytes)), m)
| Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), m) => {
let placeholder = format!("{}{}__", BINARY_PLACEHOLDER_PREFIX, bindings.len());
bindings.push(bytes);
let placeholder = next_binary_placeholder(&user_strings, &mut next_placeholder_id);
binary_bindings.insert(placeholder.clone(), bytes);
Ok(Transformed::yes(Expr::Literal(
ScalarValue::Utf8(Some(placeholder)),
m,
@@ -97,6 +287,57 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
| Expr::Literal(ScalarValue::LargeBinary(None), m) => {
Ok(Transformed::yes(Expr::Literal(ScalarValue::Null, m)))
}
Expr::Literal(ScalarValue::Decimal32(Some(value), precision, scale), _m) => {
let value = Decimal32Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal32(precision, scale),
)))
}
Expr::Literal(ScalarValue::Decimal64(Some(value), precision, scale), _m) => {
let value = Decimal64Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal64(precision, scale),
)))
}
Expr::Literal(ScalarValue::Decimal128(Some(value), precision, scale), _m) => {
let value = Decimal128Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal128(precision, scale),
)))
}
Expr::Literal(ScalarValue::Decimal256(Some(value), precision, scale), _m) => {
let value = Decimal256Type::format_decimal(value, precision, scale);
Ok(Transformed::yes(typed_string_literal(
value,
DataType::Decimal256(precision, scale),
)))
}
Expr::Literal(ScalarValue::Float16(Some(value)), _m) if !value.is_finite() => Ok(
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float16)),
),
Expr::Literal(ScalarValue::Float32(Some(value)), _m) if !value.is_finite() => Ok(
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float32)),
),
Expr::Literal(ScalarValue::Float64(Some(value)), _m) if !value.is_finite() => Ok(
Transformed::yes(typed_string_literal(value.to_string(), DataType::Float64)),
),
Expr::Cast(cast) => Ok(Transformed::yes(datafusion_arrow_cast(
*cast.expr,
Expr::Literal(
ScalarValue::Utf8(Some(cast.field.data_type().to_string())),
None,
),
))),
Expr::TryCast(cast) => Ok(Transformed::yes(datafusion_arrow_try_cast(
*cast.expr,
Expr::Literal(
ScalarValue::Utf8(Some(cast.field.data_type().to_string())),
None,
),
))),
other => Ok(Transformed::no(other)),
})
.map_err(|e| crate::Error::InvalidInput {
@@ -104,12 +345,58 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
})?
.data;
let mut sql = run_unparser(&rewritten)?;
for (i, bytes) in bindings.iter().enumerate() {
// The unparser quotes string literals with single quotes, so the
// placeholder appears as `'__lancedb_binary_placeholder_<i>__'`.
let quoted = format!("'{}{}__'", BINARY_PLACEHOLDER_PREFIX, i);
sql = sql.replace(&quoted, &bytes_to_hex_sql(bytes));
let sql = run_unparser(&rewritten)?;
if binary_bindings.is_empty() {
Ok(sql)
} else {
bind_binary_literals(&sql, binary_bindings)
}
}
#[cfg(test)]
mod tests {
use super::canonicalize_sql_predicate;
#[test]
fn normalizes_double_quoted_identifiers() {
assert_eq!(
canonicalize_sql_predicate(r#""PartyAbbrev" = 'D'"#).unwrap(),
"`PartyAbbrev` = 'D'"
);
assert_eq!(
canonicalize_sql_predicate(r#""MetaData"."userId" = 5"#).unwrap(),
"`MetaData`.`userId` = 5"
);
assert_eq!(
canonicalize_sql_predicate(r#""a""b" = 1"#).unwrap(),
"`a\"b` = 1"
);
}
#[test]
fn preserves_quotes_inside_literals_and_backticks() {
let filter = r#"name = 'Alice "Ace"' AND `quoted"field` = 1"#;
assert_eq!(canonicalize_sql_predicate(filter).unwrap(), filter);
}
#[test]
fn preserves_literals_and_comments_using_lance_dialect_rules() {
let predicate = r#"path = '\' AND "PartyAbbrev" = 'D' -- unmatched " in comment"#;
assert_eq!(
canonicalize_sql_predicate(predicate).unwrap(),
r#"path = '\' AND `PartyAbbrev` = 'D' -- unmatched " in comment"#
);
let predicate = r#"id = 1 /* unmatched " in block comment */"#;
assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate);
let predicate = r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#;
assert_eq!(canonicalize_sql_predicate(predicate).unwrap(), predicate);
}
#[test]
fn rejects_unterminated_double_quoted_identifier() {
let error = canonicalize_sql_predicate(r#""PartyAbbrev = 'D'"#).unwrap_err();
assert!(matches!(error, crate::Error::InvalidInput { .. }));
}
Ok(sql)
}
+43 -38
View File
@@ -1,10 +1,11 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Canonical values exchanged with the Enterprise Function service.
//! Canonical Function values exchanged with the Enterprise service, plus the
//! backend-neutral terminal result of a computed-column refresh.
//!
//! This module contains client/wire values only. Catalog persistence,
//! environment bake, secret resolution, and execution are owned by Sophon.
//! environment bake, and execution are owned by Sophon.
use std::collections::BTreeMap;
@@ -185,6 +186,9 @@ pub struct PythonEnvironmentSpec {
pub kind: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub packages: Vec<String>,
/// Conda channels in priority order; conda environments only.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub channels: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub path: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
@@ -194,9 +198,6 @@ pub struct PythonEnvironmentSpec {
}
/// Reproducible Python runtime definition understood by Sophon.
///
/// `env` contains non-secret values. Secret values have no client model;
/// [`FunctionVersion::required_secrets`] contains names only.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum PythonRuntimeSpec {
@@ -238,7 +239,7 @@ impl PythonRuntimeSpec {
}
}
/// Non-secret environment variables, or `None` for an unknown kind.
/// Environment variables, or `None` for an unknown kind.
pub fn env(&self) -> Option<&BTreeMap<String, String>> {
match self {
Self::Python { env, .. } => Some(env),
@@ -323,8 +324,6 @@ pub struct FunctionVersion {
runtime: PythonRuntimeSpec,
runtime_digest: String,
environment_digest: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
required_secrets: Vec<String>,
created_at: String,
}
@@ -357,11 +356,6 @@ impl FunctionVersion {
&self.environment_digest
}
/// Required secret names. Resolved values exist only inside Sophon.
pub fn required_secrets(&self) -> &[String] {
&self.required_secrets
}
pub fn created_at(&self) -> &str {
&self.created_at
}
@@ -403,18 +397,12 @@ pub struct FunctionArtifactRequest {
}
/// Stable request envelope for remote immutable Function registration.
///
/// Secret values deliberately have no field in this model. The only secret
/// material the client may send is the ordered set of names Sophon resolves
/// inside the remote runtime.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FunctionRegistrationRequest {
pub name: String,
pub artifact: FunctionArtifactRequest,
pub signature: FunctionSignature,
pub runtime: PythonRuntimeSpec,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub required_secrets: Vec<String>,
}
impl_json!(FunctionRegistrationRequest);
@@ -445,7 +433,6 @@ pub struct FunctionApplication {
function: FunctionVersionRef,
inputs: Vec<ApplicationInput>,
output: FunctionOutput,
group_id: String,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
columns: BTreeMap<String, String>,
#[serde(default, flatten, skip_serializing)]
@@ -467,10 +454,6 @@ impl FunctionApplication {
&self.output
}
pub fn group_id(&self) -> &str {
&self.group_id
}
pub fn columns(&self) -> &BTreeMap<String, String> {
&self.columns
}
@@ -512,7 +495,7 @@ pub struct InputBinding {
pub nullable: bool,
}
/// Ordered result-field to table-field mapping for a grouped binding.
/// Ordered result-field to table-field mapping for a Function binding.
///
/// Assignment state is not part of the Slice 1 client contract. During the
/// NULL transition there is no public Lance cell-flag identifier to persist.
@@ -526,20 +509,18 @@ pub struct OutputMapping {
pub nullable: bool,
}
/// Immutable grouped binding persisted by the Enterprise table service.
/// Immutable Function binding persisted by the Enterprise table service.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FunctionBinding {
binding_id: String,
revision: u64,
function: FunctionVersionRef,
group_id: String,
inputs: Vec<InputBinding>,
outputs: Vec<OutputMapping>,
/// Exact Arrow schema presented to the Function, encoded with the Lance
/// Namespace Arrow JSON representation.
#[serde(default, skip_serializing_if = "Option::is_none")]
input_schema: Option<Value>,
/// Exact physical Arrow schema of the grouped table outputs.
/// Exact physical Arrow schema of the binding's table outputs.
#[serde(default, skip_serializing_if = "Option::is_none")]
output_schema: Option<Value>,
}
@@ -549,18 +530,10 @@ impl FunctionBinding {
&self.binding_id
}
pub fn revision(&self) -> u64 {
self.revision
}
pub fn function(&self) -> &FunctionVersionRef {
&self.function
}
pub fn group_id(&self) -> &str {
&self.group_id
}
pub fn inputs(&self) -> &[InputBinding] {
&self.inputs
}
@@ -580,13 +553,22 @@ impl FunctionBinding {
impl_json!(FunctionBinding);
/// Stable terminal result of a remote Function-column refresh Job.
/// Stable terminal result of an expression-backed or Function-backed column
/// refresh [`crate::Job`].
///
/// Local refresh jobs produce this value in process. LanceDB Cloud and
/// Enterprise decode the same value from the durable job's terminal payload.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RefreshColumnResult {
/// Rows assigned a value by this refresh.
pub rows_assigned: u64,
/// Rows whose computation failed.
pub rows_failed: u64,
/// Rows that still need a value when the job completes.
pub rows_remaining: u64,
/// Exact table version the refresh read.
pub source_version: u64,
/// Table version made visible by the refresh, when one was published.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub published_version: Option<u64>,
}
@@ -604,3 +586,26 @@ impl RefreshColumnResult {
}
impl_json!(RefreshColumnResult);
#[cfg(test)]
mod conda_environment_tests {
use super::PythonEnvironmentSpec;
#[test]
fn conda_channels_round_trip_and_pip_stays_bare() {
let conda: PythonEnvironmentSpec = serde_json::from_str(
r#"{"kind":"conda","packages":["numpy"],"channels":["conda-forge"]}"#,
)
.unwrap();
assert_eq!(conda.channels, ["conda-forge"]);
assert!(
serde_json::to_string(&conda)
.unwrap()
.contains(r#""channels":["conda-forge"]"#)
);
let pip: PythonEnvironmentSpec =
serde_json::from_str(r#"{"kind":"pip","packages":["numpy"]}"#).unwrap();
assert!(!serde_json::to_string(&pip).unwrap().contains("channels"));
}
}
+1
View File
@@ -63,4 +63,5 @@ pub struct FmIndexBuilder {}
pub use lance_index::scalar::FullTextSearchQuery;
pub use lance_index::scalar::InvertedIndexParams as FtsIndexBuilder;
pub use lance_index::scalar::InvertedIndexParams;
pub use lance_index::scalar::inverted::DocumentGranularity;
pub use lance_index::scalar::inverted::query::*;
+9 -1
View File
@@ -10,7 +10,7 @@ use lance::io::WrappingObjectStore;
use object_store::{
CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result,
UploadPart, path::Path,
UploadPart, list::PaginatedListStore, path::Path,
};
use async_trait::async_trait;
@@ -187,6 +187,14 @@ impl WrappingObjectStore for MirroringObjectStoreWrapper {
secondary: self.secondary.clone(),
})
}
fn wrap_paginated(
&self,
_store_prefix: &str,
original: Arc<dyn PaginatedListStore>,
) -> Option<Arc<dyn PaginatedListStore>> {
Some(original)
}
}
// windows pathing can't be simply concatenated
@@ -12,7 +12,7 @@ use lance::io::WrappingObjectStore;
use object_store::{
CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, Result as OSResult,
UploadPart, path::Path,
UploadPart, list::PaginatedListStore, path::Path,
};
#[derive(Debug, Default)]
@@ -57,6 +57,14 @@ impl WrappingObjectStore for IoStatsHolder {
stats: self.0.clone(),
})
}
fn wrap_paginated(
&self,
_store_prefix: &str,
original: Arc<dyn PaginatedListStore>,
) -> Option<Arc<dyn PaginatedListStore>> {
Some(original)
}
}
impl IoTrackingStore {
+119 -37
View File
@@ -6,7 +6,7 @@
use std::sync::Arc;
use async_trait::async_trait;
use serde::de::DeserializeOwned;
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
use tokio::sync::watch;
use tokio::task::{AbortHandle, JoinHandle};
@@ -26,20 +26,16 @@ pub(crate) trait JobHandle: Send + Sync {
}
/// A backend-neutral successful terminal result.
///
/// Local operations do not carry a value. Remote operations may carry JSON
/// that the public [`Job`] decodes according to its result type.
#[derive(Clone)]
pub(crate) struct TerminalResult {
#[allow(dead_code)] // Typed remote submit endpoints consume this after Slice 1.
value: Option<Value>,
#[allow(dead_code)] // Preserved so typed decode errors retain request correlation.
request_id: Option<String>,
}
impl TerminalResult {
pub(crate) fn local() -> Self {
fn local(value: Value) -> Self {
Self {
value: None,
value: Some(value),
request_id: None,
}
}
@@ -51,23 +47,35 @@ impl TerminalResult {
}
}
#[allow(dead_code)] // Exercised by the remote typed-result fixtures in Slice 1.
pub(crate) fn value(&self) -> Option<&Value> {
self.value.as_ref()
}
fn decode<T: DeserializeOwned>(self) -> Result<T> {
let request_id = self.request_id.unwrap_or_default();
let value = self.value.ok_or_else(|| Error::Http {
source: "successful typed job response did not contain a result".into(),
request_id: request_id.clone(),
status_code: None,
let value = self.value.ok_or_else(|| match &self.request_id {
Some(request_id) => Error::Http {
source: "successful typed job response did not contain a result".into(),
request_id: request_id.clone(),
status_code: None,
},
None => Error::Runtime {
message: "successful typed job did not contain a result".to_string(),
},
})?;
serde_json::from_value(value).map_err(|error| Error::Http {
source: format!("failed to parse typed job result: {error}").into(),
request_id,
status_code: None,
serde_json::from_value(value).map_err(|error| match self.request_id {
Some(request_id) => Error::Http {
source: format!("failed to parse typed job result: {error}").into(),
request_id,
status_code: None,
},
None => Error::Runtime {
message: format!("failed to parse typed job result: {error}"),
},
})
}
}
type ResultDecoder<T> = fn(TerminalResult) -> Result<T>;
type ResultDecoder<T> = Arc<dyn Fn(TerminalResult) -> Result<T> + Send + Sync>;
enum JobInner<T> {
Handle {
@@ -79,7 +87,9 @@ enum JobInner<T> {
/// A handle to an operation that may still be running.
///
/// The operation may already be complete when the handle is created.
/// The operation may already be complete when the handle is created. `T` is
/// the endpoint's successful terminal result; unit-result operations use the
/// default `Job<()>`.
pub struct Job<T = ()>
where
T: Clone + Send + Sync + 'static,
@@ -111,15 +121,10 @@ impl Job<()> {
Self {
inner: JobInner::Handle {
handle,
decode: |_| Ok(()),
decode: Arc::new(|_| Ok(())),
},
}
}
/// A unit-result job running as a task in this process.
pub(crate) fn spawned(task: JoinHandle<Result<()>>) -> Self {
Self::new(Box::new(SpawnedJob::new(task)))
}
}
impl<T> Job<T>
@@ -131,12 +136,22 @@ where
Self {
inner: JobInner::Handle {
handle,
decode: TerminalResult::decode::<T>,
decode: Arc::new(TerminalResult::decode::<T>),
},
}
}
}
impl<T> Job<T>
where
T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static,
{
/// A typed job running as a task in this process.
pub(crate) fn spawned(task: JoinHandle<Result<T>>) -> Self {
Self::new_typed(Box::new(SpawnedJob::new(task)))
}
}
impl<T> Job<T>
where
T: Clone + Send + Sync + 'static,
@@ -169,11 +184,13 @@ where
/// Waits until the operation reaches a terminal state.
///
/// Returns the endpoint's typed result. Unit-result jobs return `()`.
///
/// Returns [`crate::Error::JobFailed`] if the operation failed and
/// [`crate::Error::JobCancelled`] if it was cancelled.
pub async fn wait(&self) -> Result<T> {
match &self.inner {
JobInner::Handle { handle, decode } => decode(handle.wait().await?),
JobInner::Handle { handle, decode } => (decode)(handle.wait().await?),
JobInner::Completed(result) => Ok(result.clone()),
}
}
@@ -187,21 +204,53 @@ where
JobInner::Completed(_) => Ok(()),
}
}
/// Maps a successful terminal result without changing the job lifecycle.
/// The mapping may run once for each call to [`Job::wait`], so it should
/// be deterministic and free of externally visible side effects.
///
/// ```
/// use lancedb::{Job, function::RefreshColumnResult};
///
/// # async fn rows_assigned(
/// # job: Job<RefreshColumnResult>,
/// # ) -> lancedb::Result<u64> {
/// let job = job.map(|result| result.rows_assigned);
/// job.wait().await
/// # }
/// ```
pub fn map<U, F>(self, map: F) -> Job<U>
where
U: Clone + Send + Sync + 'static,
F: Fn(T) -> U + Send + Sync + 'static,
{
match self.inner {
JobInner::Handle { handle, decode } => Job {
inner: JobInner::Handle {
handle,
decode: Arc::new(move |result| Ok(map((decode)(result)?))),
},
},
JobInner::Completed(result) => Job {
inner: JobInner::Completed(map(result)),
},
}
}
}
/// How an in-process operation ended. Cloneable so every waiter can be given
/// the outcome; [`Error`] is not, so failures share one behind an [`Arc`].
#[derive(Clone)]
enum Outcome {
Succeeded,
Succeeded(TerminalResult),
Failed(Arc<Error>),
Cancelled,
}
impl Outcome {
fn into_result(self) -> Result<()> {
fn into_result(self) -> Result<TerminalResult> {
match self {
Self::Succeeded => Ok(()),
Self::Succeeded(result) => Ok(result),
Self::Failed(source) => Err(Error::JobFailed {
job_id: None,
failure: JobFailure::from_source(source),
@@ -220,12 +269,20 @@ struct SpawnedJob {
}
impl SpawnedJob {
fn new(task: JoinHandle<Result<()>>) -> Self {
fn new<T>(task: JoinHandle<Result<T>>) -> Self
where
T: Serialize + Send + 'static,
{
let abort = task.abort_handle();
let (tx, outcome) = watch::channel(None);
tokio::spawn(async move {
let outcome = match task.await {
Ok(Ok(())) => Outcome::Succeeded,
Ok(Ok(result)) => match serde_json::to_value(result) {
Ok(value) => Outcome::Succeeded(TerminalResult::local(value)),
Err(err) => Outcome::Failed(Arc::new(Error::Runtime {
message: format!("failed to serialize job result: {err}"),
})),
},
Ok(Err(err)) => Outcome::Failed(Arc::new(err)),
Err(err) if err.is_cancelled() => Outcome::Cancelled,
Err(err) => Outcome::Failed(Arc::new(Error::Runtime {
@@ -243,7 +300,7 @@ impl JobHandle for SpawnedJob {
async fn status(&self) -> Result<String> {
let label = match &*self.outcome.borrow() {
None => "running",
Some(Outcome::Succeeded) => "finished",
Some(Outcome::Succeeded(_)) => "finished",
Some(Outcome::Failed(_)) => "failed",
Some(Outcome::Cancelled) => "cancelled",
};
@@ -256,12 +313,11 @@ impl JobHandle for SpawnedJob {
.wait_for(|outcome| outcome.is_some())
.await
.map_err(|_| Error::Runtime {
message: "index job outcome was dropped before it completed".to_string(),
message: "job outcome was dropped before it completed".to_string(),
})?
.clone()
.expect("wait_for returns once an outcome is set");
settled.into_result()?;
Ok(TerminalResult::local())
settled.into_result()
}
async fn cancel(&self) -> Result<()> {
@@ -269,3 +325,29 @@ impl JobHandle for SpawnedJob {
Ok(())
}
}
#[cfg(test)]
mod tests {
use std::future::pending;
use super::*;
#[tokio::test]
async fn mapped_spawned_job_reuses_outcome() {
let job = Job::spawned(tokio::spawn(async { Ok(41_u64) })).map(|value| value + 1);
assert_eq!(job.wait().await.unwrap(), 42);
assert_eq!(job.wait().await.unwrap(), 42);
assert_eq!(job.status().await.unwrap(), "finished");
}
#[tokio::test]
async fn mapped_spawned_job_preserves_cancellation() {
let job = Job::spawned(tokio::spawn(async { pending::<Result<u64>>().await }))
.map(|value| value.to_string());
job.cancel().await.unwrap();
assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. })));
assert_eq!(job.status().await.unwrap(), "cancelled");
}
}
+66 -5
View File
@@ -37,6 +37,13 @@ pub use refresh::{RefreshMaterializedViewResult, RefreshMode};
/// Schema metadata key holding the view definition, as kind-tagged JSON.
pub const DEFINITION_META_KEY: &str = "mv.definition";
/// Schema metadata key holding the view's incarnation: a token minted at each
/// physical creation of a view table, so a view dropped and recreated under
/// the same name and definition is still told apart from the one a caller
/// captured. A view whose metadata was replaced wholesale, or one declared
/// before tokens existed, carries none until its next refresh mints one.
pub const INCARNATION_META_KEY: &str = "mv.incarnation";
/// Schema metadata key holding the source table version the view was last
/// refreshed to. Absent until the first refresh.
pub const SOURCE_VERSION_META_KEY: &str = "mv.source_version";
@@ -163,6 +170,15 @@ pub(crate) fn plan(
filter: Option<&str>,
limit: Option<u64>,
) -> Result<(MaterializedViewDefinition, Vec<ArrowField>, Lineage)> {
let filter = filter
.map(crate::expr::canonicalize_sql_predicate)
.transpose()
.map_err(|err| match err {
Error::InvalidInput { message } => Error::InvalidInput {
message: format!("invalid view filter: {message}"),
},
err => err,
})?;
let projections: Vec<(String, String)> = if projections.is_empty() {
source_schema
.fields()
@@ -267,7 +283,7 @@ pub(crate) fn plan(
declared.push(output);
}
if let Some(filter) = filter {
if let Some(filter) = filter.as_deref() {
let expr = planner
.parse_filter(filter)
.map_err(|e| Error::InvalidInput {
@@ -307,7 +323,7 @@ pub(crate) fn plan(
.into_iter()
.map(|(output, expression)| ViewProjection { output, expression })
.collect(),
filter: filter.map(String::from),
filter,
limit,
inputs,
};
@@ -612,8 +628,17 @@ impl PreparedDeclaration {
pub async fn create(self, name: &str) -> Result<MaterializedView> {
let empty: Vec<std::result::Result<arrow_array::RecordBatch, arrow_schema::ArrowError>> =
vec![];
// Minted here, not at preparation: a declaration can be cloned and
// create more than one physical table, and each needs its own token.
let incarnation = uuid::Uuid::new_v4().to_string();
let mut metadata = self.schema.metadata().clone();
metadata.insert(INCARNATION_META_KEY.to_string(), incarnation.clone());
let schema = Arc::new(ArrowSchema::new_with_metadata(
self.schema.fields().clone(),
metadata,
));
let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
Box::new(arrow_array::RecordBatchIterator::new(empty, self.schema));
Box::new(arrow_array::RecordBatchIterator::new(empty, schema));
let mut request = CreateTableRequest::new(name.to_string(), Box::new(reader));
let write_params = request
.write_options
@@ -648,6 +673,7 @@ impl PreparedDeclaration {
Ok(MaterializedView {
table,
definition: self.definition,
incarnation: Some(incarnation),
})
}
}
@@ -878,6 +904,7 @@ impl CreateMaterializedViewBuilder {
pub struct MaterializedView {
table: Table,
definition: MaterializedViewDefinition,
incarnation: Option<String>,
}
impl MaterializedView {
@@ -893,8 +920,13 @@ impl MaterializedView {
});
}
let schema = table.schema().await?;
let incarnation = schema.metadata().get(INCARNATION_META_KEY).cloned();
match materialized_view_kind(schema.metadata())? {
Some(MaterializedViewKind::Select(definition)) => Ok(Self { table, definition }),
Some(MaterializedViewKind::Select(definition)) => Ok(Self {
table,
definition,
incarnation,
}),
Some(MaterializedViewKind::Unrecognized { kind }) => Err(Error::NotSupported {
message: format!(
"materialized view '{}' is defined by '{kind}', which this version of \
@@ -923,6 +955,13 @@ impl MaterializedView {
&self.definition
}
/// The view's incarnation token as of when this handle was opened; see
/// [`RefreshMaterializedViewBuilder::expect_incarnation`]. `None` for a
/// view that has none yet (see [`INCARNATION_META_KEY`]).
pub fn incarnation(&self) -> Option<&str> {
self.incarnation.as_deref()
}
/// Recompute the view from its source.
///
/// By default the refresh is incremental when the source's changes can be
@@ -943,6 +982,7 @@ impl MaterializedView {
view: self.clone(),
full: false,
source_version: None,
expected_incarnation: None,
}
}
}
@@ -952,6 +992,7 @@ pub struct RefreshMaterializedViewBuilder {
view: MaterializedView,
full: bool,
source_version: Option<u64>,
expected_incarnation: Option<String>,
}
impl RefreshMaterializedViewBuilder {
@@ -967,8 +1008,28 @@ impl RefreshMaterializedViewBuilder {
self
}
/// Refresh only if the view is still the incarnation that minted `token`
/// (see [`MaterializedView::incarnation`]): a refresh requested against
/// one declaration must not land in a view dropped and recreated since,
/// even under the same name and definition.
///
/// Best effort. The token is read from the latest stored manifest before
/// planning and again immediately before every commit, but it is not part
/// of the commit's own condition, so a recreation that lands between that
/// final read and the commit is not caught.
pub fn expect_incarnation(mut self, token: impl Into<String>) -> Self {
self.expected_incarnation = Some(token.into());
self
}
pub async fn execute(self) -> Result<RefreshMaterializedViewResult> {
refresh::execute_refresh(&self.view.table, self.full, self.source_version).await
refresh::execute_refresh(
&self.view.table,
self.full,
self.source_version,
self.expected_incarnation.as_deref(),
)
.await
}
}
+408 -33
View File
@@ -46,8 +46,9 @@ use lance_table::format::Fragment;
use serde::{Deserialize, Serialize};
use super::{
MaterializedViewDefinition, REFRESHED_AT_MS_META_KEY, SOURCE_ROW_ID_COLUMN,
SOURCE_VERSION_META_KEY,
DEFINITION_META_KEY, INCARNATION_META_KEY, MaterializedViewDefinition,
REFRESHED_AT_MS_META_KEY, SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY,
definition_to_metadata,
};
use crate::database::OpenTableRequest;
use crate::table::{NativeTable, NativeTableExt, Table};
@@ -108,6 +109,7 @@ pub(crate) async fn execute_refresh(
view: &Table,
full: bool,
pinned: Option<u64>,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let view_native = view.as_native().ok_or_else(|| Error::NotSupported {
message: "materialized views are supported only on local tables".into(),
@@ -122,6 +124,8 @@ pub(crate) async fn execute_refresh(
view_native.dataset.reload().await?;
let view_ds = view_native.dataset.get().await?.as_ref().clone();
ensure_incarnation(&view_ds, expected_incarnation, view.name()).await?;
// The definition a handle cached at open may since have been replaced;
// what refresh executes and what it stamps must be one generation.
let definition = match super::materialized_view_kind(&view_ds.schema().metadata)? {
@@ -194,8 +198,28 @@ pub(crate) async fn execute_refresh(
),
});
}
let definition_changed =
definition.filter != replanned.filter || definition.inputs != replanned.inputs;
let definition = &replanned;
// A watermark written for a legacy raw filter certifies the rows that
// filter produced, not the canonical predicate above. Rebuild instead of
// accepting or advancing it, and persist the migrated definition in the
// same metadata commit that certifies the replacement rows.
if definition_changed {
return rebuild(
view_native,
&view_ds,
&source_ds,
source_version,
source_ts,
definition,
true,
expected_incarnation,
)
.await;
}
let metadata = &view_ds.schema().metadata;
let watermark: Option<u64> = metadata
.get(SOURCE_VERSION_META_KEY)
@@ -240,6 +264,7 @@ pub(crate) async fn execute_refresh(
increment,
definition,
watermark,
expected_incarnation,
)
.await?;
match reconciled {
@@ -253,6 +278,8 @@ pub(crate) async fn execute_refresh(
source_version,
source_ts,
definition,
false,
expected_incarnation,
)
.await
}
@@ -266,6 +293,8 @@ pub(crate) async fn execute_refresh(
source_version,
source_ts,
definition,
false,
expected_incarnation,
)
.await
}
@@ -595,6 +624,7 @@ async fn incremental(
increment: Increment,
definition: &MaterializedViewDefinition,
watermark: Option<u64>,
expected_incarnation: Option<&str>,
) -> Result<Option<RefreshMaterializedViewResult>> {
let new_fragments = increment.appended;
let watermark_version = watermark.unwrap_or(0);
@@ -671,15 +701,37 @@ async fn incremental(
};
let nothing_to_add = (new_fragments.is_empty() && !updated_rows) || remaining == Some(0);
if nothing_to_add && eviction.is_none() {
result.version =
stamp_watermark(view_native, view_ds.clone(), source_version, source_ts).await?;
result.version = stamp_watermark(
view_native,
view_ds.clone(),
source_version,
source_ts,
None,
expected_incarnation,
)
.await?;
return Ok(Some(result));
}
// Rows left but none arrive: the removals still have to be published.
if nothing_to_add {
let filter = refresh_filter(&empty_keys(view_ds)?)?;
let published = publish(view_ds, eviction, Vec::new(), Some(filter)).await?;
result.version = stamp_watermark(view_native, published, source_version, source_ts).await?;
let published = publish(
view_ds,
eviction,
Vec::new(),
Some(filter),
expected_incarnation,
)
.await?;
result.version = stamp_watermark(
view_native,
published,
source_version,
source_ts,
None,
expected_incarnation,
)
.await?;
return Ok(Some(result));
}
@@ -737,12 +789,21 @@ async fn incremental(
eviction,
Vec::new(),
Some(refresh_filter(&empty_keys(view_ds)?)?),
expected_incarnation,
)
.await?
} else {
view_ds.clone()
};
result.version = stamp_watermark(view_native, published, source_version, source_ts).await?;
result.version = stamp_watermark(
view_native,
published,
source_version,
source_ts,
None,
expected_incarnation,
)
.await?;
return Ok(Some(result));
};
let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new(
@@ -775,12 +836,28 @@ async fn incremental(
});
};
let filter = refresh_filter(&keys)?;
let appended = publish(view_ds, eviction, new_fragments, Some(filter)).await?;
let appended = publish(
view_ds,
eviction,
new_fragments,
Some(filter),
expected_incarnation,
)
.await?;
result.rows_written = rows_written.load(Ordering::Relaxed);
result.version = stamp_watermark(view_native, appended, source_version, source_ts).await?;
result.version = stamp_watermark(
view_native,
appended,
source_version,
source_ts,
None,
expected_incarnation,
)
.await?;
Ok(Some(result))
}
#[allow(clippy::too_many_arguments)]
async fn rebuild(
view_native: &NativeTable,
view_ds: &Dataset,
@@ -788,6 +865,8 @@ async fn rebuild(
source_version: u64,
source_ts: u128,
definition: &MaterializedViewDefinition,
persist_definition: bool,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let rows_written = Arc::new(AtomicU64::new(0));
let schema = Arc::new(ArrowSchema::from(view_ds.schema()));
@@ -810,8 +889,17 @@ async fn rebuild(
// carries no schema metadata, so it cannot erase a definition update
// that raced in the way an overwrite (which adopts its stream's schema)
// durably would -- and it must land on the planned generation or abort.
let replaced = replace_retaining_indices(view_ds.clone(), stream, keys).await?;
let version = stamp_watermark(view_native, replaced, source_version, source_ts).await?;
let replaced =
replace_retaining_indices(view_ds.clone(), stream, keys, expected_incarnation).await?;
let version = stamp_watermark(
view_native,
replaced,
source_version,
source_ts,
persist_definition.then_some(definition),
expected_incarnation,
)
.await?;
Ok(RefreshMaterializedViewResult {
mode: RefreshMode::Rebuild,
rows_written: rows_written.load(Ordering::Relaxed),
@@ -828,11 +916,13 @@ async fn replace_retaining_indices(
view_ds: Dataset,
stream: SendableRecordBatchStream,
keys: Arc<StdMutex<KeyExistenceFilterBuilder>>,
expected_incarnation: Option<&str>,
) -> Result<Dataset> {
let ds = Arc::new(view_ds);
let read_version = ds.version().version;
#[cfg(test)]
tests::hold_before_publish(ds.uri()).await;
ensure_incarnation(&ds, expected_incarnation, ds.uri()).await?;
let removed_fragment_ids: Vec<u64> = ds.get_fragments().iter().map(|f| f.id() as u64).collect();
let write_txn = InsertBuilder::new(WriteDestination::Dataset(ds.clone()))
@@ -886,6 +976,32 @@ async fn replace_retaining_indices(
}
/// Record that the view now reflects `source_version`, including the view
/// Refuse to act on a view that is not `expected`'s incarnation, judged from
/// the latest stored manifest. Not a commit condition; see
/// `RefreshMaterializedViewBuilder::expect_incarnation`.
async fn ensure_incarnation(view_ds: &Dataset, expected: Option<&str>, what: &str) -> Result<()> {
let Some(expected) = expected else {
return Ok(());
};
let mut latest = view_ds.clone();
latest.checkout_latest().await?;
match latest.schema().metadata.get(INCARNATION_META_KEY) {
Some(actual) if actual == expected => Ok(()),
Some(_) => Err(Error::Runtime {
message: format!(
"materialized view '{what}' is not the incarnation this refresh was \
requested for: it was dropped and recreated"
),
}),
None => Err(Error::Runtime {
message: format!(
"materialized view '{what}' carries no incarnation token: its schema \
metadata was replaced since the token was captured"
),
}),
}
}
/// version this very commit produces. The version is predicted and then
/// verified; on a mismatch another commit raced in between, and the stamp
/// ABORTS rather than certify that commit as the refresh's own generation.
@@ -895,28 +1011,45 @@ async fn stamp_watermark(
mut dataset: Dataset,
source_version: u64,
source_ts: u128,
definition: Option<&MaterializedViewDefinition>,
expected_incarnation: Option<&str>,
) -> Result<u64> {
ensure_incarnation(&dataset, expected_incarnation, dataset.uri()).await?;
let predicted = dataset.version().version + 1;
dataset
.update_schema_metadata([
(
SOURCE_VERSION_META_KEY.to_string(),
Some(source_version.to_string()),
),
(
SOURCE_VERSION_TS_META_KEY.to_string(),
Some(source_ts.to_string()),
),
(
REFRESHED_AT_MS_META_KEY.to_string(),
Some(now_ms().to_string()),
),
(
VIEW_VERSION_META_KEY.to_string(),
Some(predicted.to_string()),
),
])
.await?;
// A view with no token (declared before tokens existed, or its metadata
// replaced wholesale) starts a new incarnation here.
let incarnation = dataset
.schema()
.metadata
.get(INCARNATION_META_KEY)
.cloned()
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let mut metadata = vec![(INCARNATION_META_KEY.to_string(), Some(incarnation))];
if let Some(definition) = definition {
metadata.push((
DEFINITION_META_KEY.to_string(),
Some(definition_to_metadata(definition)?),
));
}
metadata.extend([
(
SOURCE_VERSION_META_KEY.to_string(),
Some(source_version.to_string()),
),
(
SOURCE_VERSION_TS_META_KEY.to_string(),
Some(source_ts.to_string()),
),
(
REFRESHED_AT_MS_META_KEY.to_string(),
Some(now_ms().to_string()),
),
(
VIEW_VERSION_META_KEY.to_string(),
Some(predicted.to_string()),
),
]);
dataset.update_schema_metadata(metadata).await?;
let actual = dataset.version().version;
if actual != predicted {
return Err(Error::Runtime {
@@ -1052,12 +1185,14 @@ async fn publish(
eviction: Option<(Vec<Fragment>, Vec<u64>)>,
new_fragments: Vec<Fragment>,
keys: Option<KeyExistenceFilter>,
expected_incarnation: Option<&str>,
) -> Result<Dataset> {
let planned = view_ds.version().version;
#[cfg(test)]
tests::hold_before_publish(view_ds.uri()).await;
#[cfg(test)]
tests::hold_until_peers_planned();
ensure_incarnation(view_ds, expected_incarnation, view_ds.uri()).await?;
let (updated_fragments, removed_fragment_ids) = eviction.unwrap_or_default();
let committed = CommitBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone())))
.execute(Transaction::new(
@@ -1486,6 +1621,106 @@ mod tests {
assert_eq!(read(view.table(), "x").await, vec![20, 40]);
}
#[tokio::test]
async fn test_mixed_case_filter_is_canonicalized_for_lineage_and_refresh() {
let conn = connect("memory://").execute().await.unwrap();
let batch = record_batch!(
("id", Int32, [1, 2, 3]),
("PartyAbbrev", Utf8, ["D", "R", "D"])
)
.unwrap();
conn.create_table("src", batch)
.write_options(crate::materialized_view::tests::stable_row_ids())
.execute()
.await
.unwrap();
conn.create_materialized_view("democrats", "src")
.select([("id", "id")])
.only_if(r#""PartyAbbrev" = 'D'"#)
.execute()
.await
.unwrap();
// Reopen from schema metadata so these assertions cover the stored
// predicate and lineage, not only the declaration-time handle.
let view = conn.open_materialized_view("democrats").await.unwrap();
assert_eq!(
view.definition().filter.as_deref(),
Some("`PartyAbbrev` = 'D'")
);
assert_eq!(view.definition().inputs, ["PartyAbbrev", "id"]);
let result = view.refresh().execute().await.unwrap();
assert_eq!(result.rows_written, 2);
assert_eq!(read(view.table(), "id").await, vec![1, 3]);
}
#[tokio::test]
async fn test_legacy_raw_filter_rebuilds_and_persists_canonical_definition() {
let conn = connect("memory://").execute().await.unwrap();
let batch = record_batch!(
("id", Int32, [1, 2, 3]),
("PartyAbbrev", Utf8, ["D", "R", "D"])
)
.unwrap();
conn.create_table("legacy_src", batch)
.write_options(crate::materialized_view::tests::stable_row_ids())
.execute()
.await
.unwrap();
let view = conn
.create_materialized_view("legacy_view", "legacy_src")
.select([("id", "id")])
.only_if(r#""PartyAbbrev" = 'X'"#)
.execute()
.await
.unwrap();
assert_eq!(view.refresh().execute().await.unwrap().rows_written, 0);
// Model a definition and up-to-date watermark written before filter
// canonicalization was applied to materialized views.
let mut legacy = view.definition().clone();
legacy.filter = Some(r#""PartyAbbrev" = 'D'"#.into());
legacy.inputs = vec!["id".into()];
let native = view.table().as_native().unwrap();
let mut dataset = native.dataset.get().await.unwrap().as_ref().clone();
let predicted = dataset.version().version + 1;
dataset
.update_schema_metadata([
(
DEFINITION_META_KEY.to_string(),
Some(definition_to_metadata(&legacy).unwrap()),
),
(
VIEW_VERSION_META_KEY.to_string(),
Some(predicted.to_string()),
),
])
.await
.unwrap();
native.dataset.update(dataset);
let reopened = conn.open_materialized_view("legacy_view").await.unwrap();
let result = reopened.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Rebuild);
assert_eq!(result.rows_written, 2);
assert_eq!(read(reopened.table(), "id").await, vec![1, 3]);
// A fresh handle proves the migration was stored alongside the new
// watermark and therefore happens only once.
let migrated = conn.open_materialized_view("legacy_view").await.unwrap();
assert_eq!(
migrated.definition().filter.as_deref(),
Some("`PartyAbbrev` = 'D'")
);
assert_eq!(migrated.definition().inputs, ["PartyAbbrev", "id"]);
assert_eq!(
migrated.refresh().execute().await.unwrap().mode,
RefreshMode::NoOp
);
assert_eq!(read(migrated.table(), "id").await, vec![1, 3]);
}
#[tokio::test]
async fn test_append_refreshes_incrementally() {
let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await;
@@ -1830,7 +2065,9 @@ mod tests {
);
let staged = eviction.finish().await.unwrap();
assert!(staged.is_some(), "four ids over a chunk of two stage twice");
publish(&view_ds, staged, Vec::new(), None).await.unwrap();
publish(&view_ds, staged, Vec::new(), None, None)
.await
.unwrap();
native.dataset.reload().await.unwrap();
assert_eq!(read(view.table(), "x").await, vec![5, 6]);
@@ -2486,6 +2723,144 @@ mod tests {
assert_eq!(read(view.table(), "twice").await, vec![14]);
}
/// A refresh bound to an incarnation refuses a view dropped and recreated
/// since, even under the same name and definition; the recreated view's
/// own token is accepted, and the token survives a refresh's stamp.
#[tokio::test]
async fn test_refresh_refuses_a_recreated_view_incarnation() {
let (conn, _, view) = refreshed_doubled(vec![1]).await;
let token = view.incarnation().unwrap().to_string();
view.refresh()
.expect_incarnation(&token)
.execute()
.await
.unwrap();
let reopened = conn.open_materialized_view("doubled").await.unwrap();
assert_eq!(reopened.incarnation(), Some(token.as_str()));
conn.drop_table("doubled", &[]).await.unwrap();
let recreated = doubled_view(&conn).await;
assert_ne!(recreated.incarnation(), Some(token.as_str()));
let err = recreated
.refresh()
.expect_incarnation(&token)
.execute()
.await
.unwrap_err();
assert!(err.to_string().contains("dropped and recreated"), "{err}");
assert_eq!(read(recreated.table(), "twice").await, Vec::<i32>::new());
recreated
.refresh()
.expect_incarnation(recreated.incarnation().unwrap())
.execute()
.await
.unwrap();
assert_eq!(read(recreated.table(), "twice").await, vec![2]);
}
/// A cloned declaration creates two physical tables; each gets its own
/// token.
#[tokio::test]
async fn test_cloned_declaration_mints_a_fresh_incarnation_per_create() {
let (conn, source) = db_with_source(vec![1]).await;
let prepared = crate::materialized_view::prepare_declaration(
&source,
&[("x".into(), "x".into()), ("twice".into(), "x * 2".into())],
None,
None,
)
.await
.unwrap();
let replacement = prepared.clone();
let first = prepared.create("cloned").await.unwrap();
let first_token = first.incarnation().unwrap().to_string();
conn.drop_table("cloned", &[]).await.unwrap();
let second = replacement.create("cloned").await.unwrap();
assert_ne!(second.incarnation(), Some(first_token.as_str()));
}
/// A recreation that lands after planning but before publication is
/// caught by the pre-commit read: the stale refresh fails and the
/// replacement stays empty under its own token.
#[tokio::test(flavor = "multi_thread")]
async fn test_bound_refresh_cannot_publish_into_a_raced_recreation() {
let _serial = DRIFT_LOCK.lock().await;
let (conn, _) = db_with_source(vec![1]).await;
let view = doubled_view(&conn).await;
let token = view.incarnation().unwrap().to_string();
let uri = view
.table()
.as_native()
.unwrap()
.dataset
.get()
.await
.unwrap()
.uri()
.to_string();
*DRIFT_TARGET.lock().unwrap() = Some(uri);
let refreshing =
tokio::spawn(async move { view.refresh().expect_incarnation(token).execute().await });
tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified())
.await
.expect("refresh never reached publication");
conn.drop_table("doubled", &[]).await.unwrap();
let replacement = doubled_view(&conn).await;
let replacement_token = replacement.incarnation().unwrap().to_string();
DRIFT_RELEASED.notify_one();
let result = refreshing.await.unwrap();
assert!(result.is_err(), "the stale refresh unexpectedly succeeded");
let reopened = conn.open_materialized_view("doubled").await.unwrap();
assert_eq!(reopened.incarnation(), Some(replacement_token.as_str()));
assert_eq!(read(reopened.table(), "twice").await, Vec::<i32>::new());
}
/// Replacing the schema metadata wholesale drops the token. A refresh
/// bound to the old token is refused for that reason, not as a
/// recreation; an unbound refresh mints the view a fresh one.
#[tokio::test]
async fn test_a_view_whose_metadata_was_replaced_starts_a_new_incarnation() {
let (conn, _, view) = refreshed_doubled(vec![1]).await;
let token = view.incarnation().unwrap().to_string();
let mut metadata = HashMap::new();
metadata.insert(
crate::materialized_view::DEFINITION_META_KEY.to_string(),
crate::materialized_view::definition_to_metadata(view.definition()).unwrap(),
);
view.table()
.as_native()
.unwrap()
.replace_schema_metadata(metadata)
.await
.unwrap();
assert_eq!(
conn.open_materialized_view("doubled")
.await
.unwrap()
.incarnation(),
None
);
let err = view
.refresh()
.expect_incarnation(&token)
.execute()
.await
.unwrap_err();
assert!(err.to_string().contains("no incarnation token"), "{err}");
view.refresh().execute().await.unwrap();
let reopened = conn.open_materialized_view("doubled").await.unwrap();
assert!(reopened.incarnation().is_some());
assert_ne!(reopened.incarnation(), Some(token.as_str()));
}
/// In-process refreshes of one view serialize: the loser of the race
/// observes the winner's watermark instead of appending the same rows.
#[tokio::test(flavor = "multi_thread")]
@@ -2528,7 +2903,7 @@ mod tests {
let stale = view_native.dataset.get().await.unwrap().as_ref().clone();
view.table().delete("x = 1").await.unwrap();
let err = stamp_watermark(view_native, stale, 99, 99).await;
let err = stamp_watermark(view_native, stale, 99, 99, None, None).await;
assert!(err.is_err());
let result = view.refresh().execute().await.unwrap();
+279 -12
View File
@@ -399,6 +399,9 @@ pub trait QueryBase {
/// x > 5 OR y = 'test'
/// ```
///
/// Identifiers may be delimited with SQL-standard double quotes or
/// backticks. String literals must use single quotes.
///
/// Filtering performance can often be improved by creating a scalar index
/// on the filter column(s).
///
@@ -913,6 +916,17 @@ impl QueryRequest {
/// use different representations) the error is recorded and surfaced later
/// by [`Self::check_filter`].
pub(crate) fn add_filter(&mut self, new: QueryFilter) {
let new = match new {
QueryFilter::Sql(filter) => match crate::expr::canonicalize_sql_predicate(&filter) {
Ok(filter) => QueryFilter::Sql(filter),
Err(err) => {
self.filter_error = Some(err.to_string());
return;
}
},
other => other,
};
self.filter = Some(match self.filter.take() {
None => new,
Some(existing) => match and_filters(existing, new) {
@@ -1174,12 +1188,12 @@ impl VectorQuery {
/// Add another query vector to the search.
///
/// Multiple searches will be dispatched as part of the query.
/// This is a convenience method for adding multiple query vectors
/// to the search. It is not expected to be faster than issuing
/// multiple queries concurrently.
/// Multiple searches will be dispatched as a batch. Flat searches share
/// one table scan across the query vectors, avoiding the scan and memory
/// amplification of issuing the searches concurrently. Indexed searches
/// may still perform per-vector index work.
///
/// The output data will contain an additional columns `query_index` which
/// The output data will contain an additional column `query_index` which
/// will contain the index of the query vector that was used to generate the
/// result.
pub fn add_query_vector(mut self, vector: impl IntoQueryVector) -> Result<Self> {
@@ -1646,10 +1660,14 @@ mod tests {
use std::{collections::HashSet, sync::Arc};
use super::*;
use arrow::{array::downcast_array, compute::concat_batches, datatypes::Int32Type};
use arrow::{
array::downcast_array,
compute::concat_batches,
datatypes::{Int32Type, UInt8Type},
};
use arrow_array::{
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, StringArray, cast::AsArray,
types::Float32Type,
FixedSizeListArray, Float32Array, Int32Array, RecordBatch, RecordBatchIterator,
StringArray, cast::AsArray, types::Float32Type,
};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use futures::{StreamExt, TryStreamExt};
@@ -1774,11 +1792,14 @@ mod tests {
.postfilter();
let result = query.execute().await;
let mut stream = result.expect("should have result");
// should only have one batch
let mut num_rows = 0;
while let Some(batch) = stream.next().await {
// post filter should have removed some rows
assert!(batch.expect("should be Ok").num_rows() < 10);
let batch = batch.expect("should be Ok");
let ids: &Int32Array = batch["id"].as_primitive();
assert!(ids.iter().all(|id| id.unwrap() % 2 == 0));
num_rows += batch.num_rows();
}
assert!(num_rows <= 10);
let query = table
.query()
@@ -1875,6 +1896,157 @@ mod tests {
query.execute().await.unwrap();
}
#[tokio::test]
async fn test_double_quoted_predicates_across_table_operations() {
let tmp_dir = tempdir().unwrap();
let dataset_path = tmp_dir.path().join("test.lance");
let uri = dataset_path.to_str().unwrap();
let schema = Arc::new(ArrowSchema::new(vec![
ArrowField::new("id", DataType::Int32, false),
ArrowField::new("PartyAbbrev", DataType::Utf8, false),
ArrowField::new("path", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3, 4])),
Arc::new(StringArray::from(vec!["D", "R", "R", "D"])),
Arc::new(StringArray::from(vec!["\\", "\\", "x", "x"])),
],
)
.unwrap();
let conn = connect(uri).execute().await.unwrap();
let table = conn.create_table("parties", batch).execute().await.unwrap();
let batches = table
.query()
.only_if(r#""PartyAbbrev" = 'D'"#)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
assert_eq!(
table
.count_rows(Some(r#""PartyAbbrev" = 'D'"#.to_string()))
.await
.unwrap(),
2
);
// Public BaseTable dispatch cannot bypass canonicalization.
let query = AnyQuery::Query(QueryRequest {
filter: Some(QueryFilter::Sql(r#""PartyAbbrev" = 'D'"#.to_string())),
..Default::default()
});
let batches = table
.base_table()
.query(&query, Default::default())
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
assert_eq!(
table
.base_table()
.count_rows(Some(crate::table::Filter::Sql(
r#""PartyAbbrev" = 'D'"#.to_string(),
)))
.await
.unwrap(),
2
);
for predicate in [
r#"id = 1 -- unmatched " in a valid SQL comment"#,
r#"id = 1 /* unmatched " in a valid SQL comment */"#,
r#"id = 1 /*! OR "PartyAbbrev" = 'D' */"#,
r#"path = '\' AND "PartyAbbrev" = 'D'"#,
] {
let batches = table
.query()
.only_if(predicate)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1);
}
// The same canonical predicate contract applies to both merge filters.
let source = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec!["D", "R", "R"])),
Arc::new(StringArray::from(vec!["\\", "\\", "x"])),
],
)
.unwrap();
let mut merge = table.merge_insert(&["id"]);
merge.when_not_matched_by_source_delete(Some(r#""PartyAbbrev" = 'D'"#.to_string()));
let result = table
.base_table()
.merge_insert(
merge,
Box::new(RecordBatchIterator::new(vec![Ok(source)], schema.clone())),
)
.await
.unwrap();
assert_eq!(result.num_deleted_rows, 1);
let source = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec!["U", "U", "U"])),
Arc::new(StringArray::from(vec!["\\", "\\", "x"])),
],
)
.unwrap();
let mut merge = table.merge_insert(&["id"]);
merge.when_matched_update_all(Some(r#"target."PartyAbbrev" = 'D'"#.to_string()));
merge
.execute(Box::new(RecordBatchIterator::new(vec![Ok(source)], schema)))
.await
.unwrap();
assert_eq!(
table
.count_rows(Some(r#""PartyAbbrev" = 'U'"#.to_string()))
.await
.unwrap(),
1
);
let update = table
.update()
.only_if(r#""PartyAbbrev" = 'R'"#)
.column("PartyAbbrev", "'X'");
table.base_table().update(update).await.unwrap();
assert_eq!(
table
.count_rows(Some(r#""PartyAbbrev" = 'X'"#.to_string()))
.await
.unwrap(),
2
);
let result = table
.base_table()
.delete(crate::table::Predicate::String(r#""PartyAbbrev" = 'X'"#))
.await
.unwrap();
assert_eq!(result.num_deleted_rows, 2);
assert_eq!(table.count_rows(None).await.unwrap(), 1);
}
#[tokio::test]
async fn test_select_with_transform() {
let batches = make_non_empty_batches();
@@ -2331,7 +2503,8 @@ mod tests {
.limit(1);
let plan = query.explain_plan(true).await.unwrap();
assert!(plan.contains("UnionExec"));
assert!(plan.contains("KNNVectorDistance: queries=2"));
assert!(!plan.contains("UnionExec"));
let results = query
.execute()
@@ -2346,6 +2519,100 @@ mod tests {
// We don't guarantee order.
assert!(query_index.values().contains(&0));
assert!(query_index.values().contains(&1));
// Batch KNN does not support a per-query offset, so offset queries keep
// the legacy per-vector plan to preserve their result semantics.
let offset_query = table
.query()
.nearest_to(&[0.1, 0.2, 0.3, 0.4])
.unwrap()
.add_query_vector(&[0.5, 0.6, 0.7, 0.8])
.unwrap()
.limit(1)
.offset(1);
assert!(
offset_query
.explain_plan(true)
.await
.unwrap()
.contains("UnionExec")
);
let offset_results = offset_query
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
offset_results
.iter()
.map(RecordBatch::num_rows)
.sum::<usize>(),
2
);
}
#[tokio::test]
async fn test_multiple_binary_query_vectors() {
let vectors = FixedSizeListArray::from_iter_primitive::<UInt8Type, _, _>(
vec![
Some(vec![Some(0), Some(0)]),
Some(vec![Some(255), Some(255)]),
],
2,
);
let schema = Arc::new(ArrowSchema::new(vec![
ArrowField::new("id", DataType::Int32, false),
ArrowField::new("vector", vectors.data_type().clone(), false),
]));
let batch = RecordBatch::try_new(
schema,
vec![Arc::new(Int32Array::from(vec![0, 1])), Arc::new(vectors)],
)
.unwrap();
let conn = connect("memory://").execute().await.unwrap();
let table = conn
.create_table("binary_batch", batch)
.execute()
.await
.unwrap();
let query = table
.query()
.nearest_to(&[0.0, 0.0])
.unwrap()
.add_query_vector(&[255.0, 255.0])
.unwrap()
.distance_type(DistanceType::Hamming)
.limit(1);
// Binary queries retain the per-vector plan because Lance's binary
// nearest path requires primitive UInt8 query arrays.
assert!(
query
.explain_plan(true)
.await
.unwrap()
.contains("UnionExec")
);
let results = query
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let results = concat_batches(&results[0].schema(), &results).unwrap();
assert_eq!(results.num_rows(), 2);
let ids = results["id"].as_primitive::<Int32Type>();
assert!(ids.values().contains(&0));
assert!(ids.values().contains(&1));
let query_index = results["query_index"].as_primitive::<Int32Type>();
assert!(query_index.values().contains(&0));
assert!(query_index.values().contains(&1));
}
#[tokio::test]

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